mirror of
https://github.com/coder/coder.git
synced 2026-10-28 02:42:54 +08:00
Michael Suchacz
c8e8b21a88
## Stack Context PR 5 of 6 in a stack that migrates every Coder MCP surface from the archived `github.com/mark3labs/mcp-go` library to the official `github.com/modelcontextprotocol/go-sdk` v1.7.0. Stack: #28056 -> #28057 -> #28058 -> #28059 -> #28060 -> #28061 ## Why The aibridge injected-MCP proxy now owns an official `*mcp.Client`, `*mcp.StreamableClientTransport`, and `*mcp.ClientSession`. - The proxy constructor accepts an optional `*http.Client` instead of mark3labs options; the header-injecting wrapper shallow-copies a supplied client so its Timeout, Jar, and redirect policy survive. - Manual protocol version negotiation and the mark3labs five-second close workaround are removed; the SDK negotiates during `Connect` and fails when no mutually supported version exists. - Repeated `Init` closes the previous session, and a failed tool fetch closes the just-created session so transports do not leak. - Tool and intercept types use the official pointer content types; embedded resource blobs are re-encoded to base64 for model-facing text because the SDK decodes them into raw bytes. - `aibridge/mcpmock` is regenerated, and its stale `go:generate` source path is corrected. > Mux created this PR on Mike's behalf.
652 lines
24 KiB
Go
652 lines
24 KiB
Go
package enterprise_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"slices"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
promtest "github.com/prometheus/client_golang/prometheus/testutil"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
"go.opentelemetry.io/otel"
|
|
"go.opentelemetry.io/otel/attribute"
|
|
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
|
"go.opentelemetry.io/otel/sdk/trace/tracetest"
|
|
"golang.org/x/sync/singleflight"
|
|
|
|
"github.com/coder/coder/v2/aibridge"
|
|
"github.com/coder/coder/v2/aibridge/aibridgetest"
|
|
"github.com/coder/coder/v2/aibridge/config"
|
|
"github.com/coder/coder/v2/aibridge/keypool"
|
|
aibtracing "github.com/coder/coder/v2/aibridge/tracing"
|
|
"github.com/coder/coder/v2/coderd/aibridged"
|
|
"github.com/coder/coder/v2/coderd/aibridgedserver"
|
|
"github.com/coder/coder/v2/coderd/coderdtest"
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
|
"github.com/coder/coder/v2/coderd/database/dbtime"
|
|
"github.com/coder/coder/v2/coderd/externalauth"
|
|
"github.com/coder/coder/v2/coderd/httpmw"
|
|
"github.com/coder/coder/v2/codersdk"
|
|
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
|
|
"github.com/coder/coder/v2/testutil"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
// singleKeyPool builds a centralized key pool containing a single key.
|
|
func singleKeyPool(t *testing.T, name, key string) *keypool.Pool {
|
|
t.Helper()
|
|
pool, err := keypool.New(name, []string{key}, quartz.NewReal(), nil)
|
|
require.NoError(t, err)
|
|
return pool
|
|
}
|
|
|
|
var testTracer = otel.Tracer("aibridged_inttest")
|
|
|
|
// TestIntegration is not an exhaustive test against the upstream AI providers' SDKs (see coder/aibridge for those).
|
|
// This test validates that:
|
|
// - intercepted requests can be authenticated/authorized
|
|
// - requests can be routed to an appropriate handler
|
|
// - responses can be returned as expected
|
|
// - interceptions are logged, as well as their related prompt, token, and tool calls
|
|
// - MCP server configurations are returned as expected
|
|
// - tracing spans are properly recorded
|
|
func TestIntegration(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
|
|
sr := tracetest.NewSpanRecorder()
|
|
tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))
|
|
tracer := tp.Tracer(t.Name())
|
|
defer func() { _ = tp.Shutdown(t.Context()) }()
|
|
|
|
var mcpTokenReceived atomic.Pointer[string]
|
|
mcpHandler := sdkmcp.NewStreamableHTTPHandler(func(*http.Request) *sdkmcp.Server {
|
|
return sdkmcp.NewServer(&sdkmcp.Implementation{
|
|
Name: "test-mcp-server",
|
|
Version: "1.0.0",
|
|
}, nil)
|
|
}, &sdkmcp.StreamableHTTPOptions{Stateless: true})
|
|
mockMCPServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
t.Logf("Mock MCP server received request: %s %s", r.Method, r.URL.Path)
|
|
if auth := r.Header.Get("Authorization"); auth != "" {
|
|
mcpTokenReceived.Store(&auth)
|
|
}
|
|
mcpHandler.ServeHTTP(w, r)
|
|
}))
|
|
t.Cleanup(mockMCPServer.Close)
|
|
t.Logf("Mock MCP server running at: %s", mockMCPServer.URL)
|
|
|
|
// Set up mock OpenAI server that returns a tool call response.
|
|
mockOpenAI := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{
|
|
"id": "chatcmpl-BwkyFElDIr1egmFyfQ9z4vPBto7m2",
|
|
"object": "chat.completion",
|
|
"created": 1753343279,
|
|
"model": "gpt-4.1-2025-04-14",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": null,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_KjzAbhiZC6nk81tQzL7pwlpc",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "read_file",
|
|
"arguments": "{\"path\":\"README.md\"}"
|
|
}
|
|
}
|
|
],
|
|
"refusal": null,
|
|
"annotations": []
|
|
},
|
|
"logprobs": null,
|
|
"finish_reason": "tool_calls"
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 60,
|
|
"completion_tokens": 15,
|
|
"total_tokens": 75,
|
|
"prompt_tokens_details": {
|
|
"cached_tokens": 15,
|
|
"audio_tokens": 0
|
|
},
|
|
"completion_tokens_details": {
|
|
"reasoning_tokens": 0,
|
|
"audio_tokens": 0,
|
|
"accepted_prediction_tokens": 0,
|
|
"rejected_prediction_tokens": 0
|
|
}
|
|
},
|
|
"service_tier": "default",
|
|
"system_fingerprint": "fp_b3f1157249"
|
|
}`))
|
|
}))
|
|
t.Cleanup(mockOpenAI.Close)
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
client, _, api, firstUser := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
|
|
Options: &coderdtest.Options{
|
|
Database: db,
|
|
Pubsub: ps,
|
|
ExternalAuthConfigs: []*externalauth.Config{
|
|
{
|
|
InstrumentedOAuth2Config: &testutil.OAuth2Config{},
|
|
ID: "mock",
|
|
Type: "mock",
|
|
DisplayName: "Mock",
|
|
MCPURL: mockMCPServer.URL,
|
|
RefreshGroup: new(singleflight.Group),
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
userClient, user := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID)
|
|
|
|
// Create an API token for the user.
|
|
apiKey, err := userClient.CreateToken(ctx, "me", codersdk.CreateTokenRequest{
|
|
TokenName: fmt.Sprintf("test-key-%d", time.Now().UnixNano()),
|
|
Lifetime: time.Hour,
|
|
Scope: codersdk.APIKeyScopeAll,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create external auth link for the user.
|
|
authLink, err := db.InsertExternalAuthLink(ctx, database.InsertExternalAuthLinkParams{
|
|
ProviderID: "mock",
|
|
UserID: user.ID,
|
|
CreatedAt: dbtime.Now(),
|
|
UpdatedAt: dbtime.Now(),
|
|
OAuthAccessToken: "test-mock-token",
|
|
OAuthRefreshToken: "test-refresh-token",
|
|
OAuthExpiry: dbtime.Now().Add(time.Hour),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create aibridge server & client.
|
|
aiBridgeClient, err := api.AGPL.CreateInMemoryAIBridgeServer(ctx)
|
|
require.NoError(t, err)
|
|
|
|
logger := testutil.Logger(t)
|
|
providers := []aibridge.Provider{aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL: mockOpenAI.URL, KeyPool: singleKeyPool(t, config.ProviderOpenAI, "test-centralized-key")})}
|
|
pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, tracer)
|
|
require.NoError(t, err)
|
|
|
|
// Given: aibridged is started.
|
|
srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
|
|
return aiBridgeClient, nil
|
|
}, logger, tracer)
|
|
require.NoError(t, err, "create new aibridged")
|
|
t.Cleanup(func() {
|
|
_ = srv.Shutdown(ctx)
|
|
})
|
|
|
|
// When: a request is made to aibridged.
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "/openai/v1/chat/completions", bytes.NewBufferString(`{
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "how large is the README.md file in my current path"
|
|
}
|
|
],
|
|
"model": "gpt-4.1",
|
|
"tools": [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "read_file",
|
|
"description": "Read the contents of a file at the given path.",
|
|
"parameters": {
|
|
"properties": {
|
|
"path": {
|
|
"type": "string"
|
|
}
|
|
},
|
|
"required": [
|
|
"path"
|
|
],
|
|
"type": "object"
|
|
}
|
|
}
|
|
}
|
|
]
|
|
}`))
|
|
userAgent := "codex_cli_rs/0.87.0"
|
|
require.NoError(t, err, "make request to test server")
|
|
req.Header.Add("Authorization", "Bearer "+apiKey.Key)
|
|
req.Header.Add("Accept", "application/json")
|
|
req.Header.Add("User-Agent", userAgent)
|
|
|
|
// When: aibridged handles the request.
|
|
rec := httptest.NewRecorder()
|
|
srv.ServeHTTP(rec, req)
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
// Then: the interception & related records are stored.
|
|
interceptions, err := db.GetAIBridgeInterceptions(ctx)
|
|
require.NoError(t, err)
|
|
require.Len(t, interceptions, 1)
|
|
|
|
intc0 := interceptions[0]
|
|
keyID, _, err := httpmw.SplitAPIToken(apiKey.Key)
|
|
require.NoError(t, err)
|
|
require.Equal(t, user.ID, intc0.InitiatorID)
|
|
require.True(t, intc0.APIKeyID.Valid)
|
|
require.Equal(t, keyID, intc0.APIKeyID.String)
|
|
require.Equal(t, "openai", intc0.Provider)
|
|
require.Equal(t, "gpt-4.1", intc0.Model)
|
|
require.True(t, intc0.EndedAt.Valid)
|
|
require.False(t, intc0.EndedAt.Time.Before(intc0.StartedAt), "EndedAt should not be before StartedAt")
|
|
require.Less(t, intc0.EndedAt.Time.Sub(intc0.StartedAt), 5*time.Second)
|
|
require.True(t, intc0.Client.Valid)
|
|
require.Equal(t, string(aibridge.ClientCodex), intc0.Client.String)
|
|
require.Equal(t, database.CredentialKindCentralized, intc0.CredentialKind)
|
|
require.Equal(t, "test...-key", intc0.CredentialHint)
|
|
|
|
intc0Metadata := gjson.GetBytes(intc0.Metadata.RawMessage, aibridgedserver.MetadataUserAgentKey)
|
|
require.Equal(t, userAgent, intc0Metadata.String(), "interception metadata user agent should match request user agent")
|
|
|
|
prompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, interceptions[0].ID)
|
|
require.NoError(t, err)
|
|
require.Len(t, prompts, 1)
|
|
require.Equal(t, prompts[0].Prompt, "how large is the README.md file in my current path")
|
|
|
|
tokens, err := db.GetAIBridgeTokenUsagesByInterceptionID(ctx, interceptions[0].ID)
|
|
require.NoError(t, err)
|
|
require.Len(t, tokens, 1)
|
|
require.EqualValues(t, tokens[0].InputTokens, 45)
|
|
require.EqualValues(t, tokens[0].OutputTokens, 15)
|
|
require.EqualValues(t, 15, tokens[0].CacheReadInputTokens)
|
|
|
|
tools, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, interceptions[0].ID)
|
|
require.NoError(t, err)
|
|
require.Len(t, tools, 1)
|
|
require.False(t, tools[0].Injected)
|
|
|
|
// Then: the MCP server was initialized.
|
|
gotMCPToken := mcpTokenReceived.Load()
|
|
require.NotNil(t, gotMCPToken, "mock MCP server not requested")
|
|
require.Contains(t, *gotMCPToken, authLink.OAuthAccessToken)
|
|
|
|
// Then: verify tracing spans were recorded.
|
|
spans := sr.Ended()
|
|
require.NotEmpty(t, spans)
|
|
i := slices.IndexFunc(spans, func(s sdktrace.ReadOnlySpan) bool { return s.Name() == "CachedBridgePool.Acquire" })
|
|
require.NotEqual(t, -1, i, "span named 'CachedBridgePool.Acquire' not found")
|
|
|
|
expectAttrs := []attribute.KeyValue{
|
|
attribute.String(aibtracing.InitiatorID, user.ID.String()),
|
|
attribute.String(aibtracing.APIKeyID, keyID),
|
|
}
|
|
require.Equal(t, spans[i].Attributes(), expectAttrs)
|
|
|
|
// Check for aibridge spans.
|
|
spanNames := make(map[string]bool)
|
|
for _, span := range spans {
|
|
spanNames[span.Name()] = true
|
|
}
|
|
|
|
expectedAibridgeSpans := []string{
|
|
"CachedBridgePool.Acquire",
|
|
"ServerProxyManager.Init",
|
|
"StreamableHTTPServerProxy.Init",
|
|
"StreamableHTTPServerProxy.Init.fetchTools",
|
|
"Intercept",
|
|
"Intercept.CreateInterceptor",
|
|
"Intercept.RecordInterception",
|
|
"Intercept.ProcessRequest",
|
|
"Intercept.ProcessRequest.Upstream",
|
|
"Intercept.RecordPromptUsage",
|
|
"Intercept.RecordTokenUsage",
|
|
"Intercept.RecordToolUsage",
|
|
"Intercept.RecordInterceptionEnded",
|
|
}
|
|
|
|
for _, expectedSpan := range expectedAibridgeSpans {
|
|
require.Contains(t, spanNames, expectedSpan)
|
|
}
|
|
}
|
|
|
|
// TestIntegrationWithMetrics validates that Prometheus metrics are correctly incremented
|
|
// when requests are processed through aibridged.
|
|
func TestIntegrationWithMetrics(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
|
|
// Create prometheus registry and metrics.
|
|
registry := prometheus.NewRegistry()
|
|
metrics := aibridge.NewMetrics(registry)
|
|
|
|
// Set up mock OpenAI server.
|
|
mockOpenAI := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte(`{
|
|
"id": "chatcmpl-test",
|
|
"object": "chat.completion",
|
|
"created": 1753343279,
|
|
"model": "gpt-4.1",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "test response"
|
|
},
|
|
"finish_reason": "stop"
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"total_tokens": 15
|
|
}
|
|
}`))
|
|
}))
|
|
t.Cleanup(mockOpenAI.Close)
|
|
|
|
// Database and coderd setup.
|
|
db, ps := dbtestutil.NewDB(t)
|
|
client, _, api, firstUser := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
|
|
Options: &coderdtest.Options{
|
|
Database: db,
|
|
Pubsub: ps,
|
|
},
|
|
})
|
|
|
|
userClient, _ := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID)
|
|
|
|
// Create an API token for the user.
|
|
apiKey, err := userClient.CreateToken(ctx, "me", codersdk.CreateTokenRequest{
|
|
TokenName: fmt.Sprintf("test-key-%d", time.Now().UnixNano()),
|
|
Lifetime: time.Hour,
|
|
Scope: codersdk.APIKeyScopeCoderAll,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create aibridge client.
|
|
aiBridgeClient, err := api.AGPL.CreateInMemoryAIBridgeServer(ctx)
|
|
require.NoError(t, err)
|
|
|
|
logger := testutil.Logger(t)
|
|
providers := []aibridge.Provider{aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL: mockOpenAI.URL, KeyPool: singleKeyPool(t, config.ProviderOpenAI, "test-centralized-key")})}
|
|
|
|
// Create pool with metrics.
|
|
pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, metrics, testTracer)
|
|
require.NoError(t, err)
|
|
|
|
// Given: aibridged is started.
|
|
srv, err := aibridged.New(ctx, pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
|
|
return aiBridgeClient, nil
|
|
}, logger, testTracer)
|
|
require.NoError(t, err, "create new aibridged")
|
|
t.Cleanup(func() {
|
|
_ = srv.Shutdown(ctx)
|
|
})
|
|
|
|
// When: a request is made to aibridged.
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "/openai/v1/chat/completions", bytes.NewBufferString(`{
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "test message"
|
|
}
|
|
],
|
|
"model": "gpt-4.1"
|
|
}`))
|
|
require.NoError(t, err, "make request to test server")
|
|
req.Header.Add("Authorization", "Bearer "+apiKey.Key)
|
|
req.Header.Add("Accept", "application/json")
|
|
|
|
// When: aibridged handles the request.
|
|
rec := httptest.NewRecorder()
|
|
srv.ServeHTTP(rec, req)
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
// Then: the interceptions metric should increase to 1.
|
|
// This is not exhaustively checking the available metrics; just an indicative one to prove
|
|
// the plumbing is working.
|
|
require.Eventually(t, func() bool {
|
|
count := promtest.ToFloat64(metrics.InterceptionCount)
|
|
return count == 1
|
|
}, testutil.WaitShort, testutil.IntervalFast, "interceptions_total metric should be 1")
|
|
}
|
|
|
|
// TestIntegrationCircuitBreaker validates that the circuit breaker opens after
|
|
// consecutive failures and that the corresponding metrics are exposed.
|
|
func TestIntegrationCircuitBreaker(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
|
|
// Create prometheus registry and metrics.
|
|
registry := prometheus.NewRegistry()
|
|
metrics := aibridge.NewMetrics(registry)
|
|
|
|
// Set up mock OpenAI server that always returns 503 Service Unavailable.
|
|
mockOpenAI := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
// Disable SDK retries.
|
|
w.Header().Set("x-should-retry", "false")
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
_, _ = w.Write([]byte(`{"error":{"message":"Service Unavailable.","type":"cf_service_unavailable","code":503}}`))
|
|
}))
|
|
t.Cleanup(mockOpenAI.Close)
|
|
|
|
// Set up mock Anthropic server that always returns 529 Overloaded.
|
|
mockAnthropic := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
// Anthropic uses 529 for overloaded errors.
|
|
w.WriteHeader(529)
|
|
_, _ = w.Write([]byte(`{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`))
|
|
}))
|
|
t.Cleanup(mockAnthropic.Close)
|
|
|
|
// Database and coderd setup.
|
|
db, ps := dbtestutil.NewDB(t)
|
|
client, _, api, firstUser := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
|
|
Options: &coderdtest.Options{
|
|
Database: db,
|
|
Pubsub: ps,
|
|
},
|
|
})
|
|
|
|
userClient, _ := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID)
|
|
|
|
// Create an API token for the user.
|
|
apiKey, err := userClient.CreateToken(ctx, "me", codersdk.CreateTokenRequest{
|
|
TokenName: fmt.Sprintf("test-key-%d", time.Now().UnixNano()),
|
|
Lifetime: time.Hour,
|
|
Scope: codersdk.APIKeyScopeCoderAll,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create aibridge client.
|
|
aiBridgeClient, err := api.AGPL.CreateInMemoryAIBridgeServer(ctx)
|
|
require.NoError(t, err)
|
|
|
|
logger := testutil.Logger(t)
|
|
|
|
// Create providers with circuit breaker configured to open after 2 failures.
|
|
cbConfig := &config.CircuitBreaker{
|
|
FailureThreshold: 2,
|
|
Interval: time.Minute,
|
|
Timeout: time.Minute,
|
|
MaxRequests: 1,
|
|
}
|
|
providers := []aibridge.Provider{
|
|
aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
|
|
BaseURL: mockOpenAI.URL,
|
|
KeyPool: singleKeyPool(t, config.ProviderOpenAI, "test-key"),
|
|
CircuitBreaker: cbConfig,
|
|
}),
|
|
aibridgetest.NewAnthropicProvider(t, aibridge.AnthropicConfig{
|
|
BaseURL: mockAnthropic.URL,
|
|
KeyPool: singleKeyPool(t, config.ProviderAnthropic, "test-key"),
|
|
CircuitBreaker: cbConfig,
|
|
}, nil),
|
|
}
|
|
|
|
// Create pool with metrics.
|
|
pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, metrics, testTracer)
|
|
require.NoError(t, err)
|
|
|
|
// Given: aibridged is started.
|
|
srv, err := aibridged.New(ctx, pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
|
|
return aiBridgeClient, nil
|
|
}, logger, testTracer)
|
|
require.NoError(t, err, "create new aibridged")
|
|
t.Cleanup(func() {
|
|
_ = srv.Shutdown(ctx)
|
|
})
|
|
|
|
// Test OpenAI circuit breaker.
|
|
openaiRequestBody := `{"messages":[{"role":"user","content":"test"}],"model":"gpt-4"}`
|
|
for i := 0; i < 3; i++ {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "/openai/v1/chat/completions", bytes.NewBufferString(openaiRequestBody))
|
|
require.NoError(t, err)
|
|
req.Header.Add("Authorization", "Bearer "+apiKey.Key)
|
|
req.Header.Add("Accept", "application/json")
|
|
|
|
rec := httptest.NewRecorder()
|
|
srv.ServeHTTP(rec, req)
|
|
t.Logf("OpenAI request %d: status=%d", i+1, rec.Code)
|
|
}
|
|
|
|
// Test Anthropic circuit breaker.
|
|
anthropicRequestBody := `{"messages":[{"role":"user","content":"test"}],"model":"claude-3-5-sonnet-20241022","max_tokens":100}`
|
|
for i := 0; i < 3; i++ {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, "/anthropic/v1/messages", bytes.NewBufferString(anthropicRequestBody))
|
|
require.NoError(t, err)
|
|
req.Header.Add("Authorization", "Bearer "+apiKey.Key)
|
|
req.Header.Add("Accept", "application/json")
|
|
|
|
rec := httptest.NewRecorder()
|
|
srv.ServeHTTP(rec, req)
|
|
t.Logf("Anthropic request %d: status=%d", i+1, rec.Code)
|
|
}
|
|
|
|
// Then: the circuit breaker metrics should reflect that both circuits opened.
|
|
|
|
// OpenAI circuit breaker should have tripped (state=1 means open).
|
|
openaiTrips := promtest.ToFloat64(metrics.CircuitBreakerTrips.WithLabelValues("openai", "/v1/chat/completions", "gpt-4"))
|
|
require.Equal(t, 1.0, openaiTrips, "OpenAI CircuitBreakerTrips should be 1")
|
|
|
|
openaiState := promtest.ToFloat64(metrics.CircuitBreakerState.WithLabelValues("openai", "/v1/chat/completions", "gpt-4"))
|
|
require.Equal(t, 1.0, openaiState, "OpenAI CircuitBreakerState should be 1 (open)")
|
|
|
|
// Anthropic circuit breaker should have tripped.
|
|
anthropicTrips := promtest.ToFloat64(metrics.CircuitBreakerTrips.WithLabelValues("anthropic", "/v1/messages", "claude-3-5-sonnet-20241022"))
|
|
require.Equal(t, 1.0, anthropicTrips, "Anthropic CircuitBreakerTrips should be 1")
|
|
|
|
anthropicState := promtest.ToFloat64(metrics.CircuitBreakerState.WithLabelValues("anthropic", "/v1/messages", "claude-3-5-sonnet-20241022"))
|
|
require.Equal(t, 1.0, anthropicState, "Anthropic CircuitBreakerState should be 1 (open)")
|
|
}
|
|
|
|
// TestIntegrationRecordsUpstreamError validates that a failed interception
|
|
// persists a categorized upstream error on its database row.
|
|
func TestIntegrationRecordsUpstreamError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
|
|
// Mock upstreams that reject with 401 Unauthorized.
|
|
unauthorizedHandler := func(body string) http.HandlerFunc {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.Header().Set("x-should-retry", "false")
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
_, _ = w.Write([]byte(body))
|
|
}
|
|
}
|
|
mockOpenAI := httptest.NewServer(unauthorizedHandler(`{"error":{"message":"Incorrect API key","type":"invalid_request_error","code":"invalid_api_key"}}`))
|
|
t.Cleanup(mockOpenAI.Close)
|
|
mockAnthropic := httptest.NewServer(unauthorizedHandler(`{"type":"error","error":{"type":"authentication_error","message":"invalid x-api-key"}}`))
|
|
t.Cleanup(mockAnthropic.Close)
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
client, _, api, firstUser := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
|
|
Options: &coderdtest.Options{Database: db, Pubsub: ps},
|
|
})
|
|
|
|
userClient, _ := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID)
|
|
apiKey, err := userClient.CreateToken(ctx, "me", codersdk.CreateTokenRequest{
|
|
TokenName: fmt.Sprintf("test-key-%d", time.Now().UnixNano()),
|
|
Lifetime: time.Hour,
|
|
Scope: codersdk.APIKeyScopeCoderAll,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
aiBridgeClient, err := api.AGPL.CreateInMemoryAIBridgeServer(ctx)
|
|
require.NoError(t, err)
|
|
|
|
logger := testutil.Logger(t)
|
|
providers := []aibridge.Provider{
|
|
aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{BaseURL: mockOpenAI.URL, KeyPool: singleKeyPool(t, config.ProviderOpenAI, "test-key")}),
|
|
aibridgetest.NewAnthropicProvider(t, aibridge.AnthropicConfig{BaseURL: mockAnthropic.URL, KeyPool: singleKeyPool(t, config.ProviderAnthropic, "test-key")}, nil),
|
|
}
|
|
pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, testTracer)
|
|
require.NoError(t, err)
|
|
|
|
srv, err := aibridged.New(ctx, pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
|
|
return aiBridgeClient, nil
|
|
}, logger, testTracer)
|
|
require.NoError(t, err, "create new aibridged")
|
|
t.Cleanup(func() { _ = srv.Shutdown(ctx) })
|
|
|
|
cases := []struct {
|
|
name string
|
|
path string
|
|
body string
|
|
}{
|
|
{"openai_blocking", "/openai/v1/chat/completions", `{"messages":[{"role":"user","content":"test"}],"model":"gpt-4","stream":false}`},
|
|
{"openai_streaming", "/openai/v1/chat/completions", `{"messages":[{"role":"user","content":"test"}],"model":"gpt-4","stream":true}`},
|
|
{"anthropic_blocking", "/anthropic/v1/messages", `{"messages":[{"role":"user","content":"test"}],"model":"claude-3-5-sonnet-20241022","max_tokens":100,"stream":false}`},
|
|
{"anthropic_streaming", "/anthropic/v1/messages", `{"messages":[{"role":"user","content":"test"}],"model":"claude-3-5-sonnet-20241022","max_tokens":100,"stream":true}`},
|
|
}
|
|
|
|
// Requests are fired sequentially: they share the centralized key pool and
|
|
// database, and the aggregate assertion below depends on all of them.
|
|
for _, tc := range cases {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tc.path, bytes.NewBufferString(tc.body))
|
|
require.NoError(t, err)
|
|
req.Header.Add("Authorization", "Bearer "+apiKey.Key)
|
|
req.Header.Add("Accept", "application/json")
|
|
|
|
rec := httptest.NewRecorder()
|
|
srv.ServeHTTP(rec, req)
|
|
t.Logf("%s request: status=%d", tc.name, rec.Code)
|
|
}
|
|
|
|
// The interception rows should record the unauthorized upstream failure.
|
|
// Both blocking and streaming interceptors surface the centralized key-pool
|
|
// permanent exhaustion, which categorizes as unauthorized.
|
|
intcs, err := db.GetAIBridgeInterceptions(ctx)
|
|
require.NoError(t, err)
|
|
require.Len(t, intcs, len(cases))
|
|
for _, intc := range intcs {
|
|
require.True(t, intc.EndedAt.Valid, "interception %s should be ended", intc.ID)
|
|
require.True(t, intc.ErrorType.Valid, "interception %s should record an error type", intc.ID)
|
|
require.Equal(t, database.AibridgeInterceptionErrorTypeUnauthorized, intc.ErrorType.AIBridgeInterceptionErrorType)
|
|
require.True(t, intc.ErrorMessage.Valid, "interception %s should record an error message", intc.ID)
|
|
require.NotEmpty(t, intc.ErrorMessage.String)
|
|
}
|
|
}
|