world-model-optimizer 0.2.0__py3-none-any.whl

wmo/harness/vendor/pi-agent/src/harness/session/jsonl-storage.ts

@@ -0,0 +1,293 @@
+import type { FileSystem, JsonlSessionMetadata, LeafEntry, SessionStorage, SessionTreeEntry } from "../types.ts";
+import { SessionError, toError } from "../types.ts";
+import { getFileSystemResultOrThrow } from "./repo-utils.ts";
+import { uuidv7 } from "./uuid.ts";
+
+type JsonlSessionStorageFileSystem = Pick<FileSystem, "readTextFile" | "readTextLines" | "writeFile" | "appendFile">;
+
+interface SessionHeader {
+	type: "session";
+	version: 3;
+	id: string;
+	timestamp: string;
+	cwd: string;
+	parentSession?: string;
+}
+
+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 generateEntryId(byId: { has(id: string): boolean }): string {
+	for (let i = 0; i < 100; i++) {
+		const id = uuidv7().slice(0, 8);
+		if (!byId.has(id)) return id;
+	}
+	return uuidv7();
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null;
+}
+
+function invalidSession(filePath: string, message: string, cause?: Error): SessionError {
+	return new SessionError("invalid_session", `Invalid JSONL session file ${filePath}: ${message}`, cause);
+}
+
+function invalidEntry(filePath: string, lineNumber: number, message: string, cause?: Error): SessionError {
+	return new SessionError(
+		"invalid_entry",
+		`Invalid JSONL session file ${filePath}: line ${lineNumber} ${message}`,
+		cause,
+	);
+}
+
+function parseHeaderLine(line: string, filePath: string): SessionHeader {
+	let parsed: unknown;
+	try {
+		parsed = JSON.parse(line);
+	} catch (error) {
+		throw invalidSession(filePath, "first line is not a valid session header", toError(error));
+	}
+	if (!isRecord(parsed)) throw invalidSession(filePath, "first line is not a valid session header");
+	if (parsed.type !== "session") throw invalidSession(filePath, "first line is not a valid session header");
+	if (parsed.version !== 3) throw invalidSession(filePath, "unsupported session version");
+	if (typeof parsed.id !== "string" || !parsed.id) throw invalidSession(filePath, "session header is missing id");
+	if (typeof parsed.timestamp !== "string" || !parsed.timestamp) {
+		throw invalidSession(filePath, "session header is missing timestamp");
+	}
+	if (typeof parsed.cwd !== "string" || !parsed.cwd) throw invalidSession(filePath, "session header is missing cwd");
+	if (parsed.parentSession !== undefined && typeof parsed.parentSession !== "string") {
+		throw invalidSession(filePath, "session header parentSession must be a string");
+	}
+	return {
+		type: "session",
+		version: 3,
+		id: parsed.id,
+		timestamp: parsed.timestamp,
+		cwd: parsed.cwd,
+		parentSession: parsed.parentSession,
+	};
+}
+
+function parseEntryLine(line: string, filePath: string, lineNumber: number): SessionTreeEntry {
+	let parsed: unknown;
+	try {
+		parsed = JSON.parse(line);
+	} catch (error) {
+		throw invalidEntry(filePath, lineNumber, "is not valid JSON", toError(error));
+	}
+	if (!isRecord(parsed)) throw invalidEntry(filePath, lineNumber, "is not a valid session entry");
+	if (typeof parsed.type !== "string") throw invalidEntry(filePath, lineNumber, "is missing entry type");
+	if (typeof parsed.id !== "string" || !parsed.id) throw invalidEntry(filePath, lineNumber, "is missing entry id");
+	if (parsed.parentId !== null && typeof parsed.parentId !== "string") {
+		throw invalidEntry(filePath, lineNumber, "has invalid parentId");
+	}
+	if (typeof parsed.timestamp !== "string" || !parsed.timestamp) {
+		throw invalidEntry(filePath, lineNumber, "is missing timestamp");
+	}
+	if (parsed.type === "leaf" && parsed.targetId !== null && typeof parsed.targetId !== "string") {
+		throw invalidEntry(filePath, lineNumber, "has invalid targetId");
+	}
+	return parsed as unknown as SessionTreeEntry;
+}
+
+function leafIdAfterEntry(entry: SessionTreeEntry): string | null {
+	return entry.type === "leaf" ? entry.targetId : entry.id;
+}
+
+function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata {
+	return {
+		id: header.id,
+		createdAt: header.timestamp,
+		cwd: header.cwd,
+		path,
+		parentSessionPath: header.parentSession,
+	};
+}
+
+export async function loadJsonlSessionMetadata(
+	fs: JsonlSessionStorageFileSystem,
+	filePath: string,
+): Promise<JsonlSessionMetadata> {
+	const lines = getFileSystemResultOrThrow(
+		await fs.readTextLines(filePath, { maxLines: 1 }),
+		`Failed to read session header ${filePath}`,
+	);
+	const line = lines[0];
+	if (line?.trim()) return headerToSessionMetadata(parseHeaderLine(line, filePath), filePath);
+	throw invalidSession(filePath, "missing session header");
+}
+
+async function loadJsonlStorage(
+	fs: JsonlSessionStorageFileSystem,
+	filePath: string,
+): Promise<{
+	header: SessionHeader;
+	entries: SessionTreeEntry[];
+	leafId: string | null;
+}> {
+	const content = getFileSystemResultOrThrow(await fs.readTextFile(filePath), `Failed to read session ${filePath}`);
+	const lines = content.split("\n").filter((line) => line.trim());
+	if (lines.length === 0) {
+		throw invalidSession(filePath, "missing session header");
+	}
+
+	const header = parseHeaderLine(lines[0]!, filePath);
+	const entries: SessionTreeEntry[] = [];
+	let leafId: string | null = null;
+	for (let i = 1; i < lines.length; i++) {
+		const entry = parseEntryLine(lines[i]!, filePath, i + 1);
+		entries.push(entry);
+		leafId = leafIdAfterEntry(entry);
+	}
+	return { header, entries, leafId };
+}
+
+export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata> {
+	private readonly fs: JsonlSessionStorageFileSystem;
+	private readonly filePath: string;
+	private readonly metadata: JsonlSessionMetadata;
+	private entries: SessionTreeEntry[];
+	private byId: Map<string, SessionTreeEntry>;
+	private labelsById: Map<string, string>;
+	private currentLeafId: string | null;
+
+	private constructor(
+		fs: JsonlSessionStorageFileSystem,
+		filePath: string,
+		header: SessionHeader,
+		entries: SessionTreeEntry[],
+		leafId: string | null,
+	) {
+		this.fs = fs;
+		this.filePath = 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;
+	}
+
+	static async open(fs: JsonlSessionStorageFileSystem, filePath: string): Promise<JsonlSessionStorage> {
+		const loaded = await loadJsonlStorage(fs, filePath);
+		return new JsonlSessionStorage(fs, filePath, loaded.header, loaded.entries, loaded.leafId);
+	}
+
+	static async create(
+		fs: JsonlSessionStorageFileSystem,
+		filePath: string,
+		options: {
+			cwd: string;
+			sessionId: string;
+			parentSessionPath?: string;
+		},
+	): Promise<JsonlSessionStorage> {
+		const header: SessionHeader = {
+			type: "session",
+			version: 3,
+			id: options.sessionId,
+			timestamp: new Date().toISOString(),
+			cwd: options.cwd,
+			parentSession: options.parentSessionPath,
+		};
+		getFileSystemResultOrThrow(
+			await fs.writeFile(filePath, `${JSON.stringify(header)}\n`),
+			`Failed to create session ${filePath}`,
+		);
+		return new JsonlSessionStorage(fs, filePath, header, [], null);
+	}
+
+	async getMetadata(): Promise<JsonlSessionMetadata> {
+		return this.metadata;
+	}
+
+	async getLeafId(): Promise<string | null> {
+		if (this.currentLeafId !== null && !this.byId.has(this.currentLeafId)) {
+			throw new SessionError("invalid_session", `Entry ${this.currentLeafId} not found`);
+		}
+		return this.currentLeafId;
+	}
+
+	async setLeafId(leafId: string | null): Promise<void> {
+		if (leafId !== null && !this.byId.has(leafId)) {
+			throw new SessionError("not_found", `Entry ${leafId} not found`);
+		}
+		const entry: LeafEntry = {
+			type: "leaf",
+			id: generateEntryId(this.byId),
+			parentId: this.currentLeafId,
+			timestamp: new Date().toISOString(),
+			targetId: leafId,
+		};
+		getFileSystemResultOrThrow(
+			await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`),
+			`Failed to append session leaf ${entry.id}`,
+		);
+		this.entries.push(entry);
+		this.byId.set(entry.id, entry);
+		this.currentLeafId = leafId;
+	}
+
+	async createEntryId(): Promise<string> {
+		return generateEntryId(this.byId);
+	}
+
+	async appendEntry(entry: SessionTreeEntry): Promise<void> {
+		getFileSystemResultOrThrow(
+			await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`),
+			`Failed to append session entry ${entry.id}`,
+		);
+		this.entries.push(entry);
+		this.byId.set(entry.id, entry);
+		updateLabelCache(this.labelsById, entry);
+		this.currentLeafId = leafIdAfterEntry(entry);
+	}
+
+	async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
+		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);
+	}
+
+	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
+		if (leafId === null) return [];
+		const path: SessionTreeEntry[] = [];
+		let current = this.byId.get(leafId);
+		if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`);
+		while (current) {
+			path.unshift(current);
+			if (!current.parentId) break;
+			const parent = this.byId.get(current.parentId);
+			if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`);
+			current = parent;
+		}
+		return path;
+	}
+
+	async getEntries(): Promise<SessionTreeEntry[]> {
+		return [...this.entries];
+	}
+}

wmo/harness/vendor/pi-agent/src/harness/session/memory-repo.ts

@@ -0,0 +1,50 @@
+import { type Session, SessionError, type SessionMetadata, type SessionRepo } from "../types.ts";
+import { InMemorySessionStorage } from "./memory-storage.ts";
+import { createSessionId, createTimestamp, getEntriesToFork, toSession } from "./repo-utils.ts";
+
+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 SessionError("not_found", `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 storage = new InMemorySessionStorage({ metadata, entries: forkedEntries });
+		const session = toSession(storage);
+		this.sessions.set(metadata.id, session);
+		return session;
+	}
+}

wmo/harness/vendor/pi-agent/src/harness/session/memory-storage.ts

@@ -0,0 +1,131 @@
+import {
+	type LeafEntry,
+	SessionError,
+	type SessionMetadata,
+	type SessionStorage,
+	type SessionTreeEntry,
+} from "../types.ts";
+import { uuidv7 } from "./uuid.ts";
+
+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 generateEntryId(byId: { has(id: string): boolean }): string {
+	for (let i = 0; i < 100; i++) {
+		const id = uuidv7().slice(0, 8);
+		if (!byId.has(id)) return id;
+	}
+	return uuidv7();
+}
+
+function leafIdAfterEntry(entry: SessionTreeEntry): string | null {
+	return entry.type === "leaf" ? entry.targetId : entry.id;
+}
+
+export class InMemorySessionStorage<TMetadata extends SessionMetadata = SessionMetadata>
+	implements SessionStorage<TMetadata>
+{
+	private readonly metadata: TMetadata;
+	private entries: SessionTreeEntry[];
+	private byId: Map<string, SessionTreeEntry>;
+	private labelsById: Map<string, string>;
+	private leafId: string | null;
+
+	constructor(options?: { entries?: SessionTreeEntry[]; metadata?: TMetadata }) {
+		this.entries = options?.entries ? [...options.entries] : [];
+		this.byId = new Map(this.entries.map((entry) => [entry.id, entry]));
+		this.labelsById = buildLabelsById(this.entries);
+		this.leafId = null;
+		for (const entry of this.entries) this.leafId = leafIdAfterEntry(entry);
+		if (this.leafId !== null && !this.byId.has(this.leafId)) {
+			throw new SessionError("invalid_session", `Entry ${this.leafId} not found`);
+		}
+		this.metadata = options?.metadata ?? ({ id: uuidv7(), createdAt: new Date().toISOString() } as TMetadata);
+	}
+
+	async getMetadata(): Promise<TMetadata> {
+		return this.metadata;
+	}
+
+	async getLeafId(): Promise<string | null> {
+		if (this.leafId !== null && !this.byId.has(this.leafId)) {
+			throw new SessionError("invalid_session", `Entry ${this.leafId} not found`);
+		}
+		return this.leafId;
+	}
+
+	async setLeafId(leafId: string | null): Promise<void> {
+		if (leafId !== null && !this.byId.has(leafId)) {
+			throw new SessionError("not_found", `Entry ${leafId} not found`);
+		}
+		const entry: LeafEntry = {
+			type: "leaf",
+			id: generateEntryId(this.byId),
+			parentId: this.leafId,
+			timestamp: new Date().toISOString(),
+			targetId: leafId,
+		};
+		this.entries.push(entry);
+		this.byId.set(entry.id, entry);
+		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);
+		updateLabelCache(this.labelsById, entry);
+		this.leafId = leafIdAfterEntry(entry);
+	}
+
+	async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
+		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);
+	}
+
+	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
+		if (leafId === null) return [];
+		const path: SessionTreeEntry[] = [];
+		let current = this.byId.get(leafId);
+		if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`);
+		while (current) {
+			path.unshift(current);
+			if (!current.parentId) break;
+			const parent = this.byId.get(current.parentId);
+			if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`);
+			current = parent;
+		}
+		return path;
+	}
+
+	async getEntries(): Promise<SessionTreeEntry[]> {
+		return [...this.entries];
+	}
+}

wmo/harness/vendor/pi-agent/src/harness/session/repo-utils.ts

@@ -0,0 +1,51 @@
+import {
+	type FileError,
+	type Result,
+	SessionError,
+	type SessionMetadata,
+	type SessionStorage,
+	type SessionTreeEntry,
+} from "../types.ts";
+import { Session } from "./session.ts";
+import { uuidv7 } from "./uuid.ts";
+
+export function createSessionId(): string {
+	return uuidv7();
+}
+
+export function createTimestamp(): string {
+	return new Date().toISOString();
+}
+
+export function toSession<TMetadata extends SessionMetadata>(storage: SessionStorage<TMetadata>): Session<TMetadata> {
+	return new Session(storage);
+}
+
+export function getFileSystemResultOrThrow<TValue>(result: Result<TValue, FileError>, message: string): TValue {
+	if (!result.ok) {
+		const code = result.error.code === "not_found" ? "not_found" : "storage";
+		throw new SessionError(code, `${message}: ${result.error.message}`, result.error);
+	}
+	return result.value;
+}
+
+export async function getEntriesToFork(
+	storage: SessionStorage,
+	options: { entryId?: string; position?: "before" | "at" },
+): Promise<SessionTreeEntry[]> {
+	if (!options.entryId) return storage.getEntries();
+	const target = await storage.getEntry(options.entryId);
+	if (!target) {
+		throw new SessionError("invalid_fork_target", `Entry ${options.entryId} not found`);
+	}
+	let effectiveLeafId: string | null;
+	if ((options.position ?? "before") === "at") {
+		effectiveLeafId = target.id;
+	} else {
+		if (target.type !== "message" || target.message.role !== "user") {
+			throw new SessionError("invalid_fork_target", `Entry ${options.entryId} is not a user message`);
+		}
+		effectiveLeafId = target.parentId;
+	}
+	return storage.getPathToRoot(effectiveLeafId);
+}