Skip to content

feat(pi): Amazon Bedrock Model Provider Service support - #478

Open
bjanssen11 wants to merge 9 commits into
databricks:mainfrom
bjanssen11:feat/pi-bedrock-only
Open

feat(pi): Amazon Bedrock Model Provider Service support#478
bjanssen11 wants to merge 9 commits into
databricks:mainfrom
bjanssen11:feat/pi-bedrock-only

Conversation

@bjanssen11

@bjanssen11 bjanssen11 commented Sep 3, 2026

Copy link
Copy Markdown

What this adds

Amazon Bedrock as a model provider for Pi, routed through the Databricks AI Gateway, plus two fixes needed to make it work end to end.

Pi talks to Bedrock over Amazon's native Converse API (bedrock-converse-stream), which is the only dialect the gateway serves for the amazon_bedrock provider 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-stream returns real output. The Model Provider Service is selected per request through the Databricks-Model-Provider-Service header. Any Bedrock chat model the MPS exposes works: Claude, Nova, Llama, and the GPT-family targets all responded in testing.

The two fixes

  1. Bedrock config survives the token refresh. The launch-time refresh and the 30-minute background refresh both re-rendered models.json without the provider arguments, which dropped the databricks-bedrock block and sent the session to a system-hosted model. _refresh_token_once now reads the existing config, and when it finds a databricks-bedrock block it re-applies that block with a freshly refreshed token.

  2. No duplicate User-Agent. Pi's Converse client sets its own User-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

ucode pi --provider <your-bedrock-mps>

If the MPS declares target models, they populate Pi's /model picker. If it is allow_all_targets with none declared, ucode prompts for a Bedrock model ID, for example us.amazon.nova-lite-v1:0 or us.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

git fetch origin feat/pi-bedrock-only
git checkout feat/pi-bedrock-only
uv tool install --reinstall .
ucode pi --provider <your-bedrock-mps>

Tests for the fixes:

uv run pytest tests/test_agent_pi.py

Scope

This PR is Pi only. Codex cannot use Bedrock through the gateway: Codex speaks the OpenAI Responses API, which the gateway refuses for the amazon_bedrock provider. That is tracked in #476, and the earlier codex-plus-Bedrock branch (#455) is closed. This branch carries none of that codex launch code. It branches off main and includes the shared Model Provider Service support Pi depends on.

🤖 Generated with Claude Code

https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn

bjanssen11 and others added 8 commits September 3, 2026 12:20
@bjanssen11 @claude
feat(codex): support amazon_bedrock Model Provider Services
3601277 
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
@bjanssen11 @claude
feat: add ucode providers list/show commands
800b1bc 
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
@bjanssen11 @claude
feat: add Pi Bedrock provider support via correct gateway base URL
dc7d634 
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
@bjanssen11 @claude
fix: add --provider option to ucode pi command
60a3bcd 
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
@bjanssen11 @claude
fix: add pi to _TOOL_PROVIDER_TYPES for amazon_bedrock support
975906c 
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
@bjanssen11 @claude
fix: always prefix Bedrock selector with databricks-bedrock/ in Pi co…
743b156 
…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
@bjanssen11 @claude
fix(pi): keep Bedrock provider across token refresh and drop duplicat…
3bf054d 
…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
@bjanssen11 @claude
fix(pi): resolve provider-support test and lint after dropping codex …
77aad72 
…path

Reflect that Pi now supports anthropic/amazon_bedrock provider services, and
clean up imports left unused once the codex Bedrock launch branch is excluded.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
This was referenced Sep 3, 2026
Closed
Open
@bjanssen11 @claude
fix(pi): pin per-model output caps for Bedrock targets
92cd60d 
Some Bedrock models cap output well below Pi's default request (Nova rejects
maxTokens >= 10000). Pin maxTokens/contextWindow on a Bedrock model entry when
the model has a known limit (new `nova` entry in _MODEL_TOKEN_LIMITS); models
with no known low cap, like Claude, stay unbounded.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@bjanssen11