diff --git a/CLI/CMUXCLI+AgentHookAdmission.swift b/CLI/CMUXCLI+AgentHookAdmission.swift index ea750d3976fd..2d15b86255a0 100644 --- a/CLI/CMUXCLI+AgentHookAdmission.swift +++ b/CLI/CMUXCLI+AgentHookAdmission.swift @@ -927,7 +927,10 @@ extension CMUXCLI { } static func queuedAgentHookDataEnvironmentKeys(agent: String) -> [String] { - [ + // The routed-launch metadata is Claude's; the session-start capture + // validates it before recording (`AgentLaunchEnvironmentPolicy`). + let routedLaunchKeys = agent == "claude" ? SubrouterClaudeResumeRouting.hookCapturedEnvironmentKeys : [] + return [ "PWD", "CMUX_AGENT_HOOK_STATE_DIR", "CMUX_AGENT_HOOK_SUPPRESS_VISIBLE_MUTATIONS", "CMUX_AGENT_LAUNCH_ARGV_B64", "CMUX_AGENT_LAUNCH_CWD", @@ -935,7 +938,7 @@ extension CMUXCLI { "CMUX_AGENT_MANAGED_SUBAGENT", "CMUX_SUPPRESS_SUBAGENT_NOTIFICATIONS", "CMUX_SURFACE_ID", "CMUX_WORKSPACE_ID", agentHookRouteSnapshotEnvironmentKey, agentHookPIDEnvironmentVariable(agentName: agent), - ] + ] + routedLaunchKeys } static func agentHookPIDEnvironmentVariable(agentName: String) -> String { diff --git a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentLaunchEnvironmentPolicy.swift b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentLaunchEnvironmentPolicy.swift index 78137ab71cbc..0b328eac20c5 100644 --- a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentLaunchEnvironmentPolicy.swift +++ b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentLaunchEnvironmentPolicy.swift @@ -193,9 +193,15 @@ public struct AgentLaunchEnvironmentPolicy: Sendable { // exact command text, never a URL or credential. They cross into // the durable restore record only as an agreeing pair, so a marker // inherited from an ancestor `sr claude` session proves nothing. - selected.merge(SubrouterClaudeResumeRouting().capturedEnvironment(in: env)) { _, marker in + let router = SubrouterClaudeResumeRouting() + selected.merge(router.capturedEnvironment(in: env)) { _, marker in marker } + // The account a routed launch was pinned to, recorded by the + // wrapper. It only picks the launcher's `--account` on restore. + selected.merge(router.capturedAccountEnvironment(in: env)) { _, account in + account + } } return selected } @@ -251,6 +257,7 @@ public struct AgentLaunchEnvironmentPolicy: Sendable { ) selected.removeValue(forKey: SubrouterCodexResumeRouting.environmentKey) selected.removeValue(forKey: SubrouterCodexResumeRouting.launchBoundEnvironmentKey) + selected.removeValue(forKey: SubrouterClaudeResumeRouting.accountEnvironmentKey) return selected } diff --git a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentRestorePlanner.swift b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentRestorePlanner.swift index 1486de18e154..152a2dd1e05c 100644 --- a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentRestorePlanner.swift +++ b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentRestorePlanner.swift @@ -437,6 +437,7 @@ public struct AgentRestorePlanner: Sendable { } selected.removeValue(forKey: SubrouterClaudeResumeRouting.environmentKey) selected.removeValue(forKey: SubrouterClaudeResumeRouting.launchBoundEnvironmentKey) + selected.removeValue(forKey: SubrouterClaudeResumeRouting.accountEnvironmentKey) let keys = selected.keys.sorted().filter { Self.claudeAuthSelectionEnvironmentKeys.contains($0) } diff --git a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentSessionRecoveryPlanner.swift b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentSessionRecoveryPlanner.swift index 0391e3e06244..789beea4c419 100644 --- a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentSessionRecoveryPlanner.swift +++ b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/AgentSessionRecoveryPlanner.swift @@ -195,3 +195,55 @@ public struct AgentSessionRecoveryPlanner: Sendable { return result } } + +/// Which recovered sessions start right away and which wait until their +/// workspace is first shown. +/// +/// A heavy user can have dozens of agents running when cmux dies, and starting +/// them all on relaunch spikes CPU and memory. Only a few start now: sessions +/// from a workspace on screen, then the most recently active. The rest open +/// their workspace now and resume on its first visit, the way startup restore +/// treats background workspaces. +/// +/// A session resumed through its recorded launcher always starts now: its +/// launch claim is taken when the command is typed, and nothing would take it +/// on a later visit. Such sessions are rare. +public struct AgentRecoveryStartPlan: Equatable, Sendable { + /// How many recovered sessions start right away by default. + public static let defaultImmediateLimit = 3 + + /// Sessions whose terminal starts now, in start order. + public private(set) var startNow: [AgentRecoveryCandidate] = [] + /// Sessions whose workspace opens now and whose terminal starts on first visit. + public private(set) var startOnVisit: [AgentRecoveryCandidate] = [] + + /// - Parameters: + /// - candidates: The sessions to recover. + /// - visibleWorkspaceIds: Workspaces shown in a window right now. + /// - immediateLimit: How many sessions start right away. + public init( + candidates: [AgentRecoveryCandidate], + visibleWorkspaceIds: Set = [], + immediateLimit: Int = AgentRecoveryStartPlan.defaultImmediateLimit + ) { + func isVisible(_ candidate: AgentRecoveryCandidate) -> Bool { + candidate.workspaceId.flatMap(UUID.init(uuidString:)).map(visibleWorkspaceIds.contains) ?? false + } + let ranked = candidates.enumerated().sorted { lhs, rhs in + let lhsVisible = isVisible(lhs.element) + let rhsVisible = isVisible(rhs.element) + if lhsVisible != rhsVisible { return lhsVisible } + if lhs.element.lastActivity != rhs.element.lastActivity { + return lhs.element.lastActivity > rhs.element.lastActivity + } + return lhs.offset < rhs.offset + } + for (_, candidate) in ranked { + if startNow.count < immediateLimit || candidate.launcherResumeArguments != nil { + startNow.append(candidate) + } else { + startOnVisit.append(candidate) + } + } + } +} diff --git a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/SubrouterClaudeResumeRouting.swift b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/SubrouterClaudeResumeRouting.swift index 1050eaba3fa1..935392d2ea3a 100644 --- a/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/SubrouterClaudeResumeRouting.swift +++ b/Packages/macOS/CMUXAgentLaunch/Sources/CMUXAgentLaunch/SubrouterClaudeResumeRouting.swift @@ -31,6 +31,21 @@ public struct SubrouterClaudeResumeRouting: Sendable, Equatable { /// Wrapper-attested copy of the marker bound to the current Claude argv. public static let launchBoundEnvironmentKey = "CMUX_AGENT_LAUNCH_SUBROUTER_CLAUDE_RESUME_COMMAND" + /// The account a routed launch was pinned to, exported by + /// `cmux-claude-wrapper` from the routing headers in the launcher's + /// private `--settings` file. It is the launcher's resolved choice, so it + /// survives however the launcher was invoked or what it appended to the + /// arguments it forwarded. + public static let accountEnvironmentKey = "CMUX_AGENT_LAUNCH_ROUTED_CLAUDE_ACCOUNT" + + /// Launch metadata the wrapper exports for a routed launch, which the + /// queued Claude hooks must carry to the session-start capture. + public static let hookCapturedEnvironmentKeys = [ + environmentKey, + launchBoundEnvironmentKey, + accountEnvironmentKey, + ] + /// Directory-name prefix of the private settings directory `sr claude proxy` creates. public static let privateSettingsDirectoryPrefix = "subrouter-claude-settings-" @@ -50,6 +65,7 @@ public struct SubrouterClaudeResumeRouting: Sendable, Equatable { public static let restoreOwnedEnvironmentKeys: Set = [ environmentKey, launchBoundEnvironmentKey, + accountEnvironmentKey, "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_BASE_URL", @@ -136,6 +152,29 @@ public struct SubrouterClaudeResumeRouting: Sendable, Equatable { ] } + /// The pinned account the wrapper recorded for this launch, or `nil` when + /// the launch was pooled or the value is not a plain account id. The value + /// becomes an argument to the launcher, so anything that could read as an + /// option or carry shell or control characters is refused. + public func capturedAccount(in environment: [String: String]?) -> String? { + guard let value = environment?[Self.accountEnvironmentKey], + (1...256).contains(value.count), + !value.hasPrefix("-"), + value.unicodeScalars.allSatisfy(Self.accountScalars.contains) else { + return nil + } + return value + } + + private static let accountScalars = CharacterSet( + charactersIn: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._@+:=-" + ) + + /// The pinned account as durable launch metadata, or an empty environment. + public func capturedAccountEnvironment(in environment: [String: String]?) -> [String: String] { + capturedAccount(in: environment).map { [Self.accountEnvironmentKey: $0] } ?? [:] + } + /// Whether the captured launch proves a Subrouter-routed plain `claude` launch. /// /// cmux launchers (`claudeTeams` and friends) own their own resume shape and @@ -161,9 +200,10 @@ public struct SubrouterClaudeResumeRouting: Sendable, Equatable { /// to replay follow the session id; Subrouter's private `--settings` file is /// dropped because the launcher issues a fresh one. /// - /// When the record also carries the launcher argv that started Claude - /// (`sr claude proxy --account x`), its account pin is kept; otherwise the - /// pool picks the account. + /// A pinned launch keeps its account: the pin the wrapper recorded + /// (``accountEnvironmentKey``) when present, else one read from a captured + /// launcher argv (`sr claude proxy --account x`) on records that predate + /// it. Otherwise the pool picks the account. public func resumeArguments( launcher: String?, sessionID: String, @@ -180,7 +220,9 @@ public struct SubrouterClaudeResumeRouting: Sendable, Equatable { return nil } let markerTokens = marker.split(separator: " ").map(String.init) - let head = pinnedLauncherArguments(launcherPrefix, markerTokens: markerTokens) ?? markerTokens + let head = capturedAccount(in: environment).map { Array(markerTokens[0..<3]) + ["--account", $0, "--resume"] } + ?? pinnedLauncherArguments(launcherPrefix, markerTokens: markerTokens) + ?? markerTokens return head + [sessionID] + removingPrivateSettingsArguments(from: preserved) diff --git a/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/AgentSessionRecoveryTests.swift b/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/AgentSessionRecoveryTests.swift index dc0ecddfa331..f5506b65418f 100644 --- a/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/AgentSessionRecoveryTests.swift +++ b/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/AgentSessionRecoveryTests.swift @@ -239,3 +239,75 @@ struct AgentSessionRecoveryPlannerTests { #expect(command.launcherPrefix == nil) } } + +/// A heavy user can have dozens of agents running when cmux dies; starting +/// them all on relaunch spikes CPU and memory. Only a few start at once. +@Suite("Agent session recovery start plan") +struct AgentRecoveryStartPlanTests { + private let now = Date(timeIntervalSince1970: 1_790_428_300) + + private func candidate( + _ id: String, + minutesAgo: Double, + workspaceId: UUID? = nil, + launcherPrefix: [String]? = nil + ) -> AgentRecoveryCandidate { + AgentRecoveryCandidate( + kind: "claude", + sessionId: id, + workspaceId: workspaceId?.uuidString, + cwd: "/tmp", + launchCommand: AgentLaunchCommand(arguments: ["claude"], launcherPrefix: launcherPrefix), + lastActivity: now.addingTimeInterval(-minutesAgo * 60) + ) + } + + @Test("the most recently active sessions start now and the rest wait for a visit") + func startsMostRecentFirst() { + let plan = AgentRecoveryStartPlan( + candidates: (1...6).map { candidate("s\($0)", minutesAgo: Double(7 - $0)) }, + immediateLimit: 2 + ) + #expect(plan.startNow.map(\.sessionId) == ["s6", "s5"]) + #expect(plan.startOnVisit.map(\.sessionId) == ["s4", "s3", "s2", "s1"]) + } + + @Test("sessions from a workspace on screen start before more recent ones") + func visibleWorkspacesStartFirst() { + let visible = UUID() + let plan = AgentRecoveryStartPlan( + candidates: [ + candidate("recent", minutesAgo: 1), + candidate("on-screen", minutesAgo: 30, workspaceId: visible), + candidate("older", minutesAgo: 5), + ], + visibleWorkspaceIds: [visible], + immediateLimit: 2 + ) + #expect(plan.startNow.map(\.sessionId) == ["on-screen", "recent"]) + #expect(plan.startOnVisit.map(\.sessionId) == ["older"]) + } + + /// A session resumed through its recorded launcher takes its launch claim + /// when it is typed, and nothing would claim it on a later visit. + @Test("a session resumed through its recorded launcher always starts now") + func launcherSessionsStartNow() { + let plan = AgentRecoveryStartPlan( + candidates: [ + candidate("a", minutesAgo: 1), + candidate("b", minutesAgo: 2), + candidate("launcher", minutesAgo: 3, launcherPrefix: ["sr", "claude", "proxy", "--account", "me"]), + ], + immediateLimit: 1 + ) + #expect(plan.startNow.map(\.sessionId) == ["a", "launcher"]) + #expect(plan.startOnVisit.map(\.sessionId) == ["b"]) + } + + @Test("everything starts now when there are only a few sessions") + func fewSessionsAllStartNow() { + let plan = AgentRecoveryStartPlan(candidates: [candidate("a", minutesAgo: 1), candidate("b", minutesAgo: 2)]) + #expect(plan.startNow.map(\.sessionId) == ["a", "b"]) + #expect(plan.startOnVisit.isEmpty) + } +} diff --git a/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/SubrouterClaudeRestoreRoutingTests.swift b/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/SubrouterClaudeRestoreRoutingTests.swift index ad64576622ee..1bf9938bd97f 100644 --- a/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/SubrouterClaudeRestoreRoutingTests.swift +++ b/Packages/macOS/CMUXAgentLaunch/Tests/CMUXAgentLaunchTests/SubrouterClaudeRestoreRoutingTests.swift @@ -13,6 +13,7 @@ struct SubrouterClaudeRestoreRoutingTests { private let marker = "sr claude proxy --resume" private let markerKey = "SUBROUTER_CLAUDE_RESUME_COMMAND" private let launchBoundMarkerKey = "CMUX_AGENT_LAUNCH_SUBROUTER_CLAUDE_RESUME_COMMAND" + private let accountKey = "CMUX_AGENT_LAUNCH_ROUTED_CLAUDE_ACCOUNT" private let localPoolBaseURL = "http://127.0.0.1:31415/v1" private let capturedClaude = "/opt/homebrew/bin/claude" private let capturedConfigDir = "/Users/me/.subrouter/codex/claude-proxy/3fa7ce27b6c3bad79bd47d1a" @@ -158,6 +159,72 @@ struct SubrouterClaudeRestoreRoutingTests { } } + /// The pin recorded for the launch itself holds when the launcher argv + /// cannot be split from what it forwarded (a routed launcher that appends + /// arguments after the forwarded tail leaves no launcher prefix), and wins + /// over a prefix that names another account. + @Test("The recorded account pin survives without a launcher prefix and wins over one") + func recordedAccountPinSurvivesWithoutLauncherPrefix() throws { + for launcherPrefix in [nil, ["sr", "claude", "proxy", "--account", "stale@example.com"]] { + var environment = routedLaunchEnvironment(baseURL: localPoolBaseURL) + environment[accountKey] = "me@example.com" + let request = resumeRequest(environment: environment, launcherPrefix: launcherPrefix) + + let invocation = try #require(plannerWithSubrouterOnPath().invocation( + for: request, + ambientEnvironment: ambientEnvironment + )) + + #expect( + invocation.arguments == ["sr", "claude", "proxy", "--account", "me@example.com", "--resume", sessionID, "--model", "opus"], + "\(invocation.arguments)" + ) + // The pin is launch metadata, never replayed into the environment. + #expect(invocation.environment[accountKey] == nil) + } + } + + @Test( + "A recorded account that could read as an option or carry shell text is not a pin", + arguments: ["", "--print", "me@example.com; rm -rf /", "me example", "me\nnext"] + ) + func unusableRecordedAccountIsIgnored(account: String) throws { + var environment = routedLaunchEnvironment(baseURL: localPoolBaseURL) + environment[accountKey] = account + let invocation = try #require(plannerWithSubrouterOnPath().invocation( + for: resumeRequest(environment: environment), + ambientEnvironment: ambientEnvironment + )) + + #expect( + invocation.arguments == ["sr", "claude", "proxy", "--resume", sessionID, "--model", "opus"], + "\(invocation.arguments)" + ) + } + + /// The hook records the pin into the durable launch environment, and a + /// replayed resume command never carries it. + @Test("The recorded account pin is kept as restore metadata and dropped from replay") + func recordedAccountPinIsRestoreMetadataOnly() { + let policy = AgentLaunchEnvironmentPolicy() + let captured = [ + accountKey: "me@example.com", + "ANTHROPIC_BASE_URL": localPoolBaseURL, + ] + #expect(policy.selectedRestoreEnvironment(from: captured, kind: "claude")[accountKey] == "me@example.com") + #expect(policy.selectedRestoreEnvironment(from: captured, kind: "codex")[accountKey] == nil) + #expect(policy.selectedRestoreEnvironment( + from: [accountKey: "--print"], + kind: "claude" + )[accountKey] == nil) + #expect(policy.selectedReplayEnvironment( + from: captured, + kind: "claude", + launcher: "claude", + arguments: ["claude"] + )[accountKey] == nil) + } + /// Only the account pin is carried over: a prompt, `--settings`, or /// `--print` sr received stays out of the restore argv, and the program /// is still resolved on PATH even when the captured one moved. diff --git a/Resources/bin/cmux-claude-wrapper b/Resources/bin/cmux-claude-wrapper index 458e15de4b80..eb02d2f5eefd 100755 --- a/Resources/bin/cmux-claude-wrapper +++ b/Resources/bin/cmux-claude-wrapper @@ -1777,7 +1777,9 @@ if [[ -z "$cmux_claude_hook_cmux_bin" ]]; then fi export CMUX_CLAUDE_HOOK_CMUX_BIN="$cmux_claude_hook_cmux_bin" -claude_launch_has_subrouter_private_settings() { +# Prints the routed launcher's private `--settings` file when the argv carries +# one that still exists. +claude_launch_subrouter_private_settings_path() { local argument value expects_value=0 for argument in "$@"; do [[ "$argument" == "--" ]] && break @@ -1794,7 +1796,10 @@ claude_launch_has_subrouter_private_settings() { fi case "$value" in */subrouter-claude-settings-*/settings.json) - [[ -f "$value" ]] && return 0 + if [[ -f "$value" ]]; then + printf '%s\n' "$value" + return 0 + fi ;; esac done @@ -1807,15 +1812,33 @@ bind_subrouter_claude_resume_marker() { "sr claude proxy --resume"|"subrouter claude proxy --resume") ;; *) return 0 ;; esac - claude_launch_has_subrouter_private_settings "$@" || return 0 + claude_launch_subrouter_private_settings_path "$@" >/dev/null || return 0 export CMUX_AGENT_LAUNCH_SUBROUTER_CLAUDE_RESUME_COMMAND="$SUBROUTER_CLAUDE_RESUME_COMMAND" } +# Records the account a routed launch was pinned to, so a restore keeps it. +# The pin is read from the routing headers the launcher wrote into its private +# settings file for this launch (the file is gone once the launcher exits), so +# it holds however the launcher was invoked or what it appended to the argv it +# forwarded. A pooled launch has no pin header and records nothing. +bind_routed_claude_account() { + unset CMUX_AGENT_LAUNCH_ROUTED_CLAUDE_ACCOUNT + local settings_path account + settings_path="$(claude_launch_subrouter_private_settings_path "$@")" || return 0 + # The headers are one JSON string whose lines are joined by an escaped + # newline, so the id ends at `\n` or at the closing quote. An id with any + # other character is left unrecorded rather than cut short. + account="$(LC_ALL=C sed -nE 's/.*X-Subrouter-Account-ID: ([A-Za-z0-9._@+:=-]+)(\\n|").*/\1/p' "$settings_path" 2>/dev/null | head -n 1)" + [[ -n "$account" && "$account" != -* && ${#account} -le 256 ]] || return 0 + export CMUX_AGENT_LAUNCH_ROUTED_CLAUDE_ACCOUNT="$account" +} + if [[ "$cmux_claude_teams_wrapper_launch" != "1" ]]; then export CMUX_AGENT_LAUNCH_KIND="claude" export CMUX_AGENT_LAUNCH_EXECUTABLE="$REAL_CLAUDE" export CMUX_AGENT_LAUNCH_CWD="$PWD" bind_subrouter_claude_resume_marker "$@" + bind_routed_claude_account "$@" fi install_cmux_node_options "$REAL_CLAUDE" diff --git a/Sources/AgentHookDeliveryEvent.swift b/Sources/AgentHookDeliveryEvent.swift index 4b67af682ca3..8009d4cbfdc7 100644 --- a/Sources/AgentHookDeliveryEvent.swift +++ b/Sources/AgentHookDeliveryEvent.swift @@ -27,7 +27,7 @@ struct AgentHookDeliveryEvent: Sendable { "session-end", ] - private static let allowedHookDataEnvironmentKeys: Set = [ + private static let allowedHookDataEnvironmentKeys: Set = Set([ "PWD", "CMUX_AGENT_HOOK_STATE_DIR", "CMUX_AGENT_HOOK_SUPPRESS_VISIBLE_MUTATIONS", AgentHookDeliveryPolicy.routeSnapshotEnvironmentKey, @@ -35,12 +35,12 @@ struct AgentHookDeliveryEvent: Sendable { "CMUX_AGENT_LAUNCH_EXECUTABLE", "CMUX_AGENT_LAUNCH_KIND", "CMUX_AGENT_MANAGED_SUBAGENT", "CMUX_SUPPRESS_SUBAGENT_NOTIFICATIONS", "CMUX_SURFACE_ID", "CMUX_WORKSPACE_ID", - ] + ]).union(SubrouterClaudeResumeRouting.hookCapturedEnvironmentKeys) - private static let optionalLaunchEnvironmentKeys: Set = [ + private static let optionalLaunchEnvironmentKeys: Set = Set([ "CMUX_AGENT_LAUNCH_ARGV_B64", "CMUX_AGENT_LAUNCH_CWD", "CMUX_AGENT_LAUNCH_EXECUTABLE", "CMUX_AGENT_LAUNCH_KIND", - ] + ]).union(SubrouterClaudeResumeRouting.hookCapturedEnvironmentKeys) let agent: String let subcommand: String diff --git a/Sources/AgentSessionCloseJournal.swift b/Sources/AgentSessionCloseJournal.swift index b1834afc1854..5b3baf3587da 100644 --- a/Sources/AgentSessionCloseJournal.swift +++ b/Sources/AgentSessionCloseJournal.swift @@ -47,3 +47,75 @@ struct AgentSessionCloseJournal: Sendable { } } } + +/// Owns terminal panels that can carry agent sessions: a workspace or a Dock. +/// +/// Both close panels through their own teardown, and both call +/// ``journalClosedAgentSessions(panelId:)`` from it, so a Claude pane closed +/// anywhere is journaled the same way. +@MainActor +protocol AgentSessionPanelHost: AnyObject { + /// The workspace the journal attributes a closed panel's sessions to. + var agentSessionWorkspaceID: UUID { get } + var surfaceResumeBindingsByPanelId: [UUID: SurfaceResumeBindingSnapshot] { get } + /// Hook bindings that can name the agent session a panel carries. + func agentSessionBindingsForClose(panelId: UUID) -> [SurfaceResumeBindingSnapshot] + var restoredAgentLifecycle: RestoredAgentLifecycleCoordinator { get } + var deferredAgentResumeRestoresByPanelId: [UUID: DeferredAgentResumeRestore] { get } + var agentSessionCloseJournal: AgentSessionCloseJournal { get } +} + +extension AgentSessionPanelHost { + func agentSessionBindingsForClose(panelId: UUID) -> [SurfaceResumeBindingSnapshot] { + surfaceResumeBindingsByPanelId[panelId].map { [$0] } ?? [] + } + + /// Records the end of the recoverable agent sessions this panel carries, + /// so crash recovery never reopens a terminal the user closed. Call it + /// before the panel's bindings and restore state are discarded. Skipped + /// while the app quits: those sessions end with the app, and startup + /// restore owns them. + func journalClosedAgentSessions(panelId: UUID) { + guard AppDelegate.shared?.isTerminatingApp != true else { return } + let recoverable = Set(AgentSessionRecovery.recoverableKinds.map(\.rawValue)) + var sessions: [(kind: String, sessionID: String)] = [] + for binding in agentSessionBindingsForClose(panelId: panelId) where binding.isAgentHookBinding { + guard let kind = binding.kind?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased(), + recoverable.contains(kind), + let sessionID = binding.checkpointId else { continue } + sessions.append((kind, sessionID)) + } + let restoredAgents = [ + restoredAgentLifecycle.snapshotsByPanelId[panelId], + deferredAgentResumeRestoresByPanelId[panelId]?.restorableAgent, + ] + for agent in restoredAgents.compactMap({ $0 }) where recoverable.contains(agent.kind.rawValue) { + sessions.append((agent.kind.rawValue, agent.sessionId)) + } + guard !sessions.isEmpty else { return } + // A session another panel still carries (a restore that lost to a live + // owner, or a stale snapshot resumed elsewhere) did not end here. + let carriedElsewhere = AppDelegate.shared?.openAgentSessionIdsForRecovery(excludingPanelId: panelId) ?? [] + sessions.removeAll { carriedElsewhere.contains($0.sessionID) } + guard !sessions.isEmpty else { return } + agentSessionCloseJournal.recordClosed( + sessions: sessions, + workspaceID: agentSessionWorkspaceID, + surfaceID: panelId + ) + } +} + +extension Workspace: AgentSessionPanelHost { + var agentSessionWorkspaceID: UUID { id } +} + +extension DockSplitStore: AgentSessionPanelHost { + var agentSessionWorkspaceID: UUID { workspaceId } + + /// The effective binding, plus the agent-hook binding the Dock keeps + /// aside while process detection shows another one (a tmux binding). + func agentSessionBindingsForClose(panelId: UUID) -> [SurfaceResumeBindingSnapshot] { + [surfaceResumeBindingsByPanelId[panelId], managedAgentResumeBindingsByPanelId[panelId]].compactMap { $0 } + } +} diff --git a/Sources/AgentSessionRecovery.swift b/Sources/AgentSessionRecovery.swift index fec6530ceb7b..cd66e66d46b5 100644 --- a/Sources/AgentSessionRecovery.swift +++ b/Sources/AgentSessionRecovery.swift @@ -190,6 +190,73 @@ struct AgentSessionRecovery: Sendable { } } +extension AgentSessionRecovery { + /// Reopens each candidate in its own workspace of `tabManager` and types + /// its resume command. Sessions in `alreadyOpen` are skipped. Returns the + /// session ids that were reopened. + /// + /// Only the sessions ``AgentRecoveryStartPlan`` picks start right away; + /// the rest open their workspace now and resume on its first visit, so a + /// relaunch after a crash does not start every agent at once. + @MainActor + @discardableResult + static func reopen( + _ candidates: [AgentRecoveryCandidate], + in tabManager: TabManager, + alreadyOpen: Set, + visibleWorkspaceIds: Set = [] + ) -> [String] { + let plan = AgentRecoveryStartPlan( + candidates: candidates.filter { !alreadyOpen.contains($0.sessionId) }, + visibleWorkspaceIds: visibleWorkspaceIds + ) + var restored: [String] = [] + let ordered = plan.startNow.map { ($0, true) } + plan.startOnVisit.map { ($0, false) } + for (candidate, startsNow) in ordered { + guard let launch = launch(for: candidate) else { continue } + let directory = candidate.cwd.flatMap { FileManager.default.fileExists(atPath: $0) ? $0 : nil } + let startupAgent: SessionRestorableAgentSnapshot? + switch launch { + case .restoreVerb(_, let agent): + // `cmux restore` takes the launch claim at its admission + // boundary, so a concurrent restore cannot start the session + // twice. The panel carries the session, so a second recovery + // sees it as open, even before a deferred start. + startupAgent = agent + case .launcherCommand: + // Nothing downstream claims a launcher command; take the + // claim startup restore takes before typing it. The start + // plan always starts these now, before the claim expires. + guard AgentResumeLaunchGuard.shared.claimResumeLaunch( + kind: candidate.kind, + sessionId: candidate.sessionId + ) else { continue } + startupAgent = nil + } + guard tabManager.addWorkspaceIfActive( + title: workspaceTitle(for: candidate), + titleSource: .auto, + workingDirectory: directory, + initialTerminalInput: launch.terminalInput, + initialTerminalStartupRestoreAgent: startupAgent, + select: false, + // A session picked to start now runs without waiting for a + // visit (its hook record then carries a live pid). The rest + // hold their terminal until the workspace is first selected. + eagerLoadTerminal: startsNow, + initialTerminalStartsOnFirstVisit: !startsNow + ) != nil else { + if startupAgent == nil { + AgentResumeLaunchGuard.shared.releaseResumeLaunch(kind: candidate.kind, sessionId: candidate.sessionId) + } + continue + } + restored.append(candidate.sessionId) + } + return restored + } +} + extension AppDelegate { /// Agent session ids already carried by open panels (restored from the /// snapshot or bound since launch), which recovery must not duplicate. @@ -221,6 +288,12 @@ extension AppDelegate { bindings: dock.surfaceResumeBindingsByPanelId, deferred: dock.deferredAgentResumeRestoresByPanelId ) + ids.formUnion( + dock.managedAgentResumeBindingsByPanelId + .filter { $0.key != excludedPanelId } + .values + .compactMap(\.checkpointId) + ) } for manager in managers { for workspace in manager.tabs { @@ -237,52 +310,25 @@ extension AppDelegate { } /// Reopens each candidate in its own workspace and types its resume - /// command. Returns the session ids that were started. + /// command. Returns the session ids that were reopened. @discardableResult func restoreRecoveredAgentSessions(_ candidates: [AgentRecoveryCandidate]) -> [String] { guard let tabManager else { return [] } - let alreadyOpen = openAgentSessionIdsForRecovery() - var restored: [String] = [] - for candidate in candidates where !alreadyOpen.contains(candidate.sessionId) { - guard let launch = AgentSessionRecovery.launch(for: candidate) else { continue } - let directory = candidate.cwd.flatMap { FileManager.default.fileExists(atPath: $0) ? $0 : nil } - let startupAgent: SessionRestorableAgentSnapshot? - switch launch { - case .restoreVerb(_, let agent): - // `cmux restore` takes the launch claim at its admission - // boundary, so a concurrent restore cannot start the session - // twice. The panel carries the session, so a second recovery - // sees it as open. - startupAgent = agent - case .launcherCommand: - // Nothing downstream claims a launcher command; take the - // claim startup restore takes before typing it. - guard AgentResumeLaunchGuard.shared.claimResumeLaunch( - kind: candidate.kind, - sessionId: candidate.sessionId - ) else { continue } - startupAgent = nil - } - guard tabManager.addWorkspaceIfActive( - title: AgentSessionRecovery.workspaceTitle(for: candidate), - titleSource: .auto, - workingDirectory: directory, - initialTerminalInput: launch.terminalInput, - initialTerminalStartupRestoreAgent: startupAgent, - select: false, - // Start the terminal now so the agent is running (and its - // hook record carries a live pid) without waiting for a visit; - // a launcher command's claim would otherwise expire first. - eagerLoadTerminal: true - ) != nil else { - if startupAgent == nil { - AgentResumeLaunchGuard.shared.releaseResumeLaunch(kind: candidate.kind, sessionId: candidate.sessionId) - } - continue - } - restored.append(candidate.sessionId) + return AgentSessionRecovery.reopen( + candidates, + in: tabManager, + alreadyOpen: openAgentSessionIdsForRecovery(), + visibleWorkspaceIds: visibleWorkspaceIdsForRecovery() + ) + } + + /// Workspaces shown in a main window right now. + private func visibleWorkspaceIdsForRecovery() -> Set { + var managers = mainWindowContexts.values.map(\.tabManager) + if let tabManager, !managers.contains(where: { $0 === tabManager }) { + managers.append(tabManager) } - return restored + return Set(managers.compactMap(\.selectedTabId)) } /// After a launch that followed an unclean exit, finds agent sessions the diff --git a/Sources/DockSplitStore+PanelDestruction.swift b/Sources/DockSplitStore+PanelDestruction.swift index ce943acc3b8c..a3c771a17c5e 100644 --- a/Sources/DockSplitStore+PanelDestruction.swift +++ b/Sources/DockSplitStore+PanelDestruction.swift @@ -10,6 +10,11 @@ extension DockSplitStore { @discardableResult func discardPanelStateAndClose(panelId: UUID) -> (any Panel)? { + // A panel this Dock still owns is being closed, not handed to another + // container (a transfer detaches it first). + if panels[panelId] != nil { + journalClosedAgentSessions(panelId: panelId) + } if panels[panelId] is BrowserPanel { removeBrowserOpenTabSuggestion(panelId: panelId) } diff --git a/Sources/DockSplitStore.swift b/Sources/DockSplitStore.swift index 787701dac64e..d24ecd7728b7 100644 --- a/Sources/DockSplitStore.swift +++ b/Sources/DockSplitStore.swift @@ -92,6 +92,9 @@ final class DockSplitStore: BonsplitDelegate, FilePreviewTabMetadataHost { terminalStartupRestoreCoordinator.lifecycle } @ObservationIgnored var surfaceResumeBindingsByPanelId: [UUID: SurfaceResumeBindingSnapshot] = [:] + /// Journals agent sessions ended by closing their terminal. Tests point it + /// at a private journal. + @ObservationIgnored var agentSessionCloseJournal = AgentSessionCloseJournal() /// In-memory compare-and-claim state held while a CLI restore hands the /// validated binding to its child process. @ObservationIgnored var surfaceResumeRestoreClaimsByPanelId: [ diff --git a/Sources/TabManager.swift b/Sources/TabManager.swift index c7f2ef866a18..5452c3aa0dac 100644 --- a/Sources/TabManager.swift +++ b/Sources/TabManager.swift @@ -349,6 +349,9 @@ class TabManager: ObservableObject { if !isRestoringSessionSnapshot { workspaces.expandWorkspaceGroupForSelectionIfNeeded() } + if let selectedTabId { + workspacesById[selectedTabId]?.admitStartupRestoresAwaitingFirstVisit() + } sentryBreadcrumb("workspace.switch", data: [ "tabCount": tabs.count ]) @@ -1132,6 +1135,7 @@ class TabManager: ObservableObject { initialTerminalCommand: String?, initialTerminalInput: String? = nil, initialTerminalStartupRestoreAgent: SessionRestorableAgentSnapshot? = nil, + initialTerminalStartsOnFirstVisit: Bool = false, initialTerminalEnvironment: [String: String], initialBrowserURL: URL? = nil, initialBrowserOmnibarVisible: Bool = true, @@ -1150,6 +1154,7 @@ class TabManager: ObservableObject { initialTerminalInput: initialTerminalInput, initialTerminalStartupRestoreAgent: initialTerminalStartupRestoreAgent, initialTerminalStartupRestoreCommitOwner: .tabManagerTopology, + initialTerminalStartsOnFirstVisit: initialTerminalStartsOnFirstVisit, initialTerminalEnvironment: initialTerminalEnvironment, initialBrowserURL: initialBrowserURL, initialBrowserOmnibarVisible: initialBrowserOmnibarVisible, @@ -1350,6 +1355,7 @@ class TabManager: ObservableObject { inheritWorkingDirectory: Bool = true, select: Bool = true, eagerLoadTerminal: Bool = false, + initialTerminalStartsOnFirstVisit: Bool = false, placementOverride: WorkspacePlacement? = nil, autoWelcomeIfNeeded: Bool = true, autoRefreshMetadata: Bool = true, @@ -1445,6 +1451,7 @@ class TabManager: ObservableObject { initialTerminalCommand: initialTerminalCommand, initialTerminalInput: initialTerminalInput, initialTerminalStartupRestoreAgent: initialTerminalStartupRestoreAgent, + initialTerminalStartsOnFirstVisit: initialTerminalStartsOnFirstVisit, initialTerminalEnvironment: resolvedInitialTerminalEnvironment, initialBrowserURL: initialBrowserURL, initialBrowserOmnibarVisible: initialBrowserOmnibarVisible, diff --git a/Sources/Workspace+PanelLifecycle.swift b/Sources/Workspace+PanelLifecycle.swift index 0ab2e0ed888d..3554ad845f1f 100644 --- a/Sources/Workspace+PanelLifecycle.swift +++ b/Sources/Workspace+PanelLifecycle.swift @@ -435,35 +435,14 @@ extension Workspace { } } - /// Records the end of the recoverable agent sessions this panel carries, - /// so crash recovery never reopens a terminal the user closed. Runs before - /// the panel's bindings and restore state are discarded. Skipped while the - /// app quits: those sessions end with the app, and startup restore owns them. - func journalClosedAgentSessions(panelId: UUID) { - guard AppDelegate.shared?.isTerminatingApp != true else { return } - let recoverable = Set(AgentSessionRecovery.recoverableKinds.map(\.rawValue)) - var sessions: [(kind: String, sessionID: String)] = [] - if let binding = surfaceResumeBindingsByPanelId[panelId], - binding.isAgentHookBinding, - let kind = binding.kind?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased(), - recoverable.contains(kind), - let sessionID = binding.checkpointId { - sessions.append((kind, sessionID)) - } - let restoredAgents = [ - restoredAgentSnapshotsByPanelId[panelId], - deferredAgentResumeRestoresByPanelId[panelId]?.restorableAgent, - ] - for agent in restoredAgents.compactMap({ $0 }) where recoverable.contains(agent.kind.rawValue) { - sessions.append((agent.kind.rawValue, agent.sessionId)) - } - guard !sessions.isEmpty else { return } - // A session another panel still carries (a restore that lost to a live - // owner, or a stale snapshot resumed elsewhere) did not end here. - let carriedElsewhere = AppDelegate.shared?.openAgentSessionIdsForRecovery(excludingPanelId: panelId) ?? [] - sessions.removeAll { carriedElsewhere.contains($0.sessionID) } - guard !sessions.isEmpty else { return } - agentSessionCloseJournal.recordClosed(sessions: sessions, workspaceID: id, surfaceID: panelId) + /// Starts the restored terminals held for this workspace's first visit. + func admitStartupRestoresAwaitingFirstVisit() { + guard !startupRestorePanelIdsAwaitingFirstVisit.isEmpty else { return } + let panelIds = startupRestorePanelIdsAwaitingFirstVisit + startupRestorePanelIdsAwaitingFirstVisit.removeAll() + for panelId in panelIds { + terminalPanel(for: panelId)?.surface.admitStartupRestoreRuntime() + } } /// Discard every Workspace-owned contribution for a surface whose tab, diff --git a/Sources/Workspace.swift b/Sources/Workspace.swift index 687d4c1441ef..993c601ff7db 100644 --- a/Sources/Workspace.swift +++ b/Sources/Workspace.swift @@ -3232,6 +3232,9 @@ final class Workspace: Identifiable, ObservableObject, FilePreviewTabMetadataHos /// Journals agent sessions ended by closing their terminal. Tests point it /// at a private journal. var agentSessionCloseJournal = AgentSessionCloseJournal() + /// Restored terminals held until this workspace is first shown (crash + /// recovery defers all but a few). Their panels already carry the session. + var startupRestorePanelIdsAwaitingFirstVisit: Set = [] /// In-memory compare-and-claim state held while a CLI restore hands the /// validated binding to its child process. @ObservationIgnored var surfaceResumeRestoreClaimsByPanelId: [ @@ -3986,6 +3989,7 @@ final class Workspace: Identifiable, ObservableObject, FilePreviewTabMetadataHos initialTerminalInput: String? = nil, initialTerminalStartupRestoreAgent: SessionRestorableAgentSnapshot? = nil, initialTerminalStartupRestoreCommitOwner: WorkspaceTerminalStartupRestoreCommitOwner = .workspaceTopology, + initialTerminalStartsOnFirstVisit: Bool = false, initialTerminalEnvironment: [String: String] = [:], initialBrowserURL: URL? = nil, initialBrowserOmnibarVisible: Bool = true, @@ -4208,8 +4212,12 @@ final class Workspace: Identifiable, ObservableObject, FilePreviewTabMetadataHos snapshot: initialTerminalStartupRestoreAgent, manualResumeAvailable: true, willRunStartupInput: initialTerminalInput != nil, - resumeWorkingDirectory: initialTerminalStartupRestoreAgent.workingDirectory + resumeWorkingDirectory: initialTerminalStartupRestoreAgent.workingDirectory, + defersStartupRestoreAdmission: initialTerminalStartsOnFirstVisit ) + if initialTerminalStartsOnFirstVisit { + startupRestorePanelIdsAwaitingFirstVisit.insert(terminalPanel.id) + } if initialTerminalStartupRestoreCommitOwner == .workspaceTopology { terminalStartupRestoreCoordinator.commitPendingRestores( panelIDs: [terminalPanel.id] diff --git a/cmuxTests/AgentHookDeliveryQueueTests.swift b/cmuxTests/AgentHookDeliveryQueueTests.swift index fd55c7f16857..02df44598f22 100644 --- a/cmuxTests/AgentHookDeliveryQueueTests.swift +++ b/cmuxTests/AgentHookDeliveryQueueTests.swift @@ -1,6 +1,7 @@ import Darwin import Foundation import Testing +import CMUXAgentLaunch import CmuxRemoteWorkspace #if canImport(cmux_DEV) @@ -636,6 +637,26 @@ struct AgentHookDeliveryQueueTests { #expect(unsupportedEnvironment == nil) } + /// The routed launch's account pin rides the queued session-start hook to + /// the capture that records it; the ingress used to reject the event. + @Test("Queued Claude hooks carry the routed launch metadata") + func queuedClaudeHookCarriesRoutedLaunchMetadata() throws { + let environment = [ + "CMUX_SURFACE_ID": "surface-a", + SubrouterClaudeResumeRouting.accountEnvironmentKey: "me@example.com", + SubrouterClaudeResumeRouting.environmentKey: "sr claude proxy --resume", + SubrouterClaudeResumeRouting.launchBoundEnvironmentKey: "sr claude proxy --resume", + ] + let event = try #require(AgentHookDeliveryEvent(params: [ + "agent": "claude", + "subcommand": "session-start", + "payload": "{}", + "socket_path": "/tmp/cmux-test.sock", + "environment": environment, + ])) + #expect(event.environment == environment) + } + @Test("Every agent shares generic lifecycle queue admission") func allAgentsShareLifecycleAdmission() throws { let agents = [ diff --git a/cmuxTests/AgentSessionRecoveryAppTests.swift b/cmuxTests/AgentSessionRecoveryAppTests.swift index e5ca121184cd..9752d8906540 100644 --- a/cmuxTests/AgentSessionRecoveryAppTests.swift +++ b/cmuxTests/AgentSessionRecoveryAppTests.swift @@ -239,6 +239,215 @@ struct AgentSessionRecoveryAppTests { #expect(recovered.map(\.sessionId) == [killedID]) } + /// Recovery used to start every lost session at once, so a heavy user's + /// relaunch spawned dozens of agents together. Only a few start now; the + /// rest open their workspace and resume on its first visit. + @MainActor + @Test + func recoveryStartsOnlyAFewSessionsAtOnce() throws { + let manager = TabManager() + let existing = Set(manager.tabs.map(\.id)) + defer { + for workspace in manager.tabs where !existing.contains(workspace.id) { + workspace.teardownAllPanels() + } + } + let now = Date() + let candidates = (1...6).map { index in + AgentRecoveryCandidate( + kind: "claude", + sessionId: "6a0f2c1e-1d3b-4c55-8e0a-7b9d2f4c1a0\(index)", + workspaceId: nil, + cwd: "/tmp", + launchCommand: AgentLaunchCommand(launcher: "claude", arguments: ["claude"]), + lastActivity: now.addingTimeInterval(-Double(index) * 60) + ) + } + + let reopened = AgentSessionRecovery.reopen(candidates, in: manager, alreadyOpen: []) + + #expect(Set(reopened) == Set(candidates.map(\.sessionId))) + let recovered = manager.tabs.filter { !existing.contains($0.id) } + #expect(recovered.count == candidates.count) + #expect(manager.pendingBackgroundWorkspaceLoadIds.count == AgentRecoveryStartPlan.defaultImmediateLimit) + // Only the sessions picked to start now are released; the rest hold + // their terminal until their workspace is first shown. + func isHeld(_ workspace: Workspace) -> Bool { + workspace.focusedTerminalPanel?.surface.isAwaitingStartupRestoreAdmission == true + } + let held = recovered.filter(isHeld) + #expect(held.count == candidates.count - AgentRecoveryStartPlan.defaultImmediateLimit) + let visited = try #require(held.first) + manager.selectedTabId = visited.id + #expect(!isHeld(visited)) + #expect(visited.startupRestorePanelIdsAwaitingFirstVisit.isEmpty) + // Every recovered panel carries its session, so a second recovery + // skips the ones that have not started yet. + let carried = Set(recovered.flatMap { $0.restoredAgentSnapshotsByPanelId.values.map(\.sessionId) }) + #expect(carried == Set(candidates.map(\.sessionId))) + } + + /// Closing a Claude pane in the Dock must journal its end the same way a + /// workspace pane does, or crash recovery brings the closed pane back. + @MainActor + @Test + func closedDockClaudePaneJournalsItsEnd() async throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cmux-agent-recovery-dock-close-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + let sessionID = "5d1c7a52-0d0e-4b1f-9a4e-2f0f7a9c1e03" + // A second pane whose hook binding the Dock keeps aside while process + // detection shows a tmux binding. + let managedSessionID = "5d1c7a52-0d0e-4b1f-9a4e-2f0f7a9c1e04" + let journalURL = root.appendingPathComponent("journal.sqlite3") + let store = try AgentJournalStore(databaseURL: journalURL) + for id in [sessionID, managedSessionID] { + _ = try store.append(AgentJournalEventDraft( + kind: .sessionStarted, + occurredAtMs: Int64(Date().addingTimeInterval(-60).timeIntervalSince1970 * 1000), + source: "claude", + agentKey: "claude_code", + sessionId: id, + workspaceId: UUID().uuidString, + surfaceId: UUID().uuidString + )) + } + store.close() + + let dock = DockSplitStore(workspaceId: UUID(), baseDirectoryProvider: { nil }) + defer { dock.closeAllPanels() } + dock.agentSessionCloseJournal = AgentSessionCloseJournal( + center: AgentJournalLifecycleCenter(databaseURL: journalURL) + ) + let pane = try #require(dock.bonsplitController.allPaneIds.first) + let panelID = try #require(dock.newSurface(kind: .terminal, inPane: pane, focus: false)) + let tabID = try #require(dock.surfaceId(forPanelId: panelID)) + // What the Claude session-start hook leaves on the surface. + dock.surfaceResumeBindingsByPanelId[panelID] = SurfaceResumeBindingSnapshot( + name: "Claude Code", kind: "claude", command: "claude --resume \(sessionID)", + checkpointId: sessionID, source: "agent-hook", updatedAt: Date().timeIntervalSince1970 + ) + + let managedPanelID = try #require(dock.newSurface(kind: .terminal, inPane: pane, focus: false)) + let managedTabID = try #require(dock.surfaceId(forPanelId: managedPanelID)) + dock.managedAgentResumeBindingsByPanelId[managedPanelID] = SurfaceResumeBindingSnapshot( + name: "Claude Code", kind: "claude", command: "claude --resume \(managedSessionID)", + checkpointId: managedSessionID, source: "agent-hook", updatedAt: Date().timeIntervalSince1970 + ) + + dock.forceCloseDockTabIds.formUnion([tabID, managedTabID]) + defer { dock.forceCloseDockTabIds.subtract([tabID, managedTabID]) } + #expect(dock.bonsplitController.closeTab(tabID)) + #expect(dock.bonsplitController.closeTab(managedTabID)) + dock.reconcilePanels() + #expect(dock.panels[panelID] == nil) + #expect(dock.panels[managedPanelID] == nil) + + let reader = AgentJournalSessionTailReader(databaseURL: journalURL) + func hasEnded() -> Bool { + let tails = (try? reader.sessionTails(occurredAtOrAfterMs: 0)) ?? [] + return [sessionID, managedSessionID].allSatisfy { id in + tails.first { $0.sessionId == id }?.hasEnded == true + } + } + for _ in 0..<100 where !hasEnded() { + try await Task.sleep(for: .milliseconds(50)) + } + #expect(hasEnded()) + } + + /// A Dock keeps its agent-hook binding aside while process detection makes + /// a tmux binding effective. Closing a stale pane for that same session + /// must not end the session still carried by the other pane. + @MainActor + @Test + func closingStaleDockPanePreservesManagedSessionCarriedElsewhere() async throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cmux-agent-recovery-dock-carried-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + let carriedSessionID = "5d1c7a52-0d0e-4b1f-9a4e-2f0f7a9c1e05" + let barrierSessionID = "5d1c7a52-0d0e-4b1f-9a4e-2f0f7a9c1e06" + let journalURL = root.appendingPathComponent("journal.sqlite3") + let store = try AgentJournalStore(databaseURL: journalURL) + for id in [carriedSessionID, barrierSessionID] { + _ = try store.append(AgentJournalEventDraft( + kind: .sessionStarted, + occurredAtMs: Int64(Date().addingTimeInterval(-60).timeIntervalSince1970 * 1_000), + source: "claude", + agentKey: "claude_code", + sessionId: id, + workspaceId: UUID().uuidString, + surfaceId: UUID().uuidString + )) + } + store.close() + + let previousApp = AppDelegate.shared + let app = AppDelegate() + let manager = TabManager(autoWelcomeIfNeeded: false) + app.tabManager = manager + defer { + app.tabManager = nil + manager.finalizeAllWorkspacesForWindowClose() + AppDelegate.shared = previousApp + } + let workspace = try #require(manager.selectedWorkspace) + let dock = try #require(workspace.dockSplit) + dock.agentSessionCloseJournal = AgentSessionCloseJournal( + center: AgentJournalLifecycleCenter(databaseURL: journalURL) + ) + let pane = try #require(dock.bonsplitController.allPaneIds.first) + + let stalePanelID = try #require(dock.newSurface(kind: .terminal, inPane: pane, focus: false)) + let staleTabID = try #require(dock.surfaceId(forPanelId: stalePanelID)) + dock.managedAgentResumeBindingsByPanelId[stalePanelID] = SurfaceResumeBindingSnapshot( + name: "Claude Code", kind: "claude", command: "claude --resume \(carriedSessionID)", + checkpointId: carriedSessionID, source: "agent-hook", updatedAt: Date().timeIntervalSince1970 + ) + + let carrierPanelID = try #require(dock.newSurface(kind: .terminal, inPane: pane, focus: false)) + dock.surfaceResumeBindingsByPanelId[carrierPanelID] = SurfaceResumeBindingSnapshot( + name: "tmux", kind: "tmux", command: "tmux attach -t agent", + checkpointId: nil, source: "process-detected", updatedAt: Date().timeIntervalSince1970 + ) + dock.managedAgentResumeBindingsByPanelId[carrierPanelID] = SurfaceResumeBindingSnapshot( + name: "Claude Code", kind: "claude", command: "claude --resume \(carriedSessionID)", + checkpointId: carriedSessionID, source: "agent-hook", updatedAt: Date().timeIntervalSince1970 + ) + + let barrierPanelID = try #require(dock.newSurface(kind: .terminal, inPane: pane, focus: false)) + let barrierTabID = try #require(dock.surfaceId(forPanelId: barrierPanelID)) + dock.surfaceResumeBindingsByPanelId[barrierPanelID] = SurfaceResumeBindingSnapshot( + name: "Claude Code", kind: "claude", command: "claude --resume \(barrierSessionID)", + checkpointId: barrierSessionID, source: "agent-hook", updatedAt: Date().timeIntervalSince1970 + ) + + // Close the stale pane first. The barrier closes second, so observing + // its end proves the journal consumer already handled any stale end. + dock.forceCloseDockTabIds.formUnion([staleTabID, barrierTabID]) + defer { dock.forceCloseDockTabIds.subtract([staleTabID, barrierTabID]) } + #expect(dock.bonsplitController.closeTab(staleTabID)) + dock.reconcilePanels() + #expect(dock.bonsplitController.closeTab(barrierTabID)) + dock.reconcilePanels() + + let reader = AgentJournalSessionTailReader(databaseURL: journalURL) + func barrierHasEnded() -> Bool { + let tails = (try? reader.sessionTails(occurredAtOrAfterMs: 0)) ?? [] + return tails.first { $0.sessionId == barrierSessionID }?.hasEnded == true + } + for _ in 0..<100 where !barrierHasEnded() { + try await Task.sleep(for: .milliseconds(50)) + } + let tails = try reader.sessionTails(occurredAtOrAfterMs: 0) + #expect(tails.first { $0.sessionId == barrierSessionID }?.hasEnded == true) + #expect(tails.first { $0.sessionId == carriedSessionID }?.hasEnded != true) + } + @Test func restoreRejectsMalformedSessionIDsBeforeReadingAppState() { let invalidValues: [Any] = [ diff --git a/cmuxTests/WorkspaceCreateReviewRegressionTests.swift b/cmuxTests/WorkspaceCreateReviewRegressionTests.swift index 726ddb066260..8c4a7da77f7a 100644 --- a/cmuxTests/WorkspaceCreateReviewRegressionTests.swift +++ b/cmuxTests/WorkspaceCreateReviewRegressionTests.swift @@ -291,6 +291,7 @@ private final class RejectingWorkspaceCreationTabManager: TabManager { inheritWorkingDirectory: Bool, select: Bool, eagerLoadTerminal: Bool, + initialTerminalStartsOnFirstVisit: Bool, placementOverride: WorkspacePlacement?, autoWelcomeIfNeeded: Bool, autoRefreshMetadata: Bool, @@ -316,6 +317,7 @@ private final class RejectingWorkspaceCreationTabManager: TabManager { inheritWorkingDirectory: inheritWorkingDirectory, select: select, eagerLoadTerminal: eagerLoadTerminal, + initialTerminalStartsOnFirstVisit: initialTerminalStartsOnFirstVisit, placementOverride: placementOverride, autoWelcomeIfNeeded: autoWelcomeIfNeeded, autoRefreshMetadata: autoRefreshMetadata, diff --git a/cmuxTests/WorkspaceUnitTests.swift b/cmuxTests/WorkspaceUnitTests.swift index 476eac055c2a..d222a1350572 100644 --- a/cmuxTests/WorkspaceUnitTests.swift +++ b/cmuxTests/WorkspaceUnitTests.swift @@ -3520,6 +3520,7 @@ final class WorkspaceCreationPlacementTests: XCTestCase { initialTerminalCommand: String?, initialTerminalInput: String?, initialTerminalStartupRestoreAgent: SessionRestorableAgentSnapshot?, + initialTerminalStartsOnFirstVisit: Bool, initialTerminalEnvironment: [String: String], initialBrowserURL: URL?, initialBrowserOmnibarVisible: Bool, @@ -3538,6 +3539,7 @@ final class WorkspaceCreationPlacementTests: XCTestCase { initialTerminalCommand: initialTerminalCommand, initialTerminalInput: initialTerminalInput, initialTerminalStartupRestoreAgent: initialTerminalStartupRestoreAgent, + initialTerminalStartsOnFirstVisit: initialTerminalStartsOnFirstVisit, initialTerminalEnvironment: initialTerminalEnvironment, initialBrowserURL: initialBrowserURL, initialBrowserOmnibarVisible: initialBrowserOmnibarVisible, @@ -3836,6 +3838,7 @@ final class WorkspaceCreationConfigSanitizationTests: XCTestCase { initialTerminalCommand: String?, initialTerminalInput: String?, initialTerminalStartupRestoreAgent: SessionRestorableAgentSnapshot?, + initialTerminalStartsOnFirstVisit: Bool, initialTerminalEnvironment: [String: String], initialBrowserURL: URL?, initialBrowserOmnibarVisible: Bool, @@ -3854,6 +3857,7 @@ final class WorkspaceCreationConfigSanitizationTests: XCTestCase { initialTerminalCommand: initialTerminalCommand, initialTerminalInput: initialTerminalInput, initialTerminalStartupRestoreAgent: initialTerminalStartupRestoreAgent, + initialTerminalStartsOnFirstVisit: initialTerminalStartsOnFirstVisit, initialTerminalEnvironment: initialTerminalEnvironment, initialBrowserURL: initialBrowserURL, initialBrowserOmnibarVisible: initialBrowserOmnibarVisible, diff --git a/tests/test_claude_wrapper_subrouter_resume_marker.py b/tests/test_claude_wrapper_subrouter_resume_marker.py index dc5b3b805302..b053e7fb6c04 100644 --- a/tests/test_claude_wrapper_subrouter_resume_marker.py +++ b/tests/test_claude_wrapper_subrouter_resume_marker.py @@ -39,6 +39,7 @@ SOURCE_WRAPPER = ROOT / "Resources" / "bin" / "cmux-claude-wrapper" MARKER_KEY = "SUBROUTER_CLAUDE_RESUME_COMMAND" BOUND_KEY = "CMUX_AGENT_LAUNCH_SUBROUTER_CLAUDE_RESUME_COMMAND" +ACCOUNT_KEY = "CMUX_AGENT_LAUNCH_ROUTED_CLAUDE_ACCOUNT" SR_MARKER = "sr claude proxy --resume" SUBROUTER_MARKER = "subrouter claude proxy --resume" SESSION_ID = "0198f073-0a5b-7000-8000-000000000059" @@ -74,6 +75,7 @@ def run_wrapper( *, marker: str | None, inherited_bound: str | None = None, + inherited_account: str | None = None, ) -> tuple[int, dict[str, str], list[str], str]: """Run the wrapper against a fake claude that records its environment. @@ -152,6 +154,8 @@ def run_wrapper( env[MARKER_KEY] = marker if inherited_bound is not None: env[BOUND_KEY] = inherited_bound + if inherited_account is not None: + env[ACCOUNT_KEY] = inherited_account argv = argv_builder(tmpdir) try: @@ -171,15 +175,32 @@ def run_wrapper( return proc.returncode, read_env_log(env_log), claude_args, proc.stderr.strip() -def private_settings(tmpdir: Path, *, create: bool = True, suffix: str = "3294281412") -> Path: +def private_settings( + tmpdir: Path, + *, + create: bool = True, + suffix: str = "3294281412", + body: str = PRIVATE_SETTINGS_BODY, +) -> Path: directory = tmpdir / f"subrouter-claude-settings-{suffix}" path = directory / "settings.json" if create: directory.mkdir(parents=True, exist_ok=True) - path.write_text(PRIVATE_SETTINGS_BODY, encoding="utf-8") + path.write_text(body, encoding="utf-8") return path +def settings_with_headers(*headers: str) -> str: + """A private settings body whose routing headers are JSON-encoded, as the + launcher writes them (one string, lines joined by an escaped newline).""" + joined = "\\n".join(("X-Subrouter-Agent: claude", *headers)) + return ( + '{"env":{"ANTHROPIC_BASE_URL":"http://127.0.0.1:31415/v1",' + '"ANTHROPIC_AUTH_TOKEN":"srt_test_only_not_a_real_token",' + f'"ANTHROPIC_CUSTOM_HEADERS":"{joined}"}}}}' + ) + + def expect(condition: bool, message: str, failures: list[str]) -> None: if not condition: failures.append(message) @@ -277,6 +298,51 @@ def test_private_settings_after_option_terminator_is_not_bound(failures: list[st expect(BOUND_KEY not in env, f"after --: a literal prompt token was bound: {env.get(BOUND_KEY)!r}", failures) +def test_pinned_launch_records_its_account(failures: list[str]) -> None: + # The pin comes from the launcher's own routing headers, not its argv, so + # it holds however the launcher was invoked. It does not need the marker. + for headers, label in ( + (("X-Subrouter-Account-ID: me@example.com",), "last header"), + (("X-Subrouter-Account-ID: me@example.com", "X-Subrouter-Retry: persist"), "middle header"), + ): + code, env, _, stderr = run_wrapper( + lambda tmpdir: ["--settings", str(private_settings(tmpdir, body=settings_with_headers(*headers)))], + marker=None, + ) + expect_launched(f"pinned launch ({label})", code, env, stderr, failures) + expect( + env.get(ACCOUNT_KEY) == "me@example.com", + f"pinned launch ({label}): recorded account = {env.get(ACCOUNT_KEY)!r}", + failures, + ) + + +def test_pooled_or_unusable_account_is_not_recorded(failures: list[str]) -> None: + for headers, label in ( + ((), "pooled"), + (("X-Subrouter-Preferred-Account-ID: me@example.com",), "preferred only"), + (("X-Subrouter-Account-ID: me example.com",), "space in id"), + (("X-Subrouter-Account-ID: -me",), "option-like id"), + ): + code, env, _, stderr = run_wrapper( + lambda tmpdir: ["--settings", str(private_settings(tmpdir, body=settings_with_headers(*headers)))], + marker=None, + inherited_account="stale@example.com", + ) + expect_launched(label, code, env, stderr, failures) + expect(ACCOUNT_KEY not in env, f"{label}: recorded account = {env.get(ACCOUNT_KEY)!r}", failures) + + +def test_plain_claude_drops_an_inherited_account(failures: list[str]) -> None: + code, env, _, stderr = run_wrapper( + lambda tmpdir: ["--model", "opus"], + marker=None, + inherited_account="me@example.com", + ) + expect_launched("inherited account", code, env, stderr, failures) + expect(ACCOUNT_KEY not in env, f"inherited account survived: {env.get(ACCOUNT_KEY)!r}", failures) + + def main() -> int: if ensure_node_on_path() is None: print("SKIP: node runtime not found; the wrapper's settings merge needs node") @@ -290,6 +356,9 @@ def main() -> int: test_non_exact_marker_text_is_not_bound(failures) test_private_settings_without_marker_is_not_bound(failures) test_private_settings_after_option_terminator_is_not_bound(failures) + test_pinned_launch_records_its_account(failures) + test_pooled_or_unusable_account_is_not_recorded(failures) + test_plain_claude_drops_an_inherited_account(failures) if failures: print("FAIL: the claude wrapper does not bind the Subrouter resume marker to a proven sr launch") for failure in failures: