refactor(agent): consolidate harness session abstraction
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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
|
||||
|
||||
142
packages/agent/src/harness/session/jsonl-session-repo.ts
Normal file
142
packages/agent/src/harness/session/jsonl-session-repo.ts
Normal 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));
|
||||
}
|
||||
}
|
||||
@@ -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[] = [];
|
||||
|
||||
55
packages/agent/src/harness/session/memory-session-repo.ts
Normal file
55
packages/agent/src/harness/session/memory-session-repo.ts
Normal 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;
|
||||
}
|
||||
}
|
||||
@@ -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[] = [];
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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>;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]!;
|
||||
|
||||
Reference in New Issue
Block a user