Compare commits

...
Author SHA1 Message Date
Hari Dhanushkodiandopen-swe[bot] 6e144fb893 feat(cli): log in to LangSmith without an API key
Reuse the `langsmith login` OAuth profile (~/.langsmith/config.json) for the
deploy commands, refreshing the 5-minute access token per request, and add
`langgraph login` (device flow). With no credentials, an interactive deploy
starts the login flow instead of prompting for an API key.

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
2026-10-30 20:16:51 +00:00
Elior Nataf Lackritz 1f0cf816d5 fix(langgraph): give messages and tools events the namespace of their graph (#9458)
Fixes #6373
Fixes #5674

`messages` and `tools` events should stream under the namespace of the
graph they ran in, and a stream without `subgraphs=True` should get only
its own graph's events, plus those of a graph that one of its nodes
streams itself. langchain-ai/langgraphjs#3121 applies these rules in JS.
Python broke them in four places:

- **`@task` calls.** A task call nests its checkpoint namespace segment
under its caller's, although it runs in the caller's graph. The handlers
dropped only the last segment, so a model or tool called in a task
looked like a subgraph's: without `subgraphs=True` its events were
dropped, so a functional API agent streamed no tokens (#6373), and with
it they came under the calling node's namespace, one level deeper per
nested call. The handlers now also drop one segment per enclosing task
call, read from the task's `langgraph_path`. Send tasks and plain nodes
are unchanged.
- **A graph that runs outside the caller's task.** A graph invoked
inside a node with its own `thread_id`, `checkpoint_ns` or
`checkpoint_id`, or through a LangChain runnable called with its own
`configurable` (which replaces the caller's), starts its namespace at
the root, so a handler inherited from the caller streamed its events as
the caller's own, even with `subgraphs=False` (#5674). Such a graph now
drops the `messages` and `tools` handlers it inherits, whatever their
`subgraphs`, so its events stay out of the caller's stream like its
`updates` and `custom` output. A node that streams it still gets its
events.
- **A graph a node streams after running another.** The loop numbers the
namespace of every subgraph a task runs after its first (`node:<id>|1`),
but the handlers read their own namespace before that, so the node got
none of that graph's events. They now read it from the loop.
- **`push_message`.** It kept the pushing node's own segment, ignored
`subgraphs=False`, and reached only the first messages handler, so a
node that streamed a graph never got that graph's pushed messages. It
now follows the same rules, for every messages handler.

Reading the namespace from the loop changes one more case: a root run
given a subgraph's config, with its `checkpoint_ns` (as
`get_state(subgraphs=True)` returns it), took that namespace as its own
and streamed the subgraph's tokens with `subgraphs=False`. The loop runs
it at the root, so it now streams the subgraph node's returned messages,
like any root run.

`stream_events(version="v3")` reads the same namespaces, so a task's
tokens and tool calls and a root node's pushed messages now show in the
run's `messages` and `tool_calls` projections, where they were missing.

The new tests in `test_stream_namespaces.py` fail on `main`, except two
guards that invoke a child graph with no config. No existing test
changed.

Thanks @sdole and @Kannndev for the reports, @lxingy3 for the smaller
repro and diagnosis on #5674, and @roydonsequeira for the review that
caught the `subgraphs=True` case and the 3.10 failure.

## Release note

`messages` and `tools` events from a `@task` now stream without
`subgraphs=True`, as events of the graph that called the task, so a
functional API agent's `messages` stream carries its tokens. With
`subgraphs=True` they move from the calling node's namespace, such as
`("workflow:<id>",)`, to their graph's, `()` at the root. A graph
invoked inside a node with its own `thread_id`, `checkpoint_ns` or
`checkpoint_id`, or through a runnable called with its own
`configurable`, no longer streams into the caller's `messages` or
`tools` stream, even when the caller passes `subgraphs=True`; a node
that streams it still gets its events. A node that streams a graph after
running another one now gets that graph's events. Messages sent with
`push_message` stream under the namespace of their graph, stay out of a
stream without `subgraphs=True` when they come from a subgraph, and
reach a node that streams their graph. A root run given a subgraph's
`checkpoint_ns` with `subgraphs=False` now streams that subgraph node's
returned messages instead of its tokens. The `stream_mode` chunks that
`astream_events` carries change the same way; its own events do not.
2026-10-30 15:07:34 -04:00
Elior Nataf Lackritz 6f2df053f4 test(sdk-py): check that a reconnect loop resumes after the newest event (#9483)
Follow-up to the review on #9482: with `stream_resumable=True`, does a
loop of reconnects replay events? Each reconnect sends the newest event
id the stream got, and a failed reconnect keeps that id for the next
one, so the server replays only what the client missed. The existing
loop tests served the same events whatever the cursor, so nothing
checked this.

`test_*_reconnect_loop_resumes_after_the_newest_event` checks the
`Last-Event-ID` of every request in two loops, sync and async: progress
between drops, and a 502 between two reconnects. A stream that kept its
first cursor, or dropped it after a failed reconnect, fails it.
2026-10-29 15:08:20 -04:00
Elior Nataf Lackritz ffc3336367 fix(sdk-py): retry runs.wait, join and cancel(wait=True) reconnects with backoff (#9481)
`runs.wait`, `runs.join` and `runs.cancel(wait=True)` reconnect to the
run's `Location` when the response is cut off, but with no delay: a 429,
a 5xx or a refused connection on the reconnect raised at once, and a
dropped body reconnected again immediately, so they could not ride out a
server restart or a deploy. #9462 fixed the same for streams.

The reconnect now waits with the stream path's backoff and retries those
failures within the same five attempts; any other status and a failed
first request still raise. Each attempt is one GET with
`reconnect_limit=0`, so the loop owns every retry (the join `Location`
points back at itself, so nothing is lost by not following it).
`request_reconnect` raises plain `httpx.HTTPStatusError`, so it checks
429 and 5xx directly, the same set the stream path catches as
`RateLimitError` and `InternalServerError`.

`test_wait_retries_a_failed_reconnect_after_a_backoff` covers each
failure, sync and async, and fails on `main`; the other two new tests
check the budget and that a 404 still raises. `no_sleep` moves to
`conftest.py` so both test files can use it.
2026-10-28 18:55:31 -04:00
Elior Nataf Lackritz 5f4b28a9f7 fix(sdk-py): warn when a stream reconnects without a Last-Event-ID (#9482)
A dropped stream reconnects even when it has no event id to resume from
(a run created without `stream_resumable=True`, or one that dropped
before its first event). The reconnect can't replay what was sent while
it was disconnected, so those events were lost without a sign.

The SDK now warns when a reconnect request carries no `Last-Event-ID`,
whether the cursor would come from the stream's ids or from the caller's
own header. It still reconnects: everything sent after the reconnect
arrives, so refusing would turn a gap into losing the rest of the
stream.

The new table test covers a stream with no ids (warns; fails on `main`),
one with ids and one with a caller cursor (no warning), sync and async.
The empty-id tests now expect the warning, since they reconnect without
a cursor too.
2026-10-28 18:44:56 -04:00
16 changed files with 1117 additions and 123 deletions

No files matched your search

+8
View File
@@ -11,6 +11,7 @@ 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
@@ -233,6 +234,13 @@ 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)
+26 -12
View File
@@ -1,6 +1,7 @@
"""Deploy command and subcommands for the LangGraph CLI."""
import base64
import functools
import json as json_mod
import os
import pathlib
@@ -19,6 +20,7 @@ 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
@@ -1982,17 +1984,6 @@ 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"
)
@@ -2000,8 +1991,31 @@ 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
endpoints.control_plane_url,
resolved_api_key,
tenant_id=tenant_id,
**({"access_token": access_token} if access_token else {}),
)
+15 -7
View File
@@ -2,7 +2,7 @@
from __future__ import annotations
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Any, Literal
from urllib.parse import urlparse
@@ -130,17 +130,18 @@ class HostBackendClient:
def __init__(
self,
base_url: str,
api_key: str,
api_key: str | None,
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")
headers: dict[str, str] = {
"X-Api-Key": api_key,
"Accept": "application/json",
}
self._access_token = access_token
headers: dict[str, str] = {"Accept": "application/json"}
if api_key:
headers["X-Api-Key"] = api_key
if tenant_id:
headers["X-Tenant-ID"] = tenant_id
self._endpoints = ControlPlaneEndpoints.from_control_plane_url(base_url)
@@ -170,8 +171,15 @@ 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)
resp = self._client.request(
method, path, json=payload, params=params, headers=headers
)
resp.raise_for_status()
except httpx.HTTPStatusError as err:
detail = _error_detail(err.response)
+225
View File
@@ -0,0 +1,225 @@
"""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,12 +578,13 @@ 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="API key"):
with pytest.raises(click.ClickException, match="credentials"):
_create_host_backend_client(
host_url="https://api.example.com",
api_key=None,
@@ -0,0 +1,191 @@
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={}
)
+8 -11
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, NS_SEP
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND
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
from langgraph.pregel._messages import StreamMessagesHandler, graph_namespace
config = get_config()
message = next(x for x in convert_to_messages([message]))
@@ -359,15 +359,12 @@ def push_message(
):
handlers = config["callbacks"]
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)
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 state_key:
config[CONF][CONFIG_KEY_SEND]([(state_key, message)])
+20 -7
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
from langgraph._internal._constants import NS_SEP, PUSH
from langgraph.channels.binop import _get_overwrite
from langgraph.channels.delta import DeltaChannel
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
@@ -55,6 +55,23 @@ 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.
@@ -148,9 +165,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
**kwargs: Any,
) -> Any:
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
:-1
]
ns = graph_namespace(metadata)
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:
@@ -213,9 +228,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
and kwargs.get("name") == metadata.get("langgraph_node")
and (not tags or TAG_HIDDEN not in tags)
):
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
:-1
]
ns = graph_namespace(metadata)
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
return
self.metadata[run_id] = (ns, metadata)
+7 -12
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,12 +92,9 @@ 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`'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:
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:
- `metadata` is missing — handler is attached to a context
without Pregel routing info.
@@ -110,11 +107,9 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
return None
if tags and TAG_NOSTREAM in tags:
return None
nskey = metadata.get("langgraph_checkpoint_ns")
if not nskey:
ns: tuple[str, ...] = ()
else:
ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
ns = (
graph_namespace(metadata) if metadata.get("langgraph_checkpoint_ns") else ()
)
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
return None
return ns
+44 -33
View File
@@ -2880,6 +2880,18 @@ 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(
@@ -2898,9 +2910,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)
)
@@ -2909,24 +2921,12 @@ class Pregel(
if use_stream_messages_v2
else StreamMessagesHandler
)
run_manager.inheritable_handlers.append(
messages_handler_cls(
stream.put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
)
)
stream_handlers.append(messages_handler_cls(stream.put, subgraphs))
# set up tools stream mode
if "tools" in stream_modes:
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,
)
)
stream_handlers.append(StreamToolCallHandler(stream.put, subgraphs))
run_manager.inheritable_handlers.extend(stream_handlers)
# set up custom stream mode
if "custom" in stream_modes:
@@ -3009,6 +3009,10 @@ 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(
@@ -3291,6 +3295,18 @@ 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(
@@ -3325,10 +3341,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:
# 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)
)
@@ -3337,24 +3352,12 @@ class Pregel(
if use_stream_messages_v2
else StreamMessagesHandler
)
run_manager.inheritable_handlers.append(
messages_handler_cls(
stream_put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
)
)
stream_handlers.append(messages_handler_cls(stream_put, subgraphs))
# set up tools stream mode
if "tools" in stream_modes:
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,
)
)
stream_handlers.append(StreamToolCallHandler(stream_put, subgraphs))
run_manager.inheritable_handlers.extend(stream_handlers)
# set up custom stream mode
def stream_writer(c: Any) -> None:
@@ -3463,6 +3466,10 @@ 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(
@@ -4460,6 +4467,10 @@ 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, ...]:
@@ -0,0 +1,308 @@
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"]
+34 -8
View File
@@ -151,7 +151,8 @@ class HttpClient:
on_response: Callable[[httpx.Response], None] | None = None,
reconnect_limit: int = 5,
) -> Any:
"""Send a request that automatically reconnects to Location header."""
"""Send a request that reconnects to its Location header if the response is
cut off, with backoff between attempts."""
request_headers, content = await _aencode_json(json)
if headers:
request_headers.update(headers)
@@ -181,13 +182,27 @@ class HttpClient:
stacklevel=2,
)
await r.aclose()
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,
)
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
async def stream(
self,
@@ -230,6 +245,17 @@ 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
+34 -8
View File
@@ -152,7 +152,8 @@ class SyncHttpClient:
on_response: Callable[[httpx.Response], None] | None = None,
reconnect_limit: int = 5,
) -> Any:
"""Send a request that automatically reconnects to Location header."""
"""Send a request that reconnects to its Location header if the response is
cut off, with backoff between attempts."""
request_headers, content = _encode_json(json)
if headers:
request_headers.update(headers)
@@ -182,13 +183,27 @@ class SyncHttpClient:
stacklevel=2,
)
r.close()
return self.request_reconnect(
loc,
"GET",
headers=headers,
# don't pass on_response so it's only called once
reconnect_limit=reconnect_limit - 1,
)
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
def stream(
self,
@@ -233,6 +248,17 @@ 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
+17
View File
@@ -0,0 +1,17 @@
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
+80 -21
View File
@@ -1,7 +1,5 @@
from __future__ import annotations
import asyncio
import time
from collections.abc import Callable, Iterator, Sequence
from pathlib import Path
from typing import Any
@@ -70,19 +68,6 @@ class ListByteStream(httpx.ByteStream):
return None
@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
def iter_lines_raw(payload: list[bytes]) -> Iterator[BytesLike]:
decoder = BytesLineDecoder()
for part in payload:
@@ -501,9 +486,10 @@ _CALLER_CURSORS = pytest.mark.parametrize(
def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
empty_id: bytes, headers: dict[str, str] | None
):
reconnect = _sync_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
with pytest.warns(UserWarning, match="Last-Event-ID"):
reconnect = _sync_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
assert "last-event-id" not in reconnect.headers
@@ -514,9 +500,10 @@ def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id(
empty_id: bytes, headers: dict[str, str] | None
):
reconnect = await _async_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
with pytest.warns(UserWarning, match="Last-Event-ID"):
reconnect = await _async_reconnect(
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
)
assert "last-event-id" not in reconnect.headers
@@ -555,6 +542,45 @@ 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"}
@@ -772,6 +798,39 @@ async def test_async_long_stream_survives_many_spaced_out_drops():
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.
+98 -3
View File
@@ -1,7 +1,8 @@
from __future__ import annotations
import json
from collections.abc import AsyncIterator, Iterator
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from typing import Any
import httpx
import pytest
@@ -24,7 +25,16 @@ class _BodyDropsMidRead(httpx.SyncByteStream, httpx.AsyncByteStream):
raise httpx.ReadError("connection dropped")
def _wait_handler(requests: list[httpx.Request]):
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 handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
if request.method == "POST":
@@ -33,7 +43,7 @@ def _wait_handler(requests: list[httpx.Request]):
headers={"Location": "/threads/t1/runs/r1/join"},
stream=_BodyDropsMidRead(),
)
return httpx.Response(200, json={"done": True})
return reconnects[len(requests) - 2](request)
return handler
@@ -46,6 +56,7 @@ 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))
@@ -58,6 +69,7 @@ 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))
@@ -72,6 +84,89 @@ 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] = []