mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
Susana Ferreira
7b903cad73
## Problem Centralized requests recorded *the first available key from the pool at `CreateInterceptor` time* as `credential_hint`, so the interception could be persisted in the database with a hint that didn't match the key that actually served the request. The fix consists in storing, at end-of-interception, the hint of the key that succeeded, or the last attempted key if all keys are unavailable. ## Changes - Add `Key.Hint()` and update `credential_hint` on every failover attempt so it reflects the actually-used key. - Stop pre-populating `credential_hint` at `CreateInterceptor`. Centralized starts empty and is updated by the key failover loop. - Persist the final hint via `RecordInterceptionEnded`; SQL updates `credential_hint` only when `credential_kind = 'centralized'` so BYOK keeps its start-time value. - Log the actually-used hint on interception end/failure; start log uses a `<keypool-pending>` placeholder for centralized. > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
480 lines
16 KiB
Go
480 lines
16 KiB
Go
package messages
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.opentelemetry.io/otel"
|
|
|
|
"cdr.dev/slog/v3"
|
|
"github.com/coder/coder/v2/aibridge/config"
|
|
"github.com/coder/coder/v2/aibridge/intercept"
|
|
"github.com/coder/coder/v2/aibridge/internal/testutil"
|
|
"github.com/coder/coder/v2/aibridge/keypool"
|
|
"github.com/coder/coder/v2/aibridge/mcp"
|
|
"github.com/coder/coder/v2/aibridge/utils"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
// Common request and Anthropic-shaped response bodies.
|
|
const (
|
|
requestBody = `{"model":"claude-opus-4-5","max_tokens":1024,"messages":[{"role":"user","content":"hi"}]}`
|
|
successBody = `{"id":"msg_01","type":"message","role":"assistant","content":[{"type":"text","text":"Hello!"}],"model":"claude-opus-4-5","stop_reason":"end_turn","usage":{"input_tokens":10,"output_tokens":5}}`
|
|
toolUseBody = `{"id":"msg_01","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01","name":"test_tool","input":{}}],"model":"claude-opus-4-5","stop_reason":"tool_use","usage":{"input_tokens":10,"output_tokens":5}}`
|
|
rateLimitBody = `{"type":"error","error":{"type":"rate_limit_error","message":"rate limited"}}`
|
|
authErrorBody = `{"type":"error","error":{"type":"authentication_error","message":"invalid key"}}`
|
|
serverErrorBody = `{"type":"error","error":{"type":"api_error","message":"server error"}}`
|
|
)
|
|
|
|
type upstreamResponse struct {
|
|
statusCode int
|
|
body string
|
|
headers map[string]string
|
|
}
|
|
|
|
func TestBlockingInterception_KeyFailover(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
// Centralized pool keys. Empty when byokKey is set.
|
|
keys []string
|
|
// BYOK key. Empty when keys is set.
|
|
byokKey string
|
|
// Scripted upstream responses keyed by X-Api-Key.
|
|
responses map[string]upstreamResponse
|
|
expectedRequestCount int32
|
|
expectedStatusCode int
|
|
expectedRetryAfter string
|
|
// Expected key states after the request, by index in keys.
|
|
expectedKeyStates []keypool.KeyState
|
|
// Expected credential hint after ProcessRequest: last
|
|
// attempted key for centralized, user key from initial request for BYOK.
|
|
expectedCredentialHint string
|
|
}{
|
|
{
|
|
// Given: 1 valid key returning 200.
|
|
// Then: 1 request, 200 response, key remains valid.
|
|
name: "single_valid_key",
|
|
keys: []string{"k0-long-key"},
|
|
responses: map[string]upstreamResponse{
|
|
"k0-long-key": {statusCode: http.StatusOK, body: successBody},
|
|
},
|
|
expectedRequestCount: 1,
|
|
expectedStatusCode: http.StatusOK,
|
|
expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid},
|
|
expectedCredentialHint: utils.MaskSecret("k0-long-key"),
|
|
},
|
|
{
|
|
// Given: 2 keys; key-0 returns 429, key-1 returns 200.
|
|
// Then: 2 requests, 200 response, key-0 temporary, key-1 valid.
|
|
name: "failover_after_429",
|
|
keys: []string{"k0-long-key", "k1-long-key"},
|
|
responses: map[string]upstreamResponse{
|
|
"k0-long-key": {
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{"Retry-After": "5"},
|
|
body: rateLimitBody,
|
|
},
|
|
"k1-long-key": {statusCode: http.StatusOK, body: successBody},
|
|
},
|
|
expectedRequestCount: 2,
|
|
expectedStatusCode: http.StatusOK,
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStateTemporary,
|
|
keypool.KeyStateValid,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k1-long-key"),
|
|
},
|
|
{
|
|
// Given: 2 keys; key-0 returns 401, key-1 returns 200.
|
|
// Then: 2 requests, 200 response, key-0 permanent, key-1 valid.
|
|
name: "failover_after_401",
|
|
keys: []string{"k0-long-key", "k1-long-key"},
|
|
responses: map[string]upstreamResponse{
|
|
"k0-long-key": {statusCode: http.StatusUnauthorized, body: authErrorBody},
|
|
"k1-long-key": {statusCode: http.StatusOK, body: successBody},
|
|
},
|
|
expectedRequestCount: 2,
|
|
expectedStatusCode: http.StatusOK,
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStatePermanent,
|
|
keypool.KeyStateValid,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k1-long-key"),
|
|
},
|
|
{
|
|
// Given: 2 keys; key-0 returns 403, key-1 returns 200.
|
|
// Then: 2 requests, 200 response, key-0 permanent, key-1 valid.
|
|
name: "failover_after_403",
|
|
keys: []string{"k0-long-key", "k1-long-key"},
|
|
responses: map[string]upstreamResponse{
|
|
"k0-long-key": {statusCode: http.StatusForbidden, body: authErrorBody},
|
|
"k1-long-key": {statusCode: http.StatusOK, body: successBody},
|
|
},
|
|
expectedRequestCount: 2,
|
|
expectedStatusCode: http.StatusOK,
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStatePermanent,
|
|
keypool.KeyStateValid,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k1-long-key"),
|
|
},
|
|
{
|
|
// Given: 3 keys; all return 429 with cooldowns 5s, 3s, 10s.
|
|
// Then: 3 requests, 429 response with smallest Retry-After,
|
|
// all keys temporary.
|
|
name: "all_keys_rate_limited",
|
|
keys: []string{"k0-long-key", "k1-long-key", "k2-long-key"},
|
|
responses: map[string]upstreamResponse{
|
|
"k0-long-key": {
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{"Retry-After": "5"},
|
|
body: rateLimitBody,
|
|
},
|
|
"k1-long-key": {
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{"Retry-After": "3"},
|
|
body: rateLimitBody,
|
|
},
|
|
"k2-long-key": {
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{"Retry-After": "10"},
|
|
body: rateLimitBody,
|
|
},
|
|
},
|
|
expectedRequestCount: 3,
|
|
expectedStatusCode: http.StatusTooManyRequests,
|
|
expectedRetryAfter: "3",
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStateTemporary,
|
|
keypool.KeyStateTemporary,
|
|
keypool.KeyStateTemporary,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k2-long-key"),
|
|
},
|
|
{
|
|
// Given: 2 keys; both return 401.
|
|
// Then: 2 requests, 502 api_error response, both keys permanent.
|
|
name: "all_keys_unauthorized",
|
|
keys: []string{"k0-long-key", "k1-long-key"},
|
|
responses: map[string]upstreamResponse{
|
|
"k0-long-key": {statusCode: http.StatusUnauthorized, body: authErrorBody},
|
|
"k1-long-key": {statusCode: http.StatusUnauthorized, body: authErrorBody},
|
|
},
|
|
expectedRequestCount: 2,
|
|
expectedStatusCode: http.StatusBadGateway,
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStatePermanent,
|
|
keypool.KeyStatePermanent,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k1-long-key"),
|
|
},
|
|
{
|
|
// Given: 2 keys; key-0 returns 500.
|
|
// Then: 1 request, 500 response, both keys remain valid.
|
|
name: "server_error_no_failover",
|
|
keys: []string{"k0-long-key", "k1-long-key"},
|
|
responses: map[string]upstreamResponse{
|
|
"k0-long-key": {statusCode: http.StatusInternalServerError, body: serverErrorBody},
|
|
},
|
|
expectedRequestCount: 1,
|
|
expectedStatusCode: http.StatusInternalServerError,
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStateValid,
|
|
keypool.KeyStateValid,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k0-long-key"),
|
|
},
|
|
{
|
|
// Given: BYOK with a single key returning 429.
|
|
// Then: 1 request, 429 response, no failover, upstream
|
|
// Retry-After propagated to the client.
|
|
name: "byok_no_failover",
|
|
byokKey: "user-byok",
|
|
responses: map[string]upstreamResponse{
|
|
"user-byok": {
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{
|
|
"Retry-After": "5",
|
|
// BYOK doesn't set MaxRetries(0);
|
|
// suppress SDK retries to test a
|
|
// single attempt.
|
|
"x-should-retry": "false",
|
|
},
|
|
body: rateLimitBody,
|
|
},
|
|
},
|
|
expectedRequestCount: 1,
|
|
expectedStatusCode: http.StatusTooManyRequests,
|
|
expectedRetryAfter: "5",
|
|
expectedCredentialHint: utils.MaskSecret("user-byok"),
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
// Mock upstream: counts requests and returns
|
|
// scripted responses keyed by X-Api-Key. An unmapped
|
|
// key falls through to 500 so misconfigured cases
|
|
// surface via the status assertion.
|
|
var requestCount atomic.Int32
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
requestCount.Add(1)
|
|
_, _ = io.Copy(io.Discard, r.Body)
|
|
resp, ok := tc.responses[r.Header.Get("X-Api-Key")]
|
|
if !ok {
|
|
resp = upstreamResponse{statusCode: http.StatusInternalServerError}
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
for hk, hv := range resp.headers {
|
|
w.Header().Set(hk, hv)
|
|
}
|
|
w.WriteHeader(resp.statusCode)
|
|
_, _ = w.Write([]byte(resp.body))
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
cfg := config.Anthropic{BaseURL: upstream.URL + "/"}
|
|
var pool *keypool.Pool
|
|
credInfo := intercept.NewCredentialInfo(intercept.CredentialKindCentralized, "")
|
|
if len(tc.keys) > 0 {
|
|
var err error
|
|
pool, err = keypool.New(tc.keys, quartz.NewMock(t))
|
|
require.NoError(t, err)
|
|
cfg.KeyPool = pool
|
|
} else if tc.byokKey != "" {
|
|
cfg.Key = tc.byokKey
|
|
credInfo = intercept.NewCredentialInfo(intercept.CredentialKindBYOK, tc.byokKey)
|
|
}
|
|
|
|
payload, err := NewRequestPayload([]byte(requestBody))
|
|
require.NoError(t, err)
|
|
|
|
interceptor := NewBlockingInterceptor(
|
|
uuid.New(),
|
|
payload,
|
|
config.ProviderAnthropic,
|
|
cfg,
|
|
nil,
|
|
http.Header{},
|
|
"X-Api-Key",
|
|
otel.Tracer("blocking_test"),
|
|
credInfo,
|
|
)
|
|
interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
|
w := httptest.NewRecorder()
|
|
err = interceptor.ProcessRequest(w, req)
|
|
if tc.expectedStatusCode == http.StatusOK {
|
|
require.NoError(t, err)
|
|
} else {
|
|
require.Error(t, err)
|
|
}
|
|
|
|
assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count")
|
|
assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code")
|
|
assert.Equal(t, tc.expectedRetryAfter, w.Header().Get("Retry-After"), "Retry-After header")
|
|
assert.Equal(t, tc.expectedCredentialHint, interceptor.Credential().Hint, "credential hint")
|
|
if pool != nil {
|
|
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestBlockingInterception_AgenticLoopFailover covers the
|
|
// scenarios that span an agentic-loop continuation: the initial
|
|
// client request and the subsequent tool-call continuation can
|
|
// each fail over independently. Each iteration gets its own
|
|
// walker.
|
|
func TestBlockingInterception_AgenticLoopFailover(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
// Scripted upstream responses consumed in order of
|
|
// upstream request.
|
|
responses []upstreamResponse
|
|
expectedRequestCount int32
|
|
expectedSeenKeys []string
|
|
expectedStatusCode int
|
|
expectedRetryAfter string
|
|
expectedKeyStates []keypool.KeyState
|
|
// Expected credential hint after ProcessRequest: hint of the
|
|
// last attempted key across all agentic-loop iterations.
|
|
expectedCredentialHint string
|
|
}{
|
|
{
|
|
// Given: 2 keys; both upstream calls succeed on key-0.
|
|
// Then: 2 requests, 200 response, both keys remain valid.
|
|
name: "happy_path",
|
|
responses: []upstreamResponse{
|
|
{statusCode: http.StatusOK, body: toolUseBody},
|
|
{statusCode: http.StatusOK, body: successBody},
|
|
},
|
|
expectedRequestCount: 2,
|
|
expectedSeenKeys: []string{"k0-long-key", "k0-long-key"},
|
|
expectedStatusCode: http.StatusOK,
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStateValid,
|
|
keypool.KeyStateValid,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k0-long-key"),
|
|
},
|
|
{
|
|
// Given: 2 keys; key-0 succeeds initially, then 429s
|
|
// during the agentic continuation, key-1 succeeds.
|
|
// Then: 3 requests, 200 response, key-0 temporary,
|
|
// key-1 valid.
|
|
name: "agentic_failover_to_k1",
|
|
responses: []upstreamResponse{
|
|
{statusCode: http.StatusOK, body: toolUseBody},
|
|
{
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{"Retry-After": "5"},
|
|
body: rateLimitBody,
|
|
},
|
|
{statusCode: http.StatusOK, body: successBody},
|
|
},
|
|
expectedRequestCount: 3,
|
|
expectedSeenKeys: []string{"k0-long-key", "k0-long-key", "k1-long-key"},
|
|
expectedStatusCode: http.StatusOK,
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStateTemporary,
|
|
keypool.KeyStateValid,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k1-long-key"),
|
|
},
|
|
{
|
|
// Given: 2 keys; key-0 succeeds initially, then both
|
|
// keys 429 during the agentic continuation.
|
|
// Then: 3 requests, 429 response with smallest
|
|
// Retry-After, both keys temporary.
|
|
name: "agentic_all_keys_fail",
|
|
responses: []upstreamResponse{
|
|
{statusCode: http.StatusOK, body: toolUseBody},
|
|
{
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{"Retry-After": "5"},
|
|
body: rateLimitBody,
|
|
},
|
|
{
|
|
statusCode: http.StatusTooManyRequests,
|
|
headers: map[string]string{"Retry-After": "3"},
|
|
body: rateLimitBody,
|
|
},
|
|
},
|
|
expectedRequestCount: 3,
|
|
expectedSeenKeys: []string{"k0-long-key", "k0-long-key", "k1-long-key"},
|
|
expectedStatusCode: http.StatusTooManyRequests,
|
|
expectedRetryAfter: "3",
|
|
expectedKeyStates: []keypool.KeyState{
|
|
keypool.KeyStateTemporary,
|
|
keypool.KeyStateTemporary,
|
|
},
|
|
expectedCredentialHint: utils.MaskSecret("k1-long-key"),
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var requestCount atomic.Int32
|
|
var seenKeysMu sync.Mutex
|
|
var seenKeys []string
|
|
|
|
// Mock upstream: returns scripted responses in order,
|
|
// records each request's X-Api-Key for assertions.
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
idx := int(requestCount.Add(1)) - 1
|
|
seenKeysMu.Lock()
|
|
seenKeys = append(seenKeys, r.Header.Get("X-Api-Key"))
|
|
seenKeysMu.Unlock()
|
|
_, _ = io.Copy(io.Discard, r.Body)
|
|
|
|
if idx >= len(tc.responses) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
return
|
|
}
|
|
resp := tc.responses[idx]
|
|
w.Header().Set("Content-Type", "application/json")
|
|
for hk, hv := range resp.headers {
|
|
w.Header().Set(hk, hv)
|
|
}
|
|
w.WriteHeader(resp.statusCode)
|
|
_, _ = w.Write([]byte(resp.body))
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
|
|
pool, err := keypool.New([]string{"k0-long-key", "k1-long-key"}, quartz.NewMock(t))
|
|
require.NoError(t, err)
|
|
|
|
cfg := config.Anthropic{
|
|
BaseURL: upstream.URL + "/",
|
|
KeyPool: pool,
|
|
}
|
|
|
|
payload, err := NewRequestPayload([]byte(requestBody))
|
|
require.NoError(t, err)
|
|
|
|
interceptor := NewBlockingInterceptor(
|
|
uuid.New(),
|
|
payload,
|
|
config.ProviderAnthropic,
|
|
cfg,
|
|
nil,
|
|
http.Header{},
|
|
"X-Api-Key",
|
|
otel.Tracer("blocking_test"),
|
|
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
|
|
)
|
|
|
|
// Mock proxy with a tool the upstream's tool_use
|
|
// response will reference.
|
|
proxy := &mockServerProxier{
|
|
tools: []*mcp.Tool{
|
|
{
|
|
Client: stubToolCaller{},
|
|
ID: "test_tool",
|
|
Name: "test_tool",
|
|
ServerName: "coder",
|
|
Logger: slog.Make(),
|
|
},
|
|
},
|
|
}
|
|
interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, proxy)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
|
w := httptest.NewRecorder()
|
|
err = interceptor.ProcessRequest(w, req)
|
|
if tc.expectedStatusCode == http.StatusOK {
|
|
require.NoError(t, err)
|
|
} else {
|
|
require.Error(t, err)
|
|
}
|
|
|
|
assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count")
|
|
assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code")
|
|
assert.Equal(t, tc.expectedRetryAfter, w.Header().Get("Retry-After"), "Retry-After header")
|
|
assert.Equal(t, tc.expectedCredentialHint, interceptor.Credential().Hint, "credential hint")
|
|
|
|
seenKeysMu.Lock()
|
|
defer seenKeysMu.Unlock()
|
|
assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys")
|
|
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
|
|
})
|
|
}
|
|
}
|