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/agent_model.py

@@ -35,6 +35,23 @@
 from app.service.task import ActionCreateAgentData, Agents, get_task_lock
 from app.utils.event_loop_utils import _schedule_async_task
 
+# OpenAI chat-completions streaming only returns token usage when
+# `stream_options.include_usage` is requested. Without it the request-level
+# usage callback (on_request_usage) fires with 0 tokens, and because the
+# step-level deactivate is zeroed once request-level reporting is active,
+# streaming steps end up uncounted. These platforms use native (non-OpenAI)
+# SDKs that reject `stream_options` and surface streaming usage on their own,
+# so they are excluded from the injection below.
+_NATIVE_STREAM_USAGE_PLATFORMS = {
+    "anthropic",
+    "aws-bedrock",
+    "aws-bedrock-converse",
+    "cohere",
+    "mistral",
+    "reka",
+    "watsonx",
+}
+
 
 def agent_model(
     agent_name: str,
@@ -217,6 +234,21 @@ def build_model(force_refresh: bool = False):
             if model_config.get("max_tokens") is None:
                 model_config["max_tokens"] = 128000
 
+        # Ensure streaming steps still report token usage. OpenAI-family
+        # providers omit usage from streamed responses unless include_usage
+        # is set, which would otherwise make request-level accounting count 0.
+        # `stream_options: false` in extra_params opts out entirely, for
+        # endpoints that reject the parameter (e.g. older vLLM/Azure).
+        if model_config.get("stream_options") is False:
+            model_config.pop("stream_options")
+        elif model_config.get("stream") and (
+            effective_config["model_platform"].lower()
+            not in _NATIVE_STREAM_USAGE_PLATFORMS
+        ):
+            stream_options = model_config.setdefault("stream_options", {})
+            if isinstance(stream_options, dict):
+                stream_options.setdefault("include_usage", True)
+
         return ModelFactory.create(
             model_platform=effective_config["model_platform"],
             model_type=effective_config["model_type"],

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
@@ -44,7 +45,9 @@
     ActionBudgetNotEnough,
     ActionDeactivateAgentData,
     ActionDeactivateToolkitData,
+    ActionRequestUsageData,
     get_task_lock,
+    get_task_lock_if_exists,
     set_process_task,
 )
 from app.utils.event_loop_utils import _schedule_async_task
@@ -73,6 +76,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,
@@ -117,6 +124,12 @@ def __init__(
         ) = None,
         **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
+
         if step_timeout is None:
             step_timeout = default_step_timeout()
         super().__init__(
@@ -142,13 +155,39 @@ def __init__(
             step_timeout=step_timeout,
             **kwargs,
         )
-        self.api_task_id = api_task_id
-        self.agent_name = agent_name
         self._model_reload_callback = model_reload_callback
         self._model_reload_lock = threading.Lock()
 
     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)
+        # Lock may be gone if the task was stopped mid-request.
+        task_lock = get_task_lock_if_exists(self.api_task_id)
+        if request_tokens > 0 and task_lock is not None:
+            _schedule_async_task(
+                task_lock.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
+
     @staticmethod
     def _is_retryable_model_auth_error(error: BaseException) -> bool:
         error_text = str(error).lower()
@@ -211,7 +250,17 @@ def _send_agent_deactivate(self, message: str, tokens: int) -> None:
             message: The accumulated message content
             tokens: The total token count used
         """
-        task_lock = get_task_lock(self.api_task_id)
+        if self._camel_has_request_usage:
+            tokens = 0
+        # A missing lock (task stopped mid-step) must not fail the step.
+        task_lock = get_task_lock_if_exists(self.api_task_id)
+        if task_lock is None:
+            logger.warning(
+                "Task lock %s missing; dropping deactivate event for %s",
+                self.api_task_id,
+                self.agent_name,
+            )
+            return
         _schedule_async_task(
             task_lock.put_queue(
                 ActionDeactivateAgentData(
@@ -458,19 +507,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
@@ -585,19 +622,7 @@ async def astep(
         # Streaming responses handle deactivation in _astream_chunks
         assert message is not None
 
-        asyncio.create_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
@@ -995,6 +1020,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,
@@ -1022,6 +1051,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

@@ -1701,6 +1701,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/single_agent_service.py

@@ -175,6 +175,8 @@ def _action_to_sse(item: ActionData) -> str | None:
         return sse_json("activate_agent", item.data)
     if item.action == Action.deactivate_agent:
         return sse_json("deactivate_agent", item.data)
+    if item.action == Action.request_usage:
+        return sse_json("request_usage", item.data)
     if item.action == Action.assign_task:
         return sse_json("assign_task", item.data)
     if item.action == Action.activate_toolkit:

backend/app/service/task.py

@@ -52,6 +52,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
@@ -160,6 +161,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[
@@ -298,6 +314,7 @@ class ActionSkipTaskData(BaseModel):
     | ActionCreateAgentData
     | ActionActivateAgentData
     | ActionDeactivateAgentData
+    | ActionRequestUsageData
     | ActionAssignTaskData
     | ActionActivateToolkitData
     | ActionDeactivateToolkitData

server/tests/test_auth.py

@@ -139,8 +139,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):
@@ -181,8 +180,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: