From 197b01e985c2e3f5b9b9d576a7e1c83d9e089f54 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 27 May 2025 15:45:15 -0700
Subject: [PATCH] Update managed value usage in local_read

---
 libs/langgraph/langgraph/pregel/__init__.py | 19 +++++++++
 libs/langgraph/langgraph/pregel/algo.py     | 44 ++++++++++++---------
 2 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 73ada5aff..1cd6a225d 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -72,6 +72,7 @@ from langgraph.errors import (
 from langgraph.managed.base import ManagedValueSpec
 from langgraph.pregel.algo import (
     PregelTaskWrites,
+    _scratchpad,
     apply_writes,
     local_read,
     prepare_next_tasks,
@@ -1782,6 +1783,15 @@ class Pregel(PregelProtocol):
                             CONFIG_KEY_SEND: writes.extend,
                             CONFIG_KEY_READ: partial(
                                 local_read,
+                                _scratchpad(
+                                    None,
+                                    [],
+                                    task_id,
+                                    "",
+                                    None,
+                                    step,
+                                    step + 2,
+                                ),
                                 channels,
                                 managed,
                                 task,
@@ -2188,6 +2198,15 @@ class Pregel(PregelProtocol):
                             CONFIG_KEY_SEND: writes.extend,
                             CONFIG_KEY_READ: partial(
                                 local_read,
+                                _scratchpad(
+                                    None,
+                                    [],
+                                    task_id,
+                                    "",
+                                    None,
+                                    step,
+                                    step + 2,
+                                ),
                                 channels,
                                 managed,
                                 task,
diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index 505928867..14838de21 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -172,6 +172,7 @@ def should_interrupt(
 
 
 def local_read(
+    scratchpad: PregelScratchpad,
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
     task: WritesProtocol,
@@ -208,7 +209,7 @@ def local_read(
     else:
         values = read_channels(channels, select)
     if managed_keys:
-        values.update({k: managed[k]() for k in managed_keys})
+        values.update({k: managed[k].get(scratchpad) for k in managed_keys})
     return values
 
 
@@ -571,6 +572,15 @@ def prepare_single_task(
                 )
             else:
                 cache_key = None
+            scratchpad = _scratchpad(
+                config[CONF].get(CONFIG_KEY_SCRATCHPAD),
+                pending_writes,
+                task_id,
+                xxh3_128_hexdigest(task_checkpoint_ns.encode()),
+                config[CONF].get(CONFIG_KEY_RESUME_MAP),
+                step,
+                stop,
+            )
             return PregelExecutableTask(
                 name,
                 call.input,
@@ -587,6 +597,7 @@ def prepare_single_task(
                         CONFIG_KEY_SEND: writes.extend,
                         CONFIG_KEY_READ: partial(
                             local_read,
+                            scratchpad,
                             channels,
                             managed,
                             PregelTaskWrites(task_path, name, writes, triggers),
@@ -601,15 +612,7 @@ def prepare_single_task(
                         },
                         CONFIG_KEY_CHECKPOINT_ID: None,
                         CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                        CONFIG_KEY_SCRATCHPAD: _scratchpad(
-                            config[CONF].get(CONFIG_KEY_SCRATCHPAD),
-                            pending_writes,
-                            task_id,
-                            xxh3_128_hexdigest(task_checkpoint_ns.encode()),
-                            config[CONF].get(CONFIG_KEY_RESUME_MAP),
-                            step,
-                            stop,
-                        ),
+                        CONFIG_KEY_SCRATCHPAD: scratchpad,
                     },
                 ),
                 triggers,
@@ -695,6 +698,15 @@ def prepare_single_task(
                 )
             else:
                 cache_key = None
+            scratchpad = _scratchpad(
+                config[CONF].get(CONFIG_KEY_SCRATCHPAD),
+                pending_writes,
+                task_id,
+                xxh3_128_hexdigest(task_checkpoint_ns.encode()),
+                config[CONF].get(CONFIG_KEY_RESUME_MAP),
+                step,
+                stop,
+            )
             return PregelExecutableTask(
                 packet.node,
                 packet.arg,
@@ -712,6 +724,7 @@ def prepare_single_task(
                         CONFIG_KEY_SEND: writes.extend,
                         CONFIG_KEY_READ: partial(
                             local_read,
+                            scratchpad,
                             channels,
                             managed,
                             PregelTaskWrites(task_path, packet.node, writes, triggers),
@@ -726,15 +739,7 @@ def prepare_single_task(
                         },
                         CONFIG_KEY_CHECKPOINT_ID: None,
                         CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                        CONFIG_KEY_SCRATCHPAD: _scratchpad(
-                            config[CONF].get(CONFIG_KEY_SCRATCHPAD),
-                            pending_writes,
-                            task_id,
-                            xxh3_128_hexdigest(task_checkpoint_ns.encode()),
-                            config[CONF].get(CONFIG_KEY_RESUME_MAP),
-                            step,
-                            stop,
-                        ),
+                        CONFIG_KEY_SCRATCHPAD: scratchpad,
                         CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get(
                             PREVIOUS, None
                         ),
@@ -860,6 +865,7 @@ def prepare_single_task(
                                 CONFIG_KEY_SEND: writes.extend,
                                 CONFIG_KEY_READ: partial(
                                     local_read,
+                                    scratchpad,
                                     channels,
                                     managed,
                                     PregelTaskWrites(