How to Make SummarizationMiddleware Context-Aware When Using Custom Context Schema?

OSS Product Help LangGraph
1

Hi guys, so I’m working with create_agent from langchain.agents and I’m trying to make use of both a custom context_schema (to pass runtime data like page_context) and SummarizationMiddleware.

Setup


from langchain.agents import AgentState, create_agent
from langchain.agents.middleware import AgentMiddleware, SummarizationMiddleware

class Context(BaseModel):
    page_context: Optional[Dict[str, Any]] = None

class ToolMonitoringMiddleware(AgentMiddleware[AgentState, Context]):
    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
    ) -> ToolMessage | Command:
        print(f"Executing tool: {request.tool_call['name']}")
        print(f"Arguments: {request.tool_call['args']}")
        print(f"Runtime context: {request.runtime.context}")
        try:
            result = await handler(request)
            print(f"{request.tool_call['name']} tool completed successfully")
            return result
        except Exception as e:
            print(f"Tool failed: {e}")
            raise

# ...

agent_executor = create_agent(
    model=qwen3_model.chatModel,
    tools=tools,
    system_prompt=system_message,
    checkpointer=checkpointer,
    context_schema=Context,
    middleware=[
        ToolMonitoringMiddleware(),  # This one works fine with context
        SummarizationMiddleware(
            model=qwen3_model.chat_model,
            max_tokens_before_summary=40000,
            messages_to_keep=15,
        ),
    ],
)

Issue

However, when I specify a context_schema=Context, Pyright gives me this error:

"SummarizationMiddleware" is not assignable to "AgentMiddleware[AgentState[Unknown], Context]"
  Type parameter "ContextT@AgentMiddleware" is invariant, but "None" is not the same as "Context" Pyright (reportArgumentType)

Discussion

Is there a recommended pattern for making built-in middlewares like SummarizationMiddleware compatible with custom context schemas?

Alternatively, should SummarizationMiddleware be updated to support being parameterized with AgentMiddleware[StateT, ContextT]?

Thanks in advance!

2

Hi @ncbernar

interesting issue.

I think it is because class SummarizationMiddleware(AgentMiddleware) does not explicit pass Context type and its default is None → ContextT = TypeVar("ContextT", bound=StateLike | None, default=None). So effectively the summarization middleware becames SummarizationMiddleware(AgentMiddleware[AgentState, None])

When context_schema=Context, Pyright instantiates ContextT as Context. And Context is not assignable to None.

How to work around that?

  1. Treat it as context‑agnostic and silence the type checker locally.
from typing import cast
from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware, SummarizationMiddleware

middleware: list[AgentMiddleware[AgentState, Context]] = [
    ToolMonitoringMiddleware(),
    cast(AgentMiddleware[AgentState, Context], SummarizationMiddleware(
        model=qwen3_model.chat_model,
        max_tokens_before_summary=40000,
        messages_to_keep=15,
    )),
]

or by attaching a focused # type: ignore[arg-type] comment on the SummarizationMiddleware(...) line. This doesn’t change runtime behavior; it just tells the type checker “I know this middleware ignores context.”

  1. Create a thin, typed wrapper that delegates to SummarizationMiddleware.

If you prefer avoiding casts, you can wrap the built‑in middleware in your own AgentMiddleware[AgentState, Context] that simply forwards the hooks:

from typing import Any, cast
from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware, SummarizationMiddleware
from langgraph.runtime import Runtime

class ContextAwareSummarizationMiddleware(AgentMiddleware[AgentState, Context]):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__()
        self._inner = SummarizationMiddleware(*args, **kwargs)

    def before_model(
        self,
        state: AgentState,
        runtime: Runtime[Context],
    ) -> dict[str, Any] | None:
        # Delegate to the built-in middleware (it ignores runtime.context anyway)
        return self._inner.before_model(state, cast(Runtime, runtime))

    async def abefore_model(
        self,
        state: AgentState,
        runtime: Runtime[Context],
    ) -> dict[str, Any] | None:
        return await self._inner.abefore_model(state, cast(Runtime, runtime))

Then use ContextAwareSummarizationMiddleware instead of SummarizationMiddleware in your middleware=[...] list. This keeps the type signature aligned with context_schema=Context while preserving the existing summarization behavior.

  1. If you actually need the context inside summarization, extend the wrapper.
class ContextAwareSummarizationMiddleware(AgentMiddleware[AgentState, Context]):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__()
        self._inner = SummarizationMiddleware(*args, **kwargs)

    async def abefore_model(
        self,
        state: AgentState,
        runtime: Runtime[Context],
    ) -> dict[str, Any] | None:
        # You can inspect runtime.context here
        page_ctx = runtime.context.page_context if runtime.context else None

        # Optionally tweak state before delegating
        updates = await self._inner.abefore_model(state, cast(Runtime, runtime))
        # Potentially merge updates with something derived from page_ctx
        return updates
3

And fix(langchain): allow builtin middleware with custom context_schema by pawel-twardziak · Pull Request #34232 · langchain-ai/langchain · GitHub @sydney-runkle :slight_smile: