fix: require patched async LangGraph runtime

The old LangGraph version range allowed releases affected by a published security advisory, and the adapter called synchronous state and invocation APIs from an async path, including a state lookup with no checkpointer. This requires a patched LangGraph release and moves the adapter to the async compiled-state APIs (CompiledStateGraph, aget_state, ainvoke), skipping the state lookup when the graph has no checkpointer.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 950895002

Author: GWeale

src/google/adk/agents/langgraph_agent.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from typing import Any
 from typing import AsyncGenerator
 from typing import Union
 
@@ -23,7 +24,7 @@
 from langchain_core.messages import HumanMessage
 from langchain_core.messages import SystemMessage
 from langchain_core.runnables.config import RunnableConfig
-from langgraph.graph.graph import CompiledGraph
+from langgraph.graph.state import CompiledStateGraph
 from pydantic import ConfigDict
 from typing_extensions import override
 
@@ -53,14 +54,20 @@ def _get_last_human_messages(
 
 
 class LangGraphAgent(BaseAgent):
-  """Currently a concept implementation, supports single and multi-turn."""
+  """Adapts a compiled LangGraph state graph for single or multi-turn use.
+
+  When using a persistent checkpointer, set ``LANGGRAPH_STRICT_MSGPACK=true``
+  before importing LangGraph and compiling the graph. LangGraph's patched
+  releases provide schema-derived checkpoint allowlisting, but do not enable
+  strict deserialization by default.
+  """
 
   model_config = ConfigDict(
       arbitrary_types_allowed=True,
   )
   """The pydantic model config."""
 
-  graph: CompiledGraph
+  graph: CompiledStateGraph
 
   instruction: str = ''
 
@@ -73,13 +80,16 @@ async def _run_async_impl(
     # Needed for langgraph checkpointer (for subsequent invocations; multi-turn)
     config: RunnableConfig = {'configurable': {'thread_id': ctx.session.id}}
 
-    # Add instruction as SystemMessage if graph state is empty
-    current_graph_state = self.graph.get_state(config)
-    graph_messages = (
-        current_graph_state.values.get('messages', [])
-        if current_graph_state.values
-        else []
-    )
+    # Add instruction as SystemMessage if graph state is empty. State lookup is
+    # only valid when the compiled graph has a checkpointer.
+    graph_messages: list[Any] = []
+    if self.graph.checkpointer:
+      current_graph_state = await self.graph.aget_state(config)
+      graph_messages = (
+          current_graph_state.values.get('messages', [])
+          if current_graph_state.values
+          else []
+      )
     messages: list[BaseMessage] = (
         [SystemMessage(content=self.instruction)]
         if self.instruction and not graph_messages
@@ -89,7 +99,7 @@ async def _run_async_impl(
     messages += self._get_messages(ctx.session.events)
 
     # Use the Runnable
-    final_state = self.graph.invoke({'messages': messages}, config)
+    final_state = await self.graph.ainvoke({'messages': messages}, config)
     result = final_state['messages'][-1].content
 
     result_event = Event(

tests/unittests/agents/test_langgraph_agent.py

@@ -12,85 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import AsyncMock
 from unittest.mock import MagicMock
 
 import pytest
 
-# Skip all tests in this module if LangGraph dependencies are not available
-LANGGRAPH_AVAILABLE = True
-try:
-  from google.adk.agents.invocation_context import InvocationContext
-  from google.adk.agents.langgraph_agent import LangGraphAgent
-  from google.adk.events.event import Event
-  from google.adk.plugins.plugin_manager import PluginManager
-  from google.genai import types
-  from langchain_core.messages import AIMessage
-  from langchain_core.messages import HumanMessage
-  from langchain_core.messages import SystemMessage
-  from langgraph.graph.graph import CompiledGraph
-except ImportError:
-  LANGGRAPH_AVAILABLE = False
+pytest.importorskip("langgraph", reason="LangGraph dependencies not available")
 
-  # IMPORTANT: Dummy classes are REQUIRED in this file but NOT in A2A test files.
-  # Here's why this file is different from A2A test files:
-  #
-  # 1. MODULE-LEVEL USAGE IN DECORATORS:
-  #    This file uses @pytest.mark.parametrize decorator with complex nested structures
-  #    that directly reference imported types like Event(), types.Content(), types.Part.from_text().
-  #    These decorator expressions are evaluated during MODULE COMPILATION TIME,
-  #    not during test execution time.
-  #
-  # 2. A2A TEST FILES PATTERN:
-  #    Most A2A test files only use imported types within test method bodies:
-  #    - Inside test functions: def test_something(): Message(...)
-  #    - These are evaluated during TEST EXECUTION TIME when tests are skipped
-  #    - No NameError occurs because skipped tests don't execute their bodies
-  #
-  # 3. WHAT HAPPENS WITHOUT DUMMIES:
-  #    If we remove dummy classes from this file:
-  #    - Python tries to compile the @pytest.mark.parametrize decorator
-  #    - It encounters Event(...), types.Content(...), etc.
-  #    - These names are undefined → NameError during module compilation
-  #    - Test collection fails before pytest.mark.skipif can even run
-  #
-  # 4. WHY DUMMIES WORK:
-  #    - DummyTypes() can be called like Event() → returns DummyTypes instance
-  #    - DummyTypes.__getattr__ handles types.Content → returns DummyTypes instance
-  #    - DummyTypes.__call__ handles types.Part.from_text() → returns DummyTypes instance
-  #    - The parametrize decorator gets dummy objects instead of real ones
-  #    - Tests still get skipped due to pytestmark, so dummies never actually run
-  #
-  # 5. EXCEPTION CASES IN A2A FILES:
-  #    A few A2A files DID need dummies initially because they had:
-  #    - Type annotations: def create_helper(x: str) -> Message
-  #    - But we removed those type annotations to eliminate the need for dummies
-  #
-  # This file cannot avoid dummies because the parametrize decorator usage
-  # is fundamental to the test structure and cannot be easily refactored.
-
-  class DummyTypes:
-
-    def __getattr__(self, name):
-      return DummyTypes()
-
-    def __call__(self, *args, **kwargs):
-      return DummyTypes()
-
-  InvocationContext = DummyTypes()
-  LangGraphAgent = DummyTypes()
-  Event = DummyTypes()
-  PluginManager = DummyTypes()
-  types = (
-      DummyTypes()
-  )  # Must support chained calls like types.Content(), types.Part.from_text()
-  AIMessage = DummyTypes()
-  HumanMessage = DummyTypes()
-  SystemMessage = DummyTypes()
-  CompiledGraph = DummyTypes()
-
-pytestmark = pytest.mark.skipif(
-    not LANGGRAPH_AVAILABLE, reason="LangGraph dependencies not available"
-)
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.langgraph_agent import LangGraphAgent
+from google.adk.events.event import Event
+from google.adk.plugins.plugin_manager import PluginManager
+from google.genai import types
+from langchain_core.messages import AIMessage
+from langchain_core.messages import HumanMessage
+from langchain_core.messages import SystemMessage
+from langgraph.graph import MessagesState
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
 
 
 @pytest.mark.parametrize(
@@ -219,15 +158,15 @@ def __call__(self, *args, **kwargs):
 async def test_langgraph_agent(
     checkpointer_value, events_list, expected_messages
 ):
-  mock_graph = MagicMock(spec=CompiledGraph)
+  mock_graph = MagicMock(spec=CompiledStateGraph)
   mock_graph_state = MagicMock()
   mock_graph_state.values = {}
-  mock_graph.get_state.return_value = mock_graph_state
+  mock_graph.aget_state = AsyncMock(return_value=mock_graph_state)
 
   mock_graph.checkpointer = checkpointer_value
-  mock_graph.invoke.return_value = {
-      "messages": [AIMessage(content="test response")]
-  }
+  mock_graph.ainvoke = AsyncMock(
+      return_value={"messages": [AIMessage(content="test response")]}
+  )
 
   mock_parent_context = MagicMock(spec=InvocationContext)
   mock_parent_context._state_schema = None
@@ -257,8 +196,57 @@ async def test_langgraph_agent(
   assert result_event.author == "weather_agent"
   assert result_event.content.parts[0].text == "test response"
 
-  mock_graph.invoke.assert_called_once()
-  mock_graph.invoke.assert_called_with(
+  if checkpointer_value:
+    mock_graph.aget_state.assert_awaited_once_with(
+        {"configurable": {"thread_id": mock_session.id}}
+    )
+  else:
+    mock_graph.aget_state.assert_not_awaited()
+  mock_graph.ainvoke.assert_awaited_once_with(
       {"messages": expected_messages},
       {"configurable": {"thread_id": mock_session.id}},
   )
+
+
+@pytest.mark.asyncio
+async def test_langgraph_agent_runs_real_compiled_state_graph():
+  """A real compiled state graph runs through the asynchronous adapter."""
+  observed_messages = []
+
+  def respond(state: MessagesState) -> dict[str, list[AIMessage]]:
+    observed_messages.extend(state["messages"])
+    return {"messages": [AIMessage(content="real graph response")]}
+
+  graph_builder = StateGraph(MessagesState)
+  graph_builder.add_node("respond", respond)
+  graph_builder.set_entry_point("respond")
+  graph_builder.set_finish_point("respond")
+  graph = graph_builder.compile()
+
+  parent_context = MagicMock(spec=InvocationContext)
+  parent_context._state_schema = None
+  mock_session = MagicMock()
+  mock_session.id = "session-id"
+  mock_session.events = []
+  parent_context.session = mock_session
+  parent_context.user_content = types.Content(
+      role="user", parts=[types.Part.from_text(text="test prompt")]
+  )
+  parent_context.branch = "parent_agent"
+  parent_context.end_invocation = False
+  parent_context.invocation_id = "test_invocation_id"
+  parent_context.model_copy.return_value = parent_context
+  parent_context.plugin_manager = PluginManager(plugins=[])
+  agent = LangGraphAgent(
+      name="weather_agent",
+      instruction="test system prompt",
+      graph=graph,
+  )
+
+  events = [event async for event in agent.run_async(parent_context)]
+
+  assert len(events) == 1
+  assert events[0].content.parts[0].text == "real graph response"
+  assert len(observed_messages) == 1
+  assert isinstance(observed_messages[0], SystemMessage)
+  assert observed_messages[0].content == "test system prompt"