fix(langchain): don't allow default or optional context schemas (langchain-ai#9287)

Author: christian-bromann

.changeset/fresh-dingos-repair.md

@@ -0,0 +1,5 @@
+---
+"langchain": patch
+---
+
+fix(langchain): don't allow default or optional context schemas

libs/langchain/src/agents/middleware.ts

@@ -1,8 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import type {
   InteropZodObject,
-  InteropZodDefault,
-  InteropZodOptional,
   InferInteropZodOutput,
 } from "@langchain/core/utils/types";
 
@@ -52,11 +50,7 @@ import type {
  */
 export function createMiddleware<
   TSchema extends InteropZodObject | undefined = undefined,
-  TContextSchema extends
-    | InteropZodObject
-    | InteropZodOptional<InteropZodObject>
-    | InteropZodDefault<InteropZodObject>
-    | undefined = undefined
+  TContextSchema extends InteropZodObject | undefined = undefined
 >(config: {
   /**
    * The name of the middleware
@@ -232,15 +226,7 @@ export function createMiddleware<
 }
 
 type NormalizeContextSchema<
-  TContextSchema extends
-    | InteropZodObject
-    | InteropZodOptional<InteropZodObject>
-    | InteropZodDefault<InteropZodObject>
-    | undefined = undefined
+  TContextSchema extends InteropZodObject | undefined = undefined
 > = TContextSchema extends InteropZodObject
   ? InferInteropZodOutput<TContextSchema>
-  : TContextSchema extends InteropZodDefault<any>
-  ? InferInteropZodOutput<TContextSchema>
-  : TContextSchema extends InteropZodOptional<any>
-  ? Partial<InferInteropZodOutput<TContextSchema>>
   : never;

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

@@ -133,11 +133,9 @@ describe("middleware types", () => {
     it("doesn't require users to pass in a context if a middleware has optional context schema", async () => {
       const middleware = createMiddleware({
         name: "Middleware",
-        contextSchema: z
-          .object({
-            customOptionalContextProp: z.string().default("default value"),
-          })
-          .optional(),
+        contextSchema: z.object({
+          customOptionalContextProp: z.string().optional(),
+        }),
       });
 
       const agent = createAgent({
@@ -158,18 +156,13 @@ describe("middleware types", () => {
       );
     });
 
-    it("doesn't require users to pass in a context if a middleware has context schema with defaults", async () => {
+    it("doesn't require users to pass in a context if a middleware has context schema with defaults or optional", async () => {
       const middleware = createMiddleware({
         name: "Middleware",
-        contextSchema: z
-          .object({
-            customDefaultContextProp: z.string().default("default value"),
-            customOptionalContextProp: z.string().optional(),
-            customRequiredContextProp: z.string(),
-          })
-          .default({
-            customRequiredContextProp: "default value",
-          }),
+        contextSchema: z.object({
+          customDefaultContextProp: z.string().default("default value"),
+          customOptionalContextProp: z.string().optional(),
+        }),
         stateSchema: z.object({
           customDefaultStateProp: z.string().default("default value"),
           customOptionalStateProp: z.string().optional(),
@@ -186,7 +179,6 @@ describe("middleware types", () => {
           expectTypeOf(runtime.context).toEqualTypeOf<{
             customDefaultContextProp: string;
             customOptionalContextProp?: string;
-            customRequiredContextProp: string;
           }>();
         },
         afterModel: async (state, runtime) => {
@@ -200,7 +192,6 @@ describe("middleware types", () => {
           expectTypeOf(runtime.context).toEqualTypeOf<{
             customDefaultContextProp: string;
             customOptionalContextProp?: string;
-            customRequiredContextProp: string;
           }>();
         },
         wrapModelCall: async (request, handler) => {
@@ -210,7 +201,6 @@ describe("middleware types", () => {
           expectTypeOf(request.runtime.context).toEqualTypeOf<{
             customDefaultContextProp: string;
             customOptionalContextProp?: string;
-            customRequiredContextProp: string;
           }>();
 
           return handler({
@@ -242,40 +232,23 @@ describe("middleware types", () => {
     it("doesn't require users to pass in a context if a middleware has context schema as optional", async () => {
       const middleware = createMiddleware({
         name: "Middleware",
-        contextSchema: z
-          .object({
-            customOptionalContextProp: z.string().default("default value"),
-          })
-          .optional(),
+        contextSchema: z.object({
+          customOptionalContextProp: z.string().default("default value"),
+        }),
         beforeModel: async (_state, runtime) => {
-          expectTypeOf(runtime.context).toEqualTypeOf<
-            Partial<
-              | {
-                  customOptionalContextProp: string;
-                }
-              | undefined
-            >
-          >();
+          expectTypeOf(runtime.context).toEqualTypeOf<{
+            customOptionalContextProp: string;
+          }>();
         },
         afterModel: async (_state, runtime) => {
-          expectTypeOf(runtime.context).toEqualTypeOf<
-            Partial<
-              | {
-                  customOptionalContextProp: string;
-                }
-              | undefined
-            >
-          >();
+          expectTypeOf(runtime.context).toEqualTypeOf<{
+            customOptionalContextProp: string;
+          }>();
         },
         wrapModelCall: async (request) => {
-          expectTypeOf(request.runtime.context).toEqualTypeOf<
-            Partial<
-              | {
-                  customOptionalContextProp: string;
-                }
-              | undefined
-            >
-          >();
+          expectTypeOf(request.runtime.context).toEqualTypeOf<{
+            customOptionalContextProp: string;
+          }>();
 
           return new AIMessage("foobar");
         },