braintrust 3.6.0 → 3.7.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -19859,7 +19859,7 @@ var AnthropicPlugin = class extends BasePlugin {
19859
19859
  this.unsubscribers.push(
19860
19860
  traceStreamingChannel(anthropicChannels.betaMessagesCreate, {
19861
19861
  ...anthropicConfig,
19862
- name: "anthropic.beta.messages.create"
19862
+ name: "anthropic.messages.create"
19863
19863
  })
19864
19864
  );
19865
19865
  }
@@ -19882,9 +19882,12 @@ function parseMetricsFromUsage3(usage) {
19882
19882
  return metrics;
19883
19883
  }
19884
19884
  function aggregateAnthropicStreamChunks(chunks) {
19885
- const deltas = [];
19885
+ const fallbackTextDeltas = [];
19886
+ const contentBlocks = {};
19887
+ const contentBlockDeltas = {};
19886
19888
  let metrics = {};
19887
19889
  let metadata = {};
19890
+ let role;
19888
19891
  for (const event of chunks) {
19889
19892
  switch (event?.type) {
19890
19893
  case "message_start":
@@ -19892,15 +19895,43 @@ function aggregateAnthropicStreamChunks(chunks) {
19892
19895
  const initialMetrics = parseMetricsFromUsage3(event.message.usage);
19893
19896
  metrics = { ...metrics, ...initialMetrics };
19894
19897
  }
19898
+ if (typeof event.message?.role === "string") {
19899
+ role = event.message.role;
19900
+ }
19901
+ break;
19902
+ case "content_block_start":
19903
+ if (event.content_block) {
19904
+ contentBlocks[event.index] = event.content_block;
19905
+ contentBlockDeltas[event.index] = [];
19906
+ }
19895
19907
  break;
19896
19908
  case "content_block_delta":
19897
19909
  if (event.delta?.type === "text_delta") {
19898
19910
  const text = event.delta.text;
19899
19911
  if (text) {
19900
- deltas.push(text);
19912
+ if (contentBlocks[event.index] !== void 0 || contentBlockDeltas[event.index] !== void 0) {
19913
+ contentBlockDeltas[event.index] ??= [];
19914
+ contentBlockDeltas[event.index].push(text);
19915
+ } else {
19916
+ fallbackTextDeltas.push(text);
19917
+ }
19918
+ }
19919
+ } else if (event.delta?.type === "input_json_delta") {
19920
+ const partialJson = event.delta.partial_json;
19921
+ if (partialJson) {
19922
+ contentBlockDeltas[event.index] ??= [];
19923
+ contentBlockDeltas[event.index].push(partialJson);
19901
19924
  }
19902
19925
  }
19903
19926
  break;
19927
+ case "content_block_stop":
19928
+ finalizeContentBlock(
19929
+ event.index,
19930
+ contentBlocks,
19931
+ contentBlockDeltas,
19932
+ fallbackTextDeltas
19933
+ );
19934
+ break;
19904
19935
  case "message_delta":
19905
19936
  if (event.usage) {
19906
19937
  const finalMetrics = parseMetricsFromUsage3(event.usage);
@@ -19912,7 +19943,21 @@ function aggregateAnthropicStreamChunks(chunks) {
19912
19943
  break;
19913
19944
  }
19914
19945
  }
19915
- const output = deltas.join("");
19946
+ const orderedContent = Object.entries(contentBlocks).map(([index, block]) => ({
19947
+ block,
19948
+ index: Number(index)
19949
+ })).filter(({ block }) => block !== void 0).sort((left, right) => left.index - right.index).map(({ block }) => block);
19950
+ let output = fallbackTextDeltas.join("");
19951
+ if (orderedContent.length > 0) {
19952
+ if (orderedContent.every(isTextContentBlock)) {
19953
+ output = orderedContent.map((block) => block.text).join("");
19954
+ } else {
19955
+ output = {
19956
+ ...role ? { role } : {},
19957
+ content: orderedContent
19958
+ };
19959
+ }
19960
+ }
19916
19961
  const finalized = finalizeAnthropicTokens(metrics);
19917
19962
  const filteredMetrics = Object.fromEntries(
19918
19963
  Object.entries(finalized).filter(
@@ -19925,6 +19970,49 @@ function aggregateAnthropicStreamChunks(chunks) {
19925
19970
  metadata
19926
19971
  };
19927
19972
  }
19973
+ function finalizeContentBlock(index, contentBlocks, contentBlockDeltas, fallbackTextDeltas) {
19974
+ const contentBlock = contentBlocks[index];
19975
+ if (!contentBlock) {
19976
+ return;
19977
+ }
19978
+ const text = contentBlockDeltas[index]?.join("") ?? "";
19979
+ if (isToolUseContentBlock(contentBlock)) {
19980
+ if (!text) {
19981
+ return;
19982
+ }
19983
+ try {
19984
+ contentBlocks[index] = {
19985
+ ...contentBlock,
19986
+ input: JSON.parse(text)
19987
+ };
19988
+ } catch {
19989
+ fallbackTextDeltas.push(text);
19990
+ delete contentBlocks[index];
19991
+ }
19992
+ return;
19993
+ }
19994
+ if (isTextContentBlock(contentBlock)) {
19995
+ if (!text) {
19996
+ delete contentBlocks[index];
19997
+ return;
19998
+ }
19999
+ contentBlocks[index] = {
20000
+ ...contentBlock,
20001
+ text
20002
+ };
20003
+ return;
20004
+ }
20005
+ if (text) {
20006
+ fallbackTextDeltas.push(text);
20007
+ }
20008
+ delete contentBlocks[index];
20009
+ }
20010
+ function isTextContentBlock(contentBlock) {
20011
+ return contentBlock.type === "text";
20012
+ }
20013
+ function isToolUseContentBlock(contentBlock) {
20014
+ return contentBlock.type === "tool_use";
20015
+ }
19928
20016
  function isAnthropicBase64ContentBlock(input) {
19929
20017
  return (input.type === "image" || input.type === "document") && isObject(input.source) && input.source.type === "base64";
19930
20018
  }
@@ -20945,12 +21033,15 @@ var claudeAgentSDKChannels = defineChannels(
20945
21033
  {
20946
21034
  query: channel({
20947
21035
  channelName: "query",
20948
- kind: "async"
21036
+ kind: "sync-stream"
20949
21037
  })
20950
21038
  }
20951
21039
  );
20952
21040
 
20953
21041
  // src/instrumentation/plugins/claude-agent-sdk-plugin.ts
21042
+ function isSubAgentToolName(toolName) {
21043
+ return toolName === "Agent" || toolName === "Task";
21044
+ }
20954
21045
  function filterSerializableOptions2(options) {
20955
21046
  const allowedKeys = [
20956
21047
  "model",
@@ -21004,34 +21095,50 @@ function extractUsageFromMessage(message) {
21004
21095
  const cacheReadTokens = getNumberProperty3(usage, "cache_read_input_tokens") || 0;
21005
21096
  const cacheCreationTokens = getNumberProperty3(usage, "cache_creation_input_tokens") || 0;
21006
21097
  if (cacheReadTokens > 0 || cacheCreationTokens > 0) {
21007
- const cacheTokens = extractAnthropicCacheTokens(
21008
- cacheReadTokens,
21009
- cacheCreationTokens
21098
+ Object.assign(
21099
+ metrics,
21100
+ extractAnthropicCacheTokens(cacheReadTokens, cacheCreationTokens)
21010
21101
  );
21011
- Object.assign(metrics, cacheTokens);
21012
21102
  }
21013
21103
  if (Object.keys(metrics).length > 0) {
21014
21104
  Object.assign(metrics, finalizeAnthropicTokens(metrics));
21015
21105
  }
21016
21106
  return metrics;
21017
21107
  }
21018
- function buildLLMInput(prompt, conversationHistory) {
21019
- const promptMessage = typeof prompt === "string" ? { content: prompt, role: "user" } : void 0;
21020
- const inputParts = [
21021
- ...promptMessage ? [promptMessage] : [],
21022
- ...conversationHistory
21023
- ];
21108
+ function buildLLMInput(prompt, conversationHistory, capturedPromptMessages) {
21109
+ const promptMessages = [];
21110
+ if (typeof prompt === "string") {
21111
+ promptMessages.push({ content: prompt, role: "user" });
21112
+ } else if (capturedPromptMessages && capturedPromptMessages.length > 0) {
21113
+ for (const msg of capturedPromptMessages) {
21114
+ const role = msg.message?.role;
21115
+ const content = msg.message?.content;
21116
+ if (role && content !== void 0) {
21117
+ promptMessages.push({ content, role });
21118
+ }
21119
+ }
21120
+ }
21121
+ const inputParts = [...promptMessages, ...conversationHistory];
21024
21122
  return inputParts.length > 0 ? inputParts : void 0;
21025
21123
  }
21026
- async function createLLMSpanForMessages(messages, prompt, conversationHistory, options, startTime, parentSpan) {
21027
- if (messages.length === 0) return void 0;
21124
+ function formatCapturedMessages(messages) {
21125
+ return messages.length > 0 ? messages : [];
21126
+ }
21127
+ async function createLLMSpanForMessages(messages, prompt, conversationHistory, options, startTime, capturedPromptMessages, parentSpan) {
21128
+ if (messages.length === 0) {
21129
+ return void 0;
21130
+ }
21028
21131
  const lastMessage = messages[messages.length - 1];
21029
21132
  if (lastMessage.type !== "assistant" || !lastMessage.message?.usage) {
21030
21133
  return void 0;
21031
21134
  }
21032
21135
  const model = lastMessage.message.model || options.model;
21033
21136
  const usage = extractUsageFromMessage(lastMessage);
21034
- const input = buildLLMInput(prompt, conversationHistory);
21137
+ const input = buildLLMInput(
21138
+ prompt,
21139
+ conversationHistory,
21140
+ capturedPromptMessages
21141
+ );
21035
21142
  const outputs = messages.map(
21036
21143
  (m) => m.message?.content && m.message?.role ? { content: m.message.content, role: m.message.role } : void 0
21037
21144
  ).filter(
@@ -21039,233 +21146,552 @@ async function createLLMSpanForMessages(messages, prompt, conversationHistory, o
21039
21146
  );
21040
21147
  const span = startSpan({
21041
21148
  name: "anthropic.messages.create",
21149
+ parent: parentSpan,
21042
21150
  spanAttributes: {
21043
21151
  type: "llm" /* LLM */
21044
21152
  },
21045
- startTime,
21046
- parent: parentSpan
21153
+ startTime
21047
21154
  });
21048
21155
  span.log({
21049
21156
  input,
21050
- output: outputs,
21051
21157
  metadata: model ? { model } : void 0,
21052
- metrics: usage
21158
+ metrics: usage,
21159
+ output: outputs
21053
21160
  });
21054
21161
  await span.end();
21055
21162
  return lastMessage.message?.content && lastMessage.message?.role ? { content: lastMessage.message.content, role: lastMessage.message.role } : void 0;
21056
21163
  }
21057
- var ClaudeAgentSDKPlugin = class extends BasePlugin {
21058
- onEnable() {
21059
- this.subscribeToQuery();
21164
+ function getMcpServerMetadata2(serverName, mcpServers) {
21165
+ if (!serverName || !mcpServers) {
21166
+ return {};
21060
21167
  }
21061
- onDisable() {
21062
- for (const unsubscribe of this.unsubscribers) {
21063
- unsubscribe();
21168
+ const serverConfig = mcpServers[serverName];
21169
+ if (!serverConfig) {
21170
+ return {};
21171
+ }
21172
+ const metadata = {};
21173
+ if (serverConfig.type) {
21174
+ metadata["mcp.type"] = serverConfig.type;
21175
+ } else if (typeof serverConfig === "object" && "transport" in serverConfig) {
21176
+ metadata["mcp.type"] = "sdk";
21177
+ }
21178
+ if (serverConfig.url) {
21179
+ metadata["mcp.url"] = serverConfig.url;
21180
+ }
21181
+ if (serverConfig.command) {
21182
+ metadata["mcp.command"] = serverConfig.command;
21183
+ if (serverConfig.args) {
21184
+ metadata["mcp.args"] = serverConfig.args.join(" ");
21064
21185
  }
21065
- this.unsubscribers = [];
21066
21186
  }
21067
- /**
21068
- * Subscribe to the query channel for agent interactions.
21069
- * Handles streaming responses and traces both the top-level agent task
21070
- * and individual LLM calls.
21071
- */
21072
- subscribeToQuery() {
21073
- const channel2 = claudeAgentSDKChannels.query.tracingChannel();
21074
- const spans = /* @__PURE__ */ new WeakMap();
21075
- const handlers = {
21076
- start: (event) => {
21077
- const params = event.arguments[0];
21078
- const prompt = params?.prompt;
21079
- const options = params?.options ?? {};
21080
- const span = startSpan({
21081
- name: "Claude Agent",
21082
- spanAttributes: {
21083
- type: "task" /* TASK */
21084
- }
21085
- });
21086
- const startTime = getCurrentUnixTimestamp();
21087
- try {
21088
- span.log({
21089
- input: typeof prompt === "string" ? prompt : {
21090
- type: "streaming",
21091
- description: "AsyncIterable<ClaudeAgentSDKMessage>"
21092
- },
21093
- metadata: filterSerializableOptions2(options)
21094
- });
21095
- } catch (error) {
21096
- console.error("Error extracting input for Claude Agent SDK:", error);
21187
+ return metadata;
21188
+ }
21189
+ function parseToolName2(rawToolName) {
21190
+ const mcpMatch = rawToolName.match(/^mcp__([^_]+)__(.+)$/);
21191
+ if (mcpMatch) {
21192
+ const [, mcpServer, toolName] = mcpMatch;
21193
+ return {
21194
+ displayName: `tool: ${mcpServer}/${toolName}`,
21195
+ mcpServer,
21196
+ rawToolName,
21197
+ toolName
21198
+ };
21199
+ }
21200
+ return {
21201
+ displayName: `tool: ${rawToolName}`,
21202
+ rawToolName,
21203
+ toolName: rawToolName
21204
+ };
21205
+ }
21206
+ function createToolTracingHooks2(resolveParentSpan, activeToolSpans, mcpServers, subAgentSpans, endedSubAgentSpans) {
21207
+ const preToolUse = async (input, toolUseID) => {
21208
+ if (input.hook_event_name !== "PreToolUse" || !toolUseID) {
21209
+ return {};
21210
+ }
21211
+ if (isSubAgentToolName(input.tool_name)) {
21212
+ return {};
21213
+ }
21214
+ const parsed = parseToolName2(input.tool_name);
21215
+ const toolSpan = startSpan({
21216
+ event: {
21217
+ input: input.tool_input,
21218
+ metadata: {
21219
+ "claude_agent_sdk.cwd": input.cwd,
21220
+ "claude_agent_sdk.raw_tool_name": parsed.rawToolName,
21221
+ "claude_agent_sdk.session_id": input.session_id,
21222
+ "gen_ai.tool.call.id": toolUseID,
21223
+ "gen_ai.tool.name": parsed.toolName,
21224
+ ...parsed.mcpServer && { "mcp.server": parsed.mcpServer },
21225
+ ...getMcpServerMetadata2(parsed.mcpServer, mcpServers)
21097
21226
  }
21098
- spans.set(event, {
21099
- span,
21100
- startTime,
21101
- conversationHistory: [],
21102
- currentMessages: [],
21103
- currentMessageId: void 0,
21104
- currentMessageStartTime: startTime,
21105
- accumulatedOutputTokens: 0
21106
- });
21107
21227
  },
21108
- asyncEnd: (event) => {
21109
- const spanData = spans.get(event);
21110
- if (!spanData) {
21111
- return;
21112
- }
21113
- const eventResult = event.result;
21114
- if (eventResult === void 0) {
21115
- spanData.span.end();
21116
- spans.delete(event);
21117
- return;
21228
+ name: parsed.displayName,
21229
+ parent: await resolveParentSpan(toolUseID),
21230
+ spanAttributes: { type: "tool" /* TOOL */ }
21231
+ });
21232
+ activeToolSpans.set(toolUseID, toolSpan);
21233
+ return {};
21234
+ };
21235
+ const postToolUse = async (input, toolUseID) => {
21236
+ if (input.hook_event_name !== "PostToolUse" || !toolUseID) {
21237
+ return {};
21238
+ }
21239
+ const subAgentSpan = subAgentSpans.get(toolUseID);
21240
+ if (subAgentSpan) {
21241
+ try {
21242
+ const response = input.tool_response;
21243
+ const metadata = {};
21244
+ if (response?.status) {
21245
+ metadata["claude_agent_sdk.status"] = response.status;
21118
21246
  }
21119
- if (isAsyncIterable5(eventResult)) {
21120
- patchStreamIfNeeded(eventResult, {
21121
- onChunk: async (message) => {
21122
- const currentTime = getCurrentUnixTimestamp();
21123
- const params = event.arguments[0];
21124
- const prompt = params?.prompt;
21125
- const options = params?.options ?? {};
21126
- const messageId = message.message?.id;
21127
- if (messageId && messageId !== spanData.currentMessageId) {
21128
- if (spanData.currentMessages.length > 0) {
21129
- const finalMessage = await createLLMSpanForMessages(
21130
- spanData.currentMessages,
21131
- prompt,
21132
- spanData.conversationHistory,
21133
- options,
21134
- spanData.currentMessageStartTime,
21135
- await spanData.span.export()
21136
- );
21137
- if (finalMessage) {
21138
- spanData.conversationHistory.push(finalMessage);
21139
- }
21140
- const lastMessage = spanData.currentMessages[spanData.currentMessages.length - 1];
21141
- if (lastMessage?.message?.usage) {
21142
- const outputTokens = getNumberProperty3(
21143
- lastMessage.message.usage,
21144
- "output_tokens"
21145
- ) || 0;
21146
- spanData.accumulatedOutputTokens += outputTokens;
21147
- }
21148
- spanData.currentMessages = [];
21149
- }
21150
- spanData.currentMessageId = messageId;
21151
- spanData.currentMessageStartTime = currentTime;
21152
- }
21153
- if (message.type === "assistant" && message.message?.usage) {
21154
- spanData.currentMessages.push(message);
21155
- }
21156
- if (message.type === "result" && message.usage) {
21157
- const finalUsageMetrics = extractUsageFromMessage(message);
21158
- if (spanData.currentMessages.length > 0 && finalUsageMetrics.completion_tokens !== void 0) {
21159
- const lastMessage = spanData.currentMessages[spanData.currentMessages.length - 1];
21160
- if (lastMessage?.message?.usage) {
21161
- const adjustedTokens = finalUsageMetrics.completion_tokens - spanData.accumulatedOutputTokens;
21162
- if (adjustedTokens >= 0) {
21163
- lastMessage.message.usage.output_tokens = adjustedTokens;
21164
- }
21165
- }
21166
- }
21167
- const result_metadata = {};
21168
- if (message.num_turns !== void 0) {
21169
- result_metadata.num_turns = message.num_turns;
21170
- }
21171
- if (message.session_id !== void 0) {
21172
- result_metadata.session_id = message.session_id;
21173
- }
21174
- if (Object.keys(result_metadata).length > 0) {
21175
- spanData.span.log({
21176
- metadata: result_metadata
21177
- });
21178
- }
21179
- }
21180
- },
21181
- onComplete: async () => {
21182
- try {
21183
- const params = event.arguments[0];
21184
- const prompt = params?.prompt;
21185
- const options = params?.options ?? {};
21186
- if (spanData.currentMessages.length > 0) {
21187
- const finalMessage = await createLLMSpanForMessages(
21188
- spanData.currentMessages,
21189
- prompt,
21190
- spanData.conversationHistory,
21191
- options,
21192
- spanData.currentMessageStartTime,
21193
- await spanData.span.export()
21194
- );
21195
- if (finalMessage) {
21196
- spanData.conversationHistory.push(finalMessage);
21197
- }
21198
- }
21199
- spanData.span.log({
21200
- output: spanData.conversationHistory.length > 0 ? spanData.conversationHistory[spanData.conversationHistory.length - 1] : void 0
21201
- });
21202
- } catch (error) {
21203
- console.error(
21204
- "Error extracting output for Claude Agent SDK:",
21205
- error
21206
- );
21207
- } finally {
21208
- spanData.span.end();
21209
- spans.delete(event);
21210
- }
21211
- },
21212
- onError: (error) => {
21213
- spanData.span.log({
21214
- error: error.message
21215
- });
21216
- spanData.span.end();
21217
- spans.delete(event);
21218
- }
21219
- });
21220
- } else {
21221
- try {
21222
- spanData.span.log({
21223
- output: eventResult
21224
- });
21225
- } catch (error) {
21226
- console.error(
21227
- "Error extracting output for Claude Agent SDK:",
21228
- error
21229
- );
21230
- } finally {
21231
- spanData.span.end();
21232
- spans.delete(event);
21233
- }
21247
+ if (response?.totalDurationMs) {
21248
+ metadata["claude_agent_sdk.duration_ms"] = response.totalDurationMs;
21234
21249
  }
21235
- },
21236
- error: (event) => {
21237
- const spanData = spans.get(event);
21238
- if (!spanData || !event.error) {
21239
- return;
21250
+ if (response?.totalToolUseCount !== void 0) {
21251
+ metadata["claude_agent_sdk.tool_use_count"] = response.totalToolUseCount;
21240
21252
  }
21241
- const { span } = spanData;
21242
- span.log({
21243
- error: event.error.message
21253
+ subAgentSpan.log({
21254
+ metadata,
21255
+ output: response?.content
21244
21256
  });
21245
- span.end();
21246
- spans.delete(event);
21257
+ } finally {
21258
+ subAgentSpan.end();
21259
+ endedSubAgentSpans.add(toolUseID);
21247
21260
  }
21248
- };
21249
- channel2.subscribe(handlers);
21250
- this.unsubscribers.push(() => {
21251
- channel2.unsubscribe(handlers);
21252
- });
21253
- }
21254
- };
21255
-
21256
- // src/instrumentation/plugins/google-genai-channels.ts
21257
- var googleGenAIChannels = defineChannels("@google/genai", {
21258
- generateContent: channel({
21259
- channelName: "models.generateContent",
21260
- kind: "async"
21261
- }),
21262
- generateContentStream: channel({
21263
- channelName: "models.generateContentStream",
21264
- kind: "async"
21265
- })
21266
- });
21267
-
21268
- // src/instrumentation/plugins/google-genai-plugin.ts
21261
+ return {};
21262
+ }
21263
+ const toolSpan = activeToolSpans.get(toolUseID);
21264
+ if (!toolSpan) {
21265
+ return {};
21266
+ }
21267
+ try {
21268
+ toolSpan.log({ output: input.tool_response });
21269
+ } finally {
21270
+ toolSpan.end();
21271
+ activeToolSpans.delete(toolUseID);
21272
+ }
21273
+ return {};
21274
+ };
21275
+ const postToolUseFailure = async (input, toolUseID) => {
21276
+ if (input.hook_event_name !== "PostToolUseFailure" || !toolUseID) {
21277
+ return {};
21278
+ }
21279
+ const subAgentSpan = subAgentSpans.get(toolUseID);
21280
+ if (subAgentSpan) {
21281
+ try {
21282
+ subAgentSpan.log({ error: input.error });
21283
+ } finally {
21284
+ subAgentSpan.end();
21285
+ endedSubAgentSpans.add(toolUseID);
21286
+ }
21287
+ return {};
21288
+ }
21289
+ const toolSpan = activeToolSpans.get(toolUseID);
21290
+ if (!toolSpan) {
21291
+ return {};
21292
+ }
21293
+ const parsed = parseToolName2(input.tool_name);
21294
+ try {
21295
+ toolSpan.log({
21296
+ error: input.error,
21297
+ metadata: {
21298
+ "claude_agent_sdk.is_interrupt": input.is_interrupt,
21299
+ "claude_agent_sdk.session_id": input.session_id,
21300
+ "gen_ai.tool.call.id": toolUseID,
21301
+ "gen_ai.tool.name": parsed.toolName,
21302
+ ...parsed.mcpServer && { "mcp.server": parsed.mcpServer }
21303
+ }
21304
+ });
21305
+ } finally {
21306
+ toolSpan.end();
21307
+ activeToolSpans.delete(toolUseID);
21308
+ }
21309
+ return {};
21310
+ };
21311
+ return { postToolUse, postToolUseFailure, preToolUse };
21312
+ }
21313
+ function injectTracingHooks2(options, resolveParentSpan, activeToolSpans, subAgentSpans, endedSubAgentSpans) {
21314
+ const { preToolUse, postToolUse, postToolUseFailure } = createToolTracingHooks2(
21315
+ resolveParentSpan,
21316
+ activeToolSpans,
21317
+ options.mcpServers,
21318
+ subAgentSpans,
21319
+ endedSubAgentSpans
21320
+ );
21321
+ const existingHooks = options.hooks ?? {};
21322
+ return {
21323
+ ...options,
21324
+ hooks: {
21325
+ ...existingHooks,
21326
+ PostToolUse: [
21327
+ ...existingHooks.PostToolUse ?? [],
21328
+ { hooks: [postToolUse] }
21329
+ ],
21330
+ PostToolUseFailure: [
21331
+ ...existingHooks.PostToolUseFailure ?? [],
21332
+ {
21333
+ hooks: [postToolUseFailure]
21334
+ }
21335
+ ],
21336
+ PreToolUse: [
21337
+ ...existingHooks.PreToolUse ?? [],
21338
+ { hooks: [preToolUse] }
21339
+ ]
21340
+ }
21341
+ };
21342
+ }
21343
+ async function finalizeCurrentMessageGroup(state) {
21344
+ if (state.currentMessages.length === 0) {
21345
+ return;
21346
+ }
21347
+ const parentToolUseId = state.currentMessages[0]?.parent_tool_use_id ?? null;
21348
+ let parentSpan = await state.span.export();
21349
+ if (parentToolUseId) {
21350
+ const subAgentSpan = state.subAgentSpans.get(parentToolUseId);
21351
+ if (subAgentSpan) {
21352
+ parentSpan = await subAgentSpan.export();
21353
+ }
21354
+ }
21355
+ const finalMessage = await createLLMSpanForMessages(
21356
+ state.currentMessages,
21357
+ state.originalPrompt,
21358
+ state.finalResults,
21359
+ state.options,
21360
+ state.currentMessageStartTime,
21361
+ state.capturedPromptMessages,
21362
+ parentSpan
21363
+ );
21364
+ if (finalMessage) {
21365
+ state.finalResults.push(finalMessage);
21366
+ }
21367
+ const lastMessage = state.currentMessages[state.currentMessages.length - 1];
21368
+ if (lastMessage?.message?.usage) {
21369
+ state.accumulatedOutputTokens += getNumberProperty3(lastMessage.message.usage, "output_tokens") || 0;
21370
+ }
21371
+ state.currentMessages.length = 0;
21372
+ }
21373
+ function maybeTrackToolUseContext(state, message) {
21374
+ if (message.type !== "assistant" || !Array.isArray(message.message?.content)) {
21375
+ return;
21376
+ }
21377
+ const parentToolUseId = message.parent_tool_use_id ?? null;
21378
+ for (const block of message.message.content) {
21379
+ if (typeof block !== "object" || block === null || !("type" in block) || block.type !== "tool_use" || !("id" in block) || typeof block.id !== "string") {
21380
+ continue;
21381
+ }
21382
+ state.toolUseToParent.set(block.id, parentToolUseId);
21383
+ if (block.name === "Task" && typeof block.input === "object" && block.input !== null && "subagent_type" in block.input && typeof block.input.subagent_type === "string") {
21384
+ state.pendingSubAgentNames.set(block.id, block.input.subagent_type);
21385
+ }
21386
+ }
21387
+ }
21388
+ async function maybeStartSubAgentSpan(state, message) {
21389
+ if (!("parent_tool_use_id" in message)) {
21390
+ return;
21391
+ }
21392
+ const parentToolUseId = message.parent_tool_use_id;
21393
+ if (!parentToolUseId) {
21394
+ return;
21395
+ }
21396
+ await ensureSubAgentSpan(
21397
+ state.pendingSubAgentNames,
21398
+ state.span,
21399
+ state.subAgentSpans,
21400
+ parentToolUseId
21401
+ );
21402
+ }
21403
+ async function ensureSubAgentSpan(pendingSubAgentNames, rootSpan, subAgentSpans, parentToolUseId) {
21404
+ const existingSpan = subAgentSpans.get(parentToolUseId);
21405
+ if (existingSpan) {
21406
+ return existingSpan;
21407
+ }
21408
+ const agentName = pendingSubAgentNames.get(parentToolUseId);
21409
+ const spanName = agentName ? `Agent: ${agentName}` : "Agent: sub-agent";
21410
+ const subAgentSpan = startSpan({
21411
+ event: {
21412
+ metadata: {
21413
+ ...agentName && { "claude_agent_sdk.agent_type": agentName }
21414
+ }
21415
+ },
21416
+ name: spanName,
21417
+ parent: await rootSpan.export(),
21418
+ spanAttributes: { type: "task" /* TASK */ }
21419
+ });
21420
+ subAgentSpans.set(parentToolUseId, subAgentSpan);
21421
+ return subAgentSpan;
21422
+ }
21423
+ async function handleStreamMessage(state, message) {
21424
+ maybeTrackToolUseContext(state, message);
21425
+ await maybeStartSubAgentSpan(state, message);
21426
+ const messageId = message.message?.id;
21427
+ if (messageId && messageId !== state.currentMessageId) {
21428
+ await finalizeCurrentMessageGroup(state);
21429
+ state.currentMessageId = messageId;
21430
+ state.currentMessageStartTime = getCurrentUnixTimestamp();
21431
+ }
21432
+ if (message.type === "assistant" && message.message?.usage) {
21433
+ state.currentMessages.push(message);
21434
+ }
21435
+ if (message.type !== "result" || !message.usage) {
21436
+ return;
21437
+ }
21438
+ const finalUsageMetrics = extractUsageFromMessage(message);
21439
+ if (state.currentMessages.length > 0 && finalUsageMetrics.completion_tokens !== void 0) {
21440
+ const lastMessage = state.currentMessages[state.currentMessages.length - 1];
21441
+ if (lastMessage?.message?.usage) {
21442
+ const adjustedTokens = finalUsageMetrics.completion_tokens - state.accumulatedOutputTokens;
21443
+ if (adjustedTokens >= 0) {
21444
+ lastMessage.message.usage.output_tokens = adjustedTokens;
21445
+ }
21446
+ const resultUsage = message.usage;
21447
+ if (resultUsage && typeof resultUsage === "object") {
21448
+ const cacheReadTokens = getNumberProperty3(
21449
+ resultUsage,
21450
+ "cache_read_input_tokens"
21451
+ );
21452
+ if (cacheReadTokens !== void 0) {
21453
+ lastMessage.message.usage.cache_read_input_tokens = cacheReadTokens;
21454
+ }
21455
+ const cacheCreationTokens = getNumberProperty3(
21456
+ resultUsage,
21457
+ "cache_creation_input_tokens"
21458
+ );
21459
+ if (cacheCreationTokens !== void 0) {
21460
+ lastMessage.message.usage.cache_creation_input_tokens = cacheCreationTokens;
21461
+ }
21462
+ }
21463
+ }
21464
+ }
21465
+ const metadata = {};
21466
+ if (message.num_turns !== void 0) {
21467
+ metadata.num_turns = message.num_turns;
21468
+ }
21469
+ if (message.session_id !== void 0) {
21470
+ metadata.session_id = message.session_id;
21471
+ }
21472
+ if (Object.keys(metadata).length > 0) {
21473
+ state.span.log({ metadata });
21474
+ }
21475
+ }
21476
+ async function finalizeQuerySpan(state) {
21477
+ try {
21478
+ await finalizeCurrentMessageGroup(state);
21479
+ state.span.log({
21480
+ output: state.finalResults.length > 0 ? state.finalResults[state.finalResults.length - 1] : void 0
21481
+ });
21482
+ if (state.capturedPromptMessages) {
21483
+ if (state.promptStarted()) {
21484
+ await state.promptDone;
21485
+ }
21486
+ if (state.capturedPromptMessages.length > 0) {
21487
+ state.span.log({
21488
+ input: formatCapturedMessages(state.capturedPromptMessages)
21489
+ });
21490
+ }
21491
+ }
21492
+ } finally {
21493
+ for (const [id, subAgentSpan] of state.subAgentSpans) {
21494
+ if (!state.endedSubAgentSpans.has(id)) {
21495
+ subAgentSpan.end();
21496
+ }
21497
+ }
21498
+ state.subAgentSpans.clear();
21499
+ state.span.end();
21500
+ }
21501
+ }
21502
+ var ClaudeAgentSDKPlugin = class extends BasePlugin {
21503
+ onEnable() {
21504
+ this.subscribeToQuery();
21505
+ }
21506
+ onDisable() {
21507
+ for (const unsubscribe of this.unsubscribers) {
21508
+ unsubscribe();
21509
+ }
21510
+ this.unsubscribers = [];
21511
+ }
21512
+ subscribeToQuery() {
21513
+ const channel2 = claudeAgentSDKChannels.query.tracingChannel();
21514
+ const spans = /* @__PURE__ */ new WeakMap();
21515
+ const handlers = {
21516
+ start: (event) => {
21517
+ const params = event.arguments[0] ?? {};
21518
+ const originalPrompt = params.prompt;
21519
+ const options = params.options ?? {};
21520
+ const promptIsAsyncIterable = isAsyncIterable5(originalPrompt);
21521
+ let promptStarted = false;
21522
+ let capturedPromptMessages;
21523
+ let resolvePromptDone;
21524
+ const promptDone = new Promise((resolve) => {
21525
+ resolvePromptDone = resolve;
21526
+ });
21527
+ if (promptIsAsyncIterable) {
21528
+ capturedPromptMessages = [];
21529
+ const promptStream = originalPrompt;
21530
+ params.prompt = (async function* () {
21531
+ promptStarted = true;
21532
+ try {
21533
+ for await (const message of promptStream) {
21534
+ capturedPromptMessages.push(message);
21535
+ yield message;
21536
+ }
21537
+ } finally {
21538
+ resolvePromptDone?.();
21539
+ }
21540
+ })();
21541
+ }
21542
+ const span = startSpan({
21543
+ name: "Claude Agent",
21544
+ spanAttributes: {
21545
+ type: "task" /* TASK */
21546
+ }
21547
+ });
21548
+ const startTime = getCurrentUnixTimestamp();
21549
+ try {
21550
+ span.log({
21551
+ input: typeof originalPrompt === "string" ? originalPrompt : promptIsAsyncIterable ? void 0 : originalPrompt !== void 0 ? String(originalPrompt) : void 0,
21552
+ metadata: filterSerializableOptions2(options)
21553
+ });
21554
+ } catch (error) {
21555
+ console.error("Error extracting input for Claude Agent SDK:", error);
21556
+ }
21557
+ const activeToolSpans = /* @__PURE__ */ new Map();
21558
+ const subAgentSpans = /* @__PURE__ */ new Map();
21559
+ const endedSubAgentSpans = /* @__PURE__ */ new Set();
21560
+ const toolUseToParent = /* @__PURE__ */ new Map();
21561
+ const pendingSubAgentNames = /* @__PURE__ */ new Map();
21562
+ const optionsWithHooks = injectTracingHooks2(
21563
+ options,
21564
+ async (toolUseID) => {
21565
+ const parentToolUseId = toolUseToParent.get(toolUseID);
21566
+ if (parentToolUseId) {
21567
+ const subAgentSpan = await ensureSubAgentSpan(
21568
+ pendingSubAgentNames,
21569
+ span,
21570
+ subAgentSpans,
21571
+ parentToolUseId
21572
+ );
21573
+ return subAgentSpan.export();
21574
+ }
21575
+ return span.export();
21576
+ },
21577
+ activeToolSpans,
21578
+ subAgentSpans,
21579
+ endedSubAgentSpans
21580
+ );
21581
+ params.options = optionsWithHooks;
21582
+ event.arguments[0] = params;
21583
+ spans.set(event, {
21584
+ accumulatedOutputTokens: 0,
21585
+ activeToolSpans,
21586
+ capturedPromptMessages,
21587
+ currentMessageId: void 0,
21588
+ currentMessageStartTime: startTime,
21589
+ currentMessages: [],
21590
+ endedSubAgentSpans,
21591
+ finalResults: [],
21592
+ options: optionsWithHooks,
21593
+ originalPrompt,
21594
+ pendingSubAgentNames,
21595
+ processing: Promise.resolve(),
21596
+ promptDone,
21597
+ promptStarted: () => promptStarted,
21598
+ span,
21599
+ subAgentSpans,
21600
+ toolUseToParent
21601
+ });
21602
+ },
21603
+ end: (event) => {
21604
+ const state = spans.get(event);
21605
+ if (!state) {
21606
+ return;
21607
+ }
21608
+ const eventResult = event.result;
21609
+ if (eventResult === void 0) {
21610
+ state.span.end();
21611
+ spans.delete(event);
21612
+ return;
21613
+ }
21614
+ if (isAsyncIterable5(eventResult)) {
21615
+ patchStreamIfNeeded(eventResult, {
21616
+ onChunk: (message) => {
21617
+ maybeTrackToolUseContext(state, message);
21618
+ state.processing = state.processing.then(() => handleStreamMessage(state, message)).catch((error) => {
21619
+ console.error(
21620
+ "Error processing Claude Agent SDK stream chunk:",
21621
+ error
21622
+ );
21623
+ });
21624
+ },
21625
+ onComplete: () => {
21626
+ void state.processing.then(() => finalizeQuerySpan(state)).finally(() => {
21627
+ spans.delete(event);
21628
+ });
21629
+ },
21630
+ onError: (error) => {
21631
+ void state.processing.then(() => {
21632
+ state.span.log({
21633
+ error: error.message
21634
+ });
21635
+ }).then(() => finalizeQuerySpan(state)).finally(() => {
21636
+ spans.delete(event);
21637
+ });
21638
+ }
21639
+ });
21640
+ return;
21641
+ }
21642
+ try {
21643
+ state.span.log({ output: eventResult });
21644
+ } catch (error) {
21645
+ console.error("Error extracting output for Claude Agent SDK:", error);
21646
+ } finally {
21647
+ state.span.end();
21648
+ spans.delete(event);
21649
+ }
21650
+ },
21651
+ error: (event) => {
21652
+ const state = spans.get(event);
21653
+ if (!state || !event.error) {
21654
+ return;
21655
+ }
21656
+ state.span.log({
21657
+ error: event.error.message
21658
+ });
21659
+ state.span.end();
21660
+ spans.delete(event);
21661
+ }
21662
+ };
21663
+ channel2.subscribe(handlers);
21664
+ this.unsubscribers.push(() => {
21665
+ channel2.unsubscribe(handlers);
21666
+ });
21667
+ }
21668
+ };
21669
+
21670
+ // src/instrumentation/plugins/google-genai-channels.ts
21671
+ var googleGenAIChannels = defineChannels("@google/genai", {
21672
+ generateContent: channel({
21673
+ channelName: "models.generateContent",
21674
+ kind: "async"
21675
+ }),
21676
+ generateContentStream: channel({
21677
+ channelName: "models.generateContentStream",
21678
+ kind: "async"
21679
+ })
21680
+ });
21681
+
21682
+ // src/instrumentation/plugins/google-genai-plugin.ts
21683
+ var GOOGLE_GENAI_INTERNAL_CONTEXT = {
21684
+ caller_filename: "<node-internal>",
21685
+ caller_functionname: "<node-internal>",
21686
+ caller_lineno: 0
21687
+ };
21688
+ function createWrapperParityEvent(args) {
21689
+ return {
21690
+ context: GOOGLE_GENAI_INTERNAL_CONTEXT,
21691
+ input: args.input,
21692
+ metadata: args.metadata
21693
+ };
21694
+ }
21269
21695
  var GoogleGenAIPlugin = class extends BasePlugin {
21270
21696
  onEnable() {
21271
21697
  this.subscribeToGoogleGenAIChannels();
@@ -21274,51 +21700,282 @@ var GoogleGenAIPlugin = class extends BasePlugin {
21274
21700
  this.unsubscribers = unsubscribeAll(this.unsubscribers);
21275
21701
  }
21276
21702
  subscribeToGoogleGenAIChannels() {
21277
- this.unsubscribers.push(
21278
- traceAsyncChannel(googleGenAIChannels.generateContent, {
21279
- name: "google-genai.generateContent",
21280
- type: "llm" /* LLM */,
21281
- extractInput: ([params]) => {
21282
- const input = serializeInput2(params);
21283
- const metadata = extractMetadata2(params);
21284
- return {
21285
- input,
21286
- metadata: { ...metadata, provider: "google-genai" }
21287
- };
21288
- },
21289
- extractOutput: (result) => {
21290
- return result;
21291
- },
21292
- extractMetrics: (result, startTime) => {
21293
- return extractGenerateContentMetrics2(result, startTime);
21294
- }
21295
- })
21703
+ this.subscribeToGenerateContentChannel();
21704
+ this.subscribeToGenerateContentStreamChannel();
21705
+ }
21706
+ subscribeToGenerateContentChannel() {
21707
+ const tracingChannel = googleGenAIChannels.generateContent.tracingChannel();
21708
+ const states = /* @__PURE__ */ new WeakMap();
21709
+ const unbindCurrentSpanStore = bindCurrentSpanStoreToStart2(
21710
+ tracingChannel,
21711
+ states,
21712
+ (event) => {
21713
+ const params = event.arguments[0];
21714
+ const input = serializeInput2(params);
21715
+ const metadata = extractMetadata2(params);
21716
+ const span = startSpan({
21717
+ name: "generate_content",
21718
+ spanAttributes: {
21719
+ type: "llm" /* LLM */
21720
+ },
21721
+ event: createWrapperParityEvent({ input, metadata })
21722
+ });
21723
+ return {
21724
+ span,
21725
+ startTime: getCurrentUnixTimestamp()
21726
+ };
21727
+ }
21296
21728
  );
21297
- this.unsubscribers.push(
21298
- traceStreamingChannel(googleGenAIChannels.generateContentStream, {
21299
- name: "google-genai.generateContentStream",
21300
- type: "llm" /* LLM */,
21301
- extractInput: ([params]) => {
21729
+ const handlers = {
21730
+ start: (event) => {
21731
+ ensureSpanState(states, event, () => {
21732
+ const params = event.arguments[0];
21302
21733
  const input = serializeInput2(params);
21303
21734
  const metadata = extractMetadata2(params);
21735
+ const span = startSpan({
21736
+ name: "generate_content",
21737
+ spanAttributes: {
21738
+ type: "llm" /* LLM */
21739
+ },
21740
+ event: createWrapperParityEvent({ input, metadata })
21741
+ });
21304
21742
  return {
21305
- input,
21306
- metadata: { ...metadata, provider: "google-genai" }
21743
+ span,
21744
+ startTime: getCurrentUnixTimestamp()
21307
21745
  };
21308
- },
21309
- extractOutput: (result) => {
21310
- return result;
21311
- },
21312
- extractMetrics: () => {
21313
- return {};
21314
- },
21315
- aggregateChunks: (chunks, _result, _endEvent, startTime) => {
21316
- return aggregateGenerateContentChunks2(chunks, startTime);
21746
+ });
21747
+ },
21748
+ asyncEnd: (event) => {
21749
+ const spanState = states.get(event);
21750
+ if (!spanState) {
21751
+ return;
21317
21752
  }
21318
- })
21319
- );
21753
+ try {
21754
+ spanState.span.log({
21755
+ metrics: cleanMetrics2(
21756
+ extractGenerateContentMetrics2(
21757
+ event.result,
21758
+ spanState.startTime
21759
+ )
21760
+ ),
21761
+ output: event.result
21762
+ });
21763
+ } finally {
21764
+ spanState.span.end();
21765
+ states.delete(event);
21766
+ }
21767
+ },
21768
+ error: (event) => {
21769
+ logErrorAndEndSpan(states, event);
21770
+ }
21771
+ };
21772
+ tracingChannel.subscribe(handlers);
21773
+ this.unsubscribers.push(() => {
21774
+ unbindCurrentSpanStore?.();
21775
+ tracingChannel.unsubscribe(handlers);
21776
+ });
21777
+ }
21778
+ subscribeToGenerateContentStreamChannel() {
21779
+ const tracingChannel = googleGenAIChannels.generateContentStream.tracingChannel();
21780
+ const handlers = {
21781
+ start: (event) => {
21782
+ const streamEvent = event;
21783
+ const params = event.arguments[0];
21784
+ streamEvent.googleGenAIInput = serializeInput2(params);
21785
+ streamEvent.googleGenAIMetadata = extractMetadata2(params);
21786
+ },
21787
+ asyncEnd: (event) => {
21788
+ const streamEvent = event;
21789
+ patchGoogleGenAIStreamingResult({
21790
+ input: streamEvent.googleGenAIInput,
21791
+ metadata: streamEvent.googleGenAIMetadata,
21792
+ result: streamEvent.result
21793
+ });
21794
+ },
21795
+ error: () => {
21796
+ }
21797
+ };
21798
+ tracingChannel.subscribe(handlers);
21799
+ this.unsubscribers.push(() => {
21800
+ tracingChannel.unsubscribe(handlers);
21801
+ });
21320
21802
  }
21321
21803
  };
21804
+ function ensureSpanState(states, event, create) {
21805
+ const existing = states.get(event);
21806
+ if (existing) {
21807
+ return existing;
21808
+ }
21809
+ const created = create();
21810
+ states.set(event, created);
21811
+ return created;
21812
+ }
21813
+ function bindCurrentSpanStoreToStart2(tracingChannel, states, create) {
21814
+ const state = _internalGetGlobalState();
21815
+ const startChannel = tracingChannel.start;
21816
+ const currentSpanStore = state?.contextManager ? state.contextManager[BRAINTRUST_CURRENT_SPAN_STORE] : void 0;
21817
+ if (!startChannel?.bindStore || !currentSpanStore) {
21818
+ return void 0;
21819
+ }
21820
+ startChannel.bindStore(
21821
+ currentSpanStore,
21822
+ (event) => ensureSpanState(
21823
+ states,
21824
+ event,
21825
+ () => create(event)
21826
+ ).span
21827
+ );
21828
+ return () => {
21829
+ startChannel.unbindStore?.(currentSpanStore);
21830
+ };
21831
+ }
21832
+ function logErrorAndEndSpan(states, event) {
21833
+ const spanState = states.get(event);
21834
+ if (!spanState) {
21835
+ return;
21836
+ }
21837
+ spanState.span.log({
21838
+ error: event.error.message
21839
+ });
21840
+ spanState.span.end();
21841
+ states.delete(event);
21842
+ }
21843
+ function patchGoogleGenAIStreamingResult(args) {
21844
+ const { input, metadata, result } = args;
21845
+ if (!input || !metadata || !result || typeof result !== "object" || typeof result.next !== "function") {
21846
+ return false;
21847
+ }
21848
+ const chunks = [];
21849
+ let firstTokenTime = null;
21850
+ let finalized = false;
21851
+ let span = null;
21852
+ let startTime = null;
21853
+ const ensureSpan = () => {
21854
+ if (!span) {
21855
+ span = startSpan({
21856
+ name: "generate_content_stream",
21857
+ spanAttributes: {
21858
+ type: "llm" /* LLM */
21859
+ },
21860
+ event: {
21861
+ input,
21862
+ metadata
21863
+ }
21864
+ });
21865
+ startTime = getCurrentUnixTimestamp();
21866
+ }
21867
+ return span;
21868
+ };
21869
+ const finalize = (options) => {
21870
+ if (finalized || !span) {
21871
+ return;
21872
+ }
21873
+ finalized = true;
21874
+ if (options.result) {
21875
+ const { end, ...metricsWithoutEnd } = options.result.metrics;
21876
+ span.log({
21877
+ metrics: cleanMetrics2(metricsWithoutEnd),
21878
+ output: options.result.aggregated
21879
+ });
21880
+ span.end(typeof end === "number" ? { endTime: end } : void 0);
21881
+ return;
21882
+ }
21883
+ if (options.error !== void 0) {
21884
+ span.log({
21885
+ error: options.error instanceof Error ? options.error.message : String(options.error)
21886
+ });
21887
+ }
21888
+ span.end();
21889
+ };
21890
+ const patchIterator = (iterator) => {
21891
+ if (typeof iterator !== "object" || iterator === null || "__braintrustGoogleGenAIPatched" in iterator) {
21892
+ return iterator;
21893
+ }
21894
+ const iteratorRecord = iterator;
21895
+ const originalNext = typeof iteratorRecord.next === "function" ? iteratorRecord.next.bind(iterator) : void 0;
21896
+ const originalReturn = typeof iteratorRecord.return === "function" ? iteratorRecord.return.bind(iterator) : void 0;
21897
+ const originalThrow = typeof iteratorRecord.throw === "function" ? iteratorRecord.throw.bind(iterator) : void 0;
21898
+ const asyncIteratorMethod = iteratorRecord[Symbol.asyncIterator];
21899
+ const originalAsyncIterator = typeof asyncIteratorMethod === "function" ? asyncIteratorMethod.bind(iterator) : void 0;
21900
+ Object.defineProperty(iteratorRecord, "__braintrustGoogleGenAIPatched", {
21901
+ configurable: true,
21902
+ enumerable: false,
21903
+ value: true,
21904
+ writable: false
21905
+ });
21906
+ if (originalNext) {
21907
+ iteratorRecord.next = async (...nextArgs) => {
21908
+ ensureSpan();
21909
+ try {
21910
+ const nextResult = await originalNext(
21911
+ ...nextArgs
21912
+ );
21913
+ if (!nextResult.done && nextResult.value) {
21914
+ if (firstTokenTime === null) {
21915
+ firstTokenTime = getCurrentUnixTimestamp();
21916
+ }
21917
+ chunks.push(nextResult.value);
21918
+ }
21919
+ if (nextResult.done && startTime !== null) {
21920
+ finalize({
21921
+ result: aggregateGenerateContentChunks2(
21922
+ chunks,
21923
+ startTime,
21924
+ firstTokenTime
21925
+ )
21926
+ });
21927
+ }
21928
+ return nextResult;
21929
+ } catch (error) {
21930
+ finalize({ error });
21931
+ throw error;
21932
+ }
21933
+ };
21934
+ }
21935
+ if (originalReturn) {
21936
+ iteratorRecord.return = async (...returnArgs) => {
21937
+ ensureSpan();
21938
+ try {
21939
+ return await originalReturn(
21940
+ ...returnArgs
21941
+ );
21942
+ } finally {
21943
+ if (startTime !== null) {
21944
+ finalize({
21945
+ result: chunks.length > 0 ? aggregateGenerateContentChunks2(
21946
+ chunks,
21947
+ startTime,
21948
+ firstTokenTime
21949
+ ) : void 0
21950
+ });
21951
+ } else {
21952
+ finalize({});
21953
+ }
21954
+ }
21955
+ };
21956
+ }
21957
+ if (originalThrow) {
21958
+ iteratorRecord.throw = async (...throwArgs) => {
21959
+ ensureSpan();
21960
+ try {
21961
+ return await originalThrow(
21962
+ ...throwArgs
21963
+ );
21964
+ } catch (error) {
21965
+ finalize({ error });
21966
+ throw error;
21967
+ }
21968
+ };
21969
+ }
21970
+ iteratorRecord[Symbol.asyncIterator] = () => {
21971
+ const asyncIterator = originalAsyncIterator ? originalAsyncIterator() : iterator;
21972
+ return patchIterator(asyncIterator);
21973
+ };
21974
+ return iterator;
21975
+ };
21976
+ patchIterator(result);
21977
+ return true;
21978
+ }
21322
21979
  function serializeInput2(params) {
21323
21980
  const input = {
21324
21981
  model: params.model,
@@ -21327,11 +21984,13 @@ function serializeInput2(params) {
21327
21984
  if (params.config) {
21328
21985
  const config = tryToDict2(params.config);
21329
21986
  if (config) {
21330
- const tools = serializeTools2(params);
21331
- if (tools) {
21332
- config.tools = tools;
21333
- }
21334
- input.config = config;
21987
+ const filteredConfig = {};
21988
+ Object.keys(config).forEach((key) => {
21989
+ if (key !== "tools") {
21990
+ filteredConfig[key] = config[key];
21991
+ }
21992
+ });
21993
+ input.config = filteredConfig;
21335
21994
  }
21336
21995
  }
21337
21996
  return input;
@@ -21418,12 +22077,18 @@ function extractMetadata2(params) {
21418
22077
  });
21419
22078
  }
21420
22079
  }
22080
+ const tools = serializeTools2(params);
22081
+ if (tools) {
22082
+ metadata.tools = tools;
22083
+ }
21421
22084
  return metadata;
21422
22085
  }
21423
22086
  function extractGenerateContentMetrics2(response, startTime) {
21424
22087
  const metrics = {};
21425
- if (startTime) {
22088
+ if (startTime !== void 0) {
21426
22089
  const end = getCurrentUnixTimestamp();
22090
+ metrics.start = startTime;
22091
+ metrics.end = end;
21427
22092
  metrics.duration = end - startTime;
21428
22093
  }
21429
22094
  if (response?.usageMetadata) {
@@ -21448,19 +22113,18 @@ function populateUsageMetrics2(metrics, usage) {
21448
22113
  metrics.completion_reasoning_tokens = usage.thoughtsTokenCount;
21449
22114
  }
21450
22115
  }
21451
- function aggregateGenerateContentChunks2(chunks, startTime) {
21452
- const metrics = {};
21453
- if (startTime !== void 0) {
21454
- const end = getCurrentUnixTimestamp();
21455
- metrics.duration = end - startTime;
21456
- }
21457
- let firstTokenTime = null;
21458
- if (chunks.length > 0 && firstTokenTime === null && startTime !== void 0) {
21459
- firstTokenTime = getCurrentUnixTimestamp();
22116
+ function aggregateGenerateContentChunks2(chunks, startTime, firstTokenTime) {
22117
+ const end = getCurrentUnixTimestamp();
22118
+ const metrics = {
22119
+ start: startTime,
22120
+ end,
22121
+ duration: end - startTime
22122
+ };
22123
+ if (firstTokenTime !== null) {
21460
22124
  metrics.time_to_first_token = firstTokenTime - startTime;
21461
22125
  }
21462
22126
  if (chunks.length === 0) {
21463
- return { output: {}, metrics };
22127
+ return { aggregated: {}, metrics };
21464
22128
  }
21465
22129
  let text = "";
21466
22130
  let thoughtText = "";
@@ -21496,7 +22160,7 @@ function aggregateGenerateContentChunks2(chunks, startTime) {
21496
22160
  }
21497
22161
  }
21498
22162
  }
21499
- const output = {};
22163
+ const aggregated = {};
21500
22164
  const parts = [];
21501
22165
  if (thoughtText) {
21502
22166
  parts.push({ text: thoughtText, thought: true });
@@ -21522,16 +22186,25 @@ function aggregateGenerateContentChunks2(chunks, startTime) {
21522
22186
  }
21523
22187
  candidates.push(candidateDict);
21524
22188
  }
21525
- output.candidates = candidates;
22189
+ aggregated.candidates = candidates;
21526
22190
  }
21527
22191
  if (usageMetadata) {
21528
- output.usageMetadata = usageMetadata;
22192
+ aggregated.usageMetadata = usageMetadata;
21529
22193
  populateUsageMetrics2(metrics, usageMetadata);
21530
22194
  }
21531
22195
  if (text) {
21532
- output.text = text;
22196
+ aggregated.text = text;
21533
22197
  }
21534
- return { output, metrics };
22198
+ return { aggregated, metrics };
22199
+ }
22200
+ function cleanMetrics2(metrics) {
22201
+ const cleaned = {};
22202
+ for (const [key, value] of Object.entries(metrics)) {
22203
+ if (value !== null && value !== void 0) {
22204
+ cleaned[key] = value;
22205
+ }
22206
+ }
22207
+ return cleaned;
21535
22208
  }
21536
22209
  function tryToDict2(obj) {
21537
22210
  if (obj === null || obj === void 0) {