refactor: extract sandbox recovery logic into _recreate_sandbox util function

This commit is contained in:
aran-yogesh 2026-02-26 09:52:27 -08:00
parent 18f1801005
commit bec922b3b1

View file

@ -170,6 +170,38 @@ async def _clone_or_pull_repo_in_sandbox( # noqa: PLR0915
return repo_dir
async def _recreate_sandbox(
thread_id: str,
repo_owner: str,
repo_name: str,
github_token: str | None,
) -> tuple[SandboxBackendProtocol, str]:
"""Recreate a sandbox and clone the repo after a connection failure.
Clears the stale cache entry, sets the SANDBOX_CREATING sentinel,
creates a fresh sandbox, clones the repo, and updates repo_dir metadata.
"""
SANDBOX_BACKENDS.pop(thread_id, None)
await client.threads.update(
thread_id=thread_id,
metadata={"sandbox_id": SANDBOX_CREATING},
)
try:
sandbox_backend = await asyncio.to_thread(create_langsmith_sandbox)
repo_dir = await _clone_or_pull_repo_in_sandbox(
sandbox_backend, repo_owner, repo_name, github_token
)
await client.threads.update(
thread_id=thread_id,
metadata={"repo_dir": repo_dir},
)
except Exception:
logger.exception("Failed to recreate sandbox after connection failure")
await client.threads.update(thread_id=thread_id, metadata={"sandbox_id": None})
raise
return sandbox_backend, repo_dir
async def _wait_for_sandbox_id(thread_id: str) -> str:
"""Wait for sandbox_id to be set in thread metadata.
@ -248,24 +280,9 @@ async def get_agent(config: RunnableConfig) -> Pregel: # noqa: PLR0915
"Cached sandbox is no longer reachable for thread %s, recreating sandbox",
thread_id,
)
SANDBOX_BACKENDS.pop(thread_id, None)
await client.threads.update(
thread_id=thread_id,
metadata={"sandbox_id": SANDBOX_CREATING},
sandbox_backend, repo_dir = await _recreate_sandbox(
thread_id, repo_owner, repo_name, github_token
)
try:
sandbox_backend = await asyncio.to_thread(create_langsmith_sandbox)
repo_dir = await _clone_or_pull_repo_in_sandbox(
sandbox_backend, repo_owner, repo_name, github_token
)
await client.threads.update(
thread_id=thread_id,
metadata={"repo_dir": repo_dir},
)
except Exception:
logger.exception("Failed to recreate sandbox after connection failure")
await client.threads.update(thread_id=thread_id, metadata={"sandbox_id": None})
raise
except Exception:
logger.exception("Failed to pull repo in cached sandbox")
raise