diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index e0d2a31d6..ca7f787c3 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -915,7 +915,9 @@ class Pregel( done, inflight = concurrent.futures.wait( futures, return_when=concurrent.futures.FIRST_COMPLETED, - timeout=end_time - time.monotonic() if end_time else None, + timeout=max(0, end_time - time.monotonic()) + if end_time + else None, ) for fut in done: task = futures.pop(fut) @@ -1050,6 +1052,7 @@ class Pregel( None, ) try: + loop = asyncio.get_event_loop() bg: list[asyncio.Task] = [] if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") @@ -1231,15 +1234,15 @@ class Pregel( for task in next_tasks } end_time = ( - self.step_timeout + time.monotonic() - if self.step_timeout - else None + self.step_timeout + loop.time() if self.step_timeout else None ) while futures: done, inflight = await asyncio.wait( futures, return_when=asyncio.FIRST_COMPLETED, - timeout=end_time - time.monotonic() if end_time else None, + timeout=max(0, end_time - loop.time()) + if end_time + else None, ) for fut in done: task = futures.pop(fut) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 346f1ee88..944a6bef8 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -105,6 +105,36 @@ async def test_node_cancellation_on_other_node_exception() -> None: assert inner_task_cancelled +async def test_step_timeout_on_stream_hang() -> None: + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + try: + await asyncio.sleep(1.5) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + + async def alittlewhile(input: Any) -> None: + await asyncio.sleep(0.6) + return "1" + + builder = Graph() + builder.add_node(awhile) + builder.add_node(alittlewhile) + builder.set_conditional_entry_point(lambda _: ["awhile", "alittlewhile"], then=END) + graph = builder.compile() + graph.step_timeout = 1 + + with pytest.raises(asyncio.CancelledError): + async for chunk in graph.astream(1, stream_mode="updates"): + assert chunk == {"alittlewhile": {"alittlewhile": "1"}} + await asyncio.sleep(0.6) + + assert inner_task_cancelled + + async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")