refactor(agent): consolidate harness session abstraction

This commit is contained in:
Mario Zechner
2026-05-03 22:01:40 +02:00
parent e612149369
commit cdde2e893d
12 changed files with 599 additions and 397 deletions

View File

@@ -19,7 +19,7 @@ import type {
ExecutionEnv,
NavigateTreeResult,
PromptTemplate,
SessionTree,
Session,
Skill,
SystemPromptInputs,
} from "./types.js";
@@ -56,7 +56,7 @@ export class DefaultAgentHarness implements AgentHarness {
readonly conversation: AgentHarnessConversationState;
readonly operation: AgentHarnessOperationState;
private sessionTree: SessionTree;
private session: Session;
private promptTemplates: PromptTemplate[];
private skills: Skill[];
private requestAuth?: AgentHarnessOptions["requestAuth"];
@@ -70,7 +70,7 @@ export class DefaultAgentHarness implements AgentHarness {
constructor(options: AgentHarnessOptions) {
this.agent = options.agent;
this.env = options.env;
this.sessionTree = options.sessionTree;
this.session = options.session;
this.promptTemplates = options.promptTemplates ?? [];
this.skills = options.skills ?? [];
this.requestAuth = options.requestAuth;
@@ -78,7 +78,7 @@ export class DefaultAgentHarness implements AgentHarness {
this.toolRegistry.set(tool.name, tool);
}
this.conversation = {
sessionTree: options.sessionTree,
session: options.session,
model: options.initialModel ?? this.agent.state.model,
thinkingLevel: options.initialThinkingLevel ?? this.agent.state.thinkingLevel,
activeToolNames: options.initialActiveToolNames ?? this.agent.state.tools.map((tool) => tool.name),
@@ -191,7 +191,7 @@ export class DefaultAgentHarness implements AgentHarness {
}
private async syncFromTree(): Promise<void> {
const context = await this.sessionTree.buildContext();
const context = await this.session.buildContext();
this.agent.state.messages = context.messages;
if (context.model && this.conversation.model) {
// leave active model untouched; harness-level model is source of truth
@@ -201,7 +201,7 @@ export class DefaultAgentHarness implements AgentHarness {
private async applyPendingMutations(): Promise<void> {
for (const message of this.operation.pendingMutations.appendMessages) {
await this.sessionTree.appendMessage(message);
await this.session.appendMessage(message);
}
this.operation.pendingMutations.appendMessages = [];
@@ -210,7 +210,7 @@ export class DefaultAgentHarness implements AgentHarness {
const previousModel = this.conversation.model;
this.conversation.model = model;
this.agent.state.model = model;
await this.sessionTree.appendModelChange(model.provider, model.id);
await this.session.appendModelChange(model.provider, model.id);
await this.emitOwn({ type: "model_select", model, previousModel, source: "set" });
this.operation.pendingMutations.model = undefined;
}
@@ -220,7 +220,7 @@ export class DefaultAgentHarness implements AgentHarness {
const previousLevel = this.conversation.thinkingLevel;
this.conversation.thinkingLevel = level;
this.agent.state.thinkingLevel = level;
await this.sessionTree.appendThinkingLevelChange(level);
await this.session.appendThinkingLevelChange(level);
await this.emitOwn({ type: "thinking_level_select", level, previousLevel });
this.operation.pendingMutations.thinkingLevel = undefined;
}
@@ -258,7 +258,7 @@ export class DefaultAgentHarness implements AgentHarness {
}
}
if (event.type === "message_end") {
await this.sessionTree.appendMessage(event.message);
await this.session.appendMessage(event.message);
}
if (event.type === "turn_end") {
const hadPendingMutations =
@@ -342,7 +342,7 @@ export class DefaultAgentHarness implements AgentHarness {
async appendMessage(message: AgentMessage): Promise<void> {
if (this.operation.idle) {
await this.sessionTree.appendMessage(message);
await this.session.appendMessage(message);
await this.syncFromTree();
} else {
this.operation.pendingMutations.appendMessages.push(message);
@@ -371,7 +371,7 @@ export class DefaultAgentHarness implements AgentHarness {
if (!model) throw new Error("No model set for compaction");
const auth = await this.requestAuth?.(model);
if (!auth) throw new Error("No auth available for compaction");
const branchEntries = await this.sessionTree.getBranch();
const branchEntries = await this.session.getBranch();
const preparation = prepareCompaction(branchEntries, DEFAULT_COMPACTION_SETTINGS);
if (!preparation) throw new Error("Nothing to compact");
const hookResult = await this.emitHook("session_before_compact", {
@@ -394,14 +394,14 @@ export class DefaultAgentHarness implements AgentHarness {
undefined,
this.conversation.thinkingLevel,
));
const entryId = await this.sessionTree.appendCompaction(
const entryId = await this.session.appendCompaction(
result.summary,
result.firstKeptEntryId,
result.tokensBefore,
result.details,
provided !== undefined,
);
const entry = await this.sessionTree.getEntry(entryId);
const entry = await this.session.getEntry(entryId);
await this.syncFromTree();
if (entry?.type === "compaction") {
await this.emitOwn({ type: "session_compact", compactionEntry: entry, fromHook: provided !== undefined });
@@ -414,11 +414,11 @@ export class DefaultAgentHarness implements AgentHarness {
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
): Promise<NavigateTreeResult> {
if (!this.operation.idle) throw new Error("navigateTree() requires idle harness");
const oldLeafId = await this.sessionTree.getLeafId();
const oldLeafId = await this.session.getLeafId();
if (oldLeafId === targetId) return { cancelled: false };
const targetEntry = await this.sessionTree.getEntry(targetId);
const targetEntry = await this.session.getEntry(targetId);
if (!targetEntry) throw new Error(`Entry ${targetId} not found`);
const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.sessionTree, oldLeafId, targetId);
const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.session, oldLeafId, targetId);
const preparation = {
targetId,
oldLeafId,
@@ -469,7 +469,7 @@ export class DefaultAgentHarness implements AgentHarness {
typeof content === "string"
? content
: content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.filter((c): c is { readonly type: "text"; readonly text: string } => c.type === "text")
.map((c) => c.text)
.join("");
} else if (targetEntry.type === "custom_message") {
@@ -478,26 +478,29 @@ export class DefaultAgentHarness implements AgentHarness {
typeof targetEntry.content === "string"
? targetEntry.content
: targetEntry.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.filter((c): c is { readonly type: "text"; readonly text: string } => c.type === "text")
.map((c) => c.text)
.join("");
} else {
newLeafId = targetId;
}
await this.sessionTree.moveTo(newLeafId);
if (summaryText) {
const summaryId = await this.sessionTree.appendBranchSummary(
newLeafId ?? "root",
summaryText,
summaryDetails,
hookResult?.summary !== undefined,
);
summaryEntry = await this.sessionTree.getEntry(summaryId);
const summaryId = await this.session.moveTo(
newLeafId,
summaryText
? {
summary: summaryText,
details: summaryDetails,
fromHook: hookResult?.summary !== undefined,
}
: undefined,
);
if (summaryId) {
summaryEntry = await this.session.getEntry(summaryId);
}
await this.syncFromTree();
await this.emitOwn({
type: "session_tree",
newLeafId: await this.sessionTree.getLeafId(),
newLeafId: await this.session.getLeafId(),
oldLeafId,
summaryEntry,
fromHook: hookResult?.summary !== undefined,
@@ -510,7 +513,7 @@ export class DefaultAgentHarness implements AgentHarness {
const previousModel = this.conversation.model;
this.conversation.model = model;
this.agent.state.model = model;
await this.sessionTree.appendModelChange(model.provider, model.id);
await this.session.appendModelChange(model.provider, model.id);
await this.emitOwn({ type: "model_select", model, previousModel, source: "set" });
} else {
this.operation.pendingMutations.model = model;
@@ -522,7 +525,7 @@ export class DefaultAgentHarness implements AgentHarness {
const previousLevel = this.conversation.thinkingLevel;
this.conversation.thinkingLevel = level;
this.agent.state.thinkingLevel = level;
await this.sessionTree.appendThinkingLevelChange(level);
await this.session.appendThinkingLevelChange(level);
await this.emitOwn({ type: "thinking_level_select", level, previousLevel });
} else {
this.operation.pendingMutations.thinkingLevel = level;

View File

@@ -5,7 +5,7 @@
* a summary of the branch being left so context isn't lost.
*/
import type { Model } from "@mariozechner/pi-ai";
import type { ImageContent, Model, TextContent } from "@mariozechner/pi-ai";
import { completeSimple } from "@mariozechner/pi-ai";
import type { AgentMessage } from "../../types.js";
import {
@@ -14,7 +14,7 @@ import {
createCompactionSummaryMessage,
createCustomMessage,
} from "../messages.js";
import type { SessionTree, SessionTreeEntry } from "../types.js";
import type { Session, SessionTreeEntry } from "../types.js";
import { estimateTokens } from "./compaction.js";
import {
computeFileLists,
@@ -96,7 +96,7 @@ export interface GenerateBranchSummaryOptions {
* @returns Entries to summarize and the common ancestor
*/
export async function collectEntriesForBranchSummary(
session: SessionTree,
session: Session,
oldLeafId: string | null,
targetId: string,
): Promise<CollectEntriesResult> {
@@ -125,7 +125,7 @@ export async function collectEntriesForBranchSummary(
while (current && current !== commonAncestorId) {
const entry = await session.getEntry(current);
if (!entry) break;
entries.push(entry);
entries.push(entry as SessionTreeEntry);
current = entry.parentId;
}
@@ -148,10 +148,16 @@ function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined
case "message":
// Skip tool results - context is in assistant's tool call
if (entry.message.role === "toolResult") return undefined;
return entry.message;
return entry.message as AgentMessage;
case "custom_message":
return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
return createCustomMessage(
entry.customType,
entry.content as string | (TextContent | ImageContent)[],
entry.display,
entry.details,
entry.timestamp,
);
case "branch_summary":
return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);

View File

@@ -5,7 +5,7 @@
* and after compaction the session is reloaded.
*/
import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai";
import type { AssistantMessage, ImageContent, Model, TextContent, Usage } from "@mariozechner/pi-ai";
import { completeSimple } from "@mariozechner/pi-ai";
import type { AgentMessage, ThinkingLevel } from "../../types.js";
import {
@@ -79,10 +79,16 @@ function extractFileOperations(
*/
function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined {
if (entry.type === "message") {
return entry.message;
return entry.message as AgentMessage;
}
if (entry.type === "custom_message") {
return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
return createCustomMessage(
entry.customType,
entry.content as string | (TextContent | ImageContent)[],
entry.display,
entry.details,
entry.timestamp,
);
}
if (entry.type === "branch_summary") {
return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
@@ -158,7 +164,7 @@ export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | unde
for (let i = entries.length - 1; i >= 0; i--) {
const entry = entries[i];
if (entry.type === "message") {
const usage = getAssistantUsage(entry.message);
const usage = getAssistantUsage(entry.message as AgentMessage);
if (usage) return usage;
}
}
@@ -405,7 +411,7 @@ export function findCutPoint(
if (entry.type !== "message") continue;
// Estimate this message's size
const messageTokens = estimateTokens(entry.message);
const messageTokens = estimateTokens(entry.message as AgentMessage);
accumulatedTokens += messageTokens;
// Check if we've exceeded the budget

View File

@@ -0,0 +1,142 @@
import { constants } from "node:fs";
import { access, mkdir, readdir, rm } from "node:fs/promises";
import { join, resolve } from "node:path";
import type {
JsonlSessionCreateOptions,
JsonlSessionListQuery,
JsonlSessionMetadata,
JsonlSessionRef,
JsonlSessionRepoApi,
JsonlSessionResolveOptions,
Session,
} from "../types.js";
import { JsonlSessionStorage, loadJsonlSessionMetadata } from "./jsonl-session-storage.js";
import { createSessionId, createTimestamp, getPathEntriesToFork, toSession } from "./session-repo.js";
async function exists(path: string): Promise<boolean> {
try {
await access(path, constants.F_OK);
return true;
} catch {
return false;
}
}
function encodeCwd(cwd: string): string {
return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`;
}
export class JsonlSessionRepo implements JsonlSessionRepoApi {
private sessionsRoot: string;
constructor(options: { sessionsRoot: string }) {
this.sessionsRoot = resolve(options.sessionsRoot);
}
private getSessionDir(cwd: string): string {
return join(this.sessionsRoot, encodeCwd(cwd));
}
private createSessionFilePath(cwd: string, sessionId: string, timestamp: string): string {
return join(this.getSessionDir(cwd), `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`);
}
private refPath(ref: JsonlSessionRef): string {
return resolve(ref.path);
}
async create(options: JsonlSessionCreateOptions): Promise<Session<JsonlSessionMetadata>> {
await mkdir(this.sessionsRoot, { recursive: true });
const id = options.id ?? createSessionId();
const createdAt = createTimestamp();
const filePath = this.createSessionFilePath(options.cwd, id, createdAt);
const storage = await JsonlSessionStorage.create(filePath, {
cwd: options.cwd,
sessionId: id,
parentSessionPath: options.parentSessionPath,
});
return toSession(storage);
}
async open(ref: JsonlSessionRef): Promise<Session<JsonlSessionMetadata>> {
const filePath = this.refPath(ref);
if (!(await exists(filePath))) {
throw new Error(`Session not found: ${filePath}`);
}
const storage = await JsonlSessionStorage.open(filePath);
return toSession(storage);
}
async list(query: JsonlSessionListQuery = {}): Promise<JsonlSessionMetadata[]> {
const dirs = query.cwd ? [this.getSessionDir(query.cwd)] : await this.listSessionDirs();
const sessions: JsonlSessionMetadata[] = [];
for (const dir of dirs) {
if (!(await exists(dir))) continue;
const files = (await readdir(dir)).filter((file) => file.endsWith(".jsonl")).map((file) => join(dir, file));
for (const filePath of files) {
try {
sessions.push(await loadJsonlSessionMetadata(filePath));
} catch {
// Ignore invalid session files when listing a directory.
}
}
}
sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime());
return sessions;
}
async resolve(ref: string, options: JsonlSessionResolveOptions = {}): Promise<JsonlSessionMetadata[]> {
if (ref.includes("/") || ref.includes("\\") || ref.endsWith(".jsonl")) {
try {
return [await loadJsonlSessionMetadata(resolve(ref))];
} catch {
return [];
}
}
const local = options.cwd
? (await this.list({ cwd: options.cwd })).filter((session) => session.id.startsWith(ref))
: [];
if (local.length > 0 || !options.searchAll) return local;
return (await this.list()).filter((session) => session.id.startsWith(ref));
}
async getMostRecent(query: JsonlSessionListQuery = {}): Promise<JsonlSessionMetadata | undefined> {
return (await this.list(query))[0];
}
async delete(ref: JsonlSessionRef): Promise<void> {
const filePath = this.refPath(ref);
await rm(filePath, { force: true });
}
async fork(
ref: JsonlSessionRef,
options: JsonlSessionCreateOptions & { entryId: string; position?: "before" | "at"; id?: string },
): Promise<Session<JsonlSessionMetadata>> {
const source = await this.open(ref);
const forkedEntries = await getPathEntriesToFork(
source.getStorage(),
options.entryId,
options.position ?? "before",
);
const sourceInfo = await source.getMetadata();
const id = options.id ?? createSessionId();
const createdAt = createTimestamp();
const storage = await JsonlSessionStorage.create(this.createSessionFilePath(options.cwd, id, createdAt), {
cwd: options.cwd,
sessionId: id,
parentSessionPath: options.parentSessionPath ?? sourceInfo.path,
});
for (const entry of forkedEntries) {
await storage.appendEntry(entry);
}
return toSession(storage);
}
private async listSessionDirs(): Promise<string[]> {
if (!(await exists(this.sessionsRoot))) return [];
const entries = await readdir(this.sessionsRoot, { withFileTypes: true });
return entries.filter((entry) => entry.isDirectory()).map((entry) => join(this.sessionsRoot, entry.name));
}
}

View File

@@ -2,7 +2,7 @@ import { createReadStream } from "node:fs";
import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises";
import { dirname, resolve } from "node:path";
import { createInterface } from "node:readline";
import type { JsonlSessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js";
import type { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
interface SessionHeader {
type: "session";
@@ -13,7 +13,25 @@ interface SessionHeader {
parentSession?: string;
}
function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionInfo {
function updateLabelCache(labelsById: Map<string, string>, entry: SessionTreeEntry): void {
if (entry.type !== "label") return;
const label = entry.label?.trim();
if (label) {
labelsById.set(entry.targetId, label);
} else {
labelsById.delete(entry.targetId);
}
}
function buildLabelsById(entries: SessionTreeEntry[]): Map<string, string> {
const labelsById = new Map<string, string>();
for (const entry of entries) {
updateLabelCache(labelsById, entry);
}
return labelsById;
}
function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata {
return {
id: header.id,
createdAt: header.timestamp,
@@ -23,7 +41,7 @@ function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionI
};
}
export async function loadJsonlSessionInfo(filePath: string): Promise<JsonlSessionInfo> {
export async function loadJsonlSessionMetadata(filePath: string): Promise<JsonlSessionMetadata> {
const stream = createReadStream(filePath, { encoding: "utf8" });
const lines = createInterface({ input: stream, crlfDelay: Infinity });
try {
@@ -31,7 +49,7 @@ export async function loadJsonlSessionInfo(filePath: string): Promise<JsonlSessi
if (!line.trim()) break;
try {
const header = JSON.parse(line) as SessionHeader;
return headerToSessionInfo(header, resolve(filePath));
return headerToSessionMetadata(header, resolve(filePath));
} catch {
throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`);
}
@@ -75,35 +93,27 @@ async function loadJsonlStorage(filePath: string): Promise<{
return { header, entries, leafId };
}
export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionInfo> {
export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata> {
private readonly filePath: string;
private readonly header: SessionHeader;
private readonly sessionInfo: JsonlSessionInfo;
private readonly metadata: JsonlSessionMetadata;
private entries: SessionTreeEntry[];
private byId: Map<string, SessionTreeEntry>;
private labelsById: Map<string, string>;
private currentLeafId: string | null;
private headerWritten: boolean;
private constructor(
filePath: string,
header: SessionHeader,
entries: SessionTreeEntry[],
leafId: string | null,
headerWritten: boolean,
) {
private constructor(filePath: string, header: SessionHeader, entries: SessionTreeEntry[], leafId: string | null) {
this.filePath = resolve(filePath);
this.header = header;
this.sessionInfo = headerToSessionInfo(header, this.filePath);
this.metadata = headerToSessionMetadata(header, this.filePath);
this.entries = entries;
this.byId = new Map(entries.map((entry) => [entry.id, entry]));
this.labelsById = buildLabelsById(entries);
this.currentLeafId = leafId;
this.headerWritten = headerWritten;
}
static async open(filePath: string): Promise<JsonlSessionTreeStorage> {
static async open(filePath: string): Promise<JsonlSessionStorage> {
const resolvedPath = resolve(filePath);
const loaded = await loadJsonlStorage(resolvedPath);
return new JsonlSessionTreeStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId, true);
return new JsonlSessionStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId);
}
static async create(
@@ -113,7 +123,7 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
sessionId: string;
parentSessionPath?: string;
},
): Promise<JsonlSessionTreeStorage> {
): Promise<JsonlSessionStorage> {
const resolvedPath = resolve(filePath);
const header: SessionHeader = {
type: "session",
@@ -123,11 +133,13 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
cwd: options.cwd,
parentSession: options.parentSessionPath,
};
return new JsonlSessionTreeStorage(resolvedPath, header, [], null, false);
await mkdir(dirname(resolvedPath), { recursive: true });
await writeFile(resolvedPath, `${JSON.stringify(header)}\n`);
return new JsonlSessionStorage(resolvedPath, header, [], null);
}
async getSessionInfo(): Promise<JsonlSessionInfo> {
return this.sessionInfo;
async getMetadata(): Promise<JsonlSessionMetadata> {
return this.metadata;
}
async getLeafId(): Promise<string | null> {
@@ -142,14 +154,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
}
async appendEntry(entry: SessionTreeEntry): Promise<void> {
if (!this.headerWritten) {
await mkdir(dirname(this.filePath), { recursive: true });
await writeFile(this.filePath, `${JSON.stringify(this.header)}\n`);
this.headerWritten = true;
}
await appendFile(this.filePath, `${JSON.stringify(entry)}\n`);
this.entries.push(entry);
this.byId.set(entry.id, entry);
updateLabelCache(this.labelsById, entry);
this.currentLeafId = entry.id;
}
@@ -157,6 +165,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
return this.byId.get(id);
}
async getLabel(id: string): Promise<string | undefined> {
return this.labelsById.get(id);
}
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
if (leafId === null) return [];
const path: SessionTreeEntry[] = [];

View File

@@ -0,0 +1,55 @@
import type { Session, SessionMetadata, SessionRepo } from "../types.js";
import { InMemorySessionStorage } from "./memory-session-storage.js";
import { createSessionId, createTimestamp, getPathEntriesToFork, toSession } from "./session-repo.js";
export class InMemorySessionRepo implements SessionRepo<SessionMetadata, { id?: string }, string, void> {
private sessions = new Map<string, Session<SessionMetadata>>();
async create(options: { id?: string } = {}): Promise<Session<SessionMetadata>> {
const info: SessionMetadata = {
id: options.id ?? createSessionId(),
createdAt: createTimestamp(),
};
const storage = new InMemorySessionStorage({ metadata: info });
const session = toSession(storage);
this.sessions.set(info.id, session);
return session;
}
async open(ref: string): Promise<Session<SessionMetadata>> {
const session = this.sessions.get(ref);
if (!session) {
throw new Error(`Session not found: ${ref}`);
}
return session;
}
async list(): Promise<SessionMetadata[]> {
return Promise.all([...this.sessions.values()].map((session) => session.getMetadata()));
}
async delete(ref: string): Promise<void> {
this.sessions.delete(ref);
}
async fork(
ref: string,
options: { entryId: string; position?: "before" | "at"; id?: string },
): Promise<Session<SessionMetadata>> {
const source = await this.open(ref);
const forkedEntries = await getPathEntriesToFork(
source.getStorage(),
options.entryId,
options.position ?? "before",
);
const info: SessionMetadata = {
id: options.id ?? createSessionId(),
createdAt: createTimestamp(),
};
const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null;
const storage = new InMemorySessionStorage({ metadata: info, entries: forkedEntries, leafId });
const session = toSession(storage);
this.sessions.set(info.id, session);
return session;
}
}

View File

@@ -1,24 +1,44 @@
import { v7 as uuidv7 } from "uuid";
import type { SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js";
import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
export class InMemorySessionTreeStorage implements SessionTreeStorage {
private readonly sessionInfo: SessionInfo;
function updateLabelCache(labelsById: Map<string, string>, entry: SessionTreeEntry): void {
if (entry.type !== "label") return;
const label = entry.label?.trim();
if (label) {
labelsById.set(entry.targetId, label);
} else {
labelsById.delete(entry.targetId);
}
}
function buildLabelsById(entries: SessionTreeEntry[]): Map<string, string> {
const labelsById = new Map<string, string>();
for (const entry of entries) {
updateLabelCache(labelsById, entry);
}
return labelsById;
}
export class InMemorySessionStorage implements SessionStorage {
private readonly metadata: SessionMetadata;
private entries: SessionTreeEntry[];
private byId: Map<string, SessionTreeEntry>;
private labelsById: Map<string, string>;
private leafId: string | null;
constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; sessionInfo?: SessionInfo }) {
constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; metadata?: SessionMetadata }) {
this.entries = options?.entries ? [...options.entries] : [];
this.byId = new Map(this.entries.map((entry) => [entry.id, entry]));
this.labelsById = buildLabelsById(this.entries);
this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null;
if (this.leafId !== null && !this.byId.has(this.leafId)) {
throw new Error(`Entry ${this.leafId} not found`);
}
this.sessionInfo = options?.sessionInfo ?? { id: uuidv7(), createdAt: new Date().toISOString() };
this.metadata = options?.metadata ?? { id: uuidv7(), createdAt: new Date().toISOString() };
}
async getSessionInfo(): Promise<SessionInfo> {
return this.sessionInfo;
async getMetadata(): Promise<SessionMetadata> {
return this.metadata;
}
async getLeafId(): Promise<string | null> {
@@ -35,6 +55,7 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage {
async appendEntry(entry: SessionTreeEntry): Promise<void> {
this.entries.push(entry);
this.byId.set(entry.id, entry);
updateLabelCache(this.labelsById, entry);
this.leafId = entry.id;
}
@@ -42,6 +63,10 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage {
return this.byId.get(id);
}
async getLabel(id: string): Promise<string | undefined> {
return this.labelsById.get(id);
}
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
if (leafId === null) return [];
const path: SessionTreeEntry[] = [];

View File

@@ -1,41 +1,25 @@
import { existsSync, mkdirSync, readdirSync, rmSync } from "node:fs";
import { join, resolve } from "node:path";
import { v7 as uuidv7 } from "uuid";
import type {
JsonlSessionInfo,
JsonlSessionRepo,
Session,
SessionInfo,
SessionRepo,
SessionTreeEntry,
SessionTreeStorage,
} from "../types.js";
import { JsonlSessionTreeStorage } from "./jsonl-session-storage.js";
import { InMemorySessionTreeStorage } from "./memory-session-storage.js";
import { DefaultSessionTree } from "./session-tree.js";
import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
import { DefaultSession } from "./session-tree.js";
function createSessionId(): string {
export function createSessionId(): string {
return uuidv7();
}
function createTimestamp(): string {
export function createTimestamp(): string {
return new Date().toISOString();
}
function toSession<TInfo extends SessionInfo>(
storage: SessionTreeStorage<TInfo>,
tree: DefaultSessionTree<TInfo>,
): Session<TInfo> {
return { storage, tree };
export function toSession<TMetadata extends SessionMetadata>(storage: SessionStorage<TMetadata>): Session<TMetadata> {
return new DefaultSession(storage);
}
function getPathEntriesToFork(
entries: SessionTreeEntry[],
export async function getPathEntriesToFork(
storage: SessionStorage,
entryId: string,
position: "before" | "at",
): SessionTreeEntry[] {
const byId = new Map<string, SessionTreeEntry>(entries.map((entry) => [entry.id, entry]));
const target = byId.get(entryId);
): Promise<SessionTreeEntry[]> {
const target = await storage.getEntry(entryId);
if (!target) {
throw new Error(`Entry ${entryId} not found`);
}
@@ -48,172 +32,5 @@ function getPathEntriesToFork(
}
effectiveLeafId = target.parentId;
}
if (effectiveLeafId === null) {
return [];
}
const path: SessionTreeEntry[] = [];
let current = byId.get(effectiveLeafId);
while (current) {
path.unshift(current);
current = current.parentId ? byId.get(current.parentId) : undefined;
}
return path;
}
export class InMemorySessionRepo implements SessionRepo<string> {
private sessions = new Map<string, Session<SessionInfo>>();
async create(options?: { id?: string }): Promise<Session<SessionInfo>> {
const info: SessionInfo = {
id: options?.id ?? createSessionId(),
createdAt: createTimestamp(),
};
const storage = new InMemorySessionTreeStorage({ sessionInfo: info });
const session = toSession(storage, new DefaultSessionTree(storage));
this.sessions.set(info.id, session);
return session;
}
async open(ref: string): Promise<Session<SessionInfo>> {
const session = this.sessions.get(ref);
if (!session) {
throw new Error(`Session not found: ${ref}`);
}
return session;
}
async list(): Promise<Array<Session<SessionInfo>>> {
return [...this.sessions.values()];
}
async delete(ref: string): Promise<void> {
this.sessions.delete(ref);
}
async fork(
ref: string,
options: { entryId: string; position?: "before" | "at"; id?: string },
): Promise<Session<SessionInfo>> {
const source = await this.open(ref);
const entries = await source.tree.getEntries();
const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before");
const info: SessionInfo = {
id: options.id ?? createSessionId(),
createdAt: createTimestamp(),
};
const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null;
const storage = new InMemorySessionTreeStorage({ sessionInfo: info, entries: forkedEntries, leafId });
const session = toSession(storage, new DefaultSessionTree(storage));
this.sessions.set(info.id, session);
return session;
}
}
export class JsonlSessionFileRepo implements JsonlSessionRepo<string> {
private sessionDir: string;
private cwd: string;
constructor(options: { sessionDir: string; cwd: string }) {
this.sessionDir = resolve(options.sessionDir);
this.cwd = options.cwd;
mkdirSync(this.sessionDir, { recursive: true });
}
private createSessionFilePath(sessionId: string, timestamp: string): string {
return join(this.sessionDir, `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`);
}
async create(options?: { id?: string; parentSessionPath?: string }): Promise<Session<JsonlSessionInfo>> {
const id = options?.id ?? createSessionId();
const createdAt = createTimestamp();
const filePath = this.createSessionFilePath(id, createdAt);
const storage = await JsonlSessionTreeStorage.create(filePath, {
cwd: this.cwd,
sessionId: id,
parentSessionPath: options?.parentSessionPath,
});
return toSession(storage, new DefaultSessionTree(storage));
}
async open(ref: string): Promise<Session<JsonlSessionInfo>> {
const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref);
if (!existsSync(filePath)) {
throw new Error(`Session not found: ${ref}`);
}
const storage = await JsonlSessionTreeStorage.open(filePath);
return toSession(storage, new DefaultSessionTree(storage));
}
async list(): Promise<Array<Session<JsonlSessionInfo>>> {
if (!existsSync(this.sessionDir)) {
return [];
}
const files = readdirSync(this.sessionDir)
.filter((file) => file.endsWith(".jsonl"))
.map((file) => join(this.sessionDir, file));
const sessions: Array<Session<JsonlSessionInfo>> = [];
for (const filePath of files) {
try {
const storage = await JsonlSessionTreeStorage.open(filePath);
sessions.push(toSession(storage, new DefaultSessionTree(storage)));
} catch {
// Ignore invalid session files when listing a directory.
}
}
return sessions;
}
async listByCwd(cwd: string): Promise<Array<Session<JsonlSessionInfo>>> {
const sessions = await this.list();
const result: Array<Session<JsonlSessionInfo>> = [];
for (const session of sessions) {
if ((await session.storage.getSessionInfo()).cwd === cwd) {
result.push(session);
}
}
return result;
}
async getMostRecentByCwd(cwd: string): Promise<Session<JsonlSessionInfo> | undefined> {
const sessionsWithInfo = await Promise.all(
(await this.listByCwd(cwd)).map(async (session) => ({
session,
info: await session.storage.getSessionInfo(),
})),
);
sessionsWithInfo.sort((a, b) => new Date(b.info.createdAt).getTime() - new Date(a.info.createdAt).getTime());
return sessionsWithInfo[0]?.session;
}
async delete(ref: string): Promise<void> {
const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref);
if (existsSync(filePath)) {
rmSync(filePath, { force: true });
}
}
async fork(
ref: string,
options: { entryId: string; position?: "before" | "at"; id?: string },
): Promise<Session<JsonlSessionInfo>> {
const source = await this.open(ref);
const entries = await source.tree.getEntries();
const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before");
const id = options.id ?? createSessionId();
const createdAt = createTimestamp();
const filePath = this.createSessionFilePath(id, createdAt);
const sourceInfo = await source.storage.getSessionInfo();
const storage = await JsonlSessionTreeStorage.create(filePath, {
cwd: sourceInfo.cwd,
sessionId: id,
parentSessionPath: sourceInfo.path,
});
for (const entry of forkedEntries) {
await storage.appendEntry(entry);
}
if (forkedEntries.length === 0) {
await storage.getSessionInfo();
}
return toSession(storage, new DefaultSessionTree(storage));
}
return storage.getPathToRoot(effectiveLeafId);
}

View File

@@ -10,12 +10,12 @@ import type {
LabelEntry,
MessageEntry,
ModelChangeEntry,
Session,
SessionContext,
SessionInfo,
SessionInfoEntry,
SessionTree,
SessionMetadata,
SessionStorage,
SessionTreeEntry,
SessionTreeStorage,
ThinkingLevelChangeEntry,
} from "../types.js";
@@ -27,12 +27,12 @@ function generateId(byId: { has(id: string): boolean }): string {
return randomUUID();
}
export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext {
export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext {
let thinkingLevel = "off";
let model: { provider: string; modelId: string } | null = null;
let compaction: CompactionEntry | null = null;
for (const entry of entries) {
for (const entry of pathEntries) {
if (entry.type === "thinking_level_change") {
thinkingLevel = entry.thinkingLevel;
} else if (entry.type === "model_change") {
@@ -47,10 +47,16 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
const messages: AgentMessage[] = [];
const appendMessage = (entry: SessionTreeEntry) => {
if (entry.type === "message") {
messages.push(entry.message);
messages.push(entry.message as AgentMessage);
} else if (entry.type === "custom_message") {
messages.push(
createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp),
createCustomMessage(
entry.customType,
entry.content as string | (TextContent | ImageContent)[],
entry.display,
entry.details,
entry.timestamp,
),
);
} else if (entry.type === "branch_summary" && entry.summary) {
messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
@@ -59,18 +65,18 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
if (compaction) {
messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp));
const compactionIdx = entries.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
const compactionIdx = pathEntries.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
let foundFirstKept = false;
for (let i = 0; i < compactionIdx; i++) {
const entry = entries[i]!;
const entry = pathEntries[i]!;
if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
if (foundFirstKept) appendMessage(entry);
}
for (let i = compactionIdx + 1; i < entries.length; i++) {
appendMessage(entries[i]!);
for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
appendMessage(pathEntries[i]!);
}
} else {
for (const entry of entries) {
for (const entry of pathEntries) {
appendMessage(entry);
}
}
@@ -78,13 +84,21 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
return { messages, thinkingLevel, model };
}
export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> implements SessionTree {
private storage: SessionTreeStorage<TInfo>;
export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata> implements Session<TMetadata> {
private storage: SessionStorage<TMetadata>;
constructor(storage: SessionTreeStorage<TInfo>) {
constructor(storage: SessionStorage<TMetadata>) {
this.storage = storage;
}
getMetadata(): Promise<TMetadata> {
return this.storage.getMetadata();
}
getStorage(): SessionStorage<TMetadata> {
return this.storage;
}
getLeafId(): Promise<string | null> {
return this.storage.getLeafId();
}
@@ -106,19 +120,8 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
return buildSessionContext(await this.getBranch());
}
getSessionInfo(): Promise<TInfo> {
return this.storage.getSessionInfo();
}
async getLabel(id: string): Promise<string | undefined> {
const entries = await this.storage.getEntries();
for (let i = entries.length - 1; i >= 0; i--) {
const entry = entries[i]!;
if (entry.type === "label" && entry.targetId === id) {
return entry.label?.trim() || undefined;
}
}
return undefined;
getLabel(id: string): Promise<string | undefined> {
return this.storage.getLabel(id);
}
async getSessionName(): Promise<string | undefined> {
@@ -193,24 +196,6 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
} satisfies CompactionEntry<T>);
}
async appendBranchSummary<T = unknown>(
fromId: string,
summary: string,
details?: T,
fromHook?: boolean,
): Promise<string> {
return this.appendTypedEntry({
type: "branch_summary",
id: await this.makeEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
fromId,
summary,
details,
fromHook,
} satisfies BranchSummaryEntry<T>);
}
async appendCustomEntry(customType: string, data?: unknown): Promise<string> {
return this.appendTypedEntry({
type: "custom",
@@ -240,7 +225,10 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
} satisfies CustomMessageEntry<T>);
}
async appendLabelChange(targetId: string, label: string | undefined): Promise<string> {
async appendLabel(targetId: string, label: string | undefined): Promise<string> {
if (!(await this.storage.getEntry(targetId))) {
throw new Error(`Entry ${targetId} not found`);
}
return this.appendTypedEntry({
type: "label",
id: await this.makeEntryId(),
@@ -251,7 +239,7 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
} satisfies LabelEntry);
}
async appendSessionInfo(name: string): Promise<string> {
async appendSessionName(name: string): Promise<string> {
return this.appendTypedEntry({
type: "session_info",
id: await this.makeEntryId(),
@@ -261,10 +249,24 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
} satisfies SessionInfoEntry);
}
async moveTo(entryId: string | null): Promise<void> {
async moveTo(
entryId: string | null,
summary?: { summary: string; details?: unknown; fromHook?: boolean },
): Promise<string | undefined> {
if (entryId !== null && !(await this.storage.getEntry(entryId))) {
throw new Error(`Entry ${entryId} not found`);
}
await this.storage.setLeafId(entryId);
if (!summary) return undefined;
return this.appendTypedEntry({
type: "branch_summary",
id: await this.makeEntryId(),
parentId: entryId,
timestamp: new Date().toISOString(),
fromId: entryId ?? "root",
summary: summary.summary,
details: summary.details,
fromHook: summary.fromHook,
} satisfies BranchSummaryEntry);
}
}

View File

@@ -136,7 +136,7 @@ export interface LabelEntry extends SessionTreeEntryBase {
}
export interface SessionInfoEntry extends SessionTreeEntryBase {
type: "session_info";
type: "session_info"; // legacy name, kept for backwards compatibility
name?: string;
}
@@ -157,28 +157,32 @@ export interface SessionContext {
model: { provider: string; modelId: string } | null;
}
export interface SessionInfo {
export interface SessionMetadata {
id: string;
createdAt: string;
}
export interface JsonlSessionInfo extends SessionInfo {
export interface JsonlSessionMetadata extends SessionMetadata {
cwd: string;
path: string;
parentSessionPath?: string;
}
export interface SessionTreeStorage<TInfo extends SessionInfo = SessionInfo> {
getSessionInfo(): Promise<TInfo>;
export interface SessionStorage<TMetadata extends SessionMetadata = SessionMetadata> {
getMetadata(): Promise<TMetadata>;
getLeafId(): Promise<string | null>;
setLeafId(leafId: string | null): Promise<void>;
appendEntry(entry: SessionTreeEntry): Promise<void>;
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
getLabel(id: string): Promise<string | undefined>;
getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]>;
getEntries(): Promise<SessionTreeEntry[]>;
}
export interface SessionTree {
export interface Session<TMetadata extends SessionMetadata = SessionMetadata> {
getMetadata(): Promise<TMetadata>;
getStorage(): SessionStorage<TMetadata>;
getLeafId(): Promise<string | null>;
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
getEntries(): Promise<SessionTreeEntry[]>;
@@ -197,7 +201,6 @@ export interface SessionTree {
details?: T,
fromHook?: boolean,
): Promise<string>;
appendBranchSummary<T = unknown>(fromId: string, summary: string, details?: T, fromHook?: boolean): Promise<string>;
appendCustomEntry(customType: string, data?: unknown): Promise<string>;
appendCustomMessageEntry<T = unknown>(
customType: string,
@@ -205,15 +208,13 @@ export interface SessionTree {
display: boolean,
details?: T,
): Promise<string>;
appendLabelChange(targetId: string, label: string | undefined): Promise<string>;
appendSessionInfo(name: string): Promise<string>;
appendLabel(targetId: string, label: string | undefined): Promise<string>;
appendSessionName(name: string): Promise<string>;
moveTo(entryId: string | null): Promise<void>;
}
export interface Session<TInfo extends SessionInfo = SessionInfo> {
storage: SessionTreeStorage<TInfo>;
tree: SessionTree;
moveTo(
entryId: string | null,
summary?: { summary: string; details?: unknown; fromHook?: boolean },
): Promise<string | undefined>;
}
export interface SessionCreateOptions {
@@ -227,25 +228,38 @@ export interface SessionForkOptions {
}
export interface SessionRepo<
TRef = string,
TInfo extends SessionInfo = SessionInfo,
TMetadata extends SessionMetadata = SessionMetadata,
TCreateOptions extends SessionCreateOptions = SessionCreateOptions,
TRef = string,
TListQuery = void,
> {
create(options?: TCreateOptions): Promise<Session<TInfo>>;
open(ref: TRef): Promise<Session<TInfo>>;
list(): Promise<Array<Session<TInfo>>>;
create(options: TCreateOptions): Promise<Session<TMetadata>>;
open(ref: TRef): Promise<Session<TMetadata>>;
list(query?: TListQuery): Promise<TMetadata[]>;
delete(ref: TRef): Promise<void>;
fork(ref: TRef, options: SessionForkOptions): Promise<Session<TInfo>>;
fork(ref: TRef, options: SessionForkOptions & TCreateOptions): Promise<Session<TMetadata>>;
}
export interface JsonlSessionCreateOptions extends SessionCreateOptions {
cwd: string;
parentSessionPath?: string;
}
export interface JsonlSessionRepo<TRef = string>
extends SessionRepo<TRef, JsonlSessionInfo, JsonlSessionCreateOptions> {
listByCwd(cwd: string): Promise<Array<Session<JsonlSessionInfo>>>;
getMostRecentByCwd(cwd: string): Promise<Session<JsonlSessionInfo> | undefined>;
export type JsonlSessionRef = { path: string } | JsonlSessionMetadata;
export interface JsonlSessionListQuery {
cwd?: string;
}
export interface JsonlSessionResolveOptions {
cwd?: string;
searchAll?: boolean;
}
export interface JsonlSessionRepoApi
extends SessionRepo<JsonlSessionMetadata, JsonlSessionCreateOptions, JsonlSessionRef, JsonlSessionListQuery> {
resolve(ref: string, options?: JsonlSessionResolveOptions): Promise<JsonlSessionMetadata[]>;
getMostRecent(query?: JsonlSessionListQuery): Promise<JsonlSessionMetadata | undefined>;
}
export interface AgentHarnessPendingMutations {
@@ -257,7 +271,7 @@ export interface AgentHarnessPendingMutations {
}
export interface AgentHarnessConversationState {
sessionTree: SessionTree;
session: Session;
model: Model<any> | undefined;
thinkingLevel: ThinkingLevel;
activeToolNames: string[];
@@ -284,8 +298,8 @@ export interface SavePointSnapshot {
export interface AgentHarnessContext {
env: ExecutionEnv;
conversation: Readonly<AgentHarnessConversationState>;
operation: Readonly<AgentHarnessOperationState>;
conversation: AgentHarnessConversationState;
operation: AgentHarnessOperationState;
abortSignal?: AbortSignal;
}
@@ -549,7 +563,7 @@ export interface BranchSummaryResult {
export interface AgentHarnessOptions {
agent: Agent;
env: ExecutionEnv;
sessionTree: SessionTree;
session: Session;
promptTemplates?: PromptTemplate[];
skills?: Skill[];
requestAuth?: (model: Model<any>) => Promise<{ apiKey: string; headers?: Record<string, string> } | undefined>;
@@ -562,8 +576,8 @@ export interface AgentHarnessOptions {
export interface AgentHarness {
readonly agent: Agent;
readonly env: ExecutionEnv;
readonly conversation: Readonly<AgentHarnessConversationState>;
readonly operation: Readonly<AgentHarnessOperationState>;
readonly conversation: AgentHarnessConversationState;
readonly operation: AgentHarnessOperationState;
prompt(text: string, options?: AgentHarnessPromptOptions): Promise<void>;
skill(name: string, args?: string): Promise<void>;

View File

@@ -25,6 +25,8 @@ export {
export * from "./harness/execution-env.js";
export * from "./harness/messages.js";
export * from "./harness/prompt-templates.js";
export * from "./harness/session/jsonl-session-repo.js";
export * from "./harness/session/memory-session-repo.js";
export * from "./harness/session/session-repo.js";
export * from "./harness/session/session-tree.js";
// Harness

View File

@@ -3,10 +3,12 @@ import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { afterEach, describe, expect, it } from "vitest";
import { JsonlSessionTreeStorage, loadJsonlSessionInfo } from "../../src/harness/session/jsonl-session-storage.js";
import { InMemorySessionTreeStorage } from "../../src/harness/session/memory-session-storage.js";
import { DefaultSessionTree } from "../../src/harness/session/session-tree.js";
import type { MessageEntry, SessionInfo, SessionTreeStorage } from "../../src/harness/types.js";
import { JsonlSessionRepo } from "../../src/harness/session/jsonl-session-repo.js";
import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-session-storage.js";
import { InMemorySessionRepo } from "../../src/harness/session/memory-session-repo.js";
import { InMemorySessionStorage } from "../../src/harness/session/memory-session-storage.js";
import { DefaultSession } from "../../src/harness/session/session-tree.js";
import type { MessageEntry, SessionMetadata, SessionStorage } from "../../src/harness/types.js";
function createUserMessage(text: string): AgentMessage {
return {
@@ -56,12 +58,12 @@ afterEach(() => {
async function runSessionTreeSuite(
name: string,
createStorage: () => SessionTreeStorage | Promise<SessionTreeStorage>,
createStorage: () => SessionStorage | Promise<SessionStorage>,
inspect?: () => void,
) {
describe(name, () => {
it("appends messages and builds context in order", async () => {
const tree = new DefaultSessionTree(await createStorage());
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendMessage(createAssistantMessage("two"));
const context = await tree.buildContext();
@@ -69,7 +71,7 @@ async function runSessionTreeSuite(
});
it("tracks model and thinking level changes", async () => {
const tree = new DefaultSessionTree(await createStorage());
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendModelChange("openai", "gpt-4.1");
await tree.appendThinkingLevelChange("high");
@@ -79,7 +81,7 @@ async function runSessionTreeSuite(
});
it("supports branching by moving the leaf and appending a new branch", async () => {
const tree = new DefaultSessionTree(await createStorage());
const tree = new DefaultSession(await createStorage());
const user1 = await tree.appendMessage(createUserMessage("one"));
const assistant1 = await tree.appendMessage(createAssistantMessage("two"));
await tree.appendMessage(createUserMessage("three"));
@@ -93,7 +95,7 @@ async function runSessionTreeSuite(
});
it("supports moving the leaf to root", async () => {
const tree = new DefaultSessionTree(await createStorage());
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.moveTo(null);
expect(await tree.getLeafId()).toBeNull();
@@ -101,7 +103,7 @@ async function runSessionTreeSuite(
});
it("reconstructs compaction summaries in context", async () => {
const tree = new DefaultSessionTree(await createStorage());
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendMessage(createAssistantMessage("two"));
const user2 = await tree.appendMessage(createUserMessage("three"));
@@ -113,16 +115,19 @@ async function runSessionTreeSuite(
expect(context.messages).toHaveLength(4);
});
it("supports branch summary entries in context", async () => {
const tree = new DefaultSessionTree(await createStorage());
it("supports moving with branch summary entries in context", async () => {
const tree = new DefaultSession(await createStorage());
const user1 = await tree.appendMessage(createUserMessage("one"));
await tree.appendBranchSummary(user1, "summary text");
const summaryId = await tree.moveTo(user1, { summary: "summary text" });
expect(summaryId).toBeTruthy();
const summaryEntry = await tree.getEntry(summaryId!);
expect(summaryEntry).toMatchObject({ type: "branch_summary", parentId: user1, fromId: user1 });
const context = await tree.buildContext();
expect(context.messages[1]?.role).toBe("branchSummary");
});
it("supports custom message entries in context", async () => {
const tree = new DefaultSessionTree(await createStorage());
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendCustomMessageEntry("custom", "hello", true, { ok: true });
const context = await tree.buildContext();
@@ -130,10 +135,10 @@ async function runSessionTreeSuite(
});
it("supports labels and session info entries without affecting context", async () => {
const tree = new DefaultSessionTree(await createStorage());
const tree = new DefaultSession(await createStorage());
const user1 = await tree.appendMessage(createUserMessage("one"));
await tree.appendLabelChange(user1, "checkpoint");
await tree.appendSessionInfo("name");
await tree.appendLabel(user1, "checkpoint");
await tree.appendSessionName("name");
const entries = await tree.getEntries();
expect(entries.some((entry) => entry.type === "label")).toBe(true);
expect(entries.some((entry) => entry.type === "session_info")).toBe(true);
@@ -142,16 +147,21 @@ async function runSessionTreeSuite(
expect((await tree.buildContext()).messages).toHaveLength(1);
});
it("rejects labels for missing entries", async () => {
const tree = new DefaultSession(await createStorage());
await expect(tree.appendLabel("missing", "checkpoint")).rejects.toThrow("Entry missing not found");
});
it("persists leaf changes and appended entries via storage", async () => {
const storage = await createStorage();
const tree = new DefaultSessionTree(storage);
const tree = new DefaultSession(storage);
const user1 = await tree.appendMessage(createUserMessage("one"));
await tree.appendMessage(createAssistantMessage("two"));
await tree.appendLabelChange(user1, "checkpoint");
await tree.appendSessionInfo("name");
await tree.appendLabel(user1, "checkpoint");
await tree.appendSessionName("name");
await tree.moveTo(user1);
await tree.appendMessage(createAssistantMessage("branched"));
const tree2 = new DefaultSessionTree(storage);
const tree2 = new DefaultSession(storage);
const context = await tree2.buildContext();
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
expect(await tree2.getLabel(user1)).toBe("checkpoint");
@@ -161,11 +171,11 @@ async function runSessionTreeSuite(
});
}
describe("InMemorySessionTreeStorage", () => {
describe("InMemorySessionStorage", () => {
it("returns configured session info", async () => {
const sessionInfo: SessionInfo = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" };
const storage = new InMemorySessionTreeStorage({ sessionInfo });
expect(await storage.getSessionInfo()).toEqual(sessionInfo);
const metadata: SessionMetadata = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" };
const storage = new InMemorySessionStorage({ metadata });
expect(await storage.getMetadata()).toEqual(metadata);
});
it("copies initial entries and tracks leaf independently", async () => {
@@ -177,7 +187,7 @@ describe("InMemorySessionTreeStorage", () => {
message: createUserMessage("one"),
};
const initialEntries = [entry];
const storage = new InMemorySessionTreeStorage({ entries: initialEntries });
const storage = new InMemorySessionStorage({ entries: initialEntries });
initialEntries.push({ ...entry, id: "entry-2" });
expect((await storage.getEntries()).map((storedEntry) => storedEntry.id)).toEqual(["entry-1"]);
expect(await storage.getLeafId()).toBe("entry-1");
@@ -186,9 +196,39 @@ describe("InMemorySessionTreeStorage", () => {
});
it("rejects invalid leaf ids", async () => {
const storage = new InMemorySessionTreeStorage();
const storage = new InMemorySessionStorage();
await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found");
expect(() => new InMemorySessionTreeStorage({ leafId: "missing" })).toThrow("Entry missing not found");
expect(() => new InMemorySessionStorage({ leafId: "missing" })).toThrow("Entry missing not found");
});
it("maintains label lookup", async () => {
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
const storage = new InMemorySessionStorage({ entries: [entry] });
expect(await storage.getLabel("entry-1")).toBeUndefined();
await storage.appendEntry({
type: "label",
id: "label-1",
parentId: "entry-1",
timestamp: "2026-01-01T00:00:01.000Z",
targetId: "entry-1",
label: "checkpoint",
});
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
await storage.appendEntry({
type: "label",
id: "label-2",
parentId: "label-1",
timestamp: "2026-01-01T00:00:02.000Z",
targetId: "entry-1",
label: undefined,
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
});
it("walks paths to root", async () => {
@@ -205,26 +245,70 @@ describe("InMemorySessionTreeStorage", () => {
parentId: "root",
message: createAssistantMessage("child"),
};
const storage = new InMemorySessionTreeStorage({ entries: [root, child] });
const storage = new InMemorySessionStorage({ entries: [root, child] });
expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
expect(await storage.getPathToRoot(null)).toEqual([]);
});
});
runSessionTreeSuite("SessionTree with in-memory storage", () => new InMemorySessionTreeStorage());
runSessionTreeSuite("Session with in-memory storage", () => new InMemorySessionStorage());
describe("JsonlSessionTreeStorage", () => {
describe("InMemorySessionRepo", () => {
it("lists session infos and forks via storage path traversal", async () => {
const repo = new InMemorySessionRepo();
const session = await repo.create({ id: "session-1" });
const user1 = await session.appendMessage(createUserMessage("one"));
const assistant1 = await session.appendMessage(createAssistantMessage("two"));
const user2 = await session.appendMessage(createUserMessage("three"));
const infos = await repo.list();
expect(infos.map((info) => info.id)).toEqual(["session-1"]);
const fork = await repo.fork("session-1", { entryId: user2, id: "session-2" });
expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1]);
});
});
describe("JsonlSessionRepo", () => {
it("stores sessions below encoded cwd directories and resolves prefixes", async () => {
const root = createTempDir();
const cwd = "/tmp/my-project";
const repo = new JsonlSessionRepo({ sessionsRoot: root });
const session = await repo.create({ cwd, id: "019de8c2-de29-73e9-ae0c-e134db34c447" });
const info = await session.getMetadata();
expect(info.path).toContain("--tmp-my-project--");
expect(existsSync(info.path)).toBe(true);
expect((await repo.list({ cwd })).map((sessionInfo) => sessionInfo.id)).toEqual([info.id]);
expect((await repo.resolve("019de8c2", { cwd })).map((sessionInfo) => sessionInfo.path)).toEqual([info.path]);
});
it("forks sessions and records the parent path", async () => {
const root = createTempDir();
const repo = new JsonlSessionRepo({ sessionsRoot: root });
const source = await repo.create({ cwd: "/tmp/source", id: "source-session" });
const sourceInfo = await source.getMetadata();
const user1 = await source.appendMessage(createUserMessage("one"));
const assistant1 = await source.appendMessage(createAssistantMessage("two"));
const user2 = await source.appendMessage(createUserMessage("three"));
const fork = await repo.fork(sourceInfo, { cwd: "/tmp/target", id: "fork-session", entryId: user2 });
const forkInfo = await fork.getMetadata();
expect(forkInfo.cwd).toBe("/tmp/target");
expect(forkInfo.parentSessionPath).toBe(sourceInfo.path);
expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1]);
});
});
describe("JsonlSessionStorage", () => {
it("throws for missing files when opening", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
await expect(JsonlSessionTreeStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" });
await expect(JsonlSessionStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" });
});
it("writes the header before the first appended entry", async () => {
it("writes the header on create", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionTreeStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
expect(existsSync(filePath)).toBe(false);
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
expect(existsSync(filePath)).toBe(true);
expect(readFileSync(filePath, "utf8").trim().split("\n")).toHaveLength(1);
expect(await storage.getLeafId()).toBeNull();
expect(await storage.getEntries()).toEqual([]);
await storage.appendEntry({
@@ -245,7 +329,7 @@ describe("JsonlSessionTreeStorage", () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
writeFileSync(filePath, "not json\n");
await expect(JsonlSessionTreeStorage.open(filePath)).rejects.toThrow("first line is not a valid session header");
await expect(JsonlSessionStorage.open(filePath)).rejects.toThrow("first line is not a valid session header");
});
it("ignores malformed entry lines", async () => {
@@ -266,7 +350,7 @@ describe("JsonlSessionTreeStorage", () => {
message: createUserMessage("one"),
};
writeFileSync(filePath, `${JSON.stringify(header)}\nnot json\n${JSON.stringify(entry)}\n`);
const storage = await JsonlSessionTreeStorage.open(filePath);
const storage = await JsonlSessionStorage.open(filePath);
expect((await storage.getEntries()).map((loadedEntry) => loadedEntry.id)).toEqual(["entry-1"]);
expect(await storage.getLeafId()).toBe("entry-1");
});
@@ -274,19 +358,19 @@ describe("JsonlSessionTreeStorage", () => {
it("creates and reads session info from the header", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionTreeStorage.create(filePath, {
const storage = await JsonlSessionStorage.create(filePath, {
cwd: dir,
sessionId: "session-1",
parentSessionPath: "/tmp/parent.jsonl",
});
const info = await storage.getSessionInfo();
const info = await storage.getMetadata();
expect(info).toMatchObject({
id: "session-1",
cwd: dir,
path: filePath,
parentSessionPath: "/tmp/parent.jsonl",
});
expect(existsSync(filePath)).toBe(false);
expect(existsSync(filePath)).toBe(true);
await storage.appendEntry({
type: "message",
id: "user-1",
@@ -294,13 +378,13 @@ describe("JsonlSessionTreeStorage", () => {
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect(await loadJsonlSessionInfo(filePath)).toEqual(info);
expect(await loadJsonlSessionMetadata(filePath)).toEqual(info);
});
it("loads existing entries and reconstructs leaf", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionTreeStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
const root: MessageEntry = {
type: "message",
id: "root",
@@ -316,12 +400,46 @@ describe("JsonlSessionTreeStorage", () => {
};
await storage.appendEntry(root);
await storage.appendEntry(child);
const loaded = await JsonlSessionTreeStorage.open(filePath);
const loaded = await JsonlSessionStorage.open(filePath);
expect(await loaded.getLeafId()).toBe("child");
expect((await loaded.getEntries()).map((entry) => entry.id)).toEqual(["root", "child"]);
expect((await loaded.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
});
it("maintains label lookup", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
await storage.appendEntry({
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
await storage.appendEntry({
type: "label",
id: "label-1",
parentId: "entry-1",
timestamp: "2026-01-01T00:00:01.000Z",
targetId: "entry-1",
label: "checkpoint",
});
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
await storage.appendEntry({
type: "label",
id: "label-2",
parentId: "label-1",
timestamp: "2026-01-01T00:00:02.000Z",
targetId: "entry-1",
label: undefined,
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
const loaded = await JsonlSessionStorage.open(filePath);
expect(await loaded.getLabel("entry-1")).toBeUndefined();
});
it("reads session info from only the first JSONL line", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
@@ -334,7 +452,7 @@ describe("JsonlSessionTreeStorage", () => {
};
const malformedSecondLine = "{".repeat(10000);
writeFileSync(filePath, `${JSON.stringify(header)}\n${malformedSecondLine}\n`);
expect(await loadJsonlSessionInfo(filePath)).toEqual({
expect(await loadJsonlSessionMetadata(filePath)).toEqual({
id: "session-1",
createdAt: "2026-01-01T00:00:00.000Z",
cwd: dir,
@@ -345,10 +463,10 @@ describe("JsonlSessionTreeStorage", () => {
});
runSessionTreeSuite(
"SessionTree with JSONL storage",
"Session with JSONL storage",
async () => {
const dir = createTempDir();
return await JsonlSessionTreeStorage.create(join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" });
return await JsonlSessionStorage.create(join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" });
},
() => {
const dir = tempDirs[tempDirs.length - 1]!;