diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 762f852b3..e7456f18f 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -1,6 +1,8 @@ -### Related Issue +# Pull Request + +## Related Issue @@ -8,27 +10,29 @@ Closes # -### Description +## Description -### Testing Evidence (REQUIRED) +## Testing Evidence (REQUIRED) + + - [ ] 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? +## What is the purpose of this pull request? - [ ] 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) diff --git a/backend/app/agent/agent_model.py b/backend/app/agent/agent_model.py index acbb43685..b8eaf7e87 100644 --- a/backend/app/agent/agent_model.py +++ b/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, @@ -251,6 +268,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"], diff --git a/backend/app/agent/listen_chat_agent.py b/backend/app/agent/listen_chat_agent.py index a2d505eea..d4cf64ebf 100644 --- a/backend/app/agent/listen_chat_agent.py +++ b/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 diff --git a/backend/app/service/chat_service.py b/backend/app/service/chat_service.py index 029f83c20..2b977df8c 100644 --- a/backend/app/service/chat_service.py +++ b/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: diff --git a/backend/app/service/single_agent_service.py b/backend/app/service/single_agent_service.py index 6d963c9a7..a7994f0b2 100644 --- a/backend/app/service/single_agent_service.py +++ b/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: diff --git a/backend/app/service/task.py b/backend/app/service/task.py index eb0883bba..00d5b0af5 100644 --- a/backend/app/service/task.py +++ b/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 diff --git a/server/app/domains/remote_control/service/remote_control_service.py b/server/app/domains/remote_control/service/remote_control_service.py index bf6590dce..daa4c149e 100644 --- a/server/app/domains/remote_control/service/remote_control_service.py +++ b/server/app/domains/remote_control/service/remote_control_service.py @@ -108,6 +108,7 @@ SPACE_DISCARD_PROJECT_OVERLAYS, } NON_BRAIN_COMMANDS = {SWITCH_PROJECT_VIEW, *SPACE_COMMANDS} +REMOTE_CONTROL_HIDDEN_STEPS = ("request_usage",) def _now() -> datetime: @@ -1569,7 +1570,11 @@ def list_steps( task_ids = [legacy_history.task_id] if not task_ids: return RemoteControlStepsOut(items=[], has_more=False, next_since=since) - stmt = select(ChatStep).where(ChatStep.task_id.in_(task_ids), ChatStep.id > since) + stmt = select(ChatStep).where( + ChatStep.task_id.in_(task_ids), + ChatStep.id > since, + ChatStep.step.not_in(REMOTE_CONTROL_HIDDEN_STEPS), + ) stmt = stmt.order_by(ChatStep.id.desc() if order == "desc" else ChatStep.id.asc()).limit(limit + 1) rows = list(db.exec(stmt).all()) has_more = len(rows) > limit @@ -1593,6 +1598,8 @@ def list_steps( @staticmethod def publish_chat_step(step: ChatStep, db: Session) -> None: + if step.step in REMOTE_CONTROL_HIDDEN_STEPS: + return history = db.exec(select(ChatHistory).where(ChatHistory.task_id == step.task_id)).first() if not history: logger.warning("Skipping remote-control step publish for orphan step", extra={"task_id": step.task_id}) diff --git a/server/tests/test_auth.py b/server/tests/test_auth.py index 6ad55fd6b..b5ff63733 100644 --- a/server/tests/test_auth.py +++ b/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" ) diff --git a/server/tests/test_chat_share.py b/server/tests/test_chat_share.py index 5af72b43e..6657f3451 100644 --- a/server/tests/test_chat_share.py +++ b/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. diff --git a/server/tests/test_proxy_controller.py b/server/tests/test_proxy_controller.py index c2ca19df2..1f2fa4242 100644 --- a/server/tests/test_proxy_controller.py +++ b/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: diff --git a/server/tests/test_remote_control_security.py b/server/tests/test_remote_control_security.py index 37c129337..27fe26d3c 100644 --- a/server/tests/test_remote_control_security.py +++ b/server/tests/test_remote_control_security.py @@ -528,6 +528,98 @@ def list_steps(session_id, user_id, project_id, since, limit, order, db_session) assert result is expected +def test_list_steps_excludes_request_usage(monkeypatch): + from app.domains.remote_control.service.remote_control_service import RemoteControlService + + history = SimpleNamespace(task_id="task_1") + visible_step = SimpleNamespace( + id=12, + task_id="task_1", + step="ask", + data={"question": "Continue?"}, + timestamp=12.0, + ) + + class FakeResult: + def __init__(self, items): + self.items = items + + def all(self): + return self.items + + class FakeDb: + def __init__(self): + self.calls = 0 + + def exec(self, statement): + self.calls += 1 + if self.calls == 1: + return FakeResult([history]) + sql = str(statement.compile(compile_kwargs={"literal_binds": True})) + assert "request_usage" in sql + assert "NOT IN" in sql + return FakeResult([visible_step]) + + monkeypatch.setattr( + RemoteControlService, + "_owned_session", + staticmethod(lambda _session_id, _user_id, _db: SimpleNamespace()), + ) + monkeypatch.setattr( + RemoteControlService, + "_effective_target", + staticmethod(lambda _session: ("project_1", None, None, None)), + ) + monkeypatch.setattr( + RemoteControlService, + "_ensure_project_in_session_space", + staticmethod(lambda *_args: None), + ) + + result = RemoteControlService.list_steps( + "rcs_test", + 123, + None, + 0, + 10, + "asc", + FakeDb(), + ) + + assert [item.step for item in result.items] == ["ask"] + + +def test_publish_chat_step_excludes_request_usage(monkeypatch): + from app.domains.remote_control.service.remote_control_service import ( + RemoteControlRedis, + RemoteControlService, + ) + + class FailDb: + def exec(self, _statement): + raise AssertionError("request_usage should not query or publish") + + published = [] + monkeypatch.setattr( + RemoteControlRedis, + "publish", + lambda channel, payload: published.append((channel, payload)), + ) + + RemoteControlService.publish_chat_step( + SimpleNamespace( + id=13, + task_id="task_1", + step="request_usage", + data={"tokens": 10}, + timestamp=13.0, + ), + FailDb(), + ) + + assert published == [] + + @pytest.mark.asyncio async def test_ws_reconnect_rate_limit_is_per_identifier(monkeypatch): from app.domains.remote_control.api import remote_control_controller as controller diff --git a/src/components/ProjectPageSidebar/index.tsx b/src/components/ProjectPageSidebar/index.tsx index 6d6471f1d..5d2d1d3c4 100644 --- a/src/components/ProjectPageSidebar/index.tsx +++ b/src/components/ProjectPageSidebar/index.tsx @@ -672,14 +672,14 @@ export default function ProjectPageSidebar({