fix: bound restored agent startup input

lawrencecchen · lawrencecchen · commit 9cd3d82dfd3b · 2026-04-22T15:47:01.000-07:00
diff --git a/Sources/RestorableAgentSession.swift b/Sources/RestorableAgentSession.swift
@@ -50,6 +50,10 @@ struct AgentLaunchCommandSnapshot: Codable, Equatable, Sendable {
     var source: String?
 }
 
+fileprivate func shellSingleQuoted(_ value: String) -> String {
+    "'" + value.replacingOccurrences(of: "'", with: "'\\''") + "'"
+}
+
 private enum AgentResumeCommandBuilder {
     private static let claudeValueOptions: Set<String> = [
         "--add-dir",
@@ -647,13 +651,11 @@ private enum AgentResumeCommandBuilder {
         }
         return trimmed
     }
-
-    private static func shellSingleQuoted(_ value: String) -> String {
-        "'" + value.replacingOccurrences(of: "'", with: "'\\''") + "'"
-    }
 }
 
 struct SessionRestorableAgentSnapshot: Codable, Sendable {
+    static let maxInlineStartupInputBytes = 900
+
     var kind: RestorableAgentKind
     var sessionId: String
     var workingDirectory: String?
@@ -666,6 +668,88 @@ struct SessionRestorableAgentSnapshot: Codable, Sendable {
             workingDirectory: workingDirectory
         )
     }
+
+    func resumeStartupInput(
+        fileManager: FileManager = .default,
+        temporaryDirectory: URL = FileManager.default.temporaryDirectory
+    ) -> String? {
+        guard let command = resumeCommand else { return nil }
+
+        let inlineInput = command + "\n"
+        guard inlineInput.utf8.count > Self.maxInlineStartupInputBytes else {
+            return inlineInput
+        }
+        guard let scriptURL = AgentResumeScriptStore.writeLauncherScript(
+            command: command,
+            kind: kind,
+            sessionId: sessionId,
+            fileManager: fileManager,
+            temporaryDirectory: temporaryDirectory
+        ) else {
+            return nil
+        }
+
+        let scriptInput = "/bin/zsh \(shellSingleQuoted(scriptURL.path))\n"
+        return scriptInput.utf8.count <= Self.maxInlineStartupInputBytes ? scriptInput : nil
+    }
+}
+
+private enum AgentResumeScriptStore {
+    private static let directoryName = "cmux-agent-resume"
+    private static let scriptTTL: TimeInterval = 24 * 60 * 60
+
+    static func writeLauncherScript(
+        command: String,
+        kind: RestorableAgentKind,
+        sessionId: String,
+        fileManager: FileManager,
+        temporaryDirectory: URL
+    ) -> URL? {
+        let directoryURL = temporaryDirectory.appendingPathComponent(directoryName, isDirectory: true)
+        do {
+            try fileManager.createDirectory(at: directoryURL, withIntermediateDirectories: true)
+            try? fileManager.setAttributes([.posixPermissions: 0o700], ofItemAtPath: directoryURL.path)
+            pruneOldScripts(in: directoryURL, fileManager: fileManager)
+
+            let safeSessionPrefix = sessionId
+                .prefix(12)
+                .map { character -> Character in
+                    character.isLetter || character.isNumber || character == "-" ? character : "_"
+                }
+            let scriptURL = directoryURL.appendingPathComponent(
+                "\(kind.rawValue)-\(String(safeSessionPrefix))-\(UUID().uuidString).zsh",
+                isDirectory: false
+            )
+            let contents = """
+            #!/bin/zsh
+            rm -f -- "$0" 2>/dev/null || true
+            \(command)
+            """
+            try contents.write(to: scriptURL, atomically: true, encoding: .utf8)
+            try? fileManager.setAttributes([.posixPermissions: 0o600], ofItemAtPath: scriptURL.path)
+            return scriptURL
+        } catch {
+            return nil
+        }
+    }
+
+    private static func pruneOldScripts(in directoryURL: URL, fileManager: FileManager) {
+        guard let scriptURLs = try? fileManager.contentsOfDirectory(
+            at: directoryURL,
+            includingPropertiesForKeys: [.contentModificationDateKey],
+            options: [.skipsHiddenFiles]
+        ) else {
+            return
+        }
+
+        let cutoff = Date().addingTimeInterval(-scriptTTL)
+        for scriptURL in scriptURLs where scriptURL.pathExtension == "zsh" {
+            let values = try? scriptURL.resourceValues(forKeys: [.contentModificationDateKey])
+            if let modified = values?.contentModificationDate, modified < cutoff {
+                try? fileManager.removeItem(at: scriptURL)
+            }
+        }
+    }
 }
 
 private struct RestorableAgentHookSessionRecord: Codable, Sendable {
diff --git a/Sources/Workspace.swift b/Sources/Workspace.swift
@@ -701,8 +701,7 @@ extension Workspace {
             let shouldReplayScrollback = Self.shouldReplaySessionScrollback(
                 restorableAgent: restorableAgent
             )
-            let restoredAgentResumeCommand = restorableAgent?.resumeCommand
-            let restoredAgentResumeInput = restoredAgentResumeCommand.map { $0 + "\n" }
+            let restoredAgentResumeInput = restorableAgent?.resumeStartupInput()
 #if DEBUG
             if let restorableAgent {
                 let sessionPreview = String(restorableAgent.sessionId.prefix(8))
@@ -711,7 +710,7 @@ extension Workspace {
                     "session.restore.agent panel=\(snapshot.id.uuidString.prefix(5)) " +
                     "kind=\(restorableAgent.kind.rawValue) session=\(sessionPreview) " +
                     "hasLaunch=\(restorableAgent.launchCommand == nil ? 0 : 1) " +
-                    "launchArgc=\(launchArgc) hasResume=\(restoredAgentResumeCommand == nil ? 0 : 1) " +
+                    "launchArgc=\(launchArgc) hasResume=\(restoredAgentResumeInput == nil ? 0 : 1) " +
                     "replayScrollback=\(shouldReplayScrollback ? 1 : 0)"
                 )
             }
diff --git a/cmuxTests/SessionPersistenceTests.swift b/cmuxTests/SessionPersistenceTests.swift
@@ -1472,6 +1472,110 @@ final class SocketListenerAcceptPolicyTests: XCTestCase {
         )
     }
 
+    func testRestorableAgentStartupInputUsesInlineCommandWhenShort() {
+        let snapshot = SessionRestorableAgentSnapshot(
+            kind: .claude,
+            sessionId: "claude-session-123",
+            workingDirectory: "/tmp/cmux project",
+            launchCommand: AgentLaunchCommandSnapshot(
+                launcher: "claude",
+                executablePath: "/opt/Claude Code/bin/claude",
+                arguments: [
+                    "/opt/Claude Code/bin/claude",
+                    "--model",
+                    "sonnet"
+                ],
+                workingDirectory: "/tmp/cmux project",
+                environment: nil,
+                capturedAt: 123,
+                source: "environment"
+            )
+        )
+
+        XCTAssertEqual(snapshot.resumeStartupInput(), snapshot.resumeCommand.map { $0 + "\n" })
+    }
+
+    func testRestorableAgentStartupInputUsesLauncherScriptWhenCommandExceedsTerminalInputBudget() throws {
+        let tempDir = FileManager.default.temporaryDirectory
+            .appendingPathComponent("cmux-agent-resume-test-\(UUID().uuidString)", isDirectory: true)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: tempDir) }
+
+        let longPath = "/tmp/" + String(repeating: "nested-path-", count: 120)
+        let snapshot = SessionRestorableAgentSnapshot(
+            kind: .codex,
+            sessionId: "019dad34-d218-7943-b81a-eddac5c87951",
+            workingDirectory: "/tmp/repo",
+            launchCommand: AgentLaunchCommandSnapshot(
+                launcher: "codex",
+                executablePath: "/Users/example/.bun/bin/codex",
+                arguments: [
+                    "/Users/example/.bun/bin/codex",
+                    "--model",
+                    "gpt-5.4",
+                    "--add-dir",
+                    longPath,
+                    "initial prompt should not replay"
+                ],
+                workingDirectory: "/tmp/repo",
+                environment: ["CODEX_HOME": "/tmp/codex"],
+                capturedAt: 123,
+                source: "environment"
+            )
+        )
+
+        let input = try XCTUnwrap(snapshot.resumeStartupInput(temporaryDirectory: tempDir))
+        XCTAssertLessThanOrEqual(input.utf8.count, SessionRestorableAgentSnapshot.maxInlineStartupInputBytes)
+        XCTAssertTrue(input.hasPrefix("/bin/zsh '"))
+        XCTAssertFalse(input.contains(longPath))
+
+        let trimmedInput = input.trimmingCharacters(in: .whitespacesAndNewlines)
+        let prefix = "/bin/zsh '"
+        let scriptPath = String(trimmedInput.dropFirst(prefix.count).dropLast())
+        let scriptContents = try String(contentsOfFile: scriptPath, encoding: .utf8)
+        XCTAssertTrue(scriptContents.contains(longPath))
+        XCTAssertTrue(scriptContents.contains("'resume'"))
+        XCTAssertTrue(scriptContents.contains("'019dad34-d218-7943-b81a-eddac5c87951'"))
+
+        let attributes = try FileManager.default.attributesOfItem(atPath: scriptPath)
+        let permissions = try XCTUnwrap(attributes[.posixPermissions] as? NSNumber).intValue & 0o777
+        XCTAssertEqual(permissions, 0o600)
+    }
+
+    func testRestorableAgentStartupInputSkipsOversizedCommandWhenScriptCannotBeWritten() throws {
+        let tempDir = FileManager.default.temporaryDirectory
+            .appendingPathComponent("cmux-agent-resume-test-\(UUID().uuidString)", isDirectory: true)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: tempDir) }
+
+        let blockedDirectory = tempDir.appendingPathComponent("not-a-directory", isDirectory: false)
+        try "occupied".write(to: blockedDirectory, atomically: true, encoding: .utf8)
+        let longPath = "/tmp/" + String(repeating: "nested-path-", count: 120)
+        let snapshot = SessionRestorableAgentSnapshot(
+            kind: .codex,
+            sessionId: "019dad34-d218-7943-b81a-eddac5c87951",
+            workingDirectory: "/tmp/repo",
+            launchCommand: AgentLaunchCommandSnapshot(
+                launcher: "codex",
+                executablePath: "/Users/example/.bun/bin/codex",
+                arguments: [
+                    "/Users/example/.bun/bin/codex",
+                    "--model",
+                    "gpt-5.4",
+                    "--add-dir",
+                    longPath,
+                    "initial prompt should not replay"
+                ],
+                workingDirectory: "/tmp/repo",
+                environment: ["CODEX_HOME": "/tmp/codex"],
+                capturedAt: 123,
+                source: "environment"
+            )
+        )
+
+        XCTAssertNil(snapshot.resumeStartupInput(temporaryDirectory: blockedDirectory))
+    }
+
     func testClaudeResumeCommandPreservesDangerouslySkipPermissionsAndObservedEnvironment() {
         let snapshot = SessionRestorableAgentSnapshot(
             kind: .claude,

Original file line number	Diff line number	Diff line change
`@@ -701,8 +701,7 @@ extension Workspace {`
`701`	`701`	`let shouldReplayScrollback = Self.shouldReplaySessionScrollback(`
`702`	`702`	`restorableAgent: restorableAgent`
`703`	`703`	`)`
`704`		`- let restoredAgentResumeCommand = restorableAgent?.resumeCommand`
`705`		`- let restoredAgentResumeInput = restoredAgentResumeCommand.map { $0 + "\n" }`
	`704`	`+ let restoredAgentResumeInput = restorableAgent?.resumeStartupInput()`
`706`	`705`	`#if DEBUG`
`707`	`706`	`if let restorableAgent {`
`708`	`707`	`let sessionPreview = String(restorableAgent.sessionId.prefix(8))`
`@@ -711,7 +710,7 @@ extension Workspace {`
`711`	`710`	`"session.restore.agent panel=\(snapshot.id.uuidString.prefix(5)) " +`
`712`	`711`	`"kind=\(restorableAgent.kind.rawValue) session=\(sessionPreview) " +`
`713`	`712`	`"hasLaunch=\(restorableAgent.launchCommand == nil ? 0 : 1) " +`
`714`		`- "launchArgc=\(launchArgc) hasResume=\(restoredAgentResumeCommand == nil ? 0 : 1) " +`
	`713`	`+ "launchArgc=\(launchArgc) hasResume=\(restoredAgentResumeInput == nil ? 0 : 1) " +`
`715`	`714`	`"replayScrollback=\(shouldReplayScrollback ? 1 : 0)"`
`716`	`715`	`)`
`717`	`716`	`}`