Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
import Anthropic from '@anthropic-ai/sdk';
import { SessionToolRunner } from '@anthropic-ai/sdk/resources/beta/sessions/events';
import * as Sentry from '@sentry/node';
import express from 'express';

const STREAM_EVENTS = {
sesn_weather: [
{ type: 'agent.tool_use', id: 'sevt_weather', name: 'weather', input: { city: 'Vienna' } },
{ type: 'agent.tool_use', id: 'sevt_explode', name: 'explode', input: {} },
{ type: 'session.status_terminated', id: 'sevt_weather_end' },
],
// Served by a runner built with `new SessionToolRunner()`, as `EnvironmentWorker` does.
sesn_direct: [
{ type: 'agent.tool_use', id: 'sevt_read', name: 'read', input: { file_path: 'README.md' } },
{ type: 'session.status_terminated', id: 'sevt_direct_end' },
],
};

function startMockSessionsServer() {
const app = express();
app.use(express.json());

app.get('/anthropic/v1/sessions/:sessionId/events/stream', (req, res) => {
res.writeHead(200, { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', Connection: 'keep-alive' });

const events = STREAM_EVENTS[req.params.sessionId];
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);
});
});
app.get('/anthropic/v1/sessions/:sessionId/events', (_req, res) => {
res.json({ data: [], next_page: null });
});
app.post('/anthropic/v1/sessions/:sessionId/events', (_req, res) => {
res.json({ data: [] });
});

return new Promise(resolve => {
const server = app.listen(0, () => {
resolve(server);
});
});
}

async function drain(runner) {
for await (const _call of runner) {
// drain
}
}

async function run() {
const server = await startMockSessionsServer();

await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
const client = new Anthropic({
apiKey: 'mock-api-key',
baseURL: `http://localhost:${server.address().port}/anthropic`,
});

await drain(
client.beta.sessions.events.toolRunner('sesn_weather', {
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');
},
},
],
}),
);

await drain(
new SessionToolRunner('sesn_direct', {
client,
tools: [
{
name: 'read',
description: 'Read a file',
input_schema: { type: 'object', properties: { file_path: { type: 'string' } } },
run: () => '# Weather app',
},
],
}),
);
});

await Sentry.flush(2000);

server.close();
}

run();
Original file line number Diff line number Diff line change
Expand Up @@ -137,4 +137,67 @@ describe('Anthropic integration (0.129)', () => {
},
},
);

createEsmAndCjsTests(
__dirname,
'scenario-session-tool-runner.mjs',
'instrument.mjs',
(createRunner, test) => {
test('emits an execute_tool span for each tool a session runner runs, however it was created', 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 toolSpans = receivedSpans.filter(span => span.attributes[SENTRY_OP]?.value === GEN_AI_EXECUTE_TOOL);
expect(toolSpans).toHaveLength(3);
const weather = toolSpans.find(span => span.attributes[GEN_AI_TOOL_NAME]?.value === 'weather')!;
const explode = toolSpans.find(span => span.attributes[GEN_AI_TOOL_NAME]?.value === 'explode')!;
const read = toolSpans.find(span => span.attributes[GEN_AI_TOOL_NAME]?.value === 'read')!;

expect(weather).toMatchObject({
name: 'execute_tool weather',
parent_span_id: main.span_id,
status: 'ok',
attributes: expect.objectContaining({
[GEN_AI_TOOL_CALL_ID_ATTRIBUTE]: { type: 'string', value: 'sevt_weather' },
}),
});
expect(explode).toMatchObject({ name: 'execute_tool explode', parent_span_id: main.span_id, status: 'error' });
expect(read).toMatchObject({
name: 'execute_tool read',
parent_span_id: main.span_id,
status: 'ok',
attributes: expect.objectContaining({
[GEN_AI_TOOL_CALL_ID_ATTRIBUTE]: { type: 'string', value: 'sevt_read' },
}),
});
expect(receivedEvent).toMatchObject({
exception: {
values: [{ value: 'Tool exploded', mechanism: { type: 'auto.ai.anthropic', handled: true } }],
},
});
});
},
{
additionalDependencies: {
'@anthropic-ai/sdk': '0.129.0',
},
},
);
});
11 changes: 7 additions & 4 deletions packages/server-utils/src/ai/anthropic-ai/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import {
} from '../core/utils';
import { ANTHROPIC_AI_ORIGIN, ANTHROPIC_METHOD_REGISTRY } from './constants';
import { instrumentAsyncIterableStream, instrumentMessageStream } from './streaming';
import { instrumentToolRunnerMethod } from './tool-runner';
import { instrumentSessionToolRunner, instrumentToolRunner, instrumentToolRunnerFactory } from './tool-runner';
import type { AnthropicAiOptions, AnthropicAiResponse, AnthropicAiStreamingEvent, ContentBlock } from './types';
import { handleResponseError, messagesFromParams, setMessagesAttribute } from './utils';

