From 4b66707840b6b67f10fc8a617c819f963785c0bb Mon Sep 17 00:00:00 2001 From: Ayman Hamed Date: Fri, 21 Aug 2026 12:02:24 +0300 Subject: [PATCH] Streaming path: preserve provider token usage executeTurn's streaming branch discarded the usage on the final `.finish` chunk (`case .finish(let reason, _)`) and synthesized the response with `usage: nil`, forcing cost/context onto a local estimate. Capture the usage and thread it onto the returned response (text and tool-call paths). Adds ReplayRun.runStreaming to exercise the streaming path, and a regression test asserting the emitted response carries the provider's promptTokens/completionTokens. Co-Authored-By: Claude Opus 4.8 (1M context) --- Sources/SwiftAgentKit/Core/Agent.swift | 11 +++++-- Sources/SwiftAgentKitReplay/ReplayRun.swift | 15 ++++++++++ .../StreamingUsageTests.swift | 30 +++++++++++++++++++ 3 files changed, 54 insertions(+), 2 deletions(-) create mode 100644 Tests/SwiftAgentKitReplayTests/StreamingUsageTests.swift diff --git a/Sources/SwiftAgentKit/Core/Agent.swift b/Sources/SwiftAgentKit/Core/Agent.swift index a718944..ac5e503 100644 --- a/Sources/SwiftAgentKit/Core/Agent.swift +++ b/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 diff --git a/Sources/SwiftAgentKitReplay/ReplayRun.swift b/Sources/SwiftAgentKitReplay/ReplayRun.swift index 0f30ffd..08f57bf 100644 --- a/Sources/SwiftAgentKitReplay/ReplayRun.swift +++ b/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] { diff --git a/Tests/SwiftAgentKitReplayTests/StreamingUsageTests.swift b/Tests/SwiftAgentKitReplayTests/StreamingUsageTests.swift new file mode 100644 index 0000000..7d62924 --- /dev/null +++ b/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) +}