x

This PR adds validation to keys used in filtering logic in the SQLite store implementation.

docs/docs/cloud/reference/langgraph_server_changelog.md

@@ -4,6 +4,31 @@
 
 ---
 
+## v0.2.103 (2025-07-25)
+- Corrected the metadata endpoint to ensure accurate data retrieval.
+
+## v0.2.102 (2025-07-24)
+- Captured interrupt events in the wait method to preserve legacy behavior and stream updates by default.
+- Added support for SDK structlog in the JavaScript environment, enhancing logging capabilities.
+
+## v0.2.101 (2025-07-24)
+- Used the correct metadata endpoint for self-hosted environments, resolving an access issue.
+
+## v0.2.99 (2025-07-22)
+- Improved license validation by adding an in-memory cache and handling Redis connection errors more effectively.
+- Automatically remove agents from memory that are removed from `langgraph.json` to prevent persistence issues.
+- Ensured the UI namespace for generated UI is a valid JavaScript property name to prevent errors.
+- Raised a 422 error for improved request validation feedback.
+
+## v0.2.98 (2025-07-19)
+- Added langgraph node context for improved log filtering and trace visibility.
+
+## v0.2.97 (2025-07-19)
+- Fixed scheduling issue with ckpt ingestion worker that occurred on isolated background loops.
+- Ensured queue worker starts only after all migrations have completed.
+- Added more detailed error messages for thread state issues and improved response handling when state updates fail.
+- Exposed interrupt ID while retrieving thread state for enhanced API response details.
+
 ## v0.2.96 (2025-07-17)
 - Added a fallback mechanism for configurable header patterns to handle exclude/include settings more effectively.
 

docs/docs/concepts/langgraph_control_plane.md

@@ -119,6 +119,11 @@ These metrics are displayed as charts in the Control Plane UI.
 
 ### LangSmith Integration
 