Expand Down Expand Up @@ -351,9 +351,12 @@ 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);
const beta = (client as { beta?: { messages?: object; sessions?: { events?: object } } }).beta;
if (beta?.messages) {
instrumentToolRunnerFactory(beta.messages, instrumentToolRunner, options);
}
if (beta?.sessions?.events) {
instrumentToolRunnerFactory(beta.sessions.events, instrumentSessionToolRunner, options);
}

return client;
Expand Down
19 changes: 15 additions & 4 deletions packages/server-utils/src/ai/anthropic-ai/tool-runner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,17 @@ 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 } };
type SessionToolRunner = { 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) => {
/** Wrap the `toolRunner()` method of `owner` so the runners it returns trace their tool calls. */
export function instrumentToolRunnerFactory(
owner: object,
instrumentRunner: (runner: unknown, options: AnthropicAiOptions) => void,
options: AnthropicAiOptions,
): void {
wrapInPlace(owner, 'toolRunner', (toolRunner, thisArg, args) => {
const runner = toolRunner.apply(thisArg, args);
instrumentToolRunner(runner, options);
instrumentRunner(runner, options);
return runner;
});
}
Expand All @@ -56,6 +61,12 @@ export function instrumentToolRunner(runner: unknown, options: AnthropicAiOption
});
}

/** Give each tool call of a Managed Agents session runner a `gen_ai.execute_tool` span. */
export function instrumentSessionToolRunner(runner: unknown, options: AnthropicAiOptions): void {
// The session runner has no `addTools()`: it only runs the tools it was created with.
instrumentTools((runner as SessionToolRunner | undefined)?.tools, options);
}

function instrumentTools(tools: unknown, options: AnthropicAiOptions): void {
if (!Array.isArray(tools)) {
return;
Expand Down
7 changes: 6 additions & 1 deletion packages/server-utils/src/integrations/anthropic.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ import { wrapApiPromiseResponse } from '../ai/core/apiPromise';
import { addPrivateRequestAttributes, addResponseAttributes, extractRequestAttributes } from '../ai/anthropic-ai';
import { ANTHROPIC_AI_ORIGIN } from '../ai/anthropic-ai/constants';
import { instrumentAsyncIterableStream, instrumentMessageStream } from '../ai/anthropic-ai/streaming';
import { instrumentToolRunner } from '../ai/anthropic-ai/tool-runner';
import { instrumentSessionToolRunner, instrumentToolRunner } from '../ai/anthropic-ai/tool-runner';
import type { AnthropicAiOptions, AnthropicAiResponse } from '../ai/anthropic-ai/types';
import type { OrchestrionChannelContext } from '../orchestrion/types';
import { CHANNELS } from '../orchestrion/channels';
Expand Down Expand Up @@ -80,6 +80,11 @@ function instrumentAnthropic(options: AnthropicAiOptions): void {
.end.subscribe(message => {
instrumentToolRunner((message as OrchestrionChannelContext).result, options);
});
diagnosticsChannel
.tracingChannel<OrchestrionChannelContext>(CHANNELS.ANTHROPIC_SESSION_TOOL_RUNNER)
.end.subscribe(message => {
instrumentSessionToolRunner((message as OrchestrionChannelContext).self, options);
});
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@ export const anthropicAiConfig = [
module: { name: '@anthropic-ai/sdk', versionRange: '>=0.63.0 <1', filePath },
functionQuery: { className: 'Messages', methodName: 'toolRunner', kind: 'Sync' as const },
})),
// The constructor, not `sessions.events.toolRunner()`, so runners built directly (e.g. by `EnvironmentWorker`) are covered.
...['lib/tools/SessionToolRunner.js', 'lib/tools/SessionToolRunner.mjs'].map(filePath => ({
channelName: 'session-tool-runner',
module: { name: '@anthropic-ai/sdk', versionRange: '>=0.97.0 <1', filePath },
functionQuery: { className: 'SessionToolRunner' },
})),
] satisfies InstrumentationConfig[];

