From 85bf32064d1205b11a9f993a9eaca9513f6071df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Thu, 30 Jul 2026 16:59:06 +0800 Subject: [PATCH] fix: exclude stream replies from token accounting --- .../agent_wrapper/as_agent_wrapper.py | 9 --- .../agent_wrapper/cc_agent_wrapper.py | 1 - .../agent_wrapper/codex_agent_wrapper.py | 5 -- tests/unit/test_cc_agent_wrapper.py | 6 +- tests/unit/test_codex_agent_wrapper.py | 64 +++++++++++++++++++ tests/unit/test_token_usage.py | 32 ++++++++++ 6 files changed, 101 insertions(+), 16 deletions(-) diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 6f2171ab..d919f02f 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -498,20 +498,11 @@ class AsAgentWrapper(BaseAgentWrapper): """Stream agent events as unified StreamChunk objects.""" kwargs = self._merged_stream_kwargs(kwargs) agent, inputs = await self._build_agent(inputs, **kwargs) - usages: list[TokenUsage] = [] async for event in agent.reply_stream(inputs): - if isinstance(event, ModelCallEndEvent): - usages.append( - TokenUsage( - input_tokens=event.input_tokens, - output_tokens=event.output_tokens, - ), - ) chunk = self._event_to_chunk(event) if chunk is not None: chunk.session_id = chunk.session_id or agent.state.session_id yield chunk await self._dump_state(agent.state) - self._record_token_usage(TokenUsage.combine(usages)) diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index 3b39b30f..fa2c95f2 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -537,7 +537,6 @@ class CcAgentWrapper(BaseAgentWrapper): ) for chunk in self._result_message_to_chunks(msg): yield chunk - self._record_token_usage(self._claude_usage(msg.usage)) if reply_open or not emitted_reply_end: emitted_reply_end = True reply_open = False diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 24f939fd..5bced0b7 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -575,13 +575,10 @@ class CodexAgentWrapper(BaseAgentWrapper): turn = await thread.turn(inputs, **self._turn_kwargs(kwargs)) stream = turn.stream() completed = False - final_usage: TokenUsage | None = None try: async for event in stream: if event.method == "turn/completed": completed = True - if event.method == "thread/tokenUsage/updated": - final_usage = self._codex_usage(event.payload.token_usage.last) for chunk in self._event_to_chunks(event, thread.id): yield chunk finally: @@ -591,5 +588,3 @@ class CodexAgentWrapper(BaseAgentWrapper): except Exception as exc: # pylint: disable=broad-exception-caught self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") await stream.aclose() - if final_usage is not None: - self._record_token_usage(final_usage) diff --git a/tests/unit/test_cc_agent_wrapper.py b/tests/unit/test_cc_agent_wrapper.py index cf9ac0ee..c123faa7 100644 --- a/tests/unit/test_cc_agent_wrapper.py +++ b/tests/unit/test_cc_agent_wrapper.py @@ -473,13 +473,17 @@ async def test_reply_stream_emits_one_reply_end_for_normal_sdk_lifecycle(tmp_pat usage={"input_tokens": 1, "output_tokens": 2}, ) + wrapper = _wrapper(tmp_path) + recorded_usages = [] monkeypatch.setattr("claude_agent_sdk.query", query) - chunks = [chunk async for chunk in _wrapper(tmp_path).reply_stream("hello")] + monkeypatch.setattr(wrapper, "_record_token_usage", recorded_usages.append) + chunks = [chunk async for chunk in wrapper.reply_stream("hello")] assert sum(chunk.chunk_type == ChunkEnum.REPLY_END for chunk in chunks) == 1 delta_usage = next(chunk for chunk in chunks if chunk.metadata.get("stop_reason") == "end_turn") assert delta_usage.chunk_type == ChunkEnum.USAGE assert delta_usage.output_tokens == 2 + assert not recorded_usages @pytest.mark.asyncio diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index b32a3a13..139887c0 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -862,6 +862,70 @@ async def test_reply_stream_interrupts_turn_when_consumer_closes_early(tmp_path, assert stream_closed +@pytest.mark.asyncio +async def test_reply_stream_does_not_record_token_usage(tmp_path, monkeypatch): + wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") + recorded_usages = [] + usage = TokenUsageBreakdown( + cachedInputTokens=0, + inputTokens=3, + outputTokens=5, + reasoningOutputTokens=0, + totalTokens=8, + ) + + class FakeTurn: + id = "turn-1" + + async def stream(self): + yield SimpleNamespace( + method="thread/tokenUsage/updated", + payload=SimpleNamespace(token_usage=SimpleNamespace(last=usage)), + ) + yield SimpleNamespace( + method="turn/completed", + payload=SimpleNamespace( + turn=SimpleNamespace( + id=self.id, + status=SimpleNamespace(value="completed"), + duration_ms=1, + error=None, + ), + ), + ) + + async def interrupt(self): + raise AssertionError("Completed turns must not be interrupted") + + class FakeThread: + id = "thread-1" + + async def turn(self, _inputs, **_kwargs): + return FakeTurn() + + class FakeCodex: + def __init__(self, _config): + pass + + async def account(self): + return SimpleNamespace(account=SimpleNamespace()) + + async def close(self): + return None + + async def thread_start(self, **_kwargs): + return FakeThread() + + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) + monkeypatch.setattr(wrapper, "_record_token_usage", recorded_usages.append) + + chunks = [chunk async for chunk in wrapper.reply_stream("answer")] + await wrapper.close() + + assert any(chunk.chunk_type == ChunkEnum.USAGE for chunk in chunks) + assert not recorded_usages + + @pytest.mark.asyncio async def test_close_waits_for_active_turn(tmp_path, monkeypatch): wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") diff --git a/tests/unit/test_token_usage.py b/tests/unit/test_token_usage.py index ed36546f..6d52efa4 100644 --- a/tests/unit/test_token_usage.py +++ b/tests/unit/test_token_usage.py @@ -1,6 +1,7 @@ """Tests for unified agent token accounting.""" from agentscope.model._model_usage import ChatUsage +import pytest from reme.components.agent_wrapper import AsAgentWrapper, BaseAgentWrapper from reme.components.application_context import ApplicationContext @@ -118,3 +119,34 @@ def test_token_counter_is_a_per_agent_metric_tree(tmp_path): "cache_read_tokens_reported_calls": {"value": 1, "children": {}}, }, } + + +@pytest.mark.asyncio +async def test_agentscope_stream_reply_does_not_record_token_usage(tmp_path, monkeypatch): + """Only non-streaming AgentScope replies contribute to token accounting.""" + context = ApplicationContext(workspace_dir=str(tmp_path)) + wrapper = AsAgentWrapper(name="research", as_llm="", app_context=context) + + class FakeAgent: + """Minimal AgentScope stream double.""" + + state = type("State", (), {"session_id": "session-1"})() + + async def reply_stream(self, inputs): + """Yield no events for the supplied input.""" + if inputs is None: + yield None + + async def build_agent(inputs, **_kwargs): + """Build the minimal stream double.""" + return FakeAgent(), inputs + + async def dump_state(_state): + """Avoid durable state writes in this accounting test.""" + return None + + monkeypatch.setattr(wrapper, "_build_agent", build_agent) + monkeypatch.setattr(wrapper, "_dump_state", dump_state) + + assert [chunk async for chunk in wrapper.reply_stream("hello")] == [] + assert "__token_counter" not in context.metadata