diff --git a/.changeset/fix-concurrent-tool-spans.md b/.changeset/fix-concurrent-tool-spans.md new file mode 100644 index 000000000..4f36a9014 --- /dev/null +++ b/.changeset/fix-concurrent-tool-spans.md @@ -0,0 +1,10 @@ +--- +"@voltagent/core": patch +--- + +fix(core): resolve race condition with concurrent tool spans + +Fixed a race condition where tools running in parallel would overwrite each other's +parentToolSpan in the shared systemContext. The fix passes parentToolSpan through +execution options instead of systemContext, ensuring each tool receives its unique +span. Backward compatibility is maintained by checking both options and systemContext. diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index f54e43aae..54e61082b 100644 --- a/packages/core/src/agent/agent.ts +++ b/packages/core/src/agent/agent.ts @@ -5053,7 +5053,8 @@ export class Agent { // Push execution metadata into systemContext for tools to consume oc.systemContext.set("agentId", this.id); oc.systemContext.set("historyEntryId", oc.operationId); - oc.systemContext.set("parentToolSpan", toolSpan); + + executionOptions.parentToolSpan = toolSpan; const hasOutputOverride = ( value: unknown, @@ -5534,7 +5535,10 @@ export class Agent { topK ?? toolRouting?.topK ?? embeddingTopK ?? DEFAULT_TOOL_SEARCH_TOP_K, ); const candidates = this.buildToolSearchCandidates(); - const parentToolSpan = oc.systemContext.get("parentToolSpan") as Span | undefined; + // Check both options and systemContext for backward compatibility, prefer options + const parentToolSpan = + ((options as any).parentToolSpan as Span | undefined) || + (oc.systemContext.get("parentToolSpan") as Span | undefined); const selectionSpanAttributes = { "tool.name": TOOL_ROUTING_SEARCH_TOOL_NAME, "tool.search.name": TOOL_ROUTING_SEARCH_TOOL_NAME, diff --git a/packages/core/src/agent/concurrent-tool-spans.spec.ts b/packages/core/src/agent/concurrent-tool-spans.spec.ts new file mode 100644 index 000000000..a8a20a09a --- /dev/null +++ b/packages/core/src/agent/concurrent-tool-spans.spec.ts @@ -0,0 +1,163 @@ +import type { Span } from "@opentelemetry/api"; +import * as ai from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { z } from "zod"; +import { NodeVoltAgentObservability } from "../observability"; +import { createTool } from "../tool"; +import { Agent } from "./agent"; +import { SubAgentManager } from "./subagent"; + +// Mock the AI SDK functions +vi.mock("ai", async () => { + const actual = await vi.importActual("ai"); + return { + ...actual, + generateText: vi.fn(), + streamText: vi.fn(), + generateObject: vi.fn(), + streamObject: vi.fn(), + stepCountIs: vi.fn(() => vi.fn(() => false)), + }; +}); + +describe("Agent Concurrent Tool Spans", () => { + let observability: NodeVoltAgentObservability; + let mockModel: MockLanguageModelV3; + + beforeEach(() => { + observability = new NodeVoltAgentObservability(); + mockModel = new MockLanguageModelV3(); + }); + + it("should provide unique parent spans to tools running in parallel", async () => { + const capturedSpans: Map = new Map(); + + // Create two tools that capture their parentToolSpan + const toolA = createTool({ + name: "toolA", + description: "Tool A", + parameters: z.object({}), + execute: async (_, options) => { + // Capture the span passed in options + const span = (options as any).parentToolSpan; + capturedSpans.set("toolA", span); + // Simulate some async work to allow interleaving + await new Promise((resolve) => setTimeout(resolve, 50)); + return "resultA"; + }, + }); + + const toolB = createTool({ + name: "toolB", + description: "Tool B", + parameters: z.object({}), + execute: async (_, options) => { + // Capture the span passed in options + const span = (options as any).parentToolSpan; + capturedSpans.set("toolB", span); + await new Promise((resolve) => setTimeout(resolve, 10)); + return "resultB"; + }, + }); + + const agent = new Agent({ + name: "test-agent", + instructions: "test instructions", + model: mockModel as any, + observability, + tools: [toolA, toolB], + toolRouting: false, // Disable tool routing to ensure tools are directly available + }); + + // Mock ai.generateText to simulate parallel tool execution + vi.mocked(ai.generateText).mockImplementation(async (options: any) => { + // options.tools contains the wrapped tools from Agent.ts + const toolAWrapper = options.tools.toolA; + const toolBWrapper = options.tools.toolB; + + // Parallel tools execution simulation + await Promise.all([toolAWrapper.execute({}), toolBWrapper.execute({})]); + + return { + text: "Done", + finishReason: "stop", + usage: { inputTokens: 10, outputTokens: 10, totalTokens: 20 }, + steps: [], + toolCalls: [], + toolResults: [], + response: { id: "test", modelId: "test", timestamp: new Date(), messages: [] }, + } as any; + }); + + await agent.generateText("Run tools in parallel"); + + expect(capturedSpans.has("toolA")).toBe(true); + expect(capturedSpans.has("toolB")).toBe(true); + + const spanA = capturedSpans.get("toolA"); + const spanB = capturedSpans.get("toolB"); + + expect(spanA).toBeDefined(); + expect(spanB).toBeDefined(); + + // Spans shouldn't be the same + expect(spanA).not.toBe(spanB); + expect((spanA as any).name).toContain("tool.execution:toolA"); + expect((spanB as any).name).toContain("tool.execution:toolB"); + }); + + it("should provide unique parent spans to subagents running in parallel via delegate_task", async () => { + // 1. Setup SubAgentManager + const subAgent = new Agent({ + name: "sub", + instructions: "sub", + model: mockModel as any, + }); + + // Mock generateText on subAgent to verify parentSpan + const generateTextSpy = vi.spyOn(subAgent, "generateText"); + generateTextSpy.mockResolvedValue({ text: "ok", usage: {} } as any); + + const manager = new SubAgentManager("parent", [{ agent: subAgent, method: "generateText" }]); + + // 2. Create delegate_task tool + const delegateTool = manager.createDelegateTool({ + sourceAgent: new Agent({ + name: "parent", + instructions: "parent", + model: mockModel as any, + }), + // Simulate creation-time parentToolSpan (stale/wrong one) + parentToolSpan: { spanContext: () => ({ traceId: "stale" }) } as any, + }); + + // 3. Simulate concurrent execution with distinct parent spans + const span1 = { spanContext: () => ({ traceId: "trace1" }) } as any; + const span2 = { spanContext: () => ({ traceId: "trace2" }) } as any; + + await Promise.all([ + delegateTool.execute?.({ task: "task1", targetAgents: ["sub"] }, { + parentToolSpan: span1, + } as any), + delegateTool.execute?.({ task: "task2", targetAgents: ["sub"] }, { + parentToolSpan: span2, + } as any), + ]); + + // 4. Verify subAgent.generateText was called with correct parent spans + expect(generateTextSpy).toHaveBeenCalledTimes(2); + + const calls = generateTextSpy.mock.calls; + // Extract parentSpan from options (second argument to generateText) + const spans = calls.map((c) => (c[1] as any).parentSpan); + + // Should contain the spans passed during execution + expect(spans).toContain(span1); + expect(spans).toContain(span2); + + // Should NOT contain the stale creation-time span + const staleSpanInCalls = spans.some((s) => s && s.spanContext().traceId === "stale"); + expect(staleSpanInCalls).toBe(false); + }); +}); diff --git a/packages/core/src/agent/subagent/index.ts b/packages/core/src/agent/subagent/index.ts index 343e19cad..2ed79d573 100644 --- a/packages/core/src/agent/subagent/index.ts +++ b/packages/core/src/agent/subagent/index.ts @@ -752,6 +752,7 @@ ${task}\n\nContext: ${safeStringify(contextObj, { indentation: 2 })}`; maxSteps?: number; conversationId?: string; userId?: string; + parentToolSpan?: Span; }): Tool { const { sourceAgent, @@ -760,6 +761,7 @@ ${task}\n\nContext: ${safeStringify(contextObj, { indentation: 2 })}`; maxSteps, conversationId, userId, + parentToolSpan, } = options; return createTool({ id: "delegate_task", @@ -773,10 +775,10 @@ ${task}\n\nContext: ${safeStringify(contextObj, { indentation: 2 })}`; .optional() .describe("Additional context for the task"), }), - execute: async ({ task, targetAgents, context = {} }, options) => { + execute: async ({ task, targetAgents, context = {} }, executeOptions) => { // Extract OperationContext from options if available // Since ToolExecuteOptions extends Partial, we can cast it - const currentOperationContext = options as OperationContext | undefined; + const currentOperationContext = executeOptions as OperationContext | undefined; // Fall back to the original operation context if not available const effectiveOperationContext = currentOperationContext || operationContext; @@ -838,9 +840,10 @@ ${task}\n\nContext: ${safeStringify(contextObj, { indentation: 2 })}`; // Pass maxSteps from parent to subagents maxSteps, // Pass the parentToolSpan from executeOptions for proper span hierarchy - parentSpan: effectiveOperationContext?.systemContext?.get("parentToolSpan") as - | Span - | undefined, + parentSpan: + (executeOptions?.parentToolSpan as Span | undefined) || + parentToolSpan || + (effectiveOperationContext?.systemContext?.get("parentToolSpan") as Span | undefined), }); // Return structured results with agent names and their responses diff --git a/packages/core/src/planagent/plan-agent.ts b/packages/core/src/planagent/plan-agent.ts index 2e9bedd3f..4639eb326 100644 --- a/packages/core/src/planagent/plan-agent.ts +++ b/packages/core/src/planagent/plan-agent.ts @@ -745,7 +745,9 @@ function createTaskToolkit(options: { }), execute: async (input, executeOptions) => { const operationContext = executeOptions as OperationContext; - const toolSpan = operationContext.systemContext.get("parentToolSpan") as Span | undefined; + const toolSpan = + ((executeOptions as any).parentToolSpan as Span | undefined) || + (operationContext.systemContext.get("parentToolSpan") as Span | undefined); if (toolSpan) { toolSpan.setAttribute("voltagent.label", `task:${input.subagent_type}`); toolSpan.setAttribute("planagent.task.subagent_type", input.subagent_type); @@ -768,7 +770,7 @@ function createTaskToolkit(options: { conversationId: operationContext.conversationId, parentOperationContext: operationContext, maxSteps: taskOptions?.maxSteps, - parentSpan: operationContext.systemContext.get("parentToolSpan") as Span | undefined, + parentSpan: toolSpan, }); if (toolSpan) { diff --git a/packages/core/src/planagent/planning/index.ts b/packages/core/src/planagent/planning/index.ts index 93661380f..7c46f57a7 100644 --- a/packages/core/src/planagent/planning/index.ts +++ b/packages/core/src/planagent/planning/index.ts @@ -148,7 +148,9 @@ export function createPlanningToolkit(agent: Agent, options: PlanningToolkitOpti await resolvedBackend.setTodos(guardedTodos); operationContext.systemContext.set(PLAN_PROGRESS_CONTEXT_KEY, false); - const toolSpan = operationContext.systemContext.get("parentToolSpan") as Span | undefined; + const toolSpan = + ((executeOptions as any).parentToolSpan as Span | undefined) || + (operationContext.systemContext.get("parentToolSpan") as Span | undefined); if (toolSpan) { const pendingCount = guardedTodos.filter((todo) => todo.status === "pending").length; const inProgressCount = guardedTodos.filter((todo) => todo.status === "in_progress").length; diff --git a/packages/core/src/retriever/tools/index.ts b/packages/core/src/retriever/tools/index.ts index c4e0a6633..97558c1b6 100644 --- a/packages/core/src/retriever/tools/index.ts +++ b/packages/core/src/retriever/tools/index.ts @@ -51,9 +51,13 @@ export const createRetrieverTool = ( execute: async ({ query }, options?: ToolExecuteOptions) => { // Pass complete options to retriever for access to userId, conversationId, etc. const startTime = Date.now(); - const toolSpan = options?.systemContext?.get("parentToolSpan") as - | { setAttribute?: (key: string, value: unknown) => void } - | undefined; + const toolSpan = + ((options as any)?.parentToolSpan as + | { setAttribute?: (key: string, value: unknown) => void } + | undefined) || + (options?.systemContext?.get("parentToolSpan") as + | { setAttribute?: (key: string, value: unknown) => void } + | undefined); const normalizeAttributeValue = (value: unknown) => { if (value === null || value === undefined) return null;