feat(sdk): allow passing state_schema in create_deep_agent (#3642)

Fixes #3249

Allow callers to pass a custom `state_schema` to `create_deep_agent`,
enabling them to extend the graph state with extra fields (e.g. page
metadata, file URLs) while preserving the `DeltaChannel` message reducer
that keeps checkpoint size O(N). Defaults to `_DeepAgentState` so
existing behavior is unchanged.

## How verified

- Added two unit tests: one verifying the default `_DeepAgentState` is
used when no schema is passed, another verifying a custom subclass is
forwarded through to `create_agent`.
- All 85 existing tests in `test_graph.py` pass.
- Ruff lint checks pass with no violations.

---

**AI Disclosure:** This PR was authored with Claude Opus 4.7 assistance
for code generation and testing.

---------

Co-authored-by: Nishitha Madhu <nishitha.madhu@example.com>
Co-authored-by: Mason Daugherty <github@mdrxy.com>

libs/deepagents/deepagents/graph.py

@@ -60,7 +60,7 @@ from deepagents.profiles.harness.harness_profiles import (
 logger = logging.getLogger(__name__)
 
 
-class _DeepAgentState(AgentState):
+class DeepAgentState(AgentState):
     """AgentState with DeltaChannel on messages to reduce checkpoint growth from O(N²) to O(N)."""
 
     messages: Required[Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer, snapshot_frequency=50)]]  # ty: ignore[invalid-argument-type]
