imprint-mcp 0.6.3 → 0.6.4

package/src/imprint/codex-cli-compile.ts

@@ -10,7 +10,7 @@
 import { type ChildProcess, spawn, spawnSync } from 'node:child_process';
 import { existsSync, mkdirSync, readFileSync, unlinkSync, writeFileSync } from 'node:fs';
 import { isAbsolute as pathIsAbsolute, join as pathJoin } from 'node:path';
-import type { Span } from '@opentelemetry/api';
+import { type Span, context as otelContext } from '@opentelemetry/api';
 import type { AuthCliCompileMode } from './auth-compile-tools.ts';
 import { type SharedModuleManifestEntry, resolvePlanSliceFromFile } from './build-plan.ts';
 import type {
@@ -25,13 +25,12 @@ import { COMPILE_SENTINELS, compileDeadlineAfterVerification } from './mcp-compi
 import type { SharedCompileContext, ToolCandidate } from './tool-candidates.ts';
 import {
   endTraceSpan,
-  llmSpanAttributes,
-  resolveTraceTokenCount,
+  recordLlmUsageSpan,
   setSpanAttributes,
   startTraceSpan,
+  traceInputOutputAttributes,
   traceJsonInputOutputAttributes,
   traceLlmIoEnabled,
-  traceLlmMessages,
   traceToolIoEnabled,
   traced,
 } from './tracing.ts';
@@ -44,6 +43,80 @@ const CLI_PATH = pathJoin(REPO_ROOT, 'src', 'cli.ts');
 const MCP_SERVER_NAME = 'imprint-compile';
 const MAX_VERIFICATION_CYCLES = 5;
 const MAX_MCP_TOOL_TIMEOUT_SEC = 1800;
+const SENTINEL_USAGE_GRACE_MS = 15_000;
+
+interface SentinelGraceController {
+  observeSentinel(): void;
+  observeTurnCompleted(): void;
+  dispose(): void;
+}
+
+interface TurnActivityTracker {
+  isActive(): boolean;
+  started(): void;
+  completed(): void;
+}
+
+/** Track Codex turn lifecycle without depending on optional tracing spans. */
+export function createTurnActivityTracker(): TurnActivityTracker {
+  let active = false;
+  return {
+    isActive: () => active,
+    started: () => {
+      active = true;
+    },
+    completed: () => {
+      active = false;
+    },
+  };
+}
+
+/**
+ * Wait just long enough for the terminal Codex usage event after done/give_up.
+ * A completed turn terminates immediately; the timer is only a bounded fallback
+ * for CLI versions that never emit turn.completed.
+ */
+export function createSentinelGraceController(opts: {
+  hasActiveTurn: () => boolean;
+  terminate: () => void;
+  fallbackMs?: number;
+  schedule?: (callback: () => void, delayMs: number) => ReturnType<typeof setTimeout>;
+  cancel?: (timer: ReturnType<typeof setTimeout>) => void;
+}): SentinelGraceController {
+  const schedule = opts.schedule ?? ((callback, delayMs) => setTimeout(callback, delayMs));
+  const cancel = opts.cancel ?? clearTimeout;
+  let sentinelObserved = false;
+  let terminated = false;
+  let fallbackTimer: ReturnType<typeof setTimeout> | undefined;
+
+  const terminateOnce = (): void => {
+    if (terminated) return;
+    terminated = true;
+    if (fallbackTimer) cancel(fallbackTimer);
+    fallbackTimer = undefined;
+    opts.terminate();
+  };
+
+  return {
+    observeSentinel(): void {
+      sentinelObserved = true;
+      if (!opts.hasActiveTurn()) {
+        terminateOnce();
+        return;
+      }
+      if (fallbackTimer || terminated) return;
+      fallbackTimer = schedule(terminateOnce, opts.fallbackMs ?? SENTINEL_USAGE_GRACE_MS);
+      fallbackTimer.unref?.();
+    },
+    observeTurnCompleted(): void {
+      if (sentinelObserved) terminateOnce();
+    },
+    dispose(): void {
+      if (fallbackTimer) cancel(fallbackTimer);
+      fallbackTimer = undefined;
+    },
+  };
+}
 
 function formatRevisionMode(enabled: boolean | undefined): string {
   return enabled
@@ -127,6 +200,7 @@ export async function compileViaCodexCli(
     },
     async (span) => {
       const result = await compileViaCodexCliImpl(opts, span);
+      const model = opts.model ?? preferredAgentModel('codex-cli');
       setSpanAttributes(span, {
         'imprint.compile.outcome': result.outcome,
         'imprint.compile.success': result.success,
@@ -136,6 +210,20 @@ export async function compileViaCodexCli(
         'imprint.compile.output_tokens': result.outputTokens,
         'imprint.compile.conversation_log': result.conversationLogPath,
       });
+      // Codex reports input_tokens as the total prompt; cached_input_tokens is
+      // a subset, not an additional token bucket.
+      recordLlmUsageSpan(
+        'compile.codex_cli_usage',
+        {
+          provider: 'codex-cli',
+          model,
+          inputTokens: result.inputTokens,
+          outputTokens: result.outputTokens,
+          cacheReadTokens: result.cacheReadInputTokens,
+          cacheWriteTokens: result.cacheCreationInputTokens,
+        },
+        { 'imprint.compile.turns': result.turns },
+      );
       return result;
     },
   );
@@ -237,30 +325,16 @@ Use the imprint-compile MCP tools to inspect the session, write artifacts, run t
   }
 
   const model = opts.model ?? preferredAgentModel('codex-cli');
-  const initialTokenCount = resolveTraceTokenCount(null, initialPrompt);
   const captureLlmIo = traceLlmIoEnabled();
   // done() owns a separately bounded live-verification phase, so its transport
   // timeout cannot be derived from the compiler's remaining reasoning budget.
   const mcpToolTimeoutSec = MAX_MCP_TOOL_TIMEOUT_SEC;
   setSpanAttributes(traceSpan, {
-    ...llmSpanAttributes({
-      provider: 'codex-cli',
-      model,
-      inputTokens: initialTokenCount.tokens,
-      tokenCountsEstimated: true,
-      inputTokenSource: initialTokenCount.source,
-      inputMessages: captureLlmIo
-        ? traceLlmMessages([{ role: 'user', content: initialPrompt }])
-        : undefined,
-      inputValue: captureLlmIo ? initialPrompt : undefined,
-      invocationParameters: {
-        command: 'codex exec',
-        json: true,
-        sandbox: 'workspace-write',
-        tool_timeout_sec: mcpToolTimeoutSec,
-      },
-    }),
+    ...(captureLlmIo ? traceInputOutputAttributes('input', initialPrompt) : {}),
     'imprint.compile.initial_prompt_chars': initialPrompt.length,
+    'imprint.compile.command': 'codex exec',
+    'imprint.compile.sandbox': 'workspace-write',
+    'imprint.compile.tool_timeout_sec': mcpToolTimeoutSec,
   });
 
   const execArgs = opts.resume
@@ -352,27 +426,9 @@ Use the imprint-compile MCP tools to inspect the session, write artifacts, run t
   } finally {
     process.removeListener('exit', terminateOnParentExit);
   }
