feat: Instrument Google GenAI multimodal APIs

.changeset/instrument-google-genai-multimodal-apis.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+feat: Instrument Google GenAI multimodal APIs

e2e/scenarios/cloudflare-ai-chat-instrumentation/scenario.mjs

@@ -1,9 +1,7 @@
 import { spawn } from "node:child_process";
-import net from "node:net";
+import { stripVTControlCharacters } from "node:util";
 
 const scenarioDir = new URL(".", import.meta.url).pathname;
-const port = await reservePort();
-const baseUrl = `http://127.0.0.1:${port}`;
 const build = await runCommand([
   "exec",
   "vite",
@@ -25,7 +23,7 @@ const vite = spawn(
     "--host",
     "127.0.0.1",
     "--port",
-    String(port),
+    "0",
     "--strictPort",
   ],
   {
@@ -41,15 +39,15 @@ vite.stdout.on("data", (chunk) => (output += chunk.toString()));
 vite.stderr.on("data", (chunk) => (output += chunk.toString()));
 
 try {
-  await waitForServer();
-  const success = await run("success");
+  const baseUrl = await waitForServer();
+  const success = await run(baseUrl, "success");
   if (!JSON.stringify(success).includes("CLOUDFLARE_AI_CHAT_TOOL_OK")) {
     throw new Error(
       `Successful chat result was incomplete: ${JSON.stringify(success)}`,
     );
   }
 
-  const failure = await run("error");
+  const failure = await run(baseUrl, "error");
   if (!JSON.stringify(failure).includes("CLOUDFLARE_AI_CHAT_STREAM_ERROR")) {
     throw new Error(
       `Error chat result was incomplete: ${JSON.stringify(failure)}`,
@@ -59,7 +57,7 @@ try {
   await stopVite();
 }
 
-async function run(kind) {
+async function run(baseUrl, kind) {
   const response = await fetch(`${baseUrl}/run?kind=${kind}`, {
     signal: AbortSignal.timeout(60_000),
   });
@@ -79,14 +77,21 @@ async function waitForServer() {
         `Vite exited before startup (${vite.exitCode}):\n${output}`,
       );
     }
-    try {
-      const response = await fetch(`${baseUrl}/health`, {
-        signal: AbortSignal.timeout(1_000),
-      });
-      if (response.ok) {
-        return;
+    const baseUrl = stripVTControlCharacters(output).match(
+      /Local:\s+(http:\/\/127\.0\.0\.1:\d+)\//,
+    )?.[1];
+    if (baseUrl) {
+      try {
+        const response = await fetch(`${baseUrl}/health`, {
+          signal: AbortSignal.timeout(1_000),
+        });
+        if (response.ok) {
+          return baseUrl;
+        }
+      } catch {
+        // Continue until workerd accepts requests.
       }
-    } catch {}
+    }
     await new Promise((resolve) => setTimeout(resolve, 250));
   }
   throw new Error(`Timed out waiting for Vite:\n${output}`);
@@ -118,18 +123,6 @@ async function stopVite() {
   }
 }
 
-async function reservePort() {
-  const server = net.createServer();
-  await new Promise((resolve, reject) => {
-    server.once("error", reject);
-    server.listen(0, "127.0.0.1", resolve);
-  });
-  const address = server.address();
-  const selected = typeof address === "object" && address ? address.port : 0;
-  await new Promise((resolve) => server.close(resolve));
-  return selected;
-}
-
 async function runCommand(args) {
   return await new Promise((resolve, reject) => {
     const child = spawn("pnpm", args, {