diff --git a/CHANGELOG.md b/CHANGELOG.md index ea24218276..11bc8cca30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,10 @@ This section accumulates work toward the **2.3.0** milestone ### Fixed +- **sandbox:** AIO commands from independent subagents now run concurrently when + they carry per-call environment variables, including Lark CLI commands. + Commands within one subagent remain ordered and client cleanup drains them. + - **gateway:** Deleting a thread no longer reports success while its conversation stays readable. When the checkpoint store failed during `DELETE /api/threads/{id}`, the route logged it at DEBUG, removed the diff --git a/README.md b/README.md index ca18c0600b..32e8b66b2c 100644 --- a/README.md +++ b/README.md @@ -3064,6 +3064,23 @@ For example, independent read-only research can run concurrently when the wall-c ### Sandbox & File System +AIO runs commands from independent subagents concurrently, including Lark CLI +and skill commands with per-call environment variables. Commands within one +subagent stay ordered, and each environment-bearing call uses a fresh session. +Environment-bearing calls share a per-sandbox Bash admission limit, configured +by `sandbox.environment.AIO_BASH_MAX_SESSIONS` (positive integer, default `50`). +Excess calls wait without creating a session and can be cancelled while waiting; +the slot is retained through session cleanup. This pool is separate from +`MAX_SHELL_SESSIONS`. Docker reuse checks the persisted container environment; +K8s reuse checks the provisioner's Pod-reported Bash capacity. Upgrade the +provisioner with the Gateway: missing capacity metadata blocks admission. +Cancellation remains routed to the running command while scope release or +sandbox close waits for it to finish and clean up. +Scope release waits for in-flight execution and cleanup without a separate drain +deadline. Allow for the remaining command timeout (600 seconds by default), plus +session/transport/recovery/cleanup time when setting operator-facing release +budgets; the HTTP cleanup timeout does not bound the preceding lock wait. + With local host bash enabled, the path guard permits division in simple shell arithmetic and recognized single-quoted Python, Node.js and Perl programs. Variable-built paths remain checked even when the filename is numeric or a diff --git a/backend/packages/harness/deerflow/community/aio_sandbox/AGENTS.md b/backend/packages/harness/deerflow/community/aio_sandbox/AGENTS.md index 73a08e46ad..32dfc804e4 100644 --- a/backend/packages/harness/deerflow/community/aio_sandbox/AGENTS.md +++ b/backend/packages/harness/deerflow/community/aio_sandbox/AGENTS.md @@ -30,6 +30,26 @@ receipts fence reuse through existing quarantine. New output from a running command needs its next offset; empty polls retain the previous offsets. Markers support cooperative cleanup, not hostile-code isolation. +Scoped env calls register in the same scope table and hold its lock through +fresh bash creation, execution, polling and cleanup. Only that private path +skips the shared implicit-shell lock. Releasing entries reject admission and stay +visible to close until drained; queued calls recheck identity. Keep same-scope +ordering, fresh env sessions, cancellation ownership and closed admission. +Close also retains draining entries. Retire the command owner only after the +scope lock drains, under the registry identity check, so cancellation during +release/close reaches the active command and concurrent releases stay idempotent. +`scope_lock_held=True` requires ownership of the registered scope lock for the +whole invocation, including retries. Release has no separate drain deadline; +HTTP cleanup bounds do not bound that lock wait. A timed-out release must retain +ownership until work drains before client detach or reuse can be safe. +All env calls share a cancellable Bash semaphore, held through cleanup. Scoped +calls acquire it after the scope lock, then recheck registry admission; unscoped +calls acquire it after the shared lock. `AIO_BASH_MAX_SESSIONS` (default 50) sets +both provisioning and client limits. Docker reuse checks persisted env; remote +reuse requires Pod-attested capacity. Unknown metadata defers, never proves a +need for replacement. Keep this pool independent of persistent Shell capacity. +Regressions: `tests/test_aio_scoped_env_concurrency.py`, `tests/test_aio_bash_capacity.py`. + ## Readiness polling Build the async readiness client off the event loop. Drain construction and diff --git a/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py b/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py index 519cccc7d3..a5f3722da2 100644 --- a/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py +++ b/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox.py @@ -5,6 +5,7 @@ import math import threading import time import uuid +from contextlib import nullcontext from dataclasses import dataclass, field import httpx @@ -43,10 +44,11 @@ _BASH_EXEC_UNSUPPORTED_ERROR = ( @dataclass class _ScopedShellSession: - """One server-side shell session serialized within an agent execution.""" + """An execution's command lock and optional persistent shell session.""" lock: threading.Lock = field(default_factory=threading.Lock) session_id: str | None = None + releasing: bool = False @dataclass @@ -83,6 +85,7 @@ class AioSandbox(Sandbox): request_headers: dict[str, str] | None = None, default_command_timeout: float | None = None, lark_cli_broker: bool | None = None, + max_bash_sessions: int = 50, ): """Initialize the AIO sandbox. @@ -99,8 +102,13 @@ class AioSandbox(Sandbox): sandbox fail-closed, so a construction site that forgets to thread the attested value cannot silently authorize the plaintext credential-mount overlay. + max_bash_sessions: Bash admission limit, no greater than the + admitted runtime's AIO_BASH_MAX_SESSIONS (default 50). """ super().__init__(id) + if type(max_bash_sessions) is not int or max_bash_sessions <= 0: + raise ValueError("max_bash_sessions must be a positive integer") + self._bash_slots = threading.BoundedSemaphore(max_bash_sessions) self.lark_cli_broker = lark_cli_broker if default_command_timeout is None: self._default_command_timeout = self._DEFAULT_HARD_TIMEOUT @@ -165,15 +173,14 @@ class AioSandbox(Sandbox): cleanup is never blocked. """ # 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. + # Keep entries registered while draining so concurrent release still + # joins their locks and cancellation reaches the admitted command owner. + # Later calls fail without creating an orphaned registry entry. with self._scope_registry_lock: if self._closed: return self._closed = True scoped_sessions = list(self._scoped_shell_sessions.items()) - self._scoped_shell_sessions.clear() for scope_id, scoped in scoped_sessions: with scoped.lock: if scoped.session_id is not None and self._client is not None: @@ -184,6 +191,10 @@ class AioSandbox(Sandbox): request_options=self._bounded_cleanup_request_options(), ) scoped.session_id = None + with self._scope_registry_lock: + if self._scoped_shell_sessions.get(scope_id) is scoped: + self._command_scopes.release(scope_id) + self._scoped_shell_sessions.pop(scope_id) with self._lock: if self._recovery_session_id is not None and self._client is not None: @@ -589,13 +600,13 @@ class AioSandbox(Sandbox): *, scope_id: str | None = None, ) -> str: - """Run no-env commands in one persistent session per subagent run. + """Serialize commands within a subagent scope, independently of others. 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. + bearing commands keep the per-call fresh ``bash.exec`` session semantics. """ - if env or scope_id is None: + if scope_id is None: return self.execute_command(command, env=env, timeout=timeout) _validate_extra_env(env) @@ -607,18 +618,25 @@ class AioSandbox(Sandbox): scope_id, _ScopedShellSession(), ) - with command_lock(scoped.lock, current_command_scope(self)): + if scoped.releasing: + raise RuntimeError("sandbox command scope is no longer active") + with command_lock(scoped.lock, current_command_scope(self)), command_lock(self._bash_slots, current_command_scope(self)) if env else nullcontext(): # 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: + if self._closed or scoped.releasing 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 env: + # This registry entry also makes close()/scope release drain + # fresh bash work. Keep per-call env sessions independent of + # the shared implicit shell and of other execution scopes. + return self._execute_with_env(command, env, self._effective_command_timeout(timeout), scope_lock_held=True) if scoped.session_id is None: scoped.session_id = self._create_shell_session(client) try: @@ -696,22 +714,40 @@ class AioSandbox(Sandbox): return f"Error: {e}" def release_command_scope(self, scope_id: str) -> None: - """Release execution metadata and its explicit shell session.""" - self._command_scopes.release(scope_id) + """Fence new commands, drain in-flight work, then release the session. + + Waiting for the scope lock has no separate deadline. An active command + can hold it for its remaining command timeout (600s by default), plus + session creation, transport/recovery and cleanup time. Callers holding + a thread-key serializer must account for this in their release budget. + Returning early would require retaining ownership until work is drained. + """ 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: + scoped = self._scoped_shell_sessions.get(scope_id) + if scoped is None: + self._command_scopes.release(scope_id) return - self._cleanup_session_best_effort( - self._client, - scoped.session_id, - context=f"execution scope {scope_id}", - request_options=self._bounded_cleanup_request_options(), - ) - scoped.session_id = None + # Keep draining entries visible to close(), while fencing new + # and already-queued commands before waiting on their lock. + scoped.releasing = True + try: + 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}", + request_options=self._bounded_cleanup_request_options(), + ) + scoped.session_id = None + finally: + with self._scope_registry_lock: + if self._scoped_shell_sessions.get(scope_id) is scoped: + # Keep cancellation routed to the original owner until + # execution and cleanup drain. Only this entry may retire it. + self._command_scopes.release(scope_id) + self._scoped_shell_sessions.pop(scope_id) @staticmethod def _format_timeout_duration(timeout: float) -> str: @@ -979,11 +1015,16 @@ class AioSandbox(Sandbox): command: str, env: dict[str, str], timeout: float, + *, + scope_lock_held: bool = False, ) -> str: """Execute a command with per-call environment variables injected. The persistent-shell ``shell.exec_command`` API has no env parameter, so injected commands use the ``bash.exec`` API which accepts per-command env. + If ``scope_lock_held`` is True, the caller must hold the registered + per-scope execution lock through this call, including recovery retries; + see ``_run_bash_exec`` for the drain contract. Each call creates an explicit transient session and closes it after the command, so injected request-scoped secrets are scoped to this command, never persist across calls, and do not consume the server's session @@ -1010,10 +1051,10 @@ class AioSandbox(Sandbox): """ if self._bash_exec_unsupported: return _BASH_EXEC_UNSUPPORTED_ERROR - output, status = self._run_bash_exec(command, env, timeout) + output, status = self._run_bash_exec(command, env, timeout, scope_lock_held=scope_lock_held) if status in (None, "completed") and output and _ERROR_OBSERVATION_SIGNATURE in output: logger.warning("ErrorObservation detected in bash.exec output, retrying on a fresh session") - retried, retry_status = self._run_bash_exec(command, env, timeout) + retried, retry_status = self._run_bash_exec(command, env, timeout, scope_lock_held=scope_lock_held) if retry_status not in (None, "completed"): return retried if retried and _ERROR_OBSERVATION_SIGNATURE not in retried: @@ -1099,10 +1140,21 @@ class AioSandbox(Sandbox): command: str, env: dict[str, str], timeout: float, + *, + scope_lock_held: bool = False, ) -> tuple[str, str | None]: - """Single bash.exec invocation in an explicitly released fresh session.""" + """Single bash.exec invocation in an explicitly released fresh session. + + ``scope_lock_held=True`` requires the caller to hold its registered + ``_ScopedShellSession.lock`` and a Bash permit for this entire invocation. Release/close + join that lock to drain work before detaching the client; the shared + lock is skipped only under this precondition. + """ scope = current_command_scope(self) - with command_lock(self._lock, scope): + # Scoped callers already hold their registered execution lock, which + # close() and release_command_scope() drain before detaching the client. + # Unscoped calls retain the shared lock for that lifecycle guarantee. + with nullcontext() if scope_lock_held else command_lock(self._lock, scope), nullcontext() if scope_lock_held else command_lock(self._bash_slots, scope): for attempt in range(2): session_id: str | None = None try: diff --git a/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox_provider.py b/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox_provider.py index 07ccdfd798..cf3faab096 100644 --- a/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox_provider.py +++ b/backend/packages/harness/deerflow/community/aio_sandbox/aio_sandbox_provider.py @@ -365,6 +365,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider): provisioner_url=provisioner_url, api_key=api_key, max_shell_sessions=self._config.get("max_shell_sessions"), + max_bash_sessions=self._config.get("max_bash_sessions", 50), required_shell_sessions=self._config.get("required_shell_sessions", 0), ) @@ -405,6 +406,15 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider): AioSandbox._DEFAULT_HARD_TIMEOUT, ) max_running_subagents = int(getattr(getattr(config, "subagent_runtime", None), "max_running", 3)) + try: + max_bash_sessions = int(environment.get("AIO_BASH_MAX_SESSIONS", "50")) + except (TypeError, ValueError) as exc: + raise ValueError("sandbox.environment.AIO_BASH_MAX_SESSIONS must be a positive integer") from exc + if max_bash_sessions <= 0: + raise ValueError("sandbox.environment.AIO_BASH_MAX_SESSIONS must be a positive integer") + # Explicitly configure the same budget used by client admission. Bash + # and persistent Shell have independent server-side session pools. + environment["AIO_BASH_MAX_SESSIONS"] = str(max_bash_sessions) required_shell_sessions = max_running_subagents + _SHELL_SESSION_HEADROOM configured_shell_sessions = environment.get("MAX_SHELL_SESSIONS") if configured_shell_sessions is None: @@ -430,6 +440,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider): "thread_data_mounts": getattr(sandbox_config, "thread_data_mounts", None), "environment": environment, "max_shell_sessions": max_shell_sessions, + "max_bash_sessions": max_bash_sessions, "required_shell_sessions": required_shell_sessions, "network": sandbox_config.network.model_dump(), "ownership": getattr(sandbox_config, "ownership", None), @@ -1910,6 +1921,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider): base_url=info.sandbox_url, request_headers=info.request_headers, default_command_timeout=self._config.get("command_timeout", AioSandbox._DEFAULT_HARD_TIMEOUT), + max_bash_sessions=self._config.get("max_bash_sessions", 50), lark_cli_broker=info.lark_cli_broker, ) self._sandboxes[sandbox_id] = sandbox @@ -1963,6 +1975,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider): base_url=info.sandbox_url, request_headers=info.request_headers, default_command_timeout=self._config.get("command_timeout", AioSandbox._DEFAULT_HARD_TIMEOUT), + max_bash_sessions=self._config.get("max_bash_sessions", 50), lark_cli_broker=info.lark_cli_broker, ) # Ownership first, so a failure cannot leave a tracked-but-unowned sandbox. @@ -2024,6 +2037,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider): base_url=info.sandbox_url, request_headers=info.request_headers, default_command_timeout=self._config.get("command_timeout", AioSandbox._DEFAULT_HARD_TIMEOUT), + max_bash_sessions=self._config.get("max_bash_sessions", 50), lark_cli_broker=info.lark_cli_broker, ) key = self._thread_key(thread_id, self._effective_acquire_user_id(user_id)) if thread_id else None diff --git a/backend/packages/harness/deerflow/community/aio_sandbox/local_backend.py b/backend/packages/harness/deerflow/community/aio_sandbox/local_backend.py index 4f05b3b195..7726d38561 100644 --- a/backend/packages/harness/deerflow/community/aio_sandbox/local_backend.py +++ b/backend/packages/harness/deerflow/community/aio_sandbox/local_backend.py @@ -36,6 +36,14 @@ logger = logging.getLogger(__name__) _AIO_DEFAULT_MAX_SHELL_SESSIONS = 10 +def _parse_bash_capacity(config: dict) -> int: + value = _extract_container_environment(config, "AIO_BASH_MAX_SESSIONS") + try: + return max(0, int(value)) if value is not None else 50 + except (TypeError, ValueError): + return 0 + + class _ExistingRestrictedSandbox(RuntimeError): def __init__(self, info: SandboxInfo): super().__init__(f"restricted sandbox {info.sandbox_id} already exists") @@ -58,6 +66,7 @@ class _ContainerInspection: networks: frozenset[str] relay_token: str | None = None max_shell_sessions: int | None = None + max_bash_sessions: int | None = None container_id: str | None = None @@ -712,7 +721,13 @@ class LocalContainerBackend(SandboxBackend): except (TypeError, ValueError): return False actual = inspection.max_shell_sessions if inspection.max_shell_sessions is not None else _AIO_DEFAULT_MAX_SHELL_SESSIONS - return actual >= required + if actual < required: + return False + # A reused container may predate the current Bash environment override. + # Never admit it using the new client's larger semaphore budget. + required_bash = int(self._environment.get("AIO_BASH_MAX_SESSIONS", "50")) + actual_bash = inspection.max_bash_sessions if inspection.max_bash_sessions is not None else 50 + return actual_bash >= required_bash def _network_policy_digest(self) -> str: allow_domains = self._network_config.get("allow_domains", []) @@ -1998,6 +2013,7 @@ class LocalContainerBackend(SandboxBackend): networks=frozenset(str(value) for value in (network_settings.get("Networks") or {})), relay_token=_extract_container_environment(config, RELAY_TOKEN_ENV), max_shell_sessions=max_shell_sessions, + max_bash_sessions=_parse_bash_capacity(config), container_id=entry.get("Id") if isinstance(entry.get("Id"), str) and entry.get("Id") else None, ) return out diff --git a/backend/packages/harness/deerflow/community/aio_sandbox/remote_backend.py b/backend/packages/harness/deerflow/community/aio_sandbox/remote_backend.py index 8e2be01a1c..a276bcc3d7 100644 --- a/backend/packages/harness/deerflow/community/aio_sandbox/remote_backend.py +++ b/backend/packages/harness/deerflow/community/aio_sandbox/remote_backend.py @@ -152,6 +152,7 @@ class RemoteSandboxBackend(SandboxBackend): max_shell_sessions: int | None = None, *, required_shell_sessions: int = 0, + max_bash_sessions: int | None = None, ): """Initialize with the provisioner service URL and optional API key. @@ -163,10 +164,12 @@ class RemoteSandboxBackend(SandboxBackend): max_shell_sessions: Optional AIO shell-session capacity forwarded to each provisioned sandbox Pod. required_shell_sessions: Minimum usable capacity, even when new Pods use the image default. + max_bash_sessions: Bash capacity to provision and attest before admission. """ self._provisioner_url = provisioner_url.rstrip("/") self._api_key = api_key self._max_shell_sessions = max_shell_sessions + self._max_bash_sessions = max_bash_sessions self._required_shell_sessions = max(required_shell_sessions, max_shell_sessions or 0) @property @@ -185,7 +188,17 @@ class RemoteSandboxBackend(SandboxBackend): except (TypeError, ValueError): return True - def _sandbox_info(self, payload: dict, sandbox_id: str) -> SandboxInfo: + def _requires_bash_capacity_replacement(self, payload: dict) -> bool: + if self._max_bash_sessions is None: + return False + reported = payload.get("max_bash_sessions") + if type(reported) is not int or reported <= 0: + # Unknown metadata is not proof that a healthy Pod needs deletion. + raise RuntimeError("Provisioner did not report valid max_bash_sessions; upgrade the provisioner before admitting Bash commands") + return reported < self._max_bash_sessions + + def _sandbox_info(self, payload: dict, sandbox_id: str, *, check_bash_capacity: bool = True) -> SandboxInfo: + bash_replacement = check_bash_capacity and self._requires_bash_capacity_replacement(payload) broker = payload.get("lark_cli_broker") generation = payload.get("container_id") return SandboxInfo( @@ -193,7 +206,7 @@ class RemoteSandboxBackend(SandboxBackend): sandbox_url=payload["sandbox_url"], container_id=generation if isinstance(generation, str) and generation else None, lark_cli_broker=broker if isinstance(broker, bool) else None, - requires_replacement=self._requires_shell_capacity_replacement(payload), + requires_replacement=self._requires_shell_capacity_replacement(payload) or bash_replacement, ) # ── SandboxBackend interface ────────────────────────────────────────── @@ -339,6 +352,8 @@ class RemoteSandboxBackend(SandboxBackend): } if self._max_shell_sessions is not None: payload["max_shell_sessions"] = self._max_shell_sessions + if self._max_bash_sessions is not None: + payload["max_bash_sessions"] = self._max_bash_sessions provisioner_extra_mounts = _provisioner_extra_mounts_payload( extra_mounts, skills_container_path=normalized_skills_container_path, @@ -356,7 +371,13 @@ class RemoteSandboxBackend(SandboxBackend): ) resp.raise_for_status() data = resp.json() - info = self._sandbox_info(data, sandbox_id) + info = self._sandbox_info(data, sandbox_id, check_bash_capacity=False) + try: + insufficient_bash = self._requires_bash_capacity_replacement(data) + except RuntimeError as exc: + raise SandboxCreationError(str(exc), info=info) from exc + if insufficient_bash: + raise SandboxCreationError(f"Provisioner returned sandbox {sandbox_id} with insufficient Bash-session capacity", info=info) if self._max_shell_sessions is not None and "max_shell_sessions" not in data: raise SandboxCreationError("Provisioner did not report max_shell_sessions; Gateway/provisioner version skew prevents shell-capacity validation", info=info) if self._requires_shell_capacity_replacement(data): diff --git a/backend/tests/test_aio_bash_capacity.py b/backend/tests/test_aio_bash_capacity.py new file mode 100644 index 0000000000..ffbd51fe7c --- /dev/null +++ b/backend/tests/test_aio_bash_capacity.py @@ -0,0 +1,222 @@ +"""Bash admission includes creation, execution and cleanup, independently of PTYs.""" + +import threading +from types import SimpleNamespace + +import httpx +import pytest +from test_aio_scoped_env_concurrency import ( + _join_observed_thread, + _ObservedScopeLock, + _start_observed_thread, +) +from test_aio_scoped_env_concurrency import ( + paused_env_command as _paused_env_command, +) + +paused_env_command = _paused_env_command + + +@pytest.mark.parametrize("paused_env_command", [1], indirect=True) +@pytest.mark.parametrize("cancel", [False, True]) +@pytest.mark.parametrize("direct", [False, True]) +def test_bash_capacity_waiter_is_cancellable_and_wakes_after_cleanup(paused_env_command, monkeypatch, cancel, direct): + from deerflow.sandbox.command_scope import CommandScopeCancelled, command_owner + + case = paused_env_command + slots = _ObservedScopeLock(case.sandbox._bash_slots) + case.sandbox._bash_slots = slots + original = case.http_client._transport.handle_request + cleanup_entered = threading.Event() + cleanup_resume = threading.Event() + live_sessions = set() + + def transport(request): + path = request.url.path + if path == "/v1/bash/sessions/create": + assert not live_sessions, "Client exceeded the one-slot Bash pool" + live_sessions.add("active") + if path.startswith("/v1/bash/sessions/") and path.endswith("/close"): + cleanup_entered.set() + assert cleanup_resume.wait(5) + live_sessions.clear() + return original(request) + + monkeypatch.setattr(case.http_client._transport, "handle_request", transport) + + def run_second(): + if direct: + with command_owner(case.sandbox, "second"): + if cancel: + with pytest.raises(CommandScopeCancelled, match="cancelled"): + case.sandbox.execute_command("ordinary-command", env={"REQUEST_VALUE": "two"}) + return "cancelled" + return case.sandbox.execute_command("ordinary-command", env={"REQUEST_VALUE": "two"}) + return case.run(owner="second", value="two") + + first = _start_observed_thread(case.run, name="first-env") + queued = None + try: + assert case.entered.wait(5) + queued = _start_observed_thread(run_second, name="queued-env") + assert slots.queued.wait(5) + assert len(case.executions) == 1 + if cancel: + case.sandbox.request_command_cancel("second") + _join_observed_thread(queued) + assert "cancelled" in queued.results[0] + case.resume.set() + assert cleanup_entered.wait(5) + assert len(case.executions) == 1, "A permit must cover session cleanup" + cleanup_resume.set() + _join_observed_thread(first) + _join_observed_thread(queued) + assert first.results == ["done"] + if not cancel: + assert queued.results == ["done"] + assert len(case.executions) == 2 + assert case.run(owner="third") == "done", "Cancellation must not leak a permit" + assert not live_sessions + finally: + case.resume.set() + cleanup_resume.set() + for worker in (first, queued): + if worker is not None: + _join_observed_thread(worker) + + +@pytest.mark.parametrize("capacity", [0, -1, True, 1.5, "2"]) +def test_invalid_bash_capacity_rejected_before_client_creation(capacity): + from deerflow.community.aio_sandbox.aio_sandbox import AioSandbox + + with pytest.raises(ValueError, match="max_bash_sessions"): + AioSandbox("invalid", "http://sandbox.invalid", max_bash_sessions=capacity) + + +@pytest.mark.parametrize("configured,expected", [(None, 50), ("1", 1), ("72", 72)]) +def test_provider_configures_bash_budget_independently_of_shell(monkeypatch, configured, expected): + from deerflow.community.aio_sandbox import aio_sandbox_provider as module + from deerflow.config.sandbox_config import SandboxConfig + + config = SimpleNamespace( + sandbox=SandboxConfig(use="deerflow.community.aio_sandbox:AioSandboxProvider", environment={} if configured is None else {"AIO_BASH_MAX_SESSIONS": configured}), + subagent_runtime=SimpleNamespace(max_running=64), + stream_bridge=None, + ) + monkeypatch.setattr(module, "get_app_config", lambda: config) + provider = module.AioSandboxProvider.__new__(module.AioSandboxProvider) + loaded = provider._load_config() + assert loaded["max_bash_sessions"] == expected + assert loaded["environment"]["AIO_BASH_MAX_SESSIONS"] == str(expected) + assert loaded["max_shell_sessions"] == 65 + + +@pytest.mark.parametrize("configured", ["0", "-1", "1.5", "invalid"]) +def test_provider_rejects_invalid_bash_budget(monkeypatch, configured): + from deerflow.community.aio_sandbox import aio_sandbox_provider as module + from deerflow.config.sandbox_config import SandboxConfig + + monkeypatch.setattr(module, "get_app_config", lambda: SimpleNamespace(sandbox=SandboxConfig(use="test", environment={"AIO_BASH_MAX_SESSIONS": configured}))) + with pytest.raises(ValueError, match="AIO_BASH_MAX_SESSIONS"): + module.AioSandboxProvider.__new__(module.AioSandboxProvider)._load_config() + + +@pytest.mark.parametrize("actual,compatible", [("1", False), ("50", True), ("72", True), ("invalid", False), ("0", False), (None, True)]) +def test_local_reuse_validates_persisted_bash_budget(actual, compatible): + from deerflow.community.aio_sandbox.local_backend import LocalContainerBackend, _parse_bash_capacity + + backend = LocalContainerBackend.__new__(LocalContainerBackend) + backend._environment = {"AIO_BASH_MAX_SESSIONS": "50"} + backend._required_shell_sessions = 1 + inspection = SimpleNamespace(max_shell_sessions=10, max_bash_sessions=_parse_bash_capacity({"Env": [] if actual is None else [f"AIO_BASH_MAX_SESSIONS={actual}"]})) + assert backend._has_compatible_shell_capacity(inspection) is compatible + + +@pytest.mark.parametrize("actual", [None, "50", 0, -1, True]) +def test_remote_unknown_capacity_is_not_assumed_or_replaced(actual): + from deerflow.community.aio_sandbox.remote_backend import RemoteSandboxBackend + + backend = RemoteSandboxBackend("http://provisioner", max_bash_sessions=50) + with pytest.raises(RuntimeError, match="max_bash_sessions"): + backend._sandbox_info({"sandbox_url": "http://sandbox", "max_bash_sessions": actual}, "test") + + +@pytest.mark.parametrize("actual,replacement", [(1, True), (50, False), (72, False)]) +def test_remote_reuse_validates_attested_bash_budget(actual, replacement): + from deerflow.community.aio_sandbox.remote_backend import RemoteSandboxBackend + + backend = RemoteSandboxBackend("http://provisioner", max_bash_sessions=50) + assert backend._sandbox_info({"sandbox_url": "http://sandbox", "max_bash_sessions": actual}, "test").requires_replacement is replacement + + +@pytest.mark.parametrize("paused_env_command", [1], indirect=True) +@pytest.mark.parametrize("teardown", ["close", "release"]) +def test_capacity_waiter_rechecks_teardown_before_creation(paused_env_command, teardown): + case = paused_env_command + slots = _ObservedScopeLock(case.sandbox._bash_slots) + case.sandbox._bash_slots = slots + first = _start_observed_thread(case.run, name="first-env") + queued = stopping = None + try: + assert case.entered.wait(5) + queued = _start_observed_thread(lambda: case.run(owner="second"), name="queued-env") + assert slots.queued.wait(5) + scope = case.sandbox._scoped_shell_sessions["first" if teardown == "close" else "second"] + observed = _ObservedScopeLock(scope.lock) + scope.lock = observed + action = case.sandbox.close if teardown == "close" else lambda: case.sandbox.release_command_scope("second") + stopping = _start_observed_thread(action, name="scope-teardown") + assert observed.teardown.wait(5) + case.resume.set() + for worker in (first, queued, stopping): + _join_observed_thread(worker) + assert "no longer active" in queued.results[0] + assert len(case.executions) == 1 + finally: + case.resume.set() + for worker in (first, queued, stopping): + if worker is not None: + _join_observed_thread(worker) + + +@pytest.mark.parametrize("paused_env_command", [1], indirect=True) +@pytest.mark.parametrize("failed_path", ["/v1/bash/sessions/create", "/v1/bash/exec"]) +def test_failed_bash_request_returns_capacity(paused_env_command, monkeypatch, failed_path): + case = paused_env_command + original = case.http_client._transport.handle_request + rejected = False + + def transport(request): + nonlocal rejected + if request.url.path == failed_path and not rejected: + rejected = True + return httpx.Response(400, json={"message": "Invalid request"}) + return original(request) + + monkeypatch.setattr(case.http_client._transport, "handle_request", transport) + case.resume.set() + assert "400" in case.run() + assert case.run(owner="second") == "done" + + +@pytest.mark.parametrize("reported", [1, 50, None]) +def test_remote_create_transmits_and_verifies_bash_budget(monkeypatch, reported): + from deerflow.community.aio_sandbox import remote_backend as module + from deerflow.community.aio_sandbox.backend import SandboxCreationError + + backend = module.RemoteSandboxBackend("http://provisioner", max_bash_sessions=50) + monkeypatch.setattr(module, "user_should_see_legacy_skills", lambda _user: False) + requests = [] + + def post(url, **kwargs): + requests.append(kwargs["json"]) + return SimpleNamespace(raise_for_status=lambda: None, json=lambda: {"sandbox_url": "http://sandbox", "max_bash_sessions": reported, "container_id": "pod-uid"}) + + monkeypatch.setattr(module.requests, "post", post) + if reported == 50: + assert backend._provisioner_create("thread", "sandbox").container_id == "pod-uid" + else: + with pytest.raises(SandboxCreationError) as error: + backend._provisioner_create("thread", "sandbox") + assert error.value.info.container_id == "pod-uid" + assert requests[0]["max_bash_sessions"] == 50 diff --git a/backend/tests/test_aio_sandbox.py b/backend/tests/test_aio_sandbox.py index 5dc235e113..90fda81e48 100644 --- a/backend/tests/test_aio_sandbox.py +++ b/backend/tests/test_aio_sandbox.py @@ -1941,6 +1941,8 @@ class TestScopedShellSessions: }, ) sandbox._client.shell.create_session.assert_not_called() + assert sandbox._scoped_shell_sessions["subagent-a"].session_id is None + sandbox.release_command_scope("subagent-a") assert sandbox._scoped_shell_sessions == {} def test_env_session_cleanup_failure_does_not_mask_command_output(self, sandbox): diff --git a/backend/tests/test_aio_sandbox_provider.py b/backend/tests/test_aio_sandbox_provider.py index 616632d79b..96dfb0b2fe 100644 --- a/backend/tests/test_aio_sandbox_provider.py +++ b/backend/tests/test_aio_sandbox_provider.py @@ -104,10 +104,12 @@ def test_bash_command_timeout_rejects_non_finite_values(value): SandboxConfig(bash_command_timeout=value) -def test_register_created_sandbox_forwards_configured_command_timeout(tmp_path): +@pytest.mark.parametrize("bash_capacity", [1, 50]) +def test_register_created_sandbox_forwards_configured_command_timeout(tmp_path, bash_capacity): aio_mod = importlib.import_module("deerflow.community.aio_sandbox.aio_sandbox_provider") provider = _make_provider(tmp_path) provider._config["command_timeout"] = 42 + provider._config["max_bash_sessions"] = bash_capacity provider._warm_pool = {} provider._sandbox_infos = {} provider._thread_sandboxes = {} @@ -123,6 +125,7 @@ def test_register_created_sandbox_forwards_configured_command_timeout(tmp_path): base_url="http://sandbox", request_headers=info.request_headers, default_command_timeout=42, + max_bash_sessions=bash_capacity, lark_cli_broker=False, ) @@ -176,7 +179,7 @@ def test_backend_checks_runtime_minimum_without_overriding_image_default(monkeyp monkeypatch.setattr(local_mod.LocalContainerBackend, "_detect_runtime", lambda _self: "docker") backend = provider._create_backend() if remote: - payload = {"sandbox_id": "example", "sandbox_url": "http://sandbox:8080", "max_shell_sessions": persisted_capacity} + payload = {"sandbox_id": "example", "sandbox_url": "http://sandbox:8080", "max_shell_sessions": persisted_capacity, "max_bash_sessions": 50} def get(url, **_kwargs): data = {"sandboxes": [payload]} if url.endswith("/api/sandboxes") else payload diff --git a/backend/tests/test_aio_scoped_env_concurrency.py b/backend/tests/test_aio_scoped_env_concurrency.py new file mode 100644 index 0000000000..d64a0778d8 --- /dev/null +++ b/backend/tests/test_aio_scoped_env_concurrency.py @@ -0,0 +1,532 @@ +"""Scoped env commands run independently and drain before client teardown.""" + +import asyncio +import importlib +import json +import os +import shlex +import subprocess +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import httpx +import pytest +from agent_sandbox import Sandbox as SandboxClient +from support.shell import require_script_bash + +from deerflow.runtime.secret_context import ACTIVE_SECRETS_CONTEXT_KEY +from deerflow.sandbox.command_scope import command_owner +from deerflow.sandbox.lease import run_sync_lifecycle_operation + + +@pytest.mark.parametrize("mode", ["plain", "active-skill-env", "lark-cli"]) +def test_independent_subagent_commands_enter_distinct_sessions_concurrently(monkeypatch, tmp_path, mode): + aio = importlib.import_module("deerflow.community.aio_sandbox.aio_sandbox") + tools = importlib.import_module("deerflow.sandbox.tools") + app_config = importlib.import_module("deerflow.config.app_config") + native_bash = require_script_bash() + shim_bin = tmp_path / "bin" + shim_bin.mkdir() + shim = shim_bin / "lark-cli" + shim.write_text("#!/bin/sh\nsleep 0.4\nprintf done\n", encoding="utf-8") + shim.chmod(0o755) + bookkeeping = threading.Lock() + sessions = {} + in_flight = peak = 0 + entered = threading.Event() + requests = [] + + def complete(session): + nonlocal in_flight + if session.get("finished"): + return + try: + # Keep each service-side operation admitted until all peers have + # entered. Correctness must not depend on their sleep windows + # overlapping on a busy host. + assert entered.wait(5), "Independent commands did not all reach the service before any completion" + stdout, stderr = session["process"].communicate(timeout=5) + session.update(stdout=stdout, stderr=stderr, finished=True) + finally: + with bookkeeping: + in_flight -= 1 + + def transport(request): + nonlocal in_flight, peak + path = request.url.path + payload = json.loads(request.content) if request.content else {} + with bookkeeping: + requests.append((request.method, path, payload)) + if path == "/v1/shell/sessions/create": + return httpx.Response(200, json={"success": True, "data": {"session_id": payload["id"], "working_dir": "/mnt/user-data/workspace"}}) + if path == "/v1/bash/sessions/create": + return httpx.Response( + 200, json={"success": True, "data": {"session_id": payload["session_id"], "status": "ready", "working_dir": "/mnt/user-data/workspace", "created_at": "2026-11-03T00:00:00Z", "last_used_at": "2026-11-03T00:00:00Z"}} + ) + if path in ("/v1/shell/exec", "/v1/bash/exec"): + session_id = payload.get("id") or payload.get("session_id") + assert session_id is not None + command = payload["command"].replace("cd /mnt/user-data/workspace;", f"cd {shlex.quote(tmp_path.as_posix())};") + environment = {**os.environ, **payload.get("env", {})} + if "PATH" in payload.get("env", {}): + environment["PATH"] = environment["PATH"].replace("/mnt/integrations/lark-cli/runtime/bin", shim_bin.as_posix()) + process = subprocess.Popen([native_bash, "--noprofile", "--norc", "-c", command], cwd=tmp_path, env=environment, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8") + session = {"process": process, "command": payload["command"], "id": session_id} + with bookkeeping: + sessions[session_id] = session + in_flight += 1 + peak = max(peak, in_flight) + if in_flight == 3: + entered.set() + data = {"session_id": session_id, "command": payload["command"], "status": "running"} + if path == "/v1/bash/exec": + data["command_id"] = session_id + return httpx.Response(200, json={"success": True, "data": data}) + if path == "/v1/shell/wait": + complete(sessions[payload["id"]]) + return httpx.Response(200, json={"success": True, "data": {"status": "completed"}}) + if path == "/v1/shell/view": + session = sessions[payload["id"]] + assert session["finished"] + return httpx.Response( + 200, + json={ + "success": True, + "data": { + "session_id": session["id"], + "status": "completed", + "command": session["command"], + "output": session["stdout"], + "exit_code": session["process"].returncode, + "console": [{"ps1": "$ ", "command": session["command"], "output": session["stdout"]}], + }, + }, + ) + if path == "/v1/bash/output": + session = sessions[payload["session_id"]] + complete(session) + return httpx.Response( + 200, + json={ + "success": True, + "data": { + "session_id": session["id"], + "stdout": session["stdout"], + "stderr": session["stderr"], + "command": {"command_id": session["id"], "command": session["command"], "status": "completed", "exit_code": session["process"].returncode}, + }, + }, + ) + assert path.startswith("/v1/shell/sessions/") or (path.startswith("/v1/bash/sessions/") and path.endswith("/close")) + return httpx.Response(200, json={"success": True}) + + with httpx.Client(transport=httpx.MockTransport(transport)) as http_client: + sdk = SandboxClient(base_url="http://sandbox.invalid", httpx_client=http_client) + monkeypatch.setattr(aio, "AioSandboxClient", lambda **_kwargs: sdk) + monkeypatch.setattr(aio, "sandbox_http_trust_env", lambda _url: True) + sandbox = aio.AioSandbox("parallel-public-tool", "http://sandbox.invalid", lark_cli_broker=False) + monkeypatch.setattr(tools, "ensure_sandbox_initialized", lambda _runtime: sandbox) + monkeypatch.setattr(tools, "is_local_sandbox", lambda _runtime: False) + monkeypatch.setattr(app_config, "get_app_config", lambda: SimpleNamespace(sandbox=SimpleNamespace(bash_output_max_chars=20000, bash_command_timeout=30))) + start = threading.Barrier(3) + + def worker(index): + owner = f"subagent:{index}" + context = {"user_id": "test-user", "sandbox_lease_owner_id": owner, "sandbox_command_scope_id": owner} + if mode == "active-skill-env": + context[ACTIVE_SECRETS_CONTEXT_KEY] = {"REQUEST_TEST_VALUE": "ordinary-test-value"} + runtime = SimpleNamespace(context=context, state={}) + start.wait(timeout=5) + command = "lark-cli audit-wait" if mode == "lark-cli" else "sleep 0.4; printf done" + return tools.bash_tool.func(runtime, command) + + started = time.perf_counter() + try: + with ThreadPoolExecutor(max_workers=3) as executor: + futures = [executor.submit(worker, index) for index in range(3)] + results = [future.result(timeout=20) for future in futures] + elapsed = time.perf_counter() - started + assert results == ["done"] * 3 + assert len(sessions) == 3, "Each operation received its own real SDK session identity" + print( + "AIO_ADMISSION_REGRESSION " + + json.dumps( + { + "mode": mode, + "commands": 3, + "sleep_seconds": 0.4, + # Diagnostic only: this includes the admission gate and + # is not the separate ungated audit benchmark. + "gated_elapsed_seconds": round(elapsed, 4), + "peak_admitted_commands": peak, + "all_entered_together": entered.is_set(), + "exec_requests": sum(path in ("/v1/shell/exec", "/v1/bash/exec") for _, path, _ in requests), + } + ) + ) + assert entered.is_set() and peak == 3, "Independent fresh sessions were serialized before reaching the AIO service" + finally: + for session in sessions.values(): + if session["process"].poll() is None: + session["process"].terminate() + session["process"].communicate(timeout=5) + sandbox.close() + + +class _ObservedScopeLock: + def __init__(self, lock): + self.lock = lock + self.queued = threading.Event() + self.teardown = threading.Event() + + def acquire(self, *args, **kwargs): + if threading.current_thread().name == "queued-env": + self.queued.set() + if threading.current_thread().name == "scope-teardown": + self.teardown.set() + return self.lock.acquire(*args, **kwargs) + + def release(self): + self.lock.release() + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *_args): + self.release() + + +@pytest.fixture +def paused_env_command(monkeypatch, request): + aio = importlib.import_module("deerflow.community.aio_sandbox.aio_sandbox") + entered = threading.Event() + resume = threading.Event() + client_closed = threading.Event() + requests = [] + executions = [] + control_sessions = set() + request_lock = threading.Lock() + + def transport(request): + assert not client_closed.is_set(), "Client closed while an admitted command still owned it" + path = request.url.path + payload = json.loads(request.content) if request.content else {} + with request_lock: + requests.append((request.method, path, payload)) + if path == "/v1/bash/sessions/create": + return httpx.Response(200, json={"success": True, "data": {"session_id": payload["session_id"], "status": "ready", "working_dir": "/home/user", "created_at": "2026-11-03T00:00:00Z", "last_used_at": "2026-11-03T00:00:00Z"}}) + if path == "/v1/bash/exec": + with request_lock: + executions.append(payload) + return httpx.Response(200, json={"success": True, "data": {"session_id": payload["session_id"], "command_id": payload["session_id"], "command": payload["command"], "status": "running"}}) + if path == "/v1/bash/output": + execution = next(item for item in executions if item["session_id"] == payload["session_id"]) + if execution is executions[0]: + entered.set() + assert resume.wait(5), "Test never released the first command" + assert not client_closed.is_set() + return httpx.Response( + 200, json={"success": True, "data": {"session_id": payload["session_id"], "stdout": "done", "command": {"command_id": payload["session_id"], "command": execution["command"], "status": "completed", "exit_code": 0}}} + ) + if path.startswith("/v1/bash/sessions/") and path.endswith("/close"): + return httpx.Response(200, json={"success": True}) + if path == "/v1/shell/sessions/create": + control_sessions.add(payload["id"]) + return httpx.Response(200, json={"success": True, "data": {"session_id": payload["id"], "working_dir": "/home/user"}}) + if path == "/v1/shell/exec": + assert payload["id"] in control_sessions + assert payload["command"].startswith("python3 -c ") + return httpx.Response(200, json={"success": True, "data": {"session_id": payload["id"], "command": payload["command"], "status": "running"}}) + if path == "/v1/shell/view": + assert payload["id"] in control_sessions + return httpx.Response(200, json={"success": True, "data": {"session_id": payload["id"], "output": '{"remaining":0}', "status": "completed", "exit_code": 0}}) + assert request.method == "DELETE" and path.removeprefix("/v1/shell/sessions/") in control_sessions + return httpx.Response(200, json={"success": True}) + + http_transport = httpx.MockTransport(transport) + real_close = http_transport.close + + def close(): + client_closed.set() + real_close() + + monkeypatch.setattr(http_transport, "close", close) + http_client = httpx.Client(transport=http_transport) + sdk = SandboxClient(base_url="http://sandbox.invalid", httpx_client=http_client) + monkeypatch.setattr(aio, "AioSandboxClient", lambda **_kwargs: sdk) + monkeypatch.setattr(aio, "sandbox_http_trust_env", lambda _url: True) + options = {"max_bash_sessions": request.param} if hasattr(request, "param") else {} + sandbox = aio.AioSandbox("scoped-env-lifecycle", "http://sandbox.invalid", **options) + + def run(owner="first", value="one"): + with command_owner(sandbox, owner): + return sandbox.execute_command_in_scope("ordinary-command", env={"REQUEST_VALUE": value}, scope_id=owner) + + yield SimpleNamespace(sandbox=sandbox, entered=entered, resume=resume, client_closed=client_closed, http_client=http_client, requests=requests, executions=executions, controls=control_sessions, run=run) + resume.set() + sandbox.close() + http_client.close() + + +def _start_observed_thread(function, *, name): + results = [] + errors = [] + + def run(): + try: + results.append(function()) + except BaseException as error: + errors.append(error) + + thread = threading.Thread(target=run, name=name) + thread.start() + return SimpleNamespace(thread=thread, results=results, errors=errors) + + +def _join_observed_thread(worker): + worker.thread.join(5) + assert not worker.thread.is_alive() + assert not worker.errors + + +def test_env_commands_in_one_scope_remain_serialized(paused_env_command): + case = paused_env_command + first = _start_observed_thread(case.run, name="first-env") + second = None + try: + assert case.entered.wait(5) + scoped = case.sandbox._scoped_shell_sessions["first"] + observed = _ObservedScopeLock(scoped.lock) + scoped.lock = observed + second = _start_observed_thread(lambda: case.run(value="two"), name="queued-env") + assert observed.queued.wait(5) + assert len(case.executions) == 1 + case.resume.set() + _join_observed_thread(first) + _join_observed_thread(second) + assert first.results == second.results == ["done"] + assert [item["env"] for item in case.executions] == [{"REQUEST_VALUE": "one"}, {"REQUEST_VALUE": "two"}] + assert len({item["session_id"] for item in case.executions}) == 2 + assert not case.controls + finally: + case.resume.set() + _join_observed_thread(first) + if second is not None: + _join_observed_thread(second) + + +@pytest.mark.parametrize("teardown", ["close", "release"]) +def test_teardown_drains_env_and_rejects_already_queued_commands(paused_env_command, teardown): + case = paused_env_command + first = _start_observed_thread(case.run, name="first-env") + queued = stopping = None + try: + assert case.entered.wait(5) + scoped = case.sandbox._scoped_shell_sessions["first"] + observed = _ObservedScopeLock(scoped.lock) + scoped.lock = observed + queued = _start_observed_thread(case.run, name="queued-env") + assert observed.queued.wait(5) + action = case.sandbox.close if teardown == "close" else lambda: case.sandbox.release_command_scope("first") + stopping = _start_observed_thread(action, name="scope-teardown") + assert observed.teardown.wait(5) + assert not case.client_closed.is_set() + assert stopping.thread.is_alive() + if teardown == "close": + assert case.run(owner="new") == "Error: sandbox client is closed" + case.resume.set() + for worker in (first, queued, stopping): + _join_observed_thread(worker) + assert first.results == ["done"] + assert queued.results == ["Error: sandbox command scope is no longer active"] + assert len(case.executions) == 1 + assert case.client_closed.is_set() is (teardown == "close") + assert not case.sandbox._scoped_shell_sessions + finally: + case.resume.set() + for worker in (first, queued, stopping): + if worker is not None: + _join_observed_thread(worker) + + +def test_cancelled_env_waiter_never_submits_a_command(paused_env_command): + case = paused_env_command + first = _start_observed_thread(case.run, name="first-env") + queued = None + try: + assert case.entered.wait(5) + scoped = case.sandbox._scoped_shell_sessions["first"] + observed = _ObservedScopeLock(scoped.lock) + scoped.lock = observed + queued = _start_observed_thread(case.run, name="queued-env") + assert observed.queued.wait(5) + case.sandbox.request_command_cancel("first") + _join_observed_thread(queued) + assert "cancelled" in queued.results[0] + assert len(case.executions) == 1 + assert not case.controls + finally: + case.resume.set() + _join_observed_thread(first) + if queued is not None: + _join_observed_thread(queued) + + +@pytest.mark.parametrize("close_first", [False, True]) +@pytest.mark.parametrize("release_count", [1, 2]) +def test_cancel_reaches_env_command_while_scope_drains(paused_env_command, close_first, release_count): + case = paused_env_command + first = _start_observed_thread(case.run, name="first-env") + workers = [] + progress = {name: threading.Event() for name in ["close", "release-0", "release-1"]} + + class DrainLock(_ObservedScopeLock): + def acquire(self, *args, **kwargs): + if event := progress.get(threading.current_thread().name): + event.set() + return super().acquire(*args, **kwargs) + + try: + assert case.entered.wait(5) + owner = case.sandbox._command_scopes.find("first") + scoped = case.sandbox._scoped_shell_sessions["first"] + scoped.lock = DrainLock(scoped.lock) + if close_first: + workers.append(_start_observed_thread(case.sandbox.close, name="close")) + assert progress["close"].wait(5) + for index in range(release_count): + workers.append(_start_observed_thread(lambda: case.sandbox.release_command_scope("first"), name=f"release-{index}")) + assert progress[f"release-{index}"].wait(5), "Release must join the scope even when close is already draining it" + case.sandbox.request_command_cancel("first") + assert owner.cancelled.is_set(), "Cancellation was routed to a replacement owner instead of the running command" + assert case.sandbox._command_scopes.find("first") is owner + case.resume.set() + for worker in [first, *workers]: + _join_observed_thread(worker) + assert "cancelled" in first.results[0] + assert owner.cleaned.is_set() + assert len(case.controls) == 1 + assert sum(path.endswith("/close") for _, path, _ in case.requests) == 1 + assert case.sandbox._command_scopes.find("first") is None + assert not case.sandbox._scoped_shell_sessions + finally: + case.resume.set() + for worker in [first, *workers]: + _join_observed_thread(worker) + + +def test_cancelling_one_env_scope_preserves_another_scopes_command(paused_env_command): + case = paused_env_command + first = _start_observed_thread(case.run, name="first-env") + second = None + try: + assert case.entered.wait(5) + second = _start_observed_thread(lambda: case.run(owner="second", value="two"), name="second-env") + _join_observed_thread(second) + assert second.results == ["done"] + case.sandbox.request_command_cancel("first") + case.resume.set() + _join_observed_thread(first) + assert "cancelled" in first.results[0] + assert len(case.executions) == 2 + assert len(case.controls) == 1 + assert not case.sandbox.requires_container_recycle + finally: + case.resume.set() + _join_observed_thread(first) + if second is not None: + _join_observed_thread(second) + + +@pytest.mark.asyncio +async def test_repeated_async_cancellation_drains_scoped_env_cleanup(paused_env_command): + case = paused_env_command + task = asyncio.create_task(run_sync_lifecycle_operation(case.run, _on_cancel=lambda: case.sandbox.request_command_cancel("first"))) + try: + assert await asyncio.to_thread(case.entered.wait, 5) + task.cancel() + await asyncio.sleep(0) + task.cancel() + assert not task.done() + assert not case.client_closed.is_set() + case.resume.set() + with pytest.raises(asyncio.CancelledError): + await task + assert len(case.controls) == 1 + assert sum(path.startswith("/v1/bash/sessions/") and path.endswith("/close") for _, path, _ in case.requests) == 1 + assert case.run(owner="second", value="two") == "done" + assert len(case.executions) == 2 + assert not case.sandbox.requires_container_recycle + finally: + case.resume.set() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.parametrize("release_count", [1, 2]) +def test_releasing_env_scope_stays_visible_to_concurrent_close(paused_env_command, monkeypatch, release_count): + case = paused_env_command + first = _start_observed_thread(case.run, name="first-env") + queued = closing = other = None + releases = [] + progress = {f"release-{index}": threading.Event() for index in range(release_count)} + progress["close"] = threading.Event() + + class OverlapLock(_ObservedScopeLock): + def acquire(self, *args, **kwargs): + name = threading.current_thread().name + if name in progress: + progress[name].set() + return super().acquire(*args, **kwargs) + + real_client_close = case.http_client.close + + def client_close(): + # On the unfixed path close skips the popped releasing scope, detaches + # _client, and reaches this callback instead of waiting on its lock. + progress["close"].set() + real_client_close() + + monkeypatch.setattr(case.http_client, "close", client_close) + try: + assert case.entered.wait(5) + scoped = case.sandbox._scoped_shell_sessions["first"] + observed = OverlapLock(scoped.lock) + scoped.lock = observed + # A prior no-env call can leave a persistent shell in the same scope. + scoped.session_id = "retained-shell" + case.controls.add("retained-shell") + queued = _start_observed_thread(case.run, name="queued-env") + assert observed.queued.wait(5) + for index in range(release_count): + release = _start_observed_thread(lambda: case.sandbox.release_command_scope("first"), name=f"release-{index}") + releases.append(release) + assert progress[f"release-{index}"].wait(5) + + assert case.run() == "Error: sandbox command scope is no longer active" + other = _start_observed_thread(lambda: case.run(owner="other", value="two"), name="other-env") + _join_observed_thread(other) + assert other.results == ["done"], "Releasing one scope must not serialize an independent scope" + closing = _start_observed_thread(case.sandbox.close, name="close") + assert progress["close"].wait(5) + assert case.sandbox._client is not None, "close detached the client while a releasing scope still had an admitted command" + assert not case.client_closed.is_set() + assert closing.thread.is_alive() + + case.resume.set() + for worker in (first, queued, *releases, closing): + _join_observed_thread(worker) + assert first.results == ["done"] + assert queued.results == ["Error: sandbox command scope is no longer active"] + assert len(case.executions) == 2 + assert sum(method == "DELETE" and path.endswith("/retained-shell") for method, path, _ in case.requests) == 1 + assert sum(path.endswith("/close") for _, path, _ in case.requests) == 2 + assert not case.sandbox._scoped_shell_sessions + assert case.client_closed.is_set() + finally: + case.resume.set() + for worker in (first, queued, other, *releases, closing): + if worker is not None: + _join_observed_thread(worker) diff --git a/backend/tests/test_github_token_plumbing.py b/backend/tests/test_github_token_plumbing.py index 55c04380ff..63651c71c6 100644 --- a/backend/tests/test_github_token_plumbing.py +++ b/backend/tests/test_github_token_plumbing.py @@ -39,6 +39,7 @@ def _new_aio_sandbox_with_session_state(): from deerflow.community.aio_sandbox.aio_sandbox import AioSandbox, _SessionCreationState sbx = AioSandbox.__new__(AioSandbox) + sbx._bash_slots = threading.BoundedSemaphore(50) sbx._session_creation_state_lock = threading.Lock() sbx._shell_session_creation_state = _SessionCreationState() sbx._bash_session_creation_state = _SessionCreationState() diff --git a/backend/tests/test_provisioner_request_threading.py b/backend/tests/test_provisioner_request_threading.py index 1fb81f9735..2b73f43bbb 100644 --- a/backend/tests/test_provisioner_request_threading.py +++ b/backend/tests/test_provisioner_request_threading.py @@ -103,6 +103,37 @@ def test_provisioner_threads_shell_capacity_into_sandbox_pod( assert env["MAX_SHELL_SESSIONS"] == "13" +@pytest.mark.parametrize("capacity", [1, 50, 72]) +def test_provisioner_attests_bash_capacity_from_created_and_reused_pod(monkeypatch, provisioner_module, capacity): + core = _RecordingCoreV1(event_loop_thread_id=-1) + monkeypatch.setattr(provisioner_module, "core_v1", core) + request = provisioner_module.CreateSandboxRequest(sandbox_id="bash-capacity", max_bash_sessions=capacity) + response = provisioner_module.create_sandbox(request) + assert response.max_bash_sessions == capacity + pod = core.created_pod_specs["bash-capacity"] + assert {item.name: item.value for item in pod.spec.containers[0].env}["AIO_BASH_MAX_SESSIONS"] == str(capacity) + assert provisioner_module.create_sandbox(request).max_bash_sessions == capacity + assert len(core.created_pods) == 1 + with pytest.raises(provisioner_module.HTTPException) as error: + provisioner_module.create_sandbox(request.model_copy(update={"max_bash_sessions": capacity + 1})) + assert error.value.status_code == 409 + assert len(core.created_pods) == 1 + + +def test_bash_capacity_already_exists_race_reads_winning_pod(monkeypatch, provisioner_module): + core = _RecordingCoreV1(event_loop_thread_id=-1) + monkeypatch.setattr(provisioner_module, "core_v1", core) + # First observation misses the existing Service. The Pod creation loses + # to a peer's default-capacity Pod, so requested capacity is not evidence. + reads = iter([None, "http://sandbox"]) + monkeypatch.setattr(provisioner_module, "_sandbox_access_url", lambda *_a, **_kw: next(reads)) + monkeypatch.setattr(core, "create_namespaced_pod", MagicMock(side_effect=ApiException(status=409))) + with pytest.raises(provisioner_module.HTTPException) as error: + provisioner_module.create_sandbox(provisioner_module.CreateSandboxRequest(sandbox_id="sandbox-existing", max_bash_sessions=72)) + assert error.value.status_code == 409 + assert core.created_pods == [] + + def test_broker_request_without_image_cannot_mount_credentials(monkeypatch, provisioner_module): core = _RecordingCoreV1(event_loop_thread_id=-1) monkeypatch.setattr(provisioner_module, "core_v1", core) diff --git a/backend/tests/test_sandbox_environment_passthrough.py b/backend/tests/test_sandbox_environment_passthrough.py index 8906f6a368..4c6adabbc1 100644 --- a/backend/tests/test_sandbox_environment_passthrough.py +++ b/backend/tests/test_sandbox_environment_passthrough.py @@ -70,7 +70,11 @@ sandbox: loaded = provider._load_config() - assert loaded["environment"] == _RESOLVED + expected = dict(_RESOLVED) + if class_name == "AioSandboxProvider": + expected["AIO_BASH_MAX_SESSIONS"] = "50" + assert loaded["environment"] == expected + assert app_config.sandbox.environment == _RESOLVED @pytest.mark.usefixtures("secret_env") diff --git a/config.example.yaml b/config.example.yaml index bb82121a41..e92f9f1a93 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1853,6 +1853,7 @@ sandbox: # # An explicit value is allowed but must be at least max_running + 1. # # environment: # # MAX_SHELL_SESSIONS: "13" # Example for subagent_runtime.max_running: 12 +# # AIO_BASH_MAX_SESSIONS: "50" # Separate Bash pool; excess env calls wait (cancellable) # # NODE_ENV: production # # DEBUG: "false" # # API_KEY: $MY_API_KEY # Reads from host's MY_API_KEY env var diff --git a/docker/provisioner/app.py b/docker/provisioner/app.py index 7111702741..ba6c454e7b 100644 --- a/docker/provisioner/app.py +++ b/docker/provisioner/app.py @@ -568,6 +568,7 @@ class CreateSandboxRequest(BaseModel): # New Gateways size this from their process-wide subagent capacity. None # keeps requests from older Gateways and custom provisioner callers working. max_shell_sessions: int | None = Field(default=None, gt=0) + max_bash_sessions: int | None = Field(default=None, gt=0) class SandboxResponse(BaseModel): @@ -575,6 +576,7 @@ class SandboxResponse(BaseModel): sandbox_url: str status: str max_shell_sessions: int | None = None + max_bash_sessions: int | None = None container_id: str | None = None lark_cli_broker: bool | None = None @@ -1039,6 +1041,7 @@ def _build_pod( provision_lark_cli_runtime: bool = False, provision_lark_cli_broker: bool = False, max_shell_sessions: int | None = None, + max_bash_sessions: int | None = None, ) -> k8s_client.V1Pod: """Construct a Pod manifest for a single sandbox.""" _validate_lark_provisioning_request(provision_lark_cli_runtime, provision_lark_cli_broker) @@ -1046,6 +1049,8 @@ def _build_pod( _build_lark_cli_init_containers(provision_lark_cli_runtime, provision_lark_cli_broker) or None ) sandbox_env: list[k8s_client.V1EnvVar] = [] + if max_bash_sessions is not None: + sandbox_env.append(k8s_client.V1EnvVar(name="AIO_BASH_MAX_SESSIONS", value=str(max_bash_sessions))) if max_shell_sessions is not None: sandbox_env.append( k8s_client.V1EnvVar( @@ -1227,17 +1232,37 @@ def _pod_shell_capacity(pod) -> int: return DEFAULT_MAX_SHELL_SESSIONS +def _pod_bash_capacity(pod) -> int: + containers = getattr(getattr(pod, "spec", None), "containers", None) or [] + sandbox_container = next((container for container in containers if getattr(container, "name", None) == "sandbox"), None) + for env_var in getattr(sandbox_container, "env", None) or []: + if getattr(env_var, "name", None) != "AIO_BASH_MAX_SESSIONS": + continue + try: + value = int(env_var.value) + except (TypeError, ValueError) as exc: + raise RuntimeError("Persisted sandbox has an invalid AIO_BASH_MAX_SESSIONS value") from exc + if value <= 0: + raise RuntimeError("Persisted sandbox has an invalid AIO_BASH_MAX_SESSIONS value") + return value + return 50 + + def _sandbox_response( sandbox_id: str, sandbox_url: str, *, required_capacity: int | None = None, + required_bash_capacity: int | None = None, required_broker: bool | None = None, ) -> SandboxResponse: # Attest mode, generation and capacity from the same Pod observation, # including idempotent create and the AlreadyExists race. pod = core_v1.read_namespaced_pod(_pod_name(sandbox_id), K8S_NAMESPACE) capacity = _pod_shell_capacity(pod) + bash_capacity = _pod_bash_capacity(pod) + if required_bash_capacity is not None and bash_capacity < required_bash_capacity: + raise HTTPException(status_code=409, detail="Existing sandbox Bash capacity is below the requested value; replacement must be coordinated by the Gateway") containers = getattr(getattr(pod, "spec", None), "containers", None) or [] init_containers = getattr(getattr(pod, "spec", None), "init_containers", None) or [] sidecar = any(container.name == "lark-cli-broker" for container in containers) @@ -1267,13 +1292,14 @@ def _sandbox_response( # Terminating Pod; status/get/list callers pass no requirements and read # the phase as data — raising here would fail the whole fleet listing for # one terminating Pod (HTTPException escapes their narrower catch clause). - if (required_capacity is not None or required_broker is not None) and status == "Terminating": + if (required_capacity is not None or required_bash_capacity is not None or required_broker is not None) and status == "Terminating": raise HTTPException(status_code=409, detail="Existing sandbox is still terminating") return SandboxResponse( sandbox_id=sandbox_id, sandbox_url=sandbox_url, status=status, max_shell_sessions=capacity, + max_bash_sessions=bash_capacity, container_id=getattr(metadata, "uid", None), lark_cli_broker=broker, ) @@ -1318,6 +1344,7 @@ def create_sandbox(req: CreateSandboxRequest): provision_lark_cli_runtime = req.provision_lark_cli_runtime provision_lark_cli_broker = req.provision_lark_cli_broker max_shell_sessions = req.max_shell_sessions + max_bash_sessions = req.max_bash_sessions _validate_lark_provisioning_request(provision_lark_cli_runtime, provision_lark_cli_broker) required_broker = provision_lark_cli_broker if provision_lark_cli_runtime or provision_lark_cli_broker else None @@ -1341,6 +1368,7 @@ def create_sandbox(req: CreateSandboxRequest): sandbox_id, existing_url, required_capacity=max_shell_sessions, + required_bash_capacity=max_bash_sessions, required_broker=required_broker, ) except (ApiException, RuntimeError) as exc: @@ -1362,6 +1390,7 @@ def create_sandbox(req: CreateSandboxRequest): provision_lark_cli_runtime=provision_lark_cli_runtime, provision_lark_cli_broker=provision_lark_cli_broker, max_shell_sessions=max_shell_sessions, + max_bash_sessions=max_bash_sessions, ), ) logger.info(f"Created Pod {_pod_name(sandbox_id)}") @@ -1400,6 +1429,7 @@ def create_sandbox(req: CreateSandboxRequest): sandbox_id, sandbox_url, required_capacity=max_shell_sessions, + required_bash_capacity=max_bash_sessions, required_broker=required_broker, ) except (ApiException, RuntimeError) as exc: