feat(models): add Bedrock Grok 4.3 support

[swayamg20]

Summary

• Register bedrock_mantle/xai.grok-4.3 with its 1,000,000-token context window.
• Include bedrock_mantle/ calls in the ambient AWS credential refresh and serialization path used by AWS_USE_IMDS.
• Keep litellm.model_id scoped to classic bedrock/ runtime calls; Mantle continues to use Bedrock Projects.
• Document Grok 4.3 setup, Mantle permissions, and the cost-allocation boundary.

Context-window rationale

AWS's Grok 4.3 model card identifies xai.grok-4.3 as a bedrock-mantle model with a 1M-token context window. PR-Agent uses MAX_TOKENS as the input-context budget for prompt pruning, while operators can still impose a lower ceiling with config.max_model_tokens.

The LiteLLM route is bedrock_mantle/xai.grok-4.3, not the classic bedrock/ provider. Both providers use the AWS credential chain, so Mantle now shares the existing IMDS refresh and lock boundary while forwarding the configured model unchanged.

Classic Bedrock application inference profiles use litellm.model_id. Mantle uses Bedrock Projects instead, so model_id is intentionally omitted from Mantle completion calls.

Rebase

Rebased onto current main after the LiteLLM 1.95.0 update. The conflict resolution preserves upstream cache-control behavior alongside the Mantle credential and provider-boundary changes.

Tests

• PYTHONPATH=. ./.venv/bin/pytest -q tests/unittest/test_get_max_tokens.py tests/unittest/test_litellm_imds.py tests/unittest/test_litellm_chat_completion_core.py — 170 passed
• PYTHONPATH=. ./.venv/bin/pytest tests/unittest -q — 1,731 passed, 1 skipped, 1 xfailed
• PR-scoped pre-commit hooks — passed
• pip check — passed
• git diff --check — passed

Closes #2577

[qodo-free-for-open-source-projects]

PR Summary by Qodo

Add Amazon Bedrock Mantle Grok 4.3 model support

✨ Enhancement 📝 Documentation 🧪 Tests 🕐 20-40 Minutes

AI Description

• Register Bedrock Mantle Grok 4.3 with 1M-token context in the model token registry.
• Refresh IMDS/ambient AWS credentials for both bedrock/ and bedrock_mantle/ calls.
• Document Grok 4.3 setup and add regression tests for tokens and IMDS refresh.

Diagram

graph TD
  A["PR-Agent config (model)"] --> B["Model token registry"] --> C["get_max_tokens()"] --> D["Prompt budget / pruning"]
  A --> E["LiteLLM AI handler"] --> F["IMDS refresh + lock"] --> G{{"AWS Bedrock or Mantle"}}

Loading

High-Level Assessment

The following are alternative approaches to this PR:

1. Prefix parsing via startswith()/split() instead of substring match

• ➕ Avoids false positives if a model name ever contains 'bedrock/' as a non-prefix substring
• ➕ More explicit: reflects provider-prefix semantics used elsewhere (e.g., startswith("openrouter/"))
• ➖ Slightly less flexible if model strings can legitimately be nested/prefixed in the future (e.g., stacked routing prefixes)

2. Centralize “AWS providers requiring IMDS refresh” into a helper/constant

• ➕ Single source of truth as more AWS-backed providers are added
• ➕ Makes tests and future additions less error-prone
• ➖ Small refactor overhead for a currently simple condition

3. Ask LiteLLM for the resolved provider (provider detection API)

• ➕ Uses the same provider identification logic as the downstream client
• ➕ Reduces bespoke string matching
• ➖ Adds coupling to LiteLLM internals/APIs and may vary across LiteLLM versions

Recommendation: The PR’s approach is appropriate and minimally invasive for adding Bedrock Mantle support. If additional AWS provider prefixes are expected, consider centralizing provider detection (and using explicit prefix parsing) to keep the IMDS/lock guard accurate and maintainable.

Files changed (5) +38 / -4

Enhancement (1) +1 / -0

__init__.pyRegister bedrock_mantle/xai.grok-4.3 with 1M max tokens +1/-0

Register bedrock_mantle/xai.grok-4.3 with 1M max tokens

• Extends the central model max-token registry with an entry for Grok 4.3 on Bedrock Mantle. Records 1,000,000 tokens while noting config.max_model_tokens may cap effective usage.

pr_agent/algo/init.py

Bug fix (1) +3 / -1

litellm_ai_handler.pyRefresh IMDS credentials for both bedrock/ and bedrock_mantle/ calls +3/-1

Refresh IMDS credentials for both bedrock/ and bedrock_mantle/ calls

• Expands the AWS_USE_IMDS Bedrock guard to include the bedrock_mantle/ provider prefix. Ensures Mantle requests share the same env-serialization lock and per-call credential refresh behavior as classic Bedrock.

