@langchain/langgraph 0.2.72 → 0.2.74

package/dist/graph/messages_annotation.js

@@ -1,6 +1,8 @@
 /* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
+import { z } from "zod";
 import { Annotation } from "./annotation.js";
 import { messagesStateReducer } from "./message.js";
+import { withLangGraph } from "./zod/state.js";
 /**
  * Prebuilt state annotation that combines returned messages.
  * Can handle standard messages and special modifiers like {@link RemoveMessage}
@@ -42,4 +44,52 @@ export const MessagesAnnotation = Annotation.Root({
         default: () => [],
     }),
 });
+/**
+ * Prebuilt state object that uses Zod to combine returned messages.
+ * This utility is synonymous with the `MessagesAnnotation` annotation,
+ * but uses Zod as the way to express messages state.
+ *
+ * You can use import and use this prebuilt schema like this:
+ *
+ * @example
+ * ```ts
+ * import { MessagesZodState, StateGraph } from "@langchain/langgraph";
+ *
+ * const graph = new StateGraph(MessagesZodState)
+ *   .addNode(...)
+ *   ...
+ * ```
+ *
+ * Which is equivalent to initializing the schema object manually like this:
+ *
+ * @example
+ * ```ts
+ * import { z } from "zod";
+ * import type { BaseMessage, BaseMessageLike } from "@langchain/core/messages";
+ * import { StateGraph, messagesStateReducer } from "@langchain/langgraph";
+ * import "@langchain/langgraph/zod";
+ *
+ * const AgentState = z.object({
+ *   messages: z
+ *     .custom<BaseMessage[]>()
+ *     .default(() => [])
+ *     .langgraph.reducer(
+ *        messagesStateReducer,
+ *        z.custom<BaseMessageLike | BaseMessageLike[]>()
+ *     ),
+ * });
+ * const graph = new StateGraph(AgentState)
+ *   .addNode(...)
+ *   ...
+ * ```
+ */
+export const MessagesZodState = z.object({
+    messages: withLangGraph(z.custom(), {
+        reducer: {
+            schema: z.custom(),
+            fn: messagesStateReducer,
+        },
+        default: () => [],
+    }),
+});
 //# sourceMappingURL=messages_annotation.js.map

package/dist/graph/messages_annotation.js.map

@@ -1 +1 @@
-{"version":3,"file":"messages_annotation.js","sourceRoot":"","sources":["../../src/graph/messages_annotation.ts"],"names":[],"mappings":"AAAA,0CAA0C;AAG1C,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAC7C,OAAO,EAAY,oBAAoB,EAAE,MAAM,cAAc,CAAC;AAE9D;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAkCG;AACH,MAAM,CAAC,MAAM,kBAAkB,GAAG,UAAU,CAAC,IAAI,CAAC;IAChD,QAAQ,EAAE,UAAU,CAA0B;QAC5C,OAAO,EAAE,oBAAoB;QAC7B,OAAO,EAAE,GAAG,EAAE,CAAC,EAAE;KAClB,CAAC;CACH,CAAC,CAAC"}
+{"version":3,"file":"messages_annotation.js","sourceRoot":"","sources":["../../src/graph/messages_annotation.ts"],"names":[],"mappings":"AAAA,0CAA0C;AAG1C,OAAO,EAAE,CAAC,EAAE,MAAM,KAAK,CAAC;AACxB,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAC7C,OAAO,EAAY,oBAAoB,EAAE,MAAM,cAAc,CAAC;AAC9D,OAAO,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AAE/C;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAkCG;AACH,MAAM,CAAC,MAAM,kBAAkB,GAAG,UAAU,CAAC,IAAI,CAAC;IAChD,QAAQ,EAAE,UAAU,CAA0B;QAC5C,OAAO,EAAE,oBAAoB;QAC7B,OAAO,EAAE,GAAG,EAAE,CAAC,EAAE;KAClB,CAAC;CACH,CAAC,CAAC;AAEH;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAsCG;AACH,MAAM,CAAC,MAAM,gBAAgB,GAAG,CAAC,CAAC,MAAM,CAAC;IACvC,QAAQ,EAAE,aAAa,CAAC,CAAC,CAAC,MAAM,EAAiB,EAAE;QACjD,OAAO,EAAE;YACP,MAAM,EAAE,CAAC,CAAC,MAAM,EAAY;YAC5B,EAAE,EAAE,oBAAoB;SACzB;QACD,OAAO,EAAE,GAAG,EAAE,CAAC,EAAE;KAClB,CAAC;CACH,CAAC,CAAC"}

