feat(ai-sdk): Add AI SDK v7 support

.changeset/hip-crews-wash.md

@@ -0,0 +1,5 @@
+---
+"braintrust": minor
+---
+
+feat(ai-sdk): Add AI SDK v7 support

e2e/config/pr-comment-scenarios.json

@@ -145,7 +145,8 @@
       { "variantKey": "ai-sdk-v3", "label": "v3" },
       { "variantKey": "ai-sdk-v4", "label": "v4" },
       { "variantKey": "ai-sdk-v5", "label": "v5" },
-      { "variantKey": "ai-sdk-v6", "label": "v6" }
+      { "variantKey": "ai-sdk-v6", "label": "v6" },
+      { "variantKey": "ai-sdk-v7", "label": "v7" }
     ]
   },
   {

e2e/helpers/normalize.ts

@@ -15,6 +15,14 @@ type TokenMaps = {
   xacts: Map<string, string>;
 };
 
+export type NormalizeOptions = {
+  additionalProviderIdKeys?: Iterable<string>;
+};
+
+type ResolvedNormalizeOptions = {
+  providerIdKeys: Set<string>;
+};
+
 const ISO_DATE_REGEX = /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?Z$/;
 const ISO_DATE_SUBSTRING_REGEX =
   /(?<![A-Za-z0-9_-])\d{4}-\d{2}-\d{2}(?:[Tt ]\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:?\d{2})?)?(?![A-Za-z0-9_-])/g;
@@ -205,6 +213,7 @@ function shouldNormalizeNodeInternalStyleCaller(
 function normalizeObject(
   value: { [key: string]: Json },
   tokenMaps: TokenMaps,
+  options: ResolvedNormalizeOptions,
 ): Json {
   const callerFilename =
     typeof value.caller_filename === "string"
@@ -227,7 +236,7 @@ function normalizeObject(
         }
       }
 
-      return [key, normalizeValue(entry as Json, tokenMaps, key)];
+      return [key, normalizeValue(entry as Json, tokenMaps, options, key)];
     }),
   );
 }
