Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions Sources/RestorableAgentSession.swift
Original file line number Diff line number Diff line change
Expand Up @@ -694,7 +694,7 @@ struct SessionRestorableAgentSnapshot: Codable, Sendable {
AgentResumeCommandBuilder.resumeShellCommand(
kind: kind,
sessionId: sessionId,
launchCommand: launchCommand,
launchCommand: trustedLaunchCommandForResume,
workingDirectory: workingDirectory,
registrationOverride: registration
)
Expand All @@ -704,12 +704,21 @@ struct SessionRestorableAgentSnapshot: Codable, Sendable {
AgentResumeCommandBuilder.forkShellCommand(
kind: kind,
sessionId: sessionId,
launchCommand: launchCommand,
launchCommand: trustedLaunchCommandForResume,
workingDirectory: workingDirectory,
registrationOverride: registration
)
}

private var trustedLaunchCommandForResume: 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 resumeStartupInput(
fileManager: FileManager = .default,
temporaryDirectory: URL = FileManager.default.temporaryDirectory,
Expand Down Expand Up @@ -741,7 +750,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 ?? trustedLaunchCommandForResume?.workingDirectory)
) else {
return nil
}
Expand Down
32 changes: 27 additions & 5 deletions Sources/Workspace.swift
Original file line number Diff line number Diff line change
Expand Up @@ -1055,6 +1055,28 @@ extension Workspace {
return effectiveBinding
}

nonisolated private static func trustedSurfaceResumeBindingForRestore(
_ resumeBinding: SurfaceResumeBindingSnapshot?
) -> SurfaceResumeBindingSnapshot? {
guard let resumeBinding else { return nil }
guard !agentHookBindingLooksLikeShellWrapperResume(resumeBinding) else { return nil }
return resumeBinding
}

nonisolated private static func agentHookBindingLooksLikeShellWrapperResume(
_ binding: SurfaceResumeBindingSnapshot
) -> Bool {
guard binding.isAgentHookBinding else { return false }
let words = surfaceResumeShellWords(in: binding.command)
let commandStart = surfaceResumeCommandStartIndexAfterCwdGuard(words)
guard commandStart + 1 < words.endIndex else { return false }
let executable = (words[commandStart].value as NSString).lastPathComponent.lowercased()
let shells: Set<String> = ["sh", "bash", "zsh", "dash", "fish", "csh", "tcsh", "ksh"]
guard shells.contains(executable) else { return false }
let resumeWord = words[commandStart + 1].value
return resumeWord == "resume" || resumeWord == "--resume" || resumeWord.hasPrefix("--resume=")
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 fish included in shell set may cause false positives

fish is listed alongside traditional POSIX shells, but fish resume <session> is a plausible legitimate invocation for a user whose agent-hook command genuinely calls a Fish script named resume. Because agentHookBindingLooksLikeShellWrapperResume only fires for isAgentHookBinding bindings, the blast radius is narrow, but fish does not appear in the original poisoned binding pattern described in the PR (bash resume …) and adding it increases the heuristic surface without a concrete motivating case.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!


nonisolated private static func hermesAgentSubrouterBindingForStartup(
_ binding: SurfaceResumeBindingSnapshot
) -> SurfaceResumeBindingSnapshot {
Expand Down Expand Up @@ -1101,7 +1123,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
}
Expand Down Expand Up @@ -1137,7 +1159,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
Expand Down Expand Up @@ -1235,12 +1257,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,
Expand Down Expand Up @@ -1637,7 +1659,7 @@ extension Workspace {
) -> UUID? {
switch snapshot.type {
case .terminal:
let resumeBinding = snapshot.terminal?.resumeBinding
let resumeBinding = Self.trustedSurfaceResumeBindingForRestore(snapshot.terminal?.resumeBinding)
let restorableAgent = snapshot.terminal?.agent
let restoredHibernation = snapshot.terminal?.hibernation
let autoResumeAgentSessions = AgentSessionAutoResumeSettings.isEnabled()
Expand Down
108 changes: 108 additions & 0 deletions cmuxTests/SessionPersistenceTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -1889,6 +1889,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
Comment on lines 1994 to 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Test placed in wrong class

testPersistedAgentSnapshotDropsShellWrapperLaunchCommandWhenRenderingResume exercises SessionRestorableAgentSnapshot.resumeCommand filtering, but it lives inside SocketListenerAcceptPolicyTests. The two other new tests correctly land in SessionPersistenceTests. Moving this test keeps the file's grouping consistent and makes it easier to run the persistence test suite in isolation.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Expand Down Expand Up @@ -5796,6 +5822,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()
Expand Down
Loading