fix(langchain): make ToolNode implementation less middleware centric (#9315)

Author: christian-bromann

libs/langchain-core/src/tools/index.ts

@@ -20,6 +20,7 @@ import {
 import type { RunnableFunc } from "../runnables/base.js";
 import { isDirectToolOutput, ToolCall, ToolMessage } from "../messages/tool.js";
 import { AsyncLocalStorageProviderSingleton } from "../singletons/index.js";
+import type { RunnableToolLike } from "../runnables/base.js";
 import {
   _configHasToolCallId,
   _isToolCall,
@@ -773,3 +774,9 @@ function _stringify(content: unknown): string {
     return `${content}`;
   }
 }
+
+export type ServerTool = Record<string, unknown>;
+export type ClientTool =
+  | StructuredToolInterface
+  | DynamicTool
+  | RunnableToolLike;

libs/langchain/src/agents/ReactAgent.ts

@@ -18,9 +18,14 @@ import { ToolMessage, AIMessage } from "@langchain/core/messages";
 import { IterableReadableStream } from "@langchain/core/utils/stream";
 import type { Runnable, RunnableConfig } from "@langchain/core/runnables";
 import type { StreamEvent } from "@langchain/core/tracers/log_stream";
+import type { ClientTool, ServerTool } from "@langchain/core/tools";
 
 import { createAgentAnnotationConditional } from "./annotation.js";
-import { isClientTool, validateLLMHasNoBoundTools } from "./utils.js";
+import {
+  isClientTool,
+  validateLLMHasNoBoundTools,
+  wrapToolCall,
+} from "./utils.js";
 
 import { AgentNode } from "./nodes/AgentNode.js";
 import { ToolNode } from "./nodes/ToolNode.js";
@@ -35,7 +40,6 @@ import {
 import { StateManager } from "./state.js";
 
 import type { WithStateGraphNodes } from "./types.js";
-import type { ClientTool, ServerTool } from "./tools.js";
 
 import type {
   CreateAgentParams,
@@ -333,8 +337,7 @@ export class ReactAgent<
     if (toolClasses.filter(isClientTool).length > 0) {
       const toolNode = new ToolNode(toolClasses.filter(isClientTool), {
         signal: this.options.signal,
-        middleware,
-        stateManager: this.#stateManager,
+        wrapToolCall: wrapToolCall(middleware),
       });
       allNodeWorkflows.addNode("tools", toolNode);
     }

libs/langchain/src/agents/middleware.ts

@@ -3,8 +3,8 @@ import type {
   InteropZodObject,
   InferInteropZodOutput,
 } from "@langchain/core/utils/types";
+import type { ClientTool, ServerTool } from "@langchain/core/tools";
 
-import type { ClientTool, ServerTool } from "./tools.js";
 import type {
   AgentMiddleware,
   WrapToolCallHook,

libs/langchain/src/agents/middleware/toolRetry.ts

@@ -1,11 +1,11 @@
 /**
  * Tool retry middleware for agents.
  */
-
-import { ToolMessage } from "@langchain/core/messages";
 import { z } from "zod/v3";
+import { ToolMessage } from "@langchain/core/messages";
+import type { ClientTool, ServerTool } from "@langchain/core/tools";
+
 import { createMiddleware } from "../middleware.js";
-import type { ClientTool, ServerTool } from "../tools.js";
 import type { AgentMiddleware } from "./types.js";
 import { sleep } from "./utils.js";
 

libs/langchain/src/agents/middleware/types.ts

@@ -11,17 +11,17 @@ import type { AnnotationRoot } from "@langchain/langgraph";
 import type { AIMessage, ToolMessage } from "@langchain/core/messages";
 import type { ToolCall } from "@langchain/core/messages/tool";
 import type { Command } from "@langchain/langgraph";
+import type { ClientTool, ServerTool } from "@langchain/core/tools";
 
 import type { JumpToTarget } from "../constants.js";
-import type { ClientTool, ServerTool } from "../tools.js";
 import type { Runtime, AgentBuiltInState } from "../runtime.js";
 import type { ModelRequest } from "../nodes/types.js";
 
 type PromiseOrValue<T> = T | Promise<T>;
 
 export type AnyAnnotationRoot = AnnotationRoot<any>;
 
-type NormalizedSchemaInput<
+export type NormalizedSchemaInput<
   TSchema extends InteropZodObject | undefined | never = any
 > = [TSchema] extends [never]
   ? AgentBuiltInState

libs/langchain/src/agents/nodes/AgentNode.ts

@@ -11,6 +11,7 @@ import {
   interopZodObjectPartial,
 } from "@langchain/core/utils/types";
 import type { ToolCall } from "@langchain/core/messages/tool";
+import type { ClientTool, ServerTool } from "@langchain/core/tools";
 
 import { initChatModel } from "../../chat_models/universal.js";
 import { MultipleStructuredOutputsError } from "../errors.js";
@@ -32,7 +33,6 @@ import type {
   WrapModelCallHandler,
 } from "../middleware/types.js";
 import type { ModelRequest } from "./types.js";
-import type { ClientTool, ServerTool } from "../tools.js";
 import { withAgentName } from "../withAgentName.js";
 import {
   ToolStrategy,

libs/langchain/src/agents/nodes/ToolNode.ts

@@ -20,15 +20,13 @@ import {
 import { RunnableCallable } from "../RunnableCallable.js";
 import { PreHookAnnotation } from "../annotation.js";
 import { mergeAbortSignals } from "./utils.js";
-import { wrapToolCall } from "../utils.js";
 import { ToolInvocationError } from "../errors.js";
 import type {
-  AnyAnnotationRoot,
+  WrapToolCallHook,
   ToolCallRequest,
   ToAnnotationRoot,
 } from "../middleware/types.js";
-import type { AgentMiddleware } from "../middleware/types.js";
-import type { StateManager } from "../state.js";
+import type { AgentBuiltInState } from "../runtime.js";
 
 export interface ToolNodeOptions {
   /**
@@ -66,13 +64,11 @@ export interface ToolNodeOptions {
     | boolean
     | ((error: unknown, toolCall: ToolCall) => ToolMessage | undefined);
   /**
-   * The middleware to use for tool execution.
+   * Optional wrapper function for tool execution.
+   * Allows middleware to intercept and modify tool calls before execution.
+   * The wrapper receives the tool call request and a handler function to execute the tool.
    */
-  middleware?: readonly AgentMiddleware[];
-  /**
-   * The state manager to use for tool execution.
-   */
-  stateManager?: StateManager;
+  wrapToolCall?: WrapToolCallHook;
 }
 
 const isBaseMessageArray = (input: unknown): input is BaseMessage[] =>
@@ -165,8 +161,8 @@ function defaultHandleToolErrors(
  * ```
  */
 export class ToolNode<
-  StateSchema extends AnyAnnotationRoot | InteropZodObject = any,
-  ContextSchema extends AnyAnnotationRoot | InteropZodObject = any
+  StateSchema extends InteropZodObject = any,
+  ContextSchema extends InteropZodObject = any
 > extends RunnableCallable<StateSchema, ContextSchema> {
   tools: (StructuredToolInterface | DynamicTool | RunnableToolLike)[];
 
@@ -179,15 +175,13 @@ export class ToolNode<
     | ((error: unknown, toolCall: ToolCall) => ToolMessage | undefined) =
     defaultHandleToolErrors;
 
-  middleware: readonly AgentMiddleware[] = [];
-
-  stateManager?: StateManager;
+  wrapToolCall: WrapToolCallHook | undefined;
 
   constructor(
     tools: (StructuredToolInterface | DynamicTool | RunnableToolLike)[],
     public options?: ToolNodeOptions
   ) {
-    const { name, tags, handleToolErrors, middleware, stateManager, signal } =
+    const { name, tags, handleToolErrors, signal, wrapToolCall } =
       options ?? {};
     super({
       name,
@@ -201,9 +195,8 @@ export class ToolNode<
     });
     this.tools = tools;
     this.handleToolErrors = handleToolErrors ?? this.handleToolErrors;
-    this.middleware = middleware ?? [];
     this.signal = signal;
-    this.stateManager = stateManager;
+    this.wrapToolCall = wrapToolCall;
   }
 
   /**
@@ -279,7 +272,7 @@ export class ToolNode<
   protected async runTool(
     call: ToolCall,
     config: RunnableConfig,
-    state: ToAnnotationRoot<StateSchema>["State"] & PreHookAnnotation["State"]
+    state: AgentBuiltInState
   ): Promise<ToolMessage | Command> {
     /**
      * Define the base handler that executes the tool.
@@ -355,20 +348,12 @@ export class ToolNode<
       runtime,
     };
 
-    /**
-     * Collect and compose wrapToolCall handlers from middleware
-     * Wrap each handler with error handling and validation
-     */
-    const wrapToolCallHandler = this.stateManager
-      ? wrapToolCall(this.middleware, state)
-      : undefined;
-
     /**
      * If wrapToolCall is provided, use it to wrap the tool execution
      */
-    if (wrapToolCallHandler) {
+    if (this.wrapToolCall) {
       try {
-        return await wrapToolCallHandler(request, baseHandler);
+        return await this.wrapToolCall(request, baseHandler);
       } catch (e: unknown) {
         /**
          * Handle middleware errors

libs/langchain/src/agents/nodes/types.ts

@@ -1,6 +1,7 @@
 import type { LanguageModelLike } from "@langchain/core/language_models/base";
 import type { BaseMessage } from "@langchain/core/messages";
-import type { ServerTool, ClientTool } from "../tools.js";
+import type { ServerTool, ClientTool } from "@langchain/core/tools";
+
 import type { Runtime, AgentBuiltInState } from "../runtime.js";
 
 /**

libs/langchain/src/agents/tests/middleware.test-d.ts

@@ -2,10 +2,10 @@ import { describe, it, expectTypeOf } from "vitest";
 import { z } from "zod/v3";
 import { HumanMessage, BaseMessage, AIMessage } from "@langchain/core/messages";
 import { tool } from "@langchain/core/tools";
+import type { ServerTool, ClientTool } from "@langchain/core/tools";
 
 import { createAgent, createMiddleware } from "../index.js";
 import type { AgentBuiltInState } from "../runtime.js";
-import type { ServerTool, ClientTool } from "../tools.js";
 
 describe("middleware types", () => {
   it("a middleware can define a state schema which is propagated to the result", async () => {
@@ -202,12 +202,34 @@ describe("middleware types", () => {
             customDefaultContextProp: string;
             customOptionalContextProp?: string;
           }>();
+          expectTypeOf(request.state).toEqualTypeOf<
+            {
+              customDefaultStateProp: string;
+              customOptionalStateProp?: string;
+              customRequiredStateProp: string;
+            } & AgentBuiltInState
+          >();
 
           return handler({
             ...request,
             tools: [tool(() => "result", { name: "toolA" })],
           });
         },
+        wrapToolCall: async (request, handler) => {
+          expectTypeOf(request.runtime.context).toEqualTypeOf<{
+            customDefaultContextProp: string;
+            customOptionalContextProp?: string;
+          }>();
+          expectTypeOf(request.state).toEqualTypeOf<
+            {
+              customDefaultStateProp: string;
+              customOptionalStateProp?: string;
+              customRequiredStateProp: string;
+            } & AgentBuiltInState
+          >();
+
+          return handler(request);
+        },
       });
 
       const agent = createAgent({