langchain 1.0.5__py3-none-any.whl → 1.2.3__py3-none-any.whl

langchain/agents/middleware/_retry.py

@@ -0,0 +1,123 @@
+"""Shared retry utilities for agent middleware.
+
+This module contains common constants, utilities, and logic used by both
+model and tool retry middleware implementations.
+"""
+
+from __future__ import annotations
+
+import random
+from collections.abc import Callable
+from typing import Literal
+
+# Type aliases
+RetryOn = tuple[type[Exception], ...] | Callable[[Exception], bool]
+"""Type for specifying which exceptions to retry on.
+
+Can be either:
+- A tuple of exception types to retry on (based on `isinstance` checks)
+- A callable that takes an exception and returns `True` if it should be retried
+"""
+
+OnFailure = Literal["error", "continue"] | Callable[[Exception], str]
+"""Type for specifying failure handling behavior.
+
+Can be either:
+- A literal action string (`'error'` or `'continue'`)
+    - `'error'`: Re-raise the exception, stopping agent execution.
+    - `'continue'`: Inject a message with the error details, allowing the agent to continue.
+       For tool retries, a `ToolMessage` with the error details will be injected.
+       For model retries, an `AIMessage` with the error details will be returned.
+- A callable that takes an exception and returns a string for error message content
+"""
+
+
+def validate_retry_params(
+    max_retries: int,
+    initial_delay: float,
+    max_delay: float,
+    backoff_factor: float,
+) -> None:
+    """Validate retry parameters.
+
+    Args:
+        max_retries: Maximum number of retry attempts.
+        initial_delay: Initial delay in seconds before first retry.
+        max_delay: Maximum delay in seconds between retries.
+        backoff_factor: Multiplier for exponential backoff.
+
+    Raises:
+        ValueError: If any parameter is invalid (negative values).
+    """
+    if max_retries < 0:
+        msg = "max_retries must be >= 0"
+        raise ValueError(msg)
+    if initial_delay < 0:
+        msg = "initial_delay must be >= 0"
+        raise ValueError(msg)
+    if max_delay < 0:
+        msg = "max_delay must be >= 0"
+        raise ValueError(msg)
+    if backoff_factor < 0:
+        msg = "backoff_factor must be >= 0"
+        raise ValueError(msg)
+
+
+def should_retry_exception(
+    exc: Exception,
+    retry_on: RetryOn,
+) -> bool:
+    """Check if an exception should trigger a retry.
+
+    Args:
+        exc: The exception that occurred.
+        retry_on: Either a tuple of exception types to retry on, or a callable
+            that takes an exception and returns `True` if it should be retried.
+
+    Returns:
+        `True` if the exception should be retried, `False` otherwise.
+    """
+    if callable(retry_on):
+        return retry_on(exc)
+    return isinstance(exc, retry_on)
+
+
+def calculate_delay(
+    retry_number: int,
+    *,
+    backoff_factor: float,
+    initial_delay: float,
+    max_delay: float,
+    jitter: bool,
+) -> float:
+    """Calculate delay for a retry attempt with exponential backoff and optional jitter.
+
+    Args:
+        retry_number: The retry attempt number (0-indexed).
+        backoff_factor: Multiplier for exponential backoff.
+
+            Set to `0.0` for constant delay.
+        initial_delay: Initial delay in seconds before first retry.
+        max_delay: Maximum delay in seconds between retries.
+
+            Caps exponential backoff growth.
+        jitter: Whether to add random jitter to delay to avoid thundering herd.
+
+    Returns:
+        Delay in seconds before next retry.
+    """
+    if backoff_factor == 0.0:
+        delay = initial_delay
+    else:
+        delay = initial_delay * (backoff_factor**retry_number)
+
+    # Cap at max_delay
+    delay = min(delay, max_delay)
+
+    if jitter and delay > 0:
+        jitter_amount = delay * 0.25  # ±25% jitter
+        delay += random.uniform(-jitter_amount, jitter_amount)  # noqa: S311
+        # Ensure delay is not negative after jitter
+        delay = max(0, delay)
+
+    return delay

langchain/agents/middleware/context_editing.py

