langchain - 💡(How to fix) Fix create_agent discards state_schema subclass in inferred return type

Submission checklist

• This is a bug, not a usage question.
• I added a clear and descriptive title that summarizes this issue.
• I used the GitHub search to find a similar question and didn't find it.
• I am sure that this is a bug in LangChain rather than my code.
• The bug is not resolved by updating to the latest stable version of LangChain (or the specific integration package).
• This is not related to the langchain-community package.
• I posted a self-contained, minimal, reproducible example. A maintainer can copy it and run it AS IS.

Package (Required)

• langchain
• langchain-openai
• langchain-anthropic
• langchain-classic
• langchain-core
• langchain-model-profiles
• langchain-tests
• langchain-text-splitters
• langchain-chroma
• langchain-deepseek
• langchain-exa
• langchain-fireworks
• langchain-groq
• langchain-huggingface
• langchain-mistralai
• langchain-nomic
• langchain-ollama
• langchain-openrouter
• langchain-perplexity
• langchain-qdrant
• langchain-xai
• Other / not sure / general

Related Issues / PRs

• #33732 (closed Nov 2025 referencing #33817 as the fix). #33817 partially fixed the original request — it introduced StateT_co and rewired the middleware parameter from AgentMiddleware[AgentState[ResponseT], ContextT] to AgentMiddleware[StateT_co, ContextT], but it did not update state_schema or the return type. The original ask in #33732 was for state_schema and the return type to use StateT as well; that half was not landed.
• #34156 (related: custom state schema doesn't propagate to middleware)
• langchain-ai/langgraph#6175 (related upstream discussion)

Reproduction Steps / Example Code (Python)

from typing import Annotated, NotRequired

from langchain.agents import AgentState, create_agent
from langchain_core.language_models.fake_chat_models import FakeListChatModel


class MyState(AgentState):
    counter: NotRequired[int]


graph = create_agent(
    model=FakeListChatModel(responses=["ok"]),
    tools=[],
    state_schema=MyState,
)

# Static type checkers (mypy / pyrefly / pyright) infer `graph` as
#   CompiledStateGraph[AgentState[ResponseT], ContextT, _InputAgentState, _OutputAgentState[ResponseT]]
# The `MyState` subclass passed via `state_schema` is erased from the return type,
# even though it's the schema the graph actually runs with.
reveal_type(graph)  # noqa: F821

Error Message and Stack Trace (if applicable)

Not a runtime error — a type-inference issue. `reveal_type` (or hovering in an IDE) reports
`CompiledStateGraph[AgentState[Any], …]` instead of `CompiledStateGraph[MyState, …]`.

Downstream, anything declared as `CompiledStateGraph[MyState, …]` (e.g. a helper that wants
to receive the compiled graph with full state precision) fails to type-check against
`create_agent`'s inferred return, forcing a `cast` at every call site.

Description

langchain.agents.create_agent accepts state_schema: type[AgentState[ResponseT]] | None and returns CompiledStateGraph[AgentState[ResponseT], ContextT, _InputAgentState, _OutputAgentState[ResponseT]] (see libs/langchain_v1/langchain/agents/factory.py lines 684–702 on the installed release, lines 697–716 on master).

The state-schema subclass is therefore erased from the return type — create_agent(state_schema=MyState, …) is inferred as CompiledStateGraph[AgentState[…], …], not CompiledStateGraph[MyState, …].

This is inconsistent with the rest of the signature:

• The middleware parameter is already typed with a StateT_co TypeVar bound to AgentState[Any] (middleware/types.py:372), and that TypeVar is used on middleware (line 689) but not on state_schema or in the return type. The TypeVar exists; it's just not wired through.

Why this is a bug (not just a polish nit):

• CompiledStateGraph is invariant in its first generic, so CompiledStateGraph[MySubclass, …] is not assignable to CompiledStateGraph[AgentState, …] and vice versa. Subtyping does not rescue downstream call sites.
• Anyone holding a typed reference to the compiled graph in their own code (e.g. as a service field) is forced to either: (a) cast at every boundary, (b) downgrade their alias to CompiledStateGraph[AgentState[Any], …] and lose state precision, or (c) patch langchain locally.
• Issue #33732 raised this in Nov 2025 with a working fork (LouisAuneau/langchain@6a33ada) and was closed referencing #33817. #33817 did land part of the original fix — StateT_co was added and the middleware parameter now uses it — but the corresponding changes to state_schema and the return type (the parts that actually affect create_agent's inferred return type) were not included. As of langchain==1.3.4 (latest on PyPI) and current master, lines 704 and 714–716 still hardcode state_schema: type[AgentState[ResponseT]] and CompiledStateGraph[AgentState[ResponseT], …].

Expected behaviour

create_agent(state_schema=MyState, …) should be inferred as CompiledStateGraph[MyState, ContextT, _InputAgentState, _OutputAgentState[ResponseT]].

Proposed fix (matches the reference fork, ~30 lines, strictly backward-compatible since StateT is bound to AgentState):

# factory.py
from langchain.agents.middleware.types import StateT  # already defined there

def create_agent(
    model: str | BaseChatModel,
    tools: Sequence[BaseTool | Callable[..., Any] | dict[str, Any]] | None = None,
    *,
    system_prompt: str | SystemMessage | None = None,
    middleware: Sequence[AgentMiddleware[StateT, ContextT]] = (),
    response_format: ResponseFormat[ResponseT] | type[ResponseT] | dict[str, Any] | None = None,
    state_schema: type[StateT] | None = None,           # was: type[AgentState[ResponseT]]
    context_schema: type[ContextT] | None = None,
    ...,
) -> CompiledStateGraph[
    StateT, ContextT, _InputAgentState, _OutputAgentState[ResponseT]   # was: AgentState[ResponseT]
]:
    ...

Reference implementation: LouisAuneau/langchain@6a33ada.

Happy to open a PR with this change if there's interest.

System Info

System Information
------------------
> OS:  Darwin
> OS Version:  Darwin Kernel Version 25.1.0: Mon Oct 20 19:33:00 PDT 2025; root:xnu-12377.41.6~2/RELEASE_ARM64_T6020
> Python Version:  3.12.9 (main, Mar 17 2025, 21:36:21) [Clang 20.1.0 ]

Package Information
-------------------
> langchain_core: 1.3.2
> langchain: 1.2.15
> langsmith: 0.4.29
> langchain_cerebras: 0.8.2
> langchain_groq: 1.1.0
> langchain_mcp_adapters: 0.2.2
> langchain_openai: 1.2.1
> langchain_protocol: 0.0.13
> langchain_tavily: 0.2.13
> langgraph_sdk: 0.3.13

(Also reproduced against langchain==1.3.4, latest on PyPI — same signature.)