feat(server-utils): Emit execute_tool spans for Anthropic's beta.messages.toolRunner()

dev-packages/node-integration-tests/suites/tracing/anthropic/v0.129/scenario-tool-runner.mjs

@@ -0,0 +1,136 @@
+import Anthropic from '@anthropic-ai/sdk';
+import * as Sentry from '@sentry/node';
+import express from 'express';
+
+const MODEL = 'claude-3-haiku-20240307';
+
+function toolUseEvents() {
+  const events = [
+    {
+      type: 'message_start',
+      message: {
+        id: 'msg_tools',
+        type: 'message',
+        role: 'assistant',
+        model: MODEL,
+        content: [],
+        usage: { input_tokens: 10 },
+      },
+    },
+  ];
+  [
+    { id: 'toolu_weather', name: 'weather', input: { city: 'Vienna' } },
+    { id: 'toolu_explode', name: 'explode', input: {} },
+  ].forEach((call, index) => {
+    events.push(
+      {
+        type: 'content_block_start',
+        index,
+        content_block: { type: 'tool_use', id: call.id, name: call.name, input: {} },
+      },
+      {
+        type: 'content_block_delta',
+        index,
+        delta: { type: 'input_json_delta', partial_json: JSON.stringify(call.input) },
+      },
+      { type: 'content_block_stop', index },
+    );
+  });
+  events.push(
+    { type: 'message_delta', delta: { stop_reason: 'tool_use', stop_sequence: null }, usage: { output_tokens: 20 } },
+    { type: 'message_stop' },
+  );
+  return events;
+}
+
+function finalEvents() {
+  return [
+    {
+      type: 'message_start',
+      message: {
+        id: 'msg_final',
+        type: 'message',
+        role: 'assistant',
+        model: MODEL,
+        content: [],
+        usage: { input_tokens: 30 },
+      },
+    },
+    { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } },
+    { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Sunny.' } },
+    { type: 'content_block_stop', index: 0 },
+    { type: 'message_delta', delta: { stop_reason: 'end_turn', stop_sequence: null }, usage: { output_tokens: 5 } },
+    { type: 'message_stop' },
+  ];
+}
+
+function startMockAnthropicServer() {
+  const app = express();
+  app.use(express.json());
+
+  let requests = 0;
+  app.post('/anthropic/v1/messages', (_req, res) => {
+    res.writeHead(200, { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', Connection: 'keep-alive' });
+
+    const events = requests++ === 0 ? toolUseEvents() : finalEvents();
+    events.forEach((event, index) => {
+      setTimeout(() => {
+        res.write(`event: ${event.type}\n`);
+        res.write(`data: ${JSON.stringify(event)}\n\n`);
+        if (index === events.length - 1) {
+          res.end();
+        }
+      }, index * 10);
+    });
+  });
+
+  return new Promise(resolve => {
+    const server = app.listen(0, () => {
+      resolve(server);
+    });
+  });
+}
+
+async function run() {
+  const server = await startMockAnthropicServer();
+
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    const client = new Anthropic({
+      apiKey: 'mock-api-key',
+      baseURL: `http://localhost:${server.address().port}/anthropic`,
+    });
+
+    const runner = client.beta.messages.toolRunner({
+      model: MODEL,
+      max_tokens: 100,
+      messages: [{ role: 'user', content: 'What is the weather in Vienna?' }],
+      tools: [
+        {
+          name: 'weather',
+          description: 'Get the weather by city',
+          input_schema: { type: 'object', properties: { city: { type: 'string' } } },
+          run: () => 'sunny',
+        },
+        {
+          name: 'explode',
+          description: 'Always fails',
+          input_schema: { type: 'object', properties: {} },
+          run: () => {
+            throw new Error('Tool exploded');
+          },
+        },
+      ],
+      stream: true,
+      runToolsEagerly: true,
+    });
+    for await (const stream of runner) {
+      await stream.finalMessage();
+    }
+  });
+
+  await Sentry.flush(2000);
+
+  server.close();
+}
+
+run();

