braintrust 2.2.1-rc.0 → 2.2.2

package/dist/browser.mjs

@@ -1928,7 +1928,8 @@ var SpanType = z6.union([
     "automation",
     "facet",
     "preprocessor",
-    "classifier"
+    "classifier",
+    "review"
   ]),
   z6.null()
 ]);
@@ -5187,12 +5188,19 @@ function updateSpanImpl({
   parentObjectType,
   parentObjectId,
   id,
+  root_span_id,
+  span_id,
   event
 }) {
+  if (isEmpty2(root_span_id) !== isEmpty2(span_id)) {
+    throw new Error("both root_span_id and span_id must be set, or neither");
+  }
+  const hasExplicitSpanIds = root_span_id !== void 0 && span_id !== void 0;
   const updateEvent = deepCopyEvent(
     validateAndSanitizeExperimentLogPartialArgs({
+      ...event,
       id,
-      ...event
+      ...hasExplicitSpanIds ? { root_span_id, span_id } : {}
     })
   );
   const parentIds = async () => new SpanComponentsV3({
@@ -5224,6 +5232,8 @@ function updateSpan({
       spanComponentsToObjectIdLambda(resolvedState, components)
     ),
     id: components.data.row_id,
+    root_span_id: components.data.root_span_id,
+    span_id: components.data.span_id,
     event
   });
 }
@@ -5535,7 +5545,7 @@ var Logger = class {
    * @param event The event data to update the span with. Must include `id`. See {@link Experiment.log} for a full list of valid fields.
    */
   updateSpan(event) {
-    const { id, ...eventRest } = event;
+    const { id, root_span_id, span_id, ...eventRest } = event;
     if (!id) {
       throw new Error("Span id is required to update a span");
     }
@@ -5544,6 +5554,8 @@ var Logger = class {
       parentObjectType: this.parentObjectType(),
       parentObjectId: this.lazyId,
       id,
+      root_span_id,
+      span_id,
       event: eventRest
     });
   }
@@ -7826,7 +7838,7 @@ View complete results in Braintrust or run experiment.summarize() again.`
    * @param event The event data to update the span with. Must include `id`. See {@link Experiment.log} for a full list of valid fields.
    */
   updateSpan(event) {
-    const { id, ...eventRest } = event;
+    const { id, root_span_id, span_id, ...eventRest } = event;
     if (!id) {
       throw new Error("Span id is required to update a span");
     }
@@ -7835,6 +7847,8 @@ View complete results in Braintrust or run experiment.summarize() again.`
       parentObjectType: this.parentObjectType(),
       parentObjectId: this.lazyId,
       id,
+      root_span_id,
+      span_id,
       event: eventRest
     });
   }
@@ -12695,14 +12709,17 @@ function parseToolName(rawToolName) {
     rawToolName
   };
 }
-function createToolTracingHooks(parentSpanExportPromise, activeToolSpans, mcpServers) {
+function createToolTracingHooks(resolveParentSpan, activeToolSpans, mcpServers, subAgentSpans, endedSubAgentSpans) {
   const preToolUse = async (input, toolUseID) => {
     if (input.hook_event_name !== "PreToolUse" || !toolUseID) {
       return {};
     }
+    if (input.tool_name === "Task") {
+      return {};
+    }
     const parsed = parseToolName(input.tool_name);
     const mcpMetadata = getMcpServerMetadata(parsed.mcpServer, mcpServers);
-    const parentExport = await parentSpanExportPromise;
+    const parentExport = await resolveParentSpan(toolUseID);
     const toolSpan = startSpan({
       name: parsed.displayName,
       spanAttributes: { type: "tool" /* TOOL */ },
@@ -12730,6 +12747,30 @@ function createToolTracingHooks(parentSpanExportPromise, activeToolSpans, mcpSer
     if (input.hook_event_name !== "PostToolUse" || !toolUseID) {
       return {};
     }
+    const subAgentSpan = subAgentSpans.get(toolUseID);
+    if (subAgentSpan) {
+      try {
+        const response = input.tool_response;
+        const metadata = {};
+        if (response?.status) {
+          metadata["claude_agent_sdk.status"] = response.status;
+        }
+        if (response?.totalDurationMs) {
+          metadata["claude_agent_sdk.duration_ms"] = response.totalDurationMs;
+        }
+        if (response?.totalToolUseCount !== void 0) {
+          metadata["claude_agent_sdk.tool_use_count"] = response.totalToolUseCount;
+        }
+        subAgentSpan.log({
+          output: response?.content,
+          metadata
+        });
+      } finally {
+        subAgentSpan.end();
+        endedSubAgentSpans.add(toolUseID);
+      }
+      return {};
+    }
     const toolSpan = activeToolSpans.get(toolUseID);
     if (!toolSpan) {
       return {};
@@ -12746,6 +12787,16 @@ function createToolTracingHooks(parentSpanExportPromise, activeToolSpans, mcpSer
     if (input.hook_event_name !== "PostToolUseFailure" || !toolUseID) {
       return {};
     }
+    const subAgentSpan = subAgentSpans.get(toolUseID);
+    if (subAgentSpan) {
+      try {
+        subAgentSpan.log({ error: input.error });
+      } finally {
+        subAgentSpan.end();
+        endedSubAgentSpans.add(toolUseID);
+      }
+      return {};
+    }
     const toolSpan = activeToolSpans.get(toolUseID);
     if (!toolSpan) {
       return {};
@@ -12770,12 +12821,14 @@ function createToolTracingHooks(parentSpanExportPromise, activeToolSpans, mcpSer
   };
   return { preToolUse, postToolUse, postToolUseFailure };
 }
-function injectTracingHooks(options, parentSpanExportPromise, activeToolSpans) {
+function injectTracingHooks(options, resolveParentSpan, activeToolSpans, subAgentSpans, endedSubAgentSpans) {
   const mcpServers = options.mcpServers;
   const { preToolUse, postToolUse, postToolUseFailure } = createToolTracingHooks(
-    parentSpanExportPromise,
+    resolveParentSpan,
     activeToolSpans,
-    mcpServers
+    mcpServers,
+    subAgentSpans,
+    endedSubAgentSpans
   );
   const existingHooks = options.hooks ?? {};
   return {
@@ -12821,18 +12874,45 @@ function filterSerializableOptions(options) {
   }
   return filtered;
 }
+function isAsyncIterable(value) {
+  return value !== null && value !== void 0 && typeof value[Symbol.asyncIterator] === "function";
+}
 function wrapClaudeAgentQuery(queryFn, defaultThis) {
   const proxy = new Proxy(queryFn, {
     apply(target, thisArg, argArray) {
       const params = argArray[0] ?? {};
       const { prompt, options = {} } = params;
+      const promptIsAsyncIterable = isAsyncIterable(prompt);
+      let capturedPromptMessages;
+      let promptForQuery = prompt;
+      let promptStarted = false;
+      let resolvePromptDone;
+      const promptDone = new Promise((resolve) => {
+        resolvePromptDone = resolve;
+      });
+      if (promptIsAsyncIterable) {
+        capturedPromptMessages = [];
+        const originalPrompt = prompt;
+        const capturingPrompt = (async function* () {
+          promptStarted = true;
+          try {
+            for await (const msg of originalPrompt) {
+              capturedPromptMessages.push(msg);
+              yield msg;
+            }
+          } finally {
+            resolvePromptDone?.();
+          }
+        })();
+        promptForQuery = capturingPrompt;
+      }
       const span = startSpan({
         name: "Claude Agent",
         spanAttributes: {
           type: "task" /* TASK */
         },
         event: {
-          input: typeof prompt === "string" ? prompt : { type: "streaming", description: "AsyncIterable<SDKMessage>" },
+          input: typeof prompt === "string" ? prompt : promptIsAsyncIterable ? void 0 : prompt !== void 0 ? String(prompt) : void 0,
           metadata: filterSerializableOptions(options)
         }
       });
@@ -12843,13 +12923,22 @@ function wrapClaudeAgentQuery(queryFn, defaultThis) {
       let currentMessageStartTime = getCurrentUnixTimestamp();
       const currentMessages = [];
       const createLLMSpan = async () => {
+        const parentToolUseId = currentMessages[0]?.parent_tool_use_id ?? null;
+        let parentSpanExport;
+        if (parentToolUseId) {
+          const subAgentSpan = subAgentSpans.get(parentToolUseId);
+          parentSpanExport = subAgentSpan ? await subAgentSpan.export() : await span.export();
+        } else {
+          parentSpanExport = await span.export();
+        }
         const finalMessageContent = await _createLLMSpanForMessages(
           currentMessages,
           prompt,
           finalResults,
           options,
           currentMessageStartTime,
-          await span.export()
+          capturedPromptMessages,
+          parentSpanExport
         );
         if (finalMessageContent) {
           finalResults.push(finalMessageContent);
@@ -12863,12 +12952,34 @@ function wrapClaudeAgentQuery(queryFn, defaultThis) {
       };
       const invocationTarget = thisArg === proxy || thisArg === void 0 ? defaultThis ?? thisArg : thisArg;
       const activeToolSpans = /* @__PURE__ */ new Map();
+      const subAgentSpans = /* @__PURE__ */ new Map();
+      const endedSubAgentSpans = /* @__PURE__ */ new Set();
+      const toolUseToParent = /* @__PURE__ */ new Map();
+      const pendingSubAgentNames = /* @__PURE__ */ new Map();
+      const resolveParentSpan = async (toolUseID) => {
+        const parentToolUseId = toolUseToParent.get(toolUseID);
+        if (parentToolUseId) {
+          const subAgentSpan = subAgentSpans.get(parentToolUseId);
+          if (subAgentSpan) {
+            return subAgentSpan.export();
+          }
+        }
+        return span.export();
+      };
       const optionsWithHooks = injectTracingHooks(
         options,
-        span.export(),
-        activeToolSpans
+        resolveParentSpan,
+        activeToolSpans,
+        subAgentSpans,
+        endedSubAgentSpans
       );
-      const modifiedArgArray = [{ ...params, options: optionsWithHooks }];
+      const modifiedArgArray = [
+        {
+          ...params,
+          ...promptForQuery !== void 0 ? { prompt: promptForQuery } : {},
+          options: optionsWithHooks
+        }
+      ];
       const originalGenerator = withCurrent(
         span,
         () => Reflect.apply(target, invocationTarget, modifiedArgArray)
@@ -12877,6 +12988,41 @@ function wrapClaudeAgentQuery(queryFn, defaultThis) {
         try {
           for await (const message of originalGenerator) {
             const currentTime = getCurrentUnixTimestamp();
+            if (message.type === "assistant" && Array.isArray(message.message?.content)) {
+              const parentToolUseId = message.parent_tool_use_id ?? null;
+              for (const block of message.message.content) {
+                if (block.type === "tool_use" && block.id) {
+                  toolUseToParent.set(block.id, parentToolUseId);
+                  if (block.name === "Task" && block.input?.subagent_type) {
+                    pendingSubAgentNames.set(
+                      block.id,
+                      block.input.subagent_type
+                    );
+                  }
+                }
+              }
+            }
+            if ("parent_tool_use_id" in message) {
+              const parentToolUseId = message.parent_tool_use_id;
+              if (parentToolUseId && !subAgentSpans.has(parentToolUseId)) {
+                const agentName = pendingSubAgentNames.get(parentToolUseId);
+                const spanName = agentName ? `Agent: ${agentName}` : "Agent: sub-agent";
+                const parentExport = await span.export();
+                const subAgentSpan = startSpan({
+                  name: spanName,
+                  spanAttributes: { type: "task" /* TASK */ },
+                  event: {
+                    metadata: {
+                      ...agentName && {
+                        "claude_agent_sdk.agent_type": agentName
+                      }
+                    }
+                  },
+                  parent: parentExport
+                });
+                subAgentSpans.set(parentToolUseId, subAgentSpan);
+              }
+            }
             const messageId = message.message?.id;
             if (messageId && messageId !== currentMessageId) {
               await createLLMSpan();
@@ -12922,6 +13068,22 @@ function wrapClaudeAgentQuery(queryFn, defaultThis) {
           });
           throw error;
         } finally {
+          for (const [id, subSpan] of subAgentSpans) {
+            if (!endedSubAgentSpans.has(id)) {
+              subSpan.end();
+            }
+          }
+          subAgentSpans.clear();
+          if (capturedPromptMessages) {
+            if (promptStarted) {
+              await promptDone;
+            }
+            if (capturedPromptMessages.length > 0) {
+              span.log({
+                input: _formatCapturedMessages(capturedPromptMessages)
+              });
+            }
+          }
           span.end();
         }
       })();
@@ -12949,14 +13111,25 @@ function wrapClaudeAgentQuery(queryFn, defaultThis) {
   });
   return proxy;
 }
-function _buildLLMInput(prompt, conversationHistory) {
-  const promptMessage = typeof prompt === "string" ? { content: prompt, role: "user" } : void 0;
-  const inputParts = [
-    ...promptMessage ? [promptMessage] : [],
-    ...conversationHistory
-  ];
+function _buildLLMInput(prompt, conversationHistory, capturedPromptMessages) {
+  const promptMessages = [];
+  if (typeof prompt === "string") {
+    promptMessages.push({ content: prompt, role: "user" });
+  } else if (capturedPromptMessages && capturedPromptMessages.length > 0) {
+    for (const msg of capturedPromptMessages) {
+      const role = msg.message?.role;
+      const content = msg.message?.content;
+      if (role && content !== void 0) {
+        promptMessages.push({ content, role });
+      }
+    }
+  }
+  const inputParts = [...promptMessages, ...conversationHistory];
   return inputParts.length > 0 ? inputParts : void 0;
 }
+function _formatCapturedMessages(messages) {
+  return messages.length > 0 ? messages : [];
+}
 function _extractUsageFromMessage(message) {
   const metrics = {};
   let usage;
@@ -12990,7 +13163,7 @@ function _extractUsageFromMessage(message) {
   }
   return metrics;
 }
-async function _createLLMSpanForMessages(messages, prompt, conversationHistory, options, startTime, parentSpan) {
+async function _createLLMSpanForMessages(messages, prompt, conversationHistory, options, startTime, capturedPromptMessages, parentSpan) {
   if (messages.length === 0) return void 0;
   const lastMessage = messages[messages.length - 1];
   if (lastMessage.type !== "assistant" || !lastMessage.message?.usage) {
@@ -12998,7 +13171,11 @@ async function _createLLMSpanForMessages(messages, prompt, conversationHistory,
   }
   const model = lastMessage.message.model || options.model;
   const usage = _extractUsageFromMessage(lastMessage);
-  const input = _buildLLMInput(prompt, conversationHistory);
+  const input = _buildLLMInput(
+    prompt,
+    conversationHistory,
+    capturedPromptMessages
+  );
   const outputs = messages.map(
     (m) => m.message?.content && m.message?.role ? { content: m.message.content, role: m.message.role } : void 0
   ).filter((c) => c !== void 0);
@@ -13755,7 +13932,7 @@ function unescapePath(path) {
 }
 var graph_framework_default = { createGraph };
 
-// ../node_modules/.pnpm/async@3.2.5/node_modules/async/dist/async.mjs
+// ../node_modules/async/dist/async.mjs
 function initialParams(fn) {
   return function(...args) {
     var callback = args.pop();
@@ -13826,7 +14003,7 @@ function isAsync(fn) {
 function isAsyncGenerator2(fn) {
   return fn[Symbol.toStringTag] === "AsyncGenerator";
 }
-function isAsyncIterable(obj) {
+function isAsyncIterable2(obj) {
   return typeof obj[Symbol.asyncIterator] === "function";
 }
 function wrapAsync(asyncFn) {
@@ -13880,7 +14057,6 @@ function isArrayLike(value) {
   return value && typeof value.length === "number" && value.length >= 0 && value.length % 1 === 0;
 }
 var breakLoop = {};
-var breakLoop$1 = breakLoop;
 function once(fn) {
   function wrapper(...args) {
     if (fn === null) return;
@@ -13972,7 +14148,7 @@ function asyncEachOfLimit(generator, limit, iteratee, callback) {
       canceled = true;
       return;
     }
-    if (result === breakLoop$1 || done && running <= 0) {
+    if (result === breakLoop || done && running <= 0) {
       done = true;
       return callback(null);
     }
@@ -13998,7 +14174,7 @@ var eachOfLimit$2 = (limit) => {
     if (isAsyncGenerator2(obj)) {
       return asyncEachOfLimit(obj, limit, iteratee, callback);
     }
-    if (isAsyncIterable(obj)) {
+    if (isAsyncIterable2(obj)) {
       return asyncEachOfLimit(obj[Symbol.asyncIterator](), limit, iteratee, callback);
     }
     var nextElem = createIterator(obj);
@@ -14015,7 +14191,7 @@ var eachOfLimit$2 = (limit) => {
       } else if (err === false) {
         done = true;
         canceled = true;
-      } else if (value === breakLoop$1 || done && running <= 0) {
+      } else if (value === breakLoop || done && running <= 0) {
         done = true;
         return callback(null);
       } else if (!looping) {
@@ -14058,7 +14234,7 @@ function eachOfArrayLike(coll, iteratee, callback) {
     if (canceled === true) return;
     if (err) {
       callback(err);
-    } else if (++completed === length || value === breakLoop$1) {
+    } else if (++completed === length || value === breakLoop) {
       callback(null);
     }
   }
@@ -14454,7 +14630,7 @@ function _createTester(check, getResult) {
         if (check(result) && !testResult) {
           testPassed = true;
           testResult = getResult(true, value);
-          return callback(null, breakLoop$1);
+          return callback(null, breakLoop);
         }
         callback();
       });
@@ -15232,7 +15408,7 @@ function callEvaluatorData(data) {
     baseExperiment
   };
 }
-function isAsyncIterable2(value) {
+function isAsyncIterable3(value) {
   return typeof value === "object" && value !== null && typeof value[Symbol.asyncIterator] === "function";
 }
 function isIterable(value) {
@@ -15453,7 +15629,7 @@ async function runEvaluatorInternal(experiment, evaluator, progressReporter, fil
     }
     const resolvedDataResult = dataResult instanceof Promise ? await dataResult : dataResult;
     const dataIterable = (() => {
-      if (isAsyncIterable2(resolvedDataResult)) {
+      if (isAsyncIterable3(resolvedDataResult)) {
         return resolvedDataResult;
       }
       if (Array.isArray(resolvedDataResult) || isIterable(resolvedDataResult)) {
@@ -16305,6 +16481,7 @@ var CodeParameters = class {
     this.metadata = opts.metadata;
   }
   async toFunctionDefinition(projectNameToId) {
+    const schema = serializeEvalParameterstoParametersSchema(this.schema);
     return {
       project_id: await projectNameToId.resolve(this.project),
       name: this.name,
@@ -16313,8 +16490,8 @@ var CodeParameters = class {
       function_type: "parameters",
       function_data: {
         type: "parameters",
-        data: {},
-        __schema: serializeEvalParameterstoParametersSchema(this.schema)
+        data: getDefaultDataFromParametersSchema(schema),
+        __schema: schema
       },
       if_exists: this.ifExists,
       metadata: this.metadata
@@ -16369,6 +16546,16 @@ function serializeEvalParameterstoParametersSchema(parameters) {
     additionalProperties: true
   };
 }
+function getDefaultDataFromParametersSchema(schema) {
+  return Object.fromEntries(
+    Object.entries(schema.properties).flatMap(([name, value]) => {
+      if (!("default" in value)) {
+        return [];
+      }
+      return [[name, value.default]];
+    })
+  );
+}
 var ProjectNameIdMap = class {
   nameToId = {};
   idToName = {};