diff --git a/Sources/RestorableAgentSession.swift b/Sources/RestorableAgentSession.swift index b26e06a6292..4d0f0d8539c 100644 --- a/Sources/RestorableAgentSession.swift +++ b/Sources/RestorableAgentSession.swift @@ -694,7 +694,7 @@ struct SessionRestorableAgentSnapshot: Codable, Sendable { AgentResumeCommandBuilder.resumeShellCommand( kind: kind, sessionId: sessionId, - launchCommand: launchCommand, + launchCommand: trustedLaunchCommandForSessionRestore, workingDirectory: workingDirectory, registrationOverride: registration ) @@ -704,12 +704,50 @@ struct SessionRestorableAgentSnapshot: Codable, Sendable { AgentResumeCommandBuilder.forkShellCommand( kind: kind, sessionId: sessionId, - launchCommand: launchCommand, + launchCommand: trustedLaunchCommandForSessionRestore, workingDirectory: workingDirectory, registrationOverride: registration ) } + var trustedLaunchCommandForSessionRestore: AgentLaunchCommandSnapshot? { + guard let launchCommand else { return nil } + guard AgentLaunchCaptureTrust.launcherDescribesKind(launchCommand.launcher, kind: kind.rawValue), + !AgentLaunchCaptureTrust.argvLooksLikeShellWrapper(launchCommand.arguments) else { + return nil + } + return launchCommand + } + + func repairedForSessionRestore(fallbackWorkingDirectory: String?) -> SessionRestorableAgentSnapshot { + let trustedLaunchCommand = trustedLaunchCommandForSessionRestore + var repaired = self + repaired.launchCommand = trustedLaunchCommand + + let fallbackWorkingDirectory = Self.normalizedWorkingDirectory(fallbackWorkingDirectory) + if trustedLaunchCommand == nil { + if repaired.workingDirectory == nil + || Self.normalizedWorkingDirectory(repaired.workingDirectory) + == Self.normalizedWorkingDirectory(launchCommand?.workingDirectory) { + repaired.workingDirectory = fallbackWorkingDirectory + } + } else if repaired.workingDirectory == nil { + repaired.workingDirectory = Self.normalizedWorkingDirectory( + trustedLaunchCommand?.workingDirectory + ) ?? fallbackWorkingDirectory + } + + return repaired + } + + private static func normalizedWorkingDirectory(_ workingDirectory: String?) -> String? { + guard let trimmed = workingDirectory?.trimmingCharacters(in: .whitespacesAndNewlines), + !trimmed.isEmpty else { + return nil + } + return ((trimmed as NSString).expandingTildeInPath as NSString).standardizingPath + } + func resumeStartupInput( fileManager: FileManager = .default, temporaryDirectory: URL = FileManager.default.temporaryDirectory, @@ -741,7 +779,7 @@ struct SessionRestorableAgentSnapshot: Codable, Sendable { // the current directory (no cd), so the post-exit shell must not force the launch dir. workingDirectory: registration?.cwd == .ignore ? nil - : (workingDirectory ?? launchCommand?.workingDirectory) + : (workingDirectory ?? trustedLaunchCommandForSessionRestore?.workingDirectory) ) else { return nil } diff --git a/Sources/SessionPersistence.swift b/Sources/SessionPersistence.swift index 3717255e54b..b25ff038b04 100644 --- a/Sources/SessionPersistence.swift +++ b/Sources/SessionPersistence.swift @@ -356,6 +356,24 @@ nonisolated struct SurfaceResumeBindingSnapshot: Codable, Equatable, Sendable { autoResume == true } + var trustedForSessionRestore: SurfaceResumeBindingSnapshot? { + isPoisonedAgentHookShellWrapperResume ? nil : self + } + + var isPoisonedAgentHookShellWrapperResume: Bool { + guard isAgentHookBinding, + let tokens = SurfaceResumeCommandCanonicalizer.tokens(from: command) else { + return false + } + let commandStart = SurfaceResumeCommandCanonicalizer.commandStartIndexAfterCwdGuard(tokens) + guard commandStart + 1 < tokens.endIndex else { return false } + let executable = (tokens[commandStart] as NSString).lastPathComponent.lowercased() + let shells: Set = ["sh", "bash", "zsh", "dash", "fish", "csh", "tcsh", "ksh"] + guard shells.contains(executable) else { return false } + let resumeWord = tokens[commandStart + 1] + return resumeWord == "resume" || resumeWord == "--resume" || resumeWord.hasPrefix("--resume=") + } + func shouldYieldToDetectedSurfaceResumeBinding(_ detectedBinding: SurfaceResumeBindingSnapshot) -> Bool { detectedBinding.isProcessDetected && (isProcessDetected || isAgentHookBinding) } @@ -679,6 +697,17 @@ enum SurfaceResumeCommandCanonicalizer { return ((rawValue as NSString).expandingTildeInPath as NSString).standardizingPath } + static func commandStartIndexAfterCwdGuard(_ tokens: [String]) -> Int { + guard let first = tokens.first, + first == "{" || first == "cd" else { + return tokens.startIndex + } + guard let andIndex = tokens.firstIndex(of: "&&") else { + return tokens.startIndex + } + return tokens.index(after: andIndex) + } + static func shellQuoted(_ value: String) -> String { guard !value.isEmpty else { return "''" } let allowed = CharacterSet(charactersIn: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_+-=./:@%") @@ -1867,6 +1896,89 @@ struct AppSessionSnapshot: Codable, Sendable { var windows: [SessionWindowSnapshot] } +enum SessionSnapshotRepairer { + struct Result { + var snapshot: AppSessionSnapshot + var didRepair: Bool + } + + static func repair(_ snapshot: AppSessionSnapshot) -> Result { + var didRepair = false + var repaired = snapshot + repaired.windows = repaired.windows.map { window in + repair(window, didRepair: &didRepair) + } + return Result(snapshot: repaired, didRepair: didRepair) + } + + private static func repair( + _ window: SessionWindowSnapshot, + didRepair: inout Bool + ) -> SessionWindowSnapshot { + var repaired = window + repaired.tabManager.workspaces = repaired.tabManager.workspaces.map { workspace in + repair(workspace, didRepair: &didRepair) + } + return repaired + } + + private static func repair( + _ workspace: SessionWorkspaceSnapshot, + didRepair: inout Bool + ) -> SessionWorkspaceSnapshot { + var repaired = workspace + repaired.panels = repaired.panels.map { panel in + repair(panel, workspaceDirectory: workspace.currentDirectory, didRepair: &didRepair) + } + return repaired + } + + private static func repair( + _ panel: SessionPanelSnapshot, + workspaceDirectory: String, + didRepair: inout Bool + ) -> SessionPanelSnapshot { + guard var terminal = panel.terminal else { return panel } + let fallbackWorkingDirectory = firstNormalizedDirectory( + terminal.workingDirectory, + panel.directory, + workspaceDirectory + ) + + if let resumeBinding = terminal.resumeBinding { + let trustedBinding = resumeBinding.trustedForSessionRestore + if trustedBinding == nil { + didRepair = true + } + terminal.resumeBinding = trustedBinding + } + + if let agent = terminal.agent { + let repairedAgent = agent.repairedForSessionRestore( + fallbackWorkingDirectory: fallbackWorkingDirectory + ) + if agent.launchCommand != repairedAgent.launchCommand + || agent.workingDirectory != repairedAgent.workingDirectory { + didRepair = true + } + terminal.agent = repairedAgent + } + + var repaired = panel + repaired.terminal = terminal + return repaired + } + + private static func firstNormalizedDirectory(_ candidates: String?...) -> String? { + for candidate in candidates { + if let normalized = SurfaceResumeCommandCanonicalizer.normalizedCWD(candidate) { + return normalized + } + } + return nil + } +} + enum SessionPersistenceStore { enum SnapshotLoadOutcome { case loaded(AppSessionSnapshot) @@ -1885,7 +1997,11 @@ enum SessionPersistenceStore { guard let snapshot = try? decoder.decode(AppSessionSnapshot.self, from: data) else { return .unusable } guard snapshot.version == SessionSnapshotSchema.currentVersion else { return .unusable } guard !snapshot.windows.isEmpty else { return .unusable } - return .loaded(snapshot) + let repairResult = SessionSnapshotRepairer.repair(snapshot) + if repairResult.didRepair { + _ = save(repairResult.snapshot, fileURL: fileURL) + } + return .loaded(repairResult.snapshot) } static func load(fileURL: URL? = nil) -> AppSessionSnapshot? { diff --git a/Sources/Workspace.swift b/Sources/Workspace.swift index 501993a58db..c05f3484ae6 100644 --- a/Sources/Workspace.swift +++ b/Sources/Workspace.swift @@ -1101,7 +1101,7 @@ extension Workspace { ) -> String { let bootstrapCommand = bootstrap.joined(separator: " && ") + " && " let words = surfaceResumeShellWords(in: command) - let commandStart = hermesAgentCommandStartIndexAfterCwdGuard(words) + let commandStart = surfaceResumeCommandStartIndexAfterCwdGuard(words) guard commandStart < words.endIndex else { return bootstrapCommand + command } @@ -1137,7 +1137,7 @@ extension Workspace { nonisolated private static func hermesAgentCommandByRemovingBootstrapPrefix(_ command: String) -> String { let words = surfaceResumeShellWords(in: command) - var scanIndex = hermesAgentCommandStartIndexAfterCwdGuard(words) + var scanIndex = surfaceResumeCommandStartIndexAfterCwdGuard(words) guard scanIndex < words.endIndex else { return command } let removeStartIndex = scanIndex var removedBootstrap = false @@ -1235,12 +1235,12 @@ extension Workspace { nonisolated private static func hermesAgentWordsAfterCwdGuard( _ words: [SurfaceResumeShellWord] ) -> [SurfaceResumeShellWord] { - let commandStart = hermesAgentCommandStartIndexAfterCwdGuard(words) + let commandStart = surfaceResumeCommandStartIndexAfterCwdGuard(words) guard commandStart < words.endIndex else { return [] } return Array(words[commandStart...]) } - nonisolated private static func hermesAgentCommandStartIndexAfterCwdGuard( + nonisolated private static func surfaceResumeCommandStartIndexAfterCwdGuard( _ words: [SurfaceResumeShellWord] ) -> Int { guard let first = words.first, @@ -1637,7 +1637,7 @@ extension Workspace { ) -> UUID? { switch snapshot.type { case .terminal: - let resumeBinding = snapshot.terminal?.resumeBinding + let resumeBinding = snapshot.terminal?.resumeBinding?.trustedForSessionRestore let restorableAgent = snapshot.terminal?.agent let restoredHibernation = snapshot.terminal?.hibernation let autoResumeAgentSessions = AgentSessionAutoResumeSettings.isEnabled() diff --git a/cmuxTests/SessionPersistenceTests.swift b/cmuxTests/SessionPersistenceTests.swift index 6ffe59caaff..6efb67f7ab4 100644 --- a/cmuxTests/SessionPersistenceTests.swift +++ b/cmuxTests/SessionPersistenceTests.swift @@ -192,6 +192,94 @@ final class SessionPersistenceTests: XCTestCase { XCTAssertEqual(visibleFrame.y, 25, accuracy: 0.001) } + @MainActor + func testLoadRepairsAndPersistsPoisonedAgentHookShellWrapperResumeBinding() throws { + let tempDir = FileManager.default.temporaryDirectory + .appendingPathComponent("cmux-session-tests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: tempDir) } + + let snapshotURL = tempDir.appendingPathComponent("session.json", isDirectory: false) + let source = Workspace() + let sourcePanelId = try XCTUnwrap(source.focusedPanelId) + let sessionId = "codex-load-repair-session" + let bindingIndex = SurfaceResumeBindingIndex(bindingsByPanel: [ + SurfaceResumeBindingIndex.PanelKey(workspaceId: source.id, panelId: sourcePanelId): SurfaceResumeBindingSnapshot( + name: "Codex", + kind: "codex", + command: "{ cd -- '/tmp/right' 2>/dev/null || [ ! -d '/tmp/right' ]; } && 'bash' 'resume' '\(sessionId)'", + cwd: "/tmp/right", + checkpointId: sessionId, + source: "agent-hook", + autoResume: true, + updatedAt: 10 + ), + ]) + let workspaceSnapshot = source.sessionSnapshot( + includeScrollback: false, + surfaceResumeBindingIndex: bindingIndex + ) + let appSnapshot = makeSnapshot(workspaceSnapshot: workspaceSnapshot) + XCTAssertTrue(SessionPersistenceStore.save(appSnapshot, fileURL: snapshotURL)) + + let loaded = try XCTUnwrap(SessionPersistenceStore.load(fileURL: snapshotURL)) + let loadedBinding = loaded.windows.first?.tabManager.workspaces.first?.panels.first?.terminal?.resumeBinding + XCTAssertNil(loadedBinding) + + let persistedData = try Data(contentsOf: snapshotURL) + let persisted = try JSONDecoder().decode(AppSessionSnapshot.self, from: persistedData) + let persistedBinding = persisted.windows.first?.tabManager.workspaces.first?.panels.first?.terminal?.resumeBinding + XCTAssertNil(persistedBinding) + } + + @MainActor + func testLoadRepairsWrongForkAgentLaunchCommandAndRecoversWorkingDirectory() throws { + let tempDir = FileManager.default.temporaryDirectory + .appendingPathComponent("cmux-session-tests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: tempDir) } + + let snapshotURL = tempDir.appendingPathComponent("session.json", isDirectory: false) + let source = Workspace() + let sourcePanelId = try XCTUnwrap(source.focusedPanelId) + source.updatePanelDirectory(panelId: sourcePanelId, directory: "/tmp/right") + var workspaceSnapshot = source.sessionSnapshot(includeScrollback: false) + let panelIndex = try XCTUnwrap(workspaceSnapshot.panels.firstIndex { $0.id == sourcePanelId }) + var panel = workspaceSnapshot.panels[panelIndex] + var terminal = try XCTUnwrap(panel.terminal) + terminal.workingDirectory = "/tmp/right" + terminal.agent = SessionRestorableAgentSnapshot( + kind: .codex, + sessionId: "codex-load-repair-session", + workingDirectory: "/tmp/wrong", + launchCommand: AgentLaunchCommandSnapshot( + launcher: "claude", + executablePath: "/usr/local/bin/claude", + arguments: ["/usr/local/bin/claude", "--resume", "claude-session"], + workingDirectory: "/tmp/wrong", + environment: nil, + capturedAt: 123, + source: "process" + ) + ) + panel.directory = "/tmp/right" + panel.terminal = terminal + workspaceSnapshot.panels[panelIndex] = panel + let appSnapshot = makeSnapshot(workspaceSnapshot: workspaceSnapshot) + XCTAssertTrue(SessionPersistenceStore.save(appSnapshot, fileURL: snapshotURL)) + + let loaded = try XCTUnwrap(SessionPersistenceStore.load(fileURL: snapshotURL)) + let loadedAgent = try XCTUnwrap( + loaded.windows.first?.tabManager.workspaces.first?.panels.first?.terminal?.agent + ) + XCTAssertNil(loadedAgent.launchCommand) + XCTAssertEqual(loadedAgent.workingDirectory, "/tmp/right") + XCTAssertEqual( + loadedAgent.resumeCommand, + "{ cd -- '/tmp/right' 2>/dev/null || [ ! -d '/tmp/right' ]; } && 'codex' 'resume' 'codex-load-repair-session'" + ) + } + func testLoadReopenSessionSnapshotRequiresPreviousSnapshotFile() throws { let tempDir = FileManager.default.temporaryDirectory .appendingPathComponent("cmux-session-tests-\(UUID().uuidString)", isDirectory: true) @@ -1882,6 +1970,23 @@ final class SessionPersistenceTests: XCTestCase { ) } + private func makeSnapshot(workspaceSnapshot: SessionWorkspaceSnapshot) -> AppSessionSnapshot { + let window = SessionWindowSnapshot( + frame: SessionRectSnapshot(x: 10, y: 20, width: 900, height: 700), + display: nil, + tabManager: SessionTabManagerSnapshot( + selectedWorkspaceIndex: 0, + workspaces: [workspaceSnapshot] + ), + sidebar: SessionSidebarSnapshot(isVisible: true, selection: .tabs, width: 240) + ) + return AppSessionSnapshot( + version: SessionSnapshotSchema.currentVersion, + createdAt: Date().timeIntervalSince1970, + windows: [window] + ) + } + private func fileNumber(for fileURL: URL) throws -> Int { let attributes = try FileManager.default.attributesOfItem(atPath: fileURL.path) return try XCTUnwrap(attributes[.systemFileNumber] as? Int) @@ -1889,6 +1994,32 @@ final class SessionPersistenceTests: XCTestCase { } final class SocketListenerAcceptPolicyTests: XCTestCase { + func testPersistedAgentSnapshotDropsShellWrapperLaunchCommandWhenRenderingResume() { + let snapshot = SessionRestorableAgentSnapshot( + kind: .codex, + sessionId: "codex-session-123", + workingDirectory: "/tmp/repo", + launchCommand: AgentLaunchCommandSnapshot( + launcher: "codex", + executablePath: "bash", + arguments: [ + "bash", + "-c", + "payload=\"$1\"; shift; \"$@\" <\"$payload\" &", + ], + workingDirectory: "/tmp/dispatch-shell", + environment: ["CODEX_HOME": "/tmp/codex"], + capturedAt: 123, + source: "process" + ) + ) + + XCTAssertEqual( + snapshot.resumeCommand, + "{ cd -- '/tmp/repo' 2>/dev/null || [ ! -d '/tmp/repo' ]; } && 'codex' 'resume' 'codex-session-123'" + ) + } + func testClaudeResumeCommandRoutesThroughWrapperInsteadOfCapturedRealBinary() { // The captured launch executable is the real claude binary // (CMUX_AGENT_LAUNCH_EXECUTABLE). Resuming with it directly bypasses @@ -5796,6 +5927,88 @@ extension SessionPersistenceTests { ) } + @MainActor + func testRestoreDropsPoisonedAgentHookShellWrapperResumeBindingAndUsesAgentSnapshot() throws { + try withAutoResumeAgentSessionsEnabled { + let source = Workspace() + let sourcePanelId = try XCTUnwrap(source.focusedPanelId) + let sessionId = "019eba43-1e37-78d2-98df-d1983200273d" + let sourceIndex = try makeRestorableAgentIndex( + workspaceId: source.id, + panelId: sourcePanelId, + sessionId: sessionId, + arguments: [ + "/usr/local/bin/codex", + "--yolo", + ] + ) + let bindingIndex = SurfaceResumeBindingIndex(bindingsByPanel: [ + SurfaceResumeBindingIndex.PanelKey(workspaceId: source.id, panelId: sourcePanelId): SurfaceResumeBindingSnapshot( + name: "Codex", + kind: "codex", + command: "{ cd -- '/Users/lawrence/fun/cmuxterm-hq' 2>/dev/null || [ ! -d '/Users/lawrence/fun/cmuxterm-hq' ]; } && 'bash' 'resume' '\(sessionId)' '-c' 'payload=\"$1\"; shift; \"$@\" <\"$payload\" &'", + cwd: "/Users/lawrence/fun/cmuxterm-hq", + checkpointId: sessionId, + source: "agent-hook", + autoResume: true, + updatedAt: 10 + ), + ]) + let snapshot = source.sessionSnapshot( + includeScrollback: false, + restorableAgentIndex: sourceIndex, + surfaceResumeBindingIndex: bindingIndex + ) + + let restored = Workspace() + restored.restoreSessionSnapshot(snapshot) + let restoredPanelId = try XCTUnwrap(restored.focusedPanelId) + let restoredPanel = try XCTUnwrap(restored.terminalPanel(for: restoredPanelId)) + let payload = try restoredStartupPayload(for: restoredPanel) + + XCTAssertTrue(payload.contains("/usr/local/bin/codex"), payload) + XCTAssertTrue(payload.contains("resume"), payload) + XCTAssertTrue(payload.contains(sessionId), payload) + XCTAssertTrue(payload.contains("--yolo"), payload) + XCTAssertFalse(payload.contains("'bash' 'resume'"), payload) + XCTAssertNil(restored.sessionSnapshot(includeScrollback: false).panels.first?.terminal?.resumeBinding) + } + } + + @MainActor + func testRestoreDropsPoisonedAgentHookShellWrapperResumeBindingWithoutAgentSnapshot() throws { + try withAutoResumeAgentSessionsEnabled { + let source = Workspace() + let sourcePanelId = try XCTUnwrap(source.focusedPanelId) + let sessionId = "019eb982-d798-7123-aa2b-93326ff3bd08" + let bindingIndex = SurfaceResumeBindingIndex(bindingsByPanel: [ + SurfaceResumeBindingIndex.PanelKey(workspaceId: source.id, panelId: sourcePanelId): SurfaceResumeBindingSnapshot( + name: "Codex", + kind: "codex", + command: "{ cd -- '/Users/lawrence/fun/cmuxterm-hq' 2>/dev/null || [ ! -d '/Users/lawrence/fun/cmuxterm-hq' ]; } && 'sh' 'resume' '\(sessionId)' '-c' 'payload=\"$1\"; shift; \"$@\" <\"$payload\" &'", + cwd: "/Users/lawrence/fun/cmuxterm-hq", + checkpointId: sessionId, + source: "agent-hook", + autoResume: true, + updatedAt: 10 + ), + ]) + let snapshot = source.sessionSnapshot( + includeScrollback: false, + surfaceResumeBindingIndex: bindingIndex + ) + + let restored = Workspace() + restored.restoreSessionSnapshot(snapshot) + let restoredPanelId = try XCTUnwrap(restored.focusedPanelId) + let restoredPanel = try XCTUnwrap(restored.terminalPanel(for: restoredPanelId)) + + XCTAssertNil(restoredPanel.surface.debugInitialCommand()) + XCTAssertFalse(restoredPanel.surface.debugInitialInputMetadata().hasInitialInput) + XCTAssertNil(restored.sessionSnapshot(includeScrollback: false).panels.first?.terminal?.resumeBinding) + } + } + @MainActor func testRestoreScopesSurfaceResumeBindingEnvironmentToInitialInput() throws { let source = Workspace()