feat: trace the remaining request parameters pi puts on the wire

src/index.ts

@@ -291,6 +291,14 @@ export function toMultimodalContent(
 }
 
 const MAX_TOKENS_KEYS = ["max_tokens", "max_completion_tokens", "max_output_tokens", "maxOutputTokens", "maxTokens"];
+const MAX_PARAM_CHARS = 200;
+const THINKING_BUDGET_KEYS = [
+  "thinking_token_budget",
+  "thinking_budget",
+  "thinking_budget_tokens",
+  "budget_tokens",
+  "thinkingBudget",
+];
 
 function findNumber(bag: unknown, keys: string[], depth = 2): number | undefined {
   if (!bag || typeof bag !== "object" || Array.isArray(bag)) return undefined;
@@ -307,16 +315,53 @@ function findNumber(bag: unknown, keys: string[], depth = 2): number | undefined
   return undefined;
 }
 
+function pickCacheRetention(payload: unknown): string | undefined {
+  const bag = payload as Record<string, unknown> | undefined;
+  const flat = bag?.prompt_cache_retention;
+  if (typeof flat === "string") return flat;
+  const system = bag?.system;
+  if (!Array.isArray(system)) return undefined;
+  for (const block of system) {
+    const ttl = (block as { cache_control?: { ttl?: unknown } } | undefined)?.cache_control?.ttl;
+    if (typeof ttl === "string") return ttl;
+  }
+  return undefined;
+}
+
+function pickToolChoice(payload: unknown): string | undefined {
+  const value = (payload as Record<string, unknown> | undefined)?.tool_choice;
+  if (typeof value === "string") return value;
+  const type = (value as { type?: unknown } | undefined)?.type;
+  return typeof type === "string" ? type : undefined;
+}
+
 export function extractModelParameters(
   payload: unknown,
-  model: { reasoning?: boolean } | undefined,
+  model: { reasoning?: boolean; samplingParams?: Record<string, unknown> } | undefined,
   thinkingLevel?: string,
 ): Record<string, string | number> | undefined {
   const out: Record<string, string | number> = {};
   try {
     const maxTokens = findNumber(payload, MAX_TOKENS_KEYS);
     if (maxTokens !== undefined) out.max_tokens = maxTokens;
     if (model?.reasoning && thinkingLevel && thinkingLevel !== "off") out.thinking_level = thinkingLevel;
+    const thinkingBudget = findNumber(payload, THINKING_BUDGET_KEYS);
+    if (thinkingBudget !== undefined) out.thinking_budget_tokens = thinkingBudget;
+    const cacheRetention = pickCacheRetention(payload);
+    if (cacheRetention !== undefined) out.prompt_cache_retention = cacheRetention;
+    const serviceTier = (payload as Record<string, unknown> | undefined)?.service_tier;
+    if (typeof serviceTier === "string") out.service_tier = serviceTier;
+    const toolChoice = pickToolChoice(payload);
+    if (toolChoice !== undefined) out.tool_choice = toolChoice;
+    for (const key of Object.keys(model?.samplingParams ?? {})) {
+      const value = (payload as Record<string, unknown> | undefined)?.[key];
+      if (value === undefined) continue;
+      if (typeof value === "number") out[key] = value;
+      else {
+        const text = typeof value === "string" ? value : JSON.stringify(value);
+        if (text !== undefined && text.length <= MAX_PARAM_CHARS) out[key] = text;
+      }
+    }
   } catch {}
   return Object.keys(out).length ? out : undefined;
 }
@@ -710,7 +755,7 @@ export default function (pi: ExtensionAPI) {
         modelParameters: extractModelParameters(event.payload, ctx.model, ctx.thinkingLevel),
         metadata: {
           assistant_index: index - 1,
-          ...(ctx.model ? { provider: ctx.model.provider } : {}),
+          ...(ctx.model ? { provider: ctx.model.provider, context_window: ctx.model.contextWindow } : {}),
         },
       },
       { asType: "generation" },

test/helpers.ts

