diff --git a/backend/packages/harness/deerflow/sandbox/middleware.py b/backend/packages/harness/deerflow/sandbox/middleware.py index 7bb11c0482..924f801a0a 100644 --- a/backend/packages/harness/deerflow/sandbox/middleware.py +++ b/backend/packages/harness/deerflow/sandbox/middleware.py @@ -27,6 +27,7 @@ from deerflow.sandbox.exceptions import SandboxAuthorizationError, SandboxRuntim from deerflow.sandbox.lease import ( ensure_sandbox_lease_owner, get_sandbox_lease_manager, + run_sync_lifecycle_operation, sandbox_lease_owner, ) from deerflow.sandbox.overwrite import unwrap_sandbox @@ -352,7 +353,7 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): thread_id = (runtime.context or {}).get("thread_id") if thread_id is None: return await super().abefore_agent(state, runtime) - await asyncio.to_thread(self._apply_network_policy_response, state, runtime) + await run_sync_lifecycle_operation(self._apply_network_policy_response, state, runtime) user_id = resolve_runtime_user_id(runtime) projection = await asyncio.to_thread( self._prepare_agent_skill_projection, diff --git a/backend/tests/test_sandbox_middleware.py b/backend/tests/test_sandbox_middleware.py index fe519b2cac..fe8f8b7883 100644 --- a/backend/tests/test_sandbox_middleware.py +++ b/backend/tests/test_sandbox_middleware.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import threading from typing import get_type_hints import pytest @@ -728,6 +729,65 @@ def test_before_agent_applies_network_approval_to_same_sandbox() -> None: assert provider.decisions == [("existing", "req-1", "allow_temporary")] +@pytest.mark.anyio +async def test_abefore_agent_drains_started_network_approval_across_cancellation() -> None: + provider = _NetworkPolicyProvider() + started = threading.Event() + release = threading.Event() + + def decide(sandbox_id: str, request_id: str, decision: str) -> bool: + started.set() + assert release.wait(timeout=2) + provider.decisions.append((sandbox_id, request_id, decision)) + return True + + provider.decide_network_policy_request = decide # type: ignore[method-assign] + response = HumanMessage( + content="Allow network access for 5 minutes", + additional_kwargs={ + "hide_from_ui": True, + "human_input_response": { + "version": 1, + "kind": "human_input_response", + "source": "sandbox_network", + "request_id": "req-cancel", + "response_kind": "option", + "option_id": "allow_temporary", + "value": "Allow network access for 5 minutes", + }, + }, + ) + state = {"sandbox": {"sandbox_id": "existing"}, "messages": [response]} + task = None + set_sandbox_provider(provider) + try: + task = asyncio.create_task( + SandboxMiddleware().abefore_agent( + state, + Runtime(context={"thread_id": "thread-cancel"}), + ) + ) + assert await asyncio.to_thread(started.wait, 2) + + task.cancel() + await asyncio.sleep(0.05) + task.cancel() + await asyncio.sleep(0.05) + assert not task.done() + + release.set() + with pytest.raises(asyncio.CancelledError): + await task + finally: + release.set() + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + reset_sandbox_provider() + + assert provider.decisions == [("existing", "req-cancel", "allow_temporary")] + + def test_before_agent_does_not_reapply_network_approval_after_new_user_turn() -> None: provider = _NetworkPolicyProvider() response = HumanMessage(