fix(langchain): improved state schema typing (langchain-ai#9285)

Author: christian-bromann

.changeset/twenty-clocks-raise.md

@@ -0,0 +1,5 @@
+---
+"langchain": patch
+---
+
+fix(langchain): improved state schema typing

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

@@ -40,7 +40,7 @@ export {
 export {
   modelCallLimitMiddleware,
   type ModelCallLimitMiddlewareConfig,
-} from "./callLimit.js";
+} from "./modelCallLimit.js";
 export { modelFallbackMiddleware } from "./modelFallback.js";
 export {
   toolRetryMiddleware,

libs/langchain/src/agents/middleware/callLimit.ts renamed to libs/langchain/src/agents/middleware/modelCallLimit.ts

@@ -4,7 +4,7 @@ import { tool } from "@langchain/core/tools";
 import { MemorySaver } from "@langchain/langgraph-checkpoint";
 
 import { FakeToolCallingChatModel } from "../../tests/utils.js";
-import { modelCallLimitMiddleware } from "../callLimit.js";
+import { modelCallLimitMiddleware } from "../modelCallLimit.js";
 import { createAgent } from "../../index.js";
 
 const toolCallMessage1 = new AIMessage({

libs/langchain/src/agents/middleware/tests/callLimit.test.ts renamed to libs/langchain/src/agents/middleware/tests/modelCallLimit.test.ts

@@ -21,13 +21,24 @@ type PromiseOrValue<T> = T | Promise<T>;
 
 export type AnyAnnotationRoot = AnnotationRoot<any>;
 
-type NormalizedSchemaInput<TSchema extends InteropZodObject | undefined = any> =
-  TSchema extends InteropZodObject ? InferInteropZodInput<TSchema> : {};
+type NormalizedSchemaInput<
+  TSchema extends InteropZodObject | undefined | never = any
+> = [TSchema] extends [never]
+  ? AgentBuiltInState
+  : TSchema extends InteropZodObject
+  ? InferInteropZodOutput<TSchema> & AgentBuiltInState
+  : TSchema extends Record<string, unknown>
+  ? TSchema & AgentBuiltInState
+  : AgentBuiltInState;
 
 /**
  * Result type for middleware functions.
  */
-export type MiddlewareResult<TState> = TState | void;
+export type MiddlewareResult<TState> =
+  | (TState & {
+      jumpTo?: JumpToTarget;
+    })
+  | void;
 
 /**
  * Represents a tool call request for the wrapToolCall hook.
@@ -61,28 +72,22 @@ export interface ToolCallRequest<
  * Takes a tool call request and returns the tool result or a command.
  */
 export type ToolCallHandler<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends Record<string, unknown> = AgentBuiltInState,
   TContext = unknown
 > = (
-  request: ToolCallRequest<
-    NormalizedSchemaInput<TSchema> & AgentBuiltInState,
-    TContext
-  >
+  request: ToolCallRequest<TSchema, TContext>
 ) => PromiseOrValue<ToolMessage | Command>;
 
 /**
  * Wrapper function type for the wrapToolCall hook.
  * Allows middleware to intercept and modify tool execution.
  */
 export type WrapToolCallHook<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends InteropZodObject | undefined = undefined,
   TContext = unknown
 > = (
-  request: ToolCallRequest<
-    NormalizedSchemaInput<TSchema> & AgentBuiltInState,
-    TContext
-  >,
-  handler: ToolCallHandler<TSchema, TContext>
+  request: ToolCallRequest<NormalizedSchemaInput<TSchema>, TContext>,
+  handler: ToolCallHandler<NormalizedSchemaInput<TSchema>, TContext>
 ) => PromiseOrValue<ToolMessage | Command>;
 
 /**
@@ -93,13 +98,10 @@ export type WrapToolCallHook<
  * @returns The AI message response from the model
  */
 export type WrapModelCallHandler<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends InteropZodObject | undefined = undefined,
   TContext = unknown
 > = (
-  request: ModelRequest<
-    NormalizedSchemaInput<TSchema> & AgentBuiltInState,
-    TContext
-  >
+  request: ModelRequest<NormalizedSchemaInput<TSchema>, TContext>
 ) => PromiseOrValue<AIMessage>;
 
 /**
@@ -116,13 +118,10 @@ export type WrapModelCallHandler<
  * @returns The AI message response from the model (or a modified version)
  */
 export type WrapModelCallHook<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends InteropZodObject | undefined = undefined,
   TContext = unknown
 > = (
-  request: ModelRequest<
-    NormalizedSchemaInput<TSchema> & AgentBuiltInState,
-    TContext
-  >,
+  request: ModelRequest<NormalizedSchemaInput<TSchema>, TContext>,
   handler: WrapModelCallHandler<TSchema, TContext>
 ) => PromiseOrValue<AIMessage>;
 
@@ -134,26 +133,23 @@ export type WrapModelCallHook<
  * @param runtime - The runtime context containing metadata, signal, writer, interrupt, etc.
  * @returns A middleware result containing partial state updates or undefined to pass through
  */
-export type BeforeAgentHandler<
-  TSchema extends InteropZodObject | undefined = any,
-  TContext = unknown
-> = (
-  state: NormalizedSchemaInput<TSchema> & AgentBuiltInState,
+type BeforeAgentHandler<TSchema, TContext> = (
+  state: TSchema,
   runtime: Runtime<TContext>
-) => PromiseOrValue<MiddlewareResult<Partial<NormalizedSchemaInput<TSchema>>>>;
+) => PromiseOrValue<MiddlewareResult<Partial<TSchema>>>;
 
 /**
  * Hook type for the beforeAgent lifecycle event.
  * Can be either a handler function or an object with a handler and optional jump targets.
  * This hook is called once at the start of the agent invocation.
  */
 export type BeforeAgentHook<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends InteropZodObject | undefined = undefined,
   TContext = unknown
 > =
-  | BeforeAgentHandler<TSchema, TContext>
+  | BeforeAgentHandler<NormalizedSchemaInput<TSchema>, TContext>
   | {
-      hook: BeforeAgentHandler<TSchema, TContext>;
+      hook: BeforeAgentHandler<NormalizedSchemaInput<TSchema>, TContext>;
       canJumpTo?: JumpToTarget[];
     };
 
@@ -165,26 +161,23 @@ export type BeforeAgentHook<
  * @param runtime - The runtime context containing metadata, signal, writer, interrupt, etc.
  * @returns A middleware result containing partial state updates or undefined to pass through
  */
-export type BeforeModelHandler<
-  TSchema extends InteropZodObject | undefined = any,
-  TContext = unknown
-> = (
-  state: NormalizedSchemaInput<TSchema> & AgentBuiltInState,
+type BeforeModelHandler<TSchema, TContext> = (
+  state: TSchema,
   runtime: Runtime<TContext>
-) => PromiseOrValue<MiddlewareResult<Partial<NormalizedSchemaInput<TSchema>>>>;
+) => PromiseOrValue<MiddlewareResult<Partial<TSchema>>>;
 
 /**
  * Hook type for the beforeModel lifecycle event.
  * Can be either a handler function or an object with a handler and optional jump targets.
  * This hook is called before each model invocation.
  */
 export type BeforeModelHook<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends InteropZodObject | undefined = undefined,
   TContext = unknown
 > =
-  | BeforeModelHandler<TSchema, TContext>
+  | BeforeModelHandler<NormalizedSchemaInput<TSchema>, TContext>
   | {
-      hook: BeforeModelHandler<TSchema, TContext>;
+      hook: BeforeModelHandler<NormalizedSchemaInput<TSchema>, TContext>;
       canJumpTo?: JumpToTarget[];
     };
 
@@ -197,26 +190,23 @@ export type BeforeModelHook<
  * @param runtime - The runtime context containing metadata, signal, writer, interrupt, etc.
  * @returns A middleware result containing partial state updates or undefined to pass through
  */
-export type AfterModelHandler<
-  TSchema extends InteropZodObject | undefined = any,
-  TContext = unknown
-> = (
-  state: NormalizedSchemaInput<TSchema> & AgentBuiltInState,
+type AfterModelHandler<TSchema, TContext> = (
+  state: TSchema,
   runtime: Runtime<TContext>
-) => PromiseOrValue<MiddlewareResult<Partial<NormalizedSchemaInput<TSchema>>>>;
+) => PromiseOrValue<MiddlewareResult<Partial<TSchema>>>;
 
 /**
  * Hook type for the afterModel lifecycle event.
  * Can be either a handler function or an object with a handler and optional jump targets.
  * This hook is called after each model invocation.
  */
 export type AfterModelHook<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends InteropZodObject | undefined = undefined,
   TContext = unknown
 > =
-  | AfterModelHandler<TSchema, TContext>
+  | AfterModelHandler<NormalizedSchemaInput<TSchema>, TContext>
   | {
-      hook: AfterModelHandler<TSchema, TContext>;
+      hook: AfterModelHandler<NormalizedSchemaInput<TSchema>, TContext>;
       canJumpTo?: JumpToTarget[];
     };
 
@@ -228,26 +218,23 @@ export type AfterModelHook<
  * @param runtime - The runtime context containing metadata, signal, writer, interrupt, etc.
  * @returns A middleware result containing partial state updates or undefined to pass through
  */
-export type AfterAgentHandler<
-  TSchema extends InteropZodObject | undefined = any,
-  TContext = unknown
-> = (
-  state: NormalizedSchemaInput<TSchema> & AgentBuiltInState,
+type AfterAgentHandler<TSchema, TContext> = (
+  state: TSchema,
   runtime: Runtime<TContext>
-) => PromiseOrValue<MiddlewareResult<Partial<NormalizedSchemaInput<TSchema>>>>;
+) => PromiseOrValue<MiddlewareResult<Partial<TSchema>>>;
 
 /**
  * Hook type for the afterAgent lifecycle event.
  * Can be either a handler function or an object with a handler and optional jump targets.
  * This hook is called once at the end of the agent invocation.
  */
 export type AfterAgentHook<
-  TSchema extends InteropZodObject | undefined = any,
+  TSchema extends InteropZodObject | undefined = undefined,
   TContext = unknown
 > =
-  | AfterAgentHandler<TSchema, TContext>
+  | AfterAgentHandler<NormalizedSchemaInput<TSchema>, TContext>
   | {
-      hook: AfterAgentHandler<TSchema, TContext>;
+      hook: AfterAgentHandler<NormalizedSchemaInput<TSchema>, TContext>;
       canJumpTo?: JumpToTarget[];
     };
 

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

@@ -4,6 +4,7 @@ import { HumanMessage, BaseMessage, AIMessage } from "@langchain/core/messages";
 import { tool } 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", () => {
@@ -169,14 +170,33 @@ describe("middleware types", () => {
           .default({
             customRequiredContextProp: "default value",
           }),
-        beforeModel: async (_state, runtime) => {
+        stateSchema: z.object({
+          customDefaultStateProp: z.string().default("default value"),
+          customOptionalStateProp: z.string().optional(),
+          customRequiredStateProp: z.string(),
+        }),
+        beforeModel: async (state, runtime) => {
+          expectTypeOf(state).toEqualTypeOf<
+            {
+              customDefaultStateProp: string;
+              customOptionalStateProp?: string;
+              customRequiredStateProp: string;
+            } & AgentBuiltInState
+          >();
           expectTypeOf(runtime.context).toEqualTypeOf<{
             customDefaultContextProp: string;
             customOptionalContextProp?: string;
             customRequiredContextProp: string;
           }>();
         },
-        afterModel: async (_state, runtime) => {
+        afterModel: async (state, runtime) => {
+          expectTypeOf(state).toEqualTypeOf<
+            {
+              customDefaultStateProp: string;
+              customOptionalStateProp?: string;
+              customRequiredStateProp: string;
+            } & AgentBuiltInState
+          >();
           expectTypeOf(runtime.context).toEqualTypeOf<{
             customDefaultContextProp: string;
             customOptionalContextProp?: string;
@@ -209,6 +229,7 @@ describe("middleware types", () => {
       await agent.invoke(
         {
           messages: [new HumanMessage("Hello, world!")],
+          customRequiredStateProp: "default value",
         },
         {
           configurable: {

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

@@ -146,6 +146,7 @@ describe("middleware state management", () => {
     const model = new FakeToolCallingModel({});
     const middleware = createMiddleware({
       name: "middleware",
+      // @ts-expect-error _privateState is not an expected return type
       beforeModel: async (_, runtime) => {
         expect(runtime.threadLevelCallCount).toBe(0);
         expect(runtime.runModelCallCount).toBe(0);
@@ -165,6 +166,7 @@ describe("middleware state management", () => {
         expect(request.runtime.runModelCallCount).toBe(0);
         return handler(request);
       },
+      // @ts-expect-error _privateState is not an expected return type
       afterModel: async (_, runtime) => {
         expect(runtime.threadLevelCallCount).toBe(1);
         expect(runtime.runModelCallCount).toBe(1);