feat(langgraph): add type bag pattern for GraphNode and ConditionalEdgeRouter (#1918)

Author: hntrl

.changeset/type-bag-pattern.md

@@ -0,0 +1,36 @@
+---
+"@langchain/langgraph": patch
+---
+
+Add type bag pattern for `GraphNode` and `ConditionalEdgeRouter` type utilities.
+
+**New types:**
+- `GraphNodeTypes<InputSchema, OutputSchema, ContextSchema, Nodes>` - Type bag interface for GraphNode
+- `GraphNodeReturnValue<Update, Nodes>` - Return type helper for node functions
+- `ConditionalEdgeRouterTypes<InputSchema, ContextSchema, Nodes>` - Type bag interface for ConditionalEdgeRouter
+
+**Usage:**
+
+Both `GraphNode` and `ConditionalEdgeRouter` now support two patterns:
+
+1. **Single schema** (backward compatible):
+   ```typescript
+   const node: GraphNode<typeof AgentState, MyContext, "agent" | "tool"> = ...
+   ```
+
+2. **Type bag pattern** (new):
+   ```typescript
+   const node: GraphNode<{
+     InputSchema: typeof InputSchema;
+     OutputSchema: typeof OutputSchema;
+     ContextSchema: typeof ContextSchema;
+     Nodes: "agent" | "tool";
+   }> = (state, runtime) => {
+     // state type inferred from InputSchema
+     // return type validated against OutputSchema
+     // runtime.configurable type inferred from ContextSchema
+     return { answer: "response" };
+   };
+   ```
+
+The type bag pattern enables nodes that receive a subset of state fields and return different fields, with full type safety.

libs/langgraph-core/package.json

@@ -35,6 +35,7 @@
   "dependencies": {
     "@langchain/langgraph-checkpoint": "workspace:^",
     "@langchain/langgraph-sdk": "workspace:~",
+    "@standard-schema/spec": "1.1.0",
     "uuid": "^10.0.0"
   },
   "peerDependencies": {
@@ -55,7 +56,6 @@
     "@langchain/openai": "^1.0.0",
     "@langchain/scripts": ">=0.1.3 <0.2.0",
     "@langchain/tavily": "^1.0.0",
-    "@standard-schema/spec": "1.1.0",
     "@swc/core": "^1.3.90",
     "@testing-library/dom": "^10.4.0",
     "@tsconfig/recommended": "^1.0.3",
@@ -191,4 +191,4 @@
   "files": [
     "dist/"
   ]
-}
+}

libs/langgraph-core/src/graph/index.ts

@@ -28,5 +28,8 @@ export type {
   ExtractStateType,
   ExtractUpdateType,
   GraphNode,
+  GraphNodeTypes,
+  GraphNodeReturnValue,
   ConditionalEdgeRouter,
+  ConditionalEdgeRouterTypes,
 } from "./types.js";

libs/langgraph-core/src/graph/state.ts

@@ -416,8 +416,7 @@ export class StateGraph<
       N,
       InterruptType,
       WriterType
-    >,
-    contextSchema?: C | AnnotationRoot<ToStateDefinition<C>>
+    >
   );
 
   /** @deprecated Use `Annotation.Root`, `StateSchema`, or Zod schemas instead. */
@@ -434,7 +433,14 @@ export class StateGraph<
     options?:
       | C
       | AnnotationRoot<ToStateDefinition<C>>
-      | StateGraphOptions<I, O, C, N, InterruptType, WriterType>
+      | StateGraphOptions<
+          I,
+          O,
+          C extends ContextSchemaInit ? C : undefined,
+          N,
+          InterruptType,
+          WriterType
+        >
   ) {
     super();
 
@@ -523,13 +529,14 @@ export class StateGraph<
   ): StateGraphInit<StateDefinitionInit, I, O, C> {
     // Check if already StateGraphInit format
     if (isStateGraphInit(stateOrInit)) {
-      // Merge any 2nd arg options
+      // Second arg can be either a direct context schema or an options object
       if (isInteropZodObject(options) || AnnotationRoot.isInstance(options)) {
         return {
           ...stateOrInit,
           context: options as C,
         };
       }
+      // Merge any 2nd arg options
       const opts = options as StateGraphOptions<I, O> | undefined;
       return {
         ...stateOrInit,
@@ -547,7 +554,7 @@ export class StateGraph<
       // Second arg can be either a direct context schema or an options object
       if (isInteropZodObject(options) || AnnotationRoot.isInstance(options)) {
         return {
-          state: stateOrInit as StateDefinitionInit,
+          state: stateOrInit,
           context: options as C,
         };
       }
@@ -725,6 +732,31 @@ export class StateGraph<
     MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
   >;
 
+  override addNode<
+    K extends string,
+    InputSchema extends StateDefinitionInit,
+    NodeOutput extends U = U
+  >(
+    key: K,
+    action: NodeAction<
+      ExtractStateType<InputSchema>,
+      NodeOutput,
+      C,
+      InterruptType,
+      WriterType
+    >,
+    options: StateGraphAddNodeOptions<N | K, InputSchema>
+  ): StateGraph<
+    SD,
+    S,
+    U,
+    N | K,
+    I,
+    O,
+    C,
+    MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
+  >;
+
   override addNode<K extends string, NodeInput = S, NodeOutput extends U = U>(
     key: K,
     action: NodeAction<NodeInput, NodeOutput, C, InterruptType, WriterType>,