fix(core): concurrency cancellation, TOCTOU in callTool, defensive hook copies, and new tests

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<M extends CompletionModel = CompletionModel> {
     args: string,
     context?: ToolCallContext,
   ): Promise<NormalizedToolOutput> {
-    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 {

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";

packages/core/src/internal/concurrency.ts

@@ -2,16 +2,26 @@ export async function mapWithConcurrency<Input, Output>(
   inputs: Input[],
   concurrency: number,
   mapper: (input: Input) => Promise<Output>,
+  signal?: AbortSignal,
 ): Promise<Output[]> {
   const limit = Math.max(1, Math.trunc(concurrency));
+  if (inputs.length === 0) {
+    return [];
+  }
   const results = new Array<Output>(inputs.length);
   let nextIndex = 0;
+  let rejected = false;
 
   async function worker(): Promise<void> {
-    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;
+      }
     }
   }
 

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<ToolResult[]> {
     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 = {

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,

packages/core/src/request/prompt-request.ts

@@ -193,7 +193,7 @@ export class PromptRequest<M extends CompletionModel = CompletionModel> {
       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<M extends CompletionModel = CompletionModel> {
         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<M extends CompletionModel = CompletionModel> {
         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<M extends CompletionModel = CompletionModel> {
             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<M extends CompletionModel = CompletionModel> {
 
       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<M extends CompletionModel = CompletionModel> {
       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<AgentStreamEvent> => {
@@ -388,7 +390,7 @@ export class PromptRequest<M extends CompletionModel = CompletionModel> {
       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<M extends CompletionModel = CompletionModel> {
           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<M extends CompletionModel = CompletionModel> {
           }
         } 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<M extends CompletionModel = CompletionModel> {
         );
         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<M extends CompletionModel = CompletionModel> {
             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<M extends CompletionModel = CompletionModel> {
             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<M extends CompletionModel = CompletionModel> {
 
       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<M extends CompletionModel = CompletionModel> {
       runObservers: ActiveAgentRunObservers;
       toolDefinitions?: ToolDefinition[];
     },
+    signal?: AbortSignal,
   ): Promise<ToolResult[]> {
     const executor = new ToolCallExecutor(
       this.agent,
@@ -714,7 +722,7 @@ export class PromptRequest<M extends CompletionModel = CompletionModel> {
       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(

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<NormalizedToolOutput> {
     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);
     }
   }
 }

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);
+  });
+});