diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 9f265f80f..d7157e239 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -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, diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 8de321b84..d754db391 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -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: diff --git a/libs/langgraph/langgraph/runtime.py b/libs/langgraph/langgraph/runtime.py index b4c0afac5..9de1b65bd 100644 --- a/libs/langgraph/langgraph/runtime.py +++ b/libs/langgraph/langgraph/runtime.py @@ -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, ) diff --git a/libs/langgraph/tests/test_retry.py b/libs/langgraph/tests/test_retry.py index a7fda9740..963a3c933 100644 --- a/libs/langgraph/tests/test_retry.py +++ b/libs/langgraph/tests/test_retry.py @@ -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 diff --git a/libs/langgraph/tests/test_runtime.py b/libs/langgraph/tests/test_runtime.py index 0407b84d2..0796ddf9d 100644 --- a/libs/langgraph/tests/test_runtime.py +++ b/libs/langgraph/tests/test_runtime.py @@ -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:" (top-level, no separator) + assert main_info.checkpoint_ns.startswith("main_node:") + assert "|" not in main_info.checkpoint_ns + + # Subgraph namespace is "subgraph:|sub_node:" (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 diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 509584160..32d293248 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -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): diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index 72fd5bc7f..06b3eb70c 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/libs/prebuilt/tests/test_tool_node.py @@ -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 ---