Commit Graph
21 Commits
Author SHA1 Message Date
Susana Ferreira 01ec5e4577 feat: add key pool failover metrics to aibridge (#25901)
## Description

This PR adds Prometheus metrics for aibridge's API-key failover, giving visibility into key pool health and failover behavior per provider.

The following metrics are introduced:

- **`key_pool_state`** (gauge): number of keys currently in each state (`valid`, `temporary`, `permanent`) per provider, sampled at scrape time.
- **`key_pool_state_transitions_total`** (counter): key state transitions during failover, labeled by `reason` (`rate_limited`, `unauthorized`, `forbidden`).
- **`key_pool_exhaustions_total`** (counter): times a pool ran out of usable keys, labeled by `outcome` (`rate_limited`, `auth_failed`).
- **`key_pool_failover_attempts`** (histogram): keys attempted before success or exhaustion (per interception for bridged requests, per request for passthrough).

## Changes

- Moves `MarkKeyOnStatus` and key-pool error handling onto `*keypool.Pool`.
- Attaches metrics to each provider's key pool at install time, on construction and on provider reload.
- Adds a scrape-time state collector and a `KeyPools()` accessor on the bridge pool to feed it.
- Tracks per-request key attempts in the bridged and passthrough failover paths.
- Adds test coverage for the new metrics across the keypool unit tests, the bridged intercept failover tests, and the passthrough failover test.

Closes https://github.com/coder/internal/issues/1447
Closes https://linear.app/codercom/issue/AIGOV-198/aibridge-key-failover-observability

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-06-09 10:49:47 +01:00
Susana Ferreira f8c736f859 refactor(aibridge): consolidate key failover interceptor tests (#26032)
Consolidates the per-interceptor key-failover tests into a single table-driven `keyfailover_test.go`, parameterized over the interceptors (`messages`, `chatcompletions`, `responses`) and modes (blocking and streaming).

It keeps the two scenarios as separate tests: `TestInterception_KeyFailover` (failover within a single interception) and `TestInterception_AgenticLoopFailover` (failover across an agentic-loop continuation). Same cases and assertions as before with far less duplication, reusing the shared `testutil` mocks (`MockUpstream`, `MockServerProxier`, fixture helpers).

Closes https://linear.app/codercom/issue/AIGOV-396/consolidate-intercept-key-failover-tests-across-modes-and-providers
Closes https://linear.app/codercom/issue/AIGOV-395/share-a-single-mockupstream-helper-between-integration-tests-and

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-06-09 10:36:36 +01:00
Susana Ferreira f440cbd205 refactor(aibridge): move shared mock helpers to testutil (#25999)
Moves the shared aibridge mock test helpers (`MockUpstream`, `MockServerProxier`, `StubToolCaller`, and the `NewFixtureResponse`/`NewFixtureToolResponse` constructors) out of `aibridge/internal/integrationtest` into `aibridge/internal/testutil`, and exports them.

This lets the per-interceptor test packages (`messages`, `chatcompletions`, `responses`) reuse one set of mocks instead of each redefining its own. The symbols are exported and call sites updated, with no behavior change.

Closes: https://linear.app/codercom/issue/AIGOV-397/move-mockserverproxier-into-a-shared-testutil-package

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-06-09 10:24:33 +01:00
Susana Ferreira 18919425f9 fix(aibridge): initiate SSE stream before agentic continuation to avoid IsStreaming race (#26139)
The agentic loop has a race between the main goroutine and the `Start`
goroutine on the shared `ResponseWriter`. When an iteration's response
contains only injected-tool events (no text to relay), `Start` may not
have called `InitiateStream` by the time main reaches the `IsStreaming`
check on the next iteration. The `IsStreaming` check then returns false,
main writes a JSON error via `writeUpstreamError`, and `Start` later
writes SSE headers and events on top, producing a malformed JSON+SSE
response:

```
{\"error\":{\"message\":\"all configured keys are rate-limited\",\"type\":\"rate_limit_error\"},\"request_id\":\"\",\"type\":\"error\"}event: message_start\n..."
```

Fix: explicitly call `events.InitiateStream(w)` at the agentic
continuation point so the SSE stream is committed before the next
iteration runs. Keeps `messages` consistent with the pattern already
used in `chatcompletions/streaming.go`. `sync.Once` makes the double
call safe.

Related: coder/internal#1524
Related: coder/coder#25654
Closes:
https://linear.app/codercom/issue/AIGOV-336/flake-teststreaminginterception-agenticloopfailoveragentic-all-keys

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by
@ssncferreira
2026-06-08 20:13:13 +01:00
Danny Kopping 9afd3d0bea feat: track OpenCode sessions (#26128)
This follows #26098 by teaching AI Bridge to read OpenCode session IDs
from the X-OpenCode-Session header, so OpenCode requests are grouped
consistently in interception logs.

Adds unit and integration test coverage for the new header.

<details>
<summary>Coder Agents generated</summary>

This pull request was generated with Coder Agents assistance.
</details>
2026-06-08 10:12:03 +02:00
Danny Kopping 47a8c9572f feat: add OpenCode AI Bridge client support (#26098)
Adds OpenCode to AI Bridge client detection so requests with user agents
like `opencode/1.16.0 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.14`
show up as a first-class client instead of `Unknown`.

This also wires the existing OpenCode frontend asset into the AIBridge
UI, adds a Storybook story for the client icon, and updates the
monitoring docs list of supported client values.

<details>
<summary>Coder Agents generated</summary>

This pull request was generated by Coder Agents.

</details>
2026-06-08 06:26:30 +02:00
Susana Ferreira b7635b5036 fix(aibridge): strip proxy headers from bridge requests to fix Bedrock SigV4 signing (#26019)
## Problem

On bridge routes, aibridge acts as a client and originates new outbound
requests via the SDK. Proxy headers (`X-Forwarded-For`,
`X-Forwarded-Host`, etc.) from the inbound client request were forwarded
on the outbound request. The SigV4 signer signs all headers present, so
any in-transit modification by an egress proxy (e.g. appending an IP to
`X-Forwarded-For`) invalidated the signature, causing AWS Bedrock to
reject the request with:

> 403: "The request signature we calculated does not match the signature
you provided."

## Changes

- Strip proxy headers in `PrepareClientHeaders` on bridge routes
- Add unit test for proxy header stripping in `client_headers_test.go`
- Add integration test that verifies SigV4 signature remains valid after
an egress proxy modifies headers in transit
- Add integration test that verifies passthrough routes still set
forwarded headers correctly

Related to internal [Slack
thread](https://codercom.slack.com/archives/C096PFVBZKN/p1779919049215969).

> 🤖 Generated by Coder Agents, modified and reviewed by @ssncferreira
2026-06-04 10:15:38 +02:00
Danny KoppingandClaude Opus 4.7 5b10268827 feat: serve 503 sentinel for disabled providers (#25794)
_Disclosure: created with Coder Agents._

When providers are disabled, we should serve a sentinel error so the
requesting client (Claude Code, Coder Agents, etc) is informed. Coder
Agents can also conditionalize its display to show a helpful error
message.

---------

Signed-off-by: Danny Kopping <danny@coder.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-29 10:24:16 +02:00
DevCatsandcopilot-swe-agent[bot] 094fe971ad chore(aibridge): add AWS PRM user-agent attribution for Bedrock calls (#25221)
Adds middleware in `withAWSBedrockOptions` that appends the AWS Partner
Revenue Measurement (PRM) attribution string to the User-Agent header on
every Bedrock API call made through AI Bridge.

This is the AI Bridge counterpart to the Terraform provisioner change
merged in #23138. Together, they ensure all AWS API calls made by Coder
(both workspace infrastructure via Terraform and LLM inference via
Bedrock) include PRM attribution.

## How it works

- A middleware is added before `bedrock.WithConfig(awsCfg)` that reads
the existing `User-Agent` header and appends
`sdk-ua-app-id/APN_1.1%2Fpc_cdfmjwn8i6u8l9fwz8h82e4w3%24`
- Only affects Bedrock calls; OpenAI and direct Anthropic API calls are
unaffected
- Uses `option.WithMiddleware` rather than `option.WithHeader` because
the existing User-Agent (set by the Anthropic SDK) must be preserved and
appended to, not replaced

## Tests

- **Positive**: `TestAWSBedrockIntegration` verifies PRM attribution is
present in the User-Agent on Bedrock requests
- **Negative**: `TestAnthropicMessages` verifies PRM attribution is
absent on non-Bedrock requests

## References

- Companion Terraform provisioner PR: #23138 (merged)
- Backport: #24052 (merged)
- Preserve existing `AWS_SDK_UA_APP_ID`: #24606 (open)
- Original `coder/aibridge` PR:
https://github.com/coder/aibridge/pull/224 (superseded by this PR since
aibridge was moved into coder/coder via #24190)
- [AWS SDK Application ID
docs](https://docs.aws.amazon.com/sdkref/latest/guide/feature-appid.html)
- [AWS PRM Automated User
Agent](https://prm.partner.aws.dev/automated-user-agent.html) (partner
login required)

> Generated with [Coder Agents](https://coder.com/agents)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
2026-05-28 11:08:00 -05:00
Susana Ferreira 846aac2f74 refactor(aibridge): remove InjectAuthHeader in favor of KeyFailoverConfig (#25618)
## Description

`Provider.InjectAuthHeader` is no longer needed. With the addition of `KeyFailoverConfig` in #24920, authentication is now applied per-attempt by `KeyFailoverTransport` on passthrough routes. This PR removes the dead method from the `Provider` interface, all implementations (`Anthropic`, `OpenAI`, `Copilot`), and the test mock.

The orphaned `InjectAuthHeader` unit tests are replaced with `Test{Anthropic,OpenAI,Copilot}_KeyFailoverConfig`. `TestPassthrough_KeyFailover` is also extended to cover Copilot in the BYOK scenario.

Related to: https://linear.app/codercom/issue/AIGOV-334/aibridge-follow-ups-from-key-failover-prs

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-05-25 19:10:38 +01:00
Paweł Banaszewski 1a8a153c56 chore: fix flake in TestResponsesInjectedTool (#25630)
Fixes flake in TestResponsesInjectedTool.
See
https://github.com/coder/coder/pull/25630/changes/d9bfeb20092129127ad5e7958c5b8dbf46740527
for reproduction.
Due to AsyncRecorded token usages may be recorded in different order
then expected.

Fixes: https://github.com/coder/internal/issues/1544
2026-05-25 16:41:55 +02:00
Ethan c650aabbef chore: standardize on *_internal_test.go for white-box tests (#25601)
My agent added `//nolint:testpackage` to a test file on one of my PRs.
Again. This PR cleans it up across the entire repo and updates the
in-repo conventions so future agents stop doing it.

The repo already has a precedent for white-box tests that need to touch
unexported symbols: `*_internal_test.go` (145+ existing files). The
`testpackage` linter's default `skip-regexp` exempts that filename
suffix, so the `//nolint:testpackage` directive is unnecessary in every
case where someone reached for it. This PR renames 51 such files to
`*_internal_test.go` via `git mv` so blame and history follow, and
strips the dead directive from 2 files that were already correctly named
(`coderd/oauth2provider/authorize_internal_test.go`,
`coderd/x/chatd/advisor_internal_test.go`).

`.claude/docs/TESTING.md` now documents the rule explicitly under *Test
Package Naming*, which is imported into the root `AGENTS.md` via
`@.claude/docs/TESTING.md`. The rule: prefer `package foo_test`; if you
need internal access, rename the file to `*_internal_test.go` rather
than adding a nolint directive.
2026-05-22 20:24:38 +10:00
Paweł Banaszewski c0b4180206 fix: fix race in setupInjectedToolTest (#25455)
Fixes race condition in `setupInjectedToolTest`.

To reproduce add short sleep to AsyncRecorder.RecordToolUsage method
(inside newly spawned go-routine):
https://github.com/coder/coder/blob/46821525f7d2f7466735463898fb6e054c169f85/aibridge/recorder/recorder.go#L254-L258

Upstream request counter can reach 2 while no recording has been done
since asyncRecorder does it asynchronously:
https://github.com/coder/coder/blob/46821525f7d2f7466735463898fb6e054c169f85/aibridge/internal/integrationtest/setupbridge.go#L242-L244

Added consuming request to `setupInjectedToolTest` so
`newInterceptionProcessor` handler finishes before returning from
`setupInjectedToolTest` which guarantees that all recordings are done:
https://github.com/coder/coder/blob/46821525f7d2f7466735463898fb6e054c169f85/aibridge/bridge.go#L282

Fixes: https://github.com/coder/internal/issues/1526
2026-05-19 09:36:38 +02:00
Marcin Tojek 38772bdb7c refactor: remove cache tokens from ExtraTokenTypes (#25118)
Fixes: https://github.com/coder/aibridge/issues/243

> Generated with [Coder Agents](https://coder.com/agents)
2026-05-18 11:30:19 +02:00
Susana Ferreira 0766cc3097 feat: add automatic key failover for AI Bridge passthrough (#24920)
## Description

Adds automatic key failover for passthrough routes for the Anthropic and OpenAI providers. A new `keyFailoverTransport` wraps the reverse-proxy transport: centralized requests walk the configured key pool and retry with the next key on key-specific failures (401/403/429), reusing the same key-marking semantics as the bridged routes.

BYOK passthrough requests run as a single attempt with no failover.

## Changes

- New `keypool.KeyFailoverConfig` carrying the `Pool` to walk and the provider-specific closures (`IsBYOK`, `InjectAuthKey`, `MarkKey`, `BuildExhaustedResponse`).
- New `keypool.NewKeyFailoverTransport`: wraps an inner `http.RoundTripper`. Returns `inner` unchanged when `Pool` is nil, otherwise produces a transport that buffers the request body once, walks the pool per request, and replays each attempt with the next key.
- New `Provider.KeyFailoverConfig(logger)` interface method. Anthropic injects `X-Api-Key`; OpenAI injects `Authorization: Bearer ...`; Copilot returns an empty config.
- `passthrough.go` wires `NewKeyFailoverTransport` around the existing apidump middleware, so every retry attempt is recorded.

## Related Issues

Related to: https://github.com/coder/internal/issues/1446
Related to: https://linear.app/codercom/issue/AIGOV-197/aibridge-automatic-key-failover-for-bridged-and-passthrough-routes

## Follow-up PRs

- Remove dead `Provider.InjectAuthHeader` method now that all auth is applied per-attempt by `KeyFailoverTransport`.
- Bedrock multi-key support.
- Refactor provider vs interceptor config separation.
- Record the actually-used key in the interception credential hint after failover.

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-05-07 15:46:36 +01:00
Susana Ferreira b6dacb4a3c feat: add automatic key failover for AI Bridge OpenAI (#24847)
## Description

Adds automatic key failover for centralized OpenAI provider, covering both chat completions and responses APIs. Same shape as the Anthropic PR: each upstream call walks the configured key pool, keys are marked **temporary** on 429 (with cooldown from `Retry-After`) and **permanent** on 401/403. Each agentic-loop iteration gets its own fresh walker so a tool-call continuation can fail over independently of the initial request.

BYOK is unchanged: BYOK requests run as a single attempt with no failover.

## Changes

- `config.OpenAI` carries a `KeyPool`. `Key` remains for BYOK Authorization Bearer set per interception.
- Chat completions blocking interceptor: walks the pool via `newChatCompletionWithKeyFailover`, marks keys on key-specific failures, returns on first success or non-failover error.
- Chat completions streaming interceptor: per-iteration walker. Pre-stream failures fail over to the next key; mid-stream errors are relayed as SSE events.
- Responses blocking interceptor: extracts `newResponseWithKeyFailover` parallel to chatcompletions.
- Responses streaming interceptor: per-iteration walker, retains the existing buffer-then-forward design.

## Related Issues

Related to: https://github.com/coder/internal/issues/1446
Related to: https://linear.app/codercom/issue/AIGOV-197/aibridge-automatic-key-failover-for-bridged-and-passthrough-routes

## Follow-up PRs

- Bedrock multi-key support.
- Refactor provider vs interceptor config separation.
- Record the actually-used key in the interception credential hint after failover.

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-05-07 15:35:46 +01:00
Susana Ferreira f1155ac4d7 feat: add automatic key failover for AI Bridge Anthropic (#24836)
## Description

Adds automatic key failover for centralized Anthropic provider. When a key pool is configured, each upstream call walks the pool and tries keys in order until one succeeds or the pool is exhausted. Keys are marked **temporary** on 429 (with cooldown from `Retry-After`) and **permanent** on 401/403. Errors that aren't key-specific don't trigger failover. Each agentic-loop iteration gets its own fresh walker, so a tool-call continuation can fail over independently of the initial request.

BYOK is unchanged: BYOK requests run as a single attempt with no failover.

## Changes

- `config.Anthropic` carries a `KeyPool`. `Key` remains for BYOK X-Api-Key set per interception.
- Blocking interceptor: walks the pool, marks keys on key-specific failures, returns on first success or non-failover error.
- Streaming interceptor: per-iteration walker. Pre-stream failures fail over to the next key; mid-stream errors are relayed as SSE events.
- New `keypool` error types: `TransientExhaustionError` (carries soonest cooldown) and `ErrPermanentExhaustion`. Replace the prior `ErrAllKeysExhausted`.
- Error responses now consistently include the outer `"type": "error"` field.

## Related Issues

Related to: https://github.com/coder/internal/issues/1446
Related to: https://linear.app/codercom/issue/AIGOV-197/aibridge-automatic-key-failover-for-bridged-and-passthrough-routes

## Follow-up PRs

- Bedrock multi-key support.
- Refactor provider vs interceptor config separation.
- Record the actually-used key in the interception credential hint after failover.

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-05-07 14:57:44 +01:00
Susana Ferreira dbb50ebaaf feat: remove 429 from aibridge circuit breaker failure conditions (#24701)
## Description

Removes 429 (Too Many Requests) from the circuit breaker failure conditions. Rate limiting is now handled by automatic key failover instead of tripping the circuit breaker.

## Changes

`DefaultIsFailure` no longer treats 429 as a circuit breaker failure. The circuit breaker now only trips on server overload responses (503, 529).

Tests and integration tests updated to use 503 instead of 429 for tripping circuits. Description strings in deployment config updated to reflect the change.

Closes https://github.com/coder/internal/issues/1445

> [!NOTE]
> Initially generated by Coder Agents, modified and reviewed by @ssncferreira
2026-04-30 09:31:32 +01:00
Paweł Banaszewski a24dc19d49 chore: clean up env var usage in aibridge (#24783)
> AI tools where used when creating this PR

This PR removes environment variable parsing from `/aibridge` directory.

Added env variables/flags for dump dir as coder options.
Only added to new indexed provider options
(`CODER_AIBRIDGE_PROVIDER_<N>_*`) not to deprecated legacy env variables
(`CODER_AIBRIDGE_ANTHROPIC_*` and `CODER_AIBRIDGE_OPENAI_KEY_*`).

Reverted adding `MaxRetries` option as it will be removed soon due to
key failover work:
https://github.com/coder/coder/pull/24783#discussion_r3155544808
2026-04-29 18:28:37 +02:00
Paweł Banaszewski b8906c84a1 fix: fixes aibridge integration tests failing on windows (#24665)
Fixes aibridge responses and apidump integration tests in windows CI
job.

example failing job before the fix:
https://github.com/coder/coder/actions/runs/24873191866/job/72823942378
2026-04-24 18:28:22 +02:00
Paweł Banaszewski e00e85765b chore: move aibridge library code into coder repo (#24190)
This PR merges code from `coder/aibridge` repository into `coder/coder`.
It was split into 4 PRs for easier review but stacked PRs will need to
be merged into this PR so all checks pass.

* https://github.com/coder/coder/pull/24190 -> raw code copy (this PR,
before merging PRs on top of it, it was just 1 commit:
https://github.com/coder/coder/commit/70d33f33200c7e77df910957595715f81f9bec24)
* https://github.com/coder/coder/pull/24570 -> update imports in
`coder/coder` to use copied code
* https://github.com/coder/coder/pull/24586 -> linter fixes and CI
integration (also added README.md)
* https://github.com/coder/coder/pull/24571 -> added exclude to
scripts/check_emdash.sh check

Original PR message (before PR squash):
Moves coder/aibridge code into coder/coder repository.

Omitted files:

- `go.mod`, `go.sum`, `.gitignore`, `.github/workflows/ci.yml,`
`Makefile`, `LICENSE`, `README.md` (modified README.md is added later)
- `.github`, `example`, `buildinfo,` `scripts` directories

Simple verification script (will list omitted files)

```
tmp=$(mktemp -d)
echo "$tmp"
git clone --depth=1 https://github.com/coder/aibridge "$tmp/aibridge"
git clone --depth=1 --branch pb/aibridge-code-move https://github.com/coder/coder "$tmp/coder"
diff -rq --exclude=.git "$tmp/aibridge" "$tmp/coder/aibridge"
# rm -rf "$tmp"
```
2026-04-22 17:01:01 +02:00