@@ -1,14 +1,16 @@
 """Context editing middleware.
 
-This middleware mirrors Anthropic's context editing capabilities by clearing
-older tool results once the conversation grows beyond a configurable token
-threshold. The implementation is intentionally model-agnostic so it can be used
-with any LangChain chat model.
+Mirrors Anthropic's context editing capabilities by clearing older tool results once the
+conversation grows beyond a configurable token threshold.
+
+The implementation is intentionally model-agnostic so it can be used with any LangChain
+chat model.
 """
 
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable, Iterable, Sequence
+from copy import deepcopy
 from dataclasses import dataclass
 from typing import Literal
 
@@ -16,7 +18,6 @@ from langchain_core.messages import (
     AIMessage,
     AnyMessage,
     BaseMessage,
-    SystemMessage,
     ToolMessage,
 )
 from langchain_core.messages.utils import count_tokens_approximately
@@ -182,11 +183,13 @@ class ClearToolUsesEdit(ContextEdit):
 
 
 class ContextEditingMiddleware(AgentMiddleware):
-    """Automatically prunes tool results to manage context size.
+    """Automatically prune tool results to manage context size.
+
+    The middleware applies a sequence of edits when the total input token count exceeds
+    configured thresholds.
 
-    The middleware applies a sequence of edits when the total input token count
-    exceeds configured thresholds. Currently the `ClearToolUsesEdit` strategy is
-    supported, aligning with Anthropic's `clear_tool_uses_20250919` behaviour.
+    Currently the `ClearToolUsesEdit` strategy is supported, aligning with Anthropic's
+    `clear_tool_uses_20250919` behavior [(read more)](https://platform.claude.com/docs/en/agents-and-tools/tool-use/memory-tool).
     """
 
     edits: list[ContextEdit]
@@ -198,11 +201,12 @@ class ContextEditingMiddleware(AgentMiddleware):
         edits: Iterable[ContextEdit] | None = None,
         token_count_method: Literal["approximate", "model"] = "approximate",  # noqa: S107
     ) -> None:
-        """Initializes a context editing middleware instance.
+        """Initialize an instance of context editing middleware.
 
         Args:
-            edits: Sequence of edit strategies to apply. Defaults to a single
-                `ClearToolUsesEdit` mirroring Anthropic defaults.
+            edits: Sequence of edit strategies to apply.
+
+                Defaults to a single `ClearToolUsesEdit` mirroring Anthropic defaults.
             token_count_method: Whether to use approximate token counting
                 (faster, less accurate) or exact counting implemented by the
                 chat model (potentially slower, more accurate).
@@ -224,20 +228,20 @@ class ContextEditingMiddleware(AgentMiddleware):
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return count_tokens_approximately(messages)
+
         else:
-            system_msg = (
-                [SystemMessage(content=request.system_prompt)] if request.system_prompt else []
-            )
+            system_msg = [request.system_message] if request.system_message else []
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return request.model.get_num_tokens_from_messages(
                     system_msg + list(messages), request.tools
                 )
 
+        edited_messages = deepcopy(list(request.messages))
         for edit in self.edits:
-            edit.apply(request.messages, count_tokens=count_tokens)
+            edit.apply(edited_messages, count_tokens=count_tokens)
 
-        return handler(request)
+        return handler(request.override(messages=edited_messages))
 
     async def awrap_model_call(
         self,
@@ -252,20 +256,20 @@ class ContextEditingMiddleware(AgentMiddleware):
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return count_tokens_approximately(messages)
+
         else:
-            system_msg = (
-                [SystemMessage(content=request.system_prompt)] if request.system_prompt else []
-            )
+            system_msg = [request.system_message] if request.system_message else []
 
             def count_tokens(messages: Sequence[BaseMessage]) -> int:
                 return request.model.get_num_tokens_from_messages(
                     system_msg + list(messages), request.tools
                 )
 
+        edited_messages = deepcopy(list(request.messages))
         for edit in self.edits:
-            edit.apply(request.messages, count_tokens=count_tokens)
+            edit.apply(edited_messages, count_tokens=count_tokens)
 
-        return await handler(request)
+        return await handler(request.override(messages=edited_messages))
 
 
 __all__ = [

langchain/agents/middleware/file_search.py

@@ -21,7 +21,7 @@ from langchain.agents.middleware.types import AgentMiddleware
 
 
 def _expand_include_patterns(pattern: str) -> list[str] | None:
-    """Expand brace patterns like ``*.{py,pyi}`` into a list of globs."""
+    """Expand brace patterns like `*.{py,pyi}` into a list of globs."""
     if "}" in pattern and "{" not in pattern:
         return None
 
@@ -88,6 +88,7 @@ class FilesystemFileSearchMiddleware(AgentMiddleware):
     """Provides Glob and Grep search over filesystem files.
 
     This middleware adds two tools that search through local filesystem:
+
     - Glob: Fast file pattern matching by file path
     - Grep: Fast content search using ripgrep or Python fallback
 
@@ -100,7 +101,7 @@ class FilesystemFileSearchMiddleware(AgentMiddleware):
 
         agent = create_agent(
             model=model,
-            tools=[],
+            tools=[],  # Add tools as needed
             middleware=[
                 FilesystemFileSearchMiddleware(root_path="/workspace"),
             ],
@@ -119,9 +120,10 @@ class FilesystemFileSearchMiddleware(AgentMiddleware):
 
         Args:
             root_path: Root directory to search.
-            use_ripgrep: Whether to use ripgrep for search (default: True).
-                Falls back to Python if ripgrep unavailable.
-            max_file_size_mb: Maximum file size to search in MB (default: 10).
+            use_ripgrep: Whether to use `ripgrep` for search.
+
+                Falls back to Python if `ripgrep` unavailable.
+            max_file_size_mb: Maximum file size to search in MB.
         """
         self.root_path = Path(root_path).resolve()
         self.use_ripgrep = use_ripgrep
@@ -132,8 +134,10 @@ class FilesystemFileSearchMiddleware(AgentMiddleware):
         def glob_search(pattern: str, path: str = "/") -> str:
             """Fast file pattern matching tool that works with any codebase size.
 
-            Supports glob patterns like **/*.js or src/**/*.ts.
+            Supports glob patterns like `**/*.js` or `src/**/*.ts`.
+
             Returns matching file paths sorted by modification time.
+
             Use this tool when you need to find files by name patterns.
 
             Args:
@@ -142,7 +146,7 @@ class FilesystemFileSearchMiddleware(AgentMiddleware):
 
             Returns:
                 Newline-separated list of matching file paths, sorted by modification
-                time (most recently modified first). Returns "No files found" if no
+                time (most recently modified first). Returns `'No files found'` if no
                 matches.
             """
             try:
@@ -184,15 +188,16 @@ class FilesystemFileSearchMiddleware(AgentMiddleware):
             Args:
                 pattern: The regular expression pattern to search for in file contents.
                 path: The directory to search in. If not specified, searches from root.
-                include: File pattern to filter (e.g., "*.js", "*.{ts,tsx}").
+                include: File pattern to filter (e.g., `'*.js'`, `'*.{ts,tsx}'`).
                 output_mode: Output format:
-                    - "files_with_matches": Only file paths containing matches (default)
-                    - "content": Matching lines with file:line:content format
-                    - "count": Count of matches per file
+
+                    - `'files_with_matches'`: Only file paths containing matches
+                    - `'content'`: Matching lines with `file:line:content` format
+                    - `'count'`: Count of matches per file
 
             Returns:
-                Search results formatted according to output_mode. Returns "No matches
-                found" if no results.
+                Search results formatted according to `output_mode`.
+                    Returns `'No matches found'` if no results.
             """
             # Compile regex pattern (for validation)
             try:

langchain/agents/middleware/human_in_the_loop.py

@@ -7,17 +7,17 @@ from langgraph.runtime import Runtime
 from langgraph.types import interrupt
 from typing_extensions import NotRequired, TypedDict
 
-from langchain.agents.middleware.types import AgentMiddleware, AgentState
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, ContextT, StateT
 
 
 class Action(TypedDict):
     """Represents an action with a name and args."""
 
     name: str
-    """The type or name of action being requested (e.g., "add_numbers")."""
+    """The type or name of action being requested (e.g., `'add_numbers'`)."""
 
     args: dict[str, Any]
-    """Key-value pairs of args needed for the action (e.g., {"a": 1, "b": 2})."""
+    """Key-value pairs of args needed for the action (e.g., `{"a": 1, "b": 2}`)."""
 
 
 class ActionRequest(TypedDict):
@@ -27,7 +27,7 @@ class ActionRequest(TypedDict):
     """The name of the action being requested."""
 
     args: dict[str, Any]
-    """Key-value pairs of args needed for the action (e.g., {"a": 1, "b": 2})."""
+    """Key-value pairs of args needed for the action (e.g., `{"a": 1, "b": 2}`)."""
 
     description: NotRequired[str]
     """The description of the action to be reviewed."""
@@ -102,7 +102,7 @@ class HITLResponse(TypedDict):
 class _DescriptionFactory(Protocol):
     """Callable that generates a description for a tool call."""
 
-    def __call__(self, tool_call: ToolCall, state: AgentState, runtime: Runtime) -> str:
+    def __call__(self, tool_call: ToolCall, state: AgentState, runtime: Runtime[ContextT]) -> str:
         """Generate a description for a tool call."""
         ...
 
@@ -138,7 +138,7 @@ class InterruptOnConfig(TypedDict):
         def format_tool_description(
             tool_call: ToolCall,
             state: AgentState,
-            runtime: Runtime
+            runtime: Runtime[ContextT]
         ) -> str:
             import json
             return (
@@ -156,7 +156,7 @@ class InterruptOnConfig(TypedDict):
     """JSON schema for the args associated with the action, if edits are allowed."""
 
 
-class HumanInTheLoopMiddleware(AgentMiddleware):
+class HumanInTheLoopMiddleware(AgentMiddleware[StateT, ContextT]):
     """Human in the loop middleware."""
 
     def __init__(
@@ -169,18 +169,22 @@ class HumanInTheLoopMiddleware(AgentMiddleware):
 
         Args:
             interrupt_on: Mapping of tool name to allowed actions.
+
                 If a tool doesn't have an entry, it's auto-approved by default.
 
                 * `True` indicates all decisions are allowed: approve, edit, and reject.
                 * `False` indicates that the tool is auto-approved.
                 * `InterruptOnConfig` indicates the specific decisions allowed for this
                     tool.
-                    The InterruptOnConfig can include a `description` field (`str` or
+
+                    The `InterruptOnConfig` can include a `description` field (`str` or
                     `Callable`) for custom formatting of the interrupt description.
             description_prefix: The prefix to use when constructing action requests.
+
                 This is used to provide context about the tool call and the action being
-                requested. Not used if a tool has a `description` in its
-                `InterruptOnConfig`.
+                requested.
+
+                Not used if a tool has a `description` in its `InterruptOnConfig`.
         """
         super().__init__()
         resolved_configs: dict[str, InterruptOnConfig] = {}
@@ -200,7 +204,7 @@ class HumanInTheLoopMiddleware(AgentMiddleware):
         tool_call: ToolCall,
         config: InterruptOnConfig,
         state: AgentState,
-        runtime: Runtime,
+        runtime: Runtime[ContextT],
     ) -> tuple[ActionRequest, ReviewConfig]:
         """Create an ActionRequest and ReviewConfig for a tool call."""
         tool_name = tool_call["name"]
@@ -273,7 +277,7 @@ class HumanInTheLoopMiddleware(AgentMiddleware):
         )
         raise ValueError(msg)
 
-    def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
+    def after_model(self, state: AgentState, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
         """Trigger interrupt flows for relevant tool calls after an `AIMessage`."""
         messages = state["messages"]
         if not messages:
@@ -283,36 +287,23 @@ class HumanInTheLoopMiddleware(AgentMiddleware):
         if not last_ai_msg or not last_ai_msg.tool_calls:
             return None
 
-        # Separate tool calls that need interrupts from those that don't
-        interrupt_tool_calls: list[ToolCall] = []
-        auto_approved_tool_calls = []
-
-        for tool_call in last_ai_msg.tool_calls:
-            interrupt_tool_calls.append(tool_call) if tool_call[
-                "name"
-            ] in self.interrupt_on else auto_approved_tool_calls.append(tool_call)
-
-        # If no interrupts needed, return early
-        if not interrupt_tool_calls:
-            return None
-
-        # Process all tool calls that require interrupts
-        revised_tool_calls: list[ToolCall] = auto_approved_tool_calls.copy()
-        artificial_tool_messages: list[ToolMessage] = []
-
-        # Create action requests and review configs for all tools that need approval
+        # Create action requests and review configs for tools that need approval
         action_requests: list[ActionRequest] = []
         review_configs: list[ReviewConfig] = []
+        interrupt_indices: list[int] = []
 
-        for tool_call in interrupt_tool_calls:
-            config = self.interrupt_on[tool_call["name"]]
+        for idx, tool_call in enumerate(last_ai_msg.tool_calls):
+            if (config := self.interrupt_on.get(tool_call["name"])) is not None:
+                action_request, review_config = self._create_action_and_config(
+                    tool_call, config, state, runtime
+                )
+                action_requests.append(action_request)
+                review_configs.append(review_config)
+                interrupt_indices.append(idx)
 
-            # Create ActionRequest and ReviewConfig using helper method
-            action_request, review_config = self._create_action_and_config(
-                tool_call, config, state, runtime
-            )
-            action_requests.append(action_request)
-            review_configs.append(review_config)
+        # If no interrupts needed, return early
+        if not action_requests:
+            return None
 
         # Create single HITLRequest with all actions and configs
         hitl_request = HITLRequest(
@@ -321,31 +312,46 @@ class HumanInTheLoopMiddleware(AgentMiddleware):
         )
 
         # Send interrupt and get response
-        hitl_response: HITLResponse = interrupt(hitl_request)
-        decisions = hitl_response["decisions"]
+        decisions = interrupt(hitl_request)["decisions"]
 
         # Validate that the number of decisions matches the number of interrupt tool calls
-        if (decisions_len := len(decisions)) != (
-            interrupt_tool_calls_len := len(interrupt_tool_calls)
-        ):
+        if (decisions_len := len(decisions)) != (interrupt_count := len(interrupt_indices)):
             msg = (
                 f"Number of human decisions ({decisions_len}) does not match "
-                f"number of hanging tool calls ({interrupt_tool_calls_len})."
+                f"number of hanging tool calls ({interrupt_count})."
             )
             raise ValueError(msg)
 
-        # Process each decision using helper method
-        for i, decision in enumerate(decisions):
-            tool_call = interrupt_tool_calls[i]
-            config = self.interrupt_on[tool_call["name"]]
-
-            revised_tool_call, tool_message = self._process_decision(decision, tool_call, config)
-            if revised_tool_call:
-                revised_tool_calls.append(revised_tool_call)
-            if tool_message:
-                artificial_tool_messages.append(tool_message)
+        # Process decisions and rebuild tool calls in original order
+        revised_tool_calls: list[ToolCall] = []
+        artificial_tool_messages: list[ToolMessage] = []
+        decision_idx = 0
+
+        for idx, tool_call in enumerate(last_ai_msg.tool_calls):
+            if idx in interrupt_indices:
+                # This was an interrupt tool call - process the decision
+                config = self.interrupt_on[tool_call["name"]]
+                decision = decisions[decision_idx]
+                decision_idx += 1
+
+                revised_tool_call, tool_message = self._process_decision(
+                    decision, tool_call, config
+                )
+                if revised_tool_call is not None:
+                    revised_tool_calls.append(revised_tool_call)
+                if tool_message:
+                    artificial_tool_messages.append(tool_message)
+            else:
+                # This was auto-approved - keep original
+                revised_tool_calls.append(tool_call)
 
         # Update the AI message to only include approved tool calls
         last_ai_msg.tool_calls = revised_tool_calls
 
         return {"messages": [last_ai_msg, *artificial_tool_messages]}
+
+    async def aafter_model(
+        self, state: AgentState, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Async trigger interrupt flows for relevant tool calls after an `AIMessage`."""
+        return self.after_model(state, runtime)

langchain/agents/middleware/model_call_limit.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Literal
 
 from langchain_core.messages import AIMessage
 from langgraph.channels.untracked_value import UntrackedValue
-from typing_extensions import NotRequired
+from typing_extensions import NotRequired, override
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
@@ -20,9 +20,9 @@ if TYPE_CHECKING:
 
 
 class ModelCallLimitState(AgentState):
-    """State schema for ModelCallLimitMiddleware.
+    """State schema for `ModelCallLimitMiddleware`.
 
-    Extends AgentState with model call tracking fields.
+    Extends `AgentState` with model call tracking fields.
     """
 
     thread_model_call_count: NotRequired[Annotated[int, PrivateStateAttr]]
@@ -58,8 +58,8 @@ def _build_limit_exceeded_message(
 class ModelCallLimitExceededError(Exception):
     """Exception raised when model call limits are exceeded.
 
-    This exception is raised when the configured exit behavior is 'error'
-    and either the thread or run model call limit has been exceeded.
+    This exception is raised when the configured exit behavior is `'error'` and either
+    the thread or run model call limit has been exceeded.
     """
 
     def __init__(
@@ -127,13 +127,17 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
 
         Args:
             thread_limit: Maximum number of model calls allowed per thread.
-                None means no limit.
+
+                `None` means no limit.
             run_limit: Maximum number of model calls allowed per run.
-                None means no limit.
+
+                `None` means no limit.
             exit_behavior: What to do when limits are exceeded.
-                - "end": Jump to the end of the agent execution and
-                    inject an artificial AI message indicating that the limit was exceeded.
-                - "error": Raise a `ModelCallLimitExceededError`
+
+                - `'end'`: Jump to the end of the agent execution and
+                    inject an artificial AI message indicating that the limit was
+                    exceeded.
+                - `'error'`: Raise a `ModelCallLimitExceededError`
 
         Raises:
             ValueError: If both limits are `None` or if `exit_behavior` is invalid.
@@ -144,7 +148,7 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
             msg = "At least one limit must be specified (thread_limit or run_limit)"
             raise ValueError(msg)
 
-        if exit_behavior not in ("end", "error"):
+        if exit_behavior not in {"end", "error"}:
             msg = f"Invalid exit_behavior: {exit_behavior}. Must be 'end' or 'error'"
             raise ValueError(msg)
 
@@ -153,7 +157,8 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
         self.exit_behavior = exit_behavior
 
     @hook_config(can_jump_to=["end"])
-    def before_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+    @override
+    def before_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:
         """Check model call limits before making a model call.
 
         Args:
@@ -161,12 +166,13 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
             runtime: The langgraph runtime.
 
         Returns:
-            If limits are exceeded and exit_behavior is "end", returns
-            a Command to jump to the end with a limit exceeded message. Otherwise returns None.
+            If limits are exceeded and exit_behavior is `'end'`, returns
+                a `Command` to jump to the end with a limit exceeded message. Otherwise
+                returns `None`.
 
         Raises:
-            ModelCallLimitExceededError: If limits are exceeded and exit_behavior
-                is "error".
+            ModelCallLimitExceededError: If limits are exceeded and `exit_behavior`
+                is `'error'`.
         """
         thread_count = state.get("thread_model_call_count", 0)
         run_count = state.get("run_model_call_count", 0)
@@ -194,7 +200,31 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
 
         return None
 
-    def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+    @hook_config(can_jump_to=["end"])
+    async def abefore_model(
+        self,
+        state: ModelCallLimitState,
+        runtime: Runtime,
+    ) -> dict[str, Any] | None:
+        """Async check model call limits before making a model call.
+
+        Args:
+            state: The current agent state containing call counts.
+            runtime: The langgraph runtime.
+
+        Returns:
+            If limits are exceeded and exit_behavior is `'end'`, returns
+                a `Command` to jump to the end with a limit exceeded message. Otherwise
+                returns `None`.
+
+        Raises:
+            ModelCallLimitExceededError: If limits are exceeded and `exit_behavior`
+                is `'error'`.
+        """
+        return self.before_model(state, runtime)
+
+    @override
+    def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:
         """Increment model call counts after a model call.
 
         Args:
@@ -208,3 +238,19 @@ class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
             "thread_model_call_count": state.get("thread_model_call_count", 0) + 1,
             "run_model_call_count": state.get("run_model_call_count", 0) + 1,
         }
+
+    async def aafter_model(
+        self,
+        state: ModelCallLimitState,
+        runtime: Runtime,
+    ) -> dict[str, Any] | None:
+        """Async increment model call counts after a model call.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            State updates with incremented call counts.
+        """
+        return self.after_model(state, runtime)