Streaming path: preserve provider token usage

Sources/SwiftAgentKit/Core/Agent.swift

@@ -1252,6 +1252,11 @@ public actor Agent {
         var streamedText = ""
         var streamedToolCalls: [LLMToolCall] = []
         var sawNativeToolSignal = false
+        // Providers report real token usage on the final `.finish` chunk; capture
+        // it so the synthesized streaming response carries the model's actual
+        // consumed tokens rather than dropping them (which forced cost/context
+        // onto a local estimate).
+        var streamedUsage: LLMUsage? = nil
         for try await chunk in config.provider.stream(request) {
             switch chunk {
             case .text(let text):
@@ -1261,7 +1266,8 @@ public actor Agent {
             case .toolCall(let call):
                 streamedToolCalls.append(call)
                 sawNativeToolSignal = true
-            case .finish(let reason, _):
+            case .finish(let reason, let usage):
+                if let usage { streamedUsage = usage }
                 if reason == .toolCalls { sawNativeToolSignal = true }
             case .error(let error):
                 throw error
@@ -1277,6 +1283,7 @@ public actor Agent {
             let response = LLMResponse(
                 text: streamedText,
                 finishReason: .toolCalls,
+                usage: streamedUsage,
                 toolCalls: streamedToolCalls,
                 request: request,
                 providerName: type(of: config.provider).name
@@ -1307,7 +1314,7 @@ public actor Agent {
         let synthesized = LLMResponse(
             text: streamedText,
             finishReason: .stop,
-            usage: nil,
+            usage: streamedUsage,
             toolCalls: [],
             request: request,
             providerName: type(of: config.provider).name

Sources/SwiftAgentKitReplay/ReplayRun.swift

@@ -52,6 +52,21 @@ public final class ReplayRun: @unchecked Sendable {
         }
     }
 
+    /// Run via the STREAMING path (`onText` non-nil) to completion, returning
+    /// the concatenated streamed text. Use this to exercise streaming-only
+    /// behavior (e.g. provider usage captured from the final `.finish` chunk).
+    @discardableResult
+    public func runStreaming(_ query: String) async throws -> String {
+        let observer = agent.onEvent { [weak self] event in
+            guard let self else { return }
+            self.lock.lock(); self._events.append(event); self.lock.unlock()
+        }
+        defer { agent.removeObserver(observer) }
+        var full = ""
+        for try await chunk in agent.runStreaming(query) { full += chunk }
+        return full
+    }
+
     public var capturedRequests: [LLMRequest] { provider.capturedRequests }
 
     public var events: [AgentEvent] {

Tests/SwiftAgentKitReplayTests/StreamingUsageTests.swift

@@ -0,0 +1,30 @@
+import Testing
+import Foundation
+import LLMProviderKit
+import SwiftAgentKit
+@testable import SwiftAgentKitReplay
+
+/// Regression: the streaming path used to DROP the provider's token usage
+/// (`case .finish(let reason, _)`), synthesizing the final response with
+/// `usage: nil` — which forced cost/context onto a local estimate. The
+/// provider reports real usage on the final `.finish` chunk; it must survive
+/// onto the response the agent emits.
+@Test func streamingResponseCarriesProviderUsage() async throws {
+    let scenario = Scenario(name: "usage", turns: [
+        ScriptedTurn(
+            text: "Done.",
+            finishReason: .stop,
+            usage: LLMUsage(promptTokens: 123, completionTokens: 45, totalTokens: 168)
+        ),
+    ])
+    let run = ReplayRun(scenario: scenario)
+    _ = try await run.runStreaming("hi")
+
+    let usage: AgentTokenUsage? = run.events.compactMap { event in
+        if case .llmCallCompleted(_, let response) = event { return response.usage }
+        return nil
+    }.first ?? nil
+
+    #expect(usage?.promptTokens == 123)
+    #expect(usage?.completionTokens == 45)
+}