feat: Dynamic LLM API Pricing

src/server.ts

@@ -29,6 +29,7 @@ import {
   getAvailableLanguages,
   hasBunRuntime,
 } from "./runtime.js";
+import { refreshPricingCacheInBackground, updatePricingCache } from "./session/analytics.js";
 import { classifyNonZeroExit } from "./exit-classify.js";
 import { startLifecycleGuard, noteMcpActivity, noteRequestStart, noteRequestEnd, attachMcpActivityTap } from "./lifecycle.js";
 import { charSafePrefix } from "./truncate.js";
@@ -3933,6 +3934,31 @@ function createMinimalDb(): import("./session/analytics.js").DatabaseAdapter {
   };
 }
 
+server.registerTool(
+  "ctx_set_model",
+  {
+    title: "Set Tracking Model",
+    description: "Update the active language model used for calculating context savings.",
+    inputSchema: z.object({
+      modelName: z.string().describe("The name of the new model (e.g., 'Gemini 3.1 Pro (Low)')."),
+    }),
+  },
+  async (args: { modelName: string }) => {
+    try {
+      const modelName = String(args.modelName);
+      const result = await updatePricingCache(modelName);
+      return trackResponse("ctx_set_model", {
+        content: [{ type: "text", text: result }],
+      });
+    } catch (e: any) {
+      return trackResponse("ctx_set_model", {
+        content: [{ type: "text", text: `Failed to update model: ${e.message}` }],
+        isError: true,
+      });
+    }
+  }
+);
+
 server.registerTool(
   "ctx_stats",
   {
@@ -4966,6 +4992,9 @@ async function main() {
   // statusline staleness threshold is 30min (cliff is 30 missed ticks away).
   setInterval(() => persistStats(), 60_000).unref();
 
+  // Fire-and-forget background pricing fetch to ensure the cache stays fresh
+  refreshPricingCacheInBackground();
+
   if (process.stdin.isTTY) {
     console.error(`Context Mode MCP server v${VERSION} running on stdio`);
     console.error(`Detected runtimes:\n${getRuntimeSummary(runtimes)}`);

src/session/analytics.ts

@@ -2479,6 +2479,23 @@ function fmtNum(n: number): string {
 // Pricing (Bug #6) — Anthropic Opus input rate
 // ─────────────────────────────────────────────────────────
 
+import { resolveSessionStorageDir, resolveDefaultSessionDir } from "./db.js";
+import { readFileSync, writeFileSync, renameSync } from "node:fs";
+import { request as httpsRequest } from "node:https";
+
+function getPricingCache() {
+  try {
+    const cachePath = join(resolveSessionStorageDir(() => resolveDefaultSessionDir({ configDir: resolveClaudeConfigDir() })).path, "pricing-cache.json");
+    if (existsSync(cachePath)) {
+      const data = JSON.parse(readFileSync(cachePath, "utf8"));
+      return data;
+    }
+  } catch (e) {
+    // Ignore cache read errors
+  }
+  return null;
+}
+
 // ── Pricing (Bug #6) — per-token USD rate ─────────────────
 // Reads PI_CONTEXT_MODE_PRICE_OUTPUT_PER_TOKEN when set by a Pi host;
 // falls back to the Opus 4.7/4.8 input rate ($5/1M) for all other adapters.
@@ -2499,14 +2516,109 @@ function fmtNum(n: number): string {
  * ($5 per 1M tokens) otherwise.
  */
 export function pricePerToken(): number {
-  const env = process.env.PI_CONTEXT_MODE_PRICE_OUTPUT_PER_TOKEN;
+  const env = process.env.PI_CONTEXT_MODE_PRICE_OUTPUT_PER_TOKEN || process.env.CONTEXT_MODE_PRICE_PER_TOKEN;
   if (env !== undefined && env !== "") {
     const parsed = Number(env);
     if (Number.isFinite(parsed) && parsed > 0) return parsed;
   }
+
+  const cache = getPricingCache();
+  if (cache && typeof cache.pricePerToken === 'number') {
+    return cache.pricePerToken;
+  }
+
   return 5 / 1_000_000; // Opus 4.7/4.8 input fallback
 }
 
+/**
+ * Returns the model name for pricing labels.
+ */
+export function pricingModelName(): string {
+  if (process.env.CONTEXT_MODE_MODEL_NAME) return process.env.CONTEXT_MODE_MODEL_NAME;
+  if (process.env.PI_CONTEXT_MODE_PRICE_OUTPUT_PER_TOKEN || process.env.CONTEXT_MODE_PRICE_PER_TOKEN) return "Custom";
+
+  const cache = getPricingCache();
+  if (cache && typeof cache.modelName === 'string') {
+    return cache.modelName;
+  }
+
+  return "Opus";
+}
+
+/**
+ * Core async function to fetch pricing and update cache.
+ * Returns a confirmation string.
+ */
+export function updatePricingCache(modelNameOverride?: string): Promise<string> {
+  return new Promise((resolve, reject) => {
+    // 1. Determine model name to search for
+    const cache = getPricingCache();
+    // Use override > env var > existing cache model > fallback
+    let modelName = modelNameOverride || process.env.CONTEXT_MODE_MODEL_NAME;
+    if (!modelName && cache && cache.modelName) {
+      modelName = cache.modelName;
+    }
+
+    if (!modelName) {
+      return reject(new Error("No model name configured"));
+    }
+
+    // Fuzzy matching logic: strip (Low), (High), etc.
+    const cleanName = modelName.replace(/\s*\([^)]*\)/g, '').toLowerCase().trim();
+
+    const req = httpsRequest("https://llmpricingapi.com/api/models", (res) => {
+      let data = "";
+      res.on("data", (chunk) => data += chunk);
+      res.on("end", () => {
+        try {
+          const parsed = JSON.parse(data);
+          if (!parsed.models || !Array.isArray(parsed.models)) {
+            return reject(new Error("Invalid API response format"));
+          }
+
+          // Fuzzy match against API
+          const match = parsed.models.find((m: any) => 
+            m.name && m.name.toLowerCase().includes(cleanName)
+          );
+
+          if (match && typeof match.input_price === "number") {
+            const cacheData = {
+              modelName: match.name, // Use the API's clean name in the cache
+              pricePerToken: match.input_price / 1_000_000
+            };
+
+            const dir = resolveSessionStorageDir(() => resolveDefaultSessionDir({ configDir: resolveClaudeConfigDir() })).path;
+            const cachePath = join(dir, "pricing-cache.json");
+            const tempPath = cachePath + ".tmp";
+
+            writeFileSync(tempPath, JSON.stringify(cacheData));
+            renameSync(tempPath, cachePath);
+
+            resolve(`Successfully updated tracking to ${match.name} ($${match.input_price} / 1M tokens).`);
+          } else {
+            reject(new Error(`Model '${modelName}' (fuzzy: '${cleanName}') not found in pricing API.`));
+          }
+        } catch (e) {
+          reject(e);
+        }
+      });
+    });
+
+    req.on("error", (err) => reject(err));
+    req.end();
+  });
+}
+
+/**
+ * Asynchronously fetches pricing from llmpricingapi.com and updates the local cache.
+ * Designed to be fired-and-forget during server startup.
+ */
+export function refreshPricingCacheInBackground(): void {
+  updatePricingCache().catch(() => {
+    // silently fail for background fetch
+  });
+}
+
 /**
  * Back-compat alias for the original Opus-rate const (PR #401 architect
  * P1.1 — single source of truth). Kept as a literal so any third-party
@@ -2988,9 +3100,9 @@ export function formatReport(
   // ── Active session: visual savings dashboard ──
 
   // Line 1: Hero metric — the screenshottable number
-  // Bug #6: include Opus pricing on the hero line for credibility.
+  // Bug #6: include pricing on the hero line for credibility.
   lines.push(
-    `${fmtNum(tokensSaved)} tokens saved  ·  ${savingsPct.toFixed(1)}% reduction  ·  ${duration}  ·  ~${tokensToUsd(tokensSaved)} saved (Opus)`,
+    `${fmtNum(tokensSaved)} tokens saved  ·  ${savingsPct.toFixed(1)}% reduction  ·  ${duration}  ·  ~${tokensToUsd(tokensSaved)} saved (${pricingModelName()})`,
   );
   lines.push("");
 

tests/core/server.test.ts

@@ -5525,7 +5525,7 @@
           // `"`. Match any of those legitimate prefixes explicitly so both
           // template-literal and string-concat description shapes pass.
           const hasWhen = /(?:\\n|^|\s|")WHEN(?:\s+TO\s+USE)?:/.test(tool.description);
           expect(hasWhen, `${tool.name} description (src/server.ts:${tool.lineNo}) must contain a WHEN: section`).toBe(true);
         });
 
         // ── Canonical structure assertions (PR #683 WS3) ─────────────
@@ -5536,7 +5536,7 @@
             expect(
               flat.includes(section + ":"),
               `${tool.name} (src/server.ts:${tool.lineNo}) missing mandatory section '${section}:' per ADR-0002 canonical structure.`,
             ).toBe(true);
           }
         });
 
@@ -6684,6 +6684,7 @@
     ctx_upgrade:         { readOnlyHint: false, destructiveHint: false, idempotentHint: true,  openWorldHint: false },
     ctx_purge:           { readOnlyHint: false, destructiveHint: true,  idempotentHint: true,  openWorldHint: false },
     ctx_insight:         { readOnlyHint: false, destructiveHint: false, idempotentHint: true,  openWorldHint: true  },
+    ctx_set_model:       { readOnlyHint: false, destructiveHint: false, idempotentHint: true,  openWorldHint: false },
   };
   const tools = REGISTERED_CTX_TOOLS as Array<{ name: string; config: { annotations?: Hints } }>;
   const find = (name: string) => tools.find((t) => t.name === name);
@@ -6696,7 +6697,7 @@
     for (const [name, hints] of Object.entries(EXPECTED)) {
       const tool = find(name);
       expect(tool, `${name} not registered`).toBeDefined();
       expect(tool!.config.annotations, `${name} missing annotations`).toBeDefined();
       expect(tool!.config.annotations).toMatchObject(hints);
     }
   });