Missing State in ConditionalEdges when Using Input, Output, and State Schema Splitting

[shaharov]

langgraph version 0.2.8

I'm currently working on defining a graph with separate schemas for input, output, and additional state using StateGraph. Below is a simplified flow illustrating the issue:

START -> FirstNode -> [dummy-conditional-edge] -> LastNode -> END

The graph contains InputDefinition, OutputDefinition, and an additional state field in InnerNodesDefinition, all merged into GraphDefinition.

Expected console output:

import { END, StateGraph, Annotation, START } from "@langchain/langgraph";

export const InputDefinition = Annotation.Root({
	input: Annotation<string>,
});
export const OutputDefinition = Annotation.Root({
	output: Annotation<string>,
});

export const InnerNodesDefinition = Annotation.Root({
	innerString: Annotation<string>,
});
export const GraphDefinition = Annotation.Root({
	...InputDefinition.spec,
	...OutputDefinition.spec,
	...InnerNodesDefinition.spec,
});

const FirstNode = async (state: typeof GraphDefinition.State) => {
	console.log("FirstNode: ", state);
	return {
		innerString: "InnerNodesString",
	};
};
const route = async (state: typeof GraphDefinition.State) => {
	console.log("Route: ", state);
	return "LastNode";
};
const LastNode = async (state: typeof GraphDefinition.State) => {
	console.log("LastNode:", state);
	return {
		output: "bye",
	};
};
async function test() {
	const builder = new StateGraph({ input: InputDefinition, output: OutputDefinition, stateSchema: GraphDefinition })
		.addNode("FirstNode", FirstNode)
		.addNode("LastNode", LastNode)
		.addEdge(START, "FirstNode")
		.addConditionalEdges("FirstNode", route);
		.addEdge("LastNode", END)

	const graph = builder.compile();

	const res = await graph.invoke({ input: "InputString" });
	console.log("Response", res);
}

test();

I expect the console output to look like this:

FirstNode:  { input: 'InputString', output: null, innerString: null }
Route:  { input: 'InputString', innerString: 'InnerNodesString' }
LastNode: { input: 'InputString', output: null, innerString: 'InnerNodesString' }
Response { output: 'bye' }

However, the actual output is:

FirstNode:  { input: 'InputString', output: null, innerString: null }
Route:  {}
LastNode: { input: 'InputString', output: null, innerString: 'InnerNodesString' }
Response { output: 'bye' }

As shown above, the state in the route function is empty.

Interestingly, when I change the line:

const builder = new StateGraph({ input: InputDefinition, output: OutputDefinition, stateSchema: GraphDefinition });

to:

const builder = new StateGraph(GraphDefinition);

I get the expected output in the route node.

---

It seems like the state isn't passed correctly to the route function when splitting input, output, and state schemas. Could this be a bug, or am I missing something in the schema definition? Any insights would be appreciated!

[jacoblee93]

This appears to be a bug - will have a look!