fix(langchain): allow builtin middleware with custom context_schema

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

@@ -25,9 +25,11 @@
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
+    ContextT,
     ModelCallResult,
     ModelRequest,
     ModelResponse,
+    StateT,
 )
 
 DEFAULT_TOOL_PLACEHOLDER = "[cleared]"
@@ -182,7 +184,7 @@ def _build_cleared_tool_input_message(
         )
 
 
-class ContextEditingMiddleware(AgentMiddleware):
+class ContextEditingMiddleware(AgentMiddleware[StateT, ContextT]):
     """Automatically prune tool results to manage context size.
 
     The middleware applies a sequence of edits when the total input token count exceeds

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

@@ -17,7 +17,7 @@
 
 from langchain_core.tools import tool
 
-from langchain.agents.middleware.types import AgentMiddleware
+from langchain.agents.middleware.types import AgentMiddleware, ContextT, StateT
 
 
 def _expand_include_patterns(pattern: str) -> list[str] | None:
@@ -84,7 +84,7 @@ def _match_include_pattern(basename: str, pattern: str) -> bool:
     return any(fnmatch.fnmatch(basename, candidate) for candidate in expanded)
 
 
-class FilesystemFileSearchMiddleware(AgentMiddleware):
+class FilesystemFileSearchMiddleware(AgentMiddleware[StateT, ContextT]):
     """Provides Glob and Grep search over filesystem files.
 
     This middleware adds two tools that search through local filesystem:

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

@@ -11,6 +11,7 @@
 from langchain.agents.middleware.types import (
     AgentMiddleware,
     AgentState,
+    ContextT,
     PrivateStateAttr,
     hook_config,
 )
@@ -86,7 +87,7 @@ def __init__(
         super().__init__(msg)
 
 
-class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
+class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, ContextT]):
     """Tracks model call counts and enforces limits.
 
     This middleware monitors the number of model calls made during agent execution
@@ -158,7 +159,9 @@ def __init__(
 
     @hook_config(can_jump_to=["end"])
     @override
-    def before_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:
+    def before_model(
+        self, state: ModelCallLimitState, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
         """Check model call limits before making a model call.
 
         Args:
@@ -204,7 +207,7 @@ def before_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str
     async def abefore_model(
         self,
         state: ModelCallLimitState,
-        runtime: Runtime,
+        runtime: Runtime[ContextT],
     ) -> dict[str, Any] | None:
         """Async check model call limits before making a model call.
 
@@ -224,7 +227,9 @@ async def abefore_model(
         return self.before_model(state, runtime)
 
     @override
-    def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:
+    def after_model(
+        self, state: ModelCallLimitState, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
         """Increment model call counts after a model call.
 
         Args:
@@ -242,7 +247,7 @@ def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str,
     async def aafter_model(
         self,
         state: ModelCallLimitState,
-        runtime: Runtime,
+        runtime: Runtime[ContextT],
     ) -> dict[str, Any] | None:
         """Async increment model call counts after a model call.
 

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

@@ -6,9 +6,11 @@
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
+    ContextT,
     ModelCallResult,
     ModelRequest,
     ModelResponse,
+    StateT,
 )
 from langchain.chat_models import init_chat_model
 
@@ -18,7 +20,7 @@
     from langchain_core.language_models.chat_models import BaseChatModel
 
 
-class ModelFallbackMiddleware(AgentMiddleware):
+class ModelFallbackMiddleware(AgentMiddleware[StateT, ContextT]):
     """Automatic fallback to alternative models on errors.
 
     Retries failed model calls with alternative models in sequence until

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

@@ -15,15 +15,15 @@
     should_retry_exception,
     validate_retry_params,
 )
-from langchain.agents.middleware.types import AgentMiddleware, ModelResponse
+from langchain.agents.middleware.types import AgentMiddleware, ContextT, ModelResponse, StateT
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
 
     from langchain.agents.middleware.types import ModelRequest
 
 
-class ModelRetryMiddleware(AgentMiddleware):
+class ModelRetryMiddleware(AgentMiddleware[StateT, ContextT]):
     """Middleware that automatically retries failed model calls with configurable backoff.
 
     Supports retrying on specific exceptions and exponential backoff.

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

@@ -19,15 +19,21 @@
     detect_mac_address,
     detect_url,
 )
-from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ContextT,
+    StateT,
+    hook_config,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
     from langgraph.runtime import Runtime
 
 
-class PIIMiddleware(AgentMiddleware):
+class PIIMiddleware(AgentMiddleware[StateT, ContextT]):
     """Detect and handle Personally Identifiable Information (PII) in conversations.
 
     This middleware detects common PII types and applies configurable strategies
@@ -165,7 +171,7 @@ def _process_content(self, content: str) -> tuple[str, list[PIIMatch]]:
     def before_model(
         self,
         state: AgentState,
-        runtime: Runtime,
+        runtime: Runtime[ContextT],
     ) -> dict[str, Any] | None:
         """Check user messages and tool results for PII before model invocation.
 
@@ -260,7 +266,7 @@ def before_model(
     async def abefore_model(
         self,
         state: AgentState,
-        runtime: Runtime,
+        runtime: Runtime[ContextT],
     ) -> dict[str, Any] | None:
         """Async check user messages and tool results for PII before model invocation.
 
