From 9e0fbd60fa2d565848ece12d03fec4e5ea51eeec Mon Sep 17 00:00:00 2001 From: Aari Date: Wed, 2 Sep 2026 21:05:23 +0800 Subject: [PATCH] fix(sandbox): isolate concurrent subagent shell sessions (#5134) * fix(sandbox): isolate concurrent subagent shell sessions * fix(sandbox): make execution acquire idempotent * fix(sandbox): close execution lifecycle gaps * fix(sandbox): serialize retained client lifecycle * fix(sandbox): close remaining client lifecycle gaps * fix(sandbox): unwind failed client lookup * fix(sandbox): protect internal lease identities * fix(sandbox): make cancellation reconciliation durable * fix(sandbox): fence cancelled workers and IM uploads --- backend/app/channels/AGENTS.md | 1 + backend/app/channels/dingtalk.py | 30 +- backend/app/channels/feishu.py | 19 +- backend/app/channels/sandbox_files.py | 43 ++ backend/app/gateway/AGENTS.md | 4 +- backend/app/gateway/authz.py | 64 +- backend/app/gateway/routers/artifacts.py | 21 +- backend/app/gateway/routers/uploads.py | 229 +++--- backend/app/gateway/services.py | 32 +- backend/packages/harness/deerflow/AGENTS.md | 2 +- backend/packages/harness/deerflow/client.py | 20 +- .../community/aio_sandbox/aio_sandbox.py | 253 +++++- .../harness/deerflow/runtime/runs/worker.py | 62 +- .../harness/deerflow/sandbox/AGENTS.md | 5 +- .../harness/deerflow/sandbox/lease.py | 698 +++++++++++++++++ .../harness/deerflow/sandbox/middleware.py | 212 ++++- .../harness/deerflow/sandbox/sandbox.py | 21 + .../deerflow/sandbox/sandbox_provider.py | 11 + .../harness/deerflow/sandbox/tools.py | 179 ++++- .../harness/deerflow/subagents/AGENTS.md | 1 + .../harness/deerflow/subagents/executor.py | 24 + .../blocking_io/test_dingtalk_receive_file.py | 88 +++ .../blocking_io/test_feishu_receive_file.py | 77 ++ backend/tests/test_aio_sandbox.py | 275 ++++++- backend/tests/test_artifacts_router.py | 36 + backend/tests/test_client.py | 81 ++ .../tests/test_ensure_sandbox_initialized.py | 359 ++++++++- backend/tests/test_gateway_services.py | 18 + backend/tests/test_github_token_plumbing.py | 51 ++ backend/tests/test_run_worker_rollback.py | 148 ++++ backend/tests/test_sandbox_leases.py | 731 ++++++++++++++++++ backend/tests/test_sandbox_middleware.py | 180 ++++- backend/tests/test_subagent_executor.py | 124 +++ backend/tests/test_uploads_router.py | 79 +- 34 files changed, 3903 insertions(+), 275 deletions(-) create mode 100644 backend/app/channels/sandbox_files.py create mode 100644 backend/packages/harness/deerflow/sandbox/lease.py create mode 100644 backend/tests/test_sandbox_leases.py diff --git a/backend/app/channels/AGENTS.md b/backend/app/channels/AGENTS.md index c3bd9bc8b..4fee3fb44 100644 --- a/backend/app/channels/AGENTS.md +++ b/backend/app/channels/AGENTS.md @@ -40,6 +40,7 @@ Bridges external messaging platforms (Feishu, Slack, Telegram, Discord, DingTalk **Owner-scoped file storage**: inbound files, uploads, and output artifacts are staged under the DeerFlow owner's bucket so they land where the agent run reads/writes (`users/{user_id}/threads/{thread_id}/user-data/{uploads,outputs}`). `ChannelManager._handle_chat` resolves the storage owner once via `_channel_storage_user_id(msg)` (sanitized owner id, falling back to `safe(msg.user_id)` for unbound auth-enabled channels — mirroring `_resolve_run_params`'s run identity; `None` only when no identity is available) and threads it as the `user_id=` kwarg through the file pipeline: - `Channel.receive_file(msg, thread_id, user_id=...)` — owner-bound channels persist downloaded files under the owner's bucket instead of the default bucket - `FeishuChannel._receive_single_file(...)` / `DingTalkChannel._receive_single_file(...)` — normalize provider filenames, claim a collision-free basename and write it through `write_upload_file_no_symlink` under the same channel lock; the returned basename drives both the agent-visible virtual path and non-local sandbox sync +- `sandbox_files.py` — non-mounted Feishu/DingTalk syncs acquire unique non-releasing execution holders, drain blocking `update_file` workers across repeated cancellation, and release only after the last sandbox operation, so a parallel run cannot close the shared client mid-upload - `_ingest_inbound_files(...)` and the underlying `ensure_uploads_dir` / `get_uploads_dir` — owner-scoped via the same kwarg - `_resolve_attachments` / `_prepare_artifact_delivery` — resolve output artifacts from the bound owner's bucket The cached value is reused for both the blocking (`runs.wait`) and streaming (`_handle_streaming_chat`) paths, so uploads and artifact delivery always target the same bucket even if a channel returns a rewritten `InboundMessage` from `receive_file`. The bucket id matches the memory bucket resolved by `_resolve_memory_user_id` (both normalize through `make_safe_user_id`). diff --git a/backend/app/channels/dingtalk.py b/backend/app/channels/dingtalk.py index ceed2dae4..1c8ad48b1 100644 --- a/backend/app/channels/dingtalk.py +++ b/backend/app/channels/dingtalk.py @@ -17,6 +17,7 @@ from app.channels.base import Channel from app.channels.commands import is_known_channel_command, strip_leading_mentions from app.channels.connection_identity import attach_connection_identity from app.channels.message_bus import InboundMessage, InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment +from app.channels.sandbox_files import sync_file_to_thread_sandbox from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths from deerflow.runtime.user_context import get_effective_user_id from deerflow.sandbox.sandbox_provider import get_sandbox_provider @@ -676,20 +677,21 @@ class DingTalkChannel(Channel): virtual_path = f"{VIRTUAL_PATH_PREFIX}/uploads/{resolved_target.name}" try: - sandbox_provider = get_sandbox_provider() - # acquire_async keeps provider lifecycle work (Docker discovery, - # readiness polls) off the event loop; update_file is blocking - # transport IO on remote sandboxes, so it is offloaded too. - sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id) - if sandbox_id != "local": - sandbox = sandbox_provider.get(sandbox_id) - if sandbox is None: - # Mirror Feishu: the agent's non-local sandbox cannot see this - # file, so returning the virtual path would hand the model a - # path that reads as nothing — surface a failed-load marker. - logger.warning("[DingTalk] sandbox %s not found after acquire, dropping attachment: %s", sandbox_id, virtual_path) - return "" - await asyncio.to_thread(sandbox.update_file, virtual_path, content) + sandbox_provider = await asyncio.to_thread(get_sandbox_provider) + synced = await sync_file_to_thread_sandbox( + sandbox_provider, + thread_id=thread_id, + user_id=effective_user_id, + virtual_path=virtual_path, + content=content, + owner_prefix="dingtalk-upload", + ) + if not synced: + # Mirror Feishu: the agent's non-local sandbox cannot see this + # file, so returning the virtual path would hand the model a + # path that reads as nothing — surface a failed-load marker. + logger.warning("[DingTalk] sandbox not found after acquire, dropping attachment: %s", virtual_path) + return "" except Exception: # Same failure mode as the sandbox-is-None branch: the bytes never # reached the agent's sandbox, so the virtual path would read as diff --git a/backend/app/channels/feishu.py b/backend/app/channels/feishu.py index f75e1c920..de271341c 100644 --- a/backend/app/channels/feishu.py +++ b/backend/app/channels/feishu.py @@ -23,6 +23,7 @@ from app.channels.message_bus import ( OutboundMessage, ResolvedAttachment, ) +from app.channels.sandbox_files import sync_file_to_thread_sandbox from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths from deerflow.runtime.user_context import get_effective_user_id from deerflow.sandbox.sandbox_provider import get_sandbox_provider @@ -495,13 +496,17 @@ class FeishuChannel(Channel): try: sandbox_provider = await asyncio.to_thread(get_sandbox_provider) - if not getattr(sandbox_provider, "uses_thread_data_mounts", False): - sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id) - sandbox = sandbox_provider.get(sandbox_id) - if sandbox is None: - logger.warning("[Feishu] sandbox not found for thread_id=%s", thread_id) - return f"Failed to obtain the [{type}]" - await asyncio.to_thread(sandbox.update_file, virtual_path, content) + synced = await sync_file_to_thread_sandbox( + sandbox_provider, + thread_id=thread_id, + user_id=effective_user_id, + virtual_path=virtual_path, + content=content, + owner_prefix="feishu-upload", + ) + if not synced: + logger.warning("[Feishu] sandbox not found for thread_id=%s", thread_id) + return f"Failed to obtain the [{type}]" except Exception: logger.exception("[Feishu] failed to sync resource into non-local sandbox: %s", virtual_path) return f"Failed to obtain the [{type}]" diff --git a/backend/app/channels/sandbox_files.py b/backend/app/channels/sandbox_files.py new file mode 100644 index 000000000..eaa8a5e7f --- /dev/null +++ b/backend/app/channels/sandbox_files.py @@ -0,0 +1,43 @@ +"""Shared lifecycle-safe synchronization for inbound channel attachments.""" + +from __future__ import annotations + +from deerflow.sandbox.lease import acquire_sandbox_client_lease + + +async def sync_file_to_thread_sandbox( + sandbox_provider, + *, + thread_id: str, + user_id: str, + virtual_path: str, + content: bytes, + owner_prefix: str, +) -> bool: + """Copy one attachment while holding a non-releasing sandbox client lease. + + Thread-data mount providers already see the persisted upload. Other + providers need a unique holder so a parallel run cannot close their client + during ``update_file``. The blocking transport worker is drained even when + the channel handler is repeatedly cancelled, and only then is the holder + released. + """ + if getattr(sandbox_provider, "uses_thread_data_mounts", False): + return True + + lease = await acquire_sandbox_client_lease( + sandbox_provider, + thread_id, + user_id=user_id, + owner_prefix=owner_prefix, + release_on_last=False, + ) + try: + if lease.sandbox_id == "local" or lease.sandbox_id.startswith("local:"): + return True + if lease.sandbox is None: + return False + await lease.run_sync(lease.sandbox.update_file, virtual_path, content) + return True + finally: + await lease.release() diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index 9955d9a8b..7a820ae65 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -52,9 +52,9 @@ owner-scoped assistant version selection remains enabled. | **Subagents** (`/api/subagents`) | Admin managed-worker CRUD and listing. | | **Integrations** (`/api/integrations`) | `GET /lark/status` - inspect managed Lark/Feishu CLI integration state, including `sandbox_runtime_mode` / `sandbox_runtime_ready` (whether `lark-cli` will actually be present in the sandbox at chat time); `POST /lark/install` - admin-only install of the official `lark-*` managed skill pack; `POST /lark/config/start` and `/lark/config/complete` - internal first-time Lark connection setup; `POST /lark/config/credentials` - atomically switch the caller's per-user Lark app after validating the new `app_id`/`app_secret` through the official CLI's live tenant-token probe, revoke/remove the previous OAuth tokens, and restore the prior credential tree if the switch fails; `POST /lark/auth/start` and `/lark/auth/complete` - browser device-flow user authorization without terminal access, with optional `domains` / exact `scope` for incremental permission grants. Config and auth flows carry a server-issued, per-user generation persisted under the credential lock; a rejected direct switch leaves the current generation unchanged, stale completions return 409, and browser re-registration uses the same token-clearing/revocation transaction as direct credential switches. | | **Memory** (`/api/memory`) | `GET /` - memory data; `POST /reload` - force reload; `GET /config` - config; `GET /status` - config + data | -| **Uploads** (`/api/threads/{id}/uploads`) | `POST /` - upload files (auto-converts PDF/PPT/Excel/Word); `GET /list` - list; `DELETE /{filename}` - delete | +| **Uploads** (`/api/threads/{id}/uploads`) | `POST /` - upload files (auto-converts PDF/PPT/Excel/Word); non-mounted sandbox sync uses a non-releasing request lease; `GET /list` - list; `DELETE /{filename}` - delete | | **Threads** (`/api/threads/{id}`) | `DELETE /` - remove DeerFlow-managed local thread data after LangGraph thread deletion; `POST /branches` - branch a completed assistant turn with a replay checkpoint; inherited titles take next-free displayed sibling suffixes, including explicit/renamed ones, while explicit titles stay unchanged. Durable `branch` admission rejects races. Workspace files are not checkpointed, so the branch only best-effort copies the current workspace when branching from the **latest** turn (`workspace_clone_mode="current_thread_best_effort"`); branching from an older/historical turn skips the copy (`workspace_clone_mode="skipped_historical_turn"`) so the branch never inherits files that only exist in a later timeline. Thread-scoped runtime channels (`sandbox`, `thread_data`) are not copied onto the branch: the parent's `sandbox_id` binds path mappings and the release lifecycle to the parent's workspace, so the branch lazily acquires its own sandbox instead. Branch creation also seeds the new thread's run-event feed from the branch checkpoint's visible messages (`history_seed_mode` in the response): the thread feed reads run_events, not checkpoints, so without the seed the inherited history disappears from the UI after the branch's first run (#4380). Seeded rows are grouped into one synthetic run per inherited turn (`branch-seed-{thread_id}-{n}`, a new turn opening at every persisted human message, including an allowlisted hidden `ask_clarification` reply) because `run_id` is a turn identity to the feed's consumers, not a provenance tag: regenerating an inherited answer supersedes that row's whole `run_id` in `GET /messages/page`, so one shared id for the entire seed deleted the complete inherited history on a branch's first regenerate (#4458); `GET /goal`, `PUT /goal`, `DELETE /goal` - read, set, and clear the active thread goal; `POST /compact` - manually summarize older active context into `summary_text` and retain the recent message window, blocked while a run is in flight; unexpected failures are logged server-side and return a generic 500 detail | -| **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - stream regular text and binary artifacts with `FileResponse`, including byte-`Range` 206/416 behavior used by bounded text previews and media seeking; active content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) are always forced as download attachments to reduce XSS risk; `?download=true` still forces download for other file types. `PUT /{path}` atomically replaces an existing UTF-8 text file under `/mnt/user-data/outputs` when its expected SHA-256 still matches; active runs conflict, and non-mounted sandbox providers receive the same update explicitly. Atomic replacement applies the existing POSIX permission handling when descriptor-based APIs are available and otherwise keeps the platform-native temporary-file permissions (Windows). | +| **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - stream regular text and binary artifacts with `FileResponse`, including byte-`Range` 206/416 behavior used by bounded text previews and media seeking; active content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) are always forced as download attachments to reduce XSS risk; `?download=true` still forces download for other file types. `PUT /{path}` atomically replaces an existing UTF-8 text file under `/mnt/user-data/outputs` when its expected SHA-256 still matches; active runs conflict, and non-mounted sandbox providers receive the same update under a request lease. Atomic replacement applies the existing POSIX permission handling when descriptor-based APIs are available and otherwise keeps the platform-native temporary-file permissions (Windows). | | **Suggestions** (`/api/suggestions`) | `GET /config` - returns global suggestions config boolean; `POST /threads/{id}/suggestions` - generate follow-up questions; rich list/block model content is normalized and inline reasoning (`...`, including unclosed/truncated blocks from reasoning models like MiniMax-M3) is stripped before JSON parsing | | **Input Polish** (`/api/input-polish`) | `POST /` - rewrite a composer draft before it is sent. This is a short authenticated `runs:create` LLM request using `input_polish` config; it does not create a LangGraph run, persist a message, or modify thread state. Shares the non-graph one-shot LLM path (`deerflow.utils.oneshot_llm.run_oneshot_llm`) with the suggestions route so model build + Langfuse metadata + invoke stay in one place; validates the same stripped view of the draft it sends to the model, and preserves literal `` substrings in the rewrite (`strip_think_blocks(truncate_unclosed=False)`) | | **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block. Before the first journaled run, seed an empty feed from a checkpoint so legacy checkpoint-only history keeps its order and visibility; skip absent checkpoints or populated feeds. `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest completed or interrupted assistant answer, carrying the latest non-empty thread title in graph input so resuming an older checkpoint cannot roll back a later manual rename (#4457); `POST /edit-regenerate/prepare` - prepare a checkpoint replay from the latest editable human turn with a replacement user message and edit replay metadata; it carries the current thread title the same way, but only when the replay base already has one — an untitled base belongs to a thread the title middleware has not named yet, so pinning the current title there would keep a name generated from the prompt the edit just replaced; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/stream` hides action/wait; GET action 405 pre-owner; POST needs `runs:cancel`; `GET /{rid}/messages` - paginated per-run messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET/POST /{rid}/artifacts/archive` - receipt manifest / bounded ZIP; `GET /../messages` - legacy thread message array; `GET /../messages/page` - backward thread-global `seq` history page with middleware/subagent-AI/successful-regenerate/edit-replay filtering and page-run-scoped feedback enrichment; subagent AI callbacks remain available through run events while parent `task` ToolMessages stay visible for card restoration; `GET /../token-usage` - aggregate tokens plus an optional `context_usage` percentage. Context usage approximately counts messages from the latest materialized thread state through `build_thread_checkpoint_state_accessor`, so full and delta checkpoint modes expose the same input. The percentage uses the latest run's model and its `context_window`. | diff --git a/backend/app/gateway/authz.py b/backend/app/gateway/authz.py index 9f9104faa..269fbbb90 100644 --- a/backend/app/gateway/authz.py +++ b/backend/app/gateway/authz.py @@ -33,7 +33,9 @@ import asyncio import functools import inspect import logging +import uuid from collections.abc import Callable +from dataclasses import dataclass from types import SimpleNamespace from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar @@ -398,6 +400,27 @@ def authorize_sandbox_for_request( raise SandboxAuthorizationError(role=context.get("user_role")) from None +@dataclass(slots=True) +class SandboxRequestLease: + """One Gateway request's process-local use of a sandbox client.""" + + sandbox: object | None + sandbox_id: str | None + denied: bool + owner_id: str | None + provider: object | None + + async def release(self) -> None: + """Drop the request holder without bypassing concurrent executions.""" + if self.owner_id is None or self.provider is None: + return + from deerflow.sandbox.lease import get_sandbox_lease_manager + + owner_id = self.owner_id + self.owner_id = None + await get_sandbox_lease_manager(self.provider).release_async(owner_id) + + async def try_acquire_sandbox_for_request( request: Request, sandbox_provider, @@ -405,19 +428,21 @@ async def try_acquire_sandbox_for_request( *, user_id: str, app_config: AppConfig | None, -) -> tuple[object, str | None, bool]: + owner_prefix: str = "gateway", + release_on_last: bool = True, +) -> SandboxRequestLease: """Gate + acquire the thread sandbox for a Gateway sync path. Single entry point for the uploads/artifacts sandbox-sync paths so the deny/skip semantics live in one place: runs the ``sandbox:execute`` gate - for the request's user, then acquires the sandbox. Returns - ``(sandbox, sandbox_id, denied)``: + for the request's user, then acquires the sandbox under a unique request + holder. Callers must await :meth:`SandboxRequestLease.release` after their + last client operation. - - denied role → ``(None, None, True)``: acquisition was skipped by policy; + - denied role → no sandbox/owner and ``denied=True``: acquisition was skipped by policy; the primary operation (upload / artifact edit) proceeds without the sandbox copy. - - allowed → ``(sandbox, sandbox_id, False)``: ``sandbox`` is the acquired - instance (``sandbox_id`` for later release), or ``sandbox is None`` when + - allowed → ``sandbox`` is the acquired instance, or ``sandbox is None`` when the provider lost it right after acquiring (infrastructure error — callers surface it as 500 / RuntimeError respectively, since that is not a policy decision). @@ -434,9 +459,30 @@ async def try_acquire_sandbox_for_request( authorize_sandbox_for_request(user, is_internal=_is_internal_caller(request, user), app_config=app_config) except SandboxAuthorizationError: logger.info("Sandbox sync skipped: sandbox execution not permitted for this caller (thread_id=%s)", thread_id) - return None, None, True - sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=user_id) - return sandbox_provider.get(sandbox_id), sandbox_id, False + return SandboxRequestLease( + sandbox=None, + sandbox_id=None, + denied=True, + owner_id=None, + provider=None, + ) + + from deerflow.sandbox.lease import get_sandbox_lease_manager + + owner_id = f"{owner_prefix}:{uuid.uuid4()}" + sandbox_id = await get_sandbox_lease_manager(sandbox_provider).acquire_async( + owner_id, + thread_id, + user_id=user_id, + release_on_last=release_on_last, + ) + return SandboxRequestLease( + sandbox=sandbox_provider.get(sandbox_id), + sandbox_id=sandbox_id, + denied=False, + owner_id=owner_id, + provider=sandbox_provider, + ) async def _authenticate(request: Request) -> AuthContext: diff --git a/backend/app/gateway/routers/artifacts.py b/backend/app/gateway/routers/artifacts.py index 940c96e7b..f304af728 100644 --- a/backend/app/gateway/routers/artifacts.py +++ b/backend/app/gateway/routers/artifacts.py @@ -16,7 +16,7 @@ from fastapi import APIRouter, HTTPException, Request from fastapi.responses import FileResponse, Response from pydantic import BaseModel, Field -from app.gateway.authz import require_permission, try_acquire_sandbox_for_request +from app.gateway.authz import SandboxRequestLease, require_permission, try_acquire_sandbox_for_request from app.gateway.deps import get_run_manager from app.gateway.internal_auth import get_trusted_internal_owner_user_id from app.gateway.path_utils import resolve_thread_virtual_path @@ -511,8 +511,7 @@ async def update_artifact( raw_owner_user_id = get_trusted_internal_owner_user_id(request) effective_user_id = make_safe_user_id(raw_owner_user_id) if raw_owner_user_id else get_effective_user_id() - sandbox_provider = None - sandbox_id: str | None = None + sandbox_lease: SandboxRequestLease | None = None sandbox = None try: async with reserve_artifact_write(request, thread_id, user_id=effective_user_id): @@ -536,14 +535,16 @@ async def update_artifact( # role skips the sandbox sync; the host-side artifact update # still completes (the agent cannot consume the sandbox copy # anyway when sandbox execution is denied). - sandbox, sandbox_id, sandbox_denied = await try_acquire_sandbox_for_request( + sandbox_lease = await try_acquire_sandbox_for_request( request, sandbox_provider, thread_id, user_id=effective_user_id, app_config=safe_app_config(), + owner_prefix="gateway:artifact", ) - if not sandbox_denied and sandbox is None: + sandbox = sandbox_lease.sandbox + if not sandbox_lease.denied and sandbox is None: raise RuntimeError("Failed to acquire sandbox for artifact update") try: @@ -569,11 +570,15 @@ async def update_artifact( logger.exception("Failed to update artifact %s for thread %s", path, thread_id) raise HTTPException(status_code=500, detail="Failed to update artifact") from None finally: - if sandbox_id is not None and sandbox_provider is not None: + if sandbox_lease is not None: try: - await asyncio.to_thread(sandbox_provider.release, sandbox_id) + await sandbox_lease.release() except Exception: - logger.warning("Failed to release sandbox after artifact update: %s", sandbox_id, exc_info=True) + logger.warning( + "Failed to release sandbox request lease after artifact update: %s", + sandbox_lease.sandbox_id, + exc_info=True, + ) content_sha256 = hashlib.sha256(updated).hexdigest() return ArtifactUpdateResponse( diff --git a/backend/app/gateway/routers/uploads.py b/backend/app/gateway/routers/uploads.py index e3e8c4289..5b09d9e3e 100644 --- a/backend/app/gateway/routers/uploads.py +++ b/backend/app/gateway/routers/uploads.py @@ -11,7 +11,7 @@ from typing import BinaryIO from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile from pydantic import BaseModel, Field -from app.gateway.authz import require_permission, try_acquire_sandbox_for_request +from app.gateway.authz import SandboxRequestLease, require_permission, try_acquire_sandbox_for_request from app.gateway.deps import get_config from deerflow.config.app_config import AppConfig from deerflow.config.paths import get_paths @@ -338,126 +338,143 @@ async def upload_files( sandbox_provider = get_sandbox_provider() sync_to_sandbox = not _uses_thread_data_mounts(sandbox_provider) + sandbox_lease: SandboxRequestLease | None = None sandbox = None - if sync_to_sandbox: - # Phase 3: enforce sandbox:execute before acquiring — a role denied - # sandbox execution must not trigger sandbox allocation just by - # uploading files. Deny skips the sync; the upload itself still - # succeeds (files stay in the thread uploads dir; the agent cannot - # consume them via sandbox anyway). - sandbox, _sandbox_id, sandbox_denied = await try_acquire_sandbox_for_request( - request, - sandbox_provider, - thread_id, - user_id=effective_user_id, - app_config=config, - ) - if not sandbox_denied and sandbox is None: - raise HTTPException(status_code=500, detail="Failed to acquire sandbox") - auto_convert_documents = _auto_convert_documents_enabled(config) - - for file in files: - if not file.filename: - continue - - try: - original_filename = normalize_filename(file.filename) - safe_filename = claim_unique_filename(original_filename, seen_filenames) - except ValueError: - logger.warning(f"Skipping file with unsafe filename: {file.filename!r}") - continue - - try: - file_path, file_size, total_size = await _write_upload_file_with_limits( - file, - uploads_dir=uploads_dir, - display_filename=safe_filename, - max_single_file_size=limits.max_file_size, - max_total_size=limits.max_total_size, - total_size=total_size, + try: + if sync_to_sandbox: + # Phase 3: enforce sandbox:execute before acquiring — a role denied + # sandbox execution must not trigger sandbox allocation just by + # uploading files. Deny skips the sync; the upload itself still + # succeeds (files stay in the thread uploads dir; the agent cannot + # consume them via sandbox anyway). + sandbox_lease = await try_acquire_sandbox_for_request( + request, + sandbox_provider, + thread_id, + user_id=effective_user_id, + app_config=config, + owner_prefix="gateway:upload", + release_on_last=False, ) - written_paths.append(file_path) + sandbox = sandbox_lease.sandbox + if not sandbox_lease.denied and sandbox is None: + raise HTTPException(status_code=500, detail="Failed to acquire sandbox") - virtual_path = upload_virtual_path(safe_filename) + auto_convert_documents = _auto_convert_documents_enabled(config) - if sync_to_sandbox: - sandbox_sync_targets.append((file_path, virtual_path)) + for file in files: + if not file.filename: + continue - file_info = { - "filename": safe_filename, - "size": file_size, - "path": str(sandbox_uploads / safe_filename), - "virtual_path": virtual_path, - "artifact_url": upload_artifact_url(thread_id, safe_filename), - } - if safe_filename != original_filename: - file_info["original_filename"] = original_filename + try: + original_filename = normalize_filename(file.filename) + safe_filename = claim_unique_filename(original_filename, seen_filenames) + except ValueError: + logger.warning(f"Skipping file with unsafe filename: {file.filename!r}") + continue - logger.info(f"Saved file: {safe_filename} ({file_size} bytes) to {file_info['path']}") + try: + file_path, file_size, total_size = await _write_upload_file_with_limits( + file, + uploads_dir=uploads_dir, + display_filename=safe_filename, + max_single_file_size=limits.max_file_size, + max_total_size=limits.max_total_size, + total_size=total_size, + ) + written_paths.append(file_path) - file_ext = file_path.suffix.lower() - if auto_convert_documents and file_ext in CONVERTIBLE_EXTENSIONS: - # Reserve the companion .md name in this request's seen set - # before writing so conversion cannot silently truncate another - # uploaded or derived file (same invariant as form-part dedupe). - provisional_md_name = Path(safe_filename).with_suffix(".md").name - unique_md_name = claim_unique_filename(provisional_md_name, seen_filenames) - md_output = file_path.with_name(unique_md_name) - md_path = await convert_file_to_markdown(file_path, output_path=md_output) - if md_path: - written_paths.append(md_path) - md_virtual_path = upload_virtual_path(md_path.name) + virtual_path = upload_virtual_path(safe_filename) - if sync_to_sandbox: - sandbox_sync_targets.append((md_path, md_virtual_path)) + if sync_to_sandbox: + sandbox_sync_targets.append((file_path, virtual_path)) - file_info["markdown_file"] = md_path.name - file_info["markdown_path"] = str(sandbox_uploads / md_path.name) - file_info["markdown_virtual_path"] = md_virtual_path - file_info["markdown_artifact_url"] = upload_artifact_url(thread_id, md_path.name) - else: - # Conversion failed and wrote nothing, so release the claim; - # holding it would rename a later same-stem upload against - # a name nothing occupies. - seen_filenames.discard(unique_md_name) + file_info = { + "filename": safe_filename, + "size": file_size, + "path": str(sandbox_uploads / safe_filename), + "virtual_path": virtual_path, + "artifact_url": upload_artifact_url(thread_id, safe_filename), + } + if safe_filename != original_filename: + file_info["original_filename"] = original_filename - uploaded_files.append(file_info) + logger.info(f"Saved file: {safe_filename} ({file_size} bytes) to {file_info['path']}") - except HTTPException as e: - await run_file_io(_cleanup_uploaded_paths, written_paths) - raise e - except UnsafeUploadPathError as e: - logger.warning("Skipping upload with unsafe destination %s: %s", file.filename, e) - skipped_files.append(safe_filename) - continue - except Exception as e: - logger.error(f"Failed to upload {file.filename}: {e}") - await run_file_io(_cleanup_uploaded_paths, written_paths) - raise HTTPException(status_code=500, detail=f"Failed to upload {file.filename}: {str(e)}") + file_ext = file_path.suffix.lower() + if auto_convert_documents and file_ext in CONVERTIBLE_EXTENSIONS: + # Reserve the companion .md name in this request's seen set + # before writing so conversion cannot silently truncate another + # uploaded or derived file (same invariant as form-part dedupe). + provisional_md_name = Path(safe_filename).with_suffix(".md").name + unique_md_name = claim_unique_filename(provisional_md_name, seen_filenames) + md_output = file_path.with_name(unique_md_name) + md_path = await convert_file_to_markdown(file_path, output_path=md_output) + if md_path: + written_paths.append(md_path) + md_virtual_path = upload_virtual_path(md_path.name) - # Uploaded files are created with 0o600 permissions (owner read/write only). - # In Docker sandbox deployments the gateway writes as root but the sandbox - # process runs as a non-root user (typically UID 1000). Without group/other - # read bits the sandbox cannot access the files — whether the uploads - # directory is bind-mounted into the container or synced via - # sandbox.update_file. Always add group/other read bits so every sandbox - # configuration can read the uploaded content. - await run_file_io(_make_uploaded_paths_sandbox_readable, written_paths) + if sync_to_sandbox: + sandbox_sync_targets.append((md_path, md_virtual_path)) - if sync_to_sandbox and sandbox is not None: - for file_path, virtual_path in sandbox_sync_targets: - await run_file_io(_sync_upload_to_sandbox, sandbox, file_path, virtual_path) + file_info["markdown_file"] = md_path.name + file_info["markdown_path"] = str(sandbox_uploads / md_path.name) + file_info["markdown_virtual_path"] = md_virtual_path + file_info["markdown_artifact_url"] = upload_artifact_url(thread_id, md_path.name) + else: + # Conversion failed and wrote nothing, so release the claim; + # holding it would rename a later same-stem upload against + # a name nothing occupies. + seen_filenames.discard(unique_md_name) - message = f"Successfully uploaded {len(uploaded_files)} file(s)" - if skipped_files: - message += f"; skipped {len(skipped_files)} unsafe file(s)" + uploaded_files.append(file_info) - return UploadResponse( - success=not skipped_files, - files=uploaded_files, - message=message, - skipped_files=skipped_files, - ) + except HTTPException as e: + await run_file_io(_cleanup_uploaded_paths, written_paths) + raise e + except UnsafeUploadPathError as e: + logger.warning("Skipping upload with unsafe destination %s: %s", file.filename, e) + skipped_files.append(safe_filename) + continue + except Exception as e: + logger.error(f"Failed to upload {file.filename}: {e}") + await run_file_io(_cleanup_uploaded_paths, written_paths) + raise HTTPException(status_code=500, detail=f"Failed to upload {file.filename}: {str(e)}") + + # Uploaded files are created with 0o600 permissions (owner read/write only). + # In Docker sandbox deployments the gateway writes as root but the sandbox + # process runs as a non-root user (typically UID 1000). Without group/other + # read bits the sandbox cannot access the files — whether the uploads + # directory is bind-mounted into the container or synced via + # sandbox.update_file. Always add group/other read bits so every sandbox + # configuration can read the uploaded content. + await run_file_io(_make_uploaded_paths_sandbox_readable, written_paths) + + if sync_to_sandbox and sandbox is not None: + for file_path, virtual_path in sandbox_sync_targets: + await run_file_io(_sync_upload_to_sandbox, sandbox, file_path, virtual_path) + + message = f"Successfully uploaded {len(uploaded_files)} file(s)" + if skipped_files: + message += f"; skipped {len(skipped_files)} unsafe file(s)" + + return UploadResponse( + success=not skipped_files, + files=uploaded_files, + message=message, + skipped_files=skipped_files, + ) + + finally: + if sandbox_lease is not None: + try: + await sandbox_lease.release() + except Exception: + logger.warning( + "Failed to release sandbox request lease after upload sync: %s", + sandbox_lease.sandbox_id, + exc_info=True, + ) @router.get("/limits", response_model=UploadLimits) diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index f5c6e44d1..e6102c43c 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -77,6 +77,7 @@ from deerflow.runtime.secret_context import ( ) from deerflow.runtime.stream_modes import normalize_stream_modes from deerflow.runtime.user_context import reset_current_user, set_current_user +from deerflow.sandbox.lease import SANDBOX_SERVER_OWNED_CONTEXT_KEYS from deerflow.subagents.status_contract import SUBAGENT_ACCEPTANCE_VERDICT_KEY, SUBAGENT_RECEIPT_VERDICT_KEY, SUBAGENT_TOOL_RECEIPTS_KEY from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_context, ensure_trace_id from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY @@ -408,22 +409,27 @@ _CONTEXT_CONFIGURABLE_KEYS: frozenset[str] = frozenset( # arbitrary HTTP/IM clients must not be able to force autonomous execution. _CONTEXT_INTERNAL_CALLER_KEYS: frozenset[str] = frozenset({"non_interactive"}) -# Server-owned authorization identity fields. These must never be accepted from -# client-supplied ``body.config.context`` or ``body.config.configurable``. They +# Server-owned authorization and sandbox lifecycle identity fields. These must +# never be accepted from client-supplied ``body.config.context`` or +# ``body.config.configurable``. They # are either produced by Gateway auth state, admitted from a separately # authenticated internal request channel, or reserved for LangGraph Server. # ``is_internal`` — derived from ``request.state.auth_source`` # ``authz_attributes`` — Phase 1A has no Gateway-side producer; cleared. # ``channel_user_id`` — accepted only from trusted internal context. # ``langgraph_auth_user*`` — populated only by LangGraph Server auth. -_SERVER_OWNED_AUTHZ_CONTEXT_KEYS: frozenset[str] = frozenset( - { - "is_internal", - "authz_attributes", - "channel_user_id", - "langgraph_auth_user", - "langgraph_auth_user_id", - } +# ``sandbox_*_id`` — created only inside the run/subagent lifecycle. +_SERVER_OWNED_RUNTIME_CONTEXT_KEYS: frozenset[str] = ( + frozenset( + { + "is_internal", + "authz_attributes", + "channel_user_id", + "langgraph_auth_user", + "langgraph_auth_user_id", + } + ) + | SANDBOX_SERVER_OWNED_CONTEXT_KEYS ) # Keys forwarded from ``body.context`` into ``config['context']`` ONLY (the @@ -534,18 +540,18 @@ def inject_authenticated_user_context( Values copied through the free-form RunnableConfig are always cleared. """ - # --- Server-owned authorization identity fields --- + # --- Server-owned authorization and sandbox lifecycle identity fields --- # Clear any client-forged values from both config sections, then write the # authoritative is_internal. This runs before ALL early returns so that # even user_id-is-None paths get a defined is_internal value. runtime_context = config.setdefault("context", {}) if not isinstance(runtime_context, dict): raise TypeError("run context must be a mapping") - for key in _SERVER_OWNED_AUTHZ_CONTEXT_KEYS: + for key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: runtime_context.pop(key, None) configurable = config.get("configurable") if isinstance(configurable, dict): - for key in _SERVER_OWNED_AUTHZ_CONTEXT_KEYS: + for key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: configurable.pop(key, None) auth_source = getattr(getattr(request, "state", None), "auth_source", None) # ``user_id`` is server-owned for EXTERNAL callers: it now selects which diff --git a/backend/packages/harness/deerflow/AGENTS.md b/backend/packages/harness/deerflow/AGENTS.md index 53c3830c2..583eaeb4b 100644 --- a/backend/packages/harness/deerflow/AGENTS.md +++ b/backend/packages/harness/deerflow/AGENTS.md @@ -8,7 +8,7 @@ Entry points and binders: Gateway HTTP — `TraceMiddleware`; scheduled occurren Only the first is HTTP; the rest run outside ASGI, so the binding cannot live in middleware alone. Each scopes **one unit of work**, never a poller loop — a leaked binding on a reused worker task would tag later occurrences with the first id. `ensure_trace_context` inherits, keeping layered scheduled bindings and a manual trigger inside a Gateway request on one trace. -**Every other carrier is a derived output, never read back as an input.** `worker._bind_trace_id` stamps the runtime context and `config["metadata"]`; `services.start_run` stamps the run record; a caller-sent `deerflow_trace_id` (`body.metadata`, `body.config.context`) is replaced — honouring it would let the persisted run disagree with the header and the logs. `_SERVER_OWNED_RUNTIME_CONTEXT_KEYS` covers the embedded path, `redact_config_secrets` scrubs the kwargs echo (`runs.kwargs_json`), and `build_run_config` merges metadata onto a copy so the stamp cannot reach `body.config`. Callers pin an id with `X-Trace-Id`. +**Every other carrier is a derived output, never read back as an input.** `worker._bind_trace_id` stamps the runtime context and `config["metadata"]`; `services.start_run` stamps the run record; a caller-sent `deerflow_trace_id` (`body.metadata`, `body.config.context`) is replaced — honouring it would let the persisted run disagree with the header and the logs. `_SERVER_OWNED_RUNTIME_CONTEXT_KEYS` covers the embedded path and also rejects caller-supplied sandbox lease/scope identities, `redact_config_secrets` scrubs the kwargs echo (`runs.kwargs_json`), and `build_run_config` merges metadata onto a copy so the stamp cannot reach `body.config`. Callers pin an id with `X-Trace-Id`. Accepted divergence: a crash-recovered scheduled launch reuses its run via the idempotency key without restamping — the record keeps the first attempt's id, the retry's logs a fresh one; restamping would rewrite an existing record. Not a bug. Thread metadata omits the key entirely — a thread spans many runs. diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index 829ac509e..6450c46ad 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -23,7 +23,7 @@ import mimetypes import os import shutil import uuid -from collections.abc import Generator, Mapping, Sequence +from collections.abc import Generator, Iterator, Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path from typing import Any, Literal @@ -93,6 +93,19 @@ _EMBEDDED_AUTHORIZATION_CONTEXT_KEYS = frozenset( ) +def _stream_with_sandbox_lease_cleanup(items: Iterator[Any], context: dict[str, Any]) -> Iterator[Any]: + """Fence an embedded graph iterator with execution-lease cleanup.""" + try: + yield from items + finally: + try: + from deerflow.sandbox.lease import release_sandbox_execution_lease + + release_sandbox_execution_lease(context) + except Exception: + logger.warning("Failed to release embedded sandbox execution lease", exc_info=True) + + def _run_async_from_sync(coro): """Run an async helper from this synchronous client API.""" try: @@ -961,12 +974,13 @@ class DeerFlowClient: sent.update(delta) return delta - for item in self._agent.stream( + agent_items = self._agent.stream( state, config=config, context=context, stream_mode=["values", "messages", "custom"], - ): + ) + for item in _stream_with_sandbox_lease_cleanup(agent_items, context): if isinstance(item, tuple) and len(item) == 2: mode, chunk = item mode = str(mode) diff --git a/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py b/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py index dc6dd1da8..999f18f18 100644 --- a/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py +++ b/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py @@ -4,6 +4,7 @@ import logging import shlex import threading import uuid +from dataclasses import dataclass, field import httpx from agent_sandbox import Sandbox as AioSandboxClient @@ -35,12 +36,22 @@ _BASH_EXEC_UNSUPPORTED_ERROR = ( ) +@dataclass +class _ScopedShellSession: + """One server-side shell session serialized within an agent execution.""" + + lock: threading.Lock = field(default_factory=threading.Lock) + session_id: str | None = None + + class AioSandbox(Sandbox): """Sandbox implementation using the agent-infra/sandbox Docker container. This sandbox connects to a running AIO sandbox container via HTTP API. - A threading lock serializes shell commands to prevent concurrent requests - from corrupting the container's single persistent session (see #1433). + Lead/direct calls retain the legacy serialized shell. Delegated executions + receive separate server-side sessions, with commands serialized only inside + the same execution scope, so parallel subagents cannot corrupt one another's + shell state (see #1433 and #5128). """ #: The legacy exec path reuses one persistent shell session across calls, @@ -69,6 +80,10 @@ class AioSandbox(Sandbox): ) self._home_dir = home_dir self._lock = threading.Lock() + self._scope_registry_lock = threading.Lock() + self._scoped_shell_sessions: dict[str, _ScopedShellSession] = {} + self._recovery_session_id: str | None = None + self._default_shell_corrupted = False self._closed = False # Set to True after bash.exec answers 404 (image predates /v1/bash/*), # so later env-bearing calls fail fast instead of re-hitting HTTP (#3921). @@ -98,10 +113,34 @@ class AioSandbox(Sandbox): failures during teardown are logged and swallowed so provider/backend cleanup is never blocked. """ - with self._lock: + # Close admission for scoped commands before draining their sessions. + # A scoped call that already registered itself remains in this snapshot + # and is joined through its per-scope lock below; later calls fail + # without creating an orphaned registry entry. + with self._scope_registry_lock: if self._closed: return self._closed = True + scoped_sessions = list(self._scoped_shell_sessions.items()) + self._scoped_shell_sessions.clear() + for scope_id, scoped in scoped_sessions: + with scoped.lock: + if scoped.session_id is not None and self._client is not None: + self._cleanup_session_best_effort( + self._client, + scoped.session_id, + context=f"execution scope {scope_id}", + ) + scoped.session_id = None + + with self._lock: + if self._recovery_session_id is not None and self._client is not None: + self._cleanup_session_best_effort( + self._client, + self._recovery_session_id, + context="default recovery session", + ) + self._recovery_session_id = None client = self._client # Drop the reference under the lock for use-after-close safety: any # later command on this instance fails loudly instead of reusing a @@ -129,6 +168,157 @@ class AioSandbox(Sandbox): except Exception as e: logger.warning(f"Error closing AioSandbox client for {self.id}: {e}") + @staticmethod + def _cleanup_session_best_effort(client, session_id: str, *, context: str) -> None: + try: + client.shell.cleanup_session(session_id) + except Exception as cleanup_error: + logger.warning( + "Failed to release shell session %s (%s): %s", + session_id, + context, + cleanup_error, + ) + + @staticmethod + def _format_shell_result(result) -> tuple[str, int | None]: + data = result.data if result else None + output = data.output if data else "" + exit_code = getattr(data, "exit_code", None) if data else None + return output, exit_code + + def _create_shell_session(self, client) -> str: + session_id = str(uuid.uuid4()) + client.shell.create_session(id=session_id) + return session_id + + def _exec_shell(self, client, command: str, *, session_id: str | None) -> tuple[str, int | None]: + kwargs = { + "command": command, + "no_change_timeout": self._DEFAULT_NO_CHANGE_TIMEOUT, + } + if session_id is not None: + kwargs["id"] = session_id + return self._format_shell_result(client.shell.exec_command(**kwargs)) + + def _rotate_and_retry_shell( + self, + client, + command: str, + *, + corrupted_session_id: str | None, + context: str, + ) -> tuple[str, int | None, str | None]: + if corrupted_session_id is not None: + self._cleanup_session_best_effort( + client, + corrupted_session_id, + context=f"corrupted {context}", + ) + replacement_id = self._create_shell_session(client) + try: + output, exit_code = self._exec_shell( + client, + command, + session_id=replacement_id, + ) + except BaseException: + self._cleanup_session_best_effort( + client, + replacement_id, + context=f"abandoned replacement for {context}", + ) + raise + if output and _ERROR_OBSERVATION_SIGNATURE in output: + self._cleanup_session_best_effort( + client, + replacement_id, + context=f"failed replacement for {context}", + ) + return output, exit_code, None + return output, exit_code, replacement_id + + def execute_command_in_scope( + self, + command: str, + env: dict[str, str] | None = None, + timeout: float | None = None, + *, + scope_id: str | None = None, + ) -> str: + """Run no-env commands in one persistent session per subagent run. + + Commands within a scope remain serialized, while independent subagents + use distinct server-side sessions and can execute concurrently. Secret- + bearing commands keep the existing fresh ``bash.exec`` behavior. + """ + if env or scope_id is None: + return self.execute_command(command, env=env, timeout=timeout) + del timeout + _validate_extra_env(env) + + try: + with self._scope_registry_lock: + if self._closed: + raise RuntimeError("sandbox client is closed") + scoped = self._scoped_shell_sessions.setdefault( + scope_id, + _ScopedShellSession(), + ) + with scoped.lock: + # Registration and command execution are separated by the + # per-scope wait. Revalidate identity after that wait so a + # release/close which removed this exact scope is a hard + # lifecycle fence: queued callers cannot resurrect a session + # on an orphaned registry entry. + with self._scope_registry_lock: + if self._closed or self._scoped_shell_sessions.get(scope_id) is not scoped: + raise RuntimeError("sandbox command scope is no longer active") + client = self._client + if client is None: + raise RuntimeError("sandbox client is closed") + if scoped.session_id is None: + scoped.session_id = self._create_shell_session(client) + output, exit_code = self._exec_shell( + client, + command, + session_id=scoped.session_id, + ) + if output and _ERROR_OBSERVATION_SIGNATURE in output: + logger.warning("ErrorObservation detected in sandbox output for execution scope; rotating session") + output, exit_code, scoped.session_id = self._rotate_and_retry_shell( + client, + command, + corrupted_session_id=scoped.session_id, + context="execution scope", + ) + return self._render_shell_output(output, exit_code) + except Exception as e: + logger.error(f"Failed to execute command in sandbox: {e}") + return f"Error: {e}" + + def release_command_scope(self, scope_id: str) -> None: + """Clean up one subagent's explicit server-side shell session.""" + with self._scope_registry_lock: + scoped = self._scoped_shell_sessions.pop(scope_id, None) + if scoped is None: + return + with scoped.lock: + if scoped.session_id is None or self._client is None: + return + self._cleanup_session_best_effort( + self._client, + scoped.session_id, + context=f"execution scope {scope_id}", + ) + scoped.session_id = None + + @staticmethod + def _render_shell_output(output: str, exit_code: int | None) -> str: + if exit_code not in (0, None): + output = f"{output}\nExit Code: {exit_code}" if output else f"Command exited with code {exit_code}" + return output if output else "(no output)" + @property def home_dir(self) -> str: """Get the home directory inside the sandbox.""" @@ -161,12 +351,12 @@ class AioSandbox(Sandbox): ) -> str: """Execute a shell command in the sandbox. - Uses a lock to serialize concurrent requests. The AIO sandbox - container maintains a single persistent shell session that - corrupts when hit with concurrent exec_command calls (returns - ``ErrorObservation`` instead of real output). If corruption is - detected despite the lock (e.g. multiple processes sharing a - sandbox), the command is retried on a fresh session. + Uses a lock to serialize unscoped requests. The AIO sandbox container's + implicit persistent shell corrupts when hit with concurrent + ``exec_command`` calls (returning ``ErrorObservation`` instead of real + output). If corruption is detected despite the lock (e.g. multiple + processes sharing a sandbox), the replacement session is promoted for + subsequent calls rather than returning to the corrupted implicit one. Args: command: The command to execute. @@ -195,34 +385,31 @@ class AioSandbox(Sandbox): return self._execute_with_env(command, env) with self._lock: try: - result = self._client.shell.exec_command(command=command, no_change_timeout=self._DEFAULT_NO_CHANGE_TIMEOUT) - output = result.data.output if result.data else "" - exit_code = getattr(result.data, "exit_code", None) if result.data else None + client = self._client + if getattr(self, "_closed", False) or client is None: + raise RuntimeError("sandbox client is closed") + if self._default_shell_corrupted and self._recovery_session_id is None: + # Once the implicit session emits ErrorObservation, never + # target it again. A failed replacement is cleaned up and + # the next call starts another explicit session. + self._recovery_session_id = self._create_shell_session(client) + output, exit_code = self._exec_shell( + client, + command, + session_id=self._recovery_session_id, + ) if output and _ERROR_OBSERVATION_SIGNATURE in output: + self._default_shell_corrupted = True logger.warning("ErrorObservation detected in sandbox output, retrying on a fresh session") - # exec_command only auto-creates a session when called with - # no id, so the recovery session must be created explicitly - # before we target it on retry. - fresh_id = str(uuid.uuid4()) - self._client.shell.create_session(id=fresh_id) - try: - result = self._client.shell.exec_command(command=command, id=fresh_id, no_change_timeout=self._DEFAULT_NO_CHANGE_TIMEOUT) - output = result.data.output if result.data else "" - exit_code = getattr(result.data, "exit_code", None) if result.data else None - finally: - # Release the one-shot recovery session, best-effort, so - # repeated corruption can't accumulate sessions. - try: - self._client.shell.cleanup_session(fresh_id) - except Exception as cleanup_error: - logger.warning(f"Failed to release recovery session {fresh_id}: {cleanup_error}") + output, exit_code, self._recovery_session_id = self._rotate_and_retry_shell( + client, + command, + corrupted_session_id=self._recovery_session_id, + context="default shell", + ) - if exit_code not in (0, None): - # Mirror LocalSandbox: keep the actual shell status in the - # output text (acceptance-checklist evidence). - output = f"{output}\nExit Code: {exit_code}" if output else f"Command exited with code {exit_code}" - return output if output else "(no output)" + return self._render_shell_output(output, exit_code) except Exception as e: logger.error(f"Failed to execute command in sandbox: {e}") return f"Error: {e}" diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index d48e65439..64e8ac580 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -77,6 +77,7 @@ from deerflow.runtime.serialization import serialize from deerflow.runtime.stream_bridge import StreamBridge from deerflow.runtime.stream_modes import normalize_stream_modes, to_langgraph_stream_modes from deerflow.runtime.user_context import get_current_user, get_effective_user_id, resolve_runtime_user_id +from deerflow.sandbox.lease import SANDBOX_SERVER_OWNED_CONTEXT_KEYS from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_id from deerflow.tracing import inject_langfuse_metadata from deerflow.utils.messages import message_to_text @@ -519,11 +520,14 @@ class _LargeFileToolChunkBatcher: # strips ``__``-prefixed keys in build_run_config, but embedded harness callers # have no such filter and ``deerflow_trace_id`` carries no prefix to be caught # by it anyway. -_SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = frozenset( - { - CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY, - DEERFLOW_TRACE_METADATA_KEY, - } +_SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = ( + frozenset( + { + CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY, + DEERFLOW_TRACE_METADATA_KEY, + } + ) + | SANDBOX_SERVER_OWNED_CONTEXT_KEYS ) @@ -602,17 +606,17 @@ def _install_runtime_context(config: dict, runtime_context: dict[str, Any]) -> N if isinstance(existing_context, dict): existing_context.setdefault("thread_id", runtime_context["thread_id"]) existing_context.setdefault("run_id", runtime_context["run_id"]) - # Assigned, not setdefault: this is a server-owned key, the same rule - # _bind_trace_id applies to the runtime context and the run metadata. A - # deerflow_trace_id the caller put in body.config.context is an echo of - # a past output, not an input, and leaving it would make this one dict - # disagree with the response header and the logs. - if DEERFLOW_TRACE_METADATA_KEY in runtime_context: - existing_context[DEERFLOW_TRACE_METADATA_KEY] = runtime_context[DEERFLOW_TRACE_METADATA_KEY] + # Keep both context views authoritative. A server-owned value is + # assigned from the runtime context when present and removed otherwise, + # so an embedded caller cannot preserve a forged lifecycle identity in + # ``config['context']`` after it was rejected by _build_runtime_context. + for key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: + if key in runtime_context: + existing_context[key] = runtime_context[key] + else: + existing_context.pop(key, None) if "app_config" in runtime_context: existing_context["app_config"] = runtime_context["app_config"] - if CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY in runtime_context: - existing_context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = runtime_context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] return config["context"] = dict(runtime_context) @@ -1587,11 +1591,28 @@ async def run_agent( except Exception: logger.warning("Failed to close journal for run %s", run_id, exc_info=True) finally: - _release_run_scoped_references( - runnable_configs, - runtime_ctx, - journal, - ) + lease_cleanup_interrupt: BaseException | None = None + try: + from deerflow.sandbox.lease import release_sandbox_execution_lease_async + + await release_sandbox_execution_lease_async(runtime_ctx) + except Exception: + logger.warning("Failed to release sandbox execution lease for run %s", run_id, exc_info=True) + except BaseException as exc: + # release_async completes the underlying cleanup before it + # re-raises cancellation. Defer that interruption until the + # worker has dropped all other run-scoped references too. + lease_cleanup_interrupt = exc + logger.warning( + "Sandbox execution lease cleanup was interrupted for run %s; completing local cleanup first", + run_id, + ) + finally: + _release_run_scoped_references( + runnable_configs, + runtime_ctx, + journal, + ) # Drop graph and per-run payload references before the terminal # worker task itself becomes collectable. agent = None @@ -1618,6 +1639,9 @@ async def run_agent( _create_contextless_task(run_manager.cleanup(run_id)) _schedule_terminal_cycle_collection() + if lease_cleanup_interrupt is not None: + raise lease_cleanup_interrupt + # --------------------------------------------------------------------------- # Helpers diff --git a/backend/packages/harness/deerflow/sandbox/AGENTS.md b/backend/packages/harness/deerflow/sandbox/AGENTS.md index d15683627..840da96b5 100644 --- a/backend/packages/harness/deerflow/sandbox/AGENTS.md +++ b/backend/packages/harness/deerflow/sandbox/AGENTS.md @@ -1,9 +1,10 @@ ### Sandbox System (`packages/harness/deerflow/sandbox/`) -**Interface**: Abstract `Sandbox` with `execute_command(command, env=None)`, `read_file`, `write_file`, `list_dir`, `glob`, and `grep`. `grep` accepts either one text file or a directory tree. The optional `env` injects per-call environment variables (request-scoped secrets — see Request-Scoped Secrets below); `LocalSandbox` merges it into the host subprocess environment and `AioSandbox` routes env-bearing commands through the `bash.exec(env=...)` API on a fresh session. +**Interface**: Abstract `Sandbox` with `execute_command(command, env=None)`, additive `execute_command_in_scope(..., scope_id=...)` / `release_command_scope(scope_id)` hooks, `read_file`, `write_file`, `list_dir`, `glob`, and `grep`. Providers without server-side shell sessions inherit the scoped hook's pass-through implementation, preserving third-party subclasses. `grep` accepts either one text file or a directory tree. The optional `env` injects per-call environment variables (request-scoped secrets — see Request-Scoped Secrets below); `LocalSandbox` merges it into the host subprocess environment and `AioSandbox` routes env-bearing commands through the `bash.exec(env=...)` API on a fresh session. **Provider Pattern**: `SandboxProvider` with `acquire`, `acquire_async`, `get`, `release` lifecycle. Async agent/tool paths call async sandbox lifecycle hooks so Docker sandbox creation, discovery, cross-process locking, readiness polling, and release stay off the event loop. Providers that can enforce a lead Agent's explicit skill policy across the current Agent-accessible tool surface set `supports_agent_skill_isolation=True`; bind-mount providers observe the prepared thread roots directly, while upload providers implement `sync_agent_skills`. Host-backed providers must report the capability as false whenever an enabled shell can bypass their path mappings. The middleware fails closed before acquisition for an explicit policy on an unsupported provider. **Shared components** (RFC #4741): remote providers derive their deterministic sandbox id through `derive_sandbox_scope_token` (`sandbox/identity.py`, keyword-only; the sha256/16-hex derivation is a compatibility contract — changing it orphans existing containers), and serialize provider-selected acquire/release transitions through `AcquireSerializer` (`sandbox/acquire_serialization.py`): per-key `threading.Lock` table with holder/waiter refcount reclamation (no unbounded per-thread lock growth), a bounded dedicated executor so async waits never touch the event loop or the default executor, worker-owned cancellation cleanup that does not depend on a cancelled event loop task resuming, and idempotent `close()` called from provider `shutdown()`/`reset()`. AIO keys by `(user_id, thread_id)`; E2B keys by `(user_id, thread_id, skills_root)`; BoxLite/Tenki/OpenSandbox key by the derived sandbox id. `thread_id=None` acquires (random uuid ids) never enter the serializer. -**Authorization gate** (`sandbox:execute`, RFC #4063 Phase 3): every sandbox-backed tool call passes through the gate in `deerflow/authz/sandbox_authz.py` - a binary `authorize(principal, "sandbox", "execute", target="*")` check before either reusing a persisted sandbox id or calling `provider.acquire`. Rechecking reuse is required because authorization config and user roles can change while the sandbox remains cached. Sync tool invocations call `authorize_sandbox_execution`; async tool invocations await `authorize_sandbox_execution_async` exactly once. A task-local `ContextVar` scopes that single decision across the complete composed tool invocation, including `ReadBeforeWriteMiddleware`'s pre-write inspection, tool body, and post-read mark; the value is copied into `asyncio.to_thread` workers. Authorization denial is converted to the normal error `ToolMessage` at the composed middleware boundary and is explicitly excluded from the gate's generic fail-open handlers. Async config loading and provider class discovery/import are offloaded before `aauthorize()` so reused sandbox calls do not hash config files or import custom modules on the event loop; provider construction remains on the running event loop because async providers may initialize loop-affine clients. The gate lives at the single tool initialization entry point (`ensure_sandbox_initialized` / `ensure_sandbox_initialized_async` in `tools.py`), while `SandboxMiddleware.before_agent` / `abefore_agent` apply the matching sync/async check to eager acquisition. Deny raises `SandboxAuthorizationError` (`sandbox/exceptions.py`), which propagates out of ordinary tool execution as a friendly error `ToolMessage` ("sandbox execution is not permitted for your role") - the eager path catches it and skips acquisition instead, deferring the deny to the first sandbox-touching tool call so both paths share the same semantics. Provider errors (authorization calls and provider resolution) follow `authorization.fail_closed` / `fail_open`; no readable `config.yaml` or `authorization.enabled: false` makes the gate a no-op (`safe_app_config` tolerates missing config). Gateway auxiliary sync paths (uploads/artifacts routers) call `try_acquire_sandbox_for_request` (`app/gateway/authz.py`), which gates via `authorize_sandbox_for_request` and skips the sync on deny - the upload/artifact edit itself still succeeds. Tests: `tests/test_sandbox_authorization.py` and `tests/blocking_io/test_sandbox_authorization.py`. +**Execution leases** (`sandbox/lease.py`, #5128): cross-instance ownership decides which Gateway may reap a container; process-local `SandboxLeaseManager` tracks concurrent lead, subagent, Gateway-request, and channel-upload users of one client. Runs get ephemeral owners, persisted sandboxes are retained idempotently, and the last holder performs any pending `SandboxProvider.release`. Outer lifecycle fences repeat idempotent release after their complete graph/tool/request batch drains; per-tool terminal `Command` wrappers never release because sibling handlers may still run. Fork-restored children and upload syncs use non-releasing holders: they fence the client and own scope cleanup without themselves requesting a park; an earlier normal-owner request waits for them, and a missing fork client is replaced by a normal owner. Persisted lookup plus retention is serialized per `(user_id, thread_id)`; stale bindings fall through to acquire, and a post-acquire lookup miss rolls back before raising. Provider I/O does not hold the metadata lock. Repeated cancellation cannot interrupt acquire/rollback/release reconciliation or let a `to_thread` sandbox operation outlive its enclosing execution/request holder; failures are logged without replacing the original cancellation. Lease/scope context IDs are server-owned: Gateway and worker scrub caller values; only the internal subagent path assigns a task ID. Managers are registered by provider object identity, not hash/equality, so unhashable custom providers remain valid. Subagent owners also serve as `sandbox_command_scope_id`: AIO gives each scope one ordered persistent shell session, replaces it after `ErrorObservation`, and cleans it on lease release. Registry identity is revalidated after every scope-lock wait, preventing queued commands from resurrecting released sessions. Env-bearing commands use fresh `bash.exec` sessions so secrets do not persist. +**Authorization gate** (`sandbox:execute`, RFC #4063 Phase 3): every sandbox-backed tool call passes through the gate in `deerflow/authz/sandbox_authz.py` - a binary `authorize(principal, "sandbox", "execute", target="*")` check before either reusing a persisted sandbox id or calling `provider.acquire`. Rechecking reuse is required because authorization config and user roles can change while the sandbox remains cached. Sync tool invocations call `authorize_sandbox_execution`; async tool invocations await `authorize_sandbox_execution_async` exactly once. A task-local `ContextVar` scopes that single decision across the complete composed tool invocation, including `ReadBeforeWriteMiddleware`'s pre-write inspection, tool body, and post-read mark; the value is copied into `asyncio.to_thread` workers. Authorization denial is converted to the normal error `ToolMessage` at the composed middleware boundary and is explicitly excluded from the gate's generic fail-open handlers. Async config loading and provider class discovery/import are offloaded before `aauthorize()` so reused sandbox calls do not hash config files or import custom modules on the event loop; provider construction remains on the running event loop because async providers may initialize loop-affine clients. The gate lives at the single tool initialization entry point (`ensure_sandbox_initialized` / `ensure_sandbox_initialized_async` in `tools.py`), while `SandboxMiddleware.before_agent` / `abefore_agent` apply the matching sync/async check to eager acquisition. Deny raises `SandboxAuthorizationError` (`sandbox/exceptions.py`), which propagates out of ordinary tool execution as a friendly error `ToolMessage` ("sandbox execution is not permitted for your role") - the eager path catches it and skips acquisition instead, deferring the deny to the first sandbox-touching tool call so both paths share the same semantics. Provider errors (authorization calls and provider resolution) follow `authorization.fail_closed` / `fail_open`; no readable `config.yaml` or `authorization.enabled: false` makes the gate a no-op (`safe_app_config` tolerates missing config). Gateway upload/artifact sync calls `try_acquire_sandbox_for_request` (`app/gateway/authz.py`), which gates, returns a request lease, and skips sync on deny while preserving the primary operation. Callers release after their last sandbox operation; artifacts request normal parking, uploads do not. Tests: `tests/test_sandbox_authorization.py` and `tests/blocking_io/test_sandbox_authorization.py`. **Environment policy** (`sandbox/env_policy.py`): `execute_command` no longer inherits the full `os.environ`. `build_sandbox_env()` scrubs secret-looking names (`*KEY*`/`*SECRET*`/`*TOKEN*`/`*PASS*`/`*CREDENTIAL*`) from the inherited environment before layering injected request secrets on top, so platform credentials (e.g. `OPENAI_API_KEY`) never leak into skill subprocesses. Benign vars (`PATH`, `HOME`, `LANG`, `VIRTUAL_ENV`, ...) are preserved. **Implementations**: - `LocalSandboxProvider` - Local filesystem execution. `acquire(thread_id)` returns a per-user/thread `LocalSandbox` (id `local:{user_id}:{thread_id}`) whose `path_mappings` resolve `/mnt/user-data/{workspace,uploads,outputs}` and `/mnt/acp-workspace` to that thread's host directories, so the public `Sandbox` API honours the `/mnt/user-data` contract uniformly with AIO. `acquire()` / `acquire(None)` keeps the legacy generic singleton (id `local`) for callers without a thread context. Per-thread sandboxes are held in an LRU cache (default 256 entries) guarded by a `threading.Lock`. Shared runs use category mappings; a policy-scoped run replaces them with one `/mnt/skills` root mapping to the coherent thread view, so structured file tools resolve through one managed boundary. This is not a host filesystem security boundary: an enabled host `bash` subprocess can use canonical paths without `PathMapping`, so `supports_agent_skill_isolation` is dynamic and explicit Agent policies fail closed while host bash is enabled. Host-to-virtual output masking scans dynamic per-user/per-thread roots directly instead of compiling path-specific regexes, so evicted thread IDs do not remain in Python's global regex caches; a separate 256-entry root cache prevents repeated `realpath()` walks for every glob/grep match while bounding dynamic-path retention, and only the small process-stable skill/integration source set uses a bounded compiled cache. On Windows, Git Bash/MSYS argument-conversion exclusions are limited to safe non-root virtual path prefixes; do not restore a blanket conversion disable, because host-native CLI launchers need normal MSYS path conversion for their own installation paths. diff --git a/backend/packages/harness/deerflow/sandbox/lease.py b/backend/packages/harness/deerflow/sandbox/lease.py new file mode 100644 index 000000000..2fc3d6ff3 --- /dev/null +++ b/backend/packages/harness/deerflow/sandbox/lease.py @@ -0,0 +1,698 @@ +"""Execution-scoped leases for process-local sandbox use. + +Provider ownership stores answer which Gateway instance may reap a remote +sandbox. This module answers a different question: which concurrently running +agent executions inside one Gateway are still using the provider's active +client. The last execution lease is the only one allowed to call +``SandboxProvider.release``. +""" + +from __future__ import annotations + +import asyncio +import logging +import threading +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from deerflow.sandbox.acquire_serialization import AcquireSerializer + +if TYPE_CHECKING: + from deerflow.sandbox.sandbox_provider import SandboxProvider + +logger = logging.getLogger(__name__) + + +SANDBOX_LEASE_OWNER_CONTEXT_KEY = "sandbox_lease_owner_id" +SANDBOX_COMMAND_SCOPE_CONTEXT_KEY = "sandbox_command_scope_id" +SANDBOX_SERVER_OWNED_CONTEXT_KEYS = frozenset( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY, + SANDBOX_COMMAND_SCOPE_CONTEXT_KEY, + } +) + + +async def _drain_task_after_cancellation[T](task: asyncio.Task[T]) -> T: + """Wait for ``task`` even if the current task is cancelled again. + + Lifecycle reconciliation must keep its serializer until provider work has + finished. Repeated cancellation is therefore remembered by the caller but + cannot propagate into the reconciliation task or interrupt this drain. + """ + while True: + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + if task.done(): + return task.result() + + +async def run_sync_lifecycle_operation[T](func: Callable[..., T], /, *args: Any, **kwargs: Any) -> T: + """Run blocking client work without letting cancellation outlive it. + + ``asyncio.to_thread`` cannot stop its worker when the awaiting task is + cancelled. Sandbox cleanup must therefore wait for the worker before an + outer execution/request fence is allowed to release its client holder. + Repeated cancellation is remembered and propagated only after the worker + has terminated; a late worker failure is logged without replacing the + caller's cancellation. + """ + operation_task = asyncio.create_task(asyncio.to_thread(func, *args, **kwargs)) + try: + return await asyncio.shield(operation_task) + except asyncio.CancelledError as cancellation: + try: + await _drain_task_after_cancellation(operation_task) + except Exception: + logger.warning( + "Cancelled sandbox client operation failed while draining", + exc_info=True, + ) + raise cancellation + + +@dataclass(slots=True) +class SandboxClientLease: + """One bounded caller's process-local hold on a sandbox client.""" + + sandbox: object | None + sandbox_id: str + owner_id: str | None + provider: SandboxProvider + + async def release(self) -> None: + """Drop this holder after its final sandbox operation has drained.""" + if self.owner_id is None: + return + owner_id = self.owner_id + self.owner_id = None + await get_sandbox_lease_manager(self.provider).release_async(owner_id) + + async def run_sync[T](self, func: Callable[..., T], /, *args: Any, **kwargs: Any) -> T: + """Run one blocking client operation inside this lease boundary.""" + return await run_sync_lifecycle_operation(func, *args, **kwargs) + + +async def acquire_sandbox_client_lease( + provider: SandboxProvider, + thread_id: str, + *, + user_id: str, + owner_prefix: str, + release_on_last: bool = True, +) -> SandboxClientLease: + """Acquire a unique holder and resolve its process-local client. + + This is the non-HTTP counterpart of Gateway request leases. Callers must + keep the returned object through their last sandbox operation and release + it in ``finally``. ``release_on_last=False`` is appropriate for upload + synchronization: the upload fences a concurrent run but does not itself + request that the warm sandbox be parked. + """ + owner_id = f"{owner_prefix}:{uuid.uuid4()}" + manager = get_sandbox_lease_manager(provider) + sandbox_id = await manager.acquire_async( + owner_id, + thread_id, + user_id=user_id, + release_on_last=release_on_last, + ) + try: + sandbox = provider.get(sandbox_id) + except BaseException: + await manager.release_async(owner_id) + raise + return SandboxClientLease( + sandbox=sandbox, + sandbox_id=sandbox_id, + owner_id=owner_id, + provider=provider, + ) + + +@dataclass(frozen=True) +class _LeaseBinding: + sandbox_id: str + thread_key: tuple[str, str] + release_on_last: bool = True + + +class SandboxLeaseManager: + """Coordinate active agent users of one sandbox provider. + + Lifecycle transitions are serialized per user/thread key. Metadata is + protected separately so unrelated threads do not block each other's slow + provider operations. + """ + + def __init__(self, provider: SandboxProvider): + self._provider = provider + self._metadata_lock = threading.RLock() + self._serializer = AcquireSerializer[tuple[str, str]]( + thread_name_prefix="sandbox-execution-lease", + ) + self._bindings_by_owner: dict[str, _LeaseBinding] = {} + self._owners_by_sandbox: dict[str, set[str]] = {} + self._release_pending_by_sandbox: set[str] = set() + + @staticmethod + def _thread_key(thread_id: str, user_id: str) -> tuple[str, str]: + return user_id, thread_id + + def _remove_binding_locked( + self, + owner_id: str, + *, + request_release: bool = True, + ) -> tuple[_LeaseBinding | None, bool]: + binding = self._bindings_by_owner.pop(owner_id, None) + if binding is None: + return None, False + if request_release and binding.release_on_last: + # Closing responsibility belongs to the sandbox lifecycle, not to + # whichever holder happens to finish last. A normal owner may + # leave while a fork/upload borrower is still using the client; + # remember its close request until every holder is gone. + self._release_pending_by_sandbox.add(binding.sandbox_id) + owners = self._owners_by_sandbox.get(binding.sandbox_id) + if owners is None: + release_provider = binding.sandbox_id in self._release_pending_by_sandbox + self._release_pending_by_sandbox.discard(binding.sandbox_id) + return binding, release_provider + owners.discard(owner_id) + if owners: + return binding, False + self._owners_by_sandbox.pop(binding.sandbox_id, None) + release_provider = binding.sandbox_id in self._release_pending_by_sandbox + self._release_pending_by_sandbox.discard(binding.sandbox_id) + return binding, release_provider + + def _bind_locked( + self, + owner_id: str, + sandbox_id: str, + key: tuple[str, str], + *, + release_on_last: bool, + ) -> tuple[_LeaseBinding | None, bool]: + existing = self._bindings_by_owner.get(owner_id) + if existing is not None and existing.thread_key != key: + raise RuntimeError(f"Sandbox lease owner {owner_id!r} cannot move between thread identities") + if existing is not None and existing.sandbox_id == sandbox_id: + # Ownership is monotonic for one execution. A normal owner may + # later encounter a fork-restored view, but must not lose its + # responsibility to park the provider. A borrowed owner can be + # upgraded when it later performs a normal acquisition. + if release_on_last and not existing.release_on_last: + self._bindings_by_owner[owner_id] = _LeaseBinding( + sandbox_id=sandbox_id, + thread_key=key, + release_on_last=True, + ) + return None, False + + release_previous = False + previous: _LeaseBinding | None = None + if existing is not None: + previous, release_previous = self._remove_binding_locked(owner_id) + + self._bindings_by_owner[owner_id] = _LeaseBinding( + sandbox_id=sandbox_id, + thread_key=key, + release_on_last=release_on_last, + ) + self._owners_by_sandbox.setdefault(sandbox_id, set()).add(owner_id) + + return previous, release_previous + + def _release_unbound_acquire(self, sandbox_id: str) -> None: + """Undo a cancelled acquire when no admitted execution uses its result. + + The caller must still hold the serializer for the originating thread + key. That makes the owner check and provider release one transition: + another execution for the same thread cannot bind the sandbox between + them. + """ + with self._metadata_lock: + has_owners = bool(self._owners_by_sandbox.get(sandbox_id)) + if not has_owners: + self._provider.release(sandbox_id) + + def _active_owner_binding( + self, + owner_id: str, + key: tuple[str, str], + *, + release_on_last: bool, + ) -> str | None: + """Return one live owner binding while the caller holds ``key``. + + Middleware may retain a checkpointed id before the provider discovers + that its local client is gone. Treat that owner binding as stale so a + later acquire rebuilds it instead of preserving an unusable id. + """ + with self._metadata_lock: + existing = self._bindings_by_owner.get(owner_id) + if existing is None: + return None + if existing.thread_key != key: + raise RuntimeError(f"Sandbox lease owner {owner_id!r} cannot move between thread identities") + if self._provider.get(existing.sandbox_id) is not None: + if release_on_last and not existing.release_on_last: + with self._metadata_lock: + if self._bindings_by_owner.get(owner_id) == existing: + self._bindings_by_owner[owner_id] = _LeaseBinding( + sandbox_id=existing.sandbox_id, + thread_key=existing.thread_key, + release_on_last=True, + ) + return existing.sandbox_id + + with self._metadata_lock: + if self._bindings_by_owner.get(owner_id) == existing: + self._remove_binding_locked(owner_id, request_release=False) + return None + + def _acquire_and_bind( + self, + owner_id: str, + thread_id: str, + user_id: str, + key: tuple[str, str], + *, + release_on_last: bool, + ) -> str: + sandbox_id = self._provider.acquire(thread_id, user_id=user_id) + with self._metadata_lock: + previous, release_previous = self._bind_locked( + owner_id, + sandbox_id, + key, + release_on_last=release_on_last, + ) + if release_previous and previous is not None: + self._provider.release(previous.sandbox_id) + return sandbox_id + + async def _acquire_and_bind_async( + self, + owner_id: str, + thread_id: str, + user_id: str, + key: tuple[str, str], + *, + release_on_last: bool, + ) -> str: + acquire_task = asyncio.create_task(self._provider.acquire_async(thread_id, user_id=user_id)) + try: + sandbox_id = await asyncio.shield(acquire_task) + except asyncio.CancelledError as cancellation: + # Provider implementations commonly offload container startup to a + # worker thread, which cannot be stopped by cancelling the awaiter. + # Reconcile the result before relinquishing the serializer. + try: + sandbox_id = await _drain_task_after_cancellation(acquire_task) + except Exception: + logger.warning( + "Cancelled sandbox acquire failed during reconciliation", + exc_info=True, + ) + raise cancellation + rollback_task = asyncio.create_task( + asyncio.to_thread( + self._release_unbound_acquire, + sandbox_id, + ) + ) + try: + await _drain_task_after_cancellation(rollback_task) + except Exception: + logger.warning( + "Cancelled sandbox acquire rollback failed during reconciliation", + exc_info=True, + ) + raise cancellation + with self._metadata_lock: + previous, release_previous = self._bind_locked( + owner_id, + sandbox_id, + key, + release_on_last=release_on_last, + ) + if release_previous and previous is not None: + await asyncio.to_thread( + self._provider.release, + previous.sandbox_id, + ) + return sandbox_id + + def retain( + self, + owner_id: str, + sandbox_id: str, + *, + thread_id: str, + user_id: str, + release_on_last: bool = True, + ) -> None: + """Attach an execution to an inherited or checkpointed sandbox id. + + ``release_on_last=False`` is for a borrower: it fences the client and + owns command-scope cleanup without requesting a park itself. A normal + owner's earlier park request can still be deferred until it drains. + """ + key = self._thread_key(thread_id, user_id) + with self._serializer.hold(key): + with self._metadata_lock: + previous, release_previous = self._bind_locked( + owner_id, + sandbox_id, + key, + release_on_last=release_on_last, + ) + if release_previous and previous is not None: + self._provider.release(previous.sandbox_id) + + async def retain_async( + self, + owner_id: str, + sandbox_id: str, + *, + thread_id: str, + user_id: str, + release_on_last: bool = True, + ) -> None: + """Async attach without blocking the event loop on a lifecycle lock.""" + key = self._thread_key(thread_id, user_id) + async with self._serializer.hold_async(key): + with self._metadata_lock: + previous, release_previous = self._bind_locked( + owner_id, + sandbox_id, + key, + release_on_last=release_on_last, + ) + if release_previous and previous is not None: + await asyncio.to_thread( + self._provider.release, + previous.sandbox_id, + ) + + def acquire( + self, + owner_id: str, + thread_id: str, + *, + user_id: str, + release_on_last: bool = True, + ) -> str: + """Acquire and bind a sandbox, idempotently for one execution owner.""" + key = self._thread_key(thread_id, user_id) + with self._serializer.hold(key): + existing_sandbox_id = self._active_owner_binding( + owner_id, + key, + release_on_last=release_on_last, + ) + if existing_sandbox_id is not None: + return existing_sandbox_id + return self._acquire_and_bind( + owner_id, + thread_id, + user_id, + key, + release_on_last=release_on_last, + ) + + def reuse_or_acquire( + self, + owner_id: str, + sandbox_id: str, + *, + thread_id: str, + user_id: str, + release_on_last: bool = True, + acquire_release_on_last: bool = True, + ) -> str: + """Atomically retain a live persisted sandbox or acquire a replacement. + + A fork-restored live client is borrowed with ``release_on_last=False``; + if that persisted client is gone, its freshly acquired replacement is + owned normally unless ``acquire_release_on_last`` is also disabled. + """ + key = self._thread_key(thread_id, user_id) + with self._serializer.hold(key): + existing_sandbox_id = self._active_owner_binding( + owner_id, + key, + release_on_last=release_on_last, + ) + if existing_sandbox_id is not None: + return existing_sandbox_id + + if self._provider.get(sandbox_id) is not None: + with self._metadata_lock: + previous, release_previous = self._bind_locked( + owner_id, + sandbox_id, + key, + release_on_last=release_on_last, + ) + if release_previous and previous is not None: + self._provider.release(previous.sandbox_id) + return sandbox_id + + return self._acquire_and_bind( + owner_id, + thread_id, + user_id, + key, + release_on_last=acquire_release_on_last, + ) + + async def acquire_async( + self, + owner_id: str, + thread_id: str, + *, + user_id: str, + release_on_last: bool = True, + ) -> str: + """Async acquire while preserving the provider's own async hook.""" + key = self._thread_key(thread_id, user_id) + async with self._serializer.hold_async(key): + existing_sandbox_id = self._active_owner_binding( + owner_id, + key, + release_on_last=release_on_last, + ) + if existing_sandbox_id is not None: + return existing_sandbox_id + return await self._acquire_and_bind_async( + owner_id, + thread_id, + user_id, + key, + release_on_last=release_on_last, + ) + + async def reuse_or_acquire_async( + self, + owner_id: str, + sandbox_id: str, + *, + thread_id: str, + user_id: str, + release_on_last: bool = True, + acquire_release_on_last: bool = True, + ) -> str: + """Async atomic retain-or-replace transition for a persisted sandbox.""" + key = self._thread_key(thread_id, user_id) + async with self._serializer.hold_async(key): + existing_sandbox_id = self._active_owner_binding( + owner_id, + key, + release_on_last=release_on_last, + ) + if existing_sandbox_id is not None: + return existing_sandbox_id + + if self._provider.get(sandbox_id) is not None: + with self._metadata_lock: + previous, release_previous = self._bind_locked( + owner_id, + sandbox_id, + key, + release_on_last=release_on_last, + ) + if release_previous and previous is not None: + await asyncio.to_thread( + self._provider.release, + previous.sandbox_id, + ) + return sandbox_id + + return await self._acquire_and_bind_async( + owner_id, + thread_id, + user_id, + key, + release_on_last=acquire_release_on_last, + ) + + def release(self, owner_id: str) -> None: + """Release one execution and park the sandbox only after the last user.""" + with self._metadata_lock: + binding = self._bindings_by_owner.get(owner_id) + if binding is None: + return + + with self._serializer.hold(binding.thread_key): + with self._metadata_lock: + current = self._bindings_by_owner.get(owner_id) + if current is None: + return + binding, release_provider = self._remove_binding_locked(owner_id) + assert binding is not None + + try: + sandbox = self._provider.get(binding.sandbox_id) + if sandbox is not None: + sandbox.release_command_scope(owner_id) + finally: + if release_provider: + self._provider.release(binding.sandbox_id) + + async def release_async(self, owner_id: str) -> None: + """Release a lease without blocking the caller's event loop.""" + release_task = asyncio.create_task(asyncio.to_thread(self.release, owner_id)) + try: + await asyncio.shield(release_task) + except asyncio.CancelledError as cancellation: + # Complete lifecycle cleanup before allowing cancellation to leave + # the agent's finally block. + try: + await _drain_task_after_cancellation(release_task) + except Exception: + logger.warning( + "Cancelled sandbox release failed during reconciliation", + exc_info=True, + ) + raise cancellation + raise cancellation + + def binding_for(self, owner_id: str) -> str | None: + """Return the sandbox bound to an owner; intended for diagnostics/tests.""" + with self._metadata_lock: + binding = self._bindings_by_owner.get(owner_id) + return binding.sandbox_id if binding is not None else None + + def close(self) -> None: + """Stop accepting new transitions and release serializer workers.""" + self._serializer.close() + + +_manager_lock = threading.Lock() +_managers: dict[int, tuple[SandboxProvider, SandboxLeaseManager]] = {} + + +def get_sandbox_lease_manager(provider: SandboxProvider) -> SandboxLeaseManager: + """Return the process-local lease manager for this provider object. + + Provider implementations are not required to be hashable, and distinct + instances that compare equal must not share lifecycle state. Keep a strong + identity entry until the provider is explicitly detached; the manager + already owns the provider strongly, so a weak-key registry would not make + the lifecycle shorter. + """ + provider_id = id(provider) + with _manager_lock: + entry = _managers.get(provider_id) + if entry is not None and entry[0] is provider: + return entry[1] + manager = SandboxLeaseManager(provider) + _managers[provider_id] = (provider, manager) + return manager + + +def discard_sandbox_lease_manager(provider: SandboxProvider) -> None: + """Forget lease metadata when a provider singleton is detached.""" + provider_id = id(provider) + with _manager_lock: + entry = _managers.get(provider_id) + if entry is None or entry[0] is not provider: + manager = None + else: + _, manager = _managers.pop(provider_id) + if manager is not None: + manager.close() + + +def ensure_sandbox_lease_owner(context: Any) -> str | None: + """Create one ephemeral owner id in a mutable runtime context.""" + if not isinstance(context, dict): + return None + existing = context.get(SANDBOX_LEASE_OWNER_CONTEXT_KEY) + if isinstance(existing, str) and existing: + return existing + owner_id = f"agent:{uuid.uuid4()}" + context[SANDBOX_LEASE_OWNER_CONTEXT_KEY] = owner_id + return owner_id + + +def sandbox_lease_owner(context: Any) -> str | None: + """Read an execution owner without creating one for direct tool callers.""" + if not isinstance(context, dict): + return None + owner_id = context.get(SANDBOX_LEASE_OWNER_CONTEXT_KEY) + return owner_id if isinstance(owner_id, str) and owner_id else None + + +def sandbox_command_scope(context: Any) -> str | None: + """Read the optional shell-session scope carried by subagent execution.""" + if not isinstance(context, dict): + return None + scope_id = context.get(SANDBOX_COMMAND_SCOPE_CONTEXT_KEY) + return scope_id if isinstance(scope_id, str) and scope_id else None + + +def _sandbox_execution_lease(context: Any) -> tuple[str, str] | None: + """Return a bound execution lease without initializing a provider.""" + owner_id = sandbox_lease_owner(context) + if owner_id is None or not isinstance(context, dict): + return None + sandbox_id = context.get("sandbox_id") + if not isinstance(sandbox_id, str) or not sandbox_id: + return None + return owner_id, sandbox_id + + +def release_sandbox_execution_lease(context: Any) -> None: + """Release a lead/embedded execution lease at its outer lifecycle fence.""" + lease = _sandbox_execution_lease(context) + if lease is None: + return + + # Import lazily so runs that never touch a sandbox do not initialize the + # provider during terminal cleanup. + from deerflow.sandbox.sandbox_provider import get_sandbox_provider + + owner_id, _sandbox_id = lease + provider = get_sandbox_provider() + get_sandbox_lease_manager(provider).release(owner_id) + + +async def release_sandbox_execution_lease_async(context: Any) -> None: + """Async counterpart to :func:`release_sandbox_execution_lease`.""" + lease = _sandbox_execution_lease(context) + if lease is None: + return + + from deerflow.sandbox.sandbox_provider import get_sandbox_provider + + owner_id, _sandbox_id = lease + provider = get_sandbox_provider() + await get_sandbox_lease_manager(provider).release_async(owner_id) diff --git a/backend/packages/harness/deerflow/sandbox/middleware.py b/backend/packages/harness/deerflow/sandbox/middleware.py index f04257375..885e40f09 100644 --- a/backend/packages/harness/deerflow/sandbox/middleware.py +++ b/backend/packages/harness/deerflow/sandbox/middleware.py @@ -21,6 +21,11 @@ from deerflow.authz.sandbox_authz import ( 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 logger = logging.getLogger(__name__) @@ -39,9 +44,9 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): 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) - - Sandbox is reused across multiple turns within the same thread - - Sandbox is NOT released after each agent call to avoid wasteful recreation - - Cleanup happens at application shutdown via SandboxProvider.shutdown() + - 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 @@ -108,20 +113,105 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): 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) -> str: + def _acquire_sandbox( + self, + thread_id: str, + *, + user_id: str, + owner_id: str | None, + ) -> str: provider = get_sandbox_provider() - sandbox_id = provider.acquire(thread_id, user_id=user_id) + 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) -> str: + async def _acquire_sandbox_async( + self, + thread_id: str, + *, + user_id: str, + owner_id: str | None, + ) -> str: provider = get_sandbox_provider() - sandbox_id = await provider.acquire_async(thread_id, user_id=user_id) + 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 - async def _release_sandbox_async(self, sandbox_id: str) -> None: - await asyncio.to_thread(get_sandbox_provider().release, 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: @@ -130,11 +220,16 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): return super().before_agent(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) @@ -162,14 +257,29 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): return None provider = get_sandbox_provider() self._require_projection_support(provider, projection) - sandbox_id = self._acquire_sandbox(thread_id, user_id=user_id) - if projection is not None: - provider.sync_agent_skills( - sandbox_id, - thread_id=thread_id, - user_id=user_id, - projection=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) @@ -178,6 +288,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): "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) @override @@ -191,6 +309,7 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): 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) @@ -213,14 +332,29 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): 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) - if projection is not None: - await provider.sync_agent_skills_async( + 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, - thread_id=thread_id, - user_id=user_id, - projection=projection, + 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) @@ -229,6 +363,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): "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 @@ -242,13 +384,23 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): logger.info(f"Not releasing fork-restored sandbox {sandbox_id}") return None logger.info(f"Releasing sandbox {sandbox_id}") - get_sandbox_provider().release(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") - get_sandbox_provider().release(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 # No sandbox to release @@ -265,13 +417,19 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]): 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) + 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) + await self._release_sandbox_async( + sandbox_id, + owner_id=sandbox_lease_owner(runtime.context), + ) return None # No sandbox to release diff --git a/backend/packages/harness/deerflow/sandbox/sandbox.py b/backend/packages/harness/deerflow/sandbox/sandbox.py index 3804e7298..35ef971df 100644 --- a/backend/packages/harness/deerflow/sandbox/sandbox.py +++ b/backend/packages/harness/deerflow/sandbox/sandbox.py @@ -107,6 +107,27 @@ class Sandbox(ABC): """ pass + def execute_command_in_scope( + self, + command: str, + env: dict[str, str] | None = None, + timeout: float | None = None, + *, + scope_id: str | None = None, + ) -> str: + """Execute a command in an optional agent execution scope. + + Providers without server-side shell sessions inherit the ordinary + command behavior. Session-aware providers may isolate concurrent agent + executions while preserving serialization inside one scope. + """ + del scope_id + return self.execute_command(command, env=env, timeout=timeout) + + def release_command_scope(self, scope_id: str) -> None: + """Release provider-specific command state for one execution scope.""" + del scope_id + @abstractmethod def read_file( self, diff --git a/backend/packages/harness/deerflow/sandbox/sandbox_provider.py b/backend/packages/harness/deerflow/sandbox/sandbox_provider.py index 80c18d4f1..f4e0d5da0 100644 --- a/backend/packages/harness/deerflow/sandbox/sandbox_provider.py +++ b/backend/packages/harness/deerflow/sandbox/sandbox_provider.py @@ -178,6 +178,9 @@ def reset_sandbox_provider() -> None: provider = _default_sandbox_provider _default_sandbox_provider = None if provider is not None: + from deerflow.sandbox.lease import discard_sandbox_lease_manager + + discard_sandbox_lease_manager(provider) provider.reset() @@ -195,6 +198,9 @@ def shutdown_sandbox_provider() -> None: provider = _default_sandbox_provider _default_sandbox_provider = None if provider is not None and hasattr(provider, "shutdown"): + from deerflow.sandbox.lease import discard_sandbox_lease_manager + + discard_sandbox_lease_manager(provider) provider.shutdown() @@ -211,4 +217,9 @@ def set_sandbox_provider(provider: SandboxProvider) -> None: """ global _default_sandbox_provider with _provider_lock: + previous = _default_sandbox_provider _default_sandbox_provider = provider + if previous is not None and previous is not provider: + from deerflow.sandbox.lease import discard_sandbox_lease_manager + + discard_sandbox_lease_manager(previous) diff --git a/backend/packages/harness/deerflow/sandbox/tools.py b/backend/packages/harness/deerflow/sandbox/tools.py index 4ce3c1d0b..2116d4ec8 100644 --- a/backend/packages/harness/deerflow/sandbox/tools.py +++ b/backend/packages/harness/deerflow/sandbox/tools.py @@ -32,10 +32,16 @@ from deerflow.sandbox.exceptions import ( SandboxRuntimeError, ) from deerflow.sandbox.file_operation_lock import get_file_operation_lock +from deerflow.sandbox.lease import ( + get_sandbox_lease_manager, + run_sync_lifecycle_operation, + sandbox_command_scope, + sandbox_lease_owner, +) from deerflow.sandbox.overwrite import unwrap_sandbox from deerflow.sandbox.path_patterns import build_output_mask_pattern, replace_output_path_matches from deerflow.sandbox.sandbox import Sandbox -from deerflow.sandbox.sandbox_provider import get_sandbox_provider +from deerflow.sandbox.sandbox_provider import SandboxProvider, get_sandbox_provider from deerflow.sandbox.search import GrepMatch from deerflow.sandbox.security import LOCAL_HOST_BASH_DISABLED_MESSAGE, is_host_bash_allowed from deerflow.tools.types import Runtime @@ -1404,6 +1410,44 @@ def sandbox_from_runtime(runtime: Runtime | None = None) -> Sandbox: return sandbox +def _rollback_failed_sandbox_lookup( + provider: SandboxProvider, + sandbox_id: str, + owner_id: str | None, +) -> None: + """Undo an acquire whose active client disappeared before lookup.""" + try: + if owner_id is not None: + get_sandbox_lease_manager(provider).release(owner_id) + else: + provider.release(sandbox_id) + except Exception: + logger.warning( + "Failed to roll back sandbox after post-acquire lookup failure: %s", + sandbox_id, + exc_info=True, + ) + + +async def _rollback_failed_sandbox_lookup_async( + provider: SandboxProvider, + sandbox_id: str, + owner_id: str | None, +) -> None: + """Async rollback without blocking the event loop on provider cleanup.""" + try: + if owner_id is not None: + await get_sandbox_lease_manager(provider).release_async(owner_id) + else: + await asyncio.to_thread(provider.release, sandbox_id) + except Exception: + logger.warning( + "Failed to roll back sandbox after async post-acquire lookup failure: %s", + sandbox_id, + exc_info=True, + ) + + @contextmanager def sandbox_authorization_scope(runtime: Runtime) -> Iterator[None]: """Authorize once for one complete synchronous sandbox tool invocation.""" @@ -1440,6 +1484,14 @@ async def sandbox_authorization_scope_async(runtime: Runtime) -> AsyncIterator[N _SANDBOX_AUTHORIZATION_CHECKED.reset(token) +def _resolve_runtime_thread_id(runtime: Runtime) -> str | None: + """Resolve the thread identity consistently for reuse and acquisition.""" + thread_id = runtime.context.get("thread_id") if runtime.context else None + if thread_id is None: + thread_id = runtime.config.get("configurable", {}).get("thread_id") if runtime.config else None + return thread_id + + def ensure_sandbox_initialized(runtime: Runtime | None = None) -> Sandbox: """Ensure sandbox is initialized, acquiring lazily if needed. @@ -1473,15 +1525,28 @@ def ensure_sandbox_initialized(runtime: Runtime | None = None) -> Sandbox: app_config=safe_app_config(), ) - # Check if sandbox already exists in state - # Discarding fork_restored is safe: after_agent short-circuits on the - # still-wrapped state before the context-based release branch, so this - # reuse path never releases the parent sandbox. - sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox")) + # Check if sandbox already exists in state. A fork-restored execution keeps + # the wrapper so after_agent cannot park the parent's sandbox, but it still + # binds a non-releasing holder: parent cleanup cannot close the client under + # the child, and the child's outer fence can clean its command scope. + sandbox_state, fork_restored = unwrap_sandbox(runtime.state.get("sandbox")) if sandbox_state is not None: sandbox_id = sandbox_state.get("sandbox_id") if sandbox_id is not None: - sandbox = get_sandbox_provider().get(sandbox_id) + provider = get_sandbox_provider() + owner_id = sandbox_lease_owner(runtime.context) + thread_id = _resolve_runtime_thread_id(runtime) + if owner_id is not None and thread_id is not None: + sandbox_id = get_sandbox_lease_manager(provider).reuse_or_acquire( + owner_id, + sandbox_id, + thread_id=thread_id, + user_id=resolve_runtime_user_id(runtime), + release_on_last=not fork_restored, + ) + if not fork_restored: + runtime.state["sandbox"] = {"sandbox_id": sandbox_id} + sandbox = provider.get(sandbox_id) if sandbox is not None: if runtime.context is not None: runtime.context["sandbox_id"] = sandbox_id # Ensure sandbox_id is in context for releasing in after_agent @@ -1489,14 +1554,21 @@ def ensure_sandbox_initialized(runtime: Runtime | None = None) -> Sandbox: # Sandbox was released, fall through to acquire new one # Lazy acquisition: get thread_id and acquire sandbox - thread_id = runtime.context.get("thread_id") if runtime.context else None - if thread_id is None: - thread_id = runtime.config.get("configurable", {}).get("thread_id") if runtime.config else None + thread_id = _resolve_runtime_thread_id(runtime) if thread_id is None: raise SandboxRuntimeError("Thread ID not available in runtime context") provider = get_sandbox_provider() - sandbox_id = provider.acquire(thread_id, user_id=resolve_runtime_user_id(runtime)) + user_id = resolve_runtime_user_id(runtime) + owner_id = sandbox_lease_owner(runtime.context) + 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, + ) # Update runtime state - this persists across tool calls runtime.state["sandbox"] = {"sandbox_id": sandbox_id} @@ -1504,6 +1576,7 @@ def ensure_sandbox_initialized(runtime: Runtime | None = None) -> Sandbox: # Retrieve and return the sandbox sandbox = provider.get(sandbox_id) if sandbox is None: + _rollback_failed_sandbox_lookup(provider, sandbox_id, owner_id) raise SandboxNotFoundError("Sandbox not found after acquisition", sandbox_id=sandbox_id) if runtime.context is not None: @@ -1532,31 +1605,52 @@ async def ensure_sandbox_initialized_async(runtime: Runtime | None = None) -> Sa app_config=await safe_app_config_async(), ) - # Same discard as the sync path above: the reuse path never releases, - # because after_agent short-circuits on the still-wrapped state first. - sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox")) + # Same borrowed-holder rule as the sync path above: keep the fork wrapper + # while counting the child as an active client user. + sandbox_state, fork_restored = unwrap_sandbox(runtime.state.get("sandbox")) if sandbox_state is not None: sandbox_id = sandbox_state.get("sandbox_id") if sandbox_id is not None: - sandbox = get_sandbox_provider().get(sandbox_id) + provider = get_sandbox_provider() + owner_id = sandbox_lease_owner(runtime.context) + thread_id = _resolve_runtime_thread_id(runtime) + if owner_id is not None and thread_id is not None: + sandbox_id = await get_sandbox_lease_manager(provider).reuse_or_acquire_async( + owner_id, + sandbox_id, + thread_id=thread_id, + user_id=resolve_runtime_user_id(runtime), + release_on_last=not fork_restored, + ) + if not fork_restored: + runtime.state["sandbox"] = {"sandbox_id": sandbox_id} + sandbox = provider.get(sandbox_id) if sandbox is not None: if runtime.context is not None: runtime.context["sandbox_id"] = sandbox_id return sandbox - thread_id = runtime.context.get("thread_id") if runtime.context else None - if thread_id is None: - thread_id = runtime.config.get("configurable", {}).get("thread_id") if runtime.config else None + thread_id = _resolve_runtime_thread_id(runtime) if thread_id is None: raise SandboxRuntimeError("Thread ID not available in runtime context") provider = get_sandbox_provider() - sandbox_id = await provider.acquire_async(thread_id, user_id=resolve_runtime_user_id(runtime)) + user_id = resolve_runtime_user_id(runtime) + owner_id = sandbox_lease_owner(runtime.context) + 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, + ) runtime.state["sandbox"] = {"sandbox_id": sandbox_id} sandbox = provider.get(sandbox_id) if sandbox is None: + await _rollback_failed_sandbox_lookup_async(provider, sandbox_id, owner_id) raise SandboxNotFoundError("Sandbox not found after acquisition", sandbox_id=sandbox_id) if runtime.context is not None: @@ -1577,13 +1671,37 @@ async def _run_sync_tool_after_async_sandbox_init( if func is None: return "Error: Tool implementation not available" - return await asyncio.to_thread(func, runtime, *args) + return await run_sync_lifecycle_operation(func, runtime, *args) except SandboxError as e: return f"Error: {e}" except Exception as e: return f"Error: Unexpected error initializing sandbox: {_sanitize_error(e, runtime)}" +def _execute_bash_command( + sandbox: Sandbox, + command: str, + *, + runtime: Runtime, + env: dict[str, str] | None, + timeout: float | None = None, +) -> str: + """Route subagent bash calls through their isolated shell-session scope.""" + scope_id = sandbox_command_scope(runtime.context) + scoped_execute = getattr(sandbox, "execute_command_in_scope", None) + if scope_id is not None and callable(scoped_execute): + return scoped_execute( + command, + env=env, + timeout=timeout, + scope_id=scope_id, + ) + # Keep duck-typed custom providers and test doubles compatible: the scoped + # method is an additive Sandbox API, and ordinary/lead executions retain + # the original execute_command path. + return sandbox.execute_command(command, env=env, timeout=timeout) + + def ensure_thread_directories_exist(runtime: Runtime | None) -> None: """Ensure thread data directories (workspace, uploads, outputs) exist. @@ -1936,7 +2054,13 @@ def bash_tool(runtime: Runtime, command: str, description: str = "") -> str: except Exception: max_chars = 20000 command_timeout = None - output = sandbox.execute_command(command, env=injected_env, timeout=command_timeout) + output = _execute_bash_command( + sandbox, + command, + runtime=runtime, + env=injected_env, + timeout=command_timeout, + ) return _truncate_bash_output( mask_secret_values(mask_local_paths_in_output(output, thread_data), injected_env), max_chars, @@ -1952,7 +2076,18 @@ def bash_tool(runtime: Runtime, command: str, description: str = "") -> str: max_chars = sandbox_cfg.bash_output_max_chars if sandbox_cfg else 20000 except Exception: max_chars = 20000 - return _truncate_bash_output(mask_secret_values(sandbox.execute_command(command, env=injected_env), injected_env), max_chars) + return _truncate_bash_output( + mask_secret_values( + _execute_bash_command( + sandbox, + command, + runtime=runtime, + env=injected_env, + ), + injected_env, + ), + max_chars, + ) except SandboxError as e: return f"Error: {e}" except PermissionError as e: diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index 2855c5f92..d8489fc8d 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -6,6 +6,7 @@ **User-scoped Skills**: Subagents resolve their configured skills through `get_or_new_user_skill_storage(user_id)` using the parent runtime identity, with `DEFAULT_USER_ID` only when no identity is available. This keeps custom-skill shadowing and visibility aligned with the lead agent instead of reading the global-only catalog. **Date context (#4781)**: Every built-in subagent execution registers `SubagentDateContextMiddleware` immediately before `SystemMessageCoalescingMiddleware`. Its one-time `before_agent` hook adds a hidden framework-owned `SystemMessage` containing only `` before the first model call; it does not read `AppConfig.memory`, call the memory manager, rewrite the task `HumanMessage`, or inherit the lead agent's frozen-conversation/midnight lifecycle. The coalescer merges that reminder with the subagent's static prompt so strict providers still receive exactly one leading `SystemMessage`. The lead-only `DynamicContextMiddleware` registration and its date, optional-memory, and midnight-update behavior remain unchanged. **Execution**: Ordinary and durable-batch native subagents submit coroutines directly to one persistent isolated event loop. Gateway/embedded startup installs one process-wide async FIFO admission controller (default 3 running, bounded queue). Direct `create_deerflow_agent` callers can instead pass a caller-owned `SubagentRuntime`; reuse the same instance across graphs so its bound `task`, optional batch tools/service, middleware limits, and `SubagentExecutor` all share one controller without reading global YAML. An owned batch service must be started before graph construction and stopped at application shutdown. Waiters hold no scheduler thread, and cancellation/timeout release queue/slot ownership. +**Shared sandbox execution lifecycle** (#5128): every admitted subagent run carries a stable task-derived `sandbox_lease_owner_id` and matching `sandbox_command_scope_id` in its runtime context. Sandbox middleware retains that execution against the lead thread's active provider client, so one child finishing cannot close the sandbox while siblings still run; the final holder performs any pending provider release. A rollback/fork-restored child reusing the parent's live client binds a non-releasing holder: it fences parent cleanup and owns its command scope without requesting a park itself; a parent's earlier park request waits for the child, while a missing inherited client falls through to a normal fresh acquire. On AIO, the command scope selects one explicit persistent shell session per subagent, allowing independent scopes to run concurrently while preserving in-order shell state within one child. Sync sandbox tool bodies offloaded with `asyncio.to_thread` are shielded and drained across repeated cancellation before the outer execution can clean its holder; a cancelled worker can therefore neither re-admit an already-released owner nor run after subagent terminalization. Middleware performs the normal release, and `SubagentExecutor` repeats it idempotently in `finally` so exceptions, cooperative cancellation, and timeout unwind paths cannot leak a lease or scoped session. **Concurrency and total delegation cap**: Ordinary `task` concurrency is resolved once as the minimum of the per-run request, the startup-frozen `subagent_runtime.max_running`, and the schema safety ceiling (1-64), then shared by the lead prompt and `SubagentLimitMiddleware`. Hot reloads must not make either layer advertise more capacity than the already-created process controller; a changed startup-only value takes effect only after restart. The same middleware separately enforces `subagents.max_total_per_run` (default 6, config schema 1-50, runtime override `max_total_subagents` clamped to the same range) against current-run entries in the durable delegation ledger, so a long lead-agent run cannot bypass concurrency limits by launching repeated legal-sized batches at each planning checkpoint, but historical delegations from previous runs in the same thread do not consume the new run's budget. Explicit `batch_task` work does not consume or relax that ordinary-run ledger: its persisted total/live/running limits live under `subagent_batches`. Gateway `run_agent()` and embedded `DeerFlowClient.stream()` both provide a per-invocation `run_id` in runtime context; `DeerFlowClient.stream()` also tags its input `HumanMessage` with that same id so durable-context capture can identify the current request boundary. Gateway resume paths may not append a new `HumanMessage`, so the worker also exposes the pre-run checkpoint's message ids in runtime context; durable-context capture uses that as the current-run boundary and never re-tags older task calls as the resumed run. When no delegation slots remain, task calls are stripped, provider raw tool-call metadata is synced, `finish_reason` is forced to `stop`, and a visible "subagent delegation limit" note is appended so the agent can synthesize already-collected results. Default subagent timeout `subagents.timeout_seconds=1800` (30 min) and built-in `general-purpose` `max_turns=150`. **Flow**: Ordinary `task()` → `SubagentExecutor` → shared process slot → result polling/SSE. Explicit `batch_task()` → durable batch/item rows → lease-based batch service (`subagents/batch_service.py`, started by Gateway or an explicit direct runtime) → the same `SubagentExecutor`/process slots → bounded stored result and owner-scoped API/JSONL export. Batch mode is selected only by the explicit tool, never inferred from prompt size. Executor queue rejection/timeout occurs before model execution and therefore releases the durable lease without consuming an item attempt; real execution failure and expired leases still consume the retry budget. User cancellation terminalizes every nonterminal item immediately and clears its lease, fencing any stale worker completion. Background cancellation resolves the result/future under `_background_tasks_lock` but calls `Future.cancel()` only after releasing it, because cancellation may synchronously invoke the completion callback that reacquires the registry lock. Direct runtimes provide the tools and worker but not Gateway's HTTP/UI surface. `task_started` carries the resolved effective model name. The per-subagent `SubagentTokenCollector` publishes a cumulative usage snapshot to the shared `SubagentResult` after every completed LLM response; the next `task_running` event carries that snapshot, so collapsed workspace cards can update without re-accounting parent-run totals. Terminal ToolMessage metadata (`subagent_model_name`, `subagent_token_usage`) and the persisted `subagent.end` event retain the model/usage after reload; absent provider usage stays absent rather than being estimated as zero. The executor caches one resolved `AppConfig` snapshot (explicit or `get_app_config()` fallback) for agent assembly, deferred setup, and receipt harvesting, so `verification.receipts_enabled=false` remains authoritative on both construction paths. Terminal tool receipts are harvested before `try_set_terminal` and committed with the other payload fields under the same state lock, so status polling cannot observe a terminal result before its receipt metadata is available. Each yielded values chunk becomes the latest terminal-harvest state and immediately publishes its harvested receipts to the shared result before cooperative cancellation is checked. Tool-ended cancellation/failure evidence uses the current ToolMessage scan, but a completed result always uses the bounded ledger snapshot attached to the assistant text being returned—even when a max-turn partial ends on a later tool chunk—so omitted receipts cannot validate its citations; a missing/malformed completed snapshot fails closed with no receipts. Therefore direct task cancellation and both execution/polling timeouts retain the latest execution evidence even when cancellation interrupts before another stream boundary. **Report contract (RFC #4651 PR3)**: `report_contract.py` owns the prompt-layer text that makes Layer 1 receipt verification non-inert. `SubagentExecutor._build_initial_state` appends `build_report_contract_section(receipts_enabled=...)` to every subagent's consolidated `SystemMessage` — built-in and custom alike — requiring `[rN tool_name]` citations (from the Tool receipts ledger) for action claims, verifiable handles (absolute path, URL, ID, HTTP status) for deliverables, and explicit reporting of failures; the citation clause follows `verification.receipts_enabled`, and the citation example derives from the single-owner `format_citation`/`receipt_id` so prompt text cannot drift from the verifier. The `task` tool hands lead-supplied `acceptance_criteria` to the `SubagentExecutor` constructor, which appends them via `render_acceptance_criteria_block(...)` to the task `HumanMessage` (stripped, capped at 20 items × 500 chars, each entry neutralized) — the untrusted channel `InputSanitizationMiddleware` escapes and boundary-frames, matching their model-supplied provenance. The subagent's `SystemMessage` never carries criterion text; it gets only the framework-owned `build_acceptance_criteria_system_note(...)` pointer naming the list's location and authority, so natural-language injection inside a criterion cannot gain system-channel priority over framework instructions. Deterministic leaf checking is a separate layer. diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index 5763eefc3..66c0f4ba9 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -60,6 +60,11 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) _EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS = 3.0 +# Kept as wire keys here instead of importing ``deerflow.sandbox`` at module +# load: executor tests and extension embedders replace that package while +# breaking agent/tool import cycles. +_SANDBOX_LEASE_OWNER_CONTEXT_KEY = "sandbox_lease_owner_id" +_SANDBOX_COMMAND_SCOPE_CONTEXT_KEY = "sandbox_command_scope_id" _previous_shutdown_isolated_subagent_loop = globals().get("_shutdown_isolated_subagent_loop") @@ -1308,6 +1313,8 @@ class SubagentExecutor: status=SubagentStatus.RUNNING, started_at=datetime.now(), ) + sandbox_lease_owner_id = f"subagent:{result.task_id}" + execution_context: dict[str, Any] | None = None from deerflow_extension_api import ExtensionData, TaskInfo from deerflow.extensions import get_loaded_extensions @@ -1465,6 +1472,9 @@ class SubagentExecutor: context["authz_attributes"] = dict(self.authz_attributes) context[DEERFLOW_TRACE_METADATA_KEY] = self.deerflow_trace_id context["is_subagent"] = True + context[_SANDBOX_LEASE_OWNER_CONTEXT_KEY] = sandbox_lease_owner_id + context[_SANDBOX_COMMAND_SCOPE_CONTEXT_KEY] = sandbox_lease_owner_id + execution_context = context context["agent_id"] = self.config.name if self.loop_detection_recorder is not None: context[LOOP_DETECTION_RECORDER_CONTEXT_KEY] = self.loop_detection_recorder @@ -1623,6 +1633,20 @@ class SubagentExecutor: ) finally: + if execution_context is not None and execution_context.get("sandbox_id") is not None: + try: + from deerflow.sandbox import get_sandbox_provider + from deerflow.sandbox.lease import get_sandbox_lease_manager + + provider = get_sandbox_provider() + await get_sandbox_lease_manager(provider).release_async(sandbox_lease_owner_id) + except Exception: + logger.warning( + "[trace=%s] Failed to release sandbox execution lease for subagent %s", + self.trace_id, + self.config.name, + exc_info=True, + ) if task_info is not None and task_store is not None: try: await notify_task_stop( diff --git a/backend/tests/blocking_io/test_dingtalk_receive_file.py b/backend/tests/blocking_io/test_dingtalk_receive_file.py index e9b00d37b..1415adf20 100644 --- a/backend/tests/blocking_io/test_dingtalk_receive_file.py +++ b/backend/tests/blocking_io/test_dingtalk_receive_file.py @@ -16,6 +16,7 @@ network leg is httpx-async and not the subject here. from __future__ import annotations import asyncio +import threading from types import SimpleNamespace from unittest.mock import AsyncMock @@ -24,6 +25,40 @@ import pytest pytestmark = pytest.mark.asyncio +class _BlockingRemoteSandbox: + def __init__(self) -> None: + self.update_started = threading.Event() + self.allow_update = threading.Event() + self.closed = False + self.updates: list[tuple[str, bytes]] = [] + self.released_scopes: list[str] = [] + + def update_file(self, path: str, content: bytes) -> None: + self.update_started.set() + assert self.allow_update.wait(timeout=2) + assert not self.closed + self.updates.append((path, content)) + + def release_command_scope(self, scope_id: str) -> None: + self.released_scopes.append(scope_id) + + +class _BlockingRemoteProvider: + def __init__(self) -> None: + self.sandbox = _BlockingRemoteSandbox() + self.release_calls: list[str] = [] + + async def acquire_async(self, thread_id=None, *, user_id=None): + return "remote-sandbox" + + def get(self, sandbox_id: str): + return self.sandbox if sandbox_id == "remote-sandbox" else None + + def release(self, sandbox_id: str) -> None: + self.release_calls.append(sandbox_id) + self.sandbox.closed = True + + async def test_receive_file_persist_does_not_block_event_loop(tmp_path, monkeypatch) -> None: from app.channels.dingtalk import DingTalkChannel from app.channels.message_bus import MessageBus @@ -54,3 +89,56 @@ async def test_receive_file_persist_does_not_block_event_loop(tmp_path, monkeypa assert "/uploads/a.pdf" in out.text assert out.files == [] + + +async def test_cancelled_receive_file_holds_sandbox_lease_until_remote_sync_finishes(tmp_path, monkeypatch) -> None: + from app.channels.dingtalk import DingTalkChannel + from app.channels.message_bus import MessageBus + from deerflow.config.paths import Paths + from deerflow.sandbox.lease import discard_sandbox_lease_manager, get_sandbox_lease_manager + + paths = await asyncio.to_thread(Paths, str(tmp_path)) + provider = _BlockingRemoteProvider() + manager = get_sandbox_lease_manager(provider) + monkeypatch.setattr("app.channels.dingtalk.get_paths", lambda: paths) + monkeypatch.setattr("app.channels.dingtalk.get_sandbox_provider", lambda: provider) + + channel = DingTalkChannel(MessageBus(), config={}) + channel._download_by_code = AsyncMock(return_value=b"DATA") + + await manager.acquire_async("active-run", "thread-1", user_id="ou-user") + receive_task = asyncio.create_task( + channel._receive_single_file( + "download-code", + "file", + "report.pdf", + "thread-1", + user_id="ou-user", + ) + ) + try: + assert await asyncio.to_thread(provider.sandbox.update_started.wait, 1) + for _ in range(3): + receive_task.cancel() + await asyncio.sleep(0) + + assert not receive_task.done() + await manager.release_async("active-run") + assert provider.release_calls == [] + assert not provider.sandbox.closed + + provider.sandbox.allow_update.set() + with pytest.raises(asyncio.CancelledError): + await receive_task + + assert provider.sandbox.updates == [("/mnt/user-data/uploads/report.pdf", b"DATA")] + assert provider.release_calls == ["remote-sandbox"] + assert provider.sandbox.closed + finally: + provider.sandbox.allow_update.set() + if not receive_task.done(): + receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await receive_task + await manager.release_async("active-run") + discard_sandbox_lease_manager(provider) diff --git a/backend/tests/blocking_io/test_feishu_receive_file.py b/backend/tests/blocking_io/test_feishu_receive_file.py index 59603284c..558e9d5d0 100644 --- a/backend/tests/blocking_io/test_feishu_receive_file.py +++ b/backend/tests/blocking_io/test_feishu_receive_file.py @@ -38,11 +38,15 @@ class _RemoteSandbox: def __init__(self) -> None: self.updates: list[tuple[str, bytes]] = [] self.update_thread_id: int | None = None + self.released_scopes: list[str] = [] def update_file(self, path: str, content: bytes) -> None: self.update_thread_id = threading.get_ident() self.updates.append((path, content)) + def release_command_scope(self, scope_id: str) -> None: + self.released_scopes.append(scope_id) + class _RemoteProvider: uses_thread_data_mounts = False @@ -75,6 +79,31 @@ class _MountedProvider: raise AssertionError("mounted uploads must not look up a sandbox") +class _BlockingRemoteSandbox(_RemoteSandbox): + def __init__(self) -> None: + super().__init__() + self.update_started = threading.Event() + self.allow_update = threading.Event() + self.closed = False + + def update_file(self, path: str, content: bytes) -> None: + self.update_started.set() + assert self.allow_update.wait(timeout=2) + assert not self.closed + super().update_file(path, content) + + +class _BlockingRemoteProvider(_RemoteProvider): + def __init__(self) -> None: + super().__init__() + self.sandbox = _BlockingRemoteSandbox() + self.release_calls: list[str] = [] + + def release(self, sandbox_id: str) -> None: + self.release_calls.append(sandbox_id) + self.sandbox.closed = True + + async def test_receive_file_remote_sandbox_does_not_block_event_loop(tmp_path, monkeypatch) -> None: from deerflow.config.paths import Paths @@ -124,3 +153,51 @@ async def test_receive_file_mounted_sandbox_skips_redundant_sync(tmp_path, monke assert result == "/mnt/user-data/uploads/report.pdf" uploaded = tmp_path / "users" / "ou-user" / "threads" / "thread-1" / "user-data" / "uploads" / "report.pdf" assert await asyncio.to_thread(uploaded.read_bytes) == b"DATA" + + +async def test_cancelled_receive_file_holds_sandbox_lease_until_remote_sync_finishes(tmp_path, monkeypatch) -> None: + from deerflow.config.paths import Paths + from deerflow.sandbox.lease import discard_sandbox_lease_manager, get_sandbox_lease_manager + + paths = await asyncio.to_thread(Paths, str(tmp_path)) + provider = _BlockingRemoteProvider() + manager = get_sandbox_lease_manager(provider) + monkeypatch.setattr("app.channels.feishu.get_paths", lambda: paths) + monkeypatch.setattr("app.channels.feishu.get_sandbox_provider", lambda: provider) + + await manager.acquire_async("active-run", "thread-1", user_id="ou-user") + receive_task = asyncio.create_task( + _channel_with_file()._receive_single_file( + "message-1", + "file-key", + "file", + "thread-1", + user_id="ou-user", + ) + ) + try: + assert await asyncio.to_thread(provider.sandbox.update_started.wait, 1) + for _ in range(3): + receive_task.cancel() + await asyncio.sleep(0) + + assert not receive_task.done() + await manager.release_async("active-run") + assert provider.release_calls == [] + assert not provider.sandbox.closed + + provider.sandbox.allow_update.set() + with pytest.raises(asyncio.CancelledError): + await receive_task + + assert provider.sandbox.updates == [("/mnt/user-data/uploads/report.pdf", b"DATA")] + assert provider.release_calls == ["remote-sandbox"] + assert provider.sandbox.closed + finally: + provider.sandbox.allow_update.set() + if not receive_task.done(): + receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await receive_task + await manager.release_async("active-run") + discard_sandbox_lease_manager(provider) diff --git a/backend/tests/test_aio_sandbox.py b/backend/tests/test_aio_sandbox.py index 610adc916..1c7a99af2 100644 --- a/backend/tests/test_aio_sandbox.py +++ b/backend/tests/test_aio_sandbox.py @@ -7,6 +7,29 @@ from unittest.mock import MagicMock, patch import pytest +class _TeardownFirstScopeLock: + """Let teardown clean a scope before one already-admitted waiter runs.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self.command_waiting = threading.Event() + self.allow_command = threading.Event() + self.command_done = threading.Event() + + def __enter__(self): + if threading.current_thread().name == "queued-command": + self.command_waiting.set() + self.allow_command.wait(timeout=2) + self._lock.acquire() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self._lock.release() + if threading.current_thread().name == "scope-teardown": + self.allow_command.set() + self.command_done.wait(timeout=2) + + def test_local_sandbox_client_bypasses_environment_proxy(): """Local sandbox API calls must not inherit HTTP_PROXY (#3441).""" from deerflow.community.aio_sandbox.aio_sandbox import AioSandbox @@ -178,8 +201,13 @@ class TestErrorObservationRetry: # ...and the retry targets exactly that created session, never an # uncreated/fabricated id (which would 404). assert exec_calls[1].get("id") == created_ids[0] - # ...and that one-shot recovery session is released afterwards so a - # sandbox that keeps hitting corruption doesn't accumulate sessions. + # The recovered session is promoted: future commands must not return + # to the corrupted implicit default session. + assert cleaned_ids == [] + assert sandbox.execute_command("again") == "ok" + assert exec_calls[-1].get("id") == created_ids[0] + + sandbox.close() assert cleaned_ids == [created_ids[0]] def test_cleanup_failure_does_not_mask_successful_retry(self, sandbox): @@ -201,6 +229,33 @@ class TestErrorObservationRetry: # into an "Error: ..." result. assert sandbox.execute_command("test") == "recovered" + def test_failed_replacement_never_falls_back_to_corrupt_default(self, sandbox): + created_ids: list[str] = [] + exec_ids: list[str | None] = [] + + def mock_create_session(id, **kwargs): + created_ids.append(id) + return SimpleNamespace(data=SimpleNamespace(session_id=id)) + + def mock_exec(command, **kwargs): + session_id = kwargs.get("id") + exec_ids.append(session_id) + if len(exec_ids) <= 2: + return SimpleNamespace(data=SimpleNamespace(output="'ErrorObservation' object has no attribute 'exit_code'")) + return SimpleNamespace(data=SimpleNamespace(output="healthy")) + + sandbox._client.shell.create_session = mock_create_session + sandbox._client.shell.exec_command = mock_exec + + assert "ErrorObservation" in sandbox.execute_command("first") + assert sandbox.execute_command("second") == "healthy" + assert exec_ids[0] is None + assert exec_ids[1] == created_ids[0] + # The next call creates another explicit session instead of touching + # the already-proven-corrupt implicit default again. + assert exec_ids[2] == created_ids[1] + assert all(session_id is not None for session_id in exec_ids[1:]) + def test_no_retry_on_clean_output(self, sandbox): """Normal output should not trigger a retry.""" call_count = 0 @@ -217,6 +272,222 @@ class TestErrorObservationRetry: assert call_count == 1 +class TestScopedShellSessions: + """Concurrent subagents use independent persistent shell sessions (#5128).""" + + def test_different_scopes_execute_concurrently(self, sandbox): + active = 0 + max_active = 0 + active_lock = threading.Lock() + start_barrier = threading.Barrier(2) + session_ids: list[str] = [] + + def create_session(id, **kwargs): + session_ids.append(id) + return SimpleNamespace(data=SimpleNamespace(session_id=id)) + + def overlapping_exec(command, **kwargs): + nonlocal active, max_active + with active_lock: + active += 1 + max_active = max(max_active, active) + start_barrier.wait(timeout=1) + with active_lock: + active -= 1 + return SimpleNamespace(data=SimpleNamespace(output=command, exit_code=0)) + + sandbox._client.shell.create_session = create_session + sandbox._client.shell.exec_command = overlapping_exec + + outputs: list[str] = [] + + def worker(scope_id: str): + outputs.append( + sandbox.execute_command_in_scope( + scope_id, + scope_id=scope_id, + ) + ) + + threads = [ + threading.Thread(target=worker, args=("subagent-a",)), + threading.Thread(target=worker, args=("subagent-b",)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert sorted(outputs) == ["subagent-a", "subagent-b"] + assert max_active == 2 + assert len(set(session_ids)) == 2 + + def test_same_scope_remains_serialized(self, sandbox): + call_log: list[tuple[str, str]] = [] + start_barrier = threading.Barrier(3) + + sandbox._client.shell.create_session = lambda id, **kwargs: SimpleNamespace(data=SimpleNamespace(session_id=id)) + + def slow_exec(command, **kwargs): + call_log.append(("enter", command)) + import time + + time.sleep(0.03) + call_log.append(("exit", command)) + return SimpleNamespace(data=SimpleNamespace(output=command, exit_code=0)) + + sandbox._client.shell.exec_command = slow_exec + + def worker(command: str): + start_barrier.wait() + sandbox.execute_command_in_scope(command, scope_id="one-subagent") + + threads = [threading.Thread(target=worker, args=(f"cmd-{index}",)) for index in range(3)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + for index in range(0, len(call_log), 2): + assert call_log[index][0] == "enter" + assert call_log[index + 1] == ("exit", call_log[index][1]) + + def test_corrupt_scoped_session_is_replaced_and_reused(self, sandbox): + created_ids: list[str] = [] + cleaned_ids: list[str] = [] + exec_ids: list[str] = [] + + def create_session(id, **kwargs): + created_ids.append(id) + return SimpleNamespace(data=SimpleNamespace(session_id=id)) + + def exec_command(command, **kwargs): + exec_ids.append(kwargs["id"]) + if len(exec_ids) == 1: + return SimpleNamespace( + data=SimpleNamespace( + output="'ErrorObservation' object has no attribute 'exit_code'", + exit_code=None, + ) + ) + return SimpleNamespace(data=SimpleNamespace(output="ok", exit_code=0)) + + sandbox._client.shell.create_session = create_session + sandbox._client.shell.exec_command = exec_command + sandbox._client.shell.cleanup_session = lambda session_id, **kwargs: cleaned_ids.append(session_id) + + assert sandbox.execute_command_in_scope("first", scope_id="subagent-a") == "ok" + assert len(created_ids) == 2 + assert cleaned_ids == [created_ids[0]] + assert exec_ids == [created_ids[0], created_ids[1]] + + assert sandbox.execute_command_in_scope("second", scope_id="subagent-a") == "ok" + assert exec_ids[-1] == created_ids[1] + + sandbox.release_command_scope("subagent-a") + assert cleaned_ids == created_ids + + def test_queued_command_cannot_restart_session_after_scope_release(self, sandbox): + created_ids: list[str] = [] + executed_commands: list[str] = [] + cleaned_ids: list[str] = [] + + sandbox._client.shell.create_session = lambda id, **kwargs: created_ids.append(id) + sandbox._client.shell.exec_command = lambda command, **kwargs: executed_commands.append(command) or SimpleNamespace(data=SimpleNamespace(output="ok", exit_code=0)) + sandbox._client.shell.cleanup_session = lambda session_id, **kwargs: cleaned_ids.append(session_id) + + assert sandbox.execute_command_in_scope("initial", scope_id="subagent-a") == "ok" + scoped = sandbox._scoped_shell_sessions["subagent-a"] + controlled_lock = _TeardownFirstScopeLock() + scoped.lock = controlled_lock + queued_results: list[str] = [] + + def queued_command() -> None: + try: + queued_results.append(sandbox.execute_command_in_scope("late", scope_id="subagent-a")) + finally: + controlled_lock.command_done.set() + + command_thread = threading.Thread(target=queued_command, name="queued-command") + command_thread.start() + assert controlled_lock.command_waiting.wait(timeout=1) + teardown_thread = threading.Thread( + target=sandbox.release_command_scope, + args=("subagent-a",), + name="scope-teardown", + ) + teardown_thread.start() + command_thread.join(timeout=2) + teardown_thread.join(timeout=2) + + assert not command_thread.is_alive() + assert not teardown_thread.is_alive() + assert queued_results == ["Error: sandbox command scope is no longer active"] + assert len(created_ids) == 1 + assert executed_commands == ["initial"] + assert cleaned_ids == created_ids + + def test_queued_command_cannot_restart_session_while_sandbox_closes(self, sandbox): + created_ids: list[str] = [] + executed_commands: list[str] = [] + cleaned_ids: list[str] = [] + + sandbox._client.shell.create_session = lambda id, **kwargs: created_ids.append(id) + sandbox._client.shell.exec_command = lambda command, **kwargs: executed_commands.append(command) or SimpleNamespace(data=SimpleNamespace(output="ok", exit_code=0)) + sandbox._client.shell.cleanup_session = lambda session_id, **kwargs: cleaned_ids.append(session_id) + + assert sandbox.execute_command_in_scope("initial", scope_id="subagent-a") == "ok" + scoped = sandbox._scoped_shell_sessions["subagent-a"] + controlled_lock = _TeardownFirstScopeLock() + scoped.lock = controlled_lock + queued_results: list[str] = [] + + def queued_command() -> None: + try: + queued_results.append(sandbox.execute_command_in_scope("late", scope_id="subagent-a")) + finally: + controlled_lock.command_done.set() + + command_thread = threading.Thread(target=queued_command, name="queued-command") + command_thread.start() + assert controlled_lock.command_waiting.wait(timeout=1) + teardown_thread = threading.Thread(target=sandbox.close, name="scope-teardown") + teardown_thread.start() + command_thread.join(timeout=2) + teardown_thread.join(timeout=2) + + assert not command_thread.is_alive() + assert not teardown_thread.is_alive() + assert queued_results == ["Error: sandbox command scope is no longer active"] + assert len(created_ids) == 1 + assert executed_commands == ["initial"] + assert cleaned_ids == created_ids + + def test_env_command_keeps_fresh_bash_exec_semantics(self, sandbox): + sandbox._client.bash.exec = MagicMock(return_value=SimpleNamespace(data=SimpleNamespace(stdout="ok", stderr="", exit_code=0))) + + assert ( + sandbox.execute_command_in_scope( + "echo $TOKEN", + env={"TOKEN": "secret"}, + scope_id="subagent-a", + ) + == "ok" + ) + sandbox._client.bash.exec.assert_called_once() + sandbox._client.shell.create_session.assert_not_called() + assert sandbox._scoped_shell_sessions == {} + + def test_closed_sandbox_rejects_new_scope_without_leaking_session(self, sandbox): + client = sandbox._client + + sandbox.close() + + assert sandbox.execute_command_in_scope("echo late", scope_id="subagent-late") == "Error: sandbox client is closed" + assert sandbox._scoped_shell_sessions == {} + client.shell.create_session.assert_not_called() + + class TestBashExecUnsupportedFailFast: """Regression tests for #3921: sandbox images older than all-in-one-sandbox 1.9.x have no ``/v1/bash/exec`` route, so every env-bearing command (skills diff --git a/backend/tests/test_artifacts_router.py b/backend/tests/test_artifacts_router.py index 9a20482c2..a2f8f1b96 100644 --- a/backend/tests/test_artifacts_router.py +++ b/backend/tests/test_artifacts_router.py @@ -16,6 +16,7 @@ from starlette.responses import FileResponse import app.gateway.routers.artifacts as artifacts_router from app.gateway.internal_auth import INTERNAL_OWNER_USER_ID_HEADER_NAME, INTERNAL_SYSTEM_ROLE from deerflow.config.paths import make_safe_user_id +from deerflow.sandbox.lease import get_sandbox_lease_manager ACTIVE_ARTIFACT_CASES = [ ("poc.html", ""), @@ -66,6 +67,7 @@ class _RemoteSandbox: def __init__(self, *, fail_next_update: bool = False) -> None: self.updates: list[tuple[str, bytes]] = [] self.fail_next_update = fail_next_update + self.released_scopes: list[str] = [] def update_file(self, path: str, content: bytes) -> None: if self.fail_next_update: @@ -73,6 +75,9 @@ class _RemoteSandbox: raise RuntimeError("sandbox sync failed") self.updates.append((path, content)) + def release_command_scope(self, scope_id: str) -> None: + self.released_scopes.append(scope_id) + class _RemoteSandboxProvider: uses_thread_data_mounts = False @@ -228,6 +233,37 @@ def test_update_artifact_syncs_non_mounted_sandbox(tmp_path, monkeypatch) -> Non assert artifact_path.read_text(encoding="utf-8") == "after" +def test_update_artifact_does_not_release_under_active_execution_lease(tmp_path, monkeypatch) -> None: + artifact_path = tmp_path / "note.txt" + artifact_path.write_text("before", encoding="utf-8") + provider = _RemoteSandboxProvider() + manager = get_sandbox_lease_manager(provider) + manager.retain( + "active-agent", + "sandbox-1", + thread_id="thread-1", + user_id="default", + ) + _patch_artifact_update_dependencies(monkeypatch, artifact_path, provider) + + asyncio.run( + call_unwrapped( + artifacts_router.update_artifact, + "thread-1", + "mnt/user-data/outputs/note.txt", + artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")), + _make_request(), + ) + ) + + assert provider.sandbox.updates == [("/mnt/user-data/outputs/note.txt", b"after")] + assert manager.binding_for("active-agent") == "sandbox-1" + assert provider.released == [] + + manager.release("active-agent") + assert provider.released == ["sandbox-1"] + + def test_update_artifact_releases_sandbox_when_initial_sync_fails(tmp_path, monkeypatch) -> None: artifact_path = tmp_path / "note.txt" artifact_path.write_text("before", encoding="utf-8") diff --git a/backend/tests/test_client.py b/backend/tests/test_client.py index 392158a42..d7f7ea66c 100644 --- a/backend/tests/test_client.py +++ b/backend/tests/test_client.py @@ -27,6 +27,8 @@ from deerflow.config.authorization_config import AuthorizationConfig, Authorizat from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig from deerflow.config.paths import Paths from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig +from deerflow.sandbox.lease import ensure_sandbox_lease_owner, get_sandbox_lease_manager +from deerflow.sandbox.sandbox_provider import reset_sandbox_provider, set_sandbox_provider from deerflow.skills.types import SkillCategory from deerflow.tools.mcp_metadata import tag_mcp_tool from deerflow.uploads.manager import PathTraversalError @@ -3652,6 +3654,85 @@ class TestStreamHardening: with pytest.raises(RuntimeError, match="model quota exceeded"): list(client.stream("hi", thread_id="t-err")) + def test_agent_exception_releases_embedded_execution_lease(self, client): + provider = MagicMock() + provider.get.return_value = MagicMock() + manager = get_sandbox_lease_manager(provider) + owner_ids: list[str] = [] + + def failing_stream(state, *, config, context, stream_mode): + del state, config, stream_mode + owner_id = ensure_sandbox_lease_owner(context) + assert owner_id is not None + owner_ids.append(owner_id) + context["sandbox_id"] = "shared" + manager.retain( + owner_id, + "shared", + thread_id="t-err-lease", + user_id="anonymous", + ) + raise RuntimeError("model quota exceeded") + yield # pragma: no cover + + agent = MagicMock() + agent.stream.side_effect = failing_stream + set_sandbox_provider(provider) + try: + with ( + patch.object(client, "_ensure_agent"), + patch.object(client, "_agent", agent), + pytest.raises(RuntimeError, match="model quota exceeded"), + ): + list(client.stream("hi", thread_id="t-err-lease")) + + assert len(owner_ids) == 1 + assert manager.binding_for(owner_ids[0]) is None + provider.get.return_value.release_command_scope.assert_called_once_with(owner_ids[0]) + provider.release.assert_called_once_with("shared") + finally: + reset_sandbox_provider() + + def test_abandoned_embedded_stream_releases_execution_lease(self, client): + provider = MagicMock() + provider.get.return_value = MagicMock() + manager = get_sandbox_lease_manager(provider) + owner_ids: list[str] = [] + + def blocking_stream(state, *, config, context, stream_mode): + del state, config, stream_mode + owner_id = ensure_sandbox_lease_owner(context) + assert owner_id is not None + owner_ids.append(owner_id) + context["sandbox_id"] = "shared" + manager.retain( + owner_id, + "shared", + thread_id="t-abandoned-lease", + user_id="anonymous", + ) + yield "values", {"messages": []} + raise AssertionError("abandoned stream continued") + + agent = MagicMock() + agent.stream.side_effect = blocking_stream + set_sandbox_provider(provider) + try: + with ( + patch.object(client, "_ensure_agent"), + patch.object(client, "_agent", agent), + ): + stream = client.stream("hi", thread_id="t-abandoned-lease") + assert next(stream).type == "values" + stream.close() + + assert len(owner_ids) == 1 + assert manager.binding_for(owner_ids[0]) is None + provider.get.return_value.release_command_scope.assert_called_once_with(owner_ids[0]) + provider.release.assert_called_once_with("shared") + finally: + reset_sandbox_provider() + def test_messages_without_id(self, client): """Messages without id attribute are emitted without crashing.""" ai = AIMessage(content="no id here") diff --git a/backend/tests/test_ensure_sandbox_initialized.py b/backend/tests/test_ensure_sandbox_initialized.py index 572c6a52e..fc53cb977 100644 --- a/backend/tests/test_ensure_sandbox_initialized.py +++ b/backend/tests/test_ensure_sandbox_initialized.py @@ -2,17 +2,30 @@ from __future__ import annotations +import asyncio +import threading + import pytest from langchain.tools import ToolRuntime from langgraph.types import Overwrite +from deerflow.sandbox.exceptions import SandboxNotFoundError +from deerflow.sandbox.lease import SANDBOX_LEASE_OWNER_CONTEXT_KEY, get_sandbox_lease_manager from deerflow.sandbox.sandbox import Sandbox from deerflow.sandbox.sandbox_provider import SandboxProvider, reset_sandbox_provider, set_sandbox_provider from deerflow.sandbox.search import GrepMatch -from deerflow.sandbox.tools import ensure_sandbox_initialized, ensure_sandbox_initialized_async +from deerflow.sandbox.tools import ( + _run_sync_tool_after_async_sandbox_init, + ensure_sandbox_initialized, + ensure_sandbox_initialized_async, +) class _StubSandbox(Sandbox): + def __init__(self, sandbox_id: str) -> None: + super().__init__(sandbox_id) + self.released_scopes: list[str] = [] + def execute_command(self, command: str, env: dict | None = None, timeout: float | None = None) -> str: del env, timeout return "OK" @@ -47,10 +60,14 @@ class _StubSandbox(Sandbox): def update_file(self, path: str, content: bytes) -> None: return None + def release_command_scope(self, scope_id: str) -> None: + self.released_scopes.append(scope_id) + class _RecordingProvider(SandboxProvider): def __init__(self) -> None: self.sandbox = _StubSandbox("stub") + self.released: list[str] = [] def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: raise AssertionError("state already carries a sandbox; acquire must not run") @@ -64,7 +81,7 @@ class _RecordingProvider(SandboxProvider): return None def release(self, sandbox_id: str) -> None: - return None + self.released.append(sandbox_id) class _FallthroughProvider(SandboxProvider): @@ -73,6 +90,7 @@ class _FallthroughProvider(SandboxProvider): def __init__(self) -> None: self.sandbox = _StubSandbox("fresh") self.acquired: list[str | None] = [] + self.released: list[str] = [] def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: self.acquired.append(thread_id) @@ -88,8 +106,27 @@ class _FallthroughProvider(SandboxProvider): return None def release(self, sandbox_id: str) -> None: + self.released.append(sandbox_id) + + +class _PostAcquireLookupFailureProvider(SandboxProvider): + """Provider that binds an id but cannot return its active client.""" + + def __init__(self) -> None: + self.released: list[str] = [] + + def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: + return "lost-after-acquire" + + async def acquire_async(self, thread_id: str | None = None, *, user_id: str | None = None) -> str: + return "lost-after-acquire" + + def get(self, sandbox_id: str) -> Sandbox | None: return None + def release(self, sandbox_id: str) -> None: + self.released.append(sandbox_id) + def _make_runtime(state: dict) -> ToolRuntime: return ToolRuntime( @@ -103,6 +140,53 @@ def _make_runtime(state: dict) -> ToolRuntime: ) +def test_post_acquire_lookup_failure_unwinds_sync_execution_lease() -> None: + provider = _PostAcquireLookupFailureProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({}) + runtime.context.update( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "sync-owner", + "thread_id": "thread-1", + "user_id": "user-1", + } + ) + manager = get_sandbox_lease_manager(provider) + + with pytest.raises(SandboxNotFoundError, match="Sandbox not found after acquisition"): + ensure_sandbox_initialized(runtime) + + assert manager.binding_for("sync-owner") is None + assert provider.released == ["lost-after-acquire"] + finally: + reset_sandbox_provider() + + +@pytest.mark.anyio +async def test_post_acquire_lookup_failure_unwinds_async_execution_lease() -> None: + provider = _PostAcquireLookupFailureProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({}) + runtime.context.update( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "async-owner", + "thread_id": "thread-1", + "user_id": "user-1", + } + ) + manager = get_sandbox_lease_manager(provider) + + with pytest.raises(SandboxNotFoundError, match="Sandbox not found after acquisition"): + await ensure_sandbox_initialized_async(runtime) + + assert manager.binding_for("async-owner") is None + assert provider.released == ["lost-after-acquire"] + finally: + reset_sandbox_provider() + + def test_ensure_sandbox_initialized_unwraps_overwrite_state() -> None: """Fork-restored state must not crash on the Overwrite wrapper.""" provider = _RecordingProvider() @@ -134,6 +218,63 @@ async def test_ensure_sandbox_initialized_async_unwraps_overwrite_state() -> Non assert runtime.context["sandbox_id"] == "parent-sandbox" +def test_fork_restored_owner_holds_non_releasing_scope_lease() -> None: + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + runtime.context.update( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "fork-child", + "thread_id": "thread-1", + "user_id": "user-1", + } + ) + + sandbox = ensure_sandbox_initialized(runtime) + manager = get_sandbox_lease_manager(provider) + + assert sandbox is provider.sandbox + assert manager.binding_for("fork-child") == "parent-sandbox" + assert isinstance(runtime.state["sandbox"], Overwrite) + + manager.release("fork-child") + + assert provider.sandbox.released_scopes == ["fork-child"] + assert provider.released == [] + finally: + reset_sandbox_provider() + + +@pytest.mark.anyio +async def test_async_fork_restored_owner_holds_non_releasing_scope_lease() -> None: + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + runtime.context.update( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "fork-child", + "thread_id": "thread-1", + "user_id": "user-1", + } + ) + + sandbox = await ensure_sandbox_initialized_async(runtime) + manager = get_sandbox_lease_manager(provider) + + assert sandbox is provider.sandbox + assert manager.binding_for("fork-child") == "parent-sandbox" + assert isinstance(runtime.state["sandbox"], Overwrite) + + await manager.release_async("fork-child") + + assert provider.sandbox.released_scopes == ["fork-child"] + assert provider.released == [] + finally: + reset_sandbox_provider() + + def test_ensure_sandbox_initialized_plain_state_unchanged() -> None: provider = _RecordingProvider() set_sandbox_provider(provider) @@ -166,6 +307,59 @@ def test_ensure_sandbox_initialized_acquires_fresh_when_parent_missing() -> None assert runtime.context["sandbox_id"] == "fresh-sandbox" +def test_fork_restored_owner_normally_releases_fresh_replacement() -> None: + provider = _FallthroughProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + runtime.context.update( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "fork-child", + "thread_id": "thread-1", + "user_id": "user-1", + } + ) + + sandbox = ensure_sandbox_initialized(runtime) + manager = get_sandbox_lease_manager(provider) + + assert sandbox is provider.sandbox + assert manager.binding_for("fork-child") == "fresh-sandbox" + manager.release("fork-child") + + assert provider.sandbox.released_scopes == ["fork-child"] + assert provider.released == ["fresh-sandbox"] + finally: + reset_sandbox_provider() + + +@pytest.mark.anyio +async def test_async_fork_restored_owner_normally_releases_fresh_replacement() -> None: + provider = _FallthroughProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": Overwrite({"sandbox_id": "parent-sandbox"})}) + runtime.context.update( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "fork-child", + "thread_id": "thread-1", + "user_id": "user-1", + } + ) + + sandbox = await ensure_sandbox_initialized_async(runtime) + manager = get_sandbox_lease_manager(provider) + + assert sandbox is provider.sandbox + assert manager.binding_for("fork-child") == "fresh-sandbox" + await manager.release_async("fork-child") + + assert provider.sandbox.released_scopes == ["fork-child"] + assert provider.released == ["fresh-sandbox"] + finally: + reset_sandbox_provider() + + @pytest.mark.anyio async def test_ensure_sandbox_initialized_async_plain_state_unchanged() -> None: provider = _RecordingProvider() @@ -180,6 +374,94 @@ async def test_ensure_sandbox_initialized_async_plain_state_unchanged() -> None: assert runtime.context["sandbox_id"] == "parent-sandbox" +def test_reuse_with_config_only_thread_id_binds_execution_owner() -> None: + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": {"sandbox_id": "parent-sandbox"}}) + runtime.context[SANDBOX_LEASE_OWNER_CONTEXT_KEY] = "config-owner" + runtime.config["configurable"]["thread_id"] = "thread-from-config" + + sandbox = ensure_sandbox_initialized(runtime) + + assert sandbox is provider.sandbox + assert get_sandbox_lease_manager(provider).binding_for("config-owner") == "parent-sandbox" + finally: + reset_sandbox_provider() + + +@pytest.mark.anyio +async def test_async_reuse_with_config_only_thread_id_binds_execution_owner() -> None: + provider = _RecordingProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": {"sandbox_id": "parent-sandbox"}}) + runtime.context[SANDBOX_LEASE_OWNER_CONTEXT_KEY] = "config-owner" + runtime.config["configurable"]["thread_id"] = "thread-from-config" + + sandbox = await ensure_sandbox_initialized_async(runtime) + + assert sandbox is provider.sandbox + assert get_sandbox_lease_manager(provider).binding_for("config-owner") == "parent-sandbox" + finally: + reset_sandbox_provider() + + +def test_reuse_replaces_stale_checkpoint_and_owner_binding() -> None: + provider = _FallthroughProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": {"sandbox_id": "parent-sandbox"}}) + runtime.context[SANDBOX_LEASE_OWNER_CONTEXT_KEY] = "stale-owner" + runtime.context["thread_id"] = "thread-1" + runtime.context["user_id"] = "user-1" + manager = get_sandbox_lease_manager(provider) + manager.retain( + "stale-owner", + "parent-sandbox", + thread_id="thread-1", + user_id="user-1", + ) + + sandbox = ensure_sandbox_initialized(runtime) + + assert sandbox is provider.sandbox + assert runtime.state["sandbox"] == {"sandbox_id": "fresh-sandbox"} + assert runtime.context["sandbox_id"] == "fresh-sandbox" + assert manager.binding_for("stale-owner") == "fresh-sandbox" + assert provider.acquired == ["thread-1"] + finally: + reset_sandbox_provider() + + +@pytest.mark.anyio +async def test_async_reuse_replaces_stale_checkpoint_and_owner_binding() -> None: + provider = _FallthroughProvider() + set_sandbox_provider(provider) + try: + runtime = _make_runtime({"sandbox": {"sandbox_id": "parent-sandbox"}}) + runtime.context[SANDBOX_LEASE_OWNER_CONTEXT_KEY] = "stale-owner" + runtime.context["thread_id"] = "thread-1" + runtime.context["user_id"] = "user-1" + manager = get_sandbox_lease_manager(provider) + manager.retain( + "stale-owner", + "parent-sandbox", + thread_id="thread-1", + user_id="user-1", + ) + + sandbox = await ensure_sandbox_initialized_async(runtime) + + assert sandbox is provider.sandbox + assert runtime.state["sandbox"] == {"sandbox_id": "fresh-sandbox"} + assert runtime.context["sandbox_id"] == "fresh-sandbox" + assert manager.binding_for("stale-owner") == "fresh-sandbox" + assert provider.acquired == ["thread-1"] + finally: + reset_sandbox_provider() + + @pytest.mark.anyio async def test_ensure_sandbox_initialized_async_acquires_fresh_when_parent_missing() -> None: """Same fall-through as the sync path: the fork-restored id is gone from @@ -198,3 +480,76 @@ async def test_ensure_sandbox_initialized_async_acquires_fresh_when_parent_missi assert sandbox is provider.sandbox assert runtime.state["sandbox"] == {"sandbox_id": "fresh-sandbox"} assert runtime.context["sandbox_id"] == "fresh-sandbox" + + +@pytest.mark.anyio +async def test_cancelled_async_tool_drains_worker_before_execution_lease_cleanup(monkeypatch) -> None: + """Cancellation must not let a late sync body re-admit a released owner.""" + provider = _FallthroughProvider() + set_sandbox_provider(provider) + worker_started = threading.Event() + allow_worker = threading.Event() + worker_finished = threading.Event() + try: + runtime = _make_runtime({"sandbox": {"sandbox_id": "fresh-sandbox"}}) + runtime.context.update( + { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "cancelled-child", + "thread_id": "thread-1", + "user_id": "user-1", + } + ) + manager = get_sandbox_lease_manager(provider) + await manager.acquire_async("cancelled-child", "thread-1", user_id="user-1") + await manager.acquire_async("parallel-sibling", "thread-1", user_id="user-1") + + async def _allow_sandbox(*, context, app_config): + del context, app_config + + async def _safe_config(): + return None + + monkeypatch.setattr("deerflow.sandbox.tools.authorize_sandbox_execution_async", _allow_sandbox) + monkeypatch.setattr("deerflow.sandbox.tools.safe_app_config_async", _safe_config) + + def _blocking_tool(inner_runtime: ToolRuntime) -> str: + worker_started.set() + assert allow_worker.wait(timeout=2) + try: + sandbox = ensure_sandbox_initialized(inner_runtime) + return sandbox.execute_command("late command") + finally: + worker_finished.set() + + async def _run_then_cleanup() -> str: + try: + return await _run_sync_tool_after_async_sandbox_init(_blocking_tool, runtime) + finally: + await manager.release_async("cancelled-child") + + execution = asyncio.create_task(_run_then_cleanup()) + assert await asyncio.to_thread(worker_started.wait, 1) + + for _ in range(3): + execution.cancel() + await asyncio.sleep(0) + + assert not execution.done() + assert manager.binding_for("cancelled-child") == "fresh-sandbox" + assert provider.released == [] + + allow_worker.set() + with pytest.raises(asyncio.CancelledError): + await execution + + assert worker_finished.is_set() + assert manager.binding_for("cancelled-child") is None + assert manager.binding_for("parallel-sibling") == "fresh-sandbox" + assert provider.released == [] + + await manager.release_async("parallel-sibling") + assert provider.released == ["fresh-sandbox"] + finally: + allow_worker.set() + await asyncio.to_thread(worker_finished.wait, 2) + reset_sandbox_provider() diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index 69aae6f0a..c2fc577a1 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -3029,6 +3029,24 @@ class TestInjectAuthenticatedUserContextAuthz: assert "authz_attributes" not in config["context"] assert config["context"]["is_internal"] is True + @pytest.mark.parametrize("auth_source", ["session", AUTH_SOURCE_INTERNAL]) + @pytest.mark.parametrize("section", ["context", "configurable"]) + def test_gateway_callers_cannot_inject_sandbox_execution_identities(self, auth_source, section): + """Only the in-process lead/subagent lifecycle may assign lease identities.""" + request = _make_request_with_auth_source(auth_source) + config = _assemble_authz_run_config( + { + section: { + "sandbox_lease_owner_id": "forged-owner", + "sandbox_command_scope_id": "forged-scope", + } + }, + request, + ) + + assert "sandbox_lease_owner_id" not in config[section] + assert "sandbox_command_scope_id" not in config[section] + def test_session_body_context_cannot_inject_channel_user_id(self): request = _make_request_with_auth_source("session") config = _assemble_authz_run_config( diff --git a/backend/tests/test_github_token_plumbing.py b/backend/tests/test_github_token_plumbing.py index 1e8229e81..c3b45cd58 100644 --- a/backend/tests/test_github_token_plumbing.py +++ b/backend/tests/test_github_token_plumbing.py @@ -149,6 +149,8 @@ def test_aio_sandbox_no_env_leaves_command_unchanged() -> None: sbx._lock = __import__("threading").Lock() sbx._client = SimpleNamespace(shell=_FakeShell()) sbx._DEFAULT_NO_CHANGE_TIMEOUT = 30 + sbx._recovery_session_id = None + sbx._default_shell_corrupted = False sbx.execute_command("echo hello") @@ -394,6 +396,55 @@ def test_bash_tool_no_env_without_token(monkeypatch: pytest.MonkeyPatch) -> None assert captured["env"] is None +def test_bash_tool_routes_subagent_command_to_its_shell_scope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runtime = SimpleNamespace( + state={"sandbox": {"sandbox_id": "aio:xyz"}}, + context={ + "thread_id": "t1", + "sandbox_command_scope_id": "subagent:task-1", + }, + config={}, + ) + captured: dict = {} + + class _Sandbox: + def execute_command(self, command, env=None, timeout=None): + pytest.fail("subagent command must use its scoped shell session") + + def execute_command_in_scope( + self, + command, + env=None, + timeout=None, + *, + scope_id=None, + ): + captured.update( + command=command, + env=env, + timeout=timeout, + scope_id=scope_id, + ) + return "done" + + monkeypatch.setattr( + "deerflow.sandbox.tools.ensure_sandbox_initialized", + lambda runtime: _Sandbox(), + ) + monkeypatch.setattr( + "deerflow.sandbox.tools.ensure_thread_directories_exist", + lambda runtime: None, + ) + + result = bash_tool.func(runtime=runtime, description="ls", command="ls") + + assert result == "done" + assert captured["scope_id"] == "subagent:task-1" + assert captured["command"] == "cd /mnt/user-data/workspace; ls" + + # --------------------------------------------------------------------------- # ChannelManager._apply_channel_policy — the unified per-channel run-policy # hook. The github channel registers (is_interactive=False, diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 59e60f758..a337341d1 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -44,6 +44,13 @@ from deerflow.runtime.runs.worker import ( _try_extract_from_message, run_agent, ) +from deerflow.sandbox.lease import ( + SANDBOX_COMMAND_SCOPE_CONTEXT_KEY, + SANDBOX_LEASE_OWNER_CONTEXT_KEY, + ensure_sandbox_lease_owner, + get_sandbox_lease_manager, +) +from deerflow.sandbox.sandbox_provider import reset_sandbox_provider, set_sandbox_provider class FakeCheckpointer: @@ -53,6 +60,115 @@ class FakeCheckpointer: self.aput_writes = AsyncMock() +def _lease_test_bridge(): + return SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + + +@pytest.mark.anyio +async def test_run_agent_releases_execution_lease_when_graph_raises(): + provider = MagicMock() + provider.get.return_value = MagicMock() + manager = get_sandbox_lease_manager(provider) + run_manager = RunManager() + record = await run_manager.create("thread-lead-error") + owner_ids: list[str] = [] + + class FailingAgent: + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, stream_mode, subgraphs + context = config["configurable"]["__pregel_runtime"].context + owner_id = ensure_sandbox_lease_owner(context) + assert owner_id is not None + owner_ids.append(owner_id) + context["sandbox_id"] = "shared" + manager.retain( + owner_id, + "shared", + thread_id=record.thread_id, + user_id="anonymous", + ) + raise RuntimeError("lead model failed") + yield # pragma: no cover + + set_sandbox_provider(provider) + try: + await run_agent( + _lease_test_bridge(), + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda **_kwargs: FailingAgent(), + graph_input={}, + config={}, + ) + await asyncio.sleep(0) + + assert len(owner_ids) == 1 + assert manager.binding_for(owner_ids[0]) is None + provider.get.return_value.release_command_scope.assert_called_once_with(owner_ids[0]) + provider.release.assert_called_once_with("shared") + finally: + reset_sandbox_provider() + + +@pytest.mark.anyio +async def test_run_agent_releases_execution_lease_when_cancelled(): + provider = MagicMock() + provider.get.return_value = MagicMock() + manager = get_sandbox_lease_manager(provider) + run_manager = RunManager() + record = await run_manager.create("thread-lead-cancel") + lease_bound = asyncio.Event() + owner_ids: list[str] = [] + + class BlockingAgent: + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, stream_mode, subgraphs + context = config["configurable"]["__pregel_runtime"].context + owner_id = ensure_sandbox_lease_owner(context) + assert owner_id is not None + owner_ids.append(owner_id) + context["sandbox_id"] = "shared" + manager.retain( + owner_id, + "shared", + thread_id=record.thread_id, + user_id="anonymous", + ) + lease_bound.set() + await asyncio.Event().wait() + yield # pragma: no cover + + set_sandbox_provider(provider) + try: + task = asyncio.create_task( + run_agent( + _lease_test_bridge(), + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda **_kwargs: BlockingAgent(), + graph_input={}, + config={}, + ) + ) + await asyncio.wait_for(lease_bound.wait(), timeout=1) + task.cancel() + await task + await asyncio.sleep(0) + + assert len(owner_ids) == 1 + assert manager.binding_for(owner_ids[0]) is None + provider.get.return_value.release_command_scope.assert_called_once_with(owner_ids[0]) + provider.release.assert_called_once_with("shared") + finally: + reset_sandbox_provider() + + @pytest.mark.anyio async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled(): class CleanupTrackingRunManager(RunManager): @@ -503,6 +619,26 @@ def test_install_runtime_context_overrides_internal_pre_existing_message_ids(): assert config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"old-ai"}) +def test_install_runtime_context_removes_caller_sandbox_execution_identities(): + config = { + "context": { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "forged-owner", + SANDBOX_COMMAND_SCOPE_CONTEXT_KEY: "forged-scope", + } + } + + _install_runtime_context( + config, + { + "thread_id": "record-thread", + "run_id": "run-1", + }, + ) + + assert SANDBOX_LEASE_OWNER_CONTEXT_KEY not in config["context"] + assert SANDBOX_COMMAND_SCOPE_CONTEXT_KEY not in config["context"] + + @pytest.mark.anyio async def test_run_agent_batches_incremental_file_args_and_keeps_complete_values(): run_manager = RunManager() @@ -2195,6 +2331,18 @@ def test_build_runtime_context_ignores_caller_pre_existing_message_ids(): assert CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY not in ctx +def test_build_runtime_context_ignores_caller_sandbox_execution_identities(): + caller_context = { + SANDBOX_LEASE_OWNER_CONTEXT_KEY: "forged-owner", + SANDBOX_COMMAND_SCOPE_CONTEXT_KEY: "forged-scope", + } + + ctx = _build_runtime_context("thread-1", "run-1", caller_context) + + assert SANDBOX_LEASE_OWNER_CONTEXT_KEY not in ctx + assert SANDBOX_COMMAND_SCOPE_CONTEXT_KEY not in ctx + + def test_build_runtime_context_ignores_non_dict_caller_context(): ctx = _build_runtime_context("thread-1", "run-1", "not-a-dict") assert ctx == {"thread_id": "thread-1", "run_id": "run-1"} diff --git a/backend/tests/test_sandbox_leases.py b/backend/tests/test_sandbox_leases.py new file mode 100644 index 000000000..a5533fd44 --- /dev/null +++ b/backend/tests/test_sandbox_leases.py @@ -0,0 +1,731 @@ +from __future__ import annotations + +import asyncio +import threading +from dataclasses import dataclass, field + +import pytest + +from deerflow.sandbox.lease import ( + SandboxLeaseManager, + discard_sandbox_lease_manager, + get_sandbox_lease_manager, +) +from deerflow.sandbox.sandbox import Sandbox +from deerflow.sandbox.sandbox_provider import SandboxProvider +from deerflow.sandbox.search import GrepMatch + + +class _LeaseSandbox(Sandbox): + def __init__(self, sandbox_id: str): + super().__init__(sandbox_id) + self.released_scopes: list[str] = [] + + def execute_command(self, command, env=None, timeout=None): + return command + + def release_command_scope(self, scope_id: str) -> None: + self.released_scopes.append(scope_id) + + def read_file(self, path, start_line=None, end_line=None): + return "" + + def download_file(self, path): + return b"" + + def list_dir(self, path, max_depth=2): + return [] + + def write_file(self, path, content, append=False): + return None + + def glob(self, path, pattern, *, include_dirs=False, max_results=200): + return [], False + + def grep( + self, + path: str, + pattern: str, + *, + glob: str | None = None, + literal: bool = False, + case_sensitive: bool = False, + max_results: int = 100, + ) -> tuple[list[GrepMatch], bool]: + return [], False + + def update_file(self, path, content): + return None + + +class _LeaseProvider(SandboxProvider): + def __init__(self): + self.sandbox = _LeaseSandbox("shared") + self.acquire_calls: list[tuple[str | None, str | None]] = [] + self.release_calls: list[str] = [] + + def acquire(self, thread_id=None, *, user_id=None): + self.acquire_calls.append((thread_id, user_id)) + return self.sandbox.id + + def get(self, sandbox_id): + return self.sandbox if sandbox_id == self.sandbox.id else None + + def release(self, sandbox_id): + self.release_calls.append(sandbox_id) + + +@dataclass +class _UnhashableLeaseProvider(SandboxProvider): + """Valid value-comparable provider whose instances are not hashable.""" + + marker: str = "same" + sandbox: _LeaseSandbox = field(default_factory=lambda: _LeaseSandbox("shared"), compare=False) + release_calls: list[str] = field(default_factory=list, compare=False) + + def acquire(self, thread_id=None, *, user_id=None): + return self.sandbox.id + + def get(self, sandbox_id): + return self.sandbox if sandbox_id == self.sandbox.id else None + + def release(self, sandbox_id): + self.release_calls.append(sandbox_id) + + +def test_manager_registry_supports_unhashable_provider() -> None: + provider = _UnhashableLeaseProvider() + assert provider.__hash__ is None + + try: + manager = get_sandbox_lease_manager(provider) + assert get_sandbox_lease_manager(provider) is manager + finally: + discard_sandbox_lease_manager(provider) + + +def test_manager_registry_distinguishes_equal_provider_instances_by_identity() -> None: + first = _UnhashableLeaseProvider() + second = _UnhashableLeaseProvider() + assert first == second + assert first is not second + + try: + first_manager = get_sandbox_lease_manager(first) + second_manager = get_sandbox_lease_manager(second) + + assert first_manager is not second_manager + first_manager.retain("first-owner", "shared", thread_id="thread-1", user_id="user-1") + second_manager.retain("second-owner", "shared", thread_id="thread-1", user_id="user-1") + assert first_manager.binding_for("second-owner") is None + assert second_manager.binding_for("first-owner") is None + + discard_sandbox_lease_manager(first) + assert get_sandbox_lease_manager(second) is second_manager + finally: + discard_sandbox_lease_manager(first) + discard_sandbox_lease_manager(second) + + +def test_last_execution_lease_is_the_only_provider_releaser() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + + for owner_id in ("parent", "child-a", "child-b"): + manager.retain( + owner_id, + "shared", + thread_id="thread-1", + user_id="user-1", + ) + + manager.release("child-a") + manager.release("parent") + assert provider.release_calls == [] + + manager.release("child-b") + assert provider.release_calls == ["shared"] + assert provider.sandbox.released_scopes == ["child-a", "parent", "child-b"] + + +def test_non_releasing_holder_defers_parent_release_until_its_scope_is_clean() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "parent", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + manager.retain( + "fork-child", + "shared", + thread_id="thread-1", + user_id="user-1", + release_on_last=False, + ) + + manager.release("parent") + assert provider.release_calls == [] + + manager.release("fork-child") + + assert provider.release_calls == ["shared"] + assert provider.sandbox.released_scopes == ["parent", "fork-child"] + + +def test_lone_non_releasing_holder_does_not_park_warm_sandbox() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "upload", + "shared", + thread_id="thread-1", + user_id="user-1", + release_on_last=False, + ) + + manager.release("upload") + + assert provider.release_calls == [] + assert provider.sandbox.released_scopes == ["upload"] + + +def test_normal_acquire_upgrades_existing_non_releasing_holder() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "child", + "shared", + thread_id="thread-1", + user_id="user-1", + release_on_last=False, + ) + + sandbox_id = manager.acquire("child", "thread-1", user_id="user-1") + manager.release("child") + + assert sandbox_id == "shared" + assert provider.acquire_calls == [] + assert provider.release_calls == ["shared"] + + +def test_release_is_idempotent_for_executor_finally_safety_net() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "child", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + + manager.release("child") + manager.release("child") + + assert provider.release_calls == ["shared"] + assert provider.sandbox.released_scopes == ["child"] + + +def test_repeated_acquire_for_same_owner_does_not_reacquire_provider() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + + first = manager.acquire("child", "thread-1", user_id="user-1") + second = manager.acquire("child", "thread-1", user_id="user-1") + + assert first == second == "shared" + assert provider.acquire_calls == [("thread-1", "user-1")] + + +class _BlockingLookupProvider(_LeaseProvider): + def __init__(self) -> None: + super().__init__() + self.lookup_started = threading.Event() + self.allow_lookup = threading.Event() + self._block_next_lookup = False + self._lookup_control = threading.Lock() + + def arm_lookup(self) -> None: + with self._lookup_control: + self._block_next_lookup = True + + def get(self, sandbox_id): + with self._lookup_control: + block_lookup = self._block_next_lookup + self._block_next_lookup = False + if block_lookup: + self.lookup_started.set() + assert self.allow_lookup.wait(timeout=1) + return super().get(sandbox_id) + + +def test_reuse_lookup_and_retain_block_last_owner_release_as_one_transition() -> None: + provider = _BlockingLookupProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "previous", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + provider.arm_lookup() + + reused_ids: list[str] = [] + reuse = manager.reuse_or_acquire + reuse_thread = threading.Thread( + target=lambda: reused_ids.append( + reuse( + "next", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + ) + ) + release_thread = threading.Thread(target=manager.release, args=("previous",)) + + reuse_thread.start() + assert provider.lookup_started.wait(timeout=1) + release_thread.start() + release_thread.join(timeout=0.05) + assert release_thread.is_alive() + + provider.allow_lookup.set() + reuse_thread.join(timeout=1) + release_thread.join(timeout=1) + + assert not reuse_thread.is_alive() + assert not release_thread.is_alive() + assert reused_ids == ["shared"] + assert manager.binding_for("next") == "shared" + assert provider.release_calls == [] + + +def test_async_reuse_lookup_and_retain_block_last_owner_release_as_one_transition() -> None: + provider = _BlockingLookupProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "previous", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + provider.arm_lookup() + + reused_ids: list[str] = [] + reuse_async = manager.reuse_or_acquire_async + + def run_async_reuse() -> None: + reused_ids.append( + asyncio.run( + reuse_async( + "next", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + ) + ) + + reuse_thread = threading.Thread(target=run_async_reuse) + release_thread = threading.Thread(target=manager.release, args=("previous",)) + + reuse_thread.start() + assert provider.lookup_started.wait(timeout=1) + release_thread.start() + release_thread.join(timeout=0.05) + assert release_thread.is_alive() + + provider.allow_lookup.set() + reuse_thread.join(timeout=1) + release_thread.join(timeout=1) + + assert not reuse_thread.is_alive() + assert not release_thread.is_alive() + assert reused_ids == ["shared"] + assert manager.binding_for("next") == "shared" + assert provider.release_calls == [] + + +def test_reuse_acquires_fresh_sandbox_for_stale_owner_binding() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "next", + "stale", + thread_id="thread-1", + user_id="user-1", + ) + + sandbox_id = manager.reuse_or_acquire( + "next", + "stale", + thread_id="thread-1", + user_id="user-1", + ) + + assert sandbox_id == "shared" + assert manager.binding_for("next") == "shared" + assert provider.acquire_calls == [("thread-1", "user-1")] + + +@pytest.mark.anyio +async def test_async_reuse_acquires_fresh_sandbox_for_stale_owner_binding() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "next", + "stale", + thread_id="thread-1", + user_id="user-1", + ) + + sandbox_id = await manager.reuse_or_acquire_async( + "next", + "stale", + thread_id="thread-1", + user_id="user-1", + ) + + assert sandbox_id == "shared" + assert manager.binding_for("next") == "shared" + assert provider.acquire_calls == [("thread-1", "user-1")] + + +@pytest.mark.anyio +async def test_async_lazy_acquires_share_one_release_boundary() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + + await manager.acquire_async("child-a", "thread-1", user_id="user-1") + await manager.acquire_async("child-b", "thread-1", user_id="user-1") + + await manager.release_async("child-a") + assert provider.release_calls == [] + await manager.release_async("child-b") + assert provider.release_calls == ["shared"] + + +@pytest.mark.anyio +async def test_repeated_async_acquire_for_same_owner_does_not_reacquire_provider() -> None: + provider = _LeaseProvider() + manager = SandboxLeaseManager(provider) + + first = await manager.acquire_async("child", "thread-1", user_id="user-1") + second = await manager.acquire_async("child", "thread-1", user_id="user-1") + + assert first == second == "shared" + assert provider.acquire_calls == [("thread-1", "user-1")] + + +@pytest.mark.anyio +async def test_cancelled_async_acquire_releases_unbound_provider_result() -> None: + acquire_started = asyncio.Event() + allow_acquire = asyncio.Event() + + class _BlockingAsyncProvider(_LeaseProvider): + async def acquire_async(self, thread_id=None, *, user_id=None): + self.acquire_calls.append((thread_id, user_id)) + acquire_started.set() + await allow_acquire.wait() + return self.sandbox.id + + provider = _BlockingAsyncProvider() + manager = SandboxLeaseManager(provider) + acquire_task = asyncio.create_task(manager.acquire_async("child", "thread-1", user_id="user-1")) + await acquire_started.wait() + + acquire_task.cancel() + await asyncio.sleep(0) + assert not acquire_task.done() + + allow_acquire.set() + with pytest.raises(asyncio.CancelledError): + await acquire_task + + assert manager.binding_for("child") is None + assert provider.release_calls == ["shared"] + + +@pytest.mark.anyio +async def test_repeated_cancellation_waits_for_provider_acquire_reconciliation() -> None: + acquire_started = asyncio.Event() + allow_acquire = asyncio.Event() + + class _BlockingAsyncProvider(_LeaseProvider): + async def acquire_async(self, thread_id=None, *, user_id=None): + self.acquire_calls.append((thread_id, user_id)) + acquire_started.set() + await allow_acquire.wait() + return self.sandbox.id + + provider = _BlockingAsyncProvider() + manager = SandboxLeaseManager(provider) + acquire_task = asyncio.create_task(manager.acquire_async("cancelled", "thread-1", user_id="user-1")) + await acquire_started.wait() + + acquire_task.cancel() + await asyncio.sleep(0) + acquire_task.cancel() + await asyncio.sleep(0) + + retain_done = threading.Event() + + def retain_next_owner() -> None: + manager.retain( + "next", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + retain_done.set() + + retain_thread = threading.Thread(target=retain_next_owner) + retain_thread.start() + try: + assert not acquire_task.done() + retain_thread.join(timeout=0.05) + assert retain_thread.is_alive() + + allow_acquire.set() + with pytest.raises(asyncio.CancelledError): + await acquire_task + retain_thread.join(timeout=1) + + assert not retain_thread.is_alive() + assert retain_done.is_set() + assert provider.release_calls == ["shared"] + assert manager.binding_for("cancelled") is None + assert manager.binding_for("next") == "shared" + finally: + allow_acquire.set() + retain_thread.join(timeout=1) + manager.close() + + +@pytest.mark.anyio +async def test_cancelled_async_acquire_keeps_same_thread_serialized_through_reconciliation() -> None: + acquire_started = asyncio.Event() + allow_acquire = asyncio.Event() + release_started = threading.Event() + allow_release = threading.Event() + + class _BlockingReconciliationProvider(_LeaseProvider): + async def acquire_async(self, thread_id=None, *, user_id=None): + self.acquire_calls.append((thread_id, user_id)) + acquire_started.set() + await allow_acquire.wait() + return self.sandbox.id + + def release(self, sandbox_id): + release_started.set() + allow_release.wait(timeout=1) + super().release(sandbox_id) + + provider = _BlockingReconciliationProvider() + manager = SandboxLeaseManager(provider) + acquire_task = asyncio.create_task(manager.acquire_async("cancelled", "thread-1", user_id="user-1")) + await acquire_started.wait() + + acquire_task.cancel() + allow_acquire.set() + assert await asyncio.to_thread(release_started.wait, 1) + acquire_task.cancel() + await asyncio.sleep(0) + assert not acquire_task.done() + + retain_done = threading.Event() + + def retain_next_owner() -> None: + manager.retain( + "next", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + retain_done.set() + + retain_thread = threading.Thread(target=retain_next_owner) + retain_thread.start() + retain_thread.join(timeout=0.05) + assert retain_thread.is_alive() + + allow_release.set() + with pytest.raises(asyncio.CancelledError): + await acquire_task + retain_thread.join(timeout=1) + + assert not retain_thread.is_alive() + assert retain_done.is_set() + assert provider.release_calls == ["shared"] + assert manager.binding_for("cancelled") is None + assert manager.binding_for("next") == "shared" + + +@pytest.mark.anyio +async def test_cancelled_async_acquire_logs_reconciliation_failure(caplog) -> None: + acquire_started = asyncio.Event() + allow_acquire_failure = asyncio.Event() + + class _FailingCancelledAcquireProvider(_LeaseProvider): + async def acquire_async(self, thread_id=None, *, user_id=None): + self.acquire_calls.append((thread_id, user_id)) + acquire_started.set() + await allow_acquire_failure.wait() + raise RuntimeError("provider acquire failed") + + provider = _FailingCancelledAcquireProvider() + manager = SandboxLeaseManager(provider) + acquire_task = asyncio.create_task(manager.acquire_async("cancelled", "thread-1", user_id="user-1")) + await acquire_started.wait() + + with caplog.at_level("WARNING", logger="deerflow.sandbox.lease"): + acquire_task.cancel() + allow_acquire_failure.set() + with pytest.raises(asyncio.CancelledError): + await acquire_task + + assert "Cancelled sandbox acquire failed during reconciliation" in caplog.text + assert "provider acquire failed" in caplog.text + + +@pytest.mark.anyio +async def test_cancelled_async_acquire_preserves_cancellation_when_rollback_fails(caplog) -> None: + acquire_started = asyncio.Event() + allow_acquire = asyncio.Event() + + class _FailingRollbackProvider(_LeaseProvider): + async def acquire_async(self, thread_id=None, *, user_id=None): + self.acquire_calls.append((thread_id, user_id)) + acquire_started.set() + await allow_acquire.wait() + return self.sandbox.id + + def release(self, sandbox_id): + raise RuntimeError("rollback failed") + + provider = _FailingRollbackProvider() + manager = SandboxLeaseManager(provider) + acquire_task = asyncio.create_task(manager.acquire_async("cancelled", "thread-1", user_id="user-1")) + await acquire_started.wait() + + with caplog.at_level("WARNING", logger="deerflow.sandbox.lease"): + acquire_task.cancel() + allow_acquire.set() + with pytest.raises(asyncio.CancelledError): + await acquire_task + + assert "Cancelled sandbox acquire rollback failed during reconciliation" in caplog.text + assert "rollback failed" in caplog.text + + +@pytest.mark.anyio +async def test_cancelled_async_release_logs_reconciliation_failure(caplog) -> None: + release_started = threading.Event() + allow_release_failure = threading.Event() + + class _FailingCancelledReleaseProvider(_LeaseProvider): + def release(self, sandbox_id): + release_started.set() + assert allow_release_failure.wait(timeout=1) + raise RuntimeError("provider release failed") + + provider = _FailingCancelledReleaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "cancelled", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + release_task = asyncio.create_task(manager.release_async("cancelled")) + assert await asyncio.to_thread(release_started.wait, 1) + + with caplog.at_level("WARNING", logger="deerflow.sandbox.lease"): + release_task.cancel() + allow_release_failure.set() + with pytest.raises(asyncio.CancelledError): + await release_task + + assert "Cancelled sandbox release failed during reconciliation" in caplog.text + assert "provider release failed" in caplog.text + + +@pytest.mark.anyio +async def test_repeated_cancellation_waits_for_async_release_reconciliation() -> None: + release_started = threading.Event() + allow_release = threading.Event() + + class _BlockingReleaseProvider(_LeaseProvider): + def release(self, sandbox_id): + release_started.set() + assert allow_release.wait(timeout=1) + super().release(sandbox_id) + + provider = _BlockingReleaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "cancelled", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + release_task = asyncio.create_task(manager.release_async("cancelled")) + assert await asyncio.to_thread(release_started.wait, 1) + + for _ in range(3): + release_task.cancel() + await asyncio.sleep(0) + try: + assert not release_task.done() + allow_release.set() + with pytest.raises(asyncio.CancelledError): + await release_task + + assert manager.binding_for("cancelled") is None + assert provider.release_calls == ["shared"] + finally: + allow_release.set() + manager.close() + + +def test_new_acquire_waits_until_last_release_transition_finishes() -> None: + release_started = threading.Event() + allow_release = threading.Event() + + class _BlockingReleaseProvider(_LeaseProvider): + def release(self, sandbox_id): + release_started.set() + allow_release.wait(timeout=1) + super().release(sandbox_id) + + provider = _BlockingReleaseProvider() + manager = SandboxLeaseManager(provider) + manager.retain( + "first", + "shared", + thread_id="thread-1", + user_id="user-1", + ) + + release_thread = threading.Thread(target=manager.release, args=("first",)) + acquire_thread = threading.Thread( + target=manager.acquire, + args=("second", "thread-1"), + kwargs={"user_id": "user-1"}, + ) + release_thread.start() + assert release_started.wait(timeout=1) + acquire_thread.start() + + acquire_thread.join(timeout=0.05) + assert acquire_thread.is_alive() + + allow_release.set() + release_thread.join(timeout=1) + acquire_thread.join(timeout=1) + assert not release_thread.is_alive() + assert not acquire_thread.is_alive() + assert provider.release_calls == ["shared"] + assert manager.binding_for("second") == "shared" diff --git a/backend/tests/test_sandbox_middleware.py b/backend/tests/test_sandbox_middleware.py index 6e009c4e0..51c6e73a9 100644 --- a/backend/tests/test_sandbox_middleware.py +++ b/backend/tests/test_sandbox_middleware.py @@ -7,17 +7,23 @@ import pytest from langchain.agents.middleware import AgentMiddleware from langchain.tools import ToolRuntime from langchain_core.messages import 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.thread_state import ThreadState from deerflow.sandbox.exceptions import SandboxAuthorizationError, SandboxRuntimeError +from deerflow.sandbox.lease import ( + get_sandbox_lease_manager, + release_sandbox_execution_lease, + release_sandbox_execution_lease_async, +) from deerflow.sandbox.middleware import SandboxMiddleware, SandboxMiddlewareState from deerflow.sandbox.sandbox import Sandbox from deerflow.sandbox.sandbox_provider import SandboxProvider, reset_sandbox_provider, set_sandbox_provider from deerflow.sandbox.search import GrepMatch -from deerflow.sandbox.tools import ls_tool +from deerflow.sandbox.tools import ensure_sandbox_initialized, ls_tool class _SyncProvider(SandboxProvider): @@ -309,19 +315,68 @@ async def test_abefore_agent_delegates_to_super_when_not_acquiring( runtime: Runtime, ) -> None: calls: list[tuple[dict, Runtime]] = [] + provider = _AsyncOnlyProvider() async def fake_super_abefore_agent(self, state_arg, runtime_arg): calls.append((state_arg, runtime_arg)) return {"delegated": True} monkeypatch.setattr(AgentMiddleware, "abefore_agent", fake_super_abefore_agent) - - result = await middleware.abefore_agent(state, runtime) + set_sandbox_provider(provider) + try: + result = await middleware.abefore_agent(state, runtime) + finally: + reset_sandbox_provider() assert result == {"delegated": True} assert calls == [(state, runtime)] +def test_shared_subagents_release_provider_only_after_last_execution() -> None: + """A child finishing must not park the sandbox under a running sibling (#5128).""" + provider = _AsyncOnlyProvider() + state = {"sandbox": {"sandbox_id": "async-sandbox"}} + first_runtime = Runtime( + context={ + "thread_id": "shared-thread", + "user_id": "shared-user", + "is_subagent": True, + } + ) + second_runtime = Runtime( + context={ + "thread_id": "shared-thread", + "user_id": "shared-user", + "is_subagent": True, + } + ) + middleware = SandboxMiddleware() + set_sandbox_provider(provider) + try: + middleware.before_agent(state, first_runtime) + middleware.before_agent(state, second_runtime) + for runtime in (first_runtime, second_runtime): + ensure_sandbox_initialized( + ToolRuntime( + state=state, + context=runtime.context, + config={"configurable": {}}, + stream_writer=lambda _: None, + tools=[], + tool_call_id="call-1", + store=None, + ) + ) + + middleware.after_agent(state, first_runtime) + assert provider.released_ids == [] + + middleware.after_agent(state, second_runtime) + assert provider.released_ids == ["async-sandbox"] + finally: + reset_sandbox_provider() + + @pytest.mark.anyio async def test_default_lazy_tool_acquisition_uses_async_provider() -> None: provider = _AsyncOnlyProvider() @@ -570,6 +625,38 @@ def test_wrap_tool_call_does_not_override_non_dict_update() -> None: assert result is cmd +def test_wrap_tool_call_defers_terminal_lease_release_to_outer_run_fence() -> None: + provider = _AsyncOnlyProvider() + owner_id = "agent:terminal" + state: dict = {"sandbox": {"sandbox_id": "async-sandbox"}} + request = _make_tool_call_request(state) + request.runtime.context.update( + thread_id="thread-1", + user_id="user-1", + sandbox_lease_owner_id=owner_id, + sandbox_id="async-sandbox", + ) + set_sandbox_provider(provider) + try: + get_sandbox_lease_manager(provider).retain( + owner_id, + "async-sandbox", + thread_id="thread-1", + user_id="user-1", + ) + result = SandboxMiddleware().wrap_tool_call( + request, + lambda _: Command(goto=END), + ) + assert provider.released_ids == [] + release_sandbox_execution_lease(request.runtime.context) + finally: + reset_sandbox_provider() + + assert isinstance(result, Command) + assert provider.released_ids == ["async-sandbox"] + + @pytest.mark.anyio async def test_awrap_tool_call_emits_command_when_lazy_init_happens() -> None: middleware = SandboxMiddleware() @@ -605,6 +692,93 @@ async def test_awrap_tool_call_passthrough_when_sandbox_already_in_state() -> No assert result is original +@pytest.mark.anyio +async def test_awrap_tool_call_defers_terminal_lease_release_to_outer_run_fence() -> None: + provider = _AsyncOnlyProvider() + owner_id = "agent:async-terminal" + state: dict = {"sandbox": {"sandbox_id": "async-sandbox"}} + request = _make_tool_call_request(state) + request.runtime.context.update( + thread_id="thread-1", + user_id="user-1", + sandbox_lease_owner_id=owner_id, + sandbox_id="async-sandbox", + ) + set_sandbox_provider(provider) + try: + get_sandbox_lease_manager(provider).retain( + owner_id, + "async-sandbox", + thread_id="thread-1", + user_id="user-1", + ) + result = await SandboxMiddleware().awrap_tool_call( + request, + lambda _: asyncio.sleep(0, result=Command(goto=END)), + ) + assert provider.released_ids == [] + await release_sandbox_execution_lease_async(request.runtime.context) + finally: + reset_sandbox_provider() + + assert isinstance(result, Command) + assert provider.released_ids == ["async-sandbox"] + + +@pytest.mark.anyio +async def test_parallel_terminal_command_does_not_release_while_sibling_handler_runs() -> None: + provider = _AsyncOnlyProvider() + owner_id = "agent:parallel-terminal" + state: dict = {"sandbox": {"sandbox_id": "async-sandbox"}} + request = _make_tool_call_request(state) + request.runtime.context.update( + thread_id="thread-1", + user_id="user-1", + sandbox_lease_owner_id=owner_id, + sandbox_id="async-sandbox", + ) + sibling_started = asyncio.Event() + allow_sibling_finish = asyncio.Event() + + async def sibling_handler(_: ToolCallRequest) -> ToolMessage: + sibling_started.set() + await allow_sibling_finish.wait() + return ToolMessage(content="done", tool_call_id="call-2", name="bash") + + async def terminal_handler(_: ToolCallRequest) -> Command: + await sibling_started.wait() + return Command(goto=END) + + set_sandbox_provider(provider) + try: + get_sandbox_lease_manager(provider).retain( + owner_id, + "async-sandbox", + thread_id="thread-1", + user_id="user-1", + ) + middleware = SandboxMiddleware() + sibling_task = asyncio.create_task(middleware.awrap_tool_call(request, sibling_handler)) + terminal_task = asyncio.create_task(middleware.awrap_tool_call(request, terminal_handler)) + + terminal_result = await terminal_task + + assert isinstance(terminal_result, Command) + assert not sibling_task.done() + assert provider.released_ids == [] + + allow_sibling_finish.set() + sibling_result = await sibling_task + assert isinstance(sibling_result, ToolMessage) + assert provider.released_ids == [] + + await release_sandbox_execution_lease_async(request.runtime.context) + finally: + reset_sandbox_provider() + + assert provider.released_ids == ["async-sandbox"] + + def test_wrap_tool_call_preserves_existing_command_fields_when_merging() -> None: """Regression: when merging sandbox_update into an existing Command, all other Command fields (e.g. graph, goto, resume) must be preserved. diff --git a/backend/tests/test_subagent_executor.py b/backend/tests/test_subagent_executor.py index 1c7fd9d40..94e447190 100644 --- a/backend/tests/test_subagent_executor.py +++ b/backend/tests/test_subagent_executor.py @@ -29,6 +29,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from packaging.version import Version +from deerflow.sandbox.lease import SandboxLeaseManager from deerflow.skills.types import Skill from deerflow.subagents.capacity import SubagentCapacityRejected from deerflow.trace_context import request_trace_context @@ -1479,6 +1480,125 @@ class TestAsyncExecutionPath: assert "Agent error" in result.error assert result.completed_at is not None + @pytest.mark.anyio + async def test_aexecute_finally_releases_only_the_failing_subagent_lease( + self, + classes, + base_config, + mock_agent, + monkeypatch, + ): + """The executor's outer finally must clean up a lease on graph failure.""" + SubagentExecutor = classes["SubagentExecutor"] + SubagentStatus = classes["SubagentStatus"] + + sandbox = MagicMock() + provider = MagicMock() + provider.get.return_value = sandbox + manager = SandboxLeaseManager(provider) + manager.retain( + "lead", + "shared", + thread_id="test-thread", + user_id="default", + ) + captured_owner: list[str] = [] + + async def failing_stream(*args, context, **kwargs): + owner_id = context["sandbox_lease_owner_id"] + captured_owner.append(owner_id) + context["sandbox_id"] = "shared" + manager.retain( + owner_id, + "shared", + thread_id=context["thread_id"], + user_id=context.get("user_id") or "default", + ) + raise RuntimeError("Agent error after sandbox acquisition") + yield # pragma: no cover - make this an async generator + + mock_agent.astream = failing_stream + sys.modules["deerflow.sandbox"].get_sandbox_provider.return_value = provider + lease_module = importlib.import_module("deerflow.sandbox.lease") + monkeypatch.setattr(lease_module, "get_sandbox_lease_manager", lambda _provider: manager) + + executor = SubagentExecutor( + config=base_config, + tools=[], + thread_id="test-thread", + ) + + with patch.object(executor, "_create_agent", return_value=mock_agent): + result = await executor._aexecute("Task") + + assert result.status == SubagentStatus.FAILED + assert "Agent error after sandbox acquisition" in result.error + assert len(captured_owner) == 1 + assert manager.binding_for(captured_owner[0]) is None + assert manager.binding_for("lead") == "shared" + assert sandbox.release_command_scope.call_args_list == [((captured_owner[0],), {})] + provider.release.assert_not_called() + + manager.release("lead") + provider.release.assert_called_once_with("shared") + + @pytest.mark.anyio + async def test_aexecute_fork_restored_state_cleans_scope_without_parking_parent( + self, + classes, + base_config, + mock_agent, + monkeypatch, + ): + """A fork-restored child is a client user even though it cannot park the parent.""" + from langgraph.types import Overwrite + + SubagentExecutor = classes["SubagentExecutor"] + SubagentStatus = classes["SubagentStatus"] + + sandbox = MagicMock() + provider = MagicMock() + provider.get.return_value = sandbox + manager = SandboxLeaseManager(provider) + captured_owner: list[str] = [] + + async def failing_stream(state, *args, context, **kwargs): + assert isinstance(state["sandbox"], Overwrite) + owner_id = context["sandbox_lease_owner_id"] + captured_owner.append(owner_id) + manager.retain( + owner_id, + "shared", + thread_id=context["thread_id"], + user_id=context.get("user_id") or "default", + release_on_last=False, + ) + context["sandbox_id"] = "shared" + raise RuntimeError("forked child failed after opening a scope") + yield # pragma: no cover - make this an async generator + + mock_agent.astream = failing_stream + sys.modules["deerflow.sandbox"].get_sandbox_provider.return_value = provider + lease_module = importlib.import_module("deerflow.sandbox.lease") + monkeypatch.setattr(lease_module, "get_sandbox_lease_manager", lambda _provider: manager) + + executor = SubagentExecutor( + config=base_config, + tools=[], + thread_id="test-thread", + sandbox_state=Overwrite({"sandbox_id": "shared"}), + ) + + with patch.object(executor, "_create_agent", return_value=mock_agent): + result = await executor._aexecute("Task") + + assert result.status == SubagentStatus.FAILED + assert "forked child failed after opening a scope" in result.error + assert len(captured_owner) == 1 + assert manager.binding_for(captured_owner[0]) is None + sandbox.release_command_scope.assert_called_once_with(captured_owner[0]) + provider.release.assert_not_called() + @pytest.mark.anyio async def test_aexecute_recursion_error_with_partial_surfaces_completed_turn_capped(self, classes, base_config, mock_agent, msg): """#3875 Phase 2: ``GraphRecursionError`` (``recursion_limit`` == @@ -3720,6 +3840,10 @@ class TestSubagentGuardrailAttribution: assert context.get("oauth_id") == "subj-123" assert context.get("run_id") == "run-42" assert context.get("is_subagent") is True + lease_owner = context.get("sandbox_lease_owner_id") + assert isinstance(lease_owner, str) + assert lease_owner.startswith("subagent:") + assert context.get("sandbox_command_scope_id") == lease_owner @pytest.mark.anyio async def test_aexecute_propagates_narrow_loop_detection_recorder( diff --git a/backend/tests/test_uploads_router.py b/backend/tests/test_uploads_router.py index fc8b81bcf..ae8f98893 100644 --- a/backend/tests/test_uploads_router.py +++ b/backend/tests/test_uploads_router.py @@ -1,6 +1,7 @@ import asyncio import os import stat +import threading from io import BytesIO from pathlib import Path from types import SimpleNamespace @@ -13,6 +14,7 @@ from fastapi.testclient import TestClient from app.gateway.deps import get_config from app.gateway.routers import uploads +from deerflow.sandbox.lease import get_sandbox_lease_manager class ChunkedUpload: @@ -266,6 +268,77 @@ def test_upload_files_syncs_non_local_sandbox_and_marks_markdown_file(tmp_path): sandbox.update_file.assert_any_call("/mnt/user-data/uploads/report.md", b"converted") +def test_upload_sync_holds_non_releasing_lease_while_active_agent_finishes(tmp_path): + thread_uploads_dir = tmp_path / "uploads" + thread_uploads_dir.mkdir(parents=True) + update_started = threading.Event() + allow_update = threading.Event() + + provider = MagicMock() + provider.uses_thread_data_mounts = False + provider.acquire.side_effect = AssertionError("upload route should use acquire_async") + provider.acquire_async = AsyncMock(return_value="aio-1") + sandbox = MagicMock() + + def blocking_update(*_args) -> None: + update_started.set() + assert allow_update.wait(timeout=1) + + sandbox.update_file.side_effect = blocking_update + provider.get.return_value = sandbox + manager = get_sandbox_lease_manager(provider) + manager.retain( + "active-agent", + "aio-1", + thread_id="thread-aio", + user_id="user-1", + ) + results = [] + errors: list[BaseException] = [] + + def run_upload() -> None: + try: + file = UploadFile(filename="notes.txt", file=BytesIO(b"hello uploads")) + results.append( + asyncio.run( + call_unwrapped( + uploads.upload_files, + "thread-aio", + request=MagicMock(), + files=[file], + config=SimpleNamespace(), + ) + ) + ) + except BaseException as exc: # pragma: no cover - surfaced below + errors.append(exc) + + with ( + patch.object(uploads, "get_effective_user_id", return_value="user-1"), + patch.object(uploads, "get_uploads_dir", return_value=thread_uploads_dir), + patch.object(uploads, "ensure_uploads_dir", return_value=thread_uploads_dir), + patch.object(uploads, "get_sandbox_provider", return_value=provider), + ): + upload_thread = threading.Thread(target=run_upload) + upload_thread.start() + assert update_started.wait(timeout=1) + + manager.release("active-agent") + provider.release.assert_not_called() + + allow_update.set() + upload_thread.join(timeout=2) + + assert not upload_thread.is_alive() + assert errors == [] + assert len(results) == 1 + assert results[0].success is True + provider.release.assert_called_once_with("aio-1") + assert sandbox.release_command_scope.call_count == 2 + request_scope_id = sandbox.release_command_scope.call_args_list[-1].args[0] + assert request_scope_id.startswith("gateway:upload:") + + def test_upload_files_makes_non_local_files_sandbox_writable(tmp_path): thread_uploads_dir = tmp_path / "uploads" thread_uploads_dir.mkdir(parents=True) @@ -472,7 +545,8 @@ def test_upload_files_does_not_sync_non_local_sandbox_when_total_size_exceeds_li assert exc_info.value.status_code == 413 provider.acquire.assert_not_called() provider.acquire_async.assert_awaited_once_with("thread-aio", user_id="owner-upload") - provider.get.assert_called_once_with("aio-1") + assert provider.get.call_count == 2 + assert all(call.args == ("aio-1",) for call in provider.get.call_args_list) sandbox.update_file.assert_not_called() @@ -501,7 +575,8 @@ def test_upload_files_does_not_sync_non_local_sandbox_when_conversion_fails(tmp_ assert exc_info.value.status_code == 500 provider.acquire.assert_not_called() provider.acquire_async.assert_awaited_once_with("thread-aio", user_id="owner-upload") - provider.get.assert_called_once_with("aio-1") + assert provider.get.call_count == 2 + assert all(call.args == ("aio-1",) for call in provider.get.call_args_list) sandbox.update_file.assert_not_called() assert not (thread_uploads_dir / "report.pdf").exists()