pr_agent/algo/ai_handlers/litellm_ai_handler.py

Tests (2) +23 / -3

test_get_max_tokens.pyAdd max-token regression test for Grok 4.3 Mantle model +12/-0

Add max-token regression test for Grok 4.3 Mantle model

• Introduces a unit test asserting get_max_tokens returns 1,000,000 for bedrock_mantle/xai.grok-4.3 when no config caps are applied.

tests/unittest/test_get_max_tokens.py

test_litellm_imds.pyParametrize IMDS refresh test across Bedrock and Bedrock Mantle providers +11/-3

Parametrize IMDS refresh test across Bedrock and Bedrock Mantle providers

• Updates the IMDS refresh test to run for both bedrock/ and bedrock_mantle/ models. Also asserts the exact model string is forwarded unchanged to the LiteLLM acompletion call.

tests/unittest/test_litellm_imds.py

Documentation (1) +11 / -0

changing_a_model.mdDocument Bedrock Mantle Grok 4.3 configuration and IAM differences +11/-0

Document Bedrock Mantle Grok 4.3 configuration and IAM differences

• Adds a Grok 4.3 example using the bedrock_mantle/ provider prefix. Notes that Mantle uses the same credential sources as Bedrock but requires different IAM permissions, linking to AWS documentation.

docs/docs/usage-guide/changing_a_model.md

[qodo-free-for-open-source-projects]

Code Review by Qodo

🐞 Bugs (2) 📘 Rule violations (0) 📎 Requirement gaps (1) 🎨 UX issues (0) 📜 Skill insights (0)

1. Missing Grok in configuration.toml 📎 Requirement gap ⚙ Maintainability

Description

The configuration template does not document how to select Grok via bedrock_mantle/, so users
can’t discover/configure Grok-on-Bedrock through the standard configuration.toml defaults as
required. This creates a mismatch with the intended “template-driven” configuration flow for
enabling Grok on AWS Bedrock.

Code

pr_agent/settings/configuration.toml[R364-366]

+# Application inference profile ARN for classic `bedrock/` calls through the `bedrock-runtime` APIs.
+# Not used by `bedrock_mantle/`, which uses Projects for cost allocation.
+# model_id = ""

Evidence

Compliance ID 6 requires adding Grok-on-Bedrock configuration options to the repository
configuration template. The template’s [config] model defaults show only generic models and, even
where bedrock_mantle/ is referenced, it is only in a litellm.model_id comment—there is no
documented bedrock_mantle/xai.grok-4.3 model selection example in configuration.toml.

Add configuration support for Grok on AWS Bedrock: Add configuration support for Grok on AWS Bedrock
pr_agent/settings/configuration.toml[5-10]
pr_agent/settings/configuration.toml[364-366]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`pr_agent/settings/configuration.toml` does not include any documented example/options for selecting Grok via AWS Bedrock Mantle (e.g., `model="bedrock_mantle/xai.grok-4.3"`). The ticket-derived compliance requirement expects Grok-on-Bedrock to be discoverable via the standard configuration template.
## Issue Context
This PR adds Grok 4.3 support and documents it in the usage guide, but the repository’s configuration template still only shows generic model defaults and does not provide a Grok/Mantle example in the `[config]` models section.
## Fix Focus Areas
- pr_agent/settings/configuration.toml[5-10]
- pr_agent/settings/configuration.toml[364-366]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

2. IMDS provider substring match 🐞 Bug ➹ Performance

Description

LiteLLMAIHandler.chat_completion() enables the IMDS refresh/lock when the model string merely
contains bedrock/ or bedrock_mantle/ anywhere, so a non-Bedrock provider whose model ID includes
these substrings would be serialized behind _aws_bedrock_lock and have AWS credentials refreshed
unnecessarily. This can add avoidable latency and reduce concurrency in AWS_USE_IMDS mode.

Code

pr_agent/algo/ai_handlers/litellm_ai_handler.py[R542-544]

+        _bedrock_imds = self._aws_imds_mode and any(
+            provider in model for provider in ("bedrock/", "bedrock_mantle/")
+        )

Evidence

The IMDS gate is triggered via substring inclusion and then used to choose the Bedrock lock +
refresh path; other provider routing logic in the same function relies on strict prefixes (e.g.
openrouter/), so substring checks can misclassify model strings where bedrock_mantle/ appears
after another provider prefix.

