Add store auth types

Within `libs/langgraph`, change all `TypedDict` imports to come from
`typing_extensions` rather than `typing`, as `pydantic` doesn't like the
latter.

Additionally, add a ruff rule to ban these imports too (so this doesn't
regress).

Solves  #2909.

docs/_scripts/generate_api_reference_links.py

@@ -6,6 +6,9 @@ import re
 from typing import List, Literal, Optional
 from typing_extensions import TypedDict
 
+
+from functools import lru_cache
+
 import nbformat
 from nbconvert.preprocessors import Preprocessor
 
@@ -47,6 +50,8 @@ MANUAL_API_REFERENCES_LANGGRAPH = [
     (["langgraph.graph"], "langgraph.constants", "END", "constants"),
     (["langgraph.constants"], "langgraph.types", "Send", "types"),
     (["langgraph.constants"], "langgraph.types", "Interrupt", "types"),
+    (["langgraph.constants"], "langgraph.types", "interrupt", "types"),
+    (["langgraph.constants"], "langgraph.types", "Command", "types"),
     ([], "langgraph.types", "RetryPolicy", "types"),
     ([], "langgraph.checkpoint.base", "Checkpoint", "checkpoints"),
     ([], "langgraph.checkpoint.base", "CheckpointMetadata", "checkpoints"),
@@ -83,8 +88,11 @@ _IMPORT_LANGCHAIN_RE = _make_regular_expression("langchain")
 _IMPORT_LANGGRAPH_RE = _make_regular_expression("langgraph")
 
 
-def _get_full_module_name(module_path, class_name) -> Optional[str]:
-    """Get full module name using inspect"""
+
+
+@lru_cache(maxsize=10_000)
+def _get_full_module_name(module_path: str, class_name: str) -> Optional[str]:
+    """Get full module name using inspect, with LRU cache to memoize results."""
     try:
         module = importlib.import_module(module_path)
         class_ = getattr(module, class_name)
@@ -95,13 +103,12 @@ def _get_full_module_name(module_path, class_name) -> Optional[str]:
             return module_path
         return module.__name__
     except AttributeError as e:
-        logger.warning(f"Could not find module for {class_name}, {e}")
+        logger.warning(f"API Reference: Could not find module for {class_name}, {e}")
         return None
     except ImportError as e:
-        logger.warning(f"Failed to load for class {class_name}, {e}")
+        logger.warning(f"API Reference: Failed to load for class {class_name}, {e}")
         return None
 
-
 def _get_doc_title(data: str, file_name: str) -> str:
     try:
         return re.findall(r"^#\s*(.*)", data, re.MULTILINE)[0]
@@ -115,10 +122,10 @@ def _get_doc_title(data: str, file_name: str) -> str:
 
 
 class ImportInformation(TypedDict):
-    imported: str  # imported class name
-    source: str  # module path
-    docs: str  # URL to the documentation
-    title: str  # Title of the document
+    imported: str  # The name of the class that was imported.
+    source: str  # The full module path from which the class was imported.
+    docs: str  # The URL pointing to the class's documentation.
+    title: str  # The title of the document where the import is used.
 
 
 def _get_imports(
@@ -211,36 +218,73 @@ def _get_imports(
     return imports
 
 
-class ImportPreprocessor(Preprocessor):
-    """A preprocessor to replace imports in each Python code cell with links to their
-    documentation and append the import info in a comment."""
+def get_imports(code: str, doc_title: str) -> List[ImportInformation]:
+    """Retrieve all import references from the given code for specified ecosystems.
 
-    def preprocess(self, nb, resources):
-        self.all_imports = []
-        file_name = os.path.basename(resources.get("metadata", {}).get("name", ""))
-        _DOC_TITLE = _get_doc_title(nb.cells[0].source, file_name)
+    Args:
+        code: The source code from which to extract import references.
+        doc_title: The documentation title associated with the code.
 
-        cells = []
-        for cell in nb.cells:
-            if cell.cell_type == "code":
-                cells.append(cell)
-                imports = _get_imports(
-                    cell.source, _DOC_TITLE, "langchain"
-                ) + _get_imports(cell.source, _DOC_TITLE, "langgraph")
-                if not imports:
-                    continue
+    Returns:
+        A list of import information for each import found.
+    """
+    ecosystems = ["langchain", "langgraph"]
+    all_imports = []
+    for package_ecosystem in ecosystems:
+        all_imports.extend(_get_imports(code, doc_title, package_ecosystem))
+    return all_imports
 
-                cells.append(
-                    nbformat.v4.new_markdown_cell(
-                        source=f"""
-<div>
-<b>API Reference:</b>
-{' | '.join(f'<a href="{imp["docs"]}">{imp["imported"]}</a>' for imp in imports)}
-</div>
-                    """
-                    )
-                )
-            else:
-                cells.append(cell)
-        nb.cells = cells
-        return nb, resources
+
+def update_markdown_with_imports(markdown: str) -> str:
+    """Update markdown to include API reference links for imports in Python code blocks.
+
+    This function scans the markdown content for Python code blocks, extracts any imports, and appends links to their API documentation.
+
+    Args:
+        markdown: The markdown content to process.
+
+    Returns:
+        Updated markdown with API reference links appended to Python code blocks.
+
+    Example:
+        Given a markdown with a Python code block:
+
+        ```python
+        from langchain.nlp import TextGenerator
+        ```
+        This function will append an API reference link to the `TextGenerator` class from the `langchain.nlp` module if it's recognized.
+    """
+    code_block_pattern = re.compile(
+        r'(?P<indent>[ \t]*)```(?P<language>python|py)\n(?P<code>.*?)\n(?P=indent)```', re.DOTALL
+    )
+
+    def replace_code_block(match: re.Match) -> str:
+        """Replace the matched code block with additional API reference links if imports are found.
+
+        Args:
+            match (re.Match): The regex match object containing the code block.
+
+        Returns:
+            str: The modified code block with API reference links appended if applicable.
+        """
+        indent = match.group('indent')
+        code_block = match.group('code')
+        language = match.group('language')  # Preserve the language from the regex match
+        # Retrieve import information from the code block
+        imports = get_imports(code_block, "__unused__")
+
+        original_code_block = match.group(0)
+        # If no imports are found, return the original code block
+        if not imports:
+            return original_code_block
+
+        # Generate API reference links for each import
+        api_links = ' | '.join(
+            f'<a href="{imp["docs"]}">{imp["imported"]}</a>' for imp in imports
+        )
+        # Return the code block with appended API reference links
+        return f'{original_code_block}\n\n{indent}API Reference: {api_links}'
+
+    # Apply the replace_code_block function to all matches in the markdown
+    updated_markdown = code_block_pattern.sub(replace_code_block, markdown)
+    return updated_markdown

docs/_scripts/notebook_convert.py

@@ -6,8 +6,6 @@ import nbformat
 from nbconvert.exporters import MarkdownExporter
 from nbconvert.preprocessors import Preprocessor
 
-from generate_api_reference_links import ImportPreprocessor
-
 
 class EscapePreprocessor(Preprocessor):
     def preprocess_cell(self, cell, resources, cell_index):
@@ -107,7 +105,6 @@ exporter = MarkdownExporter(
     preprocessors=[
         EscapePreprocessor,
         ExtractAttachmentsPreprocessor,
-        ImportPreprocessor,
     ],
     template_name="mdoutput",
     extra_template_basedirs=[

docs/_scripts/notebook_hooks.py

@@ -1,10 +1,13 @@
 import logging
 import os
+import re
 from typing import Any, Dict
 
-from mkdocs.structure.pages import Page
 from mkdocs.structure.files import Files, File
+from mkdocs.structure.pages import Page
+
 from notebook_convert import convert_notebook
+from generate_api_reference_links import update_markdown_with_imports
 
 logger = logging.getLogger(__name__)
 logging.basicConfig()
@@ -35,12 +38,83 @@ def on_files(files: Files, **kwargs: Dict[str, Any]):
     return new_files
 
 
+def _highlight_code_blocks(markdown: str) -> str:
+    """Find code blocks with highlight comments and add hl_lines attribute.
+
+    Args:
+        markdown: The markdown content to process.
+
+    Returns:
+        updated Markdown code with code blocks containing highlight comments
+        updated to use the hl_lines attribute.
+    """
+    # Pattern to find code blocks with highlight comments and without
+    # existing hl_lines for Python and JavaScript
+    # Pattern to find code blocks with highlight comments, handling optional indentation
+    code_block_pattern = re.compile(
+        r"(?P<indent>[ \t]*)```(?P<language>py|python|js|javascript)(?!\s+hl_lines=)\n"
+        r"(?P<code>((?:.*\n)*?))"  # Capture the code inside the block using named group
+        r"(?P=indent)```"  # Match closing backticks with the same indentation
+    )
+
+    def replace_highlight_comments(match: re.Match) -> str:
+        indent = match.group("indent")
+        language = match.group("language")
+        code_block = match.group("code")
+        lines = code_block.split("\n")
+        highlighted_lines = []
+
+        # Skip initial empty lines
+        while lines and not lines[0].strip():
+            lines.pop(0)
+
+        lines_to_keep = []
+
+        comment_syntax = (
+            "# highlight-next-line"
+            if language in ["py", "python"]
+            else "// highlight-next-line"
+        )
+
+        for line in lines:
+            if comment_syntax in line:
+                count = len(lines_to_keep) + 1
+                highlighted_lines.append(str(count))
+            else:
+                lines_to_keep.append(line)
+
+        # Reconstruct the new code block
+        new_code_block = "\n".join(lines_to_keep)
+
+        if highlighted_lines:
+            return (
+                f'{indent}```{language} hl_lines="{" ".join(highlighted_lines)}"\n'
+                # The indent and terminating \n is already included in the code block
+                f'{new_code_block}'
+                f'{indent}```'
+            )
+        else:
+            return (
+                f"{indent}```{language}\n"
+                # The indent and terminating \n is already included in the code block
+                f"{new_code_block}"
+                f"{indent}```"
+            )
+
+    # Replace all code blocks in the markdown
+    markdown = code_block_pattern.sub(replace_highlight_comments, markdown)
+    return markdown
+
+
 def on_page_markdown(markdown: str, page: Page, **kwargs: Dict[str, Any]):
     if DISABLED:
         return markdown
     if page.file.src_path.endswith(".ipynb"):
         logger.info("Processing Jupyter notebook: %s", page.file.src_path)
-        body = convert_notebook(page.file.abs_src_path)
-        return body
+        markdown = convert_notebook(page.file.abs_src_path)
 
+    # Append API reference links to code blocks
+    markdown = update_markdown_with_imports(markdown)
+    # Apply highlight comments to code blocks
+    markdown = _highlight_code_blocks(markdown)
     return markdown

docs/docs/cloud/reference/api/api_ref_control_plane.html

@@ -0,0 +1,19 @@
+<!doctype html>
+<html>
+  <head>
+    <title>LangGraph Cloud API Reference</title>
+    <meta charset="utf-8" />
+    <meta
+      name="viewport"
+      content="width=device-width, initial-scale=1" />
+  </head>
+  <body>
+    <script id="api-reference" data-url="./openapi_control_plane.json"></script>
+    <script>
+      var configuration = {}
+      document.getElementById('api-reference').dataset.configuration =
+        JSON.stringify(configuration)
+    </script>
+    <script src="https://cdn.jsdelivr.net/npm/@scalar/api-reference"></script>
+  </body>
+</html>

docs/docs/cloud/reference/api/openapi_control_plane.json

@@ -0,0 +1,695 @@
+{
+  "openapi": "3.1.0",
+  "info": {
+    "title": "LangGraph Control Plane API (Beta)",
+    "version": "0.0.1",
+    "description": "The LangGraph Control Plane API is used to programmatically create and manage LangGraph Server deployments. For example, the APIs can be orchestrated to create custom CI/CD workflows.\n\n### Beta\nThis API is currently in beta and may change or break without notice. This API documentation may not be up-to-date with actual API functionality.\n### Host\nhttps://api.host.langchain.com/\n\n### Authentication\nTo authenticate with the LangGraph Control Plane API, set the `X-Api-Key` header to a valid LangSmith API key for each request.\n\n### Versioning\nEach endpoint path is prefixed with a version (e.g. `v1`).\n\n### Quick Start\n\n1. Call `GET /{version}/projects` to retrieve the `Project` `id`. The `Project` `id` is needed in subsequent API calls.\n2. Call `POST /{version}/projects/{project_id}/revisions` to create a new `Revision` for the `Project`.\n3. Call `GET /{version}/projects/{project_id}/revisions` to get the latest `Revision` (first element in returned list). Get the `Revision` `id`.\n4. Poll for `Revision` `status` until `status` is `DEPLOYED` by calling `GET /{version}/projects/{project_id}/revisions/{revision_id}`."
+  },
+  "servers": [
+    {
+      "url": "https://api.host.langchain.com"
+    }
+  ],
+  "tags": [
+    {
+      "name": "Projects (v1)",
+      "description": "A project corresponds to a LangGraph Server deployment and the associated LangSmith tracing project.\n\nCreating a project via API is not currently supported/documented."
+    },
+    {
+      "name": "Revisions (v1)",
+      "description": "A revision is a version of a LangGraph Server deployment. Different revisions may contain different code and/or environment variables. A project can have many revisions."
+    }
+  ],
+  "paths": {
+    "/v1/projects": {
+      "get": {
+        "tags": ["Projects (v1)"],
+        "summary": "List Projects",
+        "description": "List all projects.",
+        "operationId": "list_projects_projects_get",
+        "parameters": [
+          {
+            "required": false,
+            "schema": {
+              "type": "integer",
+              "title": "Limit",
+              "description": "Maximum number of results to return. Minimum: 1. Maximum: 100.",
+              "default": 20
+            },
+            "name": "limit",
+            "in": "query"
+          },
+          {
+            "required": false,
+            "schema": {
+              "type": "integer",
+              "title": "Offset",
+              "description": "Pagination offset value. Pass this value in subsequent requests to retrieve the next page of results. Minimum: 0.",
+              "default": 0
+            },
+            "name": "offset",
+            "in": "query"
+          },
+          {
+            "required": false,
+            "schema": {
+              "type": "string",
+              "title": "Name Contains",
+              "description": "Filter string to filter projects by `name`."
+            },
+            "name": "name_contains",
+            "in": "query"
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Success",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "type": "array",
+                  "items": {
+                    "$ref": "#/components/schemas/Project"
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/v1/projects/{project_id}": {
+      "get": {
+        "tags": ["Projects (v1)"],
+        "summary": "Get Project",
+        "description": "Get project by ID.",
+        "operationId": "get_project_projects__project_id__get",
+        "parameters": [
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Project ID"
+            },
+            "name": "project_id",
+            "in": "path"
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Success",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/Project"
+                }
+              }
+            }
+          }
+        }
+      },
+      "delete": {
+        "tags": ["Projects (v1)"],
+        "summary": "Delete Project",
+        "description": "Delete project by ID.",
+        "operationId": "delete_project_projects__project_id__delete",
+        "parameters": [
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Project ID"
+            },
+            "name": "project_id",
+            "in": "path"
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Success",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/Project"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/v1/projects/{project_id}/revisions": {
+      "get": {
+        "tags": ["Revisions (v1)"],
+        "summary": "List Revisions",
+        "description": "List revisions of a project.",
+        "operationId": "list_revisions_projects__project_id__revisions_get",
+        "parameters": [
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Project ID"
+            },
+            "name": "project_id",
+            "in": "path"
+          },
+          {
+            "required": false,
+            "schema": {
+              "type": "integer",
+              "title": "Limit",
+              "description": "Maximum number of results to return. Minimum: 1. Maximum: 100.",
+              "default": 20
+            },
+            "name": "limit",
+            "in": "query"
+          },
+          {
+            "required": false,
+            "schema": {
+              "type": "integer",
+              "title": "Offset",
+              "description": "Pagination offset value. Pass this value in subsequent requests to retrieve the next page of results. Minimum: 0.",
+              "default": 0
+            },
+            "name": "offset",
+            "in": "query"
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Success",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "type": "array",
+                  "items": {
+                    "$ref": "#/components/schemas/Revision"
+                  }
+                }
+              }
+            }
+          }
+        }
+      },
+      "post": {
+        "tags": ["Revisions (v1)"],
+        "summary": "Create Revision",
+        "description": "Create a new revision for a project.",
+        "operationId": "create_revision_projects__project_id__revisions_post",
+        "parameters": [
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Project ID"
+            },
+            "name": "project_id",
+            "in": "path"
+          }
+        ],
+        "requestBody": {
+          "content": {
+            "application/json": {
+              "schema": {
+                "$ref": "#/components/schemas/CreateRevisionRequest"
+              }
+            }
+          },
+          "required": true
+        },
+        "responses": {
+          "200": {
+            "description": "Success",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/Project"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/v1/projects/{project_id}/revisions/{revision_id}": {
+      "get": {
+        "tags": ["Revisions (v1)"],
+        "summary": "Get Revision",
+        "description": "Get revision by ID.",
+        "operationId": "get_revision_projects__project_id__revisions__revision_id__get",
+        "parameters": [
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Project ID"
+            },
+            "name": "project_id",
+            "in": "path"
+          },
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Revision ID"
+            },
+            "name": "revision_id",
+            "in": "path"
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Success",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/Revision"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/v1/projects/{project_id}/revisions/{revision_id}/interrupt": {
+      "post": {
+        "tags": ["Revisions (v1)"],
+        "summary": "Interrupt Revision",
+        "description": "Interrupt revision by ID.\n\nIf the deployment of a revision appears \"stuck\", the revision may need to be interrupted. A new revision cannot be created if the latest revision is in a non-terminal `status`. In this scenario, the revision may need to be interrupted.",
+        "operationId": "interrupt_revision_projects__project_id__revisions__revision_id__interrupt_post",
+        "parameters": [
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Project ID"
+            },
+            "name": "project_id",
+            "in": "path"
+          },
+          {
+            "required": true,
+            "schema": {
+              "type": "string",
+              "format": "uuid",
+              "title": "Revision ID"
+            },
+            "name": "revision_id",
+            "in": "path"
+          }
+        ]
+      }
+    }
+  },
+  "components": {
+    "securitySchemes": {
+      "apiKeyAuth": {
+        "type": "apiKey",
+        "in": "header",
+        "name": "X-Api-Key"
+      }
+    },
+    "schemas": {
+      "EnvVar": {
+        "type": "object",
+        "description": "An environment variable or secret.",
+        "properties": {
+          "name": {
+            "type": "string",
+            "description": "Environment variable or secret name.",
+            "required": true
+          },
+          "value": {
+            "type": "string",
+            "description": "Environment variable or secret value.",
+            "required": true
+          },
+          "type": {
+            "type": "string",
+            "enum": [
+              "default",
+              "secret"
+            ],
+            "description": "Field to designate type of the environment variable (default) or secret.",
+            "required": true
+          }
+        }
+      },
+      "ContainerSpec": {
+        "type": "object",
+        "description": "Container specification for a revision's deployment.\n\nIf any field is omitted or set to `null`, the internal default value is used depending on the deployment type (`dev` or `prod`).",
+        "properties": {
+          "min_scale": {
+            "type": ["integer", "null"],
+            "description": "Minimum number of replicas in deployment.",
+            "default": "null"
+          },
+          "max_scale": {
+            "type": ["integer", "null"],
+            "description": "Maximum number of replicas in deployment.",
+            "default": "null"
+          },
+          "cpu": {
+            "type": ["integer", "null"],
+            "description": "Number of vCPU cores per replica.",
+            "default": "null"
+          },
+          "memory_mb": {
+            "type": ["integer", "null"],
+            "description": "Amount of memory in MB per replica.",
+            "default": "null"
+          }
+        }
+      },
+      "CreateRevisionRequest": {
+        "type": "object",
+        "description": "Object for creating a new revision.",
+        "properties": {
+          "image_path": {
+            "type": ["string", "null"],
+            "description": "URI of the Docker image to deploy.\n\nIf this field is omitted or set to `null`, the previous revision's `image_path` value is used. Set this field for BYOC deployments. Omit this field if creating a new revision from a GitHub repository.",
+            "default": "null"
+          },
+          "repo_path": {
+            "type": ["string", "null"],
+            "description": "Path to `langgraph.json` configuration file. For example, `langgraph.json` or `src/langgraph.json`.\n\nIf this field is omitted or set to `null`, the previous revision's `repo_path` value is used. Set this field for deployments from a GitHub repository. Omit this field if creating a new revision from a Docker image.",
+            "default": "null"
+          },
+          "env_vars": {
+            "type": "array",
+            "description": "List of environment variables or secrets.\n\nIf this field is omitted or set to `null`, the previous revision's `env_vars` value is used.",
+            "items": {
+              "$ref": "#/components/schemas/EnvVar"
+            },
+            "default": "null"
+          },
+          "shareable": {
+            "type": ["boolean", "null"],
+            "description": "Boolean flag to configure if a deployment is shareable through LangGraph Studio.\n\nIf this field is omitted or set to `null`, the previous revision's `shareable` value is used. This field does not apply to BYOC deployments.",
+            "default": "null"
+          },
+          "container_spec": {
+            "description": "If this field is omitted or set to `null`, the previous revision's `container_spec` value is used.",
+            "$ref": "#/components/schemas/ContainerSpec",
+            "default": "null"
+          }
+        }
+      },
+      "Project": {
+        "type": "object",
+        "description": "A project corresponds to a LangGraph Server deployment and the associated LangSmith tracing project.",
+        "properties": {
+          "id": {
+            "type": "string",
+            "format": "uuid",
+            "description": "ID of the project.",
+            "required": true
+          },
+          "tool_name": {
+            "type": ["string", "null"],
+            "description": "Do not use."
+          },
+          "display_name": {
+            "type": ["string", "null"],
+            "description": "Do not use."
+          },
+          "description": {
+            "type": ["string", "null"],
+            "description": "Do not use."
+          },
+          "example_input": {
+            "type": ["object", "null"],
+            "description": "Do not use."
+          },
+          "tenant_id": {
+            "type": "string",
+            "format": "uuid",
+            "description": "ID of the tenant/workspace of the project.",
+            "required": true
+          },
+          "created_at": {
+            "type": "string",
+            "format": "date-time",
+            "description": "Timestamp of when the project was created.",
+            "required": true
+          },
+          "updated_at": {
+            "type": "string",
+            "format": "date-time",
+            "description": "Timestamp of when the project was updated.",
+            "required": true
+          },
+          "name": {
+            "type": "string",
+            "description": "Name of the project.\n\nThis is also the name of the LangSmith tracing project for the LangGraph deployment.",
+            "required": true
+          },
+          "lc_hosted": {
+            "type": "boolean",
+            "description": "Boolean flag to indicate if the deployment is hosted in LangChain's cloud or an external cloud (e.g. BYOC).",
+            "required": true
+          },
+          "repo_url": {
+            "type": ["string", "null"],
+            "description": "URL of the GitHub repository.\n\nThis field is not used for deployments from a Docker image."
+          },
+          "repo_branch": {
+            "type": ["string", "null"],
+            "description": "Branch of the GitHub repository.\n\nThis field is not used for deployments from a Docker image."
+          },
+          "tracer_session_id": {
+            "type": ["string", "null"],
+            "format": "uuid",
+            "description": "Do not use."
+          },
+          "api_key_id": {
+            "type": ["string", "null"],
+            "format": "uuid",
+            "description": "Do not use."
+          },
+          "build_on_push": {
+            "type": "boolean",
+            "description": "Boolean flag to indicate if a new revision is automatically created on push to GitHub branch (`repo_branch`).\n\nThis field does not apply for BYOC deployments."
+          },
+          "input_json_schemas": {
+            "type": ["object", "null"],
+            "description": "Do not use."
+          },
+          "output_json_schemas": {
+            "type": ["object", "null"],
+            "description": "Do not use."
+          },
+          "host_integration_id": {
+            "type": ["string", "null"],
+            "format": "uuid",
+            "description": "Do not use."
+          },
+          "metadata": {
+            "$ref": "#/components/schemas/ProjectMetadata"
+          },
+          "resource": {
+            "$ref": "#/components/schemas/ResourceService"
+          }
+        }
+      },
+      "ProjectMetadata": {
+        "type": "object",
+        "description": "Metadata associated with a `Project`.",
+        "properties": {
+          "deployment_type": {
+            "type": "string",
+            "description": "Development (`dev`) or Production (`prod`) type deployment.",
+            "enum": [
+              "dev",
+              "prod"
+            ]
+          },
+          "image_source": {
+            "type": "string",
+            "description": "Do not use.",
+            "enum": [
+              "github",
+              "internal_docker",
+              "external_docker"
+            ]
+          },
+          "shareable": {
+            "type": "boolean",
+            "description": "Boolean flag to configure if a deployment is shareable through LangGraph Studio.\n\nThis field does not apply to BYOC deployments."
+          },
+          "region": {
+            "type": "string",
+            "description": "Region of deployment.\n\nRegion value is cloud provider specific."
+          },
+          "aws_account_id": {
+            "type": "string",
+            "description": "AWS account ID of BYOC deployment.\n\nThis field does not apply to non-BYOC deployments."
+          },
+          "aws_external_id": {
+            "type": "string",
+            "description": "Do not use."
+          }
+        }
+      },
+      "ResourceId": {
+        "type": "object",
+        "description": "Internal identifier for a `ResourceRevision` or `ResourceService`.",
+        "properties": {
+          "type": {
+            "type": "string",
+            "enum": [
+              "revisions",
+              "services"
+            ]
+          },
+          "name": {
+            "type": "string"
+          }
+        }
+      },
+      "ResourceRevision": {
+        "type": "object",
+        "description": "Internal revision resource for a `ResourceService`.",
+        "properties": {
+          "id": {
+            "$ref": "#/components/schemas/ResourceId"
+          },
+          "env_vars": {
+            "type": "array",
+            "items": {
+              "$ref": "#/components/schemas/EnvVar"
+            }
+          },
+          "hosted_langserve_revision_id": {
+            "type": "string",
+            "format": "uuid",
+            "description": "References `id` of a `Revision`."
+          }
+        }
+      },
+      "ResourceService": {
+        "type": "object",
+        "description": "Internal service resource for a `Project`.",
+        "properties": {
+          "id": {
+            "$ref": "#/components/schemas/ResourceId"
+          },
+          "url": {
+            "type": ["string", "null"],
+            "description": "URL of LangGraph Server deployment."
+          },
+          "latest_revision": {
+            "description": "References latest `ResourceRevision`.\n\nThe latest `ResourceRevision` may not be active if it's currently being deployed.",
+            "$ref": "#/components/schemas/ResourceRevision"
+          },
+          "latest_active_revision": {
+            "description": "References latest active `ResourceRevision`.\n\nThe latest active `ResourceRevision` is not always the latest `ResourceRevision`.",
+            "$ref": "#/components/schemas/ResourceRevision"
+          }
+        }
+      },
+      "Revision": {
+        "type": "object",
+        "description": "A revision is a version of a LangGraph Server deployment.\n\nDifferent revisions may contain different code and/or environment variables. A project can have many revisions.",
+        "properties": {
+          "id": {
+            "type": "string",
+            "format": "uuid",
+            "description": "ID of the revision.",
+            "required": true
+          },
+          "project_id": {
+            "type": "string",
+            "format": "uuid",
+            "description": "References `id` of `Project`.",
+            "required": true
+          },
+          "created_at": {
+            "type": "string",
+            "format": "date-time",
+            "description": "Timestamp of when the revision was created.",
+            "required": true
+          },
+          "updated_at": {
+            "type": "string",
+            "format": "date-time",
+            "description": "Timestamp of when the revision was updated.",
+            "required": true
+          },
+          "repo_path": {
+            "type": ["string", "null"],
+            "description": "Path to `langgraph.json` configuration file. For example, `langgraph.json` or `src/langgraph.json`.\n\nThis field only applies to deployments from a GitHub repository.",
+            "default": "null"
+          },
+          "repo_commit": {
+            "type": ["string", "null"],
+            "description": "Git branch name of deployment.\n\nThis field only applies to deployments from a GitHub repository.",
+            "default": "null"
+          },
+          "status": {
+            "type": "string",
+            "enum": [
+              "CREATING",
+              "AWAITING_BUILD",
+              "BUILDING",
+              "AWAITING_DEPLOY",
+              "DEPLOYING",
+              "CREATE_FAILED",
+              "BUILD_FAILED",
+              "DEPLOY_FAILED",
+              "DEPLOYED",
+              "INTERRUPTED",
+              "UNKNOWN"
+            ],
+            "description": "Deployment status of the revision.\n\nNon-terminal statuses: `CREATING`, `AWAITING_BUILD`, `BUILDING`, `AWAITING_DEPLOY`, `DEPLOYING`. All other statuses are terminal."
+          },
+          "status_message": {
+            "type": "string",
+            "description": "Message associated with the `status`."
+          },
+          "gcp_build_name": {
+            "type": ["string", "null"],
+            "description": "Do not use."
+          },
+          "metadata": {
+            "$ref": "#/components/schemas/RevisionMetadata"
+          },
+          "image_path": {
+            "type": ["string", "null"],
+            "description": "URI of the Docker image to deploy.\n\nThis field does not apply to deployments from a GitHub repository.",
+            "default": "null"
+          },
+          "container_spec": {
+            "$ref": "#/components/schemas/ContainerSpec"
+          },
+          "resource": {
+            "$ref": "#/components/schemas/ResourceRevision"
+          }
+        }
+      },
+      "RevisionMetadata": {
+        "type": "object",
+        "description": "Metadata associated with a `Revision`.",
+        "properties": {
+          "created_by": {
+            "type": "object",
+            "description": "Do not use."
+          },
+          "repo_commit_sha": {
+            "type": "string",
+            "description": "Git commit SHA of the deployment.\n\nThis field only applies to deployments from a GitHub repository."
+          }
+        }
+      }
+    }
+  }
+}

docs/docs/cloud/reference/cli.md

@@ -61,6 +61,7 @@ The LangGraph CLI requires a JSON configuration file with the following keys:
 All deployments come with a DB-backed BaseStore. Adding an "index" configuration to your `langgraph.json` will enable [semantic search](../deployment/semantic_search.md) within the BaseStore of your deployment.
 
 The `fields` configuration determines which parts of your documents to embed:
+
 - If omitted or set to `["$"]`, the entire document will be embedded
 - To embed specific fields, use JSON path notation: `["metadata.title", "content.text"]`
 - Documents missing specified fields will still be stored but won't have embeddings for those fields

docs/mkdocs.yml

@@ -2,6 +2,7 @@ site_name: ""
 site_description: Build language agents as graphs
 site_url: https://langchain-ai.github.io/langgraph/
 repo_url: https://github.com/langchain-ai/langgraph
+edit_uri: edit/main/docs/docs/
 theme:
   name: material
   custom_dir: overrides
@@ -16,6 +17,7 @@ theme:
     - content.code.copy
     - content.code.select
     - content.tabs.link
+    - content.action.edit
     - content.tooltips
     - header.autohide
     - navigation.expand

libs/langgraph/bench/fanout_to_subgraph.py

@@ -1,5 +1,7 @@
 import operator
