diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 1651b0f4f..3c899c985 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), @@ -517,6 +520,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]