From 13f7d9dfc4f055167b680c2c2c324fdd25a6111b Mon Sep 17 00:00:00 2001 From: Mark Fogle Date: Thu, 23 Jul 2026 02:46:22 +0000 Subject: [PATCH] fix(langgraph): fall back A2UI tool_choice for reasoning/thinking-mode models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The A2UI render subagent binds the inner `render_a2ui` tool with a forced named `tool_choice` so the subagent is guaranteed to emit a structured surface. Reasoning / "thinking-mode" models on some OpenAI-compatible providers reject a forced/object `tool_choice` and 400 the request (observed on Alibaba DashScope qwen* with thinking enabled): The tool_choice parameter does not support being set to required or object in thinking mode That aborts the render mid-stream — the surface never paints and the client surfaces a `terminated` run error. Two changes, both fully backward compatible: 1. Runtime fallback (langgraph adapter): if binding the forced choice fails the request with that specific 400, downgrade to `tool_choice="auto"` and retry — once, and for the rest of the run so subsequent recovery attempts don't re-trip it. Providers that honor a named `tool_choice` (OpenAI, etc.) hit the success path on the first try and never enter the fallback, so their behavior is unchanged. The error match is narrow (both "tool_choice" and "thinking mode" must appear) so unrelated 400s propagate untouched, and the existing validate/retry loop still gates the result — a model that skips the tool under "auto" degrades gracefully instead of crashing. 2. `tool_choice` knob (shared toolkit `A2UIToolParams`): lets a caller pin the choice up front (e.g. "auto" for a known reasoning model). Defaults to the forced render tool, so existing callers are unaffected; every framework adapter gets the knob for free. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../python/ag_ui_langgraph/a2ui_tool.py | 55 ++++++++- .../langgraph/python/tests/test_a2ui_tool.py | 105 +++++++++++++++++- .../ag_ui_a2ui_toolkit/__init__.py | 13 +++ 3 files changed, 165 insertions(+), 8 deletions(-) diff --git a/integrations/langgraph/python/ag_ui_langgraph/a2ui_tool.py b/integrations/langgraph/python/ag_ui_langgraph/a2ui_tool.py index e6a38d911f..ca4f51c5f3 100644 --- a/integrations/langgraph/python/ag_ui_langgraph/a2ui_tool.py +++ b/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( diff --git a/integrations/langgraph/python/tests/test_a2ui_tool.py b/integrations/langgraph/python/tests/test_a2ui_tool.py index 3b2787b862..d3a9551d83 100644 --- a/integrations/langgraph/python/tests/test_a2ui_tool.py +++ b/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() diff --git a/sdks/python/a2ui_toolkit/ag_ui_a2ui_toolkit/__init__.py b/sdks/python/a2ui_toolkit/ag_ui_a2ui_toolkit/__init__.py index e369928ac3..2cf93aba74 100644 --- a/sdks/python/a2ui_toolkit/ag_ui_a2ui_toolkit/__init__.py +++ b/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"], }