shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(agent): simplify harness session repo layout

Browse Source
This commit is contained in:
Mario Zechner
2026-05-04 00:06:35 +02:00
parent cdde2e893d
commit 3d0f57187d
15 changed files with 616 additions and 581 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
packages/agent/src/harness/compaction/compaction.ts
View File

@@ -14,7 +14,7 @@ import {
createCompactionSummaryMessage,
createCustomMessage,
} from "../messages.js";
import { buildSessionContext } from "../session/session-tree.js";
import { buildSessionContext } from "../session/session.js";
import type { CompactionEntry, SessionTreeEntry } from "../types.js";
import {
computeFileLists,

55
packages/agent/src/harness/session/memory-session-repo.ts
View File

@@ -1,55 +0,0 @@
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;
}
}

67
packages/agent/src/harness/session/jsonl-session-repo.ts → packages/agent/src/harness/session/repo/jsonl.ts
View File

@@ -3,15 +3,13 @@ import { access, mkdir, readdir, rm } from "node:fs/promises";
import { join, resolve } from "node:path";
import type {
JsonlSessionCreateOptions,
JsonlSessionListQuery,
JsonlSessionListOptions,
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";
} from "../../types.js";
import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../storage/jsonl.js";
import { createSessionId, createTimestamp, getEntriesToFork, toSession } from "./shared.js";
async function exists(path: string): Promise<boolean> {
try {
@@ -41,10 +39,6 @@ export class JsonlSessionRepo implements JsonlSessionRepoApi {
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();
@@ -58,17 +52,16 @@ export class JsonlSessionRepo implements JsonlSessionRepoApi {
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}`);
async open(metadata: JsonlSessionMetadata): Promise<Session<JsonlSessionMetadata>> {
if (!(await exists(metadata.path))) {
throw new Error(`Session not found: ${metadata.path}`);
}
const storage = await JsonlSessionStorage.open(filePath);
const storage = await JsonlSessionStorage.open(metadata.path);
return toSession(storage);
}
async list(query: JsonlSessionListQuery = {}): Promise<JsonlSessionMetadata[]> {
const dirs = query.cwd ? [this.getSessionDir(query.cwd)] : await this.listSessionDirs();
async list(options: JsonlSessionListOptions = {}): Promise<JsonlSessionMetadata[]> {
const dirs = options.cwd ? [this.getSessionDir(options.cwd)] : await this.listSessionDirs();
const sessions: JsonlSessionMetadata[] = [];
for (const dir of dirs) {
if (!(await exists(dir))) continue;
@@ -85,48 +78,22 @@ export class JsonlSessionRepo implements JsonlSessionRepoApi {
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 delete(metadata: JsonlSessionMetadata): Promise<void> {
await rm(metadata.path, { force: true });
}
async fork(
ref: JsonlSessionRef,
options: JsonlSessionCreateOptions & { entryId: string; position?: "before" | "at"; id?: string },
sourceMetadata: JsonlSessionMetadata,
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 source = await this.open(sourceMetadata);
const forkedEntries = await getEntriesToFork(source.getStorage(), options);
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,
parentSessionPath: options.parentSessionPath ?? sourceMetadata.path,
});
for (const entry of forkedEntries) {
await storage.appendEntry(entry);

51
packages/agent/src/harness/session/repo/memory.ts Normal file
View File

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

18
packages/agent/src/harness/session/session-repo.ts → packages/agent/src/harness/session/repo/shared.ts
View File

@@ -1,6 +1,6 @@
import { v7 as uuidv7 } from "uuid";
import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
import { DefaultSession } from "./session-tree.js";
import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js";
import { DefaultSession } from "../session.js";
export function createSessionId(): string {
return uuidv7();
@@ -14,21 +14,21 @@ export function toSession<TMetadata extends SessionMetadata>(storage: SessionSto
return new DefaultSession(storage);
}
export async function getPathEntriesToFork(
export async function getEntriesToFork(
storage: SessionStorage,
entryId: string,
position: "before" | "at",
options: { entryId?: string; position?: "before" | "at" },
): Promise<SessionTreeEntry[]> {
const target = await storage.getEntry(entryId);
if (!options.entryId) return storage.getEntries();
const target = await storage.getEntry(options.entryId);
if (!target) {
throw new Error(`Entry ${entryId} not found`);
throw new Error(`Entry ${options.entryId} not found`);
}
let effectiveLeafId: string | null;
if (position === "at") {
if ((options.position ?? "before") === "at") {
effectiveLeafId = target.id;
} else {
if (target.type !== "message" || target.message.role !== "user") {
throw new Error(`Entry ${entryId} is not a user message`);
throw new Error(`Entry ${options.entryId} is not a user message`);
}
effectiveLeafId = target.parentId;
}

42
packages/agent/src/harness/session/session-tree.ts → packages/agent/src/harness/session/session.ts
View File

@@ -1,4 +1,3 @@
import { randomUUID } from "node:crypto";
import type { ImageContent, TextContent } from "@mariozechner/pi-ai";
import type { AgentMessage } from "../../types.js";
import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.js";
@@ -19,14 +18,6 @@ import type {
ThinkingLevelChangeEntry,
} from "../types.js";
function generateId(byId: { has(id: string): boolean }): string {
for (let i = 0; i < 100; i++) {
const id = randomUUID().slice(0, 8);
if (!byId.has(id)) return id;
}
return randomUUID();
}
export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext {
let thinkingLevel = "off";
let model: { provider: string; modelId: string } | null = null;
@@ -125,19 +116,8 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
}
async getSessionName(): 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 === "session_info") {
return entry.name?.trim() || undefined;
}
}
return undefined;
}
private async makeEntryId(): Promise<string> {
const entries = await this.storage.getEntries();
return generateId(new Set(entries.map((entry) => entry.id)));
const entries = await this.storage.findEntries("session_info");
return entries[entries.length - 1]?.name?.trim() || undefined;
}
private async appendTypedEntry<TEntry extends SessionTreeEntry>(entry: TEntry): Promise<string> {
@@ -148,7 +128,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
async appendMessage(message: AgentMessage): Promise<string> {
return this.appendTypedEntry({
type: "message",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
message,
@@ -158,7 +138,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
async appendThinkingLevelChange(thinkingLevel: string): Promise<string> {
return this.appendTypedEntry({
type: "thinking_level_change",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
thinkingLevel,
@@ -168,7 +148,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
async appendModelChange(provider: string, modelId: string): Promise<string> {
return this.appendTypedEntry({
type: "model_change",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
provider,
@@ -185,7 +165,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
): Promise<string> {
return this.appendTypedEntry({
type: "compaction",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
summary,
@@ -199,7 +179,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
async appendCustomEntry(customType: string, data?: unknown): Promise<string> {
return this.appendTypedEntry({
type: "custom",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
customType,
@@ -215,7 +195,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
): Promise<string> {
return this.appendTypedEntry({
type: "custom_message",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
customType,
@@ -231,7 +211,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
}
return this.appendTypedEntry({
type: "label",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
targetId,
@@ -242,7 +222,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
async appendSessionName(name: string): Promise<string> {
return this.appendTypedEntry({
type: "session_info",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: await this.storage.getLeafId(),
timestamp: new Date().toISOString(),
name: name.trim(),
@@ -260,7 +240,7 @@ export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata>
if (!summary) return undefined;
return this.appendTypedEntry({
type: "branch_summary",
id: await this.makeEntryId(),
id: await this.storage.createEntryId(),
parentId: entryId,
timestamp: new Date().toISOString(),
fromId: entryId ?? "root",

21
packages/agent/src/harness/session/jsonl-session-storage.ts → packages/agent/src/harness/session/storage/jsonl.ts
View File

@@ -1,8 +1,9 @@
import { randomUUID } from "node:crypto";
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 { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
import type { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js";
interface SessionHeader {
type: "session";
@@ -31,6 +32,14 @@ function buildLabelsById(entries: SessionTreeEntry[]): Map<string, string> {
return labelsById;
}
function generateEntryId(byId: { has(id: string): boolean }): string {
for (let i = 0; i < 100; i++) {
const id = randomUUID().slice(0, 8);
if (!byId.has(id)) return id;
}
return randomUUID();
}
function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata {
return {
id: header.id,
@@ -153,6 +162,10 @@ export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata>
this.currentLeafId = leafId;
}
async createEntryId(): Promise<string> {
return generateEntryId(this.byId);
}
async appendEntry(entry: SessionTreeEntry): Promise<void> {
await appendFile(this.filePath, `${JSON.stringify(entry)}\n`);
this.entries.push(entry);
@@ -165,6 +178,12 @@ export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata>
return this.byId.get(id);
}
async findEntries<TType extends SessionTreeEntry["type"]>(
type: TType,
): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>> {
return this.entries.filter((entry): entry is Extract<SessionTreeEntry, { type: TType }> => entry.type === type);
}
async getLabel(id: string): Promise<string | undefined> {
return this.labelsById.get(id);
}

21
packages/agent/src/harness/session/memory-session-storage.ts → packages/agent/src/harness/session/storage/memory.ts
View File

@@ -1,5 +1,6 @@
import { randomUUID } from "node:crypto";
import { v7 as uuidv7 } from "uuid";
import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js";
function updateLabelCache(labelsById: Map<string, string>, entry: SessionTreeEntry): void {
if (entry.type !== "label") return;
@@ -19,6 +20,14 @@ function buildLabelsById(entries: SessionTreeEntry[]): Map<string, string> {
return labelsById;
}
function generateEntryId(byId: { has(id: string): boolean }): string {
for (let i = 0; i < 100; i++) {
const id = randomUUID().slice(0, 8);
if (!byId.has(id)) return id;
}
return randomUUID();
}
export class InMemorySessionStorage implements SessionStorage {
private readonly metadata: SessionMetadata;
private entries: SessionTreeEntry[];
@@ -52,6 +61,10 @@ export class InMemorySessionStorage implements SessionStorage {
this.leafId = leafId;
}
async createEntryId(): Promise<string> {
return generateEntryId(this.byId);
}
async appendEntry(entry: SessionTreeEntry): Promise<void> {
this.entries.push(entry);
this.byId.set(entry.id, entry);
@@ -63,6 +76,12 @@ export class InMemorySessionStorage implements SessionStorage {
return this.byId.get(id);
}
async findEntries<TType extends SessionTreeEntry["type"]>(
type: TType,
): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>> {
return this.entries.filter((entry): entry is Extract<SessionTreeEntry, { type: TType }> => entry.type === type);
}
async getLabel(id: string): Promise<string | undefined> {
return this.labelsById.get(id);
}

31
packages/agent/src/harness/types.ts
View File

@@ -172,8 +172,12 @@ export interface SessionStorage<TMetadata extends SessionMetadata = SessionMetad
getMetadata(): Promise<TMetadata>;
getLeafId(): Promise<string | null>;
setLeafId(leafId: string | null): Promise<void>;
createEntryId(): Promise<string>;
appendEntry(entry: SessionTreeEntry): Promise<void>;
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
findEntries<TType extends SessionTreeEntry["type"]>(
type: TType,
): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>>;
getLabel(id: string): Promise<string | undefined>;
getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]>;
getEntries(): Promise<SessionTreeEntry[]>;
@@ -222,7 +226,7 @@ export interface SessionCreateOptions {
}
export interface SessionForkOptions {
entryId: string;
entryId?: string;
position?: "before" | "at";
id?: string;
}
@@ -230,14 +234,13 @@ export interface SessionForkOptions {
export interface SessionRepo<
TMetadata extends SessionMetadata = SessionMetadata,
TCreateOptions extends SessionCreateOptions = SessionCreateOptions,
TRef = string,
TListQuery = void,
TListOptions = void,
> {
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 & TCreateOptions): Promise<Session<TMetadata>>;
open(metadata: TMetadata): Promise<Session<TMetadata>>;
list(options?: TListOptions): Promise<TMetadata[]>;
delete(metadata: TMetadata): Promise<void>;
fork(source: TMetadata, options: SessionForkOptions & TCreateOptions): Promise<Session<TMetadata>>;
}
export interface JsonlSessionCreateOptions extends SessionCreateOptions {
@@ -245,22 +248,12 @@ export interface JsonlSessionCreateOptions extends SessionCreateOptions {
parentSessionPath?: string;
}
export type JsonlSessionRef = { path: string } | JsonlSessionMetadata;
export interface JsonlSessionListQuery {
export interface JsonlSessionListOptions {
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>;
}
extends SessionRepo<JsonlSessionMetadata, JsonlSessionCreateOptions, JsonlSessionListOptions> {}
export interface AgentHarnessPendingMutations {
appendMessages: AgentMessage[];

8
packages/agent/src/index.ts
View File

@@ -25,10 +25,10 @@ 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";
export * from "./harness/session/repo/jsonl.js";
export * from "./harness/session/repo/memory.js";
export * from "./harness/session/repo/shared.js";
export * from "./harness/session/session.js";
// Harness
export * from "./harness/types.js";
export * from "./harness/utils/shell-output.js";