diff --git a/libs/langgraph/src/graph/state.ts b/libs/langgraph/src/graph/state.ts index 541d76a2..78a76c1d 100644 --- a/libs/langgraph/src/graph/state.ts +++ b/libs/langgraph/src/graph/state.ts @@ -582,7 +582,7 @@ export class CompiledStateGraph< } else { const typeofInput = Array.isArray(input) ? "array" : typeof input; throw new InvalidUpdateError( - `Expected node "${nodeKey.toString()}" to return an object, received ${typeofInput}`, + `Expected node "${nodeKey.toString()}" to return an object or an array containing at least one Command object, received ${typeofInput}`, { lc_error_code: "INVALID_GRAPH_NODE_RETURN_VALUE", } @@ -809,13 +809,26 @@ function _controlBranch(value: any): (string | Send)[] { if (_isSend(value)) { return [value]; } - if (!isCommand(value)) { + const commands = []; + if (isCommand(value)) { + commands.push(value); + } else if (Array.isArray(value) && value.every(isCommand)) { + commands.push(...value); + } else { return []; } - if (value.graph === Command.PARENT) { - throw new ParentCommand(value); + const destinations: (string | Send)[] = []; + for (const command of commands) { + if (command.graph === Command.PARENT) { + throw new ParentCommand(command); + } + if (_isSend(command.goto) || typeof command.goto === "string") { + destinations.push(command.goto); + } else { + destinations.push(...command.goto); + } } - return Array.isArray(value.goto) ? value.goto : [value.goto]; + return destinations; } function _getControlBranch() { diff --git a/libs/langgraph/src/tests/pregel.test.ts b/libs/langgraph/src/tests/pregel.test.ts index 410d5e34..b9c51b09 100644 --- a/libs/langgraph/src/tests/pregel.test.ts +++ b/libs/langgraph/src/tests/pregel.test.ts @@ -2061,6 +2061,59 @@ graph TD; expect(await graph.invoke({ foo: "" })).toEqual({ foo: "a|c" }); }); + it("should support a simple edgeless graph", async () => { + const StateAnnotation = Annotation.Root({ + foo: Annotation, + bar: Annotation, + }); + + const nodeA = async (state: typeof StateAnnotation.State) => { + const goto = state.foo === "foo" ? "nodeB" : "nodeC"; + return [ + new Command({ + update: { + foo: "a", + }, + goto, + }), + ]; + }; + + const nodeB = async (state: typeof StateAnnotation.State) => { + return [ + { + foo: state.foo + "|b", + }, + new Command({ + update: { + bar: "test", + }, + }), + ]; + }; + + const nodeC = async (state: typeof StateAnnotation.State) => { + return { + foo: state.foo + "|c", + }; + }; + + const graph = new StateGraph(StateAnnotation) + .addNode("nodeA", nodeA, { + ends: ["nodeB", "nodeC"], + }) + .addNode("nodeB", nodeB) + .addNode("nodeC", nodeC) + .addEdge("__start__", "nodeA") + .compile(); + + expect(await graph.invoke({ foo: "foo" })).toEqual({ + foo: "a|b", + bar: "test", + }); + expect(await graph.invoke({ foo: "" })).toEqual({ foo: "a|c" }); + }); + it("should handle send sequences correctly", async () => { const StateAnnotation = Annotation.Root({ items: Annotation({