@mlx-node/agent 0.0.12 → 0.0.15

package/src/provider/events.ts

@@ -0,0 +1,341 @@
+/**
+ * `TurnEmitter` — maps one native `ChatStreamEvent` turn onto pi's
+ * `AssistantMessageEvent` protocol.
+ *
+ * Correctness contract (spike-proven):
+ *   - pi's agent loop takes the final message from `stream.result()` and
+ *     reads `stopReason` off THAT message — the `done` event's `reason`
+ *     field is discarded. Both are still emitted protocol-correct.
+ *   - Aborted native streams end with NO final event, so the caller must
+ *     invoke {@link TurnEmitter.onAborted} to synthesize the terminal
+ *     AssistantMessage (stopReason 'aborted', accumulated deltas intact).
+ *   - Every method is throw-safe: the pi StreamFn contract does not allow
+ *     the emitter to throw, so internal failures route to {@link onError}.
+ */
+
+import { randomUUID } from 'node:crypto';
+
+import type {
+  Api,
+  AssistantMessage,
+  AssistantMessageEventStream,
+  Model,
+  TextContent,
+  ThinkingContent,
+  ToolCall,
+  Usage,
+} from '@earendil-works/pi-ai';
+import type { ChatStreamDelta, ChatStreamFinal, PerformanceMetrics, ToolCallResult } from '@mlx-node/lm';
+
+import { coerceErrorMessage } from './error-coercion.js';
+import { ReasoningTagBuffer } from './reasoning-tag-buffer.js';
+import { ToolCallTagBuffer } from './tool-call-buffer.js';
+
+/**
+ * All-zero usage. Shared with the stream adapter's TurnEmitter-independent
+ * failsafe terminal, so it must stay trivially non-throwing.
+ */
+export function emptyUsage(): Usage {
+  return {
+    input: 0,
+    output: 0,
+    cacheRead: 0,
+    cacheWrite: 0,
+    totalTokens: 0,
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+  };
+}
+
+/**
+ * pi auto-compaction thresholds on `usage.totalTokens` vs
+ * `model.contextWindow`, so these numbers are load-bearing:
+ * input excludes the cache-served prefix, totalTokens is the full
+ * prompt + completion. All costs are 0 for local inference.
+ */
+function usageFromFinal(final: ChatStreamFinal): Usage {
+  const cachedTokens = final.cachedTokens ?? 0;
+  return {
+    input: Math.max(0, final.promptTokens - cachedTokens),
+    output: final.numTokens,
+    cacheRead: cachedTokens,
+    cacheWrite: 0,
+    reasoning: final.reasoningTokens,
+    totalTokens: final.promptTokens + final.numTokens,
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+  };
+}
+
+/**
+ * Native `ToolCallResult` → pi `ToolCall`.
+ *
+ * Non-ok results (invalid_json / missing_name / parse_error) become a
+ * deliberately-invalid pi ToolCall so pi's own tool validation fails it
+ * and feeds an error tool result back to the model — pi owns that retry
+ * loop (spike-proven).
+ */
+function toPiToolCall(call: ToolCallResult): ToolCall {
+  if (call.status === 'ok') {
+    const args =
+      typeof call.arguments === 'object' && call.arguments !== null ? (call.arguments as Record<string, unknown>) : {};
+    return { type: 'toolCall', id: call.id, name: call.name, arguments: args };
+  }
+  return {
+    type: 'toolCall',
+    id: call.id,
+    name: call.name || 'malformed_tool_call',
+    arguments: { raw: call.rawContent, error: call.error },
+  };
+}
+
+export class TurnEmitter {
+  /**
+   * Correlation id for this turn. Minted here, stamped on the partial message
+   * as the custom `mlxTraceId` field (survives pi's JSONL round-trip like
+   * `mlxThinkingEnabled`), and surfaced so the stream adapter can join a
+   * `MetricsTrace` record to the persisted pi message. The pi entry id is not
+   * knowable in-turn (pi emits `message_end` before persisting), so this
+   * minted id is the durable join key.
+   */
+  readonly traceId: string;
+  private readonly stream: AssistantMessageEventStream;
+  private readonly partial: AssistantMessage;
+  private readonly textBuffer = new ToolCallTagBuffer();
+  private readonly thinkingBuffer = new ReasoningTagBuffer();
+  /**
+   * Leading whitespace-only text parked before any text block exists, so
+   * a `"\n\n"` emitted right before `<tool_call>` markup never ratifies a
+   * whitespace-only text content block (mirrors the server endpoints).
+   * Joined onto the first non-whitespace text; dropped at terminal time.
+   */
+  private pendingLeadingWhitespace = '';
+  private openBlock: TextContent | ThinkingContent | null = null;
+  private finished = false;
+
+  constructor(
+    stream: AssistantMessageEventStream,
+    model: Model<Api>,
+    private readonly onPerformance?: (message: AssistantMessage, performance: PerformanceMetrics) => void,
+    thinkingEnabled?: boolean,
+  ) {
+    this.stream = stream;
+    const traceId = randomUUID();
+    this.traceId = traceId;
+    this.partial = {
+      role: 'assistant',
+      content: [],
+      api: model.api,
+      provider: model.provider,
+      model: model.id,
+      usage: emptyUsage(),
+      stopReason: 'stop',
+      timestamp: Date.now(),
+    };
+    // Pi persists provider messages as JSON without a provider-specific
+    // metadata bag. Keep this small enumerable provenance field on the
+    // assistant object so a later full-history replay can distinguish a
+    // disabled-thinking empty channel from an enabled-thinking turn that
+    // simply emitted no reasoning. Unknown fields survive Pi's session
+    // JSONL round-trip and are ignored by other providers.
+    if (thinkingEnabled !== undefined) {
+      (this.partial as AssistantMessage & { mlxThinkingEnabled?: boolean }).mlxThinkingEnabled = thinkingEnabled;
+    }
+    // Same custom-field mechanism: stamp the trace id so a dashboard can join
+    // the durable MetricsTrace record back to this persisted pi message.
+    (this.partial as AssistantMessage & { mlxTraceId?: string }).mlxTraceId = traceId;
+    stream.push({ type: 'start', partial: this.partial });
+  }
+
+  onDelta(delta: ChatStreamDelta): void {
+    if (this.finished) return;
+    try {
+      if (delta.isReasoning === true) {
+        this.appendThinking(this.thinkingBuffer.push(delta.text));
+      } else {
+        // A reasoning suffix that only resembled a partial protocol tag is
+        // ordinary text. Release it before opening the visible-text block.
+        this.appendThinking(this.thinkingBuffer.flush());
+        // Text routes through the tag buffer: partial structural markup
+        // (`<tool_call>` etc.) must never leak into pi-visible text.
+        const { safeText, tagFound, cleanPrefix } = this.textBuffer.push(delta.text);
+        if (tagFound) {
+          this.appendVisibleText(cleanPrefix);
+        } else if (safeText) {
+          this.appendVisibleText(safeText);
+        }
+      }
+    } catch (err) {
+      this.onError(err);
+    }
+  }
+
+  onFinal(final: ChatStreamFinal): void {
+    if (this.finished) return;
+    try {
+      this.partial.usage = usageFromFinal(final);
+      if (final.finishReason === 'error') {
+        this.finishWithError('error', 'model reported finishReason=error');
+        return;
+      }
+
+      // Release held-back suffixes from both structural-tag buffers, then
+      // close the open block.
+      this.appendThinking(this.thinkingBuffer.flush());
+      this.appendVisibleText(this.textBuffer.flush());
+      this.closeOpenBlock();
+
+      let sawOkToolCall = false;
+      for (const call of final.toolCalls) {
+        if (call.status === 'ok') sawOkToolCall = true;
+        const toolCall = toPiToolCall(call);
+        this.partial.content.push(toolCall);
+        const contentIndex = this.partial.content.length - 1;
+        this.stream.push({ type: 'toolcall_start', contentIndex, partial: this.partial });
+        this.stream.push({
+          type: 'toolcall_delta',
+          contentIndex,
+          delta: JSON.stringify(toolCall.arguments),
+          partial: this.partial,
+        });
+        this.stream.push({ type: 'toolcall_end', contentIndex, toolCall, partial: this.partial });
+      }
+
+      const reason = sawOkToolCall ? 'toolUse' : final.finishReason === 'length' ? 'length' : 'stop';
+      this.partial.stopReason = reason;
+      if (final.performance !== undefined && this.onPerformance !== undefined) {
+        try {
+          this.onPerformance(this.partial, final.performance);
+        } catch {
+          // Footer telemetry is best-effort and must never break inference.
+        }
+      }
+      this.finished = true;
+      this.stream.push({ type: 'done', reason, message: this.partial });
+      this.stream.end();
+    } catch (err) {
+      this.onError(err);
+    }
+  }
+
+  /**
+   * Synthesize the terminal message for an aborted native stream (which
+   * ends with no final event). Mirrors pi's provider abort pattern:
+   * `{type:'error', reason:'aborted', error: <partial message>}` with all
+   * accumulated text/thinking preserved on the message.
+   */
+  onAborted(): void {
+    if (this.finished) return;
+    this.finishWithError('aborted', 'Request was aborted');
+  }
+
+  /**
+   * Terminal for internal/adapter failures. `err` is untrusted: coercion
+   * is fully guarded (shared `coerceErrorMessage`), so even a revoked
+   * Proxy or a poisoned `message` getter cannot throw out of here.
+   */
+  onError(err: unknown): void {
+    if (this.finished) return;
+    this.finishWithError('error', coerceErrorMessage(err));
+  }
+
+  /**
+   * Shared terminal path for every non-`done` ending (abort, native
+   * finishReason=error, internal failure). Recovers any held-back buffer
+   * residue and closes the open text/thinking block BEFORE emitting the
+   * terminal event — a stream must never end `text_start/text_delta/error`
+   * with no `text_end` (pi's reference providers balance all blocks
+   * before terminals).
+   */
+  private finishWithError(reason: 'aborted' | 'error', message: string): void {
+    this.finished = true;
+    try {
+      try {
+        this.appendThinking(this.thinkingBuffer.flush());
+        this.appendVisibleText(this.textBuffer.flush());
+      } catch {
+        // Preserving buffered residue is best-effort; the block close and
+        // terminal event below must still go out.
+      }
+      this.closeOpenBlock();
+      this.partial.stopReason = reason;
+      this.partial.errorMessage = message;
+      this.stream.push({ type: 'error', reason, error: this.partial });
+      this.stream.end();
+    } catch {
+      // The StreamFn contract forbids throwing; a push/end failure here
+      // has no further recovery surface.
+    }
+  }
+
+  private appendThinking(text: string): void {
+    if (!text) return;
+    let block = this.openBlock;
+    if (block?.type !== 'thinking') {
+      this.closeOpenBlock();
+      block = { type: 'thinking', thinking: '' };
+      this.partial.content.push(block);
+      this.openBlock = block;
+      this.stream.push({ type: 'thinking_start', contentIndex: this.blockIndex(), partial: this.partial });
+    }
+    block.thinking += text;
+    this.stream.push({
+      type: 'thinking_delta',
+      contentIndex: this.blockIndex(),
+      delta: text,
+      partial: this.partial,
+    });
+  }
+
+  private appendVisibleText(text: string): void {
+    if (this.openBlock?.type === 'text') {
+      if (!text) return;
+      this.openBlock.text += text;
+      this.stream.push({ type: 'text_delta', contentIndex: this.blockIndex(), delta: text, partial: this.partial });
+      return;
+    }
+    // No text block open yet: whitespace-only text stays parked so it can
+    // front real text later or be dropped silently at tag/terminal time.
+    const combined = this.pendingLeadingWhitespace + text;
+    if (combined.trim().length === 0) {
+      this.pendingLeadingWhitespace = combined;
+      return;
+    }
+    this.pendingLeadingWhitespace = '';
+    this.closeOpenBlock();
+    const block: TextContent = { type: 'text', text: '' };
+    this.partial.content.push(block);
+    this.openBlock = block;
+    this.stream.push({ type: 'text_start', contentIndex: this.blockIndex(), partial: this.partial });
+    block.text += combined;
+    this.stream.push({
+      type: 'text_delta',
+      contentIndex: this.blockIndex(),
+      delta: combined,
+      partial: this.partial,
+    });
+  }
+
+  private closeOpenBlock(): void {
+    const block = this.openBlock;
+    if (!block) return;
+    this.openBlock = null;
+    if (block.type === 'text') {
+      this.stream.push({
+        type: 'text_end',
+        contentIndex: this.partial.content.indexOf(block),
+        content: block.text,
+        partial: this.partial,
+      });
+    } else {
+      this.stream.push({
+        type: 'thinking_end',
+        contentIndex: this.partial.content.indexOf(block),
+        content: block.thinking,
+        partial: this.partial,
+      });
+    }
+  }
+
+  private blockIndex(): number {
+    return this.partial.content.length - 1;
+  }
+}