-from typing import Annotated, TypedDict
+from typing import Annotated
+
+from typing_extensions import TypedDict
 
 from langgraph.constants import END, START, Send
 from langgraph.graph.state import StateGraph

libs/langgraph/langgraph/func/__init__.py

@@ -1,9 +1,9 @@
 import asyncio
 import concurrent
 import concurrent.futures
+import functools
 import inspect
 import types
-from functools import partial, update_wrapper
 from typing import (
     Any,
     Awaitable,
@@ -33,17 +33,17 @@ T = TypeVar("T")
 
 
 def call(
-    func: Callable[[P1], T],
-    input: P1,
-    *,
+    func: Callable[P, T],
+    *args: Any,
     retry: Optional[RetryPolicy] = None,
+    **kwargs: Any,
 ) -> concurrent.futures.Future[T]:
     from langgraph.constants import CONFIG_KEY_CALL
     from langgraph.utils.config import get_configurable
 
     conf = get_configurable()
     impl = conf[CONFIG_KEY_CALL]
-    fut = impl(func, input, retry=retry)
+    fut = impl(func, (args, kwargs), retry=retry)
     return fut
 
 
@@ -59,16 +59,51 @@ def task(  # type: ignore[overload-cannot-match]
 ) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ...
 
 
+@overload
 def task(
-    *, retry: Optional[RetryPolicy] = None
+    __func_or_none__: Callable[P, T],
+) -> Callable[P, concurrent.futures.Future[T]]: ...
+
+
+@overload
+def task(
+    __func_or_none__: Callable[P, Awaitable[T]],
+) -> Callable[P, asyncio.Future[T]]: ...
+
+
+def task(
+    __func_or_none__: Optional[Union[Callable[P, T], Callable[P, Awaitable[T]]]] = None,
+    *,
+    retry: Optional[RetryPolicy] = None,
 ) -> Union[
     Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]],
     Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]],