@@ -250,22 +259,23 @@ function tokenFor(
 function normalizeValue(
   value: Json,
   tokenMaps: TokenMaps,
+  options: ResolvedNormalizeOptions,
   currentKey?: string,
 ): Json {
   if (Array.isArray(value)) {
     if (currentKey === "span_parents") {
       return value.map((entry) =>
         typeof entry === "string"
           ? tokenFor(tokenMaps.ids, entry, "span")
-          : normalizeValue(entry, tokenMaps),
+          : normalizeValue(entry, tokenMaps, options),
       );
     }
 
-    return value.map((entry) => normalizeValue(entry, tokenMaps));
+    return value.map((entry) => normalizeValue(entry, tokenMaps, options));
   }
 
   if (value && typeof value === "object") {
-    return normalizeObject(value, tokenMaps);
+    return normalizeObject(value, tokenMaps, options);
   }
 
   if (typeof value === "number") {
@@ -339,7 +349,7 @@ function normalizeValue(
       return tokenFor(tokenMaps.runs, value, "run");
     }
 
-    if (currentKey && PROVIDER_ID_KEYS.has(currentKey)) {
+    if (currentKey && options.providerIdKeys.has(currentKey)) {
       return tokenFor(tokenMaps.ids, value, currentKey);
     }
 
@@ -400,10 +410,28 @@ function normalizeValue(
   return value;
 }
 
-export function normalizeForSnapshot(value: Json): Json {
-  return normalizeValue(value, {
-    ids: new Map(),
-    runs: new Map(),
-    xacts: new Map(),
-  });
+function resolveNormalizeOptions(
+  options: NormalizeOptions | undefined,
+): ResolvedNormalizeOptions {
+  return {
+    providerIdKeys: new Set([
+      ...PROVIDER_ID_KEYS,
+      ...(options?.additionalProviderIdKeys ?? []),
+    ]),
+  };
+}
+
+export function normalizeForSnapshot(
+  value: Json,
+  options?: NormalizeOptions,
+): Json {
+  return normalizeValue(
+    value,
+    {
+      ids: new Map(),
+      runs: new Map(),
+      xacts: new Map(),
+    },
+    resolveNormalizeOptions(options),
+  );
 }

e2e/helpers/span-tree.ts

@@ -1,6 +1,10 @@
 import { matchFileSnapshot } from "./file-snapshot";
 import type { CapturedLogEvent } from "./mock-braintrust-server";
-import { normalizeForSnapshot, type Json } from "./normalize";
+import {
+  normalizeForSnapshot,
+  type Json,
+  type NormalizeOptions,
+} from "./normalize";
 
 export type SpanTreeFields = Record<string, unknown>;
 
@@ -25,6 +29,10 @@ type SpanTreeJsonNode = {
   type?: string;
 } & Record<string, Json | SpanTreeJsonNode[] | string | undefined>;
 
+type SpanTreeSnapshotOptions = {
+  normalize?: NormalizeOptions;
+};
+
 const FIELD_ORDER = [
   "span_attributes",
   "input",
@@ -53,9 +61,14 @@ function jsonClone(value: unknown): Json | undefined {
   return JSON.parse(JSON.stringify(value)) as Json;
 }
 
-function normalizeJson(value: unknown): Json | undefined {
+function normalizeJson(
+  value: unknown,
+  options?: NormalizeOptions,
+): Json | undefined {
   const cloned = jsonClone(value);
-  return cloned === undefined ? undefined : normalizeForSnapshot(cloned);
+  return cloned === undefined
+    ? undefined
+    : normalizeForSnapshot(cloned, options);
 }
 
 function sortJsonKeys(value: Json): Json {
@@ -182,11 +195,14 @@ function shouldRenderField(value: Json | undefined): value is Json {
   return true;
 }
 
-function normalizeFields(fields: SpanTreeFields): Record<string, Json> {
+function normalizeFields(
+  fields: SpanTreeFields,
+  options?: NormalizeOptions,
+): Record<string, Json> {
   const normalized: Record<string, Json> = {};
 
   for (const [key, value] of Object.entries(fields)) {
-    const normalizedValue = normalizeJson(value);
+    const normalizedValue = normalizeJson(value, options);
     if (shouldRenderField(normalizedValue)) {
       normalized[key] = sortJsonKeys(normalizedValue);
     }
@@ -298,12 +314,13 @@ function renderEntry(
   entry: NormalizedEntry,
   prefix: string,
   isLast: boolean,
+  options?: SpanTreeSnapshotOptions,
 ): string[] {
   const connector = isLast ? "└── " : "├── ";
   const childPrefix = `${prefix}${isLast ? "    " : "│   "}`;
   const type = entry.event.span.type ? ` [${entry.event.span.type}]` : "";
   const lines = [`${prefix}${connector}${entry.name}${type}`];
-  const fields = normalizeFields(entry.fields);
+  const fields = normalizeFields(entry.fields, options?.normalize);
 
   for (const key of sortedFieldKeys(fields)) {
     lines.push(
@@ -313,7 +330,12 @@ function renderEntry(
 
   entry.children.forEach((child, index) => {
     lines.push(
-      ...renderEntry(child, childPrefix, index === entry.children.length - 1),
+      ...renderEntry(
+        child,
+        childPrefix,
+        index === entry.children.length - 1,
+        options,
+      ),
     );
   });
 
@@ -322,40 +344,45 @@ function renderEntry(
 
 export function formatSpanTreeSnapshot(
   entries: readonly (CapturedLogEvent | SpanTreeEntry)[],
+  options?: SpanTreeSnapshotOptions,
 ): string {
   const roots = buildEntries(entries);
   return [
     "span_tree:",
     ...roots.flatMap((entry, index) =>
-      renderEntry(entry, "", index === roots.length - 1),
+      renderEntry(entry, "", index === roots.length - 1, options),
     ),
     "",
   ].join("\n");
 }
 
-function jsonEntry(entry: NormalizedEntry): SpanTreeJsonNode {
+function jsonEntry(
+  entry: NormalizedEntry,
+  options?: SpanTreeSnapshotOptions,
+): SpanTreeJsonNode {
   const node: SpanTreeJsonNode = {
     name: entry.name,
     ...(entry.event.span.type ? { type: entry.event.span.type } : {}),
     children: [],
   };
-  const fields = normalizeFields(entry.fields);
+  const fields = normalizeFields(entry.fields, options?.normalize);
 
   for (const key of sortedFieldKeys(fields)) {
     node[fieldLabel(key)] = fields[key] as Json;
   }
 
-  node.children = entry.children.map((child) => jsonEntry(child));
+  node.children = entry.children.map((child) => jsonEntry(child, options));
   return node;
 }
 
 export function formatSpanTreeJsonSnapshot(
   entries: readonly (CapturedLogEvent | SpanTreeEntry)[],
+  options?: SpanTreeSnapshotOptions,
 ): string {
   const roots = buildEntries(entries);
   return `${JSON.stringify(
     {
-      span_tree: roots.map((entry) => jsonEntry(entry)),
+      span_tree: roots.map((entry) => jsonEntry(entry, options)),
     },
     null,
     2,
@@ -371,13 +398,14 @@ function txtSnapshotPath(jsonSnapshotPath: string): string {
 export async function matchSpanTreeSnapshot(
   entries: readonly (CapturedLogEvent | SpanTreeEntry)[],
   jsonSnapshotPath: string,
+  options?: SpanTreeSnapshotOptions,
 ): Promise<void> {
   await matchFileSnapshot(
-    formatSpanTreeSnapshot(entries),
+    formatSpanTreeSnapshot(entries, options),
     txtSnapshotPath(jsonSnapshotPath),
   );
   await matchFileSnapshot(
-    formatSpanTreeJsonSnapshot(entries),
+    formatSpanTreeJsonSnapshot(entries, options),
     jsonSnapshotPath,
   );
 }