pr_agent/algo/ai_handlers/litellm_ai_handler.py[536-548]
pr_agent/algo/ai_handlers/litellm_ai_handler.py[562-568]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`_bedrock_imds` is computed using substring checks (`provider in model`). This can mis-detect Bedrock IMDS mode for model strings that are not actually Bedrock/BEDROCK_MANTLE providers but contain those substrings as part of their model ID.
### Issue Context
The flag controls both credential refresh and acquisition of `_aws_bedrock_lock`, so false positives unnecessarily serialize requests and mutate AWS env vars.
### Fix Focus Areas
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[539-546]
### Proposed fix
- Compute the provider from the model prefix instead of substring matching, e.g.:
- `provider = model.split('/', 1)[0]` and check `provider in ('bedrock', 'bedrock_mantle')`, or
- `model.startswith(('bedrock/', 'bedrock_mantle/'))`.
- (Optional) Add a unit test ensuring a model like `openrouter/bedrock_mantle/xai.grok-4.3` does **not** trigger `_refresh_aws_imds_credentials()` when `AWS_USE_IMDS=true`.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

3. Mantle skips model_id 🐞 Bug ≡ Correctness

Description

LiteLLMAIHandler only attaches litellm.model_id when the model string contains 'bedrock/', so
bedrock_mantle/* requests (including bedrock_mantle/xai.grok-4.3) will not receive the
configured inference profile ID and will silently ignore it when set. This is inconsistent with the
new Mantle-as-Bedrock handling added in this PR and with the documentation’s claim that model_id
is passed to Bedrock calls.

Code

pr_agent/algo/ai_handlers/litellm_ai_handler.py[R486-488]

+        _bedrock_imds = self._aws_imds_mode and any(
+            provider in model for provider in ("bedrock/", "bedrock_mantle/")
+        )

Evidence

The PR updates IMDS/lock detection to treat bedrock_mantle/ as a Bedrock provider, but the later
model_id forwarding logic remains gated on 'bedrock/' only, so Mantle calls won’t receive
model_id. The docs also describe model_id as applying to Bedrock calls, which becomes misleading
once Mantle is a supported Bedrock route.

pr_agent/algo/ai_handlers/litellm_ai_handler.py[483-492]
pr_agent/algo/ai_handlers/litellm_ai_handler.py[686-690]
docs/docs/usage-guide/changing_a_model.md[317-336]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`litellm.model_id` (Bedrock inference profile) is only forwarded when `'bedrock/' in model`, which excludes the newly supported `bedrock_mantle/` provider.
### Issue Context
This PR adds `bedrock_mantle/` detection for the IMDS credential refresh/lock boundary, but the `model_id` forwarding logic later in `chat_completion()` still only recognizes classic `bedrock/`.
### Fix Focus Areas
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[483-492]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[686-690]
### Suggested fix
Update the `model_id` injection guard to include Mantle models as well (preferably by reusing the same provider-detection logic used for `_bedrock_imds`). For example:
- Replace `if model_id and 'bedrock/' in model:` with `if model_id and any(p in model for p in ("bedrock/", "bedrock_mantle/")):`
Optionally add/extend a unit test to verify `kwargs["model_id"]` is set for `bedrock_mantle/*` when `litellm.model_id` is configured.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

• Author self-review: I have reviewed the code review findings, and addressed the relevant ones.

Tip of the day

💡 Did you know, you can hide the parts of a finding you never read, like the evidence or the agent prompt

More tips ↗ | Customize Qodo ↗ | Qodo docs ↗

[swayamg20]

Regarding "Mantle skips model_id": this is an expected provider boundary, not a missing Mantle branch. AWS documents application inference profiles for bedrock-runtime, while Responses and Chat Completions on bedrock-mantle use Projects. With this repo's pinned LiteLLM 1.93.0, explicitly forwarding model_id would serialize it into the Mantle OpenAI-compatible request body; it would not select a Bedrock Runtime inference profile. Mantle project selection is instead mapped through aws_bedrock_project_id to the OpenAI-Project header.

Commit ede851ab clarifies the docs, configuration, and handler comments, and adds a regression test proving model_id is forwarded for classic bedrock/ calls and omitted for bedrock_mantle/ (165 focused tests and pre-commit pass). No runtime widening was made.

[qodo-free-for-open-source-projects]

Code review by qodo was updated up to the latest commit ede851a

[IsmaelMartinez]

Checked this against #2577. bedrock_mantle/xai.grok-4.3 and the 1M window both match the AWS model card, and widening the IMDS check past bedrock/ is needed since Mantle would otherwise miss the credential refresh. Shape matches the recent model additions in #2606 and #2538.

It is conflicting with main, probably from the litellm bump in #2611, so it needs a rebase.

Not for this PR: #2530 checks startswith('grok-') on the trailing path segment, which here is xai.grok-4.3, so Bedrock Grok would get no reasoning effort forwarding if both land.

[qodo-free-for-open-source-projects]

Code review by qodo was updated up to the latest commit 00b2425

[IsmaelMartinez]

Sorry for the churn here: after your quick rebase this sat two weeks on us, and this morning a merge to main added a test in the same region of test_litellm_chat_completion_core.py, so it conflicted again through no fault of yours. The resolution was purely additive (keep both tests), so I have pushed it to your branch rather than ask you twice. The substance was already verified: the model entries match the AWS card and coverage is real. Approving once CI reruns.

[IsmaelMartinez]

Green on the resolved head, approving as promised above. Thanks for the patience on this one.

[qodo-code-review]

Code Review by Qodo

🐞 Bugs (0) 📘 Rule violations (2) 📜 Skill insights (0)

1. changing_a_model.md lines too long 📘 Rule violation ⚙ Maintainability

Description

The updated documentation includes lines longer than 120 characters, which violates the repository
line-length requirement and reduces readability. This can also cause lint/formatting checks to fail
if enforced.

Code

docs/docs/usage-guide/changing_a_model.md[337]

+The `litellm.model_id` parameter applies only to classic `bedrock/` calls made through the `bedrock-runtime` APIs. It does not apply to `bedrock_mantle/`; for cost allocation with the Mantle Chat Completions and Responses APIs, use [Amazon Bedrock Projects](https://docs.aws.amazon.com/bedrock/latest/userguide/cost-mgmt-projects.html).

Relevance

●●● Strong

Recent history accepts line-length fixes; analogous overlong strings were explicitly wrapped for
compliance.

PR-#2584
PR-#2318

ⓘ Recommendations generated based on similar findings in past PRs

Evidence

PR Compliance ID 2694690 requires keeping modified source lines within 120 characters. The modified
documentation lines containing long URLs and sentences (e.g., lines 321 and 337) exceed this limit.

Rule 2694690: Enforce maximum line length of 120 characters
docs/docs/usage-guide/changing_a_model.md[321-321]
docs/docs/usage-guide/changing_a_model.md[337-337]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
Some modified Markdown lines exceed the 120-character limit.

## Issue Context
The repository enforces a 120-character maximum line length for modified source files.

## Fix Focus Areas
- docs/docs/usage-guide/changing_a_model.md[321-321]
- docs/docs/usage-guide/changing_a_model.md[337-337]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

2. Non-imperative model_id comment 📘 Rule violation ⚙ Maintainability

Description

Newly added/modified comments and a test docstring are written in descriptive form (e.g., "calls
use" / "uses" / "is called") rather than imperative mood, reducing consistency with the project's
required documentation/comment style. This conflicts with the mandated imperative phrasing
convention for new or updated comments/docstrings.

Code

pr_agent/algo/ai_handlers/litellm_ai_handler.py[R778-779]

+                # Classic `bedrock/` calls use model_id for Bedrock Runtime inference profiles.
+                # Bedrock Mantle uses Projects, so `bedrock_mantle/` intentionally omits it.

Relevance

●● Moderate

Imperative comments were accepted recently, but a closely matching descriptive test-docstring
finding was rejected.

PR-#2791
PR-#2703
PR-#2661

ⓘ Recommendations generated based on similar findings in past PRs

Evidence

PR Compliance ID 2694688 requires imperative phrasing for any new or modified comments and
docstrings. In pr_agent/algo/ai_handlers/litellm_ai_handler.py at lines 778-779, the added comments
use descriptive wording ("calls use", "uses") instead of imperative phrasing, and in
tests/unittest/test_litellm_imds.py at lines 412-413, the modified test docstring similarly uses
descriptive phrasing ("is called") rather than the required imperative mood, demonstrating
non-compliance in both locations.

Rule 2694688: Docstrings and comments must use imperative phrasing
pr_agent/algo/ai_handlers/litellm_ai_handler.py[778-779]
tests/unittest/test_litellm_imds.py[412-413]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
Newly added/modified comments and a test docstring in this PR are written in descriptive mood instead of the required imperative mood.

## Issue Context
PR Compliance ID 2694688 requires that new or changed comments/docstrings use imperative phrasing (e.g., "Use ...", "Ensure ...") rather than descriptive constructions like "uses" or "is called".

## Fix Focus Areas
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[778-779]
- tests/unittest/test_litellm_imds.py[412-413]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

• Author self-review: I have reviewed the code review findings, and addressed the relevant ones.

Context sources

✅ Compliance rules (platform): 34 rules

Review mode: ⚖️ Balanced: This changes runtime AWS credential handling and provider-specific request semantics, with real integration and compatibility risk across multiple paths, but not enough independent logic density to justify extended review.

Tip of the day

💡 Did you know, you can hide the parts of a finding you never read, like the evidence or the agent prompt

More tips ↗ | Customize Qodo ↗ | Qodo docs ↗

[swayamg20]

No worries at all, thanks for resolving the conflict getting this merged, @IsmaelMartinez!