From ab65d12356d648045ded8d19ab84478b5ae53bab Mon Sep 17 00:00:00 2001 From: Ashish Huddar <73213873+ashish921998@users.noreply.github.com> Date: Fri, 1 May 2026 21:27:17 +0530 Subject: [PATCH] Fix native restore fallback for Claude and Codex sessions (#1602) * Fix native session restore fallback for Claude and Codex * Address restore metadata review comments * Fix metadata normalization lint * Address PR metadata review feedback * Prevent fresh restore fallback for native agents --- .../__tests__/session-manager/restore.test.ts | 139 ++++++++++++++++++ packages/core/src/metadata.ts | 2 +- packages/core/src/session-manager.ts | 51 ++++++- packages/core/src/types.ts | 6 + .../agent-claude-code/src/index.test.ts | 22 +++ .../plugins/agent-claude-code/src/index.ts | 22 +-- .../plugins/agent-codex/src/index.test.ts | 14 ++ packages/plugins/agent-codex/src/index.ts | 32 ++-- 8 files changed, 261 insertions(+), 27 deletions(-) diff --git a/packages/core/src/__tests__/session-manager/restore.test.ts b/packages/core/src/__tests__/session-manager/restore.test.ts index b8022c829..f955bb509 100644 --- a/packages/core/src/__tests__/session-manager/restore.test.ts +++ b/packages/core/src/__tests__/session-manager/restore.test.ts @@ -547,6 +547,145 @@ describe("restore", () => { expect(mockAgent.getLaunchCommand).toHaveBeenCalled(); const createCall = (mockRuntime.create as ReturnType).mock.calls[0][0]; expect(createCall.launchCommand).toBe("mock-agent --start"); + + const meta = readMetadataRaw(sessionsDir, "app-1"); + expect(meta!["restoreFallbackReason"]).toBe("mock-agent.getRestoreCommand returned null"); + }); + + it("does not launch a fresh chat when a native-restore agent cannot build restore command", async () => { + const wsPath = join(tmpDir, "ws-app-native-restore-missing"); + mkdirSync(wsPath, { recursive: true }); + + const mockNativeRestoreAgent: Agent = { + ...mockAgent, + name: "codex", + getRestoreCommand: vi.fn().mockResolvedValue(null), + }; + + const registryWithNativeRestoreAgent: PluginRegistry = { + ...mockRegistry, + get: vi.fn().mockImplementation((slot: string) => { + if (slot === "runtime") return mockRuntime; + if (slot === "agent") return mockNativeRestoreAgent; + if (slot === "workspace") return mockWorkspace; + return null; + }), + }; + + writeMetadata(sessionsDir, "app-1", { + worktree: wsPath, + branch: "feat/TEST-1", + status: "killed", + project: "my-app", + runtimeHandle: makeHandle("rt-old"), + }); + + const sm = createSessionManager({ config, registry: registryWithNativeRestoreAgent }); + + await expect(sm.restore("app-1")).rejects.toThrow(SessionNotRestorableError); + expect(mockRuntime.create).not.toHaveBeenCalled(); + const meta = readMetadataRaw(sessionsDir, "app-1"); + expect(meta!["restoreFallbackReason"]).toBe("codex.getRestoreCommand returned null"); + }); + + it("clears restore fallback reason when getRestoreCommand succeeds", async () => { + const wsPath = join(tmpDir, "ws-app-restore-clears-fallback"); + mkdirSync(wsPath, { recursive: true }); + + const mockAgentWithRestore: Agent = { + ...mockAgent, + getRestoreCommand: vi.fn().mockResolvedValue("claude --resume abc123"), + }; + + const registryWithAgentRestore: PluginRegistry = { + ...mockRegistry, + get: vi.fn().mockImplementation((slot: string) => { + if (slot === "runtime") return mockRuntime; + if (slot === "agent") return mockAgentWithRestore; + if (slot === "workspace") return mockWorkspace; + return null; + }), + }; + + writeMetadata(sessionsDir, "app-1", { + worktree: wsPath, + branch: "feat/TEST-1", + status: "killed", + project: "my-app", + runtimeHandle: makeHandle("rt-old"), + restoreFallbackReason: "previous fallback", + }); + + const sm = createSessionManager({ config, registry: registryWithAgentRestore }); + await sm.restore("app-1"); + + const meta = readMetadataRaw(sessionsDir, "app-1"); + expect(meta!["restoreFallbackReason"]).toBeUndefined(); + }); + + it("normalizes agent metadata empty strings in memory like metadata persistence", async () => { + const wsPath = join(tmpDir, "ws-app-agent-metadata-normalize"); + mkdirSync(wsPath, { recursive: true }); + + const mockAgentWithMetadata: Agent = { + ...mockAgent, + getSessionInfo: vi.fn().mockResolvedValue({ + summary: null, + agentSessionId: "native-1", + metadata: { + codexThreadId: "thread-1", + restoreFallbackReason: "", + }, + }), + }; + + const registryWithAgentMetadata: PluginRegistry = { + ...mockRegistry, + get: vi.fn().mockImplementation((slot: string) => { + if (slot === "runtime") return mockRuntime; + if (slot === "agent") return mockAgentWithMetadata; + if (slot === "workspace") return mockWorkspace; + return null; + }), + }; + + writeMetadata(sessionsDir, "app-1", { + worktree: wsPath, + branch: "feat/TEST-1", + status: "killed", + project: "my-app", + runtimeHandle: makeHandle("rt-old"), + restoreFallbackReason: "previous fallback", + }); + + const sm = createSessionManager({ config, registry: registryWithAgentMetadata }); + const restored = await sm.restore("app-1"); + + expect(restored.metadata["codexThreadId"]).toBe("thread-1"); + expect(restored.metadata["restoreFallbackReason"]).toBeUndefined(); + const meta = readMetadataRaw(sessionsDir, "app-1"); + expect(meta!["codexThreadId"]).toBe("thread-1"); + expect(meta!["restoreFallbackReason"]).toBeUndefined(); + }); + + it("clears restore fallback reason when agent has no restore command", async () => { + const wsPath = join(tmpDir, "ws-app-no-restore-clears-fallback"); + mkdirSync(wsPath, { recursive: true }); + + writeMetadata(sessionsDir, "app-1", { + worktree: wsPath, + branch: "feat/TEST-1", + status: "killed", + project: "my-app", + runtimeHandle: makeHandle("rt-old"), + restoreFallbackReason: "previous fallback", + }); + + const sm = createSessionManager({ config, registry: mockRegistry }); + await sm.restore("app-1"); + + const meta = readMetadataRaw(sessionsDir, "app-1"); + expect(meta!["restoreFallbackReason"]).toBeUndefined(); }); it("does not inject OPENCODE_CONFIG when restoring OpenCode orchestrators", async () => { diff --git a/packages/core/src/metadata.ts b/packages/core/src/metadata.ts index 6e3891fb6..95d354f77 100644 --- a/packages/core/src/metadata.ts +++ b/packages/core/src/metadata.ts @@ -299,7 +299,7 @@ export function updateMetadata( }, { createIfMissing: true }); } -function applyMetadataUpdates( +export function applyMetadataUpdates( existing: Record, updates: Partial>, ): Record { diff --git a/packages/core/src/session-manager.ts b/packages/core/src/session-manager.ts index 443317dbe..5ecbd8bb4 100644 --- a/packages/core/src/session-manager.ts +++ b/packages/core/src/session-manager.ts @@ -53,6 +53,7 @@ import { readMetadataRaw, writeMetadata, updateMetadata, + applyMetadataUpdates, mutateMetadata, deleteMetadata, listMetadata, @@ -428,6 +429,15 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM return isOrchestratorSessionRecord(sessionId, metadata ?? {}, project.sessionPrefix); } + function requiresNativeRestore(agentName: string): boolean { + return ( + agentName === "claude-code" || + agentName === "codex" || + agentName === "kimicode" || + agentName === "opencode" + ); + } + function applyMetadataUpdatesToRaw( raw: Record, updates: Partial>, @@ -972,7 +982,7 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM data: {}, }; } - await enrichSessionWithRuntimeState(session, plugins, handleFromMetadata); + await enrichSessionWithRuntimeState(session, plugins, handleFromMetadata, sessionsDir); } /** @@ -985,6 +995,7 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM session: Session, plugins: ReturnType, handleFromMetadata: boolean, + sessionsDir: string, ): Promise { // Check runtime liveness first — for all statuses except "spawning". // Skip spawning sessions because tmux may not be fully initialized yet, @@ -1058,13 +1069,26 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM } // Enrich with live agent session info (summary, cost). + let info: Awaited>; try { - const info = await plugins.agent.getSessionInfo(session); - if (info) { - session.agentInfo = info; - } + info = await plugins.agent.getSessionInfo(session); } catch { // Can't get session info — keep existing values + info = null; + } + + if (info) { + session.agentInfo = info; + const metadataUpdates = info.metadata ?? {}; + if (Object.keys(metadataUpdates).length > 0) { + try { + updateMetadata(sessionsDir, session.id, metadataUpdates); + session.metadata = applyMetadataUpdates(session.metadata, metadataUpdates); + invalidateCache(); + } catch { + // Persisting agent metadata is best-effort; keep live agent info. + } + } } } } @@ -2811,7 +2835,7 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM // and isRestorable would reject it. const session = metadataToSession(sessionId, raw, projectId, project.sessionPrefix); const plugins = resolvePlugins(project, selection.agentName); - await enrichSessionWithRuntimeState(session, plugins, true); + await enrichSessionWithRuntimeState(session, plugins, true, sessionsDir); // 3. Validate restorability if (!isRestorable(session)) { @@ -2924,9 +2948,22 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM if (plugins.agent.getRestoreCommand) { const restoreCmd = await plugins.agent.getRestoreCommand(session, projectConfigForLaunch); - launchCommand = restoreCmd ?? plugins.agent.getLaunchCommand(agentLaunchConfig); + if (restoreCmd) { + launchCommand = restoreCmd; + updateMetadata(sessionsDir, sessionId, { restoreFallbackReason: "" }); + } else { + const reason = `${plugins.agent.name}.getRestoreCommand returned null`; + updateMetadata(sessionsDir, sessionId, { + restoreFallbackReason: reason, + }); + if (requiresNativeRestore(plugins.agent.name)) { + throw new SessionNotRestorableError(sessionId, reason); + } + launchCommand = plugins.agent.getLaunchCommand(agentLaunchConfig); + } } else { launchCommand = plugins.agent.getLaunchCommand(agentLaunchConfig); + updateMetadata(sessionsDir, sessionId, { restoreFallbackReason: "" }); } const environment = plugins.agent.getEnvironment(agentLaunchConfig); diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index 0393e03ea..b7eac5aea 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -601,6 +601,8 @@ export interface AgentSessionInfo { summaryIsFallback?: boolean; /** Agent's internal session ID (for resume) */ agentSessionId: string | null; + /** Agent-owned metadata worth persisting for later restore. */ + metadata?: Record; /** Estimated cost so far */ cost?: CostEstimate; } @@ -1695,6 +1697,10 @@ export interface SessionMetadata { directTerminalWsPort?: number; }; opencodeSessionId?: string; + claudeSessionUuid?: string; + codexThreadId?: string; + codexModel?: string; + restoreFallbackReason?: string; pinnedSummary?: string; // First quality summary, pinned for display stability userPrompt?: string; // Prompt used when spawning without a tracker issue /** diff --git a/packages/plugins/agent-claude-code/src/index.test.ts b/packages/plugins/agent-claude-code/src/index.test.ts index f283e681e..c05df45c5 100644 --- a/packages/plugins/agent-claude-code/src/index.test.ts +++ b/packages/plugins/agent-claude-code/src/index.test.ts @@ -566,6 +566,28 @@ describe("getSessionInfo", () => { mockJsonlFiles('{"type":"user","message":{"content":"hi"}}', ["abc-def-123.jsonl"]); const result = await agent.getSessionInfo(makeSession()); expect(result?.agentSessionId).toBe("abc-def-123"); + expect(result?.metadata?.claudeSessionUuid).toBe("abc-def-123"); + }); + }); + + describe("getRestoreCommand metadata", () => { + it("uses persisted Claude session UUID without scanning project files", async () => { + const agent = create(); + const session = makeSession({ + workspacePath: "/workspace/test-project", + metadata: { claudeSessionUuid: "persisted-uuid" }, + }); + + const command = await agent.getRestoreCommand!(session, { + name: "test-project", + repo: "owner/repo", + path: "/workspace/test-project", + defaultBranch: "main", + sessionPrefix: "test", + }); + + expect(command).toBe("claude --resume 'persisted-uuid'"); + expect(mockReaddir).not.toHaveBeenCalled(); }); }); diff --git a/packages/plugins/agent-claude-code/src/index.ts b/packages/plugins/agent-claude-code/src/index.ts index e60e761fc..7b748016d 100644 --- a/packages/plugins/agent-claude-code/src/index.ts +++ b/packages/plugins/agent-claude-code/src/index.ts @@ -832,23 +832,27 @@ function createClaudeCodeAgent(): Agent { summary: summaryResult?.summary ?? null, summaryIsFallback: summaryResult?.isFallback, agentSessionId, + metadata: { claudeSessionUuid: agentSessionId }, cost: extractCost(lines), }; }, async getRestoreCommand(session: Session, project: ProjectConfig): Promise { - if (!session.workspacePath) return null; + let sessionUuid = session.metadata?.["claudeSessionUuid"]?.trim(); + if (!sessionUuid) { + if (!session.workspacePath) return null; - // Find Claude's project directory for this workspace - const projectPath = toClaudeProjectPath(session.workspacePath); - const projectDir = join(homedir(), ".claude", "projects", projectPath); + // Find Claude's project directory for this workspace + const projectPath = toClaudeProjectPath(session.workspacePath); + const projectDir = join(homedir(), ".claude", "projects", projectPath); - // Find the latest session JSONL file - const sessionFile = await findLatestSessionFile(projectDir); - if (!sessionFile) return null; + // Find the latest session JSONL file + const sessionFile = await findLatestSessionFile(projectDir); + if (!sessionFile) return null; - // Extract session UUID from filename (e.g. "abc123-def456.jsonl" → "abc123-def456") - const sessionUuid = basename(sessionFile, ".jsonl"); + // Extract session UUID from filename (e.g. "abc123-def456.jsonl" → "abc123-def456") + sessionUuid = basename(sessionFile, ".jsonl"); + } if (!sessionUuid) return null; // Build resume command diff --git a/packages/plugins/agent-codex/src/index.test.ts b/packages/plugins/agent-codex/src/index.test.ts index 73468bfbe..0e06e0644 100644 --- a/packages/plugins/agent-codex/src/index.test.ts +++ b/packages/plugins/agent-codex/src/index.test.ts @@ -1328,6 +1328,20 @@ describe("getRestoreCommand", () => { expect(cmd).toContain("thread-abc-123"); }); + it("uses persisted Codex thread ID without scanning session files", async () => { + const session = makeSession({ + workspacePath: "/workspace/test", + metadata: { codexThreadId: "persisted-thread", codexModel: "gpt-5.3-codex" }, + }); + + const cmd = await agent.getRestoreCommand!(session, makeProjectConfig()); + + expect(cmd).toContain("'codex' resume"); + expect(cmd).toContain("--model 'gpt-5.3-codex'"); + expect(cmd).toContain("persisted-thread"); + expect(mockReaddir).not.toHaveBeenCalled(); + }); + it("builds native resume command from payload-wrapped Codex session id", async () => { const content = jsonl( { diff --git a/packages/plugins/agent-codex/src/index.ts b/packages/plugins/agent-codex/src/index.ts index 320bf5ad9..777132f29 100644 --- a/packages/plugins/agent-codex/src/index.ts +++ b/packages/plugins/agent-codex/src/index.ts @@ -703,21 +703,33 @@ function createCodexAgent(): Agent { summary: data.model ? `Codex session (${data.model})` : null, summaryIsFallback: true, agentSessionId, + metadata: data.threadId + ? { + codexThreadId: data.threadId, + ...(data.model ? { codexModel: data.model } : {}), + } + : undefined, cost, }; }, async getRestoreCommand(session: Session, project: ProjectConfig): Promise { - if (!session.workspacePath) return null; + let threadId = session.metadata?.["codexThreadId"]?.trim(); + let model: string | null = session.metadata?.["codexModel"]?.trim() || null; + if (!threadId) { + if (!session.workspacePath) return null; - // Find the Codex session file for this workspace - const sessionFile = await findCodexSessionFileCached(session.workspacePath); - if (!sessionFile) return null; + // Find the Codex session file for this workspace + const sessionFile = await findCodexSessionFileCached(session.workspacePath); + if (!sessionFile) return null; - // Stream the file line-by-line to avoid loading potentially huge - // rollout files (100 MB+) entirely into memory. - const data = await streamCodexSessionData(sessionFile); - if (!data?.threadId) return null; + // Stream the file line-by-line to avoid loading potentially huge + // rollout files (100 MB+) entirely into memory. + const data = await streamCodexSessionData(sessionFile); + if (!data?.threadId) return null; + threadId = data.threadId; + model = data.model; + } // Use Codex's native `resume` subcommand for proper conversation resume. // This restores the full thread state, not just a text prompt re-injection. @@ -727,11 +739,11 @@ function createCodexAgent(): Agent { appendNoUpdateCheckFlag(parts); appendApprovalFlags(parts, project.agentConfig?.permissions); - const effectiveModel = (project.agentConfig?.model ?? data.model) as string | undefined; + const effectiveModel = (project.agentConfig?.model ?? model) as string | undefined; appendModelFlags(parts, effectiveModel ?? undefined); // Positional threadId goes last, after all flags - parts.push(shellEscape(data.threadId)); + parts.push(shellEscape(threadId)); return parts.join(" "); },