imprint-mcp 0.6.3 → 0.6.5

package/README.md

@@ -93,6 +93,18 @@ regenerate/record it with `bun scripts/demo-teach.ts`.)*
 > [!TIP]
 > All three steps happen in a single `imprint teach` command. Credentials and PII are redacted automatically before anything reaches the LLM.
 
+When one recording yields multiple tools, all candidates are selected by default.
+Imprint automatically includes upstream lookup/producer tools required by any
+selected downstream tool. Use `--primary-tool` to compile only the primary tool
+and its prerequisites; `--all-tools` makes the default full-set selection explicit.
+
+The recorder also has a bounded fallback for Next.js React Server Component
+responses that Chromium evicts before `Network.getResponseBody` can read them.
+It streams only narrowly identified GET navigation/prefetch and POST Server
+Action responses, keeps normal body capture authoritative, and writes fallback
+bytes only when the normal read actually fails. The candidate policy and
+resource limits are documented in the [capture protocol](docs/capture-protocol.md).
+
 ---
 
 ## Why Imprint?
@@ -300,7 +312,7 @@ See [Sharing Skills](docs/credential-sharing.md) for the full flow.
 | [Sharing Skills](docs/credential-sharing.md) | Credential export/import and remote provisioning |
 | [MCP Maintenance](docs/mcp-maintenance.md) | Audit, disable, restore, and prune MCP state |
 | [Troubleshooting](docs/troubleshooting.md) | Common failures and fixes |