-A [LangSmith](https://docs.smith.langchain.com/) tracing project is automatically created for each deployment. The tracing project has the same name as the deployment. When creating a deployment, the `LANGCHAIN_TRACING` and `LANGSMITH_API_KEY`/`LANGCHAIN_API_KEY` environment variables do not need to be specified; they are set automatically by the control plane.
+A [LangSmith](https://docs.smith.langchain.com/) tracing project and LangSmith API key are automatically created for each deployment. The deployment uses the API key to automatically send traces to LangSmith.
 
-When a deployment is deleted, the traces and the tracing project are not deleted.
+- The tracing project has the same name as the deployment.
+- The API key has the description `LangGraph Platform: <deployment_name>`.
+- The API key is never revealed and cannot be deleted manually.
+- When creating a deployment, the `LANGCHAIN_TRACING` and `LANGSMITH_API_KEY`/`LANGCHAIN_API_KEY` environment variables do not need to be specified; they are set automatically by the control plane.
+
+When a deployment is deleted, the traces and the tracing project are not deleted. However, the API will be deleted when the deployment is deleted.

libs/checkpoint-sqlite/langgraph/store/sqlite/base.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import concurrent.futures
 import datetime
 import logging
+import re
 import sqlite3
 import threading
 from collections import defaultdict
@@ -107,6 +108,23 @@ def _decode_ns_text(namespace: str) -> tuple[str, ...]:
     return tuple(namespace.split("."))
 
 
+def _validate_filter_key(key: str) -> None:
+    """Validate that a filter key is safe for use in SQL queries.
+
+    Args:
+        key: The filter key to validate
+
+    Raises:
+        ValueError: If the key contains invalid characters that could enable SQL injection
+    """
+    # Allow alphanumeric characters, underscores, dots, and hyphens
+    # This covers typical JSON property names while preventing SQL injection
+    if not re.match(r"^[a-zA-Z0-9_.-]+$", key):
+        raise ValueError(
+            f"Invalid filter key: '{key}'. Filter keys must contain only alphanumeric characters, underscores, dots, and hyphens."
+        )
+
+
 def _json_loads(content: bytes | str | orjson.Fragment) -> Any:
     if isinstance(content, orjson.Fragment):
         if hasattr(content, "buf"):
@@ -372,6 +390,8 @@ class BaseSqliteStore:
             filter_conditions = []
             if op.filter:
                 for key, value in op.filter.items():
+                    _validate_filter_key(key)
+
                     if isinstance(value, dict):
                         for op_name, val in value.items():
                             condition, filter_params_ = self._get_filter_condition(
@@ -622,6 +642,8 @@ class BaseSqliteStore:
 
     def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, list]:
         """Helper to generate filter conditions."""
+        _validate_filter_key(key)
+
         # We need to properly format values for SQLite JSON extraction comparison
         if op == "$eq":
             if isinstance(value, str):
@@ -858,6 +880,8 @@ class SqliteStore(BaseSqliteStore, BaseStore):
 
     def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, list]:
         """Helper to generate filter conditions."""
+        _validate_filter_key(key)
+
         # We need to properly format values for SQLite JSON extraction comparison
         if op == "$eq":
             if isinstance(value, str):

libs/checkpoint-sqlite/tests/test_store.py

@@ -1047,3 +1047,23 @@ def test_search_items(
         for ns in test_namespaces:
             key = f"item_{ns[-1]}"
             store.delete(ns, key)
+
+
+def test_sql_injection_vulnerability(store: SqliteStore) -> None:
+    """Test that SQL injection via malicious filter keys is prevented."""
+    # Add public and private documents
+    store.put(("docs",), "public", {"access": "public", "data": "public info"})
+    store.put(
+        ("docs",), "private", {"access": "private", "data": "secret", "password": "123"}
+    )
+
+    # Normal query - returns 1 public document
+    normal = store.search(("docs",), filter={"access": "public"})
+    assert len(normal) == 1
+    assert normal[0].value["access"] == "public"
+
+    # SQL injection attempt via malicious key should raise ValueError
+    malicious_key = "access') = 'public' OR '1'='1' OR json_extract(value, '$."
+
+    with pytest.raises(ValueError, match="Invalid filter key"):
+        store.search(("docs",), filter={malicious_key: "dummy"})

libs/cli/langgraph_cli/cli.py

@@ -153,6 +153,12 @@ OPT_POSTGRES_URI = click.option(
     help="Postgres URI to use for the database. Defaults to launching a local database",
 )
 
+OPT_API_VERSION = click.option(
+    "--api-version",
+    type=str,
+    help="API server version to use for the base image. If unspecified, the latest version will be used.",
+)
+
 
 @click.group()
 @click.version_option(version=__version__, prog_name="LangGraph CLI")
@@ -170,6 +176,7 @@ def cli():
 @OPT_DEBUGGER_BASE_URL
 @OPT_WATCH
 @OPT_POSTGRES_URI
+@OPT_API_VERSION
 @click.option(
     "--image",
     type=str,
@@ -203,6 +210,7 @@ def up(
     debugger_port: Optional[int],
     debugger_base_url: Optional[str],
     postgres_uri: Optional[str],
+    api_version: Optional[str],
     image: Optional[str],
     base_image: Optional[str],
 ):
@@ -225,6 +233,7 @@ For production use, requires a license key in env var LANGGRAPH_CLOUD_LICENSE_KE
             debugger_port=debugger_port,
             debugger_base_url=debugger_base_url,
             postgres_uri=postgres_uri,
+            api_version=api_version,
             image=image,
             base_image=base_image,
         )
@@ -290,6 +299,7 @@ def _build(
     config: pathlib.Path,
     config_json: dict,
     base_image: Optional[str],
+    api_version: Optional[str],
     pull: bool,
     tag: str,
     passthrough: Sequence[str] = (),
@@ -300,7 +310,7 @@ def _build(
             subp_exec(
                 "docker",
                 "pull",
-                langgraph_cli.config.docker_tag(config_json, base_image),
+                langgraph_cli.config.docker_tag(config_json, base_image, api_version),
                 verbose=True,
             )
         )
@@ -314,7 +324,7 @@ def _build(
     ]
     # apply config
     stdin, additional_contexts = langgraph_cli.config.config_to_docker(
-        config, config_json, base_image
+        config, config_json, base_image, api_version
     )
     # add additional_contexts
     if additional_contexts:
@@ -355,6 +365,7 @@ def _build(
     "\n\n    \b\nExamples:\n    --base-image langchain/langgraph-server:0.2.18  # Pin to a specific patch version"
     "\n    --base-image langchain/langgraph-server:0.2  # Pin to a minor version (Python)",
 )
+@OPT_API_VERSION
 @click.argument("docker_build_args", nargs=-1, type=click.UNPROCESSED)
 @cli.command(
     help="📦 Build LangGraph API server Docker image.",
@@ -367,6 +378,7 @@ def build(
     config: pathlib.Path,
     docker_build_args: Sequence[str],
     base_image: Optional[str],
+    api_version: Optional[str],
     pull: bool,
     tag: str,
 ):
@@ -376,7 +388,15 @@ def build(
         config_json = langgraph_cli.config.validate_config_file(config)
         warn_non_wolfi_distro(config_json)
         _build(
-            runner, set, config, config_json, base_image, pull, tag, docker_build_args
+            runner,
+            set,
+            config,
+            config_json,
+            base_image,
+            api_version,
+            pull,
+            tag,
+            docker_build_args,
         )
 
 
@@ -456,12 +476,14 @@ tests
     "\n\n    \b\nExamples:\n    --base-image langchain/langgraph-server:0.2.18  # Pin to a specific patch version"
     "\n    --base-image langchain/langgraph-server:0.2  # Pin to a minor version (Python)",
 )
+@OPT_API_VERSION
 @log_command
 def dockerfile(
     save_path: str,
     config: pathlib.Path,
     add_docker_compose: bool,
     base_image: Optional[str] = None,
+    api_version: Optional[str] = None,
 ) -> None:
     save_path = pathlib.Path(save_path).absolute()
     secho(f"🔍 Validating configuration at path: {config}", fg="yellow")
@@ -474,6 +496,7 @@ def dockerfile(
         config,
         config_json,
         base_image=base_image,
+        api_version=api_version,
     )
     with open(str(save_path), "w", encoding="utf-8") as f:
         f.write(dockerfile)
@@ -739,6 +762,7 @@ def prepare_args_and_stdin(
     debugger_port: Optional[int] = None,
     debugger_base_url: Optional[str] = None,
     postgres_uri: Optional[str] = None,
+    api_version: Optional[str] = None,
     # Like "my-tag" (if you already built it locally)
     image: Optional[str] = None,
     # Like "langchain/langgraphjs-api" or "langchain/langgraph-api
@@ -754,6 +778,7 @@ def prepare_args_and_stdin(
         postgres_uri=postgres_uri,
         image=image,  # Pass image to compose YAML generator
         base_image=base_image,
+        api_version=api_version,
     )
     args = [
         "--project-directory",
@@ -769,6 +794,7 @@ def prepare_args_and_stdin(
         config,
         watch=watch,
         base_image=langgraph_cli.config.default_base_image(config),
+        api_version=api_version,
         image=image,
     )
     return args, stdin
@@ -787,6 +813,7 @@ def prepare(
     debugger_port: Optional[int] = None,
     debugger_base_url: Optional[str] = None,
     postgres_uri: Optional[str] = None,
+    api_version: Optional[str] = None,
     image: Optional[str] = None,
     base_image: Optional[str] = None,
 ) -> tuple[list[str], str]:
@@ -799,7 +826,7 @@ def prepare(
             subp_exec(
                 "docker",
                 "pull",
-                langgraph_cli.config.docker_tag(config_json, base_image),
+                langgraph_cli.config.docker_tag(config_json, base_image, api_version),
                 verbose=verbose,
             )
         )
@@ -814,6 +841,7 @@ def prepare(
         debugger_port=debugger_port,
         debugger_base_url=debugger_base_url or f"http://127.0.0.1:{port}",
         postgres_uri=postgres_uri,
+        api_version=api_version,
         image=image,
         base_image=base_image,
     )

libs/cli/langgraph_cli/config.py

@@ -1213,6 +1213,7 @@ def python_config_to_docker(
     config_path: pathlib.Path,
     config: Config,
     base_image: str,
+    api_version: Optional[str] = None,
 ) -> tuple[str, dict[str, str]]:
     """Generate a Dockerfile from the configuration."""
     pip_installer = config.get("pip_installer", "auto")
@@ -1360,7 +1361,7 @@ ADD {relpath} /deps/{name}
                 "# -- End of JS dependencies install --",
             ]
         )
-    image_str = docker_tag(config, base_image)
+    image_str = docker_tag(config, base_image, api_version)
     docker_file_contents = [
         f"FROM {image_str}",
         "",
@@ -1402,10 +1403,11 @@ def node_config_to_docker(
     config_path: pathlib.Path,
     config: Config,
     base_image: str,
+    api_version: Optional[str] = None,
 ) -> tuple[str, dict[str, str]]:
     faux_path = f"/deps/{config_path.parent.name}"
     install_cmd = _get_node_pm_install_cmd(config_path, config)
-    image_str = docker_tag(config, base_image)
+    image_str = docker_tag(config, base_image, api_version)
 
     env_vars: list[str] = []
 
@@ -1461,6 +1463,7 @@ def default_base_image(config: Config) -> str:
 def docker_tag(
     config: Config,
     base_image: Optional[str] = None,
+    api_version: Optional[str] = None,
 ) -> str:
     base_image = base_image or default_base_image(config)
 
@@ -1473,28 +1476,43 @@ def docker_tag(
     if "/langgraph-server" in base_image:
         return f"{base_image}-py{config['python_version']}"
 
+    # Build the standard tag format
+    language, version = None, None
     if config.get("node_version") and not config.get("python_version"):
-        return f"{base_image}:{config['node_version']}{distro_tag}"
-    return f"{base_image}:{config['python_version']}{distro_tag}"
+        language, version = "node", config["node_version"]
+    else:
+        language, version = "py", config["python_version"]
+
+    version_distro_tag = f"{version}{distro_tag}"
+
+    # Prepend API version if provided
+    if api_version:
+        full_tag = f"{api_version}-{language}{version_distro_tag}"
+    else:
+        full_tag = version_distro_tag
+
+    return f"{base_image}:{full_tag}"
 
 
 def config_to_docker(
     config_path: pathlib.Path,
     config: Config,
     base_image: Optional[str] = None,
+    api_version: Optional[str] = None,
 ) -> tuple[str, dict[str, str]]:
     base_image = base_image or default_base_image(config)
 
     if config.get("node_version") and not config.get("python_version"):
-        return node_config_to_docker(config_path, config, base_image)
+        return node_config_to_docker(config_path, config, base_image, api_version)
 
-    return python_config_to_docker(config_path, config, base_image)
+    return python_config_to_docker(config_path, config, base_image, api_version)
 
 
 def config_to_compose(
     config_path: pathlib.Path,
     config: Config,
     base_image: Optional[str] = None,
+    api_version: Optional[str] = None,
     image: Optional[str] = None,
     watch: bool = False,
 ) -> str:
@@ -1531,7 +1549,7 @@ def config_to_compose(
 
     else:
         dockerfile, additional_contexts = config_to_docker(
-            config_path, config, base_image
+            config_path, config, base_image, api_version
         )
 
         additional_contexts_str = "\n".join(

libs/cli/langgraph_cli/docker.py

@@ -147,6 +147,8 @@ def compose_as_dict(
     image: Optional[str] = None,
     # Base image to use for the LangGraph API server
     base_image: Optional[str] = None,
+    # API version of the base image
+    api_version: Optional[str] = None,
 ) -> dict:
     """Create a docker compose file as a dictionary in YML style."""
     if postgres_uri is None:
@@ -252,6 +254,7 @@ def compose(
     postgres_uri: Optional[str] = None,
     image: Optional[str] = None,
     base_image: Optional[str] = None,
+    api_version: Optional[str] = None,
 ) -> str:
     """Create a docker compose file as a string."""
     compose_content = compose_as_dict(
@@ -262,6 +265,7 @@ def compose(
         postgres_uri=postgres_uri,
         image=image,
         base_image=base_image,
+        api_version=api_version,
     )
     compose_str = dict_to_yaml(compose_content)
     return compose_str

libs/cli/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-cli"
-version = "0.3.5"
+version = "0.3.6"
 description = "CLI for interacting with LangGraph API"
 authors = []
 requires-python = ">=3.9"

libs/cli/tests/unit_tests/cli/test_cli.py

@@ -574,3 +574,248 @@ def test_build_generate_proper_build_context():
         assert len(build_contexts) == 2, (
             f"Expected 2 build contexts, but found {len(build_contexts)}"
         )
+
+
+def test_dockerfile_command_with_api_version() -> None:
+    """Test the 'dockerfile' command with --api-version flag."""
+    runner = CliRunner()
+    config_content = {
+        "python_version": "3.11",
+        "graphs": {"agent": "agent.py:graph"},
+        "dependencies": ["."],
+    }
+
+    with temporary_config_folder(config_content) as temp_dir:
+        save_path = temp_dir / "Dockerfile"
+        agent_path = temp_dir / "agent.py"
+        agent_path.touch()
+
+        result = runner.invoke(
+            cli,
+            [
+                "dockerfile",
+                str(save_path),
+                "--config",
+                str(temp_dir / "config.json"),
+                "--api-version",
+                "0.2.74",
+            ],
+        )
+
+        # Assert command was successful
+        assert result.exit_code == 0, result.output
+        assert "✅ Created: Dockerfile" in result.output
+
+        # Check if Dockerfile was created and contains correct FROM line
+        assert save_path.exists()
+        with open(save_path) as f:
+            dockerfile = f.read()
+            assert "FROM langchain/langgraph-api:0.2.74-py3.11" in dockerfile
+
+
+def test_dockerfile_command_with_api_version_and_base_image() -> None:
+    """Test the 'dockerfile' command with both --api-version and --base-image flags."""
+    runner = CliRunner()
+    config_content = {
+        "python_version": "3.12",
+        "graphs": {"agent": "agent.py:graph"},
+        "dependencies": ["."],
+        "image_distro": "wolfi",
+    }
+
+    with temporary_config_folder(config_content) as temp_dir:
+        save_path = temp_dir / "Dockerfile"
+        agent_path = temp_dir / "agent.py"
+        agent_path.touch()
+
+        result = runner.invoke(
+            cli,
+            [
+                "dockerfile",
+                str(save_path),
+                "--config",
+                str(temp_dir / "config.json"),
+                "--api-version",
+                "1.0.0",
+                "--base-image",
+                "my-registry/custom-api",
+            ],
+        )
+
+        # Assert command was successful
+        assert result.exit_code == 0, result.output
+        assert "✅ Created: Dockerfile" in result.output
+
+        # Check if Dockerfile was created and contains correct FROM line
+        assert save_path.exists()
+        with open(save_path) as f:
+            dockerfile = f.read()
+            assert "FROM my-registry/custom-api:1.0.0-py3.12-wolfi" in dockerfile
+
+
+def test_dockerfile_command_with_api_version_nodejs() -> None:
+    """Test the 'dockerfile' command with --api-version flag for Node.js config."""
+    runner = CliRunner()
+    config_content = {
+        "node_version": "20",
+        "graphs": {"agent": "agent.js:graph"},
+    }
+
+    with temporary_config_folder(config_content) as temp_dir:
+        save_path = temp_dir / "Dockerfile"
+        agent_path = temp_dir / "agent.js"
+        agent_path.touch()
+
+        result = runner.invoke(
+            cli,
+            [
+                "dockerfile",
+                str(save_path),
+                "--config",
+                str(temp_dir / "config.json"),
+                "--api-version",
+                "0.2.74",
+            ],
+        )
+
+        # Assert command was successful
+        assert result.exit_code == 0, result.output
+        assert "✅ Created: Dockerfile" in result.output
+
+        # Check if Dockerfile was created and contains correct FROM line
+        assert save_path.exists()
+        with open(save_path) as f:
+            dockerfile = f.read()
+            assert "FROM langchain/langgraphjs-api:0.2.74-node20" in dockerfile
+
+
+def test_build_command_with_api_version() -> None:
+    """Test the 'build' command with --api-version flag."""
+    runner = CliRunner()
+    config_content = {
+        "python_version": "3.11",
+        "graphs": {"agent": "agent.py:graph"},
+        "dependencies": ["."],
+        "image_distro": "wolfi",  # Use wolfi to avoid warning messages
+    }
+
+    with temporary_config_folder(config_content) as temp_dir:
+        agent_path = temp_dir / "agent.py"
+        agent_path.touch()
+
+        # Mock docker command since we don't want to actually build
+        with runner.isolated_filesystem():
+            result = runner.invoke(
+                cli,
+                [
+                    "build",
+                    "--tag",
+                    "test-image",
+                    "--config",
+                    str(temp_dir / "config.json"),
+                    "--api-version",
+                    "0.2.74",
+                    "--no-pull",  # Avoid pulling non-existent images
+                ],
+                catch_exceptions=True,
+            )
+
+        # Check that the build command is called with the correct tag
+        # The output should contain the docker build command with the api_version tag
+        assert "langchain/langgraph-api:0.2.74-py3.11-wolfi" in result.output
+
+
+def test_build_command_with_api_version_and_base_image() -> None:
+    """Test the 'build' command with both --api-version and --base-image flags."""
+    runner = CliRunner()
+    config_content = {
+        "python_version": "3.12",
+        "graphs": {"agent": "agent.py:graph"},
+        "dependencies": ["."],
+        "image_distro": "wolfi",  # Use wolfi to avoid warning messages
+    }
+
+    with temporary_config_folder(config_content) as temp_dir:
+        agent_path = temp_dir / "agent.py"
+        agent_path.touch()
+
+        # Mock docker command since we don't want to actually build
+        with runner.isolated_filesystem():
+            result = runner.invoke(
+                cli,
+                [
+                    "build",
+                    "--tag",
+                    "test-image",
+                    "--config",
+                    str(temp_dir / "config.json"),
+                    "--api-version",
+                    "1.0.0",
+                    "--base-image",
+                    "my-registry/custom-api",
+                    "--no-pull",  # Avoid pulling non-existent images
+                ],
+                catch_exceptions=True,
+            )
+
+        # Check that the build command includes the api_version
+        assert "my-registry/custom-api:1.0.0-py3.12-wolfi" in result.output
+
+
+def test_prepare_args_and_stdin_with_api_version() -> None:
+    """Test prepare_args_and_stdin function with api_version parameter."""
+    config_path = pathlib.Path(__file__).parent / "langgraph.json"
+    config = validate_config(
+        Config(dependencies=["."], graphs={"agent": "agent.py:graph"})
+    )
+    port = 8000
+    api_version = "0.2.74"
+
+    actual_args, actual_stdin = prepare_args_and_stdin(
+        capabilities=DEFAULT_DOCKER_CAPABILITIES,
+        config_path=config_path,
+        config=config,
+        docker_compose=None,
+        port=port,
+        watch=False,
+        api_version=api_version,
+    )
+
+    expected_args = [
+        "--project-directory",
+        str(pathlib.Path(__file__).parent.absolute()),
+        "-f",
+        "-",
+    ]
+
+    # Check that the args are correct
+    assert actual_args == expected_args
+
+    # Check that the stdin contains the correct FROM line with api_version
+    assert "FROM langchain/langgraph-api:0.2.74-py3.11" in actual_stdin
+
+
+def test_prepare_args_and_stdin_with_api_version_and_image() -> None:
+    """Test prepare_args_and_stdin function with both api_version and image parameters."""
+    config_path = pathlib.Path(__file__).parent / "langgraph.json"
+    config = validate_config(
+        Config(dependencies=["."], graphs={"agent": "agent.py:graph"})
+    )
+    port = 8000
+    api_version = "0.2.74"
+    image = "my-custom-image:latest"
+
+    actual_args, actual_stdin = prepare_args_and_stdin(
+        capabilities=DEFAULT_DOCKER_CAPABILITIES,
+        config_path=config_path,
+        config=config,
+        docker_compose=None,
+        port=port,
+        watch=False,
+        api_version=api_version,
+        image=image,
+    )
+
+    # When image is provided, api_version should be ignored for the image
+    # but the stdin should not contain a build section (since image is provided)
+    assert "pull_policy: build" not in actual_stdin

libs/cli/tests/unit_tests/test_config.py

@@ -1337,3 +1337,195 @@ def test_docker_tag_different_node_versions_with_distro():
         )
         tag = docker_tag(config)
         assert tag == expected_tag, f"Failed for Node.js {node_version}"
+
+
+def test_docker_tag_with_api_version():
+    """Test docker_tag function with api_version parameter."""
+
+    # Test 1: Python config with api_version and default distro
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+        }
+    )
+    tag = docker_tag(config, api_version="0.2.74")
+    assert tag == "langchain/langgraph-api:0.2.74-py3.11"
+
+    # Test 2: Python config with api_version and wolfi distro
+    config = validate_config(
+        {
+            "python_version": "3.12",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+            "image_distro": "wolfi",
+        }
+    )
+    tag = docker_tag(config, api_version="0.2.74")
+    assert tag == "langchain/langgraph-api:0.2.74-py3.12-wolfi"
+
+    # Test 3: Node.js config with api_version and default distro
+    config = validate_config(
+        {
+            "node_version": "20",
+            "graphs": {"agent": "./agent.js:graph"},
+        }
+    )
+    tag = docker_tag(config, api_version="0.2.74")
+    assert tag == "langchain/langgraphjs-api:0.2.74-node20"
+
+    # Test 4: Node.js config with api_version and wolfi distro
+    config = validate_config(
+        {
+            "node_version": "20",
+            "graphs": {"agent": "./agent.js:graph"},
+            "image_distro": "wolfi",
+        }
+    )
+    tag = docker_tag(config, api_version="0.2.74")
+    assert tag == "langchain/langgraphjs-api:0.2.74-node20-wolfi"
+
+    # Test 5: Custom base image with api_version
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+            "base_image": "my-registry/custom-image",
+        }
+    )
+    tag = docker_tag(config, base_image="my-registry/custom-image", api_version="1.0.0")
+    assert tag == "my-registry/custom-image:1.0.0-py3.11"
+
+    # Test 6: api_version with different Python versions
+    for python_version in ["3.11", "3.12", "3.13"]:
+        config = validate_config(
+            {
+                "python_version": python_version,
+                "dependencies": ["."],
+                "graphs": {"agent": "./agent.py:graph"},
+            }
+        )
+        tag = docker_tag(config, api_version="0.2.74")
+        assert tag == f"langchain/langgraph-api:0.2.74-py{python_version}"
+
+    # Test 7: Without api_version should work as before
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+        }
+    )
+    tag = docker_tag(config)
+    assert tag == "langchain/langgraph-api:3.11"
+
+    # Test 8: api_version with multiplatform config (should default to Python)
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "node_version": "20",
+            "dependencies": ["."],
+            "graphs": {"python": "./agent.py:graph", "js": "./agent.js:graph"},
+        }
+    )
+    tag = docker_tag(config, api_version="0.2.74")
+    assert tag == "langchain/langgraph-api:0.2.74-py3.11"
+
+    # Test 9: api_version with _INTERNAL_docker_tag should ignore api_version
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+            "_INTERNAL_docker_tag": "internal-tag",
+        }
+    )
+    tag = docker_tag(config, api_version="0.2.74")
+    assert tag == "langchain/langgraph-api:internal-tag"
+
+    # Test 10: api_version with langgraph-server base image should follow special format
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+        }
+    )
+    tag = docker_tag(
+        config, base_image="langchain/langgraph-server:0.2", api_version="0.2.74"
+    )
+    assert tag == "langchain/langgraph-server:0.2-py3.11"
+
+
+def test_config_to_docker_with_api_version():
+    """Test config_to_docker function with api_version parameter."""
+
+    # Test Python config with api_version
+    graphs = {"agent": "./agent.py:graph"}
+    actual_docker_stdin, additional_contexts = config_to_docker(
+        PATH_TO_CONFIG,
+        validate_config({"dependencies": ["."], "graphs": graphs}),
+        "langchain/langgraph-api",
+        api_version="0.2.74",
+    )
+
+    # Check that the FROM line uses the api_version
+    lines = actual_docker_stdin.split("\n")
+    from_line = lines[0]
+    assert from_line == "FROM langchain/langgraph-api:0.2.74-py3.11"
+
+    # Test Node.js config with api_version
+    graphs = {"agent": "./agent.js:graph"}
+    actual_docker_stdin, additional_contexts = config_to_docker(
+        PATH_TO_CONFIG,
+        validate_config({"node_version": "20", "graphs": graphs}),
+        "langchain/langgraphjs-api",
+        api_version="0.2.74",
+    )
+
+    # Check that the FROM line uses the api_version
+    lines = actual_docker_stdin.split("\n")
+    from_line = lines[0]
+    assert from_line == "FROM langchain/langgraphjs-api:0.2.74-node20"
+
+
+def test_config_to_compose_with_api_version():
+    """Test config_to_compose function with api_version parameter."""
+
+    # Test Python config with api_version
+    config = validate_config(
+        {
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+        }
+    )
+
+    actual_compose_str = config_to_compose(
+        PATH_TO_CONFIG,
+        config,
+        "langchain/langgraph-api",
+        api_version="0.2.74",
+    )
+
+    # Check that the compose file includes the correct FROM line with api_version
+    assert "FROM langchain/langgraph-api:0.2.74-py3.11" in actual_compose_str
+
+    # Test Node.js config with api_version
+    config = validate_config(
+        {
+            "node_version": "20",
+            "graphs": {"agent": "./agent.js:graph"},
+        }
+    )
+
+    actual_compose_str = config_to_compose(
+        PATH_TO_CONFIG,
+        config,
+        "langchain/langgraphjs-api",
+        api_version="0.2.74",
+    )
+
+    # Check that the compose file includes the correct FROM line with api_version
+    assert "FROM langchain/langgraphjs-api:0.2.74-node20" in actual_compose_str