+    Callable[P, asyncio.Future[T]],
+    Callable[P, concurrent.futures.Future[T]],
 ]:
-    def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]:
-        return update_wrapper(partial(call, func, retry=retry), func)
+    def decorator(
+        func: Union[Callable[P, Awaitable[T]], Callable[P, T]],
+    ) -> Callable[P, concurrent.futures.Future[T]]:
+        if asyncio.iscoroutinefunction(func):
 
-    return _task
+            @functools.wraps(func)
+            async def _tick(__allargs__: tuple) -> T:
+                return await func(*__allargs__[0], **__allargs__[1])
+
+        else:
+
+            @functools.wraps(func)
+            def _tick(__allargs__: tuple) -> T:
+                return func(*__allargs__[0], **__allargs__[1])
+
+        return functools.update_wrapper(
+            functools.partial(call, _tick, retry=retry), func
+        )
+
+    if __func_or_none__ is not None:
+        return decorator(__func_or_none__)
+
+    return decorator
 
 
 def entrypoint(

libs/langgraph/langgraph/graph/message.py

@@ -8,7 +8,6 @@ from typing import (
     Literal,
     Optional,
     Sequence,
-    TypedDict,
     Union,
     cast,
 )
@@ -22,6 +21,7 @@ from langchain_core.messages import (
     convert_to_messages,
     message_chunk_to_message,
 )
+from typing_extensions import TypedDict
 
 from langgraph.graph.state import StateGraph
 

libs/langgraph/langgraph/prebuilt/chat_agent_executor.py

@@ -381,7 +381,7 @@ def create_react_agent(
         Add complex prompt with custom graph state:
 
         ```pycon
-        >>> from typing import TypedDict
+        >>> from typing_extensions import TypedDict
         >>>
         >>> from langgraph.managed import IsLastStep
         >>> prompt = ChatPromptTemplate.from_messages(

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -601,7 +601,8 @@ def tools_condition(
         >>> from langgraph.prebuilt import ToolNode, tools_condition
         >>> from langgraph.graph.message import add_messages
         ...
-        >>> from typing import TypedDict, Annotated
+        >>> from typing import Annotated
+        >>> from typing_extensions import TypedDict
         ...
         >>> @tool
         >>> def divide(a: float, b: float) -> int:

libs/langgraph/langgraph/prebuilt/tool_validator.py

@@ -74,7 +74,8 @@ class ValidationNode(RunnableCallable):
 
     Examples:
         Example usage for re-prompting the model to generate a valid response:
-        >>> from typing import Literal, Annotated, TypedDict
+        >>> from typing import Literal, Annotated
+        >>> from typing_extensions import TypedDict
         ...
         >>> from langchain_anthropic import ChatAnthropic
         >>> from pydantic import BaseModel, validator

libs/langgraph/langgraph/pregel/debug.py

@@ -10,13 +10,13 @@ from typing import (
     Mapping,
     Optional,
     Sequence,
-    TypedDict,
     Union,
 )
 from uuid import UUID
 
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.utils.input import get_bolded_text, get_colored_text
+from typing_extensions import TypedDict
 
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite

libs/langgraph/langgraph/pregel/io.py

@@ -1,3 +1,4 @@
+from collections import Counter
 from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypeVar, Union
 from uuid import UUID
 
@@ -181,12 +182,27 @@ def map_output_updates(
                 (task.name, value) for chan, value in writes if chan == output_channels
             )
         elif any(chan in output_channels for chan, _ in writes):
-            updated.append(
-                (
-                    task.name,
-                    {chan: value for chan, value in writes if chan in output_channels},
+            counts = Counter(chan for chan, _ in writes)
+            if any(counts[chan] > 1 for chan in output_channels):
+                updated.extend(
+                    (
+                        task.name,
+                        {chan: value},
+                    )
+                    for chan, value in writes
+                    if chan in output_channels
+                )
+            else:
+                updated.append(
+                    (
+                        task.name,
+                        {
+                            chan: value
+                            for chan, value in writes
+                            if chan in output_channels
+                        },
+                    )
                 )
-            )
     grouped: dict[str, list[Any]] = {t.name: [] for t, _ in output_tasks}
     for node, value in updated:
         grouped[node].append(value)

libs/langgraph/langgraph/pregel/loop.py

@@ -1032,6 +1032,13 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
         # unwind stack
-        return await asyncio.shield(
+        exit_task = asyncio.create_task(
             self.stack.__aexit__(exc_type, exc_value, traceback)
         )
+        try:
+            return await exit_task
+        except asyncio.CancelledError as e:
+            # Bubble up the exit task upon cancellation to permit the API
+            # consumer to await it before e.g., re-using the DB connection.
+            e.args = (*e.args, exit_task)
+            raise

libs/langgraph/langgraph/types.py

@@ -13,14 +13,13 @@ from typing import (
     Optional,
     Sequence,
     Type,
-    TypedDict,
     TypeVar,
     Union,
     cast,
 )
 
 from langchain_core.runnables import Runnable, RunnableConfig
-from typing_extensions import Self
+from typing_extensions import Self, TypedDict
 
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
@@ -373,7 +372,8 @@ def interrupt(value: Any) -> Any:
     Example:
         ```python
         import uuid
-        from typing import TypedDict, Optional
+        from typing import Optional
+        from typing_extensions import TypedDict
 
         from langgraph.checkpoint.memory import MemorySaver
         from langgraph.constants import START

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.60"
+version = "0.2.61"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
@@ -38,7 +38,7 @@ py-spy = "^0.3.14"
 types-requests = "^2.32.0.20240914"
 
 [tool.ruff]
-lint.select = [ "E", "F", "I" ]
+lint.select = [ "E", "F", "I", "TID251" ]
 lint.ignore = [ "E501" ]
 line-length = 88
 indent-width = 4
@@ -52,6 +52,9 @@ line-ending = "auto"
 docstring-code-format = false
 docstring-code-line-length = "dynamic"
 
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+"typing.TypedDict".msg = "Use typing_extensions.TypedDict instead."
+
 [tool.mypy]
 # https://mypy.readthedocs.io/en/stable/config_file.html
 disallow_untyped_defs = "True"

libs/langgraph/tests/test_interruption.py

@@ -1,7 +1,6 @@
-from typing import TypedDict
-
 import pytest
 from pytest_mock import MockerFixture
+from typing_extensions import TypedDict
 
 from langgraph.graph import END, START, StateGraph
 from tests.conftest import (

libs/langgraph/tests/test_large_cases.py

@@ -4,13 +4,14 @@ import re
 import time
 from contextlib import contextmanager
 from dataclasses import replace
-from typing import Annotated, Any, Iterator, Literal, Optional, TypedDict, Union, cast
+from typing import Annotated, Any, Iterator, Literal, Optional, Union, cast
 
 import httpx
 import pytest
 from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue

libs/langgraph/tests/test_large_cases_async.py

@@ -9,7 +9,6 @@ from typing import (
     AsyncIterator,
     Literal,
     Optional,
-    TypedDict,
     Union,
     cast,
 )
@@ -21,6 +20,7 @@ from langchain_core.runnables import RunnableConfig, RunnablePick
 from pydantic import BaseModel
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue

libs/langgraph/tests/test_pregel.py

@@ -21,7 +21,6 @@ from typing import (
     Optional,
     Sequence,
     Tuple,
-    TypedDict,
     Union,
     get_type_hints,
 )
@@ -36,6 +35,7 @@ from langchain_core.runnables import (
 from langsmith import traceable
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
@@ -1515,27 +1515,32 @@ def test_imp_stream_order(
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
 
     @task()
-    def foo(state: dict) -> dict:
-        return {"a": state["a"] + "foo", "b": "bar"}
+    def foo(state: dict) -> tuple:
+        return state["a"] + "foo", "bar"
 
-    @task()
-    def bar(state: dict) -> dict:
-        return {"a": state["a"] + state["b"], "c": "bark"}
+    @task
+    def bar(a: str, b: str, c: Optional[str] = None) -> dict:
+        return {"a": a + b, "c": (c or "") + "bark"}
 
-    @task()
+    @task
     def baz(state: dict) -> dict:
         return {"a": state["a"] + "baz", "c": "something else"}
 
     @entrypoint(checkpointer=checkpointer)
     def graph(state: dict) -> dict:
         fut_foo = foo(state)
-        fut_bar = bar(fut_foo.result())
+        fut_bar = bar(*fut_foo.result())
         fut_baz = baz(fut_bar.result())
         return fut_baz.result()
 
     thread1 = {"configurable": {"thread_id": "1"}}
     assert [c for c in graph.stream({"a": "0"}, thread1)] == [
-        {"foo": {"a": "0foo", "b": "bar"}},
+        {
+            "foo": (
+                "0foo",
+                "bar",
+            )
+        },
         {"bar": {"a": "0foobar", "c": "bark"}},
         {"baz": {"a": "0foobarbaz", "c": "something else"}},
         {"graph": {"a": "0foobarbaz", "c": "something else"}},
@@ -4168,10 +4173,12 @@ def test_store_injected(
         def __call__(self, inputs: State, config: RunnableConfig, store: BaseStore):
             assert isinstance(store, BaseStore)
             store.put(
-                namespace
-                if self.i is not None
-                and config["configurable"]["thread_id"] in (thread_1, thread_2)
-                else (f"foo_{self.i}", "bar"),
+                (
+                    namespace
+                    if self.i is not None
+                    and config["configurable"]["thread_id"] in (thread_1, thread_2)
+                    else (f"foo_{self.i}", "bar")
+                ),
                 doc_id,
                 {
                     **doc,
@@ -5242,3 +5249,54 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name:
     # Verify the error was recorded in checkpoint
     failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error)
     assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error
+
+
+def test_multiple_updates_root() -> None:
+    def node_a(state):
+        return [Command(update="a1"), Command(update="a2")]
+
+    def node_b(state):
+        return "b"
+
+    graph = (
+        StateGraph(Annotated[str, operator.add])
+        .add_sequence([node_a, node_b])
+        .add_edge(START, "node_a")
+        .compile()
+    )
+
+    assert graph.invoke("") == "a1a2b"
+
+    # only streams the last update from node_a
+    assert [c for c in graph.stream("", stream_mode="updates")] == [
+        {"node_a": ["a1", "a2"]},
+        {"node_b": "b"},
+    ]
+
+
+def test_multiple_updates() -> None:
+    class State(TypedDict):
+        foo: Annotated[str, operator.add]
+
+    def node_a(state):
+        return [Command(update={"foo": "a1"}), Command(update={"foo": "a2"})]
+
+    def node_b(state):
+        return {"foo": "b"}
+
+    graph = (
+        StateGraph(State)
+        .add_sequence([node_a, node_b])
+        .add_edge(START, "node_a")
+        .compile()
+    )
+
+    assert graph.invoke({"foo": ""}) == {
+        "foo": "a1a2b",
+    }
+
+    # only streams the last update from node_a
+    assert [c for c in graph.stream({"foo": ""}, stream_mode="updates")] == [
+        {"node_a": [{"foo": "a1"}, {"foo": "a2"}]},
+        {"node_b": {"foo": "b"}},
+    ]

libs/langgraph/tests/test_pregel_async.py

@@ -19,7 +19,6 @@ from typing import (
     Literal,
     Optional,
     Tuple,
-    TypedDict,
     Union,
 )
 from uuid import UUID
@@ -34,6 +33,7 @@ from langchain_core.runnables import (
 from langchain_core.utils.aiter import aclosing
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
+from typing_extensions import TypedDict
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
@@ -180,6 +180,262 @@ async def test_checkpoint_errors() -> None:
             pass
 
 
+async def test_py_async_with_cancel_behavior() -> None:
+    """This test confirms that in all versions of Python we support, __aexit__
+    is not cancelled when the coroutine containing the async with block is cancelled."""
+
+    logs: list[str] = []
+
+    class MyContextManager:
+        async def __aenter__(self):
+            logs.append("Entering")
+            return self
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            logs.append("Starting exit")
+            try:
+                # Simulate some cleanup work
+                await asyncio.sleep(2)
+                logs.append("Cleanup completed")
+            except asyncio.CancelledError:
+                logs.append("Cleanup was cancelled!")
+                raise
+            logs.append("Exit finished")
+
+    async def main():
+        try:
+            async with MyContextManager():
+                logs.append("In context")
+                await asyncio.sleep(1)
+                logs.append("This won't print if cancelled")
+        except asyncio.CancelledError:
+            logs.append("Context was cancelled")
+            raise
+
+    # create task
+    t = asyncio.create_task(main())
+    # cancel after 0.2 seconds
+    await asyncio.sleep(0.2)
+    t.cancel()
+    # check logs before cancellation is handled
+    assert logs == [
+        "Entering",
+        "In context",
+    ], "Cancelled before cleanup started"
+    # wait for task to finish
+    try:
+        await t
+    except asyncio.CancelledError:
+        # check logs after cancellation is handled
+        assert logs == [
+            "Entering",
+            "In context",
+            "Starting exit",
+            "Cleanup completed",
+            "Exit finished",
+            "Context was cancelled",
+        ], "Cleanup started and finished after cancellation"
+    else:
+        assert False, "Task should be cancelled"
+
+
+async def test_checkpoint_put_after_cancellation() -> None:
+    logs: list[str] = []
+
+    class LongPutCheckpointer(MemorySaver):
+        async def aput(
+            self,
+            config: RunnableConfig,
+            checkpoint: Checkpoint,
+            metadata: CheckpointMetadata,
+            new_versions: ChannelVersions,
+        ) -> RunnableConfig:
+            logs.append("checkpoint.aput.start")
+            try:
+                await asyncio.sleep(1)
+                return await super().aput(config, checkpoint, metadata, new_versions)
+            finally:
+                logs.append("checkpoint.aput.end")
+
+    inner_task_cancelled = False
+
+    async def awhile(input: Any) -> None:
+        logs.append("awhile.start")
+        try:
+            await asyncio.sleep(1)
+        except asyncio.CancelledError:
+            nonlocal inner_task_cancelled
+            inner_task_cancelled = True
+            raise
+        finally:
+            logs.append("awhile.end")
+
+    builder = Graph()
+    builder.add_node("agent", awhile)
+    builder.set_entry_point("agent")
+    builder.set_finish_point("agent")
+
+    graph = builder.compile(checkpointer=LongPutCheckpointer())
+    thread1 = {"configurable": {"thread_id": "1"}}
+
+    # start the task
+    t = asyncio.create_task(graph.ainvoke(1, thread1))
+    # cancel after 0.2 seconds
+    await asyncio.sleep(0.2)
+    t.cancel()
+    # check logs before cancellation is handled
+    assert sorted(logs) == [
+        "awhile.start",
+        "checkpoint.aput.start",
+    ], "Cancelled before checkpoint put started"
+    # wait for task to finish
+    try:
+        await t
+    except asyncio.CancelledError:
+        # check logs after cancellation is handled
+        assert sorted(logs) == [
+            "awhile.end",
+            "awhile.start",
+            "checkpoint.aput.end",
+            "checkpoint.aput.start",
+        ], "Checkpoint put is not cancelled"
+    else:
+        assert False, "Task should be cancelled"
+
+
+async def test_checkpoint_put_after_cancellation_stream_anext() -> None:
+    logs: list[str] = []
+
+    class LongPutCheckpointer(MemorySaver):
+        async def aput(
+            self,
+            config: RunnableConfig,
+            checkpoint: Checkpoint,
+            metadata: CheckpointMetadata,
+            new_versions: ChannelVersions,
+        ) -> RunnableConfig:
+            logs.append("checkpoint.aput.start")
+            try:
+                await asyncio.sleep(1)
+                return await super().aput(config, checkpoint, metadata, new_versions)
+            finally:
+                logs.append("checkpoint.aput.end")
+
+    inner_task_cancelled = False
+
+    async def awhile(input: Any) -> None:
+        logs.append("awhile.start")
+        try:
+            await asyncio.sleep(1)
+        except asyncio.CancelledError:
+            nonlocal inner_task_cancelled
+            inner_task_cancelled = True
+            raise
+        finally:
+            logs.append("awhile.end")
+
+    builder = Graph()
+    builder.add_node("agent", awhile)
+    builder.set_entry_point("agent")
+    builder.set_finish_point("agent")
+
+    graph = builder.compile(checkpointer=LongPutCheckpointer())
+    thread1 = {"configurable": {"thread_id": "1"}}
+
+    # start the task
+    s = graph.astream(1, thread1)
+    t = asyncio.create_task(s.__anext__())
+    # cancel after 0.2 seconds
+    await asyncio.sleep(0.2)
+    t.cancel()
+    # check logs before cancellation is handled
+    assert sorted(logs) == [
+        "awhile.start",
+        "checkpoint.aput.start",
+    ], "Cancelled before checkpoint put started"
+    # wait for task to finish
+    try:
+        await t
+    except asyncio.CancelledError:
+        # check logs after cancellation is handled
+        assert sorted(logs) == [
+            "awhile.end",
+            "awhile.start",
+            "checkpoint.aput.end",
+            "checkpoint.aput.start",
+        ], "Checkpoint put is not cancelled"
+    else:
+        assert False, "Task should be cancelled"
+
+
+async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
+    logs: list[str] = []
+
+    class LongPutCheckpointer(MemorySaver):
+        async def aput(
+            self,
+            config: RunnableConfig,
+            checkpoint: Checkpoint,
+            metadata: CheckpointMetadata,
+            new_versions: ChannelVersions,
+        ) -> RunnableConfig:
+            logs.append("checkpoint.aput.start")
+            try:
+                await asyncio.sleep(1)
+                return await super().aput(config, checkpoint, metadata, new_versions)
+            finally:
+                logs.append("checkpoint.aput.end")
+
+    inner_task_cancelled = False
+
+    async def awhile(input: Any) -> None:
+        logs.append("awhile.start")
+        try:
+            await asyncio.sleep(1)
+        except asyncio.CancelledError:
+            nonlocal inner_task_cancelled
+            inner_task_cancelled = True
+            raise
+        finally:
+            logs.append("awhile.end")
+
+    builder = Graph()
+    builder.add_node("agent", awhile)
+    builder.set_entry_point("agent")
+    builder.set_finish_point("agent")
+
+    graph = builder.compile(checkpointer=LongPutCheckpointer())
+    thread1 = {"configurable": {"thread_id": "1"}}
+
+    # start the task
+    s = graph.astream_events(1, thread1, version="v2", include_names=["LangGraph"])
+    # skip first event (happens right away)
+    await s.__anext__()
+    # start the task for 2nd event
+    t = asyncio.create_task(s.__anext__())
+    # cancel after 0.2 seconds
+    await asyncio.sleep(0.2)
+    t.cancel()
+    # check logs before cancellation is handled
+    assert logs == [
+        "checkpoint.aput.start",
+        "awhile.start",
+    ], "Cancelled before checkpoint put started"
+    # wait for task to finish
+    try:
+        await t
+    except asyncio.CancelledError:
+        # check logs after cancellation is handled
+        assert logs == [
+            "checkpoint.aput.start",
+            "awhile.start",
+            "awhile.end",
+            "checkpoint.aput.end",
+        ], "Checkpoint put is not cancelled"
+    else:
+        assert False, "Task should be cancelled"
+
+
 async def test_node_cancellation_on_external_cancel() -> None:
     inner_task_cancelled = False
 
@@ -2315,9 +2571,9 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None:
         def foo(state: dict) -> dict:
             return {"a": state["a"] + "foo", "b": "bar"}
 
-        @task()
-        def bar(state: dict) -> dict:
-            return {"a": state["a"] + state["b"], "c": "bark"}
+        @task
+        def bar(a: str, b: str, c: Optional[str] = None) -> dict:
+            return {"a": a + b, "c": (c or "") + "bark"}
 
         @task()
         def baz(state: dict) -> dict:
@@ -2325,8 +2581,8 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None:
 
         @entrypoint(checkpointer=checkpointer)
         def graph(state: dict) -> dict:
-            fut_foo = foo(state)
-            fut_bar = bar(fut_foo.result())
+            foo_result = foo(state).result()
+            fut_bar = bar(foo_result["a"], foo_result["b"])
             fut_baz = baz(fut_bar.result())
             return fut_baz.result()
 
@@ -2351,9 +2607,9 @@ async def test_imp_stream_order(checkpointer_name: str) -> None:
         async def foo(state: dict) -> dict:
             return {"a": state["a"] + "foo", "b": "bar"}
 
-        @task()
-        async def bar(state: dict) -> dict:
-            return {"a": state["a"] + state["b"], "c": "bark"}
+        @task
+        async def bar(a: str, b: str, c: Optional[str] = None) -> dict:
+            return {"a": a + b, "c": (c or "") + "bark"}
 
         @task()
         async def baz(state: dict) -> dict:
@@ -2361,8 +2617,9 @@ async def test_imp_stream_order(checkpointer_name: str) -> None:
 
         @entrypoint(checkpointer=checkpointer)
         async def graph(state: dict) -> dict:
-            fut_foo = foo(state)
-            fut_bar = bar(await fut_foo)
+            foo_res = await foo(state)
+
+            fut_bar = bar(foo_res["a"], foo_res["b"])
             fut_baz = baz(await fut_bar)
             return await fut_baz
 
@@ -6362,3 +6619,54 @@ async def test_checkpoint_recovery_async(checkpointer_name: str):
         # Verify the error was recorded in checkpoint
         failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error)
         assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error
+
+
+async def test_multiple_updates_root() -> None:
+    def node_a(state):
+        return [Command(update="a1"), Command(update="a2")]
+
+    def node_b(state):
+        return "b"
+
+    graph = (
+        StateGraph(Annotated[str, operator.add])
+        .add_sequence([node_a, node_b])
+        .add_edge(START, "node_a")
+        .compile()
+    )
+
+    assert await graph.ainvoke("") == "a1a2b"
+
+    # only streams the last update from node_a
+    assert [c async for c in graph.astream("", stream_mode="updates")] == [
+        {"node_a": ["a1", "a2"]},
+        {"node_b": "b"},
+    ]
+
+
+async def test_multiple_updates() -> None:
+    class State(TypedDict):
+        foo: Annotated[str, operator.add]
+
+    def node_a(state):
+        return [Command(update={"foo": "a1"}), Command(update={"foo": "a2"})]
+
+    def node_b(state):
+        return {"foo": "b"}
+
+    graph = (
+        StateGraph(State)
+        .add_sequence([node_a, node_b])
+        .add_edge(START, "node_a")
+        .compile()
+    )
+
+    assert await graph.ainvoke({"foo": ""}) == {
+        "foo": "a1a2b",
+    }
+
+    # only streams the last update from node_a
+    assert [c async for c in graph.astream({"foo": ""}, stream_mode="updates")] == [
+        {"node_a": [{"foo": "a1"}, {"foo": "a2"}]},
+        {"node_b": {"foo": "b"}},
+    ]

libs/langgraph/tests/test_tracing_interops.py

@@ -1,13 +1,14 @@
 import json
 import sys
 import time
-from typing import Any, Callable, Tuple, TypedDict, TypeVar
+from typing import Any, Callable, Tuple, TypeVar
 from unittest.mock import MagicMock
 
 import langsmith as ls
 import pytest
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tracers import LangChainTracer
+from typing_extensions import TypedDict
 
 from langgraph.graph import StateGraph
 

libs/langgraph/tests/test_utils.py

@@ -9,7 +9,6 @@ from typing import (
     List,
     Literal,
     Optional,
-    TypedDict,
     TypeVar,
     Union,
 )
@@ -17,7 +16,7 @@ from unittest.mock import patch
 
 import langsmith
 import pytest
-from typing_extensions import Annotated, NotRequired, Required
+from typing_extensions import Annotated, NotRequired, Required, TypedDict
 
 from langgraph.graph import END, StateGraph
 from langgraph.graph.graph import CompiledGraph

libs/sdk-js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.32",
+  "version": "0.0.33",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",

libs/sdk-js/src/types.ts

@@ -32,7 +32,7 @@ export interface Command {
   /**
    * An object to update the thread state with.
    */
-  update?: Record<string, unknown>;
+  update?: Record<string, unknown> | [string, unknown][];
 
   /**
    * The value to return from an `interrupt` function call.

libs/sdk-py/langgraph_sdk/auth/__init__.py

@@ -461,6 +461,73 @@ class _CronsOn(
     Search = types.CronsSearch
 
 
+class _StoreOn:
+    def __init__(self, auth: Auth) -> None:
+        self._auth = auth
+
+    @typing.overload
+    def __call__(
+        self,
+        *,
+        actions: typing.Optional[
+            typing.Union[
+                typing.Literal["put", "get", "search", "list_namespaces", "delete"],
+                Sequence[
+                    typing.Literal["put", "get", "search", "list_namespaces", "delete"]
+                ],
+            ]
+        ] = None,
+    ) -> Callable[[AHO], AHO]: ...
+
+    @typing.overload
+    def __call__(self, fn: AHO) -> AHO: ...
+
+    def __call__(
+        self,
+        fn: typing.Optional[AHO] = None,
+        *,
+        actions: typing.Optional[
+            typing.Union[
+                typing.Literal["put", "get", "search", "list_namespaces", "delete"],
+                Sequence[
+                    typing.Literal["put", "get", "search", "list_namespaces", "delete"]
+                ],
+            ]
+        ] = None,
+    ) -> typing.Union[AHO, Callable[[AHO], AHO]]:
+        """Register a handler for specific resources and actions.
+
+        Can be used as a decorator or with explicit resource/action parameters:
+
+        @auth.on.store
+        async def handler(): ... # Handle all store ops
+
+        @auth.on.store(actions=("put", "get", "search", "delete"))
+        async def handler(): ... # Handle specific store ops
+
+        @auth.on.store.put
+        async def handler(): ... # Handle store.put ops
+        """
+        if fn is not None:
+            # Used as a plain decorator
+            _register_handler(self._auth, None, None, fn)
+            return fn
+
+        # Used with parameters, return a decorator
+        def decorator(
+            handler: AHO,
+        ) -> AHO:
+            if isinstance(actions, str):
+                action_list = [actions]
+            else:
+                action_list = list(actions) if actions is not None else ["*"]
+            for action in action_list:
+                _register_handler(self._auth, "store", action, handler)
+            return handler
+
+        return decorator
+
+
 AHO = typing.TypeVar("AHO", bound=_ActionHandler[dict[str, typing.Any]])
 
 
@@ -524,6 +591,7 @@ class _On:
         "threads",
         "runs",
         "crons",
+        "store",
         "value",
     )
 
@@ -532,6 +600,7 @@ class _On:
         self.assistants = _AssistantsOn(auth, "assistants")
         self.threads = _ThreadsOn(auth, "threads")
         self.crons = _CronsOn(auth, "crons")
+        self.store = _StoreOn(auth)
         self.value = dict[str, typing.Any]
 
     @typing.overload

libs/sdk-py/langgraph_sdk/auth/types.py

@@ -5,7 +5,7 @@ request handling in LangGraph. It includes user protocols, authentication contex
 and typed dictionaries for various API operations.
 
 Note:
-    All typing.TypedDict classes use total=False to make all fields optional by default.
+    All typing.TypedDict classes use total=False to make all fields typing.Optional by default.
 """
 
 import functools
@@ -157,7 +157,7 @@ class MinimalUserDict(typing.TypedDict, total=False):
     identity: typing_extensions.Required[str]
     """The required unique identifier for the user."""
     display_name: str
-    """The optional display name for the user."""
+    """The typing.Optional display name for the user."""
     is_authenticated: bool
     """Whether the user is authenticated. Defaults to True."""
     permissions: Sequence[str]
@@ -358,11 +358,34 @@ class AuthContext(BaseAuthContext):
     allowing for fine-grained access control decisions.
     """
 
-    resource: typing.Literal["runs", "threads", "crons", "assistants"]
+    resource: typing.Literal["runs", "threads", "crons", "assistants", "store"]
     """The resource being accessed."""
 
-    action: typing.Literal["create", "read", "update", "delete", "search", "create_run"]
-    """The action being performed on the resource."""
+    action: typing.Literal[
+        "create",
+        "read",
+        "update",
+        "delete",
+        "search",
+        "create_run",
+        "put",
+        "get",
+        "list_namespaces",
+    ]
+    """The action being performed on the resource.
+    
+    Most resources support the following actions:
+    - create: Create a new resource
+    - read: Read information about a resource
+    - update: Update an existing resource
+    - delete: Delete a resource
+    - search: Search for resources
+
+    The store supports the following actions:
+    - put: Add or update a document in the store
+    - get: Get a document from the store
+    - list_namespaces: List the namespaces in the store
+    """
 
 
 class ThreadsCreate(typing.TypedDict, total=False):
@@ -759,6 +782,93 @@ class CronsSearch(typing.TypedDict, total=False):
     """Offset for pagination."""
 
 
+class StoreGet(typing.TypedDict):
+    """Operation to retrieve a specific item by its namespace and key."""
+
+    namespace: tuple[str, ...]
+    """Hierarchical path that uniquely identifies the item's location."""
+
+    key: str
+    """Unique identifier for the item within its specific namespace."""
+
+
+class StoreSearch(typing.TypedDict):
+    """Operation to search for items within a specified namespace hierarchy."""
+
+    namespace_prefix: tuple[str, ...]
+    """Hierarchical path prefix defining the search scope.
+
+    ???+ example "Examples"
+
+        ```python
+        ()  # Search entire store
+        ("documents",)  # Search all documents
+        ("users", "content")  # Search within user content
+        ```
+    """
+
+    filter: typing.Optional[dict[str, typing.Any]]
+    """Key-value pairs for filtering results based on exact matches or comparison operators."""
+
+    limit: int
+    """Maximum number of items to return in the search results."""
+
+    offset: int
+    """Number of matching items to skip for pagination."""
+
+    query: typing.Optional[str]
+    """Naturalj language search query for semantic search capabilities."""
+
+
+class StoreListNamespaces(typing.TypedDict):
+    """Operation to list and filter namespaces in the store."""
+
+    prefix: typing.Optional[tuple[str, ...]]
+    """Optional conditions for filtering namespaces."""
+
+    suffix: typing.Optional[tuple[str, ...]]
+    """Optional conditions for filtering namespaces."""
+
+    max_depth: typing.Optional[int]
+    """Maximum depth of namespace hierarchy to return.
+
+    Note:
+        Namespaces deeper than this level will be truncated.
+    """
+
+    limit: int
+    """Maximum number of namespaces to return."""
+
+    offset: int
+    """Number of namespaces to skip for pagination."""
+
+
+class StorePut(typing.TypedDict):
+    """Operation to store, update, or delete an item in the store."""
+
+    namespace: tuple[str, ...]
+    """Hierarchical path that identifies the location of the item."""
+
+    key: str
+    """Unique identifier for the item within its namespace."""
+
+    value: typing.Optional[dict[str, typing.Any]]
+    """The data to store, or None to mark the item for deletion."""
+
+    index: typing.Optional[typing.Union[typing.Literal[False], list[str]]]
+    """Optional index configuration for full-text search."""
+
+
+class StoreDelete(typing.TypedDict):
+    """Operation to delete an item from the store."""
+
+    namespace: tuple[str, ...]
+    """Hierarchical path that uniquely identifies the item's location."""
+
+    key: str
+    """Unique identifier for the item within its specific namespace."""
+
+
 class on:
     """Namespace for type definitions of different API operations.
 
@@ -894,6 +1004,38 @@ class on:
 
             value = CronsSearch
 
+    class store:
+        """Types for store-related operations."""
+
+        value = typing.Union[
+            StoreGet, StoreSearch, StoreListNamespaces, StorePut, StoreDelete
+        ]
+
+        class put:
+            """Type for store put parameters."""
+
+            value = StorePut
+
+        class get:
+            """Type for store get parameters."""
+
+            value = StoreGet
+
+        class search:
+            """Type for store search parameters."""
+
+            value = StoreSearch
+
+        class delete:
+            """Type for store delete parameters."""
+
+            value = StoreDelete
+
+        class list_namespaces:
+            """Type for store list namespaces parameters."""
+
+            value = StoreListNamespaces
+
 
 __all__ = [
     "on",
@@ -909,4 +1051,9 @@ __all__ = [
     "AssistantsUpdate",
     "AssistantsDelete",
     "AssistantsSearch",
+    "StoreGet",
+    "StoreSearch",
+    "StoreListNamespaces",
+    "StorePut",
+    "StoreDelete",
 ]

libs/sdk-py/langgraph_sdk/client.py

@@ -1779,7 +1779,7 @@ class RunsClient:
 
         Args:
             thread_id: The thread ID to cancel.
-            run_id: The run ID to cancek.
+            run_id: The run ID to cancel.
             wait: Whether to wait until run has completed.
             action: Action to take when cancelling the run. Possible values
                 are `interrupt` or `rollback`. Default is `interrupt`.
@@ -3917,7 +3917,7 @@ class SyncRunsClient:
 
         Args:
             thread_id: The thread ID to cancel.
-            run_id: The run ID to cancek.
+            run_id: The run ID to cancel.
             wait: Whether to wait until run has completed.
             action: Action to take when cancelling the run. Possible values
                 are `interrupt` or `rollback`. Default is `interrupt`.

libs/sdk-py/langgraph_sdk/schema.py

@@ -1,7 +1,17 @@
 """Data models for interacting with the LangGraph API."""
 
 from datetime import datetime
-from typing import Any, Dict, Literal, NamedTuple, Optional, Sequence, TypedDict, Union
+from typing import (
+    Any,
+    Dict,
+    Literal,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    TypedDict,
+    Union,
+)
 
 Json = Optional[dict[str, Any]]
 """Represents a JSON-like structure, which can be None or a dictionary with string keys and any values."""
@@ -374,5 +384,5 @@ class Send(TypedDict):
 
 class Command(TypedDict, total=False):
     goto: Union[Send, str, Sequence[Union[Send, str]]]
-    update: dict[str, Any]
+    update: Union[dict[str, Any], Sequence[Tuple[str, Any]]]
     resume: Any

libs/sdk-py/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-sdk"
-version = "0.1.48"
+version = "0.1.50"
 description = "SDK for interacting with LangGraph API"
 authors = []
 license = "MIT"