diff --git a/README.md b/README.md index d4b71b5ba3..965e230f9c 100644 --- a/README.md +++ b/README.md @@ -2902,6 +2902,13 @@ client.clear_goal("thread-1") When continuing a thread, `values` still contains the full conversation state, while `messages-tuple` and `end.usage` cover only the current turn. +The embedded client's `end.usage` sums usage deltas across streamed chunks, +including providers such as native Gemini that report usage more than once per +reply. The completed message's usage is reconciled without counting those +chunks again; a lower completed-message count does not subtract tokens already +reported by chunks. Signed provider corrections in chunks remain supported, as +do providers that report usage only at completion. + Concurrent `upload_files()` calls preserve filenames absent at the start of each call by assigning numeric suffixes and returning the actual saved names. The same rule applies to new Markdown companions. Exact names already present retain diff --git a/backend/docs/STREAMING.md b/backend/docs/STREAMING.md index 78d391a6bb..8000312ee7 100644 --- a/backend/docs/STREAMING.md +++ b/backend/docs/STREAMING.md @@ -8,7 +8,7 @@ - DeerFlow 有**两条并行**的流式路径：**Gateway 路径**（async / HTTP SSE / JSON 序列化）服务浏览器和 IM 渠道；**DeerFlowClient 路径**（sync / in-process / 原生 LangChain 对象）服务 Jupyter、脚本、测试。它们**无法合并**——消费者模型不同。 - 两条路径都从 `create_agent()` 工厂出发，核心都是订阅 LangGraph 的 `stream_mode=["values", "messages", "custom"]`。`values` 是节点级 state 快照，`messages` 是 LLM token 级 delta，`custom` 是显式 `StreamWriter` 事件。DeerFlow 内置 custom 事件同时通过 callback dispatch 暴露为 `astream_events(version="v2")` 的 `on_custom_event`，供 AG-UI 等 callback 型消费者使用。**这些接口不是详细程度的梯度，而是独立事件源**，消费者必须订阅自己需要的接口。 -- 嵌入式 client 为每个 `stream()` 调用维护四个 `set[str]`：`seen_ids` / `streamed_ids` / `counted_usage_ids` / `historical_message_ids`。前三者分别负责去重与用量幂等；最后一个隔离续聊时的历史消息。 +- 嵌入式 client 按消息 id 分别跟踪 values 去重、已流式发送的文本、已核算的完整用量和历史消息；`streamed_usage_by_id` 累加尚未核算完整消息的 chunk 用量。 --- @@ -226,8 +226,9 @@ return "".join(chunks.get(last_id, ())) ```python seen_ids: set[str] = set() # values 路径内部 dedup streamed_ids: set[str] = set() # messages → values 跨模式 dedup -counted_usage_ids: set[str] = set() # usage_metadata 幂等计数 +counted_usage_ids: set[str] = set() # 完整消息的 usage 已核算 historical_message_ids: set[str] = set() # 续聊时隔离先前回合 +streamed_usage_by_id: dict[str, dict[str, int]] = {} # 同 id 的 chunk 用量之和 ``` 乍看用途相近，实际每个管**不同的不变式**。 @@ -236,7 +237,7 @@ historical_message_ids: set[str] = set() # 续聊时隔离先前回合 |---|---|---|---| | `seen_ids` | 连续两个 `values` 快照里同一条 message 只生成一个 `messages-tuple` 事件 | values 分支每处理一条消息就加入 | values 分支处理下一条消息前检查 | | `streamed_ids` | 如果一条消息已经通过 `messages` 模式 token 级流过，values 快照到达时**不要**再合成一次完整 `messages-tuple` | messages 分支每发一个 AI/tool 事件就加入 | values 分支看到消息时检查 | -| `counted_usage_ids` | 同一个 `usage_metadata` 在 messages 末尾 chunk 和 values 快照的 final AIMessage 里各带一份，**累计总量只算一次** | `_account_usage()` 每次接受 usage 就加入 | `_account_usage()` 每次调用时检查 | +| `counted_usage_ids` | 完整 AIMessage 的 usage 只核算一次，已累加的 chunk 部分不再重复计入 | `_account_usage()` 接受非 chunk 的 usage 时加入 | `_account_usage()` 每次调用时检查 | | `historical_message_ids` | 续聊时旧消息不再作为本轮增量事件，也不计入本轮 `end.usage` | values 快照里当前用户消息之前的 id | messages 和 values 分支发射前检查 | ### 为什么不能只用一个 set @@ -253,19 +254,20 @@ sequenceDiagram Note over M: 第一个 AI text chunk 到达 M->>SS: add(msg_id) - Note over M: 最后一个 chunk 带 usage - M->>SU: add(msg_id) + Note over M: 每个带 usage 的 chunk 都累加到 streamed_usage_by_id Note over V: snapshot 到达，包含同一条 AI message V->>SE: add(msg_id) V->>SS: 查询 → 已存在，跳过文本合成 - V->>SU: 查询 → 已存在，不重复计数 + V->>SU: 核算完整 usage 与 chunk 累计之差，add(msg_id) ``` - `seen_ids` **永远在 values 快照到达时**加入，所以它是 "values 已处理" 的标记。一条只出现在 messages 流里的消息（罕见但可能），`seen_ids` 里永远没有它。 - `streamed_ids` **在 messages 流的第一个有效事件时**加入。一条只通过 values 快照到达的非 AI 消息（HumanMessage、被 truncate 的 tool 消息），`streamed_ids` 里永远没有它。 -- `counted_usage_ids` **只在看到非空 `usage_metadata` 时**加入。一条完全没有 usage 的消息（tool message、错误消息）永远不会进去。 +- `counted_usage_ids` 只在看到完整消息的非空 `usage_metadata` 时加入。带用量的 `AIMessageChunk` 是增量，会逐段累加，不能收到第一段就将该 id 标为完成。 -**集合包含关系**：`counted_usage_ids ⊆ (streamed_ids ∪ seen_ids)` 大致成立，但**不是严格子集**，因为一条消息可以在 messages 模式流完 text 但**在最后那个带 usage 的 chunk 之前**就被 values snapshot 赶上——此时它已经在 `streamed_ids` 里，但还不在 `counted_usage_ids` 里。把它们合并成一个 dict-of-flags 会让这个微妙的时序依赖**从类型系统里消失**，变成注释里的一句话。三个独立的 set 把不变式显式化了：每个 set 名对应一个可以口头回答的问题。 +`streamed_usage_by_id` 保存每条消息已计入总量的 chunk 用量。完整消息到达时，各字段只计入 `max(完整用量 - chunk 累计, 0)`，再清除该累计值并标记 `counted_usage_ids`。中间件若调低完整消息的用量，不会扣回已经报告的 chunk 用量；这只限制完整消息的核算差值，chunk 内的有符号供应商修正仍按增量累加。重复的完整消息或晚到的同 id chunk 不会再次计数；没有完整快照时，已收到的 chunk 用量仍保留在总量中。没有 id 的消息无法跨模式匹配，维持逐事件计数。 + +例如原生 Gemini 适配器把服务端累计值转换成 chunk 增量：先 `10/1/11`，再 `0/3/3`（输入/输出/总量），最终 AIMessage 为 `10/4/14`。客户端必须累加两段并避免快照重复计数，而不是在第一段后忽略该 id。纯用量 chunk 即使没有文本也要计入；适配器返回的输入 token 修正值同样按增量累加。`tests/test_client_stream_usage.py` 用真实 Gemini、vLLM 和 OpenAI 适配器及 LangGraph、仅替换 SDK 网络调用，覆盖增量用量和仅在末尾汇总用量的路径。 --- @@ -294,11 +296,11 @@ sequenceDiagram Note over A: LLM finish_reason=stop，最后一个 chunk 带 usage A-->>C: ("messages", (AIMessageChunk(content="", usage_metadata={...}), meta)) - C->>C: counted_usage_ids.add(ai-1)
(无文本，不 yield) + C->>C: 累加 chunk usage
(无文本，不 yield) A-->>C: ("values", {messages: [..., AIMessage(complete)]}) C->>C: ai-1 in streamed_ids → 跳过合成 - C->>C: 捕获 usage (已在 counted_usage_ids，no-op) + C->>C: 核算完整 usage 与 chunk 累计之差 (0)
counted_usage_ids.add(ai-1) C-->>U: StreamEvent(type="values", ...) C-->>U: StreamEvent(type="end", data={usage:{...}}) @@ -308,7 +310,7 @@ sequenceDiagram 1. 用户看到 **35 个 messages-tuple 事件**，跨越约 476ms，每个事件带一个 token delta 和同一个 `id=ai-1`。 2. 最后一个 `values` 快照里的 `AIMessage` **不会**再触发一个完整的 `messages-tuple` 事件——因为 `ai-1 in streamed_ids` 跳过了合成。 -3. `end` 事件里的 `usage` 正好等于那一份 cumulative usage，**不是它的两倍**——`counted_usage_ids` 在 messages 末尾 chunk 上已经吸收了，values 分支的重复访问是 no-op。 +3. `end` 事件里的 `usage` 正好等于那一份 cumulative usage，**不是它的两倍**：values 分支减去已累加的 chunk 部分，之后再通过 `counted_usage_ids` 排除重复完整消息。 4. 消费者拿到的 `content` 是**增量**："ele" 只包含 3 个字符，不是 "one\ntwo\n...ele"。想要完整文本要按 `id` 累加，`chat()` 已经帮你做了。 --- diff --git a/backend/packages/harness/deerflow/AGENTS.md b/backend/packages/harness/deerflow/AGENTS.md index 3b213e2774..8390040b28 100644 --- a/backend/packages/harness/deerflow/AGENTS.md +++ b/backend/packages/harness/deerflow/AGENTS.md @@ -72,7 +72,7 @@ drift. - `"values"` — state snapshot (title, messages, artifacts, summary_text). Always forward `summary_text` (current summary or `None`), including unchanged values/resets. Never re-emit AI text delivered via `messages`; serialized `ToolMessage` entries retain non-`None` native `artifact` - `"messages-tuple"` — current-turn AI text **deltas** by `id` and each tool call/result once; excludes resumed history and preserves result `artifact` - `"custom"` — forwarded from `StreamWriter`; DeerFlow-built-in custom events are dual-emitted through `deerflow.utils.custom_events`, so `astream_events(version="v2")` consumers also receive one `on_custom_event` with `name=payload["type"]` and the unchanged payload as `data` - - `"end"` — current-turn cumulative `usage`, counted once per message id + - `"end"` — turn `usage`: signed chunks + positive full-message remainders, deduped by id - **Custom-event invariant** — use `emit_custom_event` / `aemit_custom_event`, never `StreamWriter` alone. Built-in payloads require a non-empty string `type`; typeless payloads stay writer-only, absent from `astream_events`. The writer runs first and is authoritative for Gateway/Web UI/embedded clients; best-effort callbacks must not break it. Async graph hooks must await the async helper, never dispatch synchronously on a running event loop. - Lazy graph creation uses `create_agent()` + `build_middlewares()`. - Cache graphs by storage `user_id` and the unordered set of named-agent `mcp_plugins`. `stream()` materializes `user_id` before worker/loop boundaries in every auth mode. diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index a520f23226..75f44b9e21 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -1088,39 +1088,41 @@ class DeerFlowClient: # arguments only parse once the message is complete, so their # tool_calls event is emitted from the values snapshot instead. pending_tool_call_ids: set[str] = set() - # The same message id carries identical cumulative ``usage_metadata`` - # in both the final ``messages`` chunk and the values snapshot — - # count it only on whichever arrives first. + # Chunks carry additive usage; complete messages carry the total. + # Reconcile that total once, without counting the streamed part again. + streamed_usage_by_id: dict[str, dict[str, int]] = {} counted_usage_ids: set[str] = set() sent_additional_kwargs_by_id: dict[str, dict[str, Any]] = {} cumulative_usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} - def _account_usage(msg_id: str | None, usage: Any) -> dict | None: - """Add *usage* to cumulative totals if this id has not been counted. + def _account_usage(msg_id: str | None, usage: Any, *, is_chunk: bool = False) -> dict | None: + """Accumulate chunk deltas or reconcile a complete message once. ``usage`` is a ``langchain_core.messages.UsageMetadata`` TypedDict or ``None``; typed as ``Any`` because TypedDicts are not structurally assignable to plain ``dict`` under strict type - checking. Returns the normalized usage dict (for attaching + checking. Returns the newly accounted usage (for attaching to an event) when we accepted it, otherwise ``None``. """ if not usage: return None if msg_id and msg_id in counted_usage_ids: return None + delta = {key: usage.get(key, 0) or 0 for key in cumulative_usage} if msg_id: - counted_usage_ids.add(msg_id) - input_tokens = usage.get("input_tokens", 0) or 0 - output_tokens = usage.get("output_tokens", 0) or 0 - total_tokens = usage.get("total_tokens", 0) or 0 - cumulative_usage["input_tokens"] += input_tokens - cumulative_usage["output_tokens"] += output_tokens - cumulative_usage["total_tokens"] += total_tokens - return { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": total_tokens, - } + if is_chunk: + streamed = streamed_usage_by_id.setdefault(msg_id, {}) + for key, value in delta.items(): + streamed[key] = streamed.get(key, 0) + value + else: + counted_usage_ids.add(msg_id) + streamed = streamed_usage_by_id.pop(msg_id, {}) + # State rewrites may lower totals; do not retract streamed usage. + # Signed provider corrections remain valid in the chunk path. + delta = {key: max(value - streamed.get(key, 0), 0) for key, value in delta.items()} + for key, value in delta.items(): + cumulative_usage[key] += value + return delta def _unsent_additional_kwargs(msg_id: str | None, additional_kwargs: dict[str, Any] | None) -> dict[str, Any] | None: if not additional_kwargs: @@ -1167,7 +1169,7 @@ class DeerFlowClient: if isinstance(msg_chunk, AIMessage): text = self._extract_text(msg_chunk.content) additional_kwargs = self._serialize_additional_kwargs(msg_chunk) - counted_usage = _account_usage(msg_id, msg_chunk.usage_metadata) + counted_usage = _account_usage(msg_id, msg_chunk.usage_metadata, is_chunk=isinstance(msg_chunk, AIMessageChunk)) sent_additional_kwargs = False if text: diff --git a/backend/tests/test_client_stream_usage.py b/backend/tests/test_client_stream_usage.py new file mode 100644 index 0000000000..d697a54daf --- /dev/null +++ b/backend/tests/test_client_stream_usage.py @@ -0,0 +1,220 @@ +"""Embedded stream usage must match LangChain's assembled message usage.""" + +from unittest.mock import MagicMock, patch + +import pytest +from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage +from test_client import client as client +from test_client import mock_app_config as mock_app_config + + +@pytest.mark.parametrize(("last_text", "initial_input"), [(" world", 10), ("", 10), (" world", 12)]) +def test_stream_accounts_native_gemini_usage(client, last_text, initial_input): + from google.genai.types import Candidate, Content, GenerateContentResponse, GenerateContentResponseUsageMetadata, Part + from langchain.agents import create_agent + from langchain_google_genai import ChatGoogleGenerativeAI + + model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", api_key="offline-placeholder") + responses = [ + GenerateContentResponse( + candidates=[Candidate(content=Content(role="model", parts=[Part(text=text)]), finish_reason="STOP" if index else None)], + usage_metadata=GenerateContentResponseUsageMetadata(prompt_token_count=inputs, candidates_token_count=outputs, total_token_count=inputs + outputs), + ) + for index, (text, inputs, outputs) in enumerate([("Hello", initial_input, 1), (last_text, 10, 4)]) + ] + try: + # Only replace the SDK transport; retain the real adapter and graph. + with ( + patch.object(model.client.models, "generate_content_stream", return_value=iter(responses)) as transport, + patch.object(client, "_ensure_agent"), + patch.object(client, "_agent", create_agent(model, tools=[])), + ): + events = list(client.stream("Say hello", thread_id="gemini-usage")) + transport.assert_called_once() + finally: + model.client.close() + + expected = {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} + final_values = [event for event in events if event.type == "values"][-1] + assembled = final_values.data["messages"][-1] + assert {key: assembled["usage_metadata"][key] for key in expected} == expected + assert assembled["content"] == "Hello" + last_text + assert "".join(event.data["content"] for event in events if event.type == "messages-tuple") == assembled["content"] + assert events[-1].type == "end" + assert events[-1].data["usage"] == expected + + +@pytest.mark.parametrize("provider", ["openai", "vllm"]) +def test_stream_accounts_openai_compatible_usage(client, provider): + from langchain.agents import create_agent + from langchain_openai import ChatOpenAI + + from deerflow.models.vllm_provider import VllmChatModel + + model_class = VllmChatModel if provider == "vllm" else ChatOpenAI + model = model_class( + model="offline-model", + api_key="offline-placeholder", + use_responses_api=False, + stream_usage=True, + **({"cumulative_stream_usage": True} if provider == "vllm" else {}), + ) + responses = [ + { + "id": "chatcmpl-usage", + "model": "offline-model", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": "stop" if index else None}], + "usage": {"prompt_tokens": 10, "completion_tokens": output, "total_tokens": 10 + output} if provider == "vllm" else None, + } + for index, (text, output) in enumerate([("Hello", 1), (" world", 4)]) + ] + responses.append( + { + "id": "chatcmpl-usage", + "model": "offline-model", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, + } + ) + response_stream = MagicMock() + response_stream.__enter__.return_value = iter(responses) + try: + with ( + patch.object(model.client, "create", return_value=response_stream) as transport, + patch.object(client, "_ensure_agent"), + patch.object(client, "_agent", create_agent(model, tools=[])), + ): + events = list(client.stream("Say hello", thread_id=f"{provider}-usage")) + transport.assert_called_once() + finally: + model.root_client.close() + + expected = {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} + assembled = [event for event in events if event.type == "values"][-1].data["messages"][-1] + assert {key: assembled["usage_metadata"][key] for key in expected} == expected + assert assembled["content"] == "Hello world" + assert "".join(event.data["content"] for event in events if event.type == "messages-tuple") == "Hello world" + assert events[-1].type == "end" + assert events[-1].data["usage"] == expected + + +@pytest.mark.parametrize("include_snapshots", [False, True]) +def test_stream_sums_usage_deltas_per_message(client, include_snapshots): + chunks = [ + AIMessageChunk(id="a", content="A", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}), + AIMessageChunk(id="b", content="B", usage_metadata={"input_tokens": 20, "output_tokens": 2, "total_tokens": 22}), + AIMessageChunk(id="a", content="C", usage_metadata={"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}), + AIMessageChunk(id="b", content="", usage_metadata={"input_tokens": 0, "output_tokens": 4, "total_tokens": 4}), + ] + items = [("messages", (chunk, {})) for chunk in chunks] + if include_snapshots: + messages = [ + AIMessage(id="a", content="AC", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}), + AIMessage(id="b", content="B", usage_metadata={"input_tokens": 20, "output_tokens": 6, "total_tokens": 26}), + ] + items.extend([("values", {"messages": messages}), ("values", {"messages": [message.model_copy() for message in messages]})]) + agent = MagicMock() + agent.stream.return_value = iter(items) + with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent): + events = list(client.stream("hi")) + + assert events[-1].data["usage"] == {"input_tokens": 30, "output_tokens": 10, "total_tokens": 40} + text_events = [event for event in events if event.type == "messages-tuple"] + assert [event.data["content"] for event in text_events] == ["A", "B", "C"] + assert [event.data["usage_metadata"] for event in text_events] == [chunk.usage_metadata for chunk in chunks[:3]] + + +@pytest.mark.parametrize("snapshot_mode", ["messages", "values"]) +@pytest.mark.parametrize("snapshot_first", [False, True]) +def test_stream_reconciles_full_usage_without_double_counting(client, snapshot_mode, snapshot_first): + full = AIMessage(id="a", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}) + partial = AIMessageChunk(id="a", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}) + snapshot = (snapshot_mode, (full, {}) if snapshot_mode == "messages" else {"messages": [full]}) + items = [("messages", (partial, {}))] + items.insert(0 if snapshot_first else 1, snapshot) + items.append(snapshot) + agent = MagicMock() + agent.stream.return_value = iter(items) + with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent): + events = list(client.stream("hi")) + + assert events[-1].data["usage"] == full.usage_metadata + + +def test_resumed_stream_sums_only_current_usage_deltas(client): + old = AIMessage(id="old", content="Previous reply", usage_metadata={"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}) + first = AIMessageChunk(id="new", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}) + last = AIMessageChunk(id="new", content=" again", usage_metadata={"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}) + full = AIMessage(id="new", content="Hello again", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}) + + def stream(state, **kwargs): + current = state["messages"][0] + messages = [HumanMessage(id="previous-user", content="Previous request"), old, current] + yield "values", {"messages": messages} + yield "messages", (old, {}) + yield "messages", (first, {}) + yield "messages", (last, {}) + yield "values", {"messages": [*messages, full]} + + agent = MagicMock() + agent.stream.side_effect = stream + with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent): + events = list(client.stream("hi", thread_id="existing-thread")) + + assert events[-1].data["usage"] == full.usage_metadata + assert [event.data["content"] for event in events if event.type == "messages-tuple"] == ["Hello", " again"] + + +def test_stream_keeps_no_id_usage_deltas(client): + agent = MagicMock() + agent.stream.return_value = iter( + [ + ("messages", (AIMessageChunk(content="Hi", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}), {})), + ("messages", (AIMessageChunk(content="", usage_metadata={"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}), {})), + ] + ) + with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent): + events = list(client.stream("hi")) + + assert events[-1].data["usage"] == {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14} + + +@pytest.mark.parametrize("snapshot_mode", ["messages", "values"]) +@pytest.mark.parametrize( + ("final_usage", "expected"), + [ + ({"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}), + ({"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}), + ({"input_tokens": 5, "output_tokens": 12, "total_tokens": 17}, {"input_tokens": 10, "output_tokens": 12, "total_tokens": 17}), + ], + ids=["lower", "zero", "mixed"], +) +def test_stream_snapshot_cannot_subtract_already_accounted_usage(client, snapshot_mode, final_usage, expected): + # A middleware may rewrite usage in the completed message. Reconciliation + # may add missing tokens, but must not retract usage already reported. + partial = AIMessageChunk(id="a", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}) + full = AIMessage(id="a", content="Hello again", usage_metadata=final_usage) + snapshot = (snapshot_mode, (full, {}) if snapshot_mode == "messages" else {"messages": [full]}) + duplicate = full.model_copy() + duplicate_snapshot = (snapshot_mode, (duplicate, {}) if snapshot_mode == "messages" else {"messages": [duplicate]}) + other = AIMessage(id="b", content="Another reply", usage_metadata={"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}) + agent = MagicMock() + agent.stream.return_value = iter( + [ + ("messages", (partial, {})), + snapshot, + duplicate_snapshot, + ("messages", (partial.model_copy(update={"content": ""}), {})), + ("messages", (other, {})), + ] + ) + with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent): + events = list(client.stream("hi")) + + assert events[-1].data["usage"] == {key: value + other.usage_metadata[key] for key, value in expected.items()} + # Raw values remain faithful to state; only accounting is clamped. + if snapshot_mode == "values": + assert [event for event in events if event.type == "values"][-1].data["messages"][0]["usage_metadata"] == final_usage + else: + reconciled = next(event for event in events if event.type == "messages-tuple" and event.data["content"] == "Hello again") + assert reconciled.data["usage_metadata"] == {key: max(final_usage[key] - partial.usage_metadata[key], 0) for key in expected}