Skip to content

Commit ae36796

Browse files
authored
tests(node): Add node integration tests for Vercel ToolLoopAgent (#20087)
The plan was to add support for this, but turns out we already do! Still adding some node integration tests to cover this use case explicitly. Closes #20057
1 parent 7886e12 commit ae36796

File tree

2 files changed

+138
-0
lines changed

2 files changed

+138
-0
lines changed
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import * as Sentry from '@sentry/node';
2+
import { ToolLoopAgent, stepCountIs, tool } from 'ai';
3+
import { MockLanguageModelV3 } from 'ai/test';
4+
import { z } from 'zod';
5+
6+
async function run() {
7+
await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
8+
let callCount = 0;
9+
10+
const agent = new ToolLoopAgent({
11+
experimental_telemetry: { isEnabled: true },
12+
model: new MockLanguageModelV3({
13+
doGenerate: async () => {
14+
if (callCount++ === 0) {
15+
return {
16+
finishReason: { unified: 'tool-calls', raw: 'tool_calls' },
17+
usage: {
18+
inputTokens: { total: 10, noCache: 10, cached: 0 },
19+
outputTokens: { total: 20, noCache: 20, cached: 0 },
20+
totalTokens: { total: 30, noCache: 30, cached: 0 },
21+
},
22+
content: [
23+
{
24+
type: 'tool-call',
25+
toolCallId: 'call-1',
26+
toolName: 'getWeather',
27+
input: JSON.stringify({ location: 'San Francisco' }),
28+
},
29+
],
30+
warnings: [],
31+
};
32+
}
33+
return {
34+
finishReason: { unified: 'stop', raw: 'stop' },
35+
usage: {
36+
inputTokens: { total: 15, noCache: 15, cached: 0 },
37+
outputTokens: { total: 25, noCache: 25, cached: 0 },
38+
totalTokens: { total: 40, noCache: 40, cached: 0 },
39+
},
40+
content: [{ type: 'text', text: 'The weather in San Francisco is sunny, 72°F.' }],
41+
warnings: [],
42+
};
43+
},
44+
}),
45+
tools: {
46+
getWeather: tool({
47+
description: 'Get the current weather for a location',
48+
inputSchema: z.object({ location: z.string() }),
49+
execute: async ({ location }) => `Weather in ${location}: Sunny, 72°F`,
50+
}),
51+
},
52+
stopWhen: stepCountIs(3),
53+
});
54+
55+
await agent.generate({
56+
prompt: 'What is the weather in San Francisco?',
57+
});
58+
});
59+
}
60+
61+
run();

dev-packages/node-integration-tests/suites/tracing/vercelai/v6/test.ts

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -601,4 +601,81 @@ describe('Vercel AI integration (V6)', () => {
601601
},
602602
},
603603
);
604+
605+
createEsmAndCjsTests(
606+
__dirname,
607+
'scenario-tool-loop-agent.mjs',
608+
'instrument.mjs',
609+
(createRunner, test) => {
610+
test('creates spans for ToolLoopAgent with tool calls', async () => {
611+
const expectedTransaction = {
612+
transaction: 'main',
613+
spans: expect.arrayContaining([
614+
// ToolLoopAgent outer span
615+
expect.objectContaining({
616+
data: expect.objectContaining({
617+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'invoke_agent',
618+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.invoke_agent',
619+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
620+
[GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'mock-model-id',
621+
}),
622+
op: 'gen_ai.invoke_agent',
623+
origin: 'auto.vercelai.otel',
624+
status: 'ok',
625+
}),
626+
// First doGenerate span (returns tool-calls)
627+
expect.objectContaining({
628+
data: expect.objectContaining({
629+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'generate_content',
630+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.generate_content',
631+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
632+
[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 10,
633+
[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]: 20,
634+
[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]: ['tool-calls'],
635+
}),
636+
op: 'gen_ai.generate_content',
637+
origin: 'auto.vercelai.otel',
638+
status: 'ok',
639+
}),
640+
// Tool execution span
641+
expect.objectContaining({
642+
data: expect.objectContaining({
643+
[GEN_AI_TOOL_CALL_ID_ATTRIBUTE]: 'call-1',
644+
[GEN_AI_TOOL_NAME_ATTRIBUTE]: 'getWeather',
645+
[GEN_AI_TOOL_TYPE_ATTRIBUTE]: 'function',
646+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'execute_tool',
647+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.execute_tool',
648+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
649+
}),
650+
description: 'execute_tool getWeather',
651+
op: 'gen_ai.execute_tool',
652+
origin: 'auto.vercelai.otel',
653+
status: 'ok',
654+
}),
655+
// Second doGenerate span (returns final text)
656+
expect.objectContaining({
657+
data: expect.objectContaining({
658+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'generate_content',
659+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.generate_content',
660+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
661+
[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 15,
662+
[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]: 25,
663+
[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]: ['stop'],
664+
}),
665+
op: 'gen_ai.generate_content',
666+
origin: 'auto.vercelai.otel',
667+
status: 'ok',
668+
}),
669+
]),
670+
};
671+
672+
await createRunner().expect({ transaction: expectedTransaction }).start().completed();
673+
});
674+
},
675+
{
676+
additionalDependencies: {
677+
ai: '^6.0.0',
678+
},
679+
},
680+
);
604681
});

0 commit comments

Comments
 (0)