Compare commits

..
30 changed files with 333 additions and 1880 deletions

No files matched your search

+3 -3
View File
@@ -2148,6 +2148,6 @@ yocto-queue@^0.1.0:
integrity sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==
"zod@^3.25.76 || ^4":
version "4.6.7"
resolved "https://registry.yarnpkg.com/zod/-/zod-4.6.7.tgz#e806cf64a7f4e37635ce9f2f52c04c3766f75e5b"
integrity sha512-g8W2D2kU6a8untc053fX1uVBKzM7yqAxW/3zwk83+6RyF2+SKOuV+c0wBDhUwzV1PaGapmZH+rmcUzPf+3CGwg==
version "4.3.6"
resolved "https://registry.yarnpkg.com/zod/-/zod-4.3.6.tgz#89c56e0aa7d2b05107d894412227087885ab112a"
integrity sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==
-8
View File
@@ -11,7 +11,6 @@ import click.exceptions
import langgraph_cli.config
import langgraph_cli.docker
from langgraph_cli import langsmith_auth
from langgraph_cli.analytics import log_command
from langgraph_cli.config import Config
from langgraph_cli.constants import DEFAULT_CONFIG, DEFAULT_PORT
@@ -234,13 +233,6 @@ def cli():
pass
@cli.command(help="Log in to LangSmith in the browser (no API key needed).")
@click.option("--no-browser", is_flag=True, help="Print the URL instead of opening it.")
def login(no_browser: bool) -> None:
api_url = langsmith_auth.resolve_api_url(os.environ.get("LANGSMITH_ENDPOINT"))
langsmith_auth.login(api_url, open_browser=not no_browser)
# Wire the deploy group (defined in deploy.py) into the top-level CLI.
cli.add_command(deploy)
+12 -26
View File
@@ -1,7 +1,6 @@
"""Deploy command and subcommands for the LangGraph CLI."""
import base64
import functools
import json as json_mod
import os
import pathlib
@@ -20,7 +19,6 @@ import click.exceptions
from dotenv import dotenv_values, set_key
import langgraph_cli.config
from langgraph_cli import langsmith_auth
from langgraph_cli.analytics import log_command
from langgraph_cli.config import Config
from langgraph_cli.constants import DEFAULT_CONFIG
@@ -1984,6 +1982,17 @@ def _create_host_backend_client(
if val:
resolved_api_key = val
break
if not resolved_api_key:
if effective_no_input:
raise click.ClickException(
"No LangSmith API key found. Set LANGSMITH_API_KEY in the "
"environment or .env file."
)
click.secho(
"No LangSmith API key found. Create one at Settings > API Keys in LangSmith.",
fg="yellow",
)
resolved_api_key = click.prompt("Enter LangSmith API key", hide_input=True)
tenant_id = env_vars.get("LANGSMITH_TENANT_ID") or os.environ.get(
"LANGSMITH_TENANT_ID"
)
@@ -1991,31 +2000,8 @@ def _create_host_backend_client(
"LANGSMITH_ENDPOINT"
)
endpoints = ControlPlaneEndpoints.resolve(host_url, langsmith_endpoint)
access_token: Callable[[], str] | None = None
if not resolved_api_key:
profile = langsmith_auth.load_profile()
api_url = langsmith_auth.resolve_api_url(langsmith_endpoint)
if profile.get("api_key"):
resolved_api_key = profile["api_key"]
else:
if not any((profile.get("oauth") or {}).values()):
if effective_no_input:
raise click.ClickException(
"No LangSmith credentials found. Set LANGSMITH_API_KEY "
"in the environment or .env file (an API key), or run "
"`langgraph login`."
)
click.secho("You are not logged in to LangSmith.", fg="yellow")
langsmith_auth.login(api_url)
profile = langsmith_auth.load_profile()
access_token = functools.partial(langsmith_auth.get_access_token, api_url)
tenant_id = tenant_id or profile.get("workspace_id")
return HostBackendClient(
endpoints.control_plane_url,
resolved_api_key,
tenant_id=tenant_id,
**({"access_token": access_token} if access_token else {}),
endpoints.control_plane_url, resolved_api_key, tenant_id=tenant_id
)
+7 -15
View File
@@ -2,7 +2,7 @@
from __future__ import annotations
from collections.abc import Callable, Mapping
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Literal
from urllib.parse import urlparse
@@ -130,18 +130,17 @@ class HostBackendClient:
def __init__(
self,
base_url: str,
api_key: str | None,
api_key: str,
tenant_id: str | None = None,
*,
access_token: Callable[[], str] | None = None,
transport: httpx.BaseTransport | None = None,
):
if not base_url:
raise click.UsageError("Host backend URL is required")
self._access_token = access_token
headers: dict[str, str] = {"Accept": "application/json"}
if api_key:
headers["X-Api-Key"] = api_key
headers: dict[str, str] = {
"X-Api-Key": api_key,
"Accept": "application/json",
}
if tenant_id:
headers["X-Tenant-ID"] = tenant_id
self._endpoints = ControlPlaneEndpoints.from_control_plane_url(base_url)
@@ -171,15 +170,8 @@ class HostBackendClient:
payload: dict[str, Any] | None = None,
params: dict[str, Any] | None = None,
) -> Any:
headers = (
{"Authorization": f"Bearer {self._access_token()}"}
if self._access_token
else None
)
try:
resp = self._client.request(
method, path, json=payload, params=params, headers=headers
)
resp = self._client.request(method, path, json=payload, params=params)
resp.raise_for_status()
except httpx.HTTPStatusError as err:
detail = _error_detail(err.response)
-225
View File
@@ -1,225 +0,0 @@
"""LangSmith OAuth login, shared with the ``langsmith`` CLI via ``~/.langsmith/config.json``."""
from __future__ import annotations
import datetime
import json
import os
import pathlib
import time
import webbrowser
from typing import Any
from urllib.parse import urlparse
import click
import httpx
CLIENT_ID = "langsmith-cli"
DEFAULT_API_URL = "https://api.smith.langchain.com"
_WELL_KNOWN = "/.well-known/oauth-authorization-server"
_REFRESH_LEEWAY = datetime.timedelta(seconds=60)
_DEVICE_GRANT = "urn:ietf:params:oauth:grant-type:device_code"
def _config_path() -> pathlib.Path:
override = os.environ.get("LANGSMITH_CONFIG_FILE")
if override:
return pathlib.Path(override)
return pathlib.Path.home() / ".langsmith" / "config.json"
def _read_config() -> dict[str, Any]:
try:
config = json.loads(_config_path().read_text(encoding="utf-8"))
except (OSError, ValueError):
return {}
return config if isinstance(config, dict) else {}
def _write_config(config: dict[str, Any]) -> None:
path = _config_path()
path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
tmp = path.with_name(path.name + ".tmp")
tmp.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
os.chmod(tmp, 0o600)
os.replace(tmp, path)
def _profile_name(config: dict[str, Any]) -> str:
return (
os.environ.get("LANGSMITH_PROFILE")
or config.get("current_profile")
or "default"
)
def load_profile() -> dict[str, Any]:
"""Return the active profile from the LangSmith config file, or ``{}``."""
config = _read_config()
profile = (config.get("profiles") or {}).get(_profile_name(config))
return profile if isinstance(profile, dict) else {}
def resolve_api_url(endpoint: str | None = None) -> str:
"""Pick the LangSmith API URL: explicit endpoint, then profile, then cloud."""
return endpoint or load_profile().get("api_url") or DEFAULT_API_URL
def _save_tokens(token: dict[str, Any], api_url: str | None = None) -> None:
config = _read_config()
name = _profile_name(config)
profile = config.setdefault("profiles", {}).setdefault(name, {})
if api_url:
profile["api_url"] = api_url
oauth = profile.setdefault("oauth", {})
oauth["access_token"] = token["access_token"]
if token.get("refresh_token"):
oauth["refresh_token"] = token["refresh_token"]
if token.get("expires_in"):
expires_at = datetime.datetime.now(datetime.timezone.utc) + (
datetime.timedelta(seconds=int(token["expires_in"]))
)
oauth["expires_at"] = expires_at.strftime("%Y-%m-%dT%H:%M:%SZ")
config.setdefault("current_profile", name)
_write_config(config)
def _oauth_endpoints(api_url: str) -> dict[str, str]:
"""Discover the OAuth endpoints, falling back to ``<root>/oauth/*``."""
root = api_url.rstrip("/")
for suffix in ("/api/v1", "/api"):
if root.endswith(suffix):
root = root[: -len(suffix)]
break
for base in (root, f"{root}/api"):
try:
resp = httpx.get(base + _WELL_KNOWN, timeout=10)
doc = resp.json() if resp.is_success else None
except (httpx.HTTPError, ValueError):
continue
if (
isinstance(doc, dict)
and str(doc.get("issuer", "")).rstrip("/") == base
and all(
urlparse(str(doc.get(key, ""))).netloc == urlparse(base).netloc
for key in ("token_endpoint", "device_authorization_endpoint")
)
):
return {
"issuer": base,
"token": doc["token_endpoint"],
"device": doc["device_authorization_endpoint"],
}
return {
"issuer": root,
"token": f"{root}/oauth/token",
"device": f"{root}/oauth/device/code",
}
def _expiring(oauth: dict[str, Any]) -> bool:
if not oauth.get("access_token"):
return True
expires_at = oauth.get("expires_at")
if not expires_at:
return False
try:
expiry = datetime.datetime.fromisoformat(expires_at.replace("Z", "+00:00"))
except ValueError:
return True
if expiry.tzinfo is None:
expiry = expiry.replace(tzinfo=datetime.timezone.utc)
return expiry <= datetime.datetime.now(datetime.timezone.utc) + _REFRESH_LEEWAY
def get_access_token(api_url: str) -> str:
"""Return a valid access token from the profile, refreshing it if needed."""
oauth = load_profile().get("oauth") or {}
if not _expiring(oauth):
return oauth["access_token"]
refresh_token = oauth.get("refresh_token")
if not refresh_token:
raise click.ClickException(
"LangSmith session expired. Run `langgraph login` to sign in again."
)
endpoints = _oauth_endpoints(api_url)
try:
resp = httpx.post(
endpoints["token"],
data={
"grant_type": "refresh_token",
"client_id": CLIENT_ID,
"refresh_token": refresh_token,
},
timeout=10,
)
resp.raise_for_status()
token = resp.json()
except (httpx.HTTPError, ValueError):
raise click.ClickException(
"Could not refresh your LangSmith session. "
"Run `langgraph login` to sign in again."
) from None
if not isinstance(token, dict) or not token.get("access_token"):
raise click.ClickException(
"Could not refresh your LangSmith session. "
"Run `langgraph login` to sign in again."
)
_save_tokens(token)
return token["access_token"]
def login(api_url: str, *, open_browser: bool = True) -> None:
"""Sign in with the OAuth device flow and store tokens in the profile."""
endpoints = _oauth_endpoints(api_url)
try:
resp = httpx.post(
endpoints["device"],
data={"client_id": CLIENT_ID, "resource": endpoints["issuer"]},
timeout=10,
)
resp.raise_for_status()
device = resp.json()
device_code, user_code = device["device_code"], device["user_code"]
verification_url = (
device.get("verification_uri_complete") or device["verification_uri"]
)
except (httpx.HTTPError, ValueError, KeyError):
raise click.ClickException(
f"Could not start LangSmith login against {api_url}."
) from None
click.echo(f"Open this URL to authorize the CLI:\n{verification_url}\n")
click.echo(f"Code: {user_code}\n")
if open_browser:
webbrowser.open(verification_url)
click.echo("Waiting for authorization...")
interval = max(int(device.get("interval") or 5), 5)
deadline = time.monotonic() + int(device.get("expires_in") or 600)
while time.monotonic() < deadline:
time.sleep(interval)
try:
resp = httpx.post(
endpoints["token"],
data={
"grant_type": _DEVICE_GRANT,
"client_id": CLIENT_ID,
"device_code": device_code,
"resource": endpoints["issuer"],
},
timeout=10,
)
body = resp.json()
except (httpx.HTTPError, ValueError):
continue
if resp.is_success and body.get("access_token"):
_save_tokens(body, api_url)
click.secho("Logged in to LangSmith.", fg="green")
return
error = body.get("error") if isinstance(body, dict) else None
if error == "slow_down":
interval += 5
elif error != "authorization_pending":
raise click.ClickException(f"LangSmith login failed: {error or 'error'}.")
raise click.ClickException("Timed out waiting for LangSmith login.")
@@ -578,13 +578,12 @@ def test_a_deployed_result_prints_the_deployment_url(
class TestCreateHostBackendClientNoInput:
def test_raises_when_no_api_key_and_no_input(self, monkeypatch, tmp_path):
monkeypatch.setenv("LANGSMITH_CONFIG_FILE", str(tmp_path / "missing.json"))
monkeypatch.setattr(deploy_mod, "_no_input", True)
monkeypatch.delenv("LANGSMITH_API_KEY", raising=False)
monkeypatch.delenv("LANGCHAIN_API_KEY", raising=False)
monkeypatch.delenv("LANGGRAPH_HOST_API_KEY", raising=False)
with pytest.raises(click.ClickException, match="credentials"):
with pytest.raises(click.ClickException, match="API key"):
_create_host_backend_client(
host_url="https://api.example.com",
api_key=None,
@@ -1,191 +0,0 @@
import json
from datetime import datetime, timedelta, timezone
import click
import httpx
import pytest
from langgraph_cli import deploy as deploy_mod
from langgraph_cli import langsmith_auth
from langgraph_cli.deploy import _create_host_backend_client
from langgraph_cli.host_backend import HostBackendClient
API_URL = "https://api.smith.langchain.com"
def _iso(delta: timedelta) -> str:
return (datetime.now(timezone.utc) + delta).strftime("%Y-%m-%dT%H:%M:%SZ")
@pytest.fixture
def config_file(tmp_path, monkeypatch):
path = tmp_path / "config.json"
monkeypatch.setenv("LANGSMITH_CONFIG_FILE", str(path))
monkeypatch.delenv("LANGSMITH_PROFILE", raising=False)
for name in ("LANGGRAPH_HOST_API_KEY", "LANGSMITH_API_KEY", "LANGCHAIN_API_KEY"):
monkeypatch.delenv(name, raising=False)
return path
def _write_profile(path, **oauth):
path.write_text(
json.dumps(
{
"current_profile": "default",
"profiles": {
"default": {
"api_url": API_URL,
"workspace_id": "ws-1",
"oauth": oauth,
}
},
}
)
)
def _no_discovery(monkeypatch):
def get(*_a, **_k):
raise httpx.ConnectError("no discovery")
monkeypatch.setattr(httpx, "get", get)
def test_valid_token_is_returned_without_network(config_file, monkeypatch):
_write_profile(
config_file, access_token="tok", expires_at=_iso(timedelta(minutes=5))
)
def boom(*_a, **_k):
raise AssertionError("no network expected")
monkeypatch.setattr(httpx, "get", boom)
monkeypatch.setattr(httpx, "post", boom)
assert langsmith_auth.get_access_token(API_URL) == "tok"
def test_expired_token_is_refreshed_and_saved(config_file, monkeypatch):
_write_profile(
config_file,
access_token="old",
refresh_token="refresh-1",
expires_at=_iso(timedelta(seconds=-5)),
)
_no_discovery(monkeypatch)
seen = {}
def post(url, data, timeout):
seen.update(url=url, data=data)
return httpx.Response(
200,
json={
"access_token": "new",
"refresh_token": "refresh-2",
"expires_in": 300,
},
request=httpx.Request("POST", url),
)
monkeypatch.setattr(httpx, "post", post)
assert langsmith_auth.get_access_token(API_URL) == "new"
assert seen["url"] == f"{API_URL}/oauth/token"
assert seen["data"]["refresh_token"] == "refresh-1"
saved = json.loads(config_file.read_text())["profiles"]["default"]["oauth"]
assert saved["access_token"] == "new"
assert saved["refresh_token"] == "refresh-2"
def test_expired_token_without_refresh_token_asks_to_log_in(config_file):
_write_profile(
config_file, access_token="old", expires_at=_iso(timedelta(seconds=-5))
)
with pytest.raises(click.ClickException, match="langgraph login"):
langsmith_auth.get_access_token(API_URL)
def test_login_stores_tokens_after_pending(config_file, monkeypatch):
_no_discovery(monkeypatch)
monkeypatch.setattr(langsmith_auth.time, "sleep", lambda _s: None)
monkeypatch.setattr(langsmith_auth.webbrowser, "open", lambda _u: True)
token_responses = iter(
[
{"error": "authorization_pending"},
{"access_token": "tok", "refresh_token": "r", "expires_in": 300},
]
)
def post(url, data, timeout):
request = httpx.Request("POST", url)
if url.endswith("/oauth/device/code"):
return httpx.Response(
200,
json={
"device_code": "dc",
"user_code": "ABCD",
"verification_uri": "https://example.com/activate",
"expires_in": 600,
"interval": 5,
},
request=request,
)
body = next(token_responses)
return httpx.Response(
400 if "error" in body else 200, json=body, request=request
)
monkeypatch.setattr(httpx, "post", post)
langsmith_auth.login(API_URL)
profile = json.loads(config_file.read_text())["profiles"]["default"]
assert profile["api_url"] == API_URL
assert profile["oauth"]["access_token"] == "tok"
def test_client_sends_bearer_token_per_request():
seen = []
def handler(req: httpx.Request) -> httpx.Response:
seen.append(dict(req.headers))
return httpx.Response(200, json={"resources": []})
tokens = iter(["t1", "t2"])
client = HostBackendClient(
"https://api.example.com",
None,
access_token=lambda: next(tokens),
transport=httpx.MockTransport(handler),
)
client.list_deployments()
client.list_deployments()
assert [h["authorization"] for h in seen] == ["Bearer t1", "Bearer t2"]
assert "x-api-key" not in seen[0]
def test_factory_uses_logged_in_profile(config_file, monkeypatch):
_write_profile(
config_file, access_token="tok", expires_at=_iso(timedelta(minutes=5))
)
monkeypatch.setattr(deploy_mod, "_no_input", True)
monkeypatch.delenv("LANGSMITH_TENANT_ID", raising=False)
monkeypatch.delenv("LANGSMITH_ENDPOINT", raising=False)
client = _create_host_backend_client(
host_url="https://api.example.com", api_key=None, env_vars={}
)
assert client._access_token is not None
assert client._access_token() == "tok"
assert client._client.headers["x-tenant-id"] == "ws-1"
assert "x-api-key" not in client._client.headers
def test_factory_without_credentials_and_no_input_fails(config_file, monkeypatch):
monkeypatch.setattr(deploy_mod, "_no_input", True)
with pytest.raises(click.ClickException, match="langgraph login"):
_create_host_backend_client(
host_url="https://api.example.com", api_key=None, env_vars={}
)
+11 -8
View File
@@ -22,7 +22,7 @@ from langchain_core.messages import (
)
from typing_extensions import TypedDict, deprecated
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
from langgraph.graph.state import StateGraph
from langgraph.pregel._messages import (
_messages_delta_reducer as _messages_delta_reducer,
@@ -343,7 +343,7 @@ def push_message(
)
from langgraph.config import get_config
from langgraph.pregel._messages import StreamMessagesHandler, graph_namespace
from langgraph.pregel._messages import StreamMessagesHandler
config = get_config()
message = next(x for x in convert_to_messages([message]))
@@ -359,12 +359,15 @@ def push_message(
):
handlers = config["callbacks"]
metadata = config["metadata"]
for handler in handlers:
if isinstance(handler, StreamMessagesHandler):
ns = graph_namespace(metadata)
if handler.subgraphs or not ns or ns == handler.parent_ns:
handler._emit((ns, metadata), message, dedupe=False)
if stream_handler := next(
(x for x in handlers if isinstance(x, StreamMessagesHandler)), None
):
metadata = config["metadata"]
message_meta = (
tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
metadata,
)
stream_handler._emit(message_meta, message, dedupe=False)
if state_key:
config[CONF][CONFIG_KEY_SEND]([(state_key, message)])
+34
View File
@@ -232,6 +232,8 @@ def draw_graph(
tools_by_name = getattr(node.bound, "tools_by_name", None)
if isinstance(tools_by_name, Mapping) and tools_by_name:
metadata["__tools"] = list(tools_by_name)
if subagents := _get_tool_subagents(tools_by_name):
metadata["__subagents"] = subagents
graph.add_node(node.bound, name, metadata=metadata or None)
# add start node
if START not in nodes:
@@ -281,6 +283,38 @@ def draw_graph(
return graph
def _get_tool_subagents(
tools_by_name: Mapping[str, Any],
) -> dict[str, list[dict[str, Any]]]:
"""Collect the subagents each tool can run, keyed by tool name."""
subagents: dict[str, list[dict[str, Any]]] = {}
for tool_name, tool in tools_by_name.items():
tool_subagents = getattr(tool, "subagents", None)
if (
isinstance(tool_subagents, Sequence)
and not isinstance(tool_subagents, str)
and tool_subagents
and all(
isinstance(subagent, Mapping) and isinstance(subagent.get("name"), str)
for subagent in tool_subagents
)
):
subagents[tool_name] = [
_subagent_entry(subagent) for subagent in tool_subagents
]
return subagents
def _subagent_entry(subagent: Mapping[str, Any]) -> dict[str, Any]:
"""Keep only the subagent fields graph viewers display."""
entry: dict[str, Any] = {"name": subagent["name"]}
if isinstance(description := subagent.get("description"), str):
entry["description"] = description
if subagent.get("forked") is True:
entry["forked"] = True
return entry
def add_edge(
graph: Graph,
source: str,
+8 -21
View File
@@ -23,7 +23,7 @@ from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResul
from pydantic import BaseModel
from langgraph._internal._config import filter_to_user_tags
from langgraph._internal._constants import NS_SEP, PUSH
from langgraph._internal._constants import NS_SEP
from langgraph.channels.binop import _get_overwrite
from langgraph.channels.delta import DeltaChannel
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
@@ -55,23 +55,6 @@ def _state_values(obj: Any) -> Sequence[Any]:
return ()
def graph_namespace(metadata: Mapping[str, Any]) -> tuple[str, ...]:
"""Namespace of the graph a task runs in, from the task's run metadata.
Drops the task's own segment, plus one per enclosing `@task` call: a call
runs in its caller's graph but nests its segment under the caller's. A
call's `langgraph_path` is `(PUSH, <caller's path>, <index>, True)`, while
a `Send` task's ends in `False`.
"""
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))
path = metadata.get("langgraph_path", ())
drop = 1
while len(path) == 4 and path[0] == PUSH and path[-1] is True:
drop += 1
path = path[1]
return ns[:-drop]
class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
"""A callback handler that implements stream_mode=messages.
@@ -165,7 +148,9 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
**kwargs: Any,
) -> Any:
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
ns = graph_namespace(metadata)
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
:-1
]
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
return
if (filtered_tags := filter_to_user_tags(tags)) is not None:
@@ -228,8 +213,10 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
and kwargs.get("name") == metadata.get("langgraph_node")
and (not tags or TAG_HIDDEN not in tags)
):
ns = graph_namespace(metadata)
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
:-1
]
if not self.subgraphs and len(ns) > 0:
return
self.metadata[run_id] = (ns, metadata)
for value in _state_values(inputs):
+12 -7
View File
@@ -7,9 +7,9 @@ from uuid import UUID
from langchain_core.callbacks import BaseCallbackHandler
from langgraph._internal._constants import NS_SEP
from langgraph.constants import TAG_NOSTREAM
from langgraph.errors import GraphDrained, GraphInterrupt
from langgraph.pregel._messages import graph_namespace
from langgraph.pregel.protocol import StreamChunk
try:
@@ -92,9 +92,12 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
) -> tuple[str, ...] | None:
"""Resolve the namespace this tool call should emit at, or `None` to skip.
Mirrors `StreamMessagesHandler.on_chat_model_start`: the namespace
is that of the graph the tool runs in (`graph_namespace`). Returns
`None` when the call should be silently suppressed:
Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace
derivation: parses `langgraph_checkpoint_ns` (which ends with
the `node_name:task_id` of the calling node), drops that
trailing segment, and returns the containing subgraph's own
namespace. Returns `None` when the call should be silently
suppressed:
- `metadata` is missing — handler is attached to a context
without Pregel routing info.
@@ -107,9 +110,11 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
return None
if tags and TAG_NOSTREAM in tags:
return None
ns = (
graph_namespace(metadata) if metadata.get("langgraph_checkpoint_ns") else ()
)
nskey = metadata.get("langgraph_checkpoint_ns")
if not nskey:
ns: tuple[str, ...] = ()
else:
ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
return None
return ns
+33 -44
View File
@@ -2880,18 +2880,6 @@ class Pregel(
for h in callback_manager.inheritable_handlers
if not isinstance(h, StreamMessagesHandlerV2)
]
if CONFIG_KEY_TASK_ID not in config[CONF]:
# Not run as a task (a child invoked with its own thread, say),
# this graph's namespace starts at the root, so a caller's
# handler would take this graph's events for its own graph's.
callback_manager.handlers = [
h for h in callback_manager.handlers if not _is_stream_handler(h)
]
callback_manager.inheritable_handlers = [
h
for h in callback_manager.inheritable_handlers
if not _is_stream_handler(h)
]
if "ls_integration" not in callback_manager.metadata:
callback_manager.add_metadata({"ls_integration": "langgraph"})
run_manager = callback_manager.on_chain_start(
@@ -2910,9 +2898,9 @@ class Pregel(
)
# set up subgraph checkpointing
config = self._own_checkpoint_config(config)
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
# set up messages stream mode
if "messages" in stream_modes:
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
use_stream_messages_v2 = bool(
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
)
@@ -2921,12 +2909,24 @@ class Pregel(
if use_stream_messages_v2
else StreamMessagesHandler
)
stream_handlers.append(messages_handler_cls(stream.put, subgraphs))
run_manager.inheritable_handlers.append(
messages_handler_cls(
stream.put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
)
)
# set up tools stream mode
if "tools" in stream_modes:
stream_handlers.append(StreamToolCallHandler(stream.put, subgraphs))
run_manager.inheritable_handlers.extend(stream_handlers)
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
run_manager.inheritable_handlers.append(
StreamToolCallHandler(
stream.put,
subgraphs,
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
)
)
# set up custom stream mode
if "custom" in stream_modes:
@@ -3009,10 +3009,6 @@ class Pregel(
cache_policy=self.cache_policy,
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
) as loop:
# Known only now: the loop numbers the namespace of every
# subgraph a task runs after its first.
for handler in stream_handlers:
handler.parent_ns = loop.checkpoint_ns
emit_graph_lifecycle_events(loop)
# create runner
runner = PregelRunner(
@@ -3295,18 +3291,6 @@ class Pregel(
for h in callback_manager.inheritable_handlers
if not isinstance(h, StreamMessagesHandlerV2)
]
if CONFIG_KEY_TASK_ID not in config[CONF]:
# Not run as a task (a child invoked with its own thread, say),
# this graph's namespace starts at the root, so a caller's
# handler would take this graph's events for its own graph's.
callback_manager.handlers = [
h for h in callback_manager.handlers if not _is_stream_handler(h)
]
callback_manager.inheritable_handlers = [
h
for h in callback_manager.inheritable_handlers
if not _is_stream_handler(h)
]
if "ls_integration" not in callback_manager.metadata:
callback_manager.add_metadata({"ls_integration": "langgraph"})
run_manager = await callback_manager.on_chain_start(
@@ -3341,9 +3325,10 @@ class Pregel(
)
# set up subgraph checkpointing
config = self._own_checkpoint_config(config)
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
# set up messages stream mode
if "messages" in stream_modes:
# namespace can be None in a root level graph?
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
use_stream_messages_v2 = bool(
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
)
@@ -3352,12 +3337,24 @@ class Pregel(
if use_stream_messages_v2
else StreamMessagesHandler
)
stream_handlers.append(messages_handler_cls(stream_put, subgraphs))
run_manager.inheritable_handlers.append(
messages_handler_cls(
stream_put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
)
)
# set up tools stream mode
if "tools" in stream_modes:
stream_handlers.append(StreamToolCallHandler(stream_put, subgraphs))
run_manager.inheritable_handlers.extend(stream_handlers)
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
run_manager.inheritable_handlers.append(
StreamToolCallHandler(
stream_put,
subgraphs,
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
)
)
# set up custom stream mode
def stream_writer(c: Any) -> None:
@@ -3466,10 +3463,6 @@ class Pregel(
cache_policy=self.cache_policy,
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
) as loop:
# Known only now: the loop numbers the namespace of every
# subgraph a task runs after its first.
for handler in stream_handlers:
handler.parent_ns = loop.checkpoint_ns
await aemit_graph_lifecycle_events(loop)
# create runner
runner = PregelRunner(
@@ -4467,10 +4460,6 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
payload["values"] = mapper(payload["values"])
def _is_stream_handler(handler: object) -> bool:
return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
def _resolve_parent_ns(
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
) -> tuple[str, ...]:
-10
View File
@@ -847,11 +847,6 @@ class RemoteGraph(PregelProtocol):
# filter for what was actually requested
if mode not in requested:
continue
# subgraph events may have been requested only for the parent stream.
# messages are exempt: the server dedupes them across namespaces, so
# dropping them here would also drop the root-level message
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data))
@@ -1002,11 +997,6 @@ class RemoteGraph(PregelProtocol):
# filter for what was actually requested
if mode not in requested:
continue
# subgraph events may have been requested only for the parent stream.
# messages are exempt: the server dedupes them across namespaces, so
# dropping them here would also drop the root-level message
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data))
+44 -34
View File
@@ -36,7 +36,7 @@ from langchain_core.runnables import (
RunnablePassthrough,
)
from langchain_core.runnables.graph import Edge
from langchain_core.tools import tool
from langchain_core.tools import StructuredTool, tool
from langchain_core.version import VERSION as LANGCHAIN_CORE_VERSION
from langgraph.cache.base import BaseCache
from langgraph.checkpoint.base import (
@@ -3105,6 +3105,49 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
)
def test_tool_node_lists_tool_subagents_in_graph_metadata() -> None:
class TaskTool(StructuredTool):
subagents: list[dict[str, Any]] = []
def run_task(description: str, subagent_type: str) -> str:
return subagent_type
task = TaskTool.from_function(run_task, name="task", description="Run a subagent.")
task.subagents = [
{"name": "researcher", "description": "Researches a topic."},
{
"name": "critic",
"description": "Reviews drafts.",
"forked": True,
"system_prompt": "You are a critic.",
"runnable": object(),
},
{"name": "planner", "description": None, "forked": "yes"},
]
@tool
def search(query: str) -> str:
"""Search the web."""
return query
builder = StateGraph(MessagesState)
builder.add_node("agent", lambda state: state)
builder.add_node("tools", ToolNode([search, task]))
builder.add_edge(START, "agent")
builder.add_edge("agent", "tools")
graph = builder.compile().get_graph().to_json()
nodes = {node["id"]: node for node in graph["nodes"]}
assert nodes["tools"]["metadata"]["__subagents"] == {
"task": [
{"name": "researcher", "description": "Researches a topic."},
{"name": "critic", "description": "Reviews drafts.", "forked": True},
{"name": "planner"},
]
}
assert "__subagents" not in (nodes["agent"].get("metadata") or {})
def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
@@ -8950,39 +8993,6 @@ def test_subgraph_streaming_sync() -> None:
assert result["num_chunks"] == 9
def test_subgraph_streaming_node_output_messages_sync() -> None:
"""Messages returned by a node (not streamed by an LLM) are emitted when a
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
def respond(state: MessagesState) -> dict:
return {"messages": [AIMessage("hi there", id="ai1")]}
subgraph = (
StateGraph(MessagesState)
.add_node("respond", respond)
.add_edge(START, "respond")
.compile()
)
msgs = {"messages": [("user", "hello")]}
top_level = [msg for msg, _ in subgraph.stream(msgs, stream_mode="messages")]
def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
nonlocal in_node
in_node = [
msg for msg, _ in subgraph.stream(msgs, config, stream_mode="messages")
]
return {}
in_node: list = []
StateGraph(MessagesState).add_node("parent", parent_node).add_edge(
START, "parent"
).compile().invoke(msgs)
assert [m.content for m in top_level] == ["hi there"]
assert [m.content for m in in_node] == ["hi there"]
def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
messages: list[str]
-38
View File
@@ -9503,44 +9503,6 @@ async def test_subgraph_streaming_async() -> None:
assert result["num_chunks"] == 9
async def test_subgraph_streaming_node_output_messages_async() -> None:
"""Messages returned by a node (not streamed by an LLM) are emitted when a
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
async def respond(state: MessagesState) -> dict:
return {"messages": [AIMessage("hi there", id="ai1")]}
subgraph = (
StateGraph(MessagesState)
.add_node("respond", respond)
.add_edge(START, "respond")
.compile()
)
msgs = {"messages": [("user", "hello")]}
top_level = [msg async for msg, _ in subgraph.astream(msgs, stream_mode="messages")]
async def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
nonlocal in_node
in_node = [
msg
async for msg, _ in subgraph.astream(msgs, config, stream_mode="messages")
]
return {}
in_node: list = []
await (
StateGraph(MessagesState)
.add_node("parent", parent_node)
.add_edge(START, "parent")
.compile()
.ainvoke(msgs)
)
assert [m.content for m in top_level] == ["hi there"]
assert [m.content for m in in_node] == ["hi there"]
@NEEDS_CONTEXTVARS
async def test_null_resume_disallowed_with_multiple_interrupts(
async_checkpointer: BaseCheckpointSaver,
-83
View File
@@ -16,11 +16,9 @@ from langgraph_sdk.schema import StreamPart
from pydantic import BaseModel
from typing_extensions import TypedDict
from langgraph._internal._constants import CONFIG_KEY_STREAM
from langgraph.errors import GraphInterrupt
from langgraph.graph import END, START, MessagesState, StateGraph, add_messages
from langgraph.pregel import Pregel
from langgraph.pregel.protocol import StreamProtocol
from langgraph.pregel.remote import RemoteGraph
from langgraph.types import Interrupt, StateSnapshot
from tests.any_str import AnyStr
@@ -865,87 +863,6 @@ async def test_astream_subgraphs_keeps_caller_namespace():
assert stream_parts == _CALLER_PREFIXED
# Subgraph events requested only on behalf of a streaming parent graph
_PARENT_STREAMED_PARTS = [
StreamPart(event="updates|sub:abc", data={"inner": {"chunk": "inner"}}),
StreamPart(event="messages|sub:abc", data=[{"content": "token"}, {}]),
StreamPart(event="updates", data={"sub": {"chunk": "root"}}),
]
def _parent_streamed_config(parent_parts: list) -> RunnableConfig:
return {
"configurable": {
"thread_id": "thread_1",
"checkpoint_ns": "outer:123",
CONFIG_KEY_STREAM: StreamProtocol(
parent_parts.append, {"updates", "messages"}
),
}
}
def test_stream_without_subgraphs_skips_parent_requested_subgraph_events():
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = _PARENT_STREAMED_PARTS
remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
for stream_mode, expected in [
("updates", [{"sub": {"chunk": "root"}}]),
# messages are not filtered: the server dedupes them across namespaces
("messages", [({"content": "token"}, {})]),
]:
parent_parts: list = []
stream_parts = list(
remote_pregel.stream(
{"input": "data"},
_parent_streamed_config(parent_parts),
stream_mode=stream_mode,
)
)
assert stream_parts == expected
# subgraph events are still requested and forwarded to the parent
assert mock_sync_client.runs.stream.call_args.kwargs["stream_subgraphs"]
assert [ns for ns, _, _ in parent_parts] == [
("outer:123", "sub:abc"),
("outer:123", "sub:abc"),
("outer:123",),
]
async def test_astream_without_subgraphs_skips_parent_requested_subgraph_events():
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = _PARENT_STREAMED_PARTS
mock_async_client.runs.stream.return_value = async_iter
remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
for stream_mode, expected in [
("updates", [{"sub": {"chunk": "root"}}]),
# messages are not filtered: the server dedupes them across namespaces
("messages", [({"content": "token"}, {})]),
]:
parent_parts: list = []
stream_parts = [
part
async for part in remote_pregel.astream(
{"input": "data"},
_parent_streamed_config(parent_parts),
stream_mode=stream_mode,
)
]
assert stream_parts == expected
# subgraph events are still requested and forwarded to the parent
assert mock_async_client.runs.stream.call_args.kwargs["stream_subgraphs"]
assert [ns for ns, _, _ in parent_parts] == [
("outer:123", "sub:abc"),
("outer:123", "sub:abc"),
("outer:123",),
]
def test_invoke():
# set up test
mock_sync_client = MagicMock()
@@ -1,308 +0,0 @@
import sys
from typing import Any
import pytest
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langgraph.func import entrypoint, task
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.graph.message import push_message
from langgraph.graph.state import CompiledStateGraph
pytestmark = pytest.mark.anyio
NEEDS_CONTEXTVARS = pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
@tool
def echo(text: str) -> str:
"""Echo the text."""
return text
def _one_node_graph(node: Any) -> CompiledStateGraph:
return (
StateGraph(MessagesState)
.add_node(node)
.add_edge(START, node.__name__)
.compile()
)
def _chat_and_echo(state: MessagesState) -> dict:
echo.invoke({"text": "done"})
return {"messages": [GenericFakeChatModel(messages=iter(["one two"])).invoke("hi")]}
async def _achat_and_echo(state: MessagesState) -> dict:
await echo.ainvoke({"text": "done"})
model = GenericFakeChatModel(messages=iter(["one two"]))
return {"messages": [await model.ainvoke("hi")]}
def _segment_names(ns: tuple[str, ...]) -> tuple[str, ...]:
return tuple(segment.split(":")[0] for segment in ns)
def _summary(part: dict) -> tuple[tuple[str, ...], str]:
if part["type"] == "messages":
return _segment_names(part["ns"]), part["data"][0].content
return _segment_names(part["ns"]), part["data"]["event"]
_CHAT_AND_ECHO_EVENTS = [
((), "tool-started"),
((), "tool-finished"),
((), "one"),
((), " "),
((), "two"),
]
_CHILD_CONFIGS = pytest.mark.parametrize(
"child_config",
[
None,
{"configurable": {"thread_id": "child"}},
{"configurable": {"thread_id": "parent"}},
],
ids=["no_config", "own_thread", "caller_thread"],
)
_TASK_CALLERS = pytest.mark.parametrize(
"nested", [False, True], ids=["from_entrypoint", "from_task_in_task"]
)
@_TASK_CALLERS
@pytest.mark.parametrize("subgraphs", [False, True])
def test_task_events_stream_under_the_graph_namespace(
nested: bool, subgraphs: bool
) -> None:
@task
def chat_and_echo() -> None:
_chat_and_echo({"messages": []})
@task
def outer() -> None:
chat_and_echo().result()
@entrypoint()
def workflow(inputs: dict) -> None:
(outer() if nested else chat_and_echo()).result()
parts = workflow.stream(
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
)
assert [_summary(part) for part in parts] == _CHAT_AND_ECHO_EVENTS
@NEEDS_CONTEXTVARS
@_TASK_CALLERS
@pytest.mark.parametrize("subgraphs", [False, True])
async def test_task_events_stream_under_the_graph_namespace_async(
nested: bool, subgraphs: bool
) -> None:
@task
async def chat_and_echo() -> None:
await _achat_and_echo({"messages": []})
@task
async def outer() -> None:
await chat_and_echo()
@entrypoint()
async def workflow(inputs: dict) -> None:
await (outer() if nested else chat_and_echo())
parts = workflow.astream(
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
)
assert [_summary(part) async for part in parts] == _CHAT_AND_ECHO_EVENTS
def test_task_events_stream_under_the_namespace_of_their_subgraph() -> None:
@task
def chat_and_echo() -> None:
_chat_and_echo({"messages": []})
def node(state: MessagesState) -> dict:
chat_and_echo().result()
return {}
sub = _one_node_graph(node)
parent = StateGraph(MessagesState).add_node("sub", sub).add_edge(START, "sub")
parts = parent.compile().stream(
{"messages": []},
stream_mode=["messages", "tools"],
subgraphs=True,
version="v2",
)
assert [_summary(part) for part in parts] == [
(("sub",), event) for _, event in _CHAT_AND_ECHO_EVENTS
]
@_CHILD_CONFIGS
def test_child_graph_streams_only_to_the_node_streaming_it(
child_config: dict | None,
) -> None:
child = _one_node_graph(_chat_and_echo)
received = []
def node(state: MessagesState) -> dict:
received.extend(
child.stream(
{"messages": []}, child_config, stream_mode=["messages", "tools"]
)
)
return {}
parts = _one_node_graph(node).stream(
{"messages": []},
{"configurable": {"thread_id": "parent"}},
stream_mode=["messages", "tools"],
)
assert list(parts) == []
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
@NEEDS_CONTEXTVARS
@_CHILD_CONFIGS
async def test_child_graph_streams_only_to_the_node_streaming_it_async(
child_config: dict | None,
) -> None:
child = _one_node_graph(_achat_and_echo)
received = []
async def node(state: MessagesState) -> dict:
stream = child.astream(
{"messages": []}, child_config, stream_mode=["messages", "tools"]
)
received.extend([part async for part in stream])
return {}
parts = _one_node_graph(node).astream(
{"messages": []},
{"configurable": {"thread_id": "parent"}},
stream_mode=["messages", "tools"],
)
assert [part async for part in parts] == []
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
@pytest.mark.parametrize(
"child_config",
[
{"configurable": {"thread_id": "child"}},
{"configurable": {"thread_id": "parent"}},
],
ids=["own_thread", "caller_thread"],
)
def test_child_graph_on_its_own_thread_skips_a_caller_streaming_subgraphs(
child_config: dict,
) -> None:
child = _one_node_graph(_chat_and_echo)
def node(state: MessagesState) -> dict:
child.invoke({"messages": []}, child_config)
return {}
parts = _one_node_graph(node).stream(
{"messages": []},
{"configurable": {"thread_id": "parent"}},
stream_mode=["messages", "tools"],
subgraphs=True,
)
assert list(parts) == []
def test_node_gets_the_events_of_every_subgraph_it_streams() -> None:
sub = _one_node_graph(_chat_and_echo)
counts = []
def node(state: MessagesState) -> dict:
for _ in range(2):
stream = sub.stream({"messages": []}, stream_mode=["messages", "tools"])
counts.append(len(list(stream)))
return {}
_one_node_graph(node).invoke({"messages": []})
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
@NEEDS_CONTEXTVARS
async def test_node_gets_the_events_of_every_subgraph_it_streams_async() -> None:
sub = _one_node_graph(_achat_and_echo)
counts = []
async def node(state: MessagesState) -> dict:
for _ in range(2):
stream = sub.astream({"messages": []}, stream_mode=["messages", "tools"])
counts.append(len([part async for part in stream]))
return {}
await _one_node_graph(node).ainvoke({"messages": []})
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
def _push(state: MessagesState) -> dict:
push_message(AIMessage(id="pushed", content="hi"), state_key=None)
return {}
def _namespaces(graph: CompiledStateGraph, subgraphs: bool) -> list[tuple[str, ...]]:
parts = graph.stream(
{"messages": []}, stream_mode="messages", subgraphs=subgraphs, version="v2"
)
return [_segment_names(part["ns"]) for part in parts]
def _streams_pushing_graph(received: list) -> Any:
pushing = _one_node_graph(_push)
def node(state: MessagesState) -> dict:
for message, _ in pushing.stream({"messages": []}, stream_mode="messages"):
received.append(message.id)
return {}
return node
def test_pushed_message_streams_under_the_namespace_of_its_graph() -> None:
pushing = _one_node_graph(_push)
parent = StateGraph(MessagesState).add_node("sub", pushing).add_edge(START, "sub")
assert _namespaces(pushing, subgraphs=True) == [()]
assert _namespaces(parent.compile(), subgraphs=False) == []
assert _namespaces(parent.compile(), subgraphs=True) == [("sub",)]
def test_pushed_message_streams_to_the_node_that_streams_its_graph() -> None:
received: list[str] = []
graph = _one_node_graph(_streams_pushing_graph(received))
assert _namespaces(graph, subgraphs=False) == []
assert received == ["pushed"]
def test_pushed_message_streams_to_every_stream_that_takes_it() -> None:
received: list[str] = []
graph = _one_node_graph(_streams_pushing_graph(received))
assert _namespaces(graph, subgraphs=True) == [("node",)]
assert received == ["pushed"]
+3 -3
View File
@@ -3419,11 +3419,11 @@ wheels = [
[[package]]
name = "setuptools"
version = "84.1.0"
version = "83.0.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/73/eb/7e5f628ee2eb168fecbf8026de7856b2543d6d609cc85f086b9a91ac93d6/setuptools-84.1.0.tar.gz", hash = "sha256:23657205b63d4d31395ab917815247f694aafad28cd21316bf5fd45a51048ea0", size = 1171783, upload-time = "2026-10-19T20:11:43.489Z" }
sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/18/02/b4bbb90f8f29e334898ff45c6d9b1fffaa6f378017052b3023a06262b5f7/setuptools-84.1.0-py3-none-any.whl", hash = "sha256:5231689c6ac8156cd155cfb98c3fd6fb252c3e7e9bd3d78b64d3c55f132624c9", size = 819380, upload-time = "2026-10-19T20:11:41.684Z" },
{ url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" },
]
[[package]]
+59 -105
View File
@@ -14,14 +14,9 @@ import orjson
from langgraph_sdk._shared.utilities import (
_orjson_default,
_reconnect_delay,
_validate_reconnect_location,
)
from langgraph_sdk.errors import (
InternalServerError,
RateLimitError,
_araise_for_status_typed,
)
from langgraph_sdk.errors import _araise_for_status_typed
from langgraph_sdk.schema import QueryParamTypes, StreamPart
from langgraph_sdk.sse import SSEDecoder, aiter_lines_raw
@@ -151,8 +146,7 @@ class HttpClient:
on_response: Callable[[httpx.Response], None] | None = None,
reconnect_limit: int = 5,
) -> Any:
"""Send a request that reconnects to its Location header if the response is
cut off, with backoff between attempts."""
"""Send a request that automatically reconnects to Location header."""
request_headers, content = await _aencode_json(json)
if headers:
request_headers.update(headers)
@@ -182,27 +176,13 @@ class HttpClient:
stacklevel=2,
)
await r.aclose()
for attempt in range(reconnect_limit):
await asyncio.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
# A restart or a deploy looks like a refused connection, a
# dropped body, a 429 or a 5xx; other errors won't go away.
try:
return await self.request_reconnect(
loc,
"GET",
headers=headers,
# don't pass on_response so it's only called once
reconnect_limit=0,
)
except httpx.HTTPStatusError as err:
status = err.response.status_code
if attempt == reconnect_limit - 1 or (
status != 429 and status < 500
):
raise
except httpx.TransportError:
if attempt == reconnect_limit - 1:
raise
return await self.request_reconnect(
loc,
"GET",
headers=headers,
# don't pass on_response so it's only called once
reconnect_limit=reconnect_limit - 1,
)
async def stream(
self,
@@ -245,96 +225,70 @@ class HttpClient:
}
if last_event_id:
current_headers["Last-Event-ID"] = last_event_id
if reconnect_path is not None and not any(
key.lower() == "last-event-id" and value
for key, value in current_headers.items()
):
warnings.warn(
"Reconnecting without a Last-Event-ID, so events sent while "
"disconnected are not replayed. Only runs created with "
"stream_resumable=True send the event IDs a reconnect resumes "
"from.",
stacklevel=2,
)
current_method = method if reconnect_path is None else "GET"
current_content = content if reconnect_path is None else None
current_params = params if reconnect_path is None else None
reconnect_cause: httpx.HTTPError | None = None
try:
async with self.client.stream(
current_method,
reconnect_path or path,
headers=current_headers,
content=current_content,
params=current_params,
) as res:
if reconnect_path is None and on_response:
on_response(res)
# check status
await _araise_for_status_typed(res)
# check content type
content_type = res.headers.get("content-type", "").partition(";")[0]
if "text/event-stream" not in content_type:
raise httpx.TransportError(
"Expected response header Content-Type to contain 'text/event-stream', "
f"got {content_type!r}"
)
retry = False
async with self.client.stream(
current_method,
reconnect_path or path,
headers=current_headers,
content=current_content,
params=current_params,
) as res:
if reconnect_path is None and on_response:
on_response(res)
# check status
await _araise_for_status_typed(res)
# check content type
content_type = res.headers.get("content-type", "").partition(";")[0]
if "text/event-stream" not in content_type:
raise httpx.TransportError(
"Expected response header Content-Type to contain 'text/event-stream', "
f"got {content_type!r}"
)
reconnect_location = res.headers.get("location")
if reconnect_location:
_validate_reconnect_location(
self.client.base_url, reconnect_location
)
reconnect_path = reconnect_location
reconnect_location = res.headers.get("location")
if reconnect_location:
_validate_reconnect_location(
self.client.base_url, reconnect_location
)
reconnect_path = reconnect_location
# parse SSE
decoder = SSEDecoder(last_event_id=last_event_id)
try:
async for line in aiter_lines_raw(res):
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
if sse is not None:
if decoder.last_event_id is not None:
last_event_id = decoder.last_event_id
if sse.event or decoder.has_data:
# Progress since the last drop earns a fresh
# set of reconnect attempts.
reconnect_attempts = 0
yield sse
except httpx.HTTPError as err:
# httpx.TransportError inherits from HTTPError, so transient
# disconnects during streaming land here.
if reconnect_path is None:
raise
reconnect_cause = err
else:
if sse := decoder.decode(b""):
# parse SSE
decoder = SSEDecoder()
try:
async for line in aiter_lines_raw(res):
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
if sse is not None:
if decoder.last_event_id is not None:
last_event_id = decoder.last_event_id
if sse.event or decoder.has_data:
# decoder.decode(b"") flushes the in-flight event and may
# return an empty placeholder when there is no pending
# message. Skip these no-op events so the stream doesn't
# emit a trailing blank item after reconnects.
if sse.event or sse.data is not None:
yield sse
except (httpx.TransportError, RateLimitError, InternalServerError) as err:
# A reconnect that can't open (connection refused, 429, 5xx or a
# non-SSE body) usually means the server is restarting or
# mid-deploy, so retry it like a drop. Before the server sends a
# Location there is nothing to reconnect to, so raise.
if reconnect_path is None:
raise
reconnect_cause = err
if reconnect_cause is not None:
except httpx.HTTPError:
# httpx.TransportError inherits from HTTPError, so transient
# disconnects during streaming land here.
if reconnect_path is None:
raise
retry = True
else:
if sse := decoder.decode(b""):
if decoder.last_event_id is not None:
last_event_id = decoder.last_event_id
if sse.event or sse.data is not None:
# decoder.decode(b"") flushes the in-flight event and may
# return an empty placeholder when there is no pending
# message. Skip these no-op events so the stream doesn't
# emit a trailing blank item after reconnects.
yield sse
if retry:
reconnect_attempts += 1
if reconnect_attempts > max_reconnect_attempts:
raise httpx.TransportError(
"Exceeded maximum SSE reconnection attempts"
) from reconnect_cause
await asyncio.sleep(
_reconnect_delay(reconnect_attempts - 1, base=1.0, cap=5.0)
)
)
continue
break
+7 -8
View File
@@ -16,6 +16,7 @@ from __future__ import annotations
import asyncio
import contextlib
import logging
import random
from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping
from dataclasses import dataclass, field
from typing import Any, Literal, TypedDict, cast
@@ -24,7 +25,7 @@ from langchain_core.language_models.chat_model_stream import AsyncChatModelStrea
from langchain_protocol import Event, SubscribeParams
from langgraph_sdk._async.http import HttpClient
from langgraph_sdk._shared.utilities import _quote_path_param, _reconnect_delay
from langgraph_sdk._shared.utilities import _quote_path_param
from langgraph_sdk.schema import LangSmithTracing, QueryParamTypes
from langgraph_sdk.stream.controller import _SeenEventIds
from langgraph_sdk.stream.decoders import (
@@ -1690,13 +1691,11 @@ class AsyncThreadStream:
async def _reconnect_sleep(self, attempt: int) -> None:
"""Sleep with exponential backoff and jitter for reconnect attempt `attempt`."""
await asyncio.sleep(
_reconnect_delay(
attempt,
base=self._shared_reconnect_backoff_base,
cap=self._shared_reconnect_backoff_cap,
)
)
base = self._shared_reconnect_backoff_base
cap = self._shared_reconnect_backoff_cap
delay = min(cap, base * (2**attempt))
jitter = random.uniform(0, delay * 0.25)
await asyncio.sleep(delay + jitter)
async def _reconnect_shared_stream(self) -> bool:
"""Attempt to reopen the shared stream after a post-ready transport drop.
@@ -4,7 +4,6 @@ from __future__ import annotations
import functools
import os
import random
import re
from collections.abc import Mapping
from datetime import tzinfo
@@ -195,15 +194,6 @@ def _validate_reconnect_location(base_url: httpx.URL, location: str) -> str:
return location
def _reconnect_delay(attempt: int, *, base: float, cap: float) -> float:
"""Seconds to wait before reconnect attempt `attempt` (0-based).
Exponential backoff capped at `cap`, plus up to 25% jitter.
"""
delay = min(cap, base * (2**attempt))
return delay + random.uniform(0, delay * 0.25)
def _provided_vals(d: Mapping[str, Any]) -> dict[str, Any]:
return {k: v for k, v in d.items() if v is not None}
+56 -101
View File
@@ -4,7 +4,6 @@ from __future__ import annotations
import logging
import sys
import time
import warnings
from collections.abc import Callable, Iterator, Mapping
from typing import Any, cast
@@ -14,14 +13,9 @@ import orjson
from langgraph_sdk._shared.utilities import (
_orjson_default,
_reconnect_delay,
_validate_reconnect_location,
)
from langgraph_sdk.errors import (
InternalServerError,
RateLimitError,
_raise_for_status_typed,
)
from langgraph_sdk.errors import _raise_for_status_typed
from langgraph_sdk.schema import QueryParamTypes, StreamPart
from langgraph_sdk.sse import SSEDecoder, iter_lines_raw
@@ -152,8 +146,7 @@ class SyncHttpClient:
on_response: Callable[[httpx.Response], None] | None = None,
reconnect_limit: int = 5,
) -> Any:
"""Send a request that reconnects to its Location header if the response is
cut off, with backoff between attempts."""
"""Send a request that automatically reconnects to Location header."""
request_headers, content = _encode_json(json)
if headers:
request_headers.update(headers)
@@ -183,27 +176,13 @@ class SyncHttpClient:
stacklevel=2,
)
r.close()
for attempt in range(reconnect_limit):
time.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
# A restart or a deploy looks like a refused connection, a
# dropped body, a 429 or a 5xx; other errors won't go away.
try:
return self.request_reconnect(
loc,
"GET",
headers=headers,
# don't pass on_response so it's only called once
reconnect_limit=0,
)
except httpx.HTTPStatusError as err:
status = err.response.status_code
if attempt == reconnect_limit - 1 or (
status != 429 and status < 500
):
raise
except httpx.TransportError:
if attempt == reconnect_limit - 1:
raise
return self.request_reconnect(
loc,
"GET",
headers=headers,
# don't pass on_response so it's only called once
reconnect_limit=reconnect_limit - 1,
)
def stream(
self,
@@ -248,91 +227,67 @@ class SyncHttpClient:
}
if last_event_id:
current_headers["Last-Event-ID"] = last_event_id
if reconnect_path is not None and not any(
key.lower() == "last-event-id" and value
for key, value in current_headers.items()
):
warnings.warn(
"Reconnecting without a Last-Event-ID, so events sent while "
"disconnected are not replayed. Only runs created with "
"stream_resumable=True send the event IDs a reconnect resumes "
"from.",
stacklevel=2,
)
current_method = method if reconnect_path is None else "GET"
current_content = content if reconnect_path is None else None
current_params = params if reconnect_path is None else None
reconnect_cause: httpx.HTTPError | None = None
try:
with self.client.stream(
current_method,
reconnect_path or path,
headers=current_headers,
content=current_content,
params=current_params,
) as res:
if reconnect_path is None and on_response:
on_response(res)
# check status
_raise_for_status_typed(res)
# check content type
content_type = res.headers.get("content-type", "").partition(";")[0]
if "text/event-stream" not in content_type:
raise httpx.TransportError(
"Expected response header Content-Type to contain 'text/event-stream', "
f"got {content_type!r}"
)
retry = False
with self.client.stream(
current_method,
reconnect_path or path,
headers=current_headers,
content=current_content,
params=current_params,
) as res:
if reconnect_path is None and on_response:
on_response(res)
# check status
_raise_for_status_typed(res)
# check content type
content_type = res.headers.get("content-type", "").partition(";")[0]
if "text/event-stream" not in content_type:
raise httpx.TransportError(
"Expected response header Content-Type to contain 'text/event-stream', "
f"got {content_type!r}"
)
reconnect_location = res.headers.get("location")
if reconnect_location:
_validate_reconnect_location(
self.client.base_url, reconnect_location
)
reconnect_path = reconnect_location
reconnect_location = res.headers.get("location")
if reconnect_location:
_validate_reconnect_location(
self.client.base_url, reconnect_location
)
reconnect_path = reconnect_location
decoder = SSEDecoder(last_event_id=last_event_id)
try:
for line in iter_lines_raw(res):
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
if sse is not None:
if decoder.last_event_id is not None:
last_event_id = decoder.last_event_id
if sse.event or decoder.has_data:
# Progress since the last drop earns a fresh
# set of reconnect attempts.
reconnect_attempts = 0
yield sse
except httpx.HTTPError as err:
# httpx.TransportError inherits from HTTPError, so transient
# disconnects during streaming land here.
if reconnect_path is None:
raise
reconnect_cause = err
else:
if sse := decoder.decode(b""):
decoder = SSEDecoder()
try:
for line in iter_lines_raw(res):
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
if sse is not None:
if decoder.last_event_id is not None:
last_event_id = decoder.last_event_id
if sse.event or decoder.has_data:
# See async stream implementation for rationale on
# skipping empty flush events.
if sse.event or sse.data is not None:
yield sse
except (httpx.TransportError, RateLimitError, InternalServerError) as err:
# A reconnect that can't open (connection refused, 429, 5xx or a
# non-SSE body) usually means the server is restarting or
# mid-deploy, so retry it like a drop. Before the server sends a
# Location there is nothing to reconnect to, so raise.
if reconnect_path is None:
raise
reconnect_cause = err
if reconnect_cause is not None:
except httpx.HTTPError:
# httpx.TransportError inherits from HTTPError, so transient
# disconnects during streaming land here.
if reconnect_path is None:
raise
retry = True
else:
if sse := decoder.decode(b""):
if decoder.last_event_id is not None:
last_event_id = decoder.last_event_id
if sse.event or sse.data is not None:
# See async stream implementation for rationale on
# skipping empty flush events.
yield sse
if retry:
reconnect_attempts += 1
if reconnect_attempts > max_reconnect_attempts:
raise httpx.TransportError(
"Exceeded maximum SSE reconnection attempts"
) from reconnect_cause
time.sleep(_reconnect_delay(reconnect_attempts - 1, base=1.0, cap=5.0))
)
continue
break
+1 -1
View File
@@ -811,7 +811,7 @@ class _On:
def __call__(
self,
*,
resources: str | Sequence[str] | None = None,
resources: str | Sequence[str],
actions: str | Sequence[str] | None = None,
) -> Callable[[AHO], AHO]: ...
+20 -33
View File
@@ -30,14 +30,18 @@ class BytesLineDecoder:
# See https://docs.python.org/3/glossary.html#term-universal-newlines
NEWLINE_CHARS = b"\n\r"
if not text:
return []
if self.trailing_cr and text.startswith(b"\n"):
text = text[1:]
self.trailing_cr = text.endswith(b"\r")
# We always push a trailing `\r` into the next decode iteration.
if self.trailing_cr:
text = b"\r" + text
self.trailing_cr = False
if text.endswith(b"\r"):
self.trailing_cr = True
text = text[:-1]
if not text:
return []
# NOTE: the edge case input of empty text doesn't occur in practice,
# because other httpx internals filter out this value
return [] # pragma: no cover
trailing_newline = text[-1] in NEWLINE_CHARS
lines = cast(list[BytesLike], text.splitlines())
@@ -62,7 +66,7 @@ class BytesLineDecoder:
return lines
def flush(self) -> list[BytesLike]:
if not self.buffer:
if not self.buffer and not self.trailing_cr:
return []
lines: list[BytesLike] = [self.buffer]
@@ -72,25 +76,18 @@ class BytesLineDecoder:
class SSEDecoder:
def __init__(self, *, last_event_id: str | None = None) -> None:
def __init__(self) -> None:
self._event = ""
self._data = bytearray()
self._last_event_id = last_event_id
self._last_event_id: str | None = None
self._retry: int | None = None
self._at_stream_start = True
self._has_data = False
self._id_changed = False
@property
def has_data(self) -> bool:
"""Whether the last decoded block contained a data field."""
return self._has_data
@property
def last_event_id(self) -> str | None:
"""Return the last event identifier that was seen.
`None` until an ID is supplied or received, and `""` after an empty
`None` until the stream sends an `id` field, and `""` after an empty
one, which resets the reconnect cursor.
"""
@@ -104,26 +101,17 @@ class SSEDecoder:
line = line.removeprefix(b"\xef\xbb\xbf")
if not line:
self._has_data = bool(self._data)
if (
not self._event
and not self._data
and not self._id_changed
and self._last_event_id is None
and self._retry is None
):
return None
body = self._data[:-1]
data = None
if body and not body.isspace():
try:
data = orjson.loads(body)
except orjson.JSONDecodeError:
data = orjson.loads(body.decode(errors="replace"))
sse = StreamPart(
event=self._event,
data=data, # ty: ignore[invalid-argument-type]
data=orjson.loads(self._data) if self._data else None, # ty: ignore[invalid-argument-type]
id=self._last_event_id or None,
)
@@ -131,7 +119,6 @@ class SSEDecoder:
self._event = ""
self._data = bytearray()
self._retry = None
self._id_changed = False
return sse
@@ -144,16 +131,16 @@ class SSEDecoder:
value = value[1:]
if fieldname == b"event":
self._event = value.decode(errors="replace")
self._event = value.decode()
elif fieldname == b"data":
if self._data:
self._data.extend(b"\n")
self._data.extend(value)
self._data.extend(b"\n")
elif fieldname == b"id":
if b"\0" in value:
pass
else:
self._last_event_id = value.decode(errors="replace")
self._id_changed = True
self._last_event_id = value.decode()
elif fieldname == b"retry":
with contextlib.suppress(TypeError, ValueError):
self._retry = int(value)
@@ -15,6 +15,7 @@ from __future__ import annotations
import asyncio
import contextlib
import logging
import random
from collections import OrderedDict
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
from dataclasses import dataclass, field
@@ -22,7 +23,6 @@ from typing import Any
from langchain_protocol import Event, SubscribeParams
from langgraph_sdk._shared.utilities import _reconnect_delay
from langgraph_sdk.stream.transport import AsyncProtocolTransport, EventStreamHandle
# ---------------------------------------------------------------------------
@@ -275,13 +275,11 @@ class StreamController:
async def _reconnect_sleep(self, attempt: int) -> None:
"""Sleep with exponential backoff and jitter for reconnect attempt *attempt*."""
await asyncio.sleep(
_reconnect_delay(
attempt,
base=self._reconnect_backoff_base,
cap=self._reconnect_backoff_cap,
)
)
base = self._reconnect_backoff_base
cap = self._reconnect_backoff_cap
delay = min(cap, base * (2**attempt))
jitter = random.uniform(0, delay * 0.25)
await asyncio.sleep(delay + jitter)
async def _reconnect_shared_stream(self) -> bool:
"""Attempt to reopen the shared stream after a transport drop.
@@ -4,6 +4,7 @@ from __future__ import annotations
import contextlib
import logging
import random
import threading
import time
from dataclasses import dataclass, field
@@ -12,7 +13,6 @@ from typing import Any
from langchain_protocol import Event, SubscribeParams
from langgraph_sdk._shared.utilities import _reconnect_delay
from langgraph_sdk.stream.errors import (
is_non_retryable_stream_error,
stream_error_frame,
@@ -282,13 +282,11 @@ class SyncStreamController:
def _reconnect_sleep(self, attempt: int) -> None:
"""Sleep with exponential backoff + jitter before a reconnect attempt."""
time.sleep(
_reconnect_delay(
attempt,
base=self._reconnect_backoff_base,
cap=self._reconnect_backoff_cap,
)
)
base = self._reconnect_backoff_base
cap = self._reconnect_backoff_cap
delay = min(cap, base * (2**attempt))
jitter = random.uniform(0, delay * 0.25)
time.sleep(delay + jitter)
def _reconnect_shared_stream(self) -> bool:
"""Attempt to reopen the shared stream after a transport drop.
-17
View File
@@ -1,17 +0,0 @@
import asyncio
import time
import pytest
@pytest.fixture
def no_sleep(monkeypatch):
"""Record backoff sleeps without actually waiting."""
sleeps: list[float] = []
async def fake_async_sleep(d):
sleeps.append(d)
monkeypatch.setattr(asyncio, "sleep", fake_async_sleep)
monkeypatch.setattr(time, "sleep", lambda d: sleeps.append(d))
return sleeps
-41
View File
@@ -1,49 +1,8 @@
import ast
import inspect
import pytest
from langgraph_sdk import Auth
@pytest.mark.parametrize(
"kwargs",
[{}, {"actions": "read"}, {"resources": None}],
)
def test_global_handler_factory_accepts_optional_resources(kwargs) -> None:
auth = Auth()
tree = ast.parse(inspect.getsource(type(auth.on)))
cls = tree.body[0]
assert isinstance(cls, ast.ClassDef)
factory_overload = next(
node
for node in cls.body
if isinstance(node, ast.FunctionDef)
and node.name == "__call__"
and node.decorator_list
and any(arg.arg == "resources" for arg in node.args.kwonlyargs)
)
resources = next(
arg for arg in factory_overload.args.kwonlyargs if arg.arg == "resources"
)
assert resources.annotation is not None
assert ast.unparse(resources.annotation) == "str | Sequence[str] | None"
default = factory_overload.args.kw_defaults[0]
assert isinstance(default, ast.Constant) and default.value is None
@auth.on(**kwargs)
async def handler(ctx, value):
del ctx, value
return None
if kwargs.get("actions"):
assert auth._handlers == {("*", "read"): [handler]}
assert auth._global_handlers == []
else:
assert auth._global_handlers == [handler]
assert auth._handlers == {}
def test_handler_multiple_resources_and_actions() -> None:
auth = Auth()
+7 -424
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from collections.abc import Callable, Iterator, Sequence
from collections.abc import Iterator, Sequence
from pathlib import Path
from typing import Any
@@ -13,7 +13,6 @@ from langgraph_sdk._async.runs import _wrap_stream_v2
from langgraph_sdk._shared.utilities import _sse_to_v2_dict
from langgraph_sdk._sync.runs import _wrap_stream_v2_sync
from langgraph_sdk.client import HttpClient, SyncHttpClient
from langgraph_sdk.errors import InternalServerError, NotFoundError
from langgraph_sdk.schema import (
CheckpointPayload,
CheckpointsStreamPart,
@@ -131,110 +130,6 @@ def test_sse_decoder_strips_a_leading_bom():
assert decoder.decode(b"") == StreamPart(event="values", data=1, id=None) # ty: ignore[invalid-argument-type]
def test_control_blocks_do_not_repeat_after_a_cursor_update():
decoder = SSEDecoder()
decoder.decode(b"id: cursor")
assert decoder.decode(b"") == StreamPart(event="", data=None, id="cursor") # ty: ignore[invalid-argument-type]
decoder.decode(b": heartbeat")
assert decoder.decode(b"") is None
assert decoder.decode(b"") is None
decoder.decode(b"id:")
assert decoder.decode(b"") == StreamPart(event="", data=None, id=None) # ty: ignore[invalid-argument-type]
assert decoder.last_event_id == ""
assert decoder.decode(b"") is None
def test_cr_lines_are_emitted_without_waiting_for_another_chunk():
decoder = BytesLineDecoder()
assert decoder.decode(b"data: 1\r\r") == [b"data: 1", b""]
assert decoder.decode(b"\n") == []
assert decoder.decode(b"data: 2\r") == [b"data: 2"]
assert decoder.decode(b"") == []
assert decoder.decode(b"\n\r") == [b""]
assert decoder.flush() == []
def test_empty_data_lines_preserve_their_newlines():
decoder = SSEDecoder()
for line in (b"data", b"data: 1", b"data:"):
decoder.decode(line)
assert decoder._data == b"\n1\n\n"
assert decoder.decode(b"") == StreamPart(event="", data=1, id=None) # ty: ignore[invalid-argument-type]
decoder.decode(b"data:")
decoder.decode(b"data:")
assert decoder.decode(b"") == StreamPart(event="", data=None, id=None) # ty: ignore[invalid-argument-type]
assert decoder.has_data
assert decoder.decode(b"") is None
assert not decoder.has_data
def test_malformed_utf8_is_replaced_in_fields_and_json_strings():
decoder = SSEDecoder()
for line in (b"event: \xff", b"id: \xff", b'data: "\xff"'):
decoder.decode(line)
assert decoder.decode(b"") == StreamPart(event="\ufffd", data="\ufffd", id="\ufffd") # ty: ignore[invalid-argument-type]
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_null_and_empty_data_are_not_dropped(asynchronous: bool):
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=b"data: null\n\ndata:\n\n",
)
transport = httpx.MockTransport(handler)
if asynchronous:
async with httpx.AsyncClient(
transport=transport, base_url="https://example.com"
) as client:
parts = [part async for part in HttpClient(client).stream("/stream", "GET")]
else:
with httpx.Client(
transport=transport, base_url="https://example.com"
) as client:
parts = list(SyncHttpClient(client).stream("/stream", "GET"))
assert parts == [StreamPart(event="", data=None), StreamPart(event="", data=None)] # ty: ignore[invalid-argument-type]
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_cursor_is_inherited_after_reconnect(asynchronous: bool):
requests: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
if len(requests) == 1:
stream_cls = AsyncListByteStream if asynchronous else ListByteStream
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream", "Location": "/reconnect"},
stream=stream_cls(
[b"id: cursor\nevent: values\ndata: 1\n\n"],
httpx.RemoteProtocolError("incomplete chunked read"),
),
)
assert (b"Last-Event-ID", b"cursor") in request.headers.raw
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=b"event: values\ndata: 2\n\n",
)
transport = httpx.MockTransport(handler)
if asynchronous:
async with httpx.AsyncClient(
transport=transport, base_url="https://example.com"
) as client:
parts = [part async for part in HttpClient(client).stream("/stream", "GET")]
else:
with httpx.Client(
transport=transport, base_url="https://example.com"
) as client:
parts = list(SyncHttpClient(client).stream("/stream", "GET"))
assert [part.id for part in parts] == ["cursor", "cursor"]
# --- HTTP client streaming ---
@@ -281,7 +176,6 @@ def test_sync_http_client_stream_flushes_trailing_event():
assert parts == [StreamPart(event="foo", data={"bar": 1})]
@pytest.mark.usefixtures("no_sleep")
def test_sync_http_client_stream_recovers_after_disconnect():
reconnect_path = "/reconnect"
first_chunks = [
@@ -347,7 +241,6 @@ def test_sync_http_client_stream_recovers_after_disconnect():
@pytest.mark.asyncio
@pytest.mark.usefixtures("no_sleep")
async def test_http_client_stream_recovers_after_disconnect():
reconnect_path = "/reconnect"
first_chunks = [
@@ -482,33 +375,28 @@ _CALLER_CURSORS = pytest.mark.parametrize(
@_EMPTY_IDS
@_CALLER_CURSORS
@pytest.mark.usefixtures("no_sleep")
def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
empty_id: bytes, headers: dict[str, str] | None
):
with pytest.warns(UserWarning, match="Last-Event-ID"):
reconnect = _sync_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
reconnect = _sync_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
assert "last-event-id" not in reconnect.headers
@_EMPTY_IDS
@_CALLER_CURSORS
@pytest.mark.usefixtures("no_sleep")
async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id(
empty_id: bytes, headers: dict[str, str] | None
):
with pytest.warns(UserWarning, match="Last-Event-ID"):
reconnect = await _async_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
reconnect = await _async_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
assert "last-event-id" not in reconnect.headers
@pytest.mark.usefixtures("no_sleep")
def test_sync_stream_id_replaces_a_lowercase_caller_cursor():
reconnect = _sync_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n"], {"last-event-id": "5"}
@@ -517,7 +405,6 @@ def test_sync_stream_id_replaces_a_lowercase_caller_cursor():
assert reconnect.headers.get_list("last-event-id") == ["1"]
@pytest.mark.usefixtures("no_sleep")
async def test_async_stream_id_replaces_a_lowercase_caller_cursor():
reconnect = await _async_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n"], {"last-event-id": "5"}
@@ -526,14 +413,12 @@ async def test_async_stream_id_replaces_a_lowercase_caller_cursor():
assert reconnect.headers.get_list("last-event-id") == ["1"]
@pytest.mark.usefixtures("no_sleep")
def test_sync_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id():
reconnect = _sync_reconnect([b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"})
assert reconnect.headers.get_list("last-event-id") == ["5"]
@pytest.mark.usefixtures("no_sleep")
async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id():
reconnect = await _async_reconnect(
[b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"}
@@ -542,308 +427,6 @@ async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id
assert reconnect.headers.get_list("last-event-id") == ["5"]
# A reconnect that sends no Last-Event-ID can't replay what was sent while it was
# disconnected, so it warns.
_RECONNECT_CURSORS = pytest.mark.parametrize(
("chunks", "headers", "warns"),
[
([b"event: values\ndata: 1\n\n"], None, True),
([b"id: 1\nevent: values\ndata: 1\n\n"], None, False),
([b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"}, False),
],
ids=["no-cursor", "stream-cursor", "caller-cursor"],
)
@_RECONNECT_CURSORS
@pytest.mark.usefixtures("no_sleep")
def test_sync_reconnect_warns_only_without_a_cursor(
chunks: list[bytes],
headers: dict[str, str] | None,
warns: bool,
recwarn: pytest.WarningsRecorder,
):
_sync_reconnect(chunks, headers)
assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
@_RECONNECT_CURSORS
@pytest.mark.usefixtures("no_sleep")
async def test_async_reconnect_warns_only_without_a_cursor(
chunks: list[bytes],
headers: dict[str, str] | None,
warns: bool,
recwarn: pytest.WarningsRecorder,
):
await _async_reconnect(chunks, headers)
assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
# --- reconnects survive a server restart or a deploy ---
SSE = {"Content-Type": "text/event-stream"}
DROP = httpx.RemoteProtocolError("incomplete chunked read")
def _event(i: int) -> bytes:
return f"id: {i}\nevent: values\ndata: {i}\n\n".encode()
def _scripted(steps: list[Callable[[httpx.Request], httpx.Response]]):
"""Handler that plays one step per request and records each request."""
seen: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
seen.append(request)
if len(seen) > len(steps):
raise AssertionError("unexpected request")
return steps[len(seen) - 1](request)
return handler, seen
def _ok(stream_cls, chunks, *, drop=False, location=None):
headers = dict(SSE)
if location:
headers["Location"] = location
def step(_request):
return httpx.Response(
200, headers=headers, stream=stream_cls(chunks, DROP if drop else None)
)
return step
def _status(code):
return lambda request: httpx.Response(
code, json={"detail": "upstream"}, request=request
)
def _connect_error(request):
raise httpx.ConnectError("[Errno 61] Connection refused", request=request)
def _not_sse(request):
return httpx.Response(
200, headers={"Content-Type": "text/html"}, text="<html>", request=request
)
def _transient_steps(stream_cls, failure):
"""POST drops after event 1; first reconnect fails; second succeeds."""
return [
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
failure,
_ok(stream_cls, [_event(2), b"event: end\ndata: null\n\n"]),
]
def _long_stream_steps(stream_cls, drops):
"""POST + `drops` reconnects, each delivering one event then dropping."""
steps = [_ok(stream_cls, [_event(0)], drop=True, location="/reconnect")]
for i in range(1, drops):
steps.append(_ok(stream_cls, [_event(i)], drop=True))
steps.append(_ok(stream_cls, [_event(drops), b"event: end\ndata: null\n\n"]))
return steps
def _run_sync(steps):
handler, seen = _scripted(steps)
with httpx.Client(
transport=httpx.MockTransport(handler), base_url="https://example.com"
) as client:
parts = list(SyncHttpClient(client).stream("/stream", "POST", json={}))
return parts, seen
async def _run_async(steps):
sync_handler, seen = _scripted(steps)
async def handler(request):
return sync_handler(request)
async with httpx.AsyncClient(
transport=httpx.MockTransport(handler), base_url="https://example.com"
) as client:
parts = [p async for p in HttpClient(client).stream("/stream", "POST", json={})]
return parts, seen
FAILURES = {
"429": _status(429),
"502": _status(502),
"503": _status(503),
"connect_error": _connect_error,
"not_sse": _not_sse,
}
# A reconnect that hits a 429 / 502 / 503 / connect error / non-SSE body, then
# succeeds, should recover.
@pytest.mark.usefixtures("no_sleep")
@pytest.mark.parametrize("failure", FAILURES.values(), ids=FAILURES.keys())
def test_sync_reconnect_recovers_from_transient_failure(failure):
parts, seen = _run_sync(_transient_steps(ListByteStream, failure))
assert [p.event for p in parts] == ["values", "values", "end"]
assert len(seen) == 3
@pytest.mark.usefixtures("no_sleep")
@pytest.mark.parametrize("failure", FAILURES.values(), ids=FAILURES.keys())
async def test_async_reconnect_recovers_from_transient_failure(failure):
parts, seen = await _run_async(_transient_steps(AsyncListByteStream, failure))
assert [p.event for p in parts] == ["values", "values", "end"]
assert len(seen) == 3
# Any other 4xx on a reconnect won't go away by retrying, so it raises.
def _client_error_steps(stream_cls):
return [
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
_status(404),
]
@pytest.mark.usefixtures("no_sleep")
def test_sync_reconnect_does_not_retry_a_client_error():
with pytest.raises(NotFoundError):
_run_sync(_client_error_steps(ListByteStream))
@pytest.mark.usefixtures("no_sleep")
async def test_async_reconnect_does_not_retry_a_client_error():
with pytest.raises(NotFoundError):
await _run_async(_client_error_steps(AsyncListByteStream))
# The first request isn't a reconnect, so it raises even when it names a Location.
def _first_request_503(request):
return httpx.Response(
503,
headers={"Location": "/reconnect"},
json={"detail": "upstream"},
request=request,
)
@pytest.mark.usefixtures("no_sleep")
def test_sync_first_request_failure_is_not_retried():
with pytest.raises(InternalServerError):
_run_sync([_first_request_503])
@pytest.mark.usefixtures("no_sleep")
async def test_async_first_request_failure_is_not_retried():
with pytest.raises(InternalServerError):
await _run_async([_first_request_503])
# A server that stays down gets five reconnects, backing off 1, 2, 4, 5 and 5s
# (each plus up to 25% jitter), so about 20s, before the stream gives up. The
# error it gives up with is caused by the last failure, so callers see why.
GIVE_UP = "Exceeded maximum SSE reconnection attempts"
def _server_down_steps(stream_cls):
return [
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
*[_status(502)] * 5,
]
def _assert_backed_off_for_about_20s(sleeps: list[float]) -> None:
assert len(sleeps) == 5
for slept, delay in zip(sleeps, [1, 2, 4, 5, 5], strict=True):
assert delay <= slept <= delay * 1.25
def test_sync_reconnect_gives_up_after_five_backed_off_attempts(no_sleep):
with pytest.raises(httpx.TransportError, match=GIVE_UP) as exc_info:
_run_sync(_server_down_steps(ListByteStream))
assert isinstance(exc_info.value.__cause__, InternalServerError)
_assert_backed_off_for_about_20s(no_sleep)
async def test_async_reconnect_gives_up_after_five_backed_off_attempts(no_sleep):
with pytest.raises(httpx.TransportError, match=GIVE_UP) as exc_info:
await _run_async(_server_down_steps(AsyncListByteStream))
assert isinstance(exc_info.value.__cause__, InternalServerError)
_assert_backed_off_for_about_20s(no_sleep)
# A long stream that drops 6 times, each reconnect succeeding and making
# progress, should not run out of attempts.
@pytest.mark.usefixtures("no_sleep")
def test_sync_long_stream_survives_many_spaced_out_drops():
parts, _ = _run_sync(_long_stream_steps(ListByteStream, drops=6))
assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
@pytest.mark.usefixtures("no_sleep")
async def test_async_long_stream_survives_many_spaced_out_drops():
parts, _ = await _run_async(_long_stream_steps(AsyncListByteStream, drops=6))
assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
# Each reconnect resumes after the newest event the stream got, so a reconnect
# loop on a resumable run replays only what it missed. A failed reconnect keeps
# the cursor for the next one.
RECONNECT_LOOPS = {
"progress_between_drops": (
lambda stream_cls: _long_stream_steps(stream_cls, drops=3),
[None, "0", "1", "2"],
),
"failed_reconnect": (
lambda stream_cls: _transient_steps(stream_cls, _status(502)),
[None, "1", "1"],
),
}
@pytest.mark.usefixtures("no_sleep")
@pytest.mark.parametrize(
("steps", "cursors"), RECONNECT_LOOPS.values(), ids=RECONNECT_LOOPS.keys()
)
def test_sync_reconnect_loop_resumes_after_the_newest_event(steps, cursors):
_, seen = _run_sync(steps(ListByteStream))
assert [r.headers.get("last-event-id") for r in seen] == cursors
@pytest.mark.usefixtures("no_sleep")
@pytest.mark.parametrize(
("steps", "cursors"), RECONNECT_LOOPS.values(), ids=RECONNECT_LOOPS.keys()
)
async def test_async_reconnect_loop_resumes_after_the_newest_event(steps, cursors):
_, seen = await _run_async(steps(AsyncListByteStream))
assert [r.headers.get("last-event-id") for r in seen] == cursors
# Reconnects should back off rather than hammer the server.
def test_sync_reconnect_backs_off(no_sleep):
_run_sync(_long_stream_steps(ListByteStream, drops=2))
assert len(no_sleep) == 2 and all(d > 0 for d in no_sleep)
async def test_async_reconnect_backs_off(no_sleep):
await _run_async(_long_stream_steps(AsyncListByteStream, drops=2))
assert len(no_sleep) == 2 and all(d > 0 for d in no_sleep)
# --- _sse_to_v2_dict conversion ---
+3 -98
View File
@@ -1,8 +1,7 @@
from __future__ import annotations
import json
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from typing import Any
from collections.abc import AsyncIterator, Iterator
import httpx
import pytest
@@ -25,16 +24,7 @@ class _BodyDropsMidRead(httpx.SyncByteStream, httpx.AsyncByteStream):
raise httpx.ReadError("connection dropped")
def _done(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"done": True})
def _wait_handler(
requests: list[httpx.Request],
reconnects: Sequence[Callable[[httpx.Request], httpx.Response]] = (_done,),
):
"""POST drops mid-body; the reconnects answer from `reconnects` in order."""
def _wait_handler(requests: list[httpx.Request]):
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
if request.method == "POST":
@@ -43,7 +33,7 @@ def _wait_handler(
headers={"Location": "/threads/t1/runs/r1/join"},
stream=_BodyDropsMidRead(),
)
return reconnects[len(requests) - 2](request)
return httpx.Response(200, json={"done": True})
return handler
@@ -56,7 +46,6 @@ def _assert_bodyless_get(request: httpx.Request) -> None:
assert request.read() == b""
@pytest.mark.usefixtures("no_sleep")
def test_sync_wait_reconnects_with_a_bodyless_get():
requests: list[httpx.Request] = []
transport = httpx.MockTransport(_wait_handler(requests))
@@ -69,7 +58,6 @@ def test_sync_wait_reconnects_with_a_bodyless_get():
_assert_bodyless_get(requests[1])
@pytest.mark.usefixtures("no_sleep")
async def test_async_wait_reconnects_with_a_bodyless_get():
requests: list[httpx.Request] = []
transport = httpx.MockTransport(_wait_handler(requests))
@@ -84,89 +72,6 @@ async def test_async_wait_reconnects_with_a_bodyless_get():
_assert_bodyless_get(requests[1])
def _connect_error(request: httpx.Request) -> httpx.Response:
raise httpx.ConnectError("[Errno 61] Connection refused", request=request)
async def _wait(
handler: Callable[[httpx.Request], httpx.Response], *, asynchronous: bool
) -> Any:
transport = httpx.MockTransport(handler)
if asynchronous:
async with httpx.AsyncClient(
transport=transport, base_url="https://example.com"
) as client:
return await RunsClient(HttpClient(client)).wait("t1", "agent", input={})
with httpx.Client(transport=transport, base_url="https://example.com") as client:
return SyncRunsClient(SyncHttpClient(client)).wait("t1", "agent", input={})
# What a server restart or a deploy looks like to the reconnect GET: each one is
# retried after a backoff.
RECONNECT_FAILURES = {
"429": lambda request: httpx.Response(429, request=request),
"503": lambda request: httpx.Response(503, request=request),
"connect_error": _connect_error,
"body_drops": lambda _request: httpx.Response(
200,
headers={"Location": "/threads/t1/runs/r1/join"},
stream=_BodyDropsMidRead(),
),
}
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize(
"failure", RECONNECT_FAILURES.values(), ids=RECONNECT_FAILURES.keys()
)
async def test_wait_retries_a_failed_reconnect_after_a_backoff(
failure, asynchronous, no_sleep
):
requests: list[httpx.Request] = []
with pytest.warns(UserWarning, match="attempting reconnect"):
result = await _wait(
_wait_handler(requests, [failure, _done]), asynchronous=asynchronous
)
assert result == {"done": True}
assert len(requests) == 3
assert len(no_sleep) == 2 and min(no_sleep) >= 1, no_sleep
@pytest.mark.usefixtures("no_sleep")
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_wait_raises_after_five_failed_reconnects(asynchronous):
requests: list[httpx.Request] = []
with (
pytest.warns(UserWarning, match="attempting reconnect"),
pytest.raises(httpx.HTTPStatusError, match="503"),
):
await _wait(
_wait_handler(requests, [RECONNECT_FAILURES["503"]] * 5),
asynchronous=asynchronous,
)
assert len(requests) == 6
@pytest.mark.usefixtures("no_sleep")
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_wait_does_not_retry_a_client_error(asynchronous):
requests: list[httpx.Request] = []
with (
pytest.warns(UserWarning, match="attempting reconnect"),
pytest.raises(httpx.HTTPStatusError, match="404"),
):
await _wait(
_wait_handler(
requests, [lambda request: httpx.Response(404, request=request)]
),
asynchronous=asynchronous,
)
assert len(requests) == 2
def test_sync_wait_raises_the_run_error():
requests: list[httpx.Request] = []