dev-packages/node-integration-tests/suites/tracing/anthropic/v0.129/test.ts

@@ -1,5 +1,13 @@
-import { GEN_AI_RESPONSE_ID, GEN_AI_RESPONSE_STREAMING, SENTRY_OP } from '@sentry/conventions/attributes';
+import type { Event, SerializedStreamedSpanContainer } from '@sentry/core';
+import {
+  GEN_AI_RESPONSE_ID,
+  GEN_AI_RESPONSE_STREAMING,
+  GEN_AI_TOOL_NAME,
+  SENTRY_OP,
+} from '@sentry/conventions/attributes';
+import { GEN_AI_EXECUTE_TOOL } from '@sentry/conventions/op';
 import { afterAll, describe, expect } from 'vitest';
+import { GEN_AI_TOOL_CALL_ID_ATTRIBUTE } from '../../../../../../packages/server-utils/src/ai/core/gen-ai-attributes';
 import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../../utils/runner';
 
 // The attribute extraction is version-independent and covered by the suite pinned to 0.63. What a
@@ -72,4 +80,61 @@ describe('Anthropic integration (0.129)', () => {
       },
     },
   );
+
+  createEsmAndCjsTests(
+    __dirname,
+    'scenario-tool-runner.mjs',
+    'instrument.mjs',
+    (createRunner, test) => {
+      test('emits an execute_tool span for each tool the eager runner starts', async () => {
+        let receivedEvent: Event | undefined;
+        let receivedSpans: SerializedStreamedSpanContainer['items'] = [];
+
+        await createRunner()
+          .unordered()
+          .expect({
+            span: container => {
+              expect(container.items.find(span => span.is_segment && span.name === 'main')).toBeDefined();
+              receivedSpans = container.items;
+            },
+          })
+          .expect({
+            event: event => {
+              receivedEvent = event;
+            },
+          })
+          .start()
+          .completed();
+
+        const main = receivedSpans.find(span => span.is_segment && span.name === 'main')!;
+        const chat = receivedSpans.find(span => span.attributes[GEN_AI_RESPONSE_ID]?.value === 'msg_tools')!;
+        const weather = receivedSpans.find(span => span.attributes[GEN_AI_TOOL_NAME]?.value === 'weather')!;
+        const explode = receivedSpans.find(span => span.attributes[GEN_AI_TOOL_NAME]?.value === 'explode')!;
+
+        expect(weather).toMatchObject({
+          name: 'execute_tool weather',
+          parent_span_id: main.span_id,
+          status: 'ok',
+          attributes: expect.objectContaining({
+            [SENTRY_OP]: { type: 'string', value: GEN_AI_EXECUTE_TOOL },
+            [GEN_AI_TOOL_CALL_ID_ATTRIBUTE]: { type: 'string', value: 'toolu_weather' },
+          }),
+        });
+        // The eager runner starts the first tool while the reply that called it is still streaming.
+        expect(weather.start_timestamp).toBeLessThan(chat.end_timestamp);
+
+        expect(explode).toMatchObject({ name: 'execute_tool explode', parent_span_id: main.span_id, status: 'error' });
+        expect(receivedEvent).toMatchObject({
+          exception: {
+            values: [{ value: 'Tool exploded', mechanism: { type: 'auto.ai.anthropic', handled: true } }],
+          },
+        });
+      });
+    },
+    {
+      additionalDependencies: {
+        '@anthropic-ai/sdk': '0.129.0',
+      },
+    },
+  );
 });

packages/server-utils/src/ai/anthropic-ai/constants.ts

@@ -9,3 +9,5 @@ export const ANTHROPIC_METHOD_REGISTRY = {
   'completions.create': { operation: 'chat' },
   'beta.messages.create': { operation: 'chat' },
 } as const satisfies InstrumentedMethodRegistry;
+
+export const ANTHROPIC_AI_ORIGIN = 'auto.ai.anthropic';

packages/server-utils/src/ai/anthropic-ai/index.ts

