pi-subagent-feather 0.0.0-stage → 0.2.0

package/src/collector.ts

@@ -0,0 +1,698 @@
+// Isolated, bounded JSON event collector for the pi-subagent-lite fork.
+//
+// Design goals:
+//  - NO Pi dependencies (this module must be importable stand-alone).
+//  - Never retain raw events, message objects, tool args, images or thinking.
+//  - Bound every retained string (bytes + lines) while appending, never store
+//    an unbounded suffix and slice afterwards.
+//  - Correctly decode chunked UTF-8 (stdout NDJSON and stderr), including
+//    multibyte characters split across chunks and buffers ending without a
+//    trailing newline.
+//  - Recover from oversized NDJSON lines without poisoning the next line.
+
+import { StringDecoder } from "node:string_decoder";
+
+/** Numeric byte/line budgets enforced by {@link ChildCollector}. */
+export interface CollectorLimits {
+	resultBytes: number;
+	resultLines: number;
+	progressBytes: number;
+	stderrBytes: number;
+	eventBytes: number;
+}
+
+/** Default budgets, identical in spirit to upstream but chosen for a safe fork. */
+/** Where an oversized completed answer keeps content: head (default) or tail. */
+export type KeepMode = "head" | "tail";
+
+export const DEFAULT_LIMITS: CollectorLimits = {
+	resultBytes: 8192,
+	resultLines: 200,
+	progressBytes: 2048,
+	stderrBytes: 4096,
+	eventBytes: 1024 * 1024,
+};
+
+/** Accumulated token/cost accounting exposed by {@link ChildCollector.usage}. */
+export interface CollectorUsage {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	totalTokens: number;
+	cost: {
+		input: number;
+		output: number;
+		cacheRead: number;
+		cacheWrite: number;
+		total: number;
+	};
+	turns: number;
+}
+
+function emptyUsage(): CollectorUsage {
+	return {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens: 0,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		turns: 0,
+	};
+}
+
+function num(value: unknown): number {
+	return typeof value === "number" && Number.isFinite(value) ? value : 0;
+}
+
+function utf8Length(text: string): number {
+	return Buffer.byteLength(text, "utf8");
+}
+
+/**
+ * Longest prefix of `text` whose UTF-8 byte length does not exceed `maxBytes`.
+ * Never splits a surrogate pair.
+ */
+function utf8Prefix(text: string, maxBytes: number): string {
+	if (maxBytes <= 0) return "";
+	let bytes = 0;
+	let i = 0;
+	while (i < text.length) {
+		const cp = text.codePointAt(i) as number;
+		const size = cp <= 0x7f ? 1 : cp <= 0x7ff ? 2 : cp <= 0xffff ? 3 : 4;
+		if (bytes + size > maxBytes) break;
+		bytes += size;
+		i += cp > 0xffff ? 2 : 1;
+	}
+	return text.slice(0, i);
+}
+
+/**
+ * Longest suffix of `text` whose UTF-8 byte length does not exceed `maxBytes`.
+ * Never splits a surrogate pair.
+ */
+function utf8Tail(text: string, maxBytes: number): string {
+	if (maxBytes <= 0) return "";
+	let bytes = 0;
+	let i = text.length;
+	while (i > 0) {
+		let start = i - 1;
+		let cp = text.charCodeAt(start);
+		if (cp >= 0xdc00 && cp <= 0xdfff && start > 0) {
+			start -= 1;
+			cp = text.codePointAt(start) as number;
+		}
+		const size = cp <= 0x7f ? 1 : cp <= 0x7ff ? 2 : cp <= 0xffff ? 3 : 4;
+		if (bytes + size > maxBytes) break;
+		bytes += size;
+		i = start;
+	}
+	return text.slice(i);
+}
+
+/** Byte caps for retained scalar metadata strings. */
+const SCALAR_LIMITS = {
+	toolName: 128,
+	sessionId: 256,
+	model: 512,
+	stopReason: 64,
+	errorMessage: 2048,
+} as const;
+
+/** Scalar token/cost snapshot copied out of an event (no event reference kept). */
+interface TurnUsage {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	totalTokens: number;
+	cost: { input: number; output: number; cacheRead: number; cacheWrite: number; total: number };
+}
+
+/** Force a fresh V8 string so a bounded slice cannot retain a huge backing store. */
+function detach(text: string): string {
+	return Buffer.from(text, "utf8").toString("utf8");
+}
+
+/** Detached UTF-8 prefix of `value`, capped at `maxBytes`. */
+function capScalar(value: unknown, maxBytes: number): string | undefined {
+	if (value === undefined || value === null) return undefined;
+	const text = typeof value === "string" ? value : String(value);
+	if (text.length === 0) return undefined;
+	return detach(utf8Prefix(text, maxBytes));
+}
+
+/** Copy only finite numeric token/cost fields into a fresh small object. */
+function copyUsage(usage: unknown): TurnUsage | undefined {
+	if (!usage || typeof usage !== "object") return undefined;
+	const u = usage as Record<string, unknown>;
+	const rawCost = u.cost;
+	const c = rawCost && typeof rawCost === "object" ? (rawCost as Record<string, unknown>) : {};
+	return {
+		input: num(u.input),
+		output: num(u.output),
+		cacheRead: num(u.cacheRead),
+		cacheWrite: num(u.cacheWrite),
+		totalTokens: num(u.totalTokens),
+		cost: {
+			input: num(c.input),
+			output: num(c.output),
+			cacheRead: num(c.cacheRead),
+			cacheWrite: num(c.cacheWrite),
+			total: num(c.total),
+		},
+	};
+}
+
+/**
+ * Incremental, budget-bounded text accumulator.
+ *
+ * Appends are never stored unbounded: once `maxBytes` or `maxLines` is reached
+ * the accumulator sets {@link truncated} and silently ignores the rest.
+ */
+class BoundedText {
+	private readonly chunks: string[] = [];
+	private bytes = 0;
+	private newlines = 0;
+	private hasContent = false;
+	truncated = false;
+
+	private readonly maxBytes: number;
+	private readonly maxLines: number;
+	private readonly separator: string;
+
+	constructor(maxBytes: number, maxLines: number, separator: string) {
+		this.maxBytes = maxBytes;
+		this.maxLines = maxLines;
+		this.separator = separator;
+	}
+
+	append(input: string): void {
+		if (this.truncated) return;
+		if (input.length === 0) return;
+		if (this.maxBytes <= 0 || this.maxLines <= 0) {
+			this.truncated = true;
+			return;
+		}
+
+		if (this.hasContent && this.separator.length > 0) {
+			if (this.bytes + 1 > this.maxBytes || this.newlines >= this.maxLines - 1) {
+				this.truncated = true;
+				return;
+			}
+			this.chunks.push(this.separator);
+			this.bytes += 1;
+			this.newlines += 1;
+		}
+
+		let i = 0;
+		while (i < input.length) {
+			const nl = input.indexOf("\n", i);
+			const end = nl === -1 ? input.length : nl;
+			const segment = input.slice(i, end);
+			const room = this.maxBytes - this.bytes;
+
+			// Every code unit is at least one UTF-8 byte, so a segment longer
+			// than the remaining room is guaranteed over budget: avoid scanning
+			// giant segments in full.
+			if (segment.length > room) {
+				if (room > 0) {
+					const partial = utf8Prefix(segment, room);
+					if (partial.length > 0) {
+						this.chunks.push(detach(partial));
+						this.bytes += utf8Length(partial);
+						this.hasContent = true;
+					}
+				}
+				this.truncated = true;
+				return;
+			}
+
+			const segBytes = utf8Length(segment);
+			if (this.bytes + segBytes > this.maxBytes) {
+				const partial = utf8Prefix(segment, room);
+				if (partial.length > 0) {
+					this.chunks.push(detach(partial));
+					this.bytes += utf8Length(partial);
+					this.hasContent = true;
+				}
+				this.truncated = true;
+				return;
+			}
+
+			if (segment.length > 0) {
+				this.chunks.push(detach(segment));
+				this.bytes += segBytes;
+				this.hasContent = true;
+			}
+
+			if (nl === -1) return;
+
+			if (this.newlines >= this.maxLines - 1) {
+				this.truncated = true;
+				return;
+			}
+			if (this.bytes + 1 > this.maxBytes) {
+				this.truncated = true;
+				return;
+			}
+			this.chunks.push("\n");
+			this.bytes += 1;
+			this.newlines += 1;
+			i = nl + 1;
+		}
+	}
+
+	value(): string {
+		return this.chunks.length === 0 ? "" : this.chunks.join("");
+	}
+}
+
+/** Anything BoundedText/TailText can append into. */
+interface TextSink {
+	append(input: string): void;
+}
+
+/**
+ * Raw-substring markers for the event types this collector reads. Pi writes compact
+ * JSON (`"type":"x"`, no spaces), so a line without any needle cannot be one of them;
+ * such lines skip JSON.parse entirely. Key order does not matter (scan, not prefix match).
+ */
+const EVENT_TYPE_NEEDLES = [
+	"session",
+	"message_start",
+	"message_update",
+	"message_end",
+	"tool_execution_start",
+	"tool_execution_end",
+].map(type => `"type":"${type}"`);
+
+/**
+ * Bounded accumulator that keeps the most recent content (tail) instead of the head.
+ * Compaction re-slices only when the budget is exceeded; at the default few-KiB budgets
+ * the amortized cost per appended byte is O(1).
+ */
+class TailText {
+	private readonly maxBytes: number;
+	private readonly maxLines: number;
+	private readonly separator: string;
+	private buf = "";
+	private segments = 0;
+	truncated = false;
+
+	constructor(maxBytes: number, maxLines: number, separator: string) {
+		this.maxBytes = maxBytes;
+		this.maxLines = maxLines;
+		this.separator = separator;
+	}
+
+	append(input: string): void {
+		if (input.length === 0) return;
+		if (this.maxBytes <= 0 || this.maxLines <= 0) {
+			this.truncated = true;
+			return;
+		}
+		if (this.segments > 0 && this.separator.length > 0) this.buf += this.separator;
+		this.segments += 1;
+		this.buf += input;
+		this.compact();
+	}
+
+	private compact(): void {
+		if (utf8Length(this.buf) > this.maxBytes) {
+			this.buf = detach(utf8Tail(this.buf, this.maxBytes));
+			this.truncated = true;
+		}
+		let newlines = 0;
+		for (let i = this.buf.indexOf("\n"); i !== -1; i = this.buf.indexOf("\n", i + 1)) newlines += 1;
+		if (newlines >= this.maxLines) {
+			const lines = this.buf.split("\n");
+			this.buf = detach(lines.slice(lines.length - this.maxLines).join("\n"));
+			this.truncated = true;
+		}
+	}
+
+	value(): string {
+		return this.buf;
+	}
+}
+
+function normalizeLimits(overrides?: Partial<CollectorLimits>): CollectorLimits {
+	const out: CollectorLimits = { ...DEFAULT_LIMITS };
+	if (overrides) {
+		for (const key of Object.keys(DEFAULT_LIMITS) as (keyof CollectorLimits)[]) {
+			const value = overrides[key];
+			if (typeof value === "number" && Number.isFinite(value) && value >= 0) {
+				out[key] = Math.floor(value);
+			}
+		}
+	}
+	return out;
+}
+
+/** Best-effort short tool *name* only, e.g. `mcp__server__search` -> `search`. No tool id is tracked or matched. */
+function shortToolName(name: unknown): string {
+	const raw = typeof name === "string" && name.length > 0 ? name : "unknown";
+	const slash = raw.lastIndexOf("/");
+	const base = slash >= 0 ? raw.slice(slash + 1) : raw;
+	const dunder = base.lastIndexOf("__");
+	return dunder >= 0 ? base.slice(dunder + 2) : base;
+}
+
+function appendTextContent(acc: TextSink, content: unknown): void {
+	if (typeof content === "string") {
+		acc.append(content);
+		return;
+	}
+	if (!Array.isArray(content)) return;
+	for (const part of content) {
+		if (!part || typeof part !== "object") continue;
+		const record = part as { type?: unknown; text?: unknown };
+		if (record.type === "text" && typeof record.text === "string") {
+			acc.append(record.text);
+		}
+		// Thinking, tool-args, tool results and images are intentionally ignored.
+	}
+}
+
+function addUsage(target: CollectorUsage, usage: TurnUsage): void {
+	target.input += usage.input;
+	target.output += usage.output;
+	target.cacheRead += usage.cacheRead;
+	target.cacheWrite += usage.cacheWrite;
+	if (usage.totalTokens !== 0) target.totalTokens = usage.totalTokens;
+	target.cost.input += usage.cost.input;
+	target.cost.output += usage.cost.output;
+	target.cost.cacheRead += usage.cost.cacheRead;
+	target.cost.cacheWrite += usage.cost.cacheWrite;
+	target.cost.total += usage.cost.total;
+}
+
+/**
+ * Bounded, isolated collector for a single child Pi run's JSON event stream.
+ *
+ * Feed it either structured events via {@link accept} or raw child stdout via
+ * {@link pushStdout}. Only scalar metadata and bounded strings are retained.
+ */
+export class ChildCollector {
+	readonly limits: CollectorLimits;
+	/** Whether oversized completed answers keep their head or tail. */
+	readonly resultKeep: KeepMode;
+
+	/** Most recent completed non-empty assistant answer (bounded). */
+	finalText = "";
+	/** Current in-flight assistant answer (bounded). */
+	partialText = "";
+	/** Bounded UTF-8 tail of stderr. */
+	stderr = "";
+	errorMessage?: string;
+	model?: string;
+	sessionId?: string;
+	stopReason?: string;
+
+	readonly usage: CollectorUsage = emptyUsage();
+	toolCount = 0;
+	/** Best-effort current tool short name (no tool id retained). */
+	currentTool?: string;
+	/** Number of oversized NDJSON events dropped. */
+	droppedEvents = 0;
+
+	private partialBuffer: BoundedText;
+	private finalTruncated = false;
+	private currentTurnUsage: TurnUsage | undefined;
+	/** Set when an oversized event was dropped; may mean the answer is incomplete. */
+	private droppedTruncation = false;
+	private readonly stdoutDecoder: StringDecoder;
+	private readonly stderrDecoder: StringDecoder;
+	private buffer = "";
+	private bufferBytes = 0;
+	private dropping = false;
+
+	constructor(limits?: Partial<CollectorLimits>, options?: { resultKeep?: KeepMode }) {
+		this.limits = normalizeLimits(limits);
+		this.resultKeep = options?.resultKeep === "tail" ? "tail" : "head";
+		this.partialBuffer = new BoundedText(this.limits.resultBytes, this.limits.resultLines, "");
+		this.stdoutDecoder = new StringDecoder("utf8");
+		this.stderrDecoder = new StringDecoder("utf8");
+	}
+
+	// ---------------------------------------------------------------------
+	// Answer / progress
+	// ---------------------------------------------------------------------
+
+	/**
+	 * Prefers the current partial answer when present, otherwise the most
+	 * recent completed non-empty assistant text.
+	 */
+	answer(): string {
+		if (this.partialText.trim().length > 0) return this.partialText.trim();
+		if (this.finalText.trim().length > 0) return this.finalText.trim();
+		return "";
+	}
+
+	/** Bounded progress view (≤ progressBytes) of the current answer. */
+	progress(): string {
+		const source = this.partialText.trim().length > 0 ? this.partialText : this.finalText;
+		if (source.length === 0) return "";
+		return utf8Prefix(source, this.limits.progressBytes);
+	}
+
+	/** Truncation flag for the answer {@link answer} would currently return. */
+	get truncated(): boolean {
+		if (this.partialText.trim().length > 0) {
+			return this.partialBuffer.truncated || this.droppedTruncation;
+		}
+		return this.finalTruncated;
+	}
+
+	private resetPartial(): void {
+		this.partialBuffer = new BoundedText(this.limits.resultBytes, this.limits.resultLines, "");
+		this.partialText = "";
+		this.droppedTruncation = false;
+	}
+
+	// ---------------------------------------------------------------------
+	// Structured event ingestion
+	// ---------------------------------------------------------------------
+
+	accept(event: unknown): void {
+		if (!event || typeof event !== "object") return;
+		const e = event as Record<string, any>;
+
+		switch (e.type) {
+			case "session": {
+				if (e.id !== undefined && e.id !== null) this.sessionId = capScalar(e.id, SCALAR_LIMITS.sessionId);
+				return;
+			}
+			case "message_start": {
+				if (e.message?.role === "assistant") {
+					// Flush usage left unresolved by an oversized/missing prior message_end
+					// before the new turn's accounting begins (no double count if it was handled).
+					if (this.currentTurnUsage) {
+						addUsage(this.usage, this.currentTurnUsage);
+						this.currentTurnUsage = undefined;
+					}
+					this.resetPartial();
+					this.errorMessage = undefined;
+					this.stopReason = undefined;
+					this.usage.turns += 1;
+				}
+				return;
+			}
+			case "message_update": {
+				const streamedUsage = copyUsage(e.usage ?? e.message?.usage ?? e.assistantMessageEvent?.partial?.usage);
+				if (streamedUsage) this.currentTurnUsage = streamedUsage;
+				const delta = e.assistantMessageEvent;
+				if (delta?.type === "text_delta" && typeof delta.delta === "string") {
+					this.partialBuffer.append(delta.delta);
+					this.partialText = this.partialBuffer.value();
+				}
+				return;
+			}
+			case "message_end": {
+				this.handleMessageEnd(e.message);
+				return;
+			}
+			case "tool_execution_start": {
+				this.toolCount += 1;
+				this.currentTool = capScalar(shortToolName(e.toolName), SCALAR_LIMITS.toolName);
+				return;
+			}
+			case "tool_execution_end": {
+				const name = capScalar(shortToolName(e.toolName), SCALAR_LIMITS.toolName);
+				if (this.currentTool === name) this.currentTool = undefined;
+				return;
+			}
+			default:
+				return;
+		}
+	}
+
+	private handleMessageEnd(message: unknown): void {
+		if (!message || typeof message !== "object") return;
+		const m = message as Record<string, any>;
+		if (m.role !== "assistant") return;
+
+		const acc = this.resultKeep === "tail"
+			? new TailText(this.limits.resultBytes, this.limits.resultLines, "\n")
+			: new BoundedText(this.limits.resultBytes, this.limits.resultLines, "\n");
+		appendTextContent(acc, m.content);
+		const text = acc.value();
+		if (text.trim().length > 0) {
+			this.finalText = detach(text);
+			this.finalTruncated = acc.truncated;
+		}
+
+		// A completed assistant turn ends the in-flight partial.
+		this.resetPartial();
+
+		// Accumulate usage exactly once, preferring the finalized message usage.
+		const endUsage = copyUsage(m.usage);
+		if (endUsage) {
+			addUsage(this.usage, endUsage);
+			this.currentTurnUsage = undefined;
+		} else if (this.currentTurnUsage) {
+			addUsage(this.usage, this.currentTurnUsage);
+			this.currentTurnUsage = undefined;
+		}
+
+		const reportedModel = capScalar(m.model, SCALAR_LIMITS.model);
+		if (reportedModel) {
+			const provider = capScalar(m.provider, SCALAR_LIMITS.model);
+			this.model = capScalar(provider ? `${provider}/${reportedModel}` : reportedModel, SCALAR_LIMITS.model);
+		}
+		const stopReason = capScalar(m.stopReason, SCALAR_LIMITS.stopReason);
+		if (stopReason) this.stopReason = stopReason;
+		const errorMessage = capScalar(m.errorMessage, SCALAR_LIMITS.errorMessage);
+		if (errorMessage) this.errorMessage = errorMessage;
+	}
+
+	/** Flush any usage recorded for a turn that never received a message_end. */
+	finishUsage(): void {
+		if (this.currentTurnUsage) {
+			addUsage(this.usage, this.currentTurnUsage);
+			this.currentTurnUsage = undefined;
+		}
+	}
+
+	// ---------------------------------------------------------------------
+	// Raw stdout NDJSON ingestion
+	// ---------------------------------------------------------------------
+
+	pushStdout(chunk: Buffer): void {
+		const text = this.stdoutDecoder.write(chunk as any);
+		if (text.length > 0) this.appendStdout(text);
+	}
+
+	endStdout(): void {
+		const rest = this.stdoutDecoder.end();
+		if (rest.length > 0) this.appendStdout(rest);
+		if (this.dropping) {
+			this.dropping = false;
+			return;
+		}
+		if (this.buffer.length > 0) {
+			const line = this.buffer;
+			this.buffer = "";
+			this.bufferBytes = 0;
+			this.processLine(line);
+		}
+	}
+
+	private appendStdout(text: string): void {
+		let data = text;
+		const limit = this.limits.eventBytes;
+
+		while (data.length > 0) {
+			if (this.dropping) {
+				const nl = data.indexOf("\n");
+				if (nl === -1) return;
+				data = data.slice(nl + 1);
+				this.dropping = false;
+				continue;
+			}
+
+			const nl = data.indexOf("\n");
+			if (nl === -1) {
+				// Incomplete line: buffer it, capping while we buffer.
+				const room = limit - this.bufferBytes;
+				if (data.length > room || this.bufferBytes + utf8Length(data) > limit) {
+					this.buffer = "";
+					this.bufferBytes = 0;
+					this.dropping = true;
+					this.droppedEvents += 1;
+					this.droppedTruncation = true;
+				} else {
+					this.buffer += detach(data);
+					this.bufferBytes += utf8Length(data);
+				}
+				return;
+			}
+
+			const head = data.slice(0, nl);
+			data = data.slice(nl + 1);
+			const room = limit - this.bufferBytes;
+			let over = false;
+			if (head.length > room) {
+				over = true;
+			} else if (this.bufferBytes + utf8Length(head) > limit) {
+				over = true;
+			}
+			if (over) {
+				this.buffer = "";
+				this.bufferBytes = 0;
+				this.droppedEvents += 1;
+				this.droppedTruncation = true;
+				continue;
+			}
+			const line = this.buffer + head;
+			this.buffer = "";
+			this.bufferBytes = 0;
+			this.processLine(line);
+		}
+	}
+
+	private processLine(line: string): void {
+		// Cheap pre-filter before JSON.parse: skip lines whose event type we never read.
+		let interesting = false;
+		for (const needle of EVENT_TYPE_NEEDLES) {
+			if (line.includes(needle)) {
+				interesting = true;
+				break;
+			}
+		}
+		if (!interesting) return;
+		const trimmed = line.trim();
+		if (trimmed.length === 0) return;
+		let event: unknown;
+		try {
+			event = JSON.parse(trimmed);
+		} catch {
+			return; // malformed JSON is ignored, never retained
+		}
+		this.accept(event);
+	}
+
+	// ---------------------------------------------------------------------
+	// Raw stderr ingestion
+	// ---------------------------------------------------------------------
+
+	pushStderr(chunk: Buffer): void {
+		this.appendStderr(this.stderrDecoder.write(chunk as any));
+	}
+
+	/** Flush any pending bytes from the stderr decoder and re-cap the tail. */
+	endStderr(): void {
+		this.appendStderr(this.stderrDecoder.end());
+	}
+
+	private appendStderr(text: string): void {
+		if (text.length === 0) return;
+		this.stderr += text;
+		const max = this.limits.stderrBytes;
+		if (utf8Length(this.stderr) > max) {
+			this.stderr = detach(utf8Tail(this.stderr, max));
+		}
+	}
+}