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

wmo/harness/pi_entry/runner_live.ts

@@ -0,0 +1,587 @@
+/**
+ * Live-session pi runner: the in-sandbox peer of wmo/harness/live_session.py (LiveSession).
+ *
+ * Where runner_stdio.ts runs ONE fire-and-forget episode, this runner hosts ONE long-lived pi
+ * Agent for an interactive multi-turn session: the host sends `session_start` once (materializing
+ * the champion's code surfaces), then `user_message` / `abort` / `ping` frames over the session's
+ * life. Interactive sessions retain one transcript; filesystem-backed projects can instead scope
+ * conversation to each outer turn while reusing the same Agent, runner, and project filesystem.
+ * The transport, the localhost LLM bridge (credentials stay host-side), and the env-tools-as-
+ * `tool_request` contract are identical to runner_stdio.ts — deliberately re-mirrored rather than
+ * imported, because runner_stdio boots per-episode helpers and this runner's lifecycle differs.
+ *
+ * stdout is the frame stream and NOTHING else may write to it: console.* is rebound to stderr
+ * before any agent code loads. One base64(JSON) frame per line.
+ *
+ * Run (inside the sandbox, workdir holding node_modules + package.json {"type":"module"}):
+ *   cd /home/user/pi-run && node --experimental-strip-types runner_live.ts
+ */
+import fs from "node:fs";
+import http from "node:http";
+import path from "node:path";
+import util from "node:util";
+import { pathToFileURL } from "node:url";
+import type { AddressInfo } from "node:net";
+
+// CRITICAL — FIRST statements: rebind every stdout-defaulting console channel to stderr before
+// anything else runs, so stray prints can never corrupt the frame stream.
+const toStderr = (...args: unknown[]): void => {
+	process.stderr.write(util.format(...args) + "\n");
+};
+console.log = toStderr;
+console.info = toStderr;
+console.warn = toStderr;
+console.debug = toStderr;
+
+const AGENT_MODEL = process.env.PI_AGENT_MODEL ?? "worker";
+// Last-resort model context window when session_start carries none. The host resolves the REAL
+// served window (provider/SDK model info) and sends it as context_window; never assume a size here.
+const DEFAULT_CONTEXT_WINDOW = 128000;
+const TRANSPORT_KEEPALIVE_MS = 30_000;
+const LIVE_OUTBOX = process.env.WMO_LIVE_OUTBOX?.trim() || null;
+
+type Frame = Record<string, any>;
+type TransportEnvelope = { transport_seq: number; frame: Frame };
+
+function encodeFrame(frame: Frame): string {
+	return Buffer.from(JSON.stringify(frame), "utf8").toString("base64") + "\n";
+}
+
+function decodeFrame(line: string): Frame | null {
+	const text = line.trim();
+	if (!text) return null;
+	try {
+		const frame = JSON.parse(Buffer.from(text, "base64").toString("utf8"));
+		return frame && typeof frame === "object" && !Array.isArray(frame) ? (frame as Frame) : null;
+	} catch {
+		return null;
+	}
+}
+
+/**
+ * Durable, replayable copy of the semantic output stream.
+ *
+ * A frame file is published before the head watermark advances, and both writes use a temporary
+ * sibling plus rename. Readers may therefore trust every sequence through `head` without ever
+ * observing partial JSON. The decimal watermark also lets a replacement runner continue a
+ * pre-existing outbox without reusing committed sequence numbers.
+ */
+class DurableOutbox {
+	private readonly framesDir: string;
+	private readonly headPath: string;
+	private readonly tmpNamespace = `${process.pid}-${Date.now().toString(36)}-${process.hrtime.bigint().toString(36)}`;
+	private transportSeq = 0;
+	private tmpSeq = 0;
+
+	constructor(root: string) {
+		const resolvedRoot = path.resolve(root);
+		this.framesDir = path.join(resolvedRoot, "frames");
+		this.headPath = path.join(resolvedRoot, "head");
+		fs.mkdirSync(this.framesDir, { recursive: true });
+
+		if (fs.existsSync(this.headPath)) {
+			const current = fs.readFileSync(this.headPath, "utf8").trim();
+			const parsed = Number(current);
+			if (!Number.isSafeInteger(parsed) || parsed < 0) {
+				throw new Error(`invalid live outbox head at ${this.headPath}: ${JSON.stringify(current)}`);
+			}
+			this.transportSeq = parsed;
+		} else {
+			this.atomicWrite(this.headPath, "0\n");
+		}
+	}
+
+	publish(frame: Frame): TransportEnvelope {
+		const transport_seq = this.transportSeq + 1;
+		const envelope: TransportEnvelope = { transport_seq, frame };
+		const filename = `${String(transport_seq).padStart(20, "0")}.json`;
+		this.atomicWrite(path.join(this.framesDir, filename), JSON.stringify(envelope) + "\n");
+		this.atomicWrite(this.headPath, `${transport_seq}\n`);
+		this.transportSeq = transport_seq;
+		return envelope;
+	}
+
+	private atomicWrite(target: string, content: string): void {
+		const tmp = path.join(
+			path.dirname(target),
+			`.${path.basename(target)}.tmp-${this.tmpNamespace}-${++this.tmpSeq}`,
+		);
+		try {
+			fs.writeFileSync(tmp, content, { encoding: "utf8", flag: "wx" });
+			fs.renameSync(tmp, target);
+		} catch (error) {
+			try {
+				fs.unlinkSync(tmp);
+			} catch {
+				// The temp file may not have been created, or rename may already have consumed it.
+			}
+			throw error;
+		}
+	}
+}
+
+/** The stdio twin of runner_frames.FrameConn: `request` awaits a matching response by req_id;
+ *  host-pushed frames (session_start, user_message, abort, ping, shutdown) fire handlers. */
+class StdioConn {
+	private buf = "";
+	private waiters = new Map<number, (f: Frame) => void>();
+	private handlers = new Map<string, (f: Frame) => void>();
+	private reqSeq = 0;
+	private readonly outbox: DurableOutbox | null;
+	private lastInboundSeq = 0;
+
+	constructor() {
+		// Initialize the complete durable layout before `main` can emit its ready/hello frame.
+		this.outbox = LIVE_OUTBOX ? new DurableOutbox(LIVE_OUTBOX) : null;
+		process.stdin.setEncoding("utf8");
+		process.stdin.on("data", (chunk: string) => this.onData(chunk));
+		process.stdin.on("end", () => process.exit(0));
+	}
+
+	on(type: string, handler: (f: Frame) => void): void {
+		this.handlers.set(type, handler);
+	}
+
+	send(frame: Frame): void {
+		if (!this.outbox) {
+			// Keep the original wire format byte-for-byte when durable transport is not requested.
+			process.stdout.write(encodeFrame(frame));
+			return;
+		}
+		const envelope = this.outbox.publish(frame);
+		process.stdout.write(encodeFrame(envelope));
+	}
+
+	request(type: string, payload: Frame): Promise<Frame> {
+		const req_id = ++this.reqSeq;
+		return new Promise((resolve) => {
+			this.waiters.set(req_id, resolve);
+			this.send({ type, req_id, ...payload });
+		});
+	}
+
+	/** Keep the command stream active across the persistent session; timeout remains host-owned. */
+	startTransportKeepalive(): () => void {
+		const timer = setInterval(
+			// Liveness ticks carry no semantic state: leave them on the legacy wire and out of the
+			// replay log so they neither advance the sequence nor create unbounded tiny files.
+			() => process.stdout.write(encodeFrame({ type: "transport_keepalive" })),
+			TRANSPORT_KEEPALIVE_MS,
+		);
+		timer.unref();
+		return () => clearInterval(timer);
+	}
+
+	private onData(chunk: string): void {
+		this.buf += chunk;
+		let nl = this.buf.indexOf("\n");
+		while (nl >= 0) {
+			const line = this.buf.slice(0, nl);
+			this.buf = this.buf.slice(nl + 1);
+			const frame = decodeFrame(line);
+			if (frame) this.dispatchInbound(frame);
+			nl = this.buf.indexOf("\n");
+		}
+	}
+
+	private dispatchInbound(value: Frame): void {
+		if (!this.outbox) {
+			this.dispatch(value);
+			return;
+		}
+
+		const inboundSeq = value.transport_in_seq;
+		const frame = value.frame;
+		if (
+			typeof inboundSeq !== "number" ||
+			!Number.isSafeInteger(inboundSeq) ||
+			inboundSeq <= 0 ||
+			!frame ||
+			typeof frame !== "object" ||
+			Array.isArray(frame)
+		) {
+			this.send({
+				type: "transport_nack",
+				transport_in_seq: inboundSeq,
+				expected_transport_in_seq: this.lastInboundSeq + 1,
+			});
+			return;
+		}
+		if (inboundSeq <= this.lastInboundSeq) {
+			// The original dispatch succeeded but its ack notification was lost. Re-publish only
+			// the acknowledgement: repeating a response, prompt, or tool result is unsafe.
+			this.send({ type: "transport_ack", transport_in_seq: inboundSeq });
+			return;
+		}
+		if (inboundSeq !== this.lastInboundSeq + 1) {
+			this.send({
+				type: "transport_nack",
+				transport_in_seq: inboundSeq,
+				expected_transport_in_seq: this.lastInboundSeq + 1,
+			});
+			return;
+		}
+
+		this.lastInboundSeq = inboundSeq;
+		if (frame.type === "shutdown") {
+			// shutdown exits synchronously, so persist acceptance before dispatching it.
+			this.send({ type: "transport_ack", transport_in_seq: inboundSeq });
+			this.dispatch(frame);
+			return;
+		}
+		this.dispatch(frame);
+		this.send({ type: "transport_ack", transport_in_seq: inboundSeq });
+	}
+
+	private dispatch(frame: Frame): void {
+		const rid = frame.req_id;
+		if (typeof rid === "number" && this.waiters.has(rid)) {
+			const resolve = this.waiters.get(rid);
+			this.waiters.delete(rid);
+			resolve?.(frame);
+			return;
+		}
+		const handler = this.handlers.get(frame.type);
+		if (handler) handler(frame);
+	}
+}
+
+function assistantText(msg: any): string {
+	if (!msg || msg.role !== "assistant" || !Array.isArray(msg.content)) return "";
+	return msg.content
+		.filter((c: any) => c?.type === "text")
+		.map((c: any) => String(c.text ?? ""))
+		.join("")
+		.trim();
+}
+
+interface Bridge {
+	url: string;
+	close: () => void;
+}
+
+/** Localhost endpoint pi's openai-completions transport POSTs to; frames each request to the host
+ *  and renders the returned completion back as the SSE pi's parser expects. Creds stay host-side. */
+function startLlmBridge(conn: StdioConn): Promise<Bridge> {
+	return new Promise((resolve) => {
+		const server = http.createServer((req, res) => {
+			const chunks: Buffer[] = [];
+			req.on("data", (c: Buffer) => chunks.push(c));
+			req.on("end", async () => {
+				res.writeHead(200, { "Content-Type": "text/event-stream", Connection: "close" });
+				try {
+					const body = JSON.parse(Buffer.concat(chunks).toString("utf8") || "{}");
+					const reply = await conn.request("llm_request", { openai_body: body });
+					if (reply.error) {
+						res.end(`data: ${JSON.stringify({ error: { message: reply.error } })}\n\ndata: [DONE]\n\n`);
+						return;
+					}
+					const choice = reply.completion?.choices?.[0] ?? {};
+					const msg = choice.message ?? {};
+					const delta: any = { role: "assistant", content: msg.content ?? "" };
+					if (msg.reasoning_details) delta.reasoning_details = msg.reasoning_details;
+					if (msg.tool_calls) {
+						delta.tool_calls = msg.tool_calls.map((tc: any, i: number) => ({
+							index: i,
+							id: tc.id,
+							type: tc.type ?? "function",
+							function: tc.function ?? {},
+						}));
+					}
+					const first = { choices: [{ index: 0, delta, finish_reason: null }] };
+					const last = {
+						choices: [{ index: 0, delta: {}, finish_reason: choice.finish_reason ?? "stop" }],
+						// Pi uses the latest assistant usage to estimate occupied context. Without this,
+						// it falls back to chars/4 and can prematurely clamp the next output budget.
+						usage: reply.completion?.usage,
+					};
+					res.write(`data: ${JSON.stringify(first)}\n\n`);
+					res.write(`data: ${JSON.stringify(last)}\n\n`);
+					res.end("data: [DONE]\n\n");
+				} catch (e) {
+					res.end(`data: ${JSON.stringify({ error: { message: String(e) } })}\n\ndata: [DONE]\n\n`);
+				}
+			});
+		});
+		server.listen(0, "127.0.0.1", () => {
+			const addr = server.address() as AddressInfo;
+			resolve({ url: `http://127.0.0.1:${addr.port}/v1`, close: () => server.close() });
+		});
+	});
+}
+
+/** Materialize session_start.files (the champion's code surfaces) into ./live/ once and import the
+ *  Agent. Unlike the episode runner there is no per-run dir — the session reuses one Agent. */
+async function loadAgent(start: Frame): Promise<any> {
+	const files: Record<string, string> = start.files ?? {};
+	if (!files["src/agent.ts"]) {
+		throw new Error("session_start carried no src/agent.ts (the live runner has no static fallback)");
+	}
+	const base = path.join(process.cwd(), "live");
+	const basePrefix = base.endsWith(path.sep) ? base : base + path.sep;
+	for (const [rel, content] of Object.entries(files)) {
+		// Keep every materialized file under ./live: `path.join`/`resolve` let a
+		// `../` or absolute manifest key escape and overwrite arbitrary sandbox
+		// files, so reject anything that resolves outside the base.
+		const dst = path.resolve(base, rel);
+		if (dst !== base && !dst.startsWith(basePrefix)) {
+			throw new Error(`session_start file path escapes the live directory: ${rel}`);
+		}
+		fs.mkdirSync(path.dirname(dst), { recursive: true });
+		fs.writeFileSync(dst, content);
+	}
+	const mod = await import(pathToFileURL(path.join(base, "src/agent.ts")).href);
+	if (typeof mod.Agent !== "function") {
+		throw new Error("champion src/agent.ts does not export an Agent class");
+	}
+	return mod.Agent;
+}
+
+const REQUIRED_AGENT_METHODS = ["prompt", "steer", "abort", "subscribe"] as const;
+
+function assertInteractive(AgentCtor: any): void {
+	for (const method of REQUIRED_AGENT_METHODS) {
+		if (typeof AgentCtor.prototype?.[method] !== "function") {
+			throw new Error(
+				`champion Agent is not steerable: missing ${method}() — this harness cannot run a live session`,
+			);
+		}
+	}
+}
+
+function userMessage(text: string): any {
+	return { role: "user", content: [{ type: "text", text }], timestamp: Date.now() };
+}
+
+/**
+ * After an aborted run, the transcript tail can hold an assistant message whose toolCalls never got
+ * results (the loop returned on the abort signal before executing/finishing them). Left as-is, the
+ * NEXT provider request is a 400 (OpenAI) / validation error (Bedrock Converse): an assistant
+ * tool-call with no matching tool result. Append a synthetic "cancelled by user" result for every
+ * orphaned toolCall so the next turn is well-formed. (Vendored eval runners never hit this — they
+ * only ever run one prompt; interactivity is what surfaces it.)
+ */
+function repairOrphanedToolCalls(agent: any): void {
+	const messages: any[] = agent.state?.messages ?? [];
+	const resolved = new Set<string>();
+	for (const m of messages) {
+		if (m?.role === "toolResult" && typeof m.toolCallId === "string") resolved.add(m.toolCallId);
+	}
+	const repairs: any[] = [];
+	for (const m of messages) {
+		if (m?.role !== "assistant" || !Array.isArray(m.content)) continue;
+		for (const block of m.content) {
+			if (block?.type === "toolCall" && typeof block.id === "string" && !resolved.has(block.id)) {
+				resolved.add(block.id);
+				repairs.push({
+					role: "toolResult",
+					toolCallId: block.id,
+					toolName: block.name ?? "",
+					content: [{ type: "text", text: "cancelled by user" }],
+					details: {},
+					isError: true,
+					timestamp: Date.now(),
+				});
+			}
+		}
+	}
+	if (repairs.length > 0) agent.state.messages = [...messages, ...repairs];
+}
+
+/** Owns the single Agent and serializes prompt/steer/abort across the session's frames. */
+class Session {
+	// NOTE: fields declared explicitly (not via constructor parameter properties) — node's
+	// --experimental-strip-types rejects `constructor(private x)` in strip-only mode.
+	private readonly conn: StdioConn;
+	private agent: any = null;
+	private bridge: Bridge | null = null;
+	private running = false;
+	private turns = 0;
+	private turnCap = 60;
+	private conversationScope: "session" | "turn" = "session";
+	private interrupted = false;
+	private hitTurnCap = false;
+
+	constructor(conn: StdioConn) {
+		this.conn = conn;
+	}
+
+	async start(frame: Frame): Promise<void> {
+		this.turnCap = Number(frame.turn_cap ?? 60);
+		this.conversationScope = frame.conversation_scope === "turn" ? "turn" : "session";
+		const maxOutputTokens =
+			Number.isInteger(frame.max_output_tokens) && frame.max_output_tokens >= 1
+				? frame.max_output_tokens
+				: 4096;
+		const contextWindow =
+			Number.isInteger(frame.context_window) && frame.context_window >= 1024
+				? frame.context_window
+				: DEFAULT_CONTEXT_WINDOW;
+		const AgentCtor = await loadAgent(frame);
+		assertInteractive(AgentCtor);
+		this.bridge = await startLlmBridge(this.conn);
+
+		const model = {
+			id: AGENT_MODEL,
+			name: AGENT_MODEL,
+			api: "openai-completions",
+			provider: "link", // non-builtin -> uses model.baseUrl (our localhost bridge) directly
+			baseUrl: this.bridge.url,
+			reasoning: false,
+			input: ["text"],
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+			contextWindow,
+			maxTokens: maxOutputTokens,
+		};
+
+		const tools = this.buildTools(frame.tools ?? []);
+		this.agent = new AgentCtor({
+			initialState: { systemPrompt: frame.system ?? "", model, tools },
+			getApiKey: () => "x",
+			// Real filesystem tools may race if run in parallel; keep them ordered. Steers drain
+			// as a batch ("all") so a burst of user messages doesn't each cost an extra turn.
+			toolExecution: "sequential",
+			steeringMode: "all",
+		});
+		this.agent.subscribe((event: any) => {
+			if (event.type === "turn_end") {
+				this.turns += 1;
+				if (this.turns >= this.turnCap && this.running) {
+					this.hitTurnCap = true;
+					this.agent.abort();
+				}
+			}
+		});
+		this.sendState("idle");
+	}
+
+	private buildTools(specs: any[]): any[] {
+		const envTools = specs
+			.filter((t: any) => t.name !== "submit")
+			.map((t: any) => ({
+				name: t.name,
+				label: t.name,
+				description: t.description,
+				parameters: t.parameters,
+				execute: async (_id: string, params: any, signal?: AbortSignal) => {
+					if (signal?.aborted) {
+						return { content: [{ type: "text", text: "interrupted" }], details: {}, terminate: false };
+					}
+					const r = await this.conn.request("tool_request", { name: t.name, arguments: params });
+					// A host-side failure (budget exhausted, unknown tool, executor error) must reach
+					// the agent AS a failure: throw so pi records an error tool result, rather than
+					// letting the model reason from a failed action as if it succeeded.
+					if (r.is_error) {
+						throw new Error(String(r.content ?? "tool failed"));
+					}
+					return { content: [{ type: "text", text: String(r.content ?? "") }], details: r, terminate: false };
+				},
+			}));
+		const submit = {
+			name: "submit",
+			label: "submit",
+			description: "Finish the task and submit your answer/result summary. This ends the run.",
+			parameters: { type: "object", properties: { answer: { type: "string" } }, required: ["answer"] },
+			execute: async (_id: string, params: { answer?: string }, signal?: AbortSignal) => {
+				// Honor an interrupt racing a submit: don't emit a final submit for an aborted turn.
+				if (signal?.aborted) {
+					return { content: [{ type: "text", text: "interrupted" }], details: {}, terminate: false };
+				}
+				// The host emits the submit event; the run then ends (but the SESSION stays alive
+				// awaiting the next user message).
+				await this.conn.request("tool_request", { name: "submit", arguments: { answer: params.answer ?? "" } });
+				// Re-check after the round-trip: if the interrupt landed while the request was in
+				// flight, end via the abort path rather than terminating the run as a clean submit.
+				if (signal?.aborted) {
+					return { content: [{ type: "text", text: "interrupted" }], details: {}, terminate: false };
+				}
+				return { content: [{ type: "text", text: "submitted" }], details: {}, terminate: true };
+			},
+		};
+		return [...envTools, submit];
+	}
+
+	handleUserMessage(frame: Frame): void {
+		const text = String(frame.text ?? "");
+		if (!this.agent) return;
+		if (this.running) {
+			this.agent.steer(userMessage(text));
+			return;
+		}
+		void this.runPrompt(text);
+	}
+
+	private async runPrompt(text: string): Promise<void> {
+		this.running = true;
+		this.turns = 0;
+		this.interrupted = false;
+		this.hitTurnCap = false;
+		if (this.conversationScope === "turn") {
+			// AgentProject stores durable memory in its filesystem and gives every outer task a
+			// self-contained request. Retaining all prior task transcripts duplicates that state,
+			// grows input cost without bound, and makes pi clamp maxTokens to 1 near its context cap.
+			// Reset only here, while idle: tool calls within this logical turn still share context.
+			this.agent.state.messages = [];
+		}
+		this.sendState("running");
+		try {
+			await this.agent.prompt(text);
+		} catch (e) {
+			this.conn.send({ type: "episode_error", note: String(e) });
+		} finally {
+			this.running = false;
+			repairOrphanedToolCalls(this.agent);
+			const reason = this.hitTurnCap ? "turn_limit" : this.interrupted ? "aborted" : "completed";
+			this.sendState("idle", reason);
+		}
+	}
+
+	handleAbort(_frame: Frame): void {
+		if (this.agent && this.running) {
+			this.interrupted = true;
+			// Interrupt cancels the WHOLE pending turn: clear any messages the user queued (via
+			// steer) before pressing Stop, so a follow-up typed just before the interrupt is not
+			// silently drained into the next prompt. New messages after this start a fresh turn.
+			this.agent.clearAllQueues?.();
+			this.agent.abort();
+		}
+	}
+
+	handlePing(frame: Frame): void {
+		this.conn.send({ type: "pong", nonce: frame.nonce });
+	}
+
+	private sendState(status: "idle" | "running", reason?: string): void {
+		const frame: Frame = { type: "state", status, turns: this.turns };
+		if (reason) frame.reason = reason;
+		this.conn.send(frame);
+	}
+}
+
+function main(): void {
+	const conn = new StdioConn();
+	// AgentProject deliberately keeps this ordinary live session between proposal turns. Keep the
+	// output stream attached while idle too; host-side E2B timeout/idle-suspend policy is separate.
+	conn.startTransportKeepalive();
+	const session = new Session(conn);
+	conn.on("shutdown", () => process.exit(0));
+	conn.on("session_start", (start) => {
+		session.start(start).catch((e) => {
+			conn.send({ type: "episode_error", note: String(e) });
+			process.stderr.write(`[runner-live] start fatal ${e}\n`);
+		});
+	});
+	conn.on("user_message", (f) => session.handleUserMessage(f));
+	conn.on("abort", (f) => session.handleAbort(f));
+	conn.on("ping", (f) => session.handlePing(f));
+	conn.send({
+		type: "hello",
+		node_version: process.version,
+		pi_version: "0.80.3",
+		mode: "session",
+		transport: "stdio",
+	});
+	process.stderr.write("[runner-live] ready\n");
+}
+
+main();