langgraph-executor 0.0.1a5__py3-none-any.whl → 0.0.1a6__py3-none-any.whl

langgraph_executor/common.py

@@ -1,6 +1,7 @@
 import traceback
 from collections.abc import Mapping, Sequence
 from collections.abc import Sequence as SequenceType
+from contextvars import ContextVar
 from typing import Any, cast
 
 from google.protobuf.json_format import MessageToDict
@@ -19,7 +20,12 @@ from langgraph._internal._constants import (
 from langgraph._internal._scratchpad import PregelScratchpad
 from langgraph._internal._typing import MISSING
 from langgraph.channels.base import BaseChannel, EmptyChannelError
-from langgraph.checkpoint.base import Checkpoint, PendingWrite
+from langgraph.checkpoint.base import (
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    PendingWrite,
+)
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.errors import GraphBubbleUp, GraphInterrupt
 from langgraph.managed.base import ManagedValue, ManagedValueMapping
@@ -30,6 +36,10 @@ from langgraph.types import Command, Interrupt, Send
 
 from langgraph_executor.pb import types_pb2
 
+var_child_runnable_config: ContextVar[RunnableConfig | None] = ContextVar(
+    "child_runnable_config", default=None
+)
+
 
 def map_reserved_configurable(
     reserved_configurable: types_pb2.ReservedConfigurable,
@@ -108,13 +118,18 @@ def reconstruct_checkpoint(request_checkpoint: types_pb2.Checkpoint) -> Checkpoi
     versions_seen = {
         k: dict(v.channel_versions) for k, v in request_checkpoint.versions_seen.items()
     }
-    # channel_values = unwrap_channel_values(request_checkpoint.channel_values)
+
+    channel_values = {}
+    if request_checkpoint.channel_values:
+        channel_values = {
+            k: pb_to_val(v) for k, v in request_checkpoint.channel_values.items()
+        }
 
     return Checkpoint(
         v=request_checkpoint.v,
         id=request_checkpoint.id,
         channel_versions=channel_versions,
-        # channel_values=channel_values,
+        channel_values=channel_values,
         versions_seen=versions_seen,
         ts=request_checkpoint.ts,
     )
@@ -249,7 +264,8 @@ def command_to_pb(cmd: Command) -> types_pb2.Value:
         cmd_pb.graph = cmd.graph
     if cmd.update:
         if isinstance(cmd.update, dict):
-            cmd_pb.update.update({k: val_to_pb(None, v) for k, v in cmd.update.items()})
+            for k, v in cmd.update.items():
+                cmd_pb.update[k].CopyFrom(val_to_pb(None, v))
         else:
             cmd_pb.update.update({"__root__": val_to_pb(None, cmd.update)})
     if cmd.resume:
@@ -393,3 +409,40 @@ def pb_to_pending_writes(
         return None
 
     return [(pw.task_id, pw.channel, pb_to_val(pw.value)) for pw in pb]
+
+
+def reconstruct_pending_writes(
+    pb: SequenceType[types_pb2.PendingWrite],
+) -> list[PendingWrite] | None:
+    if not pb:
+        return None
+
+    return [(pw.task_id, pw.channel, pb_to_val(pw.value)) for pw in pb]
+
+
+def reconstruct_checkpoint_tuple(
+    checkpoint_tuple_pb: types_pb2.CheckpointTuple,
+) -> CheckpointTuple | None:
+    if not checkpoint_tuple_pb:
+        return None
+
+    return CheckpointTuple(
+        config=reconstruct_config(checkpoint_tuple_pb.config),
+        checkpoint=reconstruct_checkpoint(checkpoint_tuple_pb.checkpoint),
+        metadata=reconstruct_checkpoint_metadata(checkpoint_tuple_pb.metadata),
+        parent_config=reconstruct_config(checkpoint_tuple_pb.parent_config),
+        pending_writes=reconstruct_pending_writes(checkpoint_tuple_pb.pending_writes),
+    )
+
+
+def reconstruct_checkpoint_metadata(
+    metadata_pb: types_pb2.CheckpointMetadata,
+) -> CheckpointMetadata | None:
+    if not metadata_pb:
+        return None
+
+    return CheckpointMetadata(
+        source=metadata_pb.source,
+        step=metadata_pb.step,
+        parents=metadata_pb.parents,
+    )

langgraph_executor/executor.py

@@ -16,8 +16,14 @@ from langgraph_executor.pb.executor_pb2_grpc import (
 # Internal helpers
 LOGGER = logging.getLogger(__name__)
 
+DEFAULT_RUNTIME_ADDRESS = "127.0.0.1:50051"
 
-def create_server(graphs: dict[str, Pregel], address: str) -> grpc.aio.Server:
+
+def create_server(
+    graphs: dict[str, Pregel],
+    address: str,
+    runtime_address: str = DEFAULT_RUNTIME_ADDRESS,
+) -> grpc.aio.Server:
     graphs, subgraph_map = _load_graphs(graphs)
     server = grpc.aio.server(
         # Be permissive: allow client pings without active RPCs and accept intervals
@@ -30,8 +36,16 @@ def create_server(graphs: dict[str, Pregel], address: str) -> grpc.aio.Server:
         ]
     )
     getter = functools.partial(get_graph, graphs=graphs)
+    # lazy connection, doesn't matter if runtime service is ready
+    # also thread-safe, so we can process multiple subgraphs concurrently
+    runtime_channel = grpc.insecure_channel(runtime_address)
     add_LangGraphExecutorServicer_to_server(
-        LangGraphExecutorServicer(graphs, subgraph_map=subgraph_map, get_graph=getter),
+        LangGraphExecutorServicer(
+            graphs,
+            subgraph_map=subgraph_map,
+            get_graph=getter,
+            runtime_channel=runtime_channel,
+        ),
         server,
     )
     server.add_insecure_port(address)
@@ -55,14 +69,14 @@ def _load_graphs(graphs: dict[str, Pregel]) -> tuple[dict[str, Pregel], dict[str
 
     # Then, collect all subgraphs and mappings
     all_subgraphs: dict[str, Pregel] = {}
-    subgraph_to_parent: dict[str, str] = {}
+    # subgraph_to_parent: dict[str, str] = {}
 
-    for root_graph in graphs.values():
-        subgraphs, mappings = _collect_subgraphs(root_graph, root_graph.name)
-        all_subgraphs.update(subgraphs)
-        subgraph_to_parent.update(mappings)
+    # for root_graph in graphs.values():
+    #     subgraphs, mappings = _collect_subgraphs(root_graph, root_graph.name)
+    #     all_subgraphs.update(subgraphs)
+    #     subgraph_to_parent.update(mappings)
 
-    subgraph_map.update(subgraph_to_parent)
+    # subgraph_map.update(subgraph_to_parent)
 
     # Now build self.graphs in hierarchical order (parents before children)
     for root_name in sorted(graphs.keys()):

langgraph_executor/executor_base.py

@@ -3,6 +3,7 @@ import contextlib
 import functools
 import logging
 import uuid
+from collections import deque
 from collections.abc import (
     AsyncIterator,
     Awaitable,
@@ -11,22 +12,36 @@ from collections.abc import (
     Iterator,
     Sequence,
 )
+from functools import partial
 from typing import Any, Protocol, cast
 
 import grpc
 import grpc.aio
 from google.protobuf.struct_pb2 import Struct  # type: ignore[import-untyped]
 from langchain_core.messages import BaseMessage, BaseMessageChunk
-from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables import RunnableConfig, RunnableSequence
+from langgraph._internal._config import patch_config
+from langgraph._internal._constants import (
+    CONFIG_KEY_READ,
+    CONFIG_KEY_SEND,
+    CONFIG_KEY_TASK_ID,
+    INTERRUPT,
+)
 from langgraph.checkpoint.base import Checkpoint
-from langgraph.errors import GraphBubbleUp, GraphInterrupt
+from langgraph.errors import GraphBubbleUp, GraphInterrupt, InvalidUpdateError
 from langgraph.pregel import Pregel
-from langgraph.pregel._algo import apply_writes
+from langgraph.pregel._algo import (
+    PregelTaskWrites,
+    _scratchpad,
+    apply_writes,
+    local_read,
+)
 from langgraph.pregel._checkpoint import channels_from_checkpoint
 from langgraph.pregel._retry import arun_with_retry
 from langgraph.store.base import BaseStore
 from langgraph.types import PregelExecutableTask
 
+from langgraph_executor.client.patch import _patch_pregel
 from langgraph_executor.common import (
     checkpoint_to_proto,
     exception_to_pb,
@@ -37,6 +52,7 @@ from langgraph_executor.common import (
     reconstruct_config,
     reconstruct_task_writes,
     updates_to_proto,
+    var_child_runnable_config,
 )
 from langgraph_executor.execute_task import (
     extract_writes,
@@ -44,6 +60,7 @@ from langgraph_executor.execute_task import (
 )
 from langgraph_executor.extract_graph import extract_graph
 from langgraph_executor.pb import executor_pb2, executor_pb2_grpc, types_pb2
+from langgraph_executor.pb.runtime_pb2_grpc import LangGraphRuntimeStub
 from langgraph_executor.stream_utils import ExecutorStreamHandler
 
 
@@ -72,6 +89,7 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         *,
         subgraph_map: dict[str, str],
         get_graph: GetGraph,
+        runtime_channel: grpc.Channel | None = None,
         logger: Logger | None = None,
         on_message: Callable[
             [
@@ -98,7 +116,15 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
         self.graph_names = sorted(self.graphs)
         self.subgraph_map = subgraph_map
         self.get_graph = get_graph
+        self.runtime_channel = runtime_channel
+        self.runtime_client: LangGraphRuntimeStub | None = (
+            LangGraphRuntimeStub(runtime_channel)
+            if runtime_channel is not None
+            else None
+        )
         _patch_base_message_with_ids()
+        if self.runtime_client is not None:
+            _patch_pregel(self.runtime_client, self.logger)
         self._graph_definition_cache: dict[str, executor_pb2.GetGraphResponse] = {}
         self.on_message = on_message
         self.on_custom = on_custom
@@ -351,6 +377,66 @@ class LangGraphExecutorServicer(executor_pb2_grpc.LangGraphExecutorServicer):
             self.logger.exception(f"ApplyWrites error: {e}")
             await context.abort(grpc.StatusCode.INTERNAL, str(e))
 
+    async def StateUpdate(
+        self,
+        request: executor_pb2.StateUpdateRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> executor_pb2.TaskResult | None:
+        self.logger.debug("StateUpdate called")
+
+        try:
+            async with self.get_graph(request.graph_name, RunnableConfig()) as graph:
+                config = reconstruct_config(request.config)
+                channels, managed = reconstruct_channels(
+                    request.channels.channels,
+                    graph,
+                    scratchpad=None,  # type: ignore[invalid-arg-type]
+                )
+                writes: deque[tuple[str, Any]] = deque()
+                writers = graph.nodes[request.node_name].flat_writers
+                if not writers:
+                    raise InvalidUpdateError(f"Node {request.node_name} has no writers")
+                task = PregelTaskWrites((), request.node_name, writes, [INTERRUPT])
+                task_id = request.task_id
+
+                run = RunnableSequence(*writers) if len(writers) > 1 else writers[0]
+
+                input = pb_to_val(request.values)
+
+                run.invoke(
+                    input,
+                    patch_config(
+                        config,
+                        run_name=graph.name + "UpdateState",
+                        configurable={
+                            # deque.extend is thread-safe
+                            CONFIG_KEY_SEND: writes.extend,
+                            CONFIG_KEY_TASK_ID: task_id,
+                            CONFIG_KEY_READ: partial(
+                                local_read,
+                                _scratchpad(
+                                    None,
+                                    [],
+                                    task_id,
+                                    "",
+                                    None,
+                                    request.step,
+                                    request.step + 2,
+                                ),
+                                channels,
+                                managed,
+                                task,
+                            ),
+                        },
+                    ),
+                )
+
+                return executor_pb2.TaskResult(writes=extract_writes(task.writes))
+
+        except Exception as e:
+            self.logger.exception(f"StateUpdate error: {e}")
+            await context.abort(grpc.StatusCode.INTERNAL, str(e))
+
     async def GenerateCacheKey(
         self,
         request: executor_pb2.GenerateCacheKeyRequest,
@@ -369,6 +455,8 @@ async def _run_task(
     logger: Logger,
     stream_queue: asyncio.Queue[executor_pb2.ExecuteTaskResponse],
 ) -> types_pb2.ExecutorError | None:
+    # Set config
+    var_child_runnable_config.set(task.config)
     try:
         await arun_with_retry(
             task,

langgraph_executor/extract_graph.py

@@ -117,6 +117,18 @@ def extract_trigger_to_nodes(
     return trigger_map
 
 
+def extract_string_or_slice(
+    val: str | Sequence[str] | None,
+) -> types_pb2.StringOrSlice | None:
+    if val is None:
+        return None
+    if isinstance(val, str):
+        return types_pb2.StringOrSlice(values=[val], is_string=True)
+    if isinstance(val, list):
+        return types_pb2.StringOrSlice(values=val, is_string=False)
+    raise NotImplementedError(f"Cannot extract field value {val} as string or slice")
+
+
 def extract_graph(graph: Pregel, name: str | None = None) -> graph_pb2.GraphDefinition:
     """Extract graph information from a compiled LangGraph graph.
 
@@ -134,13 +146,7 @@ def extract_graph(graph: Pregel, name: str | None = None) -> graph_pb2.GraphDefi
             else graph.stream_mode
         ),
         stream_eager=bool(graph.stream_eager),
-        stream_channels=(
-            [graph.stream_channels]
-            if isinstance(graph.stream_channels, str)
-            else list(graph.stream_channels)
-            if graph.stream_channels
-            else []
-        ),
+        stream_channels=extract_string_or_slice(graph.stream_channels),
         step_timeout=float(graph.step_timeout) if graph.step_timeout else 0.0,
         debug=bool(graph.debug),
         # TODO retry policy
@@ -150,24 +156,11 @@ def extract_graph(graph: Pregel, name: str | None = None) -> graph_pb2.GraphDefi
         config=extract_config(graph.config) if graph.config else None,
         nodes=extract_nodes(graph.nodes),
         trigger_to_nodes=extract_trigger_to_nodes(graph.trigger_to_nodes),
+        stream_channels_asis=extract_string_or_slice(graph.stream_channels_asis),
+        input_channels=extract_string_or_slice(graph.input_channels),
+        output_channels=extract_string_or_slice(graph.output_channels),
     )
 
-    # Set input_channels and input_channels_is_list flag based on type
-    if isinstance(graph.input_channels, str):
-        graph_def.input_channels.extend([graph.input_channels])
-        graph_def.input_channels_is_list = False
-    else:
-        graph_def.input_channels.extend(graph.input_channels)
-        graph_def.input_channels_is_list = True
-
-    # Set output_channels and output_channels_is_list flag based on type
-    if isinstance(graph.output_channels, str):
-        graph_def.output_channels.extend([graph.output_channels])
-        graph_def.output_channels_is_list = False
-    else:
-        graph_def.output_channels.extend(graph.output_channels)
-        graph_def.output_channels_is_list = True
-
     return graph_def
 
 

langgraph_executor/pb/executor_pb2.py

@@ -27,7 +27,7 @@ from . import graph_pb2 as graph__pb2
 from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0e\x65xecutor.proto\x12\x08\x65xecutor\x1a\x0btypes.proto\x1a\x0bgraph.proto\x1a\x1cgoogle/protobuf/struct.proto\"\x13\n\x11ListGraphsRequest\")\n\x12ListGraphsResponse\x12\x13\n\x0bgraph_names\x18\x01 \x03(\t\"%\n\x0fGetGraphRequest\x12\x12\n\ngraph_name\x18\x01 \x01(\t\"\x9a\x01\n\x10GetGraphResponse\x12\x30\n\x10graph_definition\x18\x01 \x01(\x0b\x32\x16.graph.GraphDefinition\x12\x18\n\x0bparent_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0c\x63heckpointer\x18\x03 \x01(\x08H\x01\x88\x01\x01\x42\x0e\n\x0c_parent_nameB\x0f\n\r_checkpointer\"\x15\n\x13GetAllGraphsRequest\"\x8a\x02\n\x1d\x43hannelsFromCheckpointRequest\x12\x12\n\ngraph_name\x18\x01 \x01(\t\x12\x1e\n\x05specs\x18\x02 \x01(\x0b\x32\x0f.types.Channels\x12g\n\x19\x63heckpoint_channel_values\x18\x03 \x03(\x0b\x32\x44.executor.ChannelsFromCheckpointRequest.CheckpointChannelValuesEntry\x1aL\n\x1c\x43heckpointChannelValuesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x1b\n\x05value\x18\x02 \x01(\x0b\x32\x0c.types.Value:\x02\x38\x01\"C\n\x1e\x43hannelsFromCheckpointResponse\x12!\n\x08\x63hannels\x18\x01 \x01(\x0b\x32\x0f.types.Channels\"\x88\x01\n\x12\x45xecuteTaskRequest\x12)\n\x04init\x18\x01 \x01(\x0b\x32\x19.executor.ExecuteTaskInitH\x00\x12<\n\x14\x63\x61\x63he_check_response\x18\x02 \x01(\x0b\x32\x1c.executor.CacheCheckResponseH\x00\x42\t\n\x07message\"\xd6\x01\n\x0f\x45xecuteTaskInit\x12\x19\n\x04task\x18\x01 \x01(\x0b\x32\x0b.types.Task\x12!\n\x08\x63hannels\x18\x02 \x01(\x0b\x32\x0f.types.Channels\x12\x14\n\x0cstream_modes\x18\x03 \x03(\t\x12\x0c\n\x04step\x18\x04 \x01(\x05\x12\x15\n\rcheckpoint_ns\x18\x05 \x03(\t\x12\x13\n\x0boutput_keys\x18\x06 \x03(\t\x12\x13\n\x0bstream_keys\x18\x07 \x03(\t\x12\x12\n\ngraph_name\x18\x08 \x01(\t\x12\x0c\n\x04stop\x18\t \x01(\x05\"\xf9\x01\n\x13\x45xecuteTaskResponse\x12+\n\x0btask_result\x18\x01 \x01(\x0b\x32\x14.executor.TaskResultH\x00\x12\x32\n\x18message_or_message_chunk\x18\x02 \x01(\x0b\x32\x0e.types.MessageH\x00\x12:\n\x13\x63\x61\x63he_check_request\x18\x03 \x01(\x0b\x32\x1b.executor.CacheCheckRequestH\x00\x12:\n\x13\x63ustom_stream_event\x18\x04 \x01(\x0b\x32\x1b.executor.CustomStreamEventH\x00\x42\t\n\x07message\"^\n\nTaskResult\x12(\n\x05\x65rror\x18\x01 \x01(\x0b\x32\x14.types.ExecutorErrorH\x00\x88\x01\x01\x12\x1c\n\x06writes\x18\x02 \x03(\x0b\x32\x0c.types.WriteB\x08\n\x06_error\"=\n\x11\x43ustomStreamEvent\x12(\n\x07payload\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\"\xa4\x01\n\x12\x41pplyWritesRequest\x12%\n\ncheckpoint\x18\x01 \x01(\x0b\x32\x11.types.Checkpoint\x12!\n\x08\x63hannels\x18\x02 \x01(\x0b\x32\x0f.types.Channels\x12\x1a\n\x05tasks\x18\x03 \x03(\x0b\x32\x0b.types.Task\x12\x14\n\x0cnext_version\x18\x04 \x01(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\"b\n\x13\x41pplyWritesResponse\x12!\n\x07updates\x18\x01 \x01(\x0b\x32\x0e.types.UpdatesH\x00\x12\x1d\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x0c.types.ValueH\x00\x42\t\n\x07message\"\xc8\x01\n\x17GenerateCacheKeyRequest\x12\x11\n\tnode_name\x18\x01 \x01(\t\x12\x16\n\x0einput_channels\x18\x02 \x03(\t\x12L\n\x0e\x63hannel_values\x18\x03 \x03(\x0b\x32\x34.executor.GenerateCacheKeyRequest.ChannelValuesEntry\x1a\x34\n\x12\x43hannelValuesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c:\x02\x38\x01\"-\n\x18GenerateCacheKeyResponse\x12\x11\n\tcache_key\x18\x01 \x01(\t\"Y\n\x11\x43\x61\x63heCheckRequest\x12\x17\n\x0f\x63\x61\x63he_namespace\x18\x01 \x03(\t\x12\x11\n\tcache_key\x18\x02 \x01(\t\x12\x10\n\x03ttl\x18\x03 \x01(\x05H\x00\x88\x01\x01\x42\x06\n\x04_ttl\"\'\n\x12\x43\x61\x63heCheckResponse\x12\x11\n\tcache_hit\x18\x01 \x01(\x08\x32\xd0\x04\n\x11LangGraphExecutor\x12G\n\nListGraphs\x12\x1b.executor.ListGraphsRequest\x1a\x1c.executor.ListGraphsResponse\x12\x41\n\x08GetGraph\x12\x19.executor.GetGraphRequest\x1a\x1a.executor.GetGraphResponse\x12K\n\x0cGetAllGraphs\x12\x1d.executor.GetAllGraphsRequest\x1a\x1a.executor.GetGraphResponse0\x01\x12k\n\x16\x43hannelsFromCheckpoint\x12\'.executor.ChannelsFromCheckpointRequest\x1a(.executor.ChannelsFromCheckpointResponse\x12N\n\x0b\x45xecuteTask\x12\x1c.executor.ExecuteTaskRequest\x1a\x1d.executor.ExecuteTaskResponse(\x01\x30\x01\x12J\n\x0b\x41pplyWrites\x12\x1c.executor.ApplyWritesRequest\x1a\x1d.executor.ApplyWritesResponse\x12Y\n\x10GenerateCacheKey\x12!.executor.GenerateCacheKeyRequest\x1a\".executor.GenerateCacheKeyResponseB1Z/github.com/langchain-ai/langgraph-go/runtime/pbb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0e\x65xecutor.proto\x12\x08\x65xecutor\x1a\x0btypes.proto\x1a\x0bgraph.proto\x1a\x1cgoogle/protobuf/struct.proto\"\xc2\x01\n\x12StateUpdateRequest\x12\x11\n\tnode_name\x18\x01 \x01(\t\x12\x0f\n\x07task_id\x18\x02 \x01(\t\x12\x12\n\ngraph_name\x18\x03 \x01(\t\x12\x0c\n\x04step\x18\x04 \x01(\x05\x12!\n\x08\x63hannels\x18\x05 \x01(\x0b\x32\x0f.types.Channels\x12\x1c\n\x06values\x18\x06 \x01(\x0b\x32\x0c.types.Value\x12%\n\x06\x63onfig\x18\x07 \x01(\x0b\x32\x15.types.RunnableConfig\"\x13\n\x11ListGraphsRequest\")\n\x12ListGraphsResponse\x12\x13\n\x0bgraph_names\x18\x01 \x03(\t\"%\n\x0fGetGraphRequest\x12\x12\n\ngraph_name\x18\x01 \x01(\t\"\x9a\x01\n\x10GetGraphResponse\x12\x30\n\x10graph_definition\x18\x01 \x01(\x0b\x32\x16.graph.GraphDefinition\x12\x18\n\x0bparent_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0c\x63heckpointer\x18\x03 \x01(\x08H\x01\x88\x01\x01\x42\x0e\n\x0c_parent_nameB\x0f\n\r_checkpointer\"\x15\n\x13GetAllGraphsRequest\"\x8a\x02\n\x1d\x43hannelsFromCheckpointRequest\x12\x12\n\ngraph_name\x18\x01 \x01(\t\x12\x1e\n\x05specs\x18\x02 \x01(\x0b\x32\x0f.types.Channels\x12g\n\x19\x63heckpoint_channel_values\x18\x03 \x03(\x0b\x32\x44.executor.ChannelsFromCheckpointRequest.CheckpointChannelValuesEntry\x1aL\n\x1c\x43heckpointChannelValuesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x1b\n\x05value\x18\x02 \x01(\x0b\x32\x0c.types.Value:\x02\x38\x01\"C\n\x1e\x43hannelsFromCheckpointResponse\x12!\n\x08\x63hannels\x18\x01 \x01(\x0b\x32\x0f.types.Channels\"\x88\x01\n\x12\x45xecuteTaskRequest\x12)\n\x04init\x18\x01 \x01(\x0b\x32\x19.executor.ExecuteTaskInitH\x00\x12<\n\x14\x63\x61\x63he_check_response\x18\x02 \x01(\x0b\x32\x1c.executor.CacheCheckResponseH\x00\x42\t\n\x07message\"\xd6\x01\n\x0f\x45xecuteTaskInit\x12\x19\n\x04task\x18\x01 \x01(\x0b\x32\x0b.types.Task\x12!\n\x08\x63hannels\x18\x02 \x01(\x0b\x32\x0f.types.Channels\x12\x14\n\x0cstream_modes\x18\x03 \x03(\t\x12\x0c\n\x04step\x18\x04 \x01(\x05\x12\x15\n\rcheckpoint_ns\x18\x05 \x03(\t\x12\x13\n\x0boutput_keys\x18\x06 \x03(\t\x12\x13\n\x0bstream_keys\x18\x07 \x03(\t\x12\x12\n\ngraph_name\x18\x08 \x01(\t\x12\x0c\n\x04stop\x18\t \x01(\x05\"\xf9\x01\n\x13\x45xecuteTaskResponse\x12+\n\x0btask_result\x18\x01 \x01(\x0b\x32\x14.executor.TaskResultH\x00\x12\x32\n\x18message_or_message_chunk\x18\x02 \x01(\x0b\x32\x0e.types.MessageH\x00\x12:\n\x13\x63\x61\x63he_check_request\x18\x03 \x01(\x0b\x32\x1b.executor.CacheCheckRequestH\x00\x12:\n\x13\x63ustom_stream_event\x18\x04 \x01(\x0b\x32\x1b.executor.CustomStreamEventH\x00\x42\t\n\x07message\"^\n\nTaskResult\x12(\n\x05\x65rror\x18\x01 \x01(\x0b\x32\x14.types.ExecutorErrorH\x00\x88\x01\x01\x12\x1c\n\x06writes\x18\x02 \x03(\x0b\x32\x0c.types.WriteB\x08\n\x06_error\"=\n\x11\x43ustomStreamEvent\x12(\n\x07payload\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\"\xa4\x01\n\x12\x41pplyWritesRequest\x12%\n\ncheckpoint\x18\x01 \x01(\x0b\x32\x11.types.Checkpoint\x12!\n\x08\x63hannels\x18\x02 \x01(\x0b\x32\x0f.types.Channels\x12\x1a\n\x05tasks\x18\x03 \x03(\x0b\x32\x0b.types.Task\x12\x14\n\x0cnext_version\x18\x04 \x01(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\"b\n\x13\x41pplyWritesResponse\x12!\n\x07updates\x18\x01 \x01(\x0b\x32\x0e.types.UpdatesH\x00\x12\x1d\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x0c.types.ValueH\x00\x42\t\n\x07message\"\xc8\x01\n\x17GenerateCacheKeyRequest\x12\x11\n\tnode_name\x18\x01 \x01(\t\x12\x16\n\x0einput_channels\x18\x02 \x03(\t\x12L\n\x0e\x63hannel_values\x18\x03 \x03(\x0b\x32\x34.executor.GenerateCacheKeyRequest.ChannelValuesEntry\x1a\x34\n\x12\x43hannelValuesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c:\x02\x38\x01\"-\n\x18GenerateCacheKeyResponse\x12\x11\n\tcache_key\x18\x01 \x01(\t\"Y\n\x11\x43\x61\x63heCheckRequest\x12\x17\n\x0f\x63\x61\x63he_namespace\x18\x01 \x03(\t\x12\x11\n\tcache_key\x18\x02 \x01(\t\x12\x10\n\x03ttl\x18\x03 \x01(\x05H\x00\x88\x01\x01\x42\x06\n\x04_ttl\"\'\n\x12\x43\x61\x63heCheckResponse\x12\x11\n\tcache_hit\x18\x01 \x01(\x08\x32\x93\x05\n\x11LangGraphExecutor\x12G\n\nListGraphs\x12\x1b.executor.ListGraphsRequest\x1a\x1c.executor.ListGraphsResponse\x12\x41\n\x08GetGraph\x12\x19.executor.GetGraphRequest\x1a\x1a.executor.GetGraphResponse\x12K\n\x0cGetAllGraphs\x12\x1d.executor.GetAllGraphsRequest\x1a\x1a.executor.GetGraphResponse0\x01\x12k\n\x16\x43hannelsFromCheckpoint\x12\'.executor.ChannelsFromCheckpointRequest\x1a(.executor.ChannelsFromCheckpointResponse\x12N\n\x0b\x45xecuteTask\x12\x1c.executor.ExecuteTaskRequest\x1a\x1d.executor.ExecuteTaskResponse(\x01\x30\x01\x12J\n\x0b\x41pplyWrites\x12\x1c.executor.ApplyWritesRequest\x1a\x1d.executor.ApplyWritesResponse\x12Y\n\x10GenerateCacheKey\x12!.executor.GenerateCacheKeyRequest\x1a\".executor.GenerateCacheKeyResponse\x12\x41\n\x0bStateUpdate\x12\x1c.executor.StateUpdateRequest\x1a\x14.executor.TaskResultB1Z/github.com/langchain-ai/langgraph-go/runtime/pbb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -39,46 +39,48 @@ if not _descriptor._USE_C_DESCRIPTORS:
   _globals['_CHANNELSFROMCHECKPOINTREQUEST_CHECKPOINTCHANNELVALUESENTRY']._serialized_options = b'8\001'
   _globals['_GENERATECACHEKEYREQUEST_CHANNELVALUESENTRY']._loaded_options = None
   _globals['_GENERATECACHEKEYREQUEST_CHANNELVALUESENTRY']._serialized_options = b'8\001'
-  _globals['_LISTGRAPHSREQUEST']._serialized_start=84
-  _globals['_LISTGRAPHSREQUEST']._serialized_end=103
-  _globals['_LISTGRAPHSRESPONSE']._serialized_start=105
-  _globals['_LISTGRAPHSRESPONSE']._serialized_end=146
-  _globals['_GETGRAPHREQUEST']._serialized_start=148
-  _globals['_GETGRAPHREQUEST']._serialized_end=185
-  _globals['_GETGRAPHRESPONSE']._serialized_start=188
-  _globals['_GETGRAPHRESPONSE']._serialized_end=342
-  _globals['_GETALLGRAPHSREQUEST']._serialized_start=344
-  _globals['_GETALLGRAPHSREQUEST']._serialized_end=365
-  _globals['_CHANNELSFROMCHECKPOINTREQUEST']._serialized_start=368
-  _globals['_CHANNELSFROMCHECKPOINTREQUEST']._serialized_end=634
-  _globals['_CHANNELSFROMCHECKPOINTREQUEST_CHECKPOINTCHANNELVALUESENTRY']._serialized_start=558
-  _globals['_CHANNELSFROMCHECKPOINTREQUEST_CHECKPOINTCHANNELVALUESENTRY']._serialized_end=634
-  _globals['_CHANNELSFROMCHECKPOINTRESPONSE']._serialized_start=636
-  _globals['_CHANNELSFROMCHECKPOINTRESPONSE']._serialized_end=703
-  _globals['_EXECUTETASKREQUEST']._serialized_start=706
-  _globals['_EXECUTETASKREQUEST']._serialized_end=842
-  _globals['_EXECUTETASKINIT']._serialized_start=845
-  _globals['_EXECUTETASKINIT']._serialized_end=1059
-  _globals['_EXECUTETASKRESPONSE']._serialized_start=1062
-  _globals['_EXECUTETASKRESPONSE']._serialized_end=1311
-  _globals['_TASKRESULT']._serialized_start=1313
-  _globals['_TASKRESULT']._serialized_end=1407
-  _globals['_CUSTOMSTREAMEVENT']._serialized_start=1409
-  _globals['_CUSTOMSTREAMEVENT']._serialized_end=1470
-  _globals['_APPLYWRITESREQUEST']._serialized_start=1473
-  _globals['_APPLYWRITESREQUEST']._serialized_end=1637
-  _globals['_APPLYWRITESRESPONSE']._serialized_start=1639
-  _globals['_APPLYWRITESRESPONSE']._serialized_end=1737
-  _globals['_GENERATECACHEKEYREQUEST']._serialized_start=1740
-  _globals['_GENERATECACHEKEYREQUEST']._serialized_end=1940
-  _globals['_GENERATECACHEKEYREQUEST_CHANNELVALUESENTRY']._serialized_start=1888
-  _globals['_GENERATECACHEKEYREQUEST_CHANNELVALUESENTRY']._serialized_end=1940
-  _globals['_GENERATECACHEKEYRESPONSE']._serialized_start=1942
-  _globals['_GENERATECACHEKEYRESPONSE']._serialized_end=1987
-  _globals['_CACHECHECKREQUEST']._serialized_start=1989
-  _globals['_CACHECHECKREQUEST']._serialized_end=2078
-  _globals['_CACHECHECKRESPONSE']._serialized_start=2080
-  _globals['_CACHECHECKRESPONSE']._serialized_end=2119
-  _globals['_LANGGRAPHEXECUTOR']._serialized_start=2122
-  _globals['_LANGGRAPHEXECUTOR']._serialized_end=2714
+  _globals['_STATEUPDATEREQUEST']._serialized_start=85
+  _globals['_STATEUPDATEREQUEST']._serialized_end=279
+  _globals['_LISTGRAPHSREQUEST']._serialized_start=281
+  _globals['_LISTGRAPHSREQUEST']._serialized_end=300
+  _globals['_LISTGRAPHSRESPONSE']._serialized_start=302
+  _globals['_LISTGRAPHSRESPONSE']._serialized_end=343
+  _globals['_GETGRAPHREQUEST']._serialized_start=345
+  _globals['_GETGRAPHREQUEST']._serialized_end=382
+  _globals['_GETGRAPHRESPONSE']._serialized_start=385
+  _globals['_GETGRAPHRESPONSE']._serialized_end=539
+  _globals['_GETALLGRAPHSREQUEST']._serialized_start=541
+  _globals['_GETALLGRAPHSREQUEST']._serialized_end=562
+  _globals['_CHANNELSFROMCHECKPOINTREQUEST']._serialized_start=565
+  _globals['_CHANNELSFROMCHECKPOINTREQUEST']._serialized_end=831
+  _globals['_CHANNELSFROMCHECKPOINTREQUEST_CHECKPOINTCHANNELVALUESENTRY']._serialized_start=755
+  _globals['_CHANNELSFROMCHECKPOINTREQUEST_CHECKPOINTCHANNELVALUESENTRY']._serialized_end=831
+  _globals['_CHANNELSFROMCHECKPOINTRESPONSE']._serialized_start=833
+  _globals['_CHANNELSFROMCHECKPOINTRESPONSE']._serialized_end=900
+  _globals['_EXECUTETASKREQUEST']._serialized_start=903
+  _globals['_EXECUTETASKREQUEST']._serialized_end=1039
+  _globals['_EXECUTETASKINIT']._serialized_start=1042
+  _globals['_EXECUTETASKINIT']._serialized_end=1256
+  _globals['_EXECUTETASKRESPONSE']._serialized_start=1259
+  _globals['_EXECUTETASKRESPONSE']._serialized_end=1508
+  _globals['_TASKRESULT']._serialized_start=1510
+  _globals['_TASKRESULT']._serialized_end=1604
+  _globals['_CUSTOMSTREAMEVENT']._serialized_start=1606
+  _globals['_CUSTOMSTREAMEVENT']._serialized_end=1667
+  _globals['_APPLYWRITESREQUEST']._serialized_start=1670
+  _globals['_APPLYWRITESREQUEST']._serialized_end=1834
+  _globals['_APPLYWRITESRESPONSE']._serialized_start=1836
+  _globals['_APPLYWRITESRESPONSE']._serialized_end=1934
+  _globals['_GENERATECACHEKEYREQUEST']._serialized_start=1937
+  _globals['_GENERATECACHEKEYREQUEST']._serialized_end=2137
+  _globals['_GENERATECACHEKEYREQUEST_CHANNELVALUESENTRY']._serialized_start=2085
+  _globals['_GENERATECACHEKEYREQUEST_CHANNELVALUESENTRY']._serialized_end=2137
+  _globals['_GENERATECACHEKEYRESPONSE']._serialized_start=2139
+  _globals['_GENERATECACHEKEYRESPONSE']._serialized_end=2184
+  _globals['_CACHECHECKREQUEST']._serialized_start=2186
+  _globals['_CACHECHECKREQUEST']._serialized_end=2275
+  _globals['_CACHECHECKRESPONSE']._serialized_start=2277
+  _globals['_CACHECHECKRESPONSE']._serialized_end=2316
+  _globals['_LANGGRAPHEXECUTOR']._serialized_start=2319
+  _globals['_LANGGRAPHEXECUTOR']._serialized_end=2978
 # @@protoc_insertion_point(module_scope)

langgraph_executor/pb/executor_pb2.pyi

@@ -15,6 +15,45 @@ import typing
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
+@typing.final
+class StateUpdateRequest(google.protobuf.message.Message):
+    """StateUpdate"""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    NODE_NAME_FIELD_NUMBER: builtins.int
+    TASK_ID_FIELD_NUMBER: builtins.int
+    GRAPH_NAME_FIELD_NUMBER: builtins.int
+    STEP_FIELD_NUMBER: builtins.int
+    CHANNELS_FIELD_NUMBER: builtins.int
+    VALUES_FIELD_NUMBER: builtins.int
+    CONFIG_FIELD_NUMBER: builtins.int
+    node_name: builtins.str
+    task_id: builtins.str
+    graph_name: builtins.str
+    step: builtins.int
+    @property
+    def channels(self) -> types_pb2.Channels: ...
+    @property
+    def values(self) -> types_pb2.Value: ...
+    @property
+    def config(self) -> types_pb2.RunnableConfig: ...
+    def __init__(
+        self,
+        *,
+        node_name: builtins.str = ...,
+        task_id: builtins.str = ...,
+        graph_name: builtins.str = ...,
+        step: builtins.int = ...,
+        channels: types_pb2.Channels | None = ...,
+        values: types_pb2.Value | None = ...,
+        config: types_pb2.RunnableConfig | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["channels", b"channels", "config", b"config", "values", b"values"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["channels", b"channels", "config", b"config", "graph_name", b"graph_name", "node_name", b"node_name", "step", b"step", "task_id", b"task_id", "values", b"values"]) -> None: ...
+
+global___StateUpdateRequest = StateUpdateRequest
+
 @typing.final
 class ListGraphsRequest(google.protobuf.message.Message):
     """ListGraphs"""

langgraph_executor/pb/executor_pb2_grpc.py

@@ -70,6 +70,11 @@ class LangGraphExecutorStub(object):
                 request_serializer=executor__pb2.GenerateCacheKeyRequest.SerializeToString,
                 response_deserializer=executor__pb2.GenerateCacheKeyResponse.FromString,
                 _registered_method=True)
+        self.StateUpdate = channel.unary_unary(
+                '/executor.LangGraphExecutor/StateUpdate',
+                request_serializer=executor__pb2.StateUpdateRequest.SerializeToString,
+                response_deserializer=executor__pb2.TaskResult.FromString,
+                _registered_method=True)
 
 
 class LangGraphExecutorServicer(object):
@@ -125,6 +130,13 @@ class LangGraphExecutorServicer(object):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def StateUpdate(self, request, context):
+        """State update 
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
 
 def add_LangGraphExecutorServicer_to_server(servicer, server):
     rpc_method_handlers = {
@@ -163,6 +175,11 @@ def add_LangGraphExecutorServicer_to_server(servicer, server):
                     request_deserializer=executor__pb2.GenerateCacheKeyRequest.FromString,
                     response_serializer=executor__pb2.GenerateCacheKeyResponse.SerializeToString,
             ),
+            'StateUpdate': grpc.unary_unary_rpc_method_handler(
+                    servicer.StateUpdate,
+                    request_deserializer=executor__pb2.StateUpdateRequest.FromString,
+                    response_serializer=executor__pb2.TaskResult.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'executor.LangGraphExecutor', rpc_method_handlers)
@@ -363,3 +380,30 @@ class LangGraphExecutor(object):
             timeout,
             metadata,
             _registered_method=True)
+
+    @staticmethod
+    def StateUpdate(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/executor.LangGraphExecutor/StateUpdate',
+            executor__pb2.StateUpdateRequest.SerializeToString,
+            executor__pb2.TaskResult.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

langgraph_executor/pb/executor_pb2_grpc.pyi

@@ -63,6 +63,12 @@ class LangGraphExecutorStub:
     ]
     """Generate cache key for a node execution"""
 
+    StateUpdate: grpc.UnaryUnaryMultiCallable[
+        executor_pb2.StateUpdateRequest,
+        executor_pb2.TaskResult,
+    ]
+    """State update"""
+
 class LangGraphExecutorAsyncStub:
     """Main service for LangGraph runtime operations"""
 
@@ -108,6 +114,12 @@ class LangGraphExecutorAsyncStub:
     ]
     """Generate cache key for a node execution"""
 
+    StateUpdate: grpc.aio.UnaryUnaryMultiCallable[
+        executor_pb2.StateUpdateRequest,
+        executor_pb2.TaskResult,
+    ]
+    """State update"""
+
 class LangGraphExecutorServicer(metaclass=abc.ABCMeta):
     """Main service for LangGraph runtime operations"""
 
@@ -167,4 +179,12 @@ class LangGraphExecutorServicer(metaclass=abc.ABCMeta):
     ) -> typing.Union[executor_pb2.GenerateCacheKeyResponse, collections.abc.Awaitable[executor_pb2.GenerateCacheKeyResponse]]:
         """Generate cache key for a node execution"""
 
+    @abc.abstractmethod
+    def StateUpdate(
+        self,
+        request: executor_pb2.StateUpdateRequest,
+        context: _ServicerContext,
+    ) -> typing.Union[executor_pb2.TaskResult, collections.abc.Awaitable[executor_pb2.TaskResult]]:
+        """State update"""
+
 def add_LangGraphExecutorServicer_to_server(servicer: LangGraphExecutorServicer, server: typing.Union[grpc.Server, grpc.aio.Server]) -> None: ...