refactor: Replace timeout wrapper with tool call arg (#976)

2026-09-30 19:43:15 +00:00 · 2026-02-24 08:14:31 -08:00 · 2026-02-24 08:14:31 -08:00 · 7972ebf481
commit 7972ebf481
parent ff7c52e649
8 changed files with 35 additions and 119 deletions
--- a/apps/agent/agent/integrations/langsmith.py
+++ b/apps/agent/agent/integrations/langsmith.py
@ -43,7 +43,7 @@ def _get_sandbox_template_config() -> tuple[str | None, str | None]:
    return template_name, template_image


-def _create_langsmith_sandbox(
+def create_langsmith_sandbox(
    sandbox_id: str | None = None,
 ) -> SandboxBackendProtocol:
    """Create or connect to a LangSmith sandbox without automatic cleanup.
@ -108,17 +108,28 @@ class LangSmithBackend(BaseSandbox):

    def __init__(self, sandbox: Sandbox) -> None:
        self._sandbox = sandbox
-        self._timeout: int = 30 * 60  # 30 mins default
+        self._default_timeout: int = 30 * 5  # 5 minute default

    @property
    def id(self) -> str:
        """Unique identifier for the sandbox backend."""
        return self._sandbox.name

-    def execute(self, command: str) -> ExecuteResponse:
-        """Execute a command in the sandbox and return ExecuteResponse."""
-        result = self._sandbox.run(command, timeout=self._timeout)
+    def execute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse:
+        """Execute a command in the sandbox and return ExecuteResponse.

+        Args:
+            command: Full shell command string to execute.
+            timeout: Maximum time in seconds to wait for the command to complete.
+                If None, uses the default timeout of 5 minutes.
+
+        Returns:
+            ExecuteResponse with combined output, exit code, and truncation flag.
+        """
+        effective_timeout = timeout if timeout is not None else self._default_timeout
+        result = self._sandbox.run(command, timeout=effective_timeout)
+
+        # Combine stdout and stderr (matching other backends' approach)
        output = result.stdout or ""
        if result.stderr:
            output += "\n" + result.stderr if output else result.stderr
--- a/apps/agent/agent/middleware/init.py
+++ b/apps/agent/agent/middleware/init.py
@ -1,11 +1,9 @@
 from .check_message_queue import check_message_queue_before_model
 from .open_pr import open_pr_if_needed
 from .post_to_linear import post_to_linear_after_model
-from .timeout_execute_tool import TimeoutExecuteToolMiddleware
 from .tool_error_handler import ToolErrorMiddleware

 __all__ = [
-    "TimeoutExecuteToolMiddleware",
    "ToolErrorMiddleware",
    "check_message_queue_before_model",
    "open_pr_if_needed",
--- a/apps/agent/agent/middleware/timeout_execute_tool.py
+++ b/apps/agent/agent/middleware/timeout_execute_tool.py
@ -1,83 +0,0 @@
-"""Tool middleware that wraps execute commands with a timeout."""
-
-from __future__ import annotations
-
-import logging
-import re
-import shlex
-from collections.abc import Awaitable, Callable
-
-from langchain.agents.middleware.types import AgentMiddleware, AgentState
-from langchain_core.messages import ToolMessage
-from langgraph.prebuilt.tool_node import ToolCallRequest
-from langgraph.types import Command
-
-logger = logging.getLogger(__name__)
-
-DEFAULT_TIMEOUT_SECONDS = 300
-TIMEOUT_REGEX = re.compile(r"\btimeout\s+\d+(?:\.\d+)?\s*[smhd]?\b", re.IGNORECASE)
-
-
-def _get_tool_name(request: ToolCallRequest) -> str | None:
-    tool_call = request.tool_call
-    if isinstance(tool_call, dict):
-        return tool_call.get("name")
-    return None
-
-
-def _get_command_arg(request: ToolCallRequest) -> str | None:
-    tool_call = request.tool_call
-    if not isinstance(tool_call, dict):
-        return None
-    args = tool_call.get("args")
-    if not isinstance(args, dict):
-        return None
-    command = args.get("command")
-    return command if isinstance(command, str) else None
-
-
-def _wrap_command(command: str) -> str:
-    if TIMEOUT_REGEX.search(command):
-        return command
-    quoted = shlex.quote(command)
-    return f"timeout {DEFAULT_TIMEOUT_SECONDS}s sh -c {quoted}"
-
-
-def _overwrite_request_if_needed(request: ToolCallRequest) -> ToolCallRequest:
-    if _get_tool_name(request) != "execute":
-        return request
-
-    command = _get_command_arg(request)
-    if not command:
-        return request
-
-    wrapped = _wrap_command(command)
-    if wrapped == command:
-        return request
-
-    tool_call = dict(request.tool_call)
-    args = dict(tool_call.get("args", {}))
-    args["command"] = wrapped
-    tool_call["args"] = args
-    logger.debug("Wrapped execute command with timeout")
-    return request.override(tool_call=tool_call)
-
-
-class TimeoutExecuteToolMiddleware(AgentMiddleware):
-    """Ensure execute tool calls are wrapped with a timeout."""
-
-    state_schema = AgentState
-
-    def wrap_tool_call(
-        self,
-        request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], ToolMessage | Command],
-    ) -> ToolMessage | Command:
-        return handler(_overwrite_request_if_needed(request))
-
-    async def awrap_tool_call(
-        self,
-        request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
-    ) -> ToolMessage | Command:
-        return await handler(_overwrite_request_if_needed(request))
--- a/apps/agent/agent/prompt.py
+++ b/apps/agent/agent/prompt.py
@ -8,8 +8,8 @@ All code execution and file operations happen in this sandbox environment.

 **Important:**
 - Use `{working_dir}` as your working directory for all operations
- The `execute` tool enforces a 5-minute timeout by default (`timeout 300s`)
- If a command times out and needs longer, rerun it by explicitly appending `timeout Ns`
+- The `execute` tool enforces a 5-minute timeout by default
+- If a command times out and needs longer, rerun it explicitly passing the `timeout` argument to the `execute` tool with a higher value in seconds.


 ---
--- a/apps/agent/agent/server.py
+++ b/apps/agent/agent/server.py
@ -26,9 +26,8 @@ from deepagents.backends.protocol import SandboxBackendProtocol
 from langchain_openai import ChatOpenAI

 from .encryption import decrypt_token
-from .integrations.langsmith import _create_langsmith_sandbox
+from .integrations.langsmith import create_langsmith_sandbox
 from .middleware import (
-    TimeoutExecuteToolMiddleware,
    ToolErrorMiddleware,
    check_message_queue_before_model,
    open_pr_if_needed,
@ -253,7 +252,7 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915

        try:
            # Create sandbox without context manager cleanup (sandbox persists)
-            sandbox_backend = await asyncio.to_thread(_create_langsmith_sandbox)
+            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
@ -287,7 +286,7 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915
        logger.info("Connecting to existing sandbox %s", sandbox_id)
        try:
            # Connect to existing sandbox without context manager cleanup
-            sandbox_backend = await asyncio.to_thread(_create_langsmith_sandbox, sandbox_id)
+            sandbox_backend = await asyncio.to_thread(create_langsmith_sandbox, sandbox_id)
            logger.info("Connected to existing sandbox %s", sandbox_id)
        except Exception:
            logger.warning("Failed to connect to existing sandbox %s, creating new one", sandbox_id)
@ -298,7 +297,7 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915
            )

            try:
-                sandbox_backend = await asyncio.to_thread(_create_langsmith_sandbox)
+                sandbox_backend = await asyncio.to_thread(create_langsmith_sandbox)
                logger.info("New sandbox created: %s", sandbox_backend.id)

                await client.threads.update(
@ -346,7 +345,6 @@ async def get_agent(config: RunnableConfig) -> Pregel:  # noqa: PLR0915
        tools=[http_request, fetch_url, commit_and_open_pr],
        backend=sandbox_backend,
        middleware=[
-            TimeoutExecuteToolMiddleware(),
            ToolErrorMiddleware(),
            check_message_queue_before_model,
            post_to_linear_after_model,
--- a/apps/agent/agent/utils/sandbox_state.py
+++ b/apps/agent/agent/utils/sandbox_state.py
@ -8,7 +8,7 @@ from typing import Any

 from langgraph.config import get_config

-from ..integrations.langsmith import _create_langsmith_sandbox
+from ..integrations.langsmith import create_langsmith_sandbox

 logger = logging.getLogger(__name__)

@ -36,7 +36,7 @@ async def get_sandbox_backend(thread_id: str) -> Any | None:
    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_backend = await asyncio.to_thread(create_langsmith_sandbox, sandbox_id)
    SANDBOX_BACKENDS[thread_id] = sandbox_backend
    return sandbox_backend

--- a/apps/agent/pyproject.toml
+++ b/apps/agent/pyproject.toml
@ -6,26 +6,18 @@ readme = "README.md"
 requires-python = ">=3.11"
 license = { text = "MIT" }
 dependencies = [
-    # Core deepagents library
-    "deepagents>=0.4.0",
-    # FastAPI for webhook handling
+    "deepagents>=0.4.3",
    "fastapi>=0.104.0",
    "uvicorn>=0.24.0",
-    # HTTP client
    "httpx>=0.25.0",
-    # JWT for service authentication
    "PyJWT>=2.8.0",
-    # Encryption
    "cryptography>=41.0.0",
-    # LangGraph SDK for thread management
    "langgraph-sdk>=0.1.0",
-    # LangChain dependencies (will be pulled in by deepagents but listing for clarity)
    "langchain>=1.2.9",
    "langgraph>=1.0.8",
    "markdownify>=1.2.2",
    "langchain-anthropic>1.1.0",
    "langgraph-cli[inmem]>=0.4.12",
-    # LangSmith SDK for sandbox management
    "langsmith>=0.7.1",
    "langchain-openai==1.1.10",
 ]
--- a/apps/agent/uv.lock
+++ b/apps/agent/uv.lock
@ -345,7 +345,7 @@ wheels = [

 [[package]]
 name = "deepagents"
-version = "0.4.0"
+version = "0.4.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "langchain" },
@ -354,9 +354,9 @@ dependencies = [
    { name = "langchain-google-genai" },
    { name = "wcmatch" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/b2/ac/895c5efe77ee64f38af64146509b97220867fc149a9a376b2c74f266fd45/deepagents-0.4.0.tar.gz", hash = "sha256:ccfbb2394d2c50a3cf6f61457c5c1d3868354beee2177a57b8497c72c38f8e9b", size = 77614 }
+sdist = { url = "https://files.pythonhosted.org/packages/b4/30/5bba09d1c196a9e6e2e3a3406cd131bdf01e84ec67c4b6233f68a903978f/deepagents-0.4.3.tar.gz", hash = "sha256:88033c616c5ea481f2620dbb2d05533bc8fdcd48f376d713f9dba49a8157b6f8", size = 83210 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/c4/c8/cbedac42e011889f151047cf22dbecef9a38330c2e90e517c2bd1b1e636d/deepagents-0.4.0-py3-none-any.whl", hash = "sha256:475af99429c7b6abe4c53d476d6b48f01960dc1bcaec7fa21093189acdfc2864", size = 87831 },
+    { url = "https://files.pythonhosted.org/packages/58/f8/c076a841b68cc13d89c395cc97965b37751ed008691a304119efa0f5717e/deepagents-0.4.3-py3-none-any.whl", hash = "sha256:298d19c5c0b4c6fc6a74b68049a7bfea0ba481aece7201ab21e7172b71ee61b9", size = 94882 },
 ]

 [[package]]
@ -787,21 +787,21 @@ wheels = [

 [[package]]
 name = "langchain-anthropic"
-version = "1.3.2"
+version = "1.3.4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "anthropic" },
    { name = "langchain-core" },
    { name = "pydantic" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/e7/dd/c5e094079bdd748ca3f0bd0a09189ed2fa46bba56b5a8351198dc7c19e1f/langchain_anthropic-1.3.2.tar.gz", hash = "sha256:e551726a6ebf20229bde06022b5149d33bd48d28e34bd002a744953667b8ad48", size = 686239 }
+sdist = { url = "https://files.pythonhosted.org/packages/30/4e/7c1ffac126f5e62b0b9066f331f91ae69361e73476fd3ca1b19f8d8a3cc3/langchain_anthropic-1.3.4.tar.gz", hash = "sha256:000ed4c2d6fb8842b4ffeed22a74a3e84f9e9bcb63638e4abbb4a1d8ffa07211", size = 671858 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/c9/6b/2da16c32308f79bb4588cec7095edbc770722ae4b3c3a1c135e05b0bdc2e/langchain_anthropic-1.3.2-py3-none-any.whl", hash = "sha256:35bc30862696a493680b898eb76bd6c866841f8e48a57d5eca1420a4fd807ac0", size = 46751 },
+    { url = "https://files.pythonhosted.org/packages/9b/cf/b7c7b7270efbb3db2edbf14b09ba9110a41628f3a85a11cae9527a35641c/langchain_anthropic-1.3.4-py3-none-any.whl", hash = "sha256:cd112dcc8049aef09f58b3c4338b2c9db5ee98105e08664954a4e40d8bf120b9", size = 47454 },
 ]

 [[package]]
 name = "langchain-core"
-version = "1.2.14"
+version = "1.2.15"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "jsonpatch" },
@ -813,9 +813,9 @@ dependencies = [
    { name = "typing-extensions" },
    { name = "uuid-utils" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/3f/ff/c5e3da8eca8a18719b300ef6c29e28208ee4e9da7f9749022b96292b6541/langchain_core-1.2.14.tar.gz", hash = "sha256:09549d838a2672781da3a9502f3b9c300863284b77b27e2a6dac4e6e650acfed", size = 833399 }
+sdist = { url = "https://files.pythonhosted.org/packages/cc/db/693d81b6c229aceb7c6e6809939c6ab1b023554227a25de438f00c0389c6/langchain_core-1.2.15.tar.gz", hash = "sha256:7d5f5d2daa8ddbe4054a96101dc5d509926f831b9914808c24640987d499758c", size = 835280 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/71/41/fe6ae9065b866b1397adbfc98db5e1648e8dcd78126b8e1266fcbe2d6395/langchain_core-1.2.14-py3-none-any.whl", hash = "sha256:b349ca28c057ac1f9b5280ea091bddb057db24d0f1c3c89bbb590713e1715838", size = 501411 },
+    { url = "https://files.pythonhosted.org/packages/48/e0/a6a83dde94400b43d9b091ecbb41a50d6f86c4fecacb81b13d8452a7712b/langchain_core-1.2.15-py3-none-any.whl", hash = "sha256:8d920d8a31d8c223966a3993d8c79fd6093b9665f2222fc878812f3a52072ab7", size = 502213 },
 ]

 [[package]]
@ -1050,7 +1050,7 @@ dev = [
 [package.metadata]
 requires-dist = [
    { name = "cryptography", specifier = ">=41.0.0" },
-    { name = "deepagents", specifier = ">=0.4.0" },
+    { name = "deepagents", specifier = ">=0.4.3" },
    { name = "fastapi", specifier = ">=0.104.0" },
    { name = "httpx", specifier = ">=0.25.0" },
    { name = "langchain", specifier = ">=1.2.9" },