From 49122e252b73641ab805d7774e7099348dd1d019 Mon Sep 17 00:00:00 2001 From: Jiayuan Date: Fri, 30 Jan 2026 01:29:21 +0800 Subject: [PATCH 1/2] feat(agent): add session persistence and compaction --- src/agent/cli.ts | 8 ++- src/agent/runner.ts | 56 +++++++++++++-- src/agent/session/compaction.ts | 15 ++++ src/agent/session/session-manager.ts | 102 +++++++++++++++++++++++++++ src/agent/session/storage.ts | 65 +++++++++++++++++ src/agent/session/types.ts | 12 ++++ src/agent/types.ts | 1 + src/hub/agent.ts | 1 + 8 files changed, 255 insertions(+), 5 deletions(-) create mode 100644 src/agent/session/compaction.ts create mode 100644 src/agent/session/session-manager.ts create mode 100644 src/agent/session/storage.ts create mode 100644 src/agent/session/types.ts diff --git a/src/agent/cli.ts b/src/agent/cli.ts index 18c9ad9f..f7a3d343 100644 --- a/src/agent/cli.ts +++ b/src/agent/cli.ts @@ -7,11 +7,12 @@ type CliOptions = { system?: string; thinking?: string; cwd?: string; + session?: string; help?: boolean; }; function printUsage() { - console.log("Usage: pnpm agent:cli [--provider PROVIDER] [--model MODEL] [--system TEXT] [--thinking LEVEL] [--cwd DIR] "); + console.log("Usage: pnpm agent:cli [--provider PROVIDER] [--model MODEL] [--system TEXT] [--thinking LEVEL] [--cwd DIR] [--session ID] "); console.log(" echo \"your prompt\" | pnpm agent:cli"); } @@ -47,6 +48,10 @@ function parseArgs(argv: string[]) { opts.cwd = args.shift(); continue; } + if (arg === "--session") { + opts.session = args.shift(); + continue; + } if (arg === "--") { promptParts.push(...args); break; @@ -88,6 +93,7 @@ async function main() { systemPrompt: opts.system, thinkingLevel: opts.thinking as any, cwd: opts.cwd, + sessionId: opts.session, }); const result = await agent.run(finalPrompt); diff --git a/src/agent/runner.ts b/src/agent/runner.ts index 1b3013ee..e6cd56be 100644 --- a/src/agent/runner.ts +++ b/src/agent/runner.ts @@ -1,11 +1,13 @@ -import { Agent as PiAgentCore, type AgentEvent } from "@mariozechner/pi-agent-core"; +import { Agent as PiAgentCore, type AgentEvent, type AgentMessage } from "@mariozechner/pi-agent-core"; import type { AgentOptions, AgentRunResult } from "./types.js"; import { createAgentOutput } from "./output.js"; import { resolveModel, resolveTools } from "./tools.js"; +import { SessionManager } from "./session/session-manager.js"; export class Agent { private readonly agent: PiAgentCore; private readonly output; + private readonly session: SessionManager; constructor(options: AgentOptions = {}) { const stdout = options.logger?.stdout ?? process.stdout; @@ -14,11 +16,39 @@ export class Agent { this.agent = new PiAgentCore(); if (options.systemPrompt) this.agent.setSystemPrompt(options.systemPrompt); - if (options.thinkingLevel) this.agent.setThinkingLevel(options.thinkingLevel); - this.agent.setModel(resolveModel(options)); + const sessionId = options.sessionId ?? "default"; + this.session = new SessionManager({ sessionId }); + const storedMeta = this.session.getMeta(); + if (!options.thinkingLevel && storedMeta?.thinkingLevel) { + this.agent.setThinkingLevel(storedMeta.thinkingLevel as any); + } else if (options.thinkingLevel) { + this.agent.setThinkingLevel(options.thinkingLevel); + } + + const model = options.provider && options.model ? resolveModel(options) : resolveModel({ + ...options, + provider: storedMeta?.provider, + model: storedMeta?.model, + }); + this.agent.setModel(model); this.agent.setTools(resolveTools(options)); - this.agent.subscribe((event: AgentEvent) => this.output.handleEvent(event)); + + const restoredMessages = this.session.loadMessages(); + if (restoredMessages.length > 0) { + this.agent.replaceMessages(restoredMessages); + } + + this.session.saveMeta({ + provider: this.agent.state.model?.provider, + model: this.agent.state.model?.id, + thinkingLevel: this.agent.state.thinkingLevel, + }); + + this.agent.subscribe((event: AgentEvent) => { + this.output.handleEvent(event); + this.handleSessionEvent(event); + }); } async run(prompt: string): Promise { @@ -26,4 +56,22 @@ export class Agent { await this.agent.prompt(prompt); return { text: this.output.state.lastAssistantText, error: this.agent.state.error }; } + + private handleSessionEvent(event: AgentEvent) { + if (event.type === "message_end") { + const message = event.message as AgentMessage; + this.session.saveMessage(message); + if (message.role === "assistant") { + void this.maybeCompact(); + } + } + } + + private async maybeCompact() { + const messages = this.agent.state.messages.slice(); + const result = await this.session.maybeCompact(messages); + if (result?.kept) { + this.agent.replaceMessages(result.kept); + } + } } diff --git a/src/agent/session/compaction.ts b/src/agent/session/compaction.ts new file mode 100644 index 00000000..dec56363 --- /dev/null +++ b/src/agent/session/compaction.ts @@ -0,0 +1,15 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; + +export type CompactionResult = { + kept: AgentMessage[]; + removedCount: number; +}; + +export function compactMessages(messages: AgentMessage[], maxMessages: number, keepLast: number) { + if (messages.length <= maxMessages) return null; + const kept = messages.slice(-keepLast); + return { + kept, + removedCount: messages.length - kept.length, + } satisfies CompactionResult; +} diff --git a/src/agent/session/session-manager.ts b/src/agent/session/session-manager.ts new file mode 100644 index 00000000..c62e52e3 --- /dev/null +++ b/src/agent/session/session-manager.ts @@ -0,0 +1,102 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import type { SessionEntry, SessionMeta } from "./types.js"; +import { appendEntry, readEntries, writeEntries } from "./storage.js"; +import { compactMessages } from "./compaction.js"; + +export type SessionManagerOptions = { + sessionId: string; + baseDir?: string; + maxMessages?: number; + keepLast?: number; +}; + +export class SessionManager { + private readonly sessionId: string; + private readonly baseDir?: string; + private readonly maxMessages: number; + private readonly keepLast: number; + private queue: Promise = Promise.resolve(); + private meta: SessionMeta | undefined; + + constructor(options: SessionManagerOptions) { + this.sessionId = options.sessionId; + this.baseDir = options.baseDir; + this.maxMessages = options.maxMessages ?? 80; + this.keepLast = options.keepLast ?? 60; + this.meta = this.loadMeta(); + } + + loadEntries(): SessionEntry[] { + return readEntries(this.sessionId, { baseDir: this.baseDir }); + } + + loadMessages(): AgentMessage[] { + const entries = this.loadEntries(); + return entries + .filter((entry) => entry.type === "message") + .map((entry) => entry.message); + } + + loadMeta(): SessionMeta | undefined { + const entries = this.loadEntries(); + let meta: SessionMeta | undefined; + for (const entry of entries) { + if (entry.type === "meta") { + meta = entry.meta; + } + } + return meta; + } + + getMeta(): SessionMeta | undefined { + return this.meta; + } + + saveMeta(meta: SessionMeta) { + this.meta = meta; + void this.enqueue(() => + appendEntry( + this.sessionId, + { type: "meta", meta, timestamp: Date.now() }, + { baseDir: this.baseDir }, + ), + ); + } + + saveMessage(message: AgentMessage) { + void this.enqueue(() => + appendEntry( + this.sessionId, + { type: "message", message, timestamp: Date.now() }, + { baseDir: this.baseDir }, + ), + ); + } + + async maybeCompact(messages: AgentMessage[]) { + const result = compactMessages(messages, this.maxMessages, this.keepLast); + if (!result) return null; + const entries: SessionEntry[] = []; + if (this.meta) { + entries.push({ type: "meta", meta: this.meta, timestamp: Date.now() }); + } + for (const message of result.kept) { + entries.push({ type: "message", message, timestamp: Date.now() }); + } + entries.push({ + type: "compaction", + removed: result.removedCount, + kept: result.kept.length, + timestamp: Date.now(), + }); + await this.enqueue(() => + writeEntries(this.sessionId, entries, { baseDir: this.baseDir }), + ); + return result; + } + + private enqueue(task: () => Promise) { + this.queue = this.queue.then(task, task); + return this.queue; + } +} diff --git a/src/agent/session/storage.ts b/src/agent/session/storage.ts new file mode 100644 index 00000000..b4629acb --- /dev/null +++ b/src/agent/session/storage.ts @@ -0,0 +1,65 @@ +import { homedir } from "os"; +import { join } from "path"; +import { existsSync, mkdirSync, readFileSync } from "fs"; +import { appendFile, writeFile } from "fs/promises"; +import type { SessionEntry } from "./types.js"; + +export type SessionStorageOptions = { + baseDir?: string; +}; + +export function resolveBaseDir(options?: SessionStorageOptions) { + return options?.baseDir ?? join(homedir(), ".super-multica", "sessions"); +} + +export function resolveSessionDir(sessionId: string, options?: SessionStorageOptions) { + return join(resolveBaseDir(options), sessionId); +} + +export function resolveSessionPath(sessionId: string, options?: SessionStorageOptions) { + return join(resolveSessionDir(sessionId, options), "session.jsonl"); +} + +export function ensureSessionDir(sessionId: string, options?: SessionStorageOptions) { + const dir = resolveSessionDir(sessionId, options); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } +} + +export function readEntries(sessionId: string, options?: SessionStorageOptions): SessionEntry[] { + const filePath = resolveSessionPath(sessionId, options); + if (!existsSync(filePath)) return []; + const content = readFileSync(filePath, "utf8"); + const lines = content.split("\n").filter(Boolean); + const entries: SessionEntry[] = []; + for (const line of lines) { + try { + entries.push(JSON.parse(line) as SessionEntry); + } catch { + // Skip malformed lines + } + } + return entries; +} + +export async function appendEntry( + sessionId: string, + entry: SessionEntry, + options?: SessionStorageOptions, +) { + ensureSessionDir(sessionId, options); + const filePath = resolveSessionPath(sessionId, options); + await appendFile(filePath, `${JSON.stringify(entry)}\n`, "utf8"); +} + +export async function writeEntries( + sessionId: string, + entries: SessionEntry[], + options?: SessionStorageOptions, +) { + ensureSessionDir(sessionId, options); + const filePath = resolveSessionPath(sessionId, options); + const content = entries.map((entry) => JSON.stringify(entry)).join("\n"); + await writeFile(filePath, content ? `${content}\n` : "", "utf8"); +} diff --git a/src/agent/session/types.ts b/src/agent/session/types.ts new file mode 100644 index 00000000..cbc82685 --- /dev/null +++ b/src/agent/session/types.ts @@ -0,0 +1,12 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; + +export type SessionMeta = { + provider?: string; + model?: string; + thinkingLevel?: string; +}; + +export type SessionEntry = + | { type: "message"; message: AgentMessage; timestamp: number } + | { type: "meta"; meta: SessionMeta; timestamp: number } + | { type: "compaction"; removed: number; kept: number; timestamp: number }; diff --git a/src/agent/types.ts b/src/agent/types.ts index 8261886c..58769a39 100644 --- a/src/agent/types.ts +++ b/src/agent/types.ts @@ -16,5 +16,6 @@ export type AgentOptions = { systemPrompt?: string; thinkingLevel?: ThinkingLevel; cwd?: string; + sessionId?: string; logger?: AgentLogger; }; diff --git a/src/hub/agent.ts b/src/hub/agent.ts index 334cbad0..f2cda839 100644 --- a/src/hub/agent.ts +++ b/src/hub/agent.ts @@ -21,6 +21,7 @@ export class Agent { stdout: this.createChannelStream("[assistant] "), stderr: this.createChannelStream("[tool] "), }, + sessionId: this.id, }); } From 674a88c61e50392109c78156777b23b90110d6af Mon Sep 17 00:00:00 2001 From: Jiayuan Date: Fri, 30 Jan 2026 01:33:08 +0800 Subject: [PATCH 2/2] docs: add agent CLI usage --- README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/README.md b/README.md index 00f5b4b9..27ae7934 100644 --- a/README.md +++ b/README.md @@ -19,9 +19,25 @@ pnpm install pnpm dev ``` +## Agent CLI + +Use the agent module directly from the CLI for isolated testing. + +```bash +pnpm agent:cli "hello" + +# Persist a session under ~/.super-multica/sessions//session.jsonl +pnpm agent:cli --session demo "remember my name is Alice" +pnpm agent:cli --session demo "what's my name?" + +# Override provider/model +pnpm agent:cli --provider openai --model gpt-4o-mini "hi" +``` + ## Scripts - `pnpm dev` - Run in development mode +- `pnpm agent:cli` - Run the agent CLI for module-level testing - `pnpm build` - Build for production - `pnpm start` - Run production build - `pnpm typecheck` - Type check without emitting