cloud-setup: connect Codex broker to app lifetime

Packages/macOS/CmuxNext/Sources/CmuxNextApp/Cloud/CloudService.swift

@@ -24,6 +24,8 @@ final class CloudService {
     /// The machines a `cloud deliver-secrets` runs for now: one delivery per machine at a time.
     @ObservationIgnored var deliveringSecrets: Set<String> = []
     @ObservationIgnored let paths: CloudPaths
+    /// One app-lifetime owner of confirmed, per-account Codex delivery.
+    let codexBroker: CloudCodexBrokerAdapter
     @ObservationIgnored private let work: BackgroundWork
     @ObservationIgnored private(set) var hub: CloudTunnelHub?
     /// This Mac's Cloud install principal (cx-wb5.64): install tokens for the
@@ -112,6 +114,29 @@ final class CloudService {
             tokens: { try await auth.tokens() },
             teamID: { await auth.teamID }
         )
+        let api = api
+        let build = LaunchIdentity.current().build
+        let bin = Bundle.main.resourceURL?.appendingPathComponent("bin", isDirectory: true)
+        let environment = ProcessInfo.processInfo.environment
+        let userHome = FileManager.default.homeDirectoryForCurrentUser
+        codexBroker = CloudCodexBrokerAdapter(root: paths.root, registry: CloudCodexBrokerRegistry(),
+            currentScope: {
+                guard auth.isSignedIn, let user = auth.user else { return nil }
+                return CloudCodexBrokerScope(userID: user.id, teamID: auth.teamID)
+            },
+            launch: {
+                try await work.run(.userVisible, mayBlock: .files) {
+                    guard let acpmux = AcpmuxEnvironment.resolve(build: build, bundledBinDirectory: bin,
+                        environment: environment, userHome: userHome) else { throw CloudCodexBrokerError.brokerUnavailable }
+                    let search = (environment["PATH"] ?? "").split(separator: ":").map(String.init)
+                        + [userHome.appendingPathComponent(".local/bin").path, "/opt/homebrew/bin", "/usr/local/bin"]
+                    guard let codex = search.filter { $0.hasPrefix("/") }.map({ URL(fileURLWithPath: $0).appendingPathComponent("codex") })
+                        .first(where: { FileManager.checked.isExecutableFile(atPath: $0.path) }) else {
+                        throw CloudCodexBrokerError.brokerUnavailable
+                    }
+                    return CloudCodexBrokerAdapter.Launch(executable: acpmux.executable, prefix: [], codex: codex)
+                }
+            }, target: { scope, machineID in CloudCodexBrokerTarget(api: api, machineID: machineID, scope: scope) })
         installIdentity = MacInstallIdentity(
             store: .forApp(directory: paths.root.appendingPathComponent("install", isDirectory: true),
                            service: "\(configuration.bundleID ?? "cmux").install-key.\(configuration.ownerAPIBaseURL().host ?? "unknown")",
@@ -259,6 +284,7 @@ final class CloudService {
             var last = account
             for await account in ObservationStream({ self.account }) where account != last {
                 last = account
+                await codexBroker.scopeDidChange()
                 startSessionRuns.forgetAll()
                 startSessionAsks.endAll()
             }
@@ -274,6 +300,10 @@ final class CloudService {
     func applyPolicy(disabled: Bool) {
         guard disabled != policyDisabled else { return }
         policyDisabled = disabled
+        // Fence Codex broker requests synchronously with the policy flip. The
+        // adapter tears down holders asynchronously and retains their private
+        // homes for a later confirmed run after policy is restored.
+        codexBroker.setAvailable(!disabled)
         if disabled {
             for session in machines.cloud {
                 session.daemon.policyBlock.set(true)
@@ -309,6 +339,10 @@ final class CloudService {
         if let hub { Task { await hub.stop() } }
     }
 
+    /// The quit coordinator awaits this before AppKit may terminate the
+    /// process, while Cloud transports can still erase delivered access.
+    func prepareForTermination() async { await codexBroker.shutdown() }
+
     /// Serves the `credential` family on the local daemon: answers its
     /// calls on the connection that sent them, and registers again after
     /// every handshake of the verified app connection (a reconnected
@@ -423,6 +457,12 @@ final class CloudService {
         }
         let ids = Set(visible.map(\.id))
         for session in machines.cloud where !ids.contains(session.machineID) {
+            if let user = auth.user {
+                let scope = CloudCodexBrokerScope(userID: user.id, teamID: auth.teamID)
+                let machineID = session.machineID
+                // task-owner: server-driven removal; the adapter stops renewal before cleanup.
+                Task { try? await codexBroker.revoke(machineID: machineID, scope: scope) }
+            }
             creations.machineRemoved(session.machineID)
             machines.remove(session.machineID)?.disconnect()
         }
@@ -560,7 +600,10 @@ final class CloudService {
     static let inProgressDelay: Duration = .seconds(15)
 
     func deleteMachine(_ machineID: String) async throws {
+        let brokerScope = auth.user.map { CloudCodexBrokerScope(userID: $0.id, teamID: auth.teamID) }
+        if let brokerScope { try? await codexBroker.revoke(machineID: machineID, scope: brokerScope) }
         try await api.deleteMachine(machineID)
+        if let brokerScope { await codexBroker.machineWasDeleted(machineID, scope: brokerScope) }
         // A list in flight was asked while this machine still existed.
         listGeneration += 1
         creations.machineRemoved(machineID)
@@ -615,6 +658,7 @@ final class CloudService {
         startSessionRuns.forgetAll()
         startSessionAsks.endAll()
         for hook in willSignOut { hook() }
+        await codexBroker.signOut()
         AgentPaneTailCache.shared.clear() // the panes' chat tails: memory only, forgotten whole
         // No ensure starts from here until the session's final endLink (the
         // signed-out branch above, after auth.signOut).

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

@@ -164,6 +164,7 @@ final class QuitCoordinator {
         // Awaited (bounded) so `quitting` is on disk before the quit can end
         // the process; the write itself runs off the main thread.
         await services.crashRecovery.quitBegan()
+        await services.cloud.prepareForTermination()
         await QuitCompletion.run(choice, remember: remember, QuitSteps(
             remember: { behavior in
                 guard let settings = services.settings,

Packages/macOS/CmuxNext/Sources/CmuxNextCloud/Broker/CloudCodexBroker.swift

@@ -15,6 +15,7 @@ public actor CloudCodexBroker {
     private let codexExecutable: URL
     private let home: URL
     private var session: CodexBrokerSession?
+    private var stopped = false
     private var cached: CloudCodexAccessGrant?
     private var previousModeByMachine: [String: Data?] = [:]
     private var loadedPreviousMode: Set<String> = []
@@ -25,11 +26,15 @@ public actor CloudCodexBroker {
     private var targets: [String: CloudCodexBrokerTarget] = [:]
     private var targetGenerations: [String: Int] = [:]
     private var pendingRevocations: Set<String> = []
+    private var revocationTargets: [String: CloudCodexBrokerTarget] = [:]
     private var targetLocks: Set<String> = []
     private var targetWaiters: [String: [CheckedContinuation<Void, Never>]] = [:]
     private var refreshTask: Task<Void, Never>?
     private var refreshCycle = 0
     private var refreshInFlight: Task<CloudCodexAccessGrant, any Error>?
+    private var quitWaiters: [UUID: CheckedContinuation<Void, any Error>] = [:]
+    private var quitTimers: [UUID: Task<Void, Never>] = [:]
+    private var quitWork: [UUID: Task<Void, Never>] = [:]
 
     public init(scope: CloudCodexBrokerScope,
                 currentScope: @escaping @Sendable () async -> CloudCodexBrokerScope?,
@@ -51,7 +56,9 @@ public actor CloudCodexBroker {
     /// cmux and then calls `waitForLogin` without opening a VM terminal.
     public func login(target: CloudCodexBrokerTarget, confirmed: Bool) async throws -> CloudCodexLoginPrompt {
         try await validate(target: target, confirmed: confirmed)
+        try await ensureScope()
         let response = try await request(method: "login")
+        try await ensureScope()
         guard let urlString = response.verificationURL,
               let url = URL(string: urlString), url.scheme == "https",
               let code = response.userCode, !code.isEmpty, code.utf8.count < 64 else {
@@ -124,6 +131,7 @@ public actor CloudCodexBroker {
 
     private func performRefresh() async throws -> CloudCodexAccessGrant {
         let response = try await request(method: "access", force: true, previousAccountID: cached?.chatgptAccountID)
+        try await ensureScope()
         guard let token = response.accessToken, !token.isEmpty,
               let account = response.chatgptAccountID, !account.isEmpty,
               let expiry = response.expiresAt else { throw CloudCodexBrokerError.invalidReply }
@@ -141,6 +149,7 @@ public actor CloudCodexBroker {
         targetGenerations[target.machineID, default: 0] += 1
         guard registered else { return }
         pendingRevocations.insert(target.machineID)
+        revocationTargets[target.machineID] = target
         await lockTarget(target.machineID)
         defer { unlockTarget(target.machineID) }
         if let stages = pendingStageFiles[target.machineID] {
@@ -155,6 +164,7 @@ public actor CloudCodexBroker {
             loadedPreviousMode.remove(target.machineID)
             accessRemovalAttempted.remove(target.machineID)
             pendingRevocations.remove(target.machineID)
+            revocationTargets.removeValue(forKey: target.machineID)
             stopRenewalIfEmpty()
             if signOut, targets.isEmpty { _ = try? await request(method: "logout") }
             return
@@ -189,6 +199,7 @@ public actor CloudCodexBroker {
         previousModeByMachine.removeValue(forKey: target.machineID)
         loadedPreviousMode.remove(target.machineID)
         pendingRevocations.remove(target.machineID)
+        revocationTargets.removeValue(forKey: target.machineID)
         if targets.isEmpty {
             cached = nil
             refreshTask?.cancel()
@@ -202,7 +213,24 @@ public actor CloudCodexBroker {
         cached = nil; refreshTask?.cancel(); refreshTask = nil
     }
 
+    /// The server has confirmed deletion; there is no VM file left to restore.
+    /// This is not a substitute for revocation of an existing machine.
+    public func discardDeletedTarget(_ machineID: String) {
+        targetGenerations[machineID, default: 0] += 1
+        targets.removeValue(forKey: machineID)
+        revocationTargets.removeValue(forKey: machineID)
+        pendingRevocations.remove(machineID)
+        previousModeByMachine.removeValue(forKey: machineID)
+        installedAccess.removeValue(forKey: machineID)
+        uncertainInstallations.removeValue(forKey: machineID)
+        pendingStageFiles.removeValue(forKey: machineID)
+        accessRemovalAttempted.remove(machineID)
+        loadedPreviousMode.remove(machineID)
+        if targets.isEmpty { cached = nil; refreshTask?.cancel(); refreshTask = nil }
+    }
+
     private func validate(target: CloudCodexBrokerTarget, confirmed: Bool) async throws {
+        guard !stopped else { throw CloudCodexBrokerError.cancelled }
         guard confirmed else { throw CloudCodexBrokerError.consentRequired }
         guard !pendingRevocations.contains(target.machineID) else { throw CloudCodexBrokerError.cancelled }
         guard await currentScope() == scope else { throw CloudCodexBrokerError.scopeChanged }
@@ -215,6 +243,7 @@ public actor CloudCodexBroker {
     }
 
     private func request(method: String, force: Bool? = nil, previousAccountID: String? = nil) async throws -> BrokerResult {
+        guard !stopped || method == "logout" else { throw CloudCodexBrokerError.cancelled }
         if session == nil {
             session = try CodexBrokerSession(executable: executable, executablePrefix: executablePrefix, codexExecutable: codexExecutable, home: home)
         }
@@ -330,6 +359,7 @@ public actor CloudCodexBroker {
                 targets.removeValue(forKey: target.machineID)
                 targetGenerations[target.machineID, default: 0] += 1
                 pendingRevocations.insert(target.machineID)
+                revocationTargets[target.machineID] = target
             }
             let previous = previousModeByMachine[target.machineID] ?? nil
             if committed { accessRemovalAttempted.insert(target.machineID) }
@@ -375,12 +405,68 @@ public actor CloudCodexBroker {
     }
 
     private func ensureScope() async throws {
+        guard !stopped else { throw CloudCodexBrokerError.cancelled }
         guard await currentScope() == scope else { throw CloudCodexBrokerError.scopeChanged }
         guard !Task.isCancelled else { throw CloudCodexBrokerError.cancelled }
     }
 
     func registeredMachineIDs() -> Set<String> { Set(targets.keys) }
 
+    /// Stops the local process and renewal before erasing registered VM files.
+    /// Returns failed cleanup IDs; retained local credentials allow reuse later.
+    public func shutdown() async -> Set<String> {
+        stopped = true
+        refreshTask?.cancel()
+        refreshInFlight?.cancel()
+        let endingTargets = Array(targets.merging(revocationTargets) { _, pending in pending }.values)
+        if let session { await invalidate(session) }
+        var failed = pendingRevocations
+        await withTaskGroup(of: (String, Bool).self) { group in
+            for target in endingTargets {
+                group.addTask { [weak self] in
+                    guard let self else { return (target.machineID, false) }
+                    let succeeded: Bool
+                    do { try await self.revokeWithinQuitBudget(target); succeeded = true }
+                    catch { succeeded = false }
+                    return (target.machineID, succeeded)
+                }
+            }
+            for await (machineID, succeeded) in group {
+                if succeeded { failed.remove(machineID) } else { failed.insert(machineID) }
+            }
+        }
+        cached = nil
+        return failed
+    }
+
+    private func revokeWithinQuitBudget(_ target: CloudCodexBrokerTarget) async throws {
+        let id = UUID()
+        try await withCheckedThrowingContinuation { continuation in
+            quitWaiters[id] = continuation
+            // A timed-out transport may finish later. Its broker stays retained,
+            // preventing replacement until that cleanup has actually completed.
+            quitWork[id] = Task {
+                do { try await revoke(target: target); finishQuitCleanup(id, .success(())) }
+                catch { finishQuitCleanup(id, .failure(error)) }
+                quitWork.removeValue(forKey: id)
+            }
+            quitTimers[id] = Task {
+                // wakeup-allow: one bounded quit-cleanup deadline per target.
+                do { try await Task.sleep(for: .seconds(3)) } catch { return }
+                // Bound the caller's wait, not the remote mutation: cancellation
+                // could release its lock before the server finishes erasing.
+                finishQuitCleanup(id, .failure(CloudCodexBrokerError.remoteDeliveryFailed))
+            }
+        }
+    }
+
+    @discardableResult private func finishQuitCleanup(_ id: UUID, _ result: Result<Void, any Error>) -> Bool {
+        guard let waiter = quitWaiters.removeValue(forKey: id) else { return false }
+        quitTimers.removeValue(forKey: id)?.cancel()
+        waiter.resume(with: result)
+        return true
+    }
+
     private func targetList() -> [CloudCodexBrokerTarget] { targets.values.sorted { $0.machineID < $1.machineID } }
 
     private func scheduleRefresh(grant: CloudCodexAccessGrant) {
@@ -451,6 +537,7 @@ public actor CloudCodexBroker {
 /// passed to transfer per call, while the rotating local refresh cache remains one per scope.
 public actor CloudCodexBrokerRegistry {
     private var brokers: [CloudCodexBrokerScope: CloudCodexBroker] = [:]
+    private var pendingCleanup: [CloudCodexBrokerScope: CloudCodexBroker] = [:]
 
     public init() {}
 
@@ -459,16 +546,36 @@ public actor CloudCodexBrokerRegistry {
                        executable: URL,
                        executablePrefix: [String] = [],
                        codexExecutable: URL,
-                       home: URL) -> CloudCodexBroker {
+                       home: URL) async throws -> CloudCodexBroker {
+        if let previous = pendingCleanup[scope] {
+            guard await previous.shutdown().isEmpty else { throw CloudCodexBrokerError.remoteDeliveryFailed }
+            pendingCleanup.removeValue(forKey: scope)
+        }
         if let broker = brokers[scope] { return broker }
         let broker = CloudCodexBroker(scope: scope, currentScope: currentScope, executable: executable,
                                       executablePrefix: executablePrefix, codexExecutable: codexExecutable, home: home)
         brokers[scope] = broker
         return broker
     }
 
-    public func remove(scope: CloudCodexBrokerScope) {
-        brokers.removeValue(forKey: scope)
+    /// Returns the existing holder without creating or signing in to one.
+    public func existing(for scope: CloudCodexBrokerScope) -> CloudCodexBroker? { brokers[scope] }
+
+    func hasPendingCleanup(for scope: CloudCodexBrokerScope) -> Bool { pendingCleanup[scope] != nil }
+
+    /// Called only after the service successfully deletes this VM.
+    public func discardDeletedTarget(_ machineID: String, scope: CloudCodexBrokerScope) async {
+        await brokers[scope]?.discardDeletedTarget(machineID)
+        await pendingCleanup[scope]?.discardDeletedTarget(machineID)
+    }
+
+    /// Stops a scoped holder before releasing it. Credentials remain on disk.
+    public func remove(scope: CloudCodexBrokerScope) async -> Set<String> {
+        guard let broker = brokers.removeValue(forKey: scope) ?? pendingCleanup[scope] else { return [] }
+        pendingCleanup[scope] = broker
+        let failed = await broker.shutdown()
+        if failed.isEmpty { pendingCleanup.removeValue(forKey: scope) }
+        return failed
     }
 }