refactor(agent): model structural terminal finish reasons

Author: Tom-Ryder

README.md

@@ -40,9 +40,13 @@ let weatherTool = try Tool<WeatherParams, String, EmptyContext>(
 
 let agent = Agent(client: client, tools: [weatherTool])
 let result = try await agent.run(userMessage: "What's the weather in SF?", context: EmptyContext())
-print(result.content)
+if let content = result.content {
+    print(content)
+}
 ```
 
+`result.content` is optional. Completed runs return finish-tool content, while structural terminal reasons such as max iterations or token budget exhaustion surface through `result.finishReason` with no final content.
+
 ---
 
 ## Documentation

Sources/AgentRunKit/Core/Agent.swift

@@ -178,7 +178,12 @@ extension Agent {
             }
 
             if let tokenBudget = options.tokenBudget, totalUsage.total > tokenBudget {
-                throw AgentError.tokenBudgetExceeded(budget: tokenBudget, used: totalUsage.total)
+                return makeTerminalResult(
+                    reason: .tokenBudgetExceeded(budget: tokenBudget, used: totalUsage.total),
+                    tokenUsage: totalUsage,
+                    iterations: iteration,
+                    history: messages
+                )
             }
 
             let pruneCalls = response.toolCalls.filter { $0.name == "prune_context" }
@@ -194,7 +199,12 @@ extension Agent {
             }
         }
 
-        throw AgentError.maxIterationsReached(iterations: configuration.maxIterations)
+        return makeTerminalResult(
+            reason: .maxIterationsReached(limit: configuration.maxIterations),
+            tokenUsage: totalUsage,
+            iterations: configuration.maxIterations,
+            history: messages
+        )
     }
 
     func stream(
@@ -216,7 +226,7 @@ extension Agent {
     }
 }
 
-private extension Agent {
+extension Agent {
     func stream(
         userMessage: ChatMessage,
         history: [ChatMessage],
@@ -258,17 +268,14 @@ private extension Agent {
         continuation: AsyncThrowingStream<StreamEvent, Error>.Continuation
     ) async throws {
         var messages = buildInitialMessages(
-            userMessage: userMessage, history: history,
-            systemPromptOverride: options.systemPromptOverride
+            userMessage: userMessage, history: history, systemPromptOverride: options.systemPromptOverride
         )
         var totalUsage = TokenUsage()
         var lastTotalTokens: Int?
         var sessionAllowlist: Set<String> = []
         let policy = StreamPolicy.agent
         let processor = StreamProcessor(client: client, toolDefinitions: toolDefinitions, policy: policy)
-        var compactor = ContextCompactor(
-            client: client, toolDefinitions: toolDefinitions, configuration: configuration
-        )
+        var compactor = ContextCompactor(client: client, toolDefinitions: toolDefinitions, configuration: configuration)
         var budgetPhase = try makeBudgetPhase()
 
         for iterationNumber in 1 ... configuration.maxIterations {
@@ -277,63 +284,57 @@ private extension Agent {
             let compacted = await compactor.compactOrTruncateIfNeeded(
                 &messages, lastTotalTokens: lastTotalTokens, totalUsage: &totalUsage
             )
-            if compacted, let totalTokens = lastTotalTokens, let windowSize = client.contextWindowSize {
-                continuation.yield(.make(.compacted(totalTokens: totalTokens, windowSize: windowSize)))
-            }
+            emitCompactionEventIfNeeded(compacted, lastTotalTokens: lastTotalTokens, continuation: continuation)
             let iteration = try await processor.process(
-                messages: messages,
-                totalUsage: &totalUsage,
-                continuation: continuation,
+                messages: messages, totalUsage: &totalUsage, continuation: continuation,
                 requestContext: options.requestContext
             )
 
             if let usage = iteration.usage {
-                lastTotalTokens = usage.total
                 continuation.yield(.make(.iterationCompleted(usage: usage, iteration: iterationNumber)))
             }
-
             messages.append(.assistant(iteration.toAssistantMessage()))
 
-            let filteredTools = policy.executableToolCalls(from: iteration.toolCalls)
+            let filteredTools = StreamPolicy.agent.executableToolCalls(from: iteration.toolCalls)
             let pruneCalls = filteredTools.filter { $0.name == "prune_context" }
             let regularCalls = filteredTools.filter { $0.name != "prune_context" }
-            let shouldTerminate = policy.shouldTerminateAfterIteration(toolCalls: iteration.toolCalls)
             let budgetUsage = try requireBudgetUsage(iteration.usage, budgetPhase: budgetPhase)
 
             if let budgetUsage {
-                applyBudgetPhase(
-                    &budgetPhase,
-                    usage: budgetUsage,
-                    messages: &messages,
-                    continuation: continuation
-                )
+                applyBudgetPhase(&budgetPhase, usage: budgetUsage, messages: &messages, continuation: continuation)
             }
 
-            if shouldTerminate {
-                let finishEvent = try parseFinishEvent(
-                    from: iteration.toolCalls, tokenUsage: totalUsage, history: messages
-                )
-                continuation.yield(finishEvent)
-                continuation.finish()
+            if try finishIfTerminated(iteration, usage: totalUsage, history: messages, continuation: continuation) {
                 return
             }
 
             executePruneCalls(pruneCalls, messages: &messages, continuation: continuation)
             try await executeStreamingAndAppendResults(
-                regularCalls, context: context, messages: &messages,
+                regularCalls,
+                context: context,
+                messages: &messages,
                 continuation: continuation,
-                approvalHandler: options.approvalHandler, allowlist: &sessionAllowlist
+                approvalHandler: options.approvalHandler,
+                allowlist: &sessionAllowlist
             )
 
-            if let tokenBudget = options.tokenBudget, totalUsage.total > tokenBudget {
-                continuation.finish(
-                    throwing: AgentError.tokenBudgetExceeded(budget: tokenBudget, used: totalUsage.total)
-                )
+            if finishIfOverBudget(
+                options.tokenBudget, totalUsage: totalUsage, history: messages, continuation: continuation
+            ) {
                 return
             }
+            lastTotalTokens = iteration.usage?.total
         }
 
-        continuation.finish(throwing: AgentError.maxIterationsReached(iterations: configuration.maxIterations))
+        finishStreaming(
+            continuation: continuation,
+            event: makeFinishedEvent(
+                tokenUsage: totalUsage,
+                content: nil,
+                reason: .maxIterationsReached(limit: configuration.maxIterations),
+                history: messages
+            )
+        )
     }
 
     func buildInitialMessages(
@@ -354,20 +355,25 @@ private extension Agent {
         from toolCalls: [ToolCall], tokenUsage: TokenUsage, history: [ChatMessage]
     ) throws -> StreamEvent {
         guard let finishCall = toolCalls.first(where: { $0.name == "finish" }) else {
-            return .make(.finished(tokenUsage: tokenUsage, content: nil, reason: nil, history: history))
+            return makeFinishedEvent(
+                tokenUsage: tokenUsage,
+                content: nil,
+                reason: nil,
+                history: history
+            )
         }
         let decoded: FinishArguments
         do {
             decoded = try JSONDecoder().decode(FinishArguments.self, from: finishCall.argumentsData)
         } catch {
             throw AgentError.finishDecodingFailed(message: String(describing: error))
         }
-        return .make(.finished(
+        return makeFinishedEvent(
             tokenUsage: tokenUsage,
             content: decoded.content,
             reason: FinishReason(decoded.reason ?? "completed"),
             history: history
-        ))
+        )
     }
 
     func parseFinishResult(
@@ -391,4 +397,89 @@ private extension Agent {
             history: history
         )
     }
+
+    private func makeFinishedEvent(
+        tokenUsage: TokenUsage,
+        content: String?,
+        reason: FinishReason?,
+        history: [ChatMessage]
+    ) -> StreamEvent {
+        .make(.finished(
+            tokenUsage: tokenUsage,
+            content: content,
+            reason: reason,
+            history: history
+        ))
+    }
+
+    private func makeTerminalResult(
+        reason: FinishReason,
+        tokenUsage: TokenUsage,
+        iterations: Int,
+        history: [ChatMessage]
+    ) -> AgentResult {
+        AgentResult(
+            finishReason: reason,
+            content: nil,
+            totalTokenUsage: tokenUsage,
+            iterations: iterations,
+            history: history
+        )
+    }
+
+    private func emitCompactionEventIfNeeded(
+        _ compacted: Bool,
+        lastTotalTokens: Int?,
+        continuation: AsyncThrowingStream<StreamEvent, Error>.Continuation
+    ) {
+        guard compacted, let totalTokens = lastTotalTokens, let windowSize = client.contextWindowSize else {
+            return
+        }
+        continuation.yield(.make(.compacted(totalTokens: totalTokens, windowSize: windowSize)))
+    }
+
+    private func finishIfTerminated(
+        _ iteration: StreamIteration,
+        usage: TokenUsage,
+        history: [ChatMessage],
+        continuation: AsyncThrowingStream<StreamEvent, Error>.Continuation
+    ) throws -> Bool {
+        guard StreamPolicy.agent.shouldTerminateAfterIteration(toolCalls: iteration.toolCalls) else {
+            return false
+        }
+        try finishStreaming(
+            continuation: continuation,
+            event: parseFinishEvent(from: iteration.toolCalls, tokenUsage: usage, history: history)
+        )
+        return true
+    }
+
+    private func finishIfOverBudget(
+        _ tokenBudget: Int?,
+        totalUsage: TokenUsage,
+        history: [ChatMessage],
+        continuation: AsyncThrowingStream<StreamEvent, Error>.Continuation
+    ) -> Bool {
+        guard let tokenBudget, totalUsage.total > tokenBudget else {
+            return false
+        }
+        finishStreaming(
+            continuation: continuation,
+            event: makeFinishedEvent(
+                tokenUsage: totalUsage,
+                content: nil,
+                reason: .tokenBudgetExceeded(budget: tokenBudget, used: totalUsage.total),
+                history: history
+            )
+        )
+        return true
+    }
+
+    private func finishStreaming(
+        continuation: AsyncThrowingStream<StreamEvent, Error>.Continuation,
+        event: StreamEvent
+    ) {
+        continuation.yield(event)
+        continuation.finish()
+    }
 }

Sources/AgentRunKit/Core/AgentError.swift

@@ -23,7 +23,6 @@ public enum MalformedStreamReason: Sendable, Equatable, CustomStringConvertible
 
 /// Errors thrown by the agent loop.
 public enum AgentError: Error, Sendable, Equatable, LocalizedError {
-    case maxIterationsReached(iterations: Int)
     case toolNotFound(name: String)
     case toolDecodingFailed(tool: String, message: String)
     case toolEncodingFailed(tool: String, message: String)
@@ -35,14 +34,11 @@ public enum AgentError: Error, Sendable, Equatable, LocalizedError {
     case malformedStream(MalformedStreamReason)
     case schemaInferenceFailed(type: String, message: String)
     case maxDepthExceeded(depth: Int)
-    case tokenBudgetExceeded(budget: Int, used: Int)
     case contextBudgetUsageUnavailable
     case contextBudgetWindowSizeUnavailable
 
     public var errorDescription: String? {
         switch self {
-        case let .maxIterationsReached(iterations):
-            "Agent reached maximum iterations (\(iterations))"
         case let .toolNotFound(name):
             "Tool '\(name)' not found"
         case let .toolDecodingFailed(tool, message):
@@ -65,8 +61,6 @@ public enum AgentError: Error, Sendable, Equatable, LocalizedError {
             "Schema inference failed for '\(type)': \(message)"
         case let .maxDepthExceeded(depth):
             "Sub-agent max depth exceeded (current depth: \(depth))"
-        case let .tokenBudgetExceeded(budget, used):
-            "Token budget exceeded (budget: \(budget), used: \(used))"
         case .contextBudgetUsageUnavailable:
             "Context budget requires provider-reported token usage for every budgeted turn"
         case .contextBudgetWindowSizeUnavailable:
@@ -81,14 +75,12 @@ public enum AgentError: Error, Sendable, Equatable, LocalizedError {
         case let .toolTimeout(tool): "Error: Tool '\(tool)' timed out."
         case let .toolExecutionFailed(tool, message): "Error: Tool '\(tool)' failed: \(message)"
         case let .toolEncodingFailed(tool, message): "Error: Failed to encode '\(tool)' output: \(message)"
-        case let .maxIterationsReached(count): "Error: Agent reached maximum iterations (\(count))."
         case let .finishDecodingFailed(message): "Error: Failed to decode finish arguments: \(message)"
         case let .structuredOutputDecodingFailed(message): "Error: Failed to decode structured output: \(message)"
         case let .llmError(transportError): "Error: LLM request failed: \(transportError)"
         case let .malformedStream(reason): "Error: Malformed stream: \(reason)"
         case let .schemaInferenceFailed(type, message): "Error: Schema inference failed for '\(type)': \(message)"
         case let .maxDepthExceeded(depth): "Error: Sub-agent max depth exceeded (current depth: \(depth))."
-        case let .tokenBudgetExceeded(budget, used): "Error: Token budget exceeded (budget: \(budget), used: \(used))."
         case .contextBudgetUsageUnavailable:
             "Error: Context budget requires provider-reported token usage for every budgeted turn."
         case .contextBudgetWindowSizeUnavailable: