mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-23 00:16:48 +00:00
fix(sandbox): avoid blocking sandbox readiness polling (#2822)
* fix(sandbox): offload async sandbox acquisition Run blocking sandbox provider acquisition through the async provider hook so eager sandbox setup does not stall the event loop. * fix(sandbox): add async readiness polling Introduce an async sandbox readiness poller using httpx and asyncio.sleep while preserving the existing synchronous API. * test(sandbox): cover async readiness polling Lock in non-blocking readiness behavior so the async helper does not regress to requests.get or time.sleep. * fix(sandbox): allow anonymous backend creation * fix(sandbox): use async readiness in provider acquisition * fix(sandbox): use async acquisition for lazy tools * test(sandbox): cover anonymous remote creation * fix(sandbox): clamp async readiness timeout budget * fix(sandbox): offload async lock file handling * fix(sandbox): delegate async middleware fallthrough * docs(sandbox): document async acquisition path * fix(sandbox): offload async sandbox release * docs(sandbox): mention async release hook * fix(sandbox): address async lock review Reduce duplicate sync/async sandbox acquisition state handling and move async thread-lock waits onto a dedicated executor with cancellation-safe cleanup. * chore: retrigger ci Retrigger GitHub Actions after upstream main fixed the stale PR merge lint failure. * test(sandbox): sync backend unit fixtures --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import NotRequired, override
|
||||
|
||||
@@ -48,6 +49,15 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
||||
logger.info(f"Acquiring sandbox {sandbox_id}")
|
||||
return sandbox_id
|
||||
|
||||
async def _acquire_sandbox_async(self, thread_id: str) -> str:
|
||||
provider = get_sandbox_provider()
|
||||
sandbox_id = await provider.acquire_async(thread_id)
|
||||
logger.info(f"Acquiring sandbox {sandbox_id}")
|
||||
return sandbox_id
|
||||
|
||||
async def _release_sandbox_async(self, sandbox_id: str) -> None:
|
||||
await asyncio.to_thread(get_sandbox_provider().release, sandbox_id)
|
||||
|
||||
@override
|
||||
def before_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
||||
# Skip acquisition if lazy_init is enabled
|
||||
@@ -64,6 +74,23 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
||||
return {"sandbox": {"sandbox_id": sandbox_id}}
|
||||
return super().before_agent(state, runtime)
|
||||
|
||||
@override
|
||||
async def abefore_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
||||
# Skip acquisition if lazy_init is enabled
|
||||
if self._lazy_init:
|
||||
return await super().abefore_agent(state, runtime)
|
||||
|
||||
# Eager initialization (original behavior), but use the async provider
|
||||
# hook so blocking sandbox startup/polling runs outside the event loop.
|
||||
if "sandbox" not in state or state["sandbox"] is None:
|
||||
thread_id = (runtime.context or {}).get("thread_id")
|
||||
if thread_id is None:
|
||||
return await super().abefore_agent(state, runtime)
|
||||
sandbox_id = await self._acquire_sandbox_async(thread_id)
|
||||
logger.info(f"Assigned sandbox {sandbox_id} to thread {thread_id}")
|
||||
return {"sandbox": {"sandbox_id": sandbox_id}}
|
||||
return await super().abefore_agent(state, runtime)
|
||||
|
||||
@override
|
||||
def after_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
||||
sandbox = state.get("sandbox")
|
||||
@@ -81,3 +108,21 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
||||
|
||||
# No sandbox to release
|
||||
return super().after_agent(state, runtime)
|
||||
|
||||
@override
|
||||
async def aafter_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
||||
sandbox = state.get("sandbox")
|
||||
if sandbox is not None:
|
||||
sandbox_id = sandbox["sandbox_id"]
|
||||
logger.info(f"Releasing sandbox {sandbox_id}")
|
||||
await self._release_sandbox_async(sandbox_id)
|
||||
return None
|
||||
|
||||
if (runtime.context or {}).get("sandbox_id") is not None:
|
||||
sandbox_id = runtime.context.get("sandbox_id")
|
||||
logger.info(f"Releasing sandbox {sandbox_id} from context")
|
||||
await self._release_sandbox_async(sandbox_id)
|
||||
return None
|
||||
|
||||
# No sandbox to release
|
||||
return await super().aafter_agent(state, runtime)
|
||||
|
||||
Reference in New Issue
Block a user