Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ let package = Package(
dependencies: ["AGUIClient"]),
.testTarget(
name: "AGUIAgentSDKTests",
dependencies: ["AGUIAgentSDK", "AGUITools"]),
dependencies: ["AGUIAgentSDK", "AGUIClient", "AGUITools"]),
.testTarget(
name: "AGUIToolsTests",
dependencies: ["AGUITools"]),
Expand Down
119 changes: 23 additions & 96 deletions Sources/AGUIAgentSDK/AgUiAgent.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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<any AGUIEvent, Error> {
/// // 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 }
Expand All @@ -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(
Expand All @@ -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<any AGUIEvent, Error> {
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<any AGUIEvent, Error> {
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<any AGUIEvent, Error> {
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 {
Expand All @@ -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()
Expand All @@ -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 {
Expand All @@ -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()
Expand Down
6 changes: 4 additions & 2 deletions Sources/AGUIAgentSDK/StatefulAgUiAgent.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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() }
}
}

Expand Down Expand Up @@ -361,7 +362,7 @@ public final class StatefulAgUiAgent: Sendable {
threadId: String
) -> AsyncThrowingStream<any AGUIEvent, Error> where S.Element == any AGUIEvent {
AsyncThrowingStream { continuation in
Task {
let task = Task {
var currentAssistantMessage: AssistantMessage?
let patchApplicator = PatchApplicator()
let messageDecoder = MessageDecoder()
Expand Down Expand Up @@ -501,6 +502,7 @@ public final class StatefulAgUiAgent: Sendable {
continuation.finish(throwing: error)
}
}
continuation.onTermination = { _ in task.cancel() }
}
}
}
Expand Down
38 changes: 29 additions & 9 deletions Sources/AGUIAgentSDK/StatefulAgUiAgentConfig.swift
Original file line number Diff line number Diff line change
Expand Up @@ -105,17 +105,12 @@ public struct StatefulAgUiAgentConfig: Sendable {

/// Bearer token for authentication.
///
/// When set, automatically adds an `Authorization: Bearer <token>` header
/// to every request.
/// When set, ``buildHeaders()`` includes `Authorization: Bearer <token>`.
/// 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.
///
Expand Down Expand Up @@ -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 <token>`
/// 3. `apiKey` → `<apiKeyHeader>: <key>`
///
/// - 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
Expand Down
2 changes: 1 addition & 1 deletion Sources/AGUIAgentSDK/Tools/ClientToolResponseHandler.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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?
Expand Down
Loading
Loading