diff --git a/Sources/RestorableAgentSession.swift b/Sources/RestorableAgentSession.swift index 5c6de9e029d1..09a9d5a46105 100644 --- a/Sources/RestorableAgentSession.swift +++ b/Sources/RestorableAgentSession.swift @@ -960,6 +960,17 @@ struct RestorableAgentSessionIndex: Sendable { let kind: RestorableAgentKind } + private struct PanelIDKindKey: Hashable { + let panelId: UUID + let kind: RestorableAgentKind + } + + private struct PanelIDKindCandidate { + let panelKey: PanelKey + let entry: Entry + let isAmbiguous: Bool + } + private let entriesByPanel: [PanelKey: Entry] private let entriesByPanelId: [UUID: Entry] @@ -1092,6 +1103,7 @@ struct RestorableAgentSessionIndex: Sendable { } var hookCandidatesBySession: [SessionKey: Entry] = [:] var hookCandidatesByPanelAndKind: [PanelKindKey: Entry] = [:] + var hookCandidatesByPanelIdAndKind: [PanelIDKindKey: PanelIDKindCandidate] = [:] for (kind, registration) in hookKinds { let fileURL = kind.hookStoreFileURL(homeDirectory: homeDirectory) @@ -1146,6 +1158,7 @@ struct RestorableAgentSessionIndex: Sendable { let key = PanelKey(workspaceId: workspaceId, panelId: panelId) let sessionKey = SessionKey(kind: kind, sessionId: normalizedSessionId) let panelKindKey = PanelKindKey(panelKey: key, kind: kind) + let panelIDKindKey = PanelIDKindKey(panelId: panelId, kind: kind) let liveProcessID = liveScopedProcessID( for: effectiveRecord, kind: kind, @@ -1170,6 +1183,25 @@ struct RestorableAgentSessionIndex: Sendable { ) { hookCandidatesByPanelAndKind[panelKindKey] = entry } + if let existingPanelIDCandidate = hookCandidatesByPanelIdAndKind[panelIDKindKey] { + let shouldReplace = shouldReplaceHookEntry( + existing: existingPanelIDCandidate.entry, + incoming: entry + ) + hookCandidatesByPanelIdAndKind[panelIDKindKey] = PanelIDKindCandidate( + panelKey: shouldReplace ? key : existingPanelIDCandidate.panelKey, + entry: shouldReplace ? entry : existingPanelIDCandidate.entry, + isAmbiguous: existingPanelIDCandidate.isAmbiguous || + existingPanelIDCandidate.panelKey != key || + existingPanelIDCandidate.entry.snapshot.sessionId != entry.snapshot.sessionId + ) + } else { + hookCandidatesByPanelIdAndKind[panelIDKindKey] = PanelIDKindCandidate( + panelKey: key, + entry: entry, + isAmbiguous: false + ) + } if shouldReplaceHookEntry( existing: hookCandidatesBySession[sessionKey], incoming: entry @@ -1200,6 +1232,17 @@ struct RestorableAgentSessionIndex: Sendable { let sameKindPanelCandidate = hookCandidatesByPanelAndKind[ PanelKindKey(panelKey: key, kind: detected.snapshot.kind) ] + let sameKindPanelIDCandidate = hookCandidatesByPanelIdAndKind[ + PanelIDKindKey(panelId: key.panelId, kind: detected.snapshot.kind) + ] + // Panel-only restore is safe only when this surface/kind maps back to exactly one + // old workspace/session pair. Stale hook stores can otherwise reuse a surface id + // across old workspaces, or record multiple sessions for the same old workspace and + // surface after an agent restart. In either case, shouldReplaceHookEntry would pick + // one session by recency, so the panel-only fallback must stay ambiguous. + let sameKindStablePanelCandidate = sameKindPanelCandidate ?? ( + sameKindPanelIDCandidate?.isAmbiguous == false ? sameKindPanelIDCandidate?.entry : nil + ) if detected.sessionIDSource == .forkParentFallback, let panelCandidate = sameKindPanelCandidate, Self.hookCandidateRepresentsDetectedProcess( @@ -1213,6 +1256,13 @@ struct RestorableAgentSessionIndex: Sendable { // A nested fork process inside another agent's pane must not displace // that pane's hook-backed identity. continue + } else if detected.sessionIDSource == .inferredLatestSessionFile, + let panelCandidate = sameKindStablePanelCandidate { + // Latest-file detection is ambiguous when multiple panels or restored workspaces share a + // cwd. Prefer the hook-store identity for this stable panel/surface while still carrying + // live process evidence for the restored panel. The workspace UUID can rotate during + // session restore, but the surface id is intentionally reused on the normal restore path. + resolved[key] = processDetectedEntry(snapshot: panelCandidate.snapshot, lifecycle: panelCandidate.lifecycle, updatedAt: panelCandidate.updatedAt, detected: detected) } else if let existing = Self.matchingHookEntry( for: detected.snapshot, resolved: resolved[key], @@ -1222,11 +1272,6 @@ struct RestorableAgentSessionIndex: Sendable { ] ) { resolved[key] = processDetectedEntry(snapshot: detected.snapshot, lifecycle: existing.lifecycle, updatedAt: existing.updatedAt, detected: detected) - } else if detected.sessionIDSource == .inferredLatestSessionFile, - let panelCandidate = sameKindPanelCandidate { - // Latest-file detection is ambiguous when multiple panels share a cwd; preserve the exact - // hook-store identity while still carrying live process evidence for this panel. - resolved[key] = processDetectedEntry(snapshot: panelCandidate.snapshot, lifecycle: panelCandidate.lifecycle, updatedAt: panelCandidate.updatedAt, detected: detected) } else { resolved[key] = processDetectedEntry(snapshot: detected.snapshot, lifecycle: nil, updatedAt: 0, detected: detected) } diff --git a/cmuxTests/RestorableAgentSessionIndexTests.swift b/cmuxTests/RestorableAgentSessionIndexTests.swift index 021ad72c2eb5..4976e58d02d7 100644 --- a/cmuxTests/RestorableAgentSessionIndexTests.swift +++ b/cmuxTests/RestorableAgentSessionIndexTests.swift @@ -695,6 +695,215 @@ final class RestorableAgentSessionIndexTests: XCTestCase { XCTAssertEqual(restoredSessionIds, sessionIds) } + func testPiDetectedLatestSessionDoesNotCollapseExactHookRecordsAfterWorkspaceIdRotation() throws { + let fm = FileManager.default + let root = fm.temporaryDirectory + .appendingPathComponent("cmux-pi-restore-workspace-rotation-\(UUID().uuidString)", isDirectory: true) + defer { try? fm.removeItem(at: root) } + + let cwd = root.appendingPathComponent("repo", isDirectory: true) + let sessionsRoot = root.appendingPathComponent("pi-sessions", isDirectory: true) + let projectDirectory = try XCTUnwrap(PiSessionLocator.projectDirectoryName(for: cwd.path)) + let projectSessions = sessionsRoot.appendingPathComponent(projectDirectory, isDirectory: true) + try fm.createDirectory(at: cwd, withIntermediateDirectories: true) + try fm.createDirectory(at: projectSessions, withIntermediateDirectories: true) + + var registration = CmuxVaultAgentRegistration.builtInPi + registration.sessionDirectory = sessionsRoot.path + let registry = CmuxVaultAgentRegistry(registrations: [registration]) + let oldWorkspaceIds = [UUID(), UUID(), UUID()] + let restoredWorkspaceIds = [UUID(), UUID(), UUID()] + let panels = [UUID(), UUID(), UUID()] + let sessionIds = ["pi-session-a", "pi-session-b", "pi-session-c"] + var hookSessions: [String: [String: Any]] = [:] + for (index, sessionId) in sessionIds.enumerated() { + let sessionFile = projectSessions.appendingPathComponent("\(sessionId).jsonl", isDirectory: false) + try "{}\n".write(to: sessionFile, atomically: true, encoding: .utf8) + try fm.setAttributes( + [.modificationDate: Date(timeIntervalSince1970: TimeInterval(1_000 + index))], + ofItemAtPath: sessionFile.path + ) + hookSessions[sessionId] = driftedAgentHookRecord( + launcher: "pi", + sessionId: sessionId, + workspaceId: oldWorkspaceIds[index], + panelId: panels[index], + recordedCwd: cwd.path, + launchCwd: cwd.path, + updatedAt: TimeInterval(10 + index) + ) + } + try writeHookStore(root: root, storeFilename: "pi-hook-sessions.json", sessions: hookSessions) + + let processes = panels.enumerated().map { index, panelId in + CmuxTopProcessInfo( + pid: 4_300 + index, + parentPID: 1, + name: "pi", + path: "/usr/local/bin/pi", + ttyDevice: nil, + cmuxWorkspaceID: restoredWorkspaceIds[index], + cmuxSurfaceID: panelId, + cmuxAttributionReason: "cmux-test", + processGroupID: nil, + terminalProcessGroupID: nil, + cpuPercent: 0, + residentBytes: 0, + virtualBytes: 0, + threadCount: 1 + ) + } + let processSnapshot = CmuxTopProcessSnapshot( + processes: processes, + sampledAt: Date(timeIntervalSince1970: 0), + includesProcessDetails: true + ) + let detectedSnapshots = RestorableAgentSessionIndex.processDetectedSnapshots( + registry: registry, + fileManager: fm, + processSnapshot: processSnapshot, + capturedAt: 42, + processArgumentsProvider: { processId in + guard processes.contains(where: { $0.pid == processId }) else { return nil } + return CmuxTopProcessArguments( + arguments: ["/usr/local/bin/pi"], + environment: [ + "PWD": cwd.path, + "PI_CODING_AGENT_SESSION_DIR": sessionsRoot.path, + ] + ) + } + ) + + let expectedDetectedKeys = Set(zip(restoredWorkspaceIds, panels).map { workspaceId, panelId in + RestorableAgentSessionIndex.PanelKey(workspaceId: workspaceId, panelId: panelId) + }) + XCTAssertEqual(Set(detectedSnapshots.keys), expectedDetectedKeys) + let detectedSessionIds = Set(detectedSnapshots.values.map { $0.snapshot.sessionId }) + XCTAssertEqual(detectedSessionIds.count, 1, "Pi latest-file detection is ambiguous for same-cwd workspaces") + + let index = RestorableAgentSessionIndex.load( + homeDirectory: root.path, + fileManager: fm, + registry: registry, + detectedSnapshots: detectedSnapshots, + processArgumentsProvider: { _ in nil } + ) + let restoredSessionIds = try zip(restoredWorkspaceIds, panels).map { workspaceId, panelId in + try XCTUnwrap(index.snapshot(workspaceId: workspaceId, panelId: panelId)).sessionId + } + + XCTAssertEqual(restoredSessionIds, sessionIds) + } + + func testPiDetectedLatestSessionDoesNotUsePanelOnlyFallbackWhenPanelIdIsAmbiguous() throws { + let fm = FileManager.default + let root = fm.temporaryDirectory + .appendingPathComponent("cmux-pi-restore-panel-id-ambiguous-\(UUID().uuidString)", isDirectory: true) + defer { try? fm.removeItem(at: root) } + + let cwd = root.appendingPathComponent("repo", isDirectory: true) + let sessionsRoot = root.appendingPathComponent("pi-sessions", isDirectory: true) + let projectDirectory = try XCTUnwrap(PiSessionLocator.projectDirectoryName(for: cwd.path)) + let projectSessions = sessionsRoot.appendingPathComponent(projectDirectory, isDirectory: true) + try fm.createDirectory(at: cwd, withIntermediateDirectories: true) + try fm.createDirectory(at: projectSessions, withIntermediateDirectories: true) + + var registration = CmuxVaultAgentRegistration.builtInPi + registration.sessionDirectory = sessionsRoot.path + let registry = CmuxVaultAgentRegistry(registrations: [registration]) + let oldWorkspaceId = UUID() + let otherOldWorkspaceId = UUID() + let restoredWorkspaceId = UUID() + let panelId = UUID() + let hookRecords = [ + (sessionId: "pi-old-workspace-a", workspaceId: oldWorkspaceId), + (sessionId: "pi-old-workspace-b", workspaceId: oldWorkspaceId), + (sessionId: "pi-other-old-workspace", workspaceId: otherOldWorkspaceId), + ] + let detectedLatestSessionId = "pi-detected-newest" + var hookSessions: [String: [String: Any]] = [:] + for (index, hookRecord) in hookRecords.enumerated() { + let sessionFile = projectSessions.appendingPathComponent("\(hookRecord.sessionId).jsonl", isDirectory: false) + try "{}\n".write(to: sessionFile, atomically: true, encoding: .utf8) + try fm.setAttributes( + [.modificationDate: Date(timeIntervalSince1970: TimeInterval(1_000 + index))], + ofItemAtPath: sessionFile.path + ) + hookSessions[hookRecord.sessionId] = driftedAgentHookRecord( + launcher: "pi", + sessionId: hookRecord.sessionId, + workspaceId: hookRecord.workspaceId, + panelId: panelId, + recordedCwd: cwd.path, + launchCwd: cwd.path, + updatedAt: TimeInterval(10 + index) + ) + } + let detectedLatestFile = projectSessions.appendingPathComponent("\(detectedLatestSessionId).jsonl", isDirectory: false) + try "{}\n".write(to: detectedLatestFile, atomically: true, encoding: .utf8) + try fm.setAttributes( + [.modificationDate: Date(timeIntervalSince1970: 2_000)], + ofItemAtPath: detectedLatestFile.path + ) + try writeHookStore(root: root, storeFilename: "pi-hook-sessions.json", sessions: hookSessions) + + let process = CmuxTopProcessInfo( + pid: 4_400, + parentPID: 1, + name: "pi", + path: "/usr/local/bin/pi", + ttyDevice: nil, + cmuxWorkspaceID: restoredWorkspaceId, + cmuxSurfaceID: panelId, + cmuxAttributionReason: "cmux-test", + processGroupID: nil, + terminalProcessGroupID: nil, + cpuPercent: 0, + residentBytes: 0, + virtualBytes: 0, + threadCount: 1 + ) + let processSnapshot = CmuxTopProcessSnapshot( + processes: [process], + sampledAt: Date(timeIntervalSince1970: 0), + includesProcessDetails: true + ) + let detectedSnapshots = RestorableAgentSessionIndex.processDetectedSnapshots( + registry: registry, + fileManager: fm, + processSnapshot: processSnapshot, + capturedAt: 42, + processArgumentsProvider: { processId in + guard processId == process.pid else { return nil } + return CmuxTopProcessArguments( + arguments: ["/usr/local/bin/pi"], + environment: [ + "PWD": cwd.path, + "PI_CODING_AGENT_SESSION_DIR": sessionsRoot.path, + ] + ) + } + ) + let restoredKey = RestorableAgentSessionIndex.PanelKey( + workspaceId: restoredWorkspaceId, + panelId: panelId + ) + let detected = try XCTUnwrap(detectedSnapshots[restoredKey]) + XCTAssertEqual(detected.snapshot.sessionId, detectedLatestSessionId) + + let index = RestorableAgentSessionIndex.load( + homeDirectory: root.path, + fileManager: fm, + registry: registry, + detectedSnapshots: detectedSnapshots, + processArgumentsProvider: { _ in nil } + ) + let snapshot = try XCTUnwrap(index.snapshot(workspaceId: restoredWorkspaceId, panelId: panelId)) + + XCTAssertEqual(snapshot.sessionId, detectedLatestSessionId) + } + func testPiInferredLatestFallbackUsesSameKindPanelHookWhenAnotherKindIsNewer() throws { let fm = FileManager.default let root = fm.temporaryDirectory