1 files changed, 118 insertions, 11 deletions
diff --git a/libpanto/src/provider_openai_chat.zig b/libpanto/src/provider_openai_chat.zig
index ca8fa4c..fcbc0bb 100644
--- a/libpanto/src/provider_openai_chat.zig
+++ b/libpanto/src/provider_openai_chat.zig
@@ -207,14 +207,17 @@ pub const OpenAIChatProvider = struct {
                     return;
                 }
                 try handleEvent(self.allocator, ev_payload, &state, receiver);
-                if (state.end_of_stream) {
-                    try state.finalize(receiver, conv);
-                    return;
-                }
+                // Note: we do NOT bail when state.end_of_stream is set.
+                // OpenAI emits the terminating `usage` chunk *after* the
+                // chunk carrying finish_reason, then sends `[DONE]`. If
+                // we returned on finish_reason we'd never capture usage.
+                // `[DONE]` is the authoritative end-of-stream marker.
             }
         }
 
-        // Stream ended without [DONE] or finish_reason. Finalize anyway.
+        // Stream ended without [DONE]. Some servers and proxies omit it
+        // (or drop the trailing usage chunk). Finalize with whatever we've
+        // got — usage will be null in that case, which is fine.
         try state.finalize(receiver, conv);
     }
 };
@@ -266,6 +269,11 @@ const StreamState = struct {
     /// whose block we've already emitted.
     closed_tool_indices: std.AutoHashMap(usize, void),
 
+    /// Token counts from the terminating chunk's `usage` block. Only
+    /// populated when the server sent `usage` (i.e. the request used
+    /// `stream_options.include_usage: true` AND the server honored it).
+    usage: ?provider_mod.Usage = null,
+
     const ToolUseInProgress = struct {
         /// Block index emitted to the receiver for this tool call's
         /// onBlockStart / onContentDelta / onBlockComplete callbacks.
@@ -542,7 +550,7 @@ const StreamState = struct {
         try conv.addAssistantMessage(moved_blocks);
 
         const msg = conv.messages.items[conv.messages.items.len - 1];
-        try receiver.onMessageComplete(msg);
+        try receiver.onMessageComplete(msg, self.usage);
     }
 };
 
@@ -556,6 +564,24 @@ fn handleEvent(
     defer parsed.deinit();
     const d = parsed.delta;
 
+    // Usage block arrives in the terminating chunk (after finish_reason,
+    // with an empty `choices` array). Capture it; `finalize` delivers it
+    // as part of `onMessageComplete`. OpenAI bills `prompt_tokens` as
+    // the *total* input including cached tokens; we split them so
+    // callers don't have to.
+    if (d.usage) |u| {
+        const prompt: u64 = u.prompt_tokens orelse 0;
+        const cached: u64 = u.cached_prompt_tokens orelse 0;
+        const fresh: u64 = if (cached > prompt) 0 else prompt - cached;
+        state.usage = .{
+            .input = fresh,
+            .output = u.completion_tokens orelse 0,
+            .cache_read = cached,
+            .cache_write = 0, // OpenAI doesn't bill a cache-write premium.
+            .reasoning = u.reasoning_tokens orelse 0,
+        };
+    }
+
     // Mid-stream provider error: some OpenAI-compatible endpoints (and
     // OpenAI itself on rare transient failures) return HTTP 200 with an
     // error embedded in the SSE stream. Treat the turn as failed.
@@ -631,7 +657,7 @@ const NoopReceiver = struct {
     fn noopToolDetails(_: *anyopaque, _: usize, _: []const u8, _: []const u8) anyerror!void {}
     fn noopDelta(_: *anyopaque, _: usize, _: []const u8) anyerror!void {}
     fn noopBlockComplete(_: *anyopaque, _: usize, _: conversation.ContentBlock) anyerror!void {}
-    fn noopMsgComplete(_: *anyopaque, _: conversation.Message) anyerror!void {}
+    fn noopMsgComplete(_: *anyopaque, _: conversation.Message, _: ?provider_mod.Usage) anyerror!void {}
     fn noopErr(_: *anyopaque, _: anyerror) void {}
 };
 
@@ -648,8 +674,10 @@ fn runStreamedTurn(
 
     for (events) |payload| {
         if (std.mem.eql(u8, payload, "[DONE]")) break;
+        // Process every chunk through to [DONE], including the
+        // post-finish_reason usage chunk. Mirrors the production loop
+        // in OpenAIChatProvider.streamStep.
         try handleEvent(allocator, payload, &state, receiver);
-        if (state.end_of_stream) break;
     }
     try state.finalize(receiver, conv);
 }
@@ -712,6 +740,76 @@ test "two streamed turns persist assistant replies in the conversation" {
     );
 }
 
+test "openai_chat: terminating usage chunk lands on message_complete with split cache_read" {
+    const allocator = testing.allocator;
+
+    var conv = conversation.Conversation.init(allocator);
+    defer conv.deinit();
+    try conv.addUserMessage("hi");
+
+    var rec = RecordingReceiver{ .allocator = allocator };
+    defer {
+        for (rec.events.items) |s| allocator.free(s);
+        rec.events.deinit(allocator);
+    }
+    var recv = rec.receiver();
+
+    const events = [_][]const u8{
+        \\{"choices":[{"delta":{"role":"assistant"}}]}
+        ,
+        \\{"choices":[{"delta":{"content":"hi"}}]}
+        ,
+        \\{"choices":[{"delta":{},"finish_reason":"stop"}]}
+        ,
+        // OpenAI's terminating chunk: empty choices, top-level usage.
+        \\{"choices":[],"usage":{"prompt_tokens":150,"completion_tokens":42,"prompt_tokens_details":{"cached_tokens":120},"completion_tokens_details":{"reasoning_tokens":18}}}
+        ,
+        "[DONE]",
+    };
+    try runStreamedTurn(allocator, &conv, &recv, &events);
+
+    var found: ?[]const u8 = null;
+    for (rec.events.items) |s| {
+        if (std.mem.startsWith(u8, s, "msg_complete[")) found = s;
+    }
+    try testing.expect(found != null);
+    // 150 prompt - 120 cached = 30 fresh input.
+    try testing.expectEqualStrings("msg_complete[usage:in=30,out=42,cr=120,cw=0,rsn=18]", found.?);
+}
+
+test "openai_chat: omitted stream usage yields null on message_complete" {
+    const allocator = testing.allocator;
+
+    var conv = conversation.Conversation.init(allocator);
+    defer conv.deinit();
+    try conv.addUserMessage("hi");
+
+    var rec = RecordingReceiver{ .allocator = allocator };
+    defer {
+        for (rec.events.items) |s| allocator.free(s);
+        rec.events.deinit(allocator);
+    }
+    var recv = rec.receiver();
+
+    const events = [_][]const u8{
+        \\{"choices":[{"delta":{"role":"assistant"}}]}
+        ,
+        \\{"choices":[{"delta":{"content":"hi"}}]}
+        ,
+        \\{"choices":[{"delta":{},"finish_reason":"stop"}]}
+        ,
+        "[DONE]",
+    };
+    try runStreamedTurn(allocator, &conv, &recv, &events);
+
+    var found: ?[]const u8 = null;
+    for (rec.events.items) |s| {
+        if (std.mem.startsWith(u8, s, "msg_complete")) found = s;
+    }
+    try testing.expect(found != null);
+    try testing.expectEqualStrings("msg_complete[usage:null]", found.?);
+}
+
 test "fragmented tool_call id and name are reassembled" {
     // Lenient OpenAI-compatible providers occasionally split `id` and
     // `function.name` across multiple deltas instead of sending them whole
@@ -817,9 +915,16 @@ const RecordingReceiver = struct {
         const self: *RecordingReceiver = @ptrCast(@alignCast(ptr));
         try self.recordFmt("block_complete[{d}]", .{idx});
     }
-    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.record("msg_complete");
+        if (usage) |u| {
+            try self.recordFmt(
+                "msg_complete[usage:in={d},out={d},cr={d},cw={d},rsn={d}]",
+                .{ u.input, u.output, u.cache_read, u.cache_write, u.reasoning },
+            );
+        } else {
+            try self.record("msg_complete[usage:null]");
+        }
     }
     fn onError(_: *anyopaque, _: anyerror) void {}
 };
@@ -878,7 +983,9 @@ test "parallel tool_calls emit one complete start/delta/complete cycle per block
         "tool_details[3]:c3:ping",
         "delta[3]:{\"host\":\"d\"}",
         "block_complete[3]",
-        "msg_complete",
+        // No usage chunk in this fixture (older test data) — record
+        // shows null.
+        "msg_complete[usage:null]",
     };
 
     // Identity arrives in the assembled ContentBlock at completion time.