feat(pi): Amazon Bedrock Model Provider Service support - #475
Closed
bjanssen11 wants to merge 8 commits into
Closed
Conversation
Codex speaks the OpenAI-compatible API, which Bedrock also exposes. `_TOOL_PROVIDER_TYPES` previously restricted codex to `openai` only, so `ucode codex --provider <bedrock-mps>` always failed with "which codex can't route to (supported: openai)." Three changes in databricks.py: - Add `amazon_bedrock` to codex's allowed provider types in `_TOOL_PROVIDER_TYPES`. - Gate the "exposes no Claude models" check in `resolve_provider_service` on `tool == "claude"` so a Bedrock MPS with OpenAI-compatible (non-Claude) targets isn't rejected when codex selects it. - Apply the same `tool == "claude"` guard in `service_usable_for_tool` so Bedrock services without Claude targets appear in the list when codex is the active tool. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
Adds two new subcommands under `ucode providers` to inspect Model Provider Services on the workspace: - `ucode providers list [--tool TOOL]` — lists all MPS services with name, provider type, and declared targets. `--tool claude|codex` filters to services the given tool can actually route through. - `ucode providers show <catalog.schema.service>` — shows full detail for one service: provider type, relay flag, allow_all_targets, and the complete targets list. Motivation: after `ucode codex --provider eng_dev.ai_gateway.amazonbedrock` launched without showing expected Bedrock models, there was no CLI to inspect what targets an MPS exposes. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
When `ucode codex --provider <bedrock-mps>` is used, Codex's built-in model picker queries OpenAI for its model list — showing gpt-5-codex and gpt-5 instead of the Bedrock targets declared on the MPS. Fix this by: - Fetching the MPS targets at launch time and offering a picker (or auto-pinning when there's only one target) - Honoring an explicit `--model` flag for codex in the provider path, which was previously a no-op - Making `codex.write_tool_config` actually use the `model` parameter when a provider is active (it was silently ignored before) Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
Wire `ucode pi --provider <bedrock-mps>` end-to-end:
- `build_pi_base_urls`: add "bedrock" key pointing at `{workspace}/ai-gateway`
(NOT `/ai-gateway/amazonbedrock` — that path maps to the Bedrock control
plane; the standard path routes to the runtime via the MPS header)
- `pi.render_overlay`: add `databricks-bedrock` provider block when
`bedrock_targets` is supplied; defaults the session to the first target
- `pi.write_tool_config`: accept `provider` and `bedrock_targets` kwargs
- `agents.__init__.configure_tool`: pass `bedrock_targets` to Pi; allow
Pi to launch without a model when a Bedrock provider + targets cover it
- `cli.py`: fetch MPS targets for Pi in the provider launch path; handle
`allow_all_targets` with a text prompt; thread `bedrock_targets` through
to `configure_tool`
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
Without it, --provider fell into ctx.args and was forwarded to Pi itself rather than being parsed by ucode, so the Bedrock target-fetching branch never ran. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
Without this entry, ucode pi --provider rejects any Bedrock MPS with "pi can't route to (supported: none)" before ever fetching targets. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
…nfig
`_resolve_model_selector` returns Bedrock model IDs (e.g.
`anthropic.claude-3-haiku-20240307-v1:0`) unprefixed because they
contain no `/`. The old `if not resolved` guard never fired since the
ID is truthy. `_write_settings` then gets an empty model half from
`partition("/")` and exits early — defaultProvider stays on
databricks-claude instead of databricks-bedrock.
Fix: unconditionally set `resolved = f"databricks-bedrock/{targets[0]}"`
when the Bedrock provider block is present.
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
…e UA
The launch-time and 30-minute token refresh re-rendered Pi's models.json
without the Bedrock provider, dropping the databricks-bedrock block and
falling back to a system-hosted model. _refresh_token_once now reads the
existing config and preserves a databricks-bedrock block, re-applying it
with a freshly refreshed token.
Also stop sending ucode's User-Agent on the Bedrock block: Pi's
bedrock-converse-stream client sets its own, and two values made the
gateway reject the request ("Header field 'user-agent' must only have a
single value").
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this adds
Amazon Bedrock as a model provider for Pi, routed through the Databricks AI Gateway, plus two fixes that were needed to make it work end to end.
Pi talks to Bedrock over Amazon's native Converse API (
bedrock-converse-stream). That is the only dialect the gateway serves for theamazon_bedrockprovider type. I verified this against the live gateway: OpenAI Responses and MLflow chat-completions both return "not supported for provider Amazon Bedrock", while/ai-gateway/model/{id}/converse-streamreturns a real reply. The Model Provider Service is picked per request through theDatabricks-Model-Provider-Serviceheader. Any Bedrock chat model the MPS exposes works: Claude, Nova, Llama, and the GPT-family targets all responded in testing.The two fixes
Bedrock config survives the token refresh. The launch-time refresh and the 30-minute background refresh both re-rendered
models.jsonwithout the provider arguments, which dropped thedatabricks-bedrockblock and sent the session to a system-hosted model instead._refresh_token_oncenow reads the existing config, and when it finds adatabricks-bedrockblock it re-applies that block with a freshly refreshed token.No duplicate
User-Agent. Pi's Converse client sets its ownUser-Agent, and ucode set one too. The gateway rejected the request with "Header field 'user-agent' must only have a single value". The Bedrock block now sends only the MPS header.How to use it
If the MPS declares target models, they populate Pi's
/modelpicker. If it isallow_all_targetswith none declared, ucode prompts for a Bedrock model ID, for exampleus.amazon.nova-lite-v1:0orus.anthropic.claude-sonnet-4-20250514-v1:0. Add or remove targets on the MPS and the next launch reflects the change.Install and try it locally
Tests for the fixes:
Notes for reviewers
Codex stays incompatible with Bedrock on this gateway. Codex speaks the OpenAI Responses API, which the gateway refuses for the
amazon_bedrockprovider. Only Converse works, and Pi speaks it.This branch builds on the shared Model Provider Service code from #455. It is a draft against
main, so until #455 merges the diff also shows that shared foundation. Once #455 lands, this collapses to the Pi provider commits plus the two fixes above.🤖 Generated with Claude Code
https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn