From 88d3659d00a62a6711f306949667eb062971099c Mon Sep 17 00:00:00 2001 From: Johannes du Plessis Date: Thu, 7 May 2026 11:26:25 -0700 Subject: [PATCH] start sandbox before proxy refresh (#1249) --- agent/server.py | 59 ++++++++++++++++++++++++-- tests/test_proxy_auth.py | 92 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 148 insertions(+), 3 deletions(-) diff --git a/agent/server.py b/agent/server.py index b201c187..12e6ad65 100644 --- a/agent/server.py +++ b/agent/server.py @@ -22,6 +22,7 @@ warnings.filterwarnings("ignore", message=".*Pydantic V1.*", category=UserWarnin # Now safe to import agent (which imports LangChain modules) from deepagents import create_deep_agent +from deepagents.backends import LangSmithSandbox from deepagents.backends.protocol import SandboxBackendProtocol from langchain.agents.middleware import ModelCallLimitMiddleware from langsmith.sandbox import SandboxClientError @@ -65,6 +66,29 @@ SANDBOX_POLL_INTERVAL = 1.0 from .utils.sandbox_state import SANDBOX_BACKENDS, get_sandbox_id_from_metadata +async def _start_langsmith_sandbox_if_needed(sandbox_backend: SandboxBackendProtocol) -> None: + """Start a LangSmith sandbox before operations that require it to be running.""" + if os.getenv("SANDBOX_TYPE", "langsmith") != "langsmith": + return + if not isinstance(sandbox_backend, LangSmithSandbox): + return + + sandbox = sandbox_backend._sandbox # noqa: SLF001 + status = await asyncio.to_thread(sandbox._client.get_sandbox_status, sandbox.name) # noqa: SLF001 + status_name = getattr(status, "status", status) + status_name = getattr(status_name, "value", status_name) + status_text = str(status_name or "").lower() + if status_text in {"running", "ready"}: + return + + logger.info( + "Starting LangSmith sandbox %s before proxy refresh (status=%s)", + sandbox_backend.id, + status_text or "unknown", + ) + await asyncio.to_thread(sandbox.start) + + async def _create_sandbox_with_proxy() -> SandboxBackendProtocol: """Create a new sandbox with GitHub proxy auth configured. @@ -80,6 +104,7 @@ async def _create_sandbox_with_proxy() -> SandboxBackendProtocol: msg = "Cannot configure proxy: GitHub App installation token is unavailable" logger.error(msg) raise ValueError(msg) + await _start_langsmith_sandbox_if_needed(sandbox_backend) await asyncio.to_thread(_configure_github_proxy, sandbox_backend.id, installation_token) return sandbox_backend @@ -100,9 +125,28 @@ async def _refresh_github_proxy( ) return + await _start_langsmith_sandbox_if_needed(sandbox_backend) await asyncio.to_thread(_configure_github_proxy, sandbox_backend.id, installation_token) +async def _refresh_github_proxy_or_recreate( + sandbox_backend: SandboxBackendProtocol, + thread_id: str, +) -> SandboxBackendProtocol: + """Refresh proxy credentials, recreating stale LangSmith sandboxes on failure.""" + try: + await _refresh_github_proxy(sandbox_backend) + except Exception: # noqa: BLE001 + logger.warning( + "Failed to refresh GitHub proxy for sandbox %s on thread %s, recreating sandbox", + sandbox_backend.id, + thread_id, + exc_info=True, + ) + return await _recreate_sandbox(thread_id) + return sandbox_backend + + async def _recreate_sandbox(thread_id: str) -> SandboxBackendProtocol: """Recreate a sandbox after a connection failure. @@ -199,8 +243,10 @@ async def ensure_sandbox_for_thread(thread_id: str) -> SandboxBackendProtocol: if sandbox_backend: logger.info("Using cached sandbox backend for thread %s", thread_id) - await _refresh_github_proxy(sandbox_backend) + original_sandbox_id = sandbox_backend.id sandbox_backend = await check_or_recreate_sandbox(sandbox_backend, thread_id) + if sandbox_backend.id == original_sandbox_id: + sandbox_backend = await _refresh_github_proxy_or_recreate(sandbox_backend, thread_id) elif sandbox_id is None: logger.info("Creating new sandbox for thread %s", thread_id) await client.threads.update(thread_id=thread_id, metadata={"sandbox_id": SANDBOX_CREATING}) @@ -216,6 +262,7 @@ async def ensure_sandbox_for_thread(thread_id: str) -> SandboxBackendProtocol: raise else: logger.info("Connecting to existing sandbox %s", sandbox_id) + created_replacement_sandbox = False try: sandbox_backend = await asyncio.to_thread(create_sandbox, sandbox_id) except Exception: @@ -225,12 +272,18 @@ async def ensure_sandbox_for_thread(thread_id: str) -> SandboxBackendProtocol: ) try: sandbox_backend = await _create_sandbox_with_proxy() + created_replacement_sandbox = True except Exception: logger.exception("Failed to create replacement sandbox") await client.threads.update(thread_id=thread_id, metadata={"sandbox_id": None}) raise - await _refresh_github_proxy(sandbox_backend) - sandbox_backend = await check_or_recreate_sandbox(sandbox_backend, thread_id) + if not created_replacement_sandbox: + original_sandbox_id = sandbox_backend.id + sandbox_backend = await check_or_recreate_sandbox(sandbox_backend, thread_id) + if sandbox_backend.id == original_sandbox_id: + sandbox_backend = await _refresh_github_proxy_or_recreate( + sandbox_backend, thread_id + ) SANDBOX_BACKENDS[thread_id] = sandbox_backend diff --git a/tests/test_proxy_auth.py b/tests/test_proxy_auth.py index 3d2c31d6..abf423ac 100644 --- a/tests/test_proxy_auth.py +++ b/tests/test_proxy_auth.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from deepagents.backends import LangSmithSandbox from agent.integrations.langsmith import _configure_github_proxy @@ -289,3 +290,94 @@ class TestRefreshProxyOnSandboxReuse: mock_create.assert_called_once_with("sandbox-existing") mock_proxy.assert_called_once_with("sandbox-existing", "ghs_fresh") + + @pytest.mark.asyncio + async def test_proxy_refresh_failure_recreates_sandbox(self) -> None: + """A stale sandbox whose proxy cannot be patched should be replaced.""" + mock_sandbox = MagicMock(id="sandbox-stale") + replacement_sandbox = MagicMock(id="sandbox-replacement") + request = httpx.Request( + "PATCH", "https://api.smith.langchain.com/v2/sandboxes/boxes/sandbox-stale" + ) + response = httpx.Response(400, request=request) + + with ( + patch( + "agent.server.get_github_app_installation_token", + new_callable=AsyncMock, + return_value="ghs_fresh", + ), + patch( + "agent.server._configure_github_proxy", + side_effect=httpx.HTTPStatusError( + "Bad request", + request=request, + response=response, + ), + ) as mock_proxy, + patch( + "agent.server._recreate_sandbox", + new_callable=AsyncMock, + return_value=replacement_sandbox, + ) as mock_recreate, + patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}), + ): + from agent.server import _refresh_github_proxy_or_recreate + + sandbox = await _refresh_github_proxy_or_recreate(mock_sandbox, "thread-123") + + assert sandbox is replacement_sandbox + mock_proxy.assert_called_once_with("sandbox-stale", "ghs_fresh") + mock_recreate.assert_awaited_once_with("thread-123") + + @pytest.mark.asyncio + async def test_starts_stopped_langsmith_sandbox_before_proxy_refresh(self) -> None: + """Proxy config requires a running LangSmith sandbox.""" + inner_sandbox = MagicMock(name="sandbox-stopped") + inner_sandbox.name = "sandbox-stopped" + inner_sandbox._client.get_sandbox_status.return_value = MagicMock(status="stopped") + sandbox_backend = object.__new__(LangSmithSandbox) + sandbox_backend._sandbox = inner_sandbox + + with ( + patch( + "agent.server.get_github_app_installation_token", + new_callable=AsyncMock, + return_value="ghs_fresh", + ), + patch("agent.server._configure_github_proxy") as mock_proxy, + patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}), + ): + from agent.server import _refresh_github_proxy + + await _refresh_github_proxy(sandbox_backend) + + inner_sandbox._client.get_sandbox_status.assert_called_once_with("sandbox-stopped") + inner_sandbox.start.assert_called_once_with() + mock_proxy.assert_called_once_with("sandbox-stopped", "ghs_fresh") + + @pytest.mark.asyncio + async def test_skips_start_for_ready_langsmith_sandbox_before_proxy_refresh(self) -> None: + """Ready sandboxes can be patched without starting again.""" + inner_sandbox = MagicMock(name="sandbox-ready") + inner_sandbox.name = "sandbox-ready" + inner_sandbox._client.get_sandbox_status.return_value = MagicMock(status="ready") + sandbox_backend = object.__new__(LangSmithSandbox) + sandbox_backend._sandbox = inner_sandbox + + with ( + patch( + "agent.server.get_github_app_installation_token", + new_callable=AsyncMock, + return_value="ghs_fresh", + ), + patch("agent.server._configure_github_proxy") as mock_proxy, + patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}), + ): + from agent.server import _refresh_github_proxy + + await _refresh_github_proxy(sandbox_backend) + + inner_sandbox._client.get_sandbox_status.assert_called_once_with("sandbox-ready") + inner_sandbox.start.assert_not_called() + mock_proxy.assert_called_once_with("sandbox-ready", "ghs_fresh")