libs/cli/tests/unit_tests/test_docker.py

@@ -146,3 +146,220 @@ services:
             REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {DEFAULT_POSTGRES_URI}"""
     assert clean_empty_lines(actual_compose_str) == expected_compose_str
+
+
+def test_compose_with_api_version():
+    """Test compose function with api_version parameter."""
+    port = 8123
+    api_version = "0.2.74"
+
+    actual_compose_str = compose(
+        DEFAULT_DOCKER_CAPABILITIES, port=port, api_version=api_version
+    )
+
+    # The compose function should generate a compose file that doesn't directly
+    # reference the api_version, since it's handled in the docker tag creation
+    # when building the image. The compose function mainly sets up services.
+    expected_compose_str = f"""volumes:
+    langgraph-data:
+        driver: local
+services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+    langgraph-postgres:
+        image: pgvector/pgvector:pg16
+        ports:
+            - "5433:5432"
+        environment:
+            POSTGRES_DB: postgres
+            POSTGRES_USER: postgres
+            POSTGRES_PASSWORD: postgres
+        command:
+            - postgres
+            - -c
+            - shared_preload_libraries=vector
+        volumes:
+            - langgraph-data:/var/lib/postgresql/data
+        healthcheck:
+            test: pg_isready -U postgres
+            start_period: 10s
+            timeout: 1s
+            retries: 5
+            interval: 5s
+    langgraph-api:
+        ports:
+            - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
+            langgraph-postgres:
+                condition: service_healthy
+        environment:
+            REDIS_URI: redis://langgraph-redis:6379
+            POSTGRES_URI: {DEFAULT_POSTGRES_URI}"""
+    assert clean_empty_lines(actual_compose_str) == expected_compose_str
+
+
+def test_compose_with_api_version_and_base_image():
+    """Test compose function with both api_version and base_image parameters."""
+    port = 8123
+    api_version = "1.0.0"
+    base_image = "my-registry/custom-api"
+
+    actual_compose_str = compose(
+        DEFAULT_DOCKER_CAPABILITIES,
+        port=port,
+        api_version=api_version,
+        base_image=base_image,
+    )
+
+    # Similar to the previous test - the compose function doesn't directly embed
+    # the api_version or base_image into the compose file since those are handled
+    # during the docker build process
+    expected_compose_str = f"""volumes:
+    langgraph-data:
+        driver: local
+services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+    langgraph-postgres:
+        image: pgvector/pgvector:pg16
+        ports:
+            - "5433:5432"
+        environment:
+            POSTGRES_DB: postgres
+            POSTGRES_USER: postgres
+            POSTGRES_PASSWORD: postgres
+        command:
+            - postgres
+            - -c
+            - shared_preload_libraries=vector
+        volumes:
+            - langgraph-data:/var/lib/postgresql/data
+        healthcheck:
+            test: pg_isready -U postgres
+            start_period: 10s
+            timeout: 1s
+            retries: 5
+            interval: 5s
+    langgraph-api:
+        ports:
+            - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
+            langgraph-postgres:
+                condition: service_healthy
+        environment:
+            REDIS_URI: redis://langgraph-redis:6379
+            POSTGRES_URI: {DEFAULT_POSTGRES_URI}"""
+    assert clean_empty_lines(actual_compose_str) == expected_compose_str
+
+
+def test_compose_with_api_version_and_custom_postgres():
+    """Test compose function with api_version and custom postgres URI."""
+    port = 8123
+    api_version = "0.2.74"
+    custom_postgres_uri = "postgresql://user:pass@external-db:5432/mydb"
+
+    actual_compose_str = compose(
+        DEFAULT_DOCKER_CAPABILITIES,
+        port=port,
+        api_version=api_version,
+        postgres_uri=custom_postgres_uri,
+    )
+
+    expected_compose_str = f"""services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+    langgraph-api:
+        ports:
+            - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
+        environment:
+            REDIS_URI: redis://langgraph-redis:6379
+            POSTGRES_URI: {custom_postgres_uri}"""
+    assert clean_empty_lines(actual_compose_str) == expected_compose_str
+
+
+def test_compose_with_api_version_and_debugger():
+    """Test compose function with api_version and debugger port."""
+    port = 8123
+    debugger_port = 8001
+    api_version = "0.2.74"
+
+    actual_compose_str = compose(
+        DEFAULT_DOCKER_CAPABILITIES,
+        port=port,
+        api_version=api_version,
+        debugger_port=debugger_port,
+    )
+
+    expected_compose_str = f"""volumes:
+    langgraph-data:
+        driver: local
+services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+    langgraph-postgres:
+        image: pgvector/pgvector:pg16
+        ports:
+            - "5433:5432"
+        environment:
+            POSTGRES_DB: postgres
+            POSTGRES_USER: postgres
+            POSTGRES_PASSWORD: postgres
+        command:
+            - postgres
+            - -c
+            - shared_preload_libraries=vector
+        volumes:
+            - langgraph-data:/var/lib/postgresql/data
+        healthcheck:
+            test: pg_isready -U postgres
+            start_period: 10s
+            timeout: 1s
+            retries: 5
+            interval: 5s
+    langgraph-debugger:
+        image: langchain/langgraph-debugger
+        restart: on-failure
+        depends_on:
+            langgraph-postgres:
+                condition: service_healthy
+        ports:
+            - "{debugger_port}:3968"
+    langgraph-api:
+        ports:
+            - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
+            langgraph-postgres:
+                condition: service_healthy
+        environment:
+            REDIS_URI: redis://langgraph-redis:6379
+            POSTGRES_URI: {DEFAULT_POSTGRES_URI}"""
+    assert clean_empty_lines(actual_compose_str) == expected_compose_str

libs/cli/uv.lock

@@ -531,7 +531,7 @@ wheels = [
 
 [[package]]
 name = "langgraph-cli"
-version = "0.3.5"
+version = "0.3.6"
 source = { editable = "." }
 dependencies = [
     { name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },

libs/langgraph/langgraph/func/__init__.py

@@ -261,10 +261,11 @@ class entrypoint(Generic[ContextT]):
             passed to the workflow.
         cache_policy: A cache policy to use for caching the results of the workflow.
         retry_policy: A retry policy (or list of policies) to use for the workflow in case of a failure.
-        config_schema: Specifies the schema for the `configurable` key in the `RunnableConfig` object.
-            !!! warning "Deprecated"
-                This parameter is deprecated and support will be removed in v2.0.0.
-                Please use `context_schema` instead.
+
+    !!! warning "`config_schema` Deprecated"
+        The `config_schema` parameter is deprecated in v0.6.0 and support will be removed in v2.0.0.
+        Please use `context_schema` instead to specify the schema for run-scoped context.
+
 
     Example: Using entrypoint and tasks
         ```python

libs/langgraph/langgraph/graph/state.py

@@ -128,10 +128,10 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
             Use this to expose immutable context data to your nodes, like user_id, db_conn, etc.
         input_schema: The schema class that defines the input to the graph.
         output_schema: The schema class that defines the output from the graph.
-        config_schema: Specifies the schema for the `configurable` key in the `RunnableConfig` object.
-            !!! warning "Deprecated"
-                This parameter is deprecated and support will be removed in v2.0.0.
-                Please use `context_schema` instead.
+
+    !!! warning "`config_schema` Deprecated"
+        The `config_schema` parameter is deprecated in v0.6.0 and support will be removed in v2.0.0.
+        Please use `context_schema` instead to specify the schema for run-scoped context.
 
     Example:
         ```python

libs/langgraph/langgraph/pregel/main.py

@@ -2439,6 +2439,8 @@ class Pregel(
         Args:
             input: The input to the graph.
             config: The configuration to use for the run.
+            context: The static context to use for the run.
+                !!! version-added "Added in version 0.6.0."
             stream_mode: The mode to stream output, defaults to `self.stream_mode`.
                 Options are:
 
@@ -2695,6 +2697,8 @@ class Pregel(
         Args:
             input: The input to the graph.
             config: The configuration to use for the run.
+            context: The static context to use for the run.
+                !!! version-added "Added in version 0.6.0."
             stream_mode: The mode to stream output, defaults to `self.stream_mode`.
                 Options are:
 
@@ -2982,6 +2986,8 @@ class Pregel(
         Args:
             input: The input data for the graph. It can be a dictionary or any other type.
             config: Optional. The configuration for the graph run.
+            context: The static context to use for the run.
+                !!! version-added "Added in version 0.6.0."
             stream_mode: Optional[str]. The stream mode for the graph run. Default is "values".
             print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way.
             output_keys: Optional. The output keys to retrieve from the graph run.
@@ -3059,6 +3065,8 @@ class Pregel(
         Args:
             input: The input data for the computation. It can be a dictionary or any other type.
             config: Optional. The configuration for the computation.
+            context: The static context to use for the run.
+                !!! version-added "Added in version 0.6.0."
             stream_mode: Optional. The stream mode for the computation. Default is "values".
             print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way.
             output_keys: Optional. The output keys to include in the result. Default is None.

libs/langgraph/langgraph/pregel/remote.py

@@ -8,6 +8,7 @@ from typing import (
     cast,
 )
 
+import langsmith as ls
 from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.graph import (
     Edge as DrawableEdge,
@@ -118,6 +119,7 @@ class RemoteGraph(PregelProtocol):
         sync_client: SyncLangGraphClient | None = None,
         config: RunnableConfig | None = None,
         name: str | None = None,
+        distributed_tracing: bool = False,
     ):
         """Specify `url`, `api_key`, and/or `headers` to create default sync and async clients.
 
