fix: exclude stream replies from token accounting

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))

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

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)

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

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")

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