From e862099dffb26e6b31bcc12a65194ecca035b7b2 Mon Sep 17 00:00:00 2001
From: Alex Fournier <afournier@nvidia.com>
Date: Thu, 23 Jul 2026 13:32:44 -0700
Subject: [PATCH] fix(relay): merge Anthropic stream usage

Signed-off-by: Alex Fournier <afournier@nvidia.com>
---
 agent/relay_llm.py            |  7 ++++++-
 tests/agent/test_relay_llm.py | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/agent/relay_llm.py b/agent/relay_llm.py
index 8b799b6e775..5aaaac58e7b 100644
--- a/agent/relay_llm.py
+++ b/agent/relay_llm.py
@@ -598,7 +598,12 @@ class AnthropicStreamAccumulator:
                     if key in delta:
                         self._message[key] = delta[key]
             if "usage" in payload:
-                self._message["usage"] = payload["usage"]
+                usage = payload["usage"]
+                current_usage = self._message.get("usage")
+                if isinstance(current_usage, dict) and isinstance(usage, dict):
+                    self._message["usage"] = {**current_usage, **usage}
+                else:
+                    self._message["usage"] = usage
 
     def finalize(self) -> dict[str, Any]:
         blocks = []
diff --git a/tests/agent/test_relay_llm.py b/tests/agent/test_relay_llm.py
index 75367fb93a1..05c040eb068 100644
--- a/tests/agent/test_relay_llm.py
+++ b/tests/agent/test_relay_llm.py
@@ -247,6 +247,39 @@ def test_direct_stream_close_reaches_original_provider_resource(monkeypatch):
     assert provider_stream.closed is True
 
 
+def test_anthropic_stream_accumulator_merges_terminal_usage():
+    accumulator = relay_llm.AnthropicStreamAccumulator()
+    accumulator.observe({
+        "type": "message_start",
+        "message": {
+            "id": "message-1",
+            "type": "message",
+            "role": "assistant",
+            "model": "claude-test",
+            "usage": {
+                "input_tokens": 100,
+                "output_tokens": 1,
+                "cache_creation_input_tokens": 20,
+                "cache_read_input_tokens": 30,
+            },
+        },
+    })
+    accumulator.observe({
+        "type": "message_delta",
+        "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+        "usage": {"output_tokens": 12},
+    })
+
+    response = accumulator.finalize()
+
+    assert response["usage"] == {
+        "input_tokens": 100,
+        "output_tokens": 12,
+        "cache_creation_input_tokens": 20,
+        "cache_read_input_tokens": 30,
+    }
+
+
 def test_non_stream_preserves_raw_provider_response_identity(relay_turn):
     _relay, _turn = relay_turn
     raw_response = SimpleNamespace(model="test-model", content="raw")