From b916a2e18338b0a3d10fa591d30db128ca7f62c1 Mon Sep 17 00:00:00 2001 From: John Kennedy <65985482+jkennedyvz@users.noreply.github.com> Date: Sat, 28 Feb 2026 11:47:45 -0800 Subject: [PATCH] refactor: extract _track helper for available_channels sync in apply_writes Replace 5 repeated inline blocks that sync available_channels with a local _track() helper that checks is_available() and updates the set, returning the availability bool for callers that also need to update updated_channels. Co-Authored-By: Claude Opus 4.6 --- libs/langgraph/langgraph/pregel/_algo.py | 48 ++++++++++-------------- libs/langgraph/langgraph/pregel/_loop.py | 1 - 2 files changed, 20 insertions(+), 29 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index ba56765b0..f4538f537 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -282,6 +282,18 @@ def apply_writes( None, ) + # Sync available_channels with channel's actual availability state. + # Returns True if the channel is available (for callers that also need + # to update updated_channels). + def _track(chan: str) -> bool: + avail = channels[chan].is_available() + if available_channels is not None: + if avail: + available_channels.add(chan) + else: + available_channels.discard(chan) + return avail + # Consume all channels that were read for chan in { chan @@ -291,11 +303,7 @@ def apply_writes( }: if channels[chan].consume() and next_version is not None: checkpoint["channel_versions"][chan] = next_version - if available_channels is not None: - if channels[chan].is_available(): - available_channels.add(chan) - else: - available_channels.discard(chan) + _track(chan) # Group writes by channel pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) @@ -325,17 +333,10 @@ def apply_writes( if channels[chan].update(vals) and next_version is not None: checkpoint["channel_versions"][chan] = next_version # unavailable channels can't trigger tasks, so don't add them - if channels[chan].is_available(): + if _track(chan): updated_channels.add(chan) - if available_channels is not None: - available_channels.add(chan) - elif available_channels is not None: - available_channels.discard(chan) - elif available_channels is not None: - if channels[chan].is_available(): - available_channels.add(chan) - else: - available_channels.discard(chan) + else: + _track(chan) # Channels that weren't updated in this step are notified of a new step if bump_step: @@ -352,10 +353,8 @@ def apply_writes( if channels[chan].update(EMPTY_SEQ) and next_version is not None: checkpoint["channel_versions"][chan] = next_version # unavailable channels can't trigger tasks, so don't add them - if channels[chan].is_available(): + if _track(chan): updated_channels.add(chan) - elif available_channels is not None: - available_channels.discard(chan) # If this is (tentatively) the last superstep, notify all channels of finish if bump_step and updated_channels.isdisjoint(trigger_to_nodes): @@ -363,17 +362,10 @@ def apply_writes( if channels[chan].finish() and next_version is not None: checkpoint["channel_versions"][chan] = next_version # unavailable channels can't trigger tasks, so don't add them - if channels[chan].is_available(): + if _track(chan): updated_channels.add(chan) - if available_channels is not None: - available_channels.add(chan) - elif available_channels is not None: - available_channels.discard(chan) - elif available_channels is not None: - if channels[chan].is_available(): - available_channels.add(chan) - else: - available_channels.discard(chan) + else: + _track(chan) # Return managed values writes to be applied externally return updated_channels diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 2edda868c..c08dce785 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -198,7 +198,6 @@ class PregelLoop: _migrate_checkpoint: Callable[[Checkpoint], None] | None submit: Submit channels: Mapping[str, BaseChannel] - _has_untracked_channels: bool _available_channels: set[str] # Futures from `checkpointer.put_writes` calls that produced delta-channel # writes. `_checkpointer_put_after_previous` drains this list (swap to a