From 6722c24dc577b116238b757618ffbd59d6c3cb40 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 17:17:23 +0800 Subject: [PATCH] Revert "fix: exclude stream replies from token accounting" This reverts commit 85bf32064d1205b11a9f993a9eaca9513f6071df. --- .../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, 16 insertions(+), 101 deletions(-) diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index d919f02f..6f2171ab 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -498,11 +498,20 @@ 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 fa2c95f2..3b39b30f 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -537,6 +537,7 @@ 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 5bced0b7..24f939fd 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -575,10 +575,13 @@ 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: @@ -588,3 +591,5 @@ 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 c123faa7..cf9ac0ee 100644 --- a/tests/unit/test_cc_agent_wrapper.py +++ b/tests/unit/test_cc_agent_wrapper.py @@ -473,17 +473,13 @@ 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) - monkeypatch.setattr(wrapper, "_record_token_usage", recorded_usages.append) - chunks = [chunk async for chunk in wrapper.reply_stream("hello")] + chunks = [chunk async for chunk in _wrapper(tmp_path).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 139887c0..b32a3a13 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -862,70 +862,6 @@ 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 6d52efa4..ed36546f 100644 --- a/tests/unit/test_token_usage.py +++ b/tests/unit/test_token_usage.py @@ -1,7 +1,6 @@ """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 @@ -119,34 +118,3 @@ 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