fix test

Author: mayinghan

tests/test_rollout_control_plane_integration.py

@@ -50,9 +50,13 @@ async def __call__(self, tool_schema, env_index, conversation_history):
         tool_calls = []
         tool_call = MCPToolCall(tool_name="lake_move", arguments={"action": action})
         tool_calls.append(tool_call)
+        if self.step_count == 3:
+            self.step_count += 1
+            no_tool_call = MCPToolCall(tool_name="_no_tool_call", arguments={})
+            return [no_tool_call], None, "stop"
 
         self.step_count += 1
-        return tool_calls, None
+        return tool_calls, None, None
 
     def add_tool_response(
         self,
@@ -285,11 +289,11 @@ def mock_step_side_effect(env_index, tool_call):
             final_cp_step = final_msg.control_plane_step
             assert final_cp_step["terminated"] == True, "Final step should be terminated"
             assert final_cp_step["reward"] == 1.0, "Final step should have correct reward"
-            assert final_cp_step["termination_reason"] == "control_plane_signal", "Should terminate via control plane"
+            assert final_cp_step["termination_reason"] == "stop", "Should terminate via control plane"
             assert final_cp_step["step"] == 2, "Should record final step"
 
             # Validate policy interaction
-            assert policy.step_count == 4, "Policy should have been called 3 times"
+            assert policy.step_count == 4, "Policy should have been called 4 times"
 
     @pytest.mark.asyncio
     async def test_rollout_trajectory_recording_with_control_plane(self):