Skip to content
This repository was archived by the owner on Aug 6, 2026. It is now read-only.

Commit cea445f

Browse files
tatoalotatoalo
authored
feat(agent): persist native session state in workspace snapshots (#3333)
1 parent 863768a commit cea445f

6 files changed

Lines changed: 295 additions & 1 deletion

File tree

‎packages/agent/src/persistent-agent-state.test.ts‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import {
2+
lstat,
3+
mkdir,
4+
mkdtemp,
5+
readFile,
6+
readlink,
7+
rm,
8+
writeFile,
9+
} from "node:fs/promises";
10+
import { tmpdir } from "node:os";
11+
import path from "node:path";
12+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
13+
import { configurePersistentAgentState } from "./persistent-agent-state";
14+
15+
describe("configurePersistentAgentState", () => {
16+
let testRoot: string;
17+
let homeDir: string;
18+
let stateRoot: string;
19+
20+
beforeEach(async () => {
21+
testRoot = await mkdtemp(path.join(tmpdir(), "persistent-agent-state-"));
22+
homeDir = path.join(testRoot, "home");
23+
stateRoot = path.join(testRoot, "workspace", ".posthog", "agent-state");
24+
vi.stubEnv("CLAUDE_CONFIG_DIR", "");
25+
vi.stubEnv("CODEX_HOME", "");
26+
});
27+
28+
afterEach(async () => {
29+
vi.unstubAllEnvs();
30+
await rm(testRoot, { recursive: true, force: true });
31+
});
32+
33+
it("routes native session state into the durable workspace root", async () => {
34+
const claudeSkillsDir = path.join(homeDir, ".claude", "skills");
35+
const claudeProjectsDir = path.join(homeDir, ".claude", "projects");
36+
const claudeSessionEnvDir = path.join(homeDir, ".claude", "session-env");
37+
const codexSessionsDir = path.join(homeDir, ".codex", "sessions");
38+
const codexConfigPath = path.join(homeDir, ".codex", "config.toml");
39+
const restoredSessionEnvDir = path.join(stateRoot, "claude", "session-env");
40+
await mkdir(claudeSkillsDir, { recursive: true });
41+
await mkdir(claudeProjectsDir, { recursive: true });
42+
await mkdir(claudeSessionEnvDir, { recursive: true });
43+
await mkdir(codexSessionsDir, { recursive: true });
44+
await mkdir(path.dirname(codexConfigPath), { recursive: true });
45+
await mkdir(restoredSessionEnvDir, { recursive: true });
46+
await writeFile(path.join(claudeSkillsDir, "catalog.txt"), "fresh");
47+
await writeFile(path.join(claudeProjectsDir, "session.jsonl"), "claude");
48+
await writeFile(path.join(restoredSessionEnvDir, "restored"), "state");
49+
await writeFile(path.join(codexSessionsDir, "rollout.jsonl"), "codex");
50+
await writeFile(codexConfigPath, "fresh = true");
51+
52+
await configurePersistentAgentState(stateRoot, homeDir);
53+
54+
const mappings = [
55+
[".claude/projects", "claude/projects"],
56+
[".claude/session-env", "claude/session-env"],
57+
[".claude/plans", "claude/plans"],
58+
[".claude/todos", "claude/todos"],
59+
[".codex/sessions", "codex/sessions"],
60+
[".codex/shell_snapshots", "codex/shell_snapshots"],
61+
];
62+
63+
for (const [sourceRelative, targetRelative] of mappings) {
64+
const source = path.join(homeDir, sourceRelative);
65+
const target = path.join(stateRoot, targetRelative);
66+
expect((await lstat(source)).isSymbolicLink()).toBe(true);
67+
expect(path.resolve(path.dirname(source), await readlink(source))).toBe(
68+
target,
69+
);
70+
await writeFile(path.join(source, "marker"), sourceRelative);
71+
await expect(
72+
readFile(path.join(target, "marker"), "utf-8"),
73+
).resolves.toBe(sourceRelative);
74+
}
75+
76+
await configurePersistentAgentState(stateRoot, homeDir);
77+
78+
await expect(
79+
readFile(
80+
path.join(stateRoot, "claude", "projects", "session.jsonl"),
81+
"utf-8",
82+
),
83+
).resolves.toBe("claude");
84+
await expect(
85+
readFile(
86+
path.join(stateRoot, "codex", "sessions", "rollout.jsonl"),
87+
"utf-8",
88+
),
89+
).resolves.toBe("codex");
90+
await expect(
91+
readFile(path.join(restoredSessionEnvDir, "restored"), "utf-8"),
92+
).resolves.toBe("state");
93+
94+
await expect(
95+
readFile(path.join(claudeSkillsDir, "catalog.txt"), "utf-8"),
96+
).resolves.toBe("fresh");
97+
await expect(readFile(codexConfigPath, "utf-8")).resolves.toBe(
98+
"fresh = true",
99+
);
100+
});
101+
});

‎packages/agent/src/persistent-agent-state.ts‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
import {
2+
cp,
3+
lstat,
4+
mkdir,
5+
readdir,
6+
readlink,
7+
rename,
8+
rm,
9+
symlink,
10+
} from "node:fs/promises";
11+
import os from "node:os";
12+
import path from "node:path";
13+
14+
interface PersistentDirectory {
15+
source: string;
16+
target: string;
17+
}
18+
19+
function isMissingPathError(error: unknown): boolean {
20+
return (
21+
error instanceof Error &&
22+
"code" in error &&
23+
(error as NodeJS.ErrnoException).code === "ENOENT"
24+
);
25+
}
26+
27+
function isCrossDeviceError(error: unknown): boolean {
28+
return (
29+
error instanceof Error &&
30+
"code" in error &&
31+
(error as NodeJS.ErrnoException).code === "EXDEV"
32+
);
33+
}
34+
35+
async function migrateExistingDirectory(
36+
source: string,
37+
target: string,
38+
): Promise<void> {
39+
const sourceEntries = await readdir(source);
40+
const targetEntries = await readdir(target);
41+
if (sourceEntries.length === 0) {
42+
await rm(source, { recursive: true, force: true });
43+
return;
44+
}
45+
if (targetEntries.length > 0) {
46+
throw new Error(
47+
`Cannot migrate persistent agent state from ${source}: ${target} is not empty`,
48+
);
49+
}
50+
51+
await rm(target, { recursive: true, force: true });
52+
try {
53+
await rename(source, target);
54+
} catch (error) {
55+
if (!isCrossDeviceError(error)) throw error;
56+
await cp(source, target, {
57+
recursive: true,
58+
errorOnExist: true,
59+
force: false,
60+
});
61+
await rm(source, { recursive: true, force: true });
62+
}
63+
}
64+
65+
async function linkPersistentDirectory({
66+
source,
67+
target,
68+
}: PersistentDirectory): Promise<void> {
69+
const resolvedSource = path.resolve(source);
70+
const resolvedTarget = path.resolve(target);
71+
if (resolvedSource === resolvedTarget) return;
72+
73+
await mkdir(resolvedTarget, { recursive: true });
74+
await mkdir(path.dirname(resolvedSource), { recursive: true });
75+
76+
try {
77+
const sourceStats = await lstat(resolvedSource);
78+
if (sourceStats.isSymbolicLink()) {
79+
const currentTarget = path.resolve(
80+
path.dirname(resolvedSource),
81+
await readlink(resolvedSource),
82+
);
83+
if (currentTarget === resolvedTarget) return;
84+
await rm(resolvedSource, { force: true });
85+
} else if (sourceStats.isDirectory()) {
86+
await migrateExistingDirectory(resolvedSource, resolvedTarget);
87+
} else {
88+
throw new Error(
89+
`Persistent agent state source must be a directory: ${resolvedSource}`,
90+
);
91+
}
92+
} catch (error) {
93+
if (!isMissingPathError(error)) throw error;
94+
}
95+
96+
await symlink(resolvedTarget, resolvedSource, "dir");
97+
}
98+
99+
export async function configurePersistentAgentState(
100+
stateRoot: string,
101+
homeDir: string = os.homedir(),
102+
): Promise<void> {
103+
if (!path.isAbsolute(stateRoot)) {
104+
throw new Error("Persistent agent state root must be an absolute path");
105+
}
106+
107+
const claudeConfigDir =
108+
process.env.CLAUDE_CONFIG_DIR || path.join(homeDir, ".claude");
109+
const codexHome = process.env.CODEX_HOME || path.join(homeDir, ".codex");
110+
const directories: PersistentDirectory[] = [
111+
{
112+
source: path.join(claudeConfigDir, "projects"),
113+
target: path.join(stateRoot, "claude", "projects"),
114+
},
115+
{
116+
source: path.join(claudeConfigDir, "session-env"),
117+
target: path.join(stateRoot, "claude", "session-env"),
118+
},
119+
{
120+
source: path.join(claudeConfigDir, "plans"),
121+
target: path.join(stateRoot, "claude", "plans"),
122+
},
123+
{
124+
source: path.join(claudeConfigDir, "todos"),
125+
target: path.join(stateRoot, "claude", "todos"),
126+
},
127+
{
128+
source: path.join(codexHome, "sessions"),
129+
target: path.join(stateRoot, "codex", "sessions"),
130+
},
131+
{
132+
source: path.join(codexHome, "shell_snapshots"),
133+
target: path.join(stateRoot, "codex", "shell_snapshots"),
134+
},
135+
];
136+
137+
// Avoid leaving other mappings mutating after one setup operation fails.
138+
for (const directory of directories) {
139+
await linkPersistentDirectory(directory);
140+
}
141+
}

‎packages/agent/src/server/agent-server.test.ts‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,13 @@
11
import { createHash } from "node:crypto";
2-
import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
2+
import {
3+
lstat,
4+
mkdir,
5+
mkdtemp,
6+
readFile,
7+
readlink,
8+
rm,
9+
writeFile,
10+
} from "node:fs/promises";
311
import { tmpdir } from "node:os";
412
import { join } from "node:path";
513
import type { ContentBlock } from "@agentclientprotocol/sdk";
@@ -442,6 +450,42 @@ describe("AgentServer HTTP Mode", () => {
442450
sessionInitMs: expect.any(Number),
443451
});
444452
}, 30000);
453+
454+
it("links native agent state before initializing the session", async () => {
455+
const originalClaudeConfigDir = process.env.CLAUDE_CONFIG_DIR;
456+
const originalCodexHome = process.env.CODEX_HOME;
457+
const claudeConfigDir = join(repo.path, ".claude-test");
458+
const codexHome = join(repo.path, ".codex-test");
459+
const agentStateDir = join(repo.path, ".posthog", "agent-state");
460+
process.env.CLAUDE_CONFIG_DIR = claudeConfigDir;
461+
process.env.CODEX_HOME = codexHome;
462+
463+
try {
464+
await createServer({ agentStateDir }).start();
465+
466+
const claudeProjects = join(claudeConfigDir, "projects");
467+
const codexSessions = join(codexHome, "sessions");
468+
expect((await lstat(claudeProjects)).isSymbolicLink()).toBe(true);
469+
expect((await lstat(codexSessions)).isSymbolicLink()).toBe(true);
470+
expect(await readlink(claudeProjects)).toBe(
471+
join(agentStateDir, "claude", "projects"),
472+
);
473+
expect(await readlink(codexSessions)).toBe(
474+
join(agentStateDir, "codex", "sessions"),
475+
);
476+
} finally {
477+
if (originalClaudeConfigDir === undefined) {
478+
delete process.env.CLAUDE_CONFIG_DIR;
479+
} else {
480+
process.env.CLAUDE_CONFIG_DIR = originalClaudeConfigDir;
481+
}
482+
if (originalCodexHome === undefined) {
483+
delete process.env.CODEX_HOME;
484+
} else {
485+
process.env.CODEX_HOME = originalCodexHome;
486+
}
487+
}
488+
}, 30000);
445489
});
446490

