fix(langgraph): fall back A2UI tool_choice for reasoning/thinking-mode models

integrations/langgraph/python/ag_ui_langgraph/a2ui_tool.py

@@ -59,6 +59,26 @@
 RENDER_A2UI_TOOL_NAME: str = RENDER_A2UI_TOOL_DEF["function"]["name"]
 
 
+def _is_thinking_mode_tool_choice_error(exc: BaseException) -> bool:
+    """True when a provider rejected a forced/named ``tool_choice`` because the
+    model is in reasoning/"thinking" mode.
+
+    Reasoning models on some OpenAI-compatible providers (observed: Alibaba
+    DashScope ``qwen*`` with thinking enabled) forbid a ``required``/object
+    ``tool_choice`` and fail the request with a 400 such as::
+
+        The tool_choice parameter does not support being set to required or
+        object in thinking mode
+
+    We match on the error text — not the exception class — so this adapter
+    doesn't hard-depend on the ``openai`` SDK's error types, and the match stays
+    narrow (both ``tool_choice`` and ``thinking mode`` must appear) so unrelated
+    400s are re-raised untouched.
+    """
+    msg = str(exc).lower()
+    return "tool_choice" in msg and "thinking mode" in msg
+
+
 # Re-export the toolkit constants/types for callers that previously imported
 # them from this package — keeps the public surface stable and lets consumers
 # type the shared params object + its guidelines without depending on the
@@ -136,6 +156,11 @@ def get_a2ui_tools(params: A2UIToolParams):
     catalog = cfg["catalog"]
     recovery = cfg["recovery"]
     on_a2ui_attempt = cfg["on_a2ui_attempt"]