-| [Tracing](docs/tracing.md) | OpenTelemetry tracing, cost rollup, and Phoenix setup |
+| [Tracing](docs/tracing.md) | OpenTelemetry tracing, Phoenix-owned cost tracking, and setup |
 
 <details>
 <summary>More docs</summary>

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "imprint-mcp",
-  "version": "0.6.3",
+  "version": "0.6.5",
   "description": "Teach an AI agent how to use any website. Once. Records a real browser session + narration; generates a deterministic MCP tool plus a DOM-replay playbook fallback.",
   "type": "module",
   "exports": {

package/src/cli.ts

@@ -7,7 +7,7 @@ import { parseArgs } from 'node:util';
 import { IS_COMPILED_BINARY } from './imprint/is-compiled.ts';
 import type { ProviderName } from './imprint/llm.ts';
 import { isDebug } from './imprint/log.ts';
-import { shutdownTracing, tracedWithCostRollup } from './imprint/tracing.ts';
+import { shutdownTracing, traced } from './imprint/tracing.ts';
 import { VERSION } from './imprint/version.ts';
 
 /** Load .env from the project root (next to src/) if present.
@@ -40,6 +40,45 @@ export function parseDuration(dur: string): number | null {
   return num;
 }
 
+export type TeachSelectionMode =
+  | 'tool'
+  | 'primary'
+  | 'all-explicit'
+  | 'all-default'
+  | 'interactive-default-all';
+
+/** Pure CLI selection-mode resolver shared by validation and telemetry. */
+export function resolveTeachSelectionMode(opts: {
+  tool?: string;
+  primaryTool?: boolean;
+  allTools?: boolean;
+  noInteractive?: boolean;
+}): { mode: TeachSelectionMode; error?: string } {
+  if (opts.tool && opts.allTools) {
+    return { mode: 'tool', error: 'error: --tool cannot be combined with --all-tools' };
+  }
+  if (opts.tool && opts.primaryTool) {
+    return { mode: 'tool', error: 'error: --tool cannot be combined with --primary-tool' };
+  }
+  if (opts.allTools && opts.primaryTool) {
+    return {
+      mode: 'primary',
+      error: 'error: --all-tools cannot be combined with --primary-tool',
+    };
+  }
+  return {
+    mode: opts.tool
+      ? 'tool'
+      : opts.primaryTool
+        ? 'primary'
+        : opts.allTools
+          ? 'all-explicit'
+          : opts.noInteractive
+            ? 'all-default'
+            : 'interactive-default-all',
+  };
+}
+
 const HELP = `imprint v${VERSION} — teach an AI agent to use any website. Once.
 
 USAGE
@@ -112,7 +151,7 @@ export const VERB_HELP: Record<string, VerbHelp> = {
     summary:
       'Record a workflow, compile both artifacts, emit the tool, and connect to your AI platform — all in one interactive flow. Supports resuming incomplete runs and multiple workflows per site.',
     usage: [
-      'imprint teach <site> [--url <url>] [--from-session <path>] [--persist-profile] [--no-interactive] [--all-tools] [--tool <toolName>] [--provider <name>] [--model <name>] [--timeout <duration>] [--keep-test] [--skip-replay] [--from-step <step>] [--to-step <step>] [--only <step>]',
+      'imprint teach <site> [--url <url>] [--from-session <path>] [--persist-profile] [--no-interactive] [--primary-tool | --all-tools] [--tool <toolName>] [--provider <name>] [--model <name>] [--timeout <duration>] [--keep-test] [--skip-replay] [--from-step <step>] [--to-step <step>] [--only <step>]',
     ],
     flags: [
       { name: '--url <url>', description: 'Starting URL (else about:blank).' },
@@ -123,13 +162,15 @@ export const VERB_HELP: Record<string, VerbHelp> = {
       { name: '--persist-profile', description: 'Reuse a stable Chrome profile for this site.' },
       {
         name: '--no-interactive',
-        description:
-          'Run without prompts; compile the primary detected tool and print integration snippets.',
+        description: 'Run without prompts; compile all detected tools by default.',
       },
       {
         name: '--all-tools',
-        description:
-          'With --no-interactive, compile every detected candidate tool instead of only the primary.',
+        description: 'Explicitly compile all detected tools (the default).',
+      },
+      {
+        name: '--primary-tool',
+        description: 'Compile only the primary detected tool and its upstream dependencies.',
       },
       {
         name: '--tool <toolName>',
@@ -1190,7 +1231,7 @@ async function main(argv: string[]): Promise<number> {
         return 2;
       }
       const outPath = values.out ?? localAuditReportPath(site);
-      const score = await tracedWithCostRollup(
+      const score = await traced(
         'cli.audit',
         'AGENT',
         {
@@ -1423,6 +1464,7 @@ async function main(argv: string[]): Promise<number> {
           'persist-profile': { type: 'boolean' },
           'no-interactive': { type: 'boolean' },
           'all-tools': { type: 'boolean' },
+          'primary-tool': { type: 'boolean' },
           tool: { type: 'string' },
           provider: { type: 'string' },
           model: { type: 'string' },
@@ -1464,8 +1506,14 @@ async function main(argv: string[]): Promise<number> {
           return 2;
         }
       }
-      if (values.tool && values['all-tools']) {
-        console.error('error: --tool cannot be combined with --all-tools');
+      const selection = resolveTeachSelectionMode({
+        tool: values.tool,
+        primaryTool: values['primary-tool'],
+        allTools: values['all-tools'],
+        noInteractive: values['no-interactive'],
+      });
+      if (selection.error) {
+        console.error(selection.error);
         return 2;
       }
 
@@ -1500,10 +1548,11 @@ async function main(argv: string[]): Promise<number> {
       const ctrl = new AbortController();
       const onSigint = (): void => ctrl.abort();
       process.once('SIGINT', onSigint);
+      const teachSelectionMode = selection.mode;
 
       try {
         const { teach } = await import('./imprint/teach.ts');
-        await tracedWithCostRollup(
+        await traced(
           'cli.teach',
           'AGENT',
           {
@@ -1513,7 +1562,10 @@ async function main(argv: string[]): Promise<number> {
             'imprint.provider': values.provider ?? 'auto',
             'imprint.model': values.model ?? 'auto',
             'imprint.timeout_ms': teachTimeoutMs ?? 'default',
-            'imprint.all_tools': values['all-tools'] ?? false,
+            'imprint.all_tools':
+              teachSelectionMode === 'all-explicit' || teachSelectionMode === 'all-default',
+            'imprint.selection_mode': teachSelectionMode,
+            'imprint.primary_tool': values['primary-tool'] ?? false,
             'imprint.tool': values.tool,
             'imprint.no_interactive': values['no-interactive'] ?? false,
             'imprint.skip_replay': values['skip-replay'] ?? false,
@@ -1530,7 +1582,7 @@ async function main(argv: string[]): Promise<number> {
               model: values.model,
               maxDurationMs: teachTimeoutMs,
               keepTest: values['keep-test'] || process.env.IMPRINT_KEEP_TEST === '1',
-              allTools: values['all-tools'],
+              primaryTool: values['primary-tool'],
               tool: values.tool,
               skipReplay: values['skip-replay'],
               fromStep: fromStepArg,

package/src/imprint/audit.ts

@@ -26,7 +26,7 @@ import { createLog } from './log.ts';
 import { imprintHomeDir } from './paths.ts';
 import { type TeachState, loadTeachState } from './teach-state.ts';
 import { type ResolvedTool, discoverTools } from './tool-loader.ts';
-import { llmSpanAttributes, setSpanAttributes, totalPromptTokens, traced } from './tracing.ts';
+import { recordLlmUsageSpan, setSpanAttributes, totalPromptTokens, traced } from './tracing.ts';
 
 const log = createLog('audit');
 
@@ -531,6 +531,20 @@ export function buildAuditMcpCommandArgs(site: string, toolNames: readonly strin
   return ['run', CLI_PATH, 'mcp-server', site, '--include-tools', toolNames.join(',')];
 }
 
+/** Normalize provider usage to OpenInference's total prompt-token contract. */
+export function auditPromptTokens(
+  provider: ProviderName,
+  inputTokens: number,
+  cacheReadTokens: number,
+  cacheWriteTokens: number,
+): number {
+  // Codex reports input_tokens as the total and cached_input_tokens as a subset.
+  // Anthropic reports input_tokens as the uncached delta, with cache buckets
+  // alongside it, so only Anthropic usage needs summing.
+  if (provider === 'codex-cli') return inputTokens;
+  return totalPromptTokens(inputTokens, cacheReadTokens, cacheWriteTokens) ?? inputTokens;
+}
+
 export async function runAudit(opts: RunAuditOptions): Promise<AuditScore> {
   return await traced(
     'audit.session',
@@ -636,6 +650,28 @@ export async function runAudit(opts: RunAuditOptions): Promise<AuditScore> {
       if (auditedToolCount === 0) {
         log('all discovered tools require human-interactive authentication; auditor not spawned');
       }
+      recordLlmUsageSpan(
+        'audit.llm_usage',
+        {
+          provider,
+          model,
+          inputTokens: auditPromptTokens(
+            provider,
+            drive.inputTokens,
+            drive.cacheReadInputTokens,
+            drive.cacheCreationInputTokens,
+          ),
+          outputTokens: drive.outputTokens,
+          cacheReadTokens: drive.cacheReadInputTokens,
+          cacheWriteTokens: drive.cacheCreationInputTokens,
+        },
+        {
+          'imprint.audit.turns': drive.turns,
+          ...(drive.totalCostUsd != null
+            ? { 'imprint.audit.provider_reported_cost_usd': drive.totalCostUsd }
+            : {}),
+        },
+      );
 
       // Cross-reference compile-time live verification with the audit grade.
       // The downgrade rule's purpose is to surface "flying blind" runs —
@@ -695,15 +731,6 @@ export async function runAudit(opts: RunAuditOptions): Promise<AuditScore> {
         }
       }
 
-      // TOTAL prompt (uncached + cache) for the cost calc; the cache split is
-      // passed to llmSpanAttributes separately. Always a number here
-      // (drive.inputTokens is non-null), so the cost-suppression happens via the
-      // `|| undefined` at the call site below.
-      const totalInputTokens = totalPromptTokens(
-        drive.inputTokens,
-        drive.cacheReadInputTokens,
-        drive.cacheCreationInputTokens,
-      );
       setSpanAttributes(span, {
         'imprint.audit.score': score.score,
         'imprint.audit.correct': score.correct,
@@ -727,16 +754,6 @@ export async function runAudit(opts: RunAuditOptions): Promise<AuditScore> {
         'imprint.audit.timed_out': drive.timedOut,
         'imprint.audit.turns': drive.turns,
         ...(drive.totalCostUsd != null ? { 'imprint.audit.cost_usd': drive.totalCostUsd } : {}),
-        ...llmSpanAttributes({
-          provider,
-          model,
-          // `|| undefined`: when no usage was captured (e.g. spawn failure → 0
-          // tokens), suppress a bogus $0 cost instead of emitting it.
-          inputTokens: totalInputTokens || undefined,
-          outputTokens: drive.outputTokens || undefined,
-          cacheReadTokens: drive.cacheReadInputTokens || undefined,
-          cacheWriteTokens: drive.cacheCreationInputTokens || undefined,
-        }),
       });
 
       // Persist the full result (deterministic score + the raw model report).
@@ -1093,7 +1110,7 @@ ${buildAuditInitialPrompt(opts, unverifiedNote)}`;
     transcript: session.transcript,
     inputTokens: session.inputTokens,
     outputTokens: session.outputTokens,
-    cacheReadInputTokens: 0,
+    cacheReadInputTokens: session.cacheReadInputTokens,
     cacheCreationInputTokens: 0,
     totalCostUsd: null,
   };
