From 21f59ea33c2019288420c686a8b9774db4736b59 Mon Sep 17 00:00:00 2001
From: worldmozara <worldmozara@163.com>
Date: Tue, 6 Oct 2026 23:36:35 +0800
Subject: [PATCH] core/agent: count the prompt a search turn really sent

The compat opens with the question alone (``message_start`` reports 165 tokens)
and puts what the search then pulled in on ``message_delta`` (12468); merging
preferred the earlier value, so the CLI's context gauge and the auto-compact
threshold ran ~12k short of the request. The larger of the two wins now (output
tokens are the delta's cumulative count either way).
---
 src/plyngent/agent/messages_dispatch.py    | 12 ++++++--
 tests/test_agent/test_messages_dispatch.py | 35 ++++++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/plyngent/agent/messages_dispatch.py b/src/plyngent/agent/messages_dispatch.py
index c373188..7ff6dfb 100644
--- a/src/plyngent/agent/messages_dispatch.py
+++ b/src/plyngent/agent/messages_dispatch.py
@@ -66,11 +66,17 @@ class _ToolBlockState:
 
 
 def _merge_usage(previous: dict[str, Any] | None, new: dict[str, Any]) -> dict[str, Any]:
-    """Merge message_start input tokens with message_delta output tokens."""
+    """Merge ``message_start`` usage with the later ``message_delta`` usage.
+
+    DeepSeek's compat reports a search turn's real prompt tokens only on the
+    delta (``message_start`` sees the question alone, the delta counts what the
+    search pulled in), so the larger of the two wins; output tokens are the
+    delta's cumulative count either way.
+    """
     if previous is None:
         return new
-    prompt = int(previous.get("prompt_tokens", 0) or new.get("prompt_tokens", 0) or 0)
-    completion = int(new.get("completion_tokens", 0) or previous.get("completion_tokens", 0) or 0)
+    prompt = max(int(previous.get("prompt_tokens", 0) or 0), int(new.get("prompt_tokens", 0) or 0))
+    completion = max(int(previous.get("completion_tokens", 0) or 0), int(new.get("completion_tokens", 0) or 0))
     return {
         "prompt_tokens": prompt,
         "completion_tokens": completion,
diff --git a/tests/test_agent/test_messages_dispatch.py b/tests/test_agent/test_messages_dispatch.py
index 0c0600a..98db65a 100644
--- a/tests/test_agent/test_messages_dispatch.py
+++ b/tests/test_agent/test_messages_dispatch.py
@@ -497,3 +497,38 @@ async def test_dispatch_messages_without_provider_tools_sends_no_server_tool() -
     param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")])
     _ = await dispatch_messages(cast("Any", client), param, provider_tools=[])
     assert client.calls[0].tools is UNSET
+
+
+async def test_search_turn_reports_the_prompt_the_delta_counts() -> None:
+    """A search turn's real prompt size only shows up on ``message_delta``.
+
+    DeepSeek's compat opens with the question alone (``message_start`` sees 165
+    tokens) and reports what the search then pulled in on the delta (12468);
+    keeping the smaller one leaves the CLI's context gauge — and the
+    auto-compact threshold — that far short of what the model was sent.
+    """
+    events: list[AnthropicStreamEvent] = [
+        AnthropicMessageStart(
+            message=AnthropicMessageResponse(
+                id="msg_s",
+                model="deepseek-flash",
+                content=[],
+                usage=AnthropicUsage(input_tokens=165, output_tokens=0),
+            )
+        ),
+        AnthropicContentBlockStart(index=0, content_block=AnthropicRawContentBlock(type="text", text="")),
+        AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="text_delta", text="done")),
+        AnthropicMessageDelta(
+            delta={"stop_reason": "end_turn"},
+            usage=AnthropicUsage(input_tokens=12468, output_tokens=198),
+        ),
+        AnthropicMessageStop(),
+    ]
+    client = ScriptedMessagesClient(stream_events=[events])
+    param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")])
+    stream = await dispatch_messages(cast("Any", client), param, stream=True)
+    chunks = [chunk async for chunk in cast("Any", stream)]
+    usage = next(getattr(chunk, "usage", None) for chunk in chunks if getattr(chunk, "usage", None))
+    assert isinstance(usage, dict)
+    assert usage["prompt_tokens"] == 12468
+    assert usage["completion_tokens"] == 198