From eaa1b3764522790a43bac306e2b91aacf5ac30d9 Mon Sep 17 00:00:00 2001
From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
Date: Sun, 1 Jun 2025 21:52:35 -0400
Subject: [PATCH] Require `state_schema` in `StateGraph.__init__` (#4897)

---
 libs/langgraph/bench/fanout_to_subgraph.py |  8 ++++++--
 libs/langgraph/langgraph/graph/state.py    | 23 ++++++----------------
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/libs/langgraph/bench/fanout_to_subgraph.py b/libs/langgraph/bench/fanout_to_subgraph.py
index fc8c7133b..107366ca3 100644
--- a/libs/langgraph/bench/fanout_to_subgraph.py
+++ b/libs/langgraph/bench/fanout_to_subgraph.py
@@ -21,6 +21,8 @@ def fanout_to_subgraph() -> StateGraph:
     class JokeOutput(TypedDict):
         jokes: list[str]
 
+    class JokeState(JokeInput, JokeOutput): ...
+
     async def bump(state: JokeOutput):
         return {"jokes": [state["jokes"][0] + " a"]}
 
@@ -35,7 +37,7 @@ def fanout_to_subgraph() -> StateGraph:
         return END if state["jokes"][0].endswith(" a" * 10) else "bump"
 
     # subgraph
-    subgraph = StateGraph(input=JokeInput, output=JokeOutput)
+    subgraph = StateGraph(JokeState, input=JokeInput, output=JokeOutput)
     subgraph.add_node("edit", edit)
     subgraph.add_node("generate", generate)
     subgraph.add_node("bump", bump)
@@ -69,6 +71,8 @@ def fanout_to_subgraph_sync() -> StateGraph:
     class JokeOutput(TypedDict):
         jokes: list[str]
 
+    class JokeState(JokeInput, JokeOutput): ...
+
     def bump(state: JokeOutput):
         return {"jokes": [state["jokes"][0] + " a"]}
 
@@ -83,7 +87,7 @@ def fanout_to_subgraph_sync() -> StateGraph:
         return END if state["jokes"][0].endswith(" a" * 10) else "bump"
 
     # subgraph
-    subgraph = StateGraph(input=JokeInput, output=JokeOutput)
+    subgraph = StateGraph(JokeState, input=JokeInput, output=JokeOutput)
     subgraph.add_node("edit", edit)
     subgraph.add_node("generate", generate)
     subgraph.add_node("bump", bump)
diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index ccd773a58..bbfdd439e 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -25,7 +25,6 @@ from langchain_core.runnables import Runnable, RunnableConfig
 from pydantic import BaseModel
 from typing_extensions import Self
 
-from langgraph._api.deprecation import LangGraphDeprecationWarning
 from langgraph.cache.base import BaseCache
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
@@ -170,27 +169,17 @@ class StateGraph:
 
     def __init__(
         self,
-        state_schema: Optional[type[Any]] = None,
+        state_schema: type[Any],
         config_schema: Optional[type[Any]] = None,
         *,
         input: Optional[type[Any]] = None,
         output: Optional[type[Any]] = None,
     ) -> None:
-        if state_schema is None:
-            if input is None or output is None:
-                raise ValueError("Must provide state_schema or input and output")
-            state_schema = input
-            warnings.warn(
-                "Initializing StateGraph without state_schema is deprecated. "
-                "Please pass in an explicit state_schema instead of just an input and output schema.",
-                LangGraphDeprecationWarning,
-                stacklevel=2,
-            )
-        else:
-            if input is None:
-                input = state_schema
-            if output is None:
-                output = state_schema
+        if input is None:
+            input = state_schema
+        if output is None:
+            output = state_schema
+
         self.nodes = {}
         self.edges = set[tuple[str, str]]()
         self.branches = defaultdict(dict)