From 0aafa04bacac8931f2735493859b4d598ba26e2e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 18:31:15 -0700 Subject: [PATCH 1/5] WIP Fix missing interrupts in stream - When multiple parallel tasks and/or subgraphs emit interrupts some were missing from stream output --- libs/langgraph/langgraph/pregel/debug.py | 3 ++- libs/langgraph/langgraph/pregel/io.py | 2 +- libs/langgraph/langgraph/pregel/loop.py | 33 ++++++++++++++++++----- libs/langgraph/langgraph/pregel/runner.py | 17 +++++++++--- 4 files changed, 43 insertions(+), 12 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 54c260682..7bd9e634e 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -293,8 +293,9 @@ def tasks_w_writes( ), tuple( v - for tid, n, v in pending_writes + for tid, n, vv in pending_writes if tid == task.id and n == INTERRUPT + for v in (vv if isinstance(vv, Sequence) else [vv]) ), states.get(task.id) if states else None, ( diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index e9963f5c8..f07745f1b 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -28,7 +28,7 @@ def is_task_id(task_id: str) -> bool: """Check if a string is a valid task id.""" try: UUID(task_id) - except ValueError: + except Exception: return False return True diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 66f0a60f5..0c007949e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -794,11 +794,14 @@ class PregelLoop(LoopProtocol): [w for t in self.tasks.values() for w in t.writes], self.channels, ) - # emit INTERRUPT event - self._emit( - "updates", - lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]), - ) + # emit INTERRUPT if exception is empty (otherwise emitted by put_writes) + if not exc_value.args or not exc_value.args[0]: + self._emit( + "updates", + lambda: iter( + [{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}] + ), + ) # save final output self.output = read_channels(self.channels, self.output_keys) # suppress interrupt @@ -829,7 +832,25 @@ class PregelLoop(LoopProtocol): "tags", EMPTY_SEQ ): return - if writes[0][0] != ERROR and writes[0][0] != INTERRUPT: + if writes[0][0] == INTERRUPT: + self._emit( + "updates", + lambda: iter( + [ + { + INTERRUPT: tuple( + v + for w in writes + if w[0] == INTERRUPT + for v in ( + w[1] if isinstance(w[1], Sequence) else (w[1],) + ) + ) + } + ] + ), + ) + elif writes[0][0] != ERROR: self._emit( "updates", map_output_updates, diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6336bc5a1..e54640ce4 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -543,10 +543,11 @@ class PregelRunner: elif exception: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer - if interrupts := [(INTERRUPT, i) for i in exception.args[0]]: + if exception.args[0]: + writes = [(INTERRUPT, exception.args[0])] if resumes := [w for w in task.writes if w[0] == RESUME]: - interrupts.extend(resumes) - self.put_writes(task.id, interrupts) + writes.extend(resumes) + self.put_writes(task.id, writes) elif isinstance(exception, GraphBubbleUp): raise exception else: @@ -608,6 +609,7 @@ def _panic_or_proceed( done.add(fut) else: inflight.add(fut) + interrupts: list[GraphInterrupt] = [] while done: # if any task failed if exc := _exception(done.pop()): @@ -616,7 +618,14 @@ def _panic_or_proceed( inflight.pop().cancel() # raise the exception if panic: - raise exc + if isinstance(exc, GraphInterrupt): + # collect interrupts + interrupts.append(exc) + else: + raise exc + # raise combined interrupts + if interrupts: + raise GraphInterrupt(tuple(i for exc in interrupts for i in exc.args[0])) if inflight: # if we got here means we timed out while inflight: From be7dee1c3b1f6f57893e5b5274cb3c985f99e489 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 18:35:53 -0700 Subject: [PATCH 2/5] Fix --- libs/langgraph/langgraph/pregel/debug.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 7bd9e634e..dc2e403a7 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -137,7 +137,12 @@ def map_debug_task_results( "result": [ w for w in writes if w[0] in stream_channels_list or w[0] == RETURN ], - "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT], + "interrupts": [ + asdict(v) + for w in writes + if w[0] == INTERRUPT + for v in (w[1] if isinstance(w[1], Sequence) else [w[1]]) + ], }, } From 59e7c63c93daf12cfb42fc85bb7783c3991dd109 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 18:38:37 -0700 Subject: [PATCH 3/5] Lint --- libs/langgraph/langgraph/pregel/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 0c007949e..21f297b9e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -795,7 +795,7 @@ class PregelLoop(LoopProtocol): self.channels, ) # emit INTERRUPT if exception is empty (otherwise emitted by put_writes) - if not exc_value.args or not exc_value.args[0]: + if exc_value is not None and not exc_value.args or not exc_value.args[0]: self._emit( "updates", lambda: iter( From e28e97d5e0ae41ea468a7597668b95cfc1249b31 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 18:39:13 -0700 Subject: [PATCH 4/5] Lint --- libs/langgraph/langgraph/pregel/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 21f297b9e..dec608b62 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -795,7 +795,7 @@ class PregelLoop(LoopProtocol): self.channels, ) # emit INTERRUPT if exception is empty (otherwise emitted by put_writes) - if exc_value is not None and not exc_value.args or not exc_value.args[0]: + if exc_value is not None and (not exc_value.args or not exc_value.args[0]): self._emit( "updates", lambda: iter( From 576aa1ca02951495fb588c1edcec0d9ddabaf62d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 18:41:36 -0700 Subject: [PATCH 5/5] Order --- libs/langgraph/tests/test_pregel_async.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 39d3d34a6..e97b8e164 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -939,9 +939,6 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: {"my_key": "value ⛰️", "market": "DE"}, thread2 ) ] == [ - { - "tool_one": {"my_key": " one"}, - }, { "__interrupt__": ( Interrupt( @@ -951,6 +948,9 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: ), ) }, + { + "tool_one": {"my_key": " one"}, + }, ] # resume with answer assert [