diff --git a/examples/rag/langgraph_crag_local.ipynb b/examples/rag/langgraph_crag_local.ipynb
index 19dc84410..c454127f2 100644
--- a/examples/rag/langgraph_crag_local.ipynb
+++ b/examples/rag/langgraph_crag_local.ipynb
@@ -707,7 +707,9 @@
     "    \"\"\"\n",
     "    Find all tool calls in the messages returned\n",
     "    \"\"\"\n",
-    "    tool_calls = [tc['name'] for m in messages['messages'] for tc in getattr(m, 'tool_calls', [])]\n",
+    "    tool_calls = [\n",
+    "        tc[\"name\"] for m in messages[\"messages\"] for tc in getattr(m, \"tool_calls\", [])\n",
+    "    ]\n",
     "    return tool_calls\n",
     "\n",
     "\n",
diff --git a/libs/cli/examples/graphs/agent.py b/libs/cli/examples/graphs/agent.py
index f39df4cae..931982061 100644
--- a/libs/cli/examples/graphs/agent.py
+++ b/libs/cli/examples/graphs/agent.py
@@ -7,6 +7,7 @@ from langchain_core.messages import BaseMessage
 from langchain_openai import ChatOpenAI
 from langgraph.graph import END, StateGraph, add_messages
 from langgraph.prebuilt import ToolNode
+from langgraph.runtime import Runtime
 
 tools = [TavilySearchResults(max_results=1)]
 
@@ -17,6 +18,10 @@ model_anth = model_anth.bind_tools(tools)
 model_oai = model_oai.bind_tools(tools)
 
 
+class AgentContext(TypedDict):
+    model: Literal["anthropic", "openai"]
+
+
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], add_messages]
 
@@ -34,8 +39,8 @@ def should_continue(state):
 
 
 # Define the function that calls the model
-def call_model(state, config):
-    if config["configurable"].get("model", "anthropic") == "anthropic":
+def call_model(state, runtime: Runtime[AgentContext]):
+    if runtime.context.get("model", "anthropic") == "anthropic":
         model = model_anth
     else:
         model = model_oai
@@ -49,12 +54,8 @@ def call_model(state, config):
 tool_node = ToolNode(tools)
 
 
-class ContextSchema(TypedDict):
-    model: Literal["anthropic", "openai"]
-
-
 # Define a new graph
-workflow = StateGraph(AgentState, context_schema=ContextSchema)
+workflow = StateGraph(AgentState, context_schema=AgentContext)
 
 # Define the two nodes we will cycle between
 workflow.add_node("agent", call_model)
diff --git a/libs/cli/examples/graphs_reqs_a/graphs_submod/agent.py b/libs/cli/examples/graphs_reqs_a/graphs_submod/agent.py
index 1073110cb..7971b3ce0 100644
--- a/libs/cli/examples/graphs_reqs_a/graphs_submod/agent.py
+++ b/libs/cli/examples/graphs_reqs_a/graphs_submod/agent.py
@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Annotated, TypedDict
+from typing import Annotated, Literal, TypedDict
 
 from langchain_anthropic import ChatAnthropic
 from langchain_community.tools.tavily_search import TavilySearchResults
@@ -8,6 +8,7 @@ from langchain_core.messages import BaseMessage
 from langchain_openai import ChatOpenAI
 from langgraph.graph import END, StateGraph, add_messages
 from langgraph.prebuilt import ToolNode
+from langgraph.runtime import Runtime
 
 tools = [TavilySearchResults(max_results=1)]
 
@@ -21,6 +22,10 @@ prompt = open(Path(__file__).parent.parent / "prompt.txt").read()
 subprompt = open(Path(__file__).parent / "subprompt.txt").read()
 
 
+class AgentContext(TypedDict):
+    model: Literal["anthropic", "openai"]
+
+
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], add_messages]
 
@@ -38,8 +43,8 @@ def should_continue(state):
 
 
 # Define the function that calls the model
-def call_model(state, config):
-    if config["configurable"].get("model", "anthropic") == "anthropic":
+def call_model(state, runtime: Runtime[AgentContext]):
+    if runtime.context.get("model", "anthropic") == "anthropic":
         model = model_anth
     else:
         model = model_oai
@@ -52,9 +57,8 @@ def call_model(state, config):
 # Define the function to execute tools
 tool_node = ToolNode(tools)
 
-
 # Define a new graph
-workflow = StateGraph(AgentState)
+workflow = StateGraph(AgentState, context_schema=AgentContext)
 
 # Define the two nodes we will cycle between
 workflow.add_node("agent", call_model)
diff --git a/libs/cli/examples/graphs_reqs_b/graphs_submod/agent.py b/libs/cli/examples/graphs_reqs_b/graphs_submod/agent.py
index 1073110cb..eb81a0a7f 100644
--- a/libs/cli/examples/graphs_reqs_b/graphs_submod/agent.py
+++ b/libs/cli/examples/graphs_reqs_b/graphs_submod/agent.py
@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Annotated, TypedDict
+from typing import Annotated, Literal, TypedDict
 
 from langchain_anthropic import ChatAnthropic
 from langchain_community.tools.tavily_search import TavilySearchResults
@@ -8,6 +8,7 @@ from langchain_core.messages import BaseMessage
 from langchain_openai import ChatOpenAI
 from langgraph.graph import END, StateGraph, add_messages
 from langgraph.prebuilt import ToolNode
+from langgraph.runtime import Runtime
 
 tools = [TavilySearchResults(max_results=1)]
 
@@ -21,6 +22,10 @@ prompt = open(Path(__file__).parent.parent / "prompt.txt").read()
 subprompt = open(Path(__file__).parent / "subprompt.txt").read()
 
 
+class AgentContext(TypedDict):
+    model: Literal["anthropic", "openai"]
+
+
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], add_messages]
 
@@ -38,8 +43,8 @@ def should_continue(state):
 
 
 # Define the function that calls the model
-def call_model(state, config):
-    if config["configurable"].get("model", "anthropic") == "anthropic":
+def call_model(state, runtime: Runtime[AgentContext]):
+    if runtime.context.get("model", "anthropic") == "anthropic":
         model = model_anth
     else:
         model = model_oai
@@ -54,7 +59,7 @@ tool_node = ToolNode(tools)
 
 
 # Define a new graph
-workflow = StateGraph(AgentState)
+workflow = StateGraph(AgentState, context_schema=AgentContext)
 
 # Define the two nodes we will cycle between
 workflow.add_node("agent", call_model)