From f864118f03d8003b424f8d48a425394fc3c4fc38 Mon Sep 17 00:00:00 2001 From: Amaad Martin Date: Wed, 24 Jun 2026 14:20:17 -0700 Subject: [PATCH 1/2] Feat: Implement tracing for tool execution --- core/src/agents/functions.ts | 309 +++++++++++------------------ core/test/agents/functions_test.ts | 88 ++++++++ package-lock.json | 1 - 3 files changed, 209 insertions(+), 189 deletions(-) diff --git a/core/src/agents/functions.ts b/core/src/agents/functions.ts index 2bba788d..6988ea66 100644 --- a/core/src/agents/functions.ts +++ b/core/src/agents/functions.ts @@ -5,6 +5,7 @@ */ import {Content, createUserContent, FunctionCall, Part} from '@google/genai'; +import {SpanStatusCode} from '@opentelemetry/api'; import {isEmpty} from 'lodash-es'; import {InvocationContext} from '../agents/invocation_context.js'; @@ -211,69 +212,6 @@ export function generateRequestConfirmationEvent({ }); } -async function callToolAsync( - tool: BaseTool, - args: Record, // eslint-disable-line @typescript-eslint/no-explicit-any - toolContext: Context, - // eslint-disable-next-line @typescript-eslint/no-explicit-any -): Promise { - return tracer.startActiveSpan(`execute_tool ${tool.name}`, async (span) => { - try { - logger.debug(`callToolAsync ${tool.name}`); - const result = await tool.runAsync({args, toolContext}); - traceToolCall({ - tool, - args, - functionResponseEvent: buildResponseEvent( - tool, - result, - toolContext, - toolContext.invocationContext, - ), - }); - return result; - } finally { - span.end(); - } - }); -} - -function buildResponseEvent( - tool: BaseTool, - functionResult: unknown, - toolContext: Context, - invocationContext: InvocationContext, -): Event { - let responseResult: Record; - if (typeof functionResult !== 'object' || functionResult == null) { - responseResult = {result: functionResult}; - } else if (Array.isArray(functionResult)) { - responseResult = {results: functionResult}; - } else { - responseResult = functionResult as Record; - } - - const partFunctionResponse: Part = { - functionResponse: { - name: tool.name, - response: responseResult, - id: toolContext.functionCallId, - }, - }; - - const content: Content = { - role: 'user', - parts: [partFunctionResponse], - }; - - return createEvent({ - invocationId: invocationContext.invocationId, - author: invocationContext.agent.name, - content: content, - actions: toolContext.actions, - branch: invocationContext.branch, - }); -} /** * Handles function calls. * Runtime behavior to pay attention to: @@ -357,142 +295,142 @@ export async function handleFunctionCallList({ toolConfirmation, }); - // TODO - b/436079721: implement [tracer.start_as_current_span] - logger.debug(`execute_tool ${tool.name}`); const functionArgs = functionCall.args ?? {}; - // Step 1: Check if plugin before_tool_callback overrides the function - // response. - let functionResponse = null; - let functionResponseError: string | unknown | undefined; - functionResponse = - await invocationContext.pluginManager.runBeforeToolCallback({ - tool: tool, - toolArgs: functionArgs, - toolContext: toolContext, - }); - - // Step 2: If no overrides are provided from the plugins, further run the - // canonical callback. - // TODO - b/425992518: validate the callback response type matches. - if (functionResponse == null) { - // Cover both null and undefined - for (const callback of beforeToolCallbacks) { - functionResponse = await callback({ - tool: tool, - args: functionArgs, - context: toolContext, - }); - if (functionResponse) { - break; - } - } - } - - // Step 3: Otherwise, proceed calling the tool normally. - if (functionResponse == null) { - // Cover both null and undefined + await tracer.startActiveSpan(`execute_tool ${tool.name}`, async (span) => { try { - functionResponse = await callToolAsync(tool, functionArgs, toolContext); - } catch (e: unknown) { - if (e instanceof Error) { - const onToolErrorResponse = - await invocationContext.pluginManager.runOnToolErrorCallback({ + // Step 1: Check if plugin before_tool_callback overrides the function + // response. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let functionResponse: any = + await invocationContext.pluginManager.runBeforeToolCallback({ + tool, + toolArgs: functionArgs, + toolContext, + }); + let functionResponseError: string | unknown | undefined; + + // Step 2: If no overrides are provided from the plugins, further run the + // canonical callback. + // TODO - b/425992518: validate the callback response type matches. + if (functionResponse == null) { + // Cover both null and undefined + for (const callback of beforeToolCallbacks) { + functionResponse = await callback({ tool: tool, - toolArgs: functionArgs, - toolContext: toolContext, - error: e, + args: functionArgs, + context: toolContext, }); + if (functionResponse) { + break; + } + } + } - // Set function response to the result of the error callback and - // continue execution, do not shortcut - if (onToolErrorResponse) { - functionResponse = onToolErrorResponse; - } else { - // If the error callback returns undefined, use the error message - // as the function response error. - functionResponseError = e.message; + // Step 3: Otherwise, proceed calling the tool normally. + if (functionResponse == null) { + // Cover both null and undefined + try { + logger.debug(`callToolAsync ${tool.name}`); + functionResponse = await tool.runAsync({ + args: functionArgs, + toolContext, + }); + } catch (e: unknown) { + const err = e instanceof Error ? e : new Error(String(e)); + span.recordException(err); + if (e instanceof Error) { + functionResponse = + await invocationContext.pluginManager.runOnToolErrorCallback({ + tool, + toolArgs: functionArgs, + toolContext, + error: e, + }); + } + if (!functionResponse) { + functionResponseError = e instanceof Error ? e.message : e; + span.setStatus({ + code: SpanStatusCode.ERROR, + message: err.message, + }); + } } - } else { - // If the error is not an Error, use the error object as the function - // response error. - functionResponseError = e; } - } - } - // Step 4: Check if plugin after_tool_callback overrides the function - // response. - let alteredFunctionResponse = - await invocationContext.pluginManager.runAfterToolCallback({ - tool: tool, - toolArgs: functionArgs, - toolContext: toolContext, - result: functionResponse, - }); - - // Step 5: If no overrides are provided from the plugins, further run the - // canonical after_tool_callbacks. - if (alteredFunctionResponse == null) { - // Cover both null and undefined - for (const callback of afterToolCallbacks) { - alteredFunctionResponse = await callback({ - tool: tool, - args: functionArgs, - context: toolContext, - response: functionResponse, - }); - if (alteredFunctionResponse) { - break; + // Step 4: Check if plugin after_tool_callback overrides the function + // response. + let alteredFunctionResponse = + await invocationContext.pluginManager.runAfterToolCallback({ + tool: tool, + toolArgs: functionArgs, + toolContext: toolContext, + result: functionResponse, + }); + + // Step 5: If no overrides are provided from the plugins, further run the + // canonical after_tool_callbacks. + if (alteredFunctionResponse == null) { + // Cover both null and undefined + for (const callback of afterToolCallbacks) { + alteredFunctionResponse = await callback({ + tool: tool, + args: functionArgs, + context: toolContext, + response: functionResponse, + }); + if (alteredFunctionResponse) { + break; + } + } } - } - } - // Step 6: If alternative response exists from after_tool_callback, use it - // instead of the original function response. - if (alteredFunctionResponse != null) { - functionResponse = alteredFunctionResponse; - } + // Step 6: If alternative response exists from after_tool_callback, use it + // instead of the original function response. + functionResponse = alteredFunctionResponse ?? functionResponse; - // TODO - b/425992518: state event polluting runtime, consider fix. - // Allow long running function to return None as response. - if (tool.isLongRunning && !functionResponse) { - continue; - } + // TODO - b/425992518: state event polluting runtime, consider fix. + // Allow long running function to return None as response. + if (tool.isLongRunning && !functionResponse) { + return; + } - if (functionResponseError) { - functionResponse = {error: functionResponseError}; - } else if ( - typeof functionResponse !== 'object' || - functionResponse == null - ) { - functionResponse = {result: functionResponse}; - } else if (Array.isArray(functionResponse)) { - functionResponse = {results: functionResponse}; - } + if (functionResponseError) { + functionResponse = {error: functionResponseError}; + } else if ( + typeof functionResponse !== 'object' || + functionResponse == null + ) { + functionResponse = {result: functionResponse}; + } else if (Array.isArray(functionResponse)) { + functionResponse = {results: functionResponse}; + } - // Builds the function response event. - const functionResponseEvent = createEvent({ - invocationId: invocationContext.invocationId, - author: invocationContext.agent.name, - content: createUserContent({ - functionResponse: { - id: toolContext.functionCallId, - name: tool.name, - response: functionResponse, - }, - }), - actions: toolContext.actions, - branch: invocationContext.branch, - }); + // Builds the function response event. + const functionResponseEvent = createEvent({ + invocationId: invocationContext.invocationId, + author: invocationContext.agent.name, + content: createUserContent({ + functionResponse: { + id: toolContext.functionCallId, + name: tool.name, + response: functionResponse, + }, + }), + actions: toolContext.actions, + branch: invocationContext.branch, + }); - // TODO - b/436079721: implement [traceToolCall] - logger.debug('traceToolCall', { - tool: tool.name, - args: functionArgs, - functionResponseEvent: functionResponseEvent.id, + traceToolCall({ + tool, + args: functionArgs, + functionResponseEvent, + }); + functionResponseEvents.push(functionResponseEvent); + } finally { + span.end(); + } }); - functionResponseEvents.push(functionResponseEvent); } if (!functionResponseEvents.length) { @@ -506,11 +444,6 @@ export async function handleFunctionCallList({ tracer.startActiveSpan('execute_tool (merged)', (span) => { try { logger.debug('execute_tool (merged)'); - // TODO - b/436079721: implement [traceMergedToolCalls] - logger.debug('traceMergedToolCalls', { - responseEventId: mergedEvent.id, - functionResponseEvent: mergedEvent.id, - }); traceMergedToolCalls({ responseEventId: mergedEvent.id, functionResponseEvent: mergedEvent, diff --git a/core/test/agents/functions_test.ts b/core/test/agents/functions_test.ts index c5b3b79e..7edfa641 100644 --- a/core/test/agents/functions_test.ts +++ b/core/test/agents/functions_test.ts @@ -20,6 +20,7 @@ import { ToolConfirmation, } from '@google/adk'; import {Content, FunctionCall} from '@google/genai'; +import {SpanStatusCode} from '@opentelemetry/api'; import {beforeEach, describe, expect, it, vi} from 'vitest'; import {z} from 'zod'; import { @@ -29,6 +30,27 @@ import { populateClientFunctionCallId, removeClientFunctionCallId, } from '../../src/agents/functions.js'; +import { + traceMergedToolCalls, + tracer, + traceToolCall, +} from '../../src/telemetry/tracing.js'; + +const mockSpan = { + end: vi.fn(), + recordException: vi.fn(), + setStatus: vi.fn(), +}; + +vi.mock('../../src/telemetry/tracing.js', () => { + return { + tracer: { + startActiveSpan: vi.fn().mockImplementation((name, cb) => cb(mockSpan)), + }, + traceToolCall: vi.fn(), + traceMergedToolCalls: vi.fn(), + }; +}); // Get the test target function const { @@ -115,6 +137,8 @@ describe('handleFunctionCallList', () => { args: {}, }; toolsDict = {'testTool': testTool}; + + vi.clearAllMocks(); }); it('should execute tool with no callbacks or plugins', async () => { @@ -363,6 +387,70 @@ describe('handleFunctionCallList', () => { }), ); }); + + it('should trace tool execution', async () => { + const event = await handleFunctionCallList({ + invocationContext, + functionCalls: [functionCall], + toolsDict, + beforeToolCallbacks: [], + afterToolCallbacks: [], + }); + + expect(tracer.startActiveSpan).toHaveBeenCalledWith( + `execute_tool ${testTool.name}`, + expect.any(Function), + ); + expect(traceToolCall).toHaveBeenCalledWith( + expect.objectContaining({ + tool: testTool, + args: functionCall.args, + functionResponseEvent: event, + }), + ); + }); + + it('should record exception on span when tool throws', async () => { + const errorFunctionCall: FunctionCall = { + id: randomIdForTestingOnly(), + name: 'errorTool', + args: {}, + }; + + await handleFunctionCallList({ + invocationContext, + functionCalls: [errorFunctionCall], + toolsDict: {'errorTool': errorTool}, + beforeToolCallbacks: [], + afterToolCallbacks: [], + }); + + expect(mockSpan.recordException).toHaveBeenCalledWith(expect.any(Error)); + expect(mockSpan.setStatus).toHaveBeenCalledWith({ + code: SpanStatusCode.ERROR, + message: expect.stringContaining('tool error message content'), + }); + }); + + it('should trace merged tool calls when multiple tools executed', async () => { + const functionCall1 = {id: '1', name: 'testTool', args: {}}; + const functionCall2 = {id: '2', name: 'testTool', args: {}}; + + const event = await handleFunctionCallList({ + invocationContext, + functionCalls: [functionCall1, functionCall2], + toolsDict, + beforeToolCallbacks: [], + afterToolCallbacks: [], + }); + + expect(traceMergedToolCalls).toHaveBeenCalledWith( + expect.objectContaining({ + responseEventId: event!.id, + functionResponseEvent: event, + }), + ); + }); }); describe('generateAuthEvent', () => { diff --git a/package-lock.json b/package-lock.json index 7fe962c0..ab5736cf 100644 --- a/package-lock.json +++ b/package-lock.json @@ -5171,7 +5171,6 @@ "version": "0.5.17", "resolved": "https://registry.npmjs.org/adm-zip/-/adm-zip-0.5.17.tgz", "integrity": "sha512-+Ut8d9LLqwEvHHJl1+PIHqoyDxFgVN847JTVM3Izi3xHDWPE4UtzzXysMZQs64DMcrJfBeS/uoEP4AD3HQHnQQ==", - "dev": true, "license": "MIT", "engines": { "node": ">=12.0" From d4acaba596a35f6ce51dbc68ea34328243a8761b Mon Sep 17 00:00:00 2001 From: Amaad Martin Date: Thu, 25 Jun 2026 10:57:54 -0700 Subject: [PATCH 2/2] Test: Use 127.0.0.1 for test API server to fix Windows CI --- tests/integration/test_api_server.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_api_server.ts b/tests/integration/test_api_server.ts index 8b6337d4..62088dfc 100644 --- a/tests/integration/test_api_server.ts +++ b/tests/integration/test_api_server.ts @@ -32,7 +32,7 @@ export class AdkTsApiServer extends BaseTestServer { private params: TestApiServerParams; constructor(params: TestApiServerParams) { - super('localhost', params.port || 0); + super('127.0.0.1', params.port || 0); this.params = params; } @@ -64,6 +64,8 @@ export class AdkTsApiServer extends BaseTestServer { cliPath, params.serveDebugUI ? 'web' : 'api_server', params.agentsDir, + '--host', + this.host, '--port', this.port.toString(), '--allow_origins',