pi-crew 0.9.53 → 0.9.55

package/src/prompt/prompt-runtime.ts

@@ -0,0 +1,380 @@
+import * as fs from "node:fs";
+import * as path from "node:path";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { startChildBrokerClient } from "../runtime/crew-broker-child.ts";
+import { logInternalError } from "../utils/internal-error.ts";
+import { resolveRealContainedPath } from "../utils/safe-paths.ts";
+
+export const PI_TEAMS_INHERIT_PROJECT_CONTEXT_ENV = "PI_TEAMS_INHERIT_PROJECT_CONTEXT";
+export const PI_TEAMS_INHERIT_SKILLS_ENV = "PI_TEAMS_INHERIT_SKILLS";
+export const PI_CREW_INHERIT_PROJECT_CONTEXT_ENV = "PI_CREW_INHERIT_PROJECT_CONTEXT";
+export const PI_CREW_INHERIT_SKILLS_ENV = "PI_CREW_INHERIT_SKILLS";
+const PI_CREW_MAX_OUTPUT_ENV = "PI_CREW_MAX_OUTPUT";
+const PI_CREW_STEERING_FILE_ENV = "PI_CREW_STEERING_FILE";
+
+const PROJECT_CONTEXT_HEADER = "\n\n# Project Context\n\nProject-specific instructions and guidelines:\n\n";
+const SKILLS_HEADER = "\n\nThe following skills provide specialized instructions for specific tasks.";
+const DATE_HEADER = "\nCurrent date:";
+
+// ── FIX-02: Steering content sanitization limits ──────────────────────────
+// Bounded to keep a malformed/malicious steer entry from blowing up the
+// worker's prompt budget or smuggling control sequences into the agent.
+const MAX_STEER_MESSAGE_LENGTH = 4096;
+const MAX_STEER_MESSAGE_NEWLINES = 50;
+// C0 control characters minus the printable whitespace (\t \n \r). These
+// are the bytes most useful for ANSI escapes, terminal-control tricks, and
+// NUL-injection attacks when steer content reaches the worker's UI.
+const STEER_CONTROL_CHAR_PATTERN = /[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]/;
+
+export interface SteerSanitizeResult {
+	valid: boolean;
+	reason?: string;
+	message?: string;
+}
+
+export interface SteerEntry {
+	type?: string;
+	message?: string;
+	/** Message id. Present in entries the broker writes; absent in legacy
+	 *  entries pre-dating the broker id-forwarding change. */
+	id?: string;
+}
+
+// ── FIX-S1: Cross-channel steer dedup state ──────────────────────────────
+// Both the live broker push (mailbox.message → onSteer) and the durable
+// file poll (pollSteering JSONL) can deliver the SAME steer to the worker
+// for two reasons:
+//   1. The broker writes to BOTH the mailbox AND the steering JSONL for
+//      durability. A connected child receives the broker push FIRST, then
+//      file-poll sees the JSONL shortly after (the broker's own write is
+//      what populates that file).
+//   2. A reconnect / catch-up from a previously-disconnected child can
+//      re-deliver the same mailbox.message id a second time.
+//
+// We dedup at the recipient (worker) by tracking steer ids across BOTH
+// channels in a single bounded FIFO set. Entries without an id (legacy
+// JSONL rows written before S1) are NOT deduped, because they lack a
+// stable identity — the file-poll path is their only delivery route.
+const SEEN_STEER_ID_CAP = 1024;
+
+/**
+ * Bounded FIFO seen-set for steer message ids.
+ *
+ * - `markOrSkip(undefined)` returns true (the file-poll path may emit
+ *   legacy id-less entries; forward them).
+ * - `markOrSkip('a')` first call returns true; the same call again returns
+ *   false (id already seen → drop the duplicate deliver).
+ * - When the cap is exceeded, the oldest id is evicted so the set stays
+ *   bounded under long-running workers with high steer churn.
+ *
+ * Factory-shaped so multiple prompt-runtime instances (tests, parallel
+ * workers) get independent sets.
+ */
+export function createSeenSteerIdSet(): { markOrSkip: (id?: string) => boolean; size: () => number } {
+	const seen: string[] = [];
+	const set = new Set<string>();
+	return {
+		markOrSkip(id?: string): boolean {
+			if (id === undefined) return true; // legacy steers have no id; only file-poll path can produce these
+			if (set.has(id)) return false;
+			set.add(id);
+			seen.push(id);
+			// FIFO eviction: when the cap is exceeded, drop the oldest entry.
+			while (seen.length > SEEN_STEER_ID_CAP) {
+				const oldest = seen.shift();
+				if (oldest !== undefined) set.delete(oldest);
+			}
+			return true;
+		},
+		size: () => set.size,
+	};
+}
+
+/**
+ * Validate a single steering-file entry before forwarding it to
+ * `pi.sendMessage`. FIX-02: reject oversized payloads, excessive newlines,
+ * or control characters that could be used to confuse the worker UI.
+ */
+export function sanitizeSteerMessage(entry: SteerEntry): SteerSanitizeResult {
+	const message = entry.message;
+	if (typeof message !== "string" || message.length === 0) {
+		return { valid: false, reason: "missing-or-empty-message" };
+	}
+	if (message.length > MAX_STEER_MESSAGE_LENGTH) {
+		return { valid: false, reason: `message-too-long:${message.length}` };
+	}
+	const newlineCount = (message.match(/\n/g) ?? []).length;
+	if (newlineCount > MAX_STEER_MESSAGE_NEWLINES) {
+		return { valid: false, reason: `too-many-newlines:${newlineCount}` };
+	}
+	if (STEER_CONTROL_CHAR_PATTERN.test(message)) {
+		return { valid: false, reason: "contains-control-characters" };
+	}
+	return { valid: true, message };
+}
+
+// ── FIX-03: Steering file path containment validation ─────────────────────
+// The steering file path is inherited from the parent via env, so we
+// defensively re-validate it before the first read to catch symlink
+// redirection or paths that escape the session's artifacts root.
+export interface SteeringFileValidation {
+	valid: boolean;
+	reason?: string;
+	resolvedPath?: string;
+}
+
+/**
+ * Validate `PI_CREW_STEERING_FILE` before first read. FIX-03:
+ *   1. `lstatSync` rejects a symlink at the steering file itself.
+ *   2. `resolveRealContainedPath` walks the ancestor chain with O_NOFOLLOW
+ *      to reject any symlinked parent (e.g. a redirected `artifactsRoot`)
+ *      and to verify the resolved path stays inside the derived artifacts
+ *      root (`<artifactsRoot>/steering/<taskId>.jsonl` → `<artifactsRoot>`).
+ *
+ * Returns `{ valid: false }` with a reason on any violation. Callers must
+ * log + skip steering on failure rather than abort the worker.
+ */
+export function validateSteeringFile(steeringFile: string): SteeringFileValidation {
+	try {
+		const lst = fs.lstatSync(steeringFile);
+		if (lst.isSymbolicLink()) {
+			return { valid: false, reason: "steering-file-is-symlink" };
+		}
+	} catch (error) {
+		const errCode = (error as NodeJS.ErrnoException).code;
+		if (errCode && errCode !== "ENOENT") {
+			return { valid: false, reason: `lstat-failed:${errCode}` };
+		}
+	}
+	// Layout invariant from task-runner.ts: `steeringFile` is built as
+	// `<artifactsRoot>/steering/<taskId>.jsonl`. We don't trust the caller
+	// to pass `artifactsRoot`, so derive it as 2 levels up. `resolveRealContainedPath`
+	// then enforces both containment AND ancestor-symlink safety in one shot.
+	const artifactsRoot = path.resolve(steeringFile, "..", "..");
+	try {
+		const resolved = resolveRealContainedPath(artifactsRoot, steeringFile);
+		return { valid: true, resolvedPath: resolved };
+	} catch (error) {
+		return {
+			valid: false,
+			reason: `path-validation-failed:${error instanceof Error ? error.message : String(error)}`,
+		};
+	}
+}
+
+function readBooleanEnv(name: string): boolean | undefined {
+	const value = process.env[name];
+	if (value === undefined) return undefined;
+	const normalized = value.trim().toLowerCase();
+	if (normalized === "1" || normalized === "true" || normalized === "yes") return true;
+	if (normalized === "0" || normalized === "false" || normalized === "no") return false;
+	// Ambiguous value — treat as undefined so callers apply their default.
+	return undefined;
+}
+
+function readBooleanEnvAny(...names: string[]): boolean | undefined {
+	for (const name of names) {
+		const value = readBooleanEnv(name);
+		if (value !== undefined) return value;
+	}
+	return undefined;
+}
+
+function findSectionEnd(prompt: string, startIndex: number, nextHeaders: string[]): number {
+	let endIndex = prompt.length;
+	for (const header of nextHeaders) {
+		const index = prompt.indexOf(header, startIndex);
+		if (index !== -1 && index < endIndex) endIndex = index;
+	}
+	return endIndex;
+}
+
+export function stripProjectContext(prompt: string): string {
+	const startIndex = prompt.indexOf(PROJECT_CONTEXT_HEADER);
+	if (startIndex === -1) return prompt;
+	const endIndex = findSectionEnd(prompt, startIndex + PROJECT_CONTEXT_HEADER.length, [SKILLS_HEADER, DATE_HEADER]);
+	return `${prompt.slice(0, startIndex)}${prompt.slice(endIndex)}`;
+}
+
+export function stripInheritedSkills(prompt: string): string {
+	const startIndex = prompt.indexOf(SKILLS_HEADER);
+	if (startIndex === -1) return prompt;
+	const endIndex = findSectionEnd(prompt, startIndex + SKILLS_HEADER.length, [DATE_HEADER]);
+	return `${prompt.slice(0, startIndex)}${prompt.slice(endIndex)}`;
+}
+
+export function rewriteTeamWorkerPrompt(prompt: string, options: { inheritProjectContext: boolean; inheritSkills: boolean }): string {
+	let rewritten = prompt;
+	if (!options.inheritProjectContext) rewritten = stripProjectContext(rewritten);
+	if (!options.inheritSkills) rewritten = stripInheritedSkills(rewritten);
+	return rewritten;
+}
+
+export default function registerPiTeamsPromptRuntime(pi: ExtensionAPI): void {
+	// ── FIX-S1: cross-channel steer dedup state ────────────────────────────
+	// Both the broker push (mailbox.message → onSteer callback below) and the
+	// file poll (pollSteering below) are wired to sendMessage. The broker
+	// writes to BOTH the mailbox observer (live fanout) AND the steering
+	// JSONL (durable fallback) so a connected worker will see the same steer
+	// twice: once via the broker, once via the next poll tick. We dedup at the
+	// consumer by tracking steer ids in a bounded FIFO set shared by both
+	// delivery paths. Without id-bearing entries (legacy producer), pollSteering
+	// is the only delivery channel and dedup is a no-op for id-less entries.
+	const seenSteers = createSeenSteerIdSet();
+	// ── Feature 1: maxTokens cap ──────────────────────────────────────────
+	// Cap output tokens per API call for background workers. Reads
+	// PI_CREW_MAX_OUTPUT_TOKENS env (set by pi-args.ts from agent.maxTokens).
+	const maxTokensEnv = process.env[PI_CREW_MAX_OUTPUT_ENV];
+	const maxTokensCap = maxTokensEnv ? Number.parseInt(maxTokensEnv, 10) : undefined;
+	if (maxTokensCap && maxTokensCap > 0) {
+		pi.on("before_provider_request", (event) => {
+			const payload = event.payload as Record<string, unknown> | undefined;
+			if (!payload || typeof payload !== "object") return;
+			// Cap both OpenAI-style max_tokens and Anthropic-style max_tokens
+			if (typeof payload.max_tokens === "number" && payload.max_tokens > maxTokensCap) {
+				payload.max_tokens = maxTokensCap;
+			}
+			// Also cap newer field names used by some providers
+			if (typeof payload.max_completion_tokens === "number" && payload.max_completion_tokens > maxTokensCap) {
+				payload.max_completion_tokens = maxTokensCap;
+			}
+			if (typeof payload.max_output_tokens === "number" && payload.max_output_tokens > maxTokensCap) {
+				payload.max_output_tokens = maxTokensCap;
+			}
+			const generationConfig = payload.generationConfig as Record<string, unknown> | undefined;
+			if (
+				generationConfig &&
+				typeof generationConfig.max_output_tokens === "number" &&
+				generationConfig.max_output_tokens > maxTokensCap
+			) {
+				generationConfig.max_output_tokens = maxTokensCap;
+			}
+		});
+	}
+
+	// ── Feature 2: real-time steering ──────────────────────────────────────
+	// Poll the steering JSONL file for new steer messages. The parent (team
+	// tool) writes steers here in real-time; this reader injects them into
+	// the active session via pi.sendMessage with deliverAs:"steer".
+	const steeringFile = process.env[PI_CREW_STEERING_FILE_ENV];
+	if (steeringFile) {
+		// FIX-03: validate the steering file path once before first read.
+		const validation = validateSteeringFile(steeringFile);
+		if (!validation.valid) {
+			logInternalError(
+				"prompt-runtime.steering-file-rejected",
+				new Error(validation.reason ?? "steering-file-validation-failed"),
+				`path=${steeringFile}`,
+				"warn",
+			);
+		} else {
+			const safeSteeringFile = validation.resolvedPath ?? steeringFile;
+			let lastOffset = 0;
+			const pollSteering = (): void => {
+				try {
+					const stat = fs.statSync(safeSteeringFile, { throwIfNoEntry: false });
+					if (!stat || stat.size <= lastOffset) return;
+					const fd = fs.openSync(safeSteeringFile, "r");
+					try {
+						const buf = Buffer.alloc(stat.size - lastOffset);
+						fs.readSync(fd, buf, 0, buf.length, lastOffset);
+						lastOffset = stat.size;
+						const lines = buf.toString("utf8").split("\n").filter(Boolean);
+						for (const line of lines) {
+							try {
+								const entry = JSON.parse(line) as SteerEntry;
+								if (entry.type !== "steer") continue;
+								// FIX-S1: cross-channel dedup. The broker writes the
+								// same steer to both the mailbox (live fanout via the
+								// onSteer callback below) and this JSONL file. A
+								// connected worker receives it via the broker first
+								// and via this poll second; the seen-id set ensures
+								// only the first arrival reaches pi.sendMessage.
+								const entryId =
+									typeof (entry as { id?: unknown }).id === "string" ? (entry as { id: string }).id : undefined;
+								if (!seenSteers.markOrSkip(entryId)) continue;
+								// FIX-02: sanitize each steer entry before forwarding
+								// to pi.sendMessage. Reject oversized payloads,
+								// excessive newlines, and control characters.
+								const sanitized = sanitizeSteerMessage(entry);
+								if (!sanitized.valid || sanitized.message === undefined) {
+									logInternalError(
+										"prompt-runtime.steer-rejected",
+										new Error(sanitized.reason ?? "steer-sanitization-failed"),
+										`line-preview=${line.slice(0, 64)}`,
+										"warn",
+									);
+									continue;
+								}
+								pi.sendMessage(
+									{ customType: "crew-steer", content: sanitized.message, display: false },
+									{ deliverAs: "steer" },
+								);
+							} catch {
+								// Malformed line — skip
+							}
+						}
+					} finally {
+						try {
+							fs.closeSync(fd);
+						} catch {
+							/* already closed */
+						}
+					}
+				} catch {
+					// File doesn't exist yet or read error — will retry next tick
+				}
+			};
+			const timer = setInterval(pollSteering, 500);
+			timer.unref?.();
+		}
+	}
+
+	// ── Feature 2b: broker push steering (opt-in, layered on the file poll) ──
+	// When the parent injected broker credentials, connect a broker client and
+	// deliver pushed steers with the SAME sanitize + pi.sendMessage path as the
+	// file poll above. The file poll remains the durable fallback; a broker
+	// connect failure is invisible to the worker.
+	const brokerHandle = startChildBrokerClient({
+		onSteer: (rawMessage, id) => {
+			// FIX-S1: cross-channel dedup. The broker persists every steer to
+			// the steering JSONL (durable) AND pushes it via the mailbox
+			// observer (live). This callback sees the live push; the
+			// pollSteering path above will see the same steer on its next
+			// tick. Keying on `id` (a stable identifier emitted by the broker)
+			// guarantees a single pi.sendMessage per steer.
+			if (!seenSteers.markOrSkip(id)) return;
+			const sanitized = sanitizeSteerMessage({ type: "steer", message: rawMessage });
+			if (!sanitized.valid || sanitized.message === undefined) {
+				logInternalError(
+					"prompt-runtime.broker-steer-rejected",
+					new Error(sanitized.reason ?? "steer-sanitization-failed"),
+					undefined,
+					"warn",
+				);
+				return;
+			}
+			pi.sendMessage({ customType: "crew-steer", content: sanitized.message, display: false }, { deliverAs: "steer" });
+		},
+	});
+	// Close the broker connection on session shutdown to avoid leaking the
+	// persistent socket / reconnect timer. Fire-and-forget: the socket is
+	// already .unref()'d so it never blocks event-loop exit; errors are
+	// swallowed because teardown failures during shutdown are harmless.
+	pi.on("session_shutdown", () => {
+		void brokerHandle.close().catch(() => {});
+	});
+
+	// ── Prompt rewriting (existing) ────────────────────────────────────────
+	pi.on("before_agent_start", (event) => {
+		const inheritProjectContext = readBooleanEnvAny(PI_CREW_INHERIT_PROJECT_CONTEXT_ENV, PI_TEAMS_INHERIT_PROJECT_CONTEXT_ENV);
+		const inheritSkills = readBooleanEnvAny(PI_CREW_INHERIT_SKILLS_ENV, PI_TEAMS_INHERIT_SKILLS_ENV);
+		if (inheritProjectContext === undefined && inheritSkills === undefined) return;
+		const rewritten = rewriteTeamWorkerPrompt(event.systemPrompt, {
+			inheritProjectContext: inheritProjectContext ?? true,
+			inheritSkills: inheritSkills ?? true,
+		});
+		if (rewritten === event.systemPrompt) return;
+		return { systemPrompt: rewritten };
+	});
+}