pi-diagnostics 0.0.0 → 0.3.0

package/src/expiry.ts

@@ -0,0 +1,267 @@
+import { appendFileSync } from "node:fs";
+import { readFile } from "node:fs/promises";
+import { join } from "node:path";
+import { getAgentDir, keyText, type ExtensionAPI, type ExtensionContext, type ModelSelectEvent, type ThinkingLevelSelectEvent } from "@earendil-works/pi-coding-agent";
+import { Box, Spacer, Text } from "@earendil-works/pi-tui";
+import {
+	DEFAULT_WARM_GRACE_MS,
+	REMINDER_TYPE,
+	cacheKeyScope,
+	type CacheWarmingMode,
+	type EntryLike,
+	type ModelLike,
+	type ReminderData,
+	type WarmingDecisionObservation,
+	buildReminder,
+	evaluate,
+	inspectPayload,
+	normalizeWarmingMode,
+	predictPiWarming,
+	prefixCacheMode,
+	reminderView,
+	resolveTtl,
+	scanBranch,
+} from "./expiry-core.ts";
+
+interface SettingsFile {
+	cacheWarming?: unknown;
+}
+
+/** Set PI_EXPIRY_REMINDER_DEBUG=/path/to/log to trace scheduling decisions. */
+function debug(message: string): void {
+	const file = process.env.PI_EXPIRY_REMINDER_DEBUG;
+	if (!file) return;
+	try {
+		appendFileSync(file, `${new Date().toISOString()} ${message}\n`);
+	} catch {
+		// Debug logging must never affect the session.
+	}
+}
+
+/** Display text for Pi's "toggle tool output" key (ctrl+o unless remapped). */
+function expandKeyText(): string {
+	try {
+		return keyText("app.tools.expand") || "ctrl+o";
+	} catch {
+		return "ctrl+o";
+	}
+}
+
+function modelKey(model: ModelLike | undefined): string | undefined {
+	return model ? `${model.provider}/${model.id}` : undefined;
+}
+
+function asModel(value: ExtensionContext["model"]): ModelLike | undefined {
+	if (!value) return undefined;
+	return value as ModelLike;
+}
+
+async function readSettings(): Promise<SettingsFile> {
+	try {
+		const text = await readFile(join(getAgentDir(), "settings.json"), "utf8");
+		const value: unknown = JSON.parse(text);
+		return typeof value === "object" && value !== null ? value as SettingsFile : {};
+	} catch {
+		return {};
+	}
+}
+
+function entryLikeBranch(ctx: ExtensionContext): EntryLike[] {
+	return ctx.sessionManager.getBranch() as unknown as EntryLike[];
+}
+
+function entryLikeAll(ctx: ExtensionContext): EntryLike[] {
+	return ctx.sessionManager.getEntries() as unknown as EntryLike[];
+}
+
+/**
+ * Show one durable transcript entry after expiry. `appendEntry`, unlike
+ * `sendMessage`, is never projected into provider context.
+ */
+export default function cacheExpiryReminder(pi: ExtensionAPI): void {
+	let settings: SettingsFile = {};
+	let warmingMode: CacheWarmingMode = "streaming";
+	let envRetention = process.env.PI_CACHE_RETENTION;
+	// Pi builds a fresh ExtensionContext object for every event, so never
+	// compare contexts by identity. Track the active session by id instead and
+	// keep the most recent context for timer callbacks.
+	let latestContext: ExtensionContext | undefined;
+	let activeSessionId: string | undefined;
+	let sessionNamespaceStartAt: number | undefined;
+	let timer: ReturnType<typeof setTimeout> | undefined;
+	let payloadHints = new Map<string, ReturnType<typeof inspectPayload>>();
+	let lastDecision: WarmingDecisionObservation | undefined;
+	let generation = 0;
+
+	pi.registerEntryRenderer<ReminderData>(REMINDER_TYPE, (entry, { expanded }, theme) => {
+		const data = entry.data;
+		if (!data) return undefined;
+		const view = reminderView(data);
+		const box = new Box(1, 1, (text) => theme.bg("customMessageBg", text));
+		// Pi's app.tools.expand (ctrl+o) toggles `expanded` on every custom entry,
+		// same as tool output. Hint the toggle in both states.
+		const hint = `  ${theme.fg("dim", expandKeyText())}${theme.fg("muted", expanded ? " collapse" : " details")}`;
+		box.addChild(new Text(`${theme.fg("warning", `⏱ ${view.title}`)}${hint}`, 0, 0));
+		if (expanded) {
+			const width = Math.max(...view.rows.map(([label]) => label.length));
+			box.addChild(new Spacer(1));
+			for (const [label, value] of view.rows) {
+				box.addChild(new Text(`${theme.fg("dim", label.padEnd(width))}  ${theme.fg("customMessageText", value)}`, 2, 0));
+			}
+		}
+		return box;
+	});
+
+	const clearTimer = (): void => {
+		if (timer) clearTimeout(timer);
+		timer = undefined;
+	};
+
+	const track = (ctx: ExtensionContext): boolean => {
+		const id = ctx.sessionManager.getSessionId();
+		if (activeSessionId !== undefined && id !== activeSessionId) return false;
+		latestContext = ctx;
+		return true;
+	};
+
+	const schedule = (delayMs: number): void => {
+		clearTimer();
+		const token = generation;
+		timer = setTimeout(() => {
+			timer = undefined;
+			if (token !== generation || !latestContext) return;
+			void check(latestContext).catch((error: unknown) => {
+				// Stale context after reload/session replacement; next session_start re-arms.
+				debug(`error: ${error instanceof Error ? error.stack ?? error.message : String(error)}`);
+			});
+		}, Math.max(100, delayMs));
+		timer.unref?.();
+	};
+
+	const check = async (ctx: ExtensionContext): Promise<void> => {
+		if (ctx.sessionManager.getSessionId() !== activeSessionId) return debug("skip: inactive session");
+		settings = await readSettings();
+		warmingMode = normalizeWarmingMode(settings.cacheWarming);
+		const model = asModel(ctx.model);
+		if (!model) return debug("skip: no model");
+		const hints = payloadHints.get(modelKey(model)!) ?? {};
+		const namespaceStartAt = cacheKeyScope(model, hints) === "session" ? sessionNamespaceStartAt : undefined;
+		const scan = scanBranch(entryLikeBranch(ctx), {
+			allEntries: entryLikeAll(ctx),
+			mode: prefixCacheMode(model, hints),
+			namespaceStartAt,
+		});
+		if (!scan.lastRequest) return debug("skip: no request on branch");
+		if (scan.inheritedFrom !== undefined) return debug("skip: fork copied request into new cache namespace");
+		if (scan.contextReset) return debug("skip: context reset after last request");
+
+		// Do not report an old provider cache after model switching. New request
+		// will establish a new cache baseline, then this loop starts again.
+		if (modelKey(model) !== `${scan.lastRequest.provider}/${scan.lastRequest.model}`) {
+			return debug(`skip: model ${modelKey(model)} != last request ${scan.lastRequest.provider}/${scan.lastRequest.model}`);
+		}
+
+		const ttl = resolveTtl({
+			model,
+			hints,
+			lastUsage: scan.lastRequest.usage,
+			envRetention,
+		});
+		if (!ttl) return debug("skip: prompt caching disabled or unknown");
+
+		const warming = predictPiWarming(warmingMode, model, ctx.thinkingLevel, envRetention);
+		const result = evaluate({
+			scan,
+			model,
+			ttl,
+			warming,
+			lastDecision,
+			graceMs: DEFAULT_WARM_GRACE_MS,
+		});
+		if (!result) return debug("skip: no evaluation");
+
+		const now = Date.now();
+		if (now < result.expiresAt) {
+			debug(`wait: expires ${new Date(result.expiresAt).toISOString()} ttl=${ttl.ttlMs} (${ttl.source}); ${result.warmingSummary}`);
+			schedule(result.checkAt - now);
+			return;
+		}
+		if (scan.remindedForTouchAt === result.lastTouchAt) return debug("skip: already reminded");
+		if (!ctx.isIdle() || ctx.hasPendingMessages()) {
+			debug("defer: agent busy");
+			schedule(5_000);
+			return;
+		}
+		debug(`remind: expired ${new Date(result.expiresAt).toISOString()}`);
+
+		pi.appendEntry<ReminderData>(REMINDER_TYPE, buildReminder(model, ttl, result));
+	};
+
+	const reset = (): void => {
+		generation++;
+		clearTimer();
+	};
+
+	pi.on("session_start", async (event, ctx) => {
+		activeSessionId = ctx.sessionManager.getSessionId();
+		latestContext = ctx;
+		const header = ctx.sessionManager.getHeader?.();
+		sessionNamespaceStartAt = event.reason === "fork"
+			? (Date.parse(header?.timestamp ?? "") || Date.now())
+			: undefined;
+		settings = await readSettings();
+		warmingMode = normalizeWarmingMode(settings.cacheWarming);
+		envRetention = process.env.PI_CACHE_RETENTION;
+		reset();
+		// Existing sessions get checked immediately; fresh sessions have no request.
+		schedule(0);
+	});
+
+	pi.on("before_provider_request", (event, ctx) => {
+		const model = asModel(ctx.model);
+		if (model) payloadHints.set(modelKey(model)!, inspectPayload(event.payload));
+		return undefined;
+	});
+
+	pi.on("cache_warming_decision", (event) => {
+		lastDecision = {
+			at: Date.now(),
+			action: event.action,
+			expectedSavings: event.continuationProbability * event.missCost - event.warmCost,
+		};
+		return undefined;
+	});
+
+	pi.on("agent_start", (_event, ctx) => {
+		if (track(ctx)) reset();
+	});
+
+	pi.on("agent_settled", async (_event, ctx) => {
+		if (!track(ctx)) return;
+		// Global setting can change mid-session via /settings.
+		settings = await readSettings();
+		warmingMode = normalizeWarmingMode(settings.cacheWarming);
+		schedule(0);
+	});
+
+	pi.on("model_select", (event: ModelSelectEvent, ctx) => {
+		if (track(ctx)) reset();
+		payloadHints.delete(modelKey(asModel(event.previousModel))!);
+	});
+
+	pi.on("thinking_level_select", (_event: ThinkingLevelSelectEvent, ctx) => {
+		if (track(ctx)) reset();
+	});
+	pi.on("session_tree", (_event, ctx) => {
+		// New leaf may resend a different prefix; re-evaluate from the new branch.
+		if (!track(ctx)) return;
+		reset();
+		schedule(0);
+	});
+	pi.on("session_shutdown", () => {
+		reset();
+		latestContext = undefined;
+		activeSessionId = undefined;
+		sessionNamespaceStartAt = undefined;
+	});
+}

package/src/index.ts

@@ -0,0 +1,333 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { JsonValue } from "@earendil-works/pi-ai";
+import { Text, type Component } from "@earendil-works/pi-tui";
+import {
+  cacheDrop,
+  createSseTap,
+  formatComplement,
+  formatNotification,
+  DIAGNOSTIC_TYPE,
+  findBaseline,
+  injectDiagnostics,
+  isExpectedMiss,
+  kindFor,
+  parseConfig,
+  sniffBody,
+  summarize,
+  type DiagnosticsKind,
+  type ModelRef,
+  type RawResult,
+} from "./core.ts";
+import { createNativeNoticeSetting, detectNativeMiss, isNativeNoticeShown, type AssistantLike } from "./native.ts";
+
+/** A request we modified whose response diagnostics should be captured. */
+interface Pending {
+  kind: DiagnosticsKind;
+  modelId: string;
+  baselineId: string;
+  baselineModelId: string;
+  afterSummary: boolean;
+  baselinePromptTokens: number;
+}
+
+type Captured = RawResult & Pending;
+
+interface Notice {
+  /** Standalone text. */
+  message: string;
+  /**
+   * Follow-up text used while Pi's native cache-miss notice is visible for the same response.
+   * Present only when Pi's own criteria for that notice are met.
+   */
+  complement?: string;
+}
+
+type Theme = Parameters<Parameters<ExtensionAPI["registerEntryRenderer"]>[1]>[2];
+
+/**
+ * Chooses standalone vs. follow-up text at render time: Pi re-derives its native notice from the
+ * current `showCacheMissNotices` value on every transcript rebuild, so this must too.
+ */
+class NoticeView implements Component {
+  private text?: string;
+  private view?: Text;
+
+  constructor(
+    private readonly notice: Notice,
+    private readonly theme: Theme,
+    private readonly nativeVisible: () => boolean,
+  ) {}
+
+  render(width: number): string[] {
+    const text = this.notice.complement && this.nativeVisible() ? this.notice.complement : this.notice.message;
+    if (!this.view || text !== this.text) {
+      this.text = text;
+      this.view = new Text(this.theme.fg("dim", text), 1, 0);
+    }
+    return this.view.render(width);
+  }
+
+  invalidate(): void {
+    this.view = undefined;
+  }
+}
+
+const NOTICE_TYPE = "pi-diagnostics";
+const MAX_TRACKED = 64;
+
+function isObject(value: unknown): value is Record<string, any> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function modelRef(value: unknown): ModelRef | undefined {
+  if (!isObject(value)) return undefined;
+  const { provider, api, id } = value;
+  if (typeof provider !== "string" || typeof api !== "string" || typeof id !== "string") return undefined;
+  return { provider, api, id };
+}
+
+function trim<T>(collection: Map<string, T> | Set<T>): void {
+  while (collection.size > MAX_TRACKED) {
+    const first = collection.keys().next().value;
+    if (first === undefined) return;
+    collection.delete(first as never);
+  }
+}
+
+/**
+ * Turns on provider prompt-cache diagnostics for Anthropic Messages and OpenAI Responses.
+ *
+ * Request side: `before_provider_request` adds the provider's opt-in field, referencing the
+ * previous assistant response on the active branch.
+ * Response side: pi's adapters drop the diagnostics field, so a narrowly scoped fetch wrapper
+ * tees only requests this extension armed, reads the diagnostics-bearing SSE event, and
+ * `message_end` attaches the result to the assistant message's `diagnostics` array.
+ */
+export default function (pi: ExtensionAPI) {
+  const config = parseConfig(process.env);
+  const pending = new Set<Pending>();
+  const captured = new Map<string, Captured>();
+  let innerFetch: typeof globalThis.fetch | undefined;
+  let wrapper: typeof globalThis.fetch | undefined;
+  const nativeSetting = createNativeNoticeSetting();
+
+  pi.registerEntryRenderer<Notice>(NOTICE_TYPE, (entry, _options, theme) => {
+    const data = entry.data;
+    if (typeof data?.message !== "string") return undefined;
+    const notice: Notice = {
+      message: data.message,
+      ...(typeof data.complement === "string" ? { complement: data.complement } : {}),
+    };
+    return new NoticeView(notice, theme, () => nativeSetting.enabled());
+  });
+
+  const persistNotice = (notice: Notice, ctx: ExtensionContext) => {
+    // Custom entries survive transcript rebuilds without entering model context or starting a turn.
+    // TUI renders entry_appended itself; only RPC needs the separate UI notification. Pi's native
+    // notice is interactive-only, so RPC always gets the standalone text.
+    pi.appendEntry<Notice>(NOTICE_TYPE, notice);
+    if (ctx.hasUI && ctx.mode === "rpc") ctx.ui.notify(notice.message, "info");
+  };
+
+  const syncSettingScope = (ctx: ExtensionContext) => {
+    nativeSetting.setScope({ cwd: ctx.cwd, projectTrusted: ctx.isProjectTrusted?.() ?? false });
+  };
+
+  /** Whether Pi's native cache-miss notice criteria hold for the persisted assistant entry. */
+  const nativeMissFor = (messageEntryId: string, message: AssistantLike, ctx: ExtensionContext): boolean => {
+    try {
+      const entries = ctx.sessionManager.getEntries();
+      const index = entries.findIndex((entry) => entry.id === messageEntryId);
+      const before = index >= 0 ? entries.slice(0, index) : entries;
+      const miss = detectNativeMiss(before as never, message, (provider, model) =>
+        ctx.modelRegistry?.find(provider, model)?.cost.cacheRead);
+      return isNativeNoticeShown(miss);
+    } catch {
+      // Best-effort mirror of Pi internals; never lose the notice itself.
+      return false;
+    }
+  };
+
+  const takePending = (body: unknown): Pending | undefined => {
+    if (pending.size === 0) return undefined;
+    const probe = sniffBody(body);
+    if (!probe) return undefined;
+    for (const item of pending) {
+      if (item.kind === probe.kind && item.baselineId === probe.comparedTo) {
+        pending.delete(item);
+        return item;
+      }
+    }
+    return undefined;
+  };
+
+  const installFetch = () => {
+    if (wrapper && globalThis.fetch === wrapper) return;
+    const inner = globalThis.fetch;
+    innerFetch = inner;
+    wrapper = async (input, init) => {
+      const match = takePending(init?.body);
+      const response = await inner(input, init);
+      if (!match || !response.body) return response;
+      const contentType = response.headers.get("content-type") ?? "";
+      if (!contentType.includes("text/event-stream")) return response;
+
+      const body = response.body.pipeThrough(
+        createSseTap(match.kind, (raw) => {
+          captured.set(raw.responseId, { ...match, ...raw });
+          trim(captured);
+        }),
+      );
+      const tapped = new Response(body, {
+        status: response.status,
+        statusText: response.statusText,
+        headers: response.headers,
+      });
+      Object.defineProperty(tapped, "url", { value: response.url });
+      return tapped;
+    };
+    globalThis.fetch = wrapper;
+  };
+
+  const uninstallFetch = () => {
+    if (wrapper && globalThis.fetch === wrapper && innerFetch) globalThis.fetch = innerFetch;
+    wrapper = undefined;
+    innerFetch = undefined;
+  };
+
+  pi.on("session_start", (_event, ctx) => {
+    syncSettingScope(ctx);
+    if (config.enabled) installFetch();
+  });
+
+  pi.on("session_shutdown", () => {
+    uninstallFetch();
+    pending.clear();
+    captured.clear();
+  });
+
+  pi.on("before_provider_request", (event, ctx) => {
+    const model = modelRef(ctx.model);
+    const kind = kindFor(model, config);
+    if (!model || !kind) return undefined;
+
+    // Nested requests on another model (e.g. compaction overrides) are not ours to label.
+    const payload = event.payload;
+    if (model.api !== "azure-openai-responses" && isObject(payload) && typeof payload.model === "string" && payload.model !== model.id) {
+      return undefined;
+    }
+
+    const baseline = findBaseline(ctx.sessionManager.getBranch(), model);
+    const next = injectDiagnostics(payload, kind, baseline?.responseId);
+    if (!next) return undefined;
+
+    if (baseline) {
+      // Re-install if pi replaced global fetch after session start (e.g. HTTP settings change).
+      installFetch();
+      pending.add({
+        kind,
+        modelId: model.id,
+        baselineId: baseline.responseId,
+        baselineModelId: baseline.modelId,
+        afterSummary: baseline.afterSummary,
+        baselinePromptTokens: baseline.promptTokens,
+      });
+      trim(pending);
+    }
+    return next;
+  });
+
+  pi.on("message_end", (event) => {
+    const message = event.message;
+    if (message.role !== "assistant" || !message.responseId) return undefined;
+    const result = captured.get(message.responseId);
+    if (!result) return undefined;
+    captured.delete(message.responseId);
+
+    const drop = cacheDrop(result.baselinePromptTokens, message.usage.cacheRead);
+    const summary = summarize(result.kind, result.raw, { cacheDropped: drop !== undefined });
+    const expected = isExpectedMiss(
+      summary,
+      { modelId: result.baselineModelId, afterSummary: result.afterSummary },
+      result.modelId,
+    );
+    // `unavailable` alone says nothing about hit/miss; only surface it when cached tokens actually
+    // dropped versus the baseline. Model switches and compaction explain drops, so skip those.
+    // `expired` is defined by the drop, so it always carries it.
+    const explainedDrop = result.afterSummary || (result.baselineModelId !== "" && result.baselineModelId !== result.modelId);
+    const droppedTokens = summary.outcome === "expired" || (summary.outcome === "unavailable" && !explainedDrop)
+      ? drop
+      : undefined;
+
+    const diagnostic = {
+      type: DIAGNOSTIC_TYPE[result.kind],
+      timestamp: Date.now(),
+      details: {
+        outcome: summary.outcome,
+        ...(summary.reason ? { reason: summary.reason } : {}),
+        ...(summary.missedTokens !== undefined ? { missedTokens: summary.missedTokens } : {}),
+        comparedTo: result.baselineId,
+        ...(expected ? { expected: true } : {}),
+        ...(droppedTokens !== undefined ? { droppedTokens } : {}),
+        cacheRead: message.usage.cacheRead,
+        input: message.usage.input,
+        // Omitted when the response had no diagnostics field, so it stays distinct from `null`.
+        ...(result.raw === undefined ? {} : { raw: result.raw as JsonValue }),
+      },
+    };
+
+    return { message: { ...message, diagnostics: [...(message.diagnostics ?? []), diagnostic] } };
+  });
+
+  pi.on("turn_end", (event, ctx) => {
+    // message_end runs before Pi persists the assistant. Wait for turn_end so notices remain
+    // after their response in both the live transcript and restored session history.
+    const message = event.message;
+    if (!config.enabled || config.notify === "off" || message.role !== "assistant") return;
+    syncSettingScope(ctx);
+    let nativeMiss: boolean | undefined;
+    for (const diagnostic of message.diagnostics ?? []) {
+      const kind = diagnostic.type === DIAGNOSTIC_TYPE.anthropic ? "anthropic"
+        : diagnostic.type === DIAGNOSTIC_TYPE.openai ? "openai" : undefined;
+      const details = diagnostic.details;
+      if (!kind || !isObject(details)) continue;
+      const droppedTokens = typeof details.droppedTokens === "number" ? details.droppedTokens : undefined;
+      const summary = summarize(kind, details.raw, { cacheDropped: droppedTokens !== undefined });
+      const unexpectedMiss = (summary.outcome === "miss" && details.expected !== true) || droppedTokens !== undefined;
+      if (config.notify === "all" || unexpectedMiss) {
+        nativeMiss ??= nativeMissFor(event.messageEntryId, message, ctx);
+        persistNotice({
+          message: formatNotification(summary, droppedTokens),
+          ...(nativeMiss ? { complement: formatComplement(summary, droppedTokens) } : {}),
+        }, ctx);
+      }
+    }
+  });
+
+  pi.registerCommand("diagnostics", {
+    description: "Show provider prompt-cache diagnostics for recent turns",
+    handler: async (_args, ctx) => {
+      const types = new Set(Object.values(DIAGNOSTIC_TYPE));
+      const lines: string[] = [];
+      for (const entry of ctx.sessionManager.getBranch()) {
+        if (entry.type !== "message" || entry.message.role !== "assistant") continue;
+        const message = entry.message;
+        for (const diagnostic of message.diagnostics ?? []) {
+          if (!types.has(diagnostic.type as never)) continue;
+          const d = (diagnostic.details ?? {}) as Record<string, unknown>;
+          const reason = d.reason ? ` ${d.reason}` : "";
+          const expected = d.expected ? " (expected)" : "";
+          const dropped = d.outcome === "unavailable" && typeof d.droppedTokens === "number"
+            ? " · cache miss, diagnostics unavailable"
+            : "";
+          lines.push(`${message.model}: ${d.outcome}${reason}${expected}${dropped}`);
+        }
+      }
+      const model = modelRef(ctx.model);
+      const active = kindFor(model, config) ? "on" : "off";
+      const header = `Provider cache diagnostics: ${active} for ${model ? `${model.provider}/${model.id}` : "no model"} (providers: ${config.providers.join(",")})`;
+      const body = lines.length ? lines.slice(-15).join("\n") : "No diagnostics recorded on this branch yet.";
+      persistNotice({ message: `${header}\n${body}` }, ctx);
+    },
+  });
+}