@@ -281,7 +287,7 @@ async def abefore_model(
     def after_model(
         self,
         state: AgentState,
-        runtime: Runtime,
+        runtime: Runtime[ContextT],
     ) -> dict[str, Any] | None:
         """Check AI messages for PII after model invocation.
 
@@ -340,7 +346,7 @@ def after_model(
     async def aafter_model(
         self,
         state: AgentState,
-        runtime: Runtime,
+        runtime: Runtime[ContextT],
     ) -> dict[str, Any] | None:
         """Async check AI messages for PII after model invocation.
 

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

@@ -20,7 +20,7 @@
 from langgraph.runtime import Runtime
 from typing_extensions import override
 
-from langchain.agents.middleware.types import AgentMiddleware, AgentState
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, ContextT, StateT
 from langchain.chat_models import BaseChatModel, init_chat_model
 
 TokenCounter = Callable[[Iterable[MessageLikeRepresentation]], int]
@@ -128,7 +128,7 @@ def _get_approximate_token_counter(model: BaseChatModel) -> TokenCounter:
     return count_tokens_approximately
 
 
-class SummarizationMiddleware(AgentMiddleware):
+class SummarizationMiddleware(AgentMiddleware[StateT, ContextT]):
     """Summarizes conversation history when token limits are approached.
 
     This middleware monitors message token counts and automatically summarizes older
@@ -264,7 +264,7 @@ def __init__(
             raise ValueError(msg)
 
     @override
-    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
+    def before_model(self, state: AgentState, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
         """Process messages before model invocation, potentially triggering summarization."""
         messages = state["messages"]
         self._ensure_message_ids(messages)
@@ -292,7 +292,9 @@ def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] |
         }
 
     @override
-    async def abefore_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
+    async def abefore_model(
+        self, state: AgentState, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
         """Process messages before model invocation, potentially triggering summarization."""
         messages = state["messages"]
         self._ensure_message_ids(messages)

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

@@ -15,6 +15,7 @@
 from langchain.agents.middleware.types import (
     AgentMiddleware,
     AgentState,
+    ContextT,
     ModelCallResult,
     ModelRequest,
     ModelResponse,
@@ -126,7 +127,7 @@ def write_todos(todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCall
     )
 
 
-class TodoListMiddleware(AgentMiddleware):
+class TodoListMiddleware(AgentMiddleware[PlanningState, ContextT]):
     """Middleware that provides todo list management capabilities to agents.
 
     This middleware adds a `write_todos` tool that allows agents to create and manage

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

@@ -6,12 +6,12 @@
 
 from langchain_core.messages import AIMessage, ToolCall, ToolMessage
 from langgraph.channels.untracked_value import UntrackedValue
-from langgraph.typing import ContextT
 from typing_extensions import NotRequired, override
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
     AgentState,
+    ContextT,
     PrivateStateAttr,
     ResponseT,
     hook_config,

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

@@ -7,7 +7,7 @@
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage, ToolMessage
 
-from langchain.agents.middleware.types import AgentMiddleware
+from langchain.agents.middleware.types import AgentMiddleware, ContextT, StateT
 from langchain.chat_models.base import init_chat_model
 
 if TYPE_CHECKING:
@@ -19,7 +19,7 @@
     from langchain.tools import BaseTool
 
 
-class LLMToolEmulator(AgentMiddleware):
+class LLMToolEmulator(AgentMiddleware[StateT, ContextT]):
     """Emulates specified tools using an LLM instead of executing them.
 
     This middleware allows selective emulation of tools for testing purposes.

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

@@ -16,7 +16,7 @@
     should_retry_exception,
     validate_retry_params,
 )
-from langchain.agents.middleware.types import AgentMiddleware
+from langchain.agents.middleware.types import AgentMiddleware, ContextT, StateT
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
@@ -27,7 +27,7 @@
     from langchain.tools import BaseTool
 
 
-class ToolRetryMiddleware(AgentMiddleware):
+class ToolRetryMiddleware(AgentMiddleware[StateT, ContextT]):
     """Middleware that automatically retries failed tool calls with configurable backoff.
 
     Supports retrying on specific exceptions and exponential backoff.

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

@@ -18,9 +18,11 @@
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
+    ContextT,
     ModelCallResult,
     ModelRequest,
     ModelResponse,
+    StateT,
 )
 from langchain.chat_models.base import init_chat_model
 
@@ -85,7 +87,7 @@ def _render_tool_list(tools: list[BaseTool]) -> str:
     return "\n".join(f"- {tool.name}: {tool.description}" for tool in tools)
 
 
-class LLMToolSelectorMiddleware(AgentMiddleware):
+class LLMToolSelectorMiddleware(AgentMiddleware[StateT, ContextT]):
     """Uses an LLM to select relevant tools before calling the main model.
 
     When an agent has many tools available, this middleware filters them down