From a3bc4ad6b4c50e1b1c89674fd923f266178c1cfa Mon Sep 17 00:00:00 2001 From: SomSamantray <> Date: Tue, 11 Aug 2026 17:45:11 +0530 Subject: [PATCH] fix(core): set explicit max_tokens for Claude thinking models on Bedrock Bedrock Converse has no fixed thinking budget to derive maxTokens from, so make_model_settings never sent an explicit max_tokens for adaptive- thinking Claude models. Bedrock's low internal default then truncated long tool calls (e.g. create_vulnerability_report), crashing the scan on the next turn. Closes #623 --- strix/config/models.py | 13 +++++- strix/core/inputs.py | 50 ++++++++++++++++++++--- tests/test_inputs.py | 92 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 7 deletions(-) diff --git a/strix/config/models.py b/strix/config/models.py index e632bb06a..1d6d43ac0 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -749,7 +749,13 @@ def uses_chat_completions_tool_schema(model_name: str, settings: Settings) -> bo return not model_supports_reasoning(model_name) -def model_supports_reasoning(model_name: str) -> bool: +def model_cost_entry(model_name: str) -> dict[str, Any] | None: + """Look up a model's ``litellm.model_cost`` entry, tolerating route prefixes. + + Tries the full (route-prefix-stripped) name first, then falls back to the + bare model name after the last ``/`` for routes ``litellm.model_cost`` + doesn't key under their prefixed form (e.g. Bedrock). + """ import litellm name = model_name.strip().lower() @@ -760,6 +766,11 @@ def model_supports_reasoning(model_name: str) -> bool: entry = litellm.model_cost.get(name) if entry is None and "/" in name: entry = litellm.model_cost.get(name.rsplit("/", 1)[1]) + return entry + + +def model_supports_reasoning(model_name: str) -> bool: + entry = model_cost_entry(model_name) return bool(entry and entry.get("supports_reasoning")) diff --git a/strix/core/inputs.py b/strix/core/inputs.py index a1106ae2d..ed16ae87a 100644 --- a/strix/core/inputs.py +++ b/strix/core/inputs.py @@ -16,6 +16,7 @@ from strix.config.models import ( is_claude_model, is_known_openai_bare_model, is_openrouter_model, + model_cost_entry, model_supports_reasoning, request_timeout_extra_args, ) @@ -26,6 +27,13 @@ if TYPE_CHECKING: from strix.config.settings import ReasoningEffort +# Floor for the explicit ``max_tokens`` set on Claude-on-Bedrock requests with +# thinking enabled (see ``_bedrock_thinking_max_tokens``): large enough that a +# long tool call (e.g. ``create_vulnerability_report``) isn't truncated by +# Bedrock Converse's low internal default when no fixed thinking budget is set. +_BEDROCK_THINKING_MAX_TOKENS = 32000 + + def _accepts_required_tool_choice(model_name: str | None) -> bool: name = (model_name or "").strip().lower() for prefix in ("litellm/", "any-llm/"): @@ -206,6 +214,13 @@ def make_model_settings( has_tools: bool = True, ) -> ModelSettings: headers = _request_headers(model_name, extra_headers) + active_reasoning_effort = ( + reasoning_effort + if reasoning_effort is not None + and reasoning_effort != "none" + and model_supports_reasoning(model_name) + else None + ) model_settings = ModelSettings( parallel_tool_calls=False if has_tools else None, retry=DEFAULT_MODEL_RETRY, @@ -213,13 +228,16 @@ def make_model_settings( extra_args=request_timeout_extra_args(request_timeout), extra_headers=headers, ) - if ( - reasoning_effort is not None - and reasoning_effort != "none" - and model_supports_reasoning(model_name) - ): + if active_reasoning_effort is not None: model_settings = model_settings.resolve( - _reasoning_settings(reasoning_effort, model_settings.extra_args), + _reasoning_settings(active_reasoning_effort, model_settings.extra_args), + ) + bedrock_thinking_max_tokens = ( + _bedrock_thinking_max_tokens(model_name) if active_reasoning_effort is not None else None + ) + if bedrock_thinking_max_tokens is not None: + model_settings = model_settings.resolve( + ModelSettings(max_tokens=bedrock_thinking_max_tokens), ) if force_required_tool_choice and _accepts_required_tool_choice(model_name): model_settings = model_settings.resolve(ModelSettings(tool_choice="required")) @@ -261,6 +279,26 @@ def _reasoning_settings( ) +def _bedrock_thinking_max_tokens(model_name: str) -> int | None: + """Explicit ``max_tokens`` for a Claude-on-Bedrock request with thinking enabled. + + Bedrock Converse has no fixed thinking budget to derive ``maxTokens`` from + (see litellm's ``supports_adaptive_thinking``), so without an explicit value + it falls back to a low internal default that truncates long tool calls. + Clamped to the model's own ``max_output_tokens`` (via the same litellm + lookup ``model_supports_reasoning`` uses) so we never request more than the + model accepts; an unmapped model still gets the floor. + """ + if not (is_claude_model(model_name) and is_bedrock_route(model_name)): + return None + + entry = model_cost_entry(model_name) + ceiling = entry.get("max_output_tokens") if entry else None + if not isinstance(ceiling, int) or ceiling <= 0: + return _BEDROCK_THINKING_MAX_TOKENS + return min(_BEDROCK_THINKING_MAX_TOKENS, ceiling) + + def _prompt_cache_extra_args(model_name: str) -> dict[str, Any] | None: """LiteLLM ``cache_control_injection_points`` for Claude prompt caching. diff --git a/tests/test_inputs.py b/tests/test_inputs.py index 2ff9a6035..7b4b7c350 100644 --- a/tests/test_inputs.py +++ b/tests/test_inputs.py @@ -145,6 +145,98 @@ def test_max_reasoning_effort_sent_as_raw_body_field() -> None: assert settings.extra_args == {"timeout": 30, "extra_body": {"reasoning_effort": "max"}} +def test_bedrock_claude_thinking_gets_explicit_max_tokens() -> None: + # Bedrock Converse has no fixed thinking budget to derive maxTokens from, so + # without an explicit value it falls back to a low internal default that + # truncates long tool calls (issue #623). + settings = make_model_settings("high", model_name="bedrock/global.anthropic.claude-opus-4-8") + assert settings.max_tokens == 32000 + + +def test_bedrock_claude_thinking_max_tokens_clamped_to_model_ceiling( + monkeypatch: Any, +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + "global.anthropic.claude-opus-4-8": { + "max_output_tokens": 8000, + "supports_reasoning": True, + } + }, + raising=False, + ) + settings = make_model_settings("high", model_name="bedrock/global.anthropic.claude-opus-4-8") + assert settings.max_tokens == 8000 + + +def test_bedrock_claude_thinking_max_tokens_falls_back_when_ceiling_unknown( + monkeypatch: Any, +) -> None: + # A model litellm knows supports reasoning but has no max_output_tokens + # entry for (e.g. a brand-new custom alias) still gets the floor rather + # than no max_tokens at all. + monkeypatch.setattr( + litellm, + "model_cost", + {"global.anthropic.claude-brand-new-9": {"supports_reasoning": True}}, + raising=False, + ) + settings = make_model_settings("high", model_name="bedrock/global.anthropic.claude-brand-new-9") + assert settings.max_tokens == 32000 + + +def test_max_tokens_unset_when_reasoning_disabled_for_bedrock_claude() -> None: + settings = make_model_settings(None, model_name="bedrock/global.anthropic.claude-opus-4-8") + assert settings.max_tokens is None + + settings = make_model_settings("none", model_name="bedrock/global.anthropic.claude-opus-4-8") + assert settings.max_tokens is None + + +def test_max_tokens_unset_for_non_bedrock_claude_with_reasoning() -> None: + settings = make_model_settings("high", model_name="anthropic/claude-sonnet-4-5") + assert settings.max_tokens is None + + +@pytest.mark.parametrize("model_name", ["gpt-5", "vertex_ai/gemini-2.5-pro"]) +def test_max_tokens_unset_for_non_claude_with_reasoning(model_name: str) -> None: + assert make_model_settings("high", model_name=model_name).max_tokens is None + + +def test_max_tokens_unset_for_non_claude_bedrock_model_with_reasoning( + monkeypatch: Any, +) -> None: + # is_claude_model must independently gate the Bedrock helper, not just + # is_bedrock_route -- a non-Anthropic Bedrock model must stay unaffected. + monkeypatch.setattr( + litellm, + "model_cost", + {"global.meta.llama-5-70b": {"supports_reasoning": True, "max_output_tokens": 8000}}, + raising=False, + ) + settings = make_model_settings("high", model_name="bedrock/global.meta.llama-5-70b") + assert settings.max_tokens is None + + +def test_bedrock_claude_thinking_max_tokens_survives_full_resolve_chain() -> None: + # Confirms .resolve() composition doesn't let one ModelSettings field clobber + # another (reasoning, max_tokens, and prompt-cache extra_args all coexist). + # force_required_tool_choice is passed too: _accepts_required_tool_choice only + # matches OpenAI routes, so it stays a no-op here, but the resolve() call for + # it still runs and must not disturb the other three fields. + settings = make_model_settings( + "high", + model_name="bedrock/global.anthropic.claude-opus-4-8", + force_required_tool_choice=True, + ) + assert settings.max_tokens == 32000 + assert settings.reasoning is not None + assert settings.tool_choice is None + assert (settings.extra_args or {}).get("cache_control_injection_points") + + def test_conversation_tail_breakpoint_moves_with_appended_transcript() -> None: # LiteLLM must place the index=-1 cache_control on the last message however # long the transcript grows.