@@ -227,6 +227,7 @@ def create_deep_agent(  # noqa: C901, PLR0912, PLR0915  # Complex graph assembly
     backend: BackendProtocol | BackendFactory | None = None,
     interrupt_on: dict[str, bool | InterruptOnConfig] | None = None,
     response_format: ResponseFormat[ResponseT] | type[ResponseT] | dict[str, Any] | None = None,
+    state_schema: type[DeepAgentState] | None = None,
     context_schema: type[ContextT] | None = None,
     checkpointer: Checkpointer | None = None,
     store: BaseStore | None = None,
@@ -437,6 +438,41 @@ def create_deep_agent(  # noqa: C901, PLR0912, PLR0915  # Complex graph assembly
             For example, `interrupt_on={"edit_file": True}` pauses before
             every edit.
         response_format: A structured output response format to use for the agent.
+        state_schema: Custom state schema for the agent graph. Must be a
+            `TypedDict` subclass of
+            [`DeepAgentState`][deepagents.graph.DeepAgentState] so the
+            built-in `DeltaChannel` reducer on `messages` is preserved.
+
+            Generally, prefer defining state extensions with middleware so
+            the extra fields stay scoped to the hooks and tools that use
+            them.
+
+            When provided, this schema is used as the base graph schema and
+            is merged with state schemas contributed by middleware. It is
+            also forwarded when compiling declarative
+            [`SubAgent`][deepagents.middleware.subagents.SubAgent] specs for
+            the `task` tool, so subagents see the same custom fields as the
+            parent.
+
+            [`CompiledSubAgent`][deepagents.middleware.subagents.CompiledSubAgent]
+            runnables do not inherit this schema because they are already
+            compiled — compile those runnables with a compatible state
+            schema if they need access to the same custom state fields.
+            Remote
+            [`AsyncSubAgent`][deepagents.middleware.async_subagents.AsyncSubAgent]
+            specs likewise use the schema configured on the remote graph.
+
+            ```python
+            from deepagents.graph import DeepAgentState
+
+
+            class MyState(DeepAgentState):
+                page_url: str
+                file_urls: list[str]
+
+
+            agent = create_deep_agent(model=..., state_schema=MyState)
+            ```
         context_schema: Schema class that defines immutable run-scoped context.
 
             Passed through to [`create_agent`][langchain.agents.create_agent].
@@ -474,6 +510,10 @@ def create_deep_agent(  # noqa: C901, PLR0912, PLR0915  # Complex graph assembly
             distinct middleware classes, or matches no entry in the assembled
             stack.
     """
+    # `DeepAgentState` is a `TypedDict`; TypedDicts disallow `issubclass`, so the
+    # subclass constraint on `state_schema` is enforced by typing alone and not
+    # validated at runtime.
+
     _model_spec: str | None = model if isinstance(model, str) else None
 
     if model is None:
@@ -692,6 +732,7 @@ def create_deep_agent(  # noqa: C901, PLR0912, PLR0915  # Complex graph assembly
             # see which subagents exist. None (default) uses the built-in
             # template. Stale keys silently no-op if the tool is renamed.
             task_description=_profile.tool_description_overrides.get("task"),
+            state_schema=state_schema,
         )
         deepagent_middleware.append(sub_agent_middleware)
     deepagent_middleware.extend(
@@ -774,7 +815,7 @@ def create_deep_agent(  # noqa: C901, PLR0912, PLR0915  # Complex graph assembly
         debug=debug,
         name=name,
         cache=cache,
-        state_schema=_DeepAgentState,
+        state_schema=state_schema if state_schema is not None else DeepAgentState,
         transformers=[_subagent_factory],
     ).with_config(
         {

libs/deepagents/deepagents/middleware/subagents.py

@@ -162,6 +162,13 @@ class CompiledSubAgent(TypedDict):
         This is required for the subagent to communicate results back to
         the main agent.
 
+    !!! note
+
+        `CompiledSubAgent` runnables are used as provided. They do not
+        inherit `create_deep_agent(state_schema=...)`; if the runnable
+        needs custom state fields, compile it with a compatible state
+        schema yourself.
+
     When the subagent completes, the parent reads the returned state:
     if `structured_response` is non-`None`, it is JSON-serialized and used as
     the `ToolMessage` content; otherwise, the last non-empty `AIMessage`
@@ -637,6 +644,11 @@ class SubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
         system_prompt: Instructions appended to main agent's system prompt
             about how to use the task tool.
         task_description: Custom description for the task tool.
+        state_schema: Base graph state schema forwarded to declarative
+            `SubAgent` specs when their runnables are compiled.
+
+            Leave unset to use `create_agent`'s default. `CompiledSubAgent`
+            entries are unaffected — callers own those runnables' schemas.
 
     Example:
         ```python
@@ -671,6 +683,7 @@ class SubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
         subagents: Sequence[SubAgent | CompiledSubAgent],
         system_prompt: str | None = TASK_SYSTEM_PROMPT,
         task_description: str | None = None,
+        state_schema: type | None = None,
     ) -> None:
         """Initialize the `SubAgentMiddleware`."""
         super().__init__()
@@ -680,6 +693,7 @@ class SubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
             raise ValueError(msg)
         self._backend = backend
         self._subagents = subagents
+        self._state_schema = state_schema
         subagent_specs = self._get_subagents()
         self.subagent_names: frozenset[str] = frozenset(spec["name"] for spec in subagent_specs)
         """Declared subagent names. Public so streamers can discover them
@@ -733,18 +747,20 @@ class SubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
             if interrupt_on:
                 middleware.append(HumanInTheLoopMiddleware(interrupt_on=interrupt_on))
 
+            create_agent_kwargs: dict[str, Any] = {
+                "system_prompt": spec["system_prompt"],
+                "tools": spec["tools"],
+                "middleware": middleware,
+                "name": spec["name"],
+                "response_format": spec.get("response_format"),
+            }
+            if self._state_schema is not None:
+                create_agent_kwargs["state_schema"] = self._state_schema
             specs.append(
                 {
                     "name": spec["name"],
                     "description": spec["description"],
-                    "runnable": create_agent(
-                        model,
-                        system_prompt=spec["system_prompt"],
-                        tools=spec["tools"],
-                        middleware=middleware,
-                        name=spec["name"],
-                        response_format=spec.get("response_format"),
-                    ),
+                    "runnable": create_agent(model, **create_agent_kwargs),
                 }
             )
 

libs/deepagents/tests/unit_tests/test_graph.py

@@ -20,6 +20,7 @@ from deepagents.graph import (
     _REQUIRED_MIDDLEWARE_CLASSES,
     _REQUIRED_MIDDLEWARE_NAMES,
     BASE_AGENT_PROMPT,
+    DeepAgentState,
     create_deep_agent,
     get_default_model,
 )
@@ -698,6 +699,60 @@ class TestExtraMiddlewareWiring:
             _HARNESS_PROFILES.update(original)
 
 
+class TestStateSchema:
+    """End-to-end tests for the `state_schema` parameter on `create_deep_agent`."""
+
+    def test_default_state_schema_uses_deep_agent_state(self) -> None:
+        fake_model = GenericFakeChatModel(messages=iter([AIMessage(content="ok")]))
+        fake_agent = MagicMock()
+        fake_agent.with_config.return_value = "compiled-agent"
+
+        with (
+            patch("deepagents.graph.resolve_model", return_value=fake_model),
+            patch("deepagents.graph.create_agent", return_value=fake_agent) as mock_create,
+        ):
+            create_deep_agent(model="testprov:some-model")
+
+        assert mock_create.call_args.kwargs["state_schema"] is DeepAgentState
+
+    def test_custom_state_schema_passed_through(self) -> None:
+        class MyState(DeepAgentState):
+            page_url: str
+            file_urls: list[str]
+
+        fake_model = GenericFakeChatModel(messages=iter([AIMessage(content="ok")]))
+        fake_agent = MagicMock()
+        fake_agent.with_config.return_value = "compiled-agent"
+
+        with (
+            patch("deepagents.graph.resolve_model", return_value=fake_model),
+            patch("deepagents.graph.create_agent", return_value=fake_agent) as mock_create,
+        ):
+            create_deep_agent(model="testprov:some-model", state_schema=MyState)
+
+        assert mock_create.call_args.kwargs["state_schema"] is MyState
+
+    def test_custom_state_schema_propagates_to_subagent_middleware(self) -> None:
+        """Custom schema reaches `SubAgentMiddleware` so declarative subagents compile with it."""
+
+        class MyState(DeepAgentState):
+            page_url: str
+
+        fake_model = GenericFakeChatModel(messages=iter([AIMessage(content="ok")]))
+        fake_agent = MagicMock()
+        fake_agent.with_config.return_value = "compiled-agent"
+
+        with (
+            patch("deepagents.graph.resolve_model", return_value=fake_model),
+            patch("deepagents.graph.create_agent", return_value=fake_agent) as mock_create,
+        ):
+            create_deep_agent(model="testprov:some-model", state_schema=MyState)
+
+        mw_stack = mock_create.call_args.kwargs["middleware"]
+        sub_mw = next(m for m in mw_stack if isinstance(m, SubAgentMiddleware))
+        assert sub_mw._state_schema is MyState
+
+
 class _OtherStubMW(AgentMiddleware[Any, Any, Any]):
     """Second stub class so exclusion tests can assert coexistence with `_StubMW`."""