From d48b25420ba7553a7154d3960fb1bf327d497e9d Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 20 Jan 2025 11:40:54 -0800
Subject: [PATCH] Fix timing issue where a sync task would finish before the
 other one was registered in futures dict

- this was not possible in async where all done callbacks are called in next tick
- in sync case this would manifest as the first task done callback seeing counter == 1 and thus setting event
- the fix is to unset the event whenever a task is scheduled
---
 libs/langgraph/langgraph/pregel/runner.py |  5 +++++
 libs/langgraph/tests/test_large_cases.py  | 11 +++++++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py
index af80ed715..14e0303e1 100644
--- a/libs/langgraph/langgraph/pregel/runner.py
+++ b/libs/langgraph/langgraph/pregel/runner.py
@@ -75,6 +75,7 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]):
         super().__setitem__(key, value)  # type: ignore[index]
         if value is not None:
             with self.lock:
+                self.event.clear()
                 self.counter += 1
             key.add_done_callback(partial(self.on_done, value))
 
@@ -296,6 +297,8 @@ class PregelRunner:
         futures.event.wait(
             timeout=(max(0, end_time - time.monotonic()) if end_time else None)
         )
+        # give control back to the caller
+        yield
         # panic on failure or timeout
         _panic_or_proceed(
             futures.done.union(f for f, t in futures.items() if t is not None),
@@ -518,6 +521,8 @@ class PregelRunner:
             futures.event.wait(),
             timeout=(max(0, end_time - loop.time()) if end_time else None),
         )
+        # give control back to the caller
+        yield
         # cancel waiter task
         for fut in futures:
             fut.cancel()
diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py
index 4d2c45234..6980b7b68 100644
--- a/libs/langgraph/tests/test_large_cases.py
+++ b/libs/langgraph/tests/test_large_cases.py
@@ -2829,9 +2829,9 @@ def test_state_graph_packets(
     # Define decision-making logic
     def should_continue(data: AgentState) -> str:
         assert isinstance(data["session"], httpx.Client)
-        assert (
-            data["something_extra"] == "hi there"
-        ), "nodes can pass extra data to their cond edges, which isn't saved in state"
+        assert data["something_extra"] == "hi there", (
+            "nodes can pass extra data to their cond edges, which isn't saved in state"
+        )
         # Logic to decide whether to continue in the loop or exit
         if tool_calls := data["messages"][-1].tool_calls:
             return [Send("tools", tool_call) for tool_call in tool_calls]
@@ -5425,7 +5425,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
         docs: Annotated[list[str], sorted_add]
 
     def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def retriever_one(data: State) -> State:
         # timer ensures stream output order is stable
@@ -7277,6 +7277,9 @@ def test_send_dedupe_on_resume(
             setattr(self, "__name__", name)
 
         def __call__(self, state):
+            time.sleep(0)
+            # sleep makes it more likely to trigger edge case where 1st task
+            # finishes before 2nd is registered in futures dict
             self.ticks += 1
             update = (
                 [self.name]