Skip to content

Commit 9552dbc

Browse files
Tom-RyderTom-Ryder
committed
fix(agent): tolerate missing provider token usage in budget tracking
1 parent e454e39 commit 9552dbc

6 files changed

Lines changed: 71 additions & 116 deletions

File tree

‎Sources/AgentRunKit/Core/Agent+ContextBudget.swift‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,6 @@ extension Agent {
1313
return ContextBudgetPhase(config: budgetConfig, windowSize: windowSize)
1414
}
1515

16-
func requireBudgetUsage(_ usage: TokenUsage?, budgetPhase: ContextBudgetPhase?) throws -> TokenUsage? {
17-
guard budgetPhase != nil else { return usage }
18-
guard let usage else {
19-
throw AgentError.contextBudgetUsageUnavailable
20-
}
21-
return usage
22-
}
23-
2416
func applyBudgetPhase(
2517
_ budgetPhase: inout ContextBudgetPhase?,
2618
usage: TokenUsage,

‎Sources/AgentRunKit/Core/Agent.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ extension Agent {
167167
historyWasRewrittenLocally: &state.historyWasRewrittenLocally,
168168
requestContext: options.requestContext
169169
)
170-
let budgetUsage = try requireBudgetUsage(response.tokenUsage, budgetPhase: state.budgetPhase)
170+
let budgetUsage = response.tokenUsage
171171

172172
if let finishCall = try exclusiveFinishCall(in: response.toolCalls) {
173173
return try parseFinishResult(
@@ -351,7 +351,7 @@ extension Agent {
351351
continuation.yield(.make(.iterationCompleted(usage: usage, iteration: iterationNumber)))
352352
}
353353
state.messages.append(.assistant(iteration.toAssistantMessage()))
354-
let budgetUsage = try requireBudgetUsage(iteration.usage, budgetPhase: state.budgetPhase)
354+
let budgetUsage = iteration.usage
355355

356356
if try tryFinishOnTerminalEvent(
357357
iteration: iteration,

‎Sources/AgentRunKit/Core/AgentError.swift‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,6 @@ public enum AgentError: Error, Sendable, Equatable, LocalizedError {
6565
case malformedStream(MalformedStreamReason)
6666
case schemaInferenceFailed(type: String, message: String)
6767
case maxDepthExceeded(depth: Int)
68-
case contextBudgetUsageUnavailable
6968
case contextBudgetWindowSizeUnavailable
7069

7170
public var errorDescription: String? {
@@ -94,8 +93,6 @@ public enum AgentError: Error, Sendable, Equatable, LocalizedError {
9493
"Schema inference failed for '\(type)': \(message)"
9594
case let .maxDepthExceeded(depth):
9695
"Sub-agent max depth exceeded (current depth: \(depth))"
97-
case .contextBudgetUsageUnavailable:
98-
"Context budget requires provider-reported token usage for every budgeted turn"
9996
case .contextBudgetWindowSizeUnavailable:
10097
"Context budget requires a client contextWindowSize for usage-based features"
10198
}
@@ -115,8 +112,6 @@ public enum AgentError: Error, Sendable, Equatable, LocalizedError {
115112
case let .malformedStream(reason): "Error: Malformed stream: \(reason)"
116113
case let .schemaInferenceFailed(type, message): "Error: Schema inference failed for '\(type)': \(message)"
117114
case let .maxDepthExceeded(depth): "Error: Sub-agent max depth exceeded (current depth: \(depth))."
118-
case .contextBudgetUsageUnavailable:
119-
"Error: Context budget requires provider-reported token usage for every budgeted turn."
120115
case .contextBudgetWindowSizeUnavailable:
121116
"Error: Context budget requires a client contextWindowSize for usage-based features."
122117
}

‎Sources/AgentRunKit/Documentation.docc/Articles/ContextManagement.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ Use ``ContextBudget/formatted(_:)`` to render a human-readable annotation, eithe
9393
| `enableVisibility` | false | Appends a token usage annotation to history after each turn |
9494
| `visibilityFormat` | `.standard` | Format for the visibility annotation |
9595

96-
Budget features that track usage require the client to report both `contextWindowSize` and per-turn `tokenUsage`.
96+
Budget features require the client to report `contextWindowSize`. Per-turn `tokenUsage` is consumed when the provider reports it; iterations that return nil usage (transient SSE chunk drops, proxies that omit the usage block) are skipped and the next iteration with reported usage resumes tracking.
9797

9898
## Streaming Events
9999

‎Tests/AgentRunKitTests/AgentContextBudgetTests.swift‎

Lines changed: 66 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -628,80 +628,68 @@ struct ContextBudgetStreamingPruneTests {
628628
}
629629
}
630630

631-
// MARK: - Usage Requirements
631+
// MARK: - Missing Usage Behavior
632632

633-
struct ContextBudgetUsageRequirementTests {
634-
@Test func missingTokenUsagePreventsToolExecutionInRun() async throws {
635-
let counter = InvocationCounter()
636-
let noopTool = try Tool<NoopParams, NoopOutput, EmptyContext>(
637-
name: "noop",
638-
description: "noop",
639-
executor: { _, _ in
640-
await counter.increment()
641-
return NoopOutput()
642-
}
643-
)
644-
let client = BudgetMockLLMClient(
645-
responses: [
646-
AssistantMessage(content: "", toolCalls: [
647-
ToolCall(id: "call_1", name: "noop", arguments: "{}"),
648-
]),
649-
],
650-
contextWindowSize: 10000
633+
struct ContextBudgetMissingUsageBehaviorTests {
634+
@Test func visibilityWithoutContextWindowSizeThrowsInStream() async throws {
635+
let client = BudgetStreamingMockLLMClient(
636+
streamSequences: [[
637+
.toolCallStart(index: 0, id: "finish_1", name: "finish", kind: .function),
638+
.toolCallDelta(index: 0, arguments: #"{"content":"done"}"#),
639+
.finished(usage: TokenUsage(input: 3000, output: 200)),
640+
]],
641+
contextWindowSize: nil
651642
)
652643
let config = AgentConfiguration(
653644
contextBudget: ContextBudgetConfig(enableVisibility: true)
654645
)
655-
let agent = Agent<EmptyContext>(client: client, tools: [noopTool], configuration: config)
646+
let agent = Agent<EmptyContext>(client: client, tools: [], configuration: config)
656647

657-
await #expect(throws: AgentError.contextBudgetUsageUnavailable) {
658-
_ = try await agent.run(userMessage: "go", context: EmptyContext())
648+
await #expect(throws: AgentError.contextBudgetWindowSizeUnavailable) {
649+
for try await _ in agent.stream(userMessage: "go", context: EmptyContext()) {}
659650
}
660-
#expect(await counter.currentValue() == 0)
661651
}
662652

663-
@Test func missingTokenUsageThrowsInRun() async throws {
653+
@Test func missingTokenUsageContinuesRunAndResumesOnNextIteration() async throws {
654+
let counter = InvocationCounter()
664655
let noopTool = try Tool<NoopParams, NoopOutput, EmptyContext>(
665-
name: "noop", description: "noop", executor: { _, _ in NoopOutput() }
656+
name: "noop",
657+
description: "noop",
658+
executor: { _, _ in
659+
await counter.increment()
660+
return NoopOutput()
661+
}
666662
)
663+
let iter2Usage = TokenUsage(input: 400, output: 20)
664+
let iter3Usage = TokenUsage(input: 500, output: 30)
667665
let client = BudgetMockLLMClient(
668666
responses: [
669667
AssistantMessage(content: "", toolCalls: [
670668
ToolCall(id: "call_1", name: "noop", arguments: "{}"),
671669
]),
670+
AssistantMessage(content: "", toolCalls: [
671+
ToolCall(id: "call_2", name: "noop", arguments: "{}"),
672+
], tokenUsage: iter2Usage),
673+
AssistantMessage(content: "", toolCalls: [
674+
ToolCall(id: "f", name: "finish", arguments: #"{"content":"done"}"#),
675+
], tokenUsage: iter3Usage),
672676
],
673677
contextWindowSize: 10000
674678
)
675679
let config = AgentConfiguration(
676-
contextBudget: ContextBudgetConfig(enableVisibility: true)
680+
contextBudget: ContextBudgetConfig(softThreshold: 0.8, enableVisibility: true)
677681
)
678682
let agent = Agent<EmptyContext>(client: client, tools: [noopTool], configuration: config)
679683

680-
await #expect(throws: AgentError.contextBudgetUsageUnavailable) {
681-
_ = try await agent.run(userMessage: "go", context: EmptyContext())
682-
}
683-
}
684-
685-
@Test func missingTokenUsageOnFinishThrowsInRun() async throws {
686-
let client = BudgetMockLLMClient(
687-
responses: [
688-
AssistantMessage(content: "", toolCalls: [
689-
ToolCall(id: "finish_1", name: "finish", arguments: #"{"content":"done"}"#),
690-
]),
691-
],
692-
contextWindowSize: 10000
693-
)
694-
let config = AgentConfiguration(
695-
contextBudget: ContextBudgetConfig(enableVisibility: true)
696-
)
697-
let agent = Agent<EmptyContext>(client: client, tools: [], configuration: config)
684+
let result = try await agent.run(userMessage: "go", context: EmptyContext())
698685

699-
await #expect(throws: AgentError.contextBudgetUsageUnavailable) {
700-
_ = try await agent.run(userMessage: "go", context: EmptyContext())
701-
}
686+
#expect(try requireContent(result) == "done")
687+
#expect(result.iterations == 3)
688+
#expect(await counter.currentValue() == 2)
689+
#expect(result.totalTokenUsage == iter2Usage + iter3Usage)
702690
}
703691

704-
@Test func missingTokenUsageThrowsInStream() async throws {
692+
@Test func missingTokenUsageContinuesStreamAndResumesOnNextIteration() async throws {
705693
let counter = InvocationCounter()
706694
let noopTool = try Tool<NoopParams, NoopOutput, EmptyContext>(
707695
name: "noop",
@@ -711,60 +699,44 @@ struct ContextBudgetUsageRequirementTests {
711699
return NoopOutput()
712700
}
713701
)
702+
let iter2Usage = TokenUsage(input: 400, output: 20)
714703
let client = BudgetStreamingMockLLMClient(
715-
streamSequences: [[
716-
.toolCallStart(index: 0, id: "call_1", name: "noop", kind: .function),
717-
.toolCallDelta(index: 0, arguments: "{}"),
718-
.finished(usage: nil),
719-
]],
704+
streamSequences: [
705+
[
706+
.toolCallStart(index: 0, id: "call_1", name: "noop", kind: .function),
707+
.toolCallDelta(index: 0, arguments: "{}"),
708+
.finished(usage: nil),
709+
],
710+
[
711+
.toolCallStart(index: 0, id: "call_2", name: "noop", kind: .function),
712+
.toolCallDelta(index: 0, arguments: "{}"),
713+
.finished(usage: iter2Usage),
714+
],
715+
[
716+
.toolCallStart(index: 0, id: "f", name: "finish", kind: .function),
717+
.toolCallDelta(index: 0, arguments: #"{"content":"done"}"#),
718+
.finished(usage: TokenUsage(input: 500, output: 30)),
719+
],
720+
],
720721
contextWindowSize: 10000
721722
)
722723
let config = AgentConfiguration(
723-
contextBudget: ContextBudgetConfig(enableVisibility: true)
724+
contextBudget: ContextBudgetConfig(softThreshold: 0.8, enableVisibility: true)
724725
)
725726
let agent = Agent<EmptyContext>(client: client, tools: [noopTool], configuration: config)
726727

727-
await #expect(throws: AgentError.contextBudgetUsageUnavailable) {
728-
for try await _ in agent.stream(userMessage: "go", context: EmptyContext()) {}
729-
}
730-
#expect(await counter.currentValue() == 0)
731-
}
732-
733-
@Test func missingTokenUsageOnFinishThrowsInStream() async throws {
734-
let client = BudgetStreamingMockLLMClient(
735-
streamSequences: [[
736-
.toolCallStart(index: 0, id: "finish_1", name: "finish", kind: .function),
737-
.toolCallDelta(index: 0, arguments: #"{"content":"done"}"#),
738-
.finished(usage: nil),
739-
]],
740-
contextWindowSize: 10000
741-
)
742-
let config = AgentConfiguration(
743-
contextBudget: ContextBudgetConfig(enableVisibility: true)
744-
)
745-
let agent = Agent<EmptyContext>(client: client, tools: [], configuration: config)
746-
747-
await #expect(throws: AgentError.contextBudgetUsageUnavailable) {
748-
for try await _ in agent.stream(userMessage: "go", context: EmptyContext()) {}
728+
var budgetUpdates: [ContextBudget] = []
729+
var events: [StreamEvent.Kind] = []
730+
for try await event in agent.stream(userMessage: "go", context: EmptyContext()) {
731+
events.append(event.kind)
732+
if case let .budgetUpdated(budget) = event.kind {
733+
budgetUpdates.append(budget)
734+
}
749735
}
750-
}
751-
752-
@Test func visibilityWithoutContextWindowSizeThrowsInStream() async throws {
753-
let client = BudgetStreamingMockLLMClient(
754-
streamSequences: [[
755-
.toolCallStart(index: 0, id: "finish_1", name: "finish", kind: .function),
756-
.toolCallDelta(index: 0, arguments: #"{"content":"done"}"#),
757-
.finished(usage: TokenUsage(input: 3000, output: 200)),
758-
]],
759-
contextWindowSize: nil
760-
)
761-
let config = AgentConfiguration(
762-
contextBudget: ContextBudgetConfig(enableVisibility: true)
763-
)
764-
let agent = Agent<EmptyContext>(client: client, tools: [], configuration: config)
765736

766-
await #expect(throws: AgentError.contextBudgetWindowSizeUnavailable) {
767-
for try await _ in agent.stream(userMessage: "go", context: EmptyContext()) {}
768-
}
737+
#expect(events.contains(where: { if case .finished = $0 { true } else { false } }))
738+
#expect(await counter.currentValue() == 2)
739+
#expect(budgetUpdates.count == 1)
740+
#expect(budgetUpdates.first?.currentUsage == iter2Usage.inputOutputTotal)
769741
}
770742
}

‎Tests/AgentRunKitTests/AgentErrorTests.swift‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -153,13 +153,9 @@ struct AgentErrorTests {
153153
#expect(error10.feedbackMessage.contains("7"))
154154
#expect(error10.feedbackMessage.contains("name"))
155155

156-
let error11 = AgentError.contextBudgetUsageUnavailable
156+
let error11 = AgentError.contextBudgetWindowSizeUnavailable
157157
#expect(error11.feedbackMessage.contains("Context budget"))
158-
#expect(error11.feedbackMessage.contains("token usage"))
159-
160-
let error12 = AgentError.contextBudgetWindowSizeUnavailable
161-
#expect(error12.feedbackMessage.contains("Context budget"))
162-
#expect(error12.feedbackMessage.contains("contextWindowSize"))
158+
#expect(error11.feedbackMessage.contains("contextWindowSize"))
163159
}
164160
}
165161

0 commit comments

Comments
 (0)