From 841aa536f5efe7adace15c09be3f83890d78b0d4 Mon Sep 17 00:00:00 2001
From: aran-yogesh <yogesh.mahendran@langchain.dev>
Date: Fri, 6 Feb 2026 10:37:28 -0800
Subject: [PATCH] feat: add configurable sandbox template and fix sandbox
 lifecycle

---
 apps/agent/agent/server.py | 86 ++++++++++++++++++++++++++++++--------
 1 file changed, 68 insertions(+), 18 deletions(-)

diff --git a/apps/agent/agent/server.py b/apps/agent/agent/server.py
index 599317bc..4f3dbd63 100644
--- a/apps/agent/agent/server.py
+++ b/apps/agent/agent/server.py
@@ -28,14 +28,10 @@ import asyncio
 warnings.filterwarnings("ignore", message=".*Pydantic V1.*", category=UserWarning)
 
 # Now safe to import agent (which imports LangChain modules)
-# Async wrapper for create_sandbox
-from contextlib import asynccontextmanager
-
 from deepagents import create_deep_agent
 from deepagents.backends.sandbox import SandboxBackendProtocol
 from deepagents_cli.agent import get_system_prompt
 from deepagents_cli.config import config, settings
-from deepagents_cli.integrations.sandbox_factory import create_sandbox
 from deepagents_cli.tools import fetch_url, http_request, web_search
 
 # Local import for encryption
@@ -51,14 +47,45 @@ def _get_langsmith_api_key() -> str | None:
     return os.environ.get("LANGSMITH_API_KEY") or os.environ.get("LANGSMITH_API_KEY_PROD")
 
 
-@asynccontextmanager
-async def create_sandbox_async(provider: str, **kwargs):
-    """Async wrapper around create_sandbox for compatibility."""
-    # For langsmith provider, pass the API key explicitly
-    if provider == "langsmith" and "api_key" not in kwargs:
-        kwargs["api_key"] = _get_langsmith_api_key()
-    with create_sandbox(provider, **kwargs) as sandbox:
-        yield sandbox
+def _get_sandbox_template_config() -> tuple[str | None, str | None]:
+    """Get sandbox template configuration from environment.
+
+    Returns:
+        Tuple of (template_name, template_image) from environment variables.
+        Values are None if not set in environment.
+    """
+    template_name = os.environ.get("DEFAULT_SANDBOX_TEMPLATE_NAME")
+    template_image = os.environ.get("DEFAULT_SANDBOX_TEMPLATE_IMAGE")
+    return template_name, template_image
+
+
+def _create_langsmith_sandbox(
+    sandbox_id: str | None = None,
+) -> SandboxBackendProtocol:
+    """Create or connect to a LangSmith sandbox without automatic cleanup.
+
+    This function directly uses the LangSmithProvider to create/connect to sandboxes
+    without the context manager cleanup, allowing sandboxes to persist across
+    multiple agent invocations.
+
+    Args:
+        sandbox_id: Optional existing sandbox ID to connect to.
+                   If None, creates a new sandbox.
+
+    Returns:
+        SandboxBackendProtocol instance
+    """
+    from deepagents_cli.integrations.langsmith import LangSmithProvider
+
+    api_key = _get_langsmith_api_key()
+    template_name, template_image = _get_sandbox_template_config()
+
+    provider = LangSmithProvider(api_key=api_key)
+    return provider.get_or_create(
+        sandbox_id=sandbox_id,
+        template=template_name,
+        template_image=template_image,
+    )
 
 
 def create_server_agent(
@@ -876,8 +903,8 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915
         await client.threads.update(thread_id=thread_id, metadata={"sandbox_id": SANDBOX_CREATING})
 
         try:
-            sandbox_cm = create_sandbox_async("langsmith")
-            sandbox_backend = await sandbox_cm.__aenter__()
+            # Create sandbox without context manager cleanup (sandbox persists)
+            sandbox_backend = await asyncio.to_thread(_create_langsmith_sandbox)
             logger.info("Sandbox created: %s", sandbox_backend.id)
 
             # Update metadata immediately after sandbox creation so other callers
@@ -910,12 +937,35 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915
     else:
         logger.info("Connecting to existing sandbox %s", sandbox_id)
         try:
-            sandbox_cm = create_sandbox_async("langsmith", sandbox_id=sandbox_id)
-            sandbox_backend = await sandbox_cm.__aenter__()
+            # Connect to existing sandbox without context manager cleanup
+            sandbox_backend = await asyncio.to_thread(
+                _create_langsmith_sandbox, sandbox_id
+            )
             logger.info("Connected to existing sandbox %s", sandbox_id)
         except Exception:
-            logger.exception("Failed to connect to existing sandbox %s", sandbox_id)
-            raise
+            logger.warning(
+                "Failed to connect to existing sandbox %s, creating new one", sandbox_id
+            )
+            # Reset sandbox_id and create a new sandbox
+            await client.threads.update(
+                thread_id=thread_id,
+                metadata={"sandbox_id": SANDBOX_CREATING},
+            )
+
+            try:
+                sandbox_backend = await asyncio.to_thread(_create_langsmith_sandbox)
+                logger.info("New sandbox created: %s", sandbox_backend.id)
+
+                await client.threads.update(
+                    thread_id=thread_id,
+                    metadata={"sandbox_id": sandbox_backend.id},
+                )
+            except Exception:
+                logger.exception("Failed to create replacement sandbox")
+                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")