Skip to content

Commit 2ba2644

Browse files
jonathanlabjonathanlab
committed
feat: add session loading
1 parent 9063af6 commit 2ba2644

File tree

7 files changed

+1222
-37
lines changed

7 files changed

+1222
-37
lines changed

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,6 @@ AGENTS.md
3434
**.car
3535

3636
.envrc
37+
38+
# Session store used for example-client.ts
39+
.session-store.json

‎packages/agent/example-client.ts‎

Lines changed: 263 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#!/usr/bin/env node
22

33
import { spawn } from "node:child_process";
4+
import { existsSync, readFileSync, writeFileSync } from "node:fs";
45
import { dirname, join } from "node:path";
56
import * as readline from "node:readline/promises";
67
import { Readable, Writable } from "node:stream";
@@ -19,8 +20,47 @@ import {
1920
type WriteTextFileRequest,
2021
type WriteTextFileResponse,
2122
} from "@agentclientprotocol/sdk";
23+
import { PostHogAPIClient } from "./src/posthog-api.js";
24+
import type { SessionPersistenceConfig } from "./src/session-store.js";
25+
26+
// PostHog configuration - set via env vars
27+
const POSTHOG_CONFIG = {
28+
apiUrl: process.env.POSTHOG_API_URL || "",
29+
apiKey: process.env.POSTHOG_API_KEY || "",
30+
projectId: parseInt(process.env.POSTHOG_PROJECT_ID || "0", 10),
31+
};
32+
33+
// Simple file-based storage for session -> persistence mapping
34+
const SESSION_STORE_PATH = join(
35+
dirname(fileURLToPath(import.meta.url)),
36+
".session-store.json",
37+
);
38+
39+
interface SessionMapping {
40+
[sessionId: string]: SessionPersistenceConfig;
41+
}
42+
43+
function loadSessionMappings(): SessionMapping {
44+
if (existsSync(SESSION_STORE_PATH)) {
45+
return JSON.parse(readFileSync(SESSION_STORE_PATH, "utf-8"));
46+
}
47+
return {};
48+
}
49+
50+
function saveSessionMapping(
51+
sessionId: string,
52+
config: SessionPersistenceConfig,
53+
): void {
54+
const mappings = loadSessionMappings();
55+
mappings[sessionId] = config;
56+
writeFileSync(SESSION_STORE_PATH, JSON.stringify(mappings, null, 2));
57+
}
2258

