Skip to content

feat(codex): Amazon Bedrock Model Provider Service support - #538

Open
bjanssen11 wants to merge 13 commits into
databricks:mainfrom
bjanssen11:feat/codex-bedrock
Open

feat(codex): Amazon Bedrock Model Provider Service support#538
bjanssen11 wants to merge 13 commits into
databricks:mainfrom
bjanssen11:feat/codex-bedrock

Conversation

@bjanssen11

@bjanssen11 bjanssen11 commented Sep 9, 2026

Copy link
Copy Markdown

What this adds

Amazon Bedrock as a model provider for Codex, routed through the Databricks AI Gateway. It completes the Bedrock provider work alongside Pi (#478) and OpenCode (#481).

Codex speaks the OpenAI Responses API. Earlier the gateway refused Responses for the amazon_bedrock provider (tracked in #476), so Codex couldn't reach Bedrock. A gateway update now serves the Codex dialect for Bedrock targets, which unblocks this. The Databricks-Model-Provider-Service header selects the MPS per request, and Bedrock's provider-side target ids (e.g. us.openai.gpt-5.6-luna) are what the gateway routes on.

The three fixes

  1. Pin the Bedrock target at launch. ucode codex --provider <bedrock-mps> used to fall through to Codex's built-in picker, which queries OpenAI and lists gpt-5-codex/gpt-5 instead of the MPS targets. Launch now fetches the MPS targets and pins one (auto-pins a single target, prompts when there are several), honors an explicit --model, and write_tool_config applies the model when a provider is active.

  2. Keep the pinned model through clear_model_preferences. It runs at the start of launch() and wipes any model not backed by codex_default_model. With an MPS active it was wiping the Bedrock target id, sending Codex back to its own picker and returning a 403 ("not in the allowed models list"). clear_model_preferences now returns early when a Databricks-Model-Provider-Service header is set.

  3. Supply model metadata the gateway can't. Codex pulls its catalog from codex/v1/models, which is disabled on this workspace, so it had no metadata for the Bedrock target ids and fell back to generic defaults. That printed a warning and dropped web search and reasoning config for the session. ucode now generates the catalog client-side. codex debug models (read from the base config, not ucode's profile, to avoid a feedback loop) yields the built-in entries; ucode clones each OpenAI-dialect target under its Bedrock slug, writes the result to ~/.codex/ucode-model-catalog.json, and pins it via model_catalog_json. Non-OpenAI targets have no Codex metadata to borrow, so ucode skips them.

Verified end to end against the live gateway: the picker shows the Bedrock targets, the 403 and the metadata warning are gone, and Codex returns Bedrock output with web search intact.

How to use it

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

If the MPS declares target models, they populate the picker. If it is allow_all_targets with none declared, ucode prompts for a Bedrock model ID, for example us.openai.gpt-5.6-luna.

Install and try it locally

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

Tests:

uv run pytest tests/test_agent_codex.py

Scope and dependency

This stacks on #478 (Pi) and #481 (OpenCode), which carry the shared amazon_bedrock provider-type support and the launch-time target resolution Codex reuses. Fork PRs can't base on a fork branch, so this PR targets main. Until #478 and #481 merge, the diff here also shows their changes; once they land, this collapses to the Codex commits.

The Codex-Bedrock limitation noted in #476 no longer applies on this workspace: the gateway now serves the Codex dialect for Bedrock.

Stack

Part of the Bedrock provider stack. Merge in this order:

  1. feat(pi): Amazon Bedrock Model Provider Service support #478 Pi Bedrock (base of the stack)
  2. feat(opencode): Amazon Bedrock Model Provider Service support #481 OpenCode Bedrock
  3. feat(codex): Amazon Bedrock Model Provider Service support #538 Codex Bedrock

Every fork PR targets main, because a fork PR can't base on a fork branch. Each PR's diff therefore includes the commits below it until those merge.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HaJ5LAQfUTfJPhTAismXmU

bjanssen11 and others added 13 commits September 9, 2026 14:03
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
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
…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
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
Route OpenCode to an Amazon Bedrock MPS through the Databricks AI Gateway,
mirroring the Pi support. OpenCode uses the @ai-sdk/amazon-bedrock provider
with a bearer apiKey (no SigV4, no region) against {workspace}/ai-gateway,
which the SDK turns into /model/{id}/converse-stream. The
Databricks-Model-Provider-Service header rides per-model, since OpenCode
clobbers provider-level headers.

_refresh_token_once now preserves an existing databricks-bedrock block across
the launch-time and 30-minute token refresh (reading the saved MPS name and
target ids back out of the per-model headers), so the session keeps routing to
Bedrock instead of dropping to a system-hosted model.

Verified end to end against the live gateway: the generated config survives a
refresh and a real `opencode run` returns Bedrock output.

Co-Authored-By: Claude Opus 4.8 (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
clear_model_preferences runs at the start of launch() and wipes any
model key that is not backed by codex_default_model in state.  When a
Bedrock MPS provider is active the model written to the config is the
Bedrock target id (e.g. us.openai.gpt-5.6-luna) that the gateway uses
for routing — clearing it causes Codex to fall back to its own model
picker, which queries OpenAI directly and returns an id the MPS does not
recognise, producing a 403.

Add _has_active_mps() and return early from clear_model_preferences when
Databricks-Model-Provider-Service is set in the provider block.  This
covers both the direct call in launch() and the indirect path through
default_model() inside _app_server_start_model().

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HaJ5LAQfUTfJPhTAismXmU
Codex refreshes its model catalog from the gateway's codex/v1/models
endpoint, which is disabled on this workspace ("not enabled for this
workspace"). Codex then has no metadata for the Bedrock provider-side
target ids it resolves to (e.g. us.openai.gpt-5.6-luna) and falls back to
generic metadata — printing a warning and, worse, losing web search and
reasoning configuration for the session.

The gateway can't serve the catalog either: Bedrock has no models-list
operation, and enabling the codex dialect is an undocumented platform-team
toggle. So generate the catalog client-side instead:

- `codex debug models` (no --profile, so it reads the base config, not
  ucode's profile where model_catalog_json lives — avoids a compounding
  feedback loop) yields the pristine built-in catalog.
- For each OpenAI-dialect MPS target, clone the matching built-in entry
  under the Bedrock slug so get_model_info resolves real metadata. Non-
  OpenAI targets (claude/grok/qwen) have no Codex metadata to borrow and
  are skipped.
- Write it to ~/.codex/ucode-model-catalog.json and pin it via
  model_catalog_json in the ucode profile. Regeneration is idempotent;
  the pin is dropped when a run isn't routing through a provider.

Verified end to end against the live gateway: the warning and the "web
search is not supported" errors are gone and codex returns Bedrock output.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant