diff --git a/src/envs/fleet_env/task_env.py b/src/envs/fleet_env/task_env.py index da2b62d0f..13c2b22b2 100644 --- a/src/envs/fleet_env/task_env.py +++ b/src/envs/fleet_env/task_env.py @@ -27,6 +27,31 @@ clear_task_context, ) +# Synthetic tool injected by the harness (not from MCP). +# Mirrors orchestrator/temporal/workflows/constants.py → ANSWER_SUBMISSION_TOOL. +SUBMIT_FINAL_ANSWER_TOOL = { + "type": "function", + "function": { + "name": "submit_final_answer", + "description": ( + "Submit your final answer to complete the task. Use this when you " + "have finished the task and want to provide your answer for " + "verification. If the requested answer asks for json, then write " + "your response in the answer field using json brackets." + ), + "parameters": { + "type": "object", + "properties": { + "answer": { + "type": "string", + "description": "Your final answer", + } + }, + "required": ["answer"], + }, + }, +} + def _is_tool_error(result: Any) -> Tuple[bool, Optional[str]]: """Check if a tool result indicates an error. @@ -132,6 +157,7 @@ def __init__( self._tools_cache: Optional[List[Dict]] = None self._reward_computed = False self.final_reward: Optional[float] = None + self._submitted_answer: Optional[str] = None # Set telemetry context so init failures are tracked with full context set_task_context( @@ -288,6 +314,7 @@ async def reset_async(self, seed: Optional[int] = None) -> Dict[str, Any]: self._done = False self._reward_computed = False self.final_reward = None + self._submitted_answer = None # Reset the environment (use short timeout to avoid blocking on broken manager APIs) # reset() failure is non-fatal — env is up, just the manager API timed out @@ -385,6 +412,13 @@ async def reset_async(self, seed: Optional[int] = None) -> Dict[str, Any]: self._rollout_completed_emitted = True raise RuntimeError("No tools available after filtering") + # Inject submit_final_answer synthetic tool for tool_use tasks whose + # prompt references it. This mirrors the harness's ANSWER_SUBMISSION_TOOL + # so that models can submit answers during SkyRL training exactly as + # they would in a Fleet harness session. + if self.modality == "tool_use" and "submit_final_answer" in self.prompt: + self._tools_cache.append(SUBMIT_FINAL_ANSWER_TOOL) + # Build observation with cached tools obs = { "prompt": self.prompt, @@ -467,7 +501,14 @@ async def step_async( tool_params = action.get("params", {}) tool_result = None - if tool_name: + if tool_name == "submit_final_answer": + # Synthetic tool — handled locally, not routed to MCP. + self._submitted_answer = tool_params.get("answer", "") + tool_result = {"status": "submitted", "message": "Answer recorded. Ending session."} + info["tool_result"] = tool_result + info["submitted_answer"] = self._submitted_answer + agent_done = True # Force episode end, same as harness behaviour + elif tool_name: try: tool_result = await self._tools.call_tool(tool_name, tool_params) info["tool_result"] = tool_result @@ -603,7 +644,12 @@ async def _compute_reward(self) -> float: # Execute verifier in a thread to avoid blocking the event loop. # verify_detailed() does sync HTTP calls internally. - response = await asyncio.to_thread(fleet_task.verify_detailed, fleet_env) + # Pass final_answer when model used submit_final_answer, + # mirroring how the harness routes the answer to the verifier. + verify_kwargs = {} + if self._submitted_answer is not None: + verify_kwargs["final_answer"] = self._submitted_answer + response = await asyncio.to_thread(fleet_task.verify_detailed, fleet_env, **verify_kwargs) # Extract result from response # response.success is bool, response.result is the verifier's return value (0.0 or 1.0) diff --git a/tests/envs/test_fleet_task_env.py b/tests/envs/test_fleet_task_env.py index 00e06744b..859bf7e14 100644 --- a/tests/envs/test_fleet_task_env.py +++ b/tests/envs/test_fleet_task_env.py @@ -520,3 +520,67 @@ async def mock_list_tools(): # Should only have computer tool assert len(env._tools_cache) == 1 assert env._tools_cache[0]["function"]["name"] == "computer" + + +class TestSubmitFinalAnswer: + """Tests for synthetic submit_final_answer tool injection.""" + + def test_submit_final_answer_tool_definition(self, mock_fleet_env_client): + """SUBMIT_FINAL_ANSWER_TOOL has correct schema.""" + from envs.fleet_env.task_env import SUBMIT_FINAL_ANSWER_TOOL + + func = SUBMIT_FINAL_ANSWER_TOOL["function"] + assert func["name"] == "submit_final_answer" + assert "answer" in func["parameters"]["properties"] + assert func["parameters"]["required"] == ["answer"] + + def test_submitted_answer_init(self, sample_task_config, mock_fleet_env_client): + """_submitted_answer should be None on init.""" + from envs.fleet_env.task_env import FleetTaskEnv + + env = FleetTaskEnv(sample_task_config, api_key="test") + assert env._submitted_answer is None + + @pytest.mark.anyio + async def test_step_submit_final_answer_stores_answer( + self, sample_task_config, mock_fleet_env_client + ): + """Calling submit_final_answer should store the answer and mark done.""" + from envs.fleet_env.task_env import FleetTaskEnv + + mock_orch, _ = mock_fleet_env_client + env = FleetTaskEnv(sample_task_config, api_key="test") + env._orch = mock_orch + env._tools = MagicMock() + env._tools_cache = [{"type": "function", "function": {"name": "bash"}}] + env._done = False + env._rollout_started = True + + action = {"tool": "submit_final_answer", "params": {"answer": '["row1", "row2"]'}} + obs, reward, done, info = await env.step_async(action) + + assert env._submitted_answer == '["row1", "row2"]' + assert done is True + assert info["submitted_answer"] == '["row1", "row2"]' + assert info["tool_result"]["status"] == "submitted" + + @pytest.mark.anyio + async def test_step_submit_final_answer_not_routed_to_mcp( + self, sample_task_config, mock_fleet_env_client + ): + """submit_final_answer should NOT call MCP tools.call_tool.""" + from envs.fleet_env.task_env import FleetTaskEnv + + mock_orch, _ = mock_fleet_env_client + mock_tools = AsyncMock() + env = FleetTaskEnv(sample_task_config, api_key="test") + env._orch = mock_orch + env._tools = mock_tools + env._tools_cache = [{"type": "function", "function": {"name": "bash"}}] + env._done = False + env._rollout_started = True + + action = {"tool": "submit_final_answer", "params": {"answer": "42"}} + await env.step_async(action) + + mock_tools.call_tool.assert_not_called()