fix(workplace-assistant): release session state after verification

resources_servers/workplace_assistant/app.py

@@ -103,27 +103,31 @@ async def route_to_python_function(self, path: str, body: WorkbenchRequest, requ
                 output=f"Error executing tool '{path}': {str(e)}"
             )  # return error to model so that it can correct itself
 
-    async def verify(self, body: WorkbenchVerifyRequest) -> WorkbenchVerifyResponse:
-        ground_truth = body.ground_truth
-        response = body.response.output
+    async def verify(self, request: Request, body: WorkbenchVerifyRequest) -> WorkbenchVerifyResponse:
+        session_id = request.session[SESSION_ID_KEY]
+        try:
+            ground_truth = body.ground_truth
+            response = body.response.output
 
-        total_score = 0.0
+            total_score = 0.0
 
-        # Convert list of ResponseFunctionToolCall objects into list of dictionaries
-        predicted_function_calls = []
+            # Convert list of ResponseFunctionToolCall objects into list of dictionaries
+            predicted_function_calls = []
 
-        for message in response:
-            if message.type == "function_call":
-                predicted_function_calls.append(message.model_dump())
+            for message in response:
+                if message.type == "function_call":
+                    predicted_function_calls.append(message.model_dump())
 
-        predicted_chat_content = []
+            predicted_chat_content = []
 
-        for message in response:
-            if message.type == "output_text":
-                predicted_chat_content.append(message.model_dump())
+            for message in response:
+                if message.type == "output_text":
+                    predicted_chat_content.append(message.model_dump())
 
-        total_score += is_correct(predicted_function_calls, ground_truth, None) * 1.0
-        return WorkbenchVerifyResponse(**body.model_dump(), reward=total_score)
+            total_score += is_correct(predicted_function_calls, ground_truth, None) * 1.0
+            return WorkbenchVerifyResponse(**body.model_dump(), reward=total_score)
+        finally:
+            self.session_id_to_tool_env.pop(session_id, None)
 
 
 if __name__ == "__main__":

resources_servers/workplace_assistant/tests/test_app.py

@@ -16,7 +16,7 @@
 
 import pandas as pd
 from fastapi import Request
-from pytest import fixture
+from pytest import fixture, raises
 
 from nemo_gym.base_resources_server import BaseSeedSessionRequest
 from nemo_gym.openai_utils import (
@@ -1254,12 +1254,9 @@ async def test_verify(self, config: WorkbenchResourcesServerConfig) -> None:
         ]
         mock_email_df = pd.DataFrame(mock_email_data)
 
-        # Seed session
-        resources_server = self.init_server(config)
         session_id = "test_session_company_dir"
         mock_scope = {"type": "http", "session": {SESSION_ID_KEY: session_id}}
         mock_request_with_session = Request(scope=mock_scope)
-        await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest())
 
         with (
             patch(
@@ -1274,6 +1271,8 @@ async def test_verify(self, config: WorkbenchResourcesServerConfig) -> None:
             mock_email_csv.return_value = mock_email_df
 
             resources_server = self.init_server(config)
+            await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest())
+            assert session_id in resources_server.session_id_to_tool_env
 
             HARDCODED_CURRENT_TIME = pd.to_datetime("2023-11-30T23:59:00")
             SYS_PROMPT = (
@@ -1381,9 +1380,29 @@ async def test_verify(self, config: WorkbenchResourcesServerConfig) -> None:
                 id="0",
             )
 
-            verification_response = await resources_server.verify(verify_request)
+            verification_response = await resources_server.verify(mock_request_with_session, verify_request)
 
             assert verification_response.reward == 1.0
+            assert session_id not in resources_server.session_id_to_tool_env
+
+    async def test_verify_cleans_up_session_when_scoring_fails(self, config: WorkbenchResourcesServerConfig) -> None:
+        resources_server = self.init_server(config)
+        session_id = "test_session_scoring_failure"
+        mock_scope = {"type": "http", "session": {SESSION_ID_KEY: session_id}}
+        mock_request_with_session = Request(scope=mock_scope)
+        await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest())
+
+        verify_request = MagicMock()
+        verify_request.ground_truth = []
+        verify_request.response.output = []
+
+        with (
+            patch("resources_servers.workplace_assistant.app.is_correct", side_effect=RuntimeError("scoring failed")),
+            raises(RuntimeError, match="scoring failed"),
+        ):
+            await resources_server.verify(mock_request_with_session, verify_request)
+
+        assert session_id not in resources_server.session_id_to_tool_env
 
     async def test_stateful_email_deletion_and_fetch(self, config: WorkbenchResourcesServerConfig) -> None:
         """
@@ -1578,13 +1597,20 @@ async def test_stateful_email_deletion_and_fetch(self, config: WorkbenchResource
                 id="1",
             )
 
-            verification_response = await resources_server.verify(verify_request)
+            session_id = "test_session_stateful_email"
+            mock_scope = {"type": "http", "session": {SESSION_ID_KEY: session_id}}
+            mock_request_with_session = Request(scope=mock_scope)
+            await resources_server.seed_session(mock_request_with_session, BaseSeedSessionRequest())
+            assert session_id in resources_server.session_id_to_tool_env
+
+            verification_response = await resources_server.verify(mock_request_with_session, verify_request)
 
             # The reward should be 1.0 because the predicted function calls in our
             # crafted `response` object exactly match the `ground_truth`.
             assert verification_response.reward == 1.0, (
                 f"Verification failed with reward {verification_response.reward}"
             )
+            assert session_id not in resources_server.session_id_to_tool_env
 
     async def test_extra_arguments_error_handling(self, config: WorkbenchResourcesServerConfig) -> None:
         """