feat:add_request_level_token_callback

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,34 +1,38 @@
 <!-- Thank you for contributing! -->
 
-### Related Issue
+# Pull Request
+
+## Related Issue
 
 <!-- REQUIRED: Link to the issue this PR resolves. PRs without a linked issue will be closed. -->
 
 <!-- Example: Closes #123 or Fixes #456 -->
 
 Closes #
 
-### Description
+## Description
 
 <!-- REQUIRED: Describe what this PR does and why. PRs without a description will not be reviewed. -->
 
-### Testing Evidence (REQUIRED)
+## Testing Evidence (REQUIRED)
 
 <!-- REQUIRED: Every PR must include human-verified testing proof (e.g., test logs, screenshots, or screen recordings). -->
+
 <!-- REQUIRED for frontend/UI changes: You MUST attach at least one screenshot or screen recording in this PR. -->
+
 <!-- Frontend/UI PRs without visual evidence will not be reviewed. -->
 
 - [ ] I have included human-verified testing evidence in this PR.
 - [ ] This PR includes frontend/UI changes, and I attached screenshot(s) or screen recording(s).
 - [ ] No frontend/UI changes in this PR.
 
-### What is the purpose of this pull request? <!-- (put an "X" next to an item) -->
+## What is the purpose of this pull request? <!-- (put an "X" next to an item) -->
 
 - [ ] Bug fix
 - [ ] New Feature
 - [ ] Documentation update
 - [ ] Other
 
-### Contribution Guidelines Acknowledgement
+## Contribution Guidelines Acknowledgement
 
 - [ ] I have read and agree to the [Eigent Contribution Guideline](https://github.com/eigent-ai/eigent/blob/main/CONTRIBUTING.md#eigent-contribution-guideline)

backend/app/agent/listen_chat_agent.py

@@ -13,6 +13,7 @@
 # ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
 
 import asyncio
+import inspect
 import json
 import logging
 import threading
@@ -43,6 +44,7 @@
     ActionBudgetNotEnough,
     ActionDeactivateAgentData,
     ActionDeactivateToolkitData,
+    ActionRequestUsageData,
     get_task_lock,
     set_process_task,
 )
@@ -57,6 +59,10 @@ class ListenChatAgent(ChatAgent):
         threading.Lock()
     )  # Protects CDP URL mutation during clone
 