@@ -75,7 +75,13 @@ function streamChunks(
 /** Usage the mock reports for a compaction/branch summarization call. */
 export const SUMMARIZATION_USAGE = { prompt: 3571, completion: 313 };
 
-export function startMockProvider(): Promise<{ port: number; close: () => void }> {
+export function startMockProvider(): Promise<{
+  port: number;
+  close: () => void;
+  /** Request bodies the provider received, in call order. */
+  payloads: () => Array<Record<string, unknown>>;
+}> {
+  const received: Array<Record<string, unknown>> = [];
   const server = http.createServer((req, res) => {
     let body = "";
     req.on("data", (c) => (body += c));
@@ -85,6 +91,7 @@ export function startMockProvider(): Promise<{ port: number; close: () => void }
         messages?: OpenAiMessage[];
         tools?: Array<{ function?: { name?: string } }>;
       };
+      received.push(payload as Record<string, unknown>);
       const messages = payload.messages ?? [];
       const model = payload.model ?? "mock-gpt-1";
       const lastUserIdx = messages.map((m) => m.role).lastIndexOf("user");
@@ -169,7 +176,7 @@ export function startMockProvider(): Promise<{ port: number; close: () => void }
   return new Promise((resolvePromise) => {
     server.listen(0, "127.0.0.1", () => {
       const port = (server.address() as { port: number }).port;
-      resolvePromise({ port, close: () => server.close() });
+      resolvePromise({ port, close: () => server.close(), payloads: () => [...received] });
     });
   });
 }

test/model-parameters.test.ts

@@ -16,6 +16,7 @@ import {
 
 describe("extractModelParameters", () => {
   const plain = { reasoning: false };
+  const PARAM_CHAR_CAP = 200;
 
   it("reads the cap the payload really carries, whatever the dialect calls it", () => {
     const payloads: Array<[string, unknown]> = [
@@ -65,6 +66,107 @@ describe("extractModelParameters", () => {
     assert.deepEqual(extractModelParameters({}, { reasoning: true }, "medium"), { thinking_level: "medium" });
   });
 
+  it("reads the thinking budget the provider was really given", () => {
+    const reasoning = { reasoning: true };
+    const payloads: Array<[string, unknown]> = [
+      ["anthropic-messages", { max_tokens: 8192, thinking: { type: "enabled", budget_tokens: 7168 } }],
+      [
+        "bedrock-converse-stream",
+        {
+          inferenceConfig: { maxTokens: 8192 },
+          additionalModelRequestFields: { thinking: { type: "enabled", budget_tokens: 7168 } },
+        },
+      ],
+      ["google", { config: { maxOutputTokens: 8192, thinkingConfig: { thinkingBudget: 7168 } } }],
+      ["openai-completions (vLLM)", { max_tokens: 8192, thinking_token_budget: 7168 }],
+      ["openai-completions (thinking_budget)", { max_tokens: 8192, thinking_budget: 7168 }],
+      ["openai-completions (thinking_budget_tokens)", { max_tokens: 8192, thinking_budget_tokens: 7168 }],
+    ];
+    for (const [dialect, payload] of payloads) {
+      assert.deepEqual(
+        extractModelParameters(payload, reasoning, "high"),
+        { max_tokens: 8192, thinking_level: "high", thinking_budget_tokens: 7168 },
+        dialect,
+      );
+    }
+  });
+
+  it("omits the budget when the provider gets none", () => {
+    assert.deepEqual(
+      extractModelParameters({ max_tokens: 8192, thinking: { type: "adaptive" }, output_config: { effort: "high" } }, { reasoning: true }, "high"),
+      { max_tokens: 8192, thinking_level: "high" },
+    );
+    assert.deepEqual(extractModelParameters({ config: { thinkingConfig: { thinkingBudget: 0 } } }, plain), undefined);
+  });
+
+  it("reports the prompt cache retention both dialects express differently", () => {
+    assert.deepEqual(extractModelParameters({ prompt_cache_retention: "24h" }, plain), {
+      prompt_cache_retention: "24h",
+    });
+    assert.deepEqual(
+      extractModelParameters({ max_tokens: 8192, system: [{ type: "text", cache_control: { type: "ephemeral", ttl: "1h" } }] }, plain),
+      { max_tokens: 8192, prompt_cache_retention: "1h" },
+    );
+    assert.deepEqual(
+      extractModelParameters({ max_tokens: 8192, system: [{ type: "text", cache_control: { type: "ephemeral" } }] }, plain),
+      { max_tokens: 8192 },
+    );
+  });
+
+  it("reports the service tier and the tool choice when they are set", () => {
+    assert.deepEqual(extractModelParameters({ service_tier: "flex", tool_choice: "auto" }, plain), {
+      service_tier: "flex",
+      tool_choice: "auto",
+    });
+    assert.deepEqual(extractModelParameters({ tool_choice: { type: "any" } }, plain), { tool_choice: "any" });
+    assert.deepEqual(extractModelParameters({ tool_choice: { type: "function", function: { name: "read" } } }, plain), {
+      tool_choice: "function",
+    });
+  });
+
+  it("reports the sampling parameters the model declares, with the values the wire carried", () => {
+    const samplingParams = { temperature: 0.2, top_p: 0.9, top_k: 40, min_p: 0.05, repetition_penalty: 1.1 };
+    const onTheWire = { temperature: 0.7, top_p: 0.5, top_k: 10, min_p: 0.01, repetition_penalty: 1.9 };
+    assert.deepEqual(extractModelParameters({ max_tokens: 8192, ...onTheWire }, { reasoning: false, samplingParams }), {
+      max_tokens: 8192,
+      ...onTheWire,
+    });
+  });
+
+  it("serializes a non-scalar sampling parameter instead of dropping it", () => {
+    const samplingParams = { stop: ["</done>"], seed: 7 };
+    assert.deepEqual(extractModelParameters({ ...samplingParams }, { reasoning: false, samplingParams }), {
+      stop: '["</done>"]',
+      seed: 7,
+    });
+  });
+
+  it("never lets a sampling parameter pull a payload structure onto the span", () => {
+    const conversation = [{ role: "user", content: "private user text ".repeat(40) }];
+    for (const key of ["messages", "system", "tools", "input", "contents"]) {
+      const out = extractModelParameters(
+        { max_tokens: 8192, [key]: conversation },
+        { reasoning: false, samplingParams: { [key]: 1 } },
+      );
+      assert.deepEqual(out, { max_tokens: 8192 }, key);
+    }
+    const long = "x".repeat(PARAM_CHAR_CAP + 1);
+    assert.deepEqual(extractModelParameters({ note: long }, { reasoning: false, samplingParams: { note: 1 } }), undefined);
+    const short = "x".repeat(PARAM_CHAR_CAP);
+    assert.deepEqual(extractModelParameters({ note: short }, { reasoning: false, samplingParams: { note: 1 } }), {
+      note: short,
+    });
+  });
+
+  it("keeps the chips it already collected when a late sampling parameter blows up", () => {
+    const cyclic: Record<string, unknown> = {};
+    cyclic.self = cyclic;
+    assert.deepEqual(
+      extractModelParameters({ max_tokens: 1477, weird: cyclic }, { reasoning: true, samplingParams: { weird: 1 } }, "high"),
+      { max_tokens: 1477, thinking_level: "high" },
+    );
+  });
+
   it("rejects a cap that is not a positive whole number", () => {
     for (const bad of [Number.NaN, Number.POSITIVE_INFINITY, -5, 1.5, "8192", null, true]) {
       assert.equal(extractModelParameters({ max_tokens: bad }, plain), undefined, String(bad));
@@ -80,7 +182,13 @@ describe("extractModelParameters", () => {
       },
     });
     assert.equal(extractModelParameters(hostile, plain), undefined);
-    assert.equal(extractModelParameters(hostile, { reasoning: true }, "high"), undefined);
+  });
+
+  it("omits a declared sampling parameter the dialect never sent", () => {
+    assert.deepEqual(
+      extractModelParameters({ max_tokens: 8192 }, { reasoning: false, samplingParams: { temperature: 0.2 } }),
+      { max_tokens: 8192 },
+    );
   });
 
   it("never descends into message arrays looking for a cap", () => {
@@ -182,4 +290,81 @@ describe("integration: model parameters", () => {
       capture.close();
     }
   });
+
+  it("records the long prompt cache retention pi put on the wire", async () => {
+    const capture = await startCaptureServer();
+    try {
+      const sandbox = createSandbox(mock.port);
+      const modelsPath = join(sandbox.agentDir, "models.json");
+      const models = JSON.parse(readFileSync(modelsPath, "utf8"));
+      models.providers.mock.models[0].compat = { supportsLongCacheRetention: true };
+      writeFileSync(modelsPath, JSON.stringify(models));
+      const result = await runPi(sandbox, "Explore this project and summarize it", {
+        env: { ...langfuseEnv(capture), PI_CACHE_RETENTION: "long" },
+      });
+      assert.equal(result.status, 0, `pi failed: ${result.stderr}`);
+      await waitForRequests(capture, 1);
+
+      const params = generationParameters(capture.spans());
+      assert.equal(params.length, 3);
+      for (const p of params) {
+        assert.deepEqual(p, { max_tokens: 8192, prompt_cache_retention: "24h" });
+      }
+    } finally {
+      capture.close();
+    }
+  });
+
+  it("reports exactly the sampling parameters the request carried", async () => {
+    const capture = await startCaptureServer();
+    const wired = await startMockProvider();
+    try {
+      const sandbox = createSandbox(wired.port);
+      const modelsPath = join(sandbox.agentDir, "models.json");
+      const models = JSON.parse(readFileSync(modelsPath, "utf8"));
+      const declared: Record<string, unknown> = { temperature: 0.2, top_k: 40, stop: ["</done>"] };
+      models.providers.mock.models[0].samplingParams = declared;
+      writeFileSync(modelsPath, JSON.stringify(models));
+      const result = await runPi(sandbox, "Explore this project and summarize it", { env: langfuseEnv(capture) });
+      assert.equal(result.status, 0, `pi failed: ${result.stderr}`);
+      await waitForRequests(capture, 1);
+
+      const sent = wired.payloads();
+      const params = generationParameters(capture.spans());
+      assert.equal(params.length, sent.length, "one generation per provider request");
+      const keys = Object.keys(declared);
+      for (const [i, p] of params.entries()) {
+        assert.equal(p?.max_tokens, 8192, `generation ${i + 1} must still carry the cap`);
+        assert.deepEqual(
+          keys.filter((k) => p?.[k] !== undefined),
+          keys.filter((k) => sent[i]?.[k] !== undefined),
+          `generation ${i + 1} must report exactly the declared parameters the wire carried`,
+        );
+        assert.equal(p?.temperature, sent[i]?.temperature, `generation ${i + 1} temperature`);
+        assert.equal(p?.top_k, sent[i]?.top_k, `generation ${i + 1} top_k`);
+      }
+    } finally {
+      wired.close();
+      capture.close();
+    }
+  });
+
+  it("puts the context window on every generation", async () => {
+    const capture = await startCaptureServer();
+    try {
+      const sandbox = createSandbox(mock.port, { contextWindow: 18000 });
+      const result = await runPi(sandbox, "Explore this project and summarize it", { env: langfuseEnv(capture) });
+      assert.equal(result.status, 0, `pi failed: ${result.stderr}`);
+      await waitForRequests(capture, 1);
+
+      const generations = capture.spans().filter((s) => s.name === "LLM Call");
+      assert.equal(generations.length, 3);
+      for (const g of generations) {
+        assert.equal(g.attrs["langfuse.observation.metadata.provider"], "mock");
+        assert.equal(Number(g.attrs["langfuse.observation.metadata.context_window"]), 18000);
+      }
+    } finally {
+      capture.close();
+    }
+  });
 });