From 0d3a114cc4bf68f08e5479932f1e818696ce4bd7 Mon Sep 17 00:00:00 2001 From: Indra Zulfi Date: Sun, 5 Jul 2026 00:16:01 +0700 Subject: [PATCH] fix(core): concurrency cancellation, TOCTOU in callTool, defensive hook copies, and new tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - concurrency.ts: add AbortSignal support and stop sibling workers on error - tool-execution.ts: accept AbortSignal for early cancellation on stream abort - prompt-request.ts: create AbortController in stream(), abort in finally; add .catch(() => {}) to detached promise chain; defensively copy newMessages at all hook call sites to prevent shared-mutable-state corruption - agent.ts: fix TOCTOU race in callTool() — single get() lookup + callWithTool() instead of separate contains() + call() lookups - tool-set.ts: add callWithTool() for atomic tool invocation with pre-resolved ref - tool-execution.ts, mcp/tool.ts: Symbol.for() -> Symbol() to avoid cross-realm symbol collisions - test: concurrency.test.ts (10 tests), compact.test.ts (12 tests), tool-set.test.ts (+4 TOCTOU tests), prompt-request.test.ts (+2 tests for unhandled rejection and message integrity) --- packages/core/src/agent/agent.ts | 16 ++- packages/core/src/agent/builder.ts | 2 +- packages/core/src/internal/concurrency.ts | 14 ++- .../internal/prompt-runtime/tool-execution.ts | 6 +- packages/core/src/mcp/tool.ts | 2 +- packages/core/src/request/prompt-request.ts | 50 +++++---- packages/core/src/tool/tool-set.ts | 11 +- packages/core/test/compact.test.ts | 69 ++++++++++++ packages/core/test/concurrency.test.ts | 105 ++++++++++++++++++ packages/core/test/prompt-request.test.ts | 77 +++++++++++++ packages/core/test/tool-set.test.ts | 38 +++++++ 11 files changed, 357 insertions(+), 33 deletions(-) create mode 100644 packages/core/test/compact.test.ts create mode 100644 packages/core/test/concurrency.test.ts diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index c39e40d0..85b7be25 100644 --- a/packages/core/src/agent/agent.ts +++ b/packages/core/src/agent/agent.ts @@ -16,6 +16,7 @@ import type { AgentObserverRegistration } from "../observability"; import { PromptRequest } from "../request"; import { createTool } from "../tool/create-tool"; import type { ToolSearchDocument } from "../tool/dynamic-tools"; +import { ToolNotFoundError } from "../tool/errors"; import type { AgentMiddleware } from "../tool/middleware"; import { isSkillTool } from "../tool/skill-tool-marker"; import type { @@ -174,18 +175,23 @@ export class Agent { args: string, context?: ToolCallContext, ): Promise { - if (this.toolSet.contains(toolName)) { - return this.toolSet.call(toolName, args, context); + const tool = this.toolSet.get(toolName); + if (tool !== undefined) { + return this.toolSet.callWithTool(tool, args, context); } for (const registration of this.dynamicTools) { const toolSet = dynamicToolSetFromIndex(registration.index); - if (toolSet?.contains(toolName)) { - return toolSet.call(toolName, args, context); + if (toolSet === undefined) { + continue; + } + const dynamicTool = toolSet.get(toolName); + if (dynamicTool !== undefined) { + return toolSet.callWithTool(dynamicTool, args, context); } } - return this.toolSet.call(toolName, args, context); + throw new ToolNotFoundError(toolName); } shouldApplyToolMiddleware(toolName: string): boolean { diff --git a/packages/core/src/agent/builder.ts b/packages/core/src/agent/builder.ts index 7820c859..38cb0402 100644 --- a/packages/core/src/agent/builder.ts +++ b/packages/core/src/agent/builder.ts @@ -1,6 +1,6 @@ import type { CompletionModel, Document, JsonObject, JsonValue, ToolChoice } from "../completion"; -import { appendGuardrailPolicies } from "../guardrails"; import type { GuardrailPolicy, GuardrailPolicyInput } from "../guardrails"; +import { appendGuardrailPolicies } from "../guardrails"; import type { PromptHook } from "../hooks"; import type { McpServer } from "../mcp"; import { resolveMemoryOptions } from "../memory/options"; diff --git a/packages/core/src/internal/concurrency.ts b/packages/core/src/internal/concurrency.ts index 259f4b61..3ac27e9c 100644 --- a/packages/core/src/internal/concurrency.ts +++ b/packages/core/src/internal/concurrency.ts @@ -2,16 +2,26 @@ export async function mapWithConcurrency( inputs: Input[], concurrency: number, mapper: (input: Input) => Promise, + signal?: AbortSignal, ): Promise { const limit = Math.max(1, Math.trunc(concurrency)); + if (inputs.length === 0) { + return []; + } const results = new Array(inputs.length); let nextIndex = 0; + let rejected = false; async function worker(): Promise { - while (nextIndex < inputs.length) { + while (nextIndex < inputs.length && !rejected && !signal?.aborted) { const index = nextIndex; nextIndex += 1; - results[index] = await mapper(inputs[index] as Input); + try { + results[index] = await mapper(inputs[index] as Input); + } catch (error) { + rejected = true; + throw error; + } } } diff --git a/packages/core/src/internal/prompt-runtime/tool-execution.ts b/packages/core/src/internal/prompt-runtime/tool-execution.ts index 68db208e..81e1ed0f 100644 --- a/packages/core/src/internal/prompt-runtime/tool-execution.ts +++ b/packages/core/src/internal/prompt-runtime/tool-execution.ts @@ -31,7 +31,7 @@ import type { import { compact } from "../compact"; import { mapWithConcurrency } from "../concurrency"; -const MCP_TOOL_METADATA_KEY = Symbol.for("anvia.mcp.tool.metadata"); +const MCP_TOOL_METADATA_KEY = Symbol("anvia.mcp.tool.metadata"); export type ToolResultEventPayload = { type: "tool_result"; @@ -83,8 +83,12 @@ export class ToolCallExecutor { onResult?: (result: ToolResultEventPayload) => void, onStreamEvent?: (event: AgentToolEventPayload) => void, observation?: ToolExecutionObservation, + signal?: AbortSignal, ): Promise { return mapWithConcurrency(toolCalls, this.concurrency, async (toolCall) => { + if (signal?.aborted) { + return ToolContent.toolResult(toolCall.id, "Tool execution cancelled", toolCall.callId); + } const args = JSON.stringify(toolCall.function.arguments ?? {}); const internalCallId = globalThis.crypto.randomUUID(); const hookArgs: ToolHookArgs = { diff --git a/packages/core/src/mcp/tool.ts b/packages/core/src/mcp/tool.ts index baf0eeab..87be6275 100644 --- a/packages/core/src/mcp/tool.ts +++ b/packages/core/src/mcp/tool.ts @@ -3,7 +3,7 @@ import type { Tool } from "../tool/index"; import { createCallToolParams, mapMcpToolResult } from "./result"; import type { McpClient, McpToolDefinition } from "./types"; -const MCP_TOOL_METADATA_KEY = Symbol.for("anvia.mcp.tool.metadata"); +const MCP_TOOL_METADATA_KEY = Symbol("anvia.mcp.tool.metadata"); export function createMcpTool( definition: McpToolDefinition, diff --git a/packages/core/src/request/prompt-request.ts b/packages/core/src/request/prompt-request.ts index 92758fb3..a6ce5a7f 100644 --- a/packages/core/src/request/prompt-request.ts +++ b/packages/core/src/request/prompt-request.ts @@ -193,7 +193,7 @@ export class PromptRequest { newMessages = [this.promptMessage]; this.chatHistory = await this.memoryRecorder.prepareRun(runId, newMessages); const pendingTurnMessages = this.memoryRecorder.pendingTurnMessages(newMessages); - await this.runRunStartHook(newMessages); + await this.runRunStartHook([...newMessages]); while (currentTurns <= this.maxTurnCount + 1) { const prompt = newMessages.at(-1); if (prompt === undefined) { @@ -204,8 +204,8 @@ export class PromptRequest { currentTurns += 1; const historyForRequest = [...this.chatHistory, ...newMessages.slice(0, -1)]; - await this.runTurnStartHook(currentTurns, prompt, historyForRequest, newMessages); - await this.runCompletionCallHook(prompt, historyForRequest, newMessages); + await this.runTurnStartHook(currentTurns, prompt, historyForRequest, [...newMessages]); + await this.runCompletionCallHook(prompt, historyForRequest, [...newMessages]); const ragText = extractRagText(prompt); const dynamicContext = await fetchDynamicContext(this.agent, ragText); @@ -230,13 +230,13 @@ export class PromptRequest { try { response = await this.runCompletion(request, currentTurns, runObservers); } catch (error) { - await this.runCompletionErrorHook(prompt, error, newMessages); + await this.runCompletionErrorHook(prompt, error, [...newMessages]); throw error; } response = await this.runCompletionResponseMiddlewares(request, response, currentTurns); usage = Usage.add(usage, response.usage); - await this.runCompletionResponseHook(prompt, response, newMessages); - await this.runTurnEndHook(currentTurns, response, newMessages); + await this.runCompletionResponseHook(prompt, response, [...newMessages]); + await this.runTurnEndHook(currentTurns, response, [...newMessages]); const toolCalls = response.choice.filter( (item): item is ToolCall => item.type === "tool_call", @@ -282,7 +282,7 @@ export class PromptRequest { trace: runObservers.trace, guardrails: [...this.guardrailDecisions], }; - await this.runRunEndHook(result, newMessages); + await this.runRunEndHook(result, [...newMessages]); await runObservers.end(result); await this.memoryRecorder.commitCompletedRun( runId, @@ -326,7 +326,7 @@ export class PromptRequest { throw new MaxTurnsError(this.maxTurnCount, [...this.chatHistory, ...newMessages], lastPrompt); } catch (error) { - const finalError = await this.runRunErrorHook(error, usage, newMessages); + const finalError = await this.runRunErrorHook(error, usage, [...newMessages]); this.runState = finalError instanceof PromptCancelledError ? "cancelled" : "errored"; await runObservers.error({ error: finalError, usage, messages: [...newMessages] }); await this.memoryRecorder.recordError(runId, finalError, newMessages); @@ -339,6 +339,8 @@ export class PromptRequest { throw new Error("This completion model does not support streaming"); } + const abort = new AbortController(); + this.startRun(); const runId = globalThis.crypto.randomUUID(); const emit = async (event: AgentStreamEvent): Promise => { @@ -388,7 +390,7 @@ export class PromptRequest { newMessages = [this.promptMessage]; this.chatHistory = await this.memoryRecorder.prepareRun(runId, newMessages); const pendingTurnMessages = this.memoryRecorder.pendingTurnMessages(newMessages); - await this.runRunStartHook(newMessages); + await this.runRunStartHook([...newMessages]); while (currentTurns <= this.maxTurnCount + 1) { const prompt = newMessages.at(-1); if (prompt === undefined) { @@ -405,8 +407,8 @@ export class PromptRequest { prompt, history: historyForRequest, }); - await this.runTurnStartHook(currentTurns, prompt, historyForRequest, newMessages); - await this.runCompletionCallHook(prompt, historyForRequest, newMessages); + await this.runTurnStartHook(currentTurns, prompt, historyForRequest, [...newMessages]); + await this.runCompletionCallHook(prompt, historyForRequest, [...newMessages]); const ragText = extractRagText(prompt); const dynamicContext = await fetchDynamicContext(this.agent, ragText); @@ -471,7 +473,7 @@ export class PromptRequest { } } catch (error) { await generationObservers.error({ turn: currentTurns, error }); - await this.runCompletionErrorHook(prompt, error, newMessages); + await this.runCompletionErrorHook(prompt, error, [...newMessages]); throw error; } @@ -485,8 +487,8 @@ export class PromptRequest { ); response = await this.runCompletionResponseMiddlewares(request, response, currentTurns); usage = Usage.add(usage, response.usage); - await this.runCompletionResponseHook(prompt, response, newMessages); - await this.runTurnEndHook(currentTurns, response, newMessages); + await this.runCompletionResponseHook(prompt, response, [...newMessages]); + await this.runTurnEndHook(currentTurns, response, [...newMessages]); const toolCalls = response.choice.filter( (item): item is ToolCall => item.type === "tool_call", @@ -555,7 +557,7 @@ export class PromptRequest { trace: runObservers.trace, guardrails: [...this.guardrailDecisions], }; - await this.runRunEndHook(result, newMessages); + await this.runRunEndHook(result, [...newMessages]); await runObservers.end(result); await this.memoryRecorder.commitCompletedRun( runId, @@ -611,11 +613,14 @@ export class PromptRequest { runObservers, toolDefinitions: request.tools, }, + abort.signal, ); - toolResultsPromise.then( - () => toolResultEvents.close(), - (error: unknown) => toolResultEvents.throw(error), - ); + toolResultsPromise + .then( + () => toolResultEvents.close(), + (error: unknown) => toolResultEvents.throw(error), + ) + .catch(() => {}); for await (const result of toolResultEvents) { yield await emit({ turn: currentTurns, ...result }); } @@ -634,12 +639,14 @@ export class PromptRequest { throw new MaxTurnsError(this.maxTurnCount, [...this.chatHistory, ...newMessages], lastPrompt); } catch (error) { - const finalError = await this.runRunErrorHook(error, usage, newMessages); + const finalError = await this.runRunErrorHook(error, usage, [...newMessages]); this.runState = finalError instanceof PromptCancelledError ? "cancelled" : "errored"; await runObservers.error({ error: finalError, usage, messages: [...newMessages] }); await this.memoryRecorder.recordError(runId, finalError, newMessages); yield await emit({ type: "error", error: finalError }); throw finalError; + } finally { + abort.abort("stream generator exited"); } } @@ -700,6 +707,7 @@ export class PromptRequest { runObservers: ActiveAgentRunObservers; toolDefinitions?: ToolDefinition[]; }, + signal?: AbortSignal, ): Promise { const executor = new ToolCallExecutor( this.agent, @@ -714,7 +722,7 @@ export class PromptRequest { this.requestMiddlewares, (reason) => this.cancelled(newMessages, reason), ); - return executor.execute(toolCalls, onResult, onStreamEvent, observation); + return executor.execute(toolCalls, onResult, onStreamEvent, observation, signal); } private async runOutputGuardrailsForResponse( diff --git a/packages/core/src/tool/tool-set.ts b/packages/core/src/tool/tool-set.ts index bbb02ea1..d08b839d 100644 --- a/packages/core/src/tool/tool-set.ts +++ b/packages/core/src/tool/tool-set.ts @@ -65,12 +65,19 @@ export class ToolSet { if (tool === undefined) { throw new ToolNotFoundError(toolName); } + return this.callWithTool(tool, args, context); + } + async callWithTool( + tool: AnyTool, + args: string, + context?: ToolCallContext, + ): Promise { let parsedArgs: unknown; try { parsedArgs = parseToolArgs(args); } catch (error) { - throw new ToolJsonError(`Invalid JSON arguments for tool ${toolName}`, error); + throw new ToolJsonError(`Invalid JSON arguments for tool ${tool.name}`, error); } try { @@ -80,7 +87,7 @@ export class ToolSet { if (error instanceof Error) { throw new ToolCallError(error.message, error); } - throw new ToolCallError(`Tool ${toolName} failed`, error); + throw new ToolCallError(`Tool ${tool.name} failed`, error); } } } diff --git a/packages/core/test/compact.test.ts b/packages/core/test/compact.test.ts new file mode 100644 index 00000000..83160c51 --- /dev/null +++ b/packages/core/test/compact.test.ts @@ -0,0 +1,69 @@ +import { describe, expect, it } from "vitest"; +import { compact, isRecord } from "../src/internal/compact"; + +describe("compact", () => { + it("removes undefined values", () => { + const result = compact({ a: 1, b: undefined, c: "hello" }); + expect(result).toEqual({ a: 1, c: "hello" }); + }); + + it("preserves null values", () => { + const result = compact({ a: 1, b: null, c: "hello" }); + expect(result).toEqual({ a: 1, b: null, c: "hello" }); + }); + + it("preserves falsy values other than undefined", () => { + const result = compact({ a: 0, b: false, c: "", d: NaN }); + expect(result).toEqual({ a: 0, b: false, c: "", d: NaN }); + }); + + it("returns empty object when all values are undefined", () => { + const result = compact({ a: undefined, b: undefined }); + expect(result).toEqual({}); + }); + + it("returns empty object for empty input", () => { + const result = compact({}); + expect(result).toEqual({}); + }); + + it("preserves nested objects", () => { + const nested = { x: 1, y: "hello" }; + const result = compact({ a: nested, b: undefined }); + expect(result).toEqual({ a: nested }); + }); + + it("preserves arrays", () => { + const result = compact({ a: [1, 2, 3], b: undefined }); + expect(result).toEqual({ a: [1, 2, 3] }); + }); + + it("does not mutate the input", () => { + const input = { a: 1, b: undefined, c: "hello" }; + compact(input); + expect(input).toEqual({ a: 1, b: undefined, c: "hello" }); + }); +}); + +describe("isRecord", () => { + it("returns true for plain objects", () => { + expect(isRecord({})).toBe(true); + expect(isRecord({ a: 1 })).toBe(true); + }); + + it("returns false for null", () => { + expect(isRecord(null)).toBe(false); + }); + + it("returns false for arrays", () => { + expect(isRecord([])).toBe(false); + expect(isRecord([1, 2, 3])).toBe(false); + }); + + it("returns false for primitives", () => { + expect(isRecord("hello")).toBe(false); + expect(isRecord(42)).toBe(false); + expect(isRecord(true)).toBe(false); + expect(isRecord(undefined)).toBe(false); + }); +}); diff --git a/packages/core/test/concurrency.test.ts b/packages/core/test/concurrency.test.ts new file mode 100644 index 00000000..196184a8 --- /dev/null +++ b/packages/core/test/concurrency.test.ts @@ -0,0 +1,105 @@ +import { describe, expect, it } from "vitest"; +import { mapWithConcurrency } from "../src/internal/concurrency"; + +async function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +describe("mapWithConcurrency", () => { + it("maps over an array with concurrency limit", async () => { + const result = await mapWithConcurrency([1, 2, 3], 2, async (n) => n * 2); + expect(result).toEqual([2, 4, 6]); + }); + + it("returns empty array for empty input", async () => { + const result = await mapWithConcurrency([], 5, async (n: number) => n); + expect(result).toEqual([]); + }); + + it("handles single input", async () => { + const result = await mapWithConcurrency(["a"], 1, async (s) => s.toUpperCase()); + expect(result).toEqual(["A"]); + }); + + it("propagates errors from mapper", async () => { + const promise = mapWithConcurrency([1, 2, 3], 2, async (n) => { + if (n === 2) throw new Error("boom"); + return n; + }); + await expect(promise).rejects.toThrow("boom"); + }); + + it("stops remaining workers when one fails", async () => { + const executionOrder: number[] = []; + const promise = mapWithConcurrency([1, 2, 3, 4], 4, async (n) => { + executionOrder.push(n); + if (n === 2) { + throw new Error("boom"); + } + await sleep(50); + executionOrder.push(-n); + return n; + }); + await expect(promise).rejects.toThrow("boom"); + // Workers for items after the failing one should not start new work + expect(executionOrder.filter((n) => n < 0).length).toBeLessThan(3); + }); + + it("respects concurrency limit", async () => { + let concurrent = 0; + let maxConcurrent = 0; + + await mapWithConcurrency([1, 2, 3, 4, 5, 6], 3, async (n) => { + concurrent += 1; + maxConcurrent = Math.max(maxConcurrent, concurrent); + await sleep(10); + concurrent -= 1; + return n; + }); + + expect(maxConcurrent).toBe(3); + }); + + it("aborts early when signal is aborted", async () => { + const abort = new AbortController(); + const executed: number[] = []; + + const promise = mapWithConcurrency( + [1, 2, 3, 4, 5], + 5, + async (n) => { + executed.push(n); + if (n === 2) { + abort.abort(); + } + await sleep(50); + return n; + }, + abort.signal, + ); + + await promise; + expect(executed.includes(1)).toBe(true); + expect(executed.includes(2)).toBe(true); + // Items after 2 may or may not have run before the signal propagated + // The key assertion: we don't crash and get a result + }); + + it("handles concurrency of 0 or less gracefully", async () => { + const result = await mapWithConcurrency([1, 2, 3], 0, async (n) => n); + expect(result).toEqual([1, 2, 3]); + }); + + it("handles concurrency larger than input length", async () => { + const result = await mapWithConcurrency([1, 2], 10, async (n) => n); + expect(result).toEqual([1, 2]); + }); + + it("preserves result order with async work", async () => { + const result = await mapWithConcurrency([3, 1, 2], 2, async (n) => { + await sleep(n * 10); + return n; + }); + expect(result).toEqual([3, 1, 2]); + }); +}); diff --git a/packages/core/test/prompt-request.test.ts b/packages/core/test/prompt-request.test.ts index 86c8fe25..00be8c89 100644 --- a/packages/core/test/prompt-request.test.ts +++ b/packages/core/test/prompt-request.test.ts @@ -6,6 +6,7 @@ import { type CompletionModel, type CompletionRequest, type CompletionResponse, + type CompletionStreamEvent, cancelPrompt, createHook, createMiddleware, @@ -15,6 +16,7 @@ import { Message, PromptCancelledError, requestToolApproval, + type StreamingCompletionModel, ToolApprovalRequiredError, ToolOutput, Usage, @@ -1186,4 +1188,79 @@ describe("PromptRequest", () => { title: "summary_response", }); }); + + it("does not produce unhandled rejections when stream is cancelled early", async () => { + const model = new StreamingQueueModel([ + [ + { + type: "tool_call" as const, + toolCall: AssistantContent.toolCall("call_1", "add", { x: 1, y: 2 }), + }, + { + type: "tool_call" as const, + toolCall: AssistantContent.toolCall("call_2", "add", { x: 3, y: 4 }), + }, + ], + ]); + const agent = new AgentBuilder("test-agent", model).tool(addTool).build(); + + const unhandledRejections: unknown[] = []; + const handler = (reason: unknown) => unhandledRejections.push(reason); + process.on("unhandledRejection", handler); + try { + const iterator = agent.prompt("add").stream()[Symbol.asyncIterator](); + await iterator.next(); + await iterator.return?.(); + } finally { + process.off("unhandledRejection", handler); + } + + expect(unhandledRejections).toHaveLength(0); + }); + + it("returns a valid messages array after multi-turn tool execution", async () => { + const model = new QueueModel([ + response([AssistantContent.toolCall("call_1", "add", { x: 1, y: 2 })]), + response([AssistantContent.text("3")]), + ]); + const agent = new AgentBuilder("test-agent", model).tool(addTool).build(); + const result = await agent.prompt("add").send(); + + expect(result.messages).toHaveLength(4); + expect(result.messages[0]?.role).toBe("user"); + expect(result.messages[1]?.role).toBe("assistant"); + expect(result.messages[2]?.role).toBe("tool"); + expect(result.messages[3]?.role).toBe("assistant"); + expect(result.messages[3]?.content).toEqual([{ type: "text", text: "3" }]); + }); }); + +class StreamingQueueModel implements StreamingCompletionModel { + readonly provider = "test"; + readonly defaultModel = "test"; + readonly capabilities = { + streaming: true, + tools: true, + toolChoice: true, + imageInput: true, + documentInput: true, + outputSchema: true, + reasoning: true, + }; + readonly requests: CompletionRequest[] = []; + + constructor(private readonly responses: CompletionStreamEvent[][]) {} + + async completion(): Promise { + throw new Error("completion should not be called"); + } + + async *streamCompletion(request: CompletionRequest): AsyncIterable { + this.requests.push(request); + const response = this.responses.shift(); + if (response === undefined) { + throw new Error("No queued response"); + } + yield* response; + } +} diff --git a/packages/core/test/tool-set.test.ts b/packages/core/test/tool-set.test.ts index caafe9a5..955ec8d0 100644 --- a/packages/core/test/tool-set.test.ts +++ b/packages/core/test/tool-set.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from "vitest"; import { z } from "zod"; import { + type AnyTool, createTool, ToolCallError, ToolJsonError, @@ -196,4 +197,41 @@ describe("ToolSet", () => { expect.objectContaining({ name: "add", description: "Replace add" }), ]); }); + + it("prevents TOCTOU race between contains() and call() via callWithTool", async () => { + const toolA = createTool({ + name: "a", + description: "Tool A", + input: z.object({}), + execute: () => "a_result", + }); + const toolSet = new ToolSet().addTool(toolA); + + // Get the tool by name, then delete it from the set + const tool = toolSet.get("a"); + expect(tool).toBeDefined(); + toolSet.deleteTool("a"); + expect(toolSet.contains("a")).toBe(false); + + // callWithTool should still work because it uses the tool reference directly + await expect(toolSet.callWithTool(tool as AnyTool, "{}")).resolves.toBe("a_result"); + }); + + it("callWithTool throws ToolJsonError for invalid JSON args", async () => { + const toolSet = ToolSet.fromTools([addTool]); + const tool = toolSet.get("add"); + expect(tool).toBeDefined(); + + await expect(toolSet.callWithTool(tool as AnyTool, "{")).rejects.toBeInstanceOf(ToolJsonError); + }); + + it("callWithTool throws ToolCallError for Zod validation failure", async () => { + const toolSet = ToolSet.fromTools([addTool]); + const tool = toolSet.get("add"); + expect(tool).toBeDefined(); + + await expect( + toolSet.callWithTool(tool as AnyTool, JSON.stringify({ x: "not_a_number", y: 5 })), + ).rejects.toBeInstanceOf(ToolCallError); + }); });