diff --git a/Package.swift b/Package.swift index 65a65b0..1866484 100644 --- a/Package.swift +++ b/Package.swift @@ -47,7 +47,7 @@ let package = Package( dependencies: ["AGUIClient"]), .testTarget( name: "AGUIAgentSDKTests", - dependencies: ["AGUIAgentSDK", "AGUITools"]), + dependencies: ["AGUIAgentSDK", "AGUIClient", "AGUITools"]), .testTarget( name: "AGUIToolsTests", dependencies: ["AGUITools"]), diff --git a/Sources/AGUIAgentSDK/AgUiAgent.swift b/Sources/AGUIAgentSDK/AgUiAgent.swift index 98d869b..03aa092 100644 --- a/Sources/AGUIAgentSDK/AgUiAgent.swift +++ b/Sources/AGUIAgentSDK/AgUiAgent.swift @@ -27,75 +27,20 @@ import AGUICore import AGUITools import Foundation -/// Stateless AG-UI agent providing a simple `sendMessage` API. -/// -/// `AgUiAgent` is a convenience wrapper around `HttpAgent` for cases where no -/// persistent conversation history is needed. Each call to ``sendMessage(_:threadId:state:includeSystemPrompt:)`` -/// builds a fresh `RunAgentInput` from scratch — callers manage history externally. -/// -/// ## Basic Usage -/// -/// ```swift -/// let agent = AgUiAgent(url: URL(string: "https://agent.example.com")!) -/// -/// let stream = agent.sendMessage("Hello!") -/// for try await event in stream { -/// if let chunk = event as? TextMessageChunkEvent { -/// print(chunk.delta ?? "", terminator: "") -/// } -/// } -/// ``` -/// -/// ## Authentication -/// -/// ```swift -/// let agent = AgUiAgent(url: agentURL) { config in -/// config.bearerToken = "sk-…" -/// } -/// ``` -/// -/// ## With Tools -/// -/// ```swift -/// let agent = AgUiAgent(url: agentURL) { config in -/// config.toolRegistry = myToolRegistry -/// config.systemPrompt = "You can use tools to help users." -/// } -/// ``` -/// -/// ## Subclassing -/// -/// Override ``run(input:)`` to customise transport (e.g. WebSocket, mock): -/// -/// ```swift -/// final class MockAgent: AgUiAgent { -/// override func run(input: RunAgentInput) -> AsyncThrowingStream { -/// // return test events -/// } -/// } -/// ``` -/// -/// - SeeAlso: ``AgUiAgentConfig``, ``StatefulAgUiAgent``, ``AgentBuilders`` -open class AgUiAgent: @unchecked Sendable { +public final class AgUiAgent: Sendable { // MARK: - Public properties - /// The resolved configuration for this agent. public let config: AgUiAgentConfig // MARK: - Private state + private let transport: any AgentTransport private let httpAgent: HttpAgent private let toolExecutionManager: ToolExecutionManager? // MARK: - Initialization - /// Creates an agent pointing at `url`, optionally configured via a closure. - /// - /// - Parameters: - /// - url: Base URL of the AG-UI agent server. - /// - configure: Closure for customising ``AgUiAgentConfig`` before the agent is created. - /// Defaults to a no-op (all defaults apply). public init( url: URL, configure: (inout AgUiAgentConfig) -> Void = { _ in } @@ -110,6 +55,7 @@ open class AgUiAgent: @unchecked Sendable { let agent = HttpAgent(configuration: httpConfig) self.httpAgent = agent + self.transport = HttpAgentTransport(configuration: httpConfig) if let registry = cfg.toolRegistry { self.toolExecutionManager = ToolExecutionManager( @@ -121,44 +67,35 @@ open class AgUiAgent: @unchecked Sendable { } } - // MARK: - Core run method (overrideable) - - /// Returns a raw event stream for the given input. - /// - /// The default implementation delegates to the internal `HttpAgent`. - /// Override this in subclasses to replace the transport (e.g. for testing). - /// - /// - Parameter input: The run agent input. - /// - Returns: An `AsyncThrowingStream` of raw AG-UI events. - open func run(input: RunAgentInput) -> AsyncThrowingStream { - httpAgent.run(input: input) + public init( + transport: any AgentTransport, + config: AgUiAgentConfig = AgUiAgentConfig(), + toolExecutionManager: ToolExecutionManager? = nil + ) { + self.config = config + self.transport = transport + let url = URL(string: "https://placeholder.local")! + self.httpAgent = HttpAgent(baseURL: url) + self.toolExecutionManager = toolExecutionManager + } + + // MARK: - Core run method + + public func run(input: RunAgentInput) -> AsyncThrowingStream { + transport.run(input: input) } // MARK: - sendMessage (primary API) - /// Sends a single message and returns the resulting event stream. - /// - /// Each call builds a fresh `RunAgentInput` — no history is carried between calls. - /// An optional system prompt is prepended as the first message when `includeSystemPrompt` - /// is `true` and ``AgUiAgentConfig/systemPrompt`` is set. - /// - /// - Parameters: - /// - message: The user message text. - /// - threadId: Conversation thread ID (default: a new UUID per call). - /// - state: Optional JSON state to include (default: `{}`). - /// - includeSystemPrompt: When `true` and a system prompt is configured, it is added - /// as the first message (default: `true`). - /// - Returns: An `AsyncThrowingStream` of AG-UI events. - open func sendMessage( + public func sendMessage( _ message: String, threadId: String = UUID().uuidString, state: State? = nil, includeSystemPrompt: Bool = true ) -> AsyncThrowingStream { AsyncThrowingStream { continuation in - Task { + let task = Task { do { - // Build message list (no history — stateless per call) var messages: [any Message] = [] if includeSystemPrompt, let prompt = self.config.systemPrompt { @@ -173,7 +110,6 @@ open class AgUiAgent: @unchecked Sendable { content: message )) - // Resolve tool definitions from the registry var tools: [Tool] = [] if let registry = self.config.toolRegistry { tools = await registry.allTools() @@ -189,7 +125,6 @@ open class AgUiAgent: @unchecked Sendable { forwardedProps: self.config.forwardedProps ) - // Route through ToolExecutionManager when a registry is configured let rawStream = self.run(input: input) if let manager = self.toolExecutionManager { @@ -212,27 +147,19 @@ open class AgUiAgent: @unchecked Sendable { continuation.finish(throwing: error) } } + continuation.onTermination = { _ in task.cancel() } } } // MARK: - Subscriber - /// Subscribes to lifecycle events from the underlying `HttpAgent`. - /// - /// - Parameter subscriber: The subscriber to register. - /// - Returns: A subscription handle that can be used to unsubscribe. public func subscribe(_ subscriber: any AgentSubscriber) async -> any AgentSubscription { await httpAgent.subscribe(subscriber) } // MARK: - Lifecycle - /// Cancels any in-flight tool executions and disposes the underlying agent. - /// - /// After calling `close()`, further calls to ``sendMessage(_:threadId:state:includeSystemPrompt:)`` - /// will start a new run (the underlying agent's `dispose()` prevents its own `runAgent` - /// pipeline from re-entering, but direct `run(input:)` calls continue to work for subclasses). - open func close() { + public func close() { Task { if let manager = self.toolExecutionManager { await manager.cancelAllExecutions() diff --git a/Sources/AGUIAgentSDK/StatefulAgUiAgent.swift b/Sources/AGUIAgentSDK/StatefulAgUiAgent.swift index 7690e84..92fef1a 100644 --- a/Sources/AGUIAgentSDK/StatefulAgUiAgent.swift +++ b/Sources/AGUIAgentSDK/StatefulAgUiAgent.swift @@ -308,12 +308,13 @@ public final class StatefulAgUiAgent: Sendable { ) } else { eventStream = AsyncThrowingStream { continuation in - Task { + let task = Task { do { for try await event in rawStream { continuation.yield(event) } continuation.finish() } catch { continuation.finish(throwing: error) } } + continuation.onTermination = { _ in task.cancel() } } } @@ -361,7 +362,7 @@ public final class StatefulAgUiAgent: Sendable { threadId: String ) -> AsyncThrowingStream where S.Element == any AGUIEvent { AsyncThrowingStream { continuation in - Task { + let task = Task { var currentAssistantMessage: AssistantMessage? let patchApplicator = PatchApplicator() let messageDecoder = MessageDecoder() @@ -501,6 +502,7 @@ public final class StatefulAgUiAgent: Sendable { continuation.finish(throwing: error) } } + continuation.onTermination = { _ in task.cancel() } } } } diff --git a/Sources/AGUIAgentSDK/StatefulAgUiAgentConfig.swift b/Sources/AGUIAgentSDK/StatefulAgUiAgentConfig.swift index bfef02d..72be929 100644 --- a/Sources/AGUIAgentSDK/StatefulAgUiAgentConfig.swift +++ b/Sources/AGUIAgentSDK/StatefulAgUiAgentConfig.swift @@ -105,17 +105,12 @@ public struct StatefulAgUiAgentConfig: Sendable { /// Bearer token for authentication. /// - /// When set, automatically adds an `Authorization: Bearer ` header - /// to every request. + /// When set, ``buildHeaders()`` includes `Authorization: Bearer `. + /// This property does **not** mutate ``headers`` directly — use + /// ``buildHeaders()`` to get the merged dictionary for request construction. /// /// Default: `nil` - public var bearerToken: String? { - didSet { - if let token = bearerToken { - headers["Authorization"] = "Bearer \(token)" - } - } - } + public var bearerToken: String? /// API key value. /// @@ -143,6 +138,31 @@ public struct StatefulAgUiAgentConfig: Sendable { /// Default: `"/run"` public var endpoint: String + // MARK: - Header builder + + /// Returns the final HTTP header dictionary, merging ``bearerToken`` and + /// ``apiKey`` into ``headers``. + /// + /// Priority (highest → lowest): + /// 1. Entries already in ``headers`` + /// 2. `bearerToken` → `Authorization: Bearer ` + /// 3. `apiKey` → `: ` + /// + /// - Returns: Merged header dictionary ready for `HttpAgentConfiguration`. + public func buildHeaders() -> [String: String] { + var result: [String: String] = [:] + if let key = apiKey { + result[apiKeyHeader] = key + } + if let token = bearerToken { + result["Authorization"] = "Bearer \(token)" + } + for (k, v) in headers { + result[k] = v + } + return result + } + /// Creates a new stateful agent configuration. /// /// - Parameter baseURL: The base URL of the AG-UI agent server diff --git a/Sources/AGUIAgentSDK/Tools/ClientToolResponseHandler.swift b/Sources/AGUIAgentSDK/Tools/ClientToolResponseHandler.swift index 5d85253..032c925 100644 --- a/Sources/AGUIAgentSDK/Tools/ClientToolResponseHandler.swift +++ b/Sources/AGUIAgentSDK/Tools/ClientToolResponseHandler.swift @@ -51,7 +51,7 @@ import Foundation /// responseHandler: handler /// ) /// ``` -public final class ClientToolResponseHandler: ToolResponseHandler, @unchecked Sendable { +public final class ClientToolResponseHandler: ToolResponseHandler, Sendable { private let httpAgent: HttpAgent private let endpoint: String? diff --git a/Sources/AGUIClient/AbstractAgent.swift b/Sources/AGUIClient/AbstractAgent.swift index 8bb5a77..42e7f44 100644 --- a/Sources/AGUIClient/AbstractAgent.swift +++ b/Sources/AGUIClient/AbstractAgent.swift @@ -27,10 +27,6 @@ import Foundation // MARK: - AgentStorage -/// Thread-safe mutable storage for an agent's runtime state. -/// -/// All reads and writes are serialized through the actor executor, -/// ensuring data-race-free access across concurrent tasks. internal actor AgentStorage { var messages: [any Message] = [] var currentState: State = Data("{}".utf8) @@ -55,48 +51,25 @@ internal extension AgentStorage { // MARK: - AbstractAgent -/// Base class for AG-UI agents. -/// -/// `AbstractAgent` provides the full lifecycle pipeline — including chunk -/// transformation, protocol verification, state application, and subscriber -/// notification — while delegating the raw event stream to subclasses via -/// `run(input:)`. -/// -/// ## Subclassing -/// -/// Override `run(input:)` to return a stream of raw AG-UI events: -/// -/// ```swift -/// final class MyAgent: AbstractAgent { -/// override func run(input: RunAgentInput) -> AsyncThrowingStream { -/// // Return your event stream here -/// } -/// } -/// ``` -/// -/// ## Thread Safety -/// -/// The `@unchecked Sendable` annotation is required because `open class` -/// with mutable stored properties cannot automatically synthesize `Sendable` -/// conformance. Thread safety is provided by the internal `AgentStorage` actor. -open class AbstractAgent: @unchecked Sendable { +public final class AbstractAgent: Sendable { // MARK: - Internal storage internal let storage: AgentStorage internal let subscriberManager: SubscriberManager + // MARK: - Transport + + private let transport: any AgentTransport + // MARK: - Configuration (immutable after init) - /// When `true`, logs verbose pipeline output to stdout. public let debug: Bool // MARK: - Initialization - /// Creates a new abstract agent. - /// - /// - Parameter debug: When `true`, enables verbose pipeline logging. - public init(debug: Bool = false) { + public init(transport: any AgentTransport, debug: Bool = false) { + self.transport = transport self.storage = AgentStorage() self.subscriberManager = SubscriberManager() self.debug = debug @@ -104,46 +77,24 @@ open class AbstractAgent: @unchecked Sendable { // MARK: - Async state accessors (cross actor boundary) - /// The current conversation message list. public var messages: [any Message] { get async { await storage.messages } } - /// The current JSON state. public var state: State { get async { await storage.currentState } } - /// The accumulated raw events received during runs. public var rawEvents: [RawEvent] { get async { await storage.rawEvents } } - /// The accumulated custom events received during runs. public var customEvents: [CustomEvent] { get async { await storage.customEvents } } - /// The current thinking/reasoning telemetry state, if any. public var thinking: ThinkingTelemetryState? { get async { await storage.thinking } } - // MARK: - Abstract run method (subclasses must override) - - /// Returns a stream of raw decoded AG-UI events for the given input. - /// - /// Subclasses **must** override this method. The default implementation - /// triggers a `fatalError`. - /// - /// - Parameter input: The run agent input. - /// - Returns: An `AsyncThrowingStream` of AG-UI events. - open func run(input: RunAgentInput) -> AsyncThrowingStream { - fatalError("AbstractAgent subclasses must implement run(input:)") + // MARK: - Run method + + public func run(input: RunAgentInput) -> AsyncThrowingStream { + transport.run(input: input) } // MARK: - Public pipeline methods - /// Executes the full pipeline and blocks until the run completes. - /// - /// The pipeline applies chunk transformation, protocol verification, and - /// state application. Registered subscribers (plus any one-off `subscriber` - /// passed here) are notified at each lifecycle hook. - /// - /// - Parameters: - /// - parameters: Optional run parameters. When `nil`, defaults are used. - /// - subscriber: Optional one-off subscriber for this run only. - /// - Throws: Rethrows any error produced by the event stream. public func runAgent( parameters: RunAgentParameters? = nil, subscriber: (any AgentSubscriber)? = nil @@ -152,12 +103,10 @@ open class AbstractAgent: @unchecked Sendable { let input = buildInput(from: parameters) - // Collect all subscribers let registeredSubscribers = await subscriberManager.allSubscribers() var allSubscribers = registeredSubscribers if let s = subscriber { allSubscribers.append(s) } - // onRunInitialized let initMutation = await runSubscribersWithMutation( subscribers: allSubscribers, messages: await storage.messages, @@ -184,7 +133,6 @@ open class AbstractAgent: @unchecked Sendable { await self.applyAgentState(agentState, input: input, subscribers: allSubscribers) } - // onRunFinalized let finalMessages = await self.storage.messages let finalState = await self.storage.currentState _ = await runSubscribersWithMutation( @@ -214,14 +162,6 @@ open class AbstractAgent: @unchecked Sendable { try await task.value } - /// Returns the processed event stream without driving the pipeline internally. - /// - /// Events pass through chunk transformation and protocol verification but - /// state is **not** applied internally. Callers are responsible for - /// consuming the stream and managing state. - /// - /// - Parameter input: The run agent input. - /// - Returns: A verified, chunk-transformed event stream. public func runAgentObservable( input: RunAgentInput ) -> AsyncThrowingStream { @@ -230,22 +170,14 @@ open class AbstractAgent: @unchecked Sendable { .verifyEvents(debug: debug) } - /// Cancels the current in-flight run task, if any. public func abortRun() { Task { await storage.currentTask?.cancel() } } - /// Prevents further runs from starting. - /// - /// Calling `runAgent` after `dispose()` is a no-op. public func dispose() { Task { await storage.setDisposed(true) } } - /// Subscribes to agent lifecycle events. - /// - /// - Parameter subscriber: The subscriber to register. - /// - Returns: A subscription handle that can be used to unsubscribe. public func subscribe(_ subscriber: any AgentSubscriber) async -> any AgentSubscription { let id = await subscriberManager.subscribe(subscriber) return DefaultAgentSubscription { diff --git a/Sources/AGUIClient/AgentConfig.swift b/Sources/AGUIClient/AgentConfig.swift index c87a875..470eb57 100644 --- a/Sources/AGUIClient/AgentConfig.swift +++ b/Sources/AGUIClient/AgentConfig.swift @@ -73,19 +73,36 @@ public struct HttpAgentConfig: Sendable { public var requestTimeout: TimeInterval /// Connection timeout in seconds. Default: 30. public var connectTimeout: TimeInterval - /// Bearer token — automatically added as `Authorization: Bearer `. - public var bearerToken: String? { - didSet { - if let token = bearerToken { - headers["Authorization"] = "Bearer \(token)" - } - } - } + /// Bearer token for authentication. + /// + /// When set, ``buildHeaders()`` includes `Authorization: Bearer `. + /// This property does **not** mutate ``headers`` — use ``buildHeaders()``. + public var bearerToken: String? /// API key value. public var apiKey: String? /// Header name for the API key. Default: "X-API-Key". public var apiKeyHeader: String + // MARK: - Header builder + + /// Returns the merged HTTP header dictionary. + /// + /// Applies `bearerToken` and `apiKey` on top of ``headers``. + /// Explicit ``headers`` entries override auto-generated auth headers. + public func buildHeaders() -> [String: String] { + var result: [String: String] = [:] + if let key = apiKey { + result[apiKeyHeader] = key + } + if let token = bearerToken { + result["Authorization"] = "Bearer \(token)" + } + for (k, v) in headers { + result[k] = v + } + return result + } + public init(url: String, base: AgentConfig = AgentConfig()) { self.base = base self.url = url diff --git a/Sources/AGUIClient/HttpAgent.swift b/Sources/AGUIClient/HttpAgent.swift index 09f01b6..8d69ccc 100644 --- a/Sources/AGUIClient/HttpAgent.swift +++ b/Sources/AGUIClient/HttpAgent.swift @@ -25,245 +25,51 @@ import AGUICore import Foundation -/// High-level HTTP client for AG-UI agent communication. -/// -/// `HttpAgent` provides a convenient, fluent API for executing agent runs -/// and streaming AG-UI events. It wraps the lower-level transport and streaming -/// infrastructure with an easy-to-use interface. -/// -/// ## Basic Usage -/// -/// ```swift -/// let agent = HttpAgent(baseURL: URL(string: "https://agent.example.com")!) -/// -/// let stream = try await agent.run(threadId: "thread-1", runId: "run-1") { builder in -/// builder.message(UserMessage( -/// id: "msg1", -/// content: [TextInputContent(text: "Hello!")] -/// )) -/// } -/// -/// for try await event in stream { -/// switch event.eventType { -/// case .textMessageChunk: -/// let chunk = event as! TextMessageChunkEvent -/// print(chunk.delta ?? "", terminator: "") -/// case .runFinished: -/// print("\nDone!") -/// default: -/// break -/// } -/// } -/// ``` -/// -/// ## Advanced Usage -/// -/// ```swift -/// var config = HttpAgentConfiguration(baseURL: agentURL) -/// config.timeout = 120.0 -/// config.headers = ["Authorization": "Bearer token"] -/// -/// let agent = HttpAgent(configuration: config) -/// -/// let input = try RunAgentInput.builder() -/// .threadId("thread-1") -/// .runId("run-1") -/// .message(DeveloperMessage( -/// id: "dev1", -/// content: [TextInputContent(text: "System prompt")] -/// )) -/// .message(UserMessage( -/// id: "user1", -/// content: [TextInputContent(text: "User query")] -/// )) -/// .tool(weatherTool) -/// .context(Context(description: "timezone", value: "UTC")) -/// .build() -/// -/// let stream = try await agent.run(input, endpoint: "/custom/run") -/// ``` -/// -/// ## Error Handling -/// -/// ```swift -/// do { -/// let stream = try await agent.run(threadId: "t1", runId: "r1") -/// for try await event in stream { -/// // Process events -/// } -/// } catch ClientError.httpError(let statusCode) { -/// print("HTTP error: \(statusCode)") -/// } catch ClientError.timeout { -/// print("Request timed out") -/// } catch { -/// print("Unexpected error: \(error)") -/// } -/// ``` -/// -/// ## Thread Safety -/// -/// `HttpAgent` is safe to use across multiple concurrent tasks. Each run -/// creates an isolated stream with its own state. -public final class HttpAgent: AbstractAgent, @unchecked Sendable { - /// The underlying HTTP transport. - private let transport: HttpTransport - - /// The AG-UI event decoder. +public final class HttpAgent: Sendable { + private let abstractAgent: AbstractAgent + private let httpTransport: HttpTransport private let decoder: AGUIEventDecoder - - /// Default endpoint for agent runs. private let defaultEndpoint: String - /// Agent configuration — stored so retry helpers can read `retryPolicy`. - private let configuration: HttpAgentConfiguration - - /// Creates a new HTTP agent with a base URL. - /// - /// This convenience initializer creates an agent with default configuration. - /// - /// - Parameter baseURL: The base URL of the AG-UI agent - /// - /// ## Example - /// - /// ```swift - /// let agent = HttpAgent(baseURL: URL(string: "https://agent.example.com")!) - /// ``` public init(baseURL: URL) { let config = HttpAgentConfiguration(baseURL: baseURL) - self.configuration = config - self.transport = HttpTransport(configuration: config) + let agentTransport = HttpAgentTransport(configuration: config) + self.abstractAgent = AbstractAgent(transport: agentTransport, debug: config.debug) + self.httpTransport = HttpTransport(configuration: config) self.decoder = AGUIEventDecoder() self.defaultEndpoint = "/run" - super.init() } - /// Creates a new HTTP agent with custom configuration. - /// - /// - Parameter configuration: The HTTP agent configuration - /// - /// ## Example - /// - /// ```swift - /// var config = HttpAgentConfiguration(baseURL: agentURL) - /// config.timeout = 120.0 - /// config.headers = ["Authorization": "Bearer token"] - /// - /// let agent = HttpAgent(configuration: config) - /// ``` public init(configuration: HttpAgentConfiguration) { - self.configuration = configuration - self.transport = HttpTransport(configuration: configuration) + let agentTransport = HttpAgentTransport(configuration: configuration) + self.abstractAgent = AbstractAgent(transport: agentTransport, debug: configuration.debug) + self.httpTransport = HttpTransport(configuration: configuration) self.decoder = AGUIEventDecoder() self.defaultEndpoint = "/run" - super.init(debug: configuration.debug) } - /// Creates a new HTTP agent with custom HTTP client. - /// - /// This initializer allows dependency injection of a custom HTTP client, - /// useful for testing or custom network implementations. - /// - /// - Parameters: - /// - configuration: The HTTP agent configuration - /// - httpClient: Custom HTTP client implementation - /// - /// ## Example - /// - /// ```swift - /// let mockClient = MockHTTPClient() - /// let agent = HttpAgent( - /// configuration: config, - /// httpClient: mockClient - /// ) - /// ``` public init( configuration: HttpAgentConfiguration, httpClient: any HTTPClient ) { - self.configuration = configuration - self.transport = HttpTransport( - configuration: configuration, - httpClient: httpClient - ) + let agentTransport = HttpAgentTransport(configuration: configuration, httpClient: httpClient) + self.abstractAgent = AbstractAgent(transport: agentTransport, debug: configuration.debug) + self.httpTransport = HttpTransport(configuration: configuration, httpClient: httpClient) self.decoder = AGUIEventDecoder() self.defaultEndpoint = "/run" - super.init() } - /// Executes an agent run with the provided input. - /// - /// This is the most explicit run method, accepting a fully-configured - /// `RunAgentInput` object. - /// - /// - Parameters: - /// - input: The run agent input - /// - endpoint: Custom endpoint (default: "/run") - /// - Returns: AsyncSequence of AG-UI events - /// - Throws: `ClientError` if the request fails - /// - /// ## Example - /// - /// ```swift - /// let input = try RunAgentInput.builder() - /// .threadId("thread-1") - /// .runId("run-1") - /// .message(UserMessage(id: "msg1", content: [TextInputContent(text: "Hi")])) - /// .build() - /// - /// let stream = try await agent.run(input) - /// for try await event in stream { - /// print(event) - /// } - /// ``` public func run( _ input: RunAgentInput, endpoint: String? = nil ) async throws -> EventStream> { - let bytes = try await transport.execute( + let bytes = try await httpTransport.execute( endpoint: endpoint ?? defaultEndpoint, input: input ) - return EventStream(bytes: bytes, decoder: decoder) } - /// Executes an agent run with builder configuration. - /// - /// This method provides a fluent interface for configuring the run input - /// using a builder pattern. - /// - /// - Parameters: - /// - threadId: The thread identifier - /// - runId: The run identifier - /// - endpoint: Custom endpoint (default: "/run") - /// - configure: Closure to configure the input builder - /// - Returns: AsyncSequence of AG-UI events - /// - Throws: `ClientError` if the request fails - /// - /// ## Example - /// - /// ```swift - /// let stream = try await agent.run( - /// threadId: "thread-1", - /// runId: "run-1" - /// ) { builder in - /// builder - /// .message(DeveloperMessage( - /// id: "dev1", - /// content: [TextInputContent(text: "You are helpful")] - /// )) - /// .message(UserMessage( - /// id: "user1", - /// content: [TextInputContent(text: "Hello!")] - /// )) - /// .tool(weatherTool) - /// .context(Context(description: "timezone", value: "UTC")) - /// } - /// - /// for try await event in stream { - /// // Process events - /// } - /// ``` public func run( threadId: String, runId: String, @@ -275,125 +81,41 @@ public final class HttpAgent: AbstractAgent, @unchecked Sendable { .threadId(threadId) .runId(runId) ).build() - return try await run(input, endpoint: endpoint) } - // MARK: - AbstractAgent override - - /// Returns a stream of raw AG-UI events with automatic retry on transient failures. - /// - /// This override bridges `AbstractAgent.run(input:)` to the underlying - /// `HttpTransport` and applies the retry policy configured in - /// `HttpAgentConfiguration.retryPolicy`. On each retry, the `Last-Event-ID` - /// header is sent with the most recent SSE event id so compliant servers can - /// resume the stream from where it dropped. - /// - /// ### Retry behaviour - /// - /// | Policy | Behaviour | - /// |--------|-----------| - /// | `.none` (default) | No retry — errors propagate immediately | - /// | `.fixed(maxAttempts:delay:)` | Up to `maxAttempts` retries, each preceded by a fixed `delay` | - /// | `.exponentialBackoff(maxAttempts:baseDelay:)` | Up to `maxAttempts` retries, delay doubles each attempt (capped at 60 s) | - /// - /// ### Retryable errors - /// - /// Only transient, network-level errors trigger a retry: - /// - `ClientError.timeout` - /// - `ClientError.networkError` - /// - /// Server errors (`httpError`, `invalidResponse`) are not retried. - /// - /// - Parameter input: The run agent input. - /// - Returns: An `AsyncThrowingStream` of AG-UI events from the server. - public override func run(input: RunAgentInput) -> AsyncThrowingStream { - AsyncThrowingStream { continuation in - Task { - var lastEventId: String? = nil - var attempt = 0 - // Hold a reference to the current EventStream so we can read - // lastEventId after a mid-stream failure (the stream is a struct but - // its lastEventIdBox is a reference type — still readable after throw). - var currentStream: EventStream>? = nil - - while true { - do { - let bytes = try await self.transport.execute( - endpoint: self.defaultEndpoint, - input: input, - lastEventId: lastEventId - ) - let eventStream = EventStream(bytes: bytes, decoder: self.decoder) - currentStream = eventStream - for try await event in eventStream { - continuation.yield(event) - } - continuation.finish() - return - } catch { - // Capture the last-event-id from the stream that just failed - // (nil when the failure happened before streaming began). - if let stream = currentStream { - lastEventId = stream.lastEventId ?? lastEventId - } - currentStream = nil - - guard self.shouldRetry(error: error, attempt: attempt) else { - continuation.finish(throwing: error) - return - } + public func run(input: RunAgentInput) -> AsyncThrowingStream { + abstractAgent.run(input: input) + } - let delay = self.retryDelay(for: attempt) - attempt += 1 + public func runAgent( + parameters: RunAgentParameters? = nil, + subscriber: (any AgentSubscriber)? = nil + ) async throws { + try await abstractAgent.runAgent(parameters: parameters, subscriber: subscriber) + } - if delay > 0 { - try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) - } - } - } - } - } + public func runAgentObservable( + input: RunAgentInput + ) -> AsyncThrowingStream { + abstractAgent.runAgentObservable(input: input) } - // MARK: - Retry helpers + public var messages: [any Message] { get async { await abstractAgent.messages } } + public var state: State { get async { await abstractAgent.state } } + public var rawEvents: [RawEvent] { get async { await abstractAgent.rawEvents } } + public var customEvents: [CustomEvent] { get async { await abstractAgent.customEvents } } + public var thinking: ThinkingTelemetryState? { get async { await abstractAgent.thinking } } - /// Returns `true` when `error` is a transient network failure and the retry - /// policy permits another attempt. - private func shouldRetry(error: Error, attempt: Int) -> Bool { - guard isRetryable(error) else { return false } - switch configuration.retryPolicy { - case .none: - return false - case .fixed(let maxAttempts, _): - return attempt < maxAttempts - case .exponentialBackoff(let maxAttempts, _): - return attempt < maxAttempts - } + public func abortRun() { + abstractAgent.abortRun() } - /// Returns `true` for errors that represent a transient network condition - /// (not a deliberate server-side rejection). - private func isRetryable(_ error: Error) -> Bool { - guard let clientError = error as? ClientError else { return false } - switch clientError { - case .timeout, .networkError: - return true - default: - return false - } + public func dispose() { + abstractAgent.dispose() } - /// Returns the sleep interval (in seconds) before attempt `attempt`. - private func retryDelay(for attempt: Int) -> TimeInterval { - switch configuration.retryPolicy { - case .none: - return 0 - case .fixed(_, let delay): - return delay - case .exponentialBackoff(_, let baseDelay): - // 2^attempt * baseDelay, capped at 60 s to avoid unbounded waits - return min(baseDelay * pow(2.0, Double(attempt)), 60.0) - } + public func subscribe(_ subscriber: any AgentSubscriber) async -> any AgentSubscription { + await abstractAgent.subscribe(subscriber) } } diff --git a/Sources/AGUIClient/Streaming/AsyncSequence+Buffering.swift b/Sources/AGUIClient/Streaming/AsyncSequence+Buffering.swift index 033a354..c2bb1ca 100644 --- a/Sources/AGUIClient/Streaming/AsyncSequence+Buffering.swift +++ b/Sources/AGUIClient/Streaming/AsyncSequence+Buffering.swift @@ -76,13 +76,6 @@ extension AsyncSequence where Self: Sendable, Element: Sendable { case .dropNewest: // Drop the new element, keep buffer as-is continue - case .suspend: - // For suspend strategy, yield current buffer first - // to apply natural backpressure - while !buffer.isEmpty { - continuation.yield(buffer.removeFirst()) - } - buffer.append(element) } } else { buffer.append(element) @@ -131,7 +124,9 @@ extension AsyncSequence where Self: Sendable { public func batched(count: Int) -> BatchedAsyncSequence { BatchedAsyncSequence(base: self, batchSize: count) } +} +extension AsyncSequence where Self: Sendable, Self.Element: Sendable { /// Batches elements by time window. /// /// Groups elements arriving within a time window. Emits batches @@ -198,8 +193,83 @@ public struct BatchedAsyncSequence: AsyncSequence where Bas } } +// MARK: - TimeBatchProducer + +/// Actor that buffers elements from a concurrent producer task so that a +/// time-windowed consumer can safely collect them without races. +/// +/// The producer task feeds elements into this actor as fast as the upstream +/// allows. The consumer calls `dequeue()` at each window boundary to drain +/// whatever arrived during that window. Because actor isolation serialises +/// all reads and writes to the buffer, no element is ever lost at window +/// edges — an element added after one `dequeue()` call simply appears in the +/// next window's batch. +fileprivate actor TimeBatchProducer { + private var buffer: [Element] = [] + private(set) var isExhausted = false + private var storedError: Error? + + func enqueue(_ element: Element) { + buffer.append(element) + } + + func markExhausted() { + isExhausted = true + } + + func markFailed(_ error: Error) { + storedError = error + isExhausted = true + } + + /// Drains the buffer and returns it with the exhaustion flag. + /// Throws if the upstream ended with an error. + func dequeue() throws -> ([Element], Bool) { + if let e = storedError { throw e } + let batch = buffer + buffer = [] + return (batch, isExhausted) + } +} + +/// Cancels its wrapped `Task` on deinit, ensuring the background producer is +/// torn down if the `TimeBatchedAsyncSequence.AsyncIterator` is dropped before +/// the upstream sequence is exhausted (e.g. `break` out of a `for await` loop). +fileprivate final class ProducerTaskHandle: Sendable { + let task: Task + init(_ task: Task) { self.task = task } + deinit { task.cancel() } +} + /// Async sequence that batches elements by time window. -public struct TimeBatchedAsyncSequence: AsyncSequence where Base: Sendable { +/// +/// A dedicated producer task pulls from the upstream sequence as fast as it +/// can, feeding every element into a `TimeBatchProducer` actor. The consumer +/// (`next()`) sleeps for `timeWindow` seconds, then drains whatever arrived +/// during that window and returns it as an array. +/// +/// Because actor isolation serialises the producer's `enqueue` and the +/// consumer's `dequeue`, no element can be lost at a window boundary: an +/// element that arrives after one drain is simply included in the next. +/// +/// Empty windows (upstream alive but idle) are silently skipped — the +/// consumer starts a new window automatically and only yields once it has +/// collected at least one element. +/// +/// ## Example +/// +/// ```swift +/// let batched = textChunks.batched(timeWindow: 0.05) // 50 ms windows +/// for try await batch in batched { +/// await render(batch.joined()) +/// } +/// ``` +/// +/// - Note: Requires `Base.Element: Sendable` because elements must cross +/// the actor boundary between the producer task and the consumer. +public struct TimeBatchedAsyncSequence: AsyncSequence + where Base: Sendable, Base.Element: Sendable +{ public typealias Element = [Base.Element] private let base: Base @@ -211,49 +281,62 @@ public struct TimeBatchedAsyncSequence: AsyncSequence where } public func makeAsyncIterator() -> AsyncIterator { - AsyncIterator(base: base.makeAsyncIterator(), timeWindow: timeWindow) + let producer = TimeBatchProducer() + let rawTask = Task { [base, producer] in + do { + for try await element in base { + await producer.enqueue(element) + } + await producer.markExhausted() + } catch { + await producer.markFailed(error) + } + } + return AsyncIterator( + producer: producer, + handle: ProducerTaskHandle(rawTask), + timeWindow: timeWindow + ) } public struct AsyncIterator: AsyncIteratorProtocol { - private var baseIterator: Base.AsyncIterator + private let producer: TimeBatchProducer + /// Class reference so `deinit` cancels the producer when the iterator + /// is dropped (e.g. after a `break` in the consumer's `for await`). + private let handle: ProducerTaskHandle private let timeWindow: TimeInterval - private var buffer: [Base.Element] = [] - private var isExhausted = false - - init(base: Base.AsyncIterator, timeWindow: TimeInterval) { - self.baseIterator = base + private var isLocallyExhausted = false + + fileprivate init( + producer: TimeBatchProducer, + handle: ProducerTaskHandle, + timeWindow: TimeInterval + ) { + self.producer = producer + self.handle = handle self.timeWindow = timeWindow } public mutating func next() async throws -> [Base.Element]? { - guard !isExhausted else { return nil } + guard !isLocallyExhausted else { return nil } - var batch: [Base.Element] = [] - let windowStart = Date() - - // Collect elements for the time window while true { - // Check if window expired - let elapsed = Date().timeIntervalSince(windowStart) - if elapsed >= timeWindow { - break - } + // Sleep for the full window duration before collecting. + try await Task.sleep(for: .seconds(timeWindow)) + + let (batch, exhausted) = try await producer.dequeue() - // Try to get next element (non-blocking check) - if let element = try await baseIterator.next() { - batch.append(element) + if exhausted { + isLocallyExhausted = true + handle.task.cancel() + return batch.isEmpty ? nil : batch + } - // Small yield to allow other tasks to run - await Task.yield() - } else { - // End of sequence - isExhausted = true - break + if !batch.isEmpty { + return batch } + // Empty window — upstream is alive but idle. Start next window. } - - // Return batch if we have elements - return batch.isEmpty ? nil : batch } } } diff --git a/Sources/AGUIClient/Streaming/BufferingStrategy.swift b/Sources/AGUIClient/Streaming/BufferingStrategy.swift index 7d7cb18..34051b8 100644 --- a/Sources/AGUIClient/Streaming/BufferingStrategy.swift +++ b/Sources/AGUIClient/Streaming/BufferingStrategy.swift @@ -67,22 +67,4 @@ public enum BufferingStrategy: Sendable { /// Result: [1, 2, 3, 4, 5] // Dropped 6 /// ``` case dropNewest - - /// Block the producer when buffer is full (natural backpressure). - /// - /// This strategy applies backpressure by suspending the producer - /// until the consumer frees space in the buffer. - /// - /// **Use case**: When all events must be processed and loss is unacceptable. - /// - /// ## Example - /// - /// ``` - /// Buffer: [1, 2, 3, 4, 5] (limit: 5) - /// New: 6 (producer suspends until consumer reads) - /// ``` - /// - /// **Note**: This is the natural behavior of AsyncSequence iteration - /// and doesn't require explicit buffering. - case suspend } diff --git a/Sources/AGUIClient/Streaming/ChunkTransformer.swift b/Sources/AGUIClient/Streaming/ChunkTransformer.swift index f9148ec..db6e6a2 100644 --- a/Sources/AGUIClient/Streaming/ChunkTransformer.swift +++ b/Sources/AGUIClient/Streaming/ChunkTransformer.swift @@ -75,10 +75,11 @@ public struct ChunkTransformer { _ events: S ) -> AsyncThrowingStream where S.Element == any AGUIEvent { AsyncThrowingStream { continuation in - Task { + let task = Task { let transformer = EventTransformer(continuation: continuation) await transformer.processEvents(events) } + continuation.onTermination = { _ in task.cancel() } } } } diff --git a/Sources/AGUIClient/Streaming/EventStream.swift b/Sources/AGUIClient/Streaming/EventStream.swift index 23afcfd..f1ba3cb 100644 --- a/Sources/AGUIClient/Streaming/EventStream.swift +++ b/Sources/AGUIClient/Streaming/EventStream.swift @@ -81,8 +81,15 @@ public struct EventStream: AsyncSequence where Bytes.Eleme /// Using a class (reference semantics) means every iterator created from the same /// stream updates the same location, so callers can read `lastEventId` on the /// stream value after an iterator throws. - final class LastEventIdBox: @unchecked Sendable { + /// + /// Declared as an `actor` so the compiler synthesises `Sendable` automatically + /// and serialises all reads and writes without manual locking. + actor LastEventIdBox { var value: String? + + func set(_ newValue: String?) { + value = newValue + } } // MARK: - Stored properties @@ -102,7 +109,7 @@ public struct EventStream: AsyncSequence where Bytes.Eleme /// /// Read this after a mid-stream failure to obtain the resume cursor for /// `Last-Event-ID` on reconnect. - public var lastEventId: String? { lastEventIdBox.value } + public var lastEventId: String? { get async { await lastEventIdBox.value } } // MARK: - Initialization @@ -200,7 +207,7 @@ public struct EventStream: AsyncSequence where Bytes.Eleme // The id is captured before decoding so it is available even // when the accompanying data fails to decode. if let id = sseEvent.id { - lastEventIdBox.value = id + await lastEventIdBox.set(id) } guard let data = sseEvent.data.data(using: .utf8) else { diff --git a/Sources/AGUIClient/Streaming/EventVerifier.swift b/Sources/AGUIClient/Streaming/EventVerifier.swift index 73694f7..0eba0c7 100644 --- a/Sources/AGUIClient/Streaming/EventVerifier.swift +++ b/Sources/AGUIClient/Streaming/EventVerifier.swift @@ -249,7 +249,7 @@ extension AsyncSequence where Element == any AGUIEvent { /// - Returns: A throwing stream of verified events. public func verifyEvents(debug: Bool = false) -> AsyncThrowingStream { AsyncThrowingStream { continuation in - Task { + let task = Task { let verifier = EventVerifier(debug: debug) do { for try await event in self { @@ -261,6 +261,7 @@ extension AsyncSequence where Element == any AGUIEvent { continuation.finish(throwing: error) } } + continuation.onTermination = { _ in task.cancel() } } } } diff --git a/Sources/AGUIClient/Subscriber/SubscriberManager.swift b/Sources/AGUIClient/Subscriber/SubscriberManager.swift index 2677022..4e206c5 100644 --- a/Sources/AGUIClient/Subscriber/SubscriberManager.swift +++ b/Sources/AGUIClient/Subscriber/SubscriberManager.swift @@ -59,8 +59,16 @@ actor DefaultAgentSubscription: AgentSubscription { /// /// This actor maintains the list of active subscribers and provides /// methods to execute them in sequence with mutation chaining. +/// +/// Subscribers are returned in **registration order** — the order in which +/// `subscribe(_:)` was called. This is guaranteed even after interleaved +/// `unsubscribe` calls. A parallel `insertionOrder` array tracks the sequence +/// separately from the dictionary so that `Dictionary.values` (undefined order) +/// is never used as an execution order. public actor SubscriberManager { private var subscribers: [UUID: any AgentSubscriber] = [:] + /// Tracks the order in which subscribers were registered. + private var insertionOrder: [UUID] = [] public init() {} @@ -71,6 +79,7 @@ public actor SubscriberManager { public func subscribe(_ subscriber: any AgentSubscriber) -> UUID { let id = UUID() subscribers[id] = subscriber + insertionOrder.append(id) return id } @@ -79,13 +88,14 @@ public actor SubscriberManager { /// - Parameter id: The subscription ID to remove public func unsubscribe(_ id: UUID) { subscribers.removeValue(forKey: id) + insertionOrder.removeAll { $0 == id } } - /// Returns all active subscribers. + /// Returns all active subscribers in registration order. /// - /// - Returns: Array of all subscribers in registration order + /// - Returns: Array of all subscribers ordered by `subscribe(_:)` call sequence public func allSubscribers() -> [any AgentSubscriber] { - Array(subscribers.values) + insertionOrder.compactMap { subscribers[$0] } } } diff --git a/Sources/AGUIClient/Transport/AgentTransport.swift b/Sources/AGUIClient/Transport/AgentTransport.swift new file mode 100644 index 0000000..8f64993 --- /dev/null +++ b/Sources/AGUIClient/Transport/AgentTransport.swift @@ -0,0 +1,30 @@ +/* + * MIT License + * + * Copyright (c) 2025 Perfect Aduh + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +import AGUICore +import Foundation + +public protocol AgentTransport: Sendable { + func run(input: RunAgentInput) -> AsyncThrowingStream +} diff --git a/Sources/AGUIClient/Transport/HttpAgentConfiguration.swift b/Sources/AGUIClient/Transport/HttpAgentConfiguration.swift index 31d30db..c12ab1e 100644 --- a/Sources/AGUIClient/Transport/HttpAgentConfiguration.swift +++ b/Sources/AGUIClient/Transport/HttpAgentConfiguration.swift @@ -45,37 +45,52 @@ public struct HttpAgentConfiguration: Sendable { /// Bearer token for authentication. /// - /// When set, automatically adds `Authorization: Bearer ` to ``headers``. - /// Setting to `nil` does not remove a manually added `Authorization` header. + /// When set, ``buildHeaders()`` includes `Authorization: Bearer `. + /// This property does **not** mutate ``headers`` directly — call + /// ``buildHeaders()`` to get the merged header dictionary for requests. /// /// Default: `nil` - public var bearerToken: String? { - didSet { - if let token = bearerToken { - headers["Authorization"] = "Bearer \(token)" - } - } - } + public var bearerToken: String? /// API key value. /// - /// When set, automatically adds the key to ``headers`` under ``apiKeyHeader``. - /// Setting to `nil` does not remove a manually added header. + /// When set, ``buildHeaders()`` includes the key under ``apiKeyHeader``. + /// This property does **not** mutate ``headers`` directly — call + /// ``buildHeaders()`` to get the merged header dictionary for requests. /// /// Default: `nil` - public var apiKey: String? { - didSet { - if let key = apiKey { - headers[apiKeyHeader] = key - } - } - } + public var apiKey: String? /// Header name used when adding the ``apiKey``. /// /// Default: `"X-API-Key"` public var apiKeyHeader: String + // MARK: - Header builder + + /// Returns the final HTTP header dictionary, merging ``bearerToken`` and + /// ``apiKey`` into ``headers``. + /// + /// Priority (highest → lowest): + /// 1. Entries already in ``headers`` (override everything) + /// 2. `bearerToken` → `Authorization: Bearer ` + /// 3. `apiKey` → `: ` + /// + /// - Returns: Merged header dictionary ready for request construction. + public func buildHeaders() -> [String: String] { + var result: [String: String] = [:] + if let key = apiKey { + result[apiKeyHeader] = key + } + if let token = bearerToken { + result["Authorization"] = "Bearer \(token)" + } + for (k, v) in headers { + result[k] = v + } + return result + } + /// Retry policy options. public enum RetryPolicy: Sendable { /// No retry on failure. diff --git a/Sources/AGUIClient/Transport/HttpAgentTransport.swift b/Sources/AGUIClient/Transport/HttpAgentTransport.swift new file mode 100644 index 0000000..0cc71ba --- /dev/null +++ b/Sources/AGUIClient/Transport/HttpAgentTransport.swift @@ -0,0 +1,130 @@ +/* + * MIT License + * + * Copyright (c) 2025 Perfect Aduh + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +import AGUICore +import Foundation + +public struct HttpAgentTransport: AgentTransport { + private let transport: HttpTransport + private let decoder: AGUIEventDecoder + private let endpoint: String + private let configuration: HttpAgentConfiguration + + public init(configuration: HttpAgentConfiguration, endpoint: String = "/run") { + self.configuration = configuration + self.transport = HttpTransport(configuration: configuration) + self.decoder = AGUIEventDecoder() + self.endpoint = endpoint + } + + public init(configuration: HttpAgentConfiguration, httpClient: any HTTPClient, endpoint: String = "/run") { + self.configuration = configuration + self.transport = HttpTransport(configuration: configuration, httpClient: httpClient) + self.decoder = AGUIEventDecoder() + self.endpoint = endpoint + } + + public func run(input: RunAgentInput) -> AsyncThrowingStream { + let transport = self.transport + let decoder = self.decoder + let endpoint = self.endpoint + let configuration = self.configuration + + return AsyncThrowingStream { continuation in + let task = Task { + var lastEventId: String? = nil + var attempt = 0 + var currentStream: EventStream>? = nil + + while true { + do { + let bytes = try await transport.execute( + endpoint: endpoint, + input: input, + lastEventId: lastEventId + ) + let eventStream = EventStream(bytes: bytes, decoder: decoder) + currentStream = eventStream + for try await event in eventStream { + continuation.yield(event) + } + continuation.finish() + return + } catch { + if let stream = currentStream { + lastEventId = await stream.lastEventId ?? lastEventId + } + currentStream = nil + + guard shouldRetry(error: error, attempt: attempt, configuration: configuration) else { + continuation.finish(throwing: error) + return + } + + let delay = retryDelay(for: attempt, configuration: configuration) + attempt += 1 + + if delay > 0 { + try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + } + } + } + } + continuation.onTermination = { _ in task.cancel() } + } + } + + private func shouldRetry(error: Error, attempt: Int, configuration: HttpAgentConfiguration) -> Bool { + guard isRetryable(error) else { return false } + switch configuration.retryPolicy { + case .none: + return false + case .fixed(let maxAttempts, _): + return attempt < maxAttempts + case .exponentialBackoff(let maxAttempts, _): + return attempt < maxAttempts + } + } + + private func isRetryable(_ error: Error) -> Bool { + guard let clientError = error as? ClientError else { return false } + switch clientError { + case .timeout, .networkError: + return true + default: + return false + } + } + + private func retryDelay(for attempt: Int, configuration: HttpAgentConfiguration) -> TimeInterval { + switch configuration.retryPolicy { + case .none: + return 0 + case .fixed(_, let delay): + return delay + case .exponentialBackoff(_, let baseDelay): + return min(baseDelay * pow(2.0, Double(attempt)), 60.0) + } + } +} diff --git a/Sources/AGUIClient/Transport/HttpTransport.swift b/Sources/AGUIClient/Transport/HttpTransport.swift index fe6f468..a7baa6d 100644 --- a/Sources/AGUIClient/Transport/HttpTransport.swift +++ b/Sources/AGUIClient/Transport/HttpTransport.swift @@ -99,8 +99,10 @@ public actor HttpTransport { sessionConfig.timeoutIntervalForRequest = max(configuration.timeout, 300) sessionConfig.timeoutIntervalForResource = max(configuration.timeout, 3600) - // Set additional headers - var headers = configuration.headers + // Merge auth + explicit headers, then add User-Agent. + // buildHeaders() unifies bearerToken, apiKey, and headers in one call + // so auth is always applied even when set after init. + var headers = configuration.buildHeaders() headers["User-Agent"] = "AGUISwift/1.0" sessionConfig.httpAdditionalHeaders = headers diff --git a/Sources/AGUICore/Decoding/AGUIEventDecoder.swift b/Sources/AGUICore/Decoding/AGUIEventDecoder.swift index e88f314..b2be905 100644 --- a/Sources/AGUICore/Decoding/AGUIEventDecoder.swift +++ b/Sources/AGUICore/Decoding/AGUIEventDecoder.swift @@ -110,7 +110,7 @@ import Foundation /// is immutable after initialization, and all configuration is `Sendable`. /// /// - SeeAlso: `AGUIEvent`, `EventType`, `EventDecodingError`, `UnknownEvent` -public struct AGUIEventDecoder { +public struct AGUIEventDecoder: Sendable { /// Handler function type for decoding a specific event type. /// diff --git a/Sources/AGUITools/AGUITools.swift b/Sources/AGUITools/AGUITools.swift index 874c4fe..badb7d3 100644 --- a/Sources/AGUITools/AGUITools.swift +++ b/Sources/AGUITools/AGUITools.swift @@ -75,9 +75,4 @@ public struct AGUITools { self.core = AGUICore() } - /// Legacy tools functionality example - @available(*, deprecated, message: "Use ToolExecutor protocol instead") - public func toolsFunction() -> String { - "AGUITools is working with \(core.coreFunction())" - } } diff --git a/Sources/AGUITools/Core/ToolExecutionManager.swift b/Sources/AGUITools/Core/ToolExecutionManager.swift index 8fea7e8..dfc9ab2 100644 --- a/Sources/AGUITools/Core/ToolExecutionManager.swift +++ b/Sources/AGUITools/Core/ToolExecutionManager.swift @@ -105,7 +105,7 @@ public actor ToolExecutionManager { runId: String? ) -> AsyncThrowingStream where S.Element == any AGUIEvent { AsyncThrowingStream { continuation in - Task { + let task = Task { do { for try await event in events { continuation.yield(event) @@ -117,6 +117,7 @@ public actor ToolExecutionManager { continuation.finish(throwing: error) } } + continuation.onTermination = { _ in task.cancel() } } } @@ -231,17 +232,19 @@ public actor ToolExecutionManager { // MARK: - ToolCallBuilder /// Builds a `ToolCall` from streaming events by accumulating argument deltas. -private final class ToolCallBuilder: @unchecked Sendable { +/// +/// Value type: all mutation happens inside the `ToolExecutionManager` actor, +private struct ToolCallBuilder { let id: String let name: String - private var arguments: String = "" + private(set) var arguments: String = "" init(id: String, name: String) { self.id = id self.name = name } - func appendArguments(_ delta: String) { + mutating func appendArguments(_ delta: String) { arguments += delta } diff --git a/Sources/AGUITools/Registry/ToolRegistry.swift b/Sources/AGUITools/Registry/ToolRegistry.swift index b2f4631..a59a317 100644 --- a/Sources/AGUITools/Registry/ToolRegistry.swift +++ b/Sources/AGUITools/Registry/ToolRegistry.swift @@ -255,8 +255,8 @@ public actor DefaultToolRegistry: ToolRegistry { } public func clearStats() async { - for (_, stat) in stats { - stat.clear() + for key in stats.keys { + stats[key]?.clear() } } @@ -269,31 +269,30 @@ public actor DefaultToolRegistry: ToolRegistry { /// Mutable version of ToolExecutionStats for internal tracking. /// -/// This class is used internally by the registry to track and update -/// statistics efficiently. It provides methods for recording successes -/// and failures while maintaining accurate averages. -private final class MutableToolExecutionStats: @unchecked Sendable { - private var executionCount: Int = 0 - private var successCount: Int = 0 - private var failureCount: Int = 0 - private var totalExecutionTime: Duration = .zero - private var averageExecutionTime: Duration = .zero - - func recordSuccess(duration: Duration) { +/// Value type: all mutation happens through the `DefaultToolRegistry` actor's +/// dictionary subscript (inout), which serialises access — no `@unchecked Sendable` needed. +private struct MutableToolExecutionStats { + private(set) var executionCount: Int = 0 + private(set) var successCount: Int = 0 + private(set) var failureCount: Int = 0 + private(set) var totalExecutionTime: Duration = .zero + private(set) var averageExecutionTime: Duration = .zero + + mutating func recordSuccess(duration: Duration) { executionCount += 1 successCount += 1 totalExecutionTime += duration averageExecutionTime = totalExecutionTime / executionCount } - func recordFailure(duration: Duration) { + mutating func recordFailure(duration: Duration) { executionCount += 1 failureCount += 1 totalExecutionTime += duration averageExecutionTime = totalExecutionTime / executionCount } - func clear() { + mutating func clear() { executionCount = 0 successCount = 0 failureCount = 0 diff --git a/Tests/AGUIAgentSDKTests/AgUiAgentTests.swift b/Tests/AGUIAgentSDKTests/AgUiAgentTests.swift index 50d008f..1f56cf1 100644 --- a/Tests/AGUIAgentSDKTests/AgUiAgentTests.swift +++ b/Tests/AGUIAgentSDKTests/AgUiAgentTests.swift @@ -22,6 +22,7 @@ * SOFTWARE. */ +import AGUIClient import AGUICore import AGUITools import XCTest @@ -29,33 +30,29 @@ import XCTest // MARK: - Test helpers -/// Subclass of `AgUiAgent` that captures every `RunAgentInput` passed to `run(input:)`. -private final class CapturingAgent: AgUiAgent, @unchecked Sendable { - private let lock = NSLock() - private var _inputs: [RunAgentInput] = [] - - var capturedInputs: [RunAgentInput] { - lock.lock() - defer { lock.unlock() } - return _inputs - } - - /// Events yielded by the mock `run(input:)`. +actor CapturingTransport: AgentTransport { + private(set) var capturedInputs: [RunAgentInput] = [] var mockEvents: [any AGUIEvent] = [] - override func run(input: RunAgentInput) -> AsyncThrowingStream { - lock.lock() - _inputs.append(input) - lock.unlock() + func setMockEvents(_ events: [any AGUIEvent]) { + mockEvents = events + } - let events = mockEvents - return AsyncThrowingStream { continuation in - for event in events { - continuation.yield(event) + nonisolated func run(input: RunAgentInput) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let task = Task { + await self.record(input) + let events = await self.mockEvents + for event in events { continuation.yield(event) } + continuation.finish() } - continuation.finish() + continuation.onTermination = { _ in task.cancel() } } } + + private func record(_ input: RunAgentInput) { + capturedInputs.append(input) + } } // MARK: - Mock ToolRegistry @@ -96,31 +93,41 @@ final class AgUiAgentTests: XCTestCase { private let agentURL = URL(string: "https://agent.example.com")! + private func makeCapturingAgent( + configure: (inout AgUiAgentConfig) -> Void = { _ in } + ) -> (AgUiAgent, CapturingTransport) { + var cfg = AgUiAgentConfig() + configure(&cfg) + let transport = CapturingTransport() + let agent = AgUiAgent(transport: transport, config: cfg) + return (agent, transport) + } + // MARK: - sendMessage constructs correct RunAgentInput func testSendMessageProducesUserMessage() async throws { - let agent = CapturingAgent(url: agentURL) + let (agent, transport) = makeCapturingAgent() let stream = agent.sendMessage("Hello!") for try await _ in stream {} - let inputs = agent.capturedInputs + let inputs = await transport.capturedInputs XCTAssertEqual(inputs.count, 1) let input = try XCTUnwrap(inputs.first) - // Only user message (no system prompt configured) XCTAssertEqual(input.messages.count, 1) XCTAssertEqual(input.messages[0].role, .user) XCTAssertEqual(input.messages[0].content, "Hello!") } func testSendMessagePrependsSystemPromptWhenConfigured() async throws { - let agent = CapturingAgent(url: agentURL) { config in + let (agent, transport) = makeCapturingAgent { config in config.systemPrompt = "Be concise." } let stream = agent.sendMessage("Hi", includeSystemPrompt: true) for try await _ in stream {} - let input = try XCTUnwrap(agent.capturedInputs.first) + let captured1 = await transport.capturedInputs + let input = try XCTUnwrap(captured1.first) XCTAssertEqual(input.messages.count, 2) XCTAssertEqual(input.messages[0].role, .system) XCTAssertEqual(input.messages[0].content, "Be concise.") @@ -128,49 +135,50 @@ final class AgUiAgentTests: XCTestCase { } func testSendMessageOmitsSystemPromptWhenDisabled() async throws { - let agent = CapturingAgent(url: agentURL) { config in + let (agent, transport) = makeCapturingAgent { config in config.systemPrompt = "Be concise." } let stream = agent.sendMessage("Hi", includeSystemPrompt: false) for try await _ in stream {} - let input = try XCTUnwrap(agent.capturedInputs.first) + let captured2 = await transport.capturedInputs + let input = try XCTUnwrap(captured2.first) XCTAssertEqual(input.messages.count, 1) XCTAssertEqual(input.messages[0].role, .user) } func testSendMessageUsesProvidedThreadId() async throws { - let agent = CapturingAgent(url: agentURL) + let (agent, transport) = makeCapturingAgent() let stream = agent.sendMessage("Hello", threadId: "my-thread") for try await _ in stream {} - let input = try XCTUnwrap(agent.capturedInputs.first) + let captured3 = await transport.capturedInputs + let input = try XCTUnwrap(captured3.first) XCTAssertEqual(input.threadId, "my-thread") } func testSendMessageUsesProvidedState() async throws { let customState = Data("{\"mode\":\"test\"}".utf8) - let agent = CapturingAgent(url: agentURL) + let (agent, transport) = makeCapturingAgent() let stream = agent.sendMessage("Hello", state: customState) for try await _ in stream {} - let input = try XCTUnwrap(agent.capturedInputs.first) + let captured4 = await transport.capturedInputs + let input = try XCTUnwrap(captured4.first) XCTAssertEqual(input.state, customState) } func testSendMessageEachCallIsFreshNoHistory() async throws { - let agent = CapturingAgent(url: agentURL) + let (agent, transport) = makeCapturingAgent() - // First call for try await _ in agent.sendMessage("Message 1") {} - // Second call for try await _ in agent.sendMessage("Message 2") {} - XCTAssertEqual(agent.capturedInputs.count, 2) + let capturedAll = await transport.capturedInputs + XCTAssertEqual(capturedAll.count, 2) - // Each call should only contain its own user message - let first = agent.capturedInputs[0] - let second = agent.capturedInputs[1] + let first = capturedAll[0] + let second = capturedAll[1] XCTAssertEqual(first.messages.count, 1) XCTAssertEqual(first.messages[0].content, "Message 1") @@ -186,25 +194,27 @@ final class AgUiAgentTests: XCTestCase { let tool2 = Tool(name: "search_web", description: "Search the web", parameters: Data("{}".utf8)) let registry = MockToolRegistry(tools: [tool1, tool2]) - let agent = CapturingAgent(url: agentURL) { config in + let (agent, transport) = makeCapturingAgent { config in config.toolRegistry = registry } let stream = agent.sendMessage("What's the weather?") for try await _ in stream {} - let input = try XCTUnwrap(agent.capturedInputs.first) + let captured5 = await transport.capturedInputs + let input = try XCTUnwrap(captured5.first) XCTAssertEqual(input.tools.count, 2) XCTAssertEqual(input.tools[0].name, "get_weather") XCTAssertEqual(input.tools[1].name, "search_web") } func testSendMessageHasEmptyToolsWhenNoRegistry() async throws { - let agent = CapturingAgent(url: agentURL) + let (agent, transport) = makeCapturingAgent() let stream = agent.sendMessage("Hello") for try await _ in stream {} - let input = try XCTUnwrap(agent.capturedInputs.first) + let captured6 = await transport.capturedInputs + let input = try XCTUnwrap(captured6.first) XCTAssertTrue(input.tools.isEmpty) } @@ -212,13 +222,14 @@ final class AgUiAgentTests: XCTestCase { func testSendMessageIncludesContext() async throws { let ctx = Context(description: "timezone", value: "America/New_York") - let agent = CapturingAgent(url: agentURL) { config in + let (agent, transport) = makeCapturingAgent { config in config.context = [ctx] } let stream = agent.sendMessage("What time is it?") for try await _ in stream {} - let input = try XCTUnwrap(agent.capturedInputs.first) + let captured7 = await transport.capturedInputs + let input = try XCTUnwrap(captured7.first) XCTAssertEqual(input.context.count, 1) XCTAssertEqual(input.context[0].description, "timezone") } @@ -226,12 +237,11 @@ final class AgUiAgentTests: XCTestCase { // MARK: - Event passthrough func testSendMessageYieldsEventsFromRun() async throws { - let agentURL = URL(string: "https://agent.example.com")! - let agent = CapturingAgent(url: agentURL) - agent.mockEvents = [ + let (agent, transport) = makeCapturingAgent() + await transport.setMockEvents([ RunStartedEvent(threadId: "t1", runId: "r1"), RunFinishedEvent(threadId: "t1", runId: "r1"), - ] + ]) var received: [any AGUIEvent] = [] for try await event in agent.sendMessage("Hi") { @@ -247,13 +257,12 @@ final class AgUiAgentTests: XCTestCase { func testCloseDoesNotCrash() { let agent = AgUiAgent(url: agentURL) - // Verify close() completes without throwing or crashing agent.close() } func testCloseCanBeCalledMultipleTimes() { let agent = AgUiAgent(url: agentURL) agent.close() - agent.close() // second call must not crash + agent.close() } } diff --git a/Tests/AGUIAgentSDKTests/EndToEndPipelineTests.swift b/Tests/AGUIAgentSDKTests/EndToEndPipelineTests.swift index af16dc4..3af87b8 100644 --- a/Tests/AGUIAgentSDKTests/EndToEndPipelineTests.swift +++ b/Tests/AGUIAgentSDKTests/EndToEndPipelineTests.swift @@ -28,35 +28,29 @@ import AGUITools import XCTest @testable import AGUIAgentSDK -// MARK: - MockAbstractAgent - -/// `AbstractAgent` subclass that emits a preset sequence of events. -/// -/// Each call to `run(input:)` pops the next event sequence from `eventSequences`. -/// When the queue is empty, subsequent calls emit an empty stream. -private final class MockAbstractAgent: AbstractAgent, @unchecked Sendable { - private let lock = NSLock() +// MARK: - MockAgentTransport + +actor MockAgentTransport: AgentTransport { private var _eventSequences: [[any AGUIEvent]] = [] - /// Enqueue event sequences to be emitted on successive `run(input:)` calls. func enqueue(_ events: [any AGUIEvent]) { - lock.lock() _eventSequences.append(events) - lock.unlock() } - override func run(input: RunAgentInput) -> AsyncThrowingStream { - lock.lock() - let events: [any AGUIEvent] = _eventSequences.isEmpty ? [] : _eventSequences.removeFirst() - lock.unlock() - - return AsyncThrowingStream { continuation in - for event in events { - continuation.yield(event) + nonisolated func run(input: RunAgentInput) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let task = Task { + let events = await self.dequeue() + for event in events { continuation.yield(event) } + continuation.finish() } - continuation.finish() + continuation.onTermination = { _ in task.cancel() } } } + + private func dequeue() -> [any AGUIEvent] { + _eventSequences.isEmpty ? [] : _eventSequences.removeFirst() + } } // MARK: - EndToEndPipelineTests @@ -65,10 +59,9 @@ final class EndToEndPipelineTests: XCTestCase { // MARK: - Text-only conversation - /// Full text-message pipeline: events accumulate into an `AssistantMessage` in `agent.messages`. func testTextOnlyConversationBuildsMessages() async throws { - let agent = MockAbstractAgent() - agent.enqueue([ + let mockTransport = MockAgentTransport() + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), TextMessageStartEvent(messageId: "msg1"), TextMessageContentEvent(messageId: "msg1", delta: "Hello"), @@ -77,6 +70,7 @@ final class EndToEndPipelineTests: XCTestCase { RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let messages = await agent.messages @@ -88,22 +82,20 @@ final class EndToEndPipelineTests: XCTestCase { // MARK: - State snapshot & delta - /// `STATE_SNAPSHOT` sets initial state; `STATE_DELTA` patches it. func testStateDeltaIsAppliedCorrectly() async throws { - let agent = MockAbstractAgent() + let mockTransport = MockAgentTransport() - // Initial state: {"count": 0} let initialJSON = Data("{\"count\":0}".utf8) - // Delta: [{"op":"replace","path":"/count","value":5}] let patchJSON = Data("[{\"op\":\"replace\",\"path\":\"/count\",\"value\":5}]".utf8) - agent.enqueue([ + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), StateSnapshotEvent(snapshot: initialJSON), StateDeltaEvent(delta: patchJSON), RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let finalState = await agent.state @@ -116,15 +108,16 @@ final class EndToEndPipelineTests: XCTestCase { } func testStateSnapshotReplacesState() async throws { - let agent = MockAbstractAgent() + let mockTransport = MockAgentTransport() let snapshotJSON = Data("{\"mode\":\"creative\"}".utf8) - agent.enqueue([ + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), StateSnapshotEvent(snapshot: snapshotJSON), RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let state = await agent.state @@ -138,10 +131,9 @@ final class EndToEndPipelineTests: XCTestCase { // MARK: - Thinking telemetry - /// Full thinking sequence builds `ThinkingTelemetryState`. func testThinkingSequenceBuildsThinkingState() async throws { - let agent = MockAbstractAgent() - agent.enqueue([ + let mockTransport = MockAgentTransport() + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), ThinkingStartEvent(title: "Step 1"), ThinkingTextMessageStartEvent(), @@ -151,6 +143,7 @@ final class EndToEndPipelineTests: XCTestCase { RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let thinking = await agent.thinking @@ -161,10 +154,9 @@ final class EndToEndPipelineTests: XCTestCase { } func testRunStartedResetsThinkingState() async throws { - let agent = MockAbstractAgent() + let mockTransport = MockAgentTransport() - // First run leaves thinking state - agent.enqueue([ + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), ThinkingStartEvent(), ThinkingTextMessageStartEvent(), @@ -174,30 +166,31 @@ final class EndToEndPipelineTests: XCTestCase { RunFinishedEvent(threadId: "t1", runId: "r1"), ]) - // Second run: thinking state should be reset on RUN_STARTED - agent.enqueue([ + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r2"), RunFinishedEvent(threadId: "t1", runId: "r2"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let thinking = await agent.thinking - XCTAssertNotNil(thinking) // was set by first run + XCTAssertNotNil(thinking) } // MARK: - Sequential multi-run (state persists) func testSequentialRunsMaintainState() async throws { - let agent = MockAbstractAgent() + let mockTransport = MockAgentTransport() let state1 = Data("{\"turn\":1}".utf8) - agent.enqueue([ + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), StateSnapshotEvent(snapshot: state1), RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let stateAfterRun1 = await agent.state @@ -209,7 +202,7 @@ final class EndToEndPipelineTests: XCTestCase { XCTAssertEqual(turn1, 1) let state2 = Data("{\"turn\":2}".utf8) - agent.enqueue([ + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r2"), StateSnapshotEvent(snapshot: state2), RunFinishedEvent(threadId: "t1", runId: "r2"), @@ -228,15 +221,14 @@ final class EndToEndPipelineTests: XCTestCase { // MARK: - Invalid event stream — EventVerifier - /// A stream that starts with a non-RUN_STARTED event must throw `AGUIProtocolError`. func testInvalidEventStreamThrowsProtocolError() async throws { - let agent = MockAbstractAgent() + let mockTransport = MockAgentTransport() - // Intentionally wrong: TextMessageStart without RUN_STARTED - agent.enqueue([ + await mockTransport.enqueue([ TextMessageStartEvent(messageId: "msg1"), ]) + let agent = AbstractAgent(transport: mockTransport) do { try await agent.runAgent() XCTFail("Expected AGUIProtocolError to be thrown") @@ -248,15 +240,15 @@ final class EndToEndPipelineTests: XCTestCase { } func testRunAfterRunErrorThrows() async throws { - let agent = MockAbstractAgent() + let mockTransport = MockAgentTransport() - agent.enqueue([ + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), RunErrorEvent(message: "fatal error", code: "FATAL"), - // Any event after RUN_ERROR should cause a protocol error TextMessageStartEvent(messageId: "msg1"), ]) + let agent = AbstractAgent(transport: mockTransport) do { try await agent.runAgent() XCTFail("Expected AGUIProtocolError to be thrown") @@ -269,10 +261,9 @@ final class EndToEndPipelineTests: XCTestCase { // MARK: - Tool call accumulation - /// Tool call events (START/ARGS/END) build an `AssistantMessage` with `toolCalls`. func testToolCallSequenceBuildsAssistantMessageWithToolCalls() async throws { - let agent = MockAbstractAgent() - agent.enqueue([ + let mockTransport = MockAgentTransport() + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), ToolCallStartEvent(toolCallId: "tc1", toolCallName: "get_weather"), ToolCallArgsEvent(toolCallId: "tc1", delta: "{\"city\":"), @@ -281,6 +272,7 @@ final class EndToEndPipelineTests: XCTestCase { RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let messages = await agent.messages @@ -300,14 +292,15 @@ final class EndToEndPipelineTests: XCTestCase { // MARK: - Custom & raw events func testRawEventsAccumulate() async throws { - let agent = MockAbstractAgent() - agent.enqueue([ + let mockTransport = MockAgentTransport() + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), RawEvent(data: Data("{\"event\":\"raw_1\"}".utf8)), RawEvent(data: Data("{\"event\":\"raw_2\"}".utf8)), RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let rawEvents = await agent.rawEvents @@ -315,14 +308,15 @@ final class EndToEndPipelineTests: XCTestCase { } func testCustomEventsAccumulate() async throws { - let agent = MockAbstractAgent() - agent.enqueue([ + let mockTransport = MockAgentTransport() + await mockTransport.enqueue([ RunStartedEvent(threadId: "t1", runId: "r1"), CustomEvent(name: "ping", value: Data("{\"ts\":1}".utf8)), CustomEvent(name: "pong", value: Data("{\"ts\":2}".utf8)), RunFinishedEvent(threadId: "t1", runId: "r1"), ]) + let agent = AbstractAgent(transport: mockTransport) try await agent.runAgent() let customEvents = await agent.customEvents @@ -363,7 +357,7 @@ final class EndToEndPipelineTests: XCTestCase { // MARK: - SimpleMockExecutor -private final class SimpleMockExecutor: ToolExecutor, @unchecked Sendable { +private final class SimpleMockExecutor: ToolExecutor, Sendable { let tool: Tool init(toolName: String, description: String) { diff --git a/Tests/AGUIAgentSDKTests/MockChatAgent.swift b/Tests/AGUIAgentSDKTests/MockChatAgent.swift index 5eaf39f..aaaac38 100644 --- a/Tests/AGUIAgentSDKTests/MockChatAgent.swift +++ b/Tests/AGUIAgentSDKTests/MockChatAgent.swift @@ -28,9 +28,11 @@ import Foundation /// Test double for `ChatAgent`. /// -/// Configured before the test (setup), then called sequentially — no concurrent -/// access — so `@unchecked Sendable` without a lock is safe. -final class MockChatAgent: ChatAgent, @unchecked Sendable { +/// Configured and read exclusively from `@MainActor` context (both test classes +/// are `@MainActor`), so `@MainActor` isolation satisfies the `Sendable` requirement +/// without locks or `@unchecked`. +@MainActor +final class MockChatAgent: ChatAgent, Sendable { // MARK: - Configuration (set before each test) diff --git a/Tests/AGUIClientTests/SseReconnectionTests.swift b/Tests/AGUIClientTests/SseReconnectionTests.swift index e373a99..71197d6 100644 --- a/Tests/AGUIClientTests/SseReconnectionTests.swift +++ b/Tests/AGUIClientTests/SseReconnectionTests.swift @@ -36,7 +36,7 @@ import XCTest /// - `.midStreamFailure(sseText:error:)` — streams bytes then throws mid-stream actor SequencedHTTPClient: HTTPClient { - enum Response: @unchecked Sendable { + enum Response: Sendable { case success(HTTPResponse) case failure(Error) /// Streams `sseText` as bytes, then throws `error` — simulates a network drop @@ -443,7 +443,8 @@ final class SseReconnectionTests: XCTestCase { for try await _ in stream {} - XCTAssertEqual(stream.lastEventId, "second-id", "Must track the LAST seen SSE id") + let lastId = await stream.lastEventId + XCTAssertEqual(lastId, "second-id", "Must track the LAST seen SSE id") } func testEventStream_lastEventIdIsNil_whenStreamHasNoIds() async throws { @@ -457,7 +458,8 @@ final class SseReconnectionTests: XCTestCase { for try await _ in stream {} - XCTAssertNil(stream.lastEventId, "lastEventId must be nil when no id: fields appear") + let lastId = await stream.lastEventId + XCTAssertNil(lastId, "lastEventId must be nil when no id: fields appear") } func testEventStream_lastEventIdUpdatesAcrossEvents() async throws { @@ -478,12 +480,13 @@ final class SseReconnectionTests: XCTestCase { for try await _ in stream { eventCount += 1 if eventCount == 1 { - idAfterFirst = stream.lastEventId + idAfterFirst = await stream.lastEventId } } XCTAssertEqual(idAfterFirst, "alpha", "After first event, id should be 'alpha'") - XCTAssertEqual(stream.lastEventId, "beta", "After all events, id should be 'beta'") + let finalId = await stream.lastEventId + XCTAssertEqual(finalId, "beta", "After all events, id should be 'beta'") } func testEventStream_lastEventIdIsSetEvenWhenEventFailsToDecode() async throws { @@ -499,8 +502,9 @@ final class SseReconnectionTests: XCTestCase { for try await _ in stream {} + let lastId = await stream.lastEventId XCTAssertEqual( - stream.lastEventId, + lastId, "orphan-id", "id must be captured even when event data fails to decode" ) diff --git a/Tests/AGUIClientTests/Streaming/BufferingTests.swift b/Tests/AGUIClientTests/Streaming/BufferingTests.swift index 7140778..69d2dca 100644 --- a/Tests/AGUIClientTests/Streaming/BufferingTests.swift +++ b/Tests/AGUIClientTests/Streaming/BufferingTests.swift @@ -378,6 +378,64 @@ final class BufferingTests: XCTestCase { XCTAssertGreaterThan(receivedBatches, 0) } + // MARK: - Time-Window Correctness Tests + + func testTimeBatchedWindowFiresEvenWhenUpstreamIsIdle() async throws { + // Upstream yields one element then pauses 500 ms before the second. + // A 50 ms window must return the first batch after ~50 ms — NOT block for 500 ms. + let source = AsyncThrowingStream { continuation in + Task { + continuation.yield("a") + try await Task.sleep(for: .milliseconds(500)) + continuation.yield("b") + continuation.finish() + } + } + + let start = ContinuousClock.now + var batches: [[String]] = [] + for try await batch in source.batched(timeWindow: 0.05) { + batches.append(batch) + // After the first batch arrives, check elapsed wall time + if batches.count == 1 { + let elapsed = ContinuousClock.now - start + // Must have returned within 200 ms (50 ms window + generous tolerance) + XCTAssertLessThan( + elapsed, + .milliseconds(200), + "First batch must not block for the full 500 ms upstream pause" + ) + } + } + + XCTAssertGreaterThanOrEqual(batches.count, 2) + XCTAssertEqual(batches.flatMap { $0 }, ["a", "b"]) + } + + func testTimeBatchedGroupsFastElementsIntoBatch() async throws { + // Upstream yields 6 elements at 10 ms intervals — well within a 50 ms window. + // At least some elements should land in the same batch. + let source = AsyncStream { continuation in + Task { + for i in 0 ..< 6 { + continuation.yield(i) + try? await Task.sleep(for: .milliseconds(10)) + } + continuation.finish() + } + } + + var batches: [[Int]] = [] + for try await batch in source.batched(timeWindow: 0.05) { + batches.append(batch) + } + + // At 10 ms per element and a 50 ms window, at least 2 elements should batch together + let hasMultiElementBatch = batches.contains { $0.count > 1 } + XCTAssertTrue(hasMultiElementBatch, "Fast elements must be grouped within the window") + XCTAssertEqual(batches.flatMap { $0 }, [0, 1, 2, 3, 4, 5]) + } + // MARK: - Cancellation Tests func testBufferedCancellation() async throws { diff --git a/Tests/AGUIClientTests/Subscriber/AgentSubscriberTests.swift b/Tests/AGUIClientTests/Subscriber/AgentSubscriberTests.swift index 92458dc..1dd935a 100644 --- a/Tests/AGUIClientTests/Subscriber/AgentSubscriberTests.swift +++ b/Tests/AGUIClientTests/Subscriber/AgentSubscriberTests.swift @@ -370,6 +370,88 @@ final class AgentSubscriberTests: XCTestCase { XCTAssertEqual(subscribers.count, 3) } + func testSubscriberExecutionOrderMatchesRegistration() async throws { + // Given: A subscriber manager and an ordered recorder + let manager = SubscriberManager() + + actor CallOrderTracker { + var order: [Int] = [] + func record(_ index: Int) { order.append(index) } + } + let tracker = CallOrderTracker() + + struct IndexedSubscriber: AgentSubscriber { + let index: Int + let tracker: CallOrderTracker + + func onRunInitialized(params: AgentSubscriberParams) async -> AgentStateMutation? { + await tracker.record(index) + return nil + } + } + + let testInput = try RunAgentInput.builder() + .threadId("order-thread") + .runId("order-run") + .build() + + // When: 5 subscribers registered in order 0..4 + for i in 0 ..< 5 { + _ = await manager.subscribe(IndexedSubscriber(index: i, tracker: tracker)) + } + + // Drive each subscriber in allSubscribers() order + let params = AgentSubscriberParams(messages: [], state: Data("{}".utf8), input: testInput) + for sub in await manager.allSubscribers() { + _ = await sub.onRunInitialized(params: params) + } + + // Then: callbacks must arrive in registration order, not dict iteration order + let callOrder = await tracker.order + XCTAssertEqual(callOrder, [0, 1, 2, 3, 4]) + } + + func testUnsubscribePreservesRemainingOrder() async throws { + // Given: 4 subscribers registered in order 0..3 + let manager = SubscriberManager() + + actor CallOrderTracker { + var order: [Int] = [] + func record(_ index: Int) { order.append(index) } + } + let tracker = CallOrderTracker() + + struct IndexedSubscriber: AgentSubscriber { + let index: Int + let tracker: CallOrderTracker + func onRunInitialized(params: AgentSubscriberParams) async -> AgentStateMutation? { + await tracker.record(index) + return nil + } + } + + let testInput = try RunAgentInput.builder() + .threadId("t").runId("r").build() + + var ids: [UUID] = [] + for i in 0 ..< 4 { + let id = await manager.subscribe(IndexedSubscriber(index: i, tracker: tracker)) + ids.append(id) + } + + // When: unsubscribe the second subscriber (index 1) + await manager.unsubscribe(ids[1]) + + let params = AgentSubscriberParams(messages: [], state: Data("{}".utf8), input: testInput) + for sub in await manager.allSubscribers() { + _ = await sub.onRunInitialized(params: params) + } + + // Then: remaining subscribers fire in insertion order: [0, 2, 3] + let callOrder = await tracker.order + XCTAssertEqual(callOrder, [0, 2, 3]) + } + func testSubscriptionHandle() async throws { // Given: A subscription with tracked state actor UnsubscribeTracker { diff --git a/Tests/AGUIClientTests/Transport/HttpAgentConfigurationTests.swift b/Tests/AGUIClientTests/Transport/HttpAgentConfigurationTests.swift index 4123d3e..9c6b612 100644 --- a/Tests/AGUIClientTests/Transport/HttpAgentConfigurationTests.swift +++ b/Tests/AGUIClientTests/Transport/HttpAgentConfigurationTests.swift @@ -138,6 +138,61 @@ final class HttpAgentConfigurationTests: XCTestCase { } } + // MARK: - buildHeaders() Tests + + func testBuildHeadersIncludesBearerToken() { + var config = HttpAgentConfiguration(baseURL: URL(string: "https://example.com")!) + config.bearerToken = "sk-test-token" + XCTAssertEqual(config.buildHeaders()["Authorization"], "Bearer sk-test-token") + } + + func testBuildHeadersIncludesApiKey() { + var config = HttpAgentConfiguration(baseURL: URL(string: "https://example.com")!) + config.apiKey = "my-api-key" + XCTAssertEqual(config.buildHeaders()["X-API-Key"], "my-api-key") + } + + func testBuildHeadersUsesCustomApiKeyHeader() { + var config = HttpAgentConfiguration(baseURL: URL(string: "https://example.com")!) + config.apiKeyHeader = "X-Custom-Key" + config.apiKey = "val" + XCTAssertEqual(config.buildHeaders()["X-Custom-Key"], "val") + XCTAssertNil(config.buildHeaders()["X-API-Key"]) + } + + func testBuildHeadersIncludesExplicitHeaders() { + var config = HttpAgentConfiguration( + baseURL: URL(string: "https://example.com")!, + headers: ["X-Trace": "abc123"] + ) + XCTAssertEqual(config.buildHeaders()["X-Trace"], "abc123") + } + + func testBearerTokenDoesNotSideEffectRawHeadersDict() { + // Setting bearerToken must NOT mutate the .headers dict directly — + // auth headers are only visible through buildHeaders(). + var config = HttpAgentConfiguration(baseURL: URL(string: "https://example.com")!) + config.bearerToken = "sk-secret" + XCTAssertNil(config.headers["Authorization"], + "bearerToken must not mutate headers via didSet; use buildHeaders()") + } + + func testApiKeyDoesNotSideEffectRawHeadersDict() { + var config = HttpAgentConfiguration(baseURL: URL(string: "https://example.com")!) + config.apiKey = "key-secret" + XCTAssertNil(config.headers["X-API-Key"], + "apiKey must not mutate headers via didSet; use buildHeaders()") + } + + func testBuildHeadersMergesTokenOverApiKey() { + var config = HttpAgentConfiguration(baseURL: URL(string: "https://example.com")!) + config.bearerToken = "token" + config.apiKey = "key" + let built = config.buildHeaders() + XCTAssertEqual(built["Authorization"], "Bearer token") + XCTAssertEqual(built["X-API-Key"], "key") + } + // MARK: - URL Construction Tests func testBaseURLWithPath() { diff --git a/Tests/AGUIToolsTests/AGUIToolsTests.swift b/Tests/AGUIToolsTests/AGUIToolsTests.swift index c8a7c26..6f50a88 100644 --- a/Tests/AGUIToolsTests/AGUIToolsTests.swift +++ b/Tests/AGUIToolsTests/AGUIToolsTests.swift @@ -26,10 +26,9 @@ import XCTest @testable import AGUITools final class AGUIToolsTests: XCTestCase { - func testToolsFunction() { + func testAGUIToolsIsInstantiable() { + // AGUITools is deprecated; verify it still compiles and instantiates. let tools = AGUITools() - let result = tools.toolsFunction() - XCTAssertTrue(result.contains("AGUITools is working")) - XCTAssertTrue(result.contains("AGUICore is working")) + _ = tools } }