AI & Agent Dev Bug Sandbox Back to Radar

@Wrap_tool_call Fails To Preserve Custom Context Type And Is Incompatible With Create_agent(Context_schema=...)

The @wrap_tool_call decorator does not accept or propagate a context_schema parameter, so the returned AgentMiddleware always types ContextT as None. When used with create_agent(context_schema=AgentContext), the invariant ContextT generic causes a type error, even though the runtime behavior may be correct.

mediumConfidence 87%Langchain Affected V1.4.3

Origin Analysis

The wrap_tool_call decorator lacks a context_schema parameter and its generic ContextT defaults to None. Because AgentMiddleware's ContextT is invariant, AgentMiddleware[State, None, Any] is not assignable to AgentMiddleware[State, AgentContext, Any], leading Pylance/mypy to reject the middleware list in create_agent.

1. Define custom State extending AgentState and custom AgentContext TypedDict. 2. Decorate a middleware function with @wrap_tool_call(state_schema=State). 3. Pass the decorated middleware to create_agent along with context_schema=AgentContext. 4. Observe type checker error: 'Argument of type "list[AgentMiddleware[AgentState[Any], None, Any]]" cannot be assigned to parameter "middleware"... ContextT invariant...'

Fixing Code Block

from collections.abc import Awaitable, Callable
from typing import Any, TypeVar, overload
from langchain.agents.middleware.types import (
    AgentMiddleware,
    ToolCallRequest,
    ToolCallHandler,
    StateT,
    ContextT,
)

@overload
def wrap_tool_call(
    func: Callable[[ToolCallRequest[StateT, ContextT], ToolCallHandler], Awaitable[Any]],
) -> AgentMiddleware[StateT, ContextT, Any]: ...

@overload
def wrap_tool_call(
    *,
    state_schema: type[StateT] | None = None,
    context_schema: type[ContextT] | None = None,
) -> Callable[
    [Callable[[ToolCallRequest[StateT, ContextT], ToolCallHandler], Awaitable[Any]]],
    AgentMiddleware[StateT, ContextT, Any],
]: ...

def wrap_tool_call(
    func: Callable[[ToolCallRequest[StateT, ContextT], ToolCallHandler], Awaitable[Any]] | None = None,
    *,
    state_schema: type[StateT] | None = None,
    context_schema: type[ContextT] | None = None,
) -> Any:
    def decorator(
        f: Callable[[ToolCallRequest[StateT, ContextT], ToolCallHandler], Awaitable[Any]]
    ) -> AgentMiddleware[StateT, ContextT, Any]:
        return AgentMiddleware(
            f,
            state_schema=state_schema,
            context_schema=context_schema,
        )
    if func is not None:
        return decorator(func)
    return decorator

Add a context_schema keyword-only parameter to wrap_tool_call and pass it to the AgentMiddleware constructor. This allows the decorator to bind ContextT to the user-provided context schema, making the returned AgentMiddleware[State, AgentContext, Any] assignable to create_agent's middleware parameter. The overloads preserve the existing call patterns while adding the new parameter.

Edge Case Audit

This is a typing-only change; runtime behavior is unchanged. Users must explicitly pass context_schema=AgentContext to the decorator to achieve type compatibility. If the parameter is omitted, ContextT remains None and the same type error persists. The added parameter is optional and backward compatible, but code relying on the previous signatures may need to update overloads if they use static type checking. Rollback suggestion: revert to class-based middleware (AgentMiddleware[State, AgentContext]) or use a cast until the patch is applied.