fix: adding missing async hooks (#33957)

* filling in missing async gaps
* using recommended tool runtime injection instead of injected state
  * updating tests to use helper function as well

Author: sydney-runkle

libs/langchain_v1/langchain/agents/middleware/human_in_the_loop.py

@@ -353,3 +353,7 @@ def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | N
         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) -> dict[str, Any] | None:
+        """Async trigger interrupt flows for relevant tool calls after an `AIMessage`."""
+        return self.after_model(state, runtime)

libs/langchain_v1/langchain/agents/middleware/model_call_limit.py

@@ -198,6 +198,29 @@ def before_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str
 
         return None
 
+    @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)
+
     def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
         """Increment model call counts after a model call.
 
@@ -212,3 +235,19 @@ def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str,
             "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)

libs/langchain_v1/langchain/agents/middleware/pii.py

@@ -252,6 +252,27 @@ def before_model(
 
         return None
 
+    @hook_config(can_jump_to=["end"])
+    async def abefore_model(
+        self,
+        state: AgentState,
+        runtime: Runtime,
+    ) -> dict[str, Any] | None:
+        """Async check user messages and tool results for PII before model invocation.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            Updated state with PII handled according to strategy, or `None` if no PII
+                detected.
+
+        Raises:
+            PIIDetectionError: If PII is detected and strategy is `'block'`.
+        """
+        return self.before_model(state, runtime)
+
     def after_model(
         self,
         state: AgentState,
@@ -311,6 +332,26 @@ def after_model(
 
         return {"messages": new_messages}
 
+    async def aafter_model(
+        self,
+        state: AgentState,
+        runtime: Runtime,
+    ) -> dict[str, Any] | None:
+        """Async check AI messages for PII after model invocation.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            Updated state with PII handled according to strategy, or None if no PII
+                detected.
+
+        Raises:
+            PIIDetectionError: If PII is detected and strategy is `'block'`.
+        """
+        return self.after_model(state, runtime)
+
 
 __all__ = [
     "PIIDetectionError",

libs/langchain_v1/langchain/agents/middleware/shell_tool.py

@@ -482,7 +482,7 @@ def before_agent(self, state: ShellToolState, runtime: Runtime) -> dict[str, Any
         return {"shell_session_resources": resources}
 
     async def abefore_agent(self, state: ShellToolState, runtime: Runtime) -> dict[str, Any] | None:
-        """Async counterpart to `before_agent`."""
+        """Async start the shell session and run startup commands."""
         return self.before_agent(state, runtime)
 
     def after_agent(self, state: ShellToolState, runtime: Runtime) -> None:  # noqa: ARG002
@@ -494,7 +494,7 @@ def after_agent(self, state: ShellToolState, runtime: Runtime) -> None:  # noqa:
             resources._finalizer()
 
     async def aafter_agent(self, state: ShellToolState, runtime: Runtime) -> None:
-        """Async counterpart to `after_agent`."""
+        """Async run shutdown commands and release resources when an agent completes."""
         return self.after_agent(state, runtime)
 
     def _ensure_resources(self, state: ShellToolState) -> _SessionResources:
@@ -689,7 +689,8 @@ async def awrap_tool_call(
         request: ToolCallRequest,
         handler: typing.Callable[[ToolCallRequest], typing.Awaitable[ToolMessage | Command]],
     ) -> ToolMessage | Command:
-        """Async interception mirroring the synchronous tool handler."""
+        """Async intercept local shell tool calls and execute them via the managed session."""
+        # The sync version already handles all the work, no need for async-specific logic
         if isinstance(request.tool, _PersistentShellTool):
             resources = self._ensure_resources(request.state)
             return self._run_shell_tool(

libs/langchain_v1/langchain/agents/middleware/tool_call_limit.py

@@ -451,3 +451,28 @@ def after_model(
             "run_tool_call_count": run_counts,
             "messages": artificial_messages,
         }
+
+    @hook_config(can_jump_to=["end"])
+    async def aafter_model(
+        self,
+        state: ToolCallLimitState[ResponseT],
+        runtime: Runtime[ContextT],
+    ) -> dict[str, Any] | None:
+        """Async increment tool call counts after a model call and check limits.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            State updates with incremented tool call counts. If limits are exceeded
+                and exit_behavior is `'end'`, also includes a jump to end with a
+                `ToolMessage` and AI message for the single exceeded tool call.
+
+        Raises:
+            ToolCallLimitExceededError: If limits are exceeded and `exit_behavior`
+                is `'error'`.
+            NotImplementedError: If limits are exceeded, `exit_behavior` is `'end'`,
+                and there are multiple tool calls.
+        """
+        return self.after_model(state, runtime)

libs/partners/anthropic/langchain_anthropic/middleware/bash.py

@@ -40,6 +40,18 @@ def wrap_model_call(
             request = request.override(tools=tools)
         return handler(request)
 
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelResponse:
+        """Async: ensure the Claude bash descriptor is available to the model."""
+        tools = request.tools
+        if all(tool is not _CLAUDE_BASH_DESCRIPTOR for tool in tools):
+            tools = [*tools, _CLAUDE_BASH_DESCRIPTOR]
+            request = request.override(tools=tools)
+        return await handler(request)
+
     def wrap_tool_call(
         self,
         request: ToolCallRequest,