fix(deep-research): source-routing guard recognizes sandbox

src/aiq_agent/agents/deep_researcher/custom_middleware.py

@@ -38,6 +38,11 @@
 # Path to this agent's prompts directory
 _PROMPTS_DIR = Path(__file__).parent / "prompts"
 _SOURCE_ROUTING_PATH = "/shared/source_routing.json"
+# When a sandbox provider is configured, CompositeBackend strips the /shared/ route
+# before delegating to StateBackend, so the router's file is stored under the
+# route-local key. The guard reads raw state, so it must accept both forms or it
+# blocks the orchestrator forever on sandboxed runs.
+_SOURCE_ROUTING_STATE_KEYS = (_SOURCE_ROUTING_PATH, "/source_routing.json")
 
 
 class SourceRoutingGuardMiddleware(AgentMiddleware):
@@ -50,7 +55,7 @@ def __init__(self, *, enabled: bool, required_subagent: str = "source-router-age
     @staticmethod
     def _routing_complete(state: object) -> bool:
         files = state.get("files", {}) if isinstance(state, dict) else getattr(state, "files", {})
-        return isinstance(files, dict) and _SOURCE_ROUTING_PATH in files
+        return isinstance(files, dict) and any(key in files for key in _SOURCE_ROUTING_STATE_KEYS)
 
     async def awrap_tool_call(self, request, handler):
         """Block out-of-order calls until the source router writes its route file."""

tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py

@@ -89,6 +89,19 @@ async def test_allows_normal_tools_after_routing_file_exists(self):
         handler.assert_awaited_once_with(request)
         assert result is expected
 
+    @pytest.mark.asyncio
+    async def test_allows_normal_tools_after_routing_file_exists_sandbox_key(self):
+        """Under a sandbox provider the /shared/ route is stripped; the route-local key must also open the gate."""
+        middleware = SourceRoutingGuardMiddleware(enabled=True)
+        expected = ToolMessage(content="[]", tool_call_id="tc1")
+        handler = AsyncMock(return_value=expected)
+        request = self._request("ls", files={"/source_routing.json": {"content": "{}"}})
+
+        result = await middleware.awrap_tool_call(request, handler)
+
+        handler.assert_awaited_once_with(request)
+        assert result is expected
+
     @pytest.mark.asyncio
     async def test_disabled_guard_is_noop(self):
         """Workflows with source routing disabled preserve their existing tool behavior."""