diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/ClientIdStore.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/ClientIdStore.swift new file mode 100644 index 000000000..38b4b6fc7 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/ClientIdStore.swift @@ -0,0 +1,56 @@ +// ClientIdStore — pluggable persistence for stable per-host `clientId`s. + +import Foundation + +/// Persistence hook for stable `clientId`s per host. +/// +/// On `MultiHostClient.add(_:)`, the multi-host client looks up `hostId` +/// in this store. If the store returns a value, that id is reused — letting +/// the server treat successive launches as the same client (which the AHP +/// `reconnect` flow needs to replay missed actions). If the store returns +/// `nil`, the multi-host client generates a fresh UUID and stores it. +/// +/// The default `InMemoryClientIdStore` is **session-stable only** — it does +/// not survive process restarts. Production multi-host apps should plug a +/// keychain/file-backed implementation in so reconnects keep working across +/// launches. +public protocol ClientIdStore: AnyObject, Sendable { + /// Look up the previously stored `clientId` for `hostId`, if any. + func load(_ hostId: HostId) async -> String? + + /// Persist `clientId` for `hostId`. Implementations should overwrite any + /// previous value. + func store(_ hostId: HostId, clientId: String) async +} + +/// In-process `ClientIdStore` backed by an actor-protected dictionary. +/// +/// Keeps assigned ids in memory. Survives reconnects within the same process +/// but **not** restarts. Fine for tests, ephemeral CLIs, and as a starting +/// point — production apps should provide a persistent implementation +/// (filesystem, keychain, secure enclave, …). +public final class InMemoryClientIdStore: ClientIdStore { + private let storage: Storage + + public init() { + self.storage = Storage() + } + + public func load(_ hostId: HostId) async -> String? { + await storage.load(hostId) + } + + public func store(_ hostId: HostId, clientId: String) async { + await storage.store(hostId, clientId: clientId) + } + + private actor Storage { + private var entries: [HostId: String] = [:] + + func load(_ hostId: HostId) -> String? { entries[hostId] } + + func store(_ hostId: HostId, clientId: String) { + entries[hostId] = clientId + } + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostClientHandle.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostClientHandle.swift new file mode 100644 index 000000000..83a5e3da4 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostClientHandle.swift @@ -0,0 +1,84 @@ +// HostClientHandle — generation-checked escape hatch onto the underlying +// `AHPClient` for a host. + +import Foundation +import AgentHostProtocol + +/// Generation-checked handle to the underlying single-host `AHPClient`. +/// +/// Issued by `MultiHostClient.client(for:)`. Methods on this handle verify +/// that the host is still on the same `generation` it was when the handle +/// was minted; if a reconnect has occurred, dispatching returns +/// `HostError.hostReconnected` instead of silently writing to the new +/// connection. +/// +/// **Race note:** generation is checked once at the start of each call, so +/// it is possible (but rare) for a reconnect to land between +/// `checkAlive()` and the actual `dispatch`/`request`. In that race the +/// dispatch goes against the stale (now-shutdown) `AHPClient` and surfaces +/// as `AHPClientError.shutdown` wrapped in `HostError.client`. Acquire a +/// fresh handle when this happens. The Rust SDK has the same semantics. +public struct HostClientHandle: Sendable { + /// Host this handle was issued for. + public let hostId: HostId + /// Generation this handle was minted at. + public let generation: UInt64 + + /// The `AHPClient` instance that was current when the handle was minted. + /// May have been shut down by a subsequent reconnect. + private let client: AHPClient + private let shared: HostShared + + internal init(hostId: HostId, generation: UInt64, client: AHPClient, shared: HostShared) { + self.hostId = hostId + self.generation = generation + self.client = client + self.shared = shared + } + + /// Validate this handle against the host's current generation. Throws + /// `HostError.hostReconnected` if a reconnect has happened. + public func checkAlive() async throws { + let current = await shared.generation() + if current != generation { + throw HostError.hostReconnected( + host: hostId, + handleGeneration: generation, + currentGeneration: current + ) + } + } + + /// Dispatch an action through this connection, refusing if the connection + /// has been replaced by a reconnect. + @discardableResult + public func dispatch(_ action: StateAction) async throws -> DispatchHandle { + try await checkAlive() + do { + return try await client.dispatch(action) + } catch let error as AHPClientError { + throw HostError.client(error) + } + } + + /// Issue an arbitrary JSON-RPC request through this connection, refusing + /// if the connection has been replaced by a reconnect. + public func request( + method: String, + params: P + ) async throws -> R { + try await checkAlive() + do { + return try await client.request(method: method, params: params) + } catch let error as AHPClientError { + throw HostError.client(error) + } + } + + /// Borrow the underlying `AHPClient` for advanced use. The caller is + /// responsible for not holding it past the next reconnect — the returned + /// reference becomes a stale handle once the host reconnects. + public func rawClient() async -> AHPClient { + client + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostConfig.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostConfig.swift new file mode 100644 index 000000000..700d02882 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostConfig.swift @@ -0,0 +1,77 @@ +// HostConfig — configuration for a host registered with `MultiHostClient`. + +import Foundation + +/// Async factory that opens (or re-opens) a transport for a host. +/// +/// The supervisor calls this on every connect attempt — including reconnects +/// — so consumers can refresh tokens, rotate URLs, or pick different backends +/// per attempt. +public typealias HostTransportFactory = @Sendable (HostId) async throws -> any AHPTransport + +/// Configuration for a single host registered with `MultiHostClient`. +/// +/// Use `HostConfig(id:label:transportFactory:)` for the common case and +/// the `with*` builders to override individual fields. +public struct HostConfig: Sendable { + /// Stable host identifier. Doubles as the `ClientIdStore` persistence key. + public var id: HostId + /// Human-readable label. Surfaced through `HostHandle.label`. + public var label: String + /// Optional override for the `clientId` sent to this host. When `nil`, the + /// multi-host client asks its `ClientIdStore` for a stable id keyed on + /// `id`. + public var clientId: String? + /// URIs to include in the `initialize` handshake. Defaults to + /// `[RootResourceURI]` so root state is always tracked. + public var initialSubscriptions: [String] + /// Configuration forwarded to the underlying `AHPClient`. + public var clientConfig: AHPClientConfig + /// Factory used to (re-)open a transport for this host. + public var transportFactory: HostTransportFactory + /// Reconnect behaviour after an unexpected drop. + public var reconnectPolicy: ReconnectPolicy + + public init( + id: HostId, + label: String, + transportFactory: @escaping HostTransportFactory + ) { + self.id = id + self.label = label + self.clientId = nil + self.initialSubscriptions = [RootResourceURI] + self.clientConfig = .default + self.transportFactory = transportFactory + self.reconnectPolicy = .exponential + } + + /// Override the explicit `clientId` for this host (skips the + /// `ClientIdStore` lookup). + public func withClientId(_ clientId: String) -> Self { + var copy = self + copy.clientId = clientId + return copy + } + + /// Replace the default `initialSubscriptions` set. + public func withInitialSubscriptions(_ uris: [String]) -> Self { + var copy = self + copy.initialSubscriptions = uris + return copy + } + + /// Override the per-host `AHPClientConfig`. + public func withClientConfig(_ config: AHPClientConfig) -> Self { + var copy = self + copy.clientConfig = config + return copy + } + + /// Override the reconnect policy. + public func withReconnectPolicy(_ policy: ReconnectPolicy) -> Self { + var copy = self + copy.reconnectPolicy = policy + return copy + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostError.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostError.swift new file mode 100644 index 000000000..a2ff9ccbc --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostError.swift @@ -0,0 +1,44 @@ +// HostError — errors specific to the multi-host SDK layer. + +import Foundation + +/// Errors specific to the multi-host SDK layer. +/// +/// Errors from the underlying single-host `AHPClient` are carried through +/// `client(AHPClientError)`. +public enum HostError: Error, Sendable { + /// No host with this id is currently registered. + case unknownHost(HostId) + + /// The `HostClientHandle` was issued for a connection that has since + /// been replaced by a reconnect. Acquire a fresh handle via + /// `MultiHostClient.client(for:)`. + case hostReconnected(host: HostId, handleGeneration: UInt64, currentGeneration: UInt64) + + /// The host's runtime task has been torn down (e.g. the host was removed + /// or the multi-host client was shut down). + case hostShutDown(HostId) + + /// `MultiHostClient.add` was called with an id that is already registered. + case duplicateHost(HostId) + + /// A request bubbled up an error from the underlying `AHPClient`. + case client(AHPClientError) +} + +extension HostError: LocalizedError { + public var errorDescription: String? { + switch self { + case .unknownHost(let id): + return "no host registered with id \(id)" + case .hostReconnected(let host, let handleGeneration, let currentGeneration): + return "host \(host) reconnected (generation \(handleGeneration) -> \(currentGeneration)); request a fresh client handle" + case .hostShutDown(let id): + return "host \(id) runtime is no longer active" + case .duplicateHost(let id): + return "host \(id) is already registered; remove it first" + case .client(let error): + return error.errorDescription + } + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostEvents.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostEvents.swift new file mode 100644 index 000000000..beafbec77 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostEvents.swift @@ -0,0 +1,67 @@ +// HostEvents — fan-in event types and aggregated view types for `MultiHostClient`. + +import Foundation +import AgentHostProtocol + +/// Inbound subscription event tagged with the host that produced it. +/// +/// Delivered by `MultiHostClient.events()`. `resource` carries the URI the +/// event is scoped to (typically derived from the underlying action's +/// `session`/`terminal` field). Protocol-level notifications (session +/// added/removed/changed, auth required) carry `resource: nil` because they +/// aren't bound to a single resource. +public struct HostSubscriptionEvent: Sendable { + public let hostId: HostId + public let resource: String? + public let event: SubscriptionEvent + + public init(hostId: HostId, resource: String?, event: SubscriptionEvent) { + self.hostId = hostId + self.resource = resource + self.event = event + } +} + +/// Connection-level event for UX, delivered by `MultiHostClient.hostEvents()`. +public enum HostEvent: Sendable { + /// A new host was registered with `MultiHostClient.add(_:)`. + case added(HostId) + /// The host's `HostState` changed. + case stateChanged(HostId, HostState, lastError: String?) + /// The host successfully (re)connected; `generation` is the new value. + case connected(HostId, generation: UInt64) + /// A host was removed from `MultiHostClient`. + case removed(HostId) +} + +/// Aggregated session summary tagged with host of origin. +/// +/// Returned by `MultiHostClient.aggregatedSessions()`. URIs are per-host +/// scoped, so two hosts can legitimately advertise the same `summary.resource`; +/// consumers should treat `(hostId, summary.resource)` as the compound key. +public struct HostedSessionSummary: Sendable { + public let hostId: HostId + public let hostLabel: String + public let summary: SessionSummary + + public init(hostId: HostId, hostLabel: String, summary: SessionSummary) { + self.hostId = hostId + self.hostLabel = hostLabel + self.summary = summary + } +} + +/// Aggregated agent descriptor tagged with host of origin. +/// +/// Returned by `MultiHostClient.aggregatedAgents()`. +public struct HostedAgent: Sendable { + public let hostId: HostId + public let hostLabel: String + public let agent: AgentInfo + + public init(hostId: HostId, hostLabel: String, agent: AgentInfo) { + self.hostId = hostId + self.hostLabel = hostLabel + self.agent = agent + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostHandle.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostHandle.swift new file mode 100644 index 000000000..969e41ccc --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostHandle.swift @@ -0,0 +1,92 @@ +// HostHandle — observable snapshot of everything `MultiHostClient` knows +// about a single host. + +import Foundation +import AgentHostProtocol + +/// Snapshot of everything the multi-host SDK knows about a single host. +/// +/// This is the value type UIs render: connection state, last error, protocol +/// version, agents pulled from root state, subscribed URIs, cached session +/// summaries, and so on. +/// +/// Snapshots are immutable; refresh by calling +/// `MultiHostClient.host(_:)`/`MultiHostClient.hosts()` again or by listening +/// to `MultiHostClient.hostEvents()`. +public struct HostHandle: Sendable { + /// Stable identifier. + public let id: HostId + /// Human-readable label from the original `HostConfig`. + public let label: String + /// `clientId` actually sent to the host on `initialize`/`reconnect`. + public let clientId: String + /// Current connection state. + public let state: HostState + /// Most recent error message, set when the supervisor enters + /// `.reconnecting` or `.failed`. Cleared on a successful connect. + public let lastError: String? + /// Wall-clock time of the most recent successful `initialize` or + /// `reconnect`. `nil` until the host first connects. + public let lastConnectedAt: Date? + /// Protocol version negotiated with the host on the most recent + /// successful `initialize`. + public let protocolVersion: String? + /// Highest `serverSeq` observed on this host. + public let serverSeq: Int + /// Optional `defaultDirectory` from the host's `InitializeResult`. + public let defaultDirectory: String? + /// Agents currently advertised by the host (mirrored from root state). + public let agents: [AgentInfo] + /// Active session count from root state, when present. + public let activeSessions: Int? + /// Lightweight terminal listing from root state, when present. + public let terminals: [TerminalInfo]? + /// URIs the supervisor will (re-)subscribe to across reconnects. + public let subscriptions: [String] + /// Trigger characters from `InitializeResult.completionTriggerCharacters`. + public let completionTriggerCharacters: [String] + /// Cached session summaries, sorted by `modifiedAt` descending. Seeded by + /// `listSessions` after each connect and kept fresh by + /// `notify/sessionAdded`/`notify/sessionRemoved`/`notify/sessionSummaryChanged`. + public let sessionSummaries: [SessionSummary] + /// Generation counter — bumped on every `connect` or `reconnect`. + /// `HostClientHandle`s carry the generation they were issued at and + /// refuse to dispatch through a stale connection. + public let generation: UInt64 + + public init( + id: HostId, + label: String, + clientId: String, + state: HostState, + lastError: String?, + lastConnectedAt: Date?, + protocolVersion: String?, + serverSeq: Int, + defaultDirectory: String?, + agents: [AgentInfo], + activeSessions: Int?, + terminals: [TerminalInfo]?, + subscriptions: [String], + completionTriggerCharacters: [String], + sessionSummaries: [SessionSummary], + generation: UInt64 + ) { + self.id = id + self.label = label + self.clientId = clientId + self.state = state + self.lastError = lastError + self.lastConnectedAt = lastConnectedAt + self.protocolVersion = protocolVersion + self.serverSeq = serverSeq + self.defaultDirectory = defaultDirectory + self.agents = agents + self.activeSessions = activeSessions + self.terminals = terminals + self.subscriptions = subscriptions + self.completionTriggerCharacters = completionTriggerCharacters + self.sessionSummaries = sessionSummaries + self.generation = generation + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostId.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostId.swift new file mode 100644 index 000000000..6e019d607 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostId.swift @@ -0,0 +1,24 @@ +// HostId — stable, opaque identifier for a host registered with `MultiHostClient`. + +import Foundation + +/// Stable identifier for a host registered with `MultiHostClient`. +/// +/// Opaque to the SDK — consumers pick the format. It's used as the persistence +/// key for `ClientIdStore`, the routing key for commands on `MultiHostClient`, +/// and the tag on every `HostSubscriptionEvent`. +public struct HostId: Hashable, Sendable, CustomStringConvertible { + public let value: String + + public init(_ value: String) { + self.value = value + } + + public var description: String { value } +} + +extension HostId: ExpressibleByStringLiteral { + public init(stringLiteral value: StringLiteralType) { + self.value = value + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostRuntime.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostRuntime.swift new file mode 100644 index 000000000..fcdb114ae --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostRuntime.swift @@ -0,0 +1,800 @@ +// HostRuntime — per-host supervisor. +// +// Owns the current `AHPClient`, the reconnect state machine, the per-host +// root-state mirror, and the session-summary cache. Receives commands over +// a single internal `AsyncStream` and forwards inbound events +// to the multi-host fan-in. + +import Foundation +import AgentHostProtocol + +/// Per-host supervisor task. Internal to the multi-host SDK. +/// +/// Constructed by `MultiHostClient.add(_:)`. The `MultiHostClient` enqueues +/// commands via the public methods (`subscribe`, `dispatch`, …) which all +/// reduce to sending a `HostCommand` into the supervisor's queue. +internal final class HostRuntime: Sendable { + /// Shared, generation-checked state. Public to the multi-host layer so + /// `HostClientHandle.checkAlive()` can read generations without going + /// through the supervisor command queue. + let shared: HostShared + + private let config: HostConfig + private let clientId: String + + /// Sink the multi-host facade fans events into. Awaited from the pump + /// so per-host event ordering is preserved (a fresh per-event Task + /// would interleave on the way to the consumer actor). + private let fanOut: @Sendable (HostSubscriptionEvent) async -> Void + /// Sink for connection-level events. Awaited for the same reason. + private let hostEventSink: @Sendable (HostEvent) async -> Void + + private let cmdContinuation: AsyncStream.Continuation + private let cmdStream: AsyncStream + + /// Long-running supervisor task. Captured so `shutdown()` can await it. + nonisolated(unsafe) private var supervisorTask: Task? + + /// Monotonic counter that mints a fresh token for each pump task and + /// each backoff sleep. Stale `.connectionEnded`/`.backoffElapsed` + /// signals from a previous cycle are filtered out by token mismatch. + private let signalTokenSource = SignalTokenSource() + + init( + config: HostConfig, + clientIdStore: ClientIdStore, + fanOut: @escaping @Sendable (HostSubscriptionEvent) async -> Void, + hostEventSink: @escaping @Sendable (HostEvent) async -> Void + ) async { + self.config = config + self.fanOut = fanOut + self.hostEventSink = hostEventSink + + let resolved: String + if let explicit = config.clientId { + resolved = explicit + } else if let stored = await clientIdStore.load(config.id) { + resolved = stored + } else { + resolved = generateClientId() + } + await clientIdStore.store(config.id, clientId: resolved) + self.clientId = resolved + + let initial = HostInternal( + id: config.id, + label: config.label, + clientId: resolved, + state: .disconnected, + lastError: nil, + lastConnectedAt: nil, + protocolVersion: nil, + serverSeq: 0, + defaultDirectory: nil, + rootState: RootState(agents: []), + subscriptions: config.initialSubscriptions, + completionTriggerCharacters: [], + sessionSummaries: [:], + generation: 0, + currentClient: nil + ) + self.shared = HostShared(initial) + + var cont: AsyncStream.Continuation! + let stream = AsyncStream(bufferingPolicy: .unbounded) { c in + cont = c + } + self.cmdContinuation = cont + self.cmdStream = stream + } + + /// Start the supervisor task. Call exactly once after `init`. + func start() { + let task = Task { [self] in + await self.run() + } + self.supervisorTask = task + } + + // MARK: - Public command surface (called by `MultiHostClient`) + + /// Snapshot the current `HostHandle` directly from `HostShared`. Bypasses + /// the command queue — `HostShared` is its own actor so this is safe and + /// won't deadlock when the supervisor is mid-await on transport I/O. + func snapshot() async -> HostHandle { + await shared.snapshot() + } + + /// Acquire a generation-checked client handle, when connected. + func clientHandle() async -> HostClientHandle? { + let state = await shared.internalState + guard let client = state.currentClient else { return nil } + return HostClientHandle( + hostId: state.id, + generation: state.generation, + client: client, + shared: shared + ) + } + + /// Send a manual reconnect signal. Returns once the supervisor has + /// observed the request (it cancels any pending backoff sleep). + func reconnect() async throws { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + cmdContinuation.yield(.manualReconnect(reply: continuation)) + } + } + + /// Subscribe to `uri` on the current connection. Tracks the URI so it + /// is replayed across reconnects. Returns `HostError.hostShutDown` if + /// the host is disconnected. + func subscribe(_ uri: String) async throws -> SubscribeResult { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + cmdContinuation.yield(.subscribe(uri: uri, reply: continuation)) + } + } + + /// Unsubscribe from `uri`. Stops replay of `uri` across reconnects. Safe + /// to call when disconnected — drops the URI from the replay set. + func unsubscribe(_ uri: String) async throws { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + cmdContinuation.yield(.unsubscribe(uri: uri, reply: continuation)) + } + } + + /// Dispatch an action through the current connection. Throws + /// `HostError.hostShutDown` if the host is disconnected. + @discardableResult + func dispatch(_ action: StateAction) async throws -> DispatchHandle { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + cmdContinuation.yield(.dispatch(action: action, reply: continuation)) + } + } + + /// Tear down the supervisor: cancel any in-flight connect/sleep, close + /// the current `AHPClient` if any, and finish the supervisor task. + func shutdown() async { + cmdContinuation.yield(.shutdown) + cmdContinuation.finish() + await supervisorTask?.value + } + + // MARK: - Supervisor loop + + private func run() async { + await hostEventSink(.added(config.id)) + + var attempt = 0 + var iter = cmdStream.makeAsyncIterator() + + // Initial state announce — the host is about to make its first + // connect attempt. Subsequent transitions to `.reconnecting` + // happen at the moment the prior attempt's connection ends, so + // `HostHandle.state` accurately reflects "we are no longer + // connected; backoff is in flight" rather than continuing to + // report `.connected` through the entire backoff sleep. + await transition(to: .connecting, error: nil) + + outer: while true { + attempt += 1 + + // Try to connect. On success we run the connection until it ends. + do { + let streams = try await connectOnce() + if config.reconnectPolicy.resetOnSuccess { + attempt = 0 + } + let outcome = await runConnection(streams: streams, iter: &iter) + await tearDownClient() + switch outcome { + case .shutdown: + return + case .manualReconnect(let reply): + reply.resume() + attempt = 0 + await transition(to: .connecting, error: nil) + continue outer + case .disconnected: + // Connection dropped. Surface that immediately so + // `HostHandle.state` doesn't keep reporting `.connected` + // through the backoff sleep. `attempt` here is either + // the count of consecutive failures since last success + // (no reset) or 0 (just reset by `resetOnSuccess`); the + // displayed attempt is clamped to ≥ 1 per the + // 1-based contract on `HostState.reconnecting`. + let lastErr = await shared.lastError() + await transition( + to: .reconnecting(attempt: max(1, attempt)), + error: lastErr + ) + } + } catch { + let reason = String(describing: error) + await shared.update { $0.lastError = reason } + // Failed connect attempt. `attempt` was just bumped at the + // top of this iteration so it is ≥ 1. + await transition(to: .reconnecting(attempt: attempt), error: reason) + } + + if config.reconnectPolicy.attemptsExhausted(attempt) { + let reason = (await shared.lastError()) ?? "reconnect attempts exhausted" + await transition(to: .failed(reason: reason), error: reason) + let outcome = await waitForManualReconnectOrShutdown(iter: &iter) + switch outcome { + case .shutdown: return + case .manualReconnect(let reply): + reply.resume() + attempt = 0 + await transition(to: .connecting, error: nil) + continue outer + case .disconnected: + return + } + } + + let delay = config.reconnectPolicy.delay(forAttempt: attempt, sample: jitterSample()) + let outcome = await sleepOrCommand(delay: delay, iter: &iter) + switch outcome { + case .shutdown: + return + case .manualReconnect(let reply): + reply.resume() + attempt = 0 + await transition(to: .connecting, error: nil) + continue outer + case .disconnected: + // Backoff elapsed; `state` is already `.reconnecting(attempt:)` + // from the disconnect / failure branch above. The next + // iteration runs the actual connect. + continue outer + } + } + } + + /// Open a transport, hand it to a fresh `AHPClient`, attach the events + /// tap *before* the handshake, and complete `initialize`/`reconnect` + /// plus an opportunistic `listSessions`. Returns both the `client.events` + /// stream and the `client.stateChanges` stream so the caller can drive + /// the per-connection event pump *and* a drop detector. + /// + /// Any failure after `client.connect()` has started the writer/receive + /// tasks is funneled through `withClientShutdownOnThrow` so the + /// orphaned `AHPClient` is torn down before the error propagates back + /// to the supervisor (which would otherwise open a fresh transport on + /// the next attempt while the previous client's tasks remain alive). + private func connectOnce() async throws -> ConnectionStreams { + let transport = try await config.transportFactory(config.id) + let client = AHPClient(transport: transport, config: config.clientConfig) + + // Attach the events and state-change taps BEFORE the + // initialize/reconnect handshake so any notifications the server + // pushes between the response and the moment we enter the run + // loop are captured rather than dropped. PR 2's + // `events_tap_captures_handshake_notifications` test exists to + // protect this contract for `events`. We also need + // `stateChanges` because `AHPClient.handleTransportFailure` + // intentionally keeps the events stream alive after a transport + // drop (so consumers can observe later state transitions); the + // only signal we get on a real drop is `connectionState` flipping + // to `.disconnected`. + let events = await client.events + let stateChanges = await client.stateChanges + + try await client.connect() + + // From here on, every `throw` must shut down `client` before + // propagating — `client.connect()` started writer/receive tasks + // that hold the transport. `withClientShutdownOnThrow` enforces it. + return try await withClientShutdownOnThrow(client) { + try await self.completeHandshake(client: client, events: events, stateChanges: stateChanges) + } + } + + /// Run the `initialize`/`reconnect` handshake, the opportunistic + /// `listSessions` seed, and atomically install the new client in + /// `HostShared`. Called inside `withClientShutdownOnThrow` so a throw + /// from any of these steps tears the client down before bubbling up. + private func completeHandshake( + client: AHPClient, + events: AsyncStream, + stateChanges: AsyncStream + ) async throws -> ConnectionStreams { + // Decide between initialize and reconnect based on prior state. + let priorSnapshot = await shared.internalState + let canReconnect = priorSnapshot.serverSeq > 0 && !priorSnapshot.subscriptions.isEmpty + let priorSubscriptions = priorSnapshot.subscriptions + let priorSeq = priorSnapshot.serverSeq + + var initResult: InitializeResult? = nil + var newSeq = priorSeq + + if canReconnect { + do { + _ = try await client.reconnect( + clientId: clientId, + lastSeenServerSeq: priorSeq, + subscriptions: priorSubscriptions + ) + // NOTE: the `ReconnectResult` is intentionally discarded + // here, mirroring the Rust `ahp::hosts` runtime. Three + // related gaps follow from this and are tracked together + // for both SDKs as a follow-up: + // 1. Replay actions returned synchronously by the server + // are not fanned out (live `notify/action` frames after + // reconnect still reach consumers normally). + // 2. `replay.missing` URIs (subscriptions the server + // cannot resume) are not pruned from the replay set, + // so the next reconnect re-asks for them. + // 3. `snapshot` results are not applied to the per-host + // root mirror / `serverSeq`, so `HostHandle` can lag + // behind the post-snapshot state until live events + // catch up. + // All three should be fixed atomically across SDKs — see the + // parent multi-host series for tracking. + } catch let error as AHPClientError { + if case .rpc = error { + let init1 = try await client.initialize( + clientId: clientId, + protocolVersions: [supportedProtocolVersion], + initialSubscriptions: priorSubscriptions + ) + initResult = init1 + newSeq = init1.serverSeq + } else { + throw error + } + } + } else { + let init1 = try await client.initialize( + clientId: clientId, + protocolVersions: [supportedProtocolVersion], + initialSubscriptions: priorSubscriptions + ) + initResult = init1 + newSeq = init1.serverSeq + } + + // Refresh session summaries from `listSessions`. Cheap on first + // connect; kept in sync by notifications afterward. Failures are + // non-fatal: the cache stays as-is. + let summaries: ListSessionsResult? = try? await client.request( + method: "listSessions", + params: ListSessionsParams() + ) + + let newGeneration: UInt64 = await { + var generation: UInt64 = 0 + await shared.update { state in + state.generation = state.generation &+ 1 + state.currentClient = client + state.lastConnectedAt = Date() + state.lastError = nil + state.serverSeq = newSeq + if let init1 = initResult { + state.protocolVersion = init1.protocolVersion + state.defaultDirectory = init1.defaultDirectory + state.completionTriggerCharacters = init1.completionTriggerCharacters ?? [] + if let snap = init1.snapshots.first(where: { $0.resource == RootResourceURI }) { + if case .root(let root) = snap.state { + state.rootState = root + } + } + } + if let list = summaries { + state.sessionSummaries.removeAll() + for summary in list.items { + state.sessionSummaries[summary.resource] = summary + } + } + generation = state.generation + } + return generation + }() + + await transition(to: .connected, error: nil) + await hostEventSink(.connected(config.id, generation: newGeneration)) + return ConnectionStreams(events: events, stateChanges: stateChanges) + } + + /// Drain commands and the event pump until the connection ends, the user + /// asks for a manual reconnect, or shutdown is requested. + private func runConnection( + streams: ConnectionStreams, + iter: inout AsyncStream.AsyncIterator + ) async -> RunOutcome { + // Per-connection token so a stale `.connectionEnded` from a prior + // pump task can't trick this drain loop into thinking the new + // connection has already failed. + let connectionToken = signalTokenSource.next() + let events = streams.events + let stateChanges = streams.stateChanges + let pumpTask = Task { [weak self, cmdContinuation] in + guard let self else { return } + for await event in events { + await self.handleEvent(event) + } + cmdContinuation.yield(.connectionEnded(token: connectionToken)) + } + // Drop detector. `AHPClient.handleTransportFailure` does not finish + // the events stream after a transport drop (it intentionally keeps + // the multicast taps alive so consumers can observe later state + // transitions), so the only signal we get on a real drop is + // `connectionState` flipping to `.disconnected`. This task converts + // that into a `.connectionEnded(token:)` sentinel. + let dropDetector = Task { [cmdContinuation] in + for await state in stateChanges { + if case .disconnected = state { + cmdContinuation.yield(.connectionEnded(token: connectionToken)) + return + } + } + } + + defer { + pumpTask.cancel() + dropDetector.cancel() + } + + while let cmd = await iter.next() { + switch cmd { + case .shutdown: + return .shutdown + case .connectionEnded(let token): + if token == connectionToken { + return .disconnected + } + continue + case .backoffElapsed: + continue + case .manualReconnect(let reply): + return .manualReconnect(reply: reply) + case .subscribe(let uri, let reply): + let result = await handleSubscribe(uri) + resumeCommand(reply: reply, with: result) + case .unsubscribe(let uri, let reply): + let result = await handleUnsubscribe(uri) + resumeCommand(reply: reply, with: result) + case .dispatch(let action, let reply): + let result = await handleDispatch(action) + resumeCommand(reply: reply, with: result) + } + } + return .shutdown + } + + /// Wait for either a manual reconnect or shutdown while in the + /// `.failed` terminal state. Subscribe/unsubscribe still mutate the + /// replay set so the next reconnect picks them up. + private func waitForManualReconnectOrShutdown( + iter: inout AsyncStream.AsyncIterator + ) async -> RunOutcome { + while let cmd = await iter.next() { + switch cmd { + case .shutdown: + return .shutdown + case .manualReconnect(let reply): + return .manualReconnect(reply: reply) + case .backoffElapsed, .connectionEnded: + continue + case .subscribe(let uri, let reply): + await shared.appendSubscription(uri) + reply.resume(throwing: HostError.hostShutDown(config.id)) + case .unsubscribe(let uri, let reply): + await shared.removeSubscription(uri) + reply.resume(returning: ()) + case .dispatch(_, let reply): + reply.resume(throwing: HostError.hostShutDown(config.id)) + } + } + return .shutdown + } + + /// Sleep for `delay` while still servicing snapshot-class commands. + /// Manual reconnect or shutdown short-circuits the sleep. + private func sleepOrCommand( + delay: Duration, + iter: inout AsyncStream.AsyncIterator + ) async -> RunOutcome { + if delay == .zero { + return .disconnected + } + let cont = cmdContinuation + let sleepToken = signalTokenSource.next() + let sleepTask = Task { + try? await Task.sleep(for: delay) + cont.yield(.backoffElapsed(token: sleepToken)) + } + defer { sleepTask.cancel() } + + while let cmd = await iter.next() { + switch cmd { + case .shutdown: + return .shutdown + case .manualReconnect(let reply): + return .manualReconnect(reply: reply) + case .backoffElapsed(let token): + if token == sleepToken { + return .disconnected + } + continue + case .connectionEnded: + continue + case .subscribe(let uri, let reply): + await shared.appendSubscription(uri) + reply.resume(throwing: HostError.hostShutDown(config.id)) + case .unsubscribe(let uri, let reply): + await shared.removeSubscription(uri) + reply.resume(returning: ()) + case .dispatch(_, let reply): + reply.resume(throwing: HostError.hostShutDown(config.id)) + } + } + return .shutdown + } + + // MARK: - Event handling + + private func handleEvent(_ event: ClientEvent) async { + // Mutate per-host mirrors before broadcasting so observers reading + // the next snapshot see the post-event state. + switch event.event { + case .action(let envelope): + await applyAction(envelope) + case .notification(let notification): + await applyNotification(notification) + } + let hostEvent = HostSubscriptionEvent( + hostId: config.id, + resource: event.resource, + event: event.event + ) + await fanOut(hostEvent) + } + + private func applyAction(_ envelope: ActionEnvelope) async { + await shared.update { state in + if envelope.serverSeq > state.serverSeq { + state.serverSeq = envelope.serverSeq + } + // Best-effort root state mirror update via the existing pure + // reducer. Non-root actions slip through without effect — that's + // the same posture as the Rust SDK. + let resource = actionResource(for: envelope.action) + if resource == RootResourceURI { + state.rootState = rootReducer(state: state.rootState, action: envelope.action) + } + } + } + + private func applyNotification(_ notification: ProtocolNotification) async { + await shared.update { state in + switch notification { + case .sessionAdded(let n): + state.sessionSummaries[n.summary.resource] = n.summary + case .sessionRemoved(let n): + state.sessionSummaries.removeValue(forKey: n.session) + case .sessionSummaryChanged(let n): + if var existing = state.sessionSummaries[n.session] { + applySummaryChanges(&existing, changes: n.changes) + state.sessionSummaries[n.session] = existing + } + case .authRequired: + break + } + } + } + + // MARK: - Active-connection command handlers + + private func handleSubscribe(_ uri: String) async -> Result { + guard let client = await shared.currentClient() else { + return .failure(.hostShutDown(config.id)) + } + do { + let (result, _) = try await client.subscribe(uri) + await shared.appendSubscription(uri) + return .success(result) + } catch let error as AHPClientError { + return .failure(.client(error)) + } catch { + return .failure(.client(.transport(.io(String(describing: error))))) + } + } + + private func handleUnsubscribe(_ uri: String) async -> Result { + let client = await shared.currentClient() + if let client { + do { + try await client.unsubscribe(uri) + } catch let error as AHPClientError { + return .failure(.client(error)) + } catch { + return .failure(.client(.transport(.io(String(describing: error))))) + } + } + await shared.removeSubscription(uri) + return .success(()) + } + + private func handleDispatch(_ action: StateAction) async -> Result { + guard let client = await shared.currentClient() else { + return .failure(.hostShutDown(config.id)) + } + do { + let handle = try await client.dispatch(action) + return .success(handle) + } catch let error as AHPClientError { + return .failure(.client(error)) + } catch { + return .failure(.client(.transport(.io(String(describing: error))))) + } + } + + // MARK: - State transitions and tear-down + + private func transition(to state: HostState, error: String?) async { + await shared.update { s in + s.state = state + if let error { + s.lastError = error + } + } + await hostEventSink(.stateChanged(config.id, state, lastError: error)) + } + + private func tearDownClient() async { + let prev: AHPClient? = await { + var captured: AHPClient? = nil + await shared.update { state in + captured = state.currentClient + state.currentClient = nil + } + return captured + }() + if let prev { + await prev.shutdown() + } + } +} + +// MARK: - Helpers + +/// Commands the supervisor consumes. All public API on `HostRuntime` reduces +/// to enqueueing one of these. +internal enum HostCommand: Sendable { + case shutdown + /// Sentinel from a connection's event-pump task signalling that the + /// underlying transport drained. The token identifies *which* + /// connection ended, so stale signals queued after `runConnection` + /// already returned (e.g. for a manual reconnect) are ignored by the + /// next runConnection cycle instead of being mistaken for an immediate + /// disconnect on the brand-new connection. + case connectionEnded(token: UInt64) + /// Sentinel from `sleepOrCommand`'s sleep task signalling that the + /// backoff delay elapsed. Tagged with a token for the same reason as + /// `connectionEnded`. + case backoffElapsed(token: UInt64) + case manualReconnect(reply: CheckedContinuation) + case subscribe(uri: String, reply: CheckedContinuation) + case unsubscribe(uri: String, reply: CheckedContinuation) + case dispatch(action: StateAction, reply: CheckedContinuation) +} + +/// Outcome of one of the supervisor's drain loops. +private enum RunOutcome { + case shutdown + case disconnected + case manualReconnect(reply: CheckedContinuation) +} + +/// Bundle of multicast streams attached to the per-connection `AHPClient`. +/// Returned by `connectOnce` and consumed by `runConnection`. +internal struct ConnectionStreams: @unchecked Sendable { + let events: AsyncStream + let stateChanges: AsyncStream +} + +/// Resume a `CheckedContinuation` with a `Result`. +private func resumeCommand( + reply: CheckedContinuation, + with result: Result +) { + switch result { + case .success(let value): reply.resume(returning: value) + case .failure(let error): reply.resume(throwing: error) + } +} + +/// Run `body`; if it throws, shut down `client` before rethrowing. Used to +/// keep handshake failures (`initialize`, `reconnect`, `listSessions`) from +/// leaking an `AHPClient` whose writer/receive tasks are already running — +/// without it, a failed connect attempt would leave the previous client's +/// tasks alive while the supervisor opens a fresh transport for the next +/// retry, holding the original transport indefinitely. +private func withClientShutdownOnThrow( + _ client: AHPClient, + _ body: () async throws -> T +) async throws -> T { + do { + return try await body() + } catch { + await client.shutdown() + throw error + } +} + +/// Mirror Rust: SipHash + atomic counter is enough randomness for jitter. +private func jitterSample() -> Double { + var hasher = Hasher() + hasher.combine(jitterCounter.bumpAndGet()) + hasher.combine(Date().timeIntervalSince1970.bitPattern) + let bits = UInt64(bitPattern: Int64(hasher.finalize())) + return Double(bits) / Double(UInt64.max) +} + +private final class JitterCounter: @unchecked Sendable { + private let lock = NSLock() + private var n: UInt64 = 0 + func bumpAndGet() -> UInt64 { + lock.lock(); defer { lock.unlock() } + n &+= 1 + return n + } +} + +private let jitterCounter = JitterCounter() + +/// Thread-safe monotonic counter producing per-connection / per-sleep tokens +/// so stale signals from prior pump tasks can't leak into a fresh +/// `runConnection`/`sleepOrCommand` cycle. +internal final class SignalTokenSource: @unchecked Sendable { + private let lock = NSLock() + private var counter: UInt64 = 0 + + func next() -> UInt64 { + lock.lock(); defer { lock.unlock() } + counter &+= 1 + return counter + } +} + +/// Generate a UUIDv4-shaped client id without taking a UUID dependency from +/// other targets. Foundation's `UUID()` is fine here. +private func generateClientId() -> String { + UUID().uuidString.lowercased() +} + +/// Mirror the resource-routing logic in `AHPClient.actionResource(for:)`. +private func actionResource(for action: StateAction) -> String? { + let encoder = JSONEncoder() + guard let data = try? encoder.encode(action), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return RootResourceURI + } + if let session = object["session"] as? String { return session } + if let terminal = object["terminal"] as? String { return terminal } + return RootResourceURI +} + +/// Apply a `PartialSessionSummary` patch in-place. Identity fields are +/// ignored per spec. +private func applySummaryChanges( + _ existing: inout SessionSummary, + changes: PartialSessionSummary +) { + if let v = changes.title { existing.title = v } + if let v = changes.status { existing.status = v } + if let v = changes.activity { existing.activity = v } + if let v = changes.modifiedAt { existing.modifiedAt = v } + if let v = changes.project { existing.project = v } + if let v = changes.model { existing.model = v } + if let v = changes.workingDirectory { existing.workingDirectory = v } + if let v = changes.diffs { existing.diffs = v } +} + +/// Protocol version offered on `initialize`. Mirrors the Rust SDK's use of +/// the canonical `PROTOCOL_VERSION` constant; the Swift types library +/// doesn't ship one yet, so this is a constant string co-located with the +/// rest of the multi-host code. TODO(codegen): source from generated types. +private let supportedProtocolVersion = "0.1.0" diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostShared.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostShared.swift new file mode 100644 index 000000000..d99c73d1c --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostShared.swift @@ -0,0 +1,109 @@ +// HostShared — internal mutable per-host state, shared between `HostRuntime` +// and `HostClientHandle`s. +// +// `HostShared` is a small actor wrapping `HostInternal`. The runtime mutates +// this state under actor isolation; `HostClientHandle` reads it to validate +// its generation and to fetch the underlying `AHPClient` reference. + +import Foundation +import AgentHostProtocol + +/// Internal mutable per-host state. Updated by the runtime task; read on the +/// snapshot path to build `HostHandle`s and by `HostClientHandle.checkAlive()` +/// to validate generation tokens. +internal struct HostInternal { + var id: HostId + var label: String + var clientId: String + var state: HostState + var lastError: String? + var lastConnectedAt: Date? + var protocolVersion: String? + var serverSeq: Int + var defaultDirectory: String? + var rootState: RootState + var subscriptions: [String] + var completionTriggerCharacters: [String] + /// Session summaries keyed by URI. Sorted on snapshot. + var sessionSummaries: [String: SessionSummary] + var generation: UInt64 + /// The currently-installed `AHPClient`, when connected. `nil` between + /// connections. + var currentClient: AHPClient? + + func snapshot() -> HostHandle { + let summaries = sessionSummaries.values + .sorted { $0.modifiedAt > $1.modifiedAt } + return HostHandle( + id: id, + label: label, + clientId: clientId, + state: state, + lastError: lastError, + lastConnectedAt: lastConnectedAt, + protocolVersion: protocolVersion, + serverSeq: serverSeq, + defaultDirectory: defaultDirectory, + agents: rootState.agents, + activeSessions: rootState.activeSessions, + terminals: rootState.terminals, + subscriptions: subscriptions, + completionTriggerCharacters: completionTriggerCharacters, + sessionSummaries: summaries, + generation: generation + ) + } +} + +/// Actor-protected wrapper around `HostInternal`. Designed to be cheap to +/// poke from outside the runtime (e.g. for `HostClientHandle.checkAlive()`) +/// without contending against the supervisor's I/O. +internal actor HostShared { + private(set) var internalState: HostInternal + + init(_ initial: HostInternal) { + self.internalState = initial + } + + /// Take an immutable snapshot. + func snapshot() -> HostHandle { + internalState.snapshot() + } + + /// Read just the generation, for `HostClientHandle.checkAlive()`. + func generation() -> UInt64 { + internalState.generation + } + + /// Borrow the current `AHPClient`, when connected. + func currentClient() -> AHPClient? { + internalState.currentClient + } + + /// Apply an arbitrary mutation under actor isolation. + func update(_ body: (inout HostInternal) -> Void) { + body(&internalState) + } + + /// Convenience: append a subscription URI if not already present. + func appendSubscription(_ uri: String) { + if !internalState.subscriptions.contains(uri) { + internalState.subscriptions.append(uri) + } + } + + /// Convenience: remove a subscription URI. + func removeSubscription(_ uri: String) { + internalState.subscriptions.removeAll { $0 == uri } + } + + /// Convenience: read the last error string. + func lastError() -> String? { + internalState.lastError + } + + /// Convenience: read the host id and label. + func identity() -> (HostId, String) { + (internalState.id, internalState.label) + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostState.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostState.swift new file mode 100644 index 000000000..22d6069c8 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostState.swift @@ -0,0 +1,35 @@ +// HostState — connection state for a single host. + +import Foundation + +/// Connection state for a single host. +public enum HostState: Sendable, Equatable { + /// The host has been added but no transport is open. + case disconnected + /// A transport is being opened or the `initialize` handshake is in flight. + case connecting + /// The host is fully connected and serving subscriptions. + case connected + /// A previous connection dropped; the supervisor is retrying with backoff. + /// + /// `attempt` is one-based and resets after a successful connect when the + /// host's `ReconnectPolicy.resetOnSuccess` is `true`. + case reconnecting(attempt: Int) + /// Reconnect attempts were exhausted (or `ReconnectPolicy.disabled` was + /// configured) and the host is no longer trying. The supervisor still + /// services `snapshot`, manual `reconnect`, and `shutdown` commands while + /// in this state. + case failed(reason: String) + + /// Convenience: is the host currently `.connected`? + public var isConnected: Bool { + if case .connected = self { return true } + return false + } + + /// Convenience: is the host in a terminal failure state? + public var isFailed: Bool { + if case .failed = self { return true } + return false + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/MultiHostClient.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/MultiHostClient.swift new file mode 100644 index 000000000..14045b46d --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/MultiHostClient.swift @@ -0,0 +1,343 @@ +// MultiHostClient — public actor facade that owns multiple `HostRuntime`s and +// fans inbound events into multicast streams. + +import Foundation +import AgentHostProtocol + +/// Default buffer size for fan-in event streams. Slow consumers see oldest +/// items dropped (matching `AHPClient.events` semantics) rather than blocking +/// the publisher. +private let defaultFanInBuffer: Int = 1024 + +/// Multi-host client. +/// +/// Wraps N independent `AHPClient`s behind a single facade with per-host +/// supervisor tasks, generation-checked client handles, multicast event +/// streams, and a session-summary cache per host. Single-host consumers can +/// use `MultiHostClient.single(_:)` to skip the registry-style API entirely. +/// +/// `MultiHostClient` is an `actor` and runs off the main thread. UI threading +/// is the consumer's concern — wrap it in your `@MainActor` `@Observable` +/// store to bind into SwiftUI. +/// +/// **Lifecycle:** call `shutdown()` when you're done. Per-host runtimes spawn +/// long-running supervisor tasks that don't stop on `deinit`. +public actor MultiHostClient { + + private let clientIdStore: ClientIdStore + private var hosts: [HostId: HostRuntime] = [:] + /// Ids that are mid-add. Reserved synchronously so two concurrent + /// `add(_:)` calls for the same id can't both pass the duplicate check + /// across the `await HostRuntime(...)` suspension. + private var pendingHostIds: Set = [] + /// Insertion order of hosts (resolved on `add`). Used for deterministic + /// secondary tie-breaking in aggregated views. + private var hostOrder: [HostId] = [] + private var didShutDown: Bool = false + + // Multicast bookkeeping. Each `events()` / `hostEvents()` call gets a + // fresh `AsyncStream`; we register its continuation on the actor and + // remove it on stream termination. + private var nextListenerId: UInt64 = 1 + private var subscriptionListeners: [UInt64: AsyncStream.Continuation] = [:] + private var hostEventListeners: [UInt64: AsyncStream.Continuation] = [:] + + public init(clientIdStore: ClientIdStore = InMemoryClientIdStore()) { + self.clientIdStore = clientIdStore + } + + /// Convenience constructor for single-host consumers. + /// + /// Builds an empty `MultiHostClient`, registers `config`, and returns + /// the resulting `HostHandle` snapshot taken immediately after the + /// supervisor task is spawned. The snapshot may be `.disconnected`, + /// `.connecting`, `.connected`, `.reconnecting`, or `.failed` depending + /// on how far the first connect attempt has progressed; consumers that + /// need to wait for `.connected` should subscribe to `hostEvents()` or + /// poll `host(_:)`. + public static func single( + _ config: HostConfig, + clientIdStore: ClientIdStore = InMemoryClientIdStore() + ) async throws -> (MultiHostClient, HostHandle) { + let multi = MultiHostClient(clientIdStore: clientIdStore) + let handle = try await multi.add(config) + return (multi, handle) + } + + // MARK: - Host registry + + /// Register a new host and start its supervisor. Throws + /// `HostError.duplicateHost` if `config.id` is already registered (or + /// is mid-`add` from a concurrent caller). Throws + /// `HostError.hostShutDown` if `MultiHostClient.shutdown()` has been + /// called. + @discardableResult + public func add(_ config: HostConfig) async throws -> HostHandle { + let id = config.id + if didShutDown { + throw HostError.hostShutDown(id) + } + if hosts[id] != nil || pendingHostIds.contains(id) { + throw HostError.duplicateHost(id) + } + // Reserve the id synchronously so a concurrent `add(_:)` for the + // same id can't slip past the duplicate check while we await on + // `HostRuntime.init`'s `clientIdStore` lookups. + pendingHostIds.insert(id) + + // Capture isolated callbacks that publish back into this actor. + // Sinks are `async` and awaited from the runtime so per-host event + // ordering is preserved (a fresh per-event Task would be racy + // because each Task's hop into the actor can interleave). + let fanOut: @Sendable (HostSubscriptionEvent) async -> Void = { [weak self] event in + guard let self else { return } + await self.broadcastSubscriptionEvent(event) + } + let hostEventSink: @Sendable (HostEvent) async -> Void = { [weak self] event in + guard let self else { return } + await self.broadcastHostEvent(event) + } + + let runtime = await HostRuntime( + config: config, + clientIdStore: clientIdStore, + fanOut: fanOut, + hostEventSink: hostEventSink + ) + + // We may have been shut down while awaiting the runtime init; bail + // before exposing the new supervisor. + if didShutDown { + pendingHostIds.remove(id) + await runtime.shutdown() + throw HostError.hostShutDown(id) + } + + hosts[id] = runtime + hostOrder.append(id) + pendingHostIds.remove(id) + runtime.start() + return await runtime.snapshot() + } + + /// Remove a host, cancelling its supervisor task and dropping its current + /// connection. Outstanding `HostClientHandle`s for this host become stale + /// and surface `HostError.hostShutDown` (or `AHPClientError.shutdown` if + /// raced). + public func remove(_ id: HostId) async throws { + guard let runtime = hosts.removeValue(forKey: id) else { + throw HostError.unknownHost(id) + } + hostOrder.removeAll { $0 == id } + await runtime.shutdown() + broadcastHostEvent(.removed(id)) + } + + /// Trigger a manual reconnect. Cancels any in-flight backoff sleep and + /// jumps to the next connect attempt. Returns once the supervisor has + /// observed the request. + public func reconnect(_ id: HostId) async throws { + guard let runtime = hosts[id] else { + throw HostError.unknownHost(id) + } + try await runtime.reconnect() + } + + /// Snapshot the current state of `id`, or `nil` if no host is registered + /// under that id. + public func host(_ id: HostId) async -> HostHandle? { + guard let runtime = hosts[id] else { return nil } + return await runtime.snapshot() + } + + /// Snapshot every registered host. Order is unspecified. + public func hosts() async -> [HostHandle] { + var out: [HostHandle] = [] + out.reserveCapacity(hosts.count) + for runtime in hosts.values { + out.append(await runtime.snapshot()) + } + return out + } + + /// Acquire a generation-checked client handle for `id`. Returns `nil` if + /// the host is not registered or has no live connection. + public func client(for id: HostId) async -> HostClientHandle? { + guard let runtime = hosts[id] else { return nil } + return await runtime.clientHandle() + } + + // MARK: - Per-host convenience wrappers + + /// Subscribe to `uri` on `host`. Tracks the URI for replay across + /// reconnects. + @discardableResult + public func subscribe(host: HostId, uri: String) async throws -> SubscribeResult { + guard let runtime = hosts[host] else { + throw HostError.unknownHost(host) + } + return try await runtime.subscribe(uri) + } + + /// Unsubscribe from `uri` on `host`. Drops the URI from the replay set. + public func unsubscribe(host: HostId, uri: String) async throws { + guard let runtime = hosts[host] else { + throw HostError.unknownHost(host) + } + try await runtime.unsubscribe(uri) + } + + /// Dispatch an action on `host`. Returns the resulting `DispatchHandle` + /// (carrying `clientSeq`) for optimistic-update correlation. + @discardableResult + public func dispatch(host: HostId, action: StateAction) async throws -> DispatchHandle { + guard let runtime = hosts[host] else { + throw HostError.unknownHost(host) + } + return try await runtime.dispatch(action) + } + + // MARK: - Event multicast + + /// Subscribe to a fan-in stream of every inbound event from every + /// registered host. + /// + /// Each call returns a fresh `AsyncStream` — multiple consumers can + /// listen independently. The stream uses + /// `.bufferingNewest(defaultFanInBuffer)`; slow consumers will lose + /// older events but the stream stays alive (matching the lossy `Lagged` + /// semantics of the Rust SDK's broadcast). + /// + /// **Ordering** is per-host only. Different hosts run independently; + /// there is no cross-host total order. + /// + /// `async` so registration completes synchronously with respect to the + /// caller — no events fired between `events()` and the next `await` are + /// missed. + public func events() async -> AsyncStream { + let id = bumpListenerId() + return AsyncStream( + bufferingPolicy: .bufferingNewest(defaultFanInBuffer) + ) { cont in + self.subscriptionListeners[id] = cont + cont.onTermination = { [weak self] _ in + guard let self else { return } + Task { await self.removeSubscriptionListener(id: id) } + } + } + } + + /// Subscribe to connection-state events for UX. Each call returns a + /// fresh stream. + public func hostEvents() async -> AsyncStream { + let id = bumpListenerId() + return AsyncStream( + bufferingPolicy: .bufferingNewest(defaultFanInBuffer) + ) { cont in + self.hostEventListeners[id] = cont + cont.onTermination = { [weak self] _ in + guard let self else { return } + Task { await self.removeHostEventListener(id: id) } + } + } + } + + private func broadcastSubscriptionEvent(_ event: HostSubscriptionEvent) { + for cont in subscriptionListeners.values { + cont.yield(event) + } + } + + private func broadcastHostEvent(_ event: HostEvent) { + for cont in hostEventListeners.values { + cont.yield(event) + } + } + + private func removeSubscriptionListener(id: UInt64) { + subscriptionListeners.removeValue(forKey: id) + } + + private func removeHostEventListener(id: UInt64) { + hostEventListeners.removeValue(forKey: id) + } + + private func bumpListenerId() -> UInt64 { + let id = nextListenerId + nextListenerId &+= 1 + return id + } + + // MARK: - Aggregated views + + /// Aggregated session summaries across every registered host, sorted by + /// `summary.modifiedAt` descending. Includes both the host id and label + /// so consumers can render a unified inbox without losing host + /// attribution. + /// + /// **Tie-breaking:** for equal `modifiedAt`, summaries are ordered by + /// host registration order, then by `summary.resource`, so the result + /// is deterministic across calls. + public func aggregatedSessions() async -> [HostedSessionSummary] { + let order = hostOrder + let orderIndex = Dictionary(uniqueKeysWithValues: order.enumerated().map { ($1, $0) }) + var out: [HostedSessionSummary] = [] + for id in order { + guard let runtime = hosts[id] else { continue } + let snap = await runtime.snapshot() + for summary in snap.sessionSummaries { + out.append(HostedSessionSummary( + hostId: snap.id, + hostLabel: snap.label, + summary: summary + )) + } + } + return out.sorted { lhs, rhs in + if lhs.summary.modifiedAt != rhs.summary.modifiedAt { + return lhs.summary.modifiedAt > rhs.summary.modifiedAt + } + let li = orderIndex[lhs.hostId] ?? Int.max + let ri = orderIndex[rhs.hostId] ?? Int.max + if li != ri { return li < ri } + return lhs.summary.resource < rhs.summary.resource + } + } + + /// Aggregated agents across every registered host, in registration order + /// per host. + public func aggregatedAgents() async -> [HostedAgent] { + var out: [HostedAgent] = [] + for id in hostOrder { + guard let runtime = hosts[id] else { continue } + let snap = await runtime.snapshot() + for agent in snap.agents { + out.append(HostedAgent( + hostId: snap.id, + hostLabel: snap.label, + agent: agent + )) + } + } + return out + } + + // MARK: - Shutdown + + /// Tear down every registered host's supervisor and finish all event + /// streams. Safe to call multiple times. + public func shutdown() async { + if didShutDown { return } + didShutDown = true + let runtimes = hosts.values.map { $0 } + hosts.removeAll() + hostOrder.removeAll() + for runtime in runtimes { + await runtime.shutdown() + } + for cont in subscriptionListeners.values { cont.finish() } + subscriptionListeners.removeAll() + for cont in hostEventListeners.values { cont.finish() } + hostEventListeners.removeAll() + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/ReconnectPolicy.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/ReconnectPolicy.swift new file mode 100644 index 000000000..58627f039 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/ReconnectPolicy.swift @@ -0,0 +1,126 @@ +// ReconnectPolicy — backoff schedule for `HostRuntime`. + +import Foundation + +/// Backoff schedule between reconnect attempts. +public enum ReconnectBackoff: Sendable, Equatable { + /// Retry immediately, with no delay between attempts. + case immediate + /// Wait a fixed amount of time between attempts. + case constant(Duration) + /// Exponential backoff: `delay = min(initial * multiplier^(attempt-1), max)`. + case exponential(initial: Duration, max: Duration, multiplier: Double) + + /// Compute the delay before the `attempt`-th retry (1-based). + public func delay(forAttempt attempt: Int) -> Duration { + switch self { + case .immediate: + return .zero + case .constant(let delay): + return delay + case .exponential(let initial, let max, let multiplier): + let attempt = Swift.max(attempt, 1) + let mult = Swift.max(multiplier, 1.0) + let scaled = initial.seconds * pow(mult, Double(attempt - 1)) + let bounded = Swift.min(scaled, max.seconds) + return .seconds(bounded) + } + } +} + +/// Reconnect behaviour for a single host. +/// +/// The supervisor enters a reconnect loop whenever an established connection +/// drops unexpectedly, or whenever a connect attempt fails. Use +/// `ReconnectPolicy.disabled` to opt out entirely (a single failure leaves +/// the host in `HostState.failed`). +public struct ReconnectPolicy: Sendable, Equatable { + /// Backoff schedule between attempts. + public var backoff: ReconnectBackoff + /// Random jitter applied to each computed backoff. The actual delay is + /// uniformly sampled from `[delay * (1 - jitter), delay * (1 + jitter)]`. + /// `0.0` disables jitter; values are clamped to `[0.0, 1.0]`. + public var jitter: Double + /// Maximum number of attempts before giving up. `nil` retries forever. + public var maxAttempts: Int? + /// When `true`, the attempt counter resets to zero after a successful + /// connection so the next reconnect starts at the initial backoff. + public var resetOnSuccess: Bool + + public init( + backoff: ReconnectBackoff, + jitter: Double, + maxAttempts: Int?, + resetOnSuccess: Bool + ) { + self.backoff = backoff + self.jitter = jitter + self.maxAttempts = maxAttempts + self.resetOnSuccess = resetOnSuccess + } + + /// Disable reconnects entirely. A single failure leaves the host in + /// `HostState.failed`; consumers can recover via manual `reconnect`. + public static let disabled = ReconnectPolicy( + backoff: .immediate, + jitter: 0.0, + maxAttempts: 0, + resetOnSuccess: true + ) + + /// Retry forever with no backoff. Almost certainly not what you want in + /// production — useful for tests. + public static let immediateForever = ReconnectPolicy( + backoff: .immediate, + jitter: 0.0, + maxAttempts: nil, + resetOnSuccess: true + ) + + /// Sensible default: exponential backoff from 250 ms up to 30 s, 25 % + /// jitter, retry forever, reset on success. + public static let exponential = ReconnectPolicy( + backoff: .exponential( + initial: .milliseconds(250), + max: .seconds(30), + multiplier: 2.0 + ), + jitter: 0.25, + maxAttempts: nil, + resetOnSuccess: true + ) + + /// Compute the delay before the `attempt`-th retry (1-based), applying + /// jitter via the supplied random sample in `[0.0, 1.0]`. + /// + /// Exposed so tests can drive it deterministically; the runtime passes a + /// real random sample. + public func delay(forAttempt attempt: Int, sample: Double) -> Duration { + let base = backoff.delay(forAttempt: attempt) + if jitter <= 0.0 || base == .zero { + return base + } + let bounded = max(0.0, min(jitter, 1.0)) + let s = max(0.0, min(sample, 1.0)) + // Map [0,1] -> [-jitter, +jitter] + let factor = 1.0 + (s * 2.0 - 1.0) * bounded + return .seconds(base.seconds * max(0.0, factor)) + } + + /// Whether `attempt` exceeds `maxAttempts`. + public func attemptsExhausted(_ attempt: Int) -> Bool { + guard let cap = maxAttempts else { return false } + return attempt > cap + } +} + +// MARK: - Duration helpers + +extension Duration { + /// Convert to seconds as a Double for arithmetic. Internal helper — + /// fractional jitter math is the only place that needs this. + fileprivate var seconds: Double { + let comps = components + return Double(comps.seconds) + Double(comps.attoseconds) / 1e18 + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/Samples/MultiHostExample.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/Samples/MultiHostExample.swift new file mode 100644 index 000000000..fd5b96704 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/Samples/MultiHostExample.swift @@ -0,0 +1,139 @@ +// MultiHostExample — small in-package sample showing two-host UX. +// +// Not a binary target; intended for documentation, manual exploration in a +// playground, and as a sanity reference for downstream consumers. The +// supporting `pairedFakeHost` builder uses `InMemoryTransport.pair()` so +// you can run `MultiHostExample.runDemo()` from a test or scratch script. + +import Foundation +import AgentHostProtocol + +/// Documentation-quality sample wiring up two hosts behind `MultiHostClient`. +/// +/// ```swift +/// // Run from a Swift script or test: +/// try await MultiHostExample.runDemo() +/// ``` +public enum MultiHostExample { + + /// Spin up two in-memory "hosts" (one with one session summary, one + /// with two), connect them through a `MultiHostClient`, and print the + /// aggregated session list before tearing down. + public static func runDemo() async throws { + let storeA = ExampleHostState(label: "Local", sessions: [ + exampleSummary("copilot:/local-1", "Local: refactor", modifiedAt: 1_700) + ]) + let storeB = ExampleHostState(label: "Tunnel", sessions: [ + exampleSummary("copilot:/remote-1", "Tunnel: feature work", modifiedAt: 2_000), + exampleSummary("copilot:/remote-2", "Tunnel: bugfix", modifiedAt: 1_500) + ]) + + let multi = MultiHostClient() + + let configA = HostConfig(id: "local", label: "Local", transportFactory: pairedFakeHost(storeA)) + let configB = HostConfig(id: "tunnel", label: "Tunnel", transportFactory: pairedFakeHost(storeB)) + + _ = try await multi.add(configA) + _ = try await multi.add(configB) + + // Wait briefly for both hosts to finish handshake. Production code + // would consume `hostEvents()` instead of polling. + try await Task.sleep(for: .milliseconds(50)) + + for hosted in await multi.aggregatedSessions() { + print("[\(hosted.hostLabel)] \(hosted.summary.title) — modified \(hosted.summary.modifiedAt)") + } + + await multi.shutdown() + } +} + +/// Fake host backing for the sample. Mirrors the test helper but trimmed +/// down to the bits the demo needs (initialize + listSessions). Lives in +/// the same module so the example can be referenced from Swift Playgrounds. +private final class ExampleHostState: @unchecked Sendable { + let label: String + let sessions: [SessionSummary] + + init(label: String, sessions: [SessionSummary]) { + self.label = label + self.sessions = sessions + } +} + +private func pairedFakeHost(_ state: ExampleHostState) -> HostTransportFactory { + { _ in + let (clientSide, serverSide) = InMemoryTransport.pair() + Task { await driveExampleHost(transport: serverSide, state: state) } + return clientSide + } +} + +private func driveExampleHost( + transport: InMemoryTransport, + state: ExampleHostState +) async { + let encoder = JSONEncoder() + while true { + let frame: TransportMessage? + do { + frame = try await transport.recv() + } catch { + return + } + guard let frame else { return } + guard case .text(let text) = frame, + let data = text.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let id = object["id"] as? Int, + let method = object["method"] as? String + else { continue } + let result: Any + switch method { + case "initialize": + let snapshotJSON: [String: Any] = [ + "resource": RootResourceURI, + "state": [ + "agents": [], + "activeSessions": state.sessions.count, + ] as [String: Any], + "fromSeq": 0, + ] + result = [ + "protocolVersion": "0.1.0", + "serverSeq": 0, + "snapshots": [snapshotJSON], + ] as [String: Any] + case "listSessions": + let items = state.sessions.compactMap { summary -> Any? in + guard let bytes = try? encoder.encode(summary), + let object = try? JSONSerialization.jsonObject(with: bytes) + else { return nil } + return object + } + result = ["items": items] + default: + result = [:] as [String: Any] + } + let response: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "result": result, + ] + guard let data = try? JSONSerialization.data(withJSONObject: response), + let text = String(data: data, encoding: .utf8) + else { continue } + try? await transport.send(.text(text)) + } +} + +private func exampleSummary(_ uri: String, _ title: String, modifiedAt: Int) -> SessionSummary { + SessionSummary( + resource: uri, + provider: "copilot", + title: title, + status: .idle, + createdAt: 0, + modifiedAt: modifiedAt + ) +} diff --git a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/MultiHostClientTests.swift b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/MultiHostClientTests.swift new file mode 100644 index 000000000..6934004e2 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/MultiHostClientTests.swift @@ -0,0 +1,670 @@ +// MultiHostClientTests — integration tests for the multi-host SDK. +// +// Each test spins up one or more `FakeHost`s over `InMemoryTransport.pair()` +// (mirroring `clients/rust/crates/ahp/tests/hosts.rs`) and exercises the +// `MultiHostClient` facade end to end. + +import XCTest +import AgentHostProtocol +@testable import AgentHostProtocolClient + +final class MultiHostClientTests: XCTestCase { + + // MARK: - single_constructor_yields_connected_handle + + func testSingleConstructorYieldsConnectedHandle() async throws { + let agent = makeAgent() + let state = FakeHostState(agents: [agent]) + let factory = makeFakeHostFactory(state: state) + let config = HostConfig(id: "local", label: "Local", transportFactory: factory) + + let (multi, _) = try await MultiHostClient.single(config) + defer { Task { await multi.shutdown() } } + + await waitForHostState(multi, id: "local") { $0.isConnected } + + let snap = await multi.host("local") + XCTAssertNotNil(snap) + XCTAssertEqual(snap?.label, "Local") + XCTAssertEqual(snap?.protocolVersion, "0.1.0") + XCTAssertEqual(snap?.agents.count, 1) + XCTAssertEqual(snap?.agents.first?.provider, "copilot") + XCTAssertNotNil(snap?.lastConnectedAt) + XCTAssertTrue(snap?.state.isConnected ?? false) + + await multi.shutdown() + } + + // MARK: - two_hosts_register_and_connect_independently + + func testTwoHostsRegisterAndConnectIndependently() async throws { + let multi = MultiHostClient() + + _ = try await multi.add(HostConfig( + id: "a", + label: "A", + transportFactory: makeFakeHostFactory(state: FakeHostState()) + )) + _ = try await multi.add(HostConfig( + id: "b", + label: "B", + transportFactory: makeFakeHostFactory(state: FakeHostState()) + )) + + await waitForHostState(multi, id: "a") { $0.isConnected } + await waitForHostState(multi, id: "b") { $0.isConnected } + + let hosts = await multi.hosts() + XCTAssertEqual(hosts.count, 2) + XCTAssertTrue(hosts.allSatisfy { $0.state.isConnected }) + let labels = Set(hosts.map(\.label)) + XCTAssertEqual(labels, ["A", "B"]) + + await multi.shutdown() + } + + // MARK: - aggregated_sessions_track_listsessions_then_notification + + func testAggregatedSessionsTrackListSessionsThenNotification() async throws { + let initial = makeSummary("copilot:/s1", "Initial title", modifiedAt: 1_000) + let added = makeSummary("copilot:/s2", "Added later", modifiedAt: 2_000) + + let factory = makeFakeHostFactory( + state: FakeHostState(sessions: [initial]), + injectAfterInit: added + ) + + let multi = MultiHostClient() + _ = try await multi.add(HostConfig(id: "local", label: "Local", transportFactory: factory)) + await waitForHostState(multi, id: "local") { $0.isConnected } + + await waitUntil { await multi.aggregatedSessions().count == 2 } + let aggregated = await multi.aggregatedSessions() + let titles = aggregated.map(\.summary.title) + XCTAssertEqual(titles, ["Added later", "Initial title"]) + XCTAssertTrue(aggregated.allSatisfy { $0.hostId == "local" }) + XCTAssertTrue(aggregated.allSatisfy { $0.hostLabel == "Local" }) + + await multi.shutdown() + } + + // MARK: - host_client_handle_invalidates_after_reconnect + + func testHostClientHandleInvalidatesAfterReconnect() async throws { + let factory = makeFakeHostFactory(state: FakeHostState()) + + let multi = MultiHostClient() + let config = HostConfig(id: "local", label: "Local", transportFactory: factory) + .withReconnectPolicy(.immediateForever) + _ = try await multi.add(config) + + await waitForHostState(multi, id: "local") { $0.isConnected } + + let handleOpt = await multi.client(for: "local") + let handle = try XCTUnwrap(handleOpt) + let initialGeneration = handle.generation + try await handle.checkAlive() + + try await multi.reconnect("local") + await waitUntil { + guard let snap = await multi.host("local") else { return false } + return snap.generation > initialGeneration && snap.state.isConnected + } + + do { + try await handle.checkAlive() + XCTFail("expected HostError.hostReconnected") + } catch let error as HostError { + switch error { + case .hostReconnected(_, let handleGen, let currentGen): + XCTAssertEqual(handleGen, initialGeneration) + XCTAssertGreaterThan(currentGen, initialGeneration) + default: + XCTFail("unexpected error: \(error)") + } + } + + let freshOpt = await multi.client(for: "local") + let fresh = try XCTUnwrap(freshOpt) + XCTAssertGreaterThan(fresh.generation, initialGeneration) + try await fresh.checkAlive() + + await multi.shutdown() + } + + // MARK: - remove_host_terminates_supervisor_and_emits_event + + func testRemoveHostTerminatesSupervisorAndEmitsEvent() async throws { + let factory = makeFakeHostFactory(state: FakeHostState()) + + let multi = MultiHostClient() + _ = try await multi.add(HostConfig(id: "temp", label: "Temporary", transportFactory: factory)) + await waitForHostState(multi, id: "temp") { $0.isConnected } + + let events = await multi.hostEvents() + + try await multi.remove("temp") + + var sawRemoved = false + let deadline = ContinuousClock.now + .milliseconds(2_000) + var iter = events.makeAsyncIterator() + while ContinuousClock.now < deadline { + guard let event = try await Self.nextWithTimeout(&iter, timeout: .milliseconds(200)) + else { break } + if case .removed(let id) = event, id == "temp" { + sawRemoved = true + break + } + } + XCTAssertTrue(sawRemoved, "expected HostEvent.removed for temp") + + let snap = await multi.host("temp") + XCTAssertNil(snap) + + await multi.shutdown() + } + + // MARK: - fan_in_events_carry_host_id_and_resource + + func testFanInEventsCarryHostIdAndResource() async throws { + let initialA = makeSummary("copilot:/a-1", "first-a", modifiedAt: 100) + let injectA = makeSummary("copilot:/added-a", "a-side", modifiedAt: 200) + let initialB = makeSummary("copilot:/b-1", "first-b", modifiedAt: 100) + let injectB = makeSummary("copilot:/added-b", "b-side", modifiedAt: 300) + + let multi = MultiHostClient() + let events = await multi.events() + + _ = try await multi.add(HostConfig( + id: "a", + label: "Host A", + transportFactory: makeFakeHostFactory( + state: FakeHostState(sessions: [initialA]), + injectAfterInit: injectA + ) + )) + _ = try await multi.add(HostConfig( + id: "b", + label: "Host B", + transportFactory: makeFakeHostFactory( + state: FakeHostState(sessions: [initialB]), + injectAfterInit: injectB + ) + )) + + var hostsSeen: Set = [] + let deadline = ContinuousClock.now + .milliseconds(3_000) + var iter = events.makeAsyncIterator() + while hostsSeen.count < 2 && ContinuousClock.now < deadline { + guard let event = try await Self.nextWithTimeout(&iter, timeout: .milliseconds(500)) + else { break } + hostsSeen.insert(event.hostId) + // Notifications carry no resource URI by design. + XCTAssertNil(event.resource) + } + XCTAssertTrue(hostsSeen.contains("a"), "missing event from host A; saw \(hostsSeen)") + XCTAssertTrue(hostsSeen.contains("b"), "missing event from host B; saw \(hostsSeen)") + + await multi.shutdown() + } + + // MARK: - transport_factory_is_called_for_each_reconnect + + func testTransportFactoryIsCalledForEachReconnect() async throws { + let counter = CallCounter() + let factory: HostTransportFactory = { _ in + await counter.bump() + let (clientSide, serverSide) = InMemoryTransport.pair() + _ = FakeHost.start(transport: serverSide, state: FakeHostState()) + return clientSide + } + + let multi = MultiHostClient() + let config = HostConfig(id: "local", label: "Local", transportFactory: factory) + .withReconnectPolicy(.immediateForever) + _ = try await multi.add(config) + + await waitForHostState(multi, id: "local") { $0.isConnected } + let count1 = await counter.value() + XCTAssertEqual(count1, 1) + + try await multi.reconnect("local") + await waitUntil { + let snap = await multi.host("local") + return await counter.value() >= 2 && (snap?.state.isConnected ?? false) + } + let count2 = await counter.value() + XCTAssertEqual(count2, 2) + + await multi.shutdown() + } + + // MARK: - duplicate_host_id_is_rejected + + func testDuplicateHostIdIsRejected() async throws { + let multi = MultiHostClient() + _ = try await multi.add(HostConfig( + id: "dup", + label: "first", + transportFactory: makeFakeHostFactory(state: FakeHostState()) + )) + + do { + _ = try await multi.add(HostConfig( + id: "dup", + label: "second", + transportFactory: makeFakeHostFactory(state: FakeHostState()) + )) + XCTFail("expected duplicate-id rejection") + } catch let error as HostError { + if case .duplicateHost(let id) = error { + XCTAssertEqual(id, "dup") + } else { + XCTFail("expected .duplicateHost, got \(error)") + } + } + + await multi.shutdown() + } + + // MARK: - subscribe_while_failed_remembers_uri_for_replay + + /// While a host is in `.failed` state, `subscribe(host:uri:)` should still + /// remember the URI so the next successful reconnect picks it up. + /// Unsubscribe in the same window should drop the URI from the replay set. + func testSubscribeWhileFailedRemembersForReplay() async throws { + // Build a transport factory that fails the first attempt and then + // succeeds on subsequent attempts. Combine with `.disabled` so the + // first failure parks the host in `.failed`. + let didFirstFail = ActorBool() + let factory: HostTransportFactory = { _ in + if !(await didFirstFail.value) { + await didFirstFail.set(true) + throw TransportError.io("intentional first-attempt failure") + } + let (clientSide, serverSide) = InMemoryTransport.pair() + _ = FakeHost.start(transport: serverSide, state: FakeHostState()) + return clientSide + } + let config = HostConfig(id: "tt", label: "T", transportFactory: factory) + .withReconnectPolicy(.disabled) + + let multi = MultiHostClient() + _ = try await multi.add(config) + + await waitForHostState(multi, id: "tt") { $0.isFailed } + + // Subscribe while disconnected. The runtime returns `hostShutDown` + // but appends the URI to the replay set. + do { + _ = try await multi.subscribe(host: "tt", uri: "copilot:/queued") + XCTFail("expected subscribe to reject while failed") + } catch let error as HostError { + if case .hostShutDown = error {} else { + XCTFail("expected .hostShutDown while failed, got \(error)") + } + } + + // Unsubscribe an unrelated URI while disconnected — should succeed + // and not throw, even though no live client exists. + try await multi.unsubscribe(host: "tt", uri: "copilot:/never-subscribed") + + var snap = await multi.host("tt") + XCTAssertEqual(snap?.subscriptions.contains("copilot:/queued"), true, + "queued subscribe URI should be recorded for replay") + + // Manually reconnect — second attempt succeeds. + try await multi.reconnect("tt") + await waitForHostState(multi, id: "tt") { $0.isConnected } + + snap = await multi.host("tt") + XCTAssertEqual(snap?.subscriptions.contains("copilot:/queued"), true, + "subscription should survive into the new connection") + + await multi.shutdown() + } + + // MARK: - shutdown_tears_down_all_hosts_and_streams + + func testShutdownTearsDownAllHostsAndStreams() async throws { + let multi = MultiHostClient() + + _ = try await multi.add(HostConfig( + id: "alpha", + label: "Alpha", + transportFactory: makeFakeHostFactory(state: FakeHostState()) + )) + _ = try await multi.add(HostConfig( + id: "beta", + label: "Beta", + transportFactory: makeFakeHostFactory(state: FakeHostState()) + )) + await waitForHostState(multi, id: "alpha") { $0.isConnected } + await waitForHostState(multi, id: "beta") { $0.isConnected } + + let events = await multi.events() + let hostEvents = await multi.hostEvents() + + await multi.shutdown() + + // After shutdown, both streams should finish so `for await` exits. + var subEvents = 0 + for await _ in events { subEvents += 1 } + var hostEventCount = 0 + for await _ in hostEvents { hostEventCount += 1 } + // We don't assert specific counts — just that the loops end. + _ = subEvents + _ = hostEventCount + + // No host snapshots should be retrievable. + let alphaSnap = await multi.host("alpha") + let betaSnap = await multi.host("beta") + XCTAssertNil(alphaSnap) + XCTAssertNil(betaSnap) + + // Subsequent `add` should reject with `hostShutDown`. + do { + _ = try await multi.add(HostConfig( + id: "gamma", + label: "Gamma", + transportFactory: makeFakeHostFactory(state: FakeHostState()) + )) + XCTFail("expected add to reject after shutdown") + } catch let error as HostError { + if case .hostShutDown(let id) = error { + XCTAssertEqual(id, "gamma") + } else { + XCTFail("expected .hostShutDown, got \(error)") + } + } + + // Idempotent. + await multi.shutdown() + } + + // MARK: - state_during_backoff_after_drop_is_reconnecting + + /// Regression: while the supervisor is sleeping in backoff after a + /// successful connection dropped, snapshots must report + /// `.reconnecting(...)` rather than `.connected`. The previous + /// implementation only transitioned at the *top* of the next iteration + /// so consumers observed `.connected` for the entire backoff window. + func testStateDuringBackoffAfterDropIsReconnecting() async throws { + // Build a transport factory whose first server side answers the + // handshake + listSessions and then waits for a "drop now" signal + // before closing — that way we can deterministically wait for + // `.connected`, then trigger the drop, then assert the state + // transitioned to `.reconnecting` during the backoff sleep. + // Subsequent connect attempts park (never reply) so the runtime + // stays in the post-drop backoff/reconnecting window. + let dropSignal = DropSignal() + let didFirstConnect = ActorBool() + let factory: HostTransportFactory = { _ in + let (clientSide, serverSide) = InMemoryTransport.pair() + if !(await didFirstConnect.value) { + await didFirstConnect.set(true) + Task { + // Answer requests until the drop signal is set. + while !dropSignal.isReady { + let frame: TransportMessage? + do { frame = try await serverSide.recv() } catch { return } + guard let frame, case .text(let text) = frame, + let data = text.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let id = object["id"] as? Int, + let method = object["method"] as? String + else { continue } + let result: Any + switch method { + case "initialize": + let snap: [String: Any] = [ + "resource": RootResourceURI, + "state": ["agents": [], "activeSessions": 0] as [String: Any], + "fromSeq": 0, + ] + result = [ + "protocolVersion": "0.1.0", + "serverSeq": 0, + "snapshots": [snap], + ] as [String: Any] + case "listSessions": + result = ["items": []] as [String: Any] + default: + result = [:] as [String: Any] + } + let resp: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "result": result, + ] + if let bytes = try? JSONSerialization.data(withJSONObject: resp), + let body = String(data: bytes, encoding: .utf8) { + try? await serverSide.send(.text(body)) + } + } + try? await serverSide.close() + } + // Sample the drop signal periodically so we close even if + // no further request arrives. This keeps the test + // deterministic when there's no in-flight request to + // unblock the recv() loop. + Task { + while !dropSignal.isReady { + try? await Task.sleep(for: .milliseconds(20)) + } + try? await serverSide.close() + } + } else { + // Subsequent attempts: park (never reply) so the runtime + // stays in `.reconnecting` while we observe. + Task { _ = try? await serverSide.recv() } + } + return clientSide + } + // Long initial backoff so we have a generous window to observe + // `.reconnecting` during sleep. + let policy = ReconnectPolicy( + backoff: .constant(.seconds(5)), + jitter: 0.0, + maxAttempts: nil, + resetOnSuccess: true + ) + let config = HostConfig(id: "drop", label: "Drop", transportFactory: factory) + .withReconnectPolicy(policy) + + let multi = MultiHostClient() + _ = try await multi.add(config) + + // Wait for the first connect to land. + await waitForHostState(multi, id: "drop") { $0.isConnected } + + // Trigger the drop, then wait for the runtime to surface + // `.reconnecting`. Without the fix, this poll would spin until + // timeout because state stayed `.connected` through the entire + // backoff sleep. + dropSignal.trigger() + await waitForHostState(multi, id: "drop", timeout: .seconds(2)) { state in + if case .reconnecting = state { return true } + return false + } + + await multi.shutdown() + } + + // MARK: - failed_handshake_shuts_down_underlying_client + + /// Regression: if `initialize`/`reconnect` throws after `client.connect()` + /// has already started the writer/receive tasks, the supervisor must + /// shut the `AHPClient` down before propagating — otherwise the + /// orphaned client's tasks keep holding the transport indefinitely + /// while the supervisor opens a fresh one for the next attempt. + /// We assert this indirectly by observing that the wrapped transport's + /// `close()` is invoked. + func testFailedHandshakeShutsDownUnderlyingClient() async throws { + let observer = ClosedObserver() + let factory: HostTransportFactory = { _ in + let (clientSide, serverSide) = InMemoryTransport.pair() + // Server returns an RPC error response to `initialize`. + _ = startFailingInitFakeHost(transport: serverSide) + return TrackingTransport(clientSide, observer: observer) + } + // `disabled` so the host bails into `.failed` after one failed + // handshake instead of looping forever. + let config = HostConfig(id: "fail", label: "Fail", transportFactory: factory) + .withReconnectPolicy(.disabled) + + let multi = MultiHostClient() + _ = try await multi.add(config) + + // Wait for the host to hit `.failed`. + await waitForHostState(multi, id: "fail", timeout: .seconds(2)) { $0.isFailed } + + // The supervisor should have shut down the AHPClient on the + // handshake failure, which closes the wrapped transport. + let closed = await observer.isClosed + XCTAssertTrue(closed, "AHPClient.shutdown() should have closed the transport on a failed handshake") + + await multi.shutdown() + } + + // MARK: - Helpers + + /// Like `nextWithTimeout` from `AHPClientTestHelpers` but typed for any + /// `AsyncStream` element. + private static func nextWithTimeout( + _ iterator: inout AsyncStream.AsyncIterator, + timeout: Duration + ) async throws -> E? where E: Sendable { + try await withThrowingTaskGroup(of: E?.self) { group in + group.addTask { [iterator = iterator] in + var iter = iterator + return await iter.next() + } + group.addTask { + try await Task.sleep(for: timeout) + throw TestTimeoutError() + } + defer { group.cancelAll() } + return try await group.next()! + } + } +} + +private actor CallCounter { + private var n: Int = 0 + func bump() { n += 1 } + func value() -> Int { n } +} + +private actor ActorBool { + private var flag: Bool = false + var value: Bool { flag } + func set(_ v: Bool) { flag = v } +} + +/// `Sendable` flag used by drop-driven tests. Using a `final class` with a +/// lock instead of an actor so the server-side `recv` loop can poll it +/// without `await`-ing into actor isolation between every frame. +private final class DropSignal: @unchecked Sendable { + private let lock = NSLock() + private var flag: Bool = false + var isReady: Bool { + lock.lock(); defer { lock.unlock() } + return flag + } + func trigger() { + lock.lock(); defer { lock.unlock() } + flag = true + } +} + +private struct TestTimeoutError: Error {} + +// MARK: - Failing-handshake fake host + +/// Fake-host driver that responds to `initialize` with a JSON-RPC error +/// instead of a result. Causes the client's `initialize` request to throw +/// `AHPClientError.rpc(...)`. Used to assert the supervisor tears the +/// `AHPClient` down on a failed handshake. +private func startFailingInitFakeHost( + transport: InMemoryTransport, + code: Int = -32000, + message: String = "init refused for test" +) -> Task { + Task { + while !Task.isCancelled { + let frame: TransportMessage? + do { + frame = try await transport.recv() + } catch { + return + } + guard let frame else { return } + guard case .text(let text) = frame, + let data = text.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let id = object["id"] as? Int, + let method = object["method"] as? String + else { continue } + if method == "initialize" || method == "reconnect" { + let resp: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "error": [ + "code": code, + "message": message, + ] as [String: Any], + ] + if let respData = try? JSONSerialization.data(withJSONObject: resp), + let respText = String(data: respData, encoding: .utf8) { + try? await transport.send(.text(respText)) + } + } else { + let resp: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "result": [:] as [String: Any], + ] + if let respData = try? JSONSerialization.data(withJSONObject: resp), + let respText = String(data: respData, encoding: .utf8) { + try? await transport.send(.text(respText)) + } + } + } + } +} + +// MARK: - Tracking transport wrapper + +/// A `Sendable` thin wrapper around `InMemoryTransport` that flips an +/// observable `isClosed` flag when `close()` runs. Used to assert that the +/// supervisor calls `AHPClient.shutdown()` (which calls `transport.close()`) +/// on a failed handshake. +private final class TrackingTransport: AHPTransport, @unchecked Sendable { + private let underlying: InMemoryTransport + private let observer: ClosedObserver + + init(_ underlying: InMemoryTransport, observer: ClosedObserver) { + self.underlying = underlying + self.observer = observer + } + + func send(_ message: TransportMessage) async throws { + try await underlying.send(message) + } + + func recv() async throws -> TransportMessage? { + try await underlying.recv() + } + + func close() async throws { + await observer.markClosed() + try await underlying.close() + } +} + +private actor ClosedObserver { + private(set) var closeCount: Int = 0 + var isClosed: Bool { closeCount > 0 } + func markClosed() { closeCount += 1 } +} diff --git a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/MultiHostTestHelpers.swift b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/MultiHostTestHelpers.swift new file mode 100644 index 000000000..c863089c9 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/MultiHostTestHelpers.swift @@ -0,0 +1,240 @@ +// MultiHostTestHelpers — shared infrastructure for `MultiHostClientTests`. +// +// Mirrors the Rust integration-test scaffolding (`crates/ahp/tests/hosts.rs`): +// a small "fake host" actor that drives the server side of an +// `InMemoryTransport.pair()` and responds to `initialize`/`reconnect`/ +// `listSessions`/`subscribe`. Optionally pushes a `notify/sessionAdded` +// after `initialize` to exercise the post-handshake notification path. + +import Foundation +import AgentHostProtocol +@testable import AgentHostProtocolClient + +/// Minimal mutable state for `FakeHost`. Conceptually equivalent to Rust's +/// `FakeHostState`. +struct FakeHostState: Sendable { + var agents: [AgentInfo] = [] + var sessions: [SessionSummary] = [] + var serverSeq: Int = 0 +} + +/// Server-side responder for one in-memory transport pair. Constructed via +/// `FakeHost.start(transport:state:injectAfterInit:)`. Drives the loop in a +/// detached `Task`; cancelled implicitly when the client closes the +/// transport (`recv` throws). +struct FakeHost { + /// Spin up a fake host driving `transport` (the *server* side of an + /// `InMemoryTransport.pair()`). When `injectAfterInit` is non-nil, the + /// fake pushes a `notify/sessionAdded` for that summary shortly after + /// answering `initialize` (or `reconnect`). + static func start( + transport: InMemoryTransport, + state: FakeHostState, + injectAfterInit: SessionSummary? = nil + ) -> Task { + Task { + await drive(transport: transport, state: state, injectAfterInit: injectAfterInit) + } + } + + private static func drive( + transport: InMemoryTransport, + state: FakeHostState, + injectAfterInit: SessionSummary? + ) async { + let encoder = JSONEncoder() + while !Task.isCancelled { + let frame: TransportMessage? + do { + frame = try await transport.recv() + } catch { + return + } + guard let frame else { return } + guard case .text(let text) = frame, + let data = text.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { continue } + + let id = object["id"] as? Int + let method = object["method"] as? String + + if let id, let method { + let result = handleRequest(method: method, params: object["params"], state: state) + let resp: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "result": result, + ] + guard let respData = try? JSONSerialization.data(withJSONObject: resp), + let respText = String(data: respData, encoding: .utf8) + else { continue } + try? await transport.send(.text(respText)) + + if let summary = injectAfterInit, method == "initialize" || method == "reconnect" { + // Tiny delay so the client's `listSessions` request has + // landed before the notification arrives. + try? await Task.sleep(for: .milliseconds(20)) + let summaryAny: Any + if let bytes = try? encoder.encode(summary), + let obj = try? JSONSerialization.jsonObject(with: bytes) { + summaryAny = obj + } else { + continue + } + let notif: [String: Any] = [ + "jsonrpc": "2.0", + "method": "notification", + "params": [ + "notification": [ + "type": "notify/sessionAdded", + "summary": summaryAny, + ] as [String: Any], + ] as [String: Any], + ] + if let notifData = try? JSONSerialization.data(withJSONObject: notif), + let notifText = String(data: notifData, encoding: .utf8) { + try? await transport.send(.text(notifText)) + } + } + } + } + } + + private static func handleRequest( + method: String, + params: Any?, + state: FakeHostState + ) -> Any { + switch method { + case "initialize": + let agentsAny = sessionSummariesToJSON(state.agents) + let snapshot: [String: Any] = [ + "resource": RootResourceURI, + "state": [ + "agents": agentsAny, + "activeSessions": state.sessions.count, + ] as [String: Any], + "fromSeq": state.serverSeq, + ] + return [ + "protocolVersion": "0.1.0", + "serverSeq": state.serverSeq, + "snapshots": [snapshot], + ] + case "reconnect": + return [ + "type": "replay", + "actions": [], + "missing": [], + ] as [String: Any] + case "listSessions": + let items = sessionSummariesToJSON(state.sessions) + return ["items": items] + case "subscribe": + let resource = (params as? [String: Any])?["resource"] as? String ?? RootResourceURI + let snap: [String: Any] = [ + "resource": resource, + "state": [ + "agents": sessionSummariesToJSON(state.agents) + ] as [String: Any], + "fromSeq": state.serverSeq, + ] + return ["snapshot": snap] + default: + return [:] as [String: Any] + } + } +} + +private func sessionSummariesToJSON(_ values: [T]) -> [Any] { + let encoder = JSONEncoder() + return values.compactMap { value -> Any? in + guard let data = try? encoder.encode(value), + let object = try? JSONSerialization.jsonObject(with: data) + else { return nil } + return object + } +} + +/// Build a transport factory that, on every call, opens a fresh +/// `InMemoryTransport.pair()` and starts a `FakeHost` driving the server +/// side. Optionally injects a `notify/sessionAdded` after init. +func makeFakeHostFactory( + state: FakeHostState, + injectAfterInit: SessionSummary? = nil, + onConnect: (@Sendable () -> Void)? = nil +) -> HostTransportFactory { + { _ in + let (clientSide, serverSide) = InMemoryTransport.pair() + onConnect?() + _ = FakeHost.start( + transport: serverSide, + state: state, + injectAfterInit: injectAfterInit + ) + return clientSide + } +} + +/// Build a `SessionSummary` with the minimal required fields, defaulting +/// optional fields to `nil` so tests stay terse. +func makeSummary( + _ uri: String, + _ title: String, + modifiedAt: Int = 0, + createdAt: Int = 0 +) -> SessionSummary { + SessionSummary( + resource: uri, + provider: "copilot", + title: title, + status: .idle, + createdAt: createdAt, + modifiedAt: modifiedAt + ) +} + +/// Build an `AgentInfo` with the minimal required fields. +func makeAgent( + provider: String = "copilot", + displayName: String = "Copilot" +) -> AgentInfo { + AgentInfo( + provider: provider, + displayName: displayName, + description: "demo", + models: [] + ) +} + +/// Poll `condition` every 10 ms until it returns true or the timeout +/// elapses. Crashes (intentionally) on timeout. +func waitUntil( + timeout: Duration = .seconds(2), + _ condition: @Sendable () async -> Bool, + file: StaticString = #file, + line: UInt = #line +) async { + let deadline = ContinuousClock.now + timeout + while ContinuousClock.now < deadline { + if await condition() { return } + try? await Task.sleep(for: .milliseconds(10)) + } + fatalError("waitUntil timed out", file: file, line: line) +} + +/// Wait for a host's `HostState` to satisfy `predicate`. +func waitForHostState( + _ multi: MultiHostClient, + id: HostId, + timeout: Duration = .seconds(2), + _ predicate: @escaping @Sendable (HostState) -> Bool, + file: StaticString = #file, + line: UInt = #line +) async { + await waitUntil(timeout: timeout, { + guard let snap = await multi.host(id) else { return false } + return predicate(snap.state) + }, file: file, line: line) +} diff --git a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/ReconnectPolicyTests.swift b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/ReconnectPolicyTests.swift new file mode 100644 index 000000000..bd86fad1e --- /dev/null +++ b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/ReconnectPolicyTests.swift @@ -0,0 +1,76 @@ +// ReconnectPolicyTests — backoff + jitter unit tests, mirroring the Rust +// `policy.rs` test module. + +import XCTest +@testable import AgentHostProtocolClient + +final class ReconnectPolicyTests: XCTestCase { + + func testExponentialBackoffCapsAtMax() { + let backoff: ReconnectBackoff = .exponential( + initial: .seconds(1), + max: .seconds(10), + multiplier: 2.0 + ) + XCTAssertEqual(backoff.delay(forAttempt: 1), .seconds(1)) + XCTAssertEqual(backoff.delay(forAttempt: 2), .seconds(2)) + XCTAssertEqual(backoff.delay(forAttempt: 3), .seconds(4)) + XCTAssertEqual(backoff.delay(forAttempt: 4), .seconds(8)) + // Capped at max. + XCTAssertEqual(backoff.delay(forAttempt: 5), .seconds(10)) + XCTAssertEqual(backoff.delay(forAttempt: 50), .seconds(10)) + } + + func testJitterZeroReturnsBaseDelay() { + let policy = ReconnectPolicy( + backoff: .constant(.seconds(5)), + jitter: 0.0, + maxAttempts: nil, + resetOnSuccess: true + ) + XCTAssertEqual(policy.delay(forAttempt: 1, sample: 0.5), .seconds(5)) + } + + func testJitterAtExtremesScalesDelay() { + let policy = ReconnectPolicy( + backoff: .constant(.seconds(10)), + jitter: 0.5, + maxAttempts: nil, + resetOnSuccess: true + ) + // sample 0 -> -50% -> 5s + XCTAssertEqual(policy.delay(forAttempt: 1, sample: 0.0), .seconds(5)) + // sample 1 -> +50% -> 15s + XCTAssertEqual(policy.delay(forAttempt: 1, sample: 1.0), .seconds(15)) + // sample 0.5 -> +0% -> 10s + XCTAssertEqual(policy.delay(forAttempt: 1, sample: 0.5), .seconds(10)) + } + + func testDisabledPolicyExhaustsImmediately() { + let policy = ReconnectPolicy.disabled + XCTAssertTrue(policy.attemptsExhausted(1)) + } + + func testUnboundedPolicyNeverExhausts() { + let policy = ReconnectPolicy.exponential + XCTAssertFalse(policy.attemptsExhausted(1_000_000)) + } + + func testImmediateBackoffIsZero() { + let backoff: ReconnectBackoff = .immediate + XCTAssertEqual(backoff.delay(forAttempt: 1), .zero) + XCTAssertEqual(backoff.delay(forAttempt: 100), .zero) + } + + func testSampleClamping() { + let policy = ReconnectPolicy( + backoff: .constant(.seconds(10)), + jitter: 0.5, + maxAttempts: nil, + resetOnSuccess: true + ) + // Samples outside [0, 1] are clamped. + XCTAssertEqual(policy.delay(forAttempt: 1, sample: -1.0), .seconds(5)) + XCTAssertEqual(policy.delay(forAttempt: 1, sample: 2.0), .seconds(15)) + } +}