+    _camel_has_request_usage: bool = (
+        "on_request_usage" in inspect.signature(ChatAgent.__init__).parameters
+    )
+
     def __init__(
         self,
         api_task_id: str,
@@ -98,6 +104,12 @@ def __init__(
         step_timeout: float | None = 1800,  # 30 minutes
         **kwargs: Any,
     ) -> None:
+        self.api_task_id = api_task_id
+        self.agent_name = agent_name
+        self._user_on_request_usage = kwargs.pop("on_request_usage", None)
+        if self._camel_has_request_usage:
+            kwargs["on_request_usage"] = self._on_request_usage
+
         super().__init__(
             system_message=system_message,
             model=model,
@@ -121,18 +133,44 @@ def __init__(
             step_timeout=step_timeout,
             **kwargs,
         )
-        self.api_task_id = api_task_id
-        self.agent_name = agent_name
 
     process_task_id: str = ""
 
+    def _on_request_usage(self, payload: dict[str, Any]) -> Any:
+        request_usage = payload.get("request_usage") or {}
+        step_usage = payload.get("step_usage") or {}
+        request_tokens = int(request_usage.get("total_tokens") or 0)
+        if request_tokens > 0:
+            _schedule_async_task(
+                get_task_lock(self.api_task_id).put_queue(
+                    ActionRequestUsageData(
+                        data={
+                            "agent_name": self.agent_name,
+                            "process_task_id": self.process_task_id,
+                            "agent_id": self.agent_id,
+                            "tokens": request_tokens,
+                            "request_index": payload.get("request_index", 0),
+                            "response_id": payload.get("response_id", ""),
+                            "step_total_tokens": int(
+                                step_usage.get("total_tokens") or 0
+                            ),
+                        }
+                    )
+                )
+            )
+        if self._user_on_request_usage is not None:
+            return self._user_on_request_usage(payload)
+        return None
+
     def _send_agent_deactivate(self, message: str, tokens: int) -> None:
         """Send agent deactivation event to the frontend.
 
         Args:
             message: The accumulated message content
             tokens: The total token count used
         """
+        if self._camel_has_request_usage:
+            tokens = 0
         task_lock = get_task_lock(self.api_task_id)
         _schedule_async_task(
             task_lock.put_queue(
@@ -299,19 +337,7 @@ def step(
 
         assert message is not None
 
-        _schedule_async_task(
-            task_lock.put_queue(
-                ActionDeactivateAgentData(
-                    data={
-                        "agent_name": self.agent_name,
-                        "process_task_id": self.process_task_id,
-                        "agent_id": self.agent_id,
-                        "message": message,
-                        "tokens": total_tokens,
-                    },
-                )
-            )
-        )
+        self._send_agent_deactivate(message, total_tokens)
 
         if error_info is not None:
             raise error_info
@@ -412,7 +438,9 @@ async def astep(
                         "process_task_id": self.process_task_id,
                         "agent_id": self.agent_id,
                         "message": message,
-                        "tokens": total_tokens,
+                        "tokens": 0
+                        if self._camel_has_request_usage
+                        else total_tokens,
                     },
                 )
             )
@@ -754,6 +782,10 @@ def clone(self, with_memory: bool = False) -> ChatAgent:
         else:
             cloned_tools, toolkits_to_register = self._clone_tools()
 
+        clone_kwargs: dict[str, Any] = {}
+        if self._user_on_request_usage is not None:
+            clone_kwargs["on_request_usage"] = self._user_on_request_usage
+
         new_agent = ListenChatAgent(
             api_task_id=self.api_task_id,
             agent_name=self.agent_name,
@@ -781,6 +813,7 @@ def clone(self, with_memory: bool = False) -> ChatAgent:
             enable_snapshot_clean=self._enable_snapshot_clean,
             step_timeout=self.step_timeout,
             stream_accumulate=self.stream_accumulate,
+            **clone_kwargs,
         )
 
         new_agent.process_task_id = self.process_task_id

backend/app/service/chat_service.py

@@ -1477,6 +1477,8 @@ def on_stream_text(chunk):
                 yield sse_json("activate_agent", item.data)
             elif item.action == Action.deactivate_agent:
                 yield sse_json("deactivate_agent", dict(item.data))
+            elif item.action == Action.request_usage:
+                yield sse_json("request_usage", dict(item.data))
             elif item.action == Action.assign_task:
                 yield sse_json("assign_task", item.data)
             elif item.action == Action.activate_toolkit:

backend/app/service/task.py

@@ -49,6 +49,7 @@ class Action(str, Enum):
     create_agent = "create_agent"  # backend -> user
     activate_agent = "activate_agent"  # backend -> user
     deactivate_agent = "deactivate_agent"  # backend -> user
+    request_usage = "request_usage"  # backend -> user
     assign_task = "assign_task"  # backend -> user
     activate_toolkit = "activate_toolkit"  # backend -> user
     deactivate_toolkit = "deactivate_toolkit"  # backend -> user
@@ -155,6 +156,21 @@ class ActionDeactivateAgentData(BaseModel):
     data: DataDict
 
 
+class RequestUsageDataDict(TypedDict):
+    agent_name: str
+    agent_id: str
+    process_task_id: str
+    tokens: int
+    request_index: int
+    response_id: str
+    step_total_tokens: int
+
+
+class ActionRequestUsageData(BaseModel):
+    action: Literal[Action.request_usage] = Action.request_usage
+    data: RequestUsageDataDict
+
+
 class ActionAssignTaskData(BaseModel):
     action: Literal[Action.assign_task] = Action.assign_task
     data: dict[
@@ -288,6 +304,7 @@ class ActionSkipTaskData(BaseModel):
     | ActionCreateAgentData
     | ActionActivateAgentData
     | ActionDeactivateAgentData
+    | ActionRequestUsageData
     | ActionAssignTaskData
     | ActionActivateToolkitData
     | ActionDeactivateToolkitData

server/tests/test_auth.py

@@ -47,7 +47,7 @@ def test_auth_must_raises_on_none_token(self):
         """auth_must should raise TokenException immediately when
         token is None, not pass it to jwt.decode()."""
         import asyncio
-        from unittest.mock import MagicMock, patch
+        from unittest.mock import MagicMock
 
         from app.component.auth import auth_must
         from app.exception.exception import TokenException
@@ -88,8 +88,7 @@ def test_list_snapshots_requires_auth_dependency(self):
         sig = inspect.signature(list_chat_snapshots)
         param_names = list(sig.parameters.keys())
         assert "auth" in param_names, (
-            "list_chat_snapshots is missing the 'auth' parameter — "
-            "unauthenticated users can list all snapshots"
+            "list_chat_snapshots is missing the 'auth' parameter — unauthenticated users can list all snapshots"
         )
 
     def test_get_snapshot_requires_auth_dependency(self):
@@ -130,8 +129,7 @@ def test_create_share_link_requires_auth_dependency():
     sig = inspect.signature(create_share_link)
     param_names = list(sig.parameters.keys())
     assert "auth" in param_names, (
-        "create_share_link is missing the 'auth' parameter — "
-        "unauthenticated users can generate share tokens"
+        "create_share_link is missing the 'auth' parameter — unauthenticated users can generate share tokens"
     )
 
 

server/tests/test_chat_share.py

@@ -15,8 +15,6 @@
 import os
 from unittest.mock import patch
 
-import pytest
-
 
 class TestChatShareSecretKey:
     """Tests for ChatShare secret key generation.

server/tests/test_proxy_controller.py

@@ -12,9 +12,7 @@
 # limitations under the License.
 # ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
 
-from urllib.parse import quote_plus, urlparse, parse_qs
-
-import pytest
+from urllib.parse import parse_qs, quote_plus, urlparse
 
 
 class TestGoogleSearchUrlEncoding:

src/store/chatStore.ts

@@ -1471,6 +1471,14 @@ const chatStore = (initial?: Partial<ChatStore>) =>
             return;
           }
 
+          // Request-level token usage updates (non-stream mode)
+          if (agentMessages.step === AgentStep.REQUEST_USAGE) {
+            if (agentMessages.data.tokens) {
+              addTokens(currentTaskId, agentMessages.data.tokens);
+            }
+            return;
+          }
+
           // Activate agent
           if (
             agentMessages.step === AgentStep.ACTIVATE_AGENT ||

src/types/chatbox.d.ts

@@ -139,6 +139,9 @@ declare global {
       agent?: string;
       file_path?: string;
       process_task_id?: string;
+      request_index?: number;
+      response_id?: string;
+      step_total_tokens?: number;
       output?: string;
       result?: string;
       tools?: string[];

src/types/constants.ts

@@ -26,6 +26,7 @@ export const AgentStep = {
   TASK_STATE: 'task_state',
   ACTIVATE_AGENT: 'activate_agent',
   DEACTIVATE_AGENT: 'deactivate_agent',
+  REQUEST_USAGE: 'request_usage',
   ASSIGN_TASK: 'assign_task',
   ACTIVATE_TOOLKIT: 'activate_toolkit',
   DEACTIVATE_TOOLKIT: 'deactivate_toolkit',

test/unit/store/chatStore.test.ts

@@ -579,6 +579,59 @@ describe('ChatStore - Core Functionality', () => {
     });
   });
 
+  describe('SSE request usage events', () => {
+    it('should accumulate tokens from request_usage event in non-stream mode', async () => {
+      vi.mocked(proxyFetchGet).mockImplementation((url: string) =>
+        url?.includes?.('snapshots')
+          ? Promise.resolve([])
+          : Promise.resolve({
+              value: '',
+              api_url: '',
+              items: [],
+              warning_code: null,
+            })
+      );
+
+      const mockFetchEventSource = vi.mocked(fetchEventSource);
+      mockFetchEventSource.mockImplementation(async (_url, opts) => {
+        opts.onmessage?.({
+          data: JSON.stringify({
+            step: 'request_usage',
+            data: { tokens: 11 },
+          }),
+        } as any);
+        opts.onmessage?.({
+          data: JSON.stringify({
+            step: 'deactivate_agent',
+            data: { tokens: 0 },
+          }),
+        } as any);
+        return Promise.resolve();
+      });
+
+      const { result } = renderHook(() => useChatStore());
+      let taskId!: string;
+      await act(async () => {
+        taskId = result.current.getState().create();
+        result.current.getState().setActiveTaskId(taskId);
+        result.current.getState().setHasMessages(taskId, true);
+        result.current.getState().addMessages(taskId, {
+          id: generateUniqueId(),
+          role: 'user',
+          content: 'Test message',
+        });
+      });
+
+      await act(async () => {
+        await result.current
+          .getState()
+          .startTask(taskId, 'replay', undefined, 0.2);
+      });
+
+      expect(result.current.getState().tasks[taskId].tokens).toBe(11);
+    });
+  });
+
   describe('Replay', () => {
     const replayProjectState = () => ({
       activeProjectId: 'proj-replay',