2359
class ExampleClient implements Client {
60+
isReplaying = false;
61+
replayCount = 0;
62+
currentSessionId?: string;
63+
2464
async requestPermission(
2565
params: RequestPermissionRequest,
2666
): Promise<RequestPermissionResponse> {
@@ -57,6 +97,12 @@ class ExampleClient implements Client {
5797
async sessionUpdate(params: SessionNotification): Promise<void> {
5898
const update = params.update;
5999

100+
if (this.isReplaying) {
101+
this.replayCount++;
102+
this.renderReplayUpdate(update);
103+
return;
104+
}
105+
60106
switch (update.sessionUpdate) {
61107
case "agent_message_chunk":
62108
if (update.content.type === "text") {
@@ -83,6 +129,41 @@ class ExampleClient implements Client {
83129
}
84130
}
85131

132+
renderReplayUpdate(update: SessionNotification["update"]): void {
133+
const dim = "\x1b[2m";
134+
const reset = "\x1b[0m";
135+
136+
switch (update.sessionUpdate) {
137+
case "agent_message_chunk":
138+
if (update.content.type === "text") {
139+
process.stdout.write(`${dim}${update.content.text}${reset}`);
140+
}
141+
break;
142+
case "user_message_chunk":
143+
if (update.content.type === "text") {
144+
process.stdout.write(
145+
`\n${dim}💬 You: ${update.content.text}${reset}`,
146+
);
147+
}
148+
break;
149+
case "tool_call":
150+
console.log(`${dim}🔧 ${update.title} (${update.status})${reset}`);
151+
break;
152+
case "tool_call_update":
153+
if (update.status === "completed" || update.status === "failed") {
154+
console.log(`${dim} └─ ${update.status}${reset}`);
155+
}
156+
break;
157+
case "agent_thought_chunk":
158+
if (update.content.type === "text") {
159+
process.stdout.write(`${dim}💭 ${update.content.text}${reset}`);
160+
}
161+
break;
162+
default:
163+
break;
164+
}
165+
}
166+
86167
async writeTextFile(
87168
params: WriteTextFileRequest,
88169
): Promise<WriteTextFileResponse> {
@@ -106,16 +187,115 @@ class ExampleClient implements Client {
106187
content: "Mock file content",
107188
};
108189
}
190+
191+
async extNotification(
192+
method: string,
193+
params: Record<string, unknown>,
194+
): Promise<void> {
195+
if (method === "_posthog/sdk_session") {
196+
const { sessionId, sdkSessionId } = params as {
197+
sessionId: string;
198+
sdkSessionId: string;
199+
};
200+
// Update the session mapping with the SDK session ID
201+
const mappings = loadSessionMappings();
202+
if (mappings[sessionId]) {
203+
mappings[sessionId].sdkSessionId = sdkSessionId;
204+
writeFileSync(SESSION_STORE_PATH, JSON.stringify(mappings, null, 2));
205+
console.log(` 🔗 SDK session ID stored: ${sdkSessionId}`);
206+
}
207+
}
208+
}
209+
}
210+
211+
async function prompt(message: string): Promise<string> {
212+
const rl = readline.createInterface({
213+
input: process.stdin,
214+
output: process.stdout,
215+
});
216+
const answer = await rl.question(message);
217+
rl.close();
218+
return answer.trim();
109219
}
110220

111221
async function main() {
112222
const __filename = fileURLToPath(import.meta.url);
113223
const __dirname = dirname(__filename);
114224
const agentPath = join(__dirname, "agent.ts");
115225

226+
// Check for session ID argument: npx tsx example-client.ts [sessionId]
227+
const existingSessionId = process.argv[2];
228+
229+
// Load existing session mappings
230+
const sessionMappings = loadSessionMappings();
231+
232+
// Check if we're reloading an existing session
233+
let persistence: SessionPersistenceConfig | undefined;
234+
235+
if (existingSessionId && sessionMappings[existingSessionId]) {
236+
// Use existing persistence config
237+
persistence = sessionMappings[existingSessionId];
238+
console.log(`🔗 Loading existing session: ${existingSessionId}`);
239+
console.log(` 📋 Task: ${persistence.taskId}`);
240+
console.log(` 🏃 Run: ${persistence.runId}`);
241+
if (persistence.sdkSessionId) {
242+
console.log(
243+
` 🧠 SDK Session: ${persistence.sdkSessionId} (context will be restored)`,
244+
);
245+
}
246+
} else if (!existingSessionId) {
247+
// Create new Task/TaskRun for new sessions (only if PostHog is configured)
248+
if (
249+
POSTHOG_CONFIG.apiUrl &&
250+
POSTHOG_CONFIG.apiKey &&
251+
POSTHOG_CONFIG.projectId
252+
) {
253+
console.log("🔗 Connecting to PostHog...");
254+
const posthogClient = new PostHogAPIClient(POSTHOG_CONFIG);
255+
256+
try {
257+
// Create a task for this session
258+
const task = await posthogClient.createTask({
259+
title: `ACP Session ${new Date().toISOString()}`,
260+
description: "Session created by example-client",
261+
});
262+
console.log(`📋 Created task: ${task.id}`);
263+
264+
// Create a task run
265+
const taskRun = await posthogClient.createTaskRun(task.id);
266+
console.log(`🏃 Created task run: ${taskRun.id}`);
267+
console.log(`📦 Log URL: ${taskRun.log_url}`);
268+
269+
persistence = {
270+
taskId: task.id,
271+
runId: taskRun.id,
272+
logUrl: taskRun.log_url,
273+
};
274+
} catch (error) {
275+
console.error("❌ Failed to create Task/TaskRun:", error);
276+
console.log(" Continuing without S3 persistence...\n");
277+
}
278+
} else {
279+
console.log(
280+
"ℹ️ PostHog not configured (set POSTHOG_API_URL, POSTHOG_API_KEY, POSTHOG_PROJECT_ID)",
281+
);
282+
console.log(" Running without persistence...\n");
283+
}
284+
} else {
285+
console.log(`⚠️ Session ${existingSessionId} not found in local store`);
286+
console.log(" Starting fresh without persistence...\n");
287+
}
288+
116289
// Spawn the agent as a subprocess using tsx
290+
// Pass PostHog config as env vars so agent can create its own SessionStore
117291
const agentProcess = spawn("npx", ["tsx", agentPath], {
118292
stdio: ["pipe", "pipe", "inherit"],
293+
env: {
294+
...process.env,
295+
POSTHOG_API_URL: POSTHOG_CONFIG.apiUrl,
296+
POSTHOG_API_KEY: POSTHOG_CONFIG.apiKey,
297+
POSTHOG_PROJECT_ID: String(POSTHOG_CONFIG.projectId),
298+
},
119299
});
120300

121301
// Create streams to communicate with the agent
@@ -144,28 +324,93 @@ async function main() {
144324
console.log(
145325
`✅ Connected to agent (protocol v${initResult.protocolVersion})`,
146326
);
327+
console.log(
328+
` Load session supported: ${initResult.agentCapabilities?.loadSession ?? false}`,
329+
);
147330

148-
// Create a new session
149-
const sessionResult = await connection.newSession({
150-
cwd: process.cwd(),
151-
mcpServers: [],
152-
});
331+
let sessionId: string;
153332

154-
console.log(`📝 Created session: ${sessionResult.sessionId}`);
155-
console.log(`💬 User: Hello, agent!\n`);
333+
if (existingSessionId) {
334+
// Load existing session
335+
console.log(`\n🔄 Loading session: ${existingSessionId}`);
336+
console.log(`${"─".repeat(50)}`);
337+
console.log(`📜 Conversation history:\n`);
156338

157-
// Send a test prompt
158-
const promptResult = await connection.prompt({
159-
sessionId: sessionResult.sessionId,
160-
prompt: [
161-
{
162-
type: "text",
163-
text: "Hello, agent!",
164-
},
165-
],
166-
});
339+
client.isReplaying = true;
340+
client.replayCount = 0;
341+
342+
await connection.loadSession({
343+
sessionId: existingSessionId,
344+
cwd: process.cwd(),
345+
mcpServers: [],
346+
_meta: persistence
347+
? { persistence, sdkSessionId: persistence.sdkSessionId }
348+
: undefined,
349+
});
350+
351+
client.isReplaying = false;
352+
sessionId = existingSessionId;
167353

168-
console.log(`\n\n✅ Agent completed with: ${promptResult.stopReason}`);
354+
console.log(`\n${"─".repeat(50)}`);
355+
console.log(`✅ Replayed ${client.replayCount} events from history\n`);
356+
} else {
357+
// Create a new session
358+
const sessionResult = await connection.newSession({
359+
cwd: process.cwd(),
360+
mcpServers: [],
361+
_meta: persistence ? { persistence } : undefined,
362+
});
363+
364+
sessionId = sessionResult.sessionId;
365+
console.log(`📝 Created session: ${sessionId}`);
366+
if (persistence) {
367+
// Save the mapping so we can reload later
368+
saveSessionMapping(sessionId, persistence);
369+
console.log(
370+
` 📦 S3 persistence enabled (task: ${persistence.taskId})`,
371+
);
372+
}
373+
console.log(
374+
` (Run with session ID to reload: npx tsx example-client.ts ${sessionId})\n`,
375+
);
376+
}
377+
378+
// Interactive prompt loop
379+
while (true) {
380+
const userInput = await prompt("\n💬 You: ");
381+
382+
if (
383+
userInput.toLowerCase() === "/quit" ||
384+
userInput.toLowerCase() === "/exit"
385+
) {
386+
console.log("\n👋 Goodbye!");
387+
break;
388+
}
389+
390+
if (userInput.toLowerCase() === "/session") {
391+
console.log(`\n📝 Current session ID: ${sessionId}`);
392+
console.log(` Reload with: npx tsx example-client.ts ${sessionId}`);
393+
continue;
394+
}
395+
396+
if (!userInput) {
397+
continue;
398+
}
399+
400+
console.log("");
401+
402+
const promptResult = await connection.prompt({
403+
sessionId,
404+
prompt: [
405+
{
406+
type: "text",
407+
text: userInput,
408+
},
409+
],
410+
});
411+
412+
console.log(`\n\n✅ Agent completed with: ${promptResult.stopReason}`);
413+
}
169414
} catch (error) {
170415
console.error("[Client] Error:", error);
171416
} finally {

‎packages/agent/index.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
// Main entry point - re-exports from src
22

3+
// Session persistence
4+
export type { SessionPersistenceConfig } from "./src/session-store.js";
5+
export { SessionStore } from "./src/session-store.js";
36
// TODO: Refactor - legacy adapter removed
47
// export { ClaudeAdapter } from "./src/adapters/claude-legacy/claude-adapter.js";
58
// export type { ProviderAdapter } from "./src/adapters/types.js";

0 commit comments

Comments
 (0)