export const anthropicAiModuleNames = getModuleNames(anthropicAiConfig);
Expand All @@ -39,4 +45,5 @@ export const anthropicAiChannels = {
ANTHROPIC_CHAT: 'orchestrion:@anthropic-ai/sdk:chat',
ANTHROPIC_MESSAGES_STREAM: 'orchestrion:@anthropic-ai/sdk:messages-stream',
ANTHROPIC_TOOL_RUNNER: 'orchestrion:@anthropic-ai/sdk:tool-runner',
ANTHROPIC_SESSION_TOOL_RUNNER: 'orchestrion:@anthropic-ai/sdk:session-tool-runner',
} as const;
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,31 @@ class FakeRunner {
}
}

// Mirrors the SDK session runner: it keeps `opts.tools` as `tools` and passes the call event as both `toolUse` and `toolUseBlock`.
class FakeSessionRunner {
public readonly tools: Tool[];
public constructor(opts: { tools: Tool[] }) {
this.tools = opts.tools;
}
public async dispatch(name: string, toolUse: { id: string; input: unknown }): Promise<{ isError: boolean }> {
const tool = this.tools.find(t => t.name === name)!;
try {
await tool.run(toolUse.input, { toolUse, toolUseBlock: toolUse, signal: new AbortController().signal });
return { isError: false };
} catch {
return { isError: true };
}
}
}

function fakeClient() {
return {
beta: {
sessions: {
events: {
toolRunner: vi.fn((_sessionId: string, opts: { tools: Tool[] }) => new FakeSessionRunner(opts)),
},
},
messages: {
create: vi.fn(),
toolRunner: vi.fn((params: Params) => {
Expand Down Expand Up @@ -263,4 +285,50 @@ describe('instrumentAnthropicAiClient tool runner', () => {

expect(() => instrumented.beta.messages.toolRunner(undefined as unknown as Params)).toThrow('params required');
});

describe('sessions.events.toolRunner()', () => {
it('wraps each tool call of the session runner in an execute_tool span', async () => {
const weather = tool('weather', async () => 'sunny');

const runner = instrumentAnthropicAiClient(fakeClient()).beta.sessions.events.toolRunner('sesn_1', {
tools: [weather],
});
const parent = await startSpan({ name: 'agent' }, async span => {
await runner.dispatch('weather', { id: 'sevt_1', input: { city: 'Vienna' } });
return spanToStaticSpanJSON(span);
});

expect(toolSpans()).toEqual([
expect.objectContaining({
description: 'execute_tool weather',
parent_span_id: parent.span_id,
origin: 'auto.ai.anthropic',
status: 'ok',
data: expect.objectContaining({
[GEN_AI_TOOL_NAME]: 'weather',
[GEN_AI_TOOL_CALL_ID_ATTRIBUTE]: 'sevt_1',
}),
}),
]);
});

it('captures the error of a failing tool, which the session runner turns into a result', async () => {
const explode = tool('explode', async () => {
throw new Error('disk full');
});

const runner = instrumentAnthropicAiClient(fakeClient()).beta.sessions.events.toolRunner('sesn_1', {
tools: [explode],
});
const outcome = await runner.dispatch('explode', { id: 'sevt_2', input: {} });

expect(outcome).toEqual({ isError: true });
expect(toolSpans()[0]!.status).toBe('internal_error');
await vi.waitFor(() => expect(events).toHaveLength(1));
expect(events[0]!.exception?.values?.[0]).toMatchObject({
value: 'disk full',
mechanism: { type: 'auto.ai.anthropic', handled: true },
});
});
});
});
Loading