@@ -136,6 +138,7 @@ class RemoteGraph(PregelProtocol):
             name: Human-readable name to attach to the RemoteGraph instance.
                 This is useful for adding `RemoteGraph` as a subgraph via `graph.add_node(remote_graph)`.
                 If not provided, defaults to the assistant ID.
+            distributed_tracing: Whether to enable sending LangSmith distributed tracing headers.
         """
         self.assistant_id = assistant_id
         if name is None:
@@ -143,6 +146,7 @@ class RemoteGraph(PregelProtocol):
         else:
             self.name = name
         self.config = config
+        self.distributed_tracing = distributed_tracing
 
         if client is None and url is not None:
             client = get_client(url=url, api_key=api_key, headers=headers)
@@ -629,6 +633,7 @@ class RemoteGraph(PregelProtocol):
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         subgraphs: bool = False,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> Iterator[dict[str, Any] | Any]:
         """Create a run and stream the results.
@@ -644,6 +649,7 @@ class RemoteGraph(PregelProtocol):
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
             subgraphs: Stream from subgraphs.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to client.runs.stream.
 
         Yields:
@@ -672,6 +678,9 @@ class RemoteGraph(PregelProtocol):
             interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs or stream is not None,
             if_not_exists="create",
+            headers=_merge_tracing_headers(headers)
+            if self.distributed_tracing
+            else headers,
             **kwargs,
         ):
             # split mode and ns