package/dist/graph/state.cjs

@@ -246,54 +246,76 @@ class StateGraph extends graph_js_1.Graph {
             }
         }
     }
-    addNode(key, action, options) {
-        if (key in this.channels) {
-            throw new Error(`${key} is already being used as a state attribute (a.k.a. a channel), cannot also be used as a node name.`);
-        }
-        for (const reservedChar of [
-            constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR,
-            constants_js_1.CHECKPOINT_NAMESPACE_END,
-        ]) {
-            if (key.includes(reservedChar)) {
-                throw new Error(`"${reservedChar}" is a reserved character and is not allowed in node names.`);
+    addNode(...args) {
+        function isMultipleNodes(args) {
+            return args.length >= 1 && typeof args[0] !== "string";
+        }
+        const nodes = (isMultipleNodes(args) // eslint-disable-line no-nested-ternary
+            ? Array.isArray(args[0])
+                ? args[0]
+                : Object.entries(args[0])
+            : [[args[0], args[1], args[2]]]);
+        if (nodes.length === 0) {
+            throw new Error("No nodes provided in `addNode`");
+        }
+        for (const [key, action, options] of nodes) {
+            if (key in this.channels) {
+                throw new Error(`${key} is already being used as a state attribute (a.k.a. a channel), cannot also be used as a node name.`);
             }
+            for (const reservedChar of [
+                constants_js_1.CHECKPOINT_NAMESPACE_SEPARATOR,
+                constants_js_1.CHECKPOINT_NAMESPACE_END,
+            ]) {
+                if (key.includes(reservedChar)) {
+                    throw new Error(`"${reservedChar}" is a reserved character and is not allowed in node names.`);
+                }
+            }
+            this.warnIfCompiled(`Adding a node to a graph that has already been compiled. This will not be reflected in the compiled graph.`);
+            if (key in this.nodes) {
+                throw new Error(`Node \`${key}\` already present.`);
+            }
+            if (key === constants_js_1.END || key === constants_js_1.START) {
+                throw new Error(`Node \`${key}\` is reserved.`);
+            }
+            let inputSpec = this._schemaDefinition;
+            if (options?.input !== undefined) {
+                if ((0, state_js_1.isAnyZodObject)(options.input)) {
+                    inputSpec = (0, state_js_1.getChannelsFromZod)(options.input);
+                }
+                else if (options.input.spec !== undefined) {
+                    inputSpec = options.input.spec;
+                }
+            }
+            if (inputSpec !== undefined) {
+                this._addSchema(inputSpec);
+            }
+            let runnable;
+            if (runnables_1.Runnable.isRunnable(action)) {
+                runnable = action;
+            }
+            else if (typeof action === "function") {
+                runnable = new utils_js_1.RunnableCallable({
+                    func: action,
+                    name: key,
+                    trace: false,
+                });
+            }
+            else {
+                runnable = (0, runnables_1._coerceToRunnable)(action);
+            }
+            const nodeSpec = {
+                runnable: runnable,
+                retryPolicy: options?.retryPolicy,
+                metadata: options?.metadata,
+                input: inputSpec ?? this._schemaDefinition,
+                subgraphs: (0, subgraph_js_1.isPregelLike)(runnable)
+                    ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                        [runnable]
+                    : options?.subgraphs,
+                ends: options?.ends,
+            };
+            this.nodes[key] = nodeSpec;
         }
-        this.warnIfCompiled(`Adding a node to a graph that has already been compiled. This will not be reflected in the compiled graph.`);
-        if (key in this.nodes) {
-            throw new Error(`Node \`${key}\` already present.`);
-        }
-        if (key === constants_js_1.END || key === constants_js_1.START) {
-            throw new Error(`Node \`${key}\` is reserved.`);
-        }
-        if (options?.input !== undefined) {
-            this._addSchema(options.input.spec);
-        }
-        let runnable;
-        if (runnables_1.Runnable.isRunnable(action)) {
-            runnable = action;
-        }
-        else if (typeof action === "function") {
-            runnable = new utils_js_1.RunnableCallable({
-                func: action,
-                name: key,
-                trace: false,
-            });
-        }
-        else {
-            runnable = (0, runnables_1._coerceToRunnable)(action);
-        }
-        const nodeSpec = {
-            runnable: runnable,
-            retryPolicy: options?.retryPolicy,
-            metadata: options?.metadata,
-            input: options?.input?.spec ?? this._schemaDefinition,
-            subgraphs: (0, subgraph_js_1.isPregelLike)(runnable)
-                ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                    [runnable]
-                : options?.subgraphs,
-            ends: options?.ends,
-        };
-        this.nodes[key] = nodeSpec;
         return this;
     }
     addEdge(startKey, endKey) {
@@ -321,6 +343,27 @@ class StateGraph extends graph_js_1.Graph {
         this.waitingEdges.add([startKey, endKey]);
         return this;
     }
+    addSequence(nodes) {
+        const parsedNodes = Array.isArray(nodes)
+            ? nodes
+            : Object.entries(nodes);
+        if (parsedNodes.length === 0) {
+            throw new Error("Sequence requires at least one node.");
+        }
+        let previousNode;
+        for (const [key, action, options] of parsedNodes) {
+            if (key in this.nodes) {
+                throw new Error(`Node names must be unique: node with the name "${key}" already exists.`);
+            }
+            const validKey = key;
+            this.addNode(validKey, action, options);
+            if (previousNode != null) {
+                this.addEdge(previousNode, validKey);
+            }
+            previousNode = validKey;
+        }
+        return this;
+    }
     compile({ checkpointer, store, interruptBefore, interruptAfter, name, } = {}) {
         // validate the graph
         this.validate([
@@ -507,15 +550,14 @@ class CompiledStateGraph extends graph_js_1.CompiledGraph {
             const inputDefinition = node?.input ?? this.builder._schemaDefinition;
             const inputValues = Object.fromEntries(Object.keys(this.builder._schemaDefinitions.get(inputDefinition)).map((k) => [k, k]));
             const isSingleInput = Object.keys(inputValues).length === 1 && ROOT in inputValues;
-            this.channels[key] = new ephemeral_value_js_1.EphemeralValue(false);
+            const branchChannel = `branch:to:${key}`;
+            this.channels[branchChannel] = new ephemeral_value_js_1.EphemeralValue(false);
             this.nodes[key] = new read_js_1.PregelNode({
-                triggers: [],
+                triggers: [branchChannel],
                 // read state keys
                 channels: isSingleInput ? Object.keys(inputValues) : inputValues,
-                // publish to this channel and state keys
-                writers: [
-                    new write_js_1.ChannelWrite(stateWriteEntries.concat({ channel: key, value: key }), [constants_js_1.TAG_HIDDEN]),
-                ],
+                // publish to state keys
+                writers: [new write_js_1.ChannelWrite(stateWriteEntries, [constants_js_1.TAG_HIDDEN])],
                 mapper: isSingleInput
                     ? undefined
                     : // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -531,10 +573,12 @@ class CompiledStateGraph extends graph_js_1.CompiledGraph {
         }
     }
     attachEdge(start, end) {
-        if (end === constants_js_1.END) {
+        if (end === constants_js_1.END)
             return;
+        if (typeof start === "string") {
+            this.nodes[start].writers.push(new write_js_1.ChannelWrite([{ channel: `branch:to:${end}`, value: null }], [constants_js_1.TAG_HIDDEN]));
         }
-        if (Array.isArray(start)) {
+        else if (Array.isArray(start)) {
             const channelName = `join:${start.join("+")}:${end}`;
             // register channel
             this.channels[channelName] =
@@ -546,34 +590,16 @@ class CompiledStateGraph extends graph_js_1.CompiledGraph {
                 this.nodes[s].writers.push(new write_js_1.ChannelWrite([{ channel: channelName, value: s }], [constants_js_1.TAG_HIDDEN]));
             }
         }
-        else if (start === constants_js_1.START) {
-            const channelName = `${constants_js_1.START}:${end}`;
-            // register channel
-            this.channels[channelName] =
-                new ephemeral_value_js_1.EphemeralValue();
-            // subscribe to channel
-            this.nodes[end].triggers.push(channelName);
-            // publish to channel
-            this.nodes[constants_js_1.START].writers.push(new write_js_1.ChannelWrite([{ channel: channelName, value: constants_js_1.START }], [constants_js_1.TAG_HIDDEN]));
-        }
-        else {
-            this.nodes[end].triggers.push(start);
-        }
     }
-    attachBranch(start, name, branch, options = { withReader: true }) {
+    attachBranch(start, _, branch, options = { withReader: true }) {
         const branchWriter = async (packets, config) => {
             const filteredPackets = packets.filter((p) => p !== constants_js_1.END);
-            if (!filteredPackets.length) {
+            if (!filteredPackets.length)
                 return;
-            }
             const writes = filteredPackets.map((p) => {
-                if ((0, constants_js_1._isSend)(p)) {
+                if ((0, constants_js_1._isSend)(p))
                     return p;
-                }
-                return {
-                    channel: `branch:${start}:${name}:${p}`,
-                    value: start,
-                };
+                return { channel: `branch:to:${p}`, value: start };
             });
             await write_js_1.ChannelWrite.doWrite({ ...config, tags: (config.tags ?? []).concat([constants_js_1.TAG_HIDDEN]) }, writes);
         };
@@ -583,19 +609,6 @@ class CompiledStateGraph extends graph_js_1.CompiledGraph {
         options.withReader
             ? (config) => read_js_1.ChannelRead.doRead(config, this.streamChannels ?? this.outputChannels, true)
             : undefined));
-        // attach branch subscribers
-        const ends = branch.ends
-            ? Object.values(branch.ends)
-            : Object.keys(this.builder.nodes);
-        for (const end of ends) {
-            if (end === constants_js_1.END) {
-                continue;
-            }
-            const channelName = `branch:${start}:${name}:${end}`;
-            this.channels[channelName] =
-                new ephemeral_value_js_1.EphemeralValue(false);
-            this.nodes[end].triggers.push(channelName);
-        }
     }
     async _validateInput(input) {
         const inputSchema = this.builder._inputRuntimeDefinition;

package/dist/graph/state.d.ts

@@ -24,7 +24,7 @@ export type StateGraphNodeSpec<RunInput, RunOutput> = NodeSpec<RunInput, RunOutp
 };
 export type StateGraphAddNodeOptions = {
     retryPolicy?: RetryPolicy;
-    input?: AnnotationRoot<any>;
+    input?: AnnotationRoot<any> | AnyZodObject;
 } & AddNodeOptions;
 export type StateGraphArgsWithStateSchema<SD extends StateDefinition, I extends StateDefinition, O extends StateDefinition> = {
     stateSchema: AnnotationRoot<SD>;
@@ -42,6 +42,8 @@ type ZodStateGraphArgsWithStateSchema<SD extends AnyZodObject, I extends SDZod,
 };
 type SDZod = StateDefinition | AnyZodObject;
 type ToStateDefinition<T> = T extends AnyZodObject ? ZodToStateDefinition<T> : T extends StateDefinition ? T : never;
+type NodeAction<S, U, C extends SDZod> = RunnableLike<S, U extends object ? U & Record<string, any> : U, // eslint-disable-line @typescript-eslint/no-explicit-any
+LangGraphRunnableConfig<StateType<ToStateDefinition<C>>>>;
 /**
  * A graph whose nodes communicate by reading and writing to a shared state.
  * Each node takes a defined `State` as input and returns a `Partial<State>`.
@@ -133,8 +135,19 @@ export declare class StateGraph<SD extends SDZod | unknown, S = SD extends SDZod
     constructor(fields: SD extends AnyZodObject ? SD | ZodStateGraphArgsWithStateSchema<SD, I, O> : never, configSchema?: C | AnnotationRoot<ToStateDefinition<C>>);
     get allEdges(): Set<[string, string]>;
     _addSchema(stateDefinition: SDZod): void;
-    addNode<K extends string, NodeInput = S>(key: K, action: RunnableLike<NodeInput, U extends object ? U & Record<string, any> : U, LangGraphRunnableConfig<StateType<ToStateDefinition<C>>>>, options?: StateGraphAddNodeOptions): StateGraph<SD, S, U, N | K, I, O, C>;
+    addNode<K extends string>(nodes: Record<K, NodeAction<S, U, C>> | [
+        key: K,
+        action: NodeAction<S, U, C>,
+        options?: StateGraphAddNodeOptions
+    ][]): StateGraph<SD, S, U, N | K, I, O, C>;
+    addNode<K extends string, NodeInput = S>(key: K, action: NodeAction<NodeInput, U, C>, options?: StateGraphAddNodeOptions): StateGraph<SD, S, U, N | K, I, O, C>;
     addEdge(startKey: typeof START | N | N[], endKey: N | typeof END): this;
+    addSequence<K extends string>(nodes: [
+        key: K,
+        action: NodeAction<S, U, C>,
+        options?: StateGraphAddNodeOptions
+    ][]): StateGraph<SD, S, U, N | K, I, O, C>;
+    addSequence<K extends string>(nodes: Record<K, NodeAction<S, U, C>>): StateGraph<SD, S, U, N | K, I, O, C>;
     compile({ checkpointer, store, interruptBefore, interruptAfter, name, }?: {
         checkpointer?: BaseCheckpointSaver | false;
         store?: BaseStore;
@@ -153,7 +166,7 @@ export declare class CompiledStateGraph<S, U, N extends string = typeof START, I
     attachNode(key: typeof START, node?: never): void;
     attachNode(key: N, node: StateGraphNodeSpec<S, U>): void;
     attachEdge(start: N | N[] | "__start__", end: N | "__end__"): void;
-    attachBranch(start: N | typeof START, name: string, branch: Branch<S, N>, options?: {
+    attachBranch(start: N | typeof START, _: string, branch: Branch<S, N>, options?: {
         withReader?: boolean;
     }): void;
     protected _validateInput(input: UpdateType<ToStateDefinition<I>>): Promise<UpdateType<ToStateDefinition<I>>>;

package/dist/graph/state.js

@@ -243,54 +243,76 @@ export class StateGraph extends Graph {
             }
         }
     }
-    addNode(key, action, options) {
-        if (key in this.channels) {
-            throw new Error(`${key} is already being used as a state attribute (a.k.a. a channel), cannot also be used as a node name.`);
-        }
-        for (const reservedChar of [
-            CHECKPOINT_NAMESPACE_SEPARATOR,
-            CHECKPOINT_NAMESPACE_END,
-        ]) {
-            if (key.includes(reservedChar)) {
-                throw new Error(`"${reservedChar}" is a reserved character and is not allowed in node names.`);
+    addNode(...args) {
+        function isMultipleNodes(args) {
+            return args.length >= 1 && typeof args[0] !== "string";
+        }
+        const nodes = (isMultipleNodes(args) // eslint-disable-line no-nested-ternary
+            ? Array.isArray(args[0])
+                ? args[0]
+                : Object.entries(args[0])
+            : [[args[0], args[1], args[2]]]);
+        if (nodes.length === 0) {
+            throw new Error("No nodes provided in `addNode`");
+        }
+        for (const [key, action, options] of nodes) {
+            if (key in this.channels) {
+                throw new Error(`${key} is already being used as a state attribute (a.k.a. a channel), cannot also be used as a node name.`);
             }
+            for (const reservedChar of [
+                CHECKPOINT_NAMESPACE_SEPARATOR,
+                CHECKPOINT_NAMESPACE_END,
+            ]) {
+                if (key.includes(reservedChar)) {
+                    throw new Error(`"${reservedChar}" is a reserved character and is not allowed in node names.`);
+                }
+            }
+            this.warnIfCompiled(`Adding a node to a graph that has already been compiled. This will not be reflected in the compiled graph.`);
+            if (key in this.nodes) {
+                throw new Error(`Node \`${key}\` already present.`);
+            }
+            if (key === END || key === START) {
+                throw new Error(`Node \`${key}\` is reserved.`);
+            }
+            let inputSpec = this._schemaDefinition;
+            if (options?.input !== undefined) {
+                if (isAnyZodObject(options.input)) {
+                    inputSpec = getChannelsFromZod(options.input);
+                }
+                else if (options.input.spec !== undefined) {
+                    inputSpec = options.input.spec;
+                }
+            }
+            if (inputSpec !== undefined) {
+                this._addSchema(inputSpec);
+            }
+            let runnable;
+            if (Runnable.isRunnable(action)) {
+                runnable = action;
+            }
+            else if (typeof action === "function") {
+                runnable = new RunnableCallable({
+                    func: action,
+                    name: key,
+                    trace: false,
+                });
+            }
+            else {
+                runnable = _coerceToRunnable(action);
+            }
+            const nodeSpec = {
+                runnable: runnable,
+                retryPolicy: options?.retryPolicy,
+                metadata: options?.metadata,
+                input: inputSpec ?? this._schemaDefinition,
+                subgraphs: isPregelLike(runnable)
+                    ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                        [runnable]
+                    : options?.subgraphs,
+                ends: options?.ends,
+            };
+            this.nodes[key] = nodeSpec;
         }
-        this.warnIfCompiled(`Adding a node to a graph that has already been compiled. This will not be reflected in the compiled graph.`);
-        if (key in this.nodes) {
-            throw new Error(`Node \`${key}\` already present.`);
-        }
-        if (key === END || key === START) {
-            throw new Error(`Node \`${key}\` is reserved.`);
-        }
-        if (options?.input !== undefined) {
-            this._addSchema(options.input.spec);
-        }
-        let runnable;
-        if (Runnable.isRunnable(action)) {
-            runnable = action;
-        }
-        else if (typeof action === "function") {
-            runnable = new RunnableCallable({
-                func: action,
-                name: key,
-                trace: false,
-            });
-        }
-        else {
-            runnable = _coerceToRunnable(action);
-        }
-        const nodeSpec = {
-            runnable: runnable,
-            retryPolicy: options?.retryPolicy,
-            metadata: options?.metadata,
-            input: options?.input?.spec ?? this._schemaDefinition,
-            subgraphs: isPregelLike(runnable)
-                ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                    [runnable]
-                : options?.subgraphs,
-            ends: options?.ends,
-        };
-        this.nodes[key] = nodeSpec;
         return this;
     }
     addEdge(startKey, endKey) {
@@ -318,6 +340,27 @@ export class StateGraph extends Graph {
         this.waitingEdges.add([startKey, endKey]);
         return this;
     }
+    addSequence(nodes) {
+        const parsedNodes = Array.isArray(nodes)
+            ? nodes
+            : Object.entries(nodes);
+        if (parsedNodes.length === 0) {
+            throw new Error("Sequence requires at least one node.");
+        }
+        let previousNode;
+        for (const [key, action, options] of parsedNodes) {
+            if (key in this.nodes) {
+                throw new Error(`Node names must be unique: node with the name "${key}" already exists.`);
+            }
+            const validKey = key;
+            this.addNode(validKey, action, options);
+            if (previousNode != null) {
+                this.addEdge(previousNode, validKey);
+            }
+            previousNode = validKey;
+        }
+        return this;
+    }
     compile({ checkpointer, store, interruptBefore, interruptAfter, name, } = {}) {
         // validate the graph
         this.validate([
@@ -503,15 +546,14 @@ export class CompiledStateGraph extends CompiledGraph {
             const inputDefinition = node?.input ?? this.builder._schemaDefinition;
             const inputValues = Object.fromEntries(Object.keys(this.builder._schemaDefinitions.get(inputDefinition)).map((k) => [k, k]));
             const isSingleInput = Object.keys(inputValues).length === 1 && ROOT in inputValues;
-            this.channels[key] = new EphemeralValue(false);
+            const branchChannel = `branch:to:${key}`;
+            this.channels[branchChannel] = new EphemeralValue(false);
             this.nodes[key] = new PregelNode({
-                triggers: [],
+                triggers: [branchChannel],
                 // read state keys
                 channels: isSingleInput ? Object.keys(inputValues) : inputValues,
-                // publish to this channel and state keys
-                writers: [
-                    new ChannelWrite(stateWriteEntries.concat({ channel: key, value: key }), [TAG_HIDDEN]),
-                ],
+                // publish to state keys
+                writers: [new ChannelWrite(stateWriteEntries, [TAG_HIDDEN])],
                 mapper: isSingleInput
                     ? undefined
                     : // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -527,10 +569,12 @@ export class CompiledStateGraph extends CompiledGraph {
         }
     }
     attachEdge(start, end) {
-        if (end === END) {
+        if (end === END)
             return;
+        if (typeof start === "string") {
+            this.nodes[start].writers.push(new ChannelWrite([{ channel: `branch:to:${end}`, value: null }], [TAG_HIDDEN]));
         }
-        if (Array.isArray(start)) {
+        else if (Array.isArray(start)) {
             const channelName = `join:${start.join("+")}:${end}`;
             // register channel
             this.channels[channelName] =
@@ -542,34 +586,16 @@ export class CompiledStateGraph extends CompiledGraph {
                 this.nodes[s].writers.push(new ChannelWrite([{ channel: channelName, value: s }], [TAG_HIDDEN]));
             }
         }
-        else if (start === START) {
-            const channelName = `${START}:${end}`;
-            // register channel
-            this.channels[channelName] =
-                new EphemeralValue();
-            // subscribe to channel
-            this.nodes[end].triggers.push(channelName);
-            // publish to channel
-            this.nodes[START].writers.push(new ChannelWrite([{ channel: channelName, value: START }], [TAG_HIDDEN]));
-        }
-        else {
-            this.nodes[end].triggers.push(start);
-        }
     }
-    attachBranch(start, name, branch, options = { withReader: true }) {
+    attachBranch(start, _, branch, options = { withReader: true }) {
         const branchWriter = async (packets, config) => {
             const filteredPackets = packets.filter((p) => p !== END);
-            if (!filteredPackets.length) {
+            if (!filteredPackets.length)
                 return;
-            }
             const writes = filteredPackets.map((p) => {
-                if (_isSend(p)) {
+                if (_isSend(p))
                     return p;
-                }
-                return {
-                    channel: `branch:${start}:${name}:${p}`,
-                    value: start,
-                };
+                return { channel: `branch:to:${p}`, value: start };
             });
             await ChannelWrite.doWrite({ ...config, tags: (config.tags ?? []).concat([TAG_HIDDEN]) }, writes);
         };
@@ -579,19 +605,6 @@ export class CompiledStateGraph extends CompiledGraph {
         options.withReader
             ? (config) => ChannelRead.doRead(config, this.streamChannels ?? this.outputChannels, true)
             : undefined));
-        // attach branch subscribers
-        const ends = branch.ends
-            ? Object.values(branch.ends)
-            : Object.keys(this.builder.nodes);
-        for (const end of ends) {
-            if (end === END) {
-                continue;
-            }
-            const channelName = `branch:${start}:${name}:${end}`;
-            this.channels[channelName] =
-                new EphemeralValue(false);
-            this.nodes[end].triggers.push(channelName);
-        }
     }
     async _validateInput(input) {
         const inputSchema = this.builder._inputRuntimeDefinition;