mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-11 14:38:38 +00:00
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
This commit is contained in:
parent
08b27aef73
commit
9e0fbd60fa
@ -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:
|
**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
|
- `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
|
- `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
|
- `_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
|
- `_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`).
|
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`).
|
||||||
|
|||||||
@ -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.commands import is_known_channel_command, strip_leading_mentions
|
||||||
from app.channels.connection_identity import attach_connection_identity
|
from app.channels.connection_identity import attach_connection_identity
|
||||||
from app.channels.message_bus import InboundMessage, InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment
|
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.config.paths import VIRTUAL_PATH_PREFIX, get_paths
|
||||||
from deerflow.runtime.user_context import get_effective_user_id
|
from deerflow.runtime.user_context import get_effective_user_id
|
||||||
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
|
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}"
|
virtual_path = f"{VIRTUAL_PATH_PREFIX}/uploads/{resolved_target.name}"
|
||||||
|
|
||||||
try:
|
try:
|
||||||
sandbox_provider = get_sandbox_provider()
|
sandbox_provider = await asyncio.to_thread(get_sandbox_provider)
|
||||||
# acquire_async keeps provider lifecycle work (Docker discovery,
|
synced = await sync_file_to_thread_sandbox(
|
||||||
# readiness polls) off the event loop; update_file is blocking
|
sandbox_provider,
|
||||||
# transport IO on remote sandboxes, so it is offloaded too.
|
thread_id=thread_id,
|
||||||
sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id)
|
user_id=effective_user_id,
|
||||||
if sandbox_id != "local":
|
virtual_path=virtual_path,
|
||||||
sandbox = sandbox_provider.get(sandbox_id)
|
content=content,
|
||||||
if sandbox is None:
|
owner_prefix="dingtalk-upload",
|
||||||
|
)
|
||||||
|
if not synced:
|
||||||
# Mirror Feishu: the agent's non-local sandbox cannot see this
|
# Mirror Feishu: the agent's non-local sandbox cannot see this
|
||||||
# file, so returning the virtual path would hand the model a
|
# file, so returning the virtual path would hand the model a
|
||||||
# path that reads as nothing — surface a failed-load marker.
|
# 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)
|
logger.warning("[DingTalk] sandbox not found after acquire, dropping attachment: %s", virtual_path)
|
||||||
return ""
|
return ""
|
||||||
await asyncio.to_thread(sandbox.update_file, virtual_path, content)
|
|
||||||
except Exception:
|
except Exception:
|
||||||
# Same failure mode as the sandbox-is-None branch: the bytes never
|
# Same failure mode as the sandbox-is-None branch: the bytes never
|
||||||
# reached the agent's sandbox, so the virtual path would read as
|
# reached the agent's sandbox, so the virtual path would read as
|
||||||
|
|||||||
@ -23,6 +23,7 @@ from app.channels.message_bus import (
|
|||||||
OutboundMessage,
|
OutboundMessage,
|
||||||
ResolvedAttachment,
|
ResolvedAttachment,
|
||||||
)
|
)
|
||||||
|
from app.channels.sandbox_files import sync_file_to_thread_sandbox
|
||||||
from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths
|
from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths
|
||||||
from deerflow.runtime.user_context import get_effective_user_id
|
from deerflow.runtime.user_context import get_effective_user_id
|
||||||
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
|
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
|
||||||
@ -495,13 +496,17 @@ class FeishuChannel(Channel):
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
sandbox_provider = await asyncio.to_thread(get_sandbox_provider)
|
sandbox_provider = await asyncio.to_thread(get_sandbox_provider)
|
||||||
if not getattr(sandbox_provider, "uses_thread_data_mounts", False):
|
synced = await sync_file_to_thread_sandbox(
|
||||||
sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id)
|
sandbox_provider,
|
||||||
sandbox = sandbox_provider.get(sandbox_id)
|
thread_id=thread_id,
|
||||||
if sandbox is None:
|
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)
|
logger.warning("[Feishu] sandbox not found for thread_id=%s", thread_id)
|
||||||
return f"Failed to obtain the [{type}]"
|
return f"Failed to obtain the [{type}]"
|
||||||
await asyncio.to_thread(sandbox.update_file, virtual_path, content)
|
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("[Feishu] failed to sync resource into non-local sandbox: %s", virtual_path)
|
logger.exception("[Feishu] failed to sync resource into non-local sandbox: %s", virtual_path)
|
||||||
return f"Failed to obtain the [{type}]"
|
return f"Failed to obtain the [{type}]"
|
||||||
|
|||||||
43
backend/app/channels/sandbox_files.py
Normal file
43
backend/app/channels/sandbox_files.py
Normal file
@ -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()
|
||||||
@ -52,9 +52,9 @@ owner-scoped assistant version selection remains enabled.
|
|||||||
| **Subagents** (`/api/subagents`) | Admin managed-worker CRUD and listing. |
|
| **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. |
|
| **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 |
|
| **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 |
|
| **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 (`<think>...</think>`, including unclosed/truncated blocks from reasoning models like MiniMax-M3) is stripped before JSON parsing |
|
| **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 (`<think>...</think>`, 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 `<think>` substrings in the rewrite (`strip_think_blocks(truncate_unclosed=False)`) |
|
| **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 `<think>` 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`. |
|
| **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`. |
|
||||||
|
|||||||
@ -33,7 +33,9 @@ import asyncio
|
|||||||
import functools
|
import functools
|
||||||
import inspect
|
import inspect
|
||||||
import logging
|
import logging
|
||||||
|
import uuid
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
|
from dataclasses import dataclass
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
|
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
|
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(
|
async def try_acquire_sandbox_for_request(
|
||||||
request: Request,
|
request: Request,
|
||||||
sandbox_provider,
|
sandbox_provider,
|
||||||
@ -405,19 +428,21 @@ async def try_acquire_sandbox_for_request(
|
|||||||
*,
|
*,
|
||||||
user_id: str,
|
user_id: str,
|
||||||
app_config: AppConfig | None,
|
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.
|
"""Gate + acquire the thread sandbox for a Gateway sync path.
|
||||||
|
|
||||||
Single entry point for the uploads/artifacts sandbox-sync paths so the
|
Single entry point for the uploads/artifacts sandbox-sync paths so the
|
||||||
deny/skip semantics live in one place: runs the ``sandbox:execute`` gate
|
deny/skip semantics live in one place: runs the ``sandbox:execute`` gate
|
||||||
for the request's user, then acquires the sandbox. Returns
|
for the request's user, then acquires the sandbox under a unique request
|
||||||
``(sandbox, sandbox_id, denied)``:
|
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
|
the primary operation (upload / artifact edit) proceeds without the
|
||||||
sandbox copy.
|
sandbox copy.
|
||||||
- allowed → ``(sandbox, sandbox_id, False)``: ``sandbox`` is the acquired
|
- allowed → ``sandbox`` is the acquired instance, or ``sandbox is None`` when
|
||||||
instance (``sandbox_id`` for later release), or ``sandbox is None`` when
|
|
||||||
the provider lost it right after acquiring (infrastructure error —
|
the provider lost it right after acquiring (infrastructure error —
|
||||||
callers surface it as 500 / RuntimeError respectively, since that is
|
callers surface it as 500 / RuntimeError respectively, since that is
|
||||||
not a policy decision).
|
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)
|
authorize_sandbox_for_request(user, is_internal=_is_internal_caller(request, user), app_config=app_config)
|
||||||
except SandboxAuthorizationError:
|
except SandboxAuthorizationError:
|
||||||
logger.info("Sandbox sync skipped: sandbox execution not permitted for this caller (thread_id=%s)", thread_id)
|
logger.info("Sandbox sync skipped: sandbox execution not permitted for this caller (thread_id=%s)", thread_id)
|
||||||
return None, None, True
|
return SandboxRequestLease(
|
||||||
sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=user_id)
|
sandbox=None,
|
||||||
return sandbox_provider.get(sandbox_id), sandbox_id, False
|
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:
|
async def _authenticate(request: Request) -> AuthContext:
|
||||||
|
|||||||
@ -16,7 +16,7 @@ from fastapi import APIRouter, HTTPException, Request
|
|||||||
from fastapi.responses import FileResponse, Response
|
from fastapi.responses import FileResponse, Response
|
||||||
from pydantic import BaseModel, Field
|
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.deps import get_run_manager
|
||||||
from app.gateway.internal_auth import get_trusted_internal_owner_user_id
|
from app.gateway.internal_auth import get_trusted_internal_owner_user_id
|
||||||
from app.gateway.path_utils import resolve_thread_virtual_path
|
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)
|
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()
|
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_lease: SandboxRequestLease | None = None
|
||||||
sandbox_id: str | None = None
|
|
||||||
sandbox = None
|
sandbox = None
|
||||||
try:
|
try:
|
||||||
async with reserve_artifact_write(request, thread_id, user_id=effective_user_id):
|
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
|
# role skips the sandbox sync; the host-side artifact update
|
||||||
# still completes (the agent cannot consume the sandbox copy
|
# still completes (the agent cannot consume the sandbox copy
|
||||||
# anyway when sandbox execution is denied).
|
# 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,
|
request,
|
||||||
sandbox_provider,
|
sandbox_provider,
|
||||||
thread_id,
|
thread_id,
|
||||||
user_id=effective_user_id,
|
user_id=effective_user_id,
|
||||||
app_config=safe_app_config(),
|
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")
|
raise RuntimeError("Failed to acquire sandbox for artifact update")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@ -569,11 +570,15 @@ async def update_artifact(
|
|||||||
logger.exception("Failed to update artifact %s for thread %s", path, thread_id)
|
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
|
raise HTTPException(status_code=500, detail="Failed to update artifact") from None
|
||||||
finally:
|
finally:
|
||||||
if sandbox_id is not None and sandbox_provider is not None:
|
if sandbox_lease is not None:
|
||||||
try:
|
try:
|
||||||
await asyncio.to_thread(sandbox_provider.release, sandbox_id)
|
await sandbox_lease.release()
|
||||||
except Exception:
|
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()
|
content_sha256 = hashlib.sha256(updated).hexdigest()
|
||||||
return ArtifactUpdateResponse(
|
return ArtifactUpdateResponse(
|
||||||
|
|||||||
@ -11,7 +11,7 @@ from typing import BinaryIO
|
|||||||
from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile
|
from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile
|
||||||
from pydantic import BaseModel, Field
|
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 app.gateway.deps import get_config
|
||||||
from deerflow.config.app_config import AppConfig
|
from deerflow.config.app_config import AppConfig
|
||||||
from deerflow.config.paths import get_paths
|
from deerflow.config.paths import get_paths
|
||||||
@ -338,22 +338,28 @@ async def upload_files(
|
|||||||
|
|
||||||
sandbox_provider = get_sandbox_provider()
|
sandbox_provider = get_sandbox_provider()
|
||||||
sync_to_sandbox = not _uses_thread_data_mounts(sandbox_provider)
|
sync_to_sandbox = not _uses_thread_data_mounts(sandbox_provider)
|
||||||
|
sandbox_lease: SandboxRequestLease | None = None
|
||||||
sandbox = None
|
sandbox = None
|
||||||
|
try:
|
||||||
if sync_to_sandbox:
|
if sync_to_sandbox:
|
||||||
# Phase 3: enforce sandbox:execute before acquiring — a role denied
|
# Phase 3: enforce sandbox:execute before acquiring — a role denied
|
||||||
# sandbox execution must not trigger sandbox allocation just by
|
# sandbox execution must not trigger sandbox allocation just by
|
||||||
# uploading files. Deny skips the sync; the upload itself still
|
# uploading files. Deny skips the sync; the upload itself still
|
||||||
# succeeds (files stay in the thread uploads dir; the agent cannot
|
# succeeds (files stay in the thread uploads dir; the agent cannot
|
||||||
# consume them via sandbox anyway).
|
# consume them via sandbox anyway).
|
||||||
sandbox, _sandbox_id, sandbox_denied = await try_acquire_sandbox_for_request(
|
sandbox_lease = await try_acquire_sandbox_for_request(
|
||||||
request,
|
request,
|
||||||
sandbox_provider,
|
sandbox_provider,
|
||||||
thread_id,
|
thread_id,
|
||||||
user_id=effective_user_id,
|
user_id=effective_user_id,
|
||||||
app_config=config,
|
app_config=config,
|
||||||
|
owner_prefix="gateway:upload",
|
||||||
|
release_on_last=False,
|
||||||
)
|
)
|
||||||
if not sandbox_denied and sandbox is None:
|
sandbox = sandbox_lease.sandbox
|
||||||
|
if not sandbox_lease.denied and sandbox is None:
|
||||||
raise HTTPException(status_code=500, detail="Failed to acquire sandbox")
|
raise HTTPException(status_code=500, detail="Failed to acquire sandbox")
|
||||||
|
|
||||||
auto_convert_documents = _auto_convert_documents_enabled(config)
|
auto_convert_documents = _auto_convert_documents_enabled(config)
|
||||||
|
|
||||||
for file in files:
|
for file in files:
|
||||||
@ -459,6 +465,17 @@ async def upload_files(
|
|||||||
skipped_files=skipped_files,
|
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)
|
@router.get("/limits", response_model=UploadLimits)
|
||||||
@require_permission("threads", "read", owner_check=True)
|
@require_permission("threads", "read", owner_check=True)
|
||||||
|
|||||||
@ -77,6 +77,7 @@ from deerflow.runtime.secret_context import (
|
|||||||
)
|
)
|
||||||
from deerflow.runtime.stream_modes import normalize_stream_modes
|
from deerflow.runtime.stream_modes import normalize_stream_modes
|
||||||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
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.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.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_context, ensure_trace_id
|
||||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||||
@ -408,15 +409,18 @@ _CONTEXT_CONFIGURABLE_KEYS: frozenset[str] = frozenset(
|
|||||||
# arbitrary HTTP/IM clients must not be able to force autonomous execution.
|
# arbitrary HTTP/IM clients must not be able to force autonomous execution.
|
||||||
_CONTEXT_INTERNAL_CALLER_KEYS: frozenset[str] = frozenset({"non_interactive"})
|
_CONTEXT_INTERNAL_CALLER_KEYS: frozenset[str] = frozenset({"non_interactive"})
|
||||||
|
|
||||||
# Server-owned authorization identity fields. These must never be accepted from
|
# Server-owned authorization and sandbox lifecycle identity fields. These must
|
||||||
# client-supplied ``body.config.context`` or ``body.config.configurable``. They
|
# 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
|
# are either produced by Gateway auth state, admitted from a separately
|
||||||
# authenticated internal request channel, or reserved for LangGraph Server.
|
# authenticated internal request channel, or reserved for LangGraph Server.
|
||||||
# ``is_internal`` — derived from ``request.state.auth_source``
|
# ``is_internal`` — derived from ``request.state.auth_source``
|
||||||
# ``authz_attributes`` — Phase 1A has no Gateway-side producer; cleared.
|
# ``authz_attributes`` — Phase 1A has no Gateway-side producer; cleared.
|
||||||
# ``channel_user_id`` — accepted only from trusted internal context.
|
# ``channel_user_id`` — accepted only from trusted internal context.
|
||||||
# ``langgraph_auth_user*`` — populated only by LangGraph Server auth.
|
# ``langgraph_auth_user*`` — populated only by LangGraph Server auth.
|
||||||
_SERVER_OWNED_AUTHZ_CONTEXT_KEYS: frozenset[str] = frozenset(
|
# ``sandbox_*_id`` — created only inside the run/subagent lifecycle.
|
||||||
|
_SERVER_OWNED_RUNTIME_CONTEXT_KEYS: frozenset[str] = (
|
||||||
|
frozenset(
|
||||||
{
|
{
|
||||||
"is_internal",
|
"is_internal",
|
||||||
"authz_attributes",
|
"authz_attributes",
|
||||||
@ -425,6 +429,8 @@ _SERVER_OWNED_AUTHZ_CONTEXT_KEYS: frozenset[str] = frozenset(
|
|||||||
"langgraph_auth_user_id",
|
"langgraph_auth_user_id",
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
| SANDBOX_SERVER_OWNED_CONTEXT_KEYS
|
||||||
|
)
|
||||||
|
|
||||||
# Keys forwarded from ``body.context`` into ``config['context']`` ONLY (the
|
# Keys forwarded from ``body.context`` into ``config['context']`` ONLY (the
|
||||||
# runtime context that becomes ``ToolRuntime.context`` / ``runtime.context``),
|
# runtime context that becomes ``ToolRuntime.context`` / ``runtime.context``),
|
||||||
@ -534,18 +540,18 @@ def inject_authenticated_user_context(
|
|||||||
Values copied through the free-form RunnableConfig are always cleared.
|
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
|
# Clear any client-forged values from both config sections, then write the
|
||||||
# authoritative is_internal. This runs before ALL early returns so that
|
# authoritative is_internal. This runs before ALL early returns so that
|
||||||
# even user_id-is-None paths get a defined is_internal value.
|
# even user_id-is-None paths get a defined is_internal value.
|
||||||
runtime_context = config.setdefault("context", {})
|
runtime_context = config.setdefault("context", {})
|
||||||
if not isinstance(runtime_context, dict):
|
if not isinstance(runtime_context, dict):
|
||||||
raise TypeError("run context must be a mapping")
|
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)
|
runtime_context.pop(key, None)
|
||||||
configurable = config.get("configurable")
|
configurable = config.get("configurable")
|
||||||
if isinstance(configurable, dict):
|
if isinstance(configurable, dict):
|
||||||
for key in _SERVER_OWNED_AUTHZ_CONTEXT_KEYS:
|
for key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS:
|
||||||
configurable.pop(key, None)
|
configurable.pop(key, None)
|
||||||
auth_source = getattr(getattr(request, "state", None), "auth_source", None)
|
auth_source = getattr(getattr(request, "state", None), "auth_source", None)
|
||||||
# ``user_id`` is server-owned for EXTERNAL callers: it now selects which
|
# ``user_id`` is server-owned for EXTERNAL callers: it now selects which
|
||||||
|
|||||||
@ -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.
|
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.
|
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.
|
||||||
|
|
||||||
|
|||||||
@ -23,7 +23,7 @@ import mimetypes
|
|||||||
import os
|
import os
|
||||||
import shutil
|
import shutil
|
||||||
import uuid
|
import uuid
|
||||||
from collections.abc import Generator, Mapping, Sequence
|
from collections.abc import Generator, Iterator, Mapping, Sequence
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Literal
|
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):
|
def _run_async_from_sync(coro):
|
||||||
"""Run an async helper from this synchronous client API."""
|
"""Run an async helper from this synchronous client API."""
|
||||||
try:
|
try:
|
||||||
@ -961,12 +974,13 @@ class DeerFlowClient:
|
|||||||
sent.update(delta)
|
sent.update(delta)
|
||||||
return delta
|
return delta
|
||||||
|
|
||||||
for item in self._agent.stream(
|
agent_items = self._agent.stream(
|
||||||
state,
|
state,
|
||||||
config=config,
|
config=config,
|
||||||
context=context,
|
context=context,
|
||||||
stream_mode=["values", "messages", "custom"],
|
stream_mode=["values", "messages", "custom"],
|
||||||
):
|
)
|
||||||
|
for item in _stream_with_sandbox_lease_cleanup(agent_items, context):
|
||||||
if isinstance(item, tuple) and len(item) == 2:
|
if isinstance(item, tuple) and len(item) == 2:
|
||||||
mode, chunk = item
|
mode, chunk = item
|
||||||
mode = str(mode)
|
mode = str(mode)
|
||||||
|
|||||||
@ -4,6 +4,7 @@ import logging
|
|||||||
import shlex
|
import shlex
|
||||||
import threading
|
import threading
|
||||||
import uuid
|
import uuid
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
from agent_sandbox import Sandbox as AioSandboxClient
|
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):
|
class AioSandbox(Sandbox):
|
||||||
"""Sandbox implementation using the agent-infra/sandbox Docker container.
|
"""Sandbox implementation using the agent-infra/sandbox Docker container.
|
||||||
|
|
||||||
This sandbox connects to a running AIO sandbox container via HTTP API.
|
This sandbox connects to a running AIO sandbox container via HTTP API.
|
||||||
A threading lock serializes shell commands to prevent concurrent requests
|
Lead/direct calls retain the legacy serialized shell. Delegated executions
|
||||||
from corrupting the container's single persistent session (see #1433).
|
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,
|
#: The legacy exec path reuses one persistent shell session across calls,
|
||||||
@ -69,6 +80,10 @@ class AioSandbox(Sandbox):
|
|||||||
)
|
)
|
||||||
self._home_dir = home_dir
|
self._home_dir = home_dir
|
||||||
self._lock = threading.Lock()
|
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
|
self._closed = False
|
||||||
# Set to True after bash.exec answers 404 (image predates /v1/bash/*),
|
# 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).
|
# 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
|
failures during teardown are logged and swallowed so provider/backend
|
||||||
cleanup is never blocked.
|
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:
|
if self._closed:
|
||||||
return
|
return
|
||||||
self._closed = True
|
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
|
client = self._client
|
||||||
# Drop the reference under the lock for use-after-close safety: any
|
# Drop the reference under the lock for use-after-close safety: any
|
||||||
# later command on this instance fails loudly instead of reusing a
|
# later command on this instance fails loudly instead of reusing a
|
||||||
@ -129,6 +168,157 @@ class AioSandbox(Sandbox):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Error closing AioSandbox client for {self.id}: {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
|
@property
|
||||||
def home_dir(self) -> str:
|
def home_dir(self) -> str:
|
||||||
"""Get the home directory inside the sandbox."""
|
"""Get the home directory inside the sandbox."""
|
||||||
@ -161,12 +351,12 @@ class AioSandbox(Sandbox):
|
|||||||
) -> str:
|
) -> str:
|
||||||
"""Execute a shell command in the sandbox.
|
"""Execute a shell command in the sandbox.
|
||||||
|
|
||||||
Uses a lock to serialize concurrent requests. The AIO sandbox
|
Uses a lock to serialize unscoped requests. The AIO sandbox container's
|
||||||
container maintains a single persistent shell session that
|
implicit persistent shell corrupts when hit with concurrent
|
||||||
corrupts when hit with concurrent exec_command calls (returns
|
``exec_command`` calls (returning ``ErrorObservation`` instead of real
|
||||||
``ErrorObservation`` instead of real output). If corruption is
|
output). If corruption is detected despite the lock (e.g. multiple
|
||||||
detected despite the lock (e.g. multiple processes sharing a
|
processes sharing a sandbox), the replacement session is promoted for
|
||||||
sandbox), the command is retried on a fresh session.
|
subsequent calls rather than returning to the corrupted implicit one.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
command: The command to execute.
|
command: The command to execute.
|
||||||
@ -195,34 +385,31 @@ class AioSandbox(Sandbox):
|
|||||||
return self._execute_with_env(command, env)
|
return self._execute_with_env(command, env)
|
||||||
with self._lock:
|
with self._lock:
|
||||||
try:
|
try:
|
||||||
result = self._client.shell.exec_command(command=command, no_change_timeout=self._DEFAULT_NO_CHANGE_TIMEOUT)
|
client = self._client
|
||||||
output = result.data.output if result.data else ""
|
if getattr(self, "_closed", False) or client is None:
|
||||||
exit_code = getattr(result.data, "exit_code", None) if result.data else 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:
|
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")
|
logger.warning("ErrorObservation detected in sandbox output, retrying on a fresh session")
|
||||||
# exec_command only auto-creates a session when called with
|
output, exit_code, self._recovery_session_id = self._rotate_and_retry_shell(
|
||||||
# no id, so the recovery session must be created explicitly
|
client,
|
||||||
# before we target it on retry.
|
command,
|
||||||
fresh_id = str(uuid.uuid4())
|
corrupted_session_id=self._recovery_session_id,
|
||||||
self._client.shell.create_session(id=fresh_id)
|
context="default shell",
|
||||||
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}")
|
|
||||||
|
|
||||||
if exit_code not in (0, None):
|
return self._render_shell_output(output, exit_code)
|
||||||
# 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)"
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to execute command in sandbox: {e}")
|
logger.error(f"Failed to execute command in sandbox: {e}")
|
||||||
return f"Error: {e}"
|
return f"Error: {e}"
|
||||||
|
|||||||
@ -77,6 +77,7 @@ from deerflow.runtime.serialization import serialize
|
|||||||
from deerflow.runtime.stream_bridge import StreamBridge
|
from deerflow.runtime.stream_bridge import StreamBridge
|
||||||
from deerflow.runtime.stream_modes import normalize_stream_modes, to_langgraph_stream_modes
|
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.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.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_id
|
||||||
from deerflow.tracing import inject_langfuse_metadata
|
from deerflow.tracing import inject_langfuse_metadata
|
||||||
from deerflow.utils.messages import message_to_text
|
from deerflow.utils.messages import message_to_text
|
||||||
@ -519,12 +520,15 @@ class _LargeFileToolChunkBatcher:
|
|||||||
# strips ``__``-prefixed keys in build_run_config, but embedded harness callers
|
# 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
|
# have no such filter and ``deerflow_trace_id`` carries no prefix to be caught
|
||||||
# by it anyway.
|
# by it anyway.
|
||||||
_SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = frozenset(
|
_SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = (
|
||||||
|
frozenset(
|
||||||
{
|
{
|
||||||
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY,
|
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY,
|
||||||
DEERFLOW_TRACE_METADATA_KEY,
|
DEERFLOW_TRACE_METADATA_KEY,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
| SANDBOX_SERVER_OWNED_CONTEXT_KEYS
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _build_runtime_context(
|
def _build_runtime_context(
|
||||||
@ -602,17 +606,17 @@ def _install_runtime_context(config: dict, runtime_context: dict[str, Any]) -> N
|
|||||||
if isinstance(existing_context, dict):
|
if isinstance(existing_context, dict):
|
||||||
existing_context.setdefault("thread_id", runtime_context["thread_id"])
|
existing_context.setdefault("thread_id", runtime_context["thread_id"])
|
||||||
existing_context.setdefault("run_id", runtime_context["run_id"])
|
existing_context.setdefault("run_id", runtime_context["run_id"])
|
||||||
# Assigned, not setdefault: this is a server-owned key, the same rule
|
# Keep both context views authoritative. A server-owned value is
|
||||||
# _bind_trace_id applies to the runtime context and the run metadata. A
|
# assigned from the runtime context when present and removed otherwise,
|
||||||
# deerflow_trace_id the caller put in body.config.context is an echo of
|
# so an embedded caller cannot preserve a forged lifecycle identity in
|
||||||
# a past output, not an input, and leaving it would make this one dict
|
# ``config['context']`` after it was rejected by _build_runtime_context.
|
||||||
# disagree with the response header and the logs.
|
for key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS:
|
||||||
if DEERFLOW_TRACE_METADATA_KEY in runtime_context:
|
if key in runtime_context:
|
||||||
existing_context[DEERFLOW_TRACE_METADATA_KEY] = runtime_context[DEERFLOW_TRACE_METADATA_KEY]
|
existing_context[key] = runtime_context[key]
|
||||||
|
else:
|
||||||
|
existing_context.pop(key, None)
|
||||||
if "app_config" in runtime_context:
|
if "app_config" in runtime_context:
|
||||||
existing_context["app_config"] = runtime_context["app_config"]
|
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
|
return
|
||||||
|
|
||||||
config["context"] = dict(runtime_context)
|
config["context"] = dict(runtime_context)
|
||||||
@ -1586,6 +1590,23 @@ async def run_agent(
|
|||||||
await journal.close(flush=not record.ownership_lost)
|
await journal.close(flush=not record.ownership_lost)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.warning("Failed to close journal for run %s", run_id, exc_info=True)
|
logger.warning("Failed to close journal for run %s", run_id, exc_info=True)
|
||||||
|
finally:
|
||||||
|
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:
|
finally:
|
||||||
_release_run_scoped_references(
|
_release_run_scoped_references(
|
||||||
runnable_configs,
|
runnable_configs,
|
||||||
@ -1618,6 +1639,9 @@ async def run_agent(
|
|||||||
_create_contextless_task(run_manager.cleanup(run_id))
|
_create_contextless_task(run_manager.cleanup(run_id))
|
||||||
_schedule_terminal_cycle_collection()
|
_schedule_terminal_cycle_collection()
|
||||||
|
|
||||||
|
if lease_cleanup_interrupt is not None:
|
||||||
|
raise lease_cleanup_interrupt
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Helpers
|
# Helpers
|
||||||
|
|||||||
@ -1,9 +1,10 @@
|
|||||||
### Sandbox System (`packages/harness/deerflow/sandbox/`)
|
### 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.
|
**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.
|
**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.
|
**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**:
|
**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.
|
- `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.
|
||||||
|
|||||||
698
backend/packages/harness/deerflow/sandbox/lease.py
Normal file
698
backend/packages/harness/deerflow/sandbox/lease.py
Normal file
@ -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)
|
||||||
@ -21,6 +21,11 @@ from deerflow.authz.sandbox_authz import (
|
|||||||
from deerflow.runtime.user_context import resolve_runtime_user_id
|
from deerflow.runtime.user_context import resolve_runtime_user_id
|
||||||
from deerflow.sandbox import get_sandbox_provider
|
from deerflow.sandbox import get_sandbox_provider
|
||||||
from deerflow.sandbox.exceptions import SandboxAuthorizationError, SandboxRuntimeError
|
from deerflow.sandbox.exceptions import SandboxAuthorizationError, SandboxRuntimeError
|
||||||
|
from deerflow.sandbox.lease import (
|
||||||
|
ensure_sandbox_lease_owner,
|
||||||
|
get_sandbox_lease_manager,
|
||||||
|
sandbox_lease_owner,
|
||||||
|
)
|
||||||
from deerflow.sandbox.overwrite import unwrap_sandbox
|
from deerflow.sandbox.overwrite import unwrap_sandbox
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@ -39,9 +44,9 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
Lifecycle Management:
|
Lifecycle Management:
|
||||||
- With lazy_init=True (default): Sandbox is acquired on first tool call
|
- 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)
|
- With lazy_init=False: Sandbox is acquired on first agent invocation (before_agent)
|
||||||
- Sandbox is reused across multiple turns within the same thread
|
- Concurrent lead/subagent executions hold independent process-local leases
|
||||||
- Sandbox is NOT released after each agent call to avoid wasteful recreation
|
- Only the final execution release parks a remote sandbox in its warm pool
|
||||||
- Cleanup happens at application shutdown via SandboxProvider.shutdown()
|
- Provider shutdown remains the terminal cleanup boundary
|
||||||
"""
|
"""
|
||||||
|
|
||||||
state_schema = SandboxMiddlewareState
|
state_schema = SandboxMiddlewareState
|
||||||
@ -108,20 +113,105 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
if projection is not None and not provider.supports_agent_skill_isolation:
|
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")
|
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()
|
provider = get_sandbox_provider()
|
||||||
|
if owner_id is None:
|
||||||
sandbox_id = provider.acquire(thread_id, user_id=user_id)
|
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}")
|
logger.info(f"Acquiring sandbox {sandbox_id}")
|
||||||
return 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()
|
provider = get_sandbox_provider()
|
||||||
|
if owner_id is None:
|
||||||
sandbox_id = await provider.acquire_async(thread_id, user_id=user_id)
|
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}")
|
logger.info(f"Acquiring sandbox {sandbox_id}")
|
||||||
return sandbox_id
|
return sandbox_id
|
||||||
|
|
||||||
async def _release_sandbox_async(self, sandbox_id: str) -> None:
|
@staticmethod
|
||||||
await asyncio.to_thread(get_sandbox_provider().release, sandbox_id)
|
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
|
@override
|
||||||
def before_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
def before_agent(self, state: SandboxMiddlewareState, runtime: Runtime) -> dict | None:
|
||||||
@ -130,11 +220,16 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
return super().before_agent(state, runtime)
|
return super().before_agent(state, runtime)
|
||||||
user_id = resolve_runtime_user_id(runtime)
|
user_id = resolve_runtime_user_id(runtime)
|
||||||
projection = self._prepare_agent_skill_projection(thread_id, user_id=user_id)
|
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.
|
# Preserve lazy initialization for threads that use the shared view.
|
||||||
# A policy-scoped view is acquired eagerly so an old shared-view
|
# A policy-scoped view is acquired eagerly so an old shared-view
|
||||||
# sandbox cannot survive into this run through checkpoint state.
|
# sandbox cannot survive into this run through checkpoint state.
|
||||||
if self._lazy_init and projection is None:
|
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)
|
return super().before_agent(state, runtime)
|
||||||
|
|
||||||
existing_sandbox_id = self._read_sandbox_id_from_state(state)
|
existing_sandbox_id = self._read_sandbox_id_from_state(state)
|
||||||
@ -162,7 +257,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
return None
|
return None
|
||||||
provider = get_sandbox_provider()
|
provider = get_sandbox_provider()
|
||||||
self._require_projection_support(provider, projection)
|
self._require_projection_support(provider, projection)
|
||||||
sandbox_id = self._acquire_sandbox(thread_id, user_id=user_id)
|
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:
|
if projection is not None:
|
||||||
provider.sync_agent_skills(
|
provider.sync_agent_skills(
|
||||||
sandbox_id,
|
sandbox_id,
|
||||||
@ -170,6 +272,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
projection=projection,
|
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}")
|
logger.info(f"Assigned sandbox {sandbox_id} to thread {thread_id}")
|
||||||
if existing_sandbox_id == sandbox_id:
|
if existing_sandbox_id == sandbox_id:
|
||||||
return super().before_agent(state, runtime)
|
return super().before_agent(state, runtime)
|
||||||
@ -178,6 +288,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
"sandbox": Overwrite({"sandbox_id": sandbox_id}),
|
"sandbox": Overwrite({"sandbox_id": sandbox_id}),
|
||||||
}
|
}
|
||||||
return {"sandbox": {"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)
|
return super().before_agent(state, runtime)
|
||||||
|
|
||||||
@override
|
@override
|
||||||
@ -191,6 +309,7 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
thread_id,
|
thread_id,
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
)
|
)
|
||||||
|
owner_id = ensure_sandbox_lease_owner(runtime.context)
|
||||||
|
|
||||||
if self._lazy_init and projection is None:
|
if self._lazy_init and projection is None:
|
||||||
return await super().abefore_agent(state, runtime)
|
return await super().abefore_agent(state, runtime)
|
||||||
@ -213,7 +332,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
return None
|
return None
|
||||||
provider = get_sandbox_provider()
|
provider = get_sandbox_provider()
|
||||||
self._require_projection_support(provider, projection)
|
self._require_projection_support(provider, projection)
|
||||||
sandbox_id = await self._acquire_sandbox_async(thread_id, user_id=user_id)
|
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:
|
if projection is not None:
|
||||||
await provider.sync_agent_skills_async(
|
await provider.sync_agent_skills_async(
|
||||||
sandbox_id,
|
sandbox_id,
|
||||||
@ -221,6 +347,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
projection=projection,
|
projection=projection,
|
||||||
)
|
)
|
||||||
|
except BaseException:
|
||||||
|
await self._release_sandbox_async(
|
||||||
|
sandbox_id,
|
||||||
|
owner_id=owner_id,
|
||||||
|
)
|
||||||
|
if runtime.context is not None:
|
||||||
|
runtime.context.pop("sandbox_id", None)
|
||||||
|
raise
|
||||||
logger.info(f"Assigned sandbox {sandbox_id} to thread {thread_id}")
|
logger.info(f"Assigned sandbox {sandbox_id} to thread {thread_id}")
|
||||||
if existing_sandbox_id == sandbox_id:
|
if existing_sandbox_id == sandbox_id:
|
||||||
return await super().abefore_agent(state, runtime)
|
return await super().abefore_agent(state, runtime)
|
||||||
@ -229,6 +363,14 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
"sandbox": Overwrite({"sandbox_id": sandbox_id}),
|
"sandbox": Overwrite({"sandbox_id": sandbox_id}),
|
||||||
}
|
}
|
||||||
return {"sandbox": {"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)
|
return await super().abefore_agent(state, runtime)
|
||||||
|
|
||||||
@override
|
@override
|
||||||
@ -242,13 +384,23 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
|
logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
|
||||||
return None
|
return None
|
||||||
logger.info(f"Releasing sandbox {sandbox_id}")
|
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
|
return None
|
||||||
|
|
||||||
if (runtime.context or {}).get("sandbox_id") is not None:
|
if (runtime.context or {}).get("sandbox_id") is not None:
|
||||||
sandbox_id = runtime.context.get("sandbox_id")
|
sandbox_id = runtime.context.get("sandbox_id")
|
||||||
logger.info(f"Releasing sandbox {sandbox_id} from context")
|
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
|
return None
|
||||||
|
|
||||||
# No sandbox to release
|
# No sandbox to release
|
||||||
@ -265,13 +417,19 @@ class SandboxMiddleware(AgentMiddleware[SandboxMiddlewareState]):
|
|||||||
logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
|
logger.info(f"Not releasing fork-restored sandbox {sandbox_id}")
|
||||||
return None
|
return None
|
||||||
logger.info(f"Releasing sandbox {sandbox_id}")
|
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
|
return None
|
||||||
|
|
||||||
if (runtime.context or {}).get("sandbox_id") is not None:
|
if (runtime.context or {}).get("sandbox_id") is not None:
|
||||||
sandbox_id = runtime.context.get("sandbox_id")
|
sandbox_id = runtime.context.get("sandbox_id")
|
||||||
logger.info(f"Releasing sandbox {sandbox_id} from context")
|
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
|
return None
|
||||||
|
|
||||||
# No sandbox to release
|
# No sandbox to release
|
||||||
|
|||||||
@ -107,6 +107,27 @@ class Sandbox(ABC):
|
|||||||
"""
|
"""
|
||||||
pass
|
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
|
@abstractmethod
|
||||||
def read_file(
|
def read_file(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@ -178,6 +178,9 @@ def reset_sandbox_provider() -> None:
|
|||||||
provider = _default_sandbox_provider
|
provider = _default_sandbox_provider
|
||||||
_default_sandbox_provider = None
|
_default_sandbox_provider = None
|
||||||
if provider is not None:
|
if provider is not None:
|
||||||
|
from deerflow.sandbox.lease import discard_sandbox_lease_manager
|
||||||
|
|
||||||
|
discard_sandbox_lease_manager(provider)
|
||||||
provider.reset()
|
provider.reset()
|
||||||
|
|
||||||
|
|
||||||
@ -195,6 +198,9 @@ def shutdown_sandbox_provider() -> None:
|
|||||||
provider = _default_sandbox_provider
|
provider = _default_sandbox_provider
|
||||||
_default_sandbox_provider = None
|
_default_sandbox_provider = None
|
||||||
if provider is not None and hasattr(provider, "shutdown"):
|
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()
|
provider.shutdown()
|
||||||
|
|
||||||
|
|
||||||
@ -211,4 +217,9 @@ def set_sandbox_provider(provider: SandboxProvider) -> None:
|
|||||||
"""
|
"""
|
||||||
global _default_sandbox_provider
|
global _default_sandbox_provider
|
||||||
with _provider_lock:
|
with _provider_lock:
|
||||||
|
previous = _default_sandbox_provider
|
||||||
_default_sandbox_provider = 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)
|
||||||
|
|||||||
@ -32,10 +32,16 @@ from deerflow.sandbox.exceptions import (
|
|||||||
SandboxRuntimeError,
|
SandboxRuntimeError,
|
||||||
)
|
)
|
||||||
from deerflow.sandbox.file_operation_lock import get_file_operation_lock
|
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.overwrite import unwrap_sandbox
|
||||||
from deerflow.sandbox.path_patterns import build_output_mask_pattern, replace_output_path_matches
|
from deerflow.sandbox.path_patterns import build_output_mask_pattern, replace_output_path_matches
|
||||||
from deerflow.sandbox.sandbox import Sandbox
|
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.search import GrepMatch
|
||||||
from deerflow.sandbox.security import LOCAL_HOST_BASH_DISABLED_MESSAGE, is_host_bash_allowed
|
from deerflow.sandbox.security import LOCAL_HOST_BASH_DISABLED_MESSAGE, is_host_bash_allowed
|
||||||
from deerflow.tools.types import Runtime
|
from deerflow.tools.types import Runtime
|
||||||
@ -1404,6 +1410,44 @@ def sandbox_from_runtime(runtime: Runtime | None = None) -> Sandbox:
|
|||||||
return 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
|
@contextmanager
|
||||||
def sandbox_authorization_scope(runtime: Runtime) -> Iterator[None]:
|
def sandbox_authorization_scope(runtime: Runtime) -> Iterator[None]:
|
||||||
"""Authorize once for one complete synchronous sandbox tool invocation."""
|
"""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)
|
_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:
|
def ensure_sandbox_initialized(runtime: Runtime | None = None) -> Sandbox:
|
||||||
"""Ensure sandbox is initialized, acquiring lazily if needed.
|
"""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(),
|
app_config=safe_app_config(),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check if sandbox already exists in state
|
# Check if sandbox already exists in state. A fork-restored execution keeps
|
||||||
# Discarding fork_restored is safe: after_agent short-circuits on the
|
# the wrapper so after_agent cannot park the parent's sandbox, but it still
|
||||||
# still-wrapped state before the context-based release branch, so this
|
# binds a non-releasing holder: parent cleanup cannot close the client under
|
||||||
# reuse path never releases the parent sandbox.
|
# the child, and the child's outer fence can clean its command scope.
|
||||||
sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox"))
|
sandbox_state, fork_restored = unwrap_sandbox(runtime.state.get("sandbox"))
|
||||||
if sandbox_state is not None:
|
if sandbox_state is not None:
|
||||||
sandbox_id = sandbox_state.get("sandbox_id")
|
sandbox_id = sandbox_state.get("sandbox_id")
|
||||||
if sandbox_id is not None:
|
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 sandbox is not None:
|
||||||
if runtime.context 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
|
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
|
# Sandbox was released, fall through to acquire new one
|
||||||
|
|
||||||
# Lazy acquisition: get thread_id and acquire sandbox
|
# Lazy acquisition: get thread_id and acquire sandbox
|
||||||
thread_id = runtime.context.get("thread_id") if runtime.context else None
|
thread_id = _resolve_runtime_thread_id(runtime)
|
||||||
if thread_id is None:
|
|
||||||
thread_id = runtime.config.get("configurable", {}).get("thread_id") if runtime.config else None
|
|
||||||
if thread_id is None:
|
if thread_id is None:
|
||||||
raise SandboxRuntimeError("Thread ID not available in runtime context")
|
raise SandboxRuntimeError("Thread ID not available in runtime context")
|
||||||
|
|
||||||
provider = get_sandbox_provider()
|
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
|
# Update runtime state - this persists across tool calls
|
||||||
runtime.state["sandbox"] = {"sandbox_id": sandbox_id}
|
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
|
# Retrieve and return the sandbox
|
||||||
sandbox = provider.get(sandbox_id)
|
sandbox = provider.get(sandbox_id)
|
||||||
if sandbox is None:
|
if sandbox is None:
|
||||||
|
_rollback_failed_sandbox_lookup(provider, sandbox_id, owner_id)
|
||||||
raise SandboxNotFoundError("Sandbox not found after acquisition", sandbox_id=sandbox_id)
|
raise SandboxNotFoundError("Sandbox not found after acquisition", sandbox_id=sandbox_id)
|
||||||
|
|
||||||
if runtime.context is not None:
|
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(),
|
app_config=await safe_app_config_async(),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Same discard as the sync path above: the reuse path never releases,
|
# Same borrowed-holder rule as the sync path above: keep the fork wrapper
|
||||||
# because after_agent short-circuits on the still-wrapped state first.
|
# while counting the child as an active client user.
|
||||||
sandbox_state, _ = unwrap_sandbox(runtime.state.get("sandbox"))
|
sandbox_state, fork_restored = unwrap_sandbox(runtime.state.get("sandbox"))
|
||||||
if sandbox_state is not None:
|
if sandbox_state is not None:
|
||||||
sandbox_id = sandbox_state.get("sandbox_id")
|
sandbox_id = sandbox_state.get("sandbox_id")
|
||||||
if sandbox_id is not None:
|
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 sandbox is not None:
|
||||||
if runtime.context is not None:
|
if runtime.context is not None:
|
||||||
runtime.context["sandbox_id"] = sandbox_id
|
runtime.context["sandbox_id"] = sandbox_id
|
||||||
return sandbox
|
return sandbox
|
||||||
|
|
||||||
thread_id = runtime.context.get("thread_id") if runtime.context else None
|
thread_id = _resolve_runtime_thread_id(runtime)
|
||||||
if thread_id is None:
|
|
||||||
thread_id = runtime.config.get("configurable", {}).get("thread_id") if runtime.config else None
|
|
||||||
if thread_id is None:
|
if thread_id is None:
|
||||||
raise SandboxRuntimeError("Thread ID not available in runtime context")
|
raise SandboxRuntimeError("Thread ID not available in runtime context")
|
||||||
|
|
||||||
provider = get_sandbox_provider()
|
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}
|
runtime.state["sandbox"] = {"sandbox_id": sandbox_id}
|
||||||
|
|
||||||
sandbox = provider.get(sandbox_id)
|
sandbox = provider.get(sandbox_id)
|
||||||
if sandbox is None:
|
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)
|
raise SandboxNotFoundError("Sandbox not found after acquisition", sandbox_id=sandbox_id)
|
||||||
|
|
||||||
if runtime.context is not None:
|
if runtime.context is not None:
|
||||||
@ -1577,13 +1671,37 @@ async def _run_sync_tool_after_async_sandbox_init(
|
|||||||
if func is None:
|
if func is None:
|
||||||
return "Error: Tool implementation not available"
|
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:
|
except SandboxError as e:
|
||||||
return f"Error: {e}"
|
return f"Error: {e}"
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return f"Error: Unexpected error initializing sandbox: {_sanitize_error(e, runtime)}"
|
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:
|
def ensure_thread_directories_exist(runtime: Runtime | None) -> None:
|
||||||
"""Ensure thread data directories (workspace, uploads, outputs) exist.
|
"""Ensure thread data directories (workspace, uploads, outputs) exist.
|
||||||
|
|
||||||
@ -1936,7 +2054,13 @@ def bash_tool(runtime: Runtime, command: str, description: str = "") -> str:
|
|||||||
except Exception:
|
except Exception:
|
||||||
max_chars = 20000
|
max_chars = 20000
|
||||||
command_timeout = None
|
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(
|
return _truncate_bash_output(
|
||||||
mask_secret_values(mask_local_paths_in_output(output, thread_data), injected_env),
|
mask_secret_values(mask_local_paths_in_output(output, thread_data), injected_env),
|
||||||
max_chars,
|
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
|
max_chars = sandbox_cfg.bash_output_max_chars if sandbox_cfg else 20000
|
||||||
except Exception:
|
except Exception:
|
||||||
max_chars = 20000
|
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:
|
except SandboxError as e:
|
||||||
return f"Error: {e}"
|
return f"Error: {e}"
|
||||||
except PermissionError as e:
|
except PermissionError as e:
|
||||||
|
|||||||
@ -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.
|
**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 `<current_date>` 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.
|
**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 `<current_date>` 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.
|
**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`.
|
**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.
|
**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.
|
**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.
|
||||||
|
|||||||
@ -60,6 +60,11 @@ if TYPE_CHECKING:
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS = 3.0
|
_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")
|
_previous_shutdown_isolated_subagent_loop = globals().get("_shutdown_isolated_subagent_loop")
|
||||||
@ -1308,6 +1313,8 @@ class SubagentExecutor:
|
|||||||
status=SubagentStatus.RUNNING,
|
status=SubagentStatus.RUNNING,
|
||||||
started_at=datetime.now(),
|
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_extension_api import ExtensionData, TaskInfo
|
||||||
|
|
||||||
from deerflow.extensions import get_loaded_extensions
|
from deerflow.extensions import get_loaded_extensions
|
||||||
@ -1465,6 +1472,9 @@ class SubagentExecutor:
|
|||||||
context["authz_attributes"] = dict(self.authz_attributes)
|
context["authz_attributes"] = dict(self.authz_attributes)
|
||||||
context[DEERFLOW_TRACE_METADATA_KEY] = self.deerflow_trace_id
|
context[DEERFLOW_TRACE_METADATA_KEY] = self.deerflow_trace_id
|
||||||
context["is_subagent"] = True
|
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
|
context["agent_id"] = self.config.name
|
||||||
if self.loop_detection_recorder is not None:
|
if self.loop_detection_recorder is not None:
|
||||||
context[LOOP_DETECTION_RECORDER_CONTEXT_KEY] = self.loop_detection_recorder
|
context[LOOP_DETECTION_RECORDER_CONTEXT_KEY] = self.loop_detection_recorder
|
||||||
@ -1623,6 +1633,20 @@ class SubagentExecutor:
|
|||||||
)
|
)
|
||||||
|
|
||||||
finally:
|
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:
|
if task_info is not None and task_store is not None:
|
||||||
try:
|
try:
|
||||||
await notify_task_stop(
|
await notify_task_stop(
|
||||||
|
|||||||
@ -16,6 +16,7 @@ network leg is httpx-async and not the subject here.
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import threading
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
from unittest.mock import AsyncMock
|
from unittest.mock import AsyncMock
|
||||||
|
|
||||||
@ -24,6 +25,40 @@ import pytest
|
|||||||
pytestmark = pytest.mark.asyncio
|
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:
|
async def test_receive_file_persist_does_not_block_event_loop(tmp_path, monkeypatch) -> None:
|
||||||
from app.channels.dingtalk import DingTalkChannel
|
from app.channels.dingtalk import DingTalkChannel
|
||||||
from app.channels.message_bus import MessageBus
|
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 "/uploads/a.pdf" in out.text
|
||||||
assert out.files == []
|
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)
|
||||||
|
|||||||
@ -38,11 +38,15 @@ class _RemoteSandbox:
|
|||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.updates: list[tuple[str, bytes]] = []
|
self.updates: list[tuple[str, bytes]] = []
|
||||||
self.update_thread_id: int | None = None
|
self.update_thread_id: int | None = None
|
||||||
|
self.released_scopes: list[str] = []
|
||||||
|
|
||||||
def update_file(self, path: str, content: bytes) -> None:
|
def update_file(self, path: str, content: bytes) -> None:
|
||||||
self.update_thread_id = threading.get_ident()
|
self.update_thread_id = threading.get_ident()
|
||||||
self.updates.append((path, content))
|
self.updates.append((path, content))
|
||||||
|
|
||||||
|
def release_command_scope(self, scope_id: str) -> None:
|
||||||
|
self.released_scopes.append(scope_id)
|
||||||
|
|
||||||
|
|
||||||
class _RemoteProvider:
|
class _RemoteProvider:
|
||||||
uses_thread_data_mounts = False
|
uses_thread_data_mounts = False
|
||||||
@ -75,6 +79,31 @@ class _MountedProvider:
|
|||||||
raise AssertionError("mounted uploads must not look up a sandbox")
|
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:
|
async def test_receive_file_remote_sandbox_does_not_block_event_loop(tmp_path, monkeypatch) -> None:
|
||||||
from deerflow.config.paths import Paths
|
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"
|
assert result == "/mnt/user-data/uploads/report.pdf"
|
||||||
uploaded = tmp_path / "users" / "ou-user" / "threads" / "thread-1" / "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"
|
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)
|
||||||
|
|||||||
@ -7,6 +7,29 @@ from unittest.mock import MagicMock, patch
|
|||||||
import pytest
|
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():
|
def test_local_sandbox_client_bypasses_environment_proxy():
|
||||||
"""Local sandbox API calls must not inherit HTTP_PROXY (#3441)."""
|
"""Local sandbox API calls must not inherit HTTP_PROXY (#3441)."""
|
||||||
from deerflow.community.aio_sandbox.aio_sandbox import AioSandbox
|
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
|
# ...and the retry targets exactly that created session, never an
|
||||||
# uncreated/fabricated id (which would 404).
|
# uncreated/fabricated id (which would 404).
|
||||||
assert exec_calls[1].get("id") == created_ids[0]
|
assert exec_calls[1].get("id") == created_ids[0]
|
||||||
# ...and that one-shot recovery session is released afterwards so a
|
# The recovered session is promoted: future commands must not return
|
||||||
# sandbox that keeps hitting corruption doesn't accumulate sessions.
|
# 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]]
|
assert cleaned_ids == [created_ids[0]]
|
||||||
|
|
||||||
def test_cleanup_failure_does_not_mask_successful_retry(self, sandbox):
|
def test_cleanup_failure_does_not_mask_successful_retry(self, sandbox):
|
||||||
@ -201,6 +229,33 @@ class TestErrorObservationRetry:
|
|||||||
# into an "Error: ..." result.
|
# into an "Error: ..." result.
|
||||||
assert sandbox.execute_command("test") == "recovered"
|
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):
|
def test_no_retry_on_clean_output(self, sandbox):
|
||||||
"""Normal output should not trigger a retry."""
|
"""Normal output should not trigger a retry."""
|
||||||
call_count = 0
|
call_count = 0
|
||||||
@ -217,6 +272,222 @@ class TestErrorObservationRetry:
|
|||||||
assert call_count == 1
|
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:
|
class TestBashExecUnsupportedFailFast:
|
||||||
"""Regression tests for #3921: sandbox images older than all-in-one-sandbox
|
"""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
|
1.9.x have no ``/v1/bash/exec`` route, so every env-bearing command (skills
|
||||||
|
|||||||
@ -16,6 +16,7 @@ from starlette.responses import FileResponse
|
|||||||
import app.gateway.routers.artifacts as artifacts_router
|
import app.gateway.routers.artifacts as artifacts_router
|
||||||
from app.gateway.internal_auth import INTERNAL_OWNER_USER_ID_HEADER_NAME, INTERNAL_SYSTEM_ROLE
|
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.config.paths import make_safe_user_id
|
||||||
|
from deerflow.sandbox.lease import get_sandbox_lease_manager
|
||||||
|
|
||||||
ACTIVE_ARTIFACT_CASES = [
|
ACTIVE_ARTIFACT_CASES = [
|
||||||
("poc.html", "<html><body><script>alert('xss')</script></body></html>"),
|
("poc.html", "<html><body><script>alert('xss')</script></body></html>"),
|
||||||
@ -66,6 +67,7 @@ class _RemoteSandbox:
|
|||||||
def __init__(self, *, fail_next_update: bool = False) -> None:
|
def __init__(self, *, fail_next_update: bool = False) -> None:
|
||||||
self.updates: list[tuple[str, bytes]] = []
|
self.updates: list[tuple[str, bytes]] = []
|
||||||
self.fail_next_update = fail_next_update
|
self.fail_next_update = fail_next_update
|
||||||
|
self.released_scopes: list[str] = []
|
||||||
|
|
||||||
def update_file(self, path: str, content: bytes) -> None:
|
def update_file(self, path: str, content: bytes) -> None:
|
||||||
if self.fail_next_update:
|
if self.fail_next_update:
|
||||||
@ -73,6 +75,9 @@ class _RemoteSandbox:
|
|||||||
raise RuntimeError("sandbox sync failed")
|
raise RuntimeError("sandbox sync failed")
|
||||||
self.updates.append((path, content))
|
self.updates.append((path, content))
|
||||||
|
|
||||||
|
def release_command_scope(self, scope_id: str) -> None:
|
||||||
|
self.released_scopes.append(scope_id)
|
||||||
|
|
||||||
|
|
||||||
class _RemoteSandboxProvider:
|
class _RemoteSandboxProvider:
|
||||||
uses_thread_data_mounts = False
|
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"
|
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:
|
def test_update_artifact_releases_sandbox_when_initial_sync_fails(tmp_path, monkeypatch) -> None:
|
||||||
artifact_path = tmp_path / "note.txt"
|
artifact_path = tmp_path / "note.txt"
|
||||||
artifact_path.write_text("before", encoding="utf-8")
|
artifact_path.write_text("before", encoding="utf-8")
|
||||||
|
|||||||
@ -27,6 +27,8 @@ from deerflow.config.authorization_config import AuthorizationConfig, Authorizat
|
|||||||
from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig
|
from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig
|
||||||
from deerflow.config.paths import Paths
|
from deerflow.config.paths import Paths
|
||||||
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
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.skills.types import SkillCategory
|
||||||
from deerflow.tools.mcp_metadata import tag_mcp_tool
|
from deerflow.tools.mcp_metadata import tag_mcp_tool
|
||||||
from deerflow.uploads.manager import PathTraversalError
|
from deerflow.uploads.manager import PathTraversalError
|
||||||
@ -3652,6 +3654,85 @@ class TestStreamHardening:
|
|||||||
with pytest.raises(RuntimeError, match="model quota exceeded"):
|
with pytest.raises(RuntimeError, match="model quota exceeded"):
|
||||||
list(client.stream("hi", thread_id="t-err"))
|
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):
|
def test_messages_without_id(self, client):
|
||||||
"""Messages without id attribute are emitted without crashing."""
|
"""Messages without id attribute are emitted without crashing."""
|
||||||
ai = AIMessage(content="no id here")
|
ai = AIMessage(content="no id here")
|
||||||
|
|||||||
@ -2,17 +2,30 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import threading
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from langchain.tools import ToolRuntime
|
from langchain.tools import ToolRuntime
|
||||||
from langgraph.types import Overwrite
|
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 import Sandbox
|
||||||
from deerflow.sandbox.sandbox_provider import SandboxProvider, reset_sandbox_provider, set_sandbox_provider
|
from deerflow.sandbox.sandbox_provider import SandboxProvider, reset_sandbox_provider, set_sandbox_provider
|
||||||
from deerflow.sandbox.search import GrepMatch
|
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):
|
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:
|
def execute_command(self, command: str, env: dict | None = None, timeout: float | None = None) -> str:
|
||||||
del env, timeout
|
del env, timeout
|
||||||
return "OK"
|
return "OK"
|
||||||
@ -47,10 +60,14 @@ class _StubSandbox(Sandbox):
|
|||||||
def update_file(self, path: str, content: bytes) -> None:
|
def update_file(self, path: str, content: bytes) -> None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
def release_command_scope(self, scope_id: str) -> None:
|
||||||
|
self.released_scopes.append(scope_id)
|
||||||
|
|
||||||
|
|
||||||
class _RecordingProvider(SandboxProvider):
|
class _RecordingProvider(SandboxProvider):
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.sandbox = _StubSandbox("stub")
|
self.sandbox = _StubSandbox("stub")
|
||||||
|
self.released: list[str] = []
|
||||||
|
|
||||||
def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> 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")
|
raise AssertionError("state already carries a sandbox; acquire must not run")
|
||||||
@ -64,7 +81,7 @@ class _RecordingProvider(SandboxProvider):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def release(self, sandbox_id: str) -> None:
|
def release(self, sandbox_id: str) -> None:
|
||||||
return None
|
self.released.append(sandbox_id)
|
||||||
|
|
||||||
|
|
||||||
class _FallthroughProvider(SandboxProvider):
|
class _FallthroughProvider(SandboxProvider):
|
||||||
@ -73,6 +90,7 @@ class _FallthroughProvider(SandboxProvider):
|
|||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.sandbox = _StubSandbox("fresh")
|
self.sandbox = _StubSandbox("fresh")
|
||||||
self.acquired: list[str | None] = []
|
self.acquired: list[str | None] = []
|
||||||
|
self.released: list[str] = []
|
||||||
|
|
||||||
def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
|
def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
|
||||||
self.acquired.append(thread_id)
|
self.acquired.append(thread_id)
|
||||||
@ -88,8 +106,27 @@ class _FallthroughProvider(SandboxProvider):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def release(self, sandbox_id: str) -> 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
|
return None
|
||||||
|
|
||||||
|
def release(self, sandbox_id: str) -> None:
|
||||||
|
self.released.append(sandbox_id)
|
||||||
|
|
||||||
|
|
||||||
def _make_runtime(state: dict) -> ToolRuntime:
|
def _make_runtime(state: dict) -> ToolRuntime:
|
||||||
return 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:
|
def test_ensure_sandbox_initialized_unwraps_overwrite_state() -> None:
|
||||||
"""Fork-restored state must not crash on the Overwrite wrapper."""
|
"""Fork-restored state must not crash on the Overwrite wrapper."""
|
||||||
provider = _RecordingProvider()
|
provider = _RecordingProvider()
|
||||||
@ -134,6 +218,63 @@ async def test_ensure_sandbox_initialized_async_unwraps_overwrite_state() -> Non
|
|||||||
assert runtime.context["sandbox_id"] == "parent-sandbox"
|
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:
|
def test_ensure_sandbox_initialized_plain_state_unchanged() -> None:
|
||||||
provider = _RecordingProvider()
|
provider = _RecordingProvider()
|
||||||
set_sandbox_provider(provider)
|
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"
|
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
|
@pytest.mark.anyio
|
||||||
async def test_ensure_sandbox_initialized_async_plain_state_unchanged() -> None:
|
async def test_ensure_sandbox_initialized_async_plain_state_unchanged() -> None:
|
||||||
provider = _RecordingProvider()
|
provider = _RecordingProvider()
|
||||||
@ -180,6 +374,94 @@ async def test_ensure_sandbox_initialized_async_plain_state_unchanged() -> None:
|
|||||||
assert runtime.context["sandbox_id"] == "parent-sandbox"
|
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
|
@pytest.mark.anyio
|
||||||
async def test_ensure_sandbox_initialized_async_acquires_fresh_when_parent_missing() -> None:
|
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
|
"""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 sandbox is provider.sandbox
|
||||||
assert runtime.state["sandbox"] == {"sandbox_id": "fresh-sandbox"}
|
assert runtime.state["sandbox"] == {"sandbox_id": "fresh-sandbox"}
|
||||||
assert runtime.context["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()
|
||||||
|
|||||||
@ -3029,6 +3029,24 @@ class TestInjectAuthenticatedUserContextAuthz:
|
|||||||
assert "authz_attributes" not in config["context"]
|
assert "authz_attributes" not in config["context"]
|
||||||
assert config["context"]["is_internal"] is True
|
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):
|
def test_session_body_context_cannot_inject_channel_user_id(self):
|
||||||
request = _make_request_with_auth_source("session")
|
request = _make_request_with_auth_source("session")
|
||||||
config = _assemble_authz_run_config(
|
config = _assemble_authz_run_config(
|
||||||
|
|||||||
@ -149,6 +149,8 @@ def test_aio_sandbox_no_env_leaves_command_unchanged() -> None:
|
|||||||
sbx._lock = __import__("threading").Lock()
|
sbx._lock = __import__("threading").Lock()
|
||||||
sbx._client = SimpleNamespace(shell=_FakeShell())
|
sbx._client = SimpleNamespace(shell=_FakeShell())
|
||||||
sbx._DEFAULT_NO_CHANGE_TIMEOUT = 30
|
sbx._DEFAULT_NO_CHANGE_TIMEOUT = 30
|
||||||
|
sbx._recovery_session_id = None
|
||||||
|
sbx._default_shell_corrupted = False
|
||||||
|
|
||||||
sbx.execute_command("echo hello")
|
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
|
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
|
# ChannelManager._apply_channel_policy — the unified per-channel run-policy
|
||||||
# hook. The github channel registers (is_interactive=False,
|
# hook. The github channel registers (is_interactive=False,
|
||||||
|
|||||||
@ -44,6 +44,13 @@ from deerflow.runtime.runs.worker import (
|
|||||||
_try_extract_from_message,
|
_try_extract_from_message,
|
||||||
run_agent,
|
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:
|
class FakeCheckpointer:
|
||||||
@ -53,6 +60,115 @@ class FakeCheckpointer:
|
|||||||
self.aput_writes = AsyncMock()
|
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
|
@pytest.mark.anyio
|
||||||
async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled():
|
async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled():
|
||||||
class CleanupTrackingRunManager(RunManager):
|
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"})
|
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
|
@pytest.mark.anyio
|
||||||
async def test_run_agent_batches_incremental_file_args_and_keeps_complete_values():
|
async def test_run_agent_batches_incremental_file_args_and_keeps_complete_values():
|
||||||
run_manager = RunManager()
|
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
|
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():
|
def test_build_runtime_context_ignores_non_dict_caller_context():
|
||||||
ctx = _build_runtime_context("thread-1", "run-1", "not-a-dict")
|
ctx = _build_runtime_context("thread-1", "run-1", "not-a-dict")
|
||||||
assert ctx == {"thread_id": "thread-1", "run_id": "run-1"}
|
assert ctx == {"thread_id": "thread-1", "run_id": "run-1"}
|
||||||
|
|||||||
731
backend/tests/test_sandbox_leases.py
Normal file
731
backend/tests/test_sandbox_leases.py
Normal file
@ -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"
|
||||||
@ -7,17 +7,23 @@ import pytest
|
|||||||
from langchain.agents.middleware import AgentMiddleware
|
from langchain.agents.middleware import AgentMiddleware
|
||||||
from langchain.tools import ToolRuntime
|
from langchain.tools import ToolRuntime
|
||||||
from langchain_core.messages import ToolMessage
|
from langchain_core.messages import ToolMessage
|
||||||
|
from langgraph.graph import END
|
||||||
from langgraph.prebuilt.tool_node import ToolCallRequest
|
from langgraph.prebuilt.tool_node import ToolCallRequest
|
||||||
from langgraph.runtime import Runtime
|
from langgraph.runtime import Runtime
|
||||||
from langgraph.types import Command, Overwrite
|
from langgraph.types import Command, Overwrite
|
||||||
|
|
||||||
from deerflow.agents.thread_state import ThreadState
|
from deerflow.agents.thread_state import ThreadState
|
||||||
from deerflow.sandbox.exceptions import SandboxAuthorizationError, SandboxRuntimeError
|
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.middleware import SandboxMiddleware, SandboxMiddlewareState
|
||||||
from deerflow.sandbox.sandbox import Sandbox
|
from deerflow.sandbox.sandbox import Sandbox
|
||||||
from deerflow.sandbox.sandbox_provider import SandboxProvider, reset_sandbox_provider, set_sandbox_provider
|
from deerflow.sandbox.sandbox_provider import SandboxProvider, reset_sandbox_provider, set_sandbox_provider
|
||||||
from deerflow.sandbox.search import GrepMatch
|
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):
|
class _SyncProvider(SandboxProvider):
|
||||||
@ -309,19 +315,68 @@ async def test_abefore_agent_delegates_to_super_when_not_acquiring(
|
|||||||
runtime: Runtime,
|
runtime: Runtime,
|
||||||
) -> None:
|
) -> None:
|
||||||
calls: list[tuple[dict, Runtime]] = []
|
calls: list[tuple[dict, Runtime]] = []
|
||||||
|
provider = _AsyncOnlyProvider()
|
||||||
|
|
||||||
async def fake_super_abefore_agent(self, state_arg, runtime_arg):
|
async def fake_super_abefore_agent(self, state_arg, runtime_arg):
|
||||||
calls.append((state_arg, runtime_arg))
|
calls.append((state_arg, runtime_arg))
|
||||||
return {"delegated": True}
|
return {"delegated": True}
|
||||||
|
|
||||||
monkeypatch.setattr(AgentMiddleware, "abefore_agent", fake_super_abefore_agent)
|
monkeypatch.setattr(AgentMiddleware, "abefore_agent", fake_super_abefore_agent)
|
||||||
|
set_sandbox_provider(provider)
|
||||||
|
try:
|
||||||
result = await middleware.abefore_agent(state, runtime)
|
result = await middleware.abefore_agent(state, runtime)
|
||||||
|
finally:
|
||||||
|
reset_sandbox_provider()
|
||||||
|
|
||||||
assert result == {"delegated": True}
|
assert result == {"delegated": True}
|
||||||
assert calls == [(state, runtime)]
|
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
|
@pytest.mark.anyio
|
||||||
async def test_default_lazy_tool_acquisition_uses_async_provider() -> None:
|
async def test_default_lazy_tool_acquisition_uses_async_provider() -> None:
|
||||||
provider = _AsyncOnlyProvider()
|
provider = _AsyncOnlyProvider()
|
||||||
@ -570,6 +625,38 @@ def test_wrap_tool_call_does_not_override_non_dict_update() -> None:
|
|||||||
assert result is cmd
|
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
|
@pytest.mark.anyio
|
||||||
async def test_awrap_tool_call_emits_command_when_lazy_init_happens() -> None:
|
async def test_awrap_tool_call_emits_command_when_lazy_init_happens() -> None:
|
||||||
middleware = SandboxMiddleware()
|
middleware = SandboxMiddleware()
|
||||||
@ -605,6 +692,93 @@ async def test_awrap_tool_call_passthrough_when_sandbox_already_in_state() -> No
|
|||||||
assert result is original
|
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:
|
def test_wrap_tool_call_preserves_existing_command_fields_when_merging() -> None:
|
||||||
"""Regression: when merging sandbox_update into an existing Command,
|
"""Regression: when merging sandbox_update into an existing Command,
|
||||||
all other Command fields (e.g. graph, goto, resume) must be preserved.
|
all other Command fields (e.g. graph, goto, resume) must be preserved.
|
||||||
|
|||||||
@ -29,6 +29,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||||||
import pytest
|
import pytest
|
||||||
from packaging.version import Version
|
from packaging.version import Version
|
||||||
|
|
||||||
|
from deerflow.sandbox.lease import SandboxLeaseManager
|
||||||
from deerflow.skills.types import Skill
|
from deerflow.skills.types import Skill
|
||||||
from deerflow.subagents.capacity import SubagentCapacityRejected
|
from deerflow.subagents.capacity import SubagentCapacityRejected
|
||||||
from deerflow.trace_context import request_trace_context
|
from deerflow.trace_context import request_trace_context
|
||||||
@ -1479,6 +1480,125 @@ class TestAsyncExecutionPath:
|
|||||||
assert "Agent error" in result.error
|
assert "Agent error" in result.error
|
||||||
assert result.completed_at is not None
|
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
|
@pytest.mark.anyio
|
||||||
async def test_aexecute_recursion_error_with_partial_surfaces_completed_turn_capped(self, classes, base_config, mock_agent, msg):
|
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`` ==
|
"""#3875 Phase 2: ``GraphRecursionError`` (``recursion_limit`` ==
|
||||||
@ -3720,6 +3840,10 @@ class TestSubagentGuardrailAttribution:
|
|||||||
assert context.get("oauth_id") == "subj-123"
|
assert context.get("oauth_id") == "subj-123"
|
||||||
assert context.get("run_id") == "run-42"
|
assert context.get("run_id") == "run-42"
|
||||||
assert context.get("is_subagent") is True
|
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
|
@pytest.mark.anyio
|
||||||
async def test_aexecute_propagates_narrow_loop_detection_recorder(
|
async def test_aexecute_propagates_narrow_loop_detection_recorder(
|
||||||
|
|||||||
@ -1,6 +1,7 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import stat
|
import stat
|
||||||
|
import threading
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
@ -13,6 +14,7 @@ from fastapi.testclient import TestClient
|
|||||||
|
|
||||||
from app.gateway.deps import get_config
|
from app.gateway.deps import get_config
|
||||||
from app.gateway.routers import uploads
|
from app.gateway.routers import uploads
|
||||||
|
from deerflow.sandbox.lease import get_sandbox_lease_manager
|
||||||
|
|
||||||
|
|
||||||
class ChunkedUpload:
|
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")
|
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):
|
def test_upload_files_makes_non_local_files_sandbox_writable(tmp_path):
|
||||||
thread_uploads_dir = tmp_path / "uploads"
|
thread_uploads_dir = tmp_path / "uploads"
|
||||||
thread_uploads_dir.mkdir(parents=True)
|
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
|
assert exc_info.value.status_code == 413
|
||||||
provider.acquire.assert_not_called()
|
provider.acquire.assert_not_called()
|
||||||
provider.acquire_async.assert_awaited_once_with("thread-aio", user_id="owner-upload")
|
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()
|
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
|
assert exc_info.value.status_code == 500
|
||||||
provider.acquire.assert_not_called()
|
provider.acquire.assert_not_called()
|
||||||
provider.acquire_async.assert_awaited_once_with("thread-aio", user_id="owner-upload")
|
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()
|
sandbox.update_file.assert_not_called()
|
||||||
assert not (thread_uploads_dir / "report.pdf").exists()
|
assert not (thread_uploads_dir / "report.pdf").exists()
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user