feat(@langchain/core): support of ToolRuntime (#9316)

Author: christian-bromann

.changeset/large-onions-attack.md

@@ -0,0 +1,6 @@
+---
+"@langchain/core": patch
+"langchain": patch
+---
+
+feat(@langchain/core): support of ToolRuntime

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

@@ -55,6 +55,7 @@ import type {
   StringInputToolSchema,
   ToolInterface,
   ToolOutputType,
+  ToolRuntime,
 } from "./types.js";
 import { type JSONSchema, validatesOnlyStrings } from "../utils/json_schema.js";
 
@@ -79,6 +80,7 @@ export {
   isRunnableToolLike,
   isStructuredTool,
   isStructuredToolParams,
+  type ToolRuntime,
 } from "./types.js";
 
 export { ToolInputParsingException };
@@ -635,6 +637,98 @@ export function tool<
 >(
   func: RunnableFunc<SchemaOutputT, ToolOutputT, ToolRunnableConfig>,
   fields: ToolWrapperParams<SchemaT>
+):
+  | DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>
+  | DynamicTool<ToolOutputT>;
+
+// Overloads with ToolRuntime as CallOptions
+export function tool<
+  SchemaT extends ZodStringV3,
+  ToolOutputT = ToolOutputType,
+  TState = unknown,
+  TContext = unknown
+>(
+  func: (
+    input: InferInteropZodOutput<SchemaT>,
+    runtime: ToolRuntime<TState, TContext>
+  ) => ToolOutputT | Promise<ToolOutputT>,
+  fields: ToolWrapperParams<SchemaT>
+): DynamicTool<ToolOutputT>;
+
+export function tool<
+  SchemaT extends ZodStringV4,
+  ToolOutputT = ToolOutputType,
+  TState = unknown,
+  TContext = unknown
+>(
+  func: (
+    input: InferInteropZodOutput<SchemaT>,
+    runtime: ToolRuntime<TState, TContext>
+  ) => ToolOutputT | Promise<ToolOutputT>,
+  fields: ToolWrapperParams<SchemaT>
+): DynamicTool<ToolOutputT>;
+
+export function tool<
+  SchemaT extends ZodObjectV3,
+  SchemaOutputT = InferInteropZodOutput<SchemaT>,
+  SchemaInputT = InferInteropZodInput<SchemaT>,
+  ToolOutputT = ToolOutputType,
+  TState = unknown,
+  TContext = unknown
+>(
+  func: (
+    input: SchemaOutputT,
+    runtime: ToolRuntime<TState, TContext>
+  ) => ToolOutputT | Promise<ToolOutputT>,
+  fields: ToolWrapperParams<SchemaT>
+): DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>;
+
+export function tool<
+  SchemaT extends ZodObjectV4,
+  SchemaOutputT = InferInteropZodOutput<SchemaT>,
+  SchemaInputT = InferInteropZodInput<SchemaT>,
+  ToolOutputT = ToolOutputType,
+  TState = unknown,
+  TContext = unknown
+>(
+  func: (
+    input: SchemaOutputT,
+    runtime: ToolRuntime<TState, TContext>
+  ) => ToolOutputT | Promise<ToolOutputT>,
+  fields: ToolWrapperParams<SchemaT>
+): DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>;
+
+export function tool<
+  SchemaT extends JSONSchema,
+  SchemaOutputT = ToolInputSchemaOutputType<SchemaT>,
+  SchemaInputT = ToolInputSchemaInputType<SchemaT>,
+  ToolOutputT = ToolOutputType,
+  TState = unknown,
+  TContext = unknown
+>(
+  func: (
+    input: Parameters<DynamicStructuredToolInput<SchemaT>["func"]>[0],
+    runtime: ToolRuntime<TState, TContext>
+  ) => ToolOutputT | Promise<ToolOutputT>,
+  fields: ToolWrapperParams<SchemaT>
+): DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>;
+
+export function tool<
+  SchemaT extends
+    | InteropZodObject
+    | InteropZodType<string>
+    | JSONSchema = InteropZodObject,
+  SchemaOutputT = ToolInputSchemaOutputType<SchemaT>,
+  SchemaInputT = ToolInputSchemaInputType<SchemaT>,
+  ToolOutputT = ToolOutputType,
+  TState = unknown,
+  TContext = unknown
+>(
+  func: (
+    input: SchemaOutputT,
+    runtime: ToolRuntime<TState, TContext>
+  ) => ToolOutputT | Promise<ToolOutputT>,
+  fields: ToolWrapperParams<SchemaT>
 ):
   | DynamicStructuredTool<SchemaT, SchemaOutputT, SchemaInputT, ToolOutputT>
   | DynamicTool<ToolOutputT> {
@@ -659,9 +753,8 @@ export function tool<
             pickRunnableConfigKeys(childConfig),
             async () => {
               try {
-                // TS doesn't restrict the type here based on the guard above
                 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                resolve(func(input as any, childConfig));
+                resolve(func(input as any, childConfig as any));
               } catch (e) {
                 reject(e);
               }
@@ -704,7 +797,8 @@ export function tool<
           pickRunnableConfigKeys(childConfig),
           async () => {
             try {
-              const result = await func(input, childConfig);
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              const result = await func(input as any, childConfig as any);
 
               /**
                * If the signal is aborted, we don't want to resolve the promise

libs/langchain-core/src/tools/tests/types.test-d.ts

@@ -0,0 +1,65 @@
+import { z } from "zod/v3";
+import { describe, it, expectTypeOf } from "vitest";
+
+import { tool } from "../index.js";
+import type { ToolRuntime } from "../types.js";
+import type { RunnableConfig } from "../../runnables/config.js";
+
+describe("ToolRuntime", () => {
+  it("allows to define runnable config argument as ToolRuntime", () => {
+    const stateSchema = z.object({
+      userId: z.string(),
+    });
+    const contextSchema = z.object({
+      db: z.object({
+        foo: z.string(),
+      }),
+    });
+
+    type State = z.infer<typeof stateSchema>;
+    type Context = z.infer<typeof contextSchema>;
+
+    tool(
+      (
+        input,
+        runtime: ToolRuntime<typeof stateSchema, typeof contextSchema>
+      ) => {
+        expectTypeOf(input).toEqualTypeOf<{
+          some: string;
+        }>();
+        expectTypeOf(runtime.state).toEqualTypeOf<State>();
+        expectTypeOf(runtime.context).toEqualTypeOf<Context>();
+        expectTypeOf(runtime.toolCallId).toEqualTypeOf<string>();
+        expectTypeOf(runtime.config).toMatchTypeOf<RunnableConfig>();
+        return `Hello, ${runtime.state.userId}!`;
+      },
+      {
+        name: "test",
+        description: "test",
+        schema: z.object({
+          some: z.string(),
+        }),
+      }
+    );
+
+    tool(
+      (input, runtime: ToolRuntime<State, Context>) => {
+        expectTypeOf(input).toEqualTypeOf<{
+          some: string;
+        }>();
+        expectTypeOf(runtime.state).toEqualTypeOf<State>();
+        expectTypeOf(runtime.context).toEqualTypeOf<Context>();
+        expectTypeOf(runtime.toolCallId).toEqualTypeOf<string>();
+        expectTypeOf(runtime.config).toMatchTypeOf<RunnableConfig>();
+        return `Hello, ${runtime.state.userId}!`;
+      },
+      {
+        name: "test",
+        description: "test",
+        schema: z.object({
+          some: z.string(),
+        }),
+      }
+    );
+  });
+});

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

@@ -21,8 +21,11 @@ import {
   type InferInteropZodOutput,
   type InteropZodType,
   isInteropZodSchema,
+  type InteropZodObject,
 } from "../utils/types/zod.js";
+
 import { JSONSchema } from "../utils/json_schema.js";
+import type { BaseStore } from "../stores.js";
 
 export type ResponseFormat = "content" | "content_and_artifact" | string;
 
@@ -425,3 +428,106 @@ export function isLangChainTool(tool?: unknown): tool is StructuredToolParams {
     isStructuredTool(tool as any)
   );
 }
+
+/**
+ * Runtime context automatically injected into tools.
+ *
+ * When a tool function has a parameter named `tool_runtime` with type hint
+ * `ToolRuntime`, the tool execution system will automatically inject an instance
+ * containing:
+ *
+ * - `state`: The current graph state
+ * - `toolCallId`: The ID of the current tool call
+ * - `config`: `RunnableConfig` for the current execution
+ * - `context`: Runtime context
+ * - `store`: `BaseStore` instance for persistent storage
+ * - `writer`: Stream writer for streaming output
+ *
+ * No `Annotated` wrapper is needed - just use `runtime: ToolRuntime`
+ * as a parameter.
+ *
+ * @example
+ * ```typescript
+ * import { tool, ToolRuntime } from "@langchain/core/tools";
+ * import { z } from "zod";
+ *
+ * const stateSchema = z.object({
+ *   messages: z.array(z.any()),
+ *   userId: z.string().optional(),
+ * });
+ *
+ * const greet = tool(
+ *   async ({ name }, runtime) => {
+ *     // Access state
+ *     const messages = runtime.state.messages;
+ *
+ *     // Access tool_call_id
+ *     console.log(`Tool call ID: ${runtime.toolCallId}`);
+ *
+ *     // Access config
+ *     console.log(`Run ID: ${runtime.config.runId}`);
+ *
+ *     // Access runtime context
+ *     const userId = runtime.context?.userId;
+ *
+ *     // Access store
+ *     await runtime.store?.mset([["key", "value"]]);
+ *
+ *     // Stream output
+ *     runtime.writer?.("Processing...");
+ *
+ *     return `Hello! User ID: ${runtime.state.userId || "unknown"} ${name}`;
+ *   },
+ *   {
+ *     name: "greet",
+ *     description: "Use this to greet the user once you found their info.",
+ *     schema: z.object({ name: z.string() }),
+ *     stateSchema,
+ *   }
+ * );
+ * ```
+ *
+ * @template StateT - The type of the state schema (inferred from stateSchema)
+ * @template ContextT - The type of the context schema (inferred from contextSchema)
+ */
+export type ToolRuntime<
+  TState = unknown,
+  TContext = unknown
+> = RunnableConfig & {
+  /**
+   * The current graph state.
+   */
+  state: TState extends InteropZodObject
+    ? InferInteropZodOutput<TState>
+    : TState extends Record<string, unknown>
+    ? TState
+    : unknown;
+  /**
+   * The ID of the current tool call.
+   */
+  toolCallId: string;
+  /**
+   * The current tool call.
+   */
+  toolCall?: ToolCall;
+  /**
+   * RunnableConfig for the current execution.
+   */
+  config: ToolRunnableConfig;
+  /**
+   * Runtime context (from langgraph `Runtime`).
+   */
+  context: TContext extends InteropZodObject
+    ? InferInteropZodOutput<TContext>
+    : TContext extends Record<string, unknown>
+    ? TContext
+    : unknown;
+  /**
+   * BaseStore instance for persistent storage (from langgraph `Runtime`).
+   */
+  store: BaseStore<string, unknown> | null;
+  /**
+   * Stream writer for streaming output (from langgraph `Runtime`).
+   */
+  writer: ((chunk: unknown) => void) | null;
+};

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

@@ -294,6 +294,12 @@ export class ToolNode<
           { ...toolCall, type: "tool_call" },
           {
             ...config,
+            /**
+             * extend to match ToolRuntime
+             */
+            config,
+            toolCallId: toolCall.id!,
+            state: config.configurable?.__pregel_scratchpad?.currentTaskInput,
             signal: mergeAbortSignals(this.signal, config.signal),
           }
         );