cmux NEXT: a chat's composer chips and shell command come back after a relaunch

Packages/macOS/CmuxNext/Sources/CmuxNextAgentPane/AgentPaneHandshakeReply.swift

@@ -64,7 +64,7 @@ enum AgentPaneHandshakeReply {
             // Only a fresh page reads it; a reconnect keeps it for the page's own ready.
             if request == .ready {
                 handshake.wake = model.pendingWake
-                model.pendingWake = nil
+                (model.handedWake, model.pendingWake) = (model.pendingWake, nil)
             }
             model.hasHandshake = true
             // A filled-in folder stays the workspace's, no root of its own (cx-9aps).

Packages/macOS/CmuxNext/Sources/CmuxNextAgentPane/AgentPaneModel.swift

@@ -90,6 +90,8 @@ public final class AgentPaneModel {
     /// What this chat's last page held when it let go (hibernation); the next handshake
     /// carries it (`wake`) and clears it.
     @ObservationIgnored public var pendingWake: AgentPaneWakeState?
+    /// The state the page's last `ready` took, kept for a quit the page does not answer in time.
+    @ObservationIgnored public internal(set) var handedWake: AgentPaneWakeState?
     /// Whether the page has asked for a handshake, so its bridge is up and
     /// a turn can be revealed through it directly.
     @ObservationIgnored public internal(set) var hasHandshake = false

Packages/macOS/CmuxNext/Sources/CmuxNextAgentPane/AgentPaneRelaunchWake.swift

@@ -0,0 +1,86 @@
+public import Foundation
+
+/// What each agent chat's page held that acpmux does not (``AgentPaneWakeState``: the composer's
+/// attachments, shell mode with its unrun command, where the transcript rested), by tab, in a JSON
+/// file in the tag's state directory. The quit's window step saves it (`QuitCoordinator`), so every
+/// relaunch that goes through the quit (Restart to Update, SIGTERM's Quit and keep sessions,
+/// Cmd-Q) keeps it; the next launch hands it to each tab's first page with its handshake, as a
+/// woken hibernated page gets it. The draft's text is acpmux's and a new chat's is
+/// ``AgentUnsentDrafts``'.
+public final class AgentPaneRelaunchWake {
+    public static let fileName = "agent-pane-wake-states.json"
+    /// Attachments carry their contents: a tab holding more is not kept, and the file stops growing
+    /// past ``maximumBytes`` (the launch reads it on the main actor; a larger file is not read).
+    nonisolated static let maximumTabBytes = 4 * 1024 * 1024
+    nonisolated static let maximumBytes = 8 * 1024 * 1024
+
+    let file: URL?
+    /// Asks a live page what it holds, within the limit (``AgentPaneView/wakeState(limit:)``).
+    public var read: @MainActor (AgentPaneView, Duration) async -> AgentPaneWakeState? = { await $0.wakeState(limit: $1) }
+
+    public init(file: URL?) {
+        self.file = file
+    }
+
+    /// The file in the tag's state directory (beside the launch snapshot).
+    public convenience init(directory: URL?) {
+        self.init(file: directory?.appending(path: Self.fileName))
+    }
+
+    /// The last quit's states, by tab, once: the file goes, so a crash later in this run never
+    /// brings back chips already sent.
+    public func take() -> [String: AgentPaneWakeState] {
+        guard let file else { return [:] }
+        defer { try? FileManager.default.removeItem(at: file) }
+        guard let size = (try? file.resourceValues(forKeys: [.fileSizeKey]))?.fileSize, size <= Self.maximumBytes,
+              // concurrency-allow: one capped read of a small local state file when the tab store is made
+              let data = try? Data(contentsOf: file),
+              let states = try? JSONDecoder().decode([String: AgentPaneWakeState].self, from: data) else { return [:] }
+        return states.filter { !$0.value.isEmpty }
+    }
+
+    /// What every chat holds for the next launch, by tab: `held` (released pages' states, by tab),
+    /// then each page in `pages`. A page that has not taken its state yet (no `ready`, only a
+    /// reconnect) keeps it unasked; every other page is asked at once, each within `limit`, and
+    /// one that does not answer keeps the tab's held state or the one its `ready` took. Tabs
+    /// `live` refuses (gone from every tree) keep nothing.
+    public func collect(held: [String: AgentPaneWakeState], pages: [(key: String, view: AgentPaneView)],
+                        live: (String) -> Bool, limit: Duration = .seconds(1)) async -> [String: AgentPaneWakeState] {
+        var states = held
+        let read = read
+        // task-owner: every ask is awaited below, before the quit goes on.
+        let asks = pages.map { page in
+            (page, page.view.model.pendingWake == nil && page.view.model.hasHandshake ? Task { await read(page.view, limit) } : nil)
+        }
+        for (page, ask) in asks {
+            let model = page.view.model
+            let answer = await ask?.value
+            states[page.key] = answer ?? model.pendingWake ?? states[page.key] ?? model.handedWake
+        }
+        return states.filter { key, state in !state.isEmpty && live(key) }
+    }
+
+    /// Replaces the saved states with `states` (by tab key); with none held the file goes.
+    public func save(_ states: [String: AgentPaneWakeState]) async {
+        guard let file else { return }
+        let held = states.filter { !$0.value.isEmpty }
+        await Task.detached { Self.write(held, to: file) }.value
+    }
+
+    private nonisolated static func write(_ states: [String: AgentPaneWakeState], to file: URL) {
+        let encoder = JSONEncoder()
+        var kept: [String: AgentPaneWakeState] = [:]
+        var total = 0
+        // Smallest first, so one tab's large attachments never cost every other tab its composer.
+        let sized = states.compactMap { key, state in (try? encoder.encode(state)).map { (key, state, $0.count) } }
+        for (key, state, size) in sized.sorted(by: { $0.2 < $1.2 }) where size <= maximumTabBytes && total + size <= maximumBytes {
+            kept[key] = state
+            total += size
+        }
+        guard !kept.isEmpty, let data = try? encoder.encode(kept) else {
+            try? FileManager.default.removeItem(at: file)
+            return
+        }
+        try? data.write(to: file, options: .atomic)
+    }
+}

Packages/macOS/CmuxNext/Sources/CmuxNextApp/AgentUnsentDrafts.swift → Packages/macOS/CmuxNext/Sources/CmuxNextAgentPane/AgentUnsentDrafts.swift

@@ -1,40 +1,38 @@
-import CmuxNextAgentPane
-import CmuxNextDaemon
-import Foundation
+public import Foundation
 
 /// A new chat's typed text from before it has an acpmux session (cx-z5vd), by tab, in a JSON file
 /// beside the launch snapshot: after a relaunch that tab opens as the chat with its draft, not as
 /// the New Tab page. Once the chat has a session acpmux keeps its draft, and this copy is gone.
-final class AgentUnsentDrafts {
-    static let fileName = "agent-unsent-drafts.json"
+public final class AgentUnsentDrafts {
+    public static let fileName = "agent-unsent-drafts.json"
     /// A longer text is not kept (acpmux's own draft limit).
     static let maximumCharacters = 1_000_000
 
     let file: URL?
     private var drafts: [String: String]
 
-    init(file: URL?) {
+    public init(file: URL?) {
         self.file = file
         // concurrency-allow: one read of a small local state file when the tab store is made
         let data = file.flatMap { try? Data(contentsOf: $0) }
         drafts = data.flatMap { try? JSONDecoder().decode([String: String].self, from: $0) } ?? [:]
     }
 
-    /// The file next to the launch snapshot of `snapshot` (the tag's state).
-    convenience init(beside snapshot: SidebarSnapshotFile?) {
-        self.init(file: snapshot?.url.deletingLastPathComponent().appending(path: Self.fileName))
+    /// The file in the tag's state directory (beside the launch snapshot).
+    public convenience init(directory: URL?) {
+        self.init(file: directory?.appending(path: Self.fileName))
     }
 
     /// Tab `key`'s kept text as a new chat's seed; nil for a tab whose chat has a session, or
     /// (`local` false) one on another machine.
-    func seed(for key: String, session: String?, local: Bool = true) -> AgentPaneSeedSource? {
+    public func seed(for key: String, session: String?, local: Bool = true) -> AgentPaneSeedSource? {
         guard local, session == nil, let text = drafts[key] else { return nil }
         return AgentPaneSeedSource(AgentPaneSeed(draft: text))
     }
 
     /// Keeps the typed text of `model`'s chat for the tab `key` names when it runs, until the chat
     /// has a session. Call it after the tab's own `onSessionChange` is set.
-    func wire(_ model: AgentPaneModel, key: @escaping () -> String) {
+    public func wire(_ model: AgentPaneModel, key: @escaping () -> String) {
         model.onUnsentDraft = { [weak self] text in self?.set(text, for: key()) }
         let started = model.onSessionChange
         model.onSessionChange = { [weak self] session in
@@ -44,10 +42,10 @@ final class AgentUnsentDrafts {
     }
 
     /// Whether tab `key` holds a new chat's typed text (its close earns an undo toast).
-    func hasDraft(_ key: String) -> Bool { drafts[key] != nil }
+    public func hasDraft(_ key: String) -> Bool { drafts[key] != nil }
 
     /// Keeps `text` as tab `key`'s draft; blank text forgets it.
-    func set(_ text: String, for key: String) {
+    public func set(_ text: String, for key: String) {
         let keep = !text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty && text.count <= Self.maximumCharacters
         let value = keep ? text : nil
         guard drafts[key] != value else { return }
@@ -56,7 +54,7 @@ final class AgentUnsentDrafts {
     }
 
     /// The store gave provisional tab `key` its id `real`.
-    func move(_ key: String, to real: String) {
+    public func move(_ key: String, to real: String) {
         guard key != real, let text = drafts.removeValue(forKey: key) else { return }
         drafts[real] = text
         save()

Packages/macOS/CmuxNext/Sources/CmuxNextApp/AgentTabs+Wiring.swift

@@ -111,7 +111,11 @@ extension AgentTabStore {
         let tabs = AgentTabStore(build: services.environment.build, registry: services.registry, environment: services.agentPaneEnvironment ?? ProcessInfo.processInfo.environment, showcase: services.environment.showcase,
                                  linkScheme: services.linkScheme, settings: services.settings)
         tabs.launchImages = AgentPaneLaunchImages(beside: services.environment.sidebarSnapshotFile)
-        tabs.unsentDrafts = AgentUnsentDrafts(beside: services.environment.sidebarSnapshotFile)
+        let state = services.environment.sidebarSnapshotFile?.url.deletingLastPathComponent()
+        tabs.unsentDrafts = AgentUnsentDrafts(directory: state)
+        // The composer chips and shell commands chat pages held at the last quit (every relaunch's quit).
+        tabs.relaunchWake = AgentPaneRelaunchWake(directory: state)
+        tabs.wakeStates = tabs.relaunchWake.take()
         // The Chief home's acpmux: the Chief host's subagent tabs attach there.
         let chief = ChiefHomeAcpmux(home: ChiefHome.resolve(tag: services.environment.tag)) { [weak services] in
             await MainActor.run { _ = services?.home.brain.ensure(asked: false) }

Packages/macOS/CmuxNext/Sources/CmuxNextApp/AgentTabs.swift

@@ -169,6 +169,10 @@ final class AgentTabStore {
     /// What each released chat page held that acpmux does not (``AgentPageHibernation``), for
     /// the page built when its tab shows again.
     var wakeStates: [String: AgentPaneWakeState] = [:]
+    /// What the chat pages held at the last quit, for each tab's first page at launch (``AgentPaneRelaunchWake``).
+    var relaunchWake = AgentPaneRelaunchWake(file: nil)
+    /// What every chat holds for the next launch (the quit's window step).
+    func relaunchWakeStates() async -> [String: AgentPaneWakeState] { await relaunchWake.collect(held: wakeStates, pages: views.map { (resolve($0.key), $0.value) }, live: isAgentTab) }
     /// The git link of tab `key`: the session host of the tab's own machine (AgentPaneGitReads.swift),
     /// so a Cloud or SSH chat reads its folder there; nil answers the page `native.not_connected`. Move to Floating Chat puts a session in Start Agent's panel.
     /// The upload for a chat tab's dropped or pasted files, nil when it takes none (``CloudUploads``).
@@ -306,8 +310,9 @@ final class AgentTabStore {
         // A New Tab page made now (Cmd-N's first page, a restored page) adopts the prewarmed spare
         // when there is one: no cold load and React mount on the open path (cx-ec99).
         // (Its folder rides on the page context, `newTab.cwd`; a first page's seed holds only that.)
+        // A page with a kept composer (``wakeStates``) gets its own page: the spare's handshake is past.
         if local, kind == .local, sessions[key] == nil, record.session == nil, let newTab = newTabPages[key],
-           seed?.holdsOnlyFolder ?? true, let spare = takeSpare?() {
+           seed?.holdsOnlyFolder ?? true, wakeStates[key] == nil, let spare = takeSpare?() {
             adoptSpare(spare, key: key, page: newTab.page)
             track(key, in: store)
             return spare

Packages/macOS/CmuxNext/Sources/CmuxNextApp/Quit/QuitCoordinator.swift

@@ -165,6 +165,9 @@ final class QuitCoordinator {
                 // Remote-terminal tabs keep their last screen for the
                 // placeholder after relaunch (data-model.md 1.4).
                 await services.remoteTerminals.saveSnapshots()
+                // Agent chats get their composer's chips, shell command and place back at the next
+                // launch (pages are asked at once, each within a second).
+                await services.agentTabs.relaunchWake.save(services.agentTabs.relaunchWakeStates())
                 // Agent panes on screen draw their last page at the next launch.
                 await services.agentTabs.launchImages.save(services.agentTabs.shownPageImages())
                 // Browser tabs reopen at their recorded page (before any