From eae2321c8f578237d63ca402a217bd565806a753 Mon Sep 17 00:00:00 2001
From: Jacob Lee <jacoblee93@gmail.com>
Date: Wed, 28 Aug 2024 11:01:13 -0700
Subject: [PATCH] langgraph[patch]: Fix typing for conditional edge array route
 maps (#408)

---
 libs/langgraph/src/graph/graph.ts       |  2 +-
 libs/langgraph/src/tests/pregel.test.ts | 14 ++++----------
 2 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/libs/langgraph/src/graph/graph.ts b/libs/langgraph/src/graph/graph.ts
index 9bc2f4b75..63752f827 100644
--- a/libs/langgraph/src/graph/graph.ts
+++ b/libs/langgraph/src/graph/graph.ts
@@ -35,7 +35,7 @@ export const END = "__end__";
 export interface BranchOptions<IO, N extends string> {
   source: N;
   path: Branch<IO, N>["condition"];
-  pathMap?: Record<string, N | typeof END> | N[];
+  pathMap?: Record<string, N | typeof END> | (N | typeof END)[];
 }
 
 export class Branch<IO, N extends string> {
diff --git a/libs/langgraph/src/tests/pregel.test.ts b/libs/langgraph/src/tests/pregel.test.ts
index b624679cb..32d4bd415 100644
--- a/libs/langgraph/src/tests/pregel.test.ts
+++ b/libs/langgraph/src/tests/pregel.test.ts
@@ -2036,9 +2036,9 @@ describe("StateGraph", () => {
     data: typeof AgentAnnotation.State
   ): Promise<string> => {
     if (data.agentOutcome && "returnValues" in data.agentOutcome) {
-      return "exit";
+      return "__end__";
     }
-    return "continue";
+    return "tools";
   };
 
   it("can invoke", async () => {
@@ -2088,10 +2088,7 @@ describe("StateGraph", () => {
       .addNode("tools", executeTools)
       .addEdge(START, "agent")
       .addEdge("agent", "passthrough")
-      .addConditionalEdges("passthrough", shouldContinue, {
-        continue: "tools",
-        exit: END,
-      })
+      .addConditionalEdges("passthrough", shouldContinue, ["tools", "__end__"])
       .addEdge("tools", "agent")
       .compile();
 
@@ -2182,10 +2179,7 @@ describe("StateGraph", () => {
       .addNode("agent", agent)
       .addNode("tools", executeTools)
       .addEdge(START, "agent")
-      .addConditionalEdges("agent", shouldContinue, {
-        continue: "tools",
-        exit: END,
-      })
+      .addConditionalEdges("agent", shouldContinue, ["tools", "__end__"])
       .addEdge("tools", "agent")
       .compile();