@@ -731,6 +740,7 @@ class RemoteGraph(PregelProtocol):
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         subgraphs: bool = False,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> AsyncIterator[dict[str, Any] | Any]:
         """Create a run and stream the results.
@@ -746,6 +756,7 @@ class RemoteGraph(PregelProtocol):
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
             subgraphs: Stream from subgraphs.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to client.runs.stream.
 
         Yields:
@@ -774,6 +785,9 @@ class RemoteGraph(PregelProtocol):
             interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs or stream is not None,
             if_not_exists="create",
+            headers=_merge_tracing_headers(headers)
+            if self.distributed_tracing
+            else headers,
             **kwargs,
         ):
             # split mode and ns
@@ -847,6 +861,7 @@ class RemoteGraph(PregelProtocol):
         *,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> dict[str, Any] | Any:
         """Create a run, wait until it finishes and return the final state.
@@ -856,6 +871,7 @@ class RemoteGraph(PregelProtocol):
             config: A `RunnableConfig` for graph invocation.
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to RemoteGraph.stream.
 
         Returns:
@@ -866,6 +882,7 @@ class RemoteGraph(PregelProtocol):
             config=config,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
+            headers=headers,
             stream_mode="values",
             **kwargs,
         ):
@@ -882,6 +899,7 @@ class RemoteGraph(PregelProtocol):
         *,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> dict[str, Any] | Any:
         """Create a run, wait until it finishes and return the final state.
@@ -891,6 +909,7 @@ class RemoteGraph(PregelProtocol):
             config: A `RunnableConfig` for graph invocation.
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to RemoteGraph.astream.
 
         Returns:
@@ -901,6 +920,7 @@ class RemoteGraph(PregelProtocol):
             config=config,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
+            headers=headers,
             stream_mode="values",
             **kwargs,
         ):
@@ -909,3 +929,17 @@ class RemoteGraph(PregelProtocol):
             return chunk
         except UnboundLocalError:
             return None
+
+
+def _merge_tracing_headers(headers: dict[str, str] | None) -> dict[str, str] | None:
+    if rt := ls.get_current_run_tree():
+        tracing_headers = rt.to_headers()
+        baggage = tracing_headers.pop("baggage")
+        if headers:
+            if "baggage" in headers:
+                baggage = headers["baggage"] + "," + baggage
+            tracing_headers["baggage"] = baggage
+            headers.update(tracing_headers)
+        else:
+            headers = tracing_headers
+    return headers

libs/langgraph/pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
     "langchain-core>=0.1",
     "langgraph-checkpoint>=2.1.0,<3.0.0",
     "langgraph-sdk>=0.2.0,<0.3.0",
-    "langgraph-prebuilt>=0.5.0,<0.6.0",
+    "langgraph-prebuilt>=0.6.0,<0.7.0",
     "xxhash>=3.5.0",
     "pydantic>=2.7.4",
 ]

libs/langgraph/tests/test_remote_graph.py

@@ -3,6 +3,7 @@ import sys
 from typing import Annotated, Union
 from unittest.mock import AsyncMock, MagicMock
 
+import langsmith as ls
 import pytest
 from langchain_core.messages import AnyMessage, BaseMessage
 from langchain_core.runnables import RunnableConfig
@@ -899,21 +900,20 @@ async def test_langgraph_cloud_integration():
     }
 
     # test invoke
-    response = app.invoke(
+    app.invoke(
         input,
         config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}},
         interrupt_before=["agent"],
     )
-    print("response:", response["messages"][-1].content)
 
     # test stream
-    async for chunk in app.astream(
+    async for _ in app.astream(
         input,
         config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
         subgraphs=True,
         stream_mode=["debug", "messages"],
     ):
-        print("chunk:", chunk)
+        pass
 
     # test stream events
     async for chunk in remote_pregel.astream_events(
@@ -923,17 +923,16 @@ async def test_langgraph_cloud_integration():
         subgraphs=True,
         stream_mode=[],
     ):
-        print("chunk:", chunk)
+        pass
 
     # test get state
-    state_snapshot = await remote_pregel.aget_state(
+    await remote_pregel.aget_state(
         config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
         subgraphs=True,
     )
-    print("state snapshot:", state_snapshot)
 
     # test update state
-    response = await remote_pregel.aupdate_state(
+    await remote_pregel.aupdate_state(
         config={"configurable": {"thread_id": "6645e002-ed50-4022-92a3-d0d186fdf812"}},
         values={
             "messages": [
@@ -944,18 +943,16 @@ async def test_langgraph_cloud_integration():
             ]
         },
     )
-    print("response:", response)
 
     # test get history
     async for state in remote_pregel.aget_state_history(
         config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
     ):
-        print("state snapshot:", state)
+        pass
 
     # test get graph
     remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca"  # must be UUID
-    graph = await remote_pregel.aget_graph(xray=True)
-    print("graph:", graph)
+    await remote_pregel.aget_graph(xray=True)
 
 
 def test_sanitize_config():
@@ -1181,3 +1178,73 @@ async def test_remote_graph_stream_messages_tuple(
     assert coerced_events == coerced_inmem_events
     # TODO: Fix the namespace matching in the next api release.
     # assert namespaces == inmem_namespaces
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("distributed_tracing", [False, True])
+@pytest.mark.parametrize("stream", [False, True])
+async def test_include_headers(distributed_tracing: bool, stream: bool):
+    mock_async_client = MagicMock()
+    async_iter = MagicMock()
+    return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+    ]
+    async_iter.__aiter__.return_value = return_value
+    astream_mock = mock_async_client.runs.stream
+    astream_mock.return_value = async_iter
+
+    mock_sync_client = MagicMock()
+    sync_iter = MagicMock()
+    sync_iter.__iter__.return_value = return_value
+    stream_mock = mock_sync_client.runs.stream
+    stream_mock.return_value = async_iter
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+        sync_client=mock_sync_client,
+        distributed_tracing=distributed_tracing,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    with ls.tracing_context(enabled=True, client=MagicMock()):
+        with ls.trace("foo"):
+            if stream:
+                async for _ in remote_pregel.astream(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                ):
+                    pass
+
+            else:
+                await remote_pregel.ainvoke(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                )
+    expected = {"foo": "bar"}
+    if distributed_tracing:
+        expected["langsmith-trace"] = AnyStr()
+        expected["baggage"] = AnyStr()
+
+    assert astream_mock.call_args.kwargs["headers"] == expected
+    stream_mock.assert_not_called()
+
+    with ls.tracing_context(enabled=True, client=MagicMock()):
+        with ls.trace("foo"):
+            if stream:
+                for _ in remote_pregel.stream(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                ):
+                    pass
+
+            else:
+                remote_pregel.invoke(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                )
+    assert stream_mock.call_args.kwargs["headers"] == expected

libs/langgraph/uv.lock

@@ -1394,7 +1394,7 @@ dev = [
 
 [[package]]
 name = "langgraph-cli"
-version = "0.3.5"
+version = "0.3.6"
 source = { editable = "../cli" }
 dependencies = [
     { name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
@@ -1433,7 +1433,7 @@ dev = [
 
 [[package]]
 name = "langgraph-prebuilt"
-version = "0.5.2"
+version = "0.6.0"
 source = { editable = "../prebuilt" }
 dependencies = [
     { name = "langchain-core" },
@@ -1481,7 +1481,7 @@ wheels = [
 
 [[package]]
 name = "langgraph-sdk"
-version = "0.2.0a1"
+version = "0.2.0"
 source = { editable = "../sdk-py" }
 dependencies = [
     { name = "httpx" },

libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py

@@ -245,6 +245,435 @@ def _validate_chat_history(
     raise ValueError(error_message)
 
 
+class _AgentBuilder:
+    """Internal builder class for constructing React agents with intuitive method-to-node mapping."""
+
+    def __init__(
+        self,
+        model: Union[str, LanguageModelLike],
+        tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode],
+        *,
+        prompt: Optional[Prompt] = None,
+        response_format: Optional[
+            Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]
+        ] = None,
+        pre_model_hook: Optional[RunnableLike] = None,
+        post_model_hook: Optional[RunnableLike] = None,
+        state_schema: Optional[StateSchemaType] = None,
+        context_schema: Optional[Type[Any]] = None,
+        version: Literal["v1", "v2"] = "v2",
+        name: Optional[str] = None,
+    ):
+        # Store all parameters
+        self.model = model
+        self.tools = tools
+        self.prompt = prompt
+        self.response_format = response_format
+        self.pre_model_hook = pre_model_hook
+        self.post_model_hook = post_model_hook
+        self.state_schema = state_schema
+        self.context_schema = context_schema
+        self.version = version
+        self.name = name
+
+        # Setup tools
+        if isinstance(self.tools, ToolNode):
+            self._tool_classes = list(self.tools.tools_by_name.values())
+            self._tool_node = self.tools
+        else:
+            self._llm_builtin_tools = [t for t in self.tools if isinstance(t, dict)]
+            self._tool_node = ToolNode(
+                [t for t in self.tools if not isinstance(t, dict)]
+            )
+            self._tool_classes = list(self._tool_node.tools_by_name.values())
+
+        self._should_return_direct: set[str] = {
+            t.name for t in self._tool_classes if t.return_direct
+        }
+
+        # Setup state schema
+        if self.state_schema is not None:
+            required_keys = {"messages", "remaining_steps"}
+            if self.response_format is not None:
+                required_keys.add("structured_response")
+
+            schema_keys = set(get_type_hints(self.state_schema))
+            if missing_keys := required_keys - schema_keys:
+                raise ValueError(
+                    f"Missing required key(s) {missing_keys} in state_schema"
+                )
+
+            self._final_state_schema = self.state_schema
+        else:
+            self._final_state_schema = (
+                AgentStateWithStructuredResponse
+                if self.response_format is not None
+                else AgentState
+            )
+
+        # Setup model
+        model = self.model
+
+        # Convert string models
+        if isinstance(model, str):
+            try:
+                from langchain.chat_models import init_chat_model  # type: ignore[import-not-found]
+            except ImportError:
+                raise ImportError(
+                    "Please install langchain (`pip install langchain`) to use '<provider>:<model>' string syntax for `model` parameter."
+                )
+            model = cast(BaseChatModel, init_chat_model(model))
+
+        # Bind tools if needed
+        if (
+            _should_bind_tools(
+                model, self._tool_classes, num_builtin=len(self._llm_builtin_tools)
+            )
+            and len(self._tool_classes + self._llm_builtin_tools) > 0
+        ):
+            model = cast(BaseChatModel, model).bind_tools(
+                self._tool_classes + self._llm_builtin_tools
+            )  # type: ignore[operator]
+
+        self._model_runnable = _get_prompt_runnable(self.prompt) | model
+
+    def create_model_node(self) -> RunnableCallable:
+        """Create the 'agent' node that calls the LLM."""
+
+        def _get_model_input_state(state: StateSchema) -> StateSchema:
+            if self.pre_model_hook is not None:
+                messages: Optional[Sequence[BaseMessage]] = (
+                    _get_state_value(state, "llm_input_messages")
+                ) or _get_state_value(state, "messages")
+                error_msg: str = f"Expected input to call_model to have 'llm_input_messages' or 'messages' key, but got {state}"
+            else:
+                messages = _get_state_value(state, "messages")
+                error_msg = f"Expected input to call_model to have 'messages' key, but got {state}"
+
+            if messages is None:
+                raise ValueError(error_msg)
+
+            _validate_chat_history(messages)
+
+            if isinstance(self._final_state_schema, type) and issubclass(
+                self._final_state_schema, BaseModel
+            ):
+                state.messages = messages  # type: ignore
+            else:
+                state["messages"] = messages  # type: ignore
+            return state
+
+        def _are_more_steps_needed(state: StateSchema, response: BaseMessage) -> bool:
+            has_tool_calls = isinstance(response, AIMessage) and response.tool_calls
+            all_tools_return_direct = (
+                all(
+                    call["name"] in self._should_return_direct
+                    for call in response.tool_calls
+                )
+                if isinstance(response, AIMessage)
+                else False
+            )
+            remaining_steps = _get_state_value(state, "remaining_steps", None)
+            is_last_step = _get_state_value(state, "is_last_step", False)
+            return (
+                (remaining_steps is None and is_last_step and has_tool_calls)
+                or (
+                    remaining_steps is not None
+                    and remaining_steps < 1
+                    and all_tools_return_direct
+                )
+                or (
+                    remaining_steps is not None
+                    and remaining_steps < 2
+                    and has_tool_calls
+                )
+            )
+
+        def call_model(state: StateSchema, config: RunnableConfig) -> StateSchema:
+            state = _get_model_input_state(state)
+            response = cast(AIMessage, self._model_runnable.invoke(state, config))  # type: ignore[union-attr]
+            response.name = self.name
+
+            if _are_more_steps_needed(state, response):
+                return {
+                    "messages": [
+                        AIMessage(
+                            id=response.id,
+                            content="Sorry, need more steps to process this request.",
+                        )
+                    ]
+                }
+            return {"messages": [response]}
+
+        async def acall_model(
+            state: StateSchema, config: RunnableConfig
+        ) -> StateSchema:
+            state = _get_model_input_state(state)
+            response = cast(
+                AIMessage, await self._model_runnable.ainvoke(state, config)
+            )  # type: ignore[union-attr]
+            response.name = self.name
+
+            if _are_more_steps_needed(state, response):
+                return {
+                    "messages": [
+                        AIMessage(
+                            id=response.id,
+                            content="Sorry, need more steps to process this request.",
+                        )
+                    ]
+                }
+            return {"messages": [response]}
+
+        # Determine input schema
+        input_schema = self._final_state_schema
+        if self.pre_model_hook is not None:
+            if isinstance(self._final_state_schema, type) and issubclass(
+                self._final_state_schema, BaseModel
+            ):
+                from pydantic import create_model
+
+                input_schema = create_model(
+                    "CallModelInputSchema",
+                    llm_input_messages=(list[AnyMessage], ...),
+                    __base__=self._final_state_schema,
+                )
+            else:
+
+                class CallModelInputSchema(self._final_state_schema):  # type: ignore
+                    llm_input_messages: list[AnyMessage]
+
+                input_schema = CallModelInputSchema
+
+        return RunnableCallable(call_model, acall_model, input_schema=input_schema)
+
+    def create_structured_response_node(self) -> Optional[RunnableCallable]:
+        """Create the 'generate_structured_response' node if configured."""
+        if self.response_format is None:
+            return None
+
+        def generate_structured_response(
+            state: StateSchema, config: RunnableConfig
+        ) -> StateSchema:
+            messages = _get_state_value(state, "messages")
+            structured_response_schema = self.response_format
+            if isinstance(self.response_format, tuple):
+                system_prompt, structured_response_schema = self.response_format
+                messages = [SystemMessage(content=system_prompt)] + list(messages)
+
+            model_with_structured_output = _get_model(
+                self._model_runnable
+            ).with_structured_output(  # type: ignore[arg-type]
+                cast(StructuredResponseSchema, structured_response_schema)
+            )
+            response = model_with_structured_output.invoke(messages, config)
+            return {"structured_response": response}
+
+        async def agenerate_structured_response(
+            state: StateSchema, config: RunnableConfig
+        ) -> StateSchema:
+            messages = _get_state_value(state, "messages")
+            structured_response_schema = self.response_format
+            if isinstance(self.response_format, tuple):
+                system_prompt, structured_response_schema = self.response_format
+                messages = [SystemMessage(content=system_prompt)] + list(messages)
+
+            model_with_structured_output = _get_model(
+                self._model_runnable
+            ).with_structured_output(  # type: ignore[arg-type]
+                cast(StructuredResponseSchema, structured_response_schema)
+            )
+            response = await model_with_structured_output.ainvoke(messages, config)
+            return {"structured_response": response}
+
+        return RunnableCallable(
+            generate_structured_response, agenerate_structured_response
+        )
+
+    def create_model_router(self) -> Callable[[StateSchema], Union[str, list[Send]]]:
+        """Create routing function for model node conditional edges."""
+
+        def should_continue(state: StateSchema) -> Union[str, list[Send]]:
+            messages = _get_state_value(state, "messages")
+            last_message = messages[-1]
+
+            if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
+                if self.post_model_hook is not None:
+                    return "post_model_hook"
+                elif self.response_format is not None:
+                    return "generate_structured_response"
+                else:
+                    return END
+            else:
+                if self.version == "v1":
+                    return "tools"
+                elif self.version == "v2":
+                    if self.post_model_hook is not None:
+                        return "post_model_hook"
+                    return [
+                        Send(
+                            "tools",
+                            ToolCallWithContext(
+                                __type="tool_call_with_context",
+                                tool_call=tool_call,
+                                state=state,
+                            ),
+                        )
+                        for tool_call in last_message.tool_calls
+                    ]
+
+        return should_continue
+
+    def post_model_hook_router(self, state: StateSchema) -> Union[str, list[Send]]:
+        """Route to the next node after post_model_hook."""
+        messages = _get_state_value(state, "messages")
+        tool_messages = [m.tool_call_id for m in messages if isinstance(m, ToolMessage)]
+        last_ai_message = next(
+            m for m in reversed(messages) if isinstance(m, AIMessage)
+        )
+        pending_tool_calls = [
+            c for c in last_ai_message.tool_calls if c["id"] not in tool_messages
+        ]
+
+        if pending_tool_calls:
+            return [
+                Send(
+                    "tools",
+                    ToolCallWithContext(
+                        __type="tool_call_with_context",
+                        tool_call=tool_call,
+                        state=state,
+                    ),
+                )
+                for tool_call in pending_tool_calls
+            ]
+        elif isinstance(messages[-1], ToolMessage):
+            return self._get_entry_point()
+        elif self.response_format is not None:
+            return "generate_structured_response"
+        else:
+            return END
+
+    def create_tools_router(self) -> Optional[Callable[[StateSchema], str]]:
+        """Create routing function for tools node conditional edges."""
+        if not self._should_return_direct:
+            return None
+
+        def route_tool_responses(state: StateSchema) -> str:
+            messages = _get_state_value(state, "messages")
+            for m in reversed(messages):
+                if not isinstance(m, ToolMessage):
+                    break
+                if m.name in self._should_return_direct:
+                    return END
+
+            if isinstance(m, AIMessage) and m.tool_calls:
+                if any(
+                    call["name"] in self._should_return_direct for call in m.tool_calls
+                ):
+                    return END
+
+            return self._get_entry_point()
+
+        return route_tool_responses
+
+    def _get_entry_point(self) -> str:
+        """Get the workflow entry point."""
+        return "pre_model_hook" if self.pre_model_hook else "agent"
+
+    def _has_tools(self) -> bool:
+        """Check if agent has tools enabled."""
+        return len(self._tool_classes) > 0
+
+    def _get_model_edges(self) -> list[str]:
+        """Get possible edge destinations from model node."""
+        edges = []
+
+        # If post_model_hook exists, we don't add edges here - we use direct edge instead
+        if not self.post_model_hook:
+            if self._has_tools():
+                edges.append("tools")
+            if self.response_format:
+                edges.append("generate_structured_response")
+            if not self._has_tools() and not self.response_format:
+                edges.append(END)
+
+        return edges
+
+    def _get_post_model_hook_edges(self) -> list[str]:
+        """Get possible edge destinations from post_model_hook node."""
+        edges = [self._get_entry_point()]
+        if self._has_tools():
+            edges.append("tools")
+        if self.response_format:
+            edges.append("generate_structured_response")
+        else:
+            edges.append(END)
+        return edges
+
+    def build(self) -> StateGraph:
+        """Build the agent workflow graph (uncompiled)."""
+        # Create workflow
+        workflow = StateGraph(
+            state_schema=self._final_state_schema,  # type: ignore[arg-type]
+            context_schema=self.context_schema,
+        )
+
+        # Add nodes
+        # Always add model node (named 'agent' for backwards compatibility)
+        workflow.add_node("agent", self.create_model_node())
+
+        # Add tools node if needed
+        if self._has_tools():
+            workflow.add_node("tools", self._tool_node)
+
+        # Add hook nodes if configured
+        if self.pre_model_hook:
+            workflow.add_node("pre_model_hook", self.pre_model_hook)  # type: ignore[arg-type]
+        if self.post_model_hook:
+            workflow.add_node("post_model_hook", self.post_model_hook)  # type: ignore[arg-type]
+
+        # Add structured response node if configured
+        structured_node = self.create_structured_response_node()
+        if structured_node:
+            workflow.add_node("generate_structured_response", structured_node)
+
+        # Add edges
+        entry_point = self._get_entry_point()
+        workflow.set_entry_point(entry_point)
+
+        # Pre-model hook edge
+        if self.pre_model_hook:
+            workflow.add_edge("pre_model_hook", "agent")
+
+        # Model node edges
+        if self.post_model_hook:
+            # Direct edge from model node to post_model_hook when post_model_hook exists
+            workflow.add_edge("agent", "post_model_hook")
+            # Post-model hook conditional edges
+            post_hook_edges = self._get_post_model_hook_edges()
+            workflow.add_conditional_edges(
+                "post_model_hook", self.post_model_hook_router, path_map=post_hook_edges
+            )  # type: ignore[arg-type]
+        else:
+            # Conditional edges from model node when no post_model_hook
+            model_router = self.create_model_router()
+            model_edges = self._get_model_edges()
+            workflow.add_conditional_edges("agent", model_router, path_map=model_edges)  # type: ignore[arg-type]
+
+        # Tools edges
+        if self._has_tools():
+            tools_router = self.create_tools_router()
+            if tools_router:
+                workflow.add_conditional_edges(
+                    "tools", tools_router, path_map=[entry_point, END]
+                )
+            else:
+                workflow.add_edge("tools", entry_point)
+
+        return workflow
+
+
 def create_react_agent(
     model: Union[str, LanguageModelLike],
     tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode],
@@ -363,10 +792,11 @@ def create_react_agent(
         name: An optional name for the CompiledStateGraph.
             This name will be automatically used when adding ReAct agent graph to another graph as a subgraph node -
             particularly useful for building multi-agent systems.
-        config_schema: Specifies the schema for the `configurable` key in the `RunnableConfig` object.
-            !!! warning "Deprecated"
-                This parameter is deprecated and support will be removed in v2.0.0.
-                Please use `context_schema` instead.
+
+    !!! warning "`config_schema` Deprecated"
+        The `config_schema` parameter is deprecated in v0.6.0 and support will be removed in v2.0.0.
+        Please use `context_schema` instead to specify the schema for run-scoped context.
+
 
     Returns:
         A compiled LangChain runnable that can be used for chat interactions.
@@ -410,6 +840,7 @@ def create_react_agent(
             print(chunk)
         ```
     """
