opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-21 12:44:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add automatic key failover for AI Bridge OpenAI (#24847)

Browse Source 
## 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
This commit is contained in:
Susana Ferreira
2026-05-07 15:35:46 +01:00
committed by GitHub
parent f1155ac4d7
commit b6dacb4a3c
20 changed files with 3112 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aibridge/config/config.go
+11
View File
@@ -50,11 +50,22 @@ type AWSBedrock struct {
BaseURL string
}
// OpenAI carries configuration for an OpenAI provider.
//
// Authentication is mutually exclusive across these two fields,
// set per interception in the provider's CreateInterceptor:
// - KeyPool: centralized requests with automatic key failover.
// - Key: BYOK with Authorization Bearer (single attempt, no
// failover).
//
// TODO(ssncferreira): consolidate the authentication fields per
// https://github.com/coder/aibridge/issues/266.
type OpenAI struct {
// Name is the provider instance name. If empty, defaults to "openai".
Name string
BaseURL string
Key string
KeyPool *keypool.Pool
APIDumpDir string
CircuitBreaker *CircuitBreaker
SendActorHeaders bool
aibridge/fixtures/openai/responses/blocking/http_error.txtar
+6 -6
View File
@@ -6,16 +6,16 @@
}
-- non-streaming --
HTTP/2.0 401 Unauthorized
Content-Length: 234
HTTP/2.0 400 Bad Request
Content-Length: 281
Content-Type: application/json
{
"error": {
"message": "Incorrect API key provided: sk-***. You can find your API key at https://platform.openai.com/account/api-keys.",
"type": "authentication_error",
"param": null,
"code": "invalid_api_key"
"message": "Input tokens exceed the configured limit of 272000 tokens. Your messages resulted in 3148588 tokens. Please reduce the length of the messages.",
"type": "invalid_request_error",
"param": "messages",
"code": "context_length_exceeded"
}
}
aibridge/fixtures/openai/responses/streaming/http_error.txtar
+6 -6
View File
@@ -6,16 +6,16 @@
}
-- streaming --
HTTP/2.0 429 Too Many Requests
Content-Length: 176
HTTP/2.0 400 Bad Request
Content-Length: 281
Content-Type: application/json
{
"error": {
"message": "Rate limit exceeded. Please try again in 20 seconds.",
"type": "rate_limit_error",
"param": null,
"code": "rate_limit_exceeded"
"message": "Input tokens exceed the configured limit of 272000 tokens. Your messages resulted in 3148588 tokens. Please reduce the length of the messages.",
"type": "invalid_request_error",
"param": "messages",
"code": "context_length_exceeded"
}
}
aibridge/intercept/chatcompletions/base.go
+74 -15
View File
@@ -5,8 +5,11 @@ import (
"encoding/json"
"errors"
"fmt"
"math"
"net/http"
"strconv"
"strings"
"time"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
@@ -20,6 +23,7 @@ import (
aibcontext "github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/apidump"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/mcp"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/tracing"
@@ -44,8 +48,19 @@ type interceptionBase struct {
credential intercept.CredentialInfo
}
// newCompletionsService builds the SDK service used for upstream
// calls. BYOK auth is set here. Centralized auth is set
// per-attempt by the failover loop.
func (i *interceptionBase) newCompletionsService() openai.ChatCompletionService {
opts := []option.RequestOption{option.WithAPIKey(i.cfg.Key), option.WithBaseURL(i.cfg.BaseURL)}
// TODO(ssncferreira): validate auth is configured per
// https://github.com/coder/aibridge/issues/266.
var opts []option.RequestOption
// BYOK auth.
if i.cfg.KeyPool == nil {
opts = append(opts, option.WithAPIKey(i.cfg.Key))
}
opts = append(opts, option.WithBaseURL(i.cfg.BaseURL))
// Add extra headers if configured.
// Some providers require additional headers that are not added by the SDK.
@@ -179,6 +194,10 @@ func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *res
}
w.Header().Set("Content-Type", "application/json")
// Set Retry-After when a cooldown is configured.
if oaiErr.RetryAfter > 0 {
w.Header().Set("Retry-After", strconv.Itoa(int(math.Ceil(oaiErr.RetryAfter.Seconds()))))
}
w.WriteHeader(oaiErr.StatusCode)
out, err := json.Marshal(oaiErr)
@@ -190,13 +209,58 @@ func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *res
"type": "error",
"message":"error marshaling upstream error",
"code": "server_error"
},
}
}`))
} else {
_, _ = w.Write(out)
}
}
// For centralized requests, markKeyOnError extracts an OpenAI
// SDK error from err and marks the key based on its status
// code. Returns true if the status was a key-specific failover
// trigger so callers can retry with the next key.
func (i *interceptionBase) markKeyOnError(ctx context.Context, key *keypool.Key, err error) bool {
if i.cfg.KeyPool == nil {
return false
}
var apiErr *openai.Error
if !errors.As(err, &apiErr) {
return false
}
return keypool.MarkKeyOnStatus(
ctx, key, apiErr.Response,
i.logger, i.providerName,
)
}
// processKeyPoolError translates a keypool exhaustion error
// into a developer-facing responseError shaped for the OpenAI
// API. Returns nil if err is not an exhaustion error.
func processKeyPoolError(err error) *responseError {
var transient *keypool.TransientKeyPoolError
switch {
case errors.As(err, &transient):
return newErrorResponse(
"all configured keys are rate-limited",
intercept.OpenAIErrTypeRateLimit,
intercept.OpenAIErrCodeRateLimit,
http.StatusTooManyRequests,
transient.RetryAfter,
)
case errors.Is(err, keypool.ErrPermanentKeyPool):
return newErrorResponse(
"all configured keys failed authentication",
intercept.OpenAIErrTypeAPI,
intercept.OpenAIErrCodeServer,
http.StatusBadGateway,
0,
)
default:
return nil
}
}
func (i *interceptionBase) hasInjectableTools() bool {
return i.mcpProxy != nil && len(i.mcpProxy.ListTools()) > 0
}
@@ -233,15 +297,7 @@ func getErrorResponse(err error) *responseError {
if !errors.As(err, &apiErr) {
return nil
}
return &responseError{
ErrorObject: &shared.ErrorObject{
Code: apiErr.Code,
Message: apiErr.Message,
Type: apiErr.Type,
},
StatusCode: apiErr.StatusCode,
}
return newErrorResponse(apiErr.Message, apiErr.Type, apiErr.Code, apiErr.StatusCode, keypool.ParseRetryAfter(apiErr.Response))
}
var _ error = &responseError{}
@@ -249,15 +305,18 @@ var _ error = &responseError{}
type responseError struct {
ErrorObject *shared.ErrorObject `json:"error"`
StatusCode int `json:"-"`
RetryAfter time.Duration `json:"-"`
}
func newErrorResponse(msg error) *responseError {
func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *responseError {
return &responseError{
ErrorObject: &shared.ErrorObject{
Code: "error",
Message: msg.Error(),
Type: "error",
Code: code,
Message: msg,
Type: errType,
},
StatusCode: status,
RetryAfter: retryAfter,
}
}
aibridge/intercept/chatcompletions/base_test.go
+200
View File
@@ -1,12 +1,22 @@
package chatcompletions //nolint:testpackage // tests unexported internals
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/openai/openai-go/v3"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/utils"
"github.com/coder/quartz"
)
func TestScanForCorrelatingToolCallID(t *testing.T) {
@@ -75,3 +85,193 @@ func TestScanForCorrelatingToolCallID(t *testing.T) {
})
}
}
func TestProcessKeyPoolError(t *testing.T) {
t.Parallel()
tests := []struct {
name string
err error
expectedNil bool
expectedStatus int
expectedRetryAfter time.Duration
}{
{
// Transient with valid keys present: 429, no Retry-After.
name: "transient_zero_retry_after",
err: &keypool.TransientKeyPoolError{},
expectedStatus: http.StatusTooManyRequests,
expectedRetryAfter: 0,
},
{
// Transient with cooldown: 429, Retry-After set.
name: "transient_with_retry_after",
err: &keypool.TransientKeyPoolError{RetryAfter: 5 * time.Second},
expectedStatus: http.StatusTooManyRequests,
expectedRetryAfter: 5 * time.Second,
},
{
// Permanent: 502 api_error.
name: "permanent_returns_502",
err: keypool.ErrPermanentKeyPool,
expectedStatus: http.StatusBadGateway,
},
{
// Anything else: not a pool-exhaustion error.
name: "non_pool_exhaustion_error_returns_nil",
err: xerrors.New("some other error"),
expectedNil: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
got := processKeyPoolError(tc.err)
if tc.expectedNil {
require.Nil(t, got)
return
}
require.NotNil(t, got)
assert.Equal(t, tc.expectedStatus, got.StatusCode)
assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter)
})
}
}
func TestMarkKeyOnError(t *testing.T) {
t.Parallel()
tests := []struct {
name string
err error
expectedReturn bool
expectedState keypool.KeyState
}{
{
// Not an *openai.Error: no status code to act on.
name: "non_api_error_returns_false",
err: xerrors.New("network failure"),
expectedReturn: false,
expectedState: keypool.KeyStateValid,
},
{
// Rate-limited: temporary cooldown.
name: "429_marks_temporary",
err: &openai.Error{StatusCode: http.StatusTooManyRequests, Response: &http.Response{StatusCode: http.StatusTooManyRequests}},
expectedReturn: true,
expectedState: keypool.KeyStateTemporary,
},
{
// Auth failure: mark permanent.
name: "401_marks_permanent",
err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}},
expectedReturn: true,
expectedState: keypool.KeyStatePermanent,
},
{
// Auth forbidden: mark permanent.
name: "403_marks_permanent",
err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}},
expectedReturn: true,
expectedState: keypool.KeyStatePermanent,
},
{
// Server errors are not key-specific.
name: "500_does_not_mark",
err: &openai.Error{StatusCode: http.StatusInternalServerError, Response: &http.Response{StatusCode: http.StatusInternalServerError}},
expectedReturn: false,
expectedState: keypool.KeyStateValid,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
pool, err := keypool.New([]string{"key-0"}, quartz.NewMock(t))
require.NoError(t, err)
key, err := pool.Walker().Next()
require.NoError(t, err)
base := &interceptionBase{cfg: config.OpenAI{KeyPool: pool}, logger: slog.Make()}
got := base.markKeyOnError(context.Background(), key, tc.err)
assert.Equal(t, tc.expectedReturn, got)
assert.Equal(t, tc.expectedState, key.State())
})
}
}
func TestWriteUpstreamError(t *testing.T) {
t.Parallel()
tests := []struct {
name string
respErr *responseError
expectStatus int
// Empty string means the header should be absent.
expectRetryAfter string
// Substring expected in the marshaled body. Empty means no body check.
expectBodyContains string
}{
{
// Standard error: status, code, and JSON body written.
name: "writes_status_and_body",
respErr: newErrorResponse("upstream failed", "api_error", "server_error", http.StatusBadGateway, 0),
expectStatus: http.StatusBadGateway,
expectBodyContains: `"upstream failed"`,
},
{
// OpenAI envelope: the code field round-trips into the body.
name: "writes_code_field",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 0),
expectStatus: http.StatusTooManyRequests,
expectBodyContains: `"rate_limit_exceeded"`,
},
{
// Whole-second retryAfter: emitted as integer seconds.
name: "retry_after_in_seconds",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 60*time.Second),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "60",
},
{
// 500ms rounds up to Retry-After: 1.
name: "retry_after_500ms_rounds_up_to_one",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 500*time.Millisecond),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "1",
},
{
// 200ms rounds up to Retry-After: 1.
name: "retry_after_200ms_rounds_up_to_one",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 200*time.Millisecond),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "1",
},
{
// Negative retryAfter: header omitted.
name: "negative_retry_after_omits_header",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, -1*time.Second),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
base := &interceptionBase{logger: slog.Make()}
w := httptest.NewRecorder()
base.writeUpstreamError(w, tc.respErr)
assert.Equal(t, tc.expectStatus, w.Code, "status code")
assert.Equal(t, "application/json", w.Header().Get("Content-Type"), "Content-Type header")
assert.Equal(t, tc.expectRetryAfter, w.Header().Get("Retry-After"), "Retry-After header")
if tc.expectBodyContains != "" {
assert.Contains(t, w.Body.String(), tc.expectBodyContains, "response body")
}
})
}
}
aibridge/intercept/chatcompletions/blocking.go
+54 -2
View File
@@ -223,6 +223,13 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
return xerrors.Errorf("upstream connection closed: %w", err)
}
// The failover loop may return a keypool exhaustion
// error. Check before the SDK-error path.
if keyErr := processKeyPoolError(err); keyErr != nil {
i.writeUpstreamError(w, keyErr)
return xerrors.Errorf("key pool exhausted: %w", err)
}
if apiErr := getErrorResponse(err); apiErr != nil {
i.writeUpstreamError(w, apiErr)
return xerrors.Errorf("openai API error: %w", err)
@@ -258,9 +265,54 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
return nil
}
func (i *BlockingInterception) newChatCompletion(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (_ *openai.ChatCompletion, outErr error) {
ctx, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
// newChatCompletion routes between BYOK (single attempt) and
// centralized failover.
func (i *BlockingInterception) newChatCompletion(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (*openai.ChatCompletion, error) {
// BYOK: single attempt, no failover.
if i.cfg.KeyPool == nil {
return i.newChatCompletionWithKey(ctx, svc, opts)
}
return i.newChatCompletionWithKeyFailover(ctx, svc, opts)
}
// newChatCompletionWithKey performs a single upstream call.
func (i *BlockingInterception) newChatCompletionWithKey(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (_ *openai.ChatCompletion, outErr error) {
_, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
defer tracing.EndSpanErr(span, &outErr)
return svc.New(ctx, i.req.ChatCompletionNewParams, opts...)
}
// newChatCompletionWithKeyFailover walks the centralized key
// pool, trying each key until one succeeds or the pool is
// exhausted. Keys are marked temporary on 429 and permanent on
// 401/403. Errors that aren't key-specific don't trigger
// failover and are returned to the caller.
func (i *BlockingInterception) newChatCompletionWithKeyFailover(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (*openai.ChatCompletion, error) {
// TODO(ssncferreira): update the interception's credential
// hint with the actually-used key (the successful key on
// success, the last tried key on failure) in the upstack PR.
walker := i.cfg.KeyPool.Walker()
for {
key, err := walker.Next()
if err != nil {
return nil, err
}
requestOpts := append([]option.RequestOption{}, opts...)
requestOpts = append(requestOpts,
option.WithAPIKey(key.Value()),
// Disable SDK retries because the failover loop
// handles retries via key rotation.
option.WithMaxRetries(0),
)
completion, err := i.newChatCompletionWithKey(ctx, svc, requestOpts)
// Key-specific failure: try the next key.
if i.markKeyOnError(ctx, key, err) {
continue
}
// Either success (completion, nil) or a non-key error
// (nil, err): nothing to retry, return as-is.
return completion, err
}
}
aibridge/intercept/chatcompletions/blocking_test.go
+502
View File
@@ -0,0 +1,502 @@
package chatcompletions //nolint:testpackage // tests unexported internals
import (
"context"
"io"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"github.com/google/uuid"
mcplib "github.com/mark3labs/mcp-go/mcp"
"github.com/openai/openai-go/v3"
"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"
)
// OpenAI-shaped response bodies.
const (
successBody = `{"id":"chatcmpl-01","object":"chat.completion","created":1234567890,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"Hello!"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`
toolUseBody = `{"id":"chatcmpl-01","object":"chat.completion","created":1234567890,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call_01","type":"function","function":{"name":"test_tool","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`
textCompleteBody = `{"id":"chatcmpl-02","object":"chat.completion","created":1234567890,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"done"},"finish_reason":"stop"}],"usage":{"prompt_tokens":15,"completion_tokens":3,"total_tokens":18}}`
rateLimitBody = `{"error":{"message":"Rate limit exceeded","type":"rate_limit_error","code":"rate_limit_exceeded"}}`
authErrorBody = `{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}`
serverErrorBody = `{"error":{"message":"Internal server error","type":"server_error","code":"internal_error"}}`
)
type upstreamResponse struct {
statusCode int
body string
headers map[string]string
}
// newRequestParams builds a minimal chat-completions request
// for tests.
func newRequestParams(stream bool) *ChatCompletionNewParamsWrapper {
return &ChatCompletionNewParamsWrapper{
ChatCompletionNewParams: openai.ChatCompletionNewParams{
Model: "gpt-4",
Messages: []openai.ChatCompletionMessageParamUnion{
openai.UserMessage("hi"),
},
},
Stream: stream,
}
}
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 bearer token.
responses map[string]upstreamResponse
expectedRequestCount int32
expectedStatusCode int
expectedRetryAfter string
// Expected key states after the request, by index in keys.
expectedKeyStates []keypool.KeyState
}{
{
// Given: 1 valid key returning 200.
// Then: 1 request, 200 response, key remains valid.
name: "single_valid_key",
keys: []string{"k0"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid},
},
{
// 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", "k1"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// 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", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// 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", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusForbidden, body: authErrorBody},
"k1": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// 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", "k1", "k2"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "3"},
body: rateLimitBody,
},
"k2": {
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,
},
},
{
// Given: 2 keys; both return 401.
// Then: 2 requests, 502 api_error response, both keys permanent.
name: "all_keys_unauthorized",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {statusCode: http.StatusUnauthorized, body: authErrorBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusBadGateway,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStatePermanent,
},
},
{
// Given: 2 keys; key-0 returns 500.
// Then: 1 request, 500 response, both keys remain valid.
name: "server_error_no_failover",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusInternalServerError,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// 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",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
// Mock upstream: counts requests and returns
// scripted responses keyed by bearer token. 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[utils.ExtractBearerToken(r.Header.Get("Authorization"))]
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.OpenAI{BaseURL: upstream.URL + "/"}
var pool *keypool.Pool
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
}
interceptor := NewBlockingInterceptor(
uuid.New(),
newRequestParams(false),
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
otel.Tracer("blocking_test"),
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
)
interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil)
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", 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")
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
expectedKeyStates []keypool.KeyState
}{
{
// 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: textCompleteBody},
},
expectedRequestCount: 2,
expectedSeenKeys: []string{"k0", "k0"},
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// 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: textCompleteBody},
},
expectedRequestCount: 3,
expectedSeenKeys: []string{"k0", "k0", "k1"},
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// 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", "k0", "k1"},
expectedStatusCode: http.StatusTooManyRequests,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateTemporary,
},
},
}
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 bearer token 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, utils.ExtractBearerToken(r.Header.Get("Authorization")))
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", "k1"}, quartz.NewMock(t))
require.NoError(t, err)
cfg := config.OpenAI{
BaseURL: upstream.URL + "/",
KeyPool: pool,
}
interceptor := NewBlockingInterceptor(
uuid.New(),
newRequestParams(false),
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
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/chat/completions", 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")
seenKeysMu.Lock()
defer seenKeysMu.Unlock()
assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys")
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
})
}
}
// mockServerProxier is a test implementation of mcp.ServerProxier.
type mockServerProxier struct {
tools []*mcp.Tool
}
func (*mockServerProxier) Init(context.Context) error {
return nil
}
func (*mockServerProxier) Shutdown(context.Context) error {
return nil
}
func (m *mockServerProxier) ListTools() []*mcp.Tool {
return m.tools
}
func (m *mockServerProxier) GetTool(id string) *mcp.Tool {
for _, t := range m.tools {
if t.ID == id {
return t
}
}
return nil
}
func (*mockServerProxier) CallTool(context.Context, string, any) (*mcplib.CallToolResult, error) {
return nil, nil //nolint:nilnil // mock: no-op implementation
}
// stubToolCaller is a minimal mcp.ToolCaller that returns a fixed
// text result, so the agentic continuation can proceed.
type stubToolCaller struct{}
func (stubToolCaller) CallTool(_ context.Context, _ mcplib.CallToolRequest) (*mcplib.CallToolResult, error) {
return mcplib.NewToolResultText("tool result"), nil
}
aibridge/intercept/chatcompletions/streaming.go
+125 -33
View File
@@ -24,6 +24,7 @@ import (
aibcontext "github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/eventstream"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/mcp"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/tracing"
@@ -126,9 +127,49 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re
lastErr error
interceptionErr error
)
for {
// TODO add outer loop span (https://github.com/coder/aibridge/issues/67)
// Per-iteration walker. An iteration is either an agentic
// continuation (sending a tool result back in a new
// stream) or a failover retry (previous key marked, try
// the next one).
var walker *keypool.Walker
if i.cfg.KeyPool != nil {
walker = i.cfg.KeyPool.Walker()
}
var opts []option.RequestOption
var currentKey *keypool.Key
if walker != nil {
key, err := walker.Next()
if respErr := processKeyPoolError(err); respErr != nil {
// Pool exhausted in this iteration. Relay the
// error to the client: as an SSE event if events
// have already been sent, or by direct write
// otherwise.
interceptionErr = respErr
if events.IsStreaming() {
payload, mErr := i.marshalErr(respErr)
if mErr != nil {
logger.Warn(ctx, "failed to marshal exhaustion error", slog.Error(mErr))
} else if sErr := events.Send(streamCtx, payload); sErr != nil {
logger.Warn(ctx, "failed to relay exhaustion error", slog.Error(sErr))
}
} else {
i.writeUpstreamError(w, respErr)
}
break
}
currentKey = key
opts = append(opts,
option.WithAPIKey(key.Value()),
// Disable SDK retries because the failover
// loop handles retries via key rotation.
option.WithMaxRetries(0),
)
}
// TODO(ssncferreira): inject actor headers directly in the client-header
// middleware instead of using SDK options.
@@ -151,7 +192,17 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re
var toolCall *openai.FinishedChatCompletionToolCall
// iterationStarted is per-iteration (reset on every
// loop): true once the upstream call has produced any
// events for this iteration. While false, a key-specific
// failure can still fail over to the next key. Distinct
// from events.IsStreaming(), which is stream-wide and
// stays true once iteration 1 has sent any event
// downstream.
var iterationStarted bool
for stream.Next() {
iterationStarted = true
chunk := stream.Current()
canRelay := processor.process(chunk)
@@ -231,40 +282,52 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re
})
}
if !events.IsStreaming() {
// response/downstream Stream has not started yet; write error response and exit.
i.writeUpstreamError(w, getErrorResponse(stream.Err()))
return stream.Err()
}
// Check if the stream encountered any errors.
if streamErr := stream.Err(); streamErr != nil {
if eventstream.IsUnrecoverableError(streamErr) {
logger.Debug(ctx, "stream terminated", slog.Error(streamErr))
// We can't reflect an error back if there's a connection error or the request context was canceled.
} else if oaiErr := getErrorResponse(streamErr); oaiErr != nil {
logger.Warn(ctx, "openai stream error", slog.Error(streamErr))
interceptionErr = oaiErr
} else {
logger.Warn(ctx, "unknown stream error", slog.Error(streamErr))
// Unfortunately, the OpenAI SDK does not support parsing errors received in the stream
// into known types (i.e. [shared.OverloadedError]).
// See https://github.com/openai/openai-go/blob/v2.7.0/packages/ssestream/ssestream.go#L171
// All it does is wrap the payload in an error - which is all we can return, currently.
interceptionErr = newErrorResponse(xerrors.Errorf("unknown stream error: %w", streamErr))
if iterationStarted {
// Mid-stream error or logical error: events have
// already streamed for this iteration, so the
// error is relayed as an SSE event.
streamErr := stream.Err()
if respErr := i.mapStreamError(ctx, logger, streamErr, lastErr); respErr != nil {
interceptionErr = respErr
payload, err := i.marshalErr(respErr)
if err != nil {
logger.Warn(ctx, "failed to marshal error", slog.Error(err), slog.F("error_payload", fmt.Sprintf("%+v", respErr)))
} else if err := events.Send(streamCtx, payload); err != nil {
logger.Warn(ctx, "failed to relay error", slog.Error(err), slog.F("payload", payload))
}
} else if streamErr != nil {
// Unrecoverable (e.g., broken pipe, context
// canceled): can't relay to the client, but record
// the error so it isn't silently swallowed.
interceptionErr = streamErr
}
} else if lastErr != nil {
// Otherwise check if any logical errors occurred during processing.
logger.Warn(ctx, "stream processing failed", slog.Error(lastErr))
interceptionErr = newErrorResponse(xerrors.Errorf("processing error: %w", lastErr))
}
if interceptionErr != nil {
payload, err := i.marshalErr(interceptionErr)
if err != nil {
logger.Warn(ctx, "failed to marshal error", slog.Error(err), slog.F("error_payload", fmt.Sprintf("%+v", interceptionErr)))
} else if err := events.Send(streamCtx, payload); err != nil {
logger.Warn(ctx, "failed to relay error", slog.Error(err), slog.F("payload", payload))
} else {
// Pre-stream failure of this iteration. For
// centralized requests, mark the key and retry with
// the next one.
if currentKey != nil && i.markKeyOnError(ctx, currentKey, stream.Err()) {
continue
}
// Non-key error: relay it. Use mapStreamError so that
// unknown upstream errors (TCP reset, DNS failure, TLS
// error, deadline exceeded) are wrapped in a generic
// response instead of producing a silent HTTP 200.
respErr := i.mapStreamError(ctx, logger, stream.Err(), lastErr)
if respErr != nil {
interceptionErr = respErr
if events.IsStreaming() {
// Prior iterations have streamed, so the SSE
// connection is open: inject as an SSE event.
payload, mErr := i.marshalErr(respErr)
if mErr != nil {
logger.Warn(ctx, "failed to marshal error", slog.Error(mErr))
} else if sErr := events.Send(streamCtx, payload); sErr != nil {
logger.Warn(ctx, "failed to relay error", slog.Error(sErr))
}
} else {
// No events streamed yet, write the response directly.
i.writeUpstreamError(w, respErr)
}
}
}
@@ -407,6 +470,35 @@ func (i *StreamingInterception) newStream(ctx context.Context, svc openai.ChatCo
return svc.NewStreaming(ctx, openai.ChatCompletionNewParams{}, opts...)
}
// mapStreamError converts a mid-stream upstream error or
// processing error into a relayable responseError. Returns nil
// when the error is unrecoverable, in which case nothing can be
// relayed back.
func (*StreamingInterception) mapStreamError(ctx context.Context, logger slog.Logger, streamErr, lastErr error) *responseError {
if streamErr != nil {
if eventstream.IsUnrecoverableError(streamErr) {
logger.Debug(ctx, "stream terminated", slog.Error(streamErr))
// We can't reflect an error back if there's a connection error or the request context was canceled.
return nil
}
if oaiErr := getErrorResponse(streamErr); oaiErr != nil {
logger.Warn(ctx, "openai stream error", slog.Error(streamErr))
return oaiErr
}
logger.Warn(ctx, "unknown stream error", slog.Error(streamErr))
// Unfortunately, the OpenAI SDK does not support parsing errors received in the stream
// into known types (i.e. [shared.OverloadedError]).
// See https://github.com/openai/openai-go/blob/v2.7.0/packages/ssestream/ssestream.go#L171
// All it does is wrap the payload in an error - which is all we can return, currently.
return newErrorResponse(fmt.Sprintf("unknown stream error: %s", streamErr), intercept.OpenAIErrTypeError, intercept.OpenAIErrTypeError, http.StatusBadGateway, 0)
}
if lastErr != nil {
logger.Warn(ctx, "stream processing failed", slog.Error(lastErr))
return newErrorResponse(fmt.Sprintf("processing error: %s", lastErr), intercept.OpenAIErrTypeError, intercept.OpenAIErrTypeError, http.StatusBadGateway, 0)
}
return nil
}
type streamProcessor struct {
ctx context.Context
logger slog.Logger
aibridge/intercept/chatcompletions/streaming_test.go
+497 -8
View File
@@ -1,9 +1,12 @@
package chatcompletions_test
package chatcompletions //nolint:testpackage // tests unexported internals
import (
"io"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"sync"
"sync/atomic"
"testing"
"github.com/google/uuid"
@@ -16,8 +19,11 @@ import (
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/chatcompletions"
"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"
)
// Test that when the upstream provider returns an error before streaming starts,
@@ -36,21 +42,21 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) {
name: "bad request error",
statusCode: http.StatusBadRequest,
responseBody: `{"error":{"message":"Invalid request","type":"invalid_request_error","code":"invalid_request"}}`,
expectedErrStr: strconv.Itoa(http.StatusBadRequest),
expectedErrStr: "Invalid request",
expectedBody: "invalid_request",
},
{
name: "rate limit error",
statusCode: http.StatusTooManyRequests,
responseBody: `{"error":{"message":"Rate limit exceeded","type":"rate_limit_error","code":"rate_limit_exceeded"}}`,
expectedErrStr: strconv.Itoa(http.StatusTooManyRequests),
expectedErrStr: "Rate limit exceeded",
expectedBody: "rate_limit",
},
{
name: "internal server error",
statusCode: http.StatusInternalServerError,
responseBody: `{"error":{"message":"Internal server error","type":"server_error","code":"internal_error"}}`,
expectedErrStr: strconv.Itoa(http.StatusInternalServerError),
expectedErrStr: "Internal server error",
expectedBody: "server_error",
},
}
@@ -74,7 +80,7 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) {
Key: "test-key",
}
req := &chatcompletions.ChatCompletionNewParamsWrapper{
req := &ChatCompletionNewParamsWrapper{
ChatCompletionNewParams: openai.ChatCompletionNewParams{
Model: "gpt-4",
Messages: []openai.ChatCompletionMessageParamUnion{
@@ -89,7 +95,7 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) {
httpReq := httptest.NewRequest(http.MethodPost, "/chat/completions", nil)
tracer := otel.Tracer("test")
interceptor := chatcompletions.NewStreamingInterceptor(uuid.New(), req, config.ProviderOpenAI, cfg, httpReq.Header, "Authorization", tracer, intercept.CredentialInfo{})
interceptor := NewStreamingInterceptor(uuid.New(), req, config.ProviderOpenAI, cfg, httpReq.Header, "Authorization", tracer, intercept.CredentialInfo{})
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
interceptor.Setup(logger, &testutil.MockRecorder{}, nil)
@@ -110,3 +116,486 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) {
})
}
}
// OpenAI-shaped SSE body for a successful streaming response.
const streamingSuccessBody = `data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]}
data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}
data: [DONE]
`
func TestStreamingInterception_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 bearer token.
responses map[string]upstreamResponse
expectedRequestCount int32
expectedStatusCode int
expectedRetryAfter string
// Expected key states after the request, by index in keys.
expectedKeyStates []keypool.KeyState
}{
{
// Given: 1 valid key returning a successful stream.
// Then: 1 request, 200 response, key remains valid.
name: "single_valid_key",
keys: []string{"k0"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid},
},
{
// Given: 2 keys; key-0 returns 429 pre-stream, key-1
// streams successfully.
// Then: 2 requests, 200 response, key-0 temporary, key-1 valid.
name: "failover_after_429",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 returns 401 pre-stream, key-1
// streams successfully.
// Then: 2 requests, 200 response, key-0 permanent, key-1 valid.
name: "failover_after_401",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 returns 403 pre-stream, key-1 streams.
// Then: 2 requests, 200 response, key-0 permanent, key-1 valid.
name: "failover_after_403",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusForbidden, body: authErrorBody},
"k1": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// Given: 3 keys; all return 429 pre-stream 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", "k1", "k2"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "3"},
body: rateLimitBody,
},
"k2": {
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,
},
},
{
// Given: 2 keys; both return 401 pre-stream.
// Then: 2 requests, 502 api_error response, both keys permanent.
name: "all_keys_unauthorized",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {statusCode: http.StatusUnauthorized, body: authErrorBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusBadGateway,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStatePermanent,
},
},
{
// Given: 2 keys; key-0 returns 500 pre-stream.
// Then: 1 request, 500 response, both keys remain valid.
name: "server_error_no_failover",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusInternalServerError,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// 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",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
// Mock upstream: counts requests and returns
// scripted responses keyed by bearer token. 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[utils.ExtractBearerToken(r.Header.Get("Authorization"))]
if !ok {
resp = upstreamResponse{statusCode: http.StatusInternalServerError}
}
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.OpenAI{BaseURL: upstream.URL + "/"}
var pool *keypool.Pool
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
}
interceptor := NewStreamingInterceptor(
uuid.New(),
newRequestParams(true),
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
otel.Tracer("streaming_test"),
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
)
interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil)
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", 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")
if pool != nil {
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
}
})
}
}
// SSE bodies covering an agentic-continuation flow.
const (
// First response: a tool_calls delta referencing the
// injected "test_tool". Triggers the agentic continuation
// loop.
toolUseStreamBody = `data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_01","type":"function","function":{"name":"test_tool","arguments":""}}]},"finish_reason":null}]}
data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":null}]}
data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}
data: [DONE]
`
// Second response (after the tool result is sent back):
// a plain text completion that ends the loop.
textStreamBody = `data: {"id":"chatcmpl-02","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":"done"},"finish_reason":null}]}
data: {"id":"chatcmpl-02","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":15,"completion_tokens":3,"total_tokens":18}}
data: [DONE]
`
)
// TestStreamingInterception_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 TestStreamingInterception_AgenticLoopFailover(t *testing.T) {
t.Parallel()
sseHeaders := map[string]string{"Content-Type": "text/event-stream"}
tests := []struct {
name string
// Scripted upstream responses consumed in order of
// upstream request.
responses []upstreamResponse
expectedRequestCount int32
expectedSeenKeys []string
// Substring expected in the response body. Either a
// success marker (e.g. "done") or an error marker
// (e.g. "rate_limit_error").
expectedBodyContains string
// True when the error must be relayed as an SSE event.
expectErrorAsSSEEvent bool
// True when ProcessRequest is expected to return an
// error (e.g. all keys exhausted).
expectedErr bool
expectedKeyStates []keypool.KeyState
}{
{
// Given: 2 keys; both upstream calls succeed on key-0.
// Then: 2 requests, success body, both keys remain valid.
name: "happy_path",
responses: []upstreamResponse{
{statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody},
{statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody},
},
expectedRequestCount: 2,
expectedSeenKeys: []string{"k0", "k0"},
expectedBodyContains: "done",
expectErrorAsSSEEvent: false,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 succeeds initially, then 429s
// during the agentic continuation, key-1 succeeds.
// Then: 3 requests, success body, key-0 temporary,
// key-1 valid.
name: "agentic_failover_to_k1",
responses: []upstreamResponse{
{statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody},
{
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
{statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody},
},
expectedRequestCount: 3,
expectedSeenKeys: []string{"k0", "k0", "k1"},
expectedBodyContains: "done",
expectErrorAsSSEEvent: false,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 succeeds initially, then both
// keys 429 during the agentic continuation.
// Then: 3 requests, error injected as SSE event, both
// keys temporary.
name: "agentic_all_keys_fail",
responses: []upstreamResponse{
{statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody},
{
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", "k0", "k1"},
expectedBodyContains: "all configured keys are rate-limited",
expectErrorAsSSEEvent: true,
expectedErr: true,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateTemporary,
},
},
}
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 bearer token 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, utils.ExtractBearerToken(r.Header.Get("Authorization")))
seenKeysMu.Unlock()
_, _ = io.Copy(io.Discard, r.Body)
if idx >= len(tc.responses) {
w.WriteHeader(http.StatusInternalServerError)
return
}
resp := tc.responses[idx]
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", "k1"}, quartz.NewMock(t))
require.NoError(t, err)
cfg := config.OpenAI{
BaseURL: upstream.URL + "/",
KeyPool: pool,
}
interceptor := NewStreamingInterceptor(
uuid.New(),
newRequestParams(true),
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
otel.Tracer("streaming_test"),
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
)
// Mock proxy with a tool the upstream's tool_calls
// chunks will reference. The stub caller returns a
// fixed text result.
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/chat/completions", nil)
w := httptest.NewRecorder()
err = interceptor.ProcessRequest(w, req)
if tc.expectedErr {
require.Error(t, err)
} else {
require.NoError(t, err)
}
assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count")
body := w.Body.String()
assert.Contains(t, body, tc.expectedBodyContains, "response body")
if tc.expectErrorAsSSEEvent {
// SSE was opened before the failure, so the body
// must start with stream chunks, not a direct
// HTTP error body.
assert.True(t, strings.HasPrefix(body, "data: "), "body must start with SSE chunks")
}
seenKeysMu.Lock()
defer seenKeysMu.Unlock()
assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys")
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
})
}
}
aibridge/intercept/openai_errors.go
+14
View File
@@ -0,0 +1,14 @@
package intercept
// OpenAI error type and code constants used by the chatcompletions
// and responses interceptors. The OpenAI Go SDK does not expose
// these as typed constants, so we define our own.
// See https://platform.openai.com/docs/guides/error-codes.
const (
OpenAIErrTypeError = "error"
OpenAIErrTypeAPI = "api_error"
OpenAIErrTypeRateLimit = "rate_limit_error"
OpenAIErrCodeServer = "server_error"
OpenAIErrCodeRateLimit = "rate_limit_exceeded"
)
aibridge/intercept/responses/base.go
+119 -1
View File
@@ -4,7 +4,10 @@ import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"math"
"net/http"
"strconv"
"strings"
@@ -13,8 +16,10 @@ import (
"time"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
"github.com/openai/openai-go/v3/option"
"github.com/openai/openai-go/v3/responses"
"github.com/openai/openai-go/v3/shared"
"github.com/openai/openai-go/v3/shared/constant"
"github.com/tidwall/gjson"
"go.opentelemetry.io/otel/attribute"
@@ -26,6 +31,7 @@ import (
aibcontext "github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/apidump"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/mcp"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/tracing"
@@ -53,8 +59,19 @@ type responsesInterceptionBase struct {
credential intercept.CredentialInfo
}
// newResponsesService builds the SDK service used for upstream
// calls. BYOK auth is set here. Centralized auth is set
// per-attempt by the failover loop.
func (i *responsesInterceptionBase) newResponsesService() responses.ResponseService {
opts := []option.RequestOption{option.WithBaseURL(i.cfg.BaseURL), option.WithAPIKey(i.cfg.Key)}
// TODO(ssncferreira): validate auth is configured per
// https://github.com/coder/aibridge/issues/266.
var opts []option.RequestOption
// BYOK auth.
if i.cfg.KeyPool == nil {
opts = append(opts, option.WithAPIKey(i.cfg.Key))
}
opts = append(opts, option.WithBaseURL(i.cfg.BaseURL))
// Add extra headers if configured.
// Some providers require additional headers that are not added by the SDK.
@@ -124,6 +141,107 @@ func (i *responsesInterceptionBase) validateRequest(ctx context.Context, w http.
return nil
}
// writeUpstreamError marshals and writes a given error.
func (i *responsesInterceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *responseError) {
if oaiErr == nil {
return
}
w.Header().Set("Content-Type", "application/json")
// Set Retry-After when a cooldown is configured.
if oaiErr.RetryAfter > 0 {
w.Header().Set("Retry-After", strconv.Itoa(int(math.Ceil(oaiErr.RetryAfter.Seconds()))))
}
w.WriteHeader(oaiErr.StatusCode)
out, err := json.Marshal(oaiErr)
if err != nil {
i.logger.Warn(context.Background(), "failed to marshal upstream error", slog.Error(err), slog.F("error_payload", fmt.Sprintf("%+v", oaiErr)))
// Response has to match expected format.
_, _ = w.Write([]byte(`{
"error": {
"type": "error",
"message":"error marshaling upstream error",
"code": "server_error"
}
}`))
} else {
_, _ = w.Write(out)
}
}
// For centralized requests, markKeyOnError extracts an OpenAI
// SDK error from err and marks the key based on its status
// code. Returns true if the status was a key-specific failover
// trigger so callers can retry with the next key.
func (i *responsesInterceptionBase) markKeyOnError(ctx context.Context, key *keypool.Key, err error) bool {
if i.cfg.KeyPool == nil {
return false
}
var apiErr *openai.Error
if !errors.As(err, &apiErr) {
return false
}
return keypool.MarkKeyOnStatus(
ctx, key, apiErr.Response,
i.logger, i.providerName,
)
}
// processKeyPoolError translates a keypool exhaustion error
// into a developer-facing responseError shaped for the OpenAI
// API. Returns nil if err is not an exhaustion error.
func processKeyPoolError(err error) *responseError {
var transient *keypool.TransientKeyPoolError
switch {
case errors.As(err, &transient):
return newErrorResponse(
"all configured keys are rate-limited",
intercept.OpenAIErrTypeRateLimit,
intercept.OpenAIErrCodeRateLimit,
http.StatusTooManyRequests,
transient.RetryAfter,
)
case errors.Is(err, keypool.ErrPermanentKeyPool):
return newErrorResponse(
"all configured keys failed authentication",
intercept.OpenAIErrTypeAPI,
intercept.OpenAIErrCodeServer,
http.StatusBadGateway,
0,
)
default:
return nil
}
}
func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *responseError {
return &responseError{
ErrorObject: &shared.ErrorObject{
Code: code,
Message: msg,
Type: errType,
},
StatusCode: status,
RetryAfter: retryAfter,
}
}
var _ error = &responseError{}
type responseError struct {
ErrorObject *shared.ErrorObject `json:"error"`
StatusCode int `json:"-"`
RetryAfter time.Duration `json:"-"`
}
func (a *responseError) Error() string {
if a.ErrorObject == nil {
return ""
}
return a.ErrorObject.Message
}
// sendCustomErr sends custom responses.Error error to the client
// it should only be called before any data is sent back to the client
func (i *responsesInterceptionBase) sendCustomErr(ctx context.Context, w http.ResponseWriter, code int, err error) {
aibridge/intercept/responses/base_test.go
+198
View File
@@ -1,17 +1,25 @@
package responses //nolint:testpackage // tests unexported internals
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
oairesponses "github.com/openai/openai-go/v3/responses"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/internal/testutil"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/quartz"
)
func TestRecordPrompt(t *testing.T) {
@@ -382,3 +390,193 @@ func TestResponseCopierDoesntSendIfNoResponseReceived(t *testing.T) {
require.True(t, mrw.writeCalled)
require.True(t, mrw.writeHeaderCalled)
}
func TestProcessKeyPoolError(t *testing.T) {
t.Parallel()
tests := []struct {
name string
err error
expectedNil bool
expectedStatus int
expectedRetryAfter time.Duration
}{
{
// Transient with valid keys present: 429, no Retry-After.
name: "transient_zero_retry_after",
err: &keypool.TransientKeyPoolError{},
expectedStatus: http.StatusTooManyRequests,
expectedRetryAfter: 0,
},
{
// Transient with cooldown: 429, Retry-After set.
name: "transient_with_retry_after",
err: &keypool.TransientKeyPoolError{RetryAfter: 5 * time.Second},
expectedStatus: http.StatusTooManyRequests,
expectedRetryAfter: 5 * time.Second,
},
{
// Permanent: 502 api_error.
name: "permanent_returns_502",
err: keypool.ErrPermanentKeyPool,
expectedStatus: http.StatusBadGateway,
},
{
// Anything else: not a pool-exhaustion error.
name: "non_pool_exhaustion_error_returns_nil",
err: xerrors.New("some other error"),
expectedNil: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
got := processKeyPoolError(tc.err)
if tc.expectedNil {
require.Nil(t, got)
return
}
require.NotNil(t, got)
assert.Equal(t, tc.expectedStatus, got.StatusCode)
assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter)
})
}
}
func TestMarkKeyOnError(t *testing.T) {
t.Parallel()
tests := []struct {
name string
err error
expectedReturn bool
expectedState keypool.KeyState
}{
{
// Not an *openai.Error: no status code to act on.
name: "non_api_error_returns_false",
err: xerrors.New("network failure"),
expectedReturn: false,
expectedState: keypool.KeyStateValid,
},
{
// Rate-limited: temporary cooldown.
name: "429_marks_temporary",
err: &openai.Error{StatusCode: http.StatusTooManyRequests, Response: &http.Response{StatusCode: http.StatusTooManyRequests}},
expectedReturn: true,
expectedState: keypool.KeyStateTemporary,
},
{
// Auth failure: mark permanent.
name: "401_marks_permanent",
err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}},
expectedReturn: true,
expectedState: keypool.KeyStatePermanent,
},
{
// Auth forbidden: mark permanent.
name: "403_marks_permanent",
err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}},
expectedReturn: true,
expectedState: keypool.KeyStatePermanent,
},
{
// Server errors are not key-specific.
name: "500_does_not_mark",
err: &openai.Error{StatusCode: http.StatusInternalServerError, Response: &http.Response{StatusCode: http.StatusInternalServerError}},
expectedReturn: false,
expectedState: keypool.KeyStateValid,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
pool, err := keypool.New([]string{"key-0"}, quartz.NewMock(t))
require.NoError(t, err)
key, err := pool.Walker().Next()
require.NoError(t, err)
base := &responsesInterceptionBase{cfg: config.OpenAI{KeyPool: pool}, logger: slog.Make()}
got := base.markKeyOnError(context.Background(), key, tc.err)
assert.Equal(t, tc.expectedReturn, got)
assert.Equal(t, tc.expectedState, key.State())
})
}
}
func TestWriteUpstreamError(t *testing.T) {
t.Parallel()
tests := []struct {
name string
respErr *responseError
expectStatus int
// Empty string means the header should be absent.
expectRetryAfter string
// Substring expected in the marshaled body. Empty means no body check.
expectBodyContains string
}{
{
// Standard error: status, code, and JSON body written.
name: "writes_status_and_body",
respErr: newErrorResponse("upstream failed", "api_error", "server_error", http.StatusBadGateway, 0),
expectStatus: http.StatusBadGateway,
expectBodyContains: `"upstream failed"`,
},
{
// OpenAI envelope: the code field round-trips into the body.
name: "writes_code_field",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 0),
expectStatus: http.StatusTooManyRequests,
expectBodyContains: `"rate_limit_exceeded"`,
},
{
// Whole-second retryAfter: emitted as integer seconds.
name: "retry_after_in_seconds",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 60*time.Second),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "60",
},
{
// 500ms rounds up to Retry-After: 1.
name: "retry_after_500ms_rounds_up_to_one",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 500*time.Millisecond),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "1",
},
{
// 200ms rounds up to Retry-After: 1.
name: "retry_after_200ms_rounds_up_to_one",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 200*time.Millisecond),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "1",
},
{
// Negative retryAfter: header omitted.
name: "negative_retry_after_omits_header",
respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, -1*time.Second),
expectStatus: http.StatusTooManyRequests,
expectRetryAfter: "",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
base := &responsesInterceptionBase{logger: slog.Make()}
w := httptest.NewRecorder()
base.writeUpstreamError(w, tc.respErr)
assert.Equal(t, tc.expectStatus, w.Code, "status code")
assert.Equal(t, "application/json", w.Header().Get("Content-Type"), "Content-Type header")
assert.Equal(t, tc.expectRetryAfter, w.Header().Get("Retry-After"), "Retry-After header")
if tc.expectBodyContains != "" {
assert.Contains(t, w.Body.String(), tc.expectBodyContains, "response body")
}
})
}
}
aibridge/intercept/responses/blocking.go
+56 -2
View File
@@ -100,6 +100,15 @@ func (i *BlockingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r *
response, upstreamErr = i.newResponse(ctx, srv, opts)
// The failover loop may return a keypool exhaustion
// error. Render it here.
if upstreamErr != nil {
if keyErr := processKeyPoolError(upstreamErr); keyErr != nil {
i.writeUpstreamError(w, keyErr)
return xerrors.Errorf("key pool exhausted: %w", upstreamErr)
}
}
if upstreamErr != nil || response == nil {
break
}
@@ -135,10 +144,55 @@ func (i *BlockingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r *
return errors.Join(upstreamErr, err)
}
func (i *BlockingResponsesInterceptor) newResponse(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (_ *responses.Response, outErr error) {
ctx, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
// newResponse routes between BYOK (single attempt) and
// centralized failover.
func (i *BlockingResponsesInterceptor) newResponse(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (*responses.Response, error) {
// BYOK: single attempt, no failover.
if i.cfg.KeyPool == nil {
return i.newResponseWithKey(ctx, srv, opts)
}
return i.newResponseWithKeyFailover(ctx, srv, opts)
}
// newResponseWithKey performs a single upstream call.
func (i *BlockingResponsesInterceptor) newResponseWithKey(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (_ *responses.Response, outErr error) {
_, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
defer tracing.EndSpanErr(span, &outErr)
// The body is overridden by option.WithRequestBody(reqPayload) in requestOptions
return srv.New(ctx, responses.ResponseNewParams{}, opts...)
}
// newResponseWithKeyFailover walks the centralized key pool,
// trying each key until one succeeds or the pool is exhausted.
// Keys are marked temporary on 429 and permanent on 401/403.
// Errors that aren't key-specific don't trigger failover and
// are returned to the caller.
func (i *BlockingResponsesInterceptor) newResponseWithKeyFailover(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (*responses.Response, error) {
// TODO(ssncferreira): update the interception's credential
// hint with the actually-used key (the successful key on
// success, the last tried key on failure) in the upstack PR.
walker := i.cfg.KeyPool.Walker()
for {
key, err := walker.Next()
if err != nil {
return nil, err
}
requestOpts := append([]option.RequestOption{}, opts...)
requestOpts = append(requestOpts,
option.WithAPIKey(key.Value()),
// Disable SDK retries because the failover loop
// handles retries via key rotation.
option.WithMaxRetries(0),
)
response, err := i.newResponseWithKey(ctx, srv, requestOpts)
// Key-specific failure: try the next key.
if i.markKeyOnError(ctx, key, err) {
continue
}
// Either success (response, nil) or a non-key error
// (nil, err): nothing to retry, return as-is.
return response, err
}
}
aibridge/intercept/responses/blocking_test.go
+492
View File
@@ -0,0 +1,492 @@
package responses //nolint:testpackage // tests unexported internals
import (
"context"
"io"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"github.com/google/uuid"
mcplib "github.com/mark3labs/mcp-go/mcp"
"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"
)
// OpenAI Responses API request and response bodies.
const (
requestBody = `{"input":"hi","model":"gpt-4o-mini"}`
successBody = `{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_01","role":"assistant","content":[{"type":"output_text","text":"Hello!"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}`
toolUseBody = `{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"function_call","id":"fc_01","call_id":"call_01","name":"test_tool","arguments":"{}","status":"completed"}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}`
textCompleteBody = `{"id":"resp_02","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_02","role":"assistant","content":[{"type":"output_text","text":"done"}]}],"usage":{"input_tokens":15,"output_tokens":3,"total_tokens":18}}`
rateLimitBody = `{"error":{"message":"Rate limit exceeded","type":"rate_limit_error","code":"rate_limit_exceeded"}}`
authErrorBody = `{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}`
serverErrorBody = `{"error":{"message":"Internal server error","type":"server_error","code":"internal_error"}}`
)
type upstreamResponse struct {
statusCode int
body string
headers map[string]string
}
func TestBlockingResponsesInterceptor_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 bearer token.
responses map[string]upstreamResponse
expectedRequestCount int32
expectedStatusCode int
expectedRetryAfter string
// Expected key states after the request, by index in keys.
expectedKeyStates []keypool.KeyState
}{
{
// Given: 1 valid key returning 200.
// Then: 1 request, 200 response, key remains valid.
name: "single_valid_key",
keys: []string{"k0"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid},
},
{
// 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", "k1"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// 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", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// 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", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusForbidden, body: authErrorBody},
"k1": {statusCode: http.StatusOK, body: successBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// 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", "k1", "k2"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "3"},
body: rateLimitBody,
},
"k2": {
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,
},
},
{
// Given: 2 keys; both return 401.
// Then: 2 requests, 502 api_error response, both keys permanent.
name: "all_keys_unauthorized",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {statusCode: http.StatusUnauthorized, body: authErrorBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusBadGateway,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStatePermanent,
},
},
{
// Given: 2 keys; key-0 returns 500.
// Then: 1 request, 500 response, both keys remain valid.
name: "server_error_no_failover",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusInternalServerError,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// Given: BYOK with a single key returning 429.
// Then: 1 request, 429 response, no failover.
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,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
// Mock upstream: counts requests and returns
// scripted responses keyed by bearer token. 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[utils.ExtractBearerToken(r.Header.Get("Authorization"))]
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.OpenAI{BaseURL: upstream.URL + "/"}
var pool *keypool.Pool
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
}
payload, err := NewRequestPayload([]byte(requestBody))
require.NoError(t, err)
interceptor := NewBlockingInterceptor(
uuid.New(),
payload,
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
otel.Tracer("blocking_test"),
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
)
interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil)
req := httptest.NewRequest(http.MethodPost, "/v1/responses", 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")
if pool != nil {
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
}
})
}
}
// TestBlockingResponsesInterceptor_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 TestBlockingResponsesInterceptor_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
expectedKeyStates []keypool.KeyState
}{
{
// 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: textCompleteBody},
},
expectedRequestCount: 2,
expectedSeenKeys: []string{"k0", "k0"},
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// 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: textCompleteBody},
},
expectedRequestCount: 3,
expectedSeenKeys: []string{"k0", "k0", "k1"},
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// 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", "k0", "k1"},
expectedStatusCode: http.StatusTooManyRequests,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateTemporary,
},
},
}
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 bearer token 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, utils.ExtractBearerToken(r.Header.Get("Authorization")))
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", "k1"}, quartz.NewMock(t))
require.NoError(t, err)
cfg := config.OpenAI{
BaseURL: upstream.URL + "/",
KeyPool: pool,
}
payload, err := NewRequestPayload([]byte(requestBody))
require.NoError(t, err)
interceptor := NewBlockingInterceptor(
uuid.New(),
payload,
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
otel.Tracer("blocking_test"),
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
)
// Mock proxy with a tool the upstream's function_call
// 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/responses", 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")
seenKeysMu.Lock()
defer seenKeysMu.Unlock()
assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys")
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
})
}
}
// mockServerProxier is a test implementation of mcp.ServerProxier.
type mockServerProxier struct {
tools []*mcp.Tool
}
func (*mockServerProxier) Init(context.Context) error {
return nil
}
func (*mockServerProxier) Shutdown(context.Context) error {
return nil
}
func (m *mockServerProxier) ListTools() []*mcp.Tool {
return m.tools
}
func (m *mockServerProxier) GetTool(id string) *mcp.Tool {
for _, t := range m.tools {
if t.ID == id {
return t
}
}
return nil
}
func (*mockServerProxier) CallTool(context.Context, string, any) (*mcplib.CallToolResult, error) {
return nil, nil //nolint:nilnil // mock: no-op implementation
}
// stubToolCaller is a minimal mcp.ToolCaller that returns a fixed
// text result, so the agentic continuation can proceed.
type stubToolCaller struct{}
func (stubToolCaller) CallTool(_ context.Context, _ mcplib.CallToolRequest) (*mcplib.CallToolResult, error) {
return mcplib.NewToolResultText("tool result"), nil
}
aibridge/intercept/responses/streaming.go
+64 -12
View File
@@ -20,6 +20,7 @@ import (
aibcontext "github.com/coder/coder/v2/aibridge/context"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/eventstream"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/mcp"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/tracing"
@@ -108,36 +109,87 @@ func (i *StreamingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r
for shouldLoop {
shouldLoop = false
respCopy = responseCopier{}
opts := i.requestOptions(&respCopy)
// Per-iteration walker. An iteration is either an agentic
// continuation (sending a tool result back in a new
// stream) or a failover retry (previous key marked, try
// the next one).
var walker *keypool.Walker
if i.cfg.KeyPool != nil {
walker = i.cfg.KeyPool.Walker()
}
// TODO(ssncferreira): inject actor headers directly in the client-header
// middleware instead of using SDK options.
if actor := aibcontext.ActorFromContext(r.Context()); actor != nil && i.cfg.SendActorHeaders {
opts = append(opts, intercept.ActorHeadersAsOpenAIOpts(actor)...)
// Failover sub-loop: try keys until a stream starts
// successfully or we hit a non-recoverable error.
var stream *ssestream.Stream[responses.ResponseStreamEventUnion]
var startErr error
for {
respCopy = responseCopier{}
opts := i.requestOptions(&respCopy)
// TODO(ssncferreira): inject actor headers directly in the client-header
// middleware instead of using SDK options.
if actor := aibcontext.ActorFromContext(r.Context()); actor != nil && i.cfg.SendActorHeaders {
opts = append(opts, intercept.ActorHeadersAsOpenAIOpts(actor)...)
}
var currentKey *keypool.Key
if walker != nil {
key, err := walker.Next()
if respErr := processKeyPoolError(err); respErr != nil {
// Pool exhausted: write the error directly. In
// agentic mode the inner loop buffers events
// instead of streaming them downstream, so the
// SSE connection has not been opened yet.
i.writeUpstreamError(w, respErr)
return xerrors.Errorf("key pool exhausted: %w", err)
}
currentKey = key
opts = append(opts,
option.WithAPIKey(key.Value()),
// Disable SDK retries because the failover
// loop handles retries via key rotation.
option.WithMaxRetries(0),
)
}
stream = i.newStream(ctx, srv, opts)
if upstreamErr := stream.Err(); upstreamErr != nil {
// Pre-stream failure of this attempt. For
// centralized requests, mark the key and
// retry with the next one.
if currentKey != nil && i.markKeyOnError(ctx, currentKey, upstreamErr) {
stream.Close()
continue
}
// Non-key error: stop trying and let the
// existing handling below report it.
startErr = upstreamErr
break
}
// Stream started successfully: commit to this key.
break
}
stream := i.newStream(ctx, srv, opts)
// func scope to defer steam.Close()
err := func() error {
defer stream.Close()
if upstreamErr := stream.Err(); upstreamErr != nil {
if startErr != nil {
// events stream should never be initialized
if events.IsStreaming() {
i.logger.Warn(ctx, "event stream was initialized when no response was received from upstream")
return upstreamErr
return startErr
}
// no response received from upstream (eg. client/connection error), return custom error
if !respCopy.responseReceived.Load() {
i.sendCustomErr(ctx, w, http.StatusInternalServerError, upstreamErr)
return upstreamErr
i.sendCustomErr(ctx, w, http.StatusInternalServerError, startErr)
return startErr
}
// forward received response as-is
err := respCopy.forwardResp(w)
return errors.Join(upstreamErr, err)
return errors.Join(startErr, err)
}
for stream.Next() {
aibridge/intercept/responses/streaming_test.go
+499
View File
@@ -0,0 +1,499 @@
package responses //nolint:testpackage // tests unexported internals
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"
)
// Streaming request body for the OpenAI Responses API.
const streamingRequestBody = `{"input":"hi","model":"gpt-4o-mini","stream":true}`
// OpenAI Responses API SSE body for a successful streaming response.
const streamingSuccessBody = `event: response.created
data: {"type":"response.created","response":{"id":"resp_01","object":"response","status":"in_progress"},"sequence_number":0}
event: response.completed
data: {"type":"response.completed","response":{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_01","role":"assistant","content":[{"type":"output_text","text":"Hello!"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}},"sequence_number":1}
`
func TestStreamingResponsesInterceptor_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 bearer token.
responses map[string]upstreamResponse
expectedRequestCount int32
expectedStatusCode int
expectedRetryAfter string
// Expected key states after the request, by index in keys.
expectedKeyStates []keypool.KeyState
}{
{
// Given: 1 valid key returning a successful stream.
// Then: 1 request, 200 response, key remains valid.
name: "single_valid_key",
keys: []string{"k0"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid},
},
{
// Given: 2 keys; key-0 returns 429 pre-stream, key-1
// streams successfully.
// Then: 2 requests, 200 response, key-0 temporary, key-1 valid.
name: "failover_after_429",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 returns 401 pre-stream, key-1
// streams successfully.
// Then: 2 requests, 200 response, key-0 permanent, key-1 valid.
name: "failover_after_401",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 returns 403 pre-stream, key-1 streams.
// Then: 2 requests, 200 response, key-0 permanent, key-1 valid.
name: "failover_after_403",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusForbidden, body: authErrorBody},
"k1": {
statusCode: http.StatusOK,
headers: map[string]string{"Content-Type": "text/event-stream"},
body: streamingSuccessBody,
},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusOK,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStateValid,
},
},
{
// Given: 3 keys; all return 429 pre-stream 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", "k1", "k2"},
responses: map[string]upstreamResponse{
"k0": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
"k1": {
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "3"},
body: rateLimitBody,
},
"k2": {
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,
},
},
{
// Given: 2 keys; both return 401 pre-stream.
// Then: 2 requests, 502 api_error response, both keys permanent.
name: "all_keys_unauthorized",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusUnauthorized, body: authErrorBody},
"k1": {statusCode: http.StatusUnauthorized, body: authErrorBody},
},
expectedRequestCount: 2,
expectedStatusCode: http.StatusBadGateway,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStatePermanent,
keypool.KeyStatePermanent,
},
},
{
// Given: 2 keys; key-0 returns 500 pre-stream.
// Then: 1 request, 500 response, both keys remain valid.
name: "server_error_no_failover",
keys: []string{"k0", "k1"},
responses: map[string]upstreamResponse{
"k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody},
},
expectedRequestCount: 1,
expectedStatusCode: http.StatusInternalServerError,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// Given: BYOK with a single key returning 429.
// Then: 1 request, 429 response, no failover.
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,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
// Mock upstream: counts requests and returns
// scripted responses keyed by bearer token. 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[utils.ExtractBearerToken(r.Header.Get("Authorization"))]
if !ok {
resp = upstreamResponse{statusCode: http.StatusInternalServerError}
}
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.OpenAI{BaseURL: upstream.URL + "/"}
var pool *keypool.Pool
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
}
payload, err := NewRequestPayload([]byte(streamingRequestBody))
require.NoError(t, err)
interceptor := NewStreamingInterceptor(
uuid.New(),
payload,
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
otel.Tracer("streaming_test"),
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
)
interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil)
req := httptest.NewRequest(http.MethodPost, "/v1/responses", 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")
if pool != nil {
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
}
})
}
}
// SSE bodies covering an agentic-continuation flow.
const (
// First response: a function_call output referencing the
// injected "test_tool". Triggers the agentic continuation
// loop.
toolUseStreamBody = `event: response.created
data: {"type":"response.created","response":{"id":"resp_01","object":"response","status":"in_progress"},"sequence_number":0}
event: response.completed
data: {"type":"response.completed","response":{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"function_call","id":"fc_01","call_id":"call_01","name":"test_tool","arguments":"{}","status":"completed"}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}},"sequence_number":1}
`
// Second response (after the tool result is sent back):
// a plain text message that ends the loop.
textStreamBody = `event: response.created
data: {"type":"response.created","response":{"id":"resp_02","object":"response","status":"in_progress"},"sequence_number":0}
event: response.completed
data: {"type":"response.completed","response":{"id":"resp_02","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_02","role":"assistant","content":[{"type":"output_text","text":"done"}]}],"usage":{"input_tokens":15,"output_tokens":3,"total_tokens":18}},"sequence_number":1}
`
)
// TestStreamingResponsesInterceptor_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 TestStreamingResponsesInterceptor_AgenticLoopFailover(t *testing.T) {
t.Parallel()
sseHeaders := map[string]string{"Content-Type": "text/event-stream"}
tests := []struct {
name string
// Scripted upstream responses consumed in order of
// upstream request.
responses []upstreamResponse
expectedRequestCount int32
expectedSeenKeys []string
// Substring expected in the response body. Either a
// success marker (e.g. "done") or an error marker
// (e.g. "rate_limit_error").
expectedBodyContains string
// True when ProcessRequest is expected to return an
// error (e.g. all keys exhausted).
expectedErr bool
expectedKeyStates []keypool.KeyState
}{
{
// Given: 2 keys; both upstream calls succeed on key-0.
// Then: 2 requests, success body, both keys remain valid.
name: "happy_path",
responses: []upstreamResponse{
{statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody},
{statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody},
},
expectedRequestCount: 2,
expectedSeenKeys: []string{"k0", "k0"},
expectedBodyContains: "done",
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateValid,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 succeeds initially, then 429s
// during the agentic continuation, key-1 succeeds.
// Then: 3 requests, success body, key-0 temporary,
// key-1 valid.
name: "agentic_failover_to_k1",
responses: []upstreamResponse{
{statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody},
{
statusCode: http.StatusTooManyRequests,
headers: map[string]string{"Retry-After": "5"},
body: rateLimitBody,
},
{statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody},
},
expectedRequestCount: 3,
expectedSeenKeys: []string{"k0", "k0", "k1"},
expectedBodyContains: "done",
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
},
},
{
// Given: 2 keys; key-0 succeeds initially, then both
// keys 429 during the agentic continuation.
// Then: 3 requests, error injected as SSE event, both
// keys temporary.
name: "agentic_all_keys_fail",
responses: []upstreamResponse{
{statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody},
{
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", "k0", "k1"},
expectedBodyContains: "all configured keys are rate-limited",
expectedErr: true,
expectedKeyStates: []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateTemporary,
},
},
}
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 bearer token 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, utils.ExtractBearerToken(r.Header.Get("Authorization")))
seenKeysMu.Unlock()
_, _ = io.Copy(io.Discard, r.Body)
if idx >= len(tc.responses) {
w.WriteHeader(http.StatusInternalServerError)
return
}
resp := tc.responses[idx]
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", "k1"}, quartz.NewMock(t))
require.NoError(t, err)
cfg := config.OpenAI{
BaseURL: upstream.URL + "/",
KeyPool: pool,
}
payload, err := NewRequestPayload([]byte(streamingRequestBody))
require.NoError(t, err)
interceptor := NewStreamingInterceptor(
uuid.New(),
payload,
config.ProviderOpenAI,
cfg,
http.Header{},
"Authorization",
otel.Tracer("streaming_test"),
intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""),
)
// Mock proxy with a tool the upstream's function_call
// response will reference. The stub caller returns a
// fixed text result.
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/responses", nil)
w := httptest.NewRecorder()
err = interceptor.ProcessRequest(w, req)
if tc.expectedErr {
require.Error(t, err)
} else {
require.NoError(t, err)
}
assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count")
body := w.Body.String()
assert.Contains(t, body, tc.expectedBodyContains, "response body")
seenKeysMu.Lock()
defer seenKeysMu.Unlock()
assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys")
assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states")
})
}
}
aibridge/internal/integrationtest/keypool_failover_test.go
+132
View File
@@ -17,9 +17,141 @@ import (
"github.com/coder/coder/v2/aibridge/fixtures"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/provider"
"github.com/coder/coder/v2/aibridge/utils"
"github.com/coder/quartz"
)
// TestOpenAI_KeyFailover verifies that a pool's key state
// persists across distinct client requests for both OpenAI APIs
// (chat completions and responses), in both blocking and
// streaming modes. A key marked temporary on request 1 is
// skipped on request 2 without a wasted upstream attempt.
func TestOpenAI_KeyFailover(t *testing.T) {
t.Parallel()
tests := []struct {
name string
fixture []byte
path string
streaming bool
successCType string
}{
{
name: "chatcompletions_blocking",
fixture: fixtures.OaiChatSimple,
path: pathOpenAIChatCompletions,
streaming: false,
successCType: "application/json",
},
{
name: "chatcompletions_streaming",
fixture: fixtures.OaiChatSimple,
path: pathOpenAIChatCompletions,
streaming: true,
successCType: "text/event-stream",
},
{
name: "responses_blocking",
fixture: fixtures.OaiResponsesBlockingSimple,
path: pathOpenAIResponses,
streaming: false,
successCType: "application/json",
},
{
name: "responses_streaming",
fixture: fixtures.OaiResponsesStreamingSimple,
path: pathOpenAIResponses,
streaming: true,
successCType: "text/event-stream",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
fix := fixtures.Parse(t, tc.fixture)
var successBody []byte
if tc.streaming {
successBody = fix.Streaming()
} else {
successBody = fix.NonStreaming()
}
pool, err := keypool.New([]string{"k0", "k1"}, quartz.NewMock(t))
require.NoError(t, err)
var requestCount atomic.Int32
var seenKeysMu sync.Mutex
var seenKeys []string
// Mock upstream: k0 always returns 429, k1 returns
// the per-test success body.
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestCount.Add(1)
key := utils.ExtractBearerToken(r.Header.Get("Authorization"))
seenKeysMu.Lock()
seenKeys = append(seenKeys, key)
seenKeysMu.Unlock()
_, _ = io.Copy(io.Discard, r.Body)
switch key {
case "k0":
w.Header().Set("Content-Type", "application/json")
w.Header().Set("Retry-After", "60")
w.WriteHeader(http.StatusTooManyRequests)
_, _ = fmt.Fprint(w, `{"error":{"type":"rate_limit_error","message":"rate limited","code":"rate_limit_exceeded"}}`)
case "k1":
w.Header().Set("Content-Type", tc.successCType)
w.WriteHeader(http.StatusOK)
_, _ = w.Write(successBody)
default:
w.WriteHeader(http.StatusInternalServerError)
}
}))
t.Cleanup(upstream.Close)
bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL,
withCustomProvider(provider.NewOpenAI(config.OpenAI{
BaseURL: upstream.URL,
KeyPool: pool,
})),
)
requestBody, err := sjson.SetBytes(fix.Request(), "stream", tc.streaming)
require.NoError(t, err)
// Request 1: walker starts at k0, fails over to k1
// after 429.
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, requestBody)
require.NoError(t, err)
_, _ = io.Copy(io.Discard, resp.Body)
require.NoError(t, resp.Body.Close())
require.Equal(t, http.StatusOK, resp.StatusCode)
// Request 2: walker skips the now-temporary k0 and
// goes straight to k1 (1 upstream call, not 2).
resp, err = bridgeServer.makeRequest(t, http.MethodPost, tc.path, requestBody)
require.NoError(t, err)
_, _ = io.Copy(io.Discard, resp.Body)
require.NoError(t, resp.Body.Close())
require.Equal(t, http.StatusOK, resp.StatusCode)
seenKeysMu.Lock()
defer seenKeysMu.Unlock()
// Request 1: 2 calls (k0 then k1). Request 2: 1 call (k1).
assert.Equal(t, int32(3), requestCount.Load(), "upstream request count")
assert.Equal(t, []string{"k0", "k1", "k1"}, seenKeys, "seen keys")
// Pool state persists: k0 temporary, k1 valid.
assert.Equal(t, []keypool.KeyState{
keypool.KeyStateTemporary,
keypool.KeyStateValid,
}, pool.PoolState(), "key states")
})
}
}
// TestAnthropic_KeyFailover verifies that a pool's key state
// persists across distinct client requests: a key marked
// temporary on request 1 is still skipped on request 2 without
aibridge/internal/integrationtest/trace_test.go
+5 -6
View File
@@ -667,13 +667,12 @@ func TestTraceOpenAIErr(t *testing.T) {
},
},
{
name: "trace_openai_responses_streaming_http_error",
fixture: fixtures.OaiResponsesStreamingHTTPErr,
streaming: true,
allowOverflow: true, // 429 error causes retries
name: "trace_openai_responses_streaming_http_error",
fixture: fixtures.OaiResponsesStreamingHTTPErr,
streaming: true,
path: pathOpenAIResponses,
expectCode: http.StatusTooManyRequests,
expectCode: http.StatusBadRequest,
expect: []expectTrace{
{"Intercept", 1, codes.Error},
{"Intercept.CreateInterceptor", 1, codes.Unset},
@@ -689,7 +688,7 @@ func TestTraceOpenAIErr(t *testing.T) {
streaming: false,
path: pathOpenAIResponses,
expectCode: http.StatusUnauthorized,
expectCode: http.StatusBadRequest,
expect: []expectTrace{
{"Intercept", 1, codes.Error},
{"Intercept.CreateInterceptor", 1, codes.Unset},
aibridge/provider/openai.go
+51 -7
View File
@@ -16,8 +16,10 @@ import (
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/chatcompletions"
"github.com/coder/coder/v2/aibridge/intercept/responses"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/tracing"
"github.com/coder/coder/v2/aibridge/utils"
"github.com/coder/quartz"
)
const (
@@ -44,6 +46,25 @@ func NewOpenAI(cfg config.OpenAI) *OpenAI {
if cfg.BaseURL == "" {
cfg.BaseURL = "https://api.openai.com/v1/"
}
// Resolve centralized key configuration into KeyPool.
// Precedence:
// 1. cfg.KeyPool (explicit, highest priority).
// 2. cfg.Key (legacy single key).
// After this block cfg.Key is empty so it can only carry a
// BYOK Authorization Bearer set per interception in
// CreateInterceptor.
// TODO(ssncferreira): simplify auth field resolution per
// https://github.com/coder/aibridge/issues/266.
if cfg.KeyPool == nil && cfg.Key != "" {
// keypool.New only fails on empty or duplicate keys,
// neither possible with a single non-empty key.
pool, err := keypool.New([]string{cfg.Key}, quartz.NewReal())
if err != nil {
panic(fmt.Sprintf("openai provider: build single-key pool: %s", err))
}
cfg.KeyPool = pool
}
cfg.Key = ""
if cfg.CircuitBreaker != nil {
cfg.CircuitBreaker.OpenErrorResponse = openAIOpenErrorResponse
}
@@ -100,20 +121,34 @@ func (p *OpenAI) CreateInterceptor(_ http.ResponseWriter, r *http.Request, trace
var interceptor intercept.Interceptor
cfg := p.cfg
// At this point the request contains only LLM provider headers. Any
// Coder-specific authentication has already been stripped.
// At this point the request contains only LLM provider headers.
// Any Coder-specific authentication has already been stripped.
//
// In centralized mode Authorization is absent, so cfg keeps the
// centralized key unchanged.
// KeyPool from provider construction and the failover loop walks
// it.
//
// In BYOK mode the user's credential is in Authorization. Replace
// the centralized key with it so it is forwarded upstream.
// In BYOK mode the user's credential is in Authorization,
// populate cfg.Key and clear cfg.KeyPool so failover is disabled.
//
// TODO(ssncferreira): consolidate auth field handling per
// https://github.com/coder/aibridge/issues/266.
credKind := intercept.CredentialKindCentralized
var credSecret string
if token := utils.ExtractBearerToken(r.Header.Get("Authorization")); token != "" {
cfg.Key = token
cfg.KeyPool = nil
credKind = intercept.CredentialKindBYOK
credSecret = token
} else if cfg.KeyPool != nil {
// Centralized: use the first key as a placeholder hint.
// TODO(ssncferreira): record the actually-used key in
// the interception record to reflect failover.
if k, err := cfg.KeyPool.Walker().Next(); err == nil {
credSecret = k.Value()
}
}
cred := intercept.NewCredentialInfo(credKind, cfg.Key)
cred := intercept.NewCredentialInfo(credKind, credSecret)
path := strings.TrimPrefix(r.URL.Path, p.RoutePrefix())
switch path {
@@ -171,7 +206,16 @@ func (p *OpenAI) InjectAuthHeader(headers *http.Header) {
return
}
headers.Set(p.AuthHeader(), "Bearer "+p.cfg.Key)
// Centralized: pull a single key from the pool. No failover
// or exhaustion handling here.
// TODO(ssncferreira): replace with RoundTripper-based auth
// in the upstack passthrough PR.
if p.cfg.KeyPool == nil {
return
}
if key, err := p.cfg.KeyPool.Walker().Next(); err == nil {
headers.Set(p.AuthHeader(), "Bearer "+key.Value())
}
}
func (p *OpenAI) CircuitBreakerConfig() *config.CircuitBreaker {
enterprise/cli/aibridged.go
+7 -4
View File
@@ -122,15 +122,18 @@ func buildProviders(cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) {
}
switch p.Type {
case aibridge.ProviderOpenAI:
// TODO(ssncferreira): pass a keypool.Pool instead.
var key string
var pool *keypool.Pool
if len(p.Keys) > 0 {
key = p.Keys[0]
var err error
pool, err = keypool.New(p.Keys, quartz.NewReal())
if err != nil {
return nil, xerrors.Errorf("create openai key pool for provider %q: %w", name, err)
}
}
providers = append(providers, aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
Name: name,
BaseURL: p.BaseURL,
Key: key,
KeyPool: pool,
APIDumpDir: p.DumpDir,
CircuitBreaker: cbConfig,
SendActorHeaders: cfg.SendActorHeaders.Value(),