From bec922b3b1aa4fae487acc4cf8097f60a8a69955 Mon Sep 17 00:00:00 2001
From: aran-yogesh <yogesh.mahendran@langchain.dev>
Date: Thu, 26 Feb 2026 09:52:27 -0800
Subject: [PATCH] refactor: extract sandbox recovery logic into
 _recreate_sandbox util function

---
 apps/agent/agent/server.py | 51 +++++++++++++++++++++++++-------------
 1 file changed, 34 insertions(+), 17 deletions(-)

diff --git a/apps/agent/agent/server.py b/apps/agent/agent/server.py
index 31d6279d..86965df4 100644
--- a/apps/agent/agent/server.py
+++ b/apps/agent/agent/server.py
@@ -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