+    # Default: force the inner render tool by name so the subagent is guaranteed
+    # to emit a structured surface. Reasoning/"thinking-mode" models reject a
+    # forced choice; the runtime fallback below downgrades to "auto"
+    # for those, and a caller may pin ``tool_choice="auto"`` up front via params.
+    tool_choice = cfg["tool_choice"]
 
     @tool(cfg["tool_name"], description=cfg["tool_description"])
     async def generate_a2ui(
@@ -171,13 +196,33 @@ async def generate_a2ui(
         if prep.get("error"):
             return wrap_error_envelope(prep["error"])
 
-        # Glue: bind the structured-output tool.
-        model_with_tool = model.bind_tools(
-            [RENDER_A2UI_TOOL_DEF], tool_choice="render_a2ui"
-        )
+        # Glue: bind the structured-output tool. ``choice_state`` is mutable so a
+        # thinking-mode fallback (below) sticks for every subsequent recovery
+        # attempt instead of re-tripping — and re-recovering from — the same 400.
+        choice_state = {"tool_choice": tool_choice}
+
+        def _bind(choice):
+            return model.bind_tools([RENDER_A2UI_TOOL_DEF], tool_choice=choice)
 
         async def _invoke_subagent(prompt, _attempt):
-            return await _stream_render_subagent(model_with_tool, prompt, messages)
+            choice = choice_state["tool_choice"]
+            try:
+                return await _stream_render_subagent(_bind(choice), prompt, messages)
+            except Exception as exc:  # noqa: BLE001 - re-raised unless it's the known 400
+                # Reasoning models reject a forced/object tool_choice. Fall back
+                # to "auto" once (and for the rest of this run) so the render can
+                # proceed; the validate->retry loop below still gates the result,
+                # so a model that skips the tool degrades gracefully rather than
+                # crashing the stream.
+                if choice != "auto" and _is_thinking_mode_tool_choice_error(exc):
+                    logger.warning(
+                        "A2UI: provider rejected forced tool_choice=%r in thinking "
+                        "mode; falling back to tool_choice='auto' for this run.",
+                        choice,
+                    )
+                    choice_state["tool_choice"] = "auto"
+                    return await _stream_render_subagent(_bind("auto"), prompt, messages)
+                raise
 
         def _build_envelope(args):
             return build_a2ui_envelope(

integrations/langgraph/python/tests/test_a2ui_tool.py

@@ -26,7 +26,10 @@
 from langchain_core.messages.tool import tool_call_chunk
 
 from ag_ui_langgraph import get_a2ui_tools
-from ag_ui_langgraph.a2ui_tool import _stream_render_subagent
+from ag_ui_langgraph.a2ui_tool import (
+    _is_thinking_mode_tool_choice_error,
+    _stream_render_subagent,
+)
 from ag_ui_a2ui_toolkit import (
     A2UI_OPERATIONS_KEY,
     DEFAULT_DESIGN_GUIDELINES,
@@ -91,11 +94,39 @@ class FakeModel:
     def __init__(self, args):
         self.args = args
         self.captured_prompts: list[str] = []
+        self.bound_tool_choices: list = []
 
     def bind_tools(self, tools, tool_choice=None):
+        self.bound_tool_choices.append(tool_choice)
         return _StreamingBoundModel(self)
 
 
+class _ThinkingModeModel(FakeModel):
+    """A provider whose thinking mode rejects a forced/named ``tool_choice`` with
+    the DashScope 400, but streams normally under ``tool_choice="auto"`` — the
+    exact reasoning/thinking-mode shape."""
+
+    def bind_tools(self, tools, tool_choice=None):
+        self.bound_tool_choices.append(tool_choice)
+        if tool_choice not in (None, "auto"):
+            return _RejectingBoundModel()
+        return _StreamingBoundModel(self)
+
+
+class _RejectingBoundModel:
+    """A bound model whose stream raises the DashScope thinking-mode 400 (raised
+    when the stream is iterated, as a real provider does)."""
+
+    async def astream(self, messages):
+        if False:  # pragma: no cover - make this an async generator
+            yield None
+        raise RuntimeError(
+            "Error code: 400 - InternalError.Algo.InvalidParameter: The "
+            "tool_choice parameter does not support being set to required or "
+            "object in thinking mode"
+        )
+
+
 class FakeRuntime:
     """Stand-in for LangGraph's ``ToolRuntime`` — the tool reads ``state`` and
     ``config`` (the latter forwarded to ``adispatch_custom_event``)."""
@@ -112,13 +143,15 @@ def _invoke_tool(tool, runtime, **kwargs) -> str:
 
 
 class TestGetA2UITools(unittest.TestCase):
-    def _make(self, guidelines=None, tool_name=None):
-        model = FakeModel(VALID_ARGS)
+    def _make(self, guidelines=None, tool_name=None, tool_choice=None, model=None):
+        model = model if model is not None else FakeModel(VALID_ARGS)
         params = {"model": model, "default_catalog_id": "cat://custom"}
         if guidelines is not None:
             params["guidelines"] = guidelines
         if tool_name is not None:
             params["tool_name"] = tool_name
+        if tool_choice is not None:
+            params["tool_choice"] = tool_choice
         return get_a2ui_tools(params), model
 
     def test_single_arg_params_produces_operations_envelope(self):
@@ -167,6 +200,54 @@ def test_tool_name_resolves(self):
         custom_tool, _ = self._make(tool_name="render_ui")
         self.assertEqual(custom_tool.name, "render_ui")
 
+    def test_default_tool_choice_forces_render(self):
+        # By default the subagent forces the inner render tool by name so it is
+        # guaranteed to emit a structured surface.
+        tool, model = self._make()
+        _invoke_tool(tool, FakeRuntime({"messages": []}), intent="create")
+        self.assertEqual(model.bound_tool_choices, ["render_a2ui"])
+
+    def test_tool_choice_override_reaches_bind(self):
+        # A caller can pin tool_choice up front (explicit escape hatch)
+        # — e.g. "auto" for a reasoning model that rejects a forced choice.
+        tool, model = self._make(tool_choice="auto")
+        _invoke_tool(tool, FakeRuntime({"messages": []}), intent="create")
+        self.assertEqual(model.bound_tool_choices, ["auto"])
+
+    def test_thinking_mode_falls_back_to_auto(self):
+        # A forced tool_choice 400s in a reasoning model's thinking mode. The adapter must
+        # catch that specific error, downgrade to "auto", and still produce a
+        # valid operations envelope — no crash, no caller change.
+        model = _ThinkingModeModel(VALID_ARGS)
+        tool, _ = self._make(model=model)
+        envelope = _invoke_tool(tool, FakeRuntime({"messages": []}), intent="create")
+        parsed = json.loads(envelope)
+        self.assertTrue(any("createSurface" in o for o in parsed[A2UI_OPERATIONS_KEY]))
+        # Forced choice tried first, then the "auto" fallback.
+        self.assertEqual(model.bound_tool_choices, ["render_a2ui", "auto"])
+
+    def test_unrelated_bad_request_is_not_swallowed(self):
+        # The fallback is narrow: a 400 that isn't the thinking-mode tool_choice
+        # rejection must propagate, not silently downgrade.
+        class _OtherError(FakeModel):
+            def bind_tools(self, tools, tool_choice=None):
+                self.bound_tool_choices.append(tool_choice)
+
+                class _Boom:
+                    async def astream(self, messages):
+                        if False:  # pragma: no cover
+                            yield None
+                        raise RuntimeError("400 - context length exceeded")
+
+                return _Boom()
+
+        model = _OtherError(VALID_ARGS)
+        tool, _ = self._make(model=model)
+        with self.assertRaises(RuntimeError):
+            _invoke_tool(tool, FakeRuntime({"messages": []}), intent="create")
+        # Only the forced choice was attempted; no "auto" fallback.
+        self.assertEqual(model.bound_tool_choices, ["render_a2ui"])
+
 
 class TestStreamRenderSubagent(unittest.TestCase):
     """The subagent STREAMS the model (``astream``) so the nested render_a2ui
@@ -201,5 +282,23 @@ async def _empty_astream(_messages):
         self.assertIsNone(captured)
 
 
+class TestThinkingModeErrorPredicate(unittest.TestCase):
+    def test_matches_dashscope_thinking_mode_400(self):
+        exc = RuntimeError(
+            "Error code: 400 - InternalError.Algo.InvalidParameter: The "
+            "tool_choice parameter does not support being set to required or "
+            "object in thinking mode"
+        )
+        self.assertTrue(_is_thinking_mode_tool_choice_error(exc))
+
+    def test_ignores_unrelated_errors(self):
+        for exc in (
+            RuntimeError("400 - context length exceeded"),
+            RuntimeError("tool_choice invalid tool name"),  # tool_choice but not thinking
+            ValueError("rate limited"),
+        ):
+            self.assertFalse(_is_thinking_mode_tool_choice_error(exc))
+
+
 if __name__ == "__main__":
     unittest.main()

sdks/python/a2ui_toolkit/ag_ui_a2ui_toolkit/__init__.py

@@ -780,6 +780,14 @@ class A2UIToolParams(TypedDict, total=False):
     catalog: Optional[dict]
     recovery: Optional[dict]
     on_a2ui_attempt: Optional[Any]
+    # ``tool_choice`` the render subagent binds the inner ``render_a2ui`` tool
+    # with. Defaults to forcing that tool by name (``"render_a2ui"``) so the
+    # subagent is guaranteed to emit a structured surface. Override to
+    # ``"auto"`` for reasoning/"thinking-mode" models (e.g. Alibaba DashScope
+    # ``qwen*`` with thinking enabled) that reject a forced/object ``tool_choice``
+    # (reasoning-model thinking mode). Adapters that honor a runtime fallback only reach for this
+    # when a caller wants to pin the behavior explicitly.
+    tool_choice: Optional[Any]
 
 
 class ResolvedA2UIToolParams(TypedDict):
@@ -796,6 +804,7 @@ class ResolvedA2UIToolParams(TypedDict):
     catalog: Optional[dict]
     recovery: Optional[dict]
     on_a2ui_attempt: Optional[Any]
+    tool_choice: Any
 
 
 def resolve_a2ui_tool_params(params: A2UIToolParams) -> ResolvedA2UIToolParams:
@@ -818,6 +827,10 @@ def resolve_a2ui_tool_params(params: A2UIToolParams) -> ResolvedA2UIToolParams:
         "catalog": params.get("catalog"),
         "recovery": params.get("recovery"),
         "on_a2ui_attempt": params.get("on_a2ui_attempt"),
+        # Forced-render by default; ``or`` (not ``is None``) so an empty-string
+        # override falls back to the forced default rather than binding a blank
+        # (invalid) tool choice.
+        "tool_choice": params.get("tool_choice") or RENDER_A2UI_TOOL_DEF["function"]["name"],
     }