@@ -1308,6 +1325,7 @@ async function collectCodexAssistantText(
   let turns = 0;
   let inputTokens = 0;
   let outputTokens = 0;
+  let cacheReadInputTokens = 0;
   const t0 = Date.now();
   const elapsedStr = (): string => {
     const s = Math.floor((Date.now() - t0) / 1000);
@@ -1357,6 +1375,7 @@ async function collectCodexAssistantText(
       if (evt.type === 'turn.completed') {
         inputTokens += evt.usage?.input_tokens ?? 0;
         outputTokens += evt.usage?.output_tokens ?? 0;
+        cacheReadInputTokens += evt.usage?.cached_input_tokens ?? 0;
         log(`[${elapsedStr()}] turn ${turns} completed`);
         continue;
       }
@@ -1411,7 +1430,7 @@ async function collectCodexAssistantText(
     turns,
     inputTokens,
     outputTokens,
-    cacheReadInputTokens: 0,
+    cacheReadInputTokens,
     cacheCreationInputTokens: 0,
     totalCostUsd: null,
   };
@@ -1461,6 +1480,8 @@ interface CodexStreamJsonEvent {
   usage?: {
     input_tokens?: number;
     output_tokens?: number;
+    cached_input_tokens?: number;
+    reasoning_output_tokens?: number;
   };
   message?: string;
   error?: { message?: string };

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

@@ -43,7 +43,7 @@ import { COMPILE_SENTINELS, compileDeadlineAfterVerification } from './mcp-compi
 import type { SharedCompileContext, ToolCandidate } from './tool-candidates.ts';
 import {
   endTraceSpan,
-  llmSpanAttributes,
+  recordLlmUsageSpan,
   setSpanAttributes,
   startTraceSpan,
   totalPromptTokens,
@@ -88,6 +88,28 @@ const USAGE_POLICY_REFUSAL =
 /** Total attempts (1 initial + retries) when a usage-policy refusal is hit. */
 const MAX_USAGE_POLICY_ATTEMPTS = 3;
 
+interface CompileUsageTotals {
+  turns: number;
+  inputTokens: number;
+  outputTokens: number;
+  cacheReadInputTokens: number;
+  cacheCreationInputTokens: number;
+}
+
+/** Add one paid CLI attempt to the aggregate usage reported for the compile. */
+export function addCompileUsageTotals(
+  totals: CompileUsageTotals,
+  attempt: CompileUsageTotals,
+): CompileUsageTotals {
+  return {
+    turns: totals.turns + attempt.turns,
+    inputTokens: totals.inputTokens + attempt.inputTokens,
+    outputTokens: totals.outputTokens + attempt.outputTokens,
+    cacheReadInputTokens: totals.cacheReadInputTokens + attempt.cacheReadInputTokens,
+    cacheCreationInputTokens: totals.cacheCreationInputTokens + attempt.cacheCreationInputTokens,
+  };
+}
+
 /** Exponential backoff with jitter between refusal retries. Spacing matters:
  *  bursts of near-identical requests raise the safety-filter trip rate. */
 function usagePolicyBackoffMs(attempt: number): number {
@@ -212,6 +234,7 @@ export async function compileViaClaudeCli(
     },
     async (span) => {
       const result = await compileViaClaudeCliImpl(opts);
+      const model = opts.model ?? preferredAgentModel('claude-cli');
       setSpanAttributes(span, {
         'imprint.compile.outcome': result.outcome,
         'imprint.compile.turns': result.turns,
@@ -220,9 +243,12 @@ export async function compileViaClaudeCli(
         'imprint.compile.output_tokens': result.outputTokens,
         'imprint.compile.cache_read_input_tokens': result.cacheReadInputTokens,
         'imprint.compile.cache_creation_input_tokens': result.cacheCreationInputTokens,
-        ...llmSpanAttributes({
+      });
+      recordLlmUsageSpan(
+        'compile.claude_cli_usage',
+        {
           provider: 'claude-cli',
-          model: opts.model ?? preferredAgentModel('claude-cli'),
+          model,
           // TOTAL prompt (uncached + cache); the cache split is passed separately
           // for cost. `result.inputTokens` alone is the uncached delta (often a
           // few hundred), which would mislabel `llm.token_count.prompt`.
@@ -234,8 +260,9 @@ export async function compileViaClaudeCli(
           outputTokens: result.outputTokens,
           cacheReadTokens: result.cacheReadInputTokens,
           cacheWriteTokens: result.cacheCreationInputTokens,
-        }),
-      });
+        },
+        { 'imprint.compile.turns': result.turns },
+      );
       return result;
     },
   );
@@ -251,11 +278,20 @@ async function compileViaClaudeCliImpl(
   opts: CompileViaClaudeCliOptions,
 ): Promise<CompileAgentResult> {
   let lastResult: CompileAgentResult | undefined;
+  let usageTotals: CompileUsageTotals = {
+    turns: 0,
+    inputTokens: 0,
+    outputTokens: 0,
+    cacheReadInputTokens: 0,
+    cacheCreationInputTokens: 0,
+  };
   for (let attempt = 1; attempt <= MAX_USAGE_POLICY_ATTEMPTS; attempt++) {
     const result = await runClaudeCliAttempt(opts);
+    usageTotals = addCompileUsageTotals(usageTotals, result);
+    const resultWithAggregateUsage = { ...result, ...usageTotals };
     const isRefusal = !result.success && USAGE_POLICY_REFUSAL.test(result.message ?? '');
-    if (!isRefusal) return result;
-    lastResult = result;
+    if (!isRefusal) return resultWithAggregateUsage;
+    lastResult = resultWithAggregateUsage;
     if (attempt < MAX_USAGE_POLICY_ATTEMPTS) {
       const backoffMs = usagePolicyBackoffMs(attempt);
       log(