@@ -25,8 +25,9 @@ import {
   resolveAIRecordingOptions,
   wrapPromiseWithMethods,
 } from '../core/utils';
-import { ANTHROPIC_METHOD_REGISTRY } from './constants';
+import { ANTHROPIC_AI_ORIGIN, ANTHROPIC_METHOD_REGISTRY } from './constants';
 import { instrumentAsyncIterableStream, instrumentMessageStream } from './streaming';
+import { instrumentToolRunnerMethod } from './tool-runner';
 import type { AnthropicAiOptions, AnthropicAiResponse, AnthropicAiStreamingEvent, ContentBlock } from './types';
 import { handleResponseError, messagesFromParams, setMessagesAttribute } from './utils';
 
@@ -50,7 +51,7 @@ export function extractRequestAttributes(
   const attributes: Record<string, unknown> = {
     [GEN_AI_PROVIDER_NAME]: 'anthropic',
     [GEN_AI_OPERATION_NAME]: operationName,
-    [SENTRY_ORIGIN]: 'auto.ai.anthropic',
+    [SENTRY_ORIGIN]: ANTHROPIC_AI_ORIGIN,
   };
 
   if (args.length > 0 && typeof args[0] === 'object' && args[0] !== null) {
@@ -350,6 +351,11 @@ function instrumentClientInPlace<T extends object>(client: T, options: Anthropic
     owner[methodName] = instrumented;
   }
 
+  const betaMessages = (client as { beta?: { messages?: object } }).beta?.messages;
+  if (betaMessages) {
+    instrumentToolRunnerMethod(betaMessages, options);
+  }
+
   return client;
 }
 

packages/server-utils/src/ai/anthropic-ai/tool-runner.ts

@@ -0,0 +1,139 @@
+import type { Span, WrappedFunction } from '@sentry/core';
+import {
+  captureException,
+  fill,
+  getOriginalFunction,
+  handleCallbackErrors,
+  SPAN_STATUS_ERROR,
+  startSpan,
+  stringify,
+} from '@sentry/core';
+import {
+  GEN_AI_OPERATION_NAME,
+  GEN_AI_TOOL_CALL_ARGUMENTS,
+  GEN_AI_TOOL_CALL_RESULT,
+  GEN_AI_TOOL_DESCRIPTION,
+  GEN_AI_TOOL_NAME,
+  SENTRY_ORIGIN,
+} from '@sentry/conventions/attributes';
+import { GEN_AI_EXECUTE_TOOL } from '@sentry/conventions/op';
+import { GEN_AI_TOOL_CALL_ID_ATTRIBUTE } from '../core/gen-ai-attributes';
+import { resolveAIRecordingOptions } from '../core/utils';
+import { ANTHROPIC_AI_ORIGIN } from './constants';
+import type { AnthropicAiOptions } from './types';
+
+type AnyFunction = (...args: unknown[]) => unknown;
+type RunnableTool = { name: string; description?: unknown };
+type RunContext = { toolUse?: { id?: unknown }; toolUseBlock?: { id?: unknown }; signal?: AbortSignal } | undefined;
+type ToolRunner = { params?: { tools?: unknown } };
+
+/** Wrap `beta.messages.toolRunner()` so the runners it returns trace their tool calls. */
+export function instrumentToolRunnerMethod(messages: object, options: AnthropicAiOptions): void {
+  wrapInPlace(messages, 'toolRunner', (toolRunner, thisArg, args) => {
+    const runner = toolRunner.apply(thisArg, args);
+    instrumentToolRunner(runner, options);
+    return runner;
+  });
+}
+
+/** Give each tool call of a tool runner a `gen_ai.execute_tool` span. */
+export function instrumentToolRunner(runner: unknown, options: AnthropicAiOptions): void {
+  if (!runner || typeof runner !== 'object') {
+    return;
+  }
+
+  instrumentTools((runner as ToolRunner).params?.tools, options);
+
+  // `addTools()` keeps its tools out of `params`.
+  wrapInPlace(runner, 'addTools', (addTools, thisArg, tools) => {
+    instrumentTools(tools, options);
+    return addTools.apply(thisArg, tools);
+  });
+  wrapInPlace(runner, 'setMessagesParams', (setMessagesParams, thisArg, args) => {
+    const result = setMessagesParams.apply(thisArg, args);
+    instrumentTools((runner as ToolRunner).params?.tools, options);
+    return result;
+  });
+}
+
+function instrumentTools(tools: unknown, options: AnthropicAiOptions): void {
+  if (!Array.isArray(tools)) {
+    return;
+  }
+
+  for (const tool of tools) {
+    // Toolsets are nameless entries whose `run` takes `(ctx, toolUse)`, so only named tools are wrapped.
+    if (tool && typeof tool === 'object' && typeof (tool as RunnableTool).name === 'string') {
+      wrapInPlace(tool, 'run', (run, thisArg, args) => runWithSpan(tool as RunnableTool, run, thisArg, args, options));
+    }
+  }
+}
+
+function wrapInPlace(
+  owner: object,
+  key: string,
+  call: (original: AnyFunction, thisArg: unknown, args: unknown[]) => unknown,
+): void {
+  const current = (owner as Record<string, unknown>)[key];
+  if (typeof current === 'function' && !getOriginalFunction(current as WrappedFunction)) {
+    fill(
+      owner,
+      key,
+      (original: AnyFunction) =>
+        function (this: unknown, ...args: unknown[]): unknown {
+          return call(original, this, args);
+        },
+    );
+  }
+}
+
+function runWithSpan(
+  tool: RunnableTool,
+  run: AnyFunction,
+  thisArg: unknown,
+  args: unknown[],
+  options: AnthropicAiOptions,
+): unknown {
+  const { recordInputs, recordOutputs } = resolveAIRecordingOptions(options);
+  const [input, context] = args as [unknown, RunContext];
+  // SDK versions before `toolUse` existed pass only `toolUseBlock`.
+  const callId = (context?.toolUse ?? context?.toolUseBlock)?.id;
+
+  return startSpan(
+    {
+      name: `execute_tool ${tool.name}`,
+      op: GEN_AI_EXECUTE_TOOL,
+      attributes: {
+        [SENTRY_ORIGIN]: ANTHROPIC_AI_ORIGIN,
+        [GEN_AI_OPERATION_NAME]: 'execute_tool',
+        [GEN_AI_TOOL_NAME]: tool.name,
+        ...(typeof tool.description === 'string' ? { [GEN_AI_TOOL_DESCRIPTION]: tool.description } : {}),
+        ...(typeof callId === 'string' ? { [GEN_AI_TOOL_CALL_ID_ATTRIBUTE]: callId } : {}),
+        ...(recordInputs ? { [GEN_AI_TOOL_CALL_ARGUMENTS]: stringify(input) } : {}),
+      },
+    },
+    span =>
+      handleCallbackErrors(
+        () => run.apply(thisArg, args),
+        error => handleToolError(span, error, context),
+        undefined,
+        result => recordOutputs && span.setAttribute(GEN_AI_TOOL_CALL_RESULT, stringify(result)),
+      ),
+  );
+}
+
+function handleToolError(span: Span, error: unknown, context: RunContext): void {
+  if (context?.signal?.aborted) {
+    span.setStatus({ code: SPAN_STATUS_ERROR, message: 'cancelled' });
+    return;
+  }
+
+  span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
+
+  // Only a runner passes `toolUse`/`toolUseBlock` (SDK 0.85+), and it hands the throw to the model as a result.
+  // A direct call rethrows to its caller, and a `ToolError` is sent on purpose, so neither is captured.
+  const fromRunner = !!(context?.toolUse ?? context?.toolUseBlock);
+  if (fromRunner && (error as Error | undefined)?.name !== 'ToolError') {
+    captureException(error, { mechanism: { handled: true, type: ANTHROPIC_AI_ORIGIN } });
+  }
+}