mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-09-30 08:03:15 +00:00
fix: enforce client-side deadline on sandbox execute (#1451)
* fix: enforce client-side deadline on sandbox execute The langsmith SDK's default execute path is now a WebSocket stream with no client-side read deadline. On a live socket where the dataplane never emits an exit/error frame, CommandHandle.result blocks forever in an uncancellable thread, wedging the run (introduced by the langsmith 0.8.3 -> 0.8.8 bump in #1385). The command `timeout` is only enforced server-side, so it doesn't fire. TimeoutLangSmithSandbox drives a non-blocking CommandHandle and kills the command if it overruns its timeout by a grace window (SANDBOX_EXECUTE_CLIENT_GRACE_SECONDS, default 30), returning a timed-out tool result instead of hanging. WS connect failures fall back to the base wait=True path, whose HTTP fallback carries its own request deadline. * fix: fall back to HTTP when WS execute connect fails run(wait=False) eagerly opens the WebSocket and reads the "started" frame, so connect/setup failures (and connect timeouts) raise from the run() call itself, not from handle.result. The previous structure left run() outside the try, so those failures bypassed the HTTP fallback and would fail every sandbox command in any environment where the WS path is unavailable. Move handle creation inside the fallback handler in both execute and aexecute, and run it via to_thread in the async path since it now blocks on connect. Add tests for connect-failure and connect-timeout fallback. --------- Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
Johannes du Plessis
GitHub
parent
ec41f138fa
commit
903c327e4c
2 changed files with 285 additions and 4 deletions
|
|
@ -2,17 +2,25 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from concurrent.futures import TimeoutError as FuturesTimeout
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from deepagents.backends import LangSmithSandbox
|
||||
from deepagents.backends.protocol import SandboxBackendProtocol
|
||||
from langsmith.sandbox import SandboxClient
|
||||
from deepagents.backends.protocol import ExecuteResponse, SandboxBackendProtocol
|
||||
from langsmith.sandbox import (
|
||||
CommandTimeoutError,
|
||||
SandboxClient,
|
||||
SandboxConnectionError,
|
||||
SandboxServerReloadError,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -47,6 +55,13 @@ def _parse_optional_int(name: str, default: int) -> int:
|
|||
raise ValueError(msg) from e
|
||||
|
||||
|
||||
def _execute_client_grace_seconds() -> int:
|
||||
"""Extra wall-clock seconds the client waits past a command's own timeout
|
||||
before giving up and killing it. The server is meant to enforce the command
|
||||
timeout; this is the client-side backstop for when it doesn't."""
|
||||
return _parse_optional_int("SANDBOX_EXECUTE_CLIENT_GRACE_SECONDS", 30)
|
||||
|
||||
|
||||
def _get_sandbox_snapshot_config() -> tuple[str | None, int, int, int, int, int]:
|
||||
"""Get sandbox snapshot configuration from environment."""
|
||||
snapshot_id = os.environ.get("DEFAULT_SANDBOX_SNAPSHOT_ID")
|
||||
|
|
@ -250,6 +265,121 @@ def _update_thread_sandbox_metadata(sandbox_id: str) -> None:
|
|||
pass
|
||||
|
||||
|
||||
class TimeoutLangSmithSandbox(LangSmithSandbox):
|
||||
"""LangSmith backend that enforces a client-side execution deadline.
|
||||
|
||||
The langsmith SDK's default execute path is now a WebSocket stream with no
|
||||
client-side read deadline: on a live socket where the dataplane never emits
|
||||
an exit/error frame, ``CommandHandle.result`` blocks forever and wedges the
|
||||
run (the blocking call sits in a thread that cancellation can't reclaim).
|
||||
|
||||
We drive a non-blocking ``CommandHandle`` ourselves and, if the command
|
||||
overruns its own timeout by the grace window, kill it and surface a
|
||||
timed-out tool result instead of hanging the graph. WebSocket connect
|
||||
failures fall back to the base wait=True path, whose HTTP fallback carries
|
||||
its own request deadline.
|
||||
"""
|
||||
|
||||
_WS_FALLBACK_ERRORS = (
|
||||
SandboxConnectionError,
|
||||
SandboxServerReloadError,
|
||||
ImportError,
|
||||
OSError,
|
||||
TypeError,
|
||||
)
|
||||
|
||||
def _deadline(self, effective_timeout: int) -> int:
|
||||
return effective_timeout + _execute_client_grace_seconds()
|
||||
|
||||
@staticmethod
|
||||
def _result_to_response(result: Any) -> ExecuteResponse:
|
||||
output = result.stdout or ""
|
||||
if result.stderr:
|
||||
output += "\n" + result.stderr if output else result.stderr
|
||||
return ExecuteResponse(output=output, exit_code=result.exit_code, truncated=False)
|
||||
|
||||
@staticmethod
|
||||
def _timeout_response(seconds: int, *, server_side: bool) -> ExecuteResponse:
|
||||
where = "on the sandbox" if server_side else "by the client and killed"
|
||||
return ExecuteResponse(
|
||||
output=f"Command timed out after {seconds}s {where}.",
|
||||
exit_code=124,
|
||||
truncated=False,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _safe_kill(handle: Any) -> None:
|
||||
try:
|
||||
handle.kill()
|
||||
except Exception: # noqa: BLE001 - best-effort cleanup of a wedged command
|
||||
logger.warning("Failed to kill timed-out sandbox command", exc_info=True)
|
||||
|
||||
def _base_execute(self, command: str, timeout: int | None) -> ExecuteResponse:
|
||||
# WS path unavailable; the base wait=True path falls back to HTTP,
|
||||
# which carries its own request deadline.
|
||||
return LangSmithSandbox.execute(self, command, timeout=timeout)
|
||||
|
||||
def execute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse:
|
||||
effective = timeout if timeout is not None else self._default_timeout
|
||||
if not effective: # 0 / None: caller opted out of any deadline
|
||||
return super().execute(command, timeout=timeout)
|
||||
# run(wait=False) eagerly opens the WS and reads the "started" frame, so
|
||||
# connect/setup failures raise here — fall back to the base path.
|
||||
try:
|
||||
handle = self._sandbox.run(command, timeout=effective, wait=False)
|
||||
except (*self._WS_FALLBACK_ERRORS, TimeoutError):
|
||||
return self._base_execute(command, timeout)
|
||||
deadline = self._deadline(effective)
|
||||
pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sbx-exec")
|
||||
try:
|
||||
future = pool.submit(lambda: handle.result)
|
||||
try:
|
||||
result = future.result(timeout=deadline)
|
||||
except FuturesTimeout:
|
||||
self._safe_kill(handle)
|
||||
return self._timeout_response(deadline, server_side=False)
|
||||
except CommandTimeoutError:
|
||||
return self._timeout_response(effective, server_side=True)
|
||||
except self._WS_FALLBACK_ERRORS:
|
||||
return self._base_execute(command, timeout)
|
||||
return self._result_to_response(result)
|
||||
finally:
|
||||
# Never join: a still-wedged worker must not block the caller.
|
||||
pool.shutdown(wait=False)
|
||||
|
||||
async def aexecute(
|
||||
self,
|
||||
command: str,
|
||||
*,
|
||||
timeout: int | None = None, # noqa: ASYNC109 - forwarded semantic timeout, not an asyncio contract
|
||||
) -> ExecuteResponse:
|
||||
effective = timeout if timeout is not None else self._default_timeout
|
||||
if not effective:
|
||||
return await super().aexecute(command, timeout=timeout)
|
||||
# run(wait=False) eagerly opens the WS and reads the "started" frame
|
||||
# (blocking, bounded by the SDK connect timeout); connect/setup failures
|
||||
# raise here — fall back to the base path.
|
||||
try:
|
||||
handle = await asyncio.to_thread(
|
||||
self._sandbox.run, command, timeout=effective, wait=False
|
||||
)
|
||||
except (*self._WS_FALLBACK_ERRORS, TimeoutError):
|
||||
return await asyncio.to_thread(self._base_execute, command, timeout)
|
||||
deadline = self._deadline(effective)
|
||||
try:
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.to_thread(lambda: handle.result), timeout=deadline
|
||||
)
|
||||
except TimeoutError:
|
||||
await asyncio.to_thread(self._safe_kill, handle)
|
||||
return self._timeout_response(deadline, server_side=False)
|
||||
except CommandTimeoutError:
|
||||
return self._timeout_response(effective, server_side=True)
|
||||
except self._WS_FALLBACK_ERRORS:
|
||||
return await asyncio.to_thread(self._base_execute, command, timeout)
|
||||
return self._result_to_response(result)
|
||||
|
||||
|
||||
class SandboxProvider(ABC):
|
||||
"""Interface for creating and deleting sandbox backends."""
|
||||
|
||||
|
|
@ -341,7 +471,7 @@ class LangSmithProvider(SandboxProvider):
|
|||
except Exception as e:
|
||||
msg = f"Failed to connect to existing sandbox '{sandbox_id}': {e}"
|
||||
raise RuntimeError(msg) from e
|
||||
return LangSmithSandbox(sandbox)
|
||||
return TimeoutLangSmithSandbox(sandbox)
|
||||
|
||||
if not snapshot_id:
|
||||
msg = "DEFAULT_SANDBOX_SNAPSHOT_ID must be set when SANDBOX_TYPE=langsmith"
|
||||
|
|
@ -361,7 +491,7 @@ class LangSmithProvider(SandboxProvider):
|
|||
msg = f"Failed to create sandbox from snapshot '{snapshot_id}': {e}"
|
||||
raise RuntimeError(msg) from e
|
||||
|
||||
return LangSmithSandbox(sandbox)
|
||||
return TimeoutLangSmithSandbox(sandbox)
|
||||
|
||||
def delete(self, *, sandbox_id: str, **kwargs: Any) -> None:
|
||||
"""Delete a LangSmith sandbox."""
|
||||
|
|
|
|||
151
tests/test_langsmith_sandbox_timeout.py
Normal file
151
tests/test_langsmith_sandbox_timeout.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
"""Client-side execution deadline for the LangSmith sandbox backend."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langsmith.sandbox import CommandTimeoutError, SandboxConnectionError
|
||||
|
||||
from agent.integrations.langsmith import TimeoutLangSmithSandbox
|
||||
|
||||
|
||||
class _FakeHandle:
|
||||
def __init__(self, *, sleep: float = 0.0, result: Any = None, raises: Exception | None = None):
|
||||
self._sleep = sleep
|
||||
self._result = result
|
||||
self._raises = raises
|
||||
self.killed = False
|
||||
|
||||
@property
|
||||
def result(self) -> Any:
|
||||
if self._sleep:
|
||||
time.sleep(self._sleep)
|
||||
if self._raises is not None:
|
||||
raise self._raises
|
||||
return self._result
|
||||
|
||||
def kill(self) -> None:
|
||||
self.killed = True
|
||||
|
||||
|
||||
class _FakeSandbox:
|
||||
def __init__(self, handle: _FakeHandle, *, run_raises: Exception | None = None):
|
||||
self._handle = handle
|
||||
self._run_raises = run_raises
|
||||
self.run_calls: list[dict[str, Any]] = []
|
||||
|
||||
def run(self, command: str, *, timeout: int, wait: bool) -> _FakeHandle:
|
||||
self.run_calls.append({"command": command, "timeout": timeout, "wait": wait})
|
||||
if self._run_raises is not None:
|
||||
raise self._run_raises
|
||||
return self._handle
|
||||
|
||||
|
||||
def _backend(
|
||||
handle: _FakeHandle, *, run_raises: Exception | None = None
|
||||
) -> TimeoutLangSmithSandbox:
|
||||
sb = TimeoutLangSmithSandbox.__new__(TimeoutLangSmithSandbox)
|
||||
sb._sandbox = _FakeSandbox(handle, run_raises=run_raises)
|
||||
sb._default_timeout = 30 * 60
|
||||
return sb
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_grace(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("SANDBOX_EXECUTE_CLIENT_GRACE_SECONDS", "0")
|
||||
|
||||
|
||||
async def test_aexecute_kills_on_client_timeout() -> None:
|
||||
handle = _FakeHandle(sleep=5.0)
|
||||
sb = _backend(handle)
|
||||
start = time.monotonic()
|
||||
resp = await sb.aexecute("sleep 999", timeout=1)
|
||||
assert time.monotonic() - start < 3.0
|
||||
assert resp.exit_code == 124
|
||||
assert "killed" in resp.output
|
||||
assert handle.killed
|
||||
assert sb._sandbox.run_calls[0]["wait"] is False
|
||||
|
||||
|
||||
async def test_aexecute_success_combines_streams() -> None:
|
||||
handle = _FakeHandle(result=SimpleNamespace(stdout="out", stderr="err", exit_code=0))
|
||||
sb = _backend(handle)
|
||||
resp = await sb.aexecute("echo hi", timeout=5)
|
||||
assert resp.exit_code == 0
|
||||
assert resp.output == "out\nerr"
|
||||
assert not handle.killed
|
||||
|
||||
|
||||
async def test_aexecute_server_timeout_not_killed() -> None:
|
||||
handle = _FakeHandle(raises=CommandTimeoutError("server enforced"))
|
||||
sb = _backend(handle)
|
||||
resp = await sb.aexecute("make hang", timeout=2)
|
||||
assert resp.exit_code == 124
|
||||
assert "on the sandbox" in resp.output
|
||||
assert not handle.killed
|
||||
|
||||
|
||||
def test_execute_kills_on_client_timeout() -> None:
|
||||
handle = _FakeHandle(sleep=5.0)
|
||||
sb = _backend(handle)
|
||||
start = time.monotonic()
|
||||
resp = sb.execute("sleep 999", timeout=1)
|
||||
assert time.monotonic() - start < 3.0
|
||||
assert resp.exit_code == 124
|
||||
assert handle.killed
|
||||
|
||||
|
||||
def _patch_base_execute(monkeypatch: pytest.MonkeyPatch, sink: dict[str, Any]) -> None:
|
||||
def fake_base_execute(self: Any, command: str, *, timeout: int | None = None) -> Any:
|
||||
sink["command"] = command
|
||||
sink["timeout"] = timeout
|
||||
return SimpleNamespace(output="via-http", exit_code=0, truncated=False)
|
||||
|
||||
monkeypatch.setattr("agent.integrations.langsmith.LangSmithSandbox.execute", fake_base_execute)
|
||||
|
||||
|
||||
async def test_aexecute_ws_connect_failure_falls_back_to_base(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# run(wait=False) connects eagerly, so a connect failure raises from run().
|
||||
sb = _backend(_FakeHandle(), run_raises=SandboxConnectionError("no ws"))
|
||||
called: dict[str, Any] = {}
|
||||
_patch_base_execute(monkeypatch, called)
|
||||
resp = await sb.aexecute("git status", timeout=5)
|
||||
assert called == {"command": "git status", "timeout": 5}
|
||||
assert resp.output == "via-http"
|
||||
|
||||
|
||||
async def test_aexecute_ws_connect_timeout_falls_back_to_base(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
sb = _backend(_FakeHandle(), run_raises=TimeoutError("connect timed out"))
|
||||
called: dict[str, Any] = {}
|
||||
_patch_base_execute(monkeypatch, called)
|
||||
resp = await sb.aexecute("git status", timeout=5)
|
||||
assert called["command"] == "git status"
|
||||
assert resp.output == "via-http"
|
||||
|
||||
|
||||
def test_execute_ws_connect_failure_falls_back_to_base(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
sb = _backend(_FakeHandle(), run_raises=SandboxConnectionError("no ws"))
|
||||
called: dict[str, Any] = {}
|
||||
_patch_base_execute(monkeypatch, called)
|
||||
resp = sb.execute("git status", timeout=5)
|
||||
assert called == {"command": "git status", "timeout": 5}
|
||||
assert resp.output == "via-http"
|
||||
|
||||
|
||||
async def test_aexecute_midstream_ws_drop_falls_back_to_base(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Connect succeeds (run returns a handle) but the stream drops while draining.
|
||||
sb = _backend(_FakeHandle(raises=SandboxConnectionError("dropped")))
|
||||
called: dict[str, Any] = {}
|
||||
_patch_base_execute(monkeypatch, called)
|
||||
resp = await sb.aexecute("git status", timeout=5)
|
||||
assert called["command"] == "git status"
|
||||
assert resp.output == "via-http"
|
||||
Loading…
Add table
Reference in a new issue