Add MultiHostClient to AgentHostProtocolClient

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
+        }
+    }
+}

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<P: Encodable & Sendable, R: Decodable & Sendable>(
+        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
+    }
+}

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
+    }
+}

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
+        }
+    }
+}

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
+    }
+}