fix: Fix Langchain anthropic token metrics

.changeset/langchain-anthropic-cache-metrics.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix: Fix Langchain anthropic token metrics

e2e/scenarios/cloudflare-agents-instrumentation/scenario.ts

@@ -1,16 +1,15 @@
 import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
 import { once } from "node:events";
 import { writeFile } from "node:fs/promises";
-import net from "node:net";
 import path from "node:path";
+import { stripVTControlCharacters } from "node:util";
 import {
   getTestRunId,
   runMain,
   scopedName,
 } from "../../helpers/scenario-runtime";
 
 async function main() {
-  const port = await getFreePort();
   const viteBin = path.join(
     process.cwd(),
     "node_modules",
@@ -21,7 +20,7 @@ async function main() {
   await buildWorker(viteBin);
   const server = spawn(
     viteBin,
-    ["preview", "--host", "127.0.0.1", "--port", String(port), "--strictPort"],
+    ["preview", "--host", "127.0.0.1", "--port", "0", "--strictPort"],
     {
       cwd: process.cwd(),
       env: process.env,
@@ -31,8 +30,7 @@ async function main() {
   const output = captureOutput(server);
 
   try {
-    const baseUrl = `http://127.0.0.1:${port}`;
-    await waitForServer(baseUrl, server, output);
+    const baseUrl = await waitForServer(server, output);
     const testRunId = getTestRunId();
     const projectName = scopedName(
       "e2e-cloudflare-agents-instrumentation",
@@ -99,22 +97,6 @@ async function buildWorker(viteBin: string): Promise<void> {
   }
 }
 
-async function getFreePort(): Promise<number> {
-  const server = net.createServer();
-  await new Promise<void>((resolve, reject) => {
-    server.once("error", reject);
-    server.listen(0, "127.0.0.1", resolve);
-  });
-  const address = server.address();
-  await new Promise<void>((resolve, reject) => {
-    server.close((error) => (error ? reject(error) : resolve()));
-  });
-  if (!address || typeof address === "string") {
-    throw new Error("Could not allocate a Vite preview-server port");
-  }
-  return address.port;
-}
-
 function captureOutput(child: ChildProcessWithoutNullStreams): () => string {
   let stdout = "";
   let stderr = "";
@@ -128,24 +110,29 @@ function captureOutput(child: ChildProcessWithoutNullStreams): () => string {
 }
 
 async function waitForServer(
-  baseUrl: string,
   server: ChildProcessWithoutNullStreams,
   output: () => string,
-): Promise<void> {
+): Promise<string> {
   const startedAt = Date.now();
   while (Date.now() - startedAt < 60_000) {
     if (server.exitCode !== null) {
       throw new Error(
         `Vite exited early with code ${server.exitCode}\n${output()}`,
       );
     }
-    try {
-      const response = await fetch(`${baseUrl}/health`);
-      if (response.ok) {
-        return;
+    // Vite binds port 0 and reports the assigned URL, avoiding a port reservation race.
+    const baseUrl = stripVTControlCharacters(output()).match(
+      /Local:\s+(http:\/\/127\.0\.0\.1:\d+)\//,
+    )?.[1];
+    if (baseUrl) {
+      try {
+        const response = await fetch(`${baseUrl}/health`);
+        if (response.ok) {
+          return baseUrl;
+        }
+      } catch {
+        // Continue until workerd accepts requests.
       }
-    } catch {
-      // Continue until workerd accepts requests.
     }
     await new Promise((resolve) => setTimeout(resolve, 250));
   }

js/src/wrappers/langchain/callback-handler.test.ts

@@ -129,6 +129,91 @@ describe("BraintrustLangChainCallbackHandler metrics", () => {
     });
   });
 
+  it.each([
+    {
+      name: "TTL buckets without an aggregate",
+      details: {
+        cache_creation: undefined,
+        ephemeral_5m_input_tokens: 4,
+        ephemeral_1h_input_tokens: 0,
+      },
+      expected: {
+        prompt_cache_creation_5m_tokens: 4,
+        prompt_cache_creation_1h_tokens: 0,
+      },
+    },
+    {
+      name: "both TTL buckets",
+      details: { ephemeral_5m_input_tokens: 4, ephemeral_1h_input_tokens: 6 },
+      expected: {
+        prompt_cache_creation_5m_tokens: 4,
+        prompt_cache_creation_1h_tokens: 6,
+      },
+    },
+    {
+      name: "a zero-valued TTL bucket",
+      details: { ephemeral_5m_input_tokens: 4, ephemeral_1h_input_tokens: 0 },
+      expected: {
+        prompt_cache_creation_5m_tokens: 4,
+        prompt_cache_creation_1h_tokens: 0,
+      },
+    },
+    {
+      name: "only the 5-minute bucket",
+      details: { ephemeral_5m_input_tokens: 4 },
+      expected: { prompt_cache_creation_5m_tokens: 4 },
+    },
+    {
+      name: "only the 1-hour bucket",
+      details: { ephemeral_1h_input_tokens: 6 },
+      expected: { prompt_cache_creation_1h_tokens: 6 },
+    },
+    {
+      name: "only a zero-valued bucket",
+      details: { ephemeral_1h_input_tokens: 0 },
+      expected: { prompt_cache_creation_1h_tokens: 0 },
+    },
+    {
+      name: "null TTL buckets",
+      details: {
+        ephemeral_5m_input_tokens: null,
+        ephemeral_1h_input_tokens: null,
+      },
+      expected: { prompt_cache_creation_tokens: 10 },
+    },
+  ])(
+    "preserves cache creation metrics with $name",
+    async ({ details, expected }) => {
+      const { endLog } = await finishChatModelRun({
+        generations: [
+          [
+            {
+              message: {
+                usage_metadata: {
+                  input_tokens: 20,
+                  output_tokens: 2,
+                  input_token_details: {
+                    cache_creation: 10,
+                    cache_read: 3,
+                    ...details,
+                  },
+                },
+              },
+            },
+          ],
+        ],
+      });
+
+      expect(endLog.metrics).toEqual({
+        prompt_tokens: 20,
+        completion_tokens: 2,
+        prompt_cached_tokens: 3,
+        tokens: 22,
+        ...expected,
+      });
+    },
+  );
+
   it("preserves reasoning metrics from message usage metadata", async () => {
     const { endLog } = await finishChatModelRun({
       generations: [

js/src/wrappers/langchain/callback-handler.ts

@@ -522,18 +522,25 @@ function getMetricsFromResponse(
       continue;
     }
 
-    const inputTokenDetails = usageMetadata.input_token_details;
+    const inputTokenDetails = isRecord(usageMetadata.input_token_details)
+      ? usageMetadata.input_token_details
+      : {};
     const outputTokenDetails = usageMetadata.output_token_details;
     return normalizeTokenMetrics({
       total_tokens: usageMetadata.total_tokens,
       prompt_tokens: usageMetadata.input_tokens,
       completion_tokens: usageMetadata.output_tokens,
-      prompt_cache_creation_tokens: isRecord(inputTokenDetails)
-        ? inputTokenDetails.cache_creation
-        : undefined,
-      prompt_cached_tokens: isRecord(inputTokenDetails)
-        ? inputTokenDetails.cache_read
-        : undefined,
+      // Prefer TTL-specific cache writes over the aggregate when available.
+      prompt_cache_creation_tokens:
+        inputTokenDetails.ephemeral_5m_input_tokens == null &&
+        inputTokenDetails.ephemeral_1h_input_tokens == null
+          ? inputTokenDetails.cache_creation
+          : undefined,
+      prompt_cache_creation_5m_tokens:
+        inputTokenDetails.ephemeral_5m_input_tokens,
+      prompt_cache_creation_1h_tokens:
+        inputTokenDetails.ephemeral_1h_input_tokens,
+      prompt_cached_tokens: inputTokenDetails.cache_read,
       completion_reasoning_tokens: isRecord(outputTokenDetails)
         ? outputTokenDetails.reasoning
         : undefined,