pi-mono/packages/agent/scripts/generate-telemetry-docs.ts

import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { dirname, resolve } from "node:path";
import { pathToFileURL } from "node:url";
import type { TelemetryAttributeDefinition, TelemetrySchemaDefinition } from "@earendil-works/pi-telemetry";
import { AI_TELEMETRY_SCHEMA, HARNESS_TELEMETRY_SCHEMA } from "../src/harness/telemetry.ts";

function escapeCell(value: string): string {
	return value.replaceAll("|", "\\|").replaceAll("\n", " ");
}

function allowedValues(definition: TelemetryAttributeDefinition): string {
	if ("values" in definition && definition.values) return definition.values.map(String).join(", ");
	if ("elementValues" in definition && definition.elementValues) {
		return `elements: ${definition.elementValues.map(String).join(", ")}`;
	}
	return "";
}

function attributeNotes(definition: TelemetryAttributeDefinition): string {
	return [definition.cardinality ? `${definition.cardinality} cardinality` : "", definition.sensitive ? "sensitive" : ""]
		.filter(Boolean)
		.join(", ");
}

function parentDescription(parent: TelemetrySchemaDefinition["spans"][string]["parents"]): string {
	switch (parent.kind) {
		case "any":
			return "root or any caller span";
		case "root_or_external":
			return "root or caller-owned external span";
		case "spans":
			return parent.spans.map((span) => `\`${span}\``).join(", ");
	}
}

function renderSchema(schema: TelemetrySchemaDefinition, title: string): string[] {
	const lines = [`## ${title}`, "", `Schema version: ${schema.version}`, ""];
	for (const [spanName, span] of Object.entries(schema.spans)) {
		lines.push(`### \`${spanName}\``, "", span.description, "");
		lines.push(`- Parents: ${parentDescription(span.parents)}`);
		lines.push(`- Default status: \`${span.status.default}\``);
		lines.push(`- Error when: ${span.status.errorWhen}`, "");
		lines.push("#### Start attributes", "");
		lines.push("| Name | Type | Required | Values | Notes | Description |");
		lines.push("|---|---|---:|---|---|---|");
		for (const [name, definition] of Object.entries(span.startAttributes)) {
			lines.push(
				`| \`${name}\` | \`${definition.type}\` | ${definition.required ? "yes" : "no"} | ${escapeCell(allowedValues(definition))} | ${escapeCell(attributeNotes(definition))} | ${escapeCell(definition.description)} |`,
			);
		}
		if (Object.keys(span.startAttributes).length === 0) lines.push("| _none_ | | | | | |");
		lines.push("", "#### End attributes", "");
		lines.push("All end attributes are optional completion enrichment.", "");
		lines.push("| Name | Type | Values | Notes | Description |");
		lines.push("|---|---|---|---|---|");
		for (const [name, definition] of Object.entries(span.endAttributes)) {
			lines.push(
				`| \`${name}\` | \`${definition.type}\` | ${escapeCell(allowedValues(definition))} | ${escapeCell(attributeNotes(definition))} | ${escapeCell(definition.description)} |`,
			);
		}
		if (Object.keys(span.endAttributes).length === 0) lines.push("| _none_ | | | | |");
		lines.push("", "#### Events", "");
		const events = Object.entries(span.events ?? {});
		if (events.length === 0) {
			lines.push("No declared span events.", "");
			continue;
		}
		for (const [eventName, event] of events) {
			lines.push(`##### \`${eventName}\``, "", event.description, "");
			lines.push("| Name | Type | Required | Values | Notes | Description |");
			lines.push("|---|---|---:|---|---|---|");
			for (const [name, definition] of Object.entries(event.attributes)) {
				lines.push(
					`| \`${name}\` | \`${definition.type}\` | ${definition.required ? "yes" : "no"} | ${escapeCell(allowedValues(definition))} | ${escapeCell(attributeNotes(definition))} | ${escapeCell(definition.description)} |`,
				);
			}
			if (Object.keys(event.attributes).length === 0) lines.push("| _none_ | | | | | |");
			lines.push("");
		}
	}
	return lines;
}

export function renderAgentTelemetrySchemaMarkdown(): string {
	const lines = [
		"# Pi Agent Telemetry Schemas",
		"",
		"<!-- Generated by generate-telemetry-docs.ts. Do not edit manually. -->",
		"",
		...renderSchema(AI_TELEMETRY_SCHEMA, "AI request schema"),
		...renderSchema(HARNESS_TELEMETRY_SCHEMA, "Harness schema"),
	];
	return `${lines.join("\n").trimEnd()}\n`;
}

export function generateTelemetryDocs(outputPath: string, check: boolean): void {
	const expected = renderAgentTelemetrySchemaMarkdown();
	if (check) {
		let actual = "";
		try {
			actual = readFileSync(outputPath, "utf8");
		} catch {
			throw new Error(`${outputPath} is missing; run the telemetry documentation generator`);
		}
		if (actual !== expected) throw new Error(`${outputPath} is stale; run the telemetry documentation generator`);
		return;
	}
	mkdirSync(dirname(outputPath), { recursive: true });
	writeFileSync(outputPath, expected);
}

if (process.argv[1] && pathToFileURL(resolve(process.argv[1])).href === import.meta.url) {
	generateTelemetryDocs(
		resolve(import.meta.dirname, "../docs/telemetry-schema.md"),
		process.argv.includes("--check"),
	);
}