From 6616f749c612092891e032dab294e9fd3664188b Mon Sep 17 00:00:00 2001
From: Salah <nh.salah@gmail.com>
Date: Sat, 25 Jul 2026 15:17:51 +0400
Subject: [PATCH 1/5] feat(adapters/langgraph): make agent structured output
 not hang or silently vanish

---
 .../src/pyagentspec/adapters/_utils.py        |  10 +
 .../adapters/langgraph/_langgraphconverter.py |  20 +-
 .../adapters/langgraph/_node_execution.py     |  28 ++-
 .../adapters/langgraph/_structured_output.py  | 103 +++++++++
 .../pyagentspec/adapters/langgraph/_types.py  |   4 +-
 .../langgraph/flows/test_agentnode.py         |  74 +++++++
 .../langgraph/test_structured_output.py       | 202 ++++++++++++++++++
 7 files changed, 428 insertions(+), 13 deletions(-)
 create mode 100644 pyagentspec/src/pyagentspec/adapters/langgraph/_structured_output.py
 create mode 100644 pyagentspec/tests/adapters/langgraph/test_structured_output.py

diff --git a/pyagentspec/src/pyagentspec/adapters/_utils.py b/pyagentspec/src/pyagentspec/adapters/_utils.py
index e221d613..a69bf07e 100644
--- a/pyagentspec/src/pyagentspec/adapters/_utils.py
+++ b/pyagentspec/src/pyagentspec/adapters/_utils.py
@@ -184,6 +184,16 @@ def _build_type_from_schema(
     return mapping.get(t, Any)
 
 
+def is_single_string_output(expected_outputs: List[AgentSpecProperty]) -> bool:
+    """Whether the declared outputs are a single string property.
+
+    Such an output is the model's free text, so an adapter can read it from the final
+    message instead of asking for structured generation. Lives here rather than in one
+    adapter because it is a property of the declared outputs.
+    """
+    return len(expected_outputs) == 1 and expected_outputs[0].type == "string"
+
+
 def create_pydantic_model_from_properties(
     model_name: str, properties: List[AgentSpecProperty]
 ) -> type[BaseModel]:
diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
index 0ca55cb0..e1cbc11b 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
@@ -35,6 +35,7 @@
     SchemaRegistry,
     _build_type_from_schema,
     create_pydantic_model_from_properties,
+    is_single_string_output,
 )
 from pyagentspec.adapters.langgraph._execution_span import patch_with_execution_span
 from pyagentspec.adapters.langgraph._managerworkers import (
@@ -50,6 +51,9 @@
     NodeExecutor,
     extract_outputs_from_invoke_result,
 )
