mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-16 17:46:20 +00:00
0f7d8709d3
* feat(sandbox): add controlled egress approvals * Apply batched suggestions from code review * fix(sandbox): harden restricted network policy * fix(sandbox): harden denied egress handling * fix(sandbox): isolate network proxy sidecar * chore: retry sandbox image smoke * fix(sandbox): close remaining network policy gaps * fix(sandbox): harden relay token rejection * fix(sandbox): fence incompatible policy replacement * fix(sandbox): replace containers across network modes * fix(sandbox): close remaining lifecycle gaps --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
648 lines
29 KiB
Python
648 lines
29 KiB
Python
import asyncio
|
|
import logging
|
|
from collections.abc import Awaitable, Callable, Mapping
|
|
from dataclasses import replace as dc_replace
|
|
from typing import NotRequired, override
|
|
|
|
from langchain.agents import AgentState
|
|
from langchain.agents.middleware import AgentMiddleware
|
|
from langchain_core.messages import HumanMessage, ToolMessage
|
|
from langgraph.graph import END
|
|
from langgraph.prebuilt.tool_node import ToolCallRequest
|
|
from langgraph.runtime import Runtime
|
|
from langgraph.types import Command, Overwrite
|
|
|
|
from deerflow.agents.human_input import read_human_input_response
|
|
from deerflow.agents.thread_state import SandboxStateField, ThreadDataState
|
|
from deerflow.authz.sandbox_authz import (
|
|
authorize_sandbox_execution,
|
|
authorize_sandbox_execution_async,
|
|
safe_app_config,
|
|
safe_app_config_async,
|
|
)
|
|
from deerflow.runtime.user_context import resolve_runtime_user_id
|
|
from deerflow.sandbox import get_sandbox_provider
|
|
from deerflow.sandbox.exceptions import SandboxAuthorizationError, SandboxRuntimeError
|
|
from deerflow.sandbox.lease import (
|
|
ensure_sandbox_lease_owner,
|
|
get_sandbox_lease_manager,
|
|
sandbox_lease_owner,
|
|
)
|
|
from deerflow.sandbox.overwrite import unwrap_sandbox
|
|
from deerflow.sandbox.sandbox_provider import get_initialized_sandbox_provider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
NETWORK_POLICY_HUMAN_INPUT_SOURCE = "sandbox_network"
|
|
_NETWORK_POLICY_DECISIONS = frozenset({"deny", "allow_temporary", "allow_sandbox"})
|
|
|
|
|
|
def _network_approval_is_non_interactive(context: Mapping[str, object]) -> bool:
|
|
return bool(context.get("disable_clarification") or context.get("non_interactive"))
|
|
|
|
|
|
class SandboxMiddlewareState(AgentState):
|
|
"""Compatible with the `ThreadState` schema."""
|
|
|
|
sandbox: SandboxStateField
|
|
thread_data: NotRequired[ThreadDataState | None]
|
|
|
|
|
|
class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|
"""Create a sandbox environment and assign it to an agent.
|
|
|
|
Lifecycle Management:
|
|
- With lazy_init=True (default): Sandbox is acquired on first tool call
|
|
- With lazy_init=False: Sandbox is acquired on first agent invocation (before_agent)
|
|
- Concurrent lead/subagent executions hold independent process-local leases
|
|
- Only the final execution release parks a remote sandbox in its warm pool
|
|
- Provider shutdown remains the terminal cleanup boundary
|
|
"""
|
|
|
|
state_schema = SandboxMiddlewareState
|
|
|
|
def __init__(
|
|
self,
|
|
lazy_init: bool = True,
|
|
*,
|
|
available_skills: set[str] | None = None,
|
|
owns_agent_skill_projection: bool = True,
|
|
):
|
|
"""Initialize sandbox middleware.
|
|
|
|
Args:
|
|
lazy_init: If True, defer sandbox acquisition until first tool call.
|
|
If False, acquire sandbox eagerly in before_agent().
|
|
Default is True for optimal performance.
|
|
owns_agent_skill_projection: Whether this middleware may create or
|
|
rebuild the thread's physical skill projection. Delegated
|
|
subagents share the lead thread sandbox and must preserve the
|
|
lead-owned view instead of applying their discovery policy to it.
|
|
"""
|
|
super().__init__()
|
|
self._lazy_init = lazy_init
|
|
self._available_skills = set(available_skills) if available_skills is not None else None
|
|
self._owns_agent_skill_projection = owns_agent_skill_projection
|
|
|
|
def _prepare_agent_skill_projection(self, thread_id: str, *, user_id: str):
|
|
"""Build the run's physical skill view before any sandbox is reused."""
|
|
if not self._owns_agent_skill_projection:
|
|
# Subagents inherit the lead's thread id and sandbox state. Their
|
|
# skill lists scope discovery/activation only; rebuilding here
|
|
# would widen or narrow the shared filesystem for every concurrent
|
|
# agent using this sandbox.
|
|
return None
|
|
|
|
from deerflow.config.paths import get_paths
|
|
|
|
# Preserve the zero-copy shared view for ordinary threads. A thread
|
|
# that previously used a restricted Agent keeps its stable mount root;
|
|
# an unrestricted run repopulates that root with all enabled skills.
|
|
if self._available_skills is None and not get_paths().thread_skills_view_dir(thread_id, user_id=user_id).exists():
|
|
return None
|
|
|
|
provider = get_sandbox_provider()
|
|
if not provider.supports_agent_skill_isolation:
|
|
if self._available_skills is not None:
|
|
raise SandboxRuntimeError(f"Sandbox provider {provider.__class__.__name__} cannot enforce per-Agent skill filesystem isolation")
|
|
# The thread projection may have been created under a different
|
|
# provider. An unrestricted run does not need that policy view and
|
|
# may safely use this provider's ordinary shared skill behavior.
|
|
return None
|
|
|
|
from deerflow.config import get_app_config
|
|
from deerflow.skills.projection import ensure_thread_skill_projection
|
|
from deerflow.skills.storage import get_or_new_user_skill_storage
|
|
|
|
app_config = get_app_config()
|
|
storage = get_or_new_user_skill_storage(user_id, app_config=app_config)
|
|
return ensure_thread_skill_projection(storage, thread_id, self._available_skills)
|
|
|
|
@staticmethod
|
|
def _require_projection_support(provider, projection) -> None:
|
|
if projection is not None and not provider.supports_agent_skill_isolation:
|
|
raise SandboxRuntimeError(f"Sandbox provider {provider.__class__.__name__} cannot enforce per-Agent skill filesystem isolation")
|
|
|
|
def _acquire_sandbox(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
user_id: str,
|
|
owner_id: str | None,
|
|
) -> str:
|
|
provider = get_sandbox_provider()
|
|
if owner_id is None:
|
|
sandbox_id = provider.acquire(thread_id, user_id=user_id)
|
|
else:
|
|
sandbox_id = get_sandbox_lease_manager(provider).acquire(
|
|
owner_id,
|
|
thread_id,
|
|
user_id=user_id,
|
|
)
|
|
logger.info(f"Acquiring sandbox {sandbox_id}")
|
|
return sandbox_id
|
|
|
|
async def _acquire_sandbox_async(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
user_id: str,
|
|
owner_id: str | None,
|
|
) -> str:
|
|
provider = get_sandbox_provider()
|
|
if owner_id is None:
|
|
sandbox_id = await provider.acquire_async(thread_id, user_id=user_id)
|
|
else:
|
|
sandbox_id = await get_sandbox_lease_manager(provider).acquire_async(
|
|
owner_id,
|
|
thread_id,
|
|
user_id=user_id,
|
|
)
|
|
logger.info(f"Acquiring sandbox {sandbox_id}")
|
|
return sandbox_id
|
|
|
|
@staticmethod
|
|
def _retain_existing_sandbox(
|
|
state: SandboxMiddlewareState,
|
|
*,
|
|
thread_id: str,
|
|
user_id: str,
|
|
owner_id: str | None,
|
|
) -> str | None:
|
|
if owner_id is None:
|
|
return None
|
|
sandbox, fork_restored = unwrap_sandbox(state.get("sandbox"))
|
|
if not isinstance(sandbox, dict) or fork_restored:
|
|
return None
|
|
sandbox_id = sandbox.get("sandbox_id")
|
|
if isinstance(sandbox_id, str):
|
|
provider = get_sandbox_provider()
|
|
get_sandbox_lease_manager(provider).retain(
|
|
owner_id,
|
|
sandbox_id,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
)
|
|
return sandbox_id
|
|
return None
|
|
|
|
@staticmethod
|
|
async def _retain_existing_sandbox_async(
|
|
state: SandboxMiddlewareState,
|
|
*,
|
|
thread_id: str,
|
|
user_id: str,
|
|
owner_id: str | None,
|
|
) -> str | None:
|
|
if owner_id is None:
|
|
return None
|
|
sandbox, fork_restored = unwrap_sandbox(state.get("sandbox"))
|
|
if not isinstance(sandbox, dict) or fork_restored:
|
|
return None
|
|
sandbox_id = sandbox.get("sandbox_id")
|
|
if isinstance(sandbox_id, str):
|
|
provider = get_sandbox_provider()
|
|
await get_sandbox_lease_manager(provider).retain_async(
|
|
owner_id,
|
|
sandbox_id,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
)
|
|
return sandbox_id
|
|
return None
|
|
|
|
async def _release_sandbox_async(
|
|
self,
|
|
sandbox_id: str,
|
|
*,
|
|
owner_id: str | None,
|
|
) -> None:
|
|
provider = get_sandbox_provider()
|
|
if owner_id is not None:
|
|
await get_sandbox_lease_manager(provider).release_async(owner_id)
|
|
return
|
|
await asyncio.to_thread(provider.release, sandbox_id)
|
|
|
|
@override
|
|
def before_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
|
thread_id = (runtime.context or {}).get("thread_id")
|
|
if thread_id is None:
|
|
return super().before_agent(state, runtime)
|
|
self._apply_network_policy_response(state, runtime)
|
|
user_id = resolve_runtime_user_id(runtime)
|
|
projection = self._prepare_agent_skill_projection(thread_id, user_id=user_id)
|
|
owner_id = ensure_sandbox_lease_owner(runtime.context)
|
|
|
|
# Preserve lazy initialization for threads that use the shared view.
|
|
# A policy-scoped view is acquired eagerly so an old shared-view
|
|
# sandbox cannot survive into this run through checkpoint state.
|
|
if self._lazy_init and projection is None:
|
|
# Bind the execution lease only when a sandbox-backed tool actually
|
|
# touches the persisted sandbox. Runs that only answer or return a
|
|
# terminal Command must not leave an unused owner behind when the
|
|
# graph bypasses after_agent.
|
|
return super().before_agent(state, runtime)
|
|
|
|
existing_sandbox_id = self._read_sandbox_id_from_state(state)
|
|
if existing_sandbox_id is None or projection is not None:
|
|
# Phase 3: enforce sandbox:execute authorization before acquiring
|
|
# (eager path). On deny, skip the eager acquisition instead of
|
|
# raising: an exception here is outside any tool call, so it would
|
|
# surface as a run-level graph error rather than the RFC §9
|
|
# friendly ToolMessage. Shared-view runs skip and defer to the lazy
|
|
# gate inside ``ensure_sandbox_initialized``. Policy-scoped runs
|
|
# abort here because retaining an older checkpointed sandbox would
|
|
# bypass the new filesystem view.
|
|
try:
|
|
authorize_sandbox_execution(
|
|
context=runtime.context or {},
|
|
app_config=safe_app_config(),
|
|
)
|
|
except SandboxAuthorizationError:
|
|
if projection is not None:
|
|
# An explicit skill policy cannot leave a checkpointed,
|
|
# previously shared sandbox reusable by downstream tools.
|
|
# Abort this run before the model can reach that state.
|
|
raise
|
|
logger.info("Sandbox execution denied for this role; skipping eager sandbox acquisition (thread_id=%s)", thread_id)
|
|
return None
|
|
provider = get_sandbox_provider()
|
|
self._require_projection_support(provider, projection)
|
|
sandbox_id = self._acquire_sandbox(
|
|
thread_id,
|
|
user_id=user_id,
|
|
owner_id=owner_id,
|
|
)
|
|
if runtime.context is not None:
|
|
runtime.context["sandbox_id"] = sandbox_id
|
|
try:
|
|
if projection is not None:
|
|
provider.sync_agent_skills(
|
|
sandbox_id,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
projection=projection,
|
|
)
|
|
except BaseException:
|
|
if owner_id is not None:
|
|
get_sandbox_lease_manager(provider).release(owner_id)
|
|
else:
|
|
provider.release(sandbox_id)
|
|
if runtime.context is not None:
|
|
runtime.context.pop("sandbox_id", None)
|
|
raise
|
|
logger.info(f"Assigned sandbox {sandbox_id} to thread {thread_id}")
|
|
if existing_sandbox_id == sandbox_id:
|
|
return super().before_agent(state, runtime)
|
|
if existing_sandbox_id is not None:
|
|
return {
|
|
"sandbox": Overwrite({"sandbox_id": sandbox_id}),
|
|
}
|
|
return {"sandbox": {"sandbox_id": sandbox_id}}
|
|
retained_id = self._retain_existing_sandbox(
|
|
state,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
owner_id=owner_id,
|
|
)
|
|
if retained_id is not None and runtime.context is not None:
|
|
runtime.context["sandbox_id"] = retained_id
|
|
return super().before_agent(state, runtime)
|
|
|
|
def _apply_network_policy_response(self, state: SandboxMiddlewareState, runtime: Runtime) -> None:
|
|
sandbox_id = self._read_sandbox_id_from_state(state)
|
|
if sandbox_id is None:
|
|
return
|
|
messages = state.get("messages", [])
|
|
response = None
|
|
for message in reversed(messages):
|
|
if not isinstance(message, HumanMessage):
|
|
continue
|
|
candidate = read_human_input_response(message.additional_kwargs)
|
|
# A network decision is actionable only when it is the current
|
|
# user turn. Stop at the newest HumanMessage so an older persisted
|
|
# card response cannot be re-applied after ordinary conversation.
|
|
if candidate is None or candidate["source"] != NETWORK_POLICY_HUMAN_INPUT_SOURCE:
|
|
return
|
|
response = candidate
|
|
break
|
|
if response is None or response["response_kind"] != "option":
|
|
return
|
|
decision = response["option_id"]
|
|
if decision not in _NETWORK_POLICY_DECISIONS:
|
|
raise SandboxRuntimeError("Invalid sandbox network approval response")
|
|
context = runtime.context or {}
|
|
applied = context.setdefault("sandbox_network_decisions_applied", set())
|
|
marker = (sandbox_id, response["request_id"], decision)
|
|
if marker in applied:
|
|
return
|
|
provider = get_sandbox_provider()
|
|
if provider.sandbox_network_mode() != "allowlist":
|
|
return
|
|
if not provider.decide_network_policy_request(sandbox_id, response["request_id"], decision):
|
|
raise SandboxRuntimeError("The sandbox network approval is stale or does not belong to this sandbox")
|
|
applied.add(marker)
|
|
|
|
@override
|
|
async def abefore_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
|
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)
|
|
user_id = resolve_runtime_user_id(runtime)
|
|
projection = await asyncio.to_thread(
|
|
self._prepare_agent_skill_projection,
|
|
thread_id,
|
|
user_id=user_id,
|
|
)
|
|
owner_id = ensure_sandbox_lease_owner(runtime.context)
|
|
|
|
if self._lazy_init and projection is None:
|
|
return await super().abefore_agent(state, runtime)
|
|
|
|
existing_sandbox_id = self._read_sandbox_id_from_state(state)
|
|
if existing_sandbox_id is None or projection is not None:
|
|
# Phase 3: enforce sandbox:execute authorization before acquiring
|
|
# (eager path, async counterpart of the gate in before_agent). On
|
|
# deny, shared-view runs skip and defer to the lazy tool gate;
|
|
# policy-scoped runs abort before an older sandbox can be reused.
|
|
try:
|
|
await authorize_sandbox_execution_async(
|
|
context=runtime.context or {},
|
|
app_config=await safe_app_config_async(),
|
|
)
|
|
except SandboxAuthorizationError:
|
|
if projection is not None:
|
|
raise
|
|
logger.info("Sandbox execution denied for this role; skipping eager sandbox acquisition (thread_id=%s)", thread_id)
|
|
return None
|
|
provider = get_sandbox_provider()
|
|
self._require_projection_support(provider, projection)
|
|
sandbox_id = await self._acquire_sandbox_async(
|
|
thread_id,
|
|
user_id=user_id,
|
|
owner_id=owner_id,
|
|
)
|
|
if runtime.context is not None:
|
|
runtime.context["sandbox_id"] = sandbox_id
|
|
try:
|
|
if projection is not None:
|
|
await provider.sync_agent_skills_async(
|
|
sandbox_id,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
projection=projection,
|
|
)
|
|
except BaseException:
|
|
await self._release_sandbox_async(
|
|
sandbox_id,
|
|
owner_id=owner_id,
|
|
)
|
|
if runtime.context is not None:
|
|
runtime.context.pop("sandbox_id", None)
|
|
raise
|
|
logger.info(f"Assigned sandbox {sandbox_id} to thread {thread_id}")
|
|
if existing_sandbox_id == sandbox_id:
|
|
return await super().abefore_agent(state, runtime)
|
|
if existing_sandbox_id is not None:
|
|
return {
|
|
"sandbox": Overwrite({"sandbox_id": sandbox_id}),
|
|
}
|
|
return {"sandbox": {"sandbox_id": sandbox_id}}
|
|
retained_id = await self._retain_existing_sandbox_async(
|
|
state,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
owner_id=owner_id,
|
|
)
|
|
if retained_id is not None and runtime.context is not None:
|
|
runtime.context["sandbox_id"] = retained_id
|
|
return await super().abefore_agent(state, runtime)
|
|
|
|
@override
|
|
def after_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
|
sandbox, fork_restored = unwrap_sandbox(state.get("sandbox"))
|
|
if sandbox is not None:
|
|
sandbox_id = sandbox["sandbox_id"]
|
|
if fork_restored:
|
|
# The wrapped value replays the parent thread's sandbox state;
|
|
# releasing it here would evict the parent's warm sandbox.
|
|
logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
|
|
return None
|
|
logger.info(f"Releasing sandbox {sandbox_id}")
|
|
provider = get_sandbox_provider()
|
|
owner_id = sandbox_lease_owner(runtime.context)
|
|
if owner_id is not None:
|
|
get_sandbox_lease_manager(provider).release(owner_id)
|
|
else:
|
|
provider.release(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")
|
|
provider = get_sandbox_provider()
|
|
owner_id = sandbox_lease_owner(runtime.context)
|
|
if owner_id is not None:
|
|
get_sandbox_lease_manager(provider).release(owner_id)
|
|
else:
|
|
provider.release(sandbox_id)
|
|
return None
|
|
|
|
# No sandbox to release
|
|
return super().after_agent(state, runtime)
|
|
|
|
@override
|
|
async def aafter_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
|
sandbox, fork_restored = unwrap_sandbox(state.get("sandbox"))
|
|
if sandbox is not None:
|
|
sandbox_id = sandbox["sandbox_id"]
|
|
if fork_restored:
|
|
# The wrapped value replays the parent thread's sandbox state;
|
|
# releasing it here would evict the parent's warm sandbox.
|
|
logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
|
|
return None
|
|
logger.info(f"Releasing sandbox {sandbox_id}")
|
|
await self._release_sandbox_async(
|
|
sandbox_id,
|
|
owner_id=sandbox_lease_owner(runtime.context),
|
|
)
|
|
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,
|
|
owner_id=sandbox_lease_owner(runtime.context),
|
|
)
|
|
return None
|
|
|
|
# No sandbox to release
|
|
return await super().aafter_agent(state, runtime)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Tool-call wrappers: persist lazily-acquired sandbox state into the
|
|
# graph state via Command(update=...).
|
|
#
|
|
# Background:
|
|
# ``ensure_sandbox_initialized*`` in ``deerflow.sandbox.tools`` mutates
|
|
# ``runtime.state["sandbox"]`` directly. That mutation is local to the
|
|
# current tool invocation and is NOT picked up by LangGraph's channel
|
|
# reducer, so subsequent graph steps (and downstream consumers such as
|
|
# ``ToolOutputBudgetMiddleware`` and the sub-agent ``task_tool``)
|
|
# cannot observe the sandbox id. Wrapping the tool call lets us detect
|
|
# a fresh lazy init by diffing the state snapshot before/after the
|
|
# handler and emit a proper state update via ``Command``.
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _read_sandbox_id_from_state(state: object) -> str | None:
|
|
if not isinstance(state, dict):
|
|
return None
|
|
sandbox_state, _ = unwrap_sandbox(state.get("sandbox"))
|
|
if not isinstance(sandbox_state, dict):
|
|
return None
|
|
sandbox_id = sandbox_state.get("sandbox_id")
|
|
return sandbox_id if isinstance(sandbox_id, str) else None
|
|
|
|
@staticmethod
|
|
def _attach_sandbox_update(result: ToolMessage | Command, sandbox_id: str) -> ToolMessage | Command:
|
|
"""Wrap or merge ``result`` so that ``sandbox.sandbox_id`` is persisted.
|
|
|
|
- ``ToolMessage`` -> ``Command(update={"sandbox": ..., "messages": [msg]})``
|
|
- ``Command`` with dict update -> merge ``sandbox`` key, preserve all
|
|
existing fields (``messages``, ``goto``, ``graph``, ``resume``, ...).
|
|
- ``Command`` with non-dict / None update -> leave it untouched to
|
|
avoid silent data loss on unknown update shapes.
|
|
"""
|
|
sandbox_update = {"sandbox": {"sandbox_id": sandbox_id}}
|
|
|
|
if isinstance(result, ToolMessage):
|
|
return Command(update={**sandbox_update, "messages": [result]})
|
|
|
|
existing_update = result.update
|
|
if isinstance(existing_update, dict):
|
|
merged_update = {**existing_update, **sandbox_update}
|
|
return dc_replace(result, update=merged_update)
|
|
return result
|
|
|
|
@staticmethod
|
|
def _read_sandbox_id_from_request(request: ToolCallRequest) -> str | None:
|
|
"""Read sandbox_id from runtime.state (where ensure_sandbox_initialized writes)."""
|
|
runtime = request.runtime
|
|
if runtime is None or runtime.state is None:
|
|
return None
|
|
return SandboxMiddleware._read_sandbox_id_from_state(runtime.state)
|
|
|
|
@override
|
|
def wrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], ToolMessage | Command],
|
|
) -> ToolMessage | Command:
|
|
prev_sandbox_id = self._read_sandbox_id_from_request(request)
|
|
result = handler(request)
|
|
curr_sandbox_id = self._read_sandbox_id_from_request(request)
|
|
if prev_sandbox_id is None and curr_sandbox_id is not None:
|
|
result = self._attach_sandbox_update(result, curr_sandbox_id)
|
|
return self._maybe_request_network_approval(request, result, curr_sandbox_id or prev_sandbox_id)
|
|
|
|
@override
|
|
async def awrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
|
|
) -> ToolMessage | Command:
|
|
prev_sandbox_id = self._read_sandbox_id_from_request(request)
|
|
result = await handler(request)
|
|
curr_sandbox_id = self._read_sandbox_id_from_request(request)
|
|
if prev_sandbox_id is None and curr_sandbox_id is not None:
|
|
result = self._attach_sandbox_update(result, curr_sandbox_id)
|
|
sandbox_id = curr_sandbox_id or prev_sandbox_id
|
|
if sandbox_id is None:
|
|
return result
|
|
context = getattr(request.runtime, "context", None) or {}
|
|
provider = get_initialized_sandbox_provider()
|
|
if provider is None:
|
|
return result
|
|
if provider.sandbox_network_mode() != "allowlist":
|
|
return result
|
|
if _network_approval_is_non_interactive(context) or context.get("is_subagent"):
|
|
if not await provider.deny_pending_network_policy_events_async(sandbox_id):
|
|
logger.warning("Failed to drain sandbox network policy events for non-interactive sandbox %s", sandbox_id)
|
|
return result
|
|
events = await provider.consume_network_policy_events_async(sandbox_id)
|
|
return self._network_approval_result(request, result, sandbox_id, events)
|
|
|
|
def _maybe_request_network_approval(
|
|
self,
|
|
request: ToolCallRequest,
|
|
result: ToolMessage | Command,
|
|
sandbox_id: str | None,
|
|
) -> ToolMessage | Command:
|
|
if sandbox_id is None:
|
|
return result
|
|
context = getattr(request.runtime, "context", None) or {}
|
|
provider = get_initialized_sandbox_provider()
|
|
if provider is None:
|
|
return result
|
|
if provider.sandbox_network_mode() != "allowlist":
|
|
return result
|
|
if _network_approval_is_non_interactive(context) or context.get("is_subagent"):
|
|
if not provider.deny_pending_network_policy_events(sandbox_id):
|
|
logger.warning("Failed to drain sandbox network policy events for non-interactive sandbox %s", sandbox_id)
|
|
return result
|
|
events = provider.consume_network_policy_events(sandbox_id)
|
|
return self._network_approval_result(request, result, sandbox_id, events)
|
|
|
|
def _network_approval_result(
|
|
self,
|
|
request: ToolCallRequest,
|
|
result: ToolMessage | Command,
|
|
sandbox_id: str,
|
|
events: list[dict[str, object]],
|
|
) -> ToolMessage | Command:
|
|
if not events:
|
|
return result
|
|
event = events[0]
|
|
request_id = event.get("request_id")
|
|
host = event.get("host")
|
|
port = event.get("port")
|
|
if not isinstance(request_id, str) or not isinstance(host, str) or not isinstance(port, int):
|
|
logger.warning("Ignoring malformed trusted sandbox network event: %r", event)
|
|
return result
|
|
tool_call_id = str(request.tool_call.get("id") or "")
|
|
tool_name = str(request.tool_call.get("name") or "sandbox")
|
|
ttl_seconds = get_sandbox_provider().sandbox_network_temporary_grant_ttl()
|
|
ttl_label = f"{ttl_seconds // 60} minutes" if ttl_seconds % 60 == 0 else f"{ttl_seconds} seconds"
|
|
message = ToolMessage(
|
|
id=f"sandbox-network:{request_id}",
|
|
content=(f"Sandbox network policy blocked {host}:{port}. The command was not retried. Choose whether this destination should be available, then ask the agent to retry if appropriate."),
|
|
tool_call_id=tool_call_id,
|
|
name=tool_name,
|
|
artifact={
|
|
"human_input": {
|
|
"version": 1,
|
|
"kind": "human_input_request",
|
|
"source": NETWORK_POLICY_HUMAN_INPUT_SOURCE,
|
|
"request_id": request_id,
|
|
"tool_call_id": tool_call_id,
|
|
"clarification_type": "risk_confirmation",
|
|
"title": "Sandbox network access",
|
|
"question": f"Allow this sandbox to connect to {host}:{port}?",
|
|
"context": "Private, loopback, link-local, multicast, and cloud metadata addresses can never be approved.",
|
|
"input_mode": "single_choice",
|
|
"options": [
|
|
{"id": "deny", "label": "Deny", "value": "Deny network access"},
|
|
{"id": "allow_temporary", "label": f"Allow for {ttl_label}", "value": f"Allow network access for {ttl_label}"},
|
|
{"id": "allow_sandbox", "label": "Allow for this sandbox", "value": "Allow network access for this sandbox"},
|
|
],
|
|
}
|
|
},
|
|
)
|
|
update: dict = {"messages": [message], "sandbox": {"sandbox_id": sandbox_id}}
|
|
if isinstance(result, Command) and isinstance(result.update, dict):
|
|
update = {**result.update, **update}
|
|
return Command(update=update, goto=END)
|