fix(ios): bound link requests and keep terminal input flowing

cmux-next-ios/CmuxMobileKit/Sources/CmuxMobileLive/AcpmuxConnection.swift

@@ -25,8 +25,9 @@ public actor AcpmuxConnection {
     private var nextID = 1
     private var pending: [Int: CheckedContinuation<JSONValue, any Error>] = [:]
     private var readTask: Task<Void, Never>?
-    private var terminalError: (any Error)?
-    var isClosed: Bool { terminalError != nil }
+    private var sends: [Int: Task<Void, Never>] = [:]
+    private var ended: (any Error)?
+    var isClosed: Bool { ended != nil }
     public nonisolated let notifications: AsyncStream<Notification>
     private let continuation: AsyncStream<Notification>.Continuation
 
@@ -36,49 +37,51 @@ public actor AcpmuxConnection {
     }
 
     public func start() {
-        guard readTask == nil, terminalError == nil else { return }
+        guard readTask == nil, ended == nil else { return }
         readTask = Task { await self.readLoop() }
     }
 
+    /// Prompts resolve when the turn ends, which can include a person answering an
+    /// approval. Other requests are short reads or mutations; both have finite bounds.
+    static func replyTimeout(for method: String) -> Duration {
+        method == "session/prompt" ? .seconds(24 * 60 * 60) : .seconds(30)
+    }
+
     public func request(_ method: String, _ params: JSONValue = .object([:])) async throws -> JSONValue {
+        try await request(method, params, timeout: Self.replyTimeout(for: method))
+    }
+
+    /// A late reply is ignored; cancellation and expiry release the request and
+    /// cancel its write without closing unrelated requests on the same channel.
+    public func request(_ method: String, _ params: JSONValue, timeout: Duration) async throws -> JSONValue {
         try Task.checkCancellation()
-        if let terminalError { throw terminalError }
+        if let ended { throw ended }
         start()
         let id = nextID
         nextID += 1
         let message: JSONValue = .object(["jsonrpc": .string("2.0"), "id": .number(Double(id)), "method": .string(method), "params": params])
         let data = try JSONEncoder().encode(message)
         rpcLog.debug("-> \(id) \(method, privacy: .public) \(data.count) bytes")
-        return try await withTaskCancellationHandler {
-            try await withCheckedThrowingContinuation { continuation in
-                guard !Task.isCancelled else {
-                    continuation.resume(throwing: CancellationError())
-                    return
-                }
-                pending[id] = continuation
-                Task {
-                    do { try await channel.send(data) } catch { self.fail(id, error) }
-                }
-            }
-        } onCancel: {
-            Task { await self.fail(id, CancellationError()) }
-        }
+        return try await LinkDeadline.run(
+            timeout: timeout,
+            timeoutError: LinkError(code: "timeout", message: "The Mac didn't answer (\(method))"),
+            operation: { try await self.sendAndWait(id: id, data: data) },
+            onCancel: { await self.fail(id, CancellationError()) }
+        )
     }
 
-    /// `request` that gives up after `timeout` (a read the person waits on: the
-    /// agent catalog), or when the task that asked is cancelled. A late reply is dropped.
-    public func request(_ method: String, _ params: JSONValue, timeout: Duration) async throws -> JSONValue {
-        let id = nextID
-        let timer = Task { [weak self] in
-            try? await Task.sleep(for: timeout)
-            guard !Task.isCancelled else { return }
-            await self?.fail(id, LinkError(code: "timeout", message: "The Mac didn't answer (\(method))"))
-        }
-        defer { timer.cancel() }
-        return try await withTaskCancellationHandler {
-            try await request(method, params)
-        } onCancel: {
-            Task { [weak self] in await self?.fail(id, CancellationError()) }
+    private func sendAndWait(id: Int, data: Data) async throws -> JSONValue {
+        try Task.checkCancellation()
+        if let ended { throw ended }
+        return try await withCheckedThrowingContinuation { continuation in
+            pending[id] = continuation
+            sends[id] = Task {
+                do {
+                    try Task.checkCancellation()
+                    try await channel.send(data)
+                    sends[id] = nil
+                } catch { fail(id, error) }
+            }
         }
     }
 
@@ -127,19 +130,20 @@ public actor AcpmuxConnection {
 
     /// A JSON-RPC notification (no id, no reply), e.g. `session/cancel`.
     public func notify(_ method: String, _ params: JSONValue) async throws {
-        if let terminalError { throw terminalError }
+        if let ended { throw ended }
         start()
         let message: JSONValue = .object(["jsonrpc": .string("2.0"), "method": .string(method), "params": params])
         try await channel.send(try JSONEncoder().encode(message))
     }
 
     public func close() async {
         readTask?.cancel()
-        await channel.close(code: "eof", reason: "client closed")
         finish(LinkError(code: "closed", message: "acpmux channel closed"))
+        await channel.close(code: "eof", reason: "client closed")
     }
 
     private func fail(_ id: Int, _ error: any Error) {
+        sends.removeValue(forKey: id)?.cancel()
         pending.removeValue(forKey: id)?.resume(throwing: error)
     }
 
@@ -176,6 +180,7 @@ public actor AcpmuxConnection {
                 rpcLog.debug("<- \(Int(id)) with no request waiting")
                 return
             }
+            sends.removeValue(forKey: Int(id))?.cancel()
             rpcLog.debug("<- \(Int(id)) \(value["error"] == nil ? "result" : "error", privacy: .public)")
             if let error = value["error"] {
                 continuation.resume(throwing: RPCError(
@@ -232,8 +237,10 @@ public actor AcpmuxConnection {
     }
 
     private func finish(_ error: any Error) {
-        guard terminalError == nil else { return }
-        terminalError = error
+        guard ended == nil else { return }
+        ended = error
+        for task in sends.values { task.cancel() }
+        sends.removeAll()
         for (_, continuation) in pending { continuation.resume(throwing: error) }
         pending.removeAll()
         continuation.finish()

cmux-next-ios/CmuxMobileKit/Sources/CmuxMobileLive/AcpmuxHub.swift

@@ -26,7 +26,7 @@ public final class AcpmuxHub {
     /// Callers of `waitUntilReady` (a Feed answer whose Mac's link is opening).
     @ObservationIgnored private var readyWaiters: [UUID: CheckedContinuation<Bool, Never>] = [:]
     @ObservationIgnored private var pumpTask: Task<Void, Never>?
-    @ObservationIgnored private var sources: [String: WeakSource] = [:]
+    @ObservationIgnored private var sources: [String: [WeakSource]] = [:]
     /// Watch-level events for the Feed: `(sessionId, notification)`.
     @ObservationIgnored public var onWatchEvent: (@MainActor (AcpmuxHub.WatchEvent) -> Void)?
     /// Signs an Allow or an answer with the phone's approval key after Face ID
@@ -85,6 +85,11 @@ public final class AcpmuxHub {
     }
 
     private struct WeakSource { weak var source: AcpmuxSessionSource? }
+    /// A chat with an expired replay cache may still have an open reader. Keep routing
+    /// to it while a new reader loads its own complete snapshot from the Mac.
+    private var liveSources: [AcpmuxSessionSource] {
+        sources.values.flatMap { $0.compactMap(\.source) }
+    }
 
     public init() {}
 
@@ -96,8 +101,7 @@ public final class AcpmuxHub {
         isReady = false
         lastError = nil
         fullAccess = false
-        lastError = nil
-        for entry in sources.values { entry.source?.linkDown() }
+        for source in liveSources { source.linkDown() }
         guard let channel else { return }
         let connection = AcpmuxConnection(channel: channel)
         self.connection = connection
@@ -113,7 +117,7 @@ public final class AcpmuxHub {
             self.isReady = false
             self.fullAccess = false
             self.lastError = self.lastError ?? "acpmux channel closed"
-            for entry in self.sources.values { entry.source?.linkDown() }
+            for source in self.liveSources { source.linkDown() }
         }
         do {
             let initialized = try await connection.request("initialize", .object([
@@ -136,7 +140,7 @@ public final class AcpmuxHub {
             readyWaiters = [:]
             for waiter in waiters.values { waiter.resume(returning: true) }
             log.info("acpmux ready, \(next.count, privacy: .public) sessions")
-            for entry in sources.values { entry.source?.linkUp() }
+            for source in liveSources { source.linkUp() }
             onWatchEvent?(.sessionsLoaded(Array(next.values)))
             await seedPending(connection, requestTimeout: requestTimeout)
             try Task.checkCancellation()
@@ -146,7 +150,7 @@ public final class AcpmuxHub {
             isReady = false
             fullAccess = false
             lastError = LinkController.reason(error)
-            for entry in sources.values { entry.source?.linkDown() }
+            for source in liveSources { source.linkDown() }
             await connection.close()
             log.error("acpmux bind failed: \(String(describing: error), privacy: .public)")
         }
@@ -231,11 +235,13 @@ public final class AcpmuxHub {
         try await connection.notify(method, params)
     }
 
-    /// The chat source for an acpmux session (one per session).
+    /// Reuse a complete replay; after its budget expires, a new reader gets a fresh
+    /// attach instead of a suffix whose first deltas have no message to append to.
     public func source(sessionID: String, summary: AgentSessionSummary) -> AcpmuxSessionSource {
-        if let existing = sources[sessionID]?.source { return existing }
+        let existing = (sources[sessionID] ?? []).compactMap(\.source)
+        if let source = existing.last(where: { !$0.replay.requiresResync }) { return source }
         let source = AcpmuxSessionSource(hub: self, sessionID: sessionID, summary: summary)
-        sources[sessionID] = WeakSource(source: source)
+        sources[sessionID] = (existing + [source]).map { WeakSource(source: $0) }
         return source
     }
 
@@ -337,15 +343,17 @@ public final class AcpmuxHub {
     /// Full access mode again.
     func fullAccessRefused() {
         fullAccess = false
-        for entry in sources.values { entry.source?.fullAccessChanged(false) }
+        for source in liveSources { source.fullAccessChanged(false) }
     }
 
     private func handle(_ notification: AcpmuxConnection.Notification) {
         let params = notification.params
         let sessionID = params["sessionId"]?.stringValue
         switch notification.method {
         case "_acpmux/event":
-            if let sessionID { sources[sessionID]?.source?.receive(record: params) }
+            if let sessionID {
+                for entry in sources[sessionID] ?? [] { entry.source?.receive(record: params) }
+            }
         case "_acpmux/session_changed":
             guard let value = params["session"], let session = AcpmuxSession(value) else { return }
             sessions[session.id] = session
@@ -361,7 +369,7 @@ public final class AcpmuxHub {
             // Missed notifications: attached chats fill their gap, and watch runs again. The
             // chats first, before the next notification: a record after the dropped ones must
             // be held for the refill, not applied ahead of it.
-            for entry in sources.values { entry.source?.refill() }
+            for source in liveSources { source.refill() }
             Task { [weak self] in
                 guard let self, let connection = self.connection else { return }
                 _ = try? await connection.request("_acpmux/watch", .object(["enabled": .bool(true)]))
@@ -380,8 +388,8 @@ public final class AcpmuxSessionSource: AgentSessionSource {
     let sessionID: String
     private weak var hub: AcpmuxHub?
     private let records: AcpmuxRecords
-    private var continuations: [AsyncStream<AgentSessionEvent>.Continuation] = []
-    private var replay: [AgentSessionEvent] = []
+    private var continuations: [UUID: AsyncStream<AgentSessionEvent>.Continuation] = [:]
+    fileprivate var replay = AcpmuxReplay<AgentSessionEvent>(limit: AcpmuxSessionSource.maxHistory)
     private var lastSeq = 0
     private var attached = false
     private var attaching = false
@@ -418,10 +426,32 @@ public final class AcpmuxSessionSource: AgentSessionSource {
 
     public func events() -> AsyncStream<AgentSessionEvent> {
         let (stream, continuation) = AsyncStream<AgentSessionEvent>.makeStream(bufferingPolicy: .unbounded)
+        if replay.requiresResync {
+            guard let hub else {
+                continuation.yield(.clientError("Reconnect to load this chat."))
+                continuation.finish()
+                return stream
+            }
+            let source = hub.source(sessionID: sessionID, summary: summary)
+            let upstream = source.events()
+            let task = Task {
+                // The proxy reader owns its fresh source until cancellation; the hub's
+                // registry remains weak, so closing a chat releases its replay too.
+                defer { withExtendedLifetime(source) {} }
+                for await event in upstream { continuation.yield(event) }
+                continuation.finish()
+            }
+            continuation.onTermination = { _ in task.cancel() }
+            return stream
+        }
         // What this source already has is history to a new reader: one change, drawn in place.
         flushOpenChunk()
-        if !replay.isEmpty { continuation.yield(.batch(replay)) }
-        continuations.append(continuation)
+        if !replay.events.isEmpty { continuation.yield(.batch(replay.events)) }
+        let id = UUID()
+        continuations[id] = continuation
+        continuation.onTermination = { [weak self] _ in
+            Task { @MainActor in self?.continuations.removeValue(forKey: id) }
+        }
         return stream
     }
 
@@ -1498,7 +1528,7 @@ public final class AcpmuxSessionSource: AgentSessionSource {
         history = nil
         guard !events.isEmpty else { return }
         for event in events { record(replay: event) }
-        for continuation in continuations { continuation.yield(.batch(events)) }
+        for continuation in continuations.values { continuation.yield(.batch(events)) }
     }
 
     private func emit(_ event: AgentSessionEvent) {
@@ -1512,15 +1542,11 @@ public final class AcpmuxSessionSource: AgentSessionSource {
             }
             record(replay: event)
         }
-        for continuation in continuations { continuation.yield(event) }
+        for continuation in continuations.values { continuation.yield(event) }
     }
 
-    /// Keeps `event` for a later reader. A streamed reply arrives as one chunk per few
-    /// tokens, and the buffer kept each one for as long as the chat stayed open (tens of
-    /// thousands of events for one long turn). Consecutive chunks of one message, or of
-    /// one thought, are kept as a single chunk with their text joined: `AgentChatStore`
-    /// appends a chunk of the same message to it, so a reader draws the same transcript
-    /// and the buffer grows with messages, not tokens.
+    /// Keeps `event` for a later reader while bounding the replay. Consecutive
+    /// streamed chunks are joined so the cache grows with messages, not tokens.
     private func record(replay event: AgentSessionEvent) {
         switch (openChunk, event) {
         case (.message(let id)?, .update(.agentMessageChunk(let text, let nextID))) where id == nextID:
@@ -1537,26 +1563,26 @@ public final class AcpmuxSessionSource: AgentSessionSource {
                 openChunk = .thought
                 openChunkText = text
             default:
-                replay.append(event)
+                replay.append([event])
             }
         }
     }
 
-    /// The chunk being joined, kept out of `replay` so its text grows in place.
     private enum OpenChunk { case message(String?), thought }
     private var openChunk: OpenChunk?
     private var openChunkText = ""
 
     private func flushOpenChunk() {
         guard let chunk = openChunk else { return }
         switch chunk {
-        case .message(let id): replay.append(.update(.agentMessageChunk(text: openChunkText, messageId: id)))
-        case .thought: replay.append(.update(.agentThoughtChunk(text: openChunkText)))
+        case .message(let id): replay.append([.update(.agentMessageChunk(text: openChunkText, messageId: id))])
+        case .thought: replay.append([.update(.agentThoughtChunk(text: openChunkText))])
         }
         openChunk = nil
         openChunkText = ""
     }
 
     /// How many events the source keeps for a later reader (tests).
-    var replayCount: Int { replay.count + (openChunk == nil ? 0 : 1) }
+    var replayCount: Int { replay.events.count + (openChunk == nil ? 0 : 1) }
+
 }

cmux-next-ios/CmuxMobileKit/Sources/CmuxMobileLive/AcpmuxReplay.swift

@@ -0,0 +1,18 @@
+/// A replay is useful only while it is complete. Once its budget is exceeded,
+/// release it and require a fresh snapshot; keeping only a tail would apply
+/// deltas without their turn, permission request or earlier message chunks.
+struct AcpmuxReplay<Event> {
+    let limit: Int
+    private(set) var events: [Event] = []
+    private(set) var requiresResync = false
+
+    mutating func append(_ batch: [Event]) {
+        guard !requiresResync else { return }
+        guard batch.count <= limit - events.count else {
+            events = []
+            requiresResync = true
+            return
+        }
+        events.append(contentsOf: batch)
+    }
+}