+    # Handle deprecated config_schema parameter
     if (
         config_schema := deprecated_kwargs.pop("config_schema", MISSING)
     ) is not MISSING:
@@ -421,400 +852,29 @@ def create_react_agent(
         if context_schema is not None:
             context_schema = config_schema
 
+    # Validate version
     if version not in ("v1", "v2"):
         raise ValueError(
             f"Invalid version {version}. Supported versions are 'v1' and 'v2'."
         )
 
-    if state_schema is not None:
-        required_keys = {"messages", "remaining_steps"}
-        if response_format is not None:
-            required_keys.add("structured_response")
-
-        schema_keys = set(get_type_hints(state_schema))
-        if missing_keys := required_keys - set(schema_keys):
-            raise ValueError(f"Missing required key(s) {missing_keys} in state_schema")
-
-    if state_schema is None:
-        state_schema = (
-            AgentStateWithStructuredResponse
-            if response_format is not None
-            else AgentState
-        )
-
-    llm_builtin_tools: list[dict] = []
-    if isinstance(tools, ToolNode):
-        tool_classes = list(tools.tools_by_name.values())
-        tool_node = tools
-    else:
-        llm_builtin_tools = [t for t in tools if isinstance(t, dict)]
-        tool_node = ToolNode([t for t in tools if not isinstance(t, dict)])
-        tool_classes = list(tool_node.tools_by_name.values())
-
-    if isinstance(model, str):
-        try:
-            from langchain.chat_models import (  # type: ignore[import-not-found]
-                init_chat_model,
-            )
-        except ImportError:
-            raise ImportError(
-                "Please install langchain (`pip install langchain`) to use '<provider>:<model>' string syntax for `model` parameter."
-            )
-
-        model = cast(BaseChatModel, init_chat_model(model))
-
-    tool_calling_enabled = len(tool_classes) > 0
-
-    if (
-        _should_bind_tools(model, tool_classes, num_builtin=len(llm_builtin_tools))
-        and len(tool_classes + llm_builtin_tools) > 0
-    ):
-        model = cast(BaseChatModel, model).bind_tools(tool_classes + llm_builtin_tools)  # type: ignore[operator]
-
-    model_runnable = _get_prompt_runnable(prompt) | model
-
-    # If any of the tools are configured to return_directly after running,
-    # our graph needs to check if these were called
-    should_return_direct = {t.name for t in tool_classes if t.return_direct}
-
-    def _are_more_steps_needed(state: StateSchema, response: BaseMessage) -> bool:
-        has_tool_calls = isinstance(response, AIMessage) and response.tool_calls
-        all_tools_return_direct = (
-            all(call["name"] in should_return_direct for call in response.tool_calls)
-            if isinstance(response, AIMessage)
-            else False
-        )
-        remaining_steps = _get_state_value(state, "remaining_steps", None)
-        is_last_step = _get_state_value(state, "is_last_step", False)
-        return (
-            (remaining_steps is None and is_last_step and has_tool_calls)
-            or (
-                remaining_steps is not None
-                and remaining_steps < 1
-                and all_tools_return_direct
-            )
-            or (remaining_steps is not None and remaining_steps < 2 and has_tool_calls)
-        )
-
-    def _get_model_input_state(state: StateSchema) -> StateSchema:
-        if pre_model_hook is not None:
-            messages = (
-                _get_state_value(state, "llm_input_messages")
-            ) or _get_state_value(state, "messages")
-            error_msg = f"Expected input to call_model to have 'llm_input_messages' or 'messages' key, but got {state}"
-        else:
-            messages = _get_state_value(state, "messages")
-            error_msg = (
-                f"Expected input to call_model to have 'messages' key, but got {state}"
-            )
-
-        if messages is None:
-            raise ValueError(error_msg)
-
-        _validate_chat_history(messages)
-        # we're passing messages under `messages` key, as this is expected by the prompt
-        if isinstance(state_schema, type) and issubclass(state_schema, BaseModel):
-            state.messages = messages  # type: ignore
-        else:
-            state["messages"] = messages  # type: ignore
-
-        return state
-
-    # Define the function that calls the model
-    def call_model(state: StateSchema, config: RunnableConfig) -> StateSchema:
-        state = _get_model_input_state(state)
-        response = cast(AIMessage, model_runnable.invoke(state, config))
-        # add agent name to the AIMessage
-        response.name = name
-
-        if _are_more_steps_needed(state, response):
-            return {
-                "messages": [
-                    AIMessage(
-                        id=response.id,
-                        content="Sorry, need more steps to process this request.",
-                    )
-                ]
-            }
-        # We return a list, because this will get added to the existing list
-        return {"messages": [response]}
-
-    async def acall_model(state: StateSchema, config: RunnableConfig) -> StateSchema:
-        state = _get_model_input_state(state)
-        response = cast(AIMessage, await model_runnable.ainvoke(state, config))
-        # add agent name to the AIMessage
-        response.name = name
-        if _are_more_steps_needed(state, response):
-            return {
-                "messages": [
-                    AIMessage(
-                        id=response.id,
-                        content="Sorry, need more steps to process this request.",
-                    )
-                ]
-            }
-        # We return a list, because this will get added to the existing list
-        return {"messages": [response]}
-
-    input_schema: StateSchemaType
-    if pre_model_hook is not None:
-        # Dynamically create a schema that inherits from state_schema and adds 'llm_input_messages'
-        if isinstance(state_schema, type) and issubclass(state_schema, BaseModel):
-            # For Pydantic schemas
-            from pydantic import create_model
-
-            input_schema = create_model(
-                "CallModelInputSchema",
-                llm_input_messages=(list[AnyMessage], ...),
-                __base__=state_schema,
-            )
-        else:
-            # For TypedDict schemas
-            class CallModelInputSchema(state_schema):  # type: ignore
-                llm_input_messages: list[AnyMessage]
-
-            input_schema = CallModelInputSchema
-    else:
-        input_schema = state_schema
-
-    def generate_structured_response(
-        state: StateSchema, config: RunnableConfig
-    ) -> StateSchema:
-        messages = _get_state_value(state, "messages")
-        structured_response_schema = response_format
-        if isinstance(response_format, tuple):
-            system_prompt, structured_response_schema = response_format
-            messages = [SystemMessage(content=system_prompt)] + list(messages)
-
-        model_with_structured_output = _get_model(model).with_structured_output(
-            cast(StructuredResponseSchema, structured_response_schema)
-        )
-        response = model_with_structured_output.invoke(messages, config)
-        return {"structured_response": response}
-
-    async def agenerate_structured_response(
-        state: StateSchema, config: RunnableConfig
-    ) -> StateSchema:
-        messages = _get_state_value(state, "messages")
-        structured_response_schema = response_format
-        if isinstance(response_format, tuple):
-            system_prompt, structured_response_schema = response_format
-            messages = [SystemMessage(content=system_prompt)] + list(messages)
-
-        model_with_structured_output = _get_model(model).with_structured_output(
-            cast(StructuredResponseSchema, structured_response_schema)
-        )
-        response = await model_with_structured_output.ainvoke(messages, config)
-        return {"structured_response": response}
-
-    if not tool_calling_enabled:
-        # Define a new graph
-        workflow = StateGraph(state_schema=state_schema, context_schema=context_schema)
-        workflow.add_node(
-            "agent",
-            RunnableCallable(call_model, acall_model),
-            input_schema=input_schema,
-        )
-        if pre_model_hook is not None:
-            workflow.add_node("pre_model_hook", pre_model_hook)  # type: ignore[arg-type]
-            workflow.add_edge("pre_model_hook", "agent")
-            entrypoint = "pre_model_hook"
-        else:
-            entrypoint = "agent"
-
-        workflow.set_entry_point(entrypoint)
-
-        if post_model_hook is not None:
-            workflow.add_node("post_model_hook", post_model_hook)  # type: ignore[arg-type]
-            workflow.add_edge("agent", "post_model_hook")
-
-        if response_format is not None:
-            workflow.add_node(
-                "generate_structured_response",
-                RunnableCallable(
-                    generate_structured_response,
-                    agenerate_structured_response,
-                ),
-            )
-            if post_model_hook is not None:
-                workflow.add_edge("post_model_hook", "generate_structured_response")
-            else:
-                workflow.add_edge("agent", "generate_structured_response")
-
-        return workflow.compile(
-            checkpointer=checkpointer,
-            store=store,
-            interrupt_before=interrupt_before,
-            interrupt_after=interrupt_after,
-            debug=debug,
-            name=name,
-        )
-
-    # Define the function that determines whether to continue or not
-    def should_continue(state: StateSchema) -> Union[str, list[Send]]:
-        messages = _get_state_value(state, "messages")
-        last_message = messages[-1]
-        # If there is no function call, then we finish
-        if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
-            if post_model_hook is not None:
-                return "post_model_hook"
-            elif response_format is not None:
-                return "generate_structured_response"
-            else:
-                return END
-        # Otherwise if there is, we continue
-        else:
-            if version == "v1":
-                return "tools"
-            elif version == "v2":
-                if post_model_hook is not None:
-                    return "post_model_hook"
-                return [
-                    Send(
-                        "tools",
-                        ToolCallWithContext(
-                            __type="tool_call_with_context",
-                            tool_call=tool_call,
-                            state=state,
-                        ),
-                    )
-                    for tool_call in last_message.tool_calls
-                ]
-
-    # Define a new graph
-    workflow = StateGraph(
-        state_schema=state_schema or AgentState, context_schema=context_schema
+    # Build the graph using the internal builder
+    builder = _AgentBuilder(
+        model=model,
+        tools=tools,
+        prompt=prompt,
+        response_format=response_format,
+        pre_model_hook=pre_model_hook,
+        post_model_hook=post_model_hook,
+        state_schema=state_schema,
+        context_schema=context_schema,
+        version=version,
+        name=name,
     )
 
-    # Define the two nodes we will cycle between
-    workflow.add_node(
-        "agent",
-        RunnableCallable(call_model, acall_model),
-        input_schema=input_schema,
-    )
-    workflow.add_node("tools", tool_node)
+    workflow = builder.build()
 
-    # Optionally add a pre-model hook node that will be called
-    # every time before the "agent" (LLM-calling node)
-    if pre_model_hook is not None:
-        workflow.add_node("pre_model_hook", pre_model_hook)  # type: ignore[arg-type]
-        workflow.add_edge("pre_model_hook", "agent")
-        entrypoint = "pre_model_hook"
-    else:
-        entrypoint = "agent"
-
-    # Set the entrypoint as `agent`
-    # This means that this node is the first one called
-    workflow.set_entry_point(entrypoint)
-
-    agent_paths = []
-    post_model_hook_paths = [entrypoint, "tools"]
-
-    # Add a post model hook node if post_model_hook is provided
-    if post_model_hook is not None:
-        workflow.add_node("post_model_hook", post_model_hook)  # type: ignore[arg-type]
-        agent_paths.append("post_model_hook")
-        workflow.add_edge("agent", "post_model_hook")
-    else:
-        agent_paths.append("tools")
-
-    # Add a structured output node if response_format is provided
-    if response_format is not None:
-        workflow.add_node(
-            "generate_structured_response",
-            RunnableCallable(
-                generate_structured_response,
-                agenerate_structured_response,
-            ),
-        )
-        if post_model_hook is not None:
-            post_model_hook_paths.append("generate_structured_response")
-        else:
-            agent_paths.append("generate_structured_response")
-    else:
-        if post_model_hook is not None:
-            post_model_hook_paths.append(END)
-        else:
-            agent_paths.append(END)
-
-    if post_model_hook is not None:
-
-        def post_model_hook_router(state: StateSchema) -> Union[str, list[Send]]:
-            """Route to the next node after post_model_hook.
-
-            Routes to one of:
-            * "tools": if there are pending tool calls without a corresponding message.
-            * "generate_structured_response": if no pending tool calls exist and response_format is specified.
-            * END: if no pending tool calls exist and no response_format is specified.
-            """
-
-            messages = _get_state_value(state, "messages")
-            tool_messages = [
-                m.tool_call_id for m in messages if isinstance(m, ToolMessage)
-            ]
-            last_ai_message = next(
-                m for m in reversed(messages) if isinstance(m, AIMessage)
-            )
-            pending_tool_calls = [
-                c for c in last_ai_message.tool_calls if c["id"] not in tool_messages
-            ]
-
-            if pending_tool_calls:
-                return [
-                    Send(
-                        "tools",
-                        ToolCallWithContext(
-                            __type="tool_call_with_context",
-                            tool_call=tool_call,
-                            state=state,
-                        ),
-                    )
-                    for tool_call in pending_tool_calls
-                ]
-            elif isinstance(messages[-1], ToolMessage):
-                return entrypoint
-            elif response_format is not None:
-                return "generate_structured_response"
-            else:
-                return END
-
-        workflow.add_conditional_edges(
-            "post_model_hook",
-            post_model_hook_router,  # type: ignore[arg-type]
-            path_map=post_model_hook_paths,
-        )
-
-    workflow.add_conditional_edges(
-        "agent",
-        should_continue,  # type: ignore[arg-type]
-        path_map=agent_paths,
-    )
-
-    def route_tool_responses(state: StateSchema) -> str:
-        for m in reversed(_get_state_value(state, "messages")):
-            if not isinstance(m, ToolMessage):
-                break
-            if m.name in should_return_direct:
-                return END
-
-        # handle a case of parallel tool calls where
-        # the tool w/ `return_direct` was executed in a different `Send`
-        if isinstance(m, AIMessage) and m.tool_calls:
-            if any(call["name"] in should_return_direct for call in m.tool_calls):
-                return END
-
-        return entrypoint
-
-    if should_return_direct:
-        workflow.add_conditional_edges(
-            "tools", route_tool_responses, path_map=[entrypoint, END]
-        )
-    else:
-        workflow.add_edge("tools", entrypoint)
-
-    # Finally, we compile it!
-    # This compiles it into a LangChain Runnable,
-    # meaning you can use it as you would any other runnable
+    # Compile and return the graph
     return workflow.compile(
         checkpointer=checkpointer,
         store=store,

libs/prebuilt/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-prebuilt"
-version = "0.5.2"
+version = "0.6.0"
 description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
 authors = []
 requires-python = ">=3.9"

libs/prebuilt/uv.lock

@@ -460,7 +460,7 @@ dev = [
 
 [[package]]
 name = "langgraph-prebuilt"
-version = "0.5.2"
+version = "0.6.0"
 source = { editable = "." }
 dependencies = [
     { name = "langchain-core" },
@@ -507,7 +507,7 @@ dev = [
 
 [[package]]
 name = "langgraph-sdk"
-version = "0.2.0a1"
+version = "0.2.0"
 source = { editable = "../sdk-py" }
 dependencies = [
     { name = "httpx" },

libs/sdk-py/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-sdk"
-version = "0.2.0a1"
+version = "0.2.0"
 description = "SDK for interacting with LangGraph API"
 authors = []
 requires-python = ">=3.9"

libs/sdk-py/uv.lock

@@ -119,7 +119,7 @@ wheels = [
 
 [[package]]
 name = "langgraph-sdk"
-version = "0.2.0a1"
+version = "0.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "httpx" },