WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-27 12:04:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: enhance runtime w/ more execution information (#7363)

Browse Source 
## Summary

Enhances `ExecutionInfo` and `Runtime` to surface richer execution
context and introduces `ServerInfo` for LangGraph Server metadata.

### `ExecutionInfo` expansion

Converted from `NamedTuple` to a frozen `dataclass`. Added identity
fields populated during task preparation in `_algo.py`:

| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `checkpoint_id` | `str` | required | Checkpoint ID for the current
execution |
| `checkpoint_ns` | `str` | required | Checkpoint namespace for the
current execution |
| `task_id` | `str` | required | Task ID for the current execution |
| `thread_id` | `str \| None` | `None` | Thread ID (None without a
checkpointer) |
| `run_id` | `str \| None` | `None` | Run ID (None when not provided in
config) |
| `node_attempt` | `int` | `1` | Current node execution attempt number
(1-indexed) |
| `node_first_attempt_time` | `float \| None` | `None` | Unix timestamp
for when the first attempt started |

`checkpoint_id`, `checkpoint_ns`, and `task_id` are required (no
defaults) — they are always populated during task preparation in
`_algo.py`. `Runtime.execution_info` is `None` until that point.

### New `ServerInfo` type

Frozen dataclass with `assistant_id: str`, `graph_id: str`, and optional
`user: BaseUser | None`. Populated from config metadata (`assistant_id`,
`graph_id`) and `configurable["langgraph_auth_user"]` via
`_build_server_info()` in `pregel/main.py`.

User detection uses `isinstance(BaseUser)` with a `hasattr("identity")`
fallback — needed because the server's `ProxyUser` provides
`permissions` via `__getattr__`, which Python's `runtime_checkable`
Protocol check doesn't see.

### `Runtime` changes

- `execution_info` is now `ExecutionInfo | None` (default `None`), set
during task prep
- Added `server_info: ServerInfo | None` field, wired through `merge()`
and `override()`

### `ToolNode` / `ToolRuntime` forwarding

`execution_info` and `server_info` are forwarded from `Runtime` to
`ToolRuntime` so tools can access execution and server context.

### New public API surface

```python
from langgraph.runtime import BaseUser, ExecutionInfo, Runtime, ServerInfo, get_runtime
```

## Test plan

- [x] `ExecutionInfo` defaults, patch, and frozen behavior
- [x] Integration tests verifying identity fields are populated in sync
and async execution
- [x] Retry tests confirming identity fields persist and `node_attempt`
increments
- [x] `ServerInfo` construction, frozen behavior, and `Runtime.merge`
precedence
- [x] `server_info` populated from config metadata and
`langgraph_auth_user` (including starlette-style proxy user)
- [x] `server_info` is `None` when no server metadata present
- [x] `ToolRuntime` forwarding of `execution_info` and `server_info`

🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
Sydney Runkle
2026-04-02 19:38:03 +00:00
committed by GitHub
parent 9f5f3fb35e
commit 39288e6111
7 changed files with 569 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/_algo.py
+28 -3
View File
@@ -45,6 +45,7 @@ from langgraph._internal._constants import (
CONFIG_KEY_SCRATCHPAD,
CONFIG_KEY_SEND,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_THREAD_ID,
ERROR,
INTERRUPT,
NO_WRITES,
@@ -70,7 +71,7 @@ from langgraph.pregel._call import get_runnable_for_task, identifier
from langgraph.pregel._io import read_channels
from langgraph.pregel._log import logger
from langgraph.pregel._read import INPUT_CACHE_KEY_TYPE, PregelNode
from langgraph.runtime import DEFAULT_RUNTIME, Runtime
from langgraph.runtime import DEFAULT_RUNTIME, ExecutionInfo, Runtime
from langgraph.types import (
All,
CacheKey,
@@ -668,6 +669,13 @@ def prepare_single_task(
runtime = runtime.override(
previous=checkpoint["channel_values"].get(PREVIOUS, None),
store=store,
execution_info=ExecutionInfo(
checkpoint_id=checkpoint["id"],
checkpoint_ns=task_checkpoint_ns,
task_id=task_id,
thread_id=configurable.get(CONFIG_KEY_THREAD_ID),
run_id=str(rid) if (rid := config.get("run_id")) else None,
),
)
additional_config = {
"metadata": metadata,
@@ -813,7 +821,16 @@ def prepare_push_task_functional(
stop,
)
runtime = cast(Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME))
runtime = runtime.override(store=store)
runtime = runtime.override(
store=store,
execution_info=ExecutionInfo(
checkpoint_id=checkpoint["id"],
checkpoint_ns=task_checkpoint_ns,
task_id=task_id,
thread_id=configurable.get(CONFIG_KEY_THREAD_ID),
run_id=str(rid) if (rid := config.get("run_id")) else None,
),
)
return PregelExecutableTask(
name,
call.input,
@@ -966,7 +983,15 @@ def prepare_push_task_send(
)
runtime = cast(Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME))
runtime = runtime.override(
store=store, previous=checkpoint["channel_values"].get(PREVIOUS, None)
store=store,
previous=checkpoint["channel_values"].get(PREVIOUS, None),
execution_info=ExecutionInfo(
checkpoint_id=checkpoint["id"],
checkpoint_ns=task_checkpoint_ns,
task_id=task_id,
thread_id=configurable.get(CONFIG_KEY_THREAD_ID),
run_id=str(rid) if (rid := config.get("run_id")) else None,
),
)
additional_config: RunnableConfig = {
"metadata": metadata,
libs/langgraph/langgraph/pregel/main.py
+50 -5
View File
@@ -136,7 +136,12 @@ from langgraph.pregel._validate import validate_graph, validate_keys
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
from langgraph.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes
from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtocol
from langgraph.runtime import DEFAULT_RUNTIME, ExecutionInfo, Runtime
from langgraph.runtime import (
DEFAULT_RUNTIME,
BaseUser,
Runtime,
ServerInfo,
)
from langgraph.types import (
All,
CachePolicy,
@@ -2645,14 +2650,18 @@ class Pregel(
if durability is not None:
config[CONF][CONFIG_KEY_DURABILITY] = durability_
# build server_info from metadata + parent runtime
parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
server_info = _build_server_info(config, parent_runtime)
runtime = Runtime(
context=_coerce_context(self.context_schema, context),
store=store,
stream_writer=stream_writer,
previous=None,
execution_info=ExecutionInfo(),
execution_info=None,
server_info=server_info,
)
parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
runtime = parent_runtime.merge(runtime)
config[CONF][CONFIG_KEY_RUNTIME] = runtime
@@ -3014,14 +3023,18 @@ class Pregel(
if durability is not None:
config[CONF][CONFIG_KEY_DURABILITY] = durability_
# build server_info from metadata + parent runtime
parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
server_info = _build_server_info(config, parent_runtime)
runtime = Runtime(
context=_coerce_context(self.context_schema, context),
store=store,
stream_writer=stream_writer,
previous=None,
execution_info=ExecutionInfo(),
execution_info=None,
server_info=server_info,
)
parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
runtime = parent_runtime.merge(runtime)
config[CONF][CONFIG_KEY_RUNTIME] = runtime
@@ -3643,6 +3656,38 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
payload["values"] = mapper(payload["values"])
def _build_server_info(
config: RunnableConfig, parent_runtime: Runtime[Any]
) -> ServerInfo | None:
"""Build ServerInfo from config metadata and configurable.
The server puts assistant_id/graph_id in config metadata and the
authenticated user dict in configurable["langgraph_auth_user"].
"""
metadata = config.get("metadata") or {}
configurable = config.get(CONF) or {}
assistant_id = metadata.get("assistant_id")
graph_id = metadata.get("graph_id")
# Read authenticated user from configurable (set by LangGraph Server).
# We prefer isinstance(BaseUser) but fall back to hasattr("identity")
# because the server's ProxyUser provides `permissions` via __getattr__,
# which Python's runtime_checkable Protocol check doesn't see.
auth_user_data = configurable.get("langgraph_auth_user")
user: BaseUser | None = None
if auth_user_data is not None:
if isinstance(auth_user_data, BaseUser) or hasattr(auth_user_data, "identity"):
user = cast(BaseUser, auth_user_data)
if assistant_id is not None or graph_id is not None or user is not None:
return ServerInfo(
assistant_id=str(assistant_id) if assistant_id else "",
graph_id=str(graph_id) if graph_id else "",
user=user,
)
return None
def _coerce_context(
context_schema: type[ContextT] | None, context: Any
) -> ContextT | None:
libs/langgraph/langgraph/runtime.py
+64 -8
View File
@@ -1,9 +1,10 @@
from __future__ import annotations
from dataclasses import dataclass, field, replace
from typing import Any, Generic, NamedTuple, cast
from typing import Any, Generic, cast
from langgraph.store.base import BaseStore
from langgraph_sdk.auth.types import BaseUser
from typing_extensions import TypedDict, Unpack
from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME
@@ -11,12 +12,38 @@ from langgraph.config import get_config
from langgraph.types import _DC_KWARGS, StreamWriter
from langgraph.typing import ContextT
__all__ = ("ExecutionInfo", "Runtime", "get_runtime")
__all__ = (
"BaseUser",
"ExecutionInfo",
"Runtime",
"ServerInfo",
"get_runtime",
)
class ExecutionInfo(NamedTuple):
@dataclass(frozen=True, slots=True)
class ExecutionInfo:
"""Read-only execution info/metadata for the execution of current thread/run/node."""
checkpoint_id: str
"""The checkpoint ID for the current execution."""
checkpoint_ns: str
"""The checkpoint namespace for the current execution."""
task_id: str
"""The task ID for the current execution."""
thread_id: str | None = None
"""The thread ID for the current execution.
None when running without a checkpointer (i.e., no persistence)."""
run_id: str | None = None
"""The run ID for the current execution.
None when `run_id` is not provided in the RunnableConfig."""
node_attempt: int = 1
"""Current node execution attempt number (1-indexed)."""
@@ -25,7 +52,26 @@ class ExecutionInfo(NamedTuple):
def patch(self, **overrides: Any) -> ExecutionInfo:
"""Return a new execution info object with selected fields replaced."""
return self._replace(**overrides)
return replace(self, **overrides)
@dataclass(frozen=True, slots=True)
class ServerInfo:
"""Metadata injected by LangGraph Server. None when running open-source LangGraph without LangSmith deployments."""
assistant_id: str
"""The assistant ID for the current execution."""
graph_id: str
"""The graph ID for the current execution."""
user: BaseUser | None = None
"""The authenticated user, if any.
This implements the `BaseUser` protocol from `langgraph_sdk.auth.types`,
which supports both attribute access (e.g. `user.identity`) and dict-like
access (e.g. `user["identity"]`).
"""
def _no_op_stream_writer(_: Any) -> None: ...
@@ -37,6 +83,7 @@ class _RuntimeOverrides(TypedDict, Generic[ContextT], total=False):
stream_writer: StreamWriter
previous: Any
execution_info: ExecutionInfo
server_info: ServerInfo | None
@dataclass(**_DC_KWARGS)
@@ -130,8 +177,13 @@ class Runtime(Generic[ContextT]):
Only available with the functional API when a checkpointer is provided.
"""
execution_info: ExecutionInfo = field(default_factory=ExecutionInfo)
"""Read-only execution information/metadata for the current node run."""
execution_info: ExecutionInfo | None = field(default=None)
"""Read-only execution information/metadata for the current node run.
None before task preparation populates it."""
server_info: ServerInfo | None = field(default=None)
"""Metadata injected by LangGraph Server. None when running open-source LangGraph without LangSmith deployments."""
def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]:
"""Merge two runtimes together.
@@ -145,7 +197,8 @@ class Runtime(Generic[ContextT]):
if other.stream_writer is not _no_op_stream_writer
else self.stream_writer,
previous=self.previous if other.previous is None else other.previous,
execution_info=other.execution_info,
execution_info=other.execution_info or self.execution_info,
server_info=other.server_info or self.server_info,
)
def override(
@@ -156,6 +209,9 @@ class Runtime(Generic[ContextT]):
def patch_execution_info(self, **overrides: Any) -> Runtime[ContextT]:
"""Return a new runtime with selected execution_info fields replaced."""
if self.execution_info is None:
msg = "Cannot patch execution_info before it has been set"
raise RuntimeError(msg)
return replace(
self,
execution_info=self.execution_info.patch(**overrides),
@@ -167,7 +223,7 @@ DEFAULT_RUNTIME = Runtime(
store=None,
stream_writer=_no_op_stream_writer,
previous=None,
execution_info=ExecutionInfo(),
execution_info=None,
)
libs/langgraph/tests/test_retry.py
+63
View File
@@ -1,6 +1,7 @@
from unittest.mock import Mock, patch
import pytest
from langgraph.checkpoint.memory import MemorySaver
from typing_extensions import TypedDict
from langgraph.graph import START, StateGraph
@@ -393,3 +394,65 @@ def test_graph_with_max_attempts_exceeded():
graph.invoke({"foo": ""})
mock_sleep.assert_called_with(0.01)
def test_execution_info_identity_fields_populated_on_retry():
"""Test that thread_id, task_id, run_id, etc. are populated in execution_info during retries."""
class State(TypedDict):
foo: str
attempt_count = 0
captured_infos: list[dict] = []
def failing_node(state: State, runtime: Runtime):
nonlocal attempt_count
attempt_count += 1
info = runtime.execution_info
captured_infos.append(
{
"thread_id": info.thread_id,
"run_id": info.run_id,
"node_attempt": info.node_attempt,
"node_first_attempt_time": info.node_first_attempt_time,
"checkpoint_ns": info.checkpoint_ns,
}
)
if attempt_count < 2:
raise ValueError("Intentional failure")
return {"foo": "success"}
retry_policy = RetryPolicy(
max_attempts=3,
initial_interval=0.01,
jitter=False,
retry_on=ValueError,
)
graph = (
StateGraph(State)
.add_node("failing_node", failing_node, retry_policy=retry_policy)
.add_edge(START, "failing_node")
.compile(checkpointer=MemorySaver())
)
with patch("time.sleep"):
result = graph.invoke(
{"foo": ""},
config={"configurable": {"thread_id": "retry-thread"}},
)
assert result["foo"] == "success"
assert len(captured_infos) == 2
# Both attempts should have the same thread_id and first_attempt_time
assert captured_infos[0]["thread_id"] == "retry-thread"
assert captured_infos[1]["thread_id"] == "retry-thread"
assert (
captured_infos[0]["node_first_attempt_time"]
== captured_infos[1]["node_first_attempt_time"]
)
# node_attempt should increment
assert captured_infos[0]["node_attempt"] == 1
assert captured_infos[1]["node_attempt"] == 2
libs/langgraph/tests/test_runtime.py
+258 -1
View File
@@ -2,11 +2,12 @@ from dataclasses import dataclass
from typing import Any
import pytest
from langgraph.checkpoint.memory import MemorySaver
from pydantic import BaseModel, ValidationError
from typing_extensions import TypedDict
from langgraph.graph import END, START, StateGraph
from langgraph.runtime import Runtime, get_runtime
from langgraph.runtime import ExecutionInfo, Runtime, ServerInfo, get_runtime
def test_injected_runtime() -> None:
@@ -389,3 +390,259 @@ def test_context_coercion_pydantic_validation_errors() -> None:
compiled.invoke(
{"message": "test"}, context={"api_key": "sk_test", "timeout": "not_an_int"}
)
# --- ExecutionInfo unit tests ---
def test_execution_info_defaults_and_patch() -> None:
info = ExecutionInfo(checkpoint_id="c1", checkpoint_ns="ns1", task_id="t1")
assert info.checkpoint_id == "c1"
assert info.checkpoint_ns == "ns1"
assert info.task_id == "t1"
assert info.thread_id is None
assert info.run_id is None
assert info.node_attempt == 1
assert info.node_first_attempt_time is None
# patch returns new instance, original unchanged
patched = info.patch(thread_id="th1", node_attempt=3, task_id="tk1")
assert patched.thread_id == "th1"
assert patched.node_attempt == 3
assert patched.task_id == "tk1"
assert info.node_attempt == 1
assert info.task_id == "t1"
# frozen
with pytest.raises(AttributeError):
info.thread_id = "t2" # type: ignore[misc]
# --- ServerInfo / Runtime unit tests ---
def test_server_info_and_runtime_merge() -> None:
si = ServerInfo(assistant_id="asst-1", graph_id="graph-1")
assert si.assistant_id == "asst-1"
assert si.user is None
# frozen
with pytest.raises(AttributeError):
si.assistant_id = "asst-2" # type: ignore[misc]
# runtime default is None
assert Runtime().server_info is None
# merge preserves server_info from self when other has None
r1 = Runtime(server_info=si)
merged = r1.merge(Runtime())
assert merged.server_info is si
# merge takes server_info from other when present
si2 = ServerInfo(assistant_id="asst-2", graph_id="graph-2")
merged2 = r1.merge(Runtime(server_info=si2))
assert merged2.server_info is si2
# --- Integration tests ---
def _make_capture_graph(
capture: dict[str, Any],
*,
checkpointer: Any = None,
) -> Any:
"""Helper: build a simple graph that captures runtime info."""
class State(TypedDict):
message: str
def capture_node(state: State, runtime: Runtime) -> dict[str, Any]:
capture["execution_info"] = runtime.execution_info
capture["server_info"] = runtime.server_info
return {"message": "done"}
graph = StateGraph(state_schema=State)
graph.add_node("capture", capture_node)
graph.add_edge(START, "capture")
graph.add_edge("capture", END)
return graph.compile(checkpointer=checkpointer)
def test_execution_info_populated_in_graph() -> None:
"""execution_info fields are populated when running with a checkpointer."""
captured: dict[str, Any] = {}
compiled = _make_capture_graph(captured, checkpointer=MemorySaver())
compiled.invoke(
{"message": "hi"},
config={"configurable": {"thread_id": "t-123"}},
)
info = captured["execution_info"]
assert info.thread_id == "t-123"
assert info.task_id is not None
assert info.checkpoint_id is not None
assert info.checkpoint_ns is not None
assert info.node_attempt == 1
assert isinstance(info.node_first_attempt_time, float)
@pytest.mark.anyio
async def test_execution_info_populated_in_graph_async() -> None:
"""execution_info fields are populated in async execution."""
captured: dict[str, Any] = {}
compiled = _make_capture_graph(captured, checkpointer=MemorySaver())
await compiled.ainvoke(
{"message": "hi"},
config={"configurable": {"thread_id": "t-xyz"}},
)
info = captured["execution_info"]
assert info.thread_id == "t-xyz"
assert info.node_attempt == 1
assert isinstance(info.node_first_attempt_time, float)
def test_server_info_from_metadata() -> None:
"""server_info is built from assistant_id/graph_id in config metadata."""
captured: dict[str, Any] = {}
compiled = _make_capture_graph(captured)
compiled.invoke(
{"message": "hi"},
config={"metadata": {"assistant_id": "asst-abc", "graph_id": "my-graph"}},
)
si = captured["server_info"]
assert si is not None
assert si.assistant_id == "asst-abc"
assert si.graph_id == "my-graph"
assert si.user is None
def test_server_info_none_without_metadata() -> None:
"""server_info is None when no assistant_id/graph_id in metadata."""
captured: dict[str, Any] = {}
compiled = _make_capture_graph(captured)
compiled.invoke({"message": "hi"})
assert captured["server_info"] is None
def test_server_info_user_from_auth_user() -> None:
"""server_info.user is populated from configurable['langgraph_auth_user'].
Tests both a proper BaseUser protocol object and a starlette-style proxy
that provides `permissions` via __getattr__ (which the Protocol isinstance
check may not see).
"""
class _ProxyUser:
"""Mimics langgraph_api's ProxyUser: identity/display_name as properties,
permissions via __getattr__."""
def __init__(self, data: dict[str, Any]) -> None:
self._data = data
@property
def identity(self) -> str:
return self._data["identity"]
@property
def display_name(self) -> str:
return self._data.get("display_name", self.identity)
@property
def is_authenticated(self) -> bool:
return True
def __getattr__(self, name: str) -> Any:
return self._data[name]
def __getitem__(self, key: str) -> Any:
return self._data[key]
def __contains__(self, key: str) -> bool:
return key in self._data
def __iter__(self) -> Any:
return iter(self._data)
proxy = _ProxyUser(
{
"identity": "proxy-user",
"display_name": "Proxy User",
"is_authenticated": True,
"permissions": ["read"],
}
)
assert not isinstance(proxy, dict)
assert hasattr(proxy, "identity")
captured: dict[str, Any] = {}
compiled = _make_capture_graph(captured)
compiled.invoke(
{"message": "hi"},
config={
"configurable": {"langgraph_auth_user": proxy},
"metadata": {"assistant_id": "asst-proxy", "graph_id": "graph-proxy"},
},
)
si = captured["server_info"]
assert si is not None
assert si.assistant_id == "asst-proxy"
assert si.user is not None
assert si.user.identity == "proxy-user"
assert si.user["display_name"] == "Proxy User"
def test_execution_info_inherited_by_subgraph() -> None:
"""execution_info is correctly populated for subgraph nodes, including namespace."""
captured_main: dict[str, Any] = {}
captured_sub: dict[str, Any] = {}
class State(TypedDict, total=False):
message: str
def subgraph_node(state: State, runtime: Runtime) -> dict[str, str]:
captured_sub["execution_info"] = runtime.execution_info
return {"message": "from_sub"}
subgraph_builder = StateGraph(State)
subgraph_builder.add_node("sub_node", subgraph_node)
subgraph_builder.add_edge(START, "sub_node")
subgraph = subgraph_builder.compile()
def main_node(state: State, runtime: Runtime) -> dict[str, str]:
captured_main["execution_info"] = runtime.execution_info
return {"message": "from_main"}
builder = StateGraph(State)
builder.add_node("main_node", main_node)
builder.add_node("subgraph", subgraph)
builder.add_edge(START, "main_node")
builder.add_edge("main_node", "subgraph")
graph = builder.compile(checkpointer=MemorySaver())
graph.invoke(
{"message": "hi"},
config={"configurable": {"thread_id": "sub-thread"}},
)
main_info = captured_main["execution_info"]
sub_info = captured_sub["execution_info"]
# Both share the same thread_id
assert main_info.thread_id == "sub-thread"
assert sub_info.thread_id == "sub-thread"
# Both have node_attempt = 1
assert main_info.node_attempt == 1
assert sub_info.node_attempt == 1
# Main namespace is "main_node:<task_id>" (top-level, no separator)
assert main_info.checkpoint_ns.startswith("main_node:")
assert "|" not in main_info.checkpoint_ns
# Subgraph namespace is "subgraph:<task_id>|sub_node:<task_id>" (nested)
assert sub_info.checkpoint_ns.startswith("subgraph:")
assert "|sub_node:" in sub_info.checkpoint_ns
# task_id appears in its own namespace segment
assert main_info.task_id in main_info.checkpoint_ns
assert sub_info.task_id in sub_info.checkpoint_ns
libs/prebuilt/langgraph/prebuilt/tool_node.py
+7
View File
@@ -85,6 +85,7 @@ from langchain_core.tools.base import (
from langgraph._internal._runnable import RunnableCallable
from langgraph.errors import GraphBubbleUp
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.runtime import ExecutionInfo, ServerInfo # noqa: TC002
from langgraph.store.base import BaseStore # noqa: TC002
from langgraph.types import Command, Send, StreamWriter
from pydantic import BaseModel, ValidationError
@@ -806,6 +807,8 @@ class ToolNode(RunnableCallable):
context=runtime.context,
store=runtime.store,
stream_writer=runtime.stream_writer,
execution_info=runtime.execution_info,
server_info=runtime.server_info,
)
tool_runtimes.append(tool_runtime)
@@ -838,6 +841,8 @@ class ToolNode(RunnableCallable):
context=runtime.context,
store=runtime.store,
stream_writer=runtime.stream_writer,
execution_info=runtime.execution_info,
server_info=runtime.server_info,
)
tool_runtimes.append(tool_runtime)
@@ -1608,6 +1613,8 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
stream_writer: StreamWriter
tool_call_id: str | None
store: BaseStore | None
execution_info: ExecutionInfo | None = None
server_info: ServerInfo | None = None
class InjectedState(InjectedToolArg):
libs/prebuilt/tests/test_tool_node.py
+99
View File
@@ -59,10 +59,16 @@ def _create_mock_runtime(store: BaseStore | None = None) -> Mock:
which is injected by RunnableCallable from config["configurable"]["__pregel_runtime"].
When testing ToolNode directly (outside a graph), we need to provide this manually.
"""
from langgraph.runtime import ExecutionInfo
mock_runtime = Mock()
mock_runtime.store = store
mock_runtime.context = None
mock_runtime.stream_writer = lambda *args, **kwargs: None
mock_runtime.execution_info = ExecutionInfo(
checkpoint_id="test-cp", checkpoint_ns="", task_id="test-task"
)
mock_runtime.server_info = None
return mock_runtime
@@ -2010,6 +2016,99 @@ async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async()
assert tool_message.tool_call_id == "call_dynamic_2"
def test_tool_runtime_forwards_execution_info_and_server_info() -> None:
"""Test that execution_info and server_info are forwarded from Runtime to ToolRuntime."""
from langgraph.runtime import ExecutionInfo, ServerInfo
exec_info = ExecutionInfo(
thread_id="t-1",
checkpoint_id="cp-1",
checkpoint_ns="",
task_id="tk-1",
run_id="r-1",
)
server_info = ServerInfo(assistant_id="asst-1", graph_id="graph-1")
mock_runtime = Mock()
mock_runtime.store = None
mock_runtime.context = None
mock_runtime.stream_writer = lambda *args, **kwargs: None
mock_runtime.execution_info = exec_info
mock_runtime.server_info = server_info
captured: dict = {}
@dec_tool
def info_tool(x: int, runtime: ToolRuntime) -> str:
"""Tool that captures runtime info."""
captured["execution_info"] = runtime.execution_info
captured["server_info"] = runtime.server_info
return "ok"
node = ToolNode([info_tool])
tool_call = {
"name": "info_tool",
"args": {"x": 1},
"id": "call-1",
"type": "tool_call",
}
msg = AIMessage("", tool_calls=[tool_call])
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
node.invoke({"messages": [msg]}, config=config)
assert captured["execution_info"] is exec_info
assert captured["execution_info"].thread_id == "t-1"
assert captured["execution_info"].task_id == "tk-1"
assert captured["server_info"] is server_info
assert captured["server_info"].assistant_id == "asst-1"
async def test_tool_runtime_forwards_execution_info_and_server_info_async() -> None:
"""Test that execution_info and server_info are forwarded in async path."""
from langgraph.runtime import ExecutionInfo, ServerInfo
exec_info = ExecutionInfo(
thread_id="t-2",
checkpoint_id="cp-2",
checkpoint_ns="",
task_id="tk-2",
run_id="r-2",
)
server_info = ServerInfo(assistant_id="asst-2", graph_id="graph-2")
mock_runtime = Mock()
mock_runtime.store = None
mock_runtime.context = None
mock_runtime.stream_writer = lambda *args, **kwargs: None
mock_runtime.execution_info = exec_info
mock_runtime.server_info = server_info
captured: dict = {}
@dec_tool
async def info_tool_async(x: int, runtime: ToolRuntime) -> str:
"""Async tool that captures runtime info."""
captured["execution_info"] = runtime.execution_info
captured["server_info"] = runtime.server_info
return "ok"
node = ToolNode([info_tool_async])
tool_call = {
"name": "info_tool_async",
"args": {"x": 1},
"id": "call-2",
"type": "tool_call",
}
msg = AIMessage("", tool_calls=[tool_call])
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
await node.ainvoke({"messages": [msg]}, config=config)
assert captured["execution_info"] is exec_info
assert captured["execution_info"].thread_id == "t-2"
assert captured["server_info"] is server_info
assert captured["server_info"].graph_id == "graph-2"
# --- InjectedToolArg security tests ---