447491
describe("turn completion", () => {

‎packages/agent/src/server/agent-server.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ import {
5151
import type { PermissionMode } from "../execution-mode";
5252
import { DEFAULT_CODEX_MODEL, fetchGatewayModels } from "../gateway-models";
5353
import { HandoffCheckpointTracker } from "../handoff-checkpoint";
54+
import { configurePersistentAgentState } from "../persistent-agent-state";
5455
import { PostHogAPIClient } from "../posthog-api";
5556
import {
5657
findPrUrls,
@@ -642,6 +643,10 @@ export class AgentServer {
642643
}
643644

644645
async start(): Promise<void> {
646+
if (this.config.agentStateDir) {
647+
await configurePersistentAgentState(this.config.agentStateDir);
648+
}
649+
645650
await new Promise<void>((resolve) => {
646651
this.server = serve(
647652
{

‎packages/agent/src/server/bin.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ const envSchema = z.object({
3232
POSTHOG_CODE_REASONING_EFFORT: z
3333
.enum(["low", "medium", "high", "xhigh", "max"])
3434
.optional(),
35+
POSTHOG_AGENT_STATE_DIR: z.string().startsWith("/").optional(),
3536
POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN: z.string().min(1).optional(),
3637
// Base URL for the event-ingest POST only; falls back to POSTHOG_API_URL when unset.
3738
POSTHOG_TASK_RUN_EVENT_INGEST_URL: z.url().optional(),
@@ -174,6 +175,7 @@ program
174175

175176
const server = new AgentServer({
176177
port: parseInt(options.port, 10),
178+
agentStateDir: env.POSTHOG_AGENT_STATE_DIR,
177179
jwtPublicKey: env.JWT_PUBLIC_KEY,
178180
eventIngestToken: env.POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN,
179181
eventIngestBaseUrl: env.POSTHOG_TASK_RUN_EVENT_INGEST_URL,

‎packages/agent/src/server/types.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ export interface ClaudeCodeConfig {
1212

1313
export interface AgentServerConfig {
1414
port: number;
15+
agentStateDir?: string;
1516
repositoryPath?: string;
1617
repoReadyFile?: string;
1718
apiUrl: string;

0 commit comments

Comments
 (0)