Keep iOS Mac avatar color slots stable during Mac switches

Packages/iOS/CmuxMobileShell/Sources/CmuxMobileShell/MobileShellComposite+MacSwitchState.swift

@@ -72,4 +72,44 @@ extension MobileShellComposite {
     /// instead of `MobileShellComposite.swift` to respect that file's length
     /// budget.
     public var isMacSwitchInFlight: Bool { macSwitchAttemptID != nil }
+
+    /// Assign any newly seen real Mac a stable in-memory color slot.
+    ///
+    /// Called from `recomputeDerivedWorkspaceState()` before deriving
+    /// previews so a Mac switch's transient single-key `workspacesByMac`
+    /// state (old foreground dropped synchronously, new foreground's
+    /// siblings re-added asynchronously) never recomputes another Mac's
+    /// existing slot. Lives here (with `stableMacColorSlots` and
+    /// `workspaceAggregation` made internal) instead of
+    /// `MobileShellComposite.swift` to respect that file's length budget.
+    func updateStableMacColorSlots() {
+        let updated = workspaceAggregation.machineColorIndex(
+            existingAssignments: stableMacColorSlots,
+            adding: Array(workspacesByMac.keys.filter { !$0.isEmpty && $0 != Self.foregroundAnonymousKey })
+        )
+        if updated != stableMacColorSlots {
+            stableMacColorSlots = updated
+        }
+    }
+
+    /// Reset all stable color slots on sign-out. Slots are additive-only (see
+    /// `updateStableMacColorSlots` above), so the anonymous-placeholder reset
+    /// in `signOut()` never prunes the previous account's real Mac→color
+    /// assignments on its own. Call this last, after every didSet-triggering
+    /// assignment in `signOut()` has had its chance to recompute from the
+    /// (still-stale) previous `workspacesByMac`, so the next account starts
+    /// with a clean slate. Lives here instead of `MobileShellComposite.swift`
+    /// to respect that file's length budget.
+    func resetStableMacColorSlotsForSignOut() {
+        stableMacColorSlots = [:]
+    }
+
+    /// Prune stable color slots to only the foreground Mac on a team switch.
+    /// Slots are additive-only, so the old team's Macs would otherwise linger
+    /// in the slot map forever across repeated team switches; the new team's
+    /// Macs get reassigned lazily as they're re-aggregated. Lives here instead
+    /// of `MobileShellComposite.swift` to respect that file's length budget.
+    func pruneStableMacColorSlots(keepingForegroundKey foregroundKey: String) {
+        stableMacColorSlots = stableMacColorSlots.filter { $0.key == foregroundKey }
+    }
 }

Packages/iOS/CmuxMobileShell/Sources/CmuxMobileShell/MobileShellComposite.swift

@@ -243,7 +243,7 @@ public final class MobileShellComposite: MobileTerminalOutputSinking {
     var workspacesByMac: [String: MacWorkspaceState] = [:] {
         didSet { recomputeDerivedWorkspaceState() }
     }
-    private let workspaceAggregation = MobileWorkspaceAggregation()
+    let workspaceAggregation = MobileWorkspaceAggregation(); var stableMacColorSlots: [String: Int] = [:]  // see MobileShellComposite+MacSwitchState.swift
     /// The flat aggregated workspace list the UI renders. A materialized
     /// derivation of ``workspacesByMac``: only ``recomputeDerivedWorkspaceState``
     /// assigns it, so it is never independently mutated.
@@ -268,7 +268,7 @@ public final class MobileShellComposite: MobileTerminalOutputSinking {
     /// workspace avatars use), so the Computers screen can color each Mac's row to
     /// match its workspaces. Keyed by `macDeviceID`.
     public var machineColorIndex: [String: Int] {
-        workspaceAggregation.machineColorIndex(statesByMac: workspacesByMac)
+        stableMacColorSlots
     }
 
     public var macConnectionStatuses: [String: MobileMacConnectionStatus] {
@@ -1188,7 +1188,7 @@ public final class MobileShellComposite: MobileTerminalOutputSinking {
             workspaces: PreviewMobileHost.workspaces,
             groups: []
         )]
-        selectedWorkspaceID = workspaces.first?.id
+        resetStableMacColorSlotsForSignOut(); selectedWorkspaceID = workspaces.first?.id
         selectedTerminalID = workspaces.first?.terminals.first?.id
         // Selection resets above are done; allow draft saving again so a
         // subsequent sign-in restores drafts normally.
@@ -1218,7 +1218,7 @@ public final class MobileShellComposite: MobileTerminalOutputSinking {
         // on the next foreground / Computers `.task` / pull-to-refresh.
         teardownSecondaryMacSubscriptions()
         let foregroundKey = foregroundMacKey
-        workspacesByMac = workspacesByMac.filter { $0.key == foregroundKey }
+        workspacesByMac = workspacesByMac.filter { $0.key == foregroundKey }; pruneStableMacColorSlots(keepingForegroundKey: foregroundKey)
         // Restore memo: invalidate so the next read re-restores for the new
         // (account, team) scope, and a suspended old-team restore can't resume.
         // Invalidate the shared boundary synchronously first; actor cleanup is
@@ -3565,7 +3565,7 @@ public final class MobileShellComposite: MobileTerminalOutputSinking {
     /// source of truth. Pure and cheap; the only place those two are assigned,
     /// called on any ``workspacesByMac`` or foreground change.
     private func recomputeDerivedWorkspaceState() {
-        let previousSelection = selectedWorkspaceID.flatMap { id in
+        updateStableMacColorSlots(); let previousSelection = selectedWorkspaceID.flatMap { id in
             workspaces.first { $0.id == id }
         }
         let foregroundKey: String?
@@ -3577,7 +3577,7 @@ public final class MobileShellComposite: MobileTerminalOutputSinking {
             foregroundKey = nil
         }
         var derived = workspaceAggregation.derivedWorkspaces(
-            statesByMac: workspacesByMac, foregroundMacDeviceID: foregroundKey)
+            statesByMac: workspacesByMac, foregroundMacDeviceID: foregroundKey, machineColorIndex: stableMacColorSlots)
         // Stamp per-Mac user color/icon overrides from pairedMacs so every
         // workspace avatar matches its computer's customization (same place the
         // aggregation already assigned the automatic color index).

Packages/iOS/CmuxMobileShell/Tests/CmuxMobileShellTests/MobileShellCompositeColorSlotResetTests.swift

@@ -0,0 +1,66 @@
+import CMUXMobileCore
+import CmuxMobilePairedMac
+import CmuxMobileRPC
+import CmuxMobileShellModel
+import Foundation
+import Testing
+@testable import CmuxMobileShell
+
+/// Regression tests for `stableMacColorSlots` resets on the account/team
+/// boundaries. Split from ``MobileShellCompositePreviewTests`` to keep that
+/// file under the Swift file length budget.
+@MainActor
+@Suite struct MobileShellCompositeColorSlotResetTests {
+    @Test func signOutClearsStableMacColorSlots() {
+        let store = MobileShellComposite.preview()
+        store.signIn()
+        store.pairingCode = "debug"
+        store.connectPreviewHost()
+
+        // Seed a real (non-anonymous) Mac's color slot so sign-out has something
+        // to leak: without an explicit reset, `stableMacColorSlots` is
+        // additive-only, so the next account's slots would start numbering past
+        // this account's stale entry instead of at zero.
+        store.workspacesByMac["mac-previous-account"] = MacWorkspaceState(
+            macDeviceID: "mac-previous-account",
+            workspaces: []
+        )
+        #expect(store.machineColorIndex["mac-previous-account"] != nil)
+
+        store.signOut()
+
+        // The previous account's Mac→color assignment must not survive
+        // sign-out into the next account's session.
+        #expect(store.machineColorIndex.isEmpty)
+    }
+
+    @Test func teamSwitchDropsBackgroundMacColorSlotsButKeepsForeground() {
+        let store = MobileShellComposite.preview()
+        store.signIn()
+        store.pairingCode = "debug"
+        store.connectPreviewHost()
+        // A real (non-anonymous) foreground Mac, plus a background Mac that
+        // belongs to the team being switched away from.
+        store.foregroundMacDeviceID = "mac-foreground"
+        store.workspacesByMac["mac-foreground"] = MacWorkspaceState(
+            macDeviceID: "mac-foreground",
+            workspaces: []
+        )
+        store.workspacesByMac["mac-old-team"] = MacWorkspaceState(
+            macDeviceID: "mac-old-team",
+            workspaces: []
+        )
+        let foregroundIndexBefore = store.machineColorIndex["mac-foreground"]
+        #expect(foregroundIndexBefore != nil)
+        #expect(store.machineColorIndex["mac-old-team"] != nil)
+
+        store.currentTeamDidChange()
+
+        // The old team's background Mac must not survive the team switch
+        // (additive-only slots would otherwise leak it into the new team's
+        // numbering forever), while the foreground Mac's own slot is left
+        // intact so its color doesn't flash/change across the switch.
+        #expect(store.machineColorIndex["mac-old-team"] == nil)
+        #expect(store.machineColorIndex["mac-foreground"] == foregroundIndexBefore)
+    }
+}

Packages/iOS/CmuxMobileShellModel/Sources/CmuxMobileShellModel/MobileWorkspaceAggregation.swift

@@ -24,13 +24,26 @@ public struct MobileWorkspaceAggregation: Sendable {
         }.map(\.macDeviceID)
     }
 
-    /// A distinct stable color index per Mac, keyed by `macDeviceID`.
+    /// Return stable color index assignments after appending any newly seen Macs.
+    ///
+    /// Existing assignments are preserved verbatim. New non-empty Mac IDs are
+    /// processed in sorted order and assigned unique slots after the currently
+    /// assigned table, so cold-start assignment stays deterministic while a Mac
+    /// that was seen earlier keeps its slot across transient live-set changes.
     public func machineColorIndex(
-        statesByMac: [String: MacWorkspaceState]
+        existingAssignments: [String: Int],
+        adding macIDs: [String]
     ) -> [String: Int] {
-        var result: [String: Int] = [:]
-        for (offset, macID) in statesByMac.keys.filter({ !$0.isEmpty }).sorted().enumerated() {
-            result[macID] = offset
+        var result = existingAssignments
+        var usedSlots = Set(result.values)
+        var nextSlot = (usedSlots.max() ?? -1) + 1
+        for macID in Set(macIDs.filter { !$0.isEmpty }).sorted() where result[macID] == nil {
+            while usedSlots.contains(nextSlot) {
+                nextSlot += 1
+            }
+            result[macID] = nextSlot
+            usedSlots.insert(nextSlot)
+            nextSlot += 1
         }
         return result
     }
@@ -50,9 +63,9 @@ public struct MobileWorkspaceAggregation: Sendable {
     /// Derive the flat, ordered workspace list across all Macs.
     public func derivedWorkspaces(
         statesByMac: [String: MacWorkspaceState],
-        foregroundMacDeviceID: String?
+        foregroundMacDeviceID: String?,
+        machineColorIndex: [String: Int]
     ) -> [MobileWorkspacePreview] {
-        let colorIndex = machineColorIndex(statesByMac: statesByMac)
         let shouldScopeRowIDs = statesByMac.keys.filter { !$0.isEmpty }.count > 1
         var result: [MobileWorkspacePreview] = []
         for macID in orderedMacIDs(statesByMac: statesByMac, foregroundMacDeviceID: foregroundMacDeviceID) {
@@ -63,7 +76,7 @@ public struct MobileWorkspaceAggregation: Sendable {
                 if !ownerID.isEmpty {
                     stamped.macDeviceID = ownerID
                     stamped.macDisplayName = state.displayName
-                    stamped.machineColorIndex = colorIndex[ownerID]
+                    stamped.machineColorIndex = machineColorIndex[ownerID]
                 }
                 let remoteID = workspace.remoteWorkspaceID ?? workspace.id
                 stamped.remoteWorkspaceID = shouldScopeRowIDs && !ownerID.isEmpty ? remoteID : workspace.remoteWorkspaceID