Files
langgraph/libs/langgraph/tests/test_runtime.py
T
open-swe[bot]GitHubopen-swe[bot] <open-swe@users.noreply.github.com>Sydney Runkle
2c98c59fca fix: populate assistant_id from config configurable instead of metadata (#7468)
## Description
The `_build_server_info` function was reading `assistant_id` and
`graph_id` from `config["metadata"]`, but the server puts these values
in `config["configurable"]`. This updates the source to read from
`configurable` consistently.

## Test Plan
- [ ] Verify `server_info.assistant_id` and `server_info.graph_id` are
correctly populated from `config["configurable"]`

_Opened collaboratively by Sydney Runkle and open-swe._

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
2026-04-13 09:08:09 -04:00

652 lines
20 KiB
Python

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 ExecutionInfo, Runtime, ServerInfo, get_runtime
def test_injected_runtime() -> None:
@dataclass
class Context:
api_key: str
class State(TypedDict):
message: str
def injected_runtime(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
return {"message": f"api key: {runtime.context.api_key}"}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("injected_runtime", injected_runtime)
graph.add_edge(START, "injected_runtime")
graph.add_edge("injected_runtime", END)
compiled = graph.compile()
result = compiled.invoke(
{"message": "hello world"}, context=Context(api_key="sk_123456")
)
assert result == {"message": "api key: sk_123456"}
def test_context_runtime() -> None:
@dataclass
class Context:
api_key: str
class State(TypedDict):
message: str
def context_runtime(state: State) -> dict[str, Any]:
runtime = get_runtime(Context)
return {"message": f"api key: {runtime.context.api_key}"}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("context_runtime", context_runtime)
graph.add_edge(START, "context_runtime")
graph.add_edge("context_runtime", END)
compiled = graph.compile()
result = compiled.invoke(
{"message": "hello world"}, context=Context(api_key="sk_123456")
)
assert result == {"message": "api key: sk_123456"}
def test_override_runtime() -> None:
@dataclass
class Context:
api_key: str
prev = Runtime(context=Context(api_key="abc"))
new = prev.override(context=Context(api_key="def"))
assert new.override(context=Context(api_key="def")).context.api_key == "def"
def test_merge_runtime() -> None:
@dataclass
class Context:
api_key: str
runtime1 = Runtime(context=Context(api_key="abc"))
runtime2 = Runtime(context=Context(api_key="def"))
runtime3 = Runtime(context=None)
assert runtime1.merge(runtime2).context.api_key == "def"
# override only applies to non-falsy values
assert runtime1.merge(runtime3).context.api_key == "abc" # type: ignore
def test_runtime_propogated_to_subgraph() -> None:
@dataclass
class Context:
username: str
class State(TypedDict, total=False):
subgraph: str
main: str
def subgraph_node_1(state: State, runtime: Runtime[Context]):
return {"subgraph": f"{runtime.context.username}!"}
subgraph_builder = StateGraph(State, context_schema=Context)
subgraph_builder.add_node(subgraph_node_1)
subgraph_builder.set_entry_point("subgraph_node_1")
subgraph = subgraph_builder.compile()
def main_node(state: State, runtime: Runtime[Context]):
return {"main": f"{runtime.context.username}!"}
builder = StateGraph(State, context_schema=Context)
builder.add_node(main_node)
builder.add_node("node_1", subgraph)
builder.set_entry_point("main_node")
builder.add_edge("main_node", "node_1")
graph = builder.compile()
context = Context(username="Alice")
result = graph.invoke({}, context=context)
assert result == {"subgraph": "Alice!", "main": "Alice!"}
def test_context_coercion_dataclass() -> None:
"""Test that dict context is coerced to dataclass."""
@dataclass
class Context:
api_key: str
timeout: int = 30
class State(TypedDict):
message: str
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
return {
"message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}"
}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", node_with_context)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
# Test dict coercion with all fields
result = compiled.invoke(
{"message": "test"}, context={"api_key": "sk_test", "timeout": 60}
)
assert result == {"message": "api_key: sk_test, timeout: 60"}
# Test dict coercion with default field
result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"})
assert result == {"message": "api_key: sk_test2, timeout: 30"}
# Test with actual dataclass instance (should still work)
result = compiled.invoke(
{"message": "test"}, context=Context(api_key="sk_test3", timeout=90)
)
assert result == {"message": "api_key: sk_test3, timeout: 90"}
def test_context_coercion_pydantic() -> None:
"""Test that dict context is coerced to Pydantic model."""
class Context(BaseModel):
api_key: str
timeout: int = 30
tags: list[str] = []
class State(TypedDict):
message: str
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
return {
"message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}, tags: {runtime.context.tags}"
}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", node_with_context)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
# Test dict coercion with all fields
result = compiled.invoke(
{"message": "test"},
context={"api_key": "sk_test", "timeout": 60, "tags": ["prod", "v2"]},
)
assert result == {"message": "api_key: sk_test, timeout: 60, tags: ['prod', 'v2']"}
# Test dict coercion with defaults
result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"})
assert result == {"message": "api_key: sk_test2, timeout: 30, tags: []"}
# Test with actual Pydantic instance (should still work)
result = compiled.invoke(
{"message": "test"},
context=Context(api_key="sk_test3", timeout=90, tags=["test"]),
)
assert result == {"message": "api_key: sk_test3, timeout: 90, tags: ['test']"}
def test_context_coercion_typeddict() -> None:
"""Test that dict context with TypedDict schema passes through as-is."""
class Context(TypedDict):
api_key: str
timeout: int
class State(TypedDict):
message: str
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
# TypedDict context is just a dict at runtime
return {
"message": f"api_key: {runtime.context['api_key']}, timeout: {runtime.context['timeout']}"
}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", node_with_context)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
# Test dict passes through for TypedDict
result = compiled.invoke(
{"message": "test"}, context={"api_key": "sk_test", "timeout": 60}
)
assert result == {"message": "api_key: sk_test, timeout: 60"}
def test_context_coercion_none() -> None:
"""Test that None context is handled properly."""
@dataclass
class Context:
api_key: str
class State(TypedDict):
message: str
def node_without_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
# Should be None when no context provided
return {"message": f"context is None: {runtime.context is None}"}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", node_without_context)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
# Test with None context
result = compiled.invoke({"message": "test"}, context=None)
assert result == {"message": "context is None: True"}
# Test without context parameter (defaults to None)
result = compiled.invoke({"message": "test"})
assert result == {"message": "context is None: True"}
def test_context_coercion_errors() -> None:
"""Test error handling for invalid context."""
@dataclass
class Context:
api_key: str # Required field
class State(TypedDict):
message: str
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
return {"message": "should not reach here"}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", node_with_context)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
# Test missing required field
with pytest.raises(TypeError):
compiled.invoke({"message": "test"}, context={"timeout": 60})
# Test invalid dict keys
with pytest.raises(TypeError):
compiled.invoke(
{"message": "test"}, context={"api_key": "test", "invalid_field": "value"}
)
@pytest.mark.anyio
async def test_context_coercion_async() -> None:
"""Test context coercion with async methods."""
@dataclass
class Context:
api_key: str
async_mode: bool = True
class State(TypedDict):
message: str
async def async_node(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
return {
"message": f"async api_key: {runtime.context.api_key}, async_mode: {runtime.context.async_mode}"
}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", async_node)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
# Test dict coercion with ainvoke
result = await compiled.ainvoke(
{"message": "test"}, context={"api_key": "sk_async", "async_mode": False}
)
assert result == {"message": "async api_key: sk_async, async_mode: False"}
# Test dict coercion with astream
chunks = []
async for chunk in compiled.astream(
{"message": "test"}, context={"api_key": "sk_stream"}
):
chunks.append(chunk)
# Find the chunk with our node output
node_output = None
for chunk in chunks:
if "node" in chunk:
node_output = chunk["node"]
break
assert node_output == {"message": "async api_key: sk_stream, async_mode: True"}
def test_context_coercion_stream() -> None:
"""Test context coercion with sync stream method."""
@dataclass
class Context:
api_key: str
stream_mode: str = "default"
class State(TypedDict):
message: str
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
return {
"message": f"stream api_key: {runtime.context.api_key}, mode: {runtime.context.stream_mode}"
}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", node_with_context)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
# Test dict coercion with stream
chunks = []
for chunk in compiled.stream(
{"message": "test"}, context={"api_key": "sk_stream", "stream_mode": "fast"}
):
chunks.append(chunk)
# Find the chunk with our node output
node_output = None
for chunk in chunks:
if "node" in chunk:
node_output = chunk["node"]
break
assert node_output == {"message": "stream api_key: sk_stream, mode: fast"}
def test_context_coercion_pydantic_validation_errors() -> None:
"""Test that Pydantic validation errors are raised."""
class Context(BaseModel):
api_key: str
timeout: int
class State(TypedDict):
message: str
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
return {
"message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}"
}
graph = StateGraph(state_schema=State, context_schema=Context)
graph.add_node("node", node_with_context)
graph.add_edge(START, "node")
graph.add_edge("node", END)
compiled = graph.compile()
with pytest.raises(ValidationError):
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_configurable() -> None:
"""server_info is built from assistant_id/graph_id in config configurable."""
captured: dict[str, Any] = {}
compiled = _make_capture_graph(captured)
compiled.invoke(
{"message": "hi"},
config={"configurable": {"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_configurable() -> None:
"""server_info is None when no assistant_id/graph_id in configurable."""
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,
"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