diff --git a/agent/integrations/langsmith.py b/agent/integrations/langsmith.py index c0bc9419..794f4e94 100644 --- a/agent/integrations/langsmith.py +++ b/agent/integrations/langsmith.py @@ -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.""" diff --git a/tests/test_langsmith_sandbox_timeout.py b/tests/test_langsmith_sandbox_timeout.py new file mode 100644 index 00000000..8179bb4e --- /dev/null +++ b/tests/test_langsmith_sandbox_timeout.py @@ -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"