@illuma-ai/agents 1.4.0-alpha.6 → 1.5.0

package/src/tools/ToolNode.ts

@@ -1,6 +1,7 @@
 import { ToolCall } from '@langchain/core/messages/tool';
 import {
   ToolMessage,
+  HumanMessage,
   isAIMessage,
   isBaseMessage,
 } from '@langchain/core/messages';
@@ -24,7 +25,30 @@ import { ExecutionContext } from './approval/constants';
 import { RunnableCallable } from '@/utils';
 import { processToolOutput } from '@/utils/toonFormat';
 import { safeDispatchCustomEvent } from '@/utils/events';
-import { Constants, GraphEvents } from '@/common';
+import { Constants, GraphEvents, CODE_EXECUTION_TOOLS } from '@/common';
+import { executeHooks } from '@/hooks';
+import type { HookRegistry, AggregatedHookResult } from '@/hooks';
+import {
+  ToolOutputReferenceRegistry,
+  buildReferenceKey,
+} from '@/tools/toolOutputReferences';
+import type { ResolvedArgsByCallId } from '@/tools/toolOutputReferences';
+
+/**
+ * Per-call batch context for `runTool`. Bundles every optional batch-scoped
+ * value the method needs so the signature stays at three positional params
+ * even as new context fields are added. Upstream-aligned (PR #117).
+ */
+type RunToolBatchContext = {
+  /** Position of this call within the parent ToolNode batch. */
+  batchIndex?: number;
+  /** Batch turn shared across every call in the batch. */
+  turn?: number;
+  /** Registry partition scope (run id or anonymous batch id). */
+  batchScopeId?: string;
+  /** Batch-local sink for post-substitution args. */
+  resolvedArgsByCallId?: ResolvedArgsByCallId;
+};
 
 /**
  * Helper to check if a value is a Send object
@@ -97,6 +121,29 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
   private toolApprovalConfig?: t.ToolApprovalConfig;
   /** Buffer for recovering truncated tool call arguments from streaming data */
   private streamingToolCallBuffer?: import('@/tools/StreamingToolCallBuffer').StreamingToolCallBuffer;
+  /** Hook registry threaded down from Run for tool-lifecycle events. */
+  private hookRegistry?: HookRegistry;
+  /**
+   * Tool output reference registry threaded down from Run / Graph
+   * (upstream PR #114). When set, dispatchToolEvents resolves
+   * `{{tool<i>turn<n>}}` placeholders in args before invoking each tool
+   * and stores successful outputs under their stable reference keys for
+   * subsequent calls in the same run to pipe through.
+   */
+  private toolOutputRegistry?: ToolOutputReferenceRegistry;
+  /**
+   * Per-batch turn counter for the registry. Increments once per
+   * `dispatchToolEvents` invocation so `tool<i>turn<n>` keys remain
+   * unique across overlapping batches in a single run.
+   */
+  private toolOutputTurn: number = 0;
+  /**
+   * Counter for synthesising scope ids when `run()` is called without a
+   * `run_id` in config. Each anonymous batch gets a unique
+   * `\0anon-<n>` scope so concurrent invocations don't collide on the
+   * shared `toolOutputRegistry`.
+   */
+  private anonBatchCounter: number = 0;
 
   constructor({
     tools,
@@ -114,6 +161,10 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     directToolNames,
     streamingToolCallBuffer,
     toolApprovalConfig,
+    hookRegistry,
+    toolOutputRegistry,
+    toolOutputReferences,
+    maxToolResultChars: _maxToolResultChars,
   }: t.ToolNodeConstructorParams) {
     super({ name, tags, func: (input, config) => this.run(input, config) });
     this.toolMap = toolMap ?? new Map(tools.map((tool) => [tool.name, tool]));
@@ -129,6 +180,33 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     this.directToolNames = directToolNames;
     this.streamingToolCallBuffer = streamingToolCallBuffer;
     this.toolApprovalConfig = toolApprovalConfig;
+    this.hookRegistry = hookRegistry;
+    /**
+     * Precedence: an explicitly passed `toolOutputRegistry` instance
+     * wins (the multi-agent case where every ToolNode shares one registry
+     * for the run); otherwise the per-ToolNode `toolOutputReferences`
+     * config builds an instance scoped to this ToolNode.
+     */
+    if (toolOutputRegistry != null) {
+      this.toolOutputRegistry = toolOutputRegistry;
+    } else if (toolOutputReferences?.enabled === true) {
+      this.toolOutputRegistry = new ToolOutputReferenceRegistry({
+        maxOutputSize: toolOutputReferences.maxOutputSize,
+        maxTotalSize: toolOutputReferences.maxTotalSize,
+      });
+    }
+  }
+
+  /**
+   * Test-observation accessor for the tool-output reference registry.
+   * @internal Exposed for test observation only. Host code should rely on
+   * `{{tool<i>turn<n>}}` substitution at tool-invocation time and not
+   * mutate the registry directly.
+   */
+  public _unsafeGetToolOutputRegistry():
+    | ToolOutputReferenceRegistry
+    | undefined {
+    return this.toolOutputRegistry;
   }
 
   /**
@@ -217,8 +295,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     // Instead, gate the bypass on the actual routing-tool name prefix.
     if (
       this.directToolNames?.has(toolName) &&
-      (toolName.startsWith(Constants.LC_TRANSFER_TO_) ||
-        toolName.startsWith(Constants.LC_HANDOFF_TO_))
+      toolName.startsWith(Constants.LC_TRANSFER_TO_)
     ) {
       return false;
     }
@@ -284,13 +361,52 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
   }
 
   /**
-   * Runs a single tool call with error handling
+   * Runs a single tool call with error handling.
+   *
+   * @param batchContext  Optional per-batch context (upstream PR #117).
+   *   Threaded from `run()` for tool output reference annotation. The
+   *   `batchScopeId` field carries an anonymous synthetic scope when the
+   *   caller has no `run_id`, so concurrent batches don't collide on the
+   *   shared registry.
    */
   protected async runTool(
     call: ToolCall,
-    config: RunnableConfig
+    config: RunnableConfig,
+    batchContext: RunToolBatchContext = {}
   ): Promise<BaseMessage | Command> {
+    const {
+      batchIndex,
+      turn: batchTurn,
+      batchScopeId,
+      resolvedArgsByCallId,
+    } = batchContext;
     const tool = this.toolMap.get(call.name);
+    const registry = this.toolOutputRegistry;
+    /**
+     * Precompute the reference key once per call — captured locally so
+     * concurrent `invoke()` calls on the same ToolNode cannot race on a
+     * shared turn field.
+     */
+    const refKey =
+      registry != null && batchIndex != null && batchTurn != null
+        ? buildReferenceKey(batchIndex, batchTurn)
+        : undefined;
+    /**
+     * Hoisted outside the try so the catch branch can stamp
+     * `_unresolvedRefs` into error ToolMessage metadata. This lets the
+     * lazy annotation transform surface a `[unresolved refs: …]` hint to
+     * the LLM, helping the model self-correct when its reference key
+     * caused the failure.
+     */
+    let unresolvedRefs: string[] = [];
+    /**
+     * Use the caller-provided `batchScopeId` when threaded from `run()`.
+     * Fall back to the config's `run_id` when runTool is invoked from a
+     * context that doesn't thread the batch context — that still
+     * preserves the runId-based partitioning for named runs.
+     */
+    const runId =
+      batchScopeId ?? (config.configurable?.run_id as string | undefined);
     try {
       if (tool === undefined) {
         throw new Error(`Tool "${call.name}" not found.`);
@@ -312,6 +428,29 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         this.streamingToolCallBuffer.clear(call.id);
       }
 
+      /**
+       * Resolve `{{tool<i>turn<n>}}` placeholders in args against the
+       * registry. Captures `unresolvedRefs` so error/success branches can
+       * surface them to the LLM via `additional_kwargs._unresolvedRefs`.
+       */
+      if (registry != null) {
+        const { resolved, unresolved } = registry.resolve(runId, args);
+        args = resolved as typeof call.args;
+        unresolvedRefs = unresolved;
+        if (
+          resolvedArgsByCallId != null &&
+          call.id != null &&
+          call.id !== '' &&
+          resolved !== call.args &&
+          typeof resolved === 'object'
+        ) {
+          resolvedArgsByCallId.set(
+            call.id,
+            resolved as Record<string, unknown>
+          );
+        }
+      }
+
       // Build invoke params - LangChain extracts non-schema fields to config.toolCall
       let invokeParams: Record<string, unknown> = {
         ...call,
@@ -345,10 +484,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
        * session_id is always injected when available (even without tracked files)
        * so the CodeExecutor can fall back to the /files endpoint for session continuity.
        */
-      if (
-        call.name === Constants.EXECUTE_CODE ||
-        call.name === Constants.PROGRAMMATIC_TOOL_CALLING
-      ) {
+      if (CODE_EXECUTION_TOOLS.has(call.name)) {
         const codeSession = this.sessions?.get(Constants.EXECUTE_CODE) as
           | t.CodeSessionContext
           | undefined;
@@ -444,21 +580,58 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         minReductionPercent: 10, // Only apply TOON when clearly beneficial
       });
 
-      // Step 4: Return ToolMessage with processed string content
+      /**
+       * Tool output reference metadata (upstream PR #117). Register the
+       * raw, untruncated content so future `{{...}}` substitutions deliver
+       * the full payload, and stamp `_refKey` / `_refScope` /
+       * `_unresolvedRefs` into `additional_kwargs` for the lazy
+       * annotation transform to apply at LLM request time. Persisted
+       * `content` stays clean.
+       *
+       * For error ToolMessages from the underlying tool: bypass
+       * registration but still stamp `_unresolvedRefs` so the model can
+       * self-correct.
+       */
       if (isBaseMessage(output) && output._getType() === 'tool') {
         const toolMsg = output as ToolMessage;
+        const isError = toolMsg.status === 'error';
+        const refMeta = isError
+          ? unresolvedRefs.length > 0
+            ? { _unresolvedRefs: unresolvedRefs }
+            : undefined
+          : this.recordOutputReference(
+              runId,
+              rawContent,
+              refKey,
+              unresolvedRefs
+            );
         return new ToolMessage({
           status: toolMsg.status,
           name: toolMsg.name,
           content: processed.content,
           tool_call_id: toolMsg.tool_call_id,
+          ...(refMeta != null && {
+            additional_kwargs: {
+              ...toolMsg.additional_kwargs,
+              ...refMeta,
+            },
+          }),
         });
       } else {
+        const refMeta = this.recordOutputReference(
+          runId,
+          rawContent,
+          refKey,
+          unresolvedRefs
+        );
         return new ToolMessage({
           status: 'success',
           name: tool.name,
           content: processed.content,
           tool_call_id: call.id!,
+          ...(refMeta != null && {
+            additional_kwargs: refMeta as Record<string, unknown>,
+          }),
         });
       }
     } catch (_e: unknown) {
@@ -510,6 +683,9 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         content: `Error: ${e.message}\n Please fix your mistakes.`,
         name: call.name,
         tool_call_id: call.id ?? '',
+        ...(unresolvedRefs.length > 0 && {
+          additional_kwargs: { _unresolvedRefs: unresolvedRefs },
+        }),
       });
     }
   }
@@ -618,6 +794,34 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     return parsedArgs;
   }
 
+  /**
+   * Stores the raw, untruncated tool output in the registry under `refKey`
+   * (when registry + refKey are both present) and returns the metadata
+   * envelope to stamp on `ToolMessage.additional_kwargs`. The lazy
+   * annotation transform reads `_refKey` / `_refScope` / `_unresolvedRefs`
+   * at LLM request time to produce the transient annotated copy.
+   */
+  private recordOutputReference(
+    runId: string | undefined,
+    registryContent: string,
+    refKey: string | undefined,
+    unresolved: string[]
+  ): t.ToolMessageRefMetadata | undefined {
+    if (this.toolOutputRegistry != null && refKey != null) {
+      this.toolOutputRegistry.set(runId, refKey, registryContent);
+    }
+    if (refKey == null && unresolved.length === 0) return undefined;
+    const meta: t.ToolMessageRefMetadata = {};
+    if (refKey != null) {
+      meta._refKey = refKey;
+      // Anonymous invocations use a synthetic scope minted in run() — stamp
+      // it onto the message so annotateMessagesForLLM can recover it later.
+      if (runId != null) meta._refScope = runId;
+    }
+    if (unresolved.length > 0) meta._unresolvedRefs = unresolved;
+    return meta;
+  }
+
   /**
    * Builds code session context for injection into event-driven tool calls.
    * Mirrors the session injection logic in runTool() for direct execution.
@@ -822,13 +1026,41 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
   }
 
   /**
-   * Dispatches tool calls to the host via ON_TOOL_EXECUTE event and returns raw ToolMessages.
-   * Core logic for event-driven execution, separated from output shaping.
+   * Converts InjectedMessage entries into LangChain HumanMessages.
+   * 'user' and 'system' both become HumanMessage to avoid provider rejections
+   * (Anthropic/Google reject non-leading SystemMessages); the original role
+   * and metadata are preserved in additional_kwargs for downstream consumers.
+   */
+  private convertInjectedMessages(
+    messages: t.InjectedMessage[]
+  ): BaseMessage[] {
+    const converted: BaseMessage[] = [];
+    for (const msg of messages) {
+      const additional_kwargs: Record<string, unknown> = {
+        role: msg.role,
+      };
+      if (msg.isMeta != null) additional_kwargs.isMeta = msg.isMeta;
+      if (msg.source != null) additional_kwargs.source = msg.source;
+      if (msg.skillName != null) additional_kwargs.skillName = msg.skillName;
+
+      converted.push(
+        new HumanMessage({ content: msg.content, additional_kwargs })
+      );
+    }
+    return converted;
+  }
+
+  /**
+   * Dispatches tool calls to the host via ON_TOOL_EXECUTE event and returns
+   * the resulting ToolMessages along with any injected messages emitted by
+   * the host (e.g. SkillTool body injection). Injected messages must be
+   * placed AFTER the ToolMessages so the AIMessage tool_calls -> ToolMessage
+   * adjacency required by some providers is preserved.
    */
   private async dispatchToolEvents(
     toolCalls: ToolCall[],
     config: RunnableConfig
-  ): Promise<ToolMessage[]> {
+  ): Promise<{ toolMessages: ToolMessage[]; injected: BaseMessage[] }> {
     // ========================================================================
     // HITL: Check approval for event-dispatched tools (browser, MCP, etc.)
     // before dispatching. Uses LangGraph interrupt() for each tool needing
@@ -863,10 +1095,132 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
 
     // If all tools were denied, return denial messages without dispatching
     if (approvedCalls.length === 0) {
-      return denialMessages;
+      return { toolMessages: denialMessages, injected: [] };
+    }
+
+    const runId = config.configurable?.run_id as string | undefined;
+    const threadId = config.configurable?.thread_id as string | undefined;
+
+    /**
+     * PreToolUse hooks may deny / ask / allow each call, and may rewrite
+     * `toolInput` via `updatedInput`. A denied/asked call short-circuits
+     * with a synthetic error ToolMessage and fires PermissionDenied so
+     * the host can audit. Survivors continue to dispatch.
+     */
+    const hookDenialMessages: ToolMessage[] = [];
+    let postHookCalls = approvedCalls;
+    if (
+      runId != null &&
+      this.hookRegistry?.hasHookFor('PreToolUse', runId) === true
+    ) {
+      const HOOK_FALLBACK: AggregatedHookResult = {
+        additionalContexts: [],
+        errors: [],
+      };
+      const surviving: ToolCall[] = [];
+      const preResults = await Promise.all(
+        approvedCalls.map((call) =>
+          executeHooks({
+            registry: this.hookRegistry!,
+            input: {
+              hook_event_name: 'PreToolUse',
+              runId,
+              threadId,
+              agentId: this.agentId,
+              toolName: call.name,
+              toolInput: call.args as Record<string, unknown>,
+              toolUseId: call.id ?? '',
+              stepId: this.toolCallStepIds?.get(call.id ?? ''),
+              turn: this.toolUsageCount.get(call.name) ?? 0,
+            },
+            sessionId: runId,
+            matchQuery: call.name,
+          }).catch((): AggregatedHookResult => HOOK_FALLBACK)
+        )
+      );
+
+      for (let i = 0; i < approvedCalls.length; i++) {
+        const call = approvedCalls[i];
+        const hookResult = preResults[i];
+        const denied =
+          hookResult.decision === 'deny' || hookResult.decision === 'ask';
+        if (denied) {
+          const reason = hookResult.reason ?? 'Blocked by hook';
+          hookDenialMessages.push(
+            new ToolMessage({
+              status: 'error',
+              content: `Blocked: ${reason}`,
+              name: call.name,
+              tool_call_id: call.id ?? '',
+            })
+          );
+          if (this.hookRegistry.hasHookFor('PermissionDenied', runId)) {
+            // Fire-and-forget — denial accounting must not block dispatch.
+            void executeHooks({
+              registry: this.hookRegistry,
+              input: {
+                hook_event_name: 'PermissionDenied',
+                runId,
+                threadId,
+                agentId: this.agentId,
+                toolName: call.name,
+                toolInput: call.args as Record<string, unknown>,
+                toolUseId: call.id ?? '',
+                reason,
+              },
+              sessionId: runId,
+              matchQuery: call.name,
+            }).catch(() => {
+              /* swallow — denial is informational */
+            });
+          }
+          continue;
+        }
+        if (
+          hookResult.updatedInput != null &&
+          typeof hookResult.updatedInput === 'object'
+        ) {
+          call.args = hookResult.updatedInput as Record<string, unknown>;
+        }
+        surviving.push(call);
+      }
+      postHookCalls = surviving;
+
+      if (postHookCalls.length === 0) {
+        return {
+          toolMessages: [...denialMessages, ...hookDenialMessages],
+          injected: [],
+        };
+      }
     }
 
-    const requests: t.ToolCallRequest[] = approvedCalls.map((call) => {
+    /**
+     * Tool output reference resolution (upstream PR #114): walk each call's
+     * args and substitute `{{tool<i>turn<n>}}` placeholders with the stored
+     * raw output. Captured here BEFORE request shaping so the substituted
+     * args flow through the rest of the dispatch unchanged.
+     *
+     * Failures are non-fatal: the registry resolver returns the original
+     * string unchanged for unknown keys, and surfaces them via `unresolved`
+     * which we currently log-only. Future work can route unresolved-key
+     * complaints into a hint message for the LLM.
+     */
+    const batchTurn = this.toolOutputTurn++;
+    if (this.toolOutputRegistry != null && runId != null) {
+      const registry = this.toolOutputRegistry;
+      for (const call of postHookCalls) {
+        const args = call.args as Record<string, unknown> | undefined;
+        if (args == null) continue;
+        try {
+          const { resolved } = registry.resolve(runId, args);
+          call.args = resolved;
+        } catch {
+          /* leave args untouched on resolver failure */
+        }
+      }
+    }
+
+    const requests: t.ToolCallRequest[] = postHookCalls.map((call) => {
       const turn = this.toolUsageCount.get(call.name) ?? 0;
       this.toolUsageCount.set(call.name, turn + 1);
 
@@ -918,7 +1272,85 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
 
     this.storeCodeSessionFromResults(results, requests);
 
-    const eventMessages = results.map((result) => {
+    const injected: BaseMessage[] = [];
+    for (const result of results) {
+      if (result.injectedMessages && result.injectedMessages.length > 0) {
+        try {
+          injected.push(
+            ...this.convertInjectedMessages(result.injectedMessages)
+          );
+        } catch (e) {
+          // eslint-disable-next-line no-console
+          console.warn(
+            `[ToolNode] Failed to convert injectedMessages for toolCallId=${result.toolCallId}:`,
+            e instanceof Error ? e.message : e
+          );
+        }
+      }
+    }
+
+    /**
+     * Fire PostToolUse / PostToolUseFailure hooks after results are in.
+     * These run in parallel and never block the return — failures are
+     * swallowed so a hook bug can't masquerade as a tool failure.
+     */
+    if (runId != null && this.hookRegistry != null) {
+      const hasPost = this.hookRegistry.hasHookFor('PostToolUse', runId);
+      const hasFail = this.hookRegistry.hasHookFor('PostToolUseFailure', runId);
+      if (hasPost || hasFail) {
+        const postPromises: Promise<unknown>[] = [];
+        for (const result of results) {
+          const request = requests.find((r) => r.id === result.toolCallId);
+          if (!request) continue;
+          if (result.status === 'success' && hasPost) {
+            postPromises.push(
+              executeHooks({
+                registry: this.hookRegistry,
+                input: {
+                  hook_event_name: 'PostToolUse',
+                  runId,
+                  threadId,
+                  agentId: this.agentId,
+                  toolName: request.name,
+                  toolInput: request.args as Record<string, unknown>,
+                  toolOutput: result.content,
+                  toolUseId: result.toolCallId,
+                  stepId: this.toolCallStepIds?.get(result.toolCallId),
+                  turn: request.turn,
+                },
+                sessionId: runId,
+                matchQuery: request.name,
+              }).catch(() => undefined)
+            );
+          } else if (result.status === 'error' && hasFail) {
+            postPromises.push(
+              executeHooks({
+                registry: this.hookRegistry,
+                input: {
+                  hook_event_name: 'PostToolUseFailure',
+                  runId,
+                  threadId,
+                  agentId: this.agentId,
+                  toolName: request.name,
+                  toolInput: request.args as Record<string, unknown>,
+                  toolUseId: result.toolCallId,
+                  error: result.errorMessage ?? 'Unknown error',
+                  stepId: this.toolCallStepIds?.get(result.toolCallId),
+                  turn: request.turn,
+                },
+                sessionId: runId,
+                matchQuery: request.name,
+              }).catch(() => undefined)
+            );
+          }
+        }
+        if (postPromises.length > 0) {
+          await Promise.all(postPromises);
+        }
+      }
+    }
+
+    const eventMessages = results.map((result, idx) => {
       const request = requests.find((r) => r.id === result.toolCallId);
       const toolName = request?.name ?? 'unknown';
       const stepId = this.toolCallStepIds?.get(result.toolCallId) ?? '';
@@ -947,12 +1379,40 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           typeof result.content === 'string'
             ? result.content
             : JSON.stringify(result.content);
+        /**
+         * Tool output reference — lazy annotation (upstream PRs #114 + #117).
+         * Register the raw output under `tool<idx>turn<turn>` in the
+         * registry but leave `ToolMessage.content` clean. Stamp the ref
+         * metadata into `additional_kwargs` instead. The lazy
+         * `annotateMessagesForLLM` transform applies the visible
+         * `[ref: …]` / `_ref` annotation only on the transient copy sent
+         * to the model. Persisted message stays unannotated so
+         * conversation exports / pruning / token accounting stay clean.
+         */
+        const additionalKwargs: Record<string, unknown> = {};
+        if (this.toolOutputRegistry != null && runId != null) {
+          try {
+            const refKey = buildReferenceKey(idx, batchTurn);
+            this.toolOutputRegistry.set(runId, refKey, contentString);
+            const meta: t.ToolMessageRefMetadata = {
+              _refKey: refKey,
+              _refScope: runId,
+            };
+            Object.assign(additionalKwargs, meta);
+          } catch {
+            /* registry store failure is non-fatal — message stays clean, no annotation */
+          }
+        }
         toolMessage = new ToolMessage({
           status: 'success',
           name: toolName,
           content: contentString,
           artifact: result.artifact,
           tool_call_id: result.toolCallId,
+          additional_kwargs:
+            Object.keys(additionalKwargs).length > 0
+              ? additionalKwargs
+              : undefined,
         });
       }
 
@@ -985,7 +1445,14 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       return toolMessage;
     });
 
-    return [...denialMessages, ...eventMessages];
+    return {
+      toolMessages: [
+        ...denialMessages,
+        ...hookDenialMessages,
+        ...eventMessages,
+      ],
+      injected,
+    };
   }
 
   /**
@@ -998,13 +1465,32 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     input: any
   ): Promise<T> {
-    const outputs = await this.dispatchToolEvents(toolCalls, config);
+    const { toolMessages, injected } = await this.dispatchToolEvents(
+      toolCalls,
+      config
+    );
+    const outputs: BaseMessage[] = [...toolMessages, ...injected];
     return (Array.isArray(input) ? outputs : { messages: outputs }) as T;
   }
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   protected async run(input: any, config: RunnableConfig): Promise<T> {
     this.toolCallTurns.clear();
+    /**
+     * Per-batch local map for resolved (post-substitution) args. Lives
+     * on the stack so concurrent `run()` calls on the same ToolNode
+     * cannot read or wipe each other's entries.
+     */
+    const resolvedArgsByCallId: ResolvedArgsByCallId = new Map();
+    /**
+     * Claim this batch's turn synchronously from the registry (or fall
+     * back to 0 when the feature is disabled). For anonymous callers
+     * (no `run_id` in config), mint a unique per-batch scope id so two
+     * concurrent anonymous invocations don't target the same bucket.
+     */
+    const incomingRunId = config.configurable?.run_id as string | undefined;
+    const batchScopeId = incomingRunId ?? `\0anon-${this.anonBatchCounter++}`;
+    const batchTurn = this.toolOutputRegistry?.nextTurn(batchScopeId) ?? 0;
     let outputs: (BaseMessage | Command)[];
 
     if (this.isSendInput(input)) {
@@ -1012,7 +1498,14 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       if (this.eventDrivenMode && isDirectTool !== true) {
         return this.executeViaEvent([input.lg_tool_call], config, input);
       }
-      outputs = [await this.runTool(input.lg_tool_call, config)];
+      outputs = [
+        await this.runTool(input.lg_tool_call, config, {
+          batchIndex: 0,
+          turn: batchTurn,
+          batchScopeId,
+          resolvedArgsByCallId,
+        }),
+      ];
       this.handleRunToolCompletions([input.lg_tool_call], outputs, config);
     } else {
       let messages: BaseMessage[];
@@ -1081,23 +1574,52 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         );
 
         // Run direct tools and event tools in parallel — they are independent
-        const [directOutputs, eventOutputs] = (await Promise.all([
+        const [directOutputs, eventDispatch] = (await Promise.all([
           directCalls.length > 0
-            ? Promise.all(directCalls.map((call) => this.runTool(call, config)))
-            : [],
+            ? Promise.all(
+                directCalls.map((call, idx) =>
+                  this.runTool(call, config, {
+                    batchIndex: idx,
+                    turn: batchTurn,
+                    batchScopeId,
+                    resolvedArgsByCallId,
+                  })
+                )
+              )
+            : ([] as (BaseMessage | Command)[]),
           eventCalls.length > 0
             ? this.dispatchToolEvents(eventCalls, config)
-            : [],
-        ])) as [(BaseMessage | Command)[], ToolMessage[]];
+            : ({ toolMessages: [], injected: [] } as {
+                toolMessages: ToolMessage[];
+                injected: BaseMessage[];
+              }),
+        ])) as [
+          (BaseMessage | Command)[],
+          { toolMessages: ToolMessage[]; injected: BaseMessage[] },
+        ];
 
         if (directCalls.length > 0 && directOutputs.length > 0) {
           this.handleRunToolCompletions(directCalls, directOutputs, config);
         }
 
-        outputs = [...directOutputs, ...eventOutputs];
+        // Injected messages MUST follow ToolMessages so the AIMessage
+        // tool_calls -> ToolMessage adjacency required by some providers
+        // is preserved.
+        outputs = [
+          ...directOutputs,
+          ...eventDispatch.toolMessages,
+          ...eventDispatch.injected,
+        ];
       } else {
         outputs = await Promise.all(
-          filteredCalls.map((call) => this.runTool(call, config))
+          filteredCalls.map((call, idx) =>
+            this.runTool(call, config, {
+              batchIndex: idx,
+              turn: batchTurn,
+              batchScopeId,
+              resolvedArgsByCallId,
+            })
+          )
         );
         this.handleRunToolCompletions(filteredCalls, outputs, config);
       }