Compare commits

..
Author SHA1 Message Date
Elior Nataf Lackritz a8e732c879 fix(langgraph): give each bulk_update_state update its own task id
An update whose node has no pending task to reuse was stored under
uuid5(checkpoint_id, INTERRUPT), so every such update in one superstep
shared a task id. Savers keep one write per (task_id, idx), so all but the
first update's writes were dropped. Plain channels were unaffected, since
their value is stored in the new checkpoint, but a DeltaChannel replays
those writes and lost every update after the first.

The ith update now gets uuid5(checkpoint_id, f"{INTERRUPT}:{i}"). The first
keeps the old id, so a single update stores exactly what it did before.
2026-09-30 12:47:55 -04:00
14 changed files with 147 additions and 174 deletions
+3 -3
View File
@@ -943,11 +943,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+3 -3
View File
@@ -1122,11 +1122,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+3 -3
View File
@@ -1057,11 +1057,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+3 -3
View File
@@ -1338,11 +1338,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+3 -3
View File
@@ -732,11 +732,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+3 -3
View File
@@ -672,11 +672,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+6 -6
View File
@@ -1777,11 +1777,11 @@ wheels = [
[[package]] [[package]]
name = "pyjwt" name = "pyjwt"
version = "2.15.0" version = "2.13.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/02/a5/5197bfd06417837ac079921c66fa6393f1dea3557272a263cebfef69e432/pyjwt-2.15.0.tar.gz", hash = "sha256:b11c5f9791d7bf51c2b39a81ed669f6b2dbbd669df2942f6c60167e9e3d1abe4", size = 120513, upload-time = "2026-09-23T16:56:00.689Z" } sdist = { url = "https://files.pythonhosted.org/packages/3b/81/58d0ac84e1ef3a3843791d6954d94c0b33d526c75eeb1efbce9d0a4c4077/pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423", size = 107515, upload-time = "2026-05-21T19:54:36.618Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/e8/55/40e45bf052ee8ee12a4dfd785519660f8effa7b065442b91646ec6828619/pyjwt-2.15.0-py3-none-any.whl", hash = "sha256:7a3742debf6b879e912dbb9819ceec1594be812452b78c5f2e2dfc56564954f8", size = 33680, upload-time = "2026-09-23T16:55:59.241Z" }, { url = "https://files.pythonhosted.org/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728", size = 31274, upload-time = "2026-05-21T19:54:35.362Z" },
] ]
[[package]] [[package]]
@@ -2243,11 +2243,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
@@ -114,26 +114,6 @@ def create_metadata_for_update_state_api(
return new_counters return new_counters
def advance_delta_counters(
channels: Mapping[str, BaseChannel],
updated_channels: set[str],
*,
prev_metadata: Mapping[str, Any] | None,
) -> dict[str, Any]:
"""The `counters_since_delta_snapshot` entry for an update_state
checkpoint saved one superstep after `prev_metadata`'s, for the paths
that skip `create_checkpoint_plan_for_update_state_api`.
Without it, the next checkpoint restarts every delta channel's snapshot
cadence from zero.
"""
counters = create_metadata_for_update_state_api(
channels, updated_channels, prev_metadata=prev_metadata
)
non_zero = {k: v for k, v in counters.items() if v != (0, 0)}
return {"counters_since_delta_snapshot": non_zero} if non_zero else {}
def create_checkpoint_plan_for_update_state_api( def create_checkpoint_plan_for_update_state_api(
channels: Mapping[str, BaseChannel], channels: Mapping[str, BaseChannel],
updated_channels: set[str], updated_channels: set[str],
+20 -57
View File
@@ -129,7 +129,6 @@ from langgraph.pregel._algo import (
from langgraph.pregel._call import identifier from langgraph.pregel._call import identifier
from langgraph.pregel._checkpoint import ( from langgraph.pregel._checkpoint import (
achannels_from_checkpoint, achannels_from_checkpoint,
advance_delta_counters,
channels_from_checkpoint, channels_from_checkpoint,
copy_checkpoint, copy_checkpoint,
create_checkpoint, create_checkpoint,
@@ -1682,7 +1681,6 @@ class Pregel(
"Cannot apply multiple updates when clearing state" "Cannot apply multiple updates when clearing state"
) )
updated_channels: set[str] = set()
if saved is not None: if saved is not None:
# tasks for this checkpoint # tasks for this checkpoint
next_tasks = prepare_next_tasks( next_tasks = prepare_next_tasks(
@@ -1705,7 +1703,7 @@ class Pregel(
for w in saved.pending_writes or [] for w in saved.pending_writes or []
if w[0] == NULL_TASK_ID if w[0] == NULL_TASK_ID
]: ]:
updated_channels |= apply_writes( apply_writes(
checkpoint, checkpoint,
channels, channels,
[PregelTaskWrites((), INPUT, null_writes, [])], [PregelTaskWrites((), INPUT, null_writes, [])],
@@ -1720,7 +1718,7 @@ class Pregel(
continue continue
next_tasks[tid].writes.append((k, v)) next_tasks[tid].writes.append((k, v))
# clear all current tasks # clear all current tasks
updated_channels |= apply_writes( apply_writes(
checkpoint, checkpoint,
channels, channels,
next_tasks.values(), next_tasks.values(),
@@ -1735,13 +1733,6 @@ class Pregel(
"source": "update", "source": "update",
"step": step + 1, "step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {}, "parents": saved.metadata.get("parents", {}) if saved else {},
**(
advance_delta_counters(
channels, updated_channels, prev_metadata=saved.metadata
)
if saved
else {}
),
}, },
get_new_channel_versions( get_new_channel_versions(
checkpoint_previous_versions, checkpoint_previous_versions,
@@ -1760,7 +1751,7 @@ class Pregel(
) )
if input_writes := deque(map_input(self.input_channels, values)): if input_writes := deque(map_input(self.input_channels, values)):
updated_channels = apply_writes( apply_writes(
checkpoint, checkpoint,
channels, channels,
[PregelTaskWrites((), INPUT, input_writes, [])], [PregelTaskWrites((), INPUT, input_writes, [])],
@@ -1783,15 +1774,6 @@ class Pregel(
"parents": saved.metadata.get("parents", {}) "parents": saved.metadata.get("parents", {})
if saved if saved
else {}, else {},
**(
advance_delta_counters(
channels,
updated_channels,
prev_metadata=saved.metadata,
)
if saved
else {}
),
}, },
get_new_channel_versions( get_new_channel_versions(
checkpoint_previous_versions, checkpoint_previous_versions,
@@ -1837,12 +1819,6 @@ class Pregel(
"source": "fork", "source": "fork",
"step": step + 1, "step": step + 1,
"parents": saved.metadata.get("parents", {}), "parents": saved.metadata.get("parents", {}),
# The copy has the same values and the same parent.
**{
k: v
for k, v in saved.metadata.items()
if k == "counters_since_delta_snapshot"
},
}, },
{}, {},
) )
@@ -1974,7 +1950,7 @@ class Pregel(
run_tasks: list[PregelTaskWrites] = [] run_tasks: list[PregelTaskWrites] = []
run_task_ids: list[str] = [] run_task_ids: list[str] = []
for as_node, values, provided_task_id in valid_updates: for i, (as_node, values, provided_task_id) in enumerate(valid_updates):
# create task to run all writers of the chosen node # create task to run all writers of the chosen node
writers = self.nodes[as_node].flat_writers writers = self.nodes[as_node].flat_writers
if not writers: if not writers:
@@ -1988,7 +1964,7 @@ class Pregel(
task_id = provided_task_id or ( task_id = provided_task_id or (
prepared_task_ids.popleft() prepared_task_ids.popleft()
if prepared_task_ids if prepared_task_ids
else str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) else _update_task_id(checkpoint["id"], i)
) )
run_tasks.append(task) run_tasks.append(task)
run_task_ids.append(task_id) run_task_ids.append(task_id)
@@ -2169,7 +2145,6 @@ class Pregel(
raise InvalidUpdateError( raise InvalidUpdateError(
"Cannot apply multiple updates when clearing state" "Cannot apply multiple updates when clearing state"
) )
updated_channels: set[str] = set()
if saved is not None: if saved is not None:
# tasks for this checkpoint # tasks for this checkpoint
next_tasks = prepare_next_tasks( next_tasks = prepare_next_tasks(
@@ -2192,7 +2167,7 @@ class Pregel(
for w in saved.pending_writes or [] for w in saved.pending_writes or []
if w[0] == NULL_TASK_ID if w[0] == NULL_TASK_ID
]: ]:
updated_channels |= apply_writes( apply_writes(
checkpoint, checkpoint,
channels, channels,
[PregelTaskWrites((), INPUT, null_writes, [])], [PregelTaskWrites((), INPUT, null_writes, [])],
@@ -2207,7 +2182,7 @@ class Pregel(
continue continue
next_tasks[tid].writes.append((k, v)) next_tasks[tid].writes.append((k, v))
# clear all current tasks # clear all current tasks
updated_channels |= apply_writes( apply_writes(
checkpoint, checkpoint,
channels, channels,
next_tasks.values(), next_tasks.values(),
@@ -2222,13 +2197,6 @@ class Pregel(
"source": "update", "source": "update",
"step": step + 1, "step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {}, "parents": saved.metadata.get("parents", {}) if saved else {},
**(
advance_delta_counters(
channels, updated_channels, prev_metadata=saved.metadata
)
if saved
else {}
),
}, },
get_new_channel_versions( get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"] checkpoint_previous_versions, checkpoint["channel_versions"]
@@ -2246,7 +2214,7 @@ class Pregel(
) )
if input_writes := deque(map_input(self.input_channels, values)): if input_writes := deque(map_input(self.input_channels, values)):
updated_channels = apply_writes( apply_writes(
checkpoint, checkpoint,
channels, channels,
[PregelTaskWrites((), INPUT, input_writes, [])], [PregelTaskWrites((), INPUT, input_writes, [])],
@@ -2269,15 +2237,6 @@ class Pregel(
"parents": saved.metadata.get("parents", {}) "parents": saved.metadata.get("parents", {})
if saved if saved
else {}, else {},
**(
advance_delta_counters(
channels,
updated_channels,
prev_metadata=saved.metadata,
)
if saved
else {}
),
}, },
get_new_channel_versions( get_new_channel_versions(
checkpoint_previous_versions, checkpoint_previous_versions,
@@ -2323,12 +2282,6 @@ class Pregel(
"source": "fork", "source": "fork",
"step": step + 1, "step": step + 1,
"parents": saved.metadata.get("parents", {}), "parents": saved.metadata.get("parents", {}),
# The copy has the same values and the same parent.
**{
k: v
for k, v in saved.metadata.items()
if k == "counters_since_delta_snapshot"
},
}, },
{}, {},
) )
@@ -2457,7 +2410,7 @@ class Pregel(
run_tasks: list[PregelTaskWrites] = [] run_tasks: list[PregelTaskWrites] = []
run_task_ids: list[str] = [] run_task_ids: list[str] = []
for as_node, values, provided_task_id in valid_updates: for i, (as_node, values, provided_task_id) in enumerate(valid_updates):
# create task to run all writers of the chosen node # create task to run all writers of the chosen node
writers = self.nodes[as_node].flat_writers writers = self.nodes[as_node].flat_writers
if not writers: if not writers:
@@ -2471,7 +2424,7 @@ class Pregel(
task_id = provided_task_id or ( task_id = provided_task_id or (
prepared_task_ids.popleft() prepared_task_ids.popleft()
if prepared_task_ids if prepared_task_ids
else str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) else _update_task_id(checkpoint["id"], i)
) )
run_tasks.append(task) run_tasks.append(task)
run_task_ids.append(task_id) run_task_ids.append(task_id)
@@ -4219,6 +4172,16 @@ class Pregel(
await self.cache.aclear(namespaces) await self.cache.aclear(namespaces)
def _update_task_id(checkpoint_id: str, i: int) -> str:
"""Task id for the `i`th update of a superstep that has no task to reuse.
Savers keep one write per `(task_id, idx)`, so updates sharing an id lose
all but the first one's writes, which a `DeltaChannel` replays from. The
first update keeps the id a lone update has always had.
"""
return str(uuid5(UUID(checkpoint_id), INTERRUPT if i == 0 else f"{INTERRUPT}:{i}"))
def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]: def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]:
"""Index from a trigger to nodes that depend on it.""" """Index from a trigger to nodes that depend on it."""
trigger_to_nodes: defaultdict[str, list[str]] = defaultdict(list) trigger_to_nodes: defaultdict[str, list[str]] = defaultdict(list)
@@ -147,60 +147,6 @@ async def test_predicate_fires_on_supersteps_overflow() -> None:
assert "x" not in result2 assert "x" not in result2
def _delta_counters(saver: InMemorySaver, config: Any) -> dict[str, list[int]]:
tup = saver.get_tuple(config)
assert tup is not None
counters = tup.metadata.get("counters_since_delta_snapshot") or {}
return {ch: list(c) for ch, c in counters.items()}
_UPDATE_PATHS_WITHOUT_THE_SNAPSHOT_PLAN = pytest.mark.parametrize(
("values", "as_node", "supersteps"),
[
(None, END, 1),
(None, "__copy__", 0),
({"a": []}, "__input__", 1),
],
ids=["clear as END", "copy", "update as input"],
)
@_UPDATE_PATHS_WITHOUT_THE_SNAPSHOT_PLAN
def test_update_state_path_keeps_delta_counters(
values: Any, as_node: str, supersteps: int
) -> None:
saver = InMemorySaver()
graph = _build_two_channel_graph(saver)
config = {"configurable": {"thread_id": "counters"}}
graph.invoke({"a": ["seed-a"], "b": ["seed-b"]}, config)
before = _delta_counters(saver, config)
assert set(before) == {"a", "b"}, f"both channels need live counters: {before}"
updated = graph.update_state(config, values, as_node=as_node)
assert _delta_counters(saver, updated) == {
ch: [u, s + supersteps] for ch, (u, s) in before.items()
}
@_UPDATE_PATHS_WITHOUT_THE_SNAPSHOT_PLAN
async def test_aupdate_state_path_keeps_delta_counters(
values: Any, as_node: str, supersteps: int
) -> None:
saver = InMemorySaver()
graph = _build_two_channel_graph(saver)
config = {"configurable": {"thread_id": "counters"}}
await graph.ainvoke({"a": ["seed-a"], "b": ["seed-b"]}, config)
before = _delta_counters(saver, config)
assert set(before) == {"a", "b"}, f"both channels need live counters: {before}"
updated = await graph.aupdate_state(config, values, as_node=as_node)
assert _delta_counters(saver, updated) == {
ch: [u, s + supersteps] for ch, (u, s) in before.items()
}
async def test_counter_reset_after_supersteps_snapshot() -> None: async def test_counter_reset_after_supersteps_snapshot() -> None:
"""After the supersteps bound triggers a snapshot, the counters for """After the supersteps bound triggers a snapshot, the counters for
that channel reset. Verify by using a bound higher than one run's that channel reset. Verify by using a bound higher than one run's
@@ -20,6 +20,7 @@ from typing import Annotated, Any
import pytest import pytest
from langchain_core.messages import HumanMessage from langchain_core.messages import HumanMessage
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.memory import InMemorySaver from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict from typing_extensions import TypedDict
@@ -27,16 +28,17 @@ from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph from langgraph.graph import START, StateGraph
from langgraph.graph.message import _messages_delta_reducer from langgraph.graph.message import _messages_delta_reducer
from langgraph.types import StateUpdate from langgraph.types import StateSnapshot, StateUpdate
pytestmark = pytest.mark.anyio pytestmark = pytest.mark.anyio
def _build_graph( def _build_graph(
checkpointer: InMemorySaver, checkpointer: BaseCheckpointSaver,
*, *,
two_nodes: bool = False, two_nodes: bool = False,
snapshot_frequency: int = 1000, snapshot_frequency: int = 1000,
interrupt_before: list[str] | None = None,
) -> Any: ) -> Any:
"""Compile a minimal DeltaChannel-backed `messages` graph. """Compile a minimal DeltaChannel-backed `messages` graph.
@@ -63,7 +65,7 @@ def _build_graph(
builder.set_finish_point("assistant") builder.set_finish_point("assistant")
else: else:
builder.set_finish_point("model") builder.set_finish_point("model")
return builder.compile(checkpointer=checkpointer) return builder.compile(checkpointer=checkpointer, interrupt_before=interrupt_before)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -273,10 +275,6 @@ def test_bulk_update_state_multi_task_per_superstep_delta_channel() -> None:
that each call `put_writes`. Guards the regression where moving that each call `put_writes`. Guards the regression where moving
`put_writes` outside the per-task loop would persist only the last `put_writes` outside the per-task loop would persist only the last
task's writes. task's writes.
Explicit `task_id`s are required to disambiguate writes belonging to
different `StateUpdate`s targeting the same node — otherwise both share
the deterministic interrupt-derived id and collide in the saver.
""" """
saver = InMemorySaver() saver = InMemorySaver()
@@ -310,6 +308,92 @@ def test_bulk_update_state_multi_task_per_superstep_delta_channel() -> None:
assert sorted(ids) == ["m1", "m2"] assert sorted(ids) == ["m1", "m2"]
def _update(content: str, as_node: str) -> StateUpdate:
return StateUpdate(
values={"messages": [HumanMessage(content=content, id=content)]},
as_node=as_node,
)
def _contents(state: StateSnapshot) -> list[str]:
return [m.content for m in state.values["messages"]]
def test_bulk_update_state_keeps_every_update_without_task_ids(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
graph = _build_graph(sync_checkpointer, two_nodes=True)
config = {"configurable": {"thread_id": "bulk-no-task-ids"}}
graph.invoke({"messages": [HumanMessage(content="hi", id="hi")]}, config)
graph.bulk_update_state(
config,
[
[
_update("first", "model"),
_update("second", "model"),
_update("third", "assistant"),
]
],
)
contents = _contents(graph.get_state(config))
assert sorted(contents) == ["first", "hi", "second", "third"], (
f"every update's writes must persist; got {contents}"
)
async def test_abulk_update_state_keeps_every_update_without_task_ids(
async_checkpointer: BaseCheckpointSaver,
) -> None:
graph = _build_graph(async_checkpointer, two_nodes=True)
config = {"configurable": {"thread_id": "bulk-no-task-ids"}}
await graph.ainvoke({"messages": [HumanMessage(content="hi", id="hi")]}, config)
await graph.abulk_update_state(
config,
[
[
_update("first", "model"),
_update("second", "model"),
_update("third", "assistant"),
]
],
)
contents = _contents(await graph.aget_state(config))
assert sorted(contents) == ["first", "hi", "second", "third"], (
f"every update's writes must persist; got {contents}"
)
def test_bulk_update_state_keeps_every_update_next_to_a_pending_task(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
graph = _build_graph(
sync_checkpointer, two_nodes=True, interrupt_before=["assistant"]
)
config = {"configurable": {"thread_id": "bulk-pending-task"}}
graph.invoke({"messages": [HumanMessage(content="hi", id="hi")]}, config)
assert graph.get_state(config).next == ("assistant",)
graph.bulk_update_state(
config,
[
[
_update("first", "assistant"),
_update("second", "model"),
_update("third", "model"),
]
],
)
contents = _contents(graph.get_state(config))
assert sorted(contents) == ["first", "hi", "second", "third"], (
f"every update's writes must persist; got {contents}"
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Public-API observation of fresh-thread checkpoint shape # Public-API observation of fresh-thread checkpoint shape
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
+6 -6
View File
@@ -2841,11 +2841,11 @@ wheels = [
[[package]] [[package]]
name = "pyjwt" name = "pyjwt"
version = "2.15.1" version = "2.13.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/43/ea/5194e52748b0da83d71e082d75496eaec6e58f419f5e184786ded517e6a9/pyjwt-2.15.1.tar.gz", hash = "sha256:4f259e80cdfb6b3fc18a7de51fd1ef9ec79652f25019bae68975ca2468a34df8", size = 121252, upload-time = "2026-09-28T18:40:42.598Z" } sdist = { url = "https://files.pythonhosted.org/packages/3b/81/58d0ac84e1ef3a3843791d6954d94c0b33d526c75eeb1efbce9d0a4c4077/pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423", size = 107515, upload-time = "2026-05-21T19:54:36.618Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/50/ca/44de4e75f8aadc457f0634be3b542815078ded46dca30efb960edeecad6e/pyjwt-2.15.1-py3-none-any.whl", hash = "sha256:42d59d631f7768a1028a64c7ff581a9bf7519804daf91fc5b6c56e30eec5e193", size = 33860, upload-time = "2026-09-28T18:40:41.429Z" }, { url = "https://files.pythonhosted.org/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728", size = 31274, upload-time = "2026-05-21T19:54:35.362Z" },
] ]
[[package]] [[package]]
@@ -3687,11 +3687,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+3 -3
View File
@@ -1363,11 +1363,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]
+3 -3
View File
@@ -1175,11 +1175,11 @@ wheels = [
[[package]] [[package]]
name = "urllib3" name = "urllib3"
version = "2.8.0" version = "2.7.0"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
] ]
[[package]] [[package]]