From d2d08ff850f6fa181b88fc0bd60fef5e52560693 Mon Sep 17 00:00:00 2001
From: "Naveen R. Iyer" <12383569+iyernaveenr@users.noreply.github.com>
Date: Tue, 14 Jul 2026 03:57:15 -0500
Subject: [PATCH] Make the agent step timeout configurable and record spend of
 failed runs (#1746)

Signed-off-by: Naveen R. Iyer <iyernaveenr@gmail.com>
Co-authored-by: Tong Chen <web_chentong@163.com>
---
 backend/app/agent/listen_chat_agent.py        | 20 ++++++++++++++++++-
 .../app/agent/toolkit/screenshot_toolkit.py   |  5 ++++-
 src/store/chatStore.ts                        | 13 ++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/backend/app/agent/listen_chat_agent.py b/backend/app/agent/listen_chat_agent.py
index 309e788e..a2d505ee 100644
--- a/backend/app/agent/listen_chat_agent.py
+++ b/backend/app/agent/listen_chat_agent.py
@@ -36,6 +36,7 @@ from camel.types import ModelPlatformType, ModelType
 from camel.types.agents import ToolCallingRecord
 from pydantic import BaseModel
 
+from app.component.environment import env
 from app.service.task import (
     Action,
     ActionActivateAgentData,
@@ -52,6 +53,21 @@ from app.utils.event_loop_utils import _schedule_async_task
 logger = logging.getLogger("agent")
 
 
+# Default 30 minutes; long agent turns (e.g. writing many chapters in one
+# run) can legitimately exceed it, so allow tuning without a rebuild.
+# A non-positive value disables the per-step timeout entirely.
+def default_step_timeout() -> float | None:
+    raw = env("AGENT_STEP_TIMEOUT_SECONDS", "1800")
+    try:
+        value = float(raw)
+    except (TypeError, ValueError):
+        logger.warning(
+            "Invalid AGENT_STEP_TIMEOUT_SECONDS value %r; using 1800", raw
+        )
+        return 1800.0
+    return value if value > 0 else None
+
+
 class ListenChatAgent(ChatAgent):
     _cdp_clone_lock = (
         threading.Lock()
@@ -95,12 +111,14 @@ class ListenChatAgent(ChatAgent):
         pause_event: asyncio.Event | None = None,
         prune_tool_calls_from_memory: bool = False,
         enable_snapshot_clean: bool = False,
-        step_timeout: float | None = 1800,  # 30 minutes
+        step_timeout: float | None = None,
         model_reload_callback: (
             Callable[[], BaseModelBackend | ModelManager] | None
         ) = None,
         **kwargs: Any,
     ) -> None:
+        if step_timeout is None:
+            step_timeout = default_step_timeout()
         super().__init__(
             system_message=system_message,
             model=model,
diff --git a/backend/app/agent/toolkit/screenshot_toolkit.py b/backend/app/agent/toolkit/screenshot_toolkit.py
index 77f6aa29..43f4459e 100644
--- a/backend/app/agent/toolkit/screenshot_toolkit.py
+++ b/backend/app/agent/toolkit/screenshot_toolkit.py
@@ -20,6 +20,7 @@ from camel.messages import BaseMessage
 from camel.toolkits import ScreenshotToolkit as BaseScreenshotToolkit
 from PIL import Image
 
+from app.agent.listen_chat_agent import default_step_timeout
 from app.agent.toolkit.abstract_toolkit import AbstractToolkit
 from app.component.environment import env
 from app.utils.listen.toolkit_listen import auto_listen_toolkit
@@ -86,7 +87,9 @@ class ScreenshotToolkit(BaseScreenshotToolkit, AbstractToolkit):
                 tools=[],
                 toolkits_to_register_agent=None,
                 external_tools=None,
-                step_timeout=getattr(self.agent, "step_timeout", 1800),
+                step_timeout=getattr(
+                    self.agent, "step_timeout", default_step_timeout()
+                ),
             )
             response = vision_agent.step(message)
             if getattr(response, "msg", None) is not None:
diff --git a/src/store/chatStore.ts b/src/store/chatStore.ts
index 6b8a112c..86b98755 100644
--- a/src/store/chatStore.ts
+++ b/src/store/chatStore.ts
@@ -3350,6 +3350,19 @@ const chatStore = (initial?: Partial<ChatStore>) =>
                 role: 'agent',
                 content: `❌ **Error**: ${errorMessage}`,
               });
+              // Record the tokens consumed before the failure so the run's
+              // spend is not lost from the history row (a failed run
+              // otherwise stays at zero tokens forever).
+              if (!type && historyId && !isProjectBusyError) {
+                const tokensSoFar = getTokens(currentTaskId);
+                if (tokensSoFar > 0) {
+                  proxyFetchPut(`/api/v1/chat/history/${historyId}`, {
+                    tokens: tokensSoFar,
+                  }).catch((err) => {
+                    console.warn('History token update failed on error:', err);
+                  });
+                }
+              }
               uploadLog(currentTaskId, type);
               // Update trigger execution status to Failed on error
               updateTriggerExecutionStatus(