-  const hasActualUsage = result.inputTokens > 0 || result.outputTokens > 0;
-  const inputTokenCount = resolveTraceTokenCount(
-    hasActualUsage ? result.inputTokens : null,
-    initialPrompt,
-  );
-  const outputTokenCount = resolveTraceTokenCount(
-    hasActualUsage ? result.outputTokens : null,
-    result.message,
-  );
   setSpanAttributes(traceSpan, {
-    ...llmSpanAttributes({
-      provider: 'codex-cli',
-      model,
-      inputTokens: inputTokenCount.tokens,
-      outputTokens: outputTokenCount.tokens,
-      tokenCountsEstimated:
-        inputTokenCount.source === 'estimated' || outputTokenCount.source === 'estimated',
-      inputTokenSource: inputTokenCount.source,
-      outputTokenSource: outputTokenCount.source,
-    }),
     'imprint.compile.message': result.message,
+    ...(captureLlmIo ? traceInputOutputAttributes('output', result.message) : {}),
   });
   return result;
 }
@@ -392,6 +448,10 @@ async function driveJsonl(
   opts: CompileViaCodexCliOptions,
   traceSpan?: Span,
 ): Promise<CompileAgentResult> {
+  // Bun's child-process event emitters do not preserve AsyncLocalStorage.
+  // Restore the compile span context inside stdout callbacks so turn/tool
+  // spans remain children of compile.codex_cli_agent in Phoenix.
+  const parentCtx = otelContext.active();
   const conversationLog: unknown[] = [];
   const conversationLogPath = pathJoin(opts.absoluteToolDir, '.compile-log.json');
   const rawStdoutPath = pathJoin(opts.absoluteToolDir, '.codex-stdout.jsonl');
@@ -414,6 +474,7 @@ async function driveJsonl(
   const rawStderrChunks: string[] = [];
   let inputTokens = 0;
   let outputTokens = 0;
+  let cacheReadInputTokens = 0;
   let turn = 0;
   let lastErrorMessage = '';
   let stderrBuf = '';
@@ -421,6 +482,8 @@ async function driveJsonl(
   let capturedSessionId: string | undefined;
   const toolSpans = new Map<string, Span>();
   let currentTurnSpan: Span | null = null;
+  const turnActivity = createTurnActivityTracker();
+  let onTurnCompletedAfterSentinel: (() => void) | undefined;
 
   const budgetMs = Math.max(0, opts.deadlineMs - Date.now());
   const fireProgress = (phase: 'thinking' | 'tool', toolName?: string): void => {
@@ -481,103 +544,106 @@ async function driveJsonl(
 
   let stdoutBuf = '';
   child.stdout?.on('data', (chunk: Buffer) => {
-    const chunkText = chunk.toString('utf8');
-    rawStdoutChunks.push(chunkText);
-    stdoutBuf += chunkText;
-    while (true) {
-      const nl = stdoutBuf.indexOf('\n');
-      if (nl < 0) break;
-      const line = stdoutBuf.slice(0, nl).trim();
-      stdoutBuf = stdoutBuf.slice(nl + 1);
-      if (!line) continue;
-
-      let evt: CodexJsonEvent;
-      try {
-        evt = JSON.parse(line) as CodexJsonEvent;
-      } catch (err) {
-        log(`unparseable jsonl line: ${errMsg(err)}`);
-        continue;
-      }
+    otelContext.with(parentCtx, () => {
+      const chunkText = chunk.toString('utf8');
+      rawStdoutChunks.push(chunkText);
+      stdoutBuf += chunkText;
+      while (true) {
+        const nl = stdoutBuf.indexOf('\n');
+        if (nl < 0) break;
+        const line = stdoutBuf.slice(0, nl).trim();
+        stdoutBuf = stdoutBuf.slice(nl + 1);
+        if (!line) continue;
+
+        let evt: CodexJsonEvent;
+        try {
+          evt = JSON.parse(line) as CodexJsonEvent;
+        } catch (err) {
+          log(`unparseable jsonl line: ${errMsg(err)}`);
+          continue;
+        }
 
-      conversationLog.push(evt);
-      // Keep the human-readable compile log current during a long agent turn.
-      // Rewriting is throttled so streamed bursts do not cause one full-file
-      // write per JSONL event, while a stalled/crashed turn still leaves the
-      // latest tool activity visible on disk.
-      scheduleLogFlush();
-
-      if (evt.type === 'thread.started') {
-        log(`thread_id=${evt.thread_id ?? '(none)'}`);
-        setSpanAttributes(traceSpan, { 'codex.thread_id': evt.thread_id });
-        if (evt.thread_id) capturedSessionId = evt.thread_id;
-        continue;
-      }
+        conversationLog.push(evt);
+        // Keep the human-readable compile log current during a long agent turn.
+        // Rewriting is throttled so streamed bursts do not cause one full-file
+        // write per JSONL event, while a stalled/crashed turn still leaves the
+        // latest tool activity visible on disk.
+        scheduleLogFlush();
+
+        if (evt.type === 'thread.started') {
+          log(`thread_id=${evt.thread_id ?? '(none)'}`);
+          setSpanAttributes(traceSpan, { 'codex.thread_id': evt.thread_id });
+          if (evt.thread_id) capturedSessionId = evt.thread_id;
+          continue;
+        }
 
-      if (evt.type === 'turn.started') {
-        if (currentTurnSpan) endTraceSpan(currentTurnSpan);
-        flushLog();
-        turn++;
-        currentTurnSpan = startTraceSpan(`agent.turn.${turn}`, 'CHAIN', {
-          'imprint.agent.turn': turn,
-          'imprint.agent.cumulative_input_tokens': inputTokens,
-          'imprint.agent.cumulative_output_tokens': outputTokens,
-        });
-        fireProgress('thinking');
-        continue;
-      }
+        if (evt.type === 'turn.started') {
+          turnActivity.started();
+          if (currentTurnSpan) endTraceSpan(currentTurnSpan);
+          flushLog();
+          turn++;
+          currentTurnSpan = startTraceSpan(`agent.turn.${turn}`, 'CHAIN', {
+            'imprint.agent.turn': turn,
+            'imprint.agent.cumulative_input_tokens': inputTokens,
+            'imprint.agent.cumulative_output_tokens': outputTokens,
+          });
+          fireProgress('thinking');
+          continue;
+        }
 
-      const normalizedToolEvt = normalizeCodexToolEvent(evt);
-      if (normalizedToolEvt) {
-        const { eventType, item } = normalizedToolEvt;
-        const toolName = codexToolName(item);
-        if (toolName) {
-          traceCodexToolEvent(toolSpans, eventType, item, toolName);
-          fireProgress(eventType === 'item.started' ? 'tool' : 'thinking', toolName);
+        const normalizedToolEvt = normalizeCodexToolEvent(evt);
+        if (normalizedToolEvt) {
+          const { eventType, item } = normalizedToolEvt;
+          const toolName = codexToolName(item);
+          if (toolName) {
+            traceCodexToolEvent(toolSpans, eventType, item, toolName);
+            fireProgress(eventType === 'item.started' ? 'tool' : 'thinking', toolName);
+          }
+          continue;
         }
-        continue;
-      }
 
-      if ((evt.type === 'item.started' || evt.type === 'item.completed') && evt.item) {
-        const agentMessage = codexAgentMessageText(evt.item);
-        if (agentMessage && evt.type === 'item.completed') {
-          agentMessageCount++;
-          setSpanAttributes(traceSpan, {
-            'imprint.codex.agent_messages': agentMessageCount,
-            'imprint.codex.last_agent_message_chars': agentMessage.length,
-            ...(traceLlmIoEnabled()
-              ? llmSpanAttributes({
-                  provider: 'codex-cli',
-                  model: opts.model ?? preferredAgentModel('codex-cli'),
-                  outputMessages: traceLlmMessages([{ role: 'assistant', content: agentMessage }]),
-                  outputValue: agentMessage,
-                })
-              : {}),
-          });
+        if ((evt.type === 'item.started' || evt.type === 'item.completed') && evt.item) {
+          const agentMessage = codexAgentMessageText(evt.item);
+          if (agentMessage && evt.type === 'item.completed') {
+            agentMessageCount++;
+            setSpanAttributes(traceSpan, {
+              'imprint.codex.agent_messages': agentMessageCount,
+              'imprint.codex.last_agent_message_chars': agentMessage.length,
+              ...(traceLlmIoEnabled()
+                ? traceInputOutputAttributes('output', agentMessage, undefined, 'agent_message')
+                : {}),
+            });
+            continue;
+          }
           continue;
         }
-        continue;
-      }
 
-      if (evt.type === 'turn.completed') {
-        const turnInput = evt.usage?.input_tokens ?? 0;
-        const turnOutput = evt.usage?.output_tokens ?? 0;
-        inputTokens += turnInput;
-        outputTokens += turnOutput;
-        if (currentTurnSpan) {
-          setSpanAttributes(currentTurnSpan, {
-            'imprint.agent.turn_input_tokens': turnInput,
-            'imprint.agent.turn_output_tokens': turnOutput,
-          });
-          endTraceSpan(currentTurnSpan);
-          currentTurnSpan = null;
+        if (evt.type === 'turn.completed') {
+          const turnInput = evt.usage?.input_tokens ?? 0;
+          const turnOutput = evt.usage?.output_tokens ?? 0;
+          const turnCacheRead = evt.usage?.cached_input_tokens ?? 0;
+          inputTokens += turnInput;
+          outputTokens += turnOutput;
+          cacheReadInputTokens += turnCacheRead;
+          if (currentTurnSpan) {
+            setSpanAttributes(currentTurnSpan, {
+              'imprint.agent.turn_input_tokens': turnInput,
+              'imprint.agent.turn_output_tokens': turnOutput,
+              'imprint.agent.turn_cache_read_input_tokens': turnCacheRead,
+            });
+            endTraceSpan(currentTurnSpan);
+            currentTurnSpan = null;
+          }
+          turnActivity.completed();
+          onTurnCompletedAfterSentinel?.();
+          continue;
         }
-        continue;
-      }
 
-      if (evt.type === 'error' || evt.type === 'turn.failed') {
-        lastErrorMessage = evt.message ?? evt.error?.message ?? JSON.stringify(evt);
+        if (evt.type === 'error' || evt.type === 'turn.failed') {
+          lastErrorMessage = evt.message ?? evt.error?.message ?? JSON.stringify(evt);
+        }
       }
-    }
+    });
   });
 
   child.stderr?.on('data', (chunk: Buffer) => {
@@ -593,6 +659,8 @@ async function driveJsonl(
     const finish = (code: number): void => {
       if (resolved) return;
       resolved = true;
+      sentinelGrace.dispose();
+      onTurnCompletedAfterSentinel = undefined;
       if (forcedKillTimer) clearTimeout(forcedKillTimer);
       clearInterval(sentinelTimer);
       resolve(code);
@@ -616,10 +684,17 @@ async function driveJsonl(
       }, 2000);
       forcedKillTimer.unref?.();
     };
+    const sentinelGrace = createSentinelGraceController({
+      // Turn lifecycle is independent of tracing. startTraceSpan() returns null
+      // when tracing is disabled, but terminal usage still must be captured.
+      hasActiveTurn: turnActivity.isActive,
+      terminate: terminateForSentinel,
+    });
+    onTurnCompletedAfterSentinel = sentinelGrace.observeTurnCompleted;
     const sentinelTimer = setInterval(() => {
       const checkpointReached = opts.authMode && existsSync(checkpointSentinel);
       if (!existsSync(doneSentinel) && !existsSync(giveUpSentinel) && !checkpointReached) return;
-      terminateForSentinel();
+      sentinelGrace.observeSentinel();
     }, 500);
     child.once('close', (code) => {
       finish(code ?? -1);
@@ -629,6 +704,7 @@ async function driveJsonl(
     });
   });
   childExited = true;
+  turnActivity.completed();
   clearTimeout(deadlineTimer);
   if (currentTurnSpan) endTraceSpan(currentTurnSpan);
   for (const span of toolSpans.values()) endTraceSpan(span);
@@ -687,7 +763,7 @@ async function driveJsonl(
     durationMs: Date.now() - opts.startTime,
     inputTokens,
     outputTokens,
-    cacheReadInputTokens: 0,
+    cacheReadInputTokens,
     cacheCreationInputTokens: 0,
     sessionId: capturedSessionId,
   };

package/src/imprint/inject-listener.ts

@@ -5,13 +5,20 @@
  * `[IMPRINT]` is exact-match — keep stable.
  */
 
-export const INJECTED_LISTENER_SOURCE = `
+export function createInjectedListenerSource(eventToken: string): string {
+  return `
 (function imprintInjector() {
   if (window.__imprint_injected__) return;
   window.__imprint_injected__ = true;
 
   const SENTINEL = '[IMPRINT]';
+  const EVENT_TOKEN = ${JSON.stringify(eventToken)};
+  const emitLog = console.log.bind(console);
+  const stringify = JSON.stringify.bind(JSON);
   const MAX_VAL = 200;
+  let lastIntentTarget = null;
+  let lastIntentAt = 0;
+  let lastIntentKind = null;
 
   function safeStr(v) {
     try {
@@ -53,6 +60,7 @@ export const INJECTED_LISTENER_SOURCE = `
         text: safeStr((el.textContent || '').trim()),
         ariaLabel: el.getAttribute && el.getAttribute('aria-label') || null,
         href: el.tagName === 'A' ? el.getAttribute('href') : null,
+        absoluteHref: el.tagName === 'A' ? el.href : null,
         selector: selectorFor(el),
       };
     } catch (e) { return {}; }
@@ -60,17 +68,33 @@ export const INJECTED_LISTENER_SOURCE = `
 
   function emit(type, payload) {
     try {
-      console.log(SENTINEL, type, JSON.stringify(payload));
+      emitLog(SENTINEL, EVENT_TOKEN, type, stringify(payload));
     } catch (e) { /* ignore */ }
   }
 
   function onClick(ev) {
     try {
-      const tgt = ev.target;
+      onIntent(ev, 'click');
+      const tgt = ev.target && ev.target.closest ? ev.target.closest('a,button,input,form') || ev.target : ev.target;
       emit('click', describe(tgt));
     } catch (e) { /* ignore */ }
   }
 
+  function onIntent(ev, intentKind) {
+    try {
+      if (!ev.isTrusted) return;
+      const tgt = ev.target && ev.target.closest ? ev.target.closest('a,button,input,form') : ev.target;
+      if (!tgt) return;
+      const now = Date.now();
+      const activation = intentKind === 'pointerdown' || intentKind === 'click';
+      if (!activation && tgt === lastIntentTarget && intentKind === lastIntentKind && now - lastIntentAt < 250) return;
+      lastIntentTarget = tgt;
+      lastIntentAt = now;
+      lastIntentKind = intentKind;
+      emit('intent', { ...describe(tgt), intentAt: now, intentKind: intentKind });
+    } catch (e) { /* ignore */ }
+  }
+
   function onInput(ev) {
     try {
       const tgt = ev.target;
@@ -128,10 +152,14 @@ export const INJECTED_LISTENER_SOURCE = `
   // Capture phase = true so we see the event before the site has a chance to
   // stopPropagation it.
   document.addEventListener('click', onClick, true);
+  document.addEventListener('pointerover', (ev) => onIntent(ev, 'pointerover'), true);
+  document.addEventListener('pointerdown', (ev) => onIntent(ev, 'pointerdown'), true);
+  document.addEventListener('focusin', (ev) => onIntent(ev, 'focusin'), true);
   document.addEventListener('input', onInput, true);
   document.addEventListener('change', onChange, true);
   document.addEventListener('submit', onSubmit, true);
 })();
 `;
+}
 
 export const IMPRINT_SENTINEL = '[IMPRINT]';

package/src/imprint/llm.ts

@@ -627,9 +627,9 @@ async function traceAnalyze(
         throw err;
       }
       // Providers report `inputTokens` as the *uncached* input only; the cached
-      // portion lives in the cache fields. `llmCostAttributes` expects the TOTAL
-      // prompt tokens (it derives uncached = total − cacheRead − cacheWrite), so
-      // sum them here. A real total is also large enough to clear the
+      // portion lives in the cache fields. Phoenix expects the TOTAL prompt plus
+      // the cache breakdown so it can calculate cost server-side, so sum them
+      // here. A real total is also large enough to clear the
       // resolveTraceTokenCount sanity check, so cache-hit calls stop falling back
       // to the chars/4 estimate.
       const cacheReadTokens = result.cacheReadInputTokens ?? undefined;

package/src/imprint/react-flight-limits.ts

@@ -0,0 +1,31 @@
+/** Bound per-row metadata for untrusted React Flight bodies below the byte cap. */
+export const MAX_REACT_FLIGHT_ROWS = 16_384;
+
+/** Bound aggregate parsed JSON width across an untrusted React Flight body. */
+export const MAX_REACT_FLIGHT_JSON_NODES = 32_768;
+
+/**
+ * Estimate JSON node count before JSON.parse allocates the object graph.
+ * For valid JSON, one root plus each container opener and comma is a
+ * conservative upper approximation of containers and values. Delimiters in
+ * strings are ignored, including escaped quotes.
+ */
+export function boundedJsonNodeCount(json: string, limit: number): number | null {
+  if (limit < 1) return null;
+  let count = 1;
+  let inString = false;
+  for (let index = 0; index < json.length; index++) {
+    const char = json.charCodeAt(index);
+    if (inString) {
+      if (char === 0x5c) index++;
+      else if (char === 0x22) inString = false;
+      continue;
+    }
+    if (char === 0x22) {
+      inString = true;
+    } else if (char === 0x7b || char === 0x5b || char === 0x2c) {
+      if (++count > limit) return null;
+    }
+  }
+  return count;
+}