One more test, handle negative timeout

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)

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")