+from pyagentspec.adapters.langgraph._structured_output import (
+    StructuredOutputGuard,
+)
 from pyagentspec.adapters.langgraph._types import (
     AgentState,
     BaseCallbackHandler,
@@ -1235,8 +1239,9 @@ def _create_react_agent_with_given_info(
         state_schema: Optional[Any] = None
         response_format: Any = None
 
-        # Build response (output) model (used for response_format)
-        if outputs:
+        # A single string output is read from the final message instead, so it needs no
+        # response_format. Same rule as LlmNodeExecutor.
+        if outputs and not is_single_string_output(outputs):
             output_model = create_pydantic_model_from_properties("AgentOutputModel", outputs)
             # Explicitly use ToolStrategy instead of letting LangChain select a provider
             # strategy. OpenAI-compatible models do not necessarily support provider-native
@@ -1266,6 +1271,17 @@ def _create_react_agent_with_given_info(
             response_format=response_format,
             state_schema=state_schema,
         )
+        if output_model is not None:
+            # Rebuild rather than append: `middleware` is shared across every agent of a
+            # swarm / manager-workers graph.
+            middleware = [
+                *middleware,
+                StructuredOutputGuard(
+                    agent_name=name,
+                    output_titles=[output.title for output in outputs],
+                    model_id=llm_config.model_id,
+                ),
+            ]
         if middleware:
             create_agent_kwargs["middleware"] = middleware
         compiled_graph: CompiledStateGraph[Any, Any, Any] = langchain_agents.create_agent(
diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_node_execution.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_node_execution.py
index a1398df9..80e1141b 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_node_execution.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_node_execution.py
@@ -16,7 +16,11 @@
     maybe_warn_about_unrestricted_templated_url,
     validate_url_against_allow_list,
 )
-from pyagentspec.adapters._utils import render_nested_object_template, render_template
+from pyagentspec.adapters._utils import (
+    is_single_string_output,
+    render_nested_object_template,
+    render_template,
+)
 from pyagentspec.adapters.langgraph._types import (
     BaseChatModel,
     BaseMessage,
@@ -606,20 +610,13 @@ def __init__(self, node: AgentSpecLlmNode, llm: BaseChatModel) -> None:
         super().__init__(node)
         if not isinstance(self.node, AgentSpecLlmNode):
             raise TypeError("LlmNodeExecutor can only be initialized with LlmNode")
-        outputs = self.node.outputs
-        if outputs is not None and len(outputs) == 1 and outputs[0].type == "string":
-            self.requires_structured_generation = False
-        else:
-            self.requires_structured_generation = True
         if not isinstance(llm, BaseChatModel):
             raise TypeError("Llm can only be initialized with a BaseChatModel")
 
         self.llm: BaseChatModel = llm
 
         node_outputs = self.node.outputs or []
-        self.requires_structured_generation = not (
-            len(node_outputs) == 1 and node_outputs[0].type == "string"
-        )
+        self.requires_structured_generation = not is_single_string_output(node_outputs)
 
         self.structured_llm: Any = None
 
@@ -933,7 +930,7 @@ def extract_outputs_from_invoke_result(
     # Extracts the outputs from the return value of an invoke call made on an agent
     # The outputs are typically exposed as part of the `structured_response`, or as entries in the result directly.
     # We give priority to the latter.
-    return {
+    outputs = {
         # Defaults if available
         **{
             output.title: output.default
@@ -949,3 +946,14 @@ def extract_outputs_from_invoke_result(
             if output.title in result
         },
     }
+    # No response_format is requested for a single string output, so read it from the
+    # final message.
+    if is_single_string_output(expected_outputs):
+        title = expected_outputs[0].title
+        if title not in outputs:
+            messages = result.get("messages")
+            if messages:
+                content = getattr(messages[-1], "content", None)
+                if content is not None:
+                    outputs[title] = content
+    return outputs
diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_structured_output.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_structured_output.py
new file mode 100644
index 00000000..9f4fbe43
--- /dev/null
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_structured_output.py
@@ -0,0 +1,103 @@
+# Copyright © 2025, 2026 Oracle and/or its affiliates.
+#
+# This software is under the Apache License 2.0
+# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License
+# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option.
+
+"""Makes an agent fail loudly when its model will not produce structured output.
+
+``create_agent(response_format=...)`` ties the agent's control flow to getting a
+parseable structured response. A model that never produces one fails in one of two
+ways, depending on whether the agent has real tools:
+
+* no tools: the agent loops. ``create_agent`` sets ``recursion_limit=9999``, so that is
+  thousands of model calls before ``GraphRecursionError``.
+* with tools: routing treats a turn with no tool calls as "done", so the agent exits
+  with no ``structured_response`` and the declared outputs come back empty, with no
+  error at all (upstream langchain issue #36349).
+
+:class:`StructuredOutputGuard` covers both. An agent whose only declared output is a
+string avoids the problem entirely, since ``response_format`` is skipped for it (see
+:func:`pyagentspec.adapters._utils.is_single_string_output`).
+"""
+
+from typing import Any, Dict, List, NoReturn, Optional
+
+from typing_extensions import NotRequired
+
+from pyagentspec.adapters.langgraph._types import AgentMiddleware, AgentState
+
+# A compliant run spends at most one turn without a structured response, so a small
+# bound separates "model is working on it" from "model will never comply".
+DEFAULT_MAX_STRUCTURED_OUTPUT_ATTEMPTS = 3
+
+# Must match the field on _StructuredOutputAttemptState.
+_ATTEMPTS_STATE_KEY = "structured_output_attempts"
+
+
+class StructuredOutputNotProducedError(RuntimeError):
+    """Raised when a model will not produce the structured response an agent declares.
+
+    Subclasses ``RuntimeError`` so existing handlers around agent execution still catch it.
+    """
+
+
+class _StructuredOutputAttemptState(AgentState):
+    # On the state rather than the instance, so the count is per run and not shared
+    # between concurrent ones.
+    structured_output_attempts: NotRequired[int]
+
+
+class StructuredOutputGuard(AgentMiddleware):
+    """Fails an agent that cannot produce the structured output it declares.
+
+    Install only when ``response_format`` is set. Without it, a turn with no tool calls
+    just means the agent is done and an empty ``structured_response`` is fine.
+    """
+
+    state_schema = _StructuredOutputAttemptState
+
+    def __init__(
+        self,
+        *,
+        agent_name: str,
+        output_titles: List[str],
+        model_id: str,
+        max_attempts: int = DEFAULT_MAX_STRUCTURED_OUTPUT_ATTEMPTS,
+    ) -> None:
+        super().__init__()
+        self.agent_name = agent_name
+        self.output_titles = output_titles
+        self.model_id = model_id
+        self.max_attempts = max_attempts
+
+    def _fail(self, detail: str) -> NoReturn:
+        raise StructuredOutputNotProducedError(
+            f"Agent {self.agent_name!r} did not produce a structured response matching its "
+            f"declared outputs ({', '.join(self.output_titles)}). {detail} Model "
+            f"{self.model_id!r} may not support structured output. Declare a single string "
+            f"output to get the model's free text instead, or use a model that supports "
+            f"structured output."
+        )
+
+    # No async variants: LangGraph routes async runs to the sync hooks, and both are pure.
+    def after_model(self, state: Any, runtime: Any) -> Optional[Dict[str, int]]:
+        messages = state.get("messages") or []
+        last = messages[-1] if messages else None
+        if last is None or getattr(last, "type", None) != "ai":
+            return None
+        if getattr(last, "tool_calls", None):
+            # A tool call is progress, so the bound applies to consecutive failures only.
+            # Under ToolStrategy the structured response is itself a tool call.
+            return {_ATTEMPTS_STATE_KEY: 0}
+        attempts = (state.get(_ATTEMPTS_STATE_KEY) or 0) + 1
+        if attempts >= self.max_attempts:
+            self._fail(f"It answered in prose {attempts} times instead.")
+        return {_ATTEMPTS_STATE_KEY: attempts}
+
+    def after_agent(self, state: Any, runtime: Any) -> None:
+        # after_model never sees the silent case: an agent with tools exits on its first
+        # prose turn, well before the bound.
+        if "structured_response" in state:
+            return None
+        self._fail("It ended its run without one.")
diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py
index 1b44d0ad..b48afc61 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py
@@ -19,7 +19,7 @@
     import langchain_openai
     import langgraph.graph as langgraph_graph
     import langgraph_swarm
-    from langchain.agents.middleware.types import AgentState
+    from langchain.agents.middleware.types import AgentMiddleware, AgentState
     from langchain_core.callbacks import BaseCallbackHandler
     from langchain_core.language_models import BaseChatModel
     from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage
@@ -59,6 +59,7 @@
     GenerationChunk = LazyType("langchain_core.outputs", "GenerationChunk")
     LLMResult = LazyType("langchain_core.outputs", "LLMResult")
     AgentState = LazyType("langchain.agents.middleware.types", "AgentState")
+    AgentMiddleware = LazyType("langchain.agents.middleware.types", "AgentMiddleware")
 
 
 LangGraphTool: TypeAlias = Union[BaseTool, Callable[..., Any]]
@@ -136,6 +137,7 @@ class FlowOutputSchema(TypedDict):
     "ToolMessage",
     "BaseChatModel",
     "AgentState",
+    "AgentMiddleware",
     "Checkpointer",
     "interrupt",
     "RunnableConfig",
diff --git a/pyagentspec/tests/adapters/langgraph/flows/test_agentnode.py b/pyagentspec/tests/adapters/langgraph/flows/test_agentnode.py
index e5b4ba12..2dcf97ea 100644
--- a/pyagentspec/tests/adapters/langgraph/flows/test_agentnode.py
+++ b/pyagentspec/tests/adapters/langgraph/flows/test_agentnode.py
@@ -96,6 +96,80 @@ def test_agentnode_can_be_imported_and_executed(agent_flow: Flow) -> None:
     assert "car" in outputs
 
 
+def test_single_string_output_taken_from_final_message_without_structured_generation() -> None:
+    """The stub model has no structured output, so if the converter still attached a
+    ``response_format`` the output would come back empty instead of holding "42".
+    """
+    from unittest.mock import patch
+
+    from langchain_core.language_models.fake_chat_models import (
+        FakeMessagesListChatModel,
+    )
+    from langchain_core.messages import AIMessage
+    from langchain_openai import ChatOpenAI
+    from langgraph.checkpoint.memory import MemorySaver
+
+    from pyagentspec.adapters.langgraph import AgentSpecLoader
+    from pyagentspec.adapters.langgraph._langgraphconverter import (
+        AgentSpecToLangGraphConverter,
+    )
+    from pyagentspec.llms.openaicompatibleconfig import OpenAiCompatibleConfig
+
+    class _FakeModel(FakeMessagesListChatModel, ChatOpenAI):
+        pass
+
+    fake_llm = _FakeModel(responses=[AIMessage(content="42")])
+
+    answer = StringProperty(title="answer")
+    agent = Agent(
+        name="agent",
+        llm_config=OpenAiCompatibleConfig(name="agent_llm", model_id="fake", url="null"),
+        system_prompt="Answer the question.",
+        outputs=[answer],
+    )
+    agent_node = AgentNode(name="agent_node", agent=agent)
+    start_node = StartNode(name="start")
+    end_node = EndNode(name="end", outputs=[answer])
+    flow = Flow(
+        name="flow",
+        start_node=start_node,
+        nodes=[start_node, agent_node, end_node],
+        control_flow_connections=[
+            ControlFlowEdge(name="start_to_node", from_node=start_node, to_node=agent_node),
+            ControlFlowEdge(name="node_to_end", from_node=agent_node, to_node=end_node),
+        ],
+        data_flow_connections=[
+            DataFlowEdge(
+                name="answer_edge",
+                source_node=agent_node,
+                source_output=answer.title,
+                destination_node=end_node,
+                destination_input=answer.title,
+            ),
+        ],
+        outputs=[answer],
+    )
+
+    loader = AgentSpecLoader(tool_registry={}, checkpointer=MemorySaver())
+    with patch.object(
+        AgentSpecToLangGraphConverter,
+        "_llm_convert_to_langgraph",
+        autospec=True,
+        side_effect=lambda self_obj, llm_config, *a, **k: fake_llm,
+    ), patch.object(
+        FakeMessagesListChatModel,
+        "bind_tools",
+        new=lambda self_obj, *a, **k: self_obj,
+    ):
+        compiled = loader.load_component(flow)
+        result = compiled.invoke(
+            {"inputs": {}, "messages": [{"role": "user", "content": "What is 6*7?"}]},
+            {"configurable": {"thread_id": "agentnode-single-string"}},
+        )
+
+    assert result["outputs"]["answer"] == "42"
+
+
 @pytest.mark.anyio
 @retry_test(max_attempts=3, wait_between_tries=2)
 async def test_agentnode_can_be_executed_async(agent_flow: Flow) -> None:
diff --git a/pyagentspec/tests/adapters/langgraph/test_structured_output.py b/pyagentspec/tests/adapters/langgraph/test_structured_output.py
new file mode 100644
index 00000000..eb6b2bdf
--- /dev/null
+++ b/pyagentspec/tests/adapters/langgraph/test_structured_output.py
@@ -0,0 +1,202 @@
+# Copyright © 2025, 2026 Oracle and/or its affiliates.
+#
+# This software is under the Apache License 2.0
+# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License
+# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option.
+
+"""How the LangGraph adapter decides on, and bounds, agent structured output.
+
+These build no LLM config, so unlike the AgentNode flow tests they run rather than skip
+under ``SKIP_LLM_TESTS=1``.
+"""
+
+from typing import Any
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+
+from pyagentspec.adapters._utils import is_single_string_output
+from pyagentspec.adapters.langgraph._structured_output import (
+    StructuredOutputGuard,
+    StructuredOutputNotProducedError,
+)
+from pyagentspec.property import IntegerProperty, StringProperty
+
+# ─── is_single_string_output ──────────────────────────────────────────────────
+
+
+def test_is_single_string_output() -> None:
+    """A lone string output is free text, not a structured field."""
+    assert is_single_string_output([StringProperty(title="x")]) is True
+    assert is_single_string_output([]) is False
+    assert is_single_string_output([IntegerProperty(title="n")]) is False
+    assert is_single_string_output([StringProperty(title="a"), StringProperty(title="b")]) is False
+
+
+# ─── structured-output guard: after_model budget ───────────────────────────────────────
+
+
+def _guard(max_attempts: int = 3) -> StructuredOutputGuard:
+    return StructuredOutputGuard(
+        agent_name="researcher",
+        output_titles=["summary", "confidence"],
+        model_id="some-model",
+        max_attempts=max_attempts,
+    )
+
+
+def _prose(attempts: int = 0) -> dict:
+    """A model turn that answered in prose, so no structured response."""
+    return {
+        "messages": [HumanMessage(content="q"), AIMessage(content="42")],
+        "structured_output_attempts": attempts,
+    }
+
+
+def _tool_turn(attempts: int = 0) -> dict:
+    return {
+        "messages": [AIMessage(content="", tool_calls=[{"name": "s", "args": {}, "id": "c1"}])],
+        "structured_output_attempts": attempts,
+    }
+
+
+def test_prose_turn_increments_until_the_limit() -> None:
+    """Prose never satisfies response_format, so the agent cannot exit on its own."""
+    guard = _guard(max_attempts=3)
+
+    assert guard.after_model(_prose(0), runtime=None) == {"structured_output_attempts": 1}
+    assert guard.after_model(_prose(1), runtime=None) == {"structured_output_attempts": 2}
+    with pytest.raises(StructuredOutputNotProducedError):
+        guard.after_model(_prose(2), runtime=None)
+
+
+def test_error_message_names_agent_fields_and_model() -> None:
+    """Failing fast is only useful if the message says what to change."""
+    with pytest.raises(StructuredOutputNotProducedError) as excinfo:
+        _guard(max_attempts=1).after_model(_prose(), runtime=None)
+
+    message = str(excinfo.value)
+    assert "researcher" in message
+    assert "summary, confidence" in message
+    assert "some-model" in message
+    # It should point at the actionable escape hatch, not just complain.
+    assert "single string output" in message
+
+
+def test_tool_calling_turn_resets_rather_than_counting() -> None:
+    """A tool call is progress. Under ToolStrategy the structured response is itself a
+    tool call, so counting these would fail every compliant agent. The reset also keeps a
+    long run of mixed prose and tool calls from adding up to a false failure."""
+    guard = _guard(max_attempts=2)
+    # max_attempts=2 with 1 attempt already banked would raise if this counted.
+    assert guard.after_model(_tool_turn(attempts=1), runtime=None) == {
+        "structured_output_attempts": 0
+    }
+
+
+def test_non_ai_or_empty_last_message_is_ignored() -> None:
+    """Only a model turn can be a failed structured response."""
+    guard = _guard(max_attempts=1)
+    assert guard.after_model({"messages": []}, runtime=None) is None
+    assert guard.after_model({}, runtime=None) is None
+    assert (
+        guard.after_model({"messages": [ToolMessage(content="r", tool_call_id="c1")]}, runtime=None)
+        is None
+    )
+
+
+@pytest.mark.anyio
+async def test_budget_applies_on_async_runs_via_the_sync_hook() -> None:
+    """The guard has no ``aafter_model`` and relies on LangGraph routing async runs to the
+    sync hook. Agents do run async, so a silent no-op there would bring the hang back."""
+    from langchain.agents import create_agent
+    from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+    from pydantic import BaseModel
+
+    class _Answer(BaseModel):
+        summary: str
+        confidence: str
+
+    class _ProseOnlyModel(GenericFakeChatModel):
+        """Accepts the structured-output tool binding, then ignores it."""
+
+        def bind_tools(self, tools: Any, **kwargs: Any) -> Any:
+            return self
+
+    # A model that only ever answers in prose, so no structured response is produced.
+    # `response_format` is what makes that a hang rather than a clean exit.
+    model = _ProseOnlyModel(messages=iter([AIMessage(content="42")] * 50))
+    agent = create_agent(
+        model=model,
+        tools=[],
+        system_prompt="Answer.",
+        response_format=_Answer,
+        middleware=[_guard(max_attempts=3)],
+    )
+
+    with pytest.raises(StructuredOutputNotProducedError):
+        await agent.ainvoke({"messages": [HumanMessage(content="6*7?")]})
+
+
+# ─── after_agent: the silent-exit case the budget cannot see ──────────────────
+
+
+def test_after_agent_accepts_a_run_that_produced_the_response() -> None:
+    """The guard has to stay out of the way of a compliant agent."""
+    assert _guard().after_agent({"structured_response": object()}, runtime=None) is None
+
+
+def test_after_agent_rejects_a_run_that_produced_nothing() -> None:
+    """An agent with tools exits after one prose turn, spending no budget, so this hook is
+    the only place that failure shows up."""
+    with pytest.raises(StructuredOutputNotProducedError) as excinfo:
+        _guard().after_agent({"messages": [AIMessage(content="42")]}, runtime=None)
+
+    message = str(excinfo.value)
+    assert "researcher" in message
+    assert "summary, confidence" in message
+    assert "some-model" in message
+    assert "single string output" in message
+
+
+def test_agent_with_tools_raises_instead_of_losing_outputs_silently() -> None:
+    """Regression for langchain#36349. With at least one real tool, routing treats a
+    turn with no tool calls as "done", so the agent used to return no
+    ``structured_response`` and the declared outputs came back empty."""
+    from langchain.agents import create_agent
+    from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+    from langchain_core.tools import tool
+    from pydantic import BaseModel
+
+    class _Answer(BaseModel):
+        summary: str
+        confidence: str
+
+    @tool
+    def search(q: str) -> str:
+        """Search for something."""
+        return "a result"
+
+    class _ProseOnlyModel(GenericFakeChatModel):
+        def bind_tools(self, tools: Any, **kwargs: Any) -> Any:
+            return self
+
+    model = _ProseOnlyModel(messages=iter([AIMessage(content="42")] * 50))
+    agent = create_agent(
+        model=model,
+        tools=[search],
+        system_prompt="Answer.",
+        response_format=_Answer,
+        middleware=[_guard()],
+    )
+
+    with pytest.raises(StructuredOutputNotProducedError, match="ended its run without one"):
+        agent.invoke({"messages": [HumanMessage(content="6*7?")]})
+
+
+def test_limiter_state_schema_declares_the_counter() -> None:
+    """The counter belongs on the state schema, not the instance, so concurrent runs do
+    not share it."""
+    schema: Any = StructuredOutputGuard.state_schema
+    assert "structured_output_attempts" in schema.__annotations__
+    assert "structured_output_attempts" in schema.__optional_keys__

From 1b2491c561fa5a7bf4d1f1827d41dcb02d3eea3c Mon Sep 17 00:00:00 2001
From: Salah <nh.salah@gmail.com>
Date: Mon, 31 Aug 2026 21:14:32 +0400
Subject: [PATCH 2/5] refactor(adapters/langgraph): apply review feedback on
 the output guard

- dunder-wrap the guard's state key (__structured_output_attempts__) so
  it cannot collide with a field the agent's own state schema declares
- document how after_model classifies each turn and why an
  unclassifiable one leaves the counter untouched instead of resetting
- rename _structured_output.py to _agent_output_guard.py and the test
  file to match; the is_single_string_output test moves next to _utils
---
 ...tured_output.py => _agent_output_guard.py} | 17 ++++++++--
 .../adapters/langgraph/_langgraphconverter.py |  2 +-
 ...d_output.py => test_agent_output_guard.py} | 33 ++++++-------------
 pyagentspec/tests/adapters/test_utils.py      | 19 +++++++++++
 4 files changed, 44 insertions(+), 27 deletions(-)
 rename pyagentspec/src/pyagentspec/adapters/langgraph/{_structured_output.py => _agent_output_guard.py} (81%)
 rename pyagentspec/tests/adapters/langgraph/{test_structured_output.py => test_agent_output_guard.py} (82%)
 create mode 100644 pyagentspec/tests/adapters/test_utils.py

diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_structured_output.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
similarity index 81%
rename from pyagentspec/src/pyagentspec/adapters/langgraph/_structured_output.py
rename to pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
index 9f4fbe43..0cdc18b2 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_structured_output.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
@@ -31,8 +31,9 @@
 # bound separates "model is working on it" from "model will never comply".
 DEFAULT_MAX_STRUCTURED_OUTPUT_ATTEMPTS = 3
 
-# Must match the field on _StructuredOutputAttemptState.
-_ATTEMPTS_STATE_KEY = "structured_output_attempts"
+# Must match the field on _StructuredOutputAttemptState. Dunder-wrapped so it cannot
+# collide with a field the agent's own state schema declares.
+_ATTEMPTS_STATE_KEY = "__structured_output_attempts__"
 
 
 class StructuredOutputNotProducedError(RuntimeError):
@@ -45,7 +46,7 @@ class StructuredOutputNotProducedError(RuntimeError):
 class _StructuredOutputAttemptState(AgentState):
     # On the state rather than the instance, so the count is per run and not shared
     # between concurrent ones.
-    structured_output_attempts: NotRequired[int]
+    __structured_output_attempts__: NotRequired[int]
 
 
 class StructuredOutputGuard(AgentMiddleware):
@@ -82,6 +83,16 @@ def _fail(self, detail: str) -> NoReturn:
 
     # No async variants: LangGraph routes async runs to the sync hooks, and both are pure.
     def after_model(self, state: Any, runtime: Any) -> Optional[Dict[str, int]]:
+        # Runs after every model turn. The last message decides between three cases:
+        #
+        # * an AIMessage with tool calls: progress, counter resets;
+        # * an AIMessage without tool calls: prose instead of the structured response,
+        #   counter increments and the run fails at the bound;
+        # * anything else: in a plain run the model's AIMessage is always last here, so
+        #   this means another middleware rewrote the history and the turn cannot be
+        #   classified. The counter is left untouched (None updates nothing) rather than
+        #   reset: unclassifiable is not evidence of progress, and resetting would let
+        #   such a middleware clear real evidence of a stuck model every turn.
         messages = state.get("messages") or []
         last = messages[-1] if messages else None
         if last is None or getattr(last, "type", None) != "ai":
diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
index e1cbc11b..fd087166 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
@@ -51,7 +51,7 @@
     NodeExecutor,
     extract_outputs_from_invoke_result,
 )
-from pyagentspec.adapters.langgraph._structured_output import (
+from pyagentspec.adapters.langgraph._agent_output_guard import (
     StructuredOutputGuard,
 )
 from pyagentspec.adapters.langgraph._types import (
diff --git a/pyagentspec/tests/adapters/langgraph/test_structured_output.py b/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
similarity index 82%
rename from pyagentspec/tests/adapters/langgraph/test_structured_output.py
rename to pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
index eb6b2bdf..7236f885 100644
--- a/pyagentspec/tests/adapters/langgraph/test_structured_output.py
+++ b/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
@@ -4,7 +4,7 @@
 # (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License
 # (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option.
 
-"""How the LangGraph adapter decides on, and bounds, agent structured output.
+"""How :class:`StructuredOutputGuard` bounds an agent that declares structured output.
 
 These build no LLM config, so unlike the AgentNode flow tests they run rather than skip
 under ``SKIP_LLM_TESTS=1``.
@@ -15,25 +15,12 @@
 import pytest
 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 
-from pyagentspec.adapters._utils import is_single_string_output
-from pyagentspec.adapters.langgraph._structured_output import (
+from pyagentspec.adapters.langgraph._agent_output_guard import (
     StructuredOutputGuard,
     StructuredOutputNotProducedError,
 )
-from pyagentspec.property import IntegerProperty, StringProperty
 
-# ─── is_single_string_output ──────────────────────────────────────────────────
-
-
-def test_is_single_string_output() -> None:
-    """A lone string output is free text, not a structured field."""
-    assert is_single_string_output([StringProperty(title="x")]) is True
-    assert is_single_string_output([]) is False
-    assert is_single_string_output([IntegerProperty(title="n")]) is False
-    assert is_single_string_output([StringProperty(title="a"), StringProperty(title="b")]) is False
-
-
-# ─── structured-output guard: after_model budget ───────────────────────────────────────
+# ─── after_model budget ───────────────────────────────────────────────────────
 
 
 def _guard(max_attempts: int = 3) -> StructuredOutputGuard:
@@ -49,14 +36,14 @@ def _prose(attempts: int = 0) -> dict:
     """A model turn that answered in prose, so no structured response."""
     return {
         "messages": [HumanMessage(content="q"), AIMessage(content="42")],
-        "structured_output_attempts": attempts,
+        "__structured_output_attempts__": attempts,
     }
 
 
 def _tool_turn(attempts: int = 0) -> dict:
     return {
         "messages": [AIMessage(content="", tool_calls=[{"name": "s", "args": {}, "id": "c1"}])],
-        "structured_output_attempts": attempts,
+        "__structured_output_attempts__": attempts,
     }
 
 
@@ -64,8 +51,8 @@ def test_prose_turn_increments_until_the_limit() -> None:
     """Prose never satisfies response_format, so the agent cannot exit on its own."""
     guard = _guard(max_attempts=3)
 
-    assert guard.after_model(_prose(0), runtime=None) == {"structured_output_attempts": 1}
-    assert guard.after_model(_prose(1), runtime=None) == {"structured_output_attempts": 2}
+    assert guard.after_model(_prose(0), runtime=None) == {"__structured_output_attempts__": 1}
+    assert guard.after_model(_prose(1), runtime=None) == {"__structured_output_attempts__": 2}
     with pytest.raises(StructuredOutputNotProducedError):
         guard.after_model(_prose(2), runtime=None)
 
@@ -90,7 +77,7 @@ def test_tool_calling_turn_resets_rather_than_counting() -> None:
     guard = _guard(max_attempts=2)
     # max_attempts=2 with 1 attempt already banked would raise if this counted.
     assert guard.after_model(_tool_turn(attempts=1), runtime=None) == {
-        "structured_output_attempts": 0
+        "__structured_output_attempts__": 0
     }
 
 
@@ -198,5 +185,5 @@ def test_limiter_state_schema_declares_the_counter() -> None:
     """The counter belongs on the state schema, not the instance, so concurrent runs do
     not share it."""
     schema: Any = StructuredOutputGuard.state_schema
-    assert "structured_output_attempts" in schema.__annotations__
-    assert "structured_output_attempts" in schema.__optional_keys__
+    assert "__structured_output_attempts__" in schema.__annotations__
+    assert "__structured_output_attempts__" in schema.__optional_keys__
diff --git a/pyagentspec/tests/adapters/test_utils.py b/pyagentspec/tests/adapters/test_utils.py
new file mode 100644
index 00000000..ab8c86a8
--- /dev/null
+++ b/pyagentspec/tests/adapters/test_utils.py
@@ -0,0 +1,19 @@
+# Copyright © 2025, 2026 Oracle and/or its affiliates.
+#
+# This software is under the Apache License 2.0
+# (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License
+# (UPL) 1.0 (LICENSE-UPL or https://oss.oracle.com/licenses/upl), at your option.
+
+"""Tests for the shared adapter helpers in ``pyagentspec.adapters._utils``."""
+
+from pyagentspec.adapters._utils import is_single_string_output
+from pyagentspec.property import IntegerProperty, StringProperty
+
+
+def test_is_single_string_output() -> None:
+    """A lone string output is free text, not a structured field, so adapters read it
+    from the final message instead of requesting structured generation."""
+    assert is_single_string_output([StringProperty(title="x")]) is True
+    assert is_single_string_output([]) is False
+    assert is_single_string_output([IntegerProperty(title="n")]) is False
+    assert is_single_string_output([StringProperty(title="a"), StringProperty(title="b")]) is False

From 38dc8e8ec0b6f5210ad108ce9478be1a57b78edc Mon Sep 17 00:00:00 2001
From: Cesare Bernardis <cesare.bernardis@oracle.com>
Date: Thu, 3 Sep 2026 15:18:11 +0200
Subject: [PATCH 3/5] Updates

---
 .../adapters/langgraph/_agent_output_guard.py | 34 +++++++++++++------
 .../adapters/langgraph/_langgraphconverter.py | 26 ++++----------
 .../pyagentspec/adapters/langgraph/_types.py  |  3 ++
 .../langgraph/test_agent_output_guard.py      | 25 +++++++++-----
 4 files changed, 48 insertions(+), 40 deletions(-)

diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
index 0cdc18b2..4fcab795 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
@@ -21,19 +21,23 @@
 :func:`pyagentspec.adapters._utils.is_single_string_output`).
 """
 
-from typing import Any, Dict, List, NoReturn, Optional
+from functools import lru_cache
+from typing import Annotated, Any, Dict, List, NoReturn, Optional
 
 from typing_extensions import NotRequired
 
-from pyagentspec.adapters.langgraph._types import AgentMiddleware, AgentState
+from pyagentspec.adapters.langgraph._types import (
+    AgentMiddleware,
+    AgentState,
+    langchain_middleware_types,
+)
 
 # A compliant run spends at most one turn without a structured response, so a small
 # bound separates "model is working on it" from "model will never comply".
 DEFAULT_MAX_STRUCTURED_OUTPUT_ATTEMPTS = 3
 
-# Must match the field on _StructuredOutputAttemptState. Dunder-wrapped so it cannot
-# collide with a field the agent's own state schema declares.
-_ATTEMPTS_STATE_KEY = "__structured_output_attempts__"
+# Must match the private field on _StructuredOutputAttemptState.
+_ATTEMPTS_STATE_KEY = "_pyagentspec_structured_output_attempts"
 
 
 class StructuredOutputNotProducedError(RuntimeError):
@@ -43,10 +47,19 @@ class StructuredOutputNotProducedError(RuntimeError):
     """
 
 
-class _StructuredOutputAttemptState(AgentState):
-    # On the state rather than the instance, so the count is per run and not shared
-    # between concurrent ones.
-    __structured_output_attempts__: NotRequired[int]
+@lru_cache
+def _structured_output_attempt_state_schema() -> type:
+    """Create the LangChain state extension when its optional dependency is available."""
+
+    class _StructuredOutputAttemptState(AgentState):
+        # On the state rather than the instance, so the count is per run and not shared
+        # between concurrent ones. LangChain omits PrivateStateAttr fields from user input
+        # and output schemas.
+        _pyagentspec_structured_output_attempts: NotRequired[
+            Annotated[int, langchain_middleware_types.PrivateStateAttr]
+        ]
+
+    return _StructuredOutputAttemptState
 
 
 class StructuredOutputGuard(AgentMiddleware):
@@ -56,8 +69,6 @@ class StructuredOutputGuard(AgentMiddleware):
     just means the agent is done and an empty ``structured_response`` is fine.
     """
 
-    state_schema = _StructuredOutputAttemptState
-
     def __init__(
         self,
         *,
@@ -67,6 +78,7 @@ def __init__(
         max_attempts: int = DEFAULT_MAX_STRUCTURED_OUTPUT_ATTEMPTS,
     ) -> None:
         super().__init__()
+        self.state_schema = _structured_output_attempt_state_schema()
         self.agent_name = agent_name
         self.output_titles = output_titles
         self.model_id = model_id
diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
index fd087166..f4ae9a98 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_langgraphconverter.py
@@ -37,6 +37,7 @@
     create_pydantic_model_from_properties,
     is_single_string_output,
 )
+from pyagentspec.adapters.langgraph._agent_output_guard import StructuredOutputGuard
 from pyagentspec.adapters.langgraph._execution_span import patch_with_execution_span
 from pyagentspec.adapters.langgraph._managerworkers import (
     _MANAGER_NODE_KEY,
@@ -51,9 +52,6 @@
     NodeExecutor,
     extract_outputs_from_invoke_result,
 )
-from pyagentspec.adapters.langgraph._agent_output_guard import (
-    StructuredOutputGuard,
-)
 from pyagentspec.adapters.langgraph._types import (
     AgentState,
     BaseCallbackHandler,
@@ -74,10 +72,7 @@
     langgraph_graph,
     langgraph_swarm,
 )
-from pyagentspec.adapters.langgraph.mcp_utils import (
-    _HttpxClientFactory,
-    run_async_in_sync,
-)
+from pyagentspec.adapters.langgraph.mcp_utils import _HttpxClientFactory, run_async_in_sync
 from pyagentspec.adapters.langgraph.tracing import (
     AgentSpecLlmCallbackHandler,
     AgentSpecToolCallbackHandler,
@@ -111,10 +106,7 @@
 )
 from pyagentspec.llms.ocigenaiconfig import OciGenAiConfig
 from pyagentspec.llms.ollamaconfig import OllamaConfig
-from pyagentspec.llms.openaicompatibleconfig import (
-    OpenAIAPIType,
-    OpenAiCompatibleConfig,
-)
+from pyagentspec.llms.openaicompatibleconfig import OpenAIAPIType, OpenAiCompatibleConfig
 from pyagentspec.llms.openaiconfig import OpenAiConfig
 from pyagentspec.llms.vllmconfig import VllmConfig
 from pyagentspec.managerworkers import ManagerWorkers as AgentSpecManagerWorkers
@@ -595,9 +587,7 @@ def _input_message_node_convert_to_langgraph(
         self,
         node: AgentSpecInputMessageNode,
     ) -> "NodeExecutor":
-        from pyagentspec.adapters.langgraph._node_execution import (
-            InputMessageNodeExecutor,
-        )
+        from pyagentspec.adapters.langgraph._node_execution import InputMessageNodeExecutor
 
         return InputMessageNodeExecutor(node)
 
@@ -605,9 +595,7 @@ def _output_message_node_convert_to_langgraph(
         self,
         node: AgentSpecOutputMessageNode,
     ) -> "NodeExecutor":
-        from pyagentspec.adapters.langgraph._node_execution import (
-            OutputMessageNodeExecutor,
-        )
+        from pyagentspec.adapters.langgraph._node_execution import OutputMessageNodeExecutor
 
         return OutputMessageNodeExecutor(node)
 
@@ -672,9 +660,7 @@ def _catch_exception_node_convert_to_langgraph(
         config: RunnableConfig,
         middleware: List[Any],
     ) -> "NodeExecutor":
-        from pyagentspec.adapters.langgraph._node_execution import (
-            CatchExceptionNodeExecutor,
-        )
+        from pyagentspec.adapters.langgraph._node_execution import CatchExceptionNodeExecutor
 
         subflow = self.convert(
             catch_node.subflow,
diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py
index b48afc61..c538dfb9 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_types.py
@@ -15,6 +15,7 @@
     # Otherwise, importing the module when they are not installed would lead to an import error.
 
     import langchain.agents as langchain_agents
+    import langchain.agents.middleware.types as langchain_middleware_types
     import langchain_ollama
     import langchain_openai
     import langgraph.graph as langgraph_graph
@@ -39,6 +40,7 @@
     langchain_openai = LazyLoader("langchain_openai")
     langgraph_graph = LazyLoader("langgraph.graph")
     langchain_agents = LazyLoader("langchain.agents")
+    langchain_middleware_types = LazyLoader("langchain.agents.middleware.types")
     BaseTool = LazyType("langchain_core.tools", "BaseTool")
     StructuredTool = LazyType("langchain_core.tools", "StructuredTool")
     Checkpointer = LazyType("langgraph.types", "Checkpointer")
@@ -109,6 +111,7 @@ class FlowOutputSchema(TypedDict):
 __all__ = [
     "langgraph_graph",
     "langchain_agents",
+    "langchain_middleware_types",
     "langchain_ollama",
     "langchain_openai",
     "LangGraphTool",
diff --git a/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py b/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
index 7236f885..a2765b23 100644
--- a/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
+++ b/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
@@ -10,9 +10,10 @@
 under ``SKIP_LLM_TESTS=1``.
 """
 
-from typing import Any
+from typing import Any, get_args
 
 import pytest
+from langchain.agents.middleware.types import PrivateStateAttr
 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 
 from pyagentspec.adapters.langgraph._agent_output_guard import (
@@ -36,14 +37,14 @@ def _prose(attempts: int = 0) -> dict:
     """A model turn that answered in prose, so no structured response."""
     return {
         "messages": [HumanMessage(content="q"), AIMessage(content="42")],
-        "__structured_output_attempts__": attempts,
+        "_pyagentspec_structured_output_attempts": attempts,
     }
 
 
 def _tool_turn(attempts: int = 0) -> dict:
     return {
         "messages": [AIMessage(content="", tool_calls=[{"name": "s", "args": {}, "id": "c1"}])],
-        "__structured_output_attempts__": attempts,
+        "_pyagentspec_structured_output_attempts": attempts,
     }
 
 
@@ -51,8 +52,12 @@ def test_prose_turn_increments_until_the_limit() -> None:
     """Prose never satisfies response_format, so the agent cannot exit on its own."""
     guard = _guard(max_attempts=3)
 
-    assert guard.after_model(_prose(0), runtime=None) == {"__structured_output_attempts__": 1}
-    assert guard.after_model(_prose(1), runtime=None) == {"__structured_output_attempts__": 2}
+    assert guard.after_model(_prose(0), runtime=None) == {
+        "_pyagentspec_structured_output_attempts": 1
+    }
+    assert guard.after_model(_prose(1), runtime=None) == {
+        "_pyagentspec_structured_output_attempts": 2
+    }
     with pytest.raises(StructuredOutputNotProducedError):
         guard.after_model(_prose(2), runtime=None)
 
@@ -77,7 +82,7 @@ def test_tool_calling_turn_resets_rather_than_counting() -> None:
     guard = _guard(max_attempts=2)
     # max_attempts=2 with 1 attempt already banked would raise if this counted.
     assert guard.after_model(_tool_turn(attempts=1), runtime=None) == {
-        "__structured_output_attempts__": 0
+        "_pyagentspec_structured_output_attempts": 0
     }
 
 
@@ -184,6 +189,8 @@ def bind_tools(self, tools: Any, **kwargs: Any) -> Any:
 def test_limiter_state_schema_declares_the_counter() -> None:
     """The counter belongs on the state schema, not the instance, so concurrent runs do
     not share it."""
-    schema: Any = StructuredOutputGuard.state_schema
-    assert "__structured_output_attempts__" in schema.__annotations__
-    assert "__structured_output_attempts__" in schema.__optional_keys__
+    schema: Any = _guard().state_schema
+    key = "_pyagentspec_structured_output_attempts"
+    assert key in schema.__annotations__
+    assert key in schema.__optional_keys__
+    assert PrivateStateAttr in get_args(get_args(schema.__annotations__[key])[0])

From 660a645979bc85e50e491d2104068f92ee2729c6 Mon Sep 17 00:00:00 2001
From: Cesare Bernardis <cesare.bernardis@oracle.com>
Date: Thu, 3 Sep 2026 15:24:00 +0200
Subject: [PATCH 4/5] Changelog

---
 docs/pyagentspec/source/changelog.rst | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/docs/pyagentspec/source/changelog.rst b/docs/pyagentspec/source/changelog.rst
index c9ceb2b5..f6008515 100644
--- a/docs/pyagentspec/source/changelog.rst
+++ b/docs/pyagentspec/source/changelog.rst
@@ -17,6 +17,14 @@ Improvements
 Bug fixes
 ^^^^^^^^^
 
+* **LangGraph structured-output guard**
+
+  LangGraph agents with declared structured outputs now fail promptly with a clear error when
+  their model repeatedly returns prose or finishes without producing the declared output, rather
+  than looping until the recursion limit or silently returning empty outputs.
+
+  We thank @spichen for the contribution!
+
 Breaking Changes
 ^^^^^^^^^^^^^^^^
 

From 24610c255fe3f9dc113a26f9b029d22b7d6f2cdd Mon Sep 17 00:00:00 2001
From: Cesare Bernardis <cesare.bernardis@oracle.com>
Date: Thu, 3 Sep 2026 15:55:34 +0200
Subject: [PATCH 5/5] Fix test imports

---
 .../adapters/langgraph/_agent_output_guard.py | 21 +++++++++++--------
 .../langgraph/test_agent_output_guard.py      | 14 +++++++++++--
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py b/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
index 4fcab795..30a8b968 100644
--- a/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
+++ b/pyagentspec/src/pyagentspec/adapters/langgraph/_agent_output_guard.py
@@ -47,17 +47,20 @@ class StructuredOutputNotProducedError(RuntimeError):
     """
 
 
+class _StructuredOutputAttemptState(AgentState):
+    # On the state rather than the instance, so the count is per run and not shared
+    # between concurrent ones. LangChain omits PrivateStateAttr fields from user input
+    # and output schemas.
+    _pyagentspec_structured_output_attempts: NotRequired[int]
+
+
 @lru_cache
 def _structured_output_attempt_state_schema() -> type:
-    """Create the LangChain state extension when its optional dependency is available."""
-
-    class _StructuredOutputAttemptState(AgentState):
-        # On the state rather than the instance, so the count is per run and not shared
-        # between concurrent ones. LangChain omits PrivateStateAttr fields from user input
-        # and output schemas.
-        _pyagentspec_structured_output_attempts: NotRequired[
-            Annotated[int, langchain_middleware_types.PrivateStateAttr]
-        ]
+    """Add LangChain's private-state marker when the guard is used."""
+
+    _StructuredOutputAttemptState.__annotations__[_ATTEMPTS_STATE_KEY] = NotRequired[
+        Annotated[int, langchain_middleware_types.PrivateStateAttr]
+    ]
 
     return _StructuredOutputAttemptState
 
diff --git a/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py b/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
index a2765b23..b03ad82e 100644
--- a/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
+++ b/pyagentspec/tests/adapters/langgraph/test_agent_output_guard.py
@@ -13,8 +13,6 @@
 from typing import Any, get_args
 
 import pytest
-from langchain.agents.middleware.types import PrivateStateAttr
-from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 
 from pyagentspec.adapters.langgraph._agent_output_guard import (
     StructuredOutputGuard,
@@ -35,6 +33,8 @@ def _guard(max_attempts: int = 3) -> StructuredOutputGuard:
 
 def _prose(attempts: int = 0) -> dict:
     """A model turn that answered in prose, so no structured response."""
+    from langchain_core.messages import AIMessage, HumanMessage
+
     return {
         "messages": [HumanMessage(content="q"), AIMessage(content="42")],
         "_pyagentspec_structured_output_attempts": attempts,
@@ -42,6 +42,8 @@ def _prose(attempts: int = 0) -> dict:
 
 
 def _tool_turn(attempts: int = 0) -> dict:
+    from langchain_core.messages import AIMessage
+
     return {
         "messages": [AIMessage(content="", tool_calls=[{"name": "s", "args": {}, "id": "c1"}])],
         "_pyagentspec_structured_output_attempts": attempts,
@@ -88,6 +90,8 @@ def test_tool_calling_turn_resets_rather_than_counting() -> None:
 
 def test_non_ai_or_empty_last_message_is_ignored() -> None:
     """Only a model turn can be a failed structured response."""
+    from langchain_core.messages import ToolMessage
+
     guard = _guard(max_attempts=1)
     assert guard.after_model({"messages": []}, runtime=None) is None
     assert guard.after_model({}, runtime=None) is None
@@ -103,6 +107,7 @@ async def test_budget_applies_on_async_runs_via_the_sync_hook() -> None:
     sync hook. Agents do run async, so a silent no-op there would bring the hang back."""
     from langchain.agents import create_agent
     from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+    from langchain_core.messages import AIMessage, HumanMessage
     from pydantic import BaseModel
 
     class _Answer(BaseModel):
@@ -141,6 +146,8 @@ def test_after_agent_accepts_a_run_that_produced_the_response() -> None:
 def test_after_agent_rejects_a_run_that_produced_nothing() -> None:
     """An agent with tools exits after one prose turn, spending no budget, so this hook is
     the only place that failure shows up."""
+    from langchain_core.messages import AIMessage
+
     with pytest.raises(StructuredOutputNotProducedError) as excinfo:
         _guard().after_agent({"messages": [AIMessage(content="42")]}, runtime=None)
 
@@ -157,6 +164,7 @@ def test_agent_with_tools_raises_instead_of_losing_outputs_silently() -> None:
     ``structured_response`` and the declared outputs came back empty."""
     from langchain.agents import create_agent
     from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+    from langchain_core.messages import AIMessage, HumanMessage
     from langchain_core.tools import tool
     from pydantic import BaseModel
 
@@ -189,6 +197,8 @@ def bind_tools(self, tools: Any, **kwargs: Any) -> Any:
 def test_limiter_state_schema_declares_the_counter() -> None:
     """The counter belongs on the state schema, not the instance, so concurrent runs do
     not share it."""
+    from langchain.agents.middleware.types import PrivateStateAttr
+
     schema: Any = _guard().state_schema
     key = "_pyagentspec_structured_output_attempts"
     assert key in schema.__annotations__