1 files changed, 119 insertions, 7 deletions

diff --git a/libpanto/src/provider_anthropic_messages.zig b/libpanto/src/provider_anthropic_messages.zig
index 96649b6..5d5520a 100644
--- a/libpanto/src/provider_anthropic_messages.zig
+++ b/libpanto/src/provider_anthropic_messages.zig
@@ -216,6 +216,15 @@ const StreamState = struct {
     /// Assembled blocks for the final message, in stream order.
     blocks: std.ArrayList(conversation.ContentBlock) = .empty,
 
+    /// Accumulated token counts. Anthropic reports the input-side counts
+    /// on `message_start.usage` and the final `output_tokens` (plus
+    /// possibly updated input-side counts) on `message_delta.usage`.
+    /// `usage_seen` distinguishes "genuinely all zero" from "never
+    /// reported" — the former stamps a real `Usage` on
+    /// `onMessageComplete`, the latter stamps `null`.
+    usage: provider_mod.Usage = .{},
+    usage_seen: bool = false,
+
     const ActiveBlock = struct {
         /// Index reported on the wire (Anthropic's content-array index).
         wire_index: usize,
@@ -252,6 +261,22 @@ const StreamState = struct {
         try receiver.onMessageStart(.assistant);
     }
 
+    /// Merge a wire-level usage snapshot into the accumulated counts.
+    /// Missing fields mean "unchanged," not "reset to zero." Marks
+    /// `usage_seen` so `finalize` delivers a non-null `Usage` to
+    /// `onMessageComplete`.
+    fn mergeUsage(self: *StreamState, partial: json_mod.StreamUsage) void {
+        if (partial.input_tokens) |v| self.usage.input = v;
+        if (partial.output_tokens) |v| self.usage.output = v;
+        if (partial.cache_creation_input_tokens) |v| self.usage.cache_write = v;
+        if (partial.cache_read_input_tokens) |v| self.usage.cache_read = v;
+        if (partial.input_tokens != null or partial.output_tokens != null or
+            partial.cache_creation_input_tokens != null or partial.cache_read_input_tokens != null)
+        {
+            self.usage_seen = true;
+        }
+    }
+
     fn openBlock(
         self: *StreamState,
         receiver: *provider_mod.Receiver,
@@ -402,7 +427,8 @@ const StreamState = struct {
         try conv.addAssistantMessage(moved_blocks);
 
         const msg = conv.messages.items[conv.messages.items.len - 1];
-        try receiver.onMessageComplete(msg);
+        const usage: ?provider_mod.Usage = if (self.usage_seen) self.usage else null;
+        try receiver.onMessageComplete(msg, usage);
     }
 };
 
@@ -416,8 +442,9 @@ fn handleEvent(
     defer parsed.deinit();
 
     switch (parsed.event) {
-        .message_start => {
+        .message_start => |s| {
             try state.ensureStarted(receiver);
+            state.mergeUsage(s.usage);
         },
         .content_block_start => |s| {
             try state.ensureStarted(receiver);
@@ -438,9 +465,10 @@ fn handleEvent(
         .content_block_stop => {
             try state.closeBlock(receiver);
         },
-        .message_delta => {
+        .message_delta => |d| {
             // We don't act on stop_reason directly; message_stop is the
             // authoritative end-of-stream signal.
+            state.mergeUsage(d.usage);
         },
         .message_stop => {
             state.end_of_stream = true;
@@ -486,7 +514,7 @@ const RecordingReceiver = struct {
             text: []const u8, // owned copy
             signature: ?[]const u8 = null, // owned copy when present
         },
-        message_complete,
+        message_complete: ?provider_mod.Usage,
         err: anyerror,
     };
 
@@ -578,9 +606,9 @@ const RecordingReceiver = struct {
             else => {},
         }
     }
-    fn onMessageComplete(ptr: *anyopaque, _: conversation.Message) anyerror!void {
+    fn onMessageComplete(ptr: *anyopaque, _: conversation.Message, usage: ?provider_mod.Usage) anyerror!void {
         const self: *RecordingReceiver = @ptrCast(@alignCast(ptr));
-        try self.events.append(self.allocator, .message_complete);
+        try self.events.append(self.allocator, .{ .message_complete = usage });
     }
     fn onError(ptr: *anyopaque, err: anyerror) void {
         const self: *RecordingReceiver = @ptrCast(@alignCast(ptr));
@@ -650,7 +678,91 @@ test "streams a text-only turn end-to-end" {
     try testing.expectEqualStrings("Hello", rec.events.items[2].delta.bytes);
     try testing.expectEqualStrings("!", rec.events.items[3].delta.bytes);
     try testing.expectEqualStrings("Hello!", rec.events.items[4].block_complete.text);
-    try testing.expectEqual(@as(@TypeOf(rec.events.items[5]), .message_complete), rec.events.items[5]);
+    try testing.expect(rec.events.items[5] == .message_complete);
+    // No usage on the wire — the assertion is structural.
+    try testing.expect(rec.events.items[5].message_complete == null);
+}
+
+test "anthropic: captures usage from message_start and message_delta on message_complete" {
+    const allocator = testing.allocator;
+
+    var conv = conversation.Conversation.init(allocator);
+    defer conv.deinit();
+    try conv.addUserMessage("hi");
+
+    var rec = RecordingReceiver.init(allocator);
+    defer rec.deinit();
+    var recv = rec.receiver();
+
+    const events = [_][]const u8{
+        // Initial input-side counts on message_start.
+        \\{"type":"message_start","message":{"id":"m","type":"message","role":"assistant","content":[],"model":"claude","usage":{"input_tokens":100,"cache_creation_input_tokens":50,"cache_read_input_tokens":200}}}
+        ,
+        \\{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+        ,
+        \\{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}
+        ,
+        \\{"type":"content_block_stop","index":0}
+        ,
+        // Final output count on message_delta.
+        \\{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":42}}
+        ,
+        \\{"type":"message_stop"}
+        ,
+    };
+
+    try runStreamedTurn(allocator, &conv, &recv, &events);
+
+    // Find the message_complete event and check its usage payload.
+    var found: ?provider_mod.Usage = null;
+    for (rec.events.items) |ev| {
+        if (ev == .message_complete) found = ev.message_complete;
+    }
+    try testing.expect(found != null);
+    const u = found.?;
+    try testing.expectEqual(@as(u64, 100), u.input);
+    try testing.expectEqual(@as(u64, 42), u.output);
+    try testing.expectEqual(@as(u64, 200), u.cache_read);
+    try testing.expectEqual(@as(u64, 50), u.cache_write);
+    try testing.expectEqual(@as(u64, 0), u.reasoning); // Anthropic doesn't split reasoning separately.
+}
+
+test "anthropic: message_complete carries null usage when wire omits it" {
+    const allocator = testing.allocator;
+
+    var conv = conversation.Conversation.init(allocator);
+    defer conv.deinit();
+    try conv.addUserMessage("hi");
+
+    var rec = RecordingReceiver.init(allocator);
+    defer rec.deinit();
+    var recv = rec.receiver();
+
+    const events = [_][]const u8{
+        \\{"type":"message_start","message":{"id":"m","role":"assistant","content":[],"model":"claude"}}
+        ,
+        \\{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+        ,
+        \\{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}
+        ,
+        \\{"type":"content_block_stop","index":0}
+        ,
+        \\{"type":"message_delta","delta":{"stop_reason":"end_turn"}}
+        ,
+        \\{"type":"message_stop"}
+        ,
+    };
+
+    try runStreamedTurn(allocator, &conv, &recv, &events);
+
+    var saw_complete = false;
+    for (rec.events.items) |ev| {
+        if (ev == .message_complete) {
+            saw_complete = true;
+            try testing.expect(ev.message_complete == null);
+        }
+    }
+    try testing.expect(saw_complete);
 }
 
 test "captures thinking signature for round-trip" {