package/src/provider/index.ts

@@ -0,0 +1,255 @@
+/**
+ * `createMlxProviderExtension` — the pi inline extension that registers
+ * the in-process `mlx` provider.
+ *
+ * Task 8's `runAgent` passes the returned extension into pi's `main()`
+ * via `extensionFactories`; pi calls the factory during extension load,
+ * and `registerProvider` makes every discovered local model resolvable
+ * as `mlx/<dir-name>` with no /login (the literal apiKey marks the
+ * models available).
+ *
+ * Import discipline (load-bearing): pi is import-order sensitive to its
+ * config env vars, so this module — which the CLI imports BEFORE those
+ * env vars are set — must not runtime-import `@earendil-works/pi-coding-agent`
+ * at module top level. Only type-only pi imports appear here; the
+ * `ExtensionAPI` value arrives as the factory argument.
+ */
+
+import type { ExtensionAPI, InlineExtension } from '@earendil-works/pi-coding-agent';
+import { coldCacheStats, coldSidecarStats, type ColdCacheStats, type ColdSidecarStats } from '@mlx-node/core';
+
+import { canonicalCacheRoot } from '../cold-tier.js';
+import { parseThinkingBudget } from './chat-config.js';
+import { MetricsTrace, type MetricsTraceRecord } from './metrics-trace.js';
+import { MLX_API, MLX_API_KEY, MLX_BASE_URL, MLX_PROVIDER_ID } from './mlx-identity.js';
+import { MlxModelHost } from './model-host.js';
+import type { MlxModelInfo } from './models.js';
+import { PerformanceStatus } from './performance-status.js';
+import { makeMlxStreamSimple, type TurnRecorder } from './stream-adapter.js';
+
+/** Read the process-wide cold-tier snapshot; the native addon may be absent (unit tests). */
+function safeColdStats(): ColdCacheStats | undefined {
+  try {
+    return coldCacheStats();
+  } catch {
+    return undefined;
+  }
+}
+
+/**
+ * Read the process-wide sidecar counters. Separate from {@link safeColdStats}
+ * because it reads a different native struct that has no `enabled` gate: the
+ * counters live in a plain static, not in the tier, so they are valid even
+ * when the tier never opened — which is exactly the case a run with zero
+ * sidecars needs distinguished.
+ */
+function safeSidecarStats(): ColdSidecarStats | undefined {
+  try {
+    return coldSidecarStats();
+  } catch {
+    return undefined;
+  }
+}
+
+/** Injectable seams for {@link createMlxProviderExtension} (unit tests). */
+export interface MlxProviderExtensionDeps {
+  /** Process-wide cold-tier reader; defaults to the native addon (absent in unit tests). */
+  coldStats?: () => ColdCacheStats | undefined;
+  /** Process-wide sidecar-counter reader; defaults to the native addon. */
+  sidecarStats?: () => ColdSidecarStats | undefined;
+  /** Durable per-turn telemetry sink; defaults to a fresh {@link MetricsTrace}. */
+  metricsTrace?: MetricsTrace;
+}
+
+/**
+ * Build the `mlx-provider` inline extension serving `models`. The host
+ * (one per process — it owns the single GPU-resident model) is created
+ * eagerly so repeated factory invocations can never spawn a second
+ * host, but stays lazy about weights: nothing loads until the first
+ * `streamSimple` call.
+ */
+export function createMlxProviderExtension(
+  models: MlxModelInfo[],
+  host?: MlxModelHost,
+  deps: MlxProviderExtensionDeps = {},
+): InlineExtension {
+  const resolvedHost = host ?? new MlxModelHost(models.map((m) => m.discovered));
+  const performanceStatus = new PerformanceStatus();
+  const metricsTrace = deps.metricsTrace ?? new MetricsTrace();
+  const readColdStats = deps.coldStats ?? safeColdStats;
+  const readSidecarStats = deps.sidecarStats ?? safeSidecarStats;
+  // This closure outlives Pi runtime replacement. Pi creates a replacement
+  // runtime for /new and /resume and reruns inline extension factories; each
+  // new factory's session_start updates the root while child sessions keep
+  // using this registered stream. The root id doubles as the cache owner and
+  // the metrics-trace root; the JSONL path travels alongside it for metrics.
+  let rootCacheOwnerId: string | undefined;
+  let rootSessionFile: string | undefined;
+
+  // Cold-tier counters are cumulative since the tier opened. Snapshot them at
+  // each turn's native start (`onTurnStart`, fired inside the serialized host
+  // closure) and diff at the success terminal, so a turn that aborted or
+  // errored — which never reaches `onTurnRecord` — can't leak its restores into
+  // the next successful turn's delta. Inference is serialized per process (host
+  // promise chain), so by the time a turn snapshots, any prior turn has fully
+  // drained. The SYNCHRONOUS counters (hits/misses/bytesRestored/enqueued/
+  // corruptions/queueDrops/restoreDeclines) are exact per-turn; bytesWritten,
+  // evictions and writeErrors advance on the async writer thread, so those
+  // deltas are approximate (documented per field on the record).
+  //
+  // The SIDECAR counters are snapshotted the same way but are exact without
+  // exception: every one of them is recorded inside the native turn finalize,
+  // on the calling thread, before the turn returns. They also have no `enabled`
+  // gate — they live in a plain process static rather than in the tier — so
+  // they read honestly on a run where the tier never opened, which is the run
+  // that most needs them.
+  let turnStartCold = readColdStats();
+  let turnStartSidecar = readSidecarStats();
+  const onTurnStart = (): void => {
+    turnStartCold = readColdStats();
+    turnStartSidecar = readSidecarStats();
+  };
+  const onTurnRecord: TurnRecorder = ({
+    traceId,
+    sessionId,
+    rootSessionId,
+    rootSessionFile,
+    model,
+    final,
+    durationMs,
+    queueMs,
+    resident,
+  }) => {
+    const rec: Omit<MetricsTraceRecord, 'v'> = {
+      traceId,
+      ts: Date.now(),
+      sessionId,
+      rootSessionId,
+      rootSessionFile,
+      model,
+      durationMs,
+      queueMs,
+      resident,
+      finishReason: final.finishReason,
+      promptTokens: final.promptTokens,
+      cachedTokens: final.cachedTokens ?? 0,
+      outputTokens: final.numTokens,
+      reasoningTokens: final.reasoningTokens,
+    };
+    const perf = final.performance;
+    if (perf) {
+      rec.ttftMs = perf.ttftMs;
+      rec.prefillTps = perf.prefillTokensPerSecond;
+      rec.decodeTps = perf.decodeTokensPerSecond;
+      rec.mtpCycles = perf.mtpCycles;
+      // mlx-vlm-comparable headline accept rate (committed tokens per cycle).
+      rec.mtpMeanAccepted = perf.mtpMeanAcceptedTokensTotal;
+    }
+    const cold = readColdStats();
+    if (cold && turnStartCold) {
+      rec.coldHits = cold.hits - turnStartCold.hits;
+      rec.coldMisses = cold.misses - turnStartCold.misses;
+      rec.coldBytesWritten = cold.bytesWritten - turnStartCold.bytesWritten;
+      rec.coldBytesRestored = cold.bytesRestored - turnStartCold.bytesRestored;
+      rec.coldEnqueued = cold.enqueued - turnStartCold.enqueued;
+      rec.coldQueueDrops = cold.queueDrops - turnStartCold.queueDrops;
+      rec.coldEvictions = cold.evictions - turnStartCold.evictions;
+      rec.coldCorruptions = cold.corruptions - turnStartCold.corruptions;
+      rec.coldWriteErrors = cold.writeErrors - turnStartCold.writeErrors;
+      rec.coldRestoreDeclines = cold.restoreDeclines - turnStartCold.restoreDeclines;
+      // Absolutes, not deltas: an aborted/errored turn never reaches this
+      // recorder, so a corruption or a dropped write during one lands in NO
+      // delta. The cumulative counter observed by the next successful turn
+      // still carries it, which is what makes "corruptions must be 0"
+      // checkable at all (`MAX(total) > 0` over any window).
+      rec.coldCorruptionsTotal = cold.corruptions;
+      rec.coldQueueDropsTotal = cold.queueDrops;
+      // Same latch, and the counter that needs it most: a write error is
+      // raised on the background writer, so the one covering the LAST turn
+      // before a crash lands in no delta at all — and "is my cache root
+      // broken?" is a question about ever, not about this turn.
+      rec.coldWriteErrorsTotal = cold.writeErrors;
+      // Cache IDENTITY comes from the END-of-turn snapshot, never the baseline:
+      // the tier opens LAZILY on first use, so on a process's first turn
+      // `turnStartCold` is the all-zero default with `enabled: false` and an
+      // empty root. Canonicalized here — in the writer — so the dashboard
+      // never has to match whatever spelling Rust happened to construct.
+      rec.coldEnabled = cold.enabled;
+      // ONE emptiness test, applied to the CANONICAL value. Gating on the raw
+      // native string used a DIFFERENT test from the one `canonicalCacheRoot`
+      // applies (it trims), so a whitespace-only root passed `length > 0` here
+      // and canonicalized to `''` — which `MetricsTrace.record` then dropped,
+      // leaving a row that says the tier was ON while carrying no root at all.
+      if (cold.enabled) {
+        const canonicalRoot = canonicalCacheRoot(cold.root);
+        if (canonicalRoot.length > 0) rec.coldRoot = canonicalRoot;
+      }
+    }
+    // Its own guard, deliberately not folded into the block above: the sidecar
+    // reader never consults the tier, so it keeps reporting when `cold` is
+    // undefined or the tier failed to open. A run with `coldHits: 0` and no
+    // tier is exactly when "did the capture even run?" is the question, and
+    // gating these on `cold` would blank them precisely then.
+    const sidecar = readSidecarStats();
+    if (sidecar && turnStartSidecar) {
+      rec.coldSidecarCaptureReached = sidecar.captureReached - turnStartSidecar.captureReached;
+      rec.coldSidecarChainEmpty = sidecar.chainEmpty - turnStartSidecar.chainEmpty;
+      rec.coldSidecarBoundarySkips = sidecar.boundarySkips - turnStartSidecar.boundarySkips;
+      rec.coldSidecarAlreadyPersisted = sidecar.alreadyPersisted - turnStartSidecar.alreadyPersisted;
+      rec.coldSidecarEnqueued = sidecar.enqueued - turnStartSidecar.enqueued;
+      rec.coldSidecarQueueDrops = sidecar.queueDrops - turnStartSidecar.queueDrops;
+      rec.coldSidecarInstalled = sidecar.installed - turnStartSidecar.installed;
+      rec.coldSidecarRestoreSuppressed = sidecar.restoreSuppressed - turnStartSidecar.restoreSuppressed;
+    }
+    metricsTrace.record(rec);
+  };
+
+  let getThinkingBudget = (): number | undefined => undefined;
+  const streamSimple = makeMlxStreamSimple(
+    resolvedHost,
+    performanceStatus.record,
+    () => rootCacheOwnerId,
+    onTurnRecord,
+    onTurnStart,
+    () => rootSessionFile,
+    () => getThinkingBudget(),
+  );
+  return {
+    name: 'mlx-provider',
+    factory: (pi: ExtensionAPI) => {
+      pi.registerFlag('thinking-budget', {
+        type: 'string',
+        description: 'Maximum reasoning tokens per model turn; 0 closes thinking immediately.',
+      });
+      getThinkingBudget = () => parseThinkingBudget(pi.getFlag('thinking-budget'));
+      pi.registerProvider(MLX_PROVIDER_ID, {
+        api: MLX_API,
+        baseUrl: MLX_BASE_URL,
+        apiKey: MLX_API_KEY,
+        streamSimple,
+        models: models.map((m) => m.piModel),
+      });
+      pi.on('session_start', (_event, ctx) => {
+        rootCacheOwnerId = ctx.sessionManager.getSessionId();
+        // Snapshot the root JSONL path so a turn submitted under this root is
+        // correlated to it at completion — even for subagent turns, which have
+        // no session file of their own. `getSessionFile` is optional-chained so
+        // a minimal test/mock session manager without it still works.
+        rootSessionFile = ctx.sessionManager.getSessionFile?.();
+      });
+      pi.on('message_end', (event, ctx) => {
+        performanceStatus.showMessage(event, ctx);
+      });
+      // Do not clear on turn_start. Pi emits a fresh turn after every tool
+      // result, before the next inference has terminal metrics to replace the
+      // completed sample; clearing here makes the footer disappear precisely
+      // while a long tool-follow-up prefill is running.
+      pi.on('model_select', (_event, ctx) => {
+        performanceStatus.clear(ctx);
+      });
+      pi.on('session_shutdown', (_event, ctx) => {
+        performanceStatus.clear(ctx);
+      });
+    },
+  };
+}