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

wmo/harness/vendor/pi-agent/src/harness/messages.ts

@@ -0,0 +1,164 @@
+import type { ImageContent, Message, TextContent } from "@earendil-works/pi-ai";
+import type { AgentMessage } from "../types.ts";
+
+export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary:
+
+<summary>
+`;
+
+export const COMPACTION_SUMMARY_SUFFIX = `
+</summary>`;
+
+export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from:
+
+<summary>
+`;
+
+export const BRANCH_SUMMARY_SUFFIX = `</summary>`;
+
+export interface BashExecutionMessage {
+	role: "bashExecution";
+	command: string;
+	output: string;
+	exitCode: number | undefined;
+	cancelled: boolean;
+	truncated: boolean;
+	fullOutputPath?: string;
+	timestamp: number;
+	excludeFromContext?: boolean;
+}
+
+export interface CustomMessage<T = unknown> {
+	role: "custom";
+	customType: string;
+	content: string | (TextContent | ImageContent)[];
+	display: boolean;
+	details?: T;
+	timestamp: number;
+}
+
+export interface BranchSummaryMessage {
+	role: "branchSummary";
+	summary: string;
+	fromId: string;
+	timestamp: number;
+}
+
+export interface CompactionSummaryMessage {
+	role: "compactionSummary";
+	summary: string;
+	tokensBefore: number;
+	timestamp: number;
+}
+
+declare module "../types.ts" {
+	interface CustomAgentMessages {
+		bashExecution: BashExecutionMessage;
+		custom: CustomMessage;
+		branchSummary: BranchSummaryMessage;
+		compactionSummary: CompactionSummaryMessage;
+	}
+}
+
+export function bashExecutionToText(msg: BashExecutionMessage): string {
+	let text = `Ran \`${msg.command}\`\n`;
+	if (msg.output) {
+		text += `\`\`\`\n${msg.output}\n\`\`\``;
+	} else {
+		text += "(no output)";
+	}
+	if (msg.cancelled) {
+		text += "\n\n(command cancelled)";
+	} else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) {
+		text += `\n\nCommand exited with code ${msg.exitCode}`;
+	}
+	if (msg.truncated && msg.fullOutputPath) {
+		text += `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]`;
+	}
+	return text;
+}
+
+export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage {
+	return {
+		role: "branchSummary",
+		summary,
+		fromId,
+		timestamp: new Date(timestamp).getTime(),
+	};
+}
+
+export function createCompactionSummaryMessage(
+	summary: string,
+	tokensBefore: number,
+	timestamp: string,
+): CompactionSummaryMessage {
+	return {
+		role: "compactionSummary",
+		summary,
+		tokensBefore,
+		timestamp: new Date(timestamp).getTime(),
+	};
+}
+
+export function createCustomMessage(
+	customType: string,
+	content: string | (TextContent | ImageContent)[],
+	display: boolean,
+	details: unknown | undefined,
+	timestamp: string,
+): CustomMessage {
+	return {
+		role: "custom",
+		customType,
+		content,
+		display,
+		details,
+		timestamp: new Date(timestamp).getTime(),
+	};
+}
+
+export function convertToLlm(messages: AgentMessage[]): Message[] {
+	return messages
+		.map((m): Message | undefined => {
+			switch (m.role) {
+				case "bashExecution":
+					if (m.excludeFromContext) {
+						return undefined;
+					}
+					return {
+						role: "user",
+						content: [{ type: "text", text: bashExecutionToText(m) }],
+						timestamp: m.timestamp,
+					};
+				case "custom": {
+					const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content;
+					return {
+						role: "user",
+						content,
+						timestamp: m.timestamp,
+					};
+				}
+				case "branchSummary":
+					return {
+						role: "user",
+						content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }],
+						timestamp: m.timestamp,
+					};
+				case "compactionSummary":
+					return {
+						role: "user",
+						content: [
+							{ type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX },
+						],
+						timestamp: m.timestamp,
+					};
+				case "user":
+				case "assistant":
+				case "toolResult":
+					return m;
+				default:
+					return undefined;
+			}
+		})
+		.filter((m): m is Message => m !== undefined);
+}

wmo/harness/vendor/pi-agent/src/harness/prompt-templates.ts

@@ -0,0 +1,267 @@
+import { parse } from "yaml";
+import { type ExecutionEnv, type FileInfo, type PromptTemplate, type Result, toError } from "./types.ts";
+
+export type PromptTemplateDiagnosticCode = "file_info_failed" | "list_failed" | "read_failed" | "parse_failed";
+
+/** Warning produced while loading prompt templates. */
+export interface PromptTemplateDiagnostic {
+	/** Diagnostic severity. Currently only warnings are emitted. */
+	type: "warning";
+	/** Stable diagnostic code. */
+	code: PromptTemplateDiagnosticCode;
+	/** Human-readable diagnostic message. */
+	message: string;
+	/** Path associated with the diagnostic. */
+	path: string;
+}
+
+interface PromptTemplateFrontmatter {
+	description?: string;
+	"argument-hint"?: string;
+	[key: string]: unknown;
+}
+
+/**
+ * Load prompt templates from one or more paths.
+ *
+ * Directory inputs load direct `.md` children non-recursively. File inputs load explicit `.md` files. Missing paths and
+ * non-markdown files are skipped. Read and parse failures are returned as diagnostics.
+ */
+export async function loadPromptTemplates(
+	env: ExecutionEnv,
+	paths: string | string[],
+): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> {
+	const promptTemplates: PromptTemplate[] = [];
+	const diagnostics: PromptTemplateDiagnostic[] = [];
+	for (const path of Array.isArray(paths) ? paths : [paths]) {
+		const infoResult = await env.fileInfo(path);
+		if (!infoResult.ok) {
+			if (infoResult.error.code !== "not_found") {
+				diagnostics.push({
+					type: "warning",
+					code: "file_info_failed",
+					message: infoResult.error.message,
+					path,
+				});
+			}
+			continue;
+		}
+		const info = infoResult.value;
+		const kind = await resolveKind(env, info, diagnostics);
+		if (kind === "directory") {
+			const result = await loadTemplatesFromDir(env, info.path);
+			promptTemplates.push(...result.promptTemplates);
+			diagnostics.push(...result.diagnostics);
+		} else if (kind === "file" && info.name.endsWith(".md")) {
+			const result = await loadTemplateFromFile(env, info.path);
+			if (result.promptTemplate) promptTemplates.push(result.promptTemplate);
+			diagnostics.push(...result.diagnostics);
+		}
+	}
+	return { promptTemplates, diagnostics };
+}
+
+/**
+ * Load prompt templates from source-tagged paths.
+ *
+ * Source values are preserved exactly and attached to every loaded prompt template and diagnostic. The agent package does
+ * not interpret source values; applications define their own provenance shape.
+ */
+export async function loadSourcedPromptTemplates<TSource, TPromptTemplate extends PromptTemplate = PromptTemplate>(
+	env: ExecutionEnv,
+	inputs: Array<{ path: string; source: TSource }>,
+	mapPromptTemplate?: (promptTemplate: PromptTemplate, source: TSource) => TPromptTemplate,
+): Promise<{
+	promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }>;
+	diagnostics: Array<PromptTemplateDiagnostic & { source: TSource }>;
+}> {
+	const promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }> = [];
+	const diagnostics: Array<PromptTemplateDiagnostic & { source: TSource }> = [];
+	for (const input of inputs) {
+		const result = await loadPromptTemplates(env, input.path);
+		for (const promptTemplate of result.promptTemplates) {
+			promptTemplates.push({
+				promptTemplate: mapPromptTemplate
+					? mapPromptTemplate(promptTemplate, input.source)
+					: (promptTemplate as TPromptTemplate),
+				source: input.source,
+			});
+		}
+		for (const diagnostic of result.diagnostics) diagnostics.push({ ...diagnostic, source: input.source });
+	}
+	return { promptTemplates, diagnostics };
+}
+
+async function loadTemplatesFromDir(
+	env: ExecutionEnv,
+	dir: string,
+): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> {
+	const promptTemplates: PromptTemplate[] = [];
+	const diagnostics: PromptTemplateDiagnostic[] = [];
+	const entriesResult = await env.listDir(dir);
+	if (!entriesResult.ok) {
+		diagnostics.push({
+			type: "warning",
+			code: "list_failed",
+			message: entriesResult.error.message,
+			path: dir,
+		});
+		return { promptTemplates, diagnostics };
+	}
+	const entries = entriesResult.value;
+
+	for (const entry of entries.sort((a, b) => a.name.localeCompare(b.name))) {
+		const kind = await resolveKind(env, entry, diagnostics);
+		if (kind !== "file" || !entry.name.endsWith(".md")) continue;
+		const result = await loadTemplateFromFile(env, entry.path);
+		if (result.promptTemplate) promptTemplates.push(result.promptTemplate);
+		diagnostics.push(...result.diagnostics);
+	}
+	return { promptTemplates, diagnostics };
+}
+
+async function loadTemplateFromFile(
+	env: ExecutionEnv,
+	filePath: string,
+): Promise<{ promptTemplate: PromptTemplate | null; diagnostics: PromptTemplateDiagnostic[] }> {
+	const diagnostics: PromptTemplateDiagnostic[] = [];
+	const rawContent = await env.readTextFile(filePath);
+	if (!rawContent.ok) {
+		diagnostics.push({
+			type: "warning",
+			code: "read_failed",
+			message: rawContent.error.message,
+			path: filePath,
+		});
+		return { promptTemplate: null, diagnostics };
+	}
+
+	const parsed = parseFrontmatter<PromptTemplateFrontmatter>(rawContent.value);
+	if (!parsed.ok) {
+		diagnostics.push({
+			type: "warning",
+			code: "parse_failed",
+			message: parsed.error.message,
+			path: filePath,
+		});
+		return { promptTemplate: null, diagnostics };
+	}
+
+	const { frontmatter, body } = parsed.value;
+	const firstLine = body.split("\n").find((line) => line.trim());
+	let description = typeof frontmatter.description === "string" ? frontmatter.description : "";
+	if (!description && firstLine) {
+		description = firstLine.slice(0, 60);
+		if (firstLine.length > 60) description += "...";
+	}
+	return {
+		promptTemplate: {
+			name: basenameEnvPath(filePath).replace(/\.md$/i, ""),
+			description,
+			content: body,
+		},
+		diagnostics,
+	};
+}
+
+async function resolveKind(
+	env: ExecutionEnv,
+	info: FileInfo,
+	diagnostics: PromptTemplateDiagnostic[],
+): Promise<"file" | "directory" | undefined> {
+	if (info.kind === "file" || info.kind === "directory") return info.kind;
+	const canonicalPath = await env.canonicalPath(info.path);
+	if (!canonicalPath.ok) {
+		if (canonicalPath.error.code !== "not_found") {
+			diagnostics.push({
+				type: "warning",
+				code: "file_info_failed",
+				message: canonicalPath.error.message,
+				path: info.path,
+			});
+		}
+		return undefined;
+	}
+	const target = await env.fileInfo(canonicalPath.value);
+	if (!target.ok) {
+		if (target.error.code !== "not_found") {
+			diagnostics.push({
+				type: "warning",
+				code: "file_info_failed",
+				message: target.error.message,
+				path: info.path,
+			});
+		}
+		return undefined;
+	}
+	return target.value.kind === "file" || target.value.kind === "directory" ? target.value.kind : undefined;
+}
+
+function parseFrontmatter<T extends Record<string, unknown>>(
+	content: string,
+): Result<{ frontmatter: T; body: string }, Error> {
+	try {
+		const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
+		if (!normalized.startsWith("---")) return { ok: true, value: { frontmatter: {} as T, body: normalized } };
+		const endIndex = normalized.indexOf("\n---", 3);
+		if (endIndex === -1) return { ok: true, value: { frontmatter: {} as T, body: normalized } };
+		const yamlString = normalized.slice(4, endIndex);
+		const body = normalized.slice(endIndex + 4).trim();
+		return { ok: true, value: { frontmatter: (parse(yamlString) ?? {}) as T, body } };
+	} catch (error) {
+		return { ok: false, error: toError(error) };
+	}
+}
+
+function basenameEnvPath(path: string): string {
+	const normalized = path.replace(/\/+$/, "");
+	const slashIndex = normalized.lastIndexOf("/");
+	return slashIndex === -1 ? normalized : normalized.slice(slashIndex + 1);
+}
+
+/** Parse an argument string using simple shell-style single and double quotes. */
+export function parseCommandArgs(argsString: string): string[] {
+	const args: string[] = [];
+	let current = "";
+	let inQuote: string | null = null;
+
+	for (let i = 0; i < argsString.length; i++) {
+		const char = argsString[i]!;
+		if (inQuote) {
+			if (char === inQuote) inQuote = null;
+			else current += char;
+		} else if (char === '"' || char === "'") {
+			inQuote = char;
+		} else if (char === " " || char === "\t") {
+			if (current) {
+				args.push(current);
+				current = "";
+			}
+		} else {
+			current += char;
+		}
+	}
+	if (current) args.push(current);
+	return args;
+}
+
+/** Substitute prompt template placeholders (`$1`, `$@`, `$ARGUMENTS`, `${@:N}`, `${@:N:L}`) with command arguments. */
+export function substituteArgs(content: string, args: string[]): string {
+	let result = content;
+	result = result.replace(/\$(\d+)/g, (_, num: string) => args[parseInt(num, 10) - 1] ?? "");
+	result = result.replace(/\$\{@:(\d+)(?::(\d+))?\}/g, (_, startStr: string, lengthStr?: string) => {
+		let start = parseInt(startStr, 10) - 1;
+		if (start < 0) start = 0;
+		if (lengthStr) return args.slice(start, start + parseInt(lengthStr, 10)).join(" ");
+		return args.slice(start).join(" ");
+	});
+	const allArgs = args.join(" ");
+	result = result.replace(/\$ARGUMENTS/g, allArgs);
+	result = result.replace(/\$@/g, allArgs);
+	return result;
+}
+
+/** Format a prompt template invocation with positional arguments. */
+export function formatPromptTemplateInvocation(template: PromptTemplate, args: string[] = []): string {
+	return substituteArgs(template.content, args);
+}

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

@@ -0,0 +1,177 @@
+import type {
+	FileSystem,
+	JsonlSessionCreateOptions,
+	JsonlSessionListOptions,
+	JsonlSessionMetadata,
+	JsonlSessionRepoApi,
+	Session,
+} from "../types.ts";
+import { SessionError, toError } from "../types.ts";
+import { JsonlSessionStorage, loadJsonlSessionMetadata } from "./jsonl-storage.ts";
+import {
+	createSessionId,
+	createTimestamp,
+	getEntriesToFork,
+	getFileSystemResultOrThrow,
+	toSession,
+} from "./repo-utils.ts";
+
+type JsonlSessionRepoFileSystem = Pick<
+	FileSystem,
+	| "cwd"
+	| "absolutePath"
+	| "joinPath"
+	| "readTextFile"
+	| "readTextLines"
+	| "writeFile"
+	| "appendFile"
+	| "listDir"
+	| "exists"
+	| "createDir"
+	| "remove"
+>;
+
+function encodeCwd(cwd: string): string {
+	return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`;
+}
+
+export class JsonlSessionRepo implements JsonlSessionRepoApi {
+	private readonly fs: JsonlSessionRepoFileSystem;
+	private readonly sessionsRootInput: string;
+	private sessionsRoot: string | undefined;
+
+	constructor(options: { fs: JsonlSessionRepoFileSystem; sessionsRoot: string }) {
+		this.fs = options.fs;
+		this.sessionsRootInput = options.sessionsRoot;
+	}
+
+	private async getSessionsRoot(): Promise<string> {
+		if (!this.sessionsRoot) {
+			this.sessionsRoot = getFileSystemResultOrThrow(
+				await this.fs.absolutePath(this.sessionsRootInput),
+				`Failed to resolve sessions root ${this.sessionsRootInput}`,
+			);
+		}
+		return this.sessionsRoot;
+	}
+
+	private async getSessionDir(cwd: string): Promise<string> {
+		return getFileSystemResultOrThrow(
+			await this.fs.joinPath([await this.getSessionsRoot(), encodeCwd(cwd)]),
+			`Failed to resolve session directory for ${cwd}`,
+		);
+	}
+
+	private async createSessionFilePath(cwd: string, sessionId: string, timestamp: string): Promise<string> {
+		return getFileSystemResultOrThrow(
+			await this.fs.joinPath([
+				await this.getSessionDir(cwd),
+				`${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`,
+			]),
+			`Failed to resolve session file path for ${sessionId}`,
+		);
+	}
+
+	async create(options: JsonlSessionCreateOptions): Promise<Session<JsonlSessionMetadata>> {
+		const id = options.id ?? createSessionId();
+		const createdAt = createTimestamp();
+		const sessionDir = await this.getSessionDir(options.cwd);
+		getFileSystemResultOrThrow(
+			await this.fs.createDir(sessionDir, { recursive: true }),
+			`Failed to create session directory ${sessionDir}`,
+		);
+		const filePath = await this.createSessionFilePath(options.cwd, id, createdAt);
+		const storage = await JsonlSessionStorage.create(this.fs, filePath, {
+			cwd: options.cwd,
+			sessionId: id,
+			parentSessionPath: options.parentSessionPath,
+		});
+		return toSession(storage);
+	}
+
+	async open(metadata: JsonlSessionMetadata): Promise<Session<JsonlSessionMetadata>> {
+		if (
+			!getFileSystemResultOrThrow(await this.fs.exists(metadata.path), `Failed to check session ${metadata.path}`)
+		) {
+			throw new SessionError("not_found", `Session not found: ${metadata.path}`);
+		}
+		const storage = await JsonlSessionStorage.open(this.fs, metadata.path);
+		return toSession(storage);
+	}
+
+	async list(options: JsonlSessionListOptions = {}): Promise<JsonlSessionMetadata[]> {
+		const dirs = options.cwd ? [await this.getSessionDir(options.cwd)] : await this.listSessionDirs();
+		const sessions: JsonlSessionMetadata[] = [];
+		for (const dir of dirs) {
+			if (!getFileSystemResultOrThrow(await this.fs.exists(dir), `Failed to check session directory ${dir}`)) {
+				continue;
+			}
+			const files = getFileSystemResultOrThrow(
+				await this.fs.listDir(dir),
+				`Failed to list sessions in ${dir}`,
+			).filter((file) => file.kind !== "directory" && file.name.endsWith(".jsonl"));
+			for (const file of files) {
+				try {
+					sessions.push(await loadJsonlSessionMetadata(this.fs, file.path));
+				} catch (error) {
+					const cause = toError(error);
+					if (!(cause instanceof SessionError) || cause.code !== "invalid_session") throw cause;
+				}
+			}
+		}
+		sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime());
+		return sessions;
+	}
+
+	async delete(metadata: JsonlSessionMetadata): Promise<void> {
+		getFileSystemResultOrThrow(
+			await this.fs.remove(metadata.path, { force: true }),
+			`Failed to delete session ${metadata.path}`,
+		);
+	}
+
+	async fork(
+		sourceMetadata: JsonlSessionMetadata,
+		options: JsonlSessionCreateOptions & { entryId?: string; position?: "before" | "at"; id?: string },
+	): Promise<Session<JsonlSessionMetadata>> {
+		const source = await this.open(sourceMetadata);
+		const forkedEntries = await getEntriesToFork(source.getStorage(), options);
+		const id = options.id ?? createSessionId();
+		const createdAt = createTimestamp();
+		const sessionDir = await this.getSessionDir(options.cwd);
+		getFileSystemResultOrThrow(
+			await this.fs.createDir(sessionDir, { recursive: true }),
+			`Failed to create session directory ${sessionDir}`,
+		);
+		const storage = await JsonlSessionStorage.create(
+			this.fs,
+			await this.createSessionFilePath(options.cwd, id, createdAt),
+			{
+				cwd: options.cwd,
+				sessionId: id,
+				parentSessionPath: options.parentSessionPath ?? sourceMetadata.path,
+			},
+		);
+		for (const entry of forkedEntries) {
+			await storage.appendEntry(entry);
+		}
+		return toSession(storage);
+	}
+
+	private async listSessionDirs(): Promise<string[]> {
+		const sessionsRoot = await this.getSessionsRoot();
+		if (
+			!getFileSystemResultOrThrow(
+				await this.fs.exists(sessionsRoot),
+				`Failed to check sessions root ${sessionsRoot}`,
+			)
+		) {
+			return [];
+		}
+		const entries = getFileSystemResultOrThrow(
+			await this.fs.listDir(sessionsRoot),
+			`Failed to list sessions root ${sessionsRoot}`,
+		);
+		return entries.filter((entry) => entry.kind === "directory").map((entry) => entry.path);
+	}
+}