mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-31 03:11:16 +01:00
Compare commits
35
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d9abeabfcd | ||
|
|
5bd231cbfd | ||
|
|
3eaf06c67c | ||
|
|
6fab323753 | ||
|
|
8d5637b2dd | ||
|
|
1f0cf816d5 | ||
|
|
6f2df053f4 | ||
|
|
ffc3336367 | ||
|
|
5f4b28a9f7 | ||
|
|
5c96bce8d3 | ||
|
|
2f46bb3ca4 | ||
|
|
4b7f2f5230 | ||
|
|
5a3e5dd739 | ||
|
|
eb8ec111b3 | ||
|
|
8c177f7388 | ||
|
|
55013146ae | ||
|
|
858891aaa1 | ||
|
|
1526b0e82d | ||
|
|
c2c73b4f91 | ||
|
|
03e73ff51b | ||
|
|
3521f08e21 | ||
|
|
24294fad37 | ||
|
|
f6317d5cb1 | ||
|
|
7e56e6c80b | ||
|
|
0fc4f35b66 | ||
|
|
58be6911de | ||
|
|
1adeafeaf3 | ||
|
|
dc1741ce96 | ||
|
|
244ca657e4 | ||
|
|
c696d4e79f | ||
|
|
34c7f58655 | ||
|
|
1b1aee689d | ||
|
|
27aeb75efd | ||
|
|
c18a42fcba | ||
|
|
b54566b620 |
No files matched your search
@@ -24,6 +24,39 @@ TEST=path/to/test.py make test
|
||||
|
||||
Other pytest arguments can also be supplied inside the `TEST` variable.
|
||||
|
||||
## Direction: LangGraph is a runtime
|
||||
|
||||
LangGraph is consolidating as the low-level runtime for agents and workflows:
|
||||
graph execution, checkpointing, streaming, and interrupt/resume. Critical
|
||||
production systems and agent harnesses run on it, so correctness and stability
|
||||
come first. It is moving away from being a friendly SDK for building agents
|
||||
directly; agent development belongs one layer up, on a harness — Deep Agents
|
||||
(github.com/langchain-ai/deepagents, a separate repo) is the recommended one.
|
||||
|
||||
Optimize contributions for runtime correctness, reliability, performance, and
|
||||
simplicity — not for the shortest agent example.
|
||||
|
||||
- One way to do each thing. A change never adds a way; removals are always
|
||||
deliberate (see below). The number of ways only goes down. No aliases,
|
||||
overloads, extra accepted types, implicit defaults, or customization hooks
|
||||
for convenience: ergonomics alone never justify new public surface.
|
||||
- The mandate applies to everything that implements runtime behavior, in
|
||||
whichever repo it lives: the graph engine (`langgraph`, `langgraph-core`),
|
||||
the checkpointers (`checkpoint*`), and the server (`langgraph-api`). The
|
||||
client SDKs and CLI stay ergonomic but are thin, faithful clients, never a
|
||||
second implementation of runtime behavior.
|
||||
- Agent-level packages (`prebuilt`, `langgraph-supervisor`, `langgraph-swarm`,
|
||||
`langgraph-cua`) are maintenance-only: fix bugs, don't add agent features.
|
||||
New examples and docs that build agents should use Deep Agents.
|
||||
- Remove deliberately, not opportunistically. Removing or changing public
|
||||
behavior must be the stated purpose of the PR and go through the deprecation
|
||||
process — replacement in place, removal version recorded, migration note.
|
||||
Never drop a default or overload as a side effect of an unrelated change.
|
||||
- Behavior changes land in Python and JS together (the implementations live in
|
||||
separate repos; note in the PR when the other one needs a matching change).
|
||||
- If a task asks for something this section rules out, say so and propose the
|
||||
runtime-level or harness-level alternative instead of quietly implementing it.
|
||||
|
||||
## Libraries
|
||||
|
||||
The repository contains several Python and JavaScript/TypeScript libraries.
|
||||
|
||||
Generated
+1
-1
@@ -279,7 +279,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Generated
+1
-1
@@ -278,7 +278,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
description = "Library with a SQLite implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
@@ -12,7 +12,7 @@ readme = "README.md"
|
||||
license = "MIT"
|
||||
license-files = ['LICENSE']
|
||||
dependencies = [
|
||||
"langgraph-checkpoint>=4.3.0,<5.0.0",
|
||||
"langgraph-checkpoint>=4.3.1,<5.0.0",
|
||||
"aiosqlite>=0.20",
|
||||
"sqlite-vec>=0.1.6",
|
||||
]
|
||||
|
||||
Generated
+2
-2
@@ -285,7 +285,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -360,7 +360,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
Generated
+1
-1
@@ -301,7 +301,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -2148,6 +2148,6 @@ yocto-queue@^0.1.0:
|
||||
integrity sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==
|
||||
|
||||
"zod@^3.25.76 || ^4":
|
||||
version "4.3.6"
|
||||
resolved "https://registry.yarnpkg.com/zod/-/zod-4.3.6.tgz#89c56e0aa7d2b05107d894412227087885ab112a"
|
||||
integrity sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==
|
||||
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==
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.4.38"
|
||||
__version__ = "0.4.40"
|
||||
+136
-128
@@ -12,7 +12,7 @@ from collections.abc import Callable, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Protocol, TypeVar
|
||||
from typing import Generic, Protocol, TypeVar
|
||||
|
||||
import click
|
||||
import click.exceptions
|
||||
@@ -26,7 +26,6 @@ from langgraph_cli.dependency_tracking import find_tracked_packages
|
||||
from langgraph_cli.docker import build_docker_image, can_build_locally
|
||||
from langgraph_cli.exec import CommandRunner, Runner, subp_exec
|
||||
from langgraph_cli.host_backend import (
|
||||
MAX_PAGE_SIZE,
|
||||
ControlPlaneEndpoints,
|
||||
HostBackendClient,
|
||||
HostBackendError,
|
||||
@@ -183,12 +182,15 @@ class Listener:
|
||||
)
|
||||
return cls(
|
||||
identifier,
|
||||
str(resource.get("compute_id", "")),
|
||||
str(resource.get("compute_id") or ""),
|
||||
tuple(str(namespace) for namespace in namespaces)
|
||||
if isinstance(namespaces, list)
|
||||
else (),
|
||||
)
|
||||
|
||||
def cells(self) -> tuple[str | None, ...]:
|
||||
return (self.id, self.compute_id, ", ".join(self.namespaces))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Unplaced:
|
||||
@@ -276,8 +278,6 @@ def _describe_listeners(listeners: Sequence[Listener]) -> str:
|
||||
]
|
||||
if len(listeners) > len(shown):
|
||||
lines.append(f" ... and {len(listeners) - len(shown)} more")
|
||||
if len(listeners) == MAX_PAGE_SIZE:
|
||||
lines.append(f" (only the first {MAX_PAGE_SIZE} listeners were read)")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@@ -499,16 +499,9 @@ def _source_of(resource: object) -> str | None:
|
||||
def find_deployment_by_name(
|
||||
client: HostBackendClient, name: str
|
||||
) -> ExistingDeployment | None:
|
||||
listed = client.list_deployments(name=name, name_contains=name, limit=MAX_PAGE_SIZE)
|
||||
for resource in listed:
|
||||
for resource in client.list_deployments(name=name, name_contains=name):
|
||||
if resource.get("name") == name and resource.get("id"):
|
||||
return ExistingDeployment(str(resource["id"]), _source_of(resource))
|
||||
if len(listed) >= MAX_PAGE_SIZE:
|
||||
raise click.ClickException(
|
||||
"This workspace has more deployments than the CLI can search, so it "
|
||||
f"cannot tell whether '{name}' already exists. Pass --deployment-id to "
|
||||
"update an existing deployment."
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
@@ -583,103 +576,96 @@ def _validate_prebuilt_image(
|
||||
_get_emitter().info(f"Image is available for {_DEPLOYMENT_PLATFORM}")
|
||||
|
||||
|
||||
def _extract_deployment_url(deployment: dict[str, object]) -> str:
|
||||
source_config = deployment.get("source_config")
|
||||
if isinstance(source_config, dict):
|
||||
custom_url = source_config.get("custom_url")
|
||||
if isinstance(custom_url, str) and custom_url:
|
||||
return custom_url
|
||||
return "-"
|
||||
Resource = Mapping[str, object]
|
||||
RowT = TypeVar("RowT")
|
||||
|
||||
|
||||
def format_deployments_table(deployments: Sequence[dict[str, object]]) -> str:
|
||||
headers = ("Deployment ID", "Deployment Name", "Deployment URL")
|
||||
rows = [
|
||||
(
|
||||
str(deployment.get("id", "-") or "-"),
|
||||
str(deployment.get("name", "-") or "-"),
|
||||
_extract_deployment_url(deployment),
|
||||
)
|
||||
for deployment in deployments
|
||||
]
|
||||
widths = [
|
||||
max(len(headers[index]), *(len(row[index]) for row in rows))
|
||||
for index in range(len(headers))
|
||||
]
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Table(Generic[RowT]):
|
||||
headers: tuple[str, ...]
|
||||
cells: Callable[[RowT], tuple[str | None, ...]]
|
||||
|
||||
def format_row(row: Sequence[str]) -> str:
|
||||
return " ".join(value.ljust(widths[index]) for index, value in enumerate(row))
|
||||
|
||||
lines = [format_row(headers), format_row(tuple("-" * width for width in widths))]
|
||||
lines.extend(format_row(row) for row in rows)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _extract_listener_namespaces(listener: dict[str, object]) -> str:
|
||||
compute_config = listener.get("compute_config")
|
||||
namespaces = (
|
||||
compute_config.get("k8s_namespaces")
|
||||
if isinstance(compute_config, dict)
|
||||
else None
|
||||
)
|
||||
if isinstance(namespaces, list) and namespaces:
|
||||
return ", ".join(str(namespace) for namespace in namespaces)
|
||||
return "-"
|
||||
|
||||
|
||||
def format_listeners_table(listeners: Sequence[dict[str, object]]) -> str:
|
||||
headers = ("Listener ID", "Compute ID", "Namespaces")
|
||||
rows = [
|
||||
(
|
||||
str(listener.get("id", "-") or "-"),
|
||||
str(listener.get("compute_id", "-") or "-"),
|
||||
_extract_listener_namespaces(listener),
|
||||
)
|
||||
for listener in listeners
|
||||
]
|
||||
widths = [
|
||||
max(len(headers[index]), *(len(row[index]) for row in rows))
|
||||
for index in range(len(headers))
|
||||
]
|
||||
|
||||
def format_row(row: Sequence[str]) -> str:
|
||||
return " ".join(value.ljust(widths[index]) for index, value in enumerate(row))
|
||||
|
||||
lines = [format_row(headers), format_row(tuple("-" * width for width in widths))]
|
||||
lines.extend(format_row(row) for row in rows)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def format_revisions_table(revisions: Sequence[dict[str, object]]) -> str:
|
||||
headers = ("Revision ID", "Status", "Created At")
|
||||
latest_deployed_seen = False
|
||||
rows = []
|
||||
for revision in revisions:
|
||||
status = str(revision.get("status", "-") or "-")
|
||||
if status == "DEPLOYED":
|
||||
if latest_deployed_seen:
|
||||
status = "REPLACED"
|
||||
else:
|
||||
latest_deployed_seen = True
|
||||
rows.append(
|
||||
(
|
||||
str(revision.get("id", "-") or "-"),
|
||||
status,
|
||||
str(revision.get("created_at", "-") or "-"),
|
||||
def render(self, rows: Sequence[RowT]) -> str:
|
||||
grid = [
|
||||
self.headers,
|
||||
*(tuple(cell or "-" for cell in self.cells(row)) for row in rows),
|
||||
]
|
||||
widths = [max(map(len, column)) for column in zip(*grid, strict=True)]
|
||||
rule = tuple("-" * width for width in widths)
|
||||
return "\n".join(
|
||||
" ".join(
|
||||
cell.ljust(width) for cell, width in zip(line, widths, strict=True)
|
||||
)
|
||||
for line in (grid[0], rule, *grid[1:])
|
||||
)
|
||||
|
||||
widths = [
|
||||
max(len(headers[index]), *(len(row[index]) for row in rows))
|
||||
for index in range(len(headers))
|
||||
|
||||
def _text(resource: Resource, key: str) -> str | None:
|
||||
value = resource.get(key)
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
def _deployment_url(deployment: Resource) -> str | None:
|
||||
source_config = deployment.get("source_config")
|
||||
custom_url = (
|
||||
source_config.get("custom_url") if isinstance(source_config, Mapping) else None
|
||||
)
|
||||
return custom_url if isinstance(custom_url, str) and custom_url else None
|
||||
|
||||
|
||||
def _with_replaced_status(revisions: Sequence[Resource]) -> list[Resource]:
|
||||
live_index = next(
|
||||
(
|
||||
index
|
||||
for index, revision in enumerate(revisions)
|
||||
if revision.get("status") == "DEPLOYED"
|
||||
),
|
||||
None,
|
||||
)
|
||||
return [
|
||||
{**revision, "status": "REPLACED"}
|
||||
if revision.get("status") == "DEPLOYED" and index != live_index
|
||||
else revision
|
||||
for index, revision in enumerate(revisions)
|
||||
]
|
||||
|
||||
def format_row(row: Sequence[str]) -> str:
|
||||
return " ".join(value.ljust(widths[index]) for index, value in enumerate(row))
|
||||
|
||||
lines = [format_row(headers), format_row(tuple("-" * width for width in widths))]
|
||||
lines.extend(format_row(row) for row in rows)
|
||||
return "\n".join(lines)
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DeploymentRow:
|
||||
id: str
|
||||
name: str | None
|
||||
url: str | None
|
||||
|
||||
@classmethod
|
||||
def from_resource(cls, deployment: Resource) -> "DeploymentRow":
|
||||
identifier = _text(deployment, "id")
|
||||
if identifier is None:
|
||||
raise HostBackendError(
|
||||
"The control plane returned a deployment without an id."
|
||||
)
|
||||
return cls(identifier, _text(deployment, "name"), _deployment_url(deployment))
|
||||
|
||||
def cells(self) -> tuple[str | None, ...]:
|
||||
return (self.id, self.name, self.url)
|
||||
|
||||
def as_json(self) -> dict[str, str | None]:
|
||||
return {"id": self.id, "name": self.name}
|
||||
|
||||
|
||||
DEPLOYMENTS_TABLE = Table(
|
||||
("Deployment ID", "Deployment Name", "Deployment URL"), DeploymentRow.cells
|
||||
)
|
||||
|
||||
LISTENERS_TABLE = Table(("Listener ID", "Compute ID", "Namespaces"), Listener.cells)
|
||||
|
||||
REVISIONS_TABLE = Table(
|
||||
("Revision ID", "Status", "Created At"),
|
||||
lambda revision: (
|
||||
_text(revision, "id"),
|
||||
_text(revision, "status"),
|
||||
_text(revision, "created_at"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def format_timestamp(ts) -> str:
|
||||
@@ -841,7 +827,6 @@ def _find_deployment(
|
||||
lambda c: c.list_deployments(
|
||||
agent_id=selector.agent_id,
|
||||
agent_environment=selector.environment,
|
||||
limit=MAX_PAGE_SIZE,
|
||||
),
|
||||
)
|
||||
if len(existing) > 1:
|
||||
@@ -1007,20 +992,15 @@ def _print_deployment_result(
|
||||
dashboard_label: str,
|
||||
status_url: str | None = None,
|
||||
) -> None:
|
||||
"""Print final deployment status and raise on failure."""
|
||||
em = _get_emitter()
|
||||
dep_info = client.get_deployment(deployment_id)
|
||||
custom_url = None
|
||||
if isinstance(dep_info, dict):
|
||||
sc = dep_info.get("source_config")
|
||||
if isinstance(sc, dict):
|
||||
custom_url = sc.get("custom_url")
|
||||
deployment = client.get_deployment(deployment_id)
|
||||
url = _deployment_url(deployment) if isinstance(deployment, Mapping) else None
|
||||
|
||||
if last_status == "DEPLOYED":
|
||||
em.result(
|
||||
"succeeded",
|
||||
deployment_id=deployment_id,
|
||||
url=custom_url,
|
||||
url=url,
|
||||
status_url=status_url,
|
||||
)
|
||||
elif last_status in ("BUILD_FAILED", "DEPLOY_FAILED", "CREATE_FAILED"):
|
||||
@@ -1081,6 +1061,8 @@ def _docker_config_for_token(registry_host: str, token: str, username: str):
|
||||
|
||||
_UPLOAD_TIMEOUT_SECONDS = 300
|
||||
_BYTES_PER_MIB = 1_048_576
|
||||
# Byte range Cloud Storage signs into upload URLs, up to the 200 MB archive limit.
|
||||
_GCS_SIZE_RANGE = f"0,{200 * _BYTES_PER_MIB}"
|
||||
|
||||
|
||||
class _ProgressReader:
|
||||
@@ -1106,12 +1088,24 @@ class _ProgressReader:
|
||||
return self._file_size
|
||||
|
||||
|
||||
def _upload_to_gcs(signed_url: str, file_path: str, file_size: int) -> None:
|
||||
"""Upload tarball to GCS via signed PUT URL with progress display."""
|
||||
def _upload_source_tarball(
|
||||
signed_url: str,
|
||||
file_path: str,
|
||||
file_size: int,
|
||||
upload_headers: dict[str, str] | None,
|
||||
) -> None:
|
||||
"""Upload tarball via signed PUT URL with progress display."""
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
em = _get_emitter()
|
||||
# Control planes that return no `upload_headers` sign for the Cloud Storage
|
||||
# size range. Content-Type is a fallback: without one, urllib sends a form
|
||||
# content type, which breaks signatures that cover it.
|
||||
headers = {
|
||||
"Content-Type": "application/gzip",
|
||||
**(upload_headers or {"X-Goog-Content-Length-Range": _GCS_SIZE_RANGE}),
|
||||
}
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
reader = _ProgressReader(f, file_size, em)
|
||||
@@ -1119,11 +1113,7 @@ def _upload_to_gcs(signed_url: str, file_path: str, file_size: int) -> None:
|
||||
signed_url,
|
||||
data=reader,
|
||||
method="PUT",
|
||||
headers={
|
||||
"Content-Type": "application/gzip",
|
||||
"Content-Length": str(file_size),
|
||||
"X-Goog-Content-Length-Range": "0,209715200",
|
||||
},
|
||||
headers={**headers, "Content-Length": str(file_size)},
|
||||
)
|
||||
try:
|
||||
urllib.request.urlopen(req, timeout=_UPLOAD_TIMEOUT_SECONDS)
|
||||
@@ -1443,7 +1433,9 @@ def _run_remote_build(
|
||||
step += 1
|
||||
|
||||
_log_deploy_step(step, "Uploading source")
|
||||
_upload_to_gcs(signed_url, archive_path, file_size)
|
||||
_upload_source_tarball(
|
||||
signed_url, archive_path, file_size, upload_data.get("upload_headers")
|
||||
)
|
||||
step += 1
|
||||
|
||||
_log_deploy_step(step, "Triggering remote build")
|
||||
@@ -2594,6 +2586,13 @@ def _deploy_cmd(
|
||||
default="",
|
||||
help="Only show deployments whose names contain this value.",
|
||||
)
|
||||
@click.option(
|
||||
"--json",
|
||||
"json_output",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Print the deployments as a JSON array of {id, name} instead of a table.",
|
||||
)
|
||||
@deploy.command("list", help="[Beta] List LangSmith Deployments.")
|
||||
def deploy_list(
|
||||
api_key: str | None,
|
||||
@@ -2601,11 +2600,13 @@ def deploy_list(
|
||||
name_contains: str,
|
||||
agent_id: str | None,
|
||||
environment: str | None,
|
||||
json_output: bool,
|
||||
) -> None:
|
||||
if agent_id is not None or environment is not None:
|
||||
click.secho(
|
||||
"Note: --agent-id and --agent-environment flags are in private beta",
|
||||
fg="yellow",
|
||||
err=True,
|
||||
)
|
||||
if agent_id is not None and not agent_id.strip():
|
||||
raise click.UsageError("--agent-id must not be empty.")
|
||||
@@ -2614,15 +2615,19 @@ def deploy_list(
|
||||
filters["agent_id"] = agent_id
|
||||
if environment is not None:
|
||||
filters["agent_environment"] = environment
|
||||
client = _create_host_backend_client(host_url, api_key)
|
||||
deployments = _call_host_backend_with_optional_tenant(
|
||||
client = _create_host_backend_client(host_url, api_key, no_input=json_output)
|
||||
resources = _call_host_backend_with_optional_tenant(
|
||||
client,
|
||||
lambda c: c.list_deployments(name_contains=name_contains, **filters),
|
||||
no_input=json_output,
|
||||
)
|
||||
if not deployments:
|
||||
deployments = [DeploymentRow.from_resource(resource) for resource in resources]
|
||||
if json_output:
|
||||
click.echo(json_mod.dumps([deployment.as_json() for deployment in deployments]))
|
||||
elif not deployments:
|
||||
click.echo("No deployments found.")
|
||||
return
|
||||
click.echo(format_deployments_table(deployments))
|
||||
else:
|
||||
click.echo(DEPLOYMENTS_TABLE.render(deployments))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -2665,11 +2670,14 @@ def deploy_listeners_list(
|
||||
)
|
||||
if json_output:
|
||||
click.echo(json_mod.dumps(listeners))
|
||||
return
|
||||
if not listeners:
|
||||
elif not listeners:
|
||||
click.echo("No listeners found for this workspace.")
|
||||
return
|
||||
click.echo(format_listeners_table(listeners))
|
||||
else:
|
||||
click.echo(
|
||||
LISTENERS_TABLE.render(
|
||||
[Listener.from_resource(listener) for listener in listeners]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -2712,7 +2720,7 @@ def deploy_revisions_list(
|
||||
if not revisions:
|
||||
click.echo(f"No revisions found for deployment {deployment_id}.")
|
||||
return
|
||||
click.echo(format_revisions_table(revisions))
|
||||
click.echo(REVISIONS_TABLE.render(_with_replaced_status(revisions)))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import urlparse
|
||||
@@ -220,27 +221,51 @@ class HostBackendClient:
|
||||
*,
|
||||
name: str | None = None,
|
||||
name_contains: str | None = None,
|
||||
limit: int | None = None,
|
||||
agent_id: str | None = None,
|
||||
agent_environment: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> list[dict[str, object]]:
|
||||
given = (
|
||||
("name", name),
|
||||
("name_contains", name_contains),
|
||||
("limit", limit),
|
||||
("agent_id", agent_id),
|
||||
("agent_environment", agent_environment),
|
||||
)
|
||||
params = {key: value for key, value in given if value is not None}
|
||||
return _resources(self._request("GET", "/v2/deployments", params=params))
|
||||
filters = {key: value for key, value in given if value is not None}
|
||||
return self._every_page("/v2/deployments", filters)
|
||||
|
||||
def _every_page(
|
||||
self, path: str, filters: Mapping[str, object]
|
||||
) -> list[dict[str, object]]:
|
||||
resources: list[dict[str, object]] = []
|
||||
seen_ids: set[object] = set()
|
||||
while True:
|
||||
page = _resources(
|
||||
self._request(
|
||||
"GET",
|
||||
path,
|
||||
params={
|
||||
**filters,
|
||||
"limit": MAX_PAGE_SIZE,
|
||||
"offset": len(resources),
|
||||
},
|
||||
)
|
||||
)
|
||||
page_ids = {resource.get("id") for resource in page}
|
||||
if page_ids and page_ids <= seen_ids:
|
||||
raise HostBackendError(
|
||||
f"GET {path} returned the same resources again: this control "
|
||||
"plane ignores 'offset', so the CLI cannot read past the first page."
|
||||
)
|
||||
seen_ids |= page_ids
|
||||
resources.extend(page)
|
||||
if len(page) < MAX_PAGE_SIZE:
|
||||
return resources
|
||||
|
||||
def get_listener(self, listener_id: str) -> dict[str, Any]:
|
||||
return self._request("GET", f"/v2/listeners/{listener_id}")
|
||||
|
||||
def list_listeners(self) -> list[dict[str, Any]]:
|
||||
return _resources(
|
||||
self._request("GET", "/v2/listeners", params={"limit": MAX_PAGE_SIZE})
|
||||
)
|
||||
def list_listeners(self) -> list[dict[str, object]]:
|
||||
return self._every_page("/v2/listeners", {})
|
||||
|
||||
def get_deployment(self, deployment_id: str) -> dict[str, Any]:
|
||||
return self._request("GET", f"/v2/deployments/{deployment_id}")
|
||||
@@ -258,10 +283,13 @@ class HostBackendClient:
|
||||
)
|
||||
|
||||
def request_upload_url(self, deployment_id: str) -> dict[str, Any]:
|
||||
"""Get a signed GCS URL for uploading the source tarball."""
|
||||
"""Get a signed URL for uploading the source tarball."""
|
||||
# Tells the control plane this client sends the returned
|
||||
# `upload_headers` with its upload, which some storage requires.
|
||||
return self._request(
|
||||
"POST",
|
||||
f"/v2/deployments/{deployment_id}/upload-url",
|
||||
{"supports_upload_headers": True},
|
||||
)
|
||||
|
||||
def update_deployment(
|
||||
|
||||
@@ -482,6 +482,177 @@ def test_deploy_list_command_no_results(monkeypatch) -> None:
|
||||
assert result.output.strip() == "No deployments found."
|
||||
|
||||
|
||||
def test_deploy_list_prints_an_aligned_table(monkeypatch) -> None:
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_deployments(self, name_contains: str = ""):
|
||||
return [
|
||||
{
|
||||
"id": "dep-123",
|
||||
"name": "alpha",
|
||||
"source_config": {"custom_url": "https://alpha.example.com"},
|
||||
},
|
||||
{"id": "dep-4567", "name": "beta-agent", "source_config": {}},
|
||||
]
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"list",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert result.output.splitlines() == [
|
||||
"Deployment ID Deployment Name Deployment URL ",
|
||||
"------------- --------------- -------------------------",
|
||||
"dep-123 alpha https://alpha.example.com",
|
||||
"dep-4567 beta-agent - ",
|
||||
]
|
||||
|
||||
|
||||
def _deploy_list_json(
|
||||
monkeypatch, deployments: list[dict], *extra_args: str
|
||||
) -> click.testing.Result:
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_deployments(
|
||||
self,
|
||||
name_contains: str = "",
|
||||
agent_id: str | None = None,
|
||||
agent_environment: str | None = None,
|
||||
):
|
||||
return deployments
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
return CliRunner().invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"list",
|
||||
"--json",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
*extra_args,
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_deploy_list_command_json(monkeypatch) -> None:
|
||||
result = _deploy_list_json(
|
||||
monkeypatch,
|
||||
[
|
||||
{
|
||||
"id": "dep-123",
|
||||
"name": "alpha",
|
||||
"source_config": {"custom_url": "https://alpha.example.com"},
|
||||
},
|
||||
{"id": "dep-456", "name": "beta", "source_config": {}},
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == [
|
||||
{"id": "dep-123", "name": "alpha"},
|
||||
{"id": "dep-456", "name": "beta"},
|
||||
]
|
||||
|
||||
|
||||
def test_deploy_list_json_never_prints_deployment_secrets(monkeypatch) -> None:
|
||||
result = _deploy_list_json(
|
||||
monkeypatch,
|
||||
[
|
||||
{
|
||||
"id": "dep-123",
|
||||
"name": "alpha",
|
||||
"source_config": {"custom_url": None},
|
||||
"secrets": [{"name": "OPENAI_API_KEY", "value": "sk-plain-value"}],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "sk-plain-value" not in result.output
|
||||
assert json.loads(result.stdout) == [{"id": "dep-123", "name": "alpha"}]
|
||||
|
||||
|
||||
def test_deploy_list_command_json_no_results(monkeypatch) -> None:
|
||||
result = _deploy_list_json(monkeypatch, [])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == []
|
||||
|
||||
|
||||
def test_deploy_list_rejects_a_deployment_without_an_id(monkeypatch) -> None:
|
||||
result = _deploy_list_json(monkeypatch, [{"name": "alpha"}])
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "deployment without an id" in result.output
|
||||
|
||||
|
||||
def test_deploy_list_json_keeps_the_agent_note_off_stdout(monkeypatch) -> None:
|
||||
result = _deploy_list_json(
|
||||
monkeypatch,
|
||||
[{"id": "dep-123", "name": "alpha"}],
|
||||
"--agent-id",
|
||||
"customer-support",
|
||||
"--agent-environment",
|
||||
"staging",
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.stdout) == [{"id": "dep-123", "name": "alpha"}]
|
||||
assert "private beta" in result.stderr
|
||||
|
||||
|
||||
def test_deploy_list_json_never_prompts_for_a_workspace(monkeypatch) -> None:
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_deployments(self, name_contains: str = ""):
|
||||
raise HostBackendError("requires workspace specification", status_code=403)
|
||||
|
||||
def set_tenant(self, tenant_id: str) -> None:
|
||||
raise AssertionError("must not retry with a prompted tenant with --json")
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
prompt_calls = []
|
||||
monkeypatch.setattr(
|
||||
click, "prompt", lambda *a, **k: prompt_calls.append((a, k)) or "unreachable"
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"list",
|
||||
"--json",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert prompt_calls == []
|
||||
assert result.exit_code != 0
|
||||
assert "LANGSMITH_TENANT_ID" in result.output
|
||||
|
||||
|
||||
def test_deploy_listeners_list_command(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
captured: dict[str, str] = {}
|
||||
@@ -564,6 +735,84 @@ def test_deploy_listeners_list_command_no_results(monkeypatch) -> None:
|
||||
assert result.output.strip() == "No listeners found for this workspace."
|
||||
|
||||
|
||||
def test_deploy_listeners_list_prints_an_aligned_table(monkeypatch) -> None:
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_listeners(self):
|
||||
return [
|
||||
{
|
||||
"id": "listener-1",
|
||||
"compute_id": "prod-cluster",
|
||||
"compute_config": {"k8s_namespaces": ["agents"]},
|
||||
},
|
||||
{
|
||||
"id": "listener-22",
|
||||
"compute_id": "multi",
|
||||
"compute_config": {"k8s_namespaces": ["agents", "agents-staging"]},
|
||||
},
|
||||
{"id": "listener-3", "compute_id": "", "compute_config": {}},
|
||||
{
|
||||
"id": "listener-4",
|
||||
"compute_id": None,
|
||||
"compute_config": {"k8s_namespaces": ["agents"]},
|
||||
},
|
||||
]
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"listeners",
|
||||
"list",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert result.output.splitlines() == [
|
||||
"Listener ID Compute ID Namespaces ",
|
||||
"----------- ------------ ----------------------",
|
||||
"listener-1 prod-cluster agents ",
|
||||
"listener-22 multi agents, agents-staging",
|
||||
"listener-3 - - ",
|
||||
"listener-4 - agents ",
|
||||
]
|
||||
|
||||
|
||||
def test_deploy_listeners_list_rejects_a_listener_without_an_id(monkeypatch) -> None:
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_listeners(self):
|
||||
return [{"compute_id": "prod-cluster"}]
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"listeners",
|
||||
"list",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "listener without an id" in result.output
|
||||
|
||||
|
||||
def test_deploy_listeners_list_command_json(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
|
||||
@@ -766,6 +1015,54 @@ def test_deploy_revisions_list_command_no_results(monkeypatch) -> None:
|
||||
assert result.output.strip() == "No revisions found for deployment dep-123."
|
||||
|
||||
|
||||
def test_deploy_revisions_list_marks_older_deployed_revisions_as_replaced(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_revisions(self, deployment_id: str, limit: int = 1):
|
||||
return [
|
||||
{
|
||||
"id": "rev-3",
|
||||
"status": "DEPLOYING",
|
||||
"created_at": "2026-10-20T10:00:00Z",
|
||||
},
|
||||
{
|
||||
"id": "rev-2",
|
||||
"status": "DEPLOYED",
|
||||
"created_at": "2026-10-19T10:00:00Z",
|
||||
},
|
||||
{"id": "rev-1", "status": "DEPLOYED", "created_at": None},
|
||||
]
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"revisions",
|
||||
"list",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
"dep-123",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert result.output.splitlines() == [
|
||||
"Revision ID Status Created At ",
|
||||
"----------- --------- --------------------",
|
||||
"rev-3 DEPLOYING 2026-10-20T10:00:00Z",
|
||||
"rev-2 DEPLOYED 2026-10-19T10:00:00Z",
|
||||
"rev-1 REPLACED - ",
|
||||
]
|
||||
|
||||
|
||||
def test_deploy_revisions_list_command_with_explicit_limit(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
@@ -96,6 +96,7 @@ class ControlPlaneDouble:
|
||||
existing_deployments: list[dict] = field(default_factory=list)
|
||||
push_token_status: int = 200
|
||||
push_token_extra: dict = field(default_factory=dict)
|
||||
upload_url_extra: dict = field(default_factory=dict)
|
||||
create_error: str | None = None
|
||||
listeners: list[dict] = field(default_factory=list)
|
||||
listeners_by_id: dict[str, dict] = field(default_factory=dict)
|
||||
@@ -111,7 +112,11 @@ class ControlPlaneDouble:
|
||||
def _respond(self, request: httpx.Request) -> httpx.Response:
|
||||
method, path = request.method, request.url.path
|
||||
if (method, path) == ("GET", "/v2/listeners"):
|
||||
return httpx.Response(200, json={"resources": self.listeners})
|
||||
offset = int(request.url.params.get("offset", 0))
|
||||
limit = int(request.url.params.get("limit", 20))
|
||||
return httpx.Response(
|
||||
200, json={"resources": self.listeners[offset : offset + limit]}
|
||||
)
|
||||
if method == "GET" and path.startswith("/v2/listeners/"):
|
||||
listener_id = path.rsplit("/", 1)[-1]
|
||||
if not _looks_like_a_uuid(listener_id):
|
||||
@@ -159,7 +164,12 @@ class ControlPlaneDouble:
|
||||
)
|
||||
if path.endswith("/upload-url"):
|
||||
return httpx.Response(
|
||||
200, json={"upload_url": SIGNED_UPLOAD_URL, "object_path": OBJECT_PATH}
|
||||
200,
|
||||
json={
|
||||
"upload_url": SIGNED_UPLOAD_URL,
|
||||
"object_path": OBJECT_PATH,
|
||||
**self.upload_url_extra,
|
||||
},
|
||||
)
|
||||
if method == "PATCH":
|
||||
return httpx.Response(200, json={"tenant_id": "tenant-1"})
|
||||
@@ -276,7 +286,7 @@ class DeployProject:
|
||||
control_plane: ControlPlaneDouble
|
||||
docker: DockerDouble
|
||||
timeline: list[str]
|
||||
uploads: list[tuple[str, str, int]]
|
||||
uploads: list[tuple[str, str, int, dict[str, str] | None]]
|
||||
|
||||
def run(self, *args: str, host_url: str = CONTROL_PLANE_URL) -> Result:
|
||||
return CliRunner().invoke(
|
||||
@@ -313,16 +323,21 @@ def deploy_project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> DeployPro
|
||||
timeline: list[str] = []
|
||||
control_plane = ControlPlaneDouble(timeline)
|
||||
docker = DockerDouble(timeline)
|
||||
uploads: list[tuple[str, str, int]] = []
|
||||
uploads: list[tuple[str, str, int, dict[str, str] | None]] = []
|
||||
|
||||
@contextmanager
|
||||
def fake_create_archive(config_path: Path, config: dict) -> Iterator[tuple]:
|
||||
timeline.append("create_archive")
|
||||
yield ARCHIVE
|
||||
|
||||
def fake_upload(signed_url: str, file_path: str, file_size: int) -> None:
|
||||
def fake_upload(
|
||||
signed_url: str,
|
||||
file_path: str,
|
||||
file_size: int,
|
||||
upload_headers: dict[str, str] | None,
|
||||
) -> None:
|
||||
timeline.append("upload_archive")
|
||||
uploads.append((signed_url, file_path, file_size))
|
||||
uploads.append((signed_url, file_path, file_size, upload_headers))
|
||||
|
||||
monkeypatch.setattr(deploy_module, "_no_input", False)
|
||||
monkeypatch.setattr(deploy_module, "_emitter", None)
|
||||
@@ -340,7 +355,7 @@ def deploy_project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> DeployPro
|
||||
)
|
||||
monkeypatch.setattr(deploy_module.platform, "machine", lambda: "x86_64")
|
||||
monkeypatch.setattr(archive_module, "create_archive", fake_create_archive)
|
||||
monkeypatch.setattr(deploy_module, "_upload_to_gcs", fake_upload)
|
||||
monkeypatch.setattr(deploy_module, "_upload_source_tarball", fake_upload)
|
||||
return DeployProject(control_plane, docker, timeline, uploads)
|
||||
|
||||
|
||||
@@ -699,6 +714,9 @@ def test_deployment_not_created_by_the_cli_gets_an_actionable_error(
|
||||
def test_remote_build_creates_an_internal_source_deployment_and_uploads_the_archive(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
upload_headers = {"Content-Type": "application/gzip", "x-ms-blob-type": "BlockBlob"}
|
||||
deploy_project.control_plane.upload_url_extra = {"upload_headers": upload_headers}
|
||||
|
||||
result = deploy_project.run("--remote", "--install-command", "yarn install")
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
@@ -713,7 +731,12 @@ def test_remote_build_creates_an_internal_source_deployment_and_uploads_the_arch
|
||||
assert deploy_project.control_plane.bodies[CREATE_DEPLOYMENT]["source"] == (
|
||||
"internal_source"
|
||||
)
|
||||
assert deploy_project.uploads == [(SIGNED_UPLOAD_URL, ARCHIVE[0], ARCHIVE[1])]
|
||||
assert deploy_project.control_plane.bodies[_upload_url(CREATED_ID)] == {
|
||||
"supports_upload_headers": True
|
||||
}
|
||||
assert deploy_project.uploads == [
|
||||
(SIGNED_UPLOAD_URL, ARCHIVE[0], ARCHIVE[1], upload_headers)
|
||||
]
|
||||
assert deploy_project.control_plane.bodies[_patch(CREATED_ID)] == {
|
||||
"revision_source": "internal_source",
|
||||
"source_revision_config": {
|
||||
@@ -1288,14 +1311,16 @@ def test_listener_flags_without_push_to_make_no_call_at_all(
|
||||
assert deploy_project.timeline == []
|
||||
|
||||
|
||||
def test_a_truncated_listener_page_says_so(deploy_project: DeployProject) -> None:
|
||||
def test_listeners_beyond_the_first_page_are_counted(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
deploy_project.control_plane.listeners = [
|
||||
{
|
||||
"id": str(uuid.UUID(int=index)),
|
||||
"compute_id": "cluster",
|
||||
"compute_config": {"k8s_namespaces": ["agents"]},
|
||||
}
|
||||
for index in range(100)
|
||||
for index in range(101)
|
||||
]
|
||||
|
||||
result = deploy_project.run(
|
||||
@@ -1303,7 +1328,8 @@ def test_a_truncated_listener_page_says_so(deploy_project: DeployProject) -> Non
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "first 100" in result.output
|
||||
assert "... and 91 more" in result.output
|
||||
assert "first 100" not in result.output
|
||||
|
||||
|
||||
def test_a_managed_build_in_a_listener_workspace_points_at_push_to(
|
||||
|
||||
@@ -75,6 +75,7 @@ def test_agent_create(deployment_api, tmp_path, monkeypatch):
|
||||
"agent_id": "customer-support",
|
||||
"agent_environment": "staging",
|
||||
"limit": "100",
|
||||
"offset": "0",
|
||||
}
|
||||
payload = json.loads(requests[1].content)
|
||||
assert payload["agent"] == {
|
||||
|
||||
@@ -33,6 +33,7 @@ from langgraph_cli.deploy import (
|
||||
_Emitter,
|
||||
_env_without_deployment_name,
|
||||
_parse_env_from_config,
|
||||
_print_deployment_result,
|
||||
_resolve_env_path,
|
||||
_resolve_pushed_image_digest,
|
||||
_resource_spec,
|
||||
@@ -536,6 +537,39 @@ class TestEmitterTextMode:
|
||||
assert "URL: https://app.test" in lines
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployment", "expected_url_lines"),
|
||||
[
|
||||
pytest.param(
|
||||
{
|
||||
"id": "dep-1",
|
||||
"source_config": {"custom_url": "https://agent.example.com"},
|
||||
},
|
||||
["URL: https://agent.example.com"],
|
||||
id="custom_url_is_printed",
|
||||
),
|
||||
pytest.param(
|
||||
{"id": "dep-1", "source_config": {}}, [], id="no_custom_url_no_line"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_deployed_result_prints_the_deployment_url(
|
||||
capsys, monkeypatch, deployment, expected_url_lines
|
||||
):
|
||||
class FakeClient:
|
||||
def get_deployment(self, deployment_id: str) -> dict:
|
||||
return deployment
|
||||
|
||||
monkeypatch.setattr(deploy_mod, "_emitter", None)
|
||||
|
||||
_print_deployment_result(
|
||||
FakeClient(), "dep-1", "DEPLOYED", dashboard_label="Deployment dashboard"
|
||||
)
|
||||
|
||||
lines = [line.strip() for line in capsys.readouterr().out.splitlines()]
|
||||
assert [line for line in lines if line.startswith("URL:")] == expected_url_lines
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# --no-input guard on _create_host_backend_client
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1210,6 +1244,7 @@ def test_finding_a_deployment_by_name_narrows_the_search_for_every_server_versio
|
||||
"name": "agent",
|
||||
"name_contains": "agent",
|
||||
"limit": "100",
|
||||
"offset": "0",
|
||||
}
|
||||
assert found == ExistingDeployment("dep-1", "github")
|
||||
|
||||
@@ -1237,20 +1272,29 @@ def test_a_server_that_ignores_the_exact_name_filter_never_matches_another_deplo
|
||||
assert find_deployment_by_name(client, "brand-new-agent") is None
|
||||
|
||||
|
||||
def test_a_full_page_without_a_match_refuses_to_claim_the_name_is_free():
|
||||
page = [
|
||||
{"id": f"dep-{index}", "name": f"other-agent-{index}"} for index in range(100)
|
||||
]
|
||||
client = HostBackendClient(
|
||||
"https://api.example.com",
|
||||
"key",
|
||||
transport=httpx.MockTransport(
|
||||
lambda req: httpx.Response(200, json={"resources": page})
|
||||
def test_a_deployment_beyond_the_first_page_is_found_by_name():
|
||||
deployments = [
|
||||
*(
|
||||
{"id": f"dep-{index}", "name": f"other-agent-{index}"}
|
||||
for index in range(100)
|
||||
),
|
||||
{"id": "dep-target", "name": "agent", "source": "external_docker"},
|
||||
]
|
||||
|
||||
def handler(req: httpx.Request) -> httpx.Response:
|
||||
offset = int(req.url.params["offset"])
|
||||
limit = int(req.url.params["limit"])
|
||||
return httpx.Response(
|
||||
200, json={"resources": deployments[offset : offset + limit]}
|
||||
)
|
||||
|
||||
client = HostBackendClient(
|
||||
"https://api.example.com", "key", transport=httpx.MockTransport(handler)
|
||||
)
|
||||
|
||||
with pytest.raises(click.ClickException, match="--deployment-id"):
|
||||
find_deployment_by_name(client, "brand-new-agent")
|
||||
assert find_deployment_by_name(client, "agent") == ExistingDeployment(
|
||||
"dep-target", "external_docker"
|
||||
)
|
||||
|
||||
|
||||
def test_a_partial_page_without_a_match_means_the_name_is_free():
|
||||
@@ -1414,3 +1458,59 @@ class TestPlacementOnAKnownListener:
|
||||
def test_refuses_and_names_the_namespaces(self, request_, listener, message):
|
||||
with pytest.raises(click.UsageError, match=message):
|
||||
request_.on(listener)
|
||||
|
||||
|
||||
DEFAULT_UPLOAD_HEADERS = {
|
||||
"content-type": "application/gzip",
|
||||
"content-length": "7",
|
||||
"x-goog-content-length-range": "0,209715200",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upload_headers", "expected"),
|
||||
[
|
||||
pytest.param(None, DEFAULT_UPLOAD_HEADERS, id="missing_sends_defaults"),
|
||||
pytest.param({}, DEFAULT_UPLOAD_HEADERS, id="empty_sends_defaults"),
|
||||
pytest.param(
|
||||
{
|
||||
"content-type": "application/x-gzip",
|
||||
"x-ms-blob-type": "BlockBlob",
|
||||
"content-length": "1",
|
||||
},
|
||||
{
|
||||
"content-type": "application/x-gzip",
|
||||
"x-ms-blob-type": "BlockBlob",
|
||||
"content-length": "7",
|
||||
},
|
||||
id="returned_headers_replace_defaults",
|
||||
),
|
||||
pytest.param(
|
||||
{"x-ms-blob-type": "BlockBlob"},
|
||||
{
|
||||
"content-type": "application/gzip",
|
||||
"x-ms-blob-type": "BlockBlob",
|
||||
"content-length": "7",
|
||||
},
|
||||
id="returned_headers_keep_a_content_type",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_source_upload_sends_the_headers_the_signed_url_requires(
|
||||
tmp_path, mocker, upload_headers, expected
|
||||
):
|
||||
archive = tmp_path / "source.tar.gz"
|
||||
archive.write_bytes(b"tarball")
|
||||
sent = []
|
||||
mocker.patch(
|
||||
"langgraph_cli.deploy._get_emitter", return_value=MagicMock(json_mode=True)
|
||||
)
|
||||
mocker.patch(
|
||||
"urllib.request.urlopen", side_effect=lambda req, timeout: sent.append(req)
|
||||
)
|
||||
|
||||
deploy_mod._upload_source_tarball(
|
||||
"https://storage.example/signed", str(archive), 7, upload_headers
|
||||
)
|
||||
|
||||
assert {name.lower(): value for name, value in sent[0].header_items()} == expected
|
||||
@@ -575,7 +575,7 @@ def test_list_listeners_asks_for_a_full_page():
|
||||
)
|
||||
|
||||
assert c.list_listeners() == [{"id": "listener-1"}]
|
||||
assert seen["url"] == "https://api.example.com/v2/listeners?limit=100"
|
||||
assert seen["url"] == "https://api.example.com/v2/listeners?limit=100&offset=0"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -602,29 +602,30 @@ def test_is_cloud_recognises_the_managed_control_plane(control_plane_url, expect
|
||||
assert endpoints.is_cloud is expected
|
||||
|
||||
|
||||
FIRST_PAGE = {"limit": "100", "offset": "0"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("call", "expected_params"),
|
||||
[
|
||||
pytest.param(
|
||||
lambda c: c.list_deployments(name="agent"),
|
||||
{"name": "agent"},
|
||||
{"name": "agent", **FIRST_PAGE},
|
||||
id="exact_name_filters_server_side",
|
||||
),
|
||||
pytest.param(
|
||||
lambda c: c.list_deployments(name_contains="age"),
|
||||
{"name_contains": "age"},
|
||||
{"name_contains": "age", **FIRST_PAGE},
|
||||
id="substring_search_keeps_its_own_parameter",
|
||||
),
|
||||
pytest.param(
|
||||
lambda c: c.list_deployments(),
|
||||
{},
|
||||
id="no_filter_sends_no_parameters",
|
||||
FIRST_PAGE,
|
||||
id="no_filter_sends_only_the_page",
|
||||
),
|
||||
pytest.param(
|
||||
lambda c: c.list_deployments(
|
||||
name="agent", name_contains="agent", limit=100
|
||||
),
|
||||
{"name": "agent", "name_contains": "agent", "limit": "100"},
|
||||
lambda c: c.list_deployments(name="agent", name_contains="agent"),
|
||||
{"name": "agent", "name_contains": "agent", **FIRST_PAGE},
|
||||
id="both_filters_travel_together_for_older_servers",
|
||||
),
|
||||
],
|
||||
@@ -645,6 +646,68 @@ def test_list_deployments_sends_one_name_filter(call, expected_params):
|
||||
assert seen == expected_params
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"list_every",
|
||||
[
|
||||
pytest.param(HostBackendClient.list_deployments, id="deployments"),
|
||||
pytest.param(HostBackendClient.list_listeners, id="listeners"),
|
||||
],
|
||||
)
|
||||
def test_list_endpoints_read_every_page(list_every):
|
||||
resources = [{"id": f"resource-{index}"} for index in range(101)]
|
||||
offsets: list[str] = []
|
||||
|
||||
def handler(req: httpx.Request) -> httpx.Response:
|
||||
offset = int(req.url.params["offset"])
|
||||
limit = int(req.url.params["limit"])
|
||||
offsets.append(req.url.params["offset"])
|
||||
return httpx.Response(
|
||||
200, json={"resources": resources[offset : offset + limit]}
|
||||
)
|
||||
|
||||
c = HostBackendClient(
|
||||
"https://api.example.com", "key", transport=httpx.MockTransport(handler)
|
||||
)
|
||||
|
||||
assert list_every(c) == resources
|
||||
assert offsets == ["0", "100"]
|
||||
|
||||
|
||||
FIRST_PAGE_OF_DEPLOYMENTS = [
|
||||
{"id": f"dep-{index}", "status": "DEPLOYED"} for index in range(100)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"repeated_page",
|
||||
[
|
||||
pytest.param(FIRST_PAGE_OF_DEPLOYMENTS, id="same_page_again"),
|
||||
pytest.param(
|
||||
[
|
||||
{"id": "dep-99", "status": "DEPLOYING"},
|
||||
*FIRST_PAGE_OF_DEPLOYMENTS[:99],
|
||||
],
|
||||
id="same_deployments_reordered_and_updated",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_list_deployments_fails_when_the_server_ignores_the_offset(repeated_page):
|
||||
pages = iter([FIRST_PAGE_OF_DEPLOYMENTS, repeated_page])
|
||||
|
||||
def handler(req: httpx.Request) -> httpx.Response:
|
||||
page = next(pages, None)
|
||||
if page is None:
|
||||
return httpx.Response(500, json={"detail": "the guard let paging run on"})
|
||||
return httpx.Response(200, json={"resources": page})
|
||||
|
||||
c = HostBackendClient(
|
||||
"https://api.example.com", "key", transport=httpx.MockTransport(handler)
|
||||
)
|
||||
|
||||
with pytest.raises(HostBackendError, match="ignores 'offset'"):
|
||||
c.list_deployments()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "expected"),
|
||||
[
|
||||
|
||||
@@ -1,11 +1,6 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
from langgraph_cli.deploy import (
|
||||
_extract_deployment_url,
|
||||
format_deployments_table,
|
||||
format_listeners_table,
|
||||
format_revisions_table,
|
||||
)
|
||||
from langgraph_cli.deploy import Table
|
||||
from langgraph_cli.util import clean_empty_lines, warn_non_wolfi_distro
|
||||
|
||||
|
||||
@@ -194,96 +189,7 @@ def test_warn_non_wolfi_distro_does_not_modify_config():
|
||||
assert config_copy == original_config # Config should remain unchanged
|
||||
|
||||
|
||||
def test_extract_deployment_url_uses_custom_url():
|
||||
deployment = {"source_config": {"custom_url": "https://example.com/custom"}}
|
||||
assert _extract_deployment_url(deployment) == "https://example.com/custom"
|
||||
def test_a_table_without_rows_still_renders_its_headers():
|
||||
table = Table(("Name", "Status"), lambda resource: (None, None))
|
||||
|
||||
|
||||
def test_extract_deployment_url_defaults_to_dash():
|
||||
assert _extract_deployment_url({"id": "dep-123"}) == "-"
|
||||
|
||||
|
||||
def test_format_deployments_table():
|
||||
output = format_deployments_table(
|
||||
[
|
||||
{
|
||||
"id": "dep-123",
|
||||
"name": "alpha",
|
||||
"source_config": {"custom_url": "https://alpha.example.com"},
|
||||
},
|
||||
{
|
||||
"id": "dep-456",
|
||||
"name": "beta",
|
||||
"url": "https://beta.example.com",
|
||||
},
|
||||
]
|
||||
)
|
||||
assert "Deployment ID" in output
|
||||
assert "Deployment Name" in output
|
||||
assert "Deployment URL" in output
|
||||
assert "dep-123" in output
|
||||
assert "alpha" in output
|
||||
assert "https://alpha.example.com" in output
|
||||
assert "dep-456" in output
|
||||
|
||||
|
||||
def test_format_revisions_table():
|
||||
output = format_revisions_table(
|
||||
[
|
||||
{
|
||||
"id": "rev-123",
|
||||
"status": "DEPLOYED",
|
||||
"created_at": "2023-11-09T10:00:00Z",
|
||||
},
|
||||
{
|
||||
"id": "rev-456",
|
||||
"status": "CREATING",
|
||||
"created_at": "2023-11-07T05:31:56Z",
|
||||
},
|
||||
{
|
||||
"id": "rev-789",
|
||||
"status": "DEPLOYED",
|
||||
"created_at": "2023-11-08T10:00:00Z",
|
||||
},
|
||||
]
|
||||
)
|
||||
assert "Revision ID" in output
|
||||
assert "Status" in output
|
||||
assert "Created At" in output
|
||||
assert "rev-123" in output
|
||||
assert "CREATING" in output
|
||||
assert "2023-11-07T05:31:56Z" in output
|
||||
assert "rev-456" in output
|
||||
assert "rev-789" in output
|
||||
assert "REPLACED" in output
|
||||
|
||||
|
||||
def test_format_listeners_table():
|
||||
output = format_listeners_table(
|
||||
[
|
||||
{
|
||||
"id": "listener-1",
|
||||
"compute_id": "prod-cluster",
|
||||
"compute_config": {"k8s_namespaces": ["agents"]},
|
||||
},
|
||||
{
|
||||
"id": "listener-2",
|
||||
"compute_id": "multi-cluster",
|
||||
"compute_config": {"k8s_namespaces": ["agents", "agents-staging"]},
|
||||
},
|
||||
{
|
||||
"id": "listener-3",
|
||||
"compute_id": "broken-cluster",
|
||||
},
|
||||
]
|
||||
)
|
||||
assert "Listener ID" in output
|
||||
assert "Compute ID" in output
|
||||
assert "Namespaces" in output
|
||||
assert "listener-1" in output
|
||||
assert "prod-cluster" in output
|
||||
assert "agents" in output
|
||||
assert "listener-2" in output
|
||||
assert "agents, agents-staging" in output
|
||||
assert "listener-3" in output
|
||||
assert "broken-cluster" in output
|
||||
assert table.render([]).splitlines() == ["Name Status", "---- ------"]
|
||||
@@ -347,6 +347,7 @@ def ensure_config(*configs: RunnableConfig | None) -> RunnableConfig:
|
||||
recursion_limit=DEFAULT_RECURSION_LIMIT,
|
||||
configurable={},
|
||||
)
|
||||
inherited_tags: set[str] = set()
|
||||
if var_config := var_child_runnable_config.get():
|
||||
empty.update(
|
||||
{
|
||||
@@ -355,6 +356,7 @@ def ensure_config(*configs: RunnableConfig | None) -> RunnableConfig:
|
||||
if _is_not_empty(v)
|
||||
},
|
||||
)
|
||||
inherited_tags.update(empty["tags"])
|
||||
# An explicit config that supplies its own checkpoint coordinate (a
|
||||
# thread_id, or any checkpoint_ns/checkpoint_id/checkpoint_map) is addressing
|
||||
# its own checkpoint lineage, so drop the inherited ambient configurable
|
||||
@@ -407,12 +409,14 @@ def ensure_config(*configs: RunnableConfig | None) -> RunnableConfig:
|
||||
elif k == "tags":
|
||||
# Concatenate tags across configs so values bound via
|
||||
# with_config(...) are preserved when later configs
|
||||
# supply additional tags. Matches merge_configs.
|
||||
# supply additional tags. Matches merge_configs. Tags the
|
||||
# ambient config already has are skipped: a subgraph gets
|
||||
# its parent's config both explicitly and as the ambient
|
||||
# one, so they would double at every nesting level.
|
||||
existing_tags: list[str] | None = empty.get("tags")
|
||||
new_tags = [t for t in cast(list, v) if t not in inherited_tags]
|
||||
empty["tags"] = (
|
||||
[*existing_tags, *cast(list, v)]
|
||||
if existing_tags
|
||||
else list(cast(list, v))
|
||||
[*existing_tags, *new_tags] if existing_tags else new_tags
|
||||
)
|
||||
else:
|
||||
empty[k] = v # type: ignore[literal-required]
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import enum
|
||||
import inspect
|
||||
import logging
|
||||
@@ -23,9 +24,16 @@ from typing import (
|
||||
Optional,
|
||||
Protocol,
|
||||
TypeGuard,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
from langchain_core.callbacks.base import BaseCallbackHandler, BaseCallbackManager
|
||||
from langchain_core.callbacks.manager import (
|
||||
AsyncCallbackManagerForChainRun,
|
||||
BaseRunManager,
|
||||
CallbackManagerForChainRun,
|
||||
)
|
||||
from langchain_core.runnables.base import (
|
||||
Runnable,
|
||||
RunnableConfig,
|
||||
@@ -41,6 +49,7 @@ from langchain_core.runnables.config import (
|
||||
var_child_runnable_config,
|
||||
)
|
||||
from langchain_core.runnables.utils import Input, Output
|
||||
from langchain_core.tracers.base import BaseTracer
|
||||
from langchain_core.tracers.langchain import LangChainTracer
|
||||
from langgraph.store.base import BaseStore
|
||||
|
||||
@@ -57,7 +66,7 @@ from langgraph._internal._constants import (
|
||||
)
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.errors import NodeError
|
||||
from langgraph.types import StreamWriter
|
||||
from langgraph.types import StreamWriter, TracePolicy
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
@@ -85,6 +94,29 @@ def _trace_payload(value: Any, transform: Callable[[Any], Any] | None) -> Any:
|
||||
return value
|
||||
|
||||
|
||||
def _is_stream_handler(handler: BaseCallbackHandler) -> bool:
|
||||
# Inheritance, not isinstance on the runtime-checkable protocol, which also
|
||||
# matches any exporter that merely has the tap methods: stream handlers get a
|
||||
# node's real payload and its streamed output.
|
||||
return _StreamingCallbackHandler in type(handler).__mro__
|
||||
|
||||
|
||||
_M = TypeVar("_M", bound=BaseCallbackManager | BaseRunManager)
|
||||
|
||||
|
||||
def _split_handlers(manager: _M) -> tuple[_M, _M]:
|
||||
# The recorded payload exists to keep tracer ingest small, so only tracers get it.
|
||||
# `astream_log`'s handler is a tracer too, but it feeds the caller's own stream.
|
||||
real, tracers = copy.copy(manager), copy.copy(manager)
|
||||
tracers.handlers = [
|
||||
h
|
||||
for h in manager.handlers
|
||||
if isinstance(h, BaseTracer) and not _is_stream_handler(h)
|
||||
]
|
||||
real.handlers = [h for h in manager.handlers if h not in tracers.handlers]
|
||||
return real, tracers
|
||||
|
||||
|
||||
def _set_config_context(
|
||||
config: RunnableConfig, run: Any = None
|
||||
) -> Token[RunnableConfig | None]:
|
||||
@@ -593,13 +625,15 @@ class RunnableSeq(Runnable):
|
||||
*steps: RunnableLike,
|
||||
name: str | None = None,
|
||||
trace_inputs: Callable[[Any], Any] | None = None,
|
||||
trace_outputs: Callable[[Any], Any] | None = None,
|
||||
trace_policy: TracePolicy | None = None,
|
||||
) -> None:
|
||||
"""Create a new RunnableSeq.
|
||||
|
||||
Args:
|
||||
steps: The steps to include in the sequence.
|
||||
name: The name of the `Runnable`.
|
||||
trace_inputs: Transforms the input every callback handler receives.
|
||||
trace_policy: Processes what tracers record, after `trace_inputs`.
|
||||
|
||||
Raises:
|
||||
ValueError: If the sequence has less than 2 steps.
|
||||
@@ -619,7 +653,67 @@ class RunnableSeq(Runnable):
|
||||
self.steps = steps_flat
|
||||
self.name = name
|
||||
self.trace_inputs = trace_inputs
|
||||
self.trace_outputs = trace_outputs
|
||||
self.trace_policy = trace_policy or TracePolicy()
|
||||
|
||||
def _start(self, input: Any, config: RunnableConfig) -> CallbackManagerForChainRun:
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
payload = _trace_payload(input, self.trace_inputs)
|
||||
recorded = _trace_payload(payload, self.trace_policy.process_inputs)
|
||||
name = config.get("run_name") or self.get_name()
|
||||
run_id = config.pop("run_id", None)
|
||||
if recorded is payload:
|
||||
return callback_manager.on_chain_start(
|
||||
None, payload, name=name, run_id=run_id
|
||||
)
|
||||
real, tracers = _split_handlers(callback_manager)
|
||||
run_manager = real.on_chain_start(None, payload, name=name, run_id=run_id)
|
||||
tracers.on_chain_start(None, recorded, name=name, run_id=run_manager.run_id)
|
||||
run_manager.handlers = callback_manager.handlers
|
||||
return run_manager
|
||||
|
||||
async def _astart(
|
||||
self, input: Any, config: RunnableConfig
|
||||
) -> AsyncCallbackManagerForChainRun:
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
payload = _trace_payload(input, self.trace_inputs)
|
||||
recorded = _trace_payload(payload, self.trace_policy.process_inputs)
|
||||
name = config.get("run_name") or self.get_name()
|
||||
run_id = config.pop("run_id", None)
|
||||
if recorded is payload:
|
||||
return await callback_manager.on_chain_start(
|
||||
None, payload, name=name, run_id=run_id
|
||||
)
|
||||
real, tracers = _split_handlers(callback_manager)
|
||||
# Tracers start last, so a cancellation before then leaves no tracer run open.
|
||||
run_manager = await real.on_chain_start(None, payload, name=name, run_id=run_id)
|
||||
await tracers.on_chain_start(
|
||||
None, recorded, name=name, run_id=run_manager.run_id
|
||||
)
|
||||
run_manager.handlers = callback_manager.handlers
|
||||
return run_manager
|
||||
|
||||
def _end(self, run_manager: CallbackManagerForChainRun, output: Any) -> None:
|
||||
recorded = _trace_payload(output, self.trace_policy.process_outputs)
|
||||
if recorded is output:
|
||||
run_manager.on_chain_end(output)
|
||||
else:
|
||||
real, tracers = _split_handlers(run_manager)
|
||||
tracers.on_chain_end(recorded)
|
||||
real.on_chain_end(output)
|
||||
|
||||
async def _aend(
|
||||
self, run_manager: AsyncCallbackManagerForChainRun, output: Any
|
||||
) -> None:
|
||||
recorded = _trace_payload(output, self.trace_policy.process_outputs)
|
||||
if recorded is output:
|
||||
await run_manager.on_chain_end(output)
|
||||
else:
|
||||
real, tracers = _split_handlers(run_manager)
|
||||
# Tracers end first: on_chain_end is shielded once it starts, so a
|
||||
# cancellation can only skip the other handlers' end, never leave a
|
||||
# tracer run open.
|
||||
await tracers.on_chain_end(recorded)
|
||||
await real.on_chain_end(output)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
@@ -676,15 +770,8 @@ class RunnableSeq(Runnable):
|
||||
) -> Any:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks and context
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = self._start(input, config)
|
||||
# invoke all steps in sequence
|
||||
try:
|
||||
for i, step in enumerate(self.steps):
|
||||
@@ -712,7 +799,7 @@ class RunnableSeq(Runnable):
|
||||
run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
run_manager.on_chain_end(_trace_payload(input, self.trace_outputs))
|
||||
self._end(run_manager, input)
|
||||
return input
|
||||
|
||||
async def ainvoke(
|
||||
@@ -723,15 +810,8 @@ class RunnableSeq(Runnable):
|
||||
) -> Any:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = await self._astart(input, config)
|
||||
|
||||
# invoke all steps in sequence
|
||||
try:
|
||||
@@ -765,7 +845,7 @@ class RunnableSeq(Runnable):
|
||||
await run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
await run_manager.on_chain_end(_trace_payload(input, self.trace_outputs))
|
||||
await self._aend(run_manager, input)
|
||||
return input
|
||||
|
||||
def stream(
|
||||
@@ -776,15 +856,8 @@ class RunnableSeq(Runnable):
|
||||
) -> Iterator[Any]:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = self._start(input, config)
|
||||
# get the run object
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, LangChainTracer):
|
||||
@@ -814,10 +887,9 @@ class RunnableSeq(Runnable):
|
||||
)
|
||||
iterator = step.transform(iterator, config)
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
if _StreamingCallbackHandler is not None:
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, _StreamingCallbackHandler):
|
||||
iterator = h.tap_output_iter(run_manager.run_id, iterator)
|
||||
for h in run_manager.handlers:
|
||||
if _is_stream_handler(h):
|
||||
iterator = h.tap_output_iter(run_manager.run_id, iterator)
|
||||
# consume into final output
|
||||
output = context.run(_consume_iter, iterator)
|
||||
# sequence doesn't emit output, yield to mark as generator
|
||||
@@ -826,7 +898,7 @@ class RunnableSeq(Runnable):
|
||||
run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
run_manager.on_chain_end(_trace_payload(output, self.trace_outputs))
|
||||
self._end(run_manager, output)
|
||||
|
||||
async def astream(
|
||||
self,
|
||||
@@ -836,15 +908,8 @@ class RunnableSeq(Runnable):
|
||||
) -> AsyncIterator[Any]:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = await self._astart(input, config)
|
||||
# stream the last steps
|
||||
# transform the input stream of each step with the next
|
||||
# steps that don't natively support transforming an input stream will
|
||||
@@ -880,12 +945,11 @@ class RunnableSeq(Runnable):
|
||||
if hasattr(aiterator, "aclose"):
|
||||
stack.push_async_callback(aiterator.aclose)
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
if _StreamingCallbackHandler is not None:
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, _StreamingCallbackHandler):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
for h in run_manager.handlers:
|
||||
if _is_stream_handler(h):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
# consume into final output
|
||||
output = await asyncio.create_task(
|
||||
_consume_aiter(aiterator), context=context
|
||||
@@ -896,9 +960,7 @@ class RunnableSeq(Runnable):
|
||||
await run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
await run_manager.on_chain_end(
|
||||
_trace_payload(output, self.trace_outputs)
|
||||
)
|
||||
await self._aend(run_manager, output)
|
||||
else:
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
@@ -914,12 +976,11 @@ class RunnableSeq(Runnable):
|
||||
if hasattr(aiterator, "aclose"):
|
||||
stack.push_async_callback(aiterator.aclose)
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
if _StreamingCallbackHandler is not None:
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, _StreamingCallbackHandler):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
for h in run_manager.handlers:
|
||||
if _is_stream_handler(h):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
# consume into final output
|
||||
output = await _consume_aiter(aiterator)
|
||||
# sequence doesn't emit output, yield to mark as generator
|
||||
@@ -928,9 +989,7 @@ class RunnableSeq(Runnable):
|
||||
await run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
await run_manager.on_chain_end(
|
||||
_trace_payload(output, self.trace_outputs)
|
||||
)
|
||||
await self._aend(run_manager, output)
|
||||
|
||||
|
||||
def _consume_iter(it: Iterator[Any]) -> Any:
|
||||
|
||||
@@ -60,7 +60,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
typ: The value type (e.g. `list`, `dict`). Inferred automatically
|
||||
from the outer type when used inside `Annotated[T, DeltaChannel(...)]`.
|
||||
snapshot_frequency: Every Nth update to this channel writes a snapshot
|
||||
blob (default `1000`). Must be a positive int.
|
||||
blob (default `100`). Must be a positive int.
|
||||
"""
|
||||
|
||||
__slots__ = ("value", "reducer", "snapshot_frequency")
|
||||
@@ -71,7 +71,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
||||
reducer: Callable[[Any, Sequence[Any]], Any],
|
||||
typ: type[Value] | None = None,
|
||||
*,
|
||||
snapshot_frequency: int = 1000,
|
||||
snapshot_frequency: int = 100,
|
||||
) -> None:
|
||||
if snapshot_frequency <= 0:
|
||||
raise ValueError(
|
||||
|
||||
@@ -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)])
|
||||
|
||||
@@ -228,6 +228,10 @@ def draw_graph(
|
||||
metadata["__interrupt"] = "before"
|
||||
elif name in interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "after"
|
||||
# Duck-typed so core doesn't import ToolNode from langgraph-prebuilt.
|
||||
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)
|
||||
graph.add_node(node.bound, name, metadata=metadata or None)
|
||||
# add start node
|
||||
if START not in nodes:
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable, Coroutine
|
||||
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
|
||||
from contextvars import copy_context
|
||||
@@ -37,6 +38,13 @@ class Submit(Protocol[P, T]):
|
||||
) -> concurrent.futures.Future[T]: ...
|
||||
|
||||
|
||||
# Futures whose exception a node's error handler took over; the executors
|
||||
# don't re-raise them on exit.
|
||||
SKIP_RERAISE_SET: weakref.WeakSet[concurrent.futures.Future | asyncio.Future] = (
|
||||
weakref.WeakSet()
|
||||
)
|
||||
|
||||
|
||||
class BackgroundExecutor(AbstractContextManager):
|
||||
"""A context manager that runs sync tasks in the background.
|
||||
Uses a thread pool executor to delegate tasks to separate threads.
|
||||
@@ -111,7 +119,7 @@ class BackgroundExecutor(AbstractContextManager):
|
||||
if exc_type is None:
|
||||
# re-raise the first exception that occurred in a task
|
||||
for task, (_, reraise) in tasks.items():
|
||||
if not reraise:
|
||||
if not reraise or task in SKIP_RERAISE_SET:
|
||||
continue
|
||||
try:
|
||||
task.result()
|
||||
@@ -202,7 +210,7 @@ class AsyncBackgroundExecutor(AbstractAsyncContextManager):
|
||||
if exc_type is None:
|
||||
# re-raise the first exception that occurred in a task
|
||||
for task, (_, reraise) in tasks.items():
|
||||
if not reraise:
|
||||
if not reraise or task in SKIP_RERAISE_SET:
|
||||
continue
|
||||
try:
|
||||
if exc := task.exception():
|
||||
|
||||
@@ -132,7 +132,7 @@ from langgraph.pregel.debug import (
|
||||
map_debug_tasks,
|
||||
)
|
||||
from langgraph.pregel.protocol import StreamChunk, StreamProtocol
|
||||
from langgraph.runtime import RunControl, Runtime
|
||||
from langgraph.runtime import RunControl, Runtime, _RunExits
|
||||
from langgraph.types import (
|
||||
All,
|
||||
CachePolicy,
|
||||
@@ -657,48 +657,11 @@ class PregelLoop:
|
||||
# check if iteration limit is reached
|
||||
if self.step > self.stop:
|
||||
self.status = "out_of_steps"
|
||||
# stream the checkpoint the last step saved, as a tick that runs would
|
||||
self._prepare_tasks_and_emit_checkpoint()
|
||||
return False
|
||||
|
||||
# prepare next tasks
|
||||
self.tasks = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
self.checkpoint_pending_writes,
|
||||
self.nodes,
|
||||
self.channels,
|
||||
self.managed,
|
||||
self.config,
|
||||
self.step,
|
||||
self.stop,
|
||||
for_execution=True,
|
||||
manager=self.manager,
|
||||
store=self.store,
|
||||
checkpointer=self.checkpointer,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
updated_channels=self.updated_channels,
|
||||
retry_policy=self.retry_policy,
|
||||
cache_policy=self.cache_policy,
|
||||
)
|
||||
|
||||
# produce debug output
|
||||
if self._checkpointer_put_after_previous is not None:
|
||||
self._emit(
|
||||
"checkpoints",
|
||||
map_debug_checkpoint,
|
||||
{
|
||||
**self.checkpoint_config,
|
||||
CONF: {
|
||||
**self.checkpoint_config[CONF],
|
||||
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
|
||||
},
|
||||
},
|
||||
self.channels,
|
||||
self.stream_keys,
|
||||
self.checkpoint_metadata,
|
||||
self.tasks.values(),
|
||||
self.checkpoint_pending_writes,
|
||||
self.prev_checkpoint_config,
|
||||
self.output_keys,
|
||||
)
|
||||
self._prepare_tasks_and_emit_checkpoint()
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not self.tasks:
|
||||
@@ -806,10 +769,53 @@ class PregelLoop:
|
||||
self.checkpoint, self.interrupt_after, self.tasks.values()
|
||||
):
|
||||
self.status = "interrupt_after"
|
||||
# the next tick, which streams this checkpoint, never runs
|
||||
self._prepare_tasks_and_emit_checkpoint()
|
||||
raise GraphInterrupt()
|
||||
# unset resuming flag
|
||||
self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
|
||||
|
||||
def _prepare_tasks_and_emit_checkpoint(self) -> None:
|
||||
self.tasks = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
self.checkpoint_pending_writes,
|
||||
self.nodes,
|
||||
self.channels,
|
||||
self.managed,
|
||||
self.config,
|
||||
self.step,
|
||||
self.stop,
|
||||
for_execution=True,
|
||||
manager=self.manager,
|
||||
store=self.store,
|
||||
checkpointer=self.checkpointer,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
updated_channels=self.updated_channels,
|
||||
retry_policy=self.retry_policy,
|
||||
cache_policy=self.cache_policy,
|
||||
)
|
||||
|
||||
# produce debug output
|
||||
if self._checkpointer_put_after_previous is not None:
|
||||
self._emit(
|
||||
"checkpoints",
|
||||
map_debug_checkpoint,
|
||||
{
|
||||
**self.checkpoint_config,
|
||||
CONF: {
|
||||
**self.checkpoint_config[CONF],
|
||||
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
|
||||
},
|
||||
},
|
||||
self.channels,
|
||||
self.stream_keys,
|
||||
self.checkpoint_metadata,
|
||||
self.tasks.values(),
|
||||
self.checkpoint_pending_writes,
|
||||
self.prev_checkpoint_config,
|
||||
self.output_keys,
|
||||
)
|
||||
|
||||
def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -1491,6 +1497,10 @@ class PregelLoop:
|
||||
if self._delta_write_futs is not None:
|
||||
self._delta_write_futs.append(fut)
|
||||
|
||||
def _run_exits(self) -> _RunExits | None:
|
||||
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
|
||||
return runtime._run_exits if isinstance(runtime, Runtime) else None
|
||||
|
||||
def _suppress_interrupt(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
@@ -1868,6 +1878,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
self._exit_delta_writes = (
|
||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||
)
|
||||
# Bottom of the stack, so run-exit callbacks run last: after the
|
||||
# executor has stopped every task and the final checkpoint is saved.
|
||||
if (exits := self._run_exits()) is not None and exits.bind(is_async=False):
|
||||
self.stack.push(exits.exit)
|
||||
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs,
|
||||
@@ -2124,6 +2138,10 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
self._exit_delta_writes = (
|
||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||
)
|
||||
# Bottom of the stack, so run-exit callbacks run last: after the
|
||||
# executor has stopped every task and the final checkpoint is saved.
|
||||
if (exits := self._run_exits()) is not None and exits.bind(is_async=True):
|
||||
self.stack.push_async_exit(exits.aexit)
|
||||
self.submit = await self.stack.enter_async_context(
|
||||
AsyncBackgroundExecutor(self.config)
|
||||
)
|
||||
|
||||
@@ -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,10 +228,8 @@ 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
|
||||
]
|
||||
if not self.subgraphs and len(ns) > 0:
|
||||
ns = graph_namespace(metadata)
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
for value in _state_values(inputs):
|
||||
|
||||
@@ -227,23 +227,14 @@ class PregelNode:
|
||||
def node(self) -> Runnable[Any, Any] | None:
|
||||
"""Get a runnable that combines `bound` and `writers`."""
|
||||
writers = self.flat_writers
|
||||
trace_inputs = self.trace_policy.process_inputs if self.trace_policy else None
|
||||
trace_outputs = self.trace_policy.process_outputs if self.trace_policy else None
|
||||
if self.bound is DEFAULT_BOUND and not writers:
|
||||
return None
|
||||
elif self.bound is DEFAULT_BOUND and len(writers) == 1:
|
||||
return writers[0]
|
||||
elif self.bound is DEFAULT_BOUND:
|
||||
return RunnableSeq(
|
||||
*writers, trace_inputs=trace_inputs, trace_outputs=trace_outputs
|
||||
)
|
||||
return RunnableSeq(*writers, trace_policy=self.trace_policy)
|
||||
elif writers:
|
||||
return RunnableSeq(
|
||||
self.bound,
|
||||
*writers,
|
||||
trace_inputs=trace_inputs,
|
||||
trace_outputs=trace_outputs,
|
||||
)
|
||||
return RunnableSeq(self.bound, *writers, trace_policy=self.trace_policy)
|
||||
else:
|
||||
return self.bound
|
||||
|
||||
|
||||
@@ -34,6 +34,7 @@ from langgraph._internal._constants import (
|
||||
CONFIG_KEY_THREAD_ID,
|
||||
CONFIG_KEY_TIMED_ATTEMPT_OBSERVER,
|
||||
NS_SEP,
|
||||
RETURN,
|
||||
)
|
||||
from langgraph._internal._runnable import create_task_in_config_context
|
||||
from langgraph._internal._timeout import sync_timeout_unsupported
|
||||
@@ -733,15 +734,14 @@ async def arun_with_retry(
|
||||
if stream:
|
||||
async for _ in task.proc.astream(task.input, config):
|
||||
pass
|
||||
break
|
||||
return _return_value(task)
|
||||
return await task.proc.ainvoke(task.input, config)
|
||||
result = await _arun_with_timeout(
|
||||
task, config, resolved_timeout, attempt_ctx, stream=stream
|
||||
)
|
||||
_finish_timed_attempt(config, attempt_ctx)
|
||||
if stream:
|
||||
# if successful, end
|
||||
break
|
||||
return _return_value(task)
|
||||
return result
|
||||
except ParentCommand as exc:
|
||||
ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
|
||||
@@ -831,6 +831,12 @@ async def arun_with_retry(
|
||||
config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
|
||||
|
||||
|
||||
def _return_value(task: PregelExecutableTask) -> Any:
|
||||
# Streaming drains the task instead of returning its result, so a caller
|
||||
# awaiting a `@task` needs the value the task wrote.
|
||||
return next((value for channel, value in task.writes if channel == RETURN), None)
|
||||
|
||||
|
||||
def _should_retry_on(retry_policy: RetryPolicy, exc: Exception) -> bool:
|
||||
"""Check if the given exception should be retried based on the retry policy."""
|
||||
if isinstance(retry_policy.retry_on, Sequence):
|
||||
|
||||
@@ -43,7 +43,7 @@ from langgraph._internal._typing import MISSING
|
||||
from langgraph.constants import TAG_HIDDEN
|
||||
from langgraph.errors import GraphBubbleUp, GraphInterrupt, is_invalid_resume
|
||||
from langgraph.pregel._algo import Call
|
||||
from langgraph.pregel._executor import Submit
|
||||
from langgraph.pregel._executor import SKIP_RERAISE_SET, Submit
|
||||
from langgraph.pregel._retry import arun_with_retry, run_with_retry
|
||||
from langgraph.pregel._task_status import CONTROL_WRITES
|
||||
from langgraph.types import (
|
||||
@@ -68,10 +68,6 @@ EXCLUDED_FRAME_FNAMES = (
|
||||
"concurrent/futures/_base.py",
|
||||
)
|
||||
|
||||
SKIP_RERAISE_SET: weakref.WeakSet[concurrent.futures.Future | asyncio.Future] = (
|
||||
weakref.WeakSet()
|
||||
)
|
||||
|
||||
|
||||
class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
|
||||
event: E
|
||||
@@ -239,6 +235,8 @@ class PregelRunner:
|
||||
fut.set_exception(exc)
|
||||
futures.done.add(fut)
|
||||
elif reraise and id(exc) not in self._handled_exception_ids:
|
||||
# let the caller stream this task's result before it raises
|
||||
yield
|
||||
if tb := exc.__traceback__:
|
||||
while tb.tb_next is not None and any(
|
||||
tb.tb_frame.f_code.co_filename.endswith(name)
|
||||
@@ -430,6 +428,8 @@ class PregelRunner:
|
||||
fut.set_exception(exc)
|
||||
futures.done.add(fut)
|
||||
elif reraise and id(exc) not in self._handled_exception_ids:
|
||||
# let the caller stream this task's result before it raises
|
||||
yield
|
||||
if tb := exc.__traceback__:
|
||||
while tb.tb_next is not None and any(
|
||||
tb.tb_frame.f_code.co_filename.endswith(name)
|
||||
|
||||
@@ -7,8 +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:
|
||||
@@ -91,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.
|
||||
@@ -109,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
|
||||
@@ -188,6 +184,10 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
return
|
||||
ns, tool_call_id, token = info
|
||||
self._reset_writer(token)
|
||||
# An interrupted or drained tool is paused, not failed: it runs again
|
||||
# on resume and reports its start and end then.
|
||||
if isinstance(error, (GraphInterrupt, GraphDrained)):
|
||||
return
|
||||
self.stream(
|
||||
(
|
||||
ns,
|
||||
|
||||
@@ -97,6 +97,7 @@ from langgraph._internal._runnable import (
|
||||
Runnable,
|
||||
RunnableLike,
|
||||
RunnableSeq,
|
||||
_is_stream_handler,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langgraph._internal._timeout import coerce_timeout_policy
|
||||
@@ -174,6 +175,7 @@ from langgraph.runtime import (
|
||||
RunControl,
|
||||
Runtime,
|
||||
ServerInfo,
|
||||
_RunExits,
|
||||
)
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import StreamTransformer
|
||||
@@ -203,11 +205,6 @@ from langgraph.types import (
|
||||
from langgraph.typing import ContextT, InputT, OutputT, StateT
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV10
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = None # type: ignore
|
||||
|
||||
__all__ = ("NodeBuilder", "Pregel")
|
||||
|
||||
_WriteValue = Callable[[Input], Output] | Any
|
||||
@@ -2880,6 +2877,20 @@ 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_mode_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_mode_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 +2909,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 +2920,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:
|
||||
@@ -2968,6 +2967,8 @@ class Pregel(
|
||||
execution_info=None,
|
||||
server_info=server_info,
|
||||
control=control or parent_runtime.control or RunControl(),
|
||||
# a subgraph shares the root run's exit callbacks
|
||||
_run_exits=parent_runtime._run_exits or _RunExits(),
|
||||
)
|
||||
runtime = parent_runtime.merge(runtime)
|
||||
config[CONF][CONFIG_KEY_RUNTIME] = runtime
|
||||
@@ -3009,6 +3010,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 +3296,20 @@ 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_mode_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_mode_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(
|
||||
@@ -3304,18 +3323,10 @@ class Pregel(
|
||||
run_id=run_manager.run_id,
|
||||
)
|
||||
# if running from astream_log() run each proc with streaming
|
||||
do_stream = (
|
||||
next(
|
||||
(
|
||||
True
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
and not isinstance(h, StreamMessagesHandler)
|
||||
),
|
||||
False,
|
||||
)
|
||||
if _StreamingCallbackHandler is not None
|
||||
else False
|
||||
do_stream = any(
|
||||
_is_stream_handler(h)
|
||||
and not isinstance(h, (StreamMessagesHandler, StreamToolCallHandler))
|
||||
for h in run_manager.handlers
|
||||
)
|
||||
if checkpointer is None and durability is not None:
|
||||
warnings.warn(
|
||||
@@ -3323,10 +3334,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)
|
||||
)
|
||||
@@ -3335,24 +3345,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:
|
||||
@@ -3409,6 +3407,8 @@ class Pregel(
|
||||
execution_info=None,
|
||||
server_info=server_info,
|
||||
control=control or parent_runtime.control or RunControl(),
|
||||
# a subgraph shares the root run's exit callbacks
|
||||
_run_exits=parent_runtime._run_exits or _RunExits(),
|
||||
)
|
||||
runtime = parent_runtime.merge(runtime)
|
||||
config[CONF][CONFIG_KEY_RUNTIME] = runtime
|
||||
@@ -3461,6 +3461,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(
|
||||
@@ -4458,6 +4462,10 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
|
||||
payload["values"] = mapper(payload["values"])
|
||||
|
||||
|
||||
def _is_stream_mode_handler(handler: object) -> bool:
|
||||
return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
|
||||
|
||||
|
||||
def _resolve_parent_ns(
|
||||
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
|
||||
) -> tuple[str, ...]:
|
||||
|
||||
@@ -847,6 +847,11 @@ 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))
|
||||
@@ -997,6 +1002,11 @@ 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))
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import inspect
|
||||
import logging
|
||||
import threading
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field, replace
|
||||
from types import TracebackType
|
||||
from typing import Any, Generic, cast
|
||||
|
||||
from langgraph.store.base import BaseStore
|
||||
@@ -104,6 +108,129 @@ class RunControl:
|
||||
return self._drain_reason
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ExitCallback = Callable[[], Awaitable[None] | None]
|
||||
"""A cleanup callback registered with `Runtime.on_exit`."""
|
||||
|
||||
|
||||
class _RunExits:
|
||||
"""Cleanup callbacks for one root graph invocation.
|
||||
|
||||
Created by the root `stream`/`astream` call and shared by every subgraph
|
||||
it runs, through `Runtime`. The root Pregel loop binds it at the bottom of
|
||||
its exit stack, so callbacks run after the loop has stopped every task,
|
||||
persisted its final checkpoint and handled any interrupt, whatever the
|
||||
reason the run ended.
|
||||
"""
|
||||
|
||||
__slots__ = ("_callbacks", "_is_async", "_lock", "_state")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._callbacks: list[ExitCallback] = []
|
||||
self._lock = threading.Lock()
|
||||
self._is_async = False
|
||||
# "new" -> "bound" (owned by a loop) -> "closed" (callbacks ran)
|
||||
self._state = "new"
|
||||
|
||||
def bind(self, *, is_async: bool) -> bool:
|
||||
"""Claim these exits for a loop. Only the first (root) loop wins."""
|
||||
with self._lock:
|
||||
if self._state != "new":
|
||||
return False
|
||||
self._state = "bound"
|
||||
self._is_async = is_async
|
||||
return True
|
||||
|
||||
def register(self, callback: ExitCallback) -> None:
|
||||
if not callable(callback):
|
||||
msg = "on_exit() expects a callable"
|
||||
raise TypeError(msg)
|
||||
with self._lock:
|
||||
if self._state == "closed":
|
||||
msg = "on_exit() called after the graph run has already exited"
|
||||
raise RuntimeError(msg)
|
||||
if not self._is_async and _is_async_callable(callback):
|
||||
msg = (
|
||||
"on_exit() got an async callback in a synchronous graph run "
|
||||
"(invoke/stream). Register a sync callback, or run the graph "
|
||||
"with ainvoke/astream."
|
||||
)
|
||||
raise TypeError(msg)
|
||||
self._callbacks.append(callback)
|
||||
|
||||
def _drain(self) -> list[ExitCallback]:
|
||||
with self._lock:
|
||||
self._state = "closed"
|
||||
callbacks = self._callbacks[::-1] # LIFO
|
||||
self._callbacks = []
|
||||
return callbacks
|
||||
|
||||
def exit(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> bool:
|
||||
errors: list[Exception] = []
|
||||
for callback in self._drain():
|
||||
try:
|
||||
result = callback()
|
||||
if inspect.isawaitable(result):
|
||||
if inspect.iscoroutine(result):
|
||||
result.close()
|
||||
msg = (
|
||||
f"on_exit() callback {callback!r} returned an awaitable "
|
||||
"in a synchronous graph run; it was not awaited."
|
||||
)
|
||||
raise TypeError(msg)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
return self._finish(errors, exc_value)
|
||||
|
||||
async def aexit(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> bool:
|
||||
errors: list[Exception] = []
|
||||
for callback in self._drain():
|
||||
try:
|
||||
result = callback()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
return self._finish(errors, exc_value)
|
||||
|
||||
@staticmethod
|
||||
def _finish(errors: list[Exception], exc_value: BaseException | None) -> bool:
|
||||
# Every callback runs even if an earlier one failed. A failing callback
|
||||
# never replaces the run's own exception (error, cancellation, stream
|
||||
# close): it is logged instead. If the run otherwise ended cleanly, the
|
||||
# first callback error is raised so cleanup bugs aren't silent.
|
||||
if not errors:
|
||||
return False
|
||||
if exc_value is None:
|
||||
for e in errors[1:]:
|
||||
logger.error("on_exit() callback failed", exc_info=e)
|
||||
raise errors[0]
|
||||
for e in errors:
|
||||
logger.error(
|
||||
"on_exit() callback failed while the graph run was exiting with %r",
|
||||
exc_value,
|
||||
exc_info=e,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _is_async_callable(callback: Any) -> bool:
|
||||
return inspect.iscoroutinefunction(callback) or inspect.iscoroutinefunction(
|
||||
getattr(callback, "__call__", None)
|
||||
)
|
||||
|
||||
|
||||
def _no_op_stream_writer(_: Any) -> None: ...
|
||||
|
||||
|
||||
@@ -237,6 +364,40 @@ class Runtime(Generic[ContextT]):
|
||||
graph runtime.
|
||||
"""
|
||||
|
||||
_run_exits: _RunExits | None = field(default=None, repr=False, compare=False)
|
||||
"""Exit callbacks of the root graph run. Use `on_exit` to register one."""
|
||||
|
||||
def on_exit(self, callback: ExitCallback) -> None:
|
||||
"""Register cleanup to run when the current graph run ends.
|
||||
|
||||
The callback runs once the root invocation (the `invoke`/`stream`/
|
||||
`ainvoke`/`astream` call the caller made) exits for any reason: it
|
||||
completed, raised, was cancelled, paused on an interrupt, or the caller
|
||||
closed the stream or broke out of it. Calls from subgraphs register on
|
||||
that same root run.
|
||||
|
||||
Callbacks take no arguments and run in reverse registration order,
|
||||
after every task of the run has stopped and the final checkpoint has
|
||||
been written. In `ainvoke`/`astream` a callback may be sync or async; in
|
||||
`invoke`/`stream` it must be sync. A failing callback never hides the
|
||||
run's own exception: it is logged, and raised only if the run otherwise
|
||||
ended without one.
|
||||
|
||||
A process that is killed does not run its callbacks, so use them to
|
||||
free resources early, not as the only guarantee they are freed.
|
||||
|
||||
Args:
|
||||
callback: A zero-argument sync or async callable.
|
||||
|
||||
Raises:
|
||||
RuntimeError: When called outside a graph run, or after it exited.
|
||||
TypeError: When given an async callback in a synchronous run.
|
||||
"""
|
||||
if self._run_exits is None:
|
||||
msg = "on_exit() can only be called during a graph run"
|
||||
raise RuntimeError(msg)
|
||||
self._run_exits.register(callback)
|
||||
|
||||
def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]:
|
||||
"""Merge two runtimes together.
|
||||
|
||||
@@ -255,6 +416,7 @@ class Runtime(Generic[ContextT]):
|
||||
execution_info=other.execution_info or self.execution_info,
|
||||
server_info=other.server_info or self.server_info,
|
||||
control=other.control or self.control,
|
||||
_run_exits=other._run_exits or self._run_exits,
|
||||
)
|
||||
|
||||
def override(
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.2.17"
|
||||
version = "1.2.19"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -133,6 +133,11 @@
|
||||
"ToolNode"
|
||||
],
|
||||
"name": "tools"
|
||||
},
|
||||
"metadata": {
|
||||
"__tools": [
|
||||
"search_api"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -221,6 +226,11 @@
|
||||
"ToolNode"
|
||||
],
|
||||
"name": "tools"
|
||||
},
|
||||
"metadata": {
|
||||
"__tools": [
|
||||
"search_api"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
|
||||
@@ -20,6 +20,7 @@ from langgraph.errors import EmptyChannelError, InvalidUpdateError
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.graph.state import _get_channel
|
||||
from langgraph.pregel._checkpoint import delta_channels_to_snapshot
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
@@ -342,6 +343,26 @@ def test_delta_channel_checkpoint_returns_missing() -> None:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frequency", [None, 50, 1000])
|
||||
def test_delta_channel_snapshot_frequency(frequency: int | None) -> None:
|
||||
spec = (
|
||||
DeltaChannel(_messages_delta_reducer)
|
||||
if frequency is None
|
||||
else DeltaChannel(_messages_delta_reducer, snapshot_frequency=frequency)
|
||||
)
|
||||
expected = 100 if frequency is None else frequency
|
||||
assert spec.snapshot_frequency == expected
|
||||
channel = spec.from_checkpoint(MISSING)
|
||||
channel.update([HumanMessage(content="hi", id="h1")])
|
||||
assert channel.copy().snapshot_frequency == expected
|
||||
for updates in (expected - 1, expected):
|
||||
assert delta_channels_to_snapshot(
|
||||
{"messages": channel},
|
||||
{"messages": (updates, updates)},
|
||||
{"messages": 1},
|
||||
) == ({"messages"} if updates == expected else set())
|
||||
|
||||
|
||||
def test_delta_channel_snapshot_version_based() -> None:
|
||||
"""Snapshots fire when a channel accumulates `snapshot_frequency` updates.
|
||||
|
||||
|
||||
@@ -4346,6 +4346,25 @@ def test_enum_node_names():
|
||||
assert graph.invoke({"foo": "hello"}) == {"foo": "hello", "bar": "hello!"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_mode", ["tasks", "debug"])
|
||||
def test_failing_task_result_is_streamed_before_the_error(stream_mode: str) -> None:
|
||||
class State(TypedDict):
|
||||
x: str
|
||||
|
||||
def failing(state: State) -> State:
|
||||
raise ValueError("boom")
|
||||
|
||||
graph = StateGraph(State).add_node("failing", failing).add_edge(START, "failing")
|
||||
errors = []
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
for chunk in graph.compile().stream({"x": ""}, stream_mode=stream_mode):
|
||||
payload = chunk["payload"] if stream_mode == "debug" else chunk
|
||||
if "error" in payload:
|
||||
errors.append(payload["error"])
|
||||
|
||||
assert [type(error) for error in errors] == [ValueError]
|
||||
|
||||
|
||||
def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[str], operator.add]
|
||||
@@ -4410,6 +4429,52 @@ def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
|
||||
assert stream_parent_conf == history_parent_conf
|
||||
|
||||
|
||||
def test_interrupt_after_streams_the_checkpoint_it_pauses_on(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
x: int
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("one", lambda state: {"x": 1})
|
||||
.add_node("two", lambda state: {"x": 2})
|
||||
.add_edge(START, "one")
|
||||
.add_edge("one", "two")
|
||||
.compile(checkpointer=sync_checkpointer, interrupt_after=["one"])
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
checkpoints = list(graph.stream({"x": 0}, config, stream_mode="checkpoints"))
|
||||
|
||||
state = graph.get_state(config)
|
||||
assert checkpoints[-1]["config"] == state.config
|
||||
assert checkpoints[-1]["next"] == ["two"]
|
||||
|
||||
|
||||
def test_recursion_limit_streams_the_last_checkpoint_saved(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
x: int
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("loop", lambda state: {"x": state["x"] + 1})
|
||||
.add_edge(START, "loop")
|
||||
.add_edge("loop", "loop")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}, "recursion_limit": 3}
|
||||
checkpoints = []
|
||||
|
||||
with pytest.raises(GraphRecursionError):
|
||||
for chunk in graph.stream({"x": 0}, config, stream_mode="checkpoints"):
|
||||
checkpoints.append(chunk)
|
||||
|
||||
assert checkpoints[-1]["config"] == graph.get_state(config).config
|
||||
|
||||
|
||||
def test_debug_subgraphs(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
):
|
||||
@@ -6880,6 +6945,22 @@ def test_named_tasks_functional() -> None:
|
||||
]
|
||||
|
||||
|
||||
def test_invoke_tags_reach_a_node_three_subgraphs_deep_once() -> None:
|
||||
class State(TypedDict):
|
||||
copies: int
|
||||
|
||||
def leaf(state: State, config: RunnableConfig) -> State:
|
||||
return {"copies": config["tags"].count("my-tag")}
|
||||
|
||||
graph = StateGraph(State).add_node("leaf", leaf).add_edge(START, "leaf").compile()
|
||||
for _ in range(3):
|
||||
graph = (
|
||||
StateGraph(State).add_node("sub", graph).add_edge(START, "sub").compile()
|
||||
)
|
||||
|
||||
assert graph.invoke({"copies": 0}, {"tags": ["my-tag"]}) == {"copies": 1}
|
||||
|
||||
|
||||
def test_tags_stream_mode_messages() -> None:
|
||||
model = GenericFakeChatModel(messages=iter(["foo"]), tags=["meow"])
|
||||
graph = (
|
||||
@@ -8869,6 +8950,39 @@ 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]
|
||||
|
||||
@@ -5809,6 +5809,60 @@ async def test_store_injected_async(
|
||||
) # still overwriting the same one
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_mode", ["tasks", "debug"])
|
||||
async def test_failing_task_result_is_streamed_before_the_error(
|
||||
stream_mode: str,
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
x: str
|
||||
|
||||
async def failing(state: State) -> State:
|
||||
raise ValueError("boom")
|
||||
|
||||
graph = StateGraph(State).add_node("failing", failing).add_edge(START, "failing")
|
||||
errors = []
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
async for chunk in graph.compile().astream({"x": ""}, stream_mode=stream_mode):
|
||||
payload = chunk["payload"] if stream_mode == "debug" else chunk
|
||||
if "error" in payload:
|
||||
errors.append(payload["error"])
|
||||
|
||||
assert [type(error) for error in errors] == [ValueError]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("parallel", "stream_mode"), [(False, "custom"), (True, "values")]
|
||||
)
|
||||
async def test_graph_error_handler_recovers_on_the_executor_path(
|
||||
parallel: bool, stream_mode: str
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
x: str
|
||||
|
||||
async def failing(state: State) -> State:
|
||||
raise ValueError("boom")
|
||||
|
||||
async def handler(state: State) -> Command:
|
||||
return Command(update={"x": "handled"}, goto=END)
|
||||
|
||||
builder = (
|
||||
StateGraph(State)
|
||||
.add_node("failing", failing, error_handler=handler)
|
||||
.add_edge(START, "failing")
|
||||
)
|
||||
if parallel:
|
||||
builder.add_node("sibling", lambda state: {}).add_edge(START, "sibling")
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in builder.compile().astream(
|
||||
{"x": ""}, stream_mode=[stream_mode, "values"]
|
||||
)
|
||||
]
|
||||
|
||||
assert [c[-1] for c in chunks if c[0] == "values"][-1] == {"x": "handled"}
|
||||
|
||||
|
||||
async def test_debug_retry(async_checkpointer: BaseCheckpointSaver):
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[str], operator.add]
|
||||
@@ -5875,6 +5929,55 @@ async def test_debug_retry(async_checkpointer: BaseCheckpointSaver):
|
||||
assert stream_parent_conf == history_parent_conf
|
||||
|
||||
|
||||
async def test_interrupt_after_streams_the_checkpoint_it_pauses_on(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
x: int
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("one", lambda state: {"x": 1})
|
||||
.add_node("two", lambda state: {"x": 2})
|
||||
.add_edge(START, "one")
|
||||
.add_edge("one", "two")
|
||||
.compile(checkpointer=async_checkpointer, interrupt_after=["one"])
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
checkpoints = [
|
||||
chunk
|
||||
async for chunk in graph.astream({"x": 0}, config, stream_mode="checkpoints")
|
||||
]
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
assert checkpoints[-1]["config"] == state.config
|
||||
assert checkpoints[-1]["next"] == ["two"]
|
||||
|
||||
|
||||
async def test_recursion_limit_streams_the_last_checkpoint_saved(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
x: int
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("loop", lambda state: {"x": state["x"] + 1})
|
||||
.add_edge(START, "loop")
|
||||
.add_edge("loop", "loop")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}, "recursion_limit": 3}
|
||||
checkpoints = []
|
||||
|
||||
with pytest.raises(GraphRecursionError):
|
||||
async for chunk in graph.astream({"x": 0}, config, stream_mode="checkpoints"):
|
||||
checkpoints.append(chunk)
|
||||
|
||||
assert checkpoints[-1]["config"] == (await graph.aget_state(config)).config
|
||||
|
||||
|
||||
async def test_debug_subgraphs(
|
||||
async_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
):
|
||||
@@ -6731,6 +6834,31 @@ async def test_multiple_updates() -> None:
|
||||
]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_task_in_subgraph_node_returns_its_value_with_tools_mode() -> None:
|
||||
class State(TypedDict):
|
||||
out: int | None
|
||||
|
||||
@task
|
||||
async def double(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
async def node(state: State) -> State:
|
||||
return {"out": await double(21)}
|
||||
|
||||
inner = StateGraph(State).add_node("node", node).add_edge(START, "node").compile()
|
||||
graph = (
|
||||
StateGraph(State).add_node("inner", inner).add_edge(START, "inner").compile()
|
||||
)
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in graph.astream({"out": None}, stream_mode=["values", "tools"])
|
||||
]
|
||||
|
||||
assert chunks[-1] == ("values", {"out": 42})
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_falsy_return_from_task(async_checkpointer: BaseCheckpointSaver) -> None:
|
||||
"""Test with a falsy return from a task."""
|
||||
@@ -7495,6 +7623,40 @@ async def test_entrypoint_from_async_generator() -> None:
|
||||
yield "b"
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("timeout", [None, 30])
|
||||
@pytest.mark.parametrize("version", ["v1", "v2"])
|
||||
async def test_awaited_task_returns_its_value_under_astream_events(
|
||||
version: Literal["v1", "v2"], timeout: int | None
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
doubled: int
|
||||
zeroed: int
|
||||
|
||||
@task(timeout=timeout)
|
||||
async def double(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
@task(timeout=timeout)
|
||||
async def zero() -> int:
|
||||
return 0
|
||||
|
||||
async def node(state: State) -> State:
|
||||
return {"doubled": await double(21), "zeroed": await zero()}
|
||||
|
||||
graph = StateGraph(State).add_node("node", node).add_edge(START, "node").compile()
|
||||
|
||||
outputs = [
|
||||
event["data"]["output"]
|
||||
async for event in graph.astream_events(
|
||||
{"doubled": 0, "zeroed": 1}, version=version
|
||||
)
|
||||
if event["event"] == "on_chain_end" and event["name"] == "node"
|
||||
]
|
||||
|
||||
assert outputs == [{"doubled": 42, "zeroed": 0}]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_named_tasks_functional() -> None:
|
||||
class Foo:
|
||||
@@ -9341,6 +9503,44 @@ 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,
|
||||
|
||||
@@ -16,9 +16,11 @@ 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
|
||||
@@ -863,6 +865,87 @@ 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()
|
||||
|
||||
@@ -2097,6 +2097,47 @@ def test_graph_error_handler_can_route_with_command():
|
||||
assert result["foo"] == "handled_next"
|
||||
|
||||
|
||||
# Each case runs the failing node on the executor path instead of the
|
||||
# single-task fast path.
|
||||
@pytest.mark.parametrize(
|
||||
("parallel", "stream_mode", "subgraphs"),
|
||||
[
|
||||
(False, "custom", False),
|
||||
(False, "messages", False),
|
||||
(False, "values", True),
|
||||
(True, "values", False),
|
||||
],
|
||||
)
|
||||
def test_graph_error_handler_recovers_on_the_executor_path(
|
||||
parallel: bool, stream_mode: str, subgraphs: bool
|
||||
) -> None:
|
||||
class State(TypedDict):
|
||||
x: str
|
||||
|
||||
def failing(state: State) -> State:
|
||||
raise ValueError("boom")
|
||||
|
||||
def handler(state: State) -> Command:
|
||||
return Command(update={"x": "handled"}, goto=END)
|
||||
|
||||
builder = (
|
||||
StateGraph(State)
|
||||
.add_node("failing", failing, error_handler=handler)
|
||||
.add_edge(START, "failing")
|
||||
)
|
||||
if parallel:
|
||||
builder.add_node("sibling", lambda state: {}).add_edge(START, "sibling")
|
||||
graph = builder.compile()
|
||||
|
||||
chunks = list(
|
||||
graph.stream(
|
||||
{"x": ""}, stream_mode=[stream_mode, "values"], subgraphs=subgraphs
|
||||
)
|
||||
)
|
||||
|
||||
assert [c[-1] for c in chunks if c[-2] == "values"][-1] == {"x": "handled"}
|
||||
|
||||
|
||||
def test_graph_error_handler_failure_fails_run():
|
||||
class State(TypedDict):
|
||||
foo: str
|
||||
|
||||
@@ -0,0 +1,339 @@
|
||||
"""Tests for `Runtime.on_exit`: cleanup that runs however a graph run ends."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.runtime import Runtime, get_runtime
|
||||
from langgraph.types import Command, interrupt
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
|
||||
|
||||
class State(TypedDict, total=False):
|
||||
value: str
|
||||
|
||||
|
||||
def _two_node_graph(first, second=None, *, checkpointer=None):
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("first", first)
|
||||
builder.add_edge(START, "first")
|
||||
if second is None:
|
||||
builder.add_edge("first", END)
|
||||
else:
|
||||
builder.add_node("second", second)
|
||||
builder.add_edge("first", "second")
|
||||
builder.add_edge("second", END)
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
# sync
|
||||
|
||||
|
||||
def test_runs_on_success_in_lifo_order() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup-a"))
|
||||
runtime.on_exit(lambda: events.append("cleanup-b"))
|
||||
events.append("first")
|
||||
return {"value": "x"}
|
||||
|
||||
def second(state: State) -> State:
|
||||
events.append("second")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first, second)
|
||||
assert graph.invoke({}) == {"value": "x"}
|
||||
assert events == ["first", "second", "cleanup-b", "cleanup-a"]
|
||||
|
||||
|
||||
def test_runs_on_error_and_keeps_the_run_exception() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def failing_cleanup() -> None:
|
||||
events.append("failing")
|
||||
raise KeyError("cleanup failed")
|
||||
|
||||
def first(state: State) -> State:
|
||||
runtime = get_runtime()
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
runtime.on_exit(failing_cleanup)
|
||||
raise ValueError("node failed")
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(ValueError, match="node failed"):
|
||||
graph.invoke({})
|
||||
# the failing callback did not stop the earlier one from running
|
||||
assert events == ["failing", "cleanup"]
|
||||
|
||||
|
||||
def test_callback_error_is_raised_when_the_run_succeeded() -> None:
|
||||
def failing_cleanup() -> None:
|
||||
raise KeyError("cleanup failed")
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(failing_cleanup)
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(KeyError, match="cleanup failed"):
|
||||
graph.invoke({})
|
||||
|
||||
|
||||
def test_runs_on_interrupt() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": interrupt("approve?")}
|
||||
|
||||
graph = _two_node_graph(first, checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
result = graph.invoke({}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert events == ["cleanup"]
|
||||
|
||||
# each resumed run is its own invocation with its own callbacks
|
||||
graph.invoke(Command(resume="ok"), config)
|
||||
assert events == ["cleanup", "cleanup"]
|
||||
|
||||
|
||||
def test_runs_when_the_stream_is_closed_early() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": "x"}
|
||||
|
||||
def second(state: State) -> State:
|
||||
events.append("second")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first, second)
|
||||
stream = graph.stream({}, stream_mode="updates")
|
||||
assert next(stream) == {"first": {"value": "x"}}
|
||||
assert events == []
|
||||
stream.close()
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
def test_runs_when_breaking_out_of_stream() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": "x"}
|
||||
|
||||
graph = _two_node_graph(first, lambda s: {})
|
||||
for _ in graph.stream({}, stream_mode="updates"):
|
||||
break
|
||||
# the dropped generator is closed by refcounting on CPython
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
def test_subgraph_registers_on_the_root_run() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def inner(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
events.append("inner")
|
||||
return {}
|
||||
|
||||
sub = StateGraph(State)
|
||||
sub.add_node("inner", inner)
|
||||
sub.add_edge(START, "inner")
|
||||
subgraph = sub.compile()
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
subgraph.invoke(state)
|
||||
events.append("subgraph returned")
|
||||
return {}
|
||||
|
||||
def after(state: State) -> State:
|
||||
events.append("after")
|
||||
return {}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("as_node", subgraph)
|
||||
builder.add_node("call_subgraph", call_subgraph)
|
||||
builder.add_node("after", after)
|
||||
builder.add_edge(START, "as_node")
|
||||
builder.add_edge("as_node", "call_subgraph")
|
||||
builder.add_edge("call_subgraph", "after")
|
||||
graph = builder.compile()
|
||||
|
||||
graph.invoke({})
|
||||
# nothing runs when either subgraph ends; both run when the root ends
|
||||
assert events == [
|
||||
"inner",
|
||||
"inner",
|
||||
"subgraph returned",
|
||||
"after",
|
||||
"cleanup",
|
||||
"cleanup",
|
||||
]
|
||||
|
||||
|
||||
def test_async_callback_rejected_in_sync_run() -> None:
|
||||
async def acleanup() -> None:
|
||||
pass
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(acleanup)
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(TypeError, match="async callback in a synchronous"):
|
||||
graph.invoke({})
|
||||
|
||||
|
||||
def test_outside_a_run_raises() -> None:
|
||||
with pytest.raises(RuntimeError, match="during a graph run"):
|
||||
Runtime().on_exit(lambda: None)
|
||||
|
||||
|
||||
def test_each_invocation_gets_its_own_callbacks() -> None:
|
||||
calls: list[int] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: calls.append(1))
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
graph.invoke({})
|
||||
graph.invoke({})
|
||||
assert calls == [1, 1]
|
||||
|
||||
|
||||
# async
|
||||
|
||||
|
||||
async def test_async_runs_on_success_sync_and_async_callbacks() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def acleanup() -> None:
|
||||
await asyncio.sleep(0)
|
||||
events.append("async-cleanup")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("sync-cleanup"))
|
||||
runtime.on_exit(acleanup)
|
||||
events.append("first")
|
||||
return {"value": "x"}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
assert await graph.ainvoke({}) == {"value": "x"}
|
||||
assert events == ["first", "async-cleanup", "sync-cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_on_error_and_keeps_the_run_exception() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def failing_cleanup() -> None:
|
||||
raise KeyError("cleanup failed")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
runtime.on_exit(failing_cleanup)
|
||||
raise ValueError("node failed")
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(ValueError, match="node failed"):
|
||||
await graph.ainvoke({})
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_on_interrupt() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": interrupt("approve?")}
|
||||
|
||||
graph = _two_node_graph(first, checkpointer=InMemorySaver())
|
||||
result = await graph.ainvoke({}, {"configurable": {"thread_id": "1"}})
|
||||
assert "__interrupt__" in result
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_on_cancel_after_the_node_unwinds() -> None:
|
||||
events: list[str] = []
|
||||
started = asyncio.Event()
|
||||
|
||||
async def acleanup() -> None:
|
||||
events.append("cleanup")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(acleanup)
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
finally:
|
||||
events.append("node finally")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
task = asyncio.create_task(graph.ainvoke({}))
|
||||
await started.wait()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
# cleanup ran after the node had stopped
|
||||
assert events == ["node finally", "cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_when_the_stream_is_closed_early() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def acleanup() -> None:
|
||||
events.append("cleanup")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(acleanup)
|
||||
return {"value": "x"}
|
||||
|
||||
async def second(state: State) -> State:
|
||||
events.append("second")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first, second)
|
||||
stream = graph.astream({}, stream_mode="updates")
|
||||
assert await stream.__anext__() == {"first": {"value": "x"}}
|
||||
assert events == []
|
||||
await stream.aclose()
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_async_subgraph_registers_on_the_root_run() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def inner(state: State) -> State:
|
||||
get_runtime().on_exit(lambda: events.append("cleanup"))
|
||||
return {}
|
||||
|
||||
sub = StateGraph(State)
|
||||
sub.add_node("inner", inner)
|
||||
sub.add_edge(START, "inner")
|
||||
subgraph = sub.compile()
|
||||
|
||||
async def call_subgraph(state: State) -> State:
|
||||
await subgraph.ainvoke(state)
|
||||
events.append("subgraph returned")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(call_subgraph, lambda s: events.append("after") or {})
|
||||
await graph.ainvoke({})
|
||||
assert events == ["subgraph returned", "after", "cleanup"]
|
||||
@@ -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"]
|
||||
@@ -13,13 +13,17 @@ from typing import Annotated, Any
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import tool
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.prebuilt import ToolNode, ToolRuntime
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.errors import GraphDrained
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.pregel._tools import _tool_call_writer
|
||||
from langgraph.runtime import RunControl
|
||||
from langgraph.types import Command, interrupt
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
@@ -54,14 +58,14 @@ def _caller_async(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1")
|
||||
return caller
|
||||
|
||||
|
||||
def _build_graph(caller, tools) -> Any:
|
||||
def _build_graph(caller, tools, checkpointer=None) -> Any:
|
||||
sg = StateGraph(_State)
|
||||
sg.add_node("caller", caller)
|
||||
sg.add_node("tools", ToolNode(tools))
|
||||
sg.add_edge(START, "caller")
|
||||
sg.add_edge("caller", "tools")
|
||||
sg.add_edge("tools", END)
|
||||
return sg.compile()
|
||||
return sg.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _tool_events(stream) -> list[tuple[tuple[str, ...], dict]]:
|
||||
@@ -146,6 +150,105 @@ class TestSyncGraphSyncTool:
|
||||
assert kinds == ["tool-started", "tool-error"]
|
||||
assert events[1][1]["message"] == "nope"
|
||||
|
||||
def test_interrupt_in_tool_is_not_a_tool_error(self) -> None:
|
||||
@tool
|
||||
def ask() -> str:
|
||||
"""asks."""
|
||||
return interrupt("ok?")
|
||||
|
||||
graph = _build_graph(_caller_sync("ask", {}), [ask], InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
paused = _tool_events(
|
||||
graph.stream(
|
||||
{"messages": []}, config, stream_mode=["tools"], subgraphs=True
|
||||
)
|
||||
)
|
||||
resumed = _tool_events(
|
||||
graph.stream(
|
||||
Command(resume="yes"), config, stream_mode=["tools"], subgraphs=True
|
||||
)
|
||||
)
|
||||
|
||||
assert [p["event"] for _, p in paused] == ["tool-started"]
|
||||
assert [(p["event"], p["tool_call_id"]) for _, p in resumed] == [
|
||||
("tool-started", "tc1"),
|
||||
("tool-finished", "tc1"),
|
||||
]
|
||||
|
||||
def test_tool_handing_off_to_the_parent_still_ends_its_call(self) -> None:
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
done: bool
|
||||
|
||||
def hand_off(state: State) -> Command:
|
||||
return Command(graph=Command.PARENT, goto="finish", update={"done": True})
|
||||
|
||||
inner = StateGraph(State).add_node("hand_off", hand_off)
|
||||
inner = inner.add_edge(START, "hand_off").compile()
|
||||
|
||||
@tool
|
||||
def handoff() -> str:
|
||||
"""hands off."""
|
||||
inner.invoke({"messages": [], "done": False})
|
||||
return "unreachable"
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("caller", _caller_sync("handoff", {}))
|
||||
.add_node("tools", ToolNode([handoff]))
|
||||
.add_node("finish", lambda state: {"done": True})
|
||||
.add_edge(START, "caller")
|
||||
.add_edge("caller", "tools")
|
||||
.add_edge("tools", END)
|
||||
.add_edge("finish", END)
|
||||
.compile()
|
||||
)
|
||||
events = [
|
||||
payload["event"]
|
||||
for _, payload in _tool_events(
|
||||
graph.stream(
|
||||
{"messages": [], "done": False},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
assert events[0] == "tool-started"
|
||||
assert events[-1] == "tool-error", events
|
||||
|
||||
def test_drained_tool_is_not_a_tool_error(self) -> None:
|
||||
control = RunControl()
|
||||
|
||||
def request_drain(state: _State) -> dict:
|
||||
control.request_drain()
|
||||
return {}
|
||||
|
||||
inner = (
|
||||
StateGraph(_State)
|
||||
.add_node("request_drain", request_drain)
|
||||
.add_node("after", lambda state: {})
|
||||
.add_edge(START, "request_drain")
|
||||
.add_edge("request_drain", "after")
|
||||
.compile()
|
||||
)
|
||||
|
||||
@tool
|
||||
def drained() -> str:
|
||||
"""drains."""
|
||||
inner.invoke({"messages": []}, control=control)
|
||||
return "unreachable"
|
||||
|
||||
graph = _build_graph(_caller_sync("drained", {}), [drained])
|
||||
chunks: list = []
|
||||
with pytest.raises(GraphDrained):
|
||||
for chunk in graph.stream(
|
||||
{"messages": []}, stream_mode=["tools"], subgraphs=True
|
||||
):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert [p["event"] for _, p in _tool_events(chunks)] == ["tool-started"]
|
||||
|
||||
def test_writer_unset_outside_tool(self) -> None:
|
||||
# Outside any tool body the ContextVar that ToolRuntime reads
|
||||
# is unset — emitting from there would be a no-op.
|
||||
@@ -196,6 +299,25 @@ class TestAsyncGraphAsyncTool:
|
||||
assert kinds == ["tool-started", "tool-output-delta", "tool-finished"]
|
||||
assert events[1][1]["delta"] == "hi"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interrupt_in_async_tool_is_not_a_tool_error(self) -> None:
|
||||
@tool
|
||||
async def aask() -> str:
|
||||
"""asks."""
|
||||
return interrupt("ok?")
|
||||
|
||||
graph = _build_graph(_caller_async("aask", {}), [aask], InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
events = []
|
||||
for graph_input in ({"messages": []}, Command(resume="yes")):
|
||||
async for ns, mode, payload in graph.astream(
|
||||
graph_input, config, stream_mode=["tools"], subgraphs=True
|
||||
):
|
||||
if mode == "tools":
|
||||
events.append(payload["event"])
|
||||
|
||||
assert events == ["tool-started", "tool-started", "tool-finished"]
|
||||
|
||||
|
||||
class TestConcurrentToolCalls:
|
||||
def test_parallel_tool_calls_do_not_bleed(self) -> None:
|
||||
|
||||
@@ -1,11 +1,24 @@
|
||||
"""End-to-end tests for `TracePolicy` input processing on node trace runs."""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph._internal._runnable import (
|
||||
ASYNCIO_ACCEPTS_CONTEXT,
|
||||
RunnableCallable,
|
||||
RunnableSeq,
|
||||
)
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
from langgraph.types import TracePolicy
|
||||
from tests.fake_tracer import FakeTracer, Run
|
||||
|
||||
@@ -147,3 +160,179 @@ async def test_trace_policy_transforms_recorded_inputs_async() -> None:
|
||||
run = _node_run(tracer, "n")
|
||||
assert run.inputs == {"scrubbed_in": True}
|
||||
assert run.outputs == {"value": 6}
|
||||
|
||||
|
||||
_RECORDING_POLICY = TracePolicy(
|
||||
process_inputs=lambda _: {"recorded": "inputs"},
|
||||
process_outputs=lambda _: {"recorded": "outputs"},
|
||||
)
|
||||
|
||||
|
||||
def _chat_graph() -> CompiledStateGraph:
|
||||
model = FakeMessagesListChatModel(
|
||||
responses=[AIMessage("model reply", id="ai-model")]
|
||||
)
|
||||
|
||||
def chat(state: MessagesState) -> dict[str, Any]:
|
||||
reply = model.invoke(state["messages"])
|
||||
return {
|
||||
"messages": [
|
||||
*state["messages"],
|
||||
reply,
|
||||
AIMessage("from node", id="ai-node"),
|
||||
]
|
||||
}
|
||||
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("chat", chat, trace_policy=_RECORDING_POLICY)
|
||||
.add_edge(START, "chat")
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
def _assert_scrubbed_node_run(tracer: FakeTracer, *, model_is_child: bool) -> None:
|
||||
run = _node_run(tracer, "chat")
|
||||
assert run.inputs == {"recorded": "inputs"}
|
||||
assert run.outputs == {"recorded": "outputs"}
|
||||
if model_is_child:
|
||||
assert [c.name for c in run.child_runs] == ["FakeMessagesListChatModel"]
|
||||
|
||||
|
||||
def test_trace_policy_scrubs_tracer_not_messages_stream() -> None:
|
||||
tracer = FakeTracer()
|
||||
streamed = [
|
||||
message.content
|
||||
for message, _ in _chat_graph().stream(
|
||||
{"messages": [HumanMessage("hello", id="h-1")]},
|
||||
{"callbacks": [tracer]},
|
||||
stream_mode="messages",
|
||||
)
|
||||
]
|
||||
|
||||
assert streamed == ["model reply", "from node"]
|
||||
_assert_scrubbed_node_run(tracer, model_is_child=True)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_trace_policy_scrubs_tracer_not_messages_stream_async() -> None:
|
||||
tracer = FakeTracer()
|
||||
streamed = [
|
||||
message.content
|
||||
async for message, _ in _chat_graph().astream(
|
||||
{"messages": [HumanMessage("hello", id="h-1")]},
|
||||
{"callbacks": [tracer]},
|
||||
stream_mode="messages",
|
||||
)
|
||||
]
|
||||
|
||||
assert streamed == ["model reply", "from node"]
|
||||
_assert_scrubbed_node_run(tracer, model_is_child=ASYNCIO_ACCEPTS_CONTEXT)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_trace_policy_scrubs_tracer_not_astream_events() -> None:
|
||||
tracer = FakeTracer()
|
||||
chat_events = [
|
||||
event
|
||||
async for event in _chat_graph().astream_events(
|
||||
{"messages": [HumanMessage("hello", id="h-1")]},
|
||||
{"callbacks": [tracer]},
|
||||
version="v2",
|
||||
)
|
||||
if event["name"] == "chat"
|
||||
and event["event"] in ("on_chain_start", "on_chain_end")
|
||||
]
|
||||
|
||||
assert [e["event"] for e in chat_events] == ["on_chain_start", "on_chain_end"]
|
||||
start, end = (e["data"] for e in chat_events)
|
||||
assert [m.content for m in start["input"]["messages"]] == ["hello"]
|
||||
assert [m.content for m in end["output"]["messages"]] == [
|
||||
"hello",
|
||||
"model reply",
|
||||
"from node",
|
||||
]
|
||||
_assert_scrubbed_node_run(tracer, model_is_child=ASYNCIO_ACCEPTS_CONTEXT)
|
||||
|
||||
|
||||
class _HandlerWithTapMethods(BaseCallbackHandler):
|
||||
def __init__(self) -> None:
|
||||
self.node_payloads: list[Any] = []
|
||||
self.node_run_ids: set[UUID] = set()
|
||||
|
||||
def _record(self, run_id: UUID, payload: Any) -> None:
|
||||
if run_id in self.node_run_ids:
|
||||
self.node_payloads.append(payload)
|
||||
|
||||
def tap_output_iter(self, run_id: UUID, output: Iterator[Any]) -> Iterator[Any]:
|
||||
for chunk in output:
|
||||
self._record(run_id, chunk)
|
||||
yield chunk
|
||||
|
||||
async def tap_output_aiter(
|
||||
self, run_id: UUID, output: AsyncIterator[Any]
|
||||
) -> AsyncIterator[Any]:
|
||||
async for chunk in output:
|
||||
self._record(run_id, chunk)
|
||||
yield chunk
|
||||
|
||||
def on_chain_start(
|
||||
self, serialized: Any, inputs: Any, *, run_id: UUID, **kwargs: Any
|
||||
) -> None:
|
||||
if kwargs.get("name") == "n":
|
||||
self.node_run_ids.add(run_id)
|
||||
self._record(run_id, inputs)
|
||||
|
||||
def on_chain_end(self, outputs: Any, *, run_id: UUID, **kwargs: Any) -> None:
|
||||
self._record(run_id, outputs)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("via_events", [False, True], ids=["ainvoke", "astream_events"])
|
||||
async def test_trace_policy_reaches_tracers_but_not_custom_handlers(
|
||||
via_events: bool,
|
||||
) -> None:
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("n", _incr, trace_policy=_RECORDING_POLICY)
|
||||
.add_edge(START, "n")
|
||||
.add_edge("n", END)
|
||||
.compile()
|
||||
)
|
||||
handler = _HandlerWithTapMethods()
|
||||
tracer = FakeTracer()
|
||||
config: RunnableConfig = {"callbacks": [handler, tracer]}
|
||||
|
||||
if via_events:
|
||||
async for _ in graph.astream_events({"value": 1}, config, version="v2"):
|
||||
pass
|
||||
else:
|
||||
assert await graph.ainvoke({"value": 1}, config) == {"value": 2}
|
||||
assert handler.node_payloads == [{"value": 1}, {"value": 2}]
|
||||
run = _node_run(tracer, "n")
|
||||
assert (run.inputs, run.outputs) == (
|
||||
{"recorded": "inputs"},
|
||||
{"recorded": "outputs"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_trace_policy_cancel_at_node_end_still_closes_tracer_run() -> None:
|
||||
def cancel_own_task(value: Any) -> Any:
|
||||
asyncio.current_task().cancel()
|
||||
return value
|
||||
|
||||
tracer = FakeTracer()
|
||||
seq = RunnableSeq(
|
||||
RunnableCallable(lambda x: x, trace=False),
|
||||
RunnableCallable(cancel_own_task, trace=False),
|
||||
trace_policy=_RECORDING_POLICY,
|
||||
)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.create_task(seq.ainvoke({}, {"callbacks": [tracer]}))
|
||||
for _ in range(100):
|
||||
if tracer.runs:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert [r.outputs for r in tracer.runs] == [{"recorded": "outputs"}]
|
||||
@@ -799,6 +799,15 @@ def test_ensure_config_tags_concat_preserves_order_and_duplicates() -> None:
|
||||
assert merged["tags"] == ["shared", "alpha", "shared", "beta"]
|
||||
|
||||
|
||||
def test_ensure_config_skips_tags_inherited_from_ambient() -> None:
|
||||
token = var_child_runnable_config.set({"tags": ["inherited"]})
|
||||
try:
|
||||
merged = ensure_config({"tags": ["inherited", "own"]})
|
||||
finally:
|
||||
var_child_runnable_config.reset(token)
|
||||
assert merged["tags"] == ["inherited", "own"]
|
||||
|
||||
|
||||
def test_ensure_config_merges_callbacks_across_configs() -> None:
|
||||
a_cb = _TrackingCB("a")
|
||||
b_cb = _TrackingCB("b")
|
||||
|
||||
Generated
+6
-6
@@ -1439,7 +1439,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.17"
|
||||
version = "1.2.19"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1625,7 +1625,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1722,7 +1722,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
@@ -3419,11 +3419,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "setuptools"
|
||||
version = "83.0.0"
|
||||
version = "84.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
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" }
|
||||
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" }
|
||||
wheels = [
|
||||
{ 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" },
|
||||
{ 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" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Generated
+3
-3
@@ -286,7 +286,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.17"
|
||||
version = "1.2.19"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -370,7 +370,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -467,7 +467,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -5,7 +5,7 @@ from langgraph_sdk.encryption.types import DecryptResult, EncryptionContext
|
||||
from langgraph_sdk.priority import Priority
|
||||
from langgraph_sdk.priority.types import PriorityContext, PrioritySettings
|
||||
|
||||
__version__ = "0.4.9"
|
||||
__version__ = "0.4.11"
|
||||
|
||||
__all__ = [
|
||||
"Auth",
|
||||
|
||||
@@ -14,9 +14,14 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import _araise_for_status_typed
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_araise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, aiter_lines_raw
|
||||
|
||||
@@ -146,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)
|
||||
@@ -176,13 +182,27 @@ class HttpClient:
|
||||
stacklevel=2,
|
||||
)
|
||||
await r.aclose()
|
||||
return await self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=request_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,
|
||||
@@ -218,71 +238,103 @@ class HttpClient:
|
||||
request_headers if reconnect_path is None else reconnect_headers
|
||||
)
|
||||
if last_event_id is not None:
|
||||
current_headers = {
|
||||
key: value
|
||||
for key, value in current_headers.items()
|
||||
if key.lower() != "last-event-id"
|
||||
}
|
||||
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
|
||||
|
||||
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_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}"
|
||||
)
|
||||
|
||||
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()
|
||||
try:
|
||||
async for line in aiter_lines_raw(res):
|
||||
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
# 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""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
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.
|
||||
yield sse
|
||||
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:
|
||||
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:
|
||||
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
|
||||
|
||||
|
||||
@@ -99,6 +99,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v1"] = "v1",
|
||||
) -> AsyncIterator[StreamPart]: ...
|
||||
|
||||
@@ -131,6 +132,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v2"],
|
||||
) -> AsyncIterator[StreamPartV2]: ...
|
||||
|
||||
@@ -147,6 +149,7 @@ class RunsClient:
|
||||
stream_resumable: bool = False,
|
||||
metadata: Mapping[str, Any] | None = None,
|
||||
config: Config | None = None,
|
||||
context: Context | None = None,
|
||||
checkpoint_during: bool | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -160,6 +163,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v1"] = "v1",
|
||||
) -> AsyncIterator[StreamPart]: ...
|
||||
|
||||
@@ -176,6 +180,7 @@ class RunsClient:
|
||||
stream_resumable: bool = False,
|
||||
metadata: Mapping[str, Any] | None = None,
|
||||
config: Config | None = None,
|
||||
context: Context | None = None,
|
||||
checkpoint_during: bool | None = None,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
@@ -189,6 +194,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v2"],
|
||||
) -> AsyncIterator[StreamPartV2]: ...
|
||||
|
||||
@@ -385,6 +391,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> Run: ...
|
||||
|
||||
@overload
|
||||
@@ -414,6 +421,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> Run: ...
|
||||
|
||||
async def create(
|
||||
@@ -647,6 +655,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> builtins.list[dict] | dict[str, Any]: ...
|
||||
|
||||
@overload
|
||||
@@ -673,6 +682,7 @@ class RunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> builtins.list[dict] | dict[str, Any]: ...
|
||||
|
||||
async def wait(
|
||||
|
||||
@@ -16,7 +16,6 @@ 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
|
||||
@@ -25,7 +24,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
|
||||
from langgraph_sdk._shared.utilities import _quote_path_param, _reconnect_delay
|
||||
from langgraph_sdk.schema import LangSmithTracing, QueryParamTypes
|
||||
from langgraph_sdk.stream.controller import _SeenEventIds
|
||||
from langgraph_sdk.stream.decoders import (
|
||||
@@ -1691,11 +1690,13 @@ class AsyncThreadStream:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt `attempt`."""
|
||||
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)
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._shared_reconnect_backoff_base,
|
||||
cap=self._shared_reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a post-ready transport drop.
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import tzinfo
|
||||
@@ -194,6 +195,15 @@ 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}
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from typing import Any, cast
|
||||
@@ -13,9 +14,14 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import _raise_for_status_typed
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_raise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, iter_lines_raw
|
||||
|
||||
@@ -146,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)
|
||||
@@ -176,13 +183,27 @@ class SyncHttpClient:
|
||||
stacklevel=2,
|
||||
)
|
||||
r.close()
|
||||
return self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=request_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,
|
||||
@@ -220,73 +241,105 @@ class SyncHttpClient:
|
||||
request_headers if reconnect_path is None else reconnect_headers
|
||||
)
|
||||
if last_event_id is not None:
|
||||
current_headers = {
|
||||
key: value
|
||||
for key, value in current_headers.items()
|
||||
if key.lower() != "last-event-id"
|
||||
}
|
||||
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
|
||||
|
||||
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_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}"
|
||||
)
|
||||
|
||||
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()
|
||||
try:
|
||||
for line in iter_lines_raw(res):
|
||||
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
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""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
if sse.event or decoder.has_data:
|
||||
# See async stream implementation for rationale on
|
||||
# skipping empty flush events.
|
||||
yield sse
|
||||
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:
|
||||
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:
|
||||
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
|
||||
|
||||
|
||||
def _encode_json(json: Any) -> tuple[dict[str, str], bytes]:
|
||||
def _encode_json(json: Any) -> tuple[dict[str, str], bytes | None]:
|
||||
if json is None:
|
||||
return {}, None
|
||||
body = orjson.dumps(
|
||||
json,
|
||||
_orjson_default,
|
||||
|
||||
@@ -80,6 +80,7 @@ class SyncRunsClient:
|
||||
command: Command | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
stream_resumable: bool = False,
|
||||
metadata: Mapping[str, Any] | None = None,
|
||||
config: Config | None = None,
|
||||
context: Context | None = None,
|
||||
@@ -98,6 +99,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v1"] = "v1",
|
||||
) -> Iterator[StreamPart]: ...
|
||||
|
||||
@@ -111,6 +113,7 @@ class SyncRunsClient:
|
||||
command: Command | None = None,
|
||||
stream_mode: StreamMode | Sequence[StreamMode] = "values",
|
||||
stream_subgraphs: bool = False,
|
||||
stream_resumable: bool = False,
|
||||
metadata: Mapping[str, Any] | None = None,
|
||||
config: Config | None = None,
|
||||
context: Context | None = None,
|
||||
@@ -129,6 +132,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v2"],
|
||||
) -> Iterator[StreamPartV2]: ...
|
||||
|
||||
@@ -159,6 +163,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v1"] = "v1",
|
||||
) -> Iterator[StreamPart]: ...
|
||||
|
||||
@@ -189,6 +194,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
version: Literal["v2"],
|
||||
) -> Iterator[StreamPartV2]: ...
|
||||
|
||||
@@ -380,6 +386,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> Run: ...
|
||||
|
||||
@overload
|
||||
@@ -409,6 +416,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> Run: ...
|
||||
|
||||
def create(
|
||||
@@ -642,6 +650,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> builtins.list[dict] | dict[str, Any]: ...
|
||||
|
||||
@overload
|
||||
@@ -668,6 +677,7 @@ class SyncRunsClient:
|
||||
headers: Mapping[str, str] | None = None,
|
||||
params: QueryParamTypes | None = None,
|
||||
on_run_created: Callable[[RunCreateMetadata], None] | None = None,
|
||||
durability: Durability | None = None,
|
||||
) -> builtins.list[dict] | dict[str, Any]: ...
|
||||
|
||||
def wait(
|
||||
@@ -815,7 +825,6 @@ class SyncRunsClient:
|
||||
"checkpoint_during": checkpoint_during,
|
||||
"on_completion": on_completion,
|
||||
"after_seconds": after_seconds,
|
||||
"raise_error": raise_error,
|
||||
"durability": durability,
|
||||
"langsmith_tracer": langsmith_tracing,
|
||||
}
|
||||
@@ -830,7 +839,7 @@ class SyncRunsClient:
|
||||
if thread_id is not None
|
||||
else "/runs/wait"
|
||||
)
|
||||
return self.http.request_reconnect(
|
||||
response = self.http.request_reconnect(
|
||||
endpoint,
|
||||
"POST",
|
||||
json={k: v for k, v in payload.items() if v is not None},
|
||||
@@ -838,6 +847,16 @@ class SyncRunsClient:
|
||||
headers=headers,
|
||||
on_response=on_response if on_run_created else None,
|
||||
)
|
||||
if (
|
||||
raise_error
|
||||
and isinstance(response, dict)
|
||||
and "__error__" in response
|
||||
and isinstance(response["__error__"], dict)
|
||||
):
|
||||
raise Exception(
|
||||
f"{response['__error__'].get('error')}: {response['__error__'].get('message')}"
|
||||
)
|
||||
return response
|
||||
|
||||
def list(
|
||||
self,
|
||||
|
||||
@@ -811,7 +811,7 @@ class _On:
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
resources: str | Sequence[str],
|
||||
resources: str | Sequence[str] | None = None,
|
||||
actions: str | Sequence[str] | None = None,
|
||||
) -> Callable[[AHO], AHO]: ...
|
||||
|
||||
|
||||
@@ -30,18 +30,14 @@ class BytesLineDecoder:
|
||||
# See https://docs.python.org/3/glossary.html#term-universal-newlines
|
||||
NEWLINE_CHARS = b"\n\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 []
|
||||
if self.trailing_cr and text.startswith(b"\n"):
|
||||
text = text[1:]
|
||||
self.trailing_cr = text.endswith(b"\r")
|
||||
|
||||
if not text:
|
||||
# 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
|
||||
return []
|
||||
|
||||
trailing_newline = text[-1] in NEWLINE_CHARS
|
||||
lines = cast(list[BytesLike], text.splitlines())
|
||||
@@ -66,7 +62,7 @@ class BytesLineDecoder:
|
||||
return lines
|
||||
|
||||
def flush(self) -> list[BytesLike]:
|
||||
if not self.buffer and not self.trailing_cr:
|
||||
if not self.buffer:
|
||||
return []
|
||||
|
||||
lines: list[BytesLike] = [self.buffer]
|
||||
@@ -76,40 +72,66 @@ class BytesLineDecoder:
|
||||
|
||||
|
||||
class SSEDecoder:
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, *, last_event_id: str | None = None) -> None:
|
||||
self._event = ""
|
||||
self._data = bytearray()
|
||||
self._last_event_id = ""
|
||||
self._last_event_id = last_event_id
|
||||
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."""
|
||||
"""Return the last event identifier that was seen.
|
||||
|
||||
return self._last_event_id or None
|
||||
`None` until an ID is supplied or received, and `""` after an empty
|
||||
one, which resets the reconnect cursor.
|
||||
"""
|
||||
|
||||
return self._last_event_id
|
||||
|
||||
def decode(self, line: bytes) -> StreamPart | None:
|
||||
# See: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation
|
||||
|
||||
if self._at_stream_start:
|
||||
self._at_stream_start = False
|
||||
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._last_event_id
|
||||
and not self._id_changed
|
||||
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=orjson.loads(self._data) if self._data else None, # ty: ignore[invalid-argument-type]
|
||||
id=self.last_event_id,
|
||||
data=data, # ty: ignore[invalid-argument-type]
|
||||
id=self._last_event_id or None,
|
||||
)
|
||||
|
||||
# NOTE: as per the SSE spec, do not reset last_event_id.
|
||||
self._event = ""
|
||||
self._data = bytearray()
|
||||
self._retry = None
|
||||
self._id_changed = False
|
||||
|
||||
return sse
|
||||
|
||||
@@ -122,14 +144,16 @@ class SSEDecoder:
|
||||
value = value[1:]
|
||||
|
||||
if fieldname == b"event":
|
||||
self._event = value.decode()
|
||||
self._event = value.decode(errors="replace")
|
||||
elif fieldname == b"data":
|
||||
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()
|
||||
self._last_event_id = value.decode(errors="replace")
|
||||
self._id_changed = True
|
||||
elif fieldname == b"retry":
|
||||
with contextlib.suppress(TypeError, ValueError):
|
||||
self._retry = int(value)
|
||||
|
||||
@@ -15,7 +15,6 @@ 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
|
||||
@@ -23,6 +22,7 @@ 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,11 +275,13 @@ class StreamController:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt *attempt*."""
|
||||
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)
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -4,7 +4,6 @@ from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
@@ -13,6 +12,7 @@ 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,11 +282,13 @@ class SyncStreamController:
|
||||
|
||||
def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff + jitter before a reconnect attempt."""
|
||||
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)
|
||||
time.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -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
|
||||
@@ -1,7 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -128,3 +130,45 @@ def test_sync_api_matches_async(async_cls, sync_cls):
|
||||
f"Return annotation mismatch for {async_cls.__name__}.{name}: "
|
||||
f"async={a_ret}, sync={s_ret}"
|
||||
)
|
||||
|
||||
|
||||
# Stateless overloads (`thread_id: None`) leave out the thread-only parameters,
|
||||
# and stateful ones (`thread_id: str`) leave out `on_completion`, which only
|
||||
# applies to stateless runs.
|
||||
_RUN_OVERLOAD_OMISSIONS = {
|
||||
"None": {"checkpoint", "checkpoint_id", "multitask_strategy"},
|
||||
"str": {"on_completion"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_cls", [RunsClient, SyncRunsClient])
|
||||
def test_run_overloads_accept_implementation_params(client_cls) -> None:
|
||||
tree = ast.parse(Path(inspect.getfile(client_cls)).read_text())
|
||||
cls = next(
|
||||
node
|
||||
for node in tree.body
|
||||
if isinstance(node, ast.ClassDef) and node.name == client_cls.__name__
|
||||
)
|
||||
defs = [
|
||||
node
|
||||
for node in cls.body
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
and node.name in {"stream", "create", "wait"}
|
||||
]
|
||||
impl_params = {
|
||||
node.name: {arg.arg for arg in node.args.kwonlyargs}
|
||||
for node in defs
|
||||
if not node.decorator_list
|
||||
}
|
||||
gaps = []
|
||||
for stub in defs:
|
||||
if not stub.decorator_list:
|
||||
continue
|
||||
thread_id = next(arg for arg in stub.args.args if arg.arg == "thread_id")
|
||||
assert thread_id.annotation is not None
|
||||
kind = ast.unparse(thread_id.annotation)
|
||||
expected = impl_params[stub.name] - _RUN_OVERLOAD_OMISSIONS[kind]
|
||||
if missing := expected - {arg.arg for arg in stub.args.kwonlyargs}:
|
||||
gaps.append(f"{stub.name}(thread_id: {kind}) line {stub.lineno}: {missing}")
|
||||
|
||||
assert not gaps, gaps
|
||||
@@ -1,8 +1,49 @@
|
||||
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()
|
||||
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator, Sequence
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import orjson
|
||||
import pytest
|
||||
from typing_extensions import assert_type
|
||||
|
||||
@@ -12,6 +13,7 @@ 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,
|
||||
@@ -110,6 +112,129 @@ def test_stream_sse():
|
||||
assert len(parts) == 79
|
||||
|
||||
|
||||
def test_sse_decoder_joins_data_lines_with_a_newline():
|
||||
decoder = SSEDecoder()
|
||||
for line in (b"data: [1,", b"data: 2]"):
|
||||
decoder.decode(line)
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=[1, 2], id=None) # ty: ignore[invalid-argument-type]
|
||||
|
||||
for line in (b"data: [1", b"data: 0]"):
|
||||
decoder.decode(line)
|
||||
with pytest.raises(orjson.JSONDecodeError):
|
||||
decoder.decode(b"")
|
||||
|
||||
|
||||
def test_sse_decoder_strips_a_leading_bom():
|
||||
decoder = SSEDecoder()
|
||||
for line in (b"\xef\xbb\xbfevent: values", b"data: 1"):
|
||||
decoder.decode(line)
|
||||
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 ---
|
||||
|
||||
|
||||
@@ -156,6 +281,7 @@ 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 = [
|
||||
@@ -221,6 +347,7 @@ 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 = [
|
||||
@@ -292,6 +419,431 @@ async def test_http_client_stream_recovers_after_disconnect():
|
||||
]
|
||||
|
||||
|
||||
def _drop_after(chunks: list[bytes], stream_cls: type, reconnects: list[httpx.Request]):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
if request.method == "POST":
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream", "Location": "/reconnect"},
|
||||
stream=stream_cls(
|
||||
chunks, httpx.RemoteProtocolError("incomplete chunked read")
|
||||
),
|
||||
)
|
||||
reconnects.append(request)
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
stream=stream_cls([b"event: end\ndata: null\n\n"]),
|
||||
)
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def _sync_reconnect(
|
||||
chunks: list[bytes], headers: dict[str, str] | None
|
||||
) -> httpx.Request:
|
||||
reconnects: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(_drop_after(chunks, ListByteStream, reconnects))
|
||||
with httpx.Client(transport=transport, base_url="https://example.com") as client:
|
||||
list(SyncHttpClient(client).stream("/stream", "POST", json={}, headers=headers))
|
||||
return reconnects[0]
|
||||
|
||||
|
||||
async def _async_reconnect(
|
||||
chunks: list[bytes], headers: dict[str, str] | None
|
||||
) -> httpx.Request:
|
||||
reconnects: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(
|
||||
_drop_after(chunks, AsyncListByteStream, reconnects)
|
||||
)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
[
|
||||
part
|
||||
async for part in HttpClient(client).stream(
|
||||
"/stream", "POST", json={}, headers=headers
|
||||
)
|
||||
]
|
||||
return reconnects[0]
|
||||
|
||||
|
||||
_EMPTY_IDS = pytest.mark.parametrize(
|
||||
"empty_id",
|
||||
[b"id:\nevent: values\ndata: 2\n\n", b"id:\n\n"],
|
||||
ids=["with-an-event", "alone"],
|
||||
)
|
||||
_CALLER_CURSORS = pytest.mark.parametrize(
|
||||
"headers",
|
||||
[None, {"Last-Event-ID": "5"}, {"last-event-id": "5"}],
|
||||
ids=["no-caller-cursor", "caller-cursor", "lowercase-caller-cursor"],
|
||||
)
|
||||
|
||||
|
||||
@_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
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
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"}
|
||||
)
|
||||
|
||||
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"}
|
||||
)
|
||||
|
||||
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"}
|
||||
)
|
||||
|
||||
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 ---
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from langgraph_sdk.client import (
|
||||
HttpClient,
|
||||
RunsClient,
|
||||
SyncHttpClient,
|
||||
SyncRunsClient,
|
||||
)
|
||||
|
||||
|
||||
class _BodyDropsMidRead(httpx.SyncByteStream, httpx.AsyncByteStream):
|
||||
def __iter__(self) -> Iterator[bytes]:
|
||||
yield b'{"partial'
|
||||
raise httpx.ReadError("connection dropped")
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
yield b'{"partial'
|
||||
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 handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if request.method == "POST":
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Location": "/threads/t1/runs/r1/join"},
|
||||
stream=_BodyDropsMidRead(),
|
||||
)
|
||||
return reconnects[len(requests) - 2](request)
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def _assert_bodyless_get(request: httpx.Request) -> None:
|
||||
assert request.method == "GET"
|
||||
assert request.url.path == "/threads/t1/runs/r1/join"
|
||||
assert "content-length" not in request.headers, request.headers
|
||||
assert "content-type" not in request.headers, request.headers
|
||||
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))
|
||||
with httpx.Client(transport=transport, base_url="https://example.com") as client:
|
||||
runs = SyncRunsClient(SyncHttpClient(client))
|
||||
with pytest.warns(UserWarning, match="attempting reconnect"):
|
||||
result = runs.wait("t1", "agent", input={"question": "hi"})
|
||||
|
||||
assert result == {"done": True}
|
||||
_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))
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
runs = RunsClient(HttpClient(client))
|
||||
with pytest.warns(UserWarning, match="attempting reconnect"):
|
||||
result = await runs.wait("t1", "agent", input={"question": "hi"})
|
||||
|
||||
assert result == {"done": True}
|
||||
_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] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
return httpx.Response(
|
||||
200, json={"__error__": {"error": "RuntimeError", "message": "boom"}}
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
with httpx.Client(transport=transport, base_url="https://example.com") as client:
|
||||
runs = SyncRunsClient(SyncHttpClient(client))
|
||||
with pytest.raises(Exception, match="RuntimeError: boom"):
|
||||
runs.wait("t1", "agent", input={}, raise_error=True)
|
||||
|
||||
assert "raise_error" not in json.loads(requests[0].content)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_wait_returns_the_run_error_without_raise_error(asynchronous):
|
||||
requests: list[httpx.Request] = []
|
||||
run_error = {"__error__": {"error": "RuntimeError", "message": "boom"}}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
return httpx.Response(200, json=run_error)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
if asynchronous:
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
result = await RunsClient(HttpClient(client)).wait(
|
||||
"t1", "agent", input={}, raise_error=False
|
||||
)
|
||||
else:
|
||||
with httpx.Client(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
result = SyncRunsClient(SyncHttpClient(client)).wait(
|
||||
"t1", "agent", input={}, raise_error=False
|
||||
)
|
||||
|
||||
assert result == run_error
|
||||
assert "raise_error" not in json.loads(requests[0].content)
|
||||
Generated
+2
-2
@@ -301,7 +301,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.17"
|
||||
version = "1.2.19"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -385,7 +385,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "4.3.0"
|
||||
version = "4.3.1"
|
||||
source = { editable = "../checkpoint" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Reference in new issue
Block a user