mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
Susana Ferreira
01ec5e4577
## 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
70 lines
1.7 KiB
Go
70 lines
1.7 KiB
Go
package keypool_test
|
|
|
|
import (
|
|
"net/http"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/xerrors"
|
|
|
|
"github.com/coder/coder/v2/aibridge/keypool"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
// errFakeRoundTripperCalled is returned by fakeRoundTripper if it
|
|
// ever gets invoked. The constructor identity tests should never
|
|
// trigger a RoundTrip call.
|
|
var errFakeRoundTripperCalled = xerrors.New("fakeRoundTripper should not be invoked")
|
|
|
|
// fakeRoundTripper is a no-op http.RoundTripper used to check
|
|
// constructor identity in tests.
|
|
type fakeRoundTripper struct{}
|
|
|
|
func (*fakeRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
|
|
return nil, errFakeRoundTripperCalled
|
|
}
|
|
|
|
func TestNewKeyFailoverTransport(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
pool, err := keypool.New("test-provider", []string{"k0"}, quartz.NewMock(t), nil)
|
|
require.NoError(t, err)
|
|
|
|
tests := []struct {
|
|
name string
|
|
// Constructor input.
|
|
config keypool.KeyFailoverConfig
|
|
// Whether the constructor returns inner unchanged.
|
|
expectSame bool
|
|
}{
|
|
{
|
|
// Pool is nil: failover is disabled, inner is returned unchanged.
|
|
name: "pool_nil_returns_inner",
|
|
config: keypool.KeyFailoverConfig{},
|
|
expectSame: true,
|
|
},
|
|
{
|
|
// Pool is set: inner is wrapped in a key-failover transport.
|
|
name: "pool_set_returns_wrapper",
|
|
config: keypool.KeyFailoverConfig{Pool: pool},
|
|
expectSame: false,
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
inner := &fakeRoundTripper{}
|
|
got := keypool.NewKeyFailoverTransport(inner, tc.config)
|
|
|
|
if tc.expectSame {
|
|
assert.Same(t, inner, got)
|
|
} else {
|
|
assert.NotSame(t, inner, got)
|
|
}
|
|
})
|
|
}
|
|
}
|