Merge pull request #954 from langchain-ai/yogesh/SANDBOX_BACKENDS-metadata

fix: restore sandbox backend via thread metadata when cache misses

apps/agent/agent/middleware/open_pr.py

@@ -32,7 +32,7 @@ from ..utils.github import (
     git_push,
 )
 from ..utils.linear import comment_on_linear_issue
-from ..utils.sandbox_state import SANDBOX_BACKENDS
+from ..utils.sandbox_state import get_sandbox_backend
 
 logger = logging.getLogger(__name__)
 
@@ -141,13 +141,13 @@ I've {action} pull request to address this issue:
 
 {last_message_content}"""
                 await comment_on_linear_issue(linear_issue_id, comment)
-            return None
+            raise ValueError("No thread_id found in config")
 
         repo_config = configurable.get("repo", {})
         repo_owner = repo_config.get("owner")
         repo_name = repo_config.get("name")
 
-        sandbox_backend = SANDBOX_BACKENDS.get(thread_id)
+        sandbox_backend = await get_sandbox_backend(thread_id)
 
         repo_dir = f"/workspace/{repo_name}"
 

apps/agent/agent/server.py

@@ -9,6 +9,7 @@ import warnings
 
 logger = logging.getLogger(__name__)
 
+from langgraph.config import get_config
 from langgraph.graph.state import RunnableConfig
 from langgraph.pregel import Pregel
 from langgraph_sdk import get_client
@@ -45,7 +46,7 @@ SANDBOX_CREATING = "__creating__"
 SANDBOX_CREATION_TIMEOUT = 180
 SANDBOX_POLL_INTERVAL = 1.0
 
-from .utils.sandbox_state import SANDBOX_BACKENDS
+from .utils.sandbox_state import SANDBOX_BACKENDS, get_sandbox_id_from_metadata
 from .utils.github import (
     git_has_uncommitted_changes,
     is_valid_git_repo,
@@ -173,12 +174,6 @@ async def _clone_or_pull_repo_in_sandbox(  # noqa: PLR0915
     return repo_dir
 
 
-async def _get_sandbox_id_from_metadata(thread_id: str) -> str | None:
-    """Get sandbox_id from thread metadata."""
-    thread = await client.threads.get(thread_id=thread_id)
-    return thread.get("metadata", {}).get("sandbox_id")
-
-
 async def _wait_for_sandbox_id(thread_id: str) -> str:
     """Wait for sandbox_id to be set in thread metadata.
 
@@ -190,7 +185,7 @@ async def _wait_for_sandbox_id(thread_id: str) -> str:
     """
     elapsed = 0.0
     while elapsed < SANDBOX_CREATION_TIMEOUT:
-        sandbox_id = await _get_sandbox_id_from_metadata(thread_id)
+        sandbox_id = await get_sandbox_id_from_metadata(thread_id)
         if sandbox_id is not None and sandbox_id != SANDBOX_CREATING:
             return sandbox_id
         await asyncio.sleep(SANDBOX_POLL_INTERVAL)
@@ -235,13 +230,29 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915
             tools=[],
         ).with_config(config)
 
-    sandbox_id = await _get_sandbox_id_from_metadata(thread_id)
+    sandbox_backend = SANDBOX_BACKENDS.get(thread_id)
+    sandbox_id = await get_sandbox_id_from_metadata(thread_id)
 
-    if sandbox_id == SANDBOX_CREATING:
+    if sandbox_id == SANDBOX_CREATING and not sandbox_backend:
         logger.info("Sandbox creation in progress, waiting...")
         sandbox_id = await _wait_for_sandbox_id(thread_id)
 
-    if sandbox_id is None:
+    if sandbox_backend:
+        logger.info("Using cached sandbox backend for thread %s", thread_id)
+        metadata = get_config().get("metadata", {})
+        repo_dir = metadata.get("repo_dir")
+
+        if repo_owner and repo_name:
+            logger.info("Pulling latest changes for repo %s/%s", repo_owner, repo_name)
+            try:
+                repo_dir = await _clone_or_pull_repo_in_sandbox(
+                    sandbox_backend, repo_owner, repo_name, github_token
+                )
+            except Exception:
+                logger.exception("Failed to pull repo in cached sandbox")
+                raise
+
+    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})
 
@@ -304,8 +315,8 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915
                 await client.threads.update(thread_id=thread_id, metadata={"sandbox_id": None})
                 raise
 
-        thread = await client.threads.get(thread_id=thread_id)
-        repo_dir = thread.get("metadata", {}).get("repo_dir")
+        metadata = get_config().get("metadata", {})
+        repo_dir = metadata.get("repo_dir")
 
         if repo_owner and repo_name:
             logger.info("Pulling latest changes for repo %s/%s", repo_owner, repo_name)

apps/agent/agent/tools/commit_and_open_pr.py

@@ -5,7 +5,6 @@ from typing import Any
 from langgraph.config import get_config
 
 from ..encryption import decrypt_token
-from ..integrations.langsmith import _create_langsmith_sandbox
 from ..utils.github import (
     create_github_pr,
     get_github_default_branch,
@@ -19,7 +18,7 @@ from ..utils.github import (
     git_has_unpushed_commits,
     git_push,
 )
-from ..utils.sandbox_state import SANDBOX_BACKENDS
+from ..utils.sandbox_state import get_sandbox_backend_sync
 
 logger = logging.getLogger(__name__)
 
@@ -128,15 +127,9 @@ def commit_and_open_pr(
                 "pr_url": None,
             }
 
-        sandbox_backend = SANDBOX_BACKENDS.get(thread_id)
+        sandbox_backend = get_sandbox_backend_sync(thread_id)
         if not sandbox_backend:
-            sandbox_id = configurable.get("sandbox_id")
-
-            if not sandbox_id:
-                return {"success": False, "error": "No sandbox found for thread", "pr_url": None}
-
-            sandbox_backend = _create_langsmith_sandbox(sandbox_id)
-            SANDBOX_BACKENDS[thread_id] = sandbox_backend
+            return {"success": False, "error": "No sandbox found for thread", "pr_url": None}
 
         repo_dir = f"/workspace/{repo_name}"
 

apps/agent/agent/utils/sandbox_state.py

@@ -2,7 +2,45 @@
 
 from __future__ import annotations
 
+import asyncio
+import logging
 from typing import Any
 
+from langgraph.config import get_config
+
+from ..integrations.langsmith import _create_langsmith_sandbox
+
+logger = logging.getLogger(__name__)
+
 # Thread ID -> SandboxBackend mapping, shared between server.py and middleware
 SANDBOX_BACKENDS: dict[str, Any] = {}
+
+
+async def get_sandbox_id_from_metadata(thread_id: str) -> str | None:
+    """Fetch sandbox_id from thread metadata."""
+    try:
+        config = get_config()
+    except Exception:
+        logger.exception("Failed to read thread metadata for sandbox")
+        return None
+    return config.get("metadata", {}).get("sandbox_id")
+
+
+async def get_sandbox_backend(thread_id: str) -> Any | None:
+    """Get sandbox backend from cache, or connect using thread metadata."""
+    sandbox_backend = SANDBOX_BACKENDS.get(thread_id)
+    if sandbox_backend:
+        return sandbox_backend
+
+    sandbox_id = await get_sandbox_id_from_metadata(thread_id)
+    if not sandbox_id:
+        raise ValueError(f"Missing sandbox_id in thread metadata for {thread_id}")
+
+    sandbox_backend = await asyncio.to_thread(_create_langsmith_sandbox, sandbox_id)
+    SANDBOX_BACKENDS[thread_id] = sandbox_backend
+    return sandbox_backend
+
+
+def get_sandbox_backend_sync(thread_id: str) -> Any | None:
+    """Sync wrapper for get_sandbox_backend."""
+    return asyncio.run(get_sandbox_backend(thread_id))