diff --git a/packages/agent/src/persistent-agent-state.test.ts b/packages/agent/src/persistent-agent-state.test.ts new file mode 100644 index 0000000000..36e854d5c9 --- /dev/null +++ b/packages/agent/src/persistent-agent-state.test.ts @@ -0,0 +1,101 @@ +import { + lstat, + mkdir, + mkdtemp, + readFile, + readlink, + rm, + writeFile, +} from "node:fs/promises"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { configurePersistentAgentState } from "./persistent-agent-state"; + +describe("configurePersistentAgentState", () => { + let testRoot: string; + let homeDir: string; + let stateRoot: string; + + beforeEach(async () => { + testRoot = await mkdtemp(path.join(tmpdir(), "persistent-agent-state-")); + homeDir = path.join(testRoot, "home"); + stateRoot = path.join(testRoot, "workspace", ".posthog", "agent-state"); + vi.stubEnv("CLAUDE_CONFIG_DIR", ""); + vi.stubEnv("CODEX_HOME", ""); + }); + + afterEach(async () => { + vi.unstubAllEnvs(); + await rm(testRoot, { recursive: true, force: true }); + }); + + it("routes native session state into the durable workspace root", async () => { + const claudeSkillsDir = path.join(homeDir, ".claude", "skills"); + const claudeProjectsDir = path.join(homeDir, ".claude", "projects"); + const claudeSessionEnvDir = path.join(homeDir, ".claude", "session-env"); + const codexSessionsDir = path.join(homeDir, ".codex", "sessions"); + const codexConfigPath = path.join(homeDir, ".codex", "config.toml"); + const restoredSessionEnvDir = path.join(stateRoot, "claude", "session-env"); + await mkdir(claudeSkillsDir, { recursive: true }); + await mkdir(claudeProjectsDir, { recursive: true }); + await mkdir(claudeSessionEnvDir, { recursive: true }); + await mkdir(codexSessionsDir, { recursive: true }); + await mkdir(path.dirname(codexConfigPath), { recursive: true }); + await mkdir(restoredSessionEnvDir, { recursive: true }); + await writeFile(path.join(claudeSkillsDir, "catalog.txt"), "fresh"); + await writeFile(path.join(claudeProjectsDir, "session.jsonl"), "claude"); + await writeFile(path.join(restoredSessionEnvDir, "restored"), "state"); + await writeFile(path.join(codexSessionsDir, "rollout.jsonl"), "codex"); + await writeFile(codexConfigPath, "fresh = true"); + + await configurePersistentAgentState(stateRoot, homeDir); + + const mappings = [ + [".claude/projects", "claude/projects"], + [".claude/session-env", "claude/session-env"], + [".claude/plans", "claude/plans"], + [".claude/todos", "claude/todos"], + [".codex/sessions", "codex/sessions"], + [".codex/shell_snapshots", "codex/shell_snapshots"], + ]; + + for (const [sourceRelative, targetRelative] of mappings) { + const source = path.join(homeDir, sourceRelative); + const target = path.join(stateRoot, targetRelative); + expect((await lstat(source)).isSymbolicLink()).toBe(true); + expect(path.resolve(path.dirname(source), await readlink(source))).toBe( + target, + ); + await writeFile(path.join(source, "marker"), sourceRelative); + await expect( + readFile(path.join(target, "marker"), "utf-8"), + ).resolves.toBe(sourceRelative); + } + + await configurePersistentAgentState(stateRoot, homeDir); + + await expect( + readFile( + path.join(stateRoot, "claude", "projects", "session.jsonl"), + "utf-8", + ), + ).resolves.toBe("claude"); + await expect( + readFile( + path.join(stateRoot, "codex", "sessions", "rollout.jsonl"), + "utf-8", + ), + ).resolves.toBe("codex"); + await expect( + readFile(path.join(restoredSessionEnvDir, "restored"), "utf-8"), + ).resolves.toBe("state"); + + await expect( + readFile(path.join(claudeSkillsDir, "catalog.txt"), "utf-8"), + ).resolves.toBe("fresh"); + await expect(readFile(codexConfigPath, "utf-8")).resolves.toBe( + "fresh = true", + ); + }); +}); diff --git a/packages/agent/src/persistent-agent-state.ts b/packages/agent/src/persistent-agent-state.ts new file mode 100644 index 0000000000..a90bfa8044 --- /dev/null +++ b/packages/agent/src/persistent-agent-state.ts @@ -0,0 +1,141 @@ +import { + cp, + lstat, + mkdir, + readdir, + readlink, + rename, + rm, + symlink, +} from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; + +interface PersistentDirectory { + source: string; + target: string; +} + +function isMissingPathError(error: unknown): boolean { + return ( + error instanceof Error && + "code" in error && + (error as NodeJS.ErrnoException).code === "ENOENT" + ); +} + +function isCrossDeviceError(error: unknown): boolean { + return ( + error instanceof Error && + "code" in error && + (error as NodeJS.ErrnoException).code === "EXDEV" + ); +} + +async function migrateExistingDirectory( + source: string, + target: string, +): Promise { + const sourceEntries = await readdir(source); + const targetEntries = await readdir(target); + if (sourceEntries.length === 0) { + await rm(source, { recursive: true, force: true }); + return; + } + if (targetEntries.length > 0) { + throw new Error( + `Cannot migrate persistent agent state from ${source}: ${target} is not empty`, + ); + } + + await rm(target, { recursive: true, force: true }); + try { + await rename(source, target); + } catch (error) { + if (!isCrossDeviceError(error)) throw error; + await cp(source, target, { + recursive: true, + errorOnExist: true, + force: false, + }); + await rm(source, { recursive: true, force: true }); + } +} + +async function linkPersistentDirectory({ + source, + target, +}: PersistentDirectory): Promise { + const resolvedSource = path.resolve(source); + const resolvedTarget = path.resolve(target); + if (resolvedSource === resolvedTarget) return; + + await mkdir(resolvedTarget, { recursive: true }); + await mkdir(path.dirname(resolvedSource), { recursive: true }); + + try { + const sourceStats = await lstat(resolvedSource); + if (sourceStats.isSymbolicLink()) { + const currentTarget = path.resolve( + path.dirname(resolvedSource), + await readlink(resolvedSource), + ); + if (currentTarget === resolvedTarget) return; + await rm(resolvedSource, { force: true }); + } else if (sourceStats.isDirectory()) { + await migrateExistingDirectory(resolvedSource, resolvedTarget); + } else { + throw new Error( + `Persistent agent state source must be a directory: ${resolvedSource}`, + ); + } + } catch (error) { + if (!isMissingPathError(error)) throw error; + } + + await symlink(resolvedTarget, resolvedSource, "dir"); +} + +export async function configurePersistentAgentState( + stateRoot: string, + homeDir: string = os.homedir(), +): Promise { + if (!path.isAbsolute(stateRoot)) { + throw new Error("Persistent agent state root must be an absolute path"); + } + + const claudeConfigDir = + process.env.CLAUDE_CONFIG_DIR || path.join(homeDir, ".claude"); + const codexHome = process.env.CODEX_HOME || path.join(homeDir, ".codex"); + const directories: PersistentDirectory[] = [ + { + source: path.join(claudeConfigDir, "projects"), + target: path.join(stateRoot, "claude", "projects"), + }, + { + source: path.join(claudeConfigDir, "session-env"), + target: path.join(stateRoot, "claude", "session-env"), + }, + { + source: path.join(claudeConfigDir, "plans"), + target: path.join(stateRoot, "claude", "plans"), + }, + { + source: path.join(claudeConfigDir, "todos"), + target: path.join(stateRoot, "claude", "todos"), + }, + { + source: path.join(codexHome, "sessions"), + target: path.join(stateRoot, "codex", "sessions"), + }, + { + source: path.join(codexHome, "shell_snapshots"), + target: path.join(stateRoot, "codex", "shell_snapshots"), + }, + ]; + + // Avoid leaving other mappings mutating after one setup operation fails. + for (const directory of directories) { + await linkPersistentDirectory(directory); + } +} diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts index 961deee65b..75dc1d2e55 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -1,5 +1,13 @@ import { createHash } from "node:crypto"; -import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { + lstat, + mkdir, + mkdtemp, + readFile, + readlink, + rm, + writeFile, +} from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { ContentBlock } from "@agentclientprotocol/sdk"; @@ -441,6 +449,42 @@ describe("AgentServer HTTP Mode", () => { sessionInitMs: expect.any(Number), }); }, 30000); + + it("links native agent state before initializing the session", async () => { + const originalClaudeConfigDir = process.env.CLAUDE_CONFIG_DIR; + const originalCodexHome = process.env.CODEX_HOME; + const claudeConfigDir = join(repo.path, ".claude-test"); + const codexHome = join(repo.path, ".codex-test"); + const agentStateDir = join(repo.path, ".posthog", "agent-state"); + process.env.CLAUDE_CONFIG_DIR = claudeConfigDir; + process.env.CODEX_HOME = codexHome; + + try { + await createServer({ agentStateDir }).start(); + + const claudeProjects = join(claudeConfigDir, "projects"); + const codexSessions = join(codexHome, "sessions"); + expect((await lstat(claudeProjects)).isSymbolicLink()).toBe(true); + expect((await lstat(codexSessions)).isSymbolicLink()).toBe(true); + expect(await readlink(claudeProjects)).toBe( + join(agentStateDir, "claude", "projects"), + ); + expect(await readlink(codexSessions)).toBe( + join(agentStateDir, "codex", "sessions"), + ); + } finally { + if (originalClaudeConfigDir === undefined) { + delete process.env.CLAUDE_CONFIG_DIR; + } else { + process.env.CLAUDE_CONFIG_DIR = originalClaudeConfigDir; + } + if (originalCodexHome === undefined) { + delete process.env.CODEX_HOME; + } else { + process.env.CODEX_HOME = originalCodexHome; + } + } + }, 30000); }); describe("turn completion", () => { diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 75fa3d1ba3..e30a3627e0 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -49,6 +49,7 @@ import { import type { PermissionMode } from "../execution-mode"; import { DEFAULT_CODEX_MODEL, fetchGatewayModels } from "../gateway-models"; import { HandoffCheckpointTracker } from "../handoff-checkpoint"; +import { configurePersistentAgentState } from "../persistent-agent-state"; import { PostHogAPIClient } from "../posthog-api"; import { findPrUrls, @@ -635,6 +636,10 @@ export class AgentServer { } async start(): Promise { + if (this.config.agentStateDir) { + await configurePersistentAgentState(this.config.agentStateDir); + } + await new Promise((resolve) => { this.server = serve( { diff --git a/packages/agent/src/server/bin.ts b/packages/agent/src/server/bin.ts index 3904517b71..46828f8e47 100644 --- a/packages/agent/src/server/bin.ts +++ b/packages/agent/src/server/bin.ts @@ -32,6 +32,7 @@ const envSchema = z.object({ POSTHOG_CODE_REASONING_EFFORT: z .enum(["low", "medium", "high", "xhigh", "max"]) .optional(), + POSTHOG_AGENT_STATE_DIR: z.string().startsWith("/").optional(), POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN: z.string().min(1).optional(), // Base URL for the event-ingest POST only; falls back to POSTHOG_API_URL when unset. POSTHOG_TASK_RUN_EVENT_INGEST_URL: z.url().optional(), @@ -174,6 +175,7 @@ program const server = new AgentServer({ port: parseInt(options.port, 10), + agentStateDir: env.POSTHOG_AGENT_STATE_DIR, jwtPublicKey: env.JWT_PUBLIC_KEY, eventIngestToken: env.POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN, eventIngestBaseUrl: env.POSTHOG_TASK_RUN_EVENT_INGEST_URL, diff --git a/packages/agent/src/server/types.ts b/packages/agent/src/server/types.ts index 12465f8cb8..6bb38efba2 100644 --- a/packages/agent/src/server/types.ts +++ b/packages/agent/src/server/types.ts @@ -11,6 +11,7 @@ export interface ClaudeCodeConfig { export interface AgentServerConfig { port: number; + agentStateDir?: string; repositoryPath?: string; repoReadyFile?: string; apiUrl: string;