Fix codex sessions restoring into the wrong surface after reload (#4920)

.github/workflows/ci.yml

@@ -410,6 +410,7 @@ jobs:
             CmuxSettings
             CmuxSettingsUI
             CmuxSocketControl
+            CMUXAgentLaunch
           )
           # Pre-existing failure unrelated to package correctness: the catalog
           # intentionally reuses some UserDefaults keys across sections (e.g. a

CLI/cmux.swift

@@ -3030,6 +3030,14 @@ struct CMUXCLI {
             return
         }
 
+        if command == "__debug-tmux-compat-env" {
+            try debugDumpTmuxCompatEnvironment(
+                socketPath: resolvedSocketPath,
+                explicitPassword: socketPasswordArg
+            )
+            return
+        }
+
         if command == "omc" {
             try runOMC(
                 commandArgs: commandArgs,
@@ -17344,13 +17352,59 @@ struct CMUXCLI {
             setenv(envVar.key, envVar.value, 1)
         }
         if let focusedContext {
-            setenv("CMUX_WORKSPACE_ID", focusedContext.workspaceId, 1)
-            if let surfaceId = focusedContext.surfaceId, !surfaceId.isEmpty {
+            // The launcher's OWN surface (its inherited env, passed in as processEnvironment) is the
+            // agent's canonical identity; the focused pane is only a fallback. Stamping the focused
+            // pane here would desync CMUX_SURFACE_ID from the inherited CMUX_PANEL_ID and make agents
+            // (codex omx/teams) record + restore into the wrong surface after reload (#4920).
+            let identity = AgentSpawnIdentity().resolve(
+                ownWorkspaceId: processEnvironment["CMUX_WORKSPACE_ID"],
+                ownSurfaceId: processEnvironment["CMUX_SURFACE_ID"],
+                focusedWorkspaceId: focusedContext.workspaceId,
+                focusedSurfaceId: focusedContext.surfaceId
+            )
+            if let workspaceId = identity.workspaceId {
+                setenv("CMUX_WORKSPACE_ID", workspaceId, 1)
+            }
+            if let surfaceId = identity.surfaceId {
                 setenv("CMUX_SURFACE_ID", surfaceId, 1)
             }
         }
     }
 
+    /// Hidden `__debug-tmux-compat-env` seam: runs the real ``configureTmuxCompatEnvironment`` against
+    /// the live socket's focused context and prints the resolved CMUX_* identity, so an integration
+    /// test can assert the launcher stamps the launch surface (its own env) rather than the focused
+    /// pane (#4920). Not user-facing. The shim/tmux params here do not affect the id resolution.
+    func debugDumpTmuxCompatEnvironment(socketPath: String, explicitPassword: String?) throws {
+        var processEnvironment = ProcessInfo.processInfo.environment
+        // Resolve the focused context from the SAME socket this command was pointed at, not whatever
+        // CMUX_SOCKET_PATH the process happened to inherit.
+        processEnvironment["CMUX_SOCKET_PATH"] = socketPath
+        let focusedContext = try tmuxCompatFocusedContext(
+            processEnvironment: processEnvironment,
+            explicitPassword: explicitPassword
+        )
+        let shimDirectory = FileManager.default.temporaryDirectory
+            .appendingPathComponent("cmux-debug-tmux-shim-\(UUID().uuidString)", isDirectory: true)
+        try? FileManager.default.createDirectory(at: shimDirectory, withIntermediateDirectories: true)
+        configureTmuxCompatEnvironment(
+            processEnvironment: processEnvironment,
+            shimDirectory: shimDirectory,
+            executablePath: processEnvironment["CMUX_BUNDLED_CLI_PATH"] ?? "cmux",
+            socketPath: socketPath,
+            explicitPassword: explicitPassword,
+            focusedContext: focusedContext,
+            tmuxPathPrefix: "cmux-debug",
+            cmuxBinEnvVar: "CMUX_BIN",
+            termOverrideEnvVar: "TERM"
+        )
+        func dump(_ key: String) -> String { getenv(key).map { String(cString: $0) } ?? "" }
+        print("CMUX_WORKSPACE_ID=\(dump("CMUX_WORKSPACE_ID"))")
+        print("CMUX_SURFACE_ID=\(dump("CMUX_SURFACE_ID"))")
+        print("CMUX_PANEL_ID=\(dump("CMUX_PANEL_ID"))")
+        print("CMUX_TAB_ID=\(dump("CMUX_TAB_ID"))")
+    }
+
     private static let claudeNodeOptionsRestoreModule = """
     const hadOriginalNodeOptions = process.env.CMUX_ORIGINAL_NODE_OPTIONS_PRESENT === "1";
     if (hadOriginalNodeOptions) {
@@ -18219,11 +18273,21 @@ struct CMUXCLI {
         guard let focusedContext = try tmuxCompatFocusedContext(
             processEnvironment: launcherEnvironment,
             explicitPassword: explicitPassword
-        ),
-              let rootSurfaceId = focusedContext.surfaceId?.trimmingCharacters(in: .whitespacesAndNewlines),
-              !rootSurfaceId.isEmpty else {
+        ) else {
+            throw CLIError(message: "cmux codex-teams must be started from a cmux terminal surface")
+        }
+        // The codex-teams root identity is the LAUNCH surface (this process's own env), not the
+        // operator's focused pane, so the watcher records the surface codex actually runs in (#4920).
+        let rootIdentity = AgentSpawnIdentity().resolve(
+            ownWorkspaceId: launcherEnvironment["CMUX_WORKSPACE_ID"],
+            ownSurfaceId: launcherEnvironment["CMUX_SURFACE_ID"],
+            focusedWorkspaceId: focusedContext.workspaceId,
+            focusedSurfaceId: focusedContext.surfaceId
+        )
+        guard let rootSurfaceId = rootIdentity.surfaceId, !rootSurfaceId.isEmpty else {
             throw CLIError(message: "cmux codex-teams must be started from a cmux terminal surface")
         }
+        let rootWorkspaceId = rootIdentity.workspaceId ?? focusedContext.workspaceId
 
         let codexExecutablePath = resolveCodexExecutable(searchPath: launcherEnvironment["PATH"])
         let codexExecutableForShell = codexExecutablePath ?? "codex"
@@ -18319,7 +18383,7 @@ struct CMUXCLI {
             socketPath,
             "__codex-teams-watch",
             "--workspace-id",
-            focusedContext.workspaceId,
+            rootWorkspaceId,
             "--surface-id",
             rootSurfaceId,
             "--app-server-url",

Packages/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentSpawnIdentity.swift

@@ -0,0 +1,55 @@
+import Foundation
+
+/// Resolves the canonical workspace/surface identity to stamp into a spawned agent's environment.
+///
+/// cmux's CLI agent launchers (the `claude-teams` / `oh-my-codex` / `codex-teams` family) run inside a
+/// terminal surface and inherit that surface's own `CMUX_SURFACE_ID` / `CMUX_WORKSPACE_ID`. They also
+/// query the daemon for the operator's currently *focused* pane (for the tmux-compat shim). The agent's
+/// canonical identity must come from the launch surface, **not** the focused pane: stamping the focused
+/// pane desyncs `CMUX_SURFACE_ID` from the inherited `CMUX_PANEL_ID`, so an agent launched in surface B
+/// while surface A is focused records surface A and later restores into the wrong surface
+/// (https://github.com/manaflow-ai/cmux/issues/4920, the codex "jumble after reload" symptom).
+///
+/// This type is a stateless value; construct one at the call site (`AgentSpawnIdentity()`).
+///
+/// ```swift
+/// // Launcher runs in surface B while the operator's focus is on surface A:
+/// AgentSpawnIdentity().resolve(
+///     ownWorkspaceId: "WS-B", ownSurfaceId: "B",
+///     focusedWorkspaceId: "WS-A", focusedSurfaceId: "A"
+/// ) // == (workspaceId: "WS-B", surfaceId: "B")  — the launch surface, not the focused pane
+/// ```
+public struct AgentSpawnIdentity: Sendable, Equatable {
+    /// Creates a spawn-identity resolver. The type holds no state.
+    public init() {}
+
+    /// The canonical `(workspaceId, surfaceId)` to stamp, preferring the launcher's own identity.
+    ///
+    /// The launching process's own identity wins whenever present; the focused pane is only a fallback
+    /// for launches that have no own identity (e.g. started outside a cmux-managed terminal). Inputs are
+    /// trimmed and empty values are treated as absent.
+    ///
+    /// - Parameters:
+    ///   - ownWorkspaceId: the launching process's inherited `CMUX_WORKSPACE_ID` (the launch surface).
+    ///   - ownSurfaceId: the launching process's inherited `CMUX_SURFACE_ID` (the launch surface).
+    ///   - focusedWorkspaceId: the operator's currently focused workspace id (fallback only).
+    ///   - focusedSurfaceId: the operator's currently focused surface id (fallback only).
+    /// - Returns: the workspace/surface ids to stamp; either element is `nil` when neither source has it.
+    public func resolve(
+        ownWorkspaceId: String?,
+        ownSurfaceId: String?,
+        focusedWorkspaceId: String?,
+        focusedSurfaceId: String?
+    ) -> (workspaceId: String?, surfaceId: String?) {
+        let workspaceId = normalized(ownWorkspaceId) ?? normalized(focusedWorkspaceId)
+        let surfaceId = normalized(ownSurfaceId) ?? normalized(focusedSurfaceId)
+        return (workspaceId, surfaceId)
+    }
+
+    private func normalized(_ value: String?) -> String? {
+        guard let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
+            return nil
+        }
+        return trimmed
+    }
+}

Packages/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/AgentSpawnIdentityTests.swift

@@ -0,0 +1,58 @@
+import CMUXAgentLaunch
+import Testing
+
+@Suite("AgentSpawnIdentity")
+struct AgentSpawnIdentityTests {
+    // The bug: a codex launched in surface B while surface A is focused must keep surface B, not the
+    // focused surface A (which would desync from CMUX_PANEL_ID and restore into the wrong surface).
+    @Test("Prefers the launcher's own surface over the focused pane")
+    func prefersOwnOverFocused() {
+        let resolved = AgentSpawnIdentity().resolve(
+            ownWorkspaceId: "WS-B", ownSurfaceId: "B",
+            focusedWorkspaceId: "WS-A", focusedSurfaceId: "A"
+        )
+        #expect(resolved.workspaceId == "WS-B")
+        #expect(resolved.surfaceId == "B")
+    }
+
+    @Test("Falls back to the focused pane only when the launcher has no own identity")
+    func fallsBackToFocusedWhenNoOwn() {
+        let resolved = AgentSpawnIdentity().resolve(
+            ownWorkspaceId: nil, ownSurfaceId: nil,
+            focusedWorkspaceId: "WS-A", focusedSurfaceId: "A"
+        )
+        #expect(resolved.workspaceId == "WS-A")
+        #expect(resolved.surfaceId == "A")
+    }
+
+    @Test("Blank own identity is treated as absent")
+    func blankOwnIsAbsent() {
+        let resolved = AgentSpawnIdentity().resolve(
+            ownWorkspaceId: "   ", ownSurfaceId: "",
+            focusedWorkspaceId: "WS-A", focusedSurfaceId: "A"
+        )
+        #expect(resolved.workspaceId == "WS-A")
+        #expect(resolved.surfaceId == "A")
+    }
+
+    @Test("Per-element independence: own workspace + focused-only surface")
+    func perElementIndependence() {
+        // Own workspace present but no own surface, focused has both: workspace from own, surface from focused.
+        let resolved = AgentSpawnIdentity().resolve(
+            ownWorkspaceId: "WS-B", ownSurfaceId: nil,
+            focusedWorkspaceId: "WS-A", focusedSurfaceId: "A"
+        )
+        #expect(resolved.workspaceId == "WS-B")
+        #expect(resolved.surfaceId == "A")
+    }
+
+    @Test("No identity anywhere yields nil")
+    func noIdentity() {
+        let resolved = AgentSpawnIdentity().resolve(
+            ownWorkspaceId: nil, ownSurfaceId: nil,
+            focusedWorkspaceId: nil, focusedSurfaceId: nil
+        )
+        #expect(resolved.workspaceId == nil)
+        #expect(resolved.surfaceId == nil)
+    }
+}

cmuxTests/CLINotifyProcessIntegrationRegressionTests.swift

@@ -8416,3 +8416,79 @@ final class CLINotifyProcessIntegrationRegressionTests: XCTestCase {
     }
 
 }
+
+extension CLINotifyProcessIntegrationRegressionTests {
+    // E2E for #4920: the REAL CLI launcher env builder (configureTmuxCompatEnvironment, exercised via
+    // the hidden __debug-tmux-compat-env seam) must stamp the LAUNCH surface (the launcher's own
+    // inherited env), not the operator's focused pane returned by system.identify. Without the fix it
+    // stamped the focused surface (A), desyncing CMUX_SURFACE_ID from CMUX_PANEL_ID and jumbling codex
+    // into the wrong surface on reload.
+    func testTmuxCompatEnvStampsLaunchSurfaceNotFocusedPane() throws {
+        let cliPath = try bundledCLIPath()
+        let tmpDir = FileManager.default.temporaryDirectory
+            .appendingPathComponent("cmux-spawn-id-e2e-\(UUID().uuidString)", isDirectory: true)
+        try FileManager.default.createDirectory(at: tmpDir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: tmpDir) }
+        let socketPath = tmpDir.appendingPathComponent("sock").path
+        let listenerFD = try bindUnixSocket(at: socketPath)
+        defer { Darwin.close(listenerFD); unlink(socketPath) }
+        let state = MockSocketServerState()
+
+        // The operator's FOCUSED pane is surface A (what system.identify returns).
+        let focusedWorkspace = "11111111-1111-1111-1111-111111111111"
+        let focusedSurface = "22222222-2222-2222-2222-222222222222"
+        let handled = startMockServer(listenerFD: listenerFD, state: state) { line in
+            guard let payload = self.jsonObject(line),
+                  let id = payload["id"] as? String,
+                  let method = payload["method"] as? String else {
+                return self.malformedRequestResponse(raw: line)
+            }
+            if method == "system.identify" {
+                return self.v2Response(id: id, ok: true, result: [
+                    "focused": [
+                        "workspace_id": focusedWorkspace,
+                        "surface_id": focusedSurface,
+                        "pane_id": "%1",
+                    ],
+                ])
+            }
+            // resolveWorkspaceId / tmuxCanonicalPaneId fail gracefully (CLI uses try?).
+            return self.v2Response(id: id, ok: false, error: ["code": "unsupported", "message": method])
+        }
+
+        // ...but the launcher RUNS in surface B (its own inherited env). Tab id is surface-scoped, so
+        // it is distinct from the workspace id.
+        let launchWorkspace = "33333333-3333-3333-3333-333333333333"
+        let launchSurface = "44444444-4444-4444-4444-444444444444"
+        let launchTab = "55555555-5555-5555-5555-555555555555"
+        let result = runProcess(
+            executablePath: cliPath,
+            arguments: ["__debug-tmux-compat-env"],
+            environment: [
+                "CMUX_SOCKET_PATH": socketPath,
+                "CMUX_WORKSPACE_ID": launchWorkspace,
+                "CMUX_SURFACE_ID": launchSurface,
+                "CMUX_PANEL_ID": launchSurface,
+                "CMUX_TAB_ID": launchTab,
+                "HOME": tmpDir.path,
+                "PATH": ProcessInfo.processInfo.environment["PATH"] ?? "/usr/bin:/bin",
+            ],
+            timeout: 30
+        )
+        wait(for: [handled], timeout: 30)
+
+        XCTAssertTrue(
+            result.stdout.contains("CMUX_SURFACE_ID=\(launchSurface)"),
+            "launcher must stamp the LAUNCH surface; stdout:\n\(result.stdout)\nstderr:\n\(result.stderr)"
+        )
+        XCTAssertFalse(
+            result.stdout.contains("CMUX_SURFACE_ID=\(focusedSurface)"),
+            "launcher must NOT stamp the focused surface; stdout:\n\(result.stdout)"
+        )
+        XCTAssertTrue(result.stdout.contains("CMUX_WORKSPACE_ID=\(launchWorkspace)"), result.stdout)
+        // Matched-pair invariant: SURFACE == PANEL (the desync is exactly the bug). The surface-scoped
+        // tab id passes through untouched.
+        XCTAssertTrue(result.stdout.contains("CMUX_PANEL_ID=\(launchSurface)"), result.stdout)
+        XCTAssertTrue(result.stdout.contains("CMUX_TAB_ID=\(launchTab)"), result.stdout)
+    }
+}