mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-09-30 09:13:14 +00:00
start sandbox before proxy refresh (#1249)
This commit is contained in:
Johannes du Plessis
GitHub
parent
da74342da4
commit
88d3659d00
2 changed files with 148 additions and 3 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue