diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index b797ce2c94..9447043b11 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -2830,9 +2830,11 @@ class ClaudeAgentSession implements AgentSession { this.startFreshConversationSession(); return; } + const currentSessionId = this.claudeSessionId; await revertClaudeConversation({ sdk: this.rewindSdk, - sessionId: this.claudeSessionId, + sessionId: currentSessionId, + transcriptPath: currentSessionId ? this.resolveHistoryPath(currentSessionId) : null, messageId: target.messageId, resolveMessageId: (messageId) => this.resolveClaudeMessageId(messageId), setSessionId: (sessionId) => { diff --git a/packages/server/src/server/agent/providers/claude/rewind.test.ts b/packages/server/src/server/agent/providers/claude/rewind.test.ts index 99243bb9c4..28a6bd38fa 100644 --- a/packages/server/src/server/agent/providers/claude/rewind.test.ts +++ b/packages/server/src/server/agent/providers/claude/rewind.test.ts @@ -1,10 +1,15 @@ -import { describe, expect, test } from "vitest"; +import { randomUUID } from "node:crypto"; +import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import { afterEach, describe, expect, test } from "vitest"; import type { Query } from "@anthropic-ai/claude-agent-sdk"; import { revertClaudeConversation, revertClaudeConversationAndFiles, revertClaudeFiles, + realClaudeRewindSdk, } from "./rewind.js"; import { FakeClaudeSdk } from "./test-rewind-claude-sdk.js"; @@ -16,6 +21,7 @@ describe("Claude rewind", () => { await revertClaudeConversation({ sdk: claude, sessionId, + transcriptPath: "/claude/projects/repo/original-session.jsonl", messageId: "user-message-1", setSessionId: (nextSessionId) => { sessionId = nextSessionId; @@ -33,6 +39,7 @@ describe("Claude rewind", () => { await revertClaudeConversation({ sdk: claude, sessionId, + transcriptPath: "/claude/projects/repo/original-session.jsonl", messageId: "timeline-message-1", resolveMessageId: () => "claude-jsonl-message-1", setSessionId: (nextSessionId) => { @@ -76,6 +83,7 @@ describe("Claude rewind", () => { sdk: claude, query: claude.createQuery() as Query, sessionId, + transcriptPath: "/claude/projects/repo/original-session.jsonl", messageId: "user-message-1", setSessionId: (nextSessionId) => { sessionId = nextSessionId; @@ -87,3 +95,84 @@ describe("Claude rewind", () => { expect(sessionId).toBe("forked-before-rehydrate"); }); }); + +interface TranscriptLine { + type: "user" | "assistant"; + uuid: string; + parentUuid: string | null; + text: string; +} + +function writeTranscript(transcriptPath: string, sessionId: string, lines: TranscriptLine[]): void { + const entries = lines.map(({ type, uuid, parentUuid, text }) => + JSON.stringify({ + type, + uuid, + parentUuid, + sessionId, + cwd: "/repo", + timestamp: new Date().toISOString(), + message: { role: type, content: text }, + }), + ); + writeFileSync(transcriptPath, `${entries.join("\n")}\n`, "utf8"); +} + +function readMessageTexts(transcriptPath: string): string[] { + const entries = readFileSync(transcriptPath, "utf8") + .trim() + .split("\n") + .map((line) => JSON.parse(line) as { message?: { content: string } }); + return entries.flatMap((entry) => (entry.message ? [entry.message.content] : [])); +} + +describe("Claude rewind against the Claude SDK", () => { + let tempRoot: string | null = null; + const previousConfigDir = process.env.CLAUDE_CONFIG_DIR; + + afterEach(() => { + if (tempRoot) rmSync(tempRoot, { recursive: true, force: true }); + tempRoot = null; + if (previousConfigDir === undefined) { + delete process.env.CLAUDE_CONFIG_DIR; + } else { + process.env.CLAUDE_CONFIG_DIR = previousConfigDir; + } + }); + + test("forks a session that lives in a provider's own CLAUDE_CONFIG_DIR", async () => { + tempRoot = mkdtempSync(path.join(tmpdir(), "claude-rewind-config-dir-")); + // The provider sets CLAUDE_CONFIG_DIR for Claude Code; the daemon's own environment does not. + delete process.env.CLAUDE_CONFIG_DIR; + const projectDir = path.join(tempRoot, "custom-claude-dir", "projects", "-repo"); + mkdirSync(projectDir, { recursive: true }); + const sessionId = randomUUID(); + const firstUser = randomUUID(); + const firstReply = randomUUID(); + const secondUser = randomUUID(); + const transcriptPath = path.join(projectDir, `${sessionId}.jsonl`); + writeTranscript(transcriptPath, sessionId, [ + { type: "user", uuid: firstUser, parentUuid: null, text: "first" }, + { type: "assistant", uuid: firstReply, parentUuid: firstUser, text: "first reply" }, + { type: "user", uuid: secondUser, parentUuid: firstReply, text: "second" }, + { type: "assistant", uuid: randomUUID(), parentUuid: secondUser, text: "second reply" }, + ]); + let currentSessionId = sessionId; + + await revertClaudeConversation({ + sdk: realClaudeRewindSdk, + sessionId, + transcriptPath, + messageId: firstReply, + setSessionId: (nextSessionId) => { + currentSessionId = nextSessionId; + }, + }); + + expect(currentSessionId).not.toBe(sessionId); + expect(readMessageTexts(path.join(projectDir, `${currentSessionId}.jsonl`))).toEqual([ + "first", + "first reply", + ]); + }); +}); diff --git a/packages/server/src/server/agent/providers/claude/rewind.ts b/packages/server/src/server/agent/providers/claude/rewind.ts index 9c3e3d708d..f3c831575a 100644 --- a/packages/server/src/server/agent/providers/claude/rewind.ts +++ b/packages/server/src/server/agent/providers/claude/rewind.ts @@ -1,29 +1,78 @@ -import { forkSession as claudeForkSession, type Query } from "@anthropic-ai/claude-agent-sdk"; +import { appendFile, readFile } from "node:fs/promises"; +import path from "node:path"; +import { + forkSession as claudeForkSession, + type Query, + type SessionStore, + type SessionStoreEntry, +} from "@anthropic-ai/claude-agent-sdk"; export interface ClaudeRewindSdk { forkSession( sessionId: string, - options: { upToMessageId: string }, + options: { upToMessageId: string; transcriptPath: string }, ): Promise<{ sessionId: string }>; } +// On its own, the SDK looks for the session under the daemon's CLAUDE_CONFIG_DIR, while a +// provider can point Claude Code at another one. Forking through a store rooted at the +// transcript's folder reads the session where Claude Code wrote it and writes the fork beside it. export const realClaudeRewindSdk: ClaudeRewindSdk = { - forkSession: claudeForkSession, + forkSession: (sessionId, options) => + claudeForkSession(sessionId, { + upToMessageId: options.upToMessageId, + sessionStore: projectFolderSessionStore(path.dirname(options.transcriptPath)), + }), }; +function projectFolderSessionStore(projectFolder: string): SessionStore { + const transcriptPath = (sessionId: string) => path.join(projectFolder, `${sessionId}.jsonl`); + return { + async load({ sessionId }) { + let text: string; + try { + text = await readFile(transcriptPath(sessionId), "utf8"); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") return null; + throw error; + } + return parseTranscriptEntries(text); + }, + async append({ sessionId }, entries) { + const lines = entries.map((entry) => `${JSON.stringify(entry)}\n`).join(""); + await appendFile(transcriptPath(sessionId), lines, "utf8"); + }, + }; +} + +function parseTranscriptEntries(text: string): SessionStoreEntry[] { + const entries: SessionStoreEntry[] = []; + for (const line of text.split("\n")) { + if (!line.trim()) continue; + try { + entries.push(JSON.parse(line) as SessionStoreEntry); + } catch { + // Skip a line that is not JSON, as the SDK's own transcript reader does. + } + } + return entries; +} + export async function revertClaudeConversation(input: { sdk: ClaudeRewindSdk; sessionId: string | null; + transcriptPath: string | null; messageId: string; resolveMessageId?: (messageId: string) => string | Promise; setSessionId: (sessionId: string) => void; }): Promise { - if (!input.sessionId) { + if (!input.sessionId || !input.transcriptPath) { throw new Error("Claude session is not ready for rewind"); } const messageId = (await input.resolveMessageId?.(input.messageId)) ?? input.messageId; const fork = await input.sdk.forkSession(input.sessionId, { upToMessageId: messageId, + transcriptPath: input.transcriptPath, }); input.setSessionId(fork.sessionId); } @@ -44,6 +93,7 @@ export async function revertClaudeConversationAndFiles(input: { sdk: ClaudeRewindSdk; query: Query; sessionId: string | null; + transcriptPath: string | null; messageId: string; resolveMessageId?: (messageId: string) => string | Promise; setSessionId: (sessionId: string) => void; diff --git a/packages/server/src/server/agent/providers/claude/test-rewind-claude-sdk.ts b/packages/server/src/server/agent/providers/claude/test-rewind-claude-sdk.ts index 496d3bc209..d607b5c063 100644 --- a/packages/server/src/server/agent/providers/claude/test-rewind-claude-sdk.ts +++ b/packages/server/src/server/agent/providers/claude/test-rewind-claude-sdk.ts @@ -12,7 +12,7 @@ export class FakeClaudeSdk implements ClaudeRewindSdk { async forkSession( _sessionId: string, - options: { upToMessageId: string }, + options: { upToMessageId: string; transcriptPath: string }, ): Promise<{ sessionId: string }> { this.recordedForks.push({ upToMessageId: options.upToMessageId }); return { sessionId: this.nextSessionId };