opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
b6033aee03025b105c31eca47f64126342e97082
coder/aibridge/internal/integrationtest/bridge_internal_test.go
T
Susana Ferreira f440cbd205 refactor(aibridge): move shared mock helpers to testutil (#25999) 
Moves the shared aibridge mock test helpers (`MockUpstream`, `MockServerProxier`, `StubToolCaller`, and the `NewFixtureResponse`/`NewFixtureToolResponse` constructors) out of `aibridge/internal/integrationtest` into `aibridge/internal/testutil`, and exports them.

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

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

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
2026-06-09 10:24:33 +01:00

2381 lines
86 KiB
Go
Raw Blame History

package integrationtest
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"slices"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/packages/ssestream"
"github.com/anthropics/anthropic-sdk-go/shared/constant"
"github.com/aws/aws-sdk-go-v2/aws"
v4signer "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
oaissestream "github.com/openai/openai-go/v3/packages/ssestream"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
"go.uber.org/goleak"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/fixtures"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/internal/testutil"
"github.com/coder/coder/v2/aibridge/mcp"
"github.com/coder/coder/v2/aibridge/provider"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/aibridge/utils"
)
func TestMain(m *testing.M) {
goleak.VerifyTestMain(m)
}
func TestAnthropicMessages(t *testing.T) {
t.Parallel()
t.Run("single builtin tool", func(t *testing.T) {
t.Parallel()
cases := []struct {
name string
streaming bool
expectedInputTokens int
expectedOutputTokens int
expectedCacheReadInputTokens int
expectedCacheWriteInputTokens int
expectedToolCallID string
}{
{
name: "streaming",
streaming: true,
expectedInputTokens: 2,
expectedOutputTokens: 66,
expectedCacheReadInputTokens: 13993,
expectedCacheWriteInputTokens: 22,
expectedToolCallID: "toolu_01RX68weRSquLx6HUTj65iBo",
},
{
name: "non-streaming",
streaming: false,
expectedInputTokens: 5,
expectedOutputTokens: 84,
expectedCacheReadInputTokens: 23490,
expectedCacheWriteInputTokens: 0,
expectedToolCallID: "toolu_01AusGgY5aKFhzWrFBv9JfHq",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, fixtures.AntSingleBuiltinTool)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
// Make API call to aibridge for Anthropic /v1/messages
reqBody, err := sjson.SetBytes(fix.Request(), "stream", tc.streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Response-specific checks.
if tc.streaming {
sp := aibridge.NewSSEParser()
require.NoError(t, sp.Parse(resp.Body))
// Ensure the message starts and completes, at a minimum.
assert.Contains(t, sp.AllEvents(), "message_start")
assert.Contains(t, sp.AllEvents(), "message_stop")
}
expectedTokenRecordings := 1
if tc.streaming {
// One for message_start, one for message_delta.
expectedTokenRecordings = 2
}
tokenUsages := bridgeServer.Recorder.RecordedTokenUsages()
require.Len(t, tokenUsages, expectedTokenRecordings)
assert.EqualValues(t, tc.expectedInputTokens, bridgeServer.Recorder.TotalInputTokens(), "input tokens miscalculated")
assert.EqualValues(t, tc.expectedOutputTokens, bridgeServer.Recorder.TotalOutputTokens(), "output tokens miscalculated")
assert.EqualValues(t, tc.expectedCacheReadInputTokens, bridgeServer.Recorder.TotalCacheReadInputTokens(), "cache read input tokens miscalculated")
assert.EqualValues(t, tc.expectedCacheWriteInputTokens, bridgeServer.Recorder.TotalCacheWriteInputTokens(), "cache write input tokens miscalculated")
toolUsages := bridgeServer.Recorder.RecordedToolUsages()
require.Len(t, toolUsages, 1)
assert.Equal(t, "Read", toolUsages[0].Tool)
assert.Equal(t, tc.expectedToolCallID, toolUsages[0].ToolCallID)
require.IsType(t, json.RawMessage{}, toolUsages[0].Args)
var args map[string]any
require.NoError(t, json.Unmarshal(toolUsages[0].Args.(json.RawMessage), &args))
require.Contains(t, args, "file_path")
assert.Equal(t, "/tmp/blah/foo", args["file_path"])
promptUsages := bridgeServer.Recorder.RecordedPromptUsages()
require.Len(t, promptUsages, 1)
assert.Equal(t, "read the foo file", promptUsages[0].Prompt)
// Verify PRM attribution is NOT present on non-Bedrock Anthropic requests.
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
ua := received[0].Header.Get("User-Agent")
assert.NotContains(t, ua, "sdk-ua-app-id",
"PRM attribution should not be present on non-Bedrock requests")
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
// When the upstream's first response is an injected tool call with no
// text preamble and the next upstream call fails, the response must
// remain a well-formed SSE stream. The upstream error is relayed as a
// well-formed SSE event.
t.Run("streaming injected tool call no preamble with upstream 500", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, fixtures.AntSingleInjectedToolNoPreamble)
upstream := testutil.NewMockUpstream(ctx, t,
testutil.NewFixtureResponse(fix),
testutil.NewErrorResponse(http.StatusInternalServerError, ""),
)
mockMCP := setupMCPForTest(t, defaultTracer)
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withMCP(mockMCP))
reqBody, err := sjson.SetBytes(fix.Request(), "stream", true)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
require.Equal(t, "text/event-stream", resp.Header.Get("Content-Type"))
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
bodyStr := string(body)
// Once iteration 1 succeeded the response is committed as SSE,
// so the iteration-2 error MUST be an SSE event and not a raw JSON body.
require.Contains(t, bodyStr, "event: error",
"iteration-2 error must be relayed as an SSE event")
// Tool was invoked despite the iteration-2 failure.
require.Len(t, mockMCP.getCallsByTool(mockToolName), 1,
"expected MCP tool to be invoked exactly once")
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
func TestAnthropicMessagesModelThoughts(t *testing.T) {
t.Parallel()
cases := []struct {
name string
streaming bool
fixture []byte
expectedThoughts []recorder.ModelThoughtRecord // nil means no model thoughts expected
}{
{
name: "single thinking block/streaming",
streaming: true,
fixture: fixtures.AntSingleBuiltinTool,
expectedThoughts: []recorder.ModelThoughtRecord{newModelThought("The user wants me to read", recorder.ThoughtSourceThinking)},
},
{
name: "single thinking block/blocking",
streaming: false,
fixture: fixtures.AntSingleBuiltinTool,
expectedThoughts: []recorder.ModelThoughtRecord{newModelThought("The user wants me to read", recorder.ThoughtSourceThinking)},
},
{
name: "multiple thinking blocks/streaming",
streaming: true,
fixture: fixtures.AntMultiThinkingBuiltinTool,
expectedThoughts: []recorder.ModelThoughtRecord{
newModelThought("The user wants me to read", recorder.ThoughtSourceThinking),
newModelThought("I should use the Read tool", recorder.ThoughtSourceThinking),
},
},
{
name: "multiple thinking blocks/blocking",
streaming: false,
fixture: fixtures.AntMultiThinkingBuiltinTool,
expectedThoughts: []recorder.ModelThoughtRecord{
newModelThought("The user wants me to read", recorder.ThoughtSourceThinking),
newModelThought("I should use the Read tool", recorder.ThoughtSourceThinking),
},
},
{
name: "parallel tool calls/streaming",
streaming: true,
fixture: fixtures.AntSingleBuiltinToolParallel,
expectedThoughts: []recorder.ModelThoughtRecord{newModelThought("The user wants me to read two files", recorder.ThoughtSourceThinking)},
},
{
name: "parallel tool calls/blocking",
streaming: false,
fixture: fixtures.AntSingleBuiltinToolParallel,
expectedThoughts: []recorder.ModelThoughtRecord{newModelThought("The user wants me to read two files", recorder.ThoughtSourceThinking)},
},
{
name: "thoughts without tool calls/streaming",
streaming: true,
fixture: fixtures.AntSimple,
expectedThoughts: []recorder.ModelThoughtRecord{newModelThought("This is a classic philosophical question about medieval scholasticism", recorder.ThoughtSourceThinking)},
},
{
name: "thoughts without tool calls/blocking",
streaming: false,
fixture: fixtures.AntSimple,
expectedThoughts: []recorder.ModelThoughtRecord{newModelThought("This is a classic philosophical question about medieval scholasticism", recorder.ThoughtSourceThinking)},
},
{
name: "no thoughts captured",
streaming: false,
fixture: fixtures.AntSingleInjectedTool,
expectedThoughts: nil,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
reqBody, err := sjson.SetBytes(fix.Request(), "stream", tc.streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
if tc.streaming {
sp := aibridge.NewSSEParser()
require.NoError(t, sp.Parse(resp.Body))
assert.Contains(t, sp.AllEvents(), "message_start")
assert.Contains(t, sp.AllEvents(), "message_stop")
}
bridgeServer.Recorder.VerifyModelThoughtsRecorded(t, tc.expectedThoughts)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
}
func TestAWSBedrockIntegration(t *testing.T) {
t.Parallel()
t.Run("invalid config", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
// Invalid bedrock config - missing region & base url
bedrockCfg := &config.AWSBedrock{
Region: "",
AccessKey: "test-key",
AccessKeySecret: "test-secret",
Model: "test-model",
SmallFastModel: "test-haiku",
}
bridgeServer := newBridgeTestServer(ctx, t, "http://unused",
withCustomProvider(provider.NewAnthropic(anthropicCfg("http://unused", apiKey), bedrockCfg)),
)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, fixtures.Request(t, fixtures.AntSingleBuiltinTool))
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
require.Contains(t, string(body), "create anthropic client")
require.Contains(t, string(body), "region or base url required")
})
t.Run("/v1/messages", func(t *testing.T) {
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("%s/streaming=%v", t.Name(), streaming), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, fixtures.AntSingleBuiltinTool)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
// We define region here to validate that with Region & BaseURL defined, the latter takes precedence.
bedrockCfg := &config.AWSBedrock{
Region: "us-west-2",
AccessKey: "test-access-key",
AccessKeySecret: "test-secret-key",
Model: "danthropic", // This model should override the request's given one.
SmallFastModel: "danthropic-mini", // Unused but needed for validation.
BaseURL: upstream.URL, // Use the mock server.
}
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withCustomProvider(provider.NewAnthropic(anthropicCfg(upstream.URL, apiKey), bedrockCfg)),
)
// Make API call to aibridge for Anthropic /v1/messages, which will be routed via AWS Bedrock.
// We override the AWS Bedrock client to route requests through our mock server.
reqBody, err := sjson.SetBytes(fix.Request(), "stream", streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
// For streaming responses, consume the body to allow the stream to complete.
if streaming {
// Read the streaming response.
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
}
// Verify that Bedrock-specific model name was used in the request to the mock server
// and the interception data.
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
// The Anthropic SDK's Bedrock middleware extracts "model" and "stream"
// from the JSON body and encodes them in the URL path.
// See: https://github.com/anthropics/anthropic-sdk-go/blob/4d669338f2041f3c60640b6dd317c4895dc71cd4/bedrock/bedrock.go#L247-L248
pathParts := strings.Split(received[0].Path, "/")
require.True(t, len(pathParts) >= 3 && pathParts[1] == "model", "unexpected path: %s", received[0].Path)
require.Equal(t, bedrockCfg.Model, pathParts[2])
require.False(t, gjson.GetBytes(received[0].Body, "model").Exists(), "model should be stripped from body")
require.False(t, gjson.GetBytes(received[0].Body, "stream").Exists(), "stream should be stripped from body")
// Verify PRM attribution is appended to the User-Agent header.
ua := received[0].Header.Get("User-Agent")
require.Contains(t, ua, "sdk-ua-app-id/APN_1.1%2Fpc_cdfmjwn8i6u8l9fwz8h82e4w3%24",
"expected AWS PRM attribution in User-Agent header")
interceptions := bridgeServer.Recorder.RecordedInterceptions()
require.Len(t, interceptions, 1)
require.Equal(t, interceptions[0].Model, bedrockCfg.Model)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
// Tests that Bedrock-incompatible fields are stripped and adaptive thinking
// is handled correctly per model. Different Bedrock model names trigger
// different behavior for beta flag filtering and field stripping.
t.Run("unsupported fields removed", func(t *testing.T) {
t.Parallel()
// All fields in the fixture request that Bedrock may strip. Fields
// listed in a test case's expectKeptFields survive; all others must
// be absent from the forwarded body.
strippableFields := []string{
"metadata", "service_tier", "container", "inference_geo", // always stripped
"output_config", "context_management", // stripped unless their beta flag survives
}
cases := []struct {
name string
model string
smallFastModel string
expectThinkingType string
expectBudgetTokens int64 // 0 means budget_tokens should not be present
expectKeptFields []string // fields from strippableFields expected to survive
expectedBetaFlags []string // values expected in the anthropic_beta array in the forwarded body
}{
// "beddel" matches no model prefix, so adaptive thinking is converted
// to enabled with budget, and all model-gated beta flags are stripped.
{
name: "beddel",
model: "beddel",
smallFastModel: "modrock",
expectThinkingType: "enabled",
expectBudgetTokens: 16000, // 32000 * 0.5 (medium effort)
expectedBetaFlags: []string{"interleaved-thinking-2025-05-14"},
},
// Opus 4.5 supports the effort beta, so output_config is kept.
{
name: "opus-4.5",
model: "anthropic.claude-opus-4-5-20250514-v1:0",
smallFastModel: "anthropic.claude-haiku-4-5-20241022-v1:0",
expectThinkingType: "enabled",
expectBudgetTokens: 16000,
expectKeptFields: []string{"output_config"},
expectedBetaFlags: []string{"interleaved-thinking-2025-05-14", "effort-2025-11-24"},
},
// Sonnet 4.5 supports context-management beta, so context_management is kept.
{
name: "sonnet-4.5",
model: "anthropic.claude-sonnet-4-5-20241022-v2:0",
smallFastModel: "anthropic.claude-haiku-4-5-20241022-v1:0",
expectThinkingType: "enabled",
expectBudgetTokens: 16000,
expectKeptFields: []string{"context_management"},
expectedBetaFlags: []string{"interleaved-thinking-2025-05-14", "context-management-2025-06-27"},
},
// Opus 4.6 supports adaptive thinking natively, so it is kept as-is.
// Neither effort nor context-management betas apply to this model.
{
name: "opus-4.6",
model: "anthropic.claude-opus-4-6-20260619-v1:0",
smallFastModel: "anthropic.claude-haiku-4-5-20241022-v1:0",
expectThinkingType: "adaptive",
expectedBetaFlags: []string{"interleaved-thinking-2025-05-14"},
},
}
for _, tc := range cases {
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("%s/streaming=%v", tc.name, streaming), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, fixtures.AntSimpleBedrock)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bCfg := &config.AWSBedrock{
Region: "us-west-2",
AccessKey: "test-access-key",
AccessKeySecret: "test-secret-key",
Model: tc.model,
SmallFastModel: tc.smallFastModel,
BaseURL: upstream.URL,
}
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withCustomProvider(provider.NewAnthropic(anthropicCfg(upstream.URL, apiKey), bCfg)),
)
reqBody, err := sjson.SetBytes(fix.Request(), "stream", streaming)
require.NoError(t, err)
// Send with Anthropic-Beta header containing flags that should be filtered.
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody, http.Header{
"Anthropic-Beta": {"interleaved-thinking-2025-05-14,effort-2025-11-24,context-management-2025-06-27,prompt-caching-scope-2026-01-05"},
})
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
body := received[0].Body
// Verify strippable fields: kept only if listed in expectKeptFields.
for _, field := range strippableFields {
assert.Equal(t, slices.Contains(tc.expectKeptFields, field), gjson.GetBytes(body, field).Exists(), "field %s", field)
}
// Verify thinking behavior.
assert.Equal(t, tc.expectThinkingType, gjson.GetBytes(body, "thinking.type").String(), "thinking type mismatch")
if tc.expectBudgetTokens > 0 {
assert.Equal(t, tc.expectBudgetTokens, gjson.GetBytes(body, "thinking.budget_tokens").Int(), "budget_tokens mismatch")
} else {
assert.False(t, gjson.GetBytes(body, "thinking.budget_tokens").Exists(), "budget_tokens should not be present")
}
// The Bedrock SDK middleware moves Anthropic-Beta from the header
// into the body as "anthropic_beta".
betaArr := gjson.GetBytes(body, "anthropic_beta").Array()
var gotFlags []string
for _, v := range betaArr {
gotFlags = append(gotFlags, v.String())
}
assert.Equal(t, tc.expectedBetaFlags, gotFlags, "beta flags mismatch")
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
}
})
// SigV4 signs all headers on the outbound Bedrock request. If any header
// is modified in transit (e.g. an egress proxy appending to X-Forwarded-For),
// the signature becomes invalid and AWS rejects the request with:
// 403: "The request signature we calculated does not match the signature
// you provided."
t.Run("SigV4 signed headers", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, fixtures.AntSingleBuiltinTool)
proxyHeaders := http.Header{
"X-Forwarded-For": {"203.0.113.50, 10.0.0.1"},
"X-Forwarded-Host": {"app.example.com"},
"X-Forwarded-Proto": {"https"},
}
// Credentials used for both the Bedrock config and the mock's
// signature re-verification.
accessKey := "test-access-key"
secretKey := "test-secret-key"
region := "us-west-2"
var signatureValid atomic.Bool
// Mock Bedrock endpoint (simulates AWS). The OnRequest callback
// re-signs the received request using only the declared
// SignedHeaders and stores whether the signatures match.
fixResp := testutil.NewFixtureResponse(fix)
fixResp.OnRequest = func(r *http.Request, body []byte) {
authHeader := r.Header.Get("Authorization")
// Passthrough requests have no SigV4 auth; skip verification.
if !strings.HasPrefix(authHeader, "AWS4-HMAC-SHA256") {
return
}
originalSig := extractSigV4Field(authHeader, "Signature=")
// Rebuild the request the way AWS would: keep only
// the declared SignedHeaders.
signedHeaders := strings.Split(extractSigV4Field(authHeader, "SignedHeaders="), ";")
verifyReq := r.Clone(r.Context())
verifyReq.Header.Del("Authorization")
for h := range verifyReq.Header {
if !slices.Contains(signedHeaders, strings.ToLower(h)) {
verifyReq.Header.Del(h)
}
}
// Restore ContentLength: Go's HTTP server parses it
// from the request but does not put it in r.Header;
// the SigV4 signer reads the struct field.
verifyReq.ContentLength = int64(len(body))
// Re-sign with the same credentials, body hash, and
// timestamp. SigV4 derives the signature from all three,
// so any difference means a header was altered in transit.
signingTime, err := time.Parse("20060102T150405Z", verifyReq.Header.Get("X-Amz-Date"))
require.NoError(t, err)
bodyHash := sha256.Sum256(body)
err = v4signer.NewSigner().SignHTTP(
ctx,
aws.Credentials{AccessKeyID: accessKey, SecretAccessKey: secretKey},
verifyReq, hex.EncodeToString(bodyHash[:]),
"bedrock", region, signingTime,
)
require.NoError(t, err)
recomputedSig := extractSigV4Field(verifyReq.Header.Get("Authorization"), "Signature=")
signatureValid.Store(originalSig == recomputedSig)
}
mockBedrock := testutil.NewMockUpstream(ctx, t, fixResp)
mockBedrock.AllowOverflow = true
// Simulated egress proxy: modifies X-Forwarded-For and
// forwards to mockBedrock, preserving the original Host.
mockEgressProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
r.Header.Set("X-Forwarded-For", xff+", 10.255.0.1")
}
proxyReq, err := http.NewRequestWithContext(r.Context(), r.Method, mockBedrock.URL+r.URL.Path, r.Body)
require.NoError(t, err)
proxyReq.Header = r.Header.Clone()
proxyReq.Host = r.Host // preserve signed Host
resp, err := http.DefaultClient.Do(proxyReq)
require.NoError(t, err)
defer resp.Body.Close()
for k, vs := range resp.Header {
for _, v := range vs {
w.Header().Add(k, v)
}
}
w.WriteHeader(resp.StatusCode)
_, _ = io.Copy(w, resp.Body)
}))
t.Cleanup(mockEgressProxy.Close)
bCfg := bedrockCfg(mockEgressProxy.URL)
bCfg.AccessKey = accessKey
bCfg.AccessKeySecret = secretKey
bCfg.Region = region
bridgeServer := newBridgeTestServer(ctx, t, mockEgressProxy.URL,
withCustomProvider(provider.NewAnthropic(anthropicCfg(mockEgressProxy.URL, apiKey), bCfg)),
)
// Sends a bridge request through a mock egress proxy that
// mutates X-Forwarded-For, then verifies the SigV4 signature
// still matches at the mock Bedrock endpoint.
t.Run("bridge SigV4 signature valid", func(t *testing.T) {
reqBody, err := sjson.SetBytes(fix.Request(), "stream", false)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody, proxyHeaders)
require.NoError(t, err)
defer resp.Body.Close()
_, _ = io.ReadAll(resp.Body)
assert.True(t, signatureValid.Load(),
"SigV4 signature mismatch: a header modified in transit "+
"was included in the signed-headers set")
})
// Passthrough routes use httputil.ReverseProxy, which forwards
// the request as-is without SigV4 signing, so proxy headers
// are safe to include. ReverseProxy sets its own X-Forwarded-*
// headers via SetXForwarded. This verifies they arrive upstream.
t.Run("passthrough proxy sets own forwarded headers", func(t *testing.T) {
resp, err := bridgeServer.makeRequest(t, http.MethodGet, "/anthropic/v1/models", nil, proxyHeaders)
require.NoError(t, err)
defer resp.Body.Close()
_, _ = io.ReadAll(resp.Body)
received := mockBedrock.ReceivedRequests()
require.NotEmpty(t, received)
last := received[len(received)-1]
assert.NotEmpty(t, last.Header.Get("X-Forwarded-For"),
"passthrough should set X-Forwarded-For via SetXForwarded")
assert.NotEmpty(t, last.Header.Get("X-Forwarded-Host"),
"passthrough should set X-Forwarded-Host via SetXForwarded")
assert.NotEmpty(t, last.Header.Get("X-Forwarded-Proto"),
"passthrough should set X-Forwarded-Proto via SetXForwarded")
})
})
}
func TestOpenAIChatCompletions(t *testing.T) {
t.Parallel()
t.Run("single builtin tool", func(t *testing.T) {
t.Parallel()
cases := []struct {
name string
streaming bool
expectedInputTokens, expectedOutputTokens int
expectedToolCallID string
}{
{
name: "streaming",
streaming: true,
expectedInputTokens: 60,
expectedOutputTokens: 15,
expectedToolCallID: "call_HjeqP7YeRkoNj0de9e3U4X4B",
},
{
name: "non-streaming",
streaming: false,
expectedInputTokens: 60,
expectedOutputTokens: 15,
expectedToolCallID: "call_KjzAbhiZC6nk81tQzL7pwlpc",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, fixtures.OaiChatSingleBuiltinTool)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
// Make API call to aibridge for OpenAI /v1/chat/completions
reqBody, err := sjson.SetBytes(fix.Request(), "stream", tc.streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIChatCompletions, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Response-specific checks.
if tc.streaming {
sp := aibridge.NewSSEParser()
require.NoError(t, sp.Parse(resp.Body))
// OpenAI sends all events under the same type.
messageEvents := sp.MessageEvents()
assert.NotEmpty(t, messageEvents)
// OpenAI streaming ends with [DONE]
lastEvent := messageEvents[len(messageEvents)-1]
assert.Equal(t, "[DONE]", lastEvent.Data)
}
tokenUsages := bridgeServer.Recorder.RecordedTokenUsages()
require.Len(t, tokenUsages, 1)
assert.EqualValues(t, tc.expectedInputTokens, bridgeServer.Recorder.TotalInputTokens(), "input tokens miscalculated")
assert.EqualValues(t, tc.expectedOutputTokens, bridgeServer.Recorder.TotalOutputTokens(), "output tokens miscalculated")
toolUsages := bridgeServer.Recorder.RecordedToolUsages()
require.Len(t, toolUsages, 1)
assert.Equal(t, "read_file", toolUsages[0].Tool)
assert.Equal(t, tc.expectedToolCallID, toolUsages[0].ToolCallID)
require.IsType(t, map[string]any{}, toolUsages[0].Args)
require.Contains(t, toolUsages[0].Args, "path")
assert.Equal(t, "README.md", toolUsages[0].Args.(map[string]any)["path"])
promptUsages := bridgeServer.Recorder.RecordedPromptUsages()
require.Len(t, promptUsages, 1)
assert.Equal(t, "how large is the README.md file in my current path", promptUsages[0].Prompt)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
t.Run("streaming injected tool call edge cases", func(t *testing.T) {
t.Parallel()
cases := []struct {
name string
fixture []byte
expectedArgs map[string]any
}{
{
name: "tool call no preamble",
fixture: fixtures.OaiChatStreamingInjectedToolNoPreamble,
expectedArgs: map[string]any{"owner": "me"},
},
{
name: "tool call with non-zero index",
fixture: fixtures.OaiChatStreamingInjectedToolNonzeroIndex,
expectedArgs: nil, // No arguments in this fixture
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
// Setup mock server for multi-turn interaction.
// First request → tool call response, second → tool response.
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix), testutil.NewFixtureToolResponse(fix))
// Setup MCP proxies with the tool from the fixture
mockMCP := setupMCPForTest(t, defaultTracer)
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withMCP(mockMCP),
)
// Add the stream param to the request.
reqBody, err := sjson.SetBytes(fix.Request(), "stream", true)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIChatCompletions, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Verify SSE headers are sent correctly
require.Equal(t, "text/event-stream", resp.Header.Get("Content-Type"))
require.Equal(t, "no-cache", resp.Header.Get("Cache-Control"))
require.Equal(t, "keep-alive", resp.Header.Get("Connection"))
// Consume the full response body to ensure the interception completes
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
// Verify the MCP tool was actually invoked
invocations := mockMCP.getCallsByTool(mockToolName)
require.Len(t, invocations, 1, "expected MCP tool to be invoked")
// Verify tool was invoked with the expected args (if specified)
if tc.expectedArgs != nil {
expected, err := json.Marshal(tc.expectedArgs)
require.NoError(t, err)
actual, err := json.Marshal(invocations[0])
require.NoError(t, err)
require.EqualValues(t, expected, actual)
}
// Verify tool usage was recorded
toolUsages := bridgeServer.Recorder.RecordedToolUsages()
require.Len(t, toolUsages, 1)
assert.Equal(t, mockToolName, toolUsages[0].Tool)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
}
func TestSimple(t *testing.T) {
t.Parallel()
getAnthropicResponseID := func(streaming bool, resp *http.Response) (string, error) {
if streaming {
decoder := ssestream.NewDecoder(resp)
stream := ssestream.NewStream[anthropic.MessageStreamEventUnion](decoder, nil)
var message anthropic.Message
for stream.Next() {
event := stream.Current()
if err := message.Accumulate(event); err != nil {
return "", xerrors.Errorf("accumulate event: %w", err)
}
}
if stream.Err() != nil {
return "", xerrors.Errorf("stream error: %w", stream.Err())
}
return message.ID, nil
}
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", xerrors.Errorf("read body: %w", err)
}
var message anthropic.Message
if err := json.Unmarshal(body, &message); err != nil {
return "", xerrors.Errorf("unmarshal response: %w", err)
}
return message.ID, nil
}
getOpenAIResponseID := func(streaming bool, resp *http.Response) (string, error) {
if streaming {
// Parse the response stream.
decoder := oaissestream.NewDecoder(resp)
stream := oaissestream.NewStream[openai.ChatCompletionChunk](decoder, nil)
var message openai.ChatCompletionAccumulator
for stream.Next() {
chunk := stream.Current()
message.AddChunk(chunk)
}
if stream.Err() != nil {
return "", xerrors.Errorf("stream error: %w", stream.Err())
}
return message.ID, nil
}
// Parse & unmarshal the response.
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", xerrors.Errorf("read body: %w", err)
}
var message openai.ChatCompletion
if err := json.Unmarshal(body, &message); err != nil {
return "", xerrors.Errorf("unmarshal response: %w", err)
}
return message.ID, nil
}
testCases := []struct {
name string
fixture []byte
basePath string
expectedPath string
getResponseIDFunc func(streaming bool, resp *http.Response) (string, error)
path string
expectedMsgID string
userAgent string
expectedClient aibridge.Client
}{
{
name: config.ProviderAnthropic,
fixture: fixtures.AntSimple,
basePath: "",
expectedPath: "/v1/messages",
getResponseIDFunc: getAnthropicResponseID,
path: pathAnthropicMessages,
expectedMsgID: "msg_01Pvyf26bY17RcjmWfJsXGBn",
userAgent: "claude-cli/2.0.67 (external, cli)",
expectedClient: aibridge.ClientClaudeCode,
},
{
name: config.ProviderAnthropic + "_haiku_prompt_capture",
fixture: fixtures.AntHaikuSimple,
basePath: "",
expectedPath: "/v1/messages",
getResponseIDFunc: getAnthropicResponseID,
path: pathAnthropicMessages,
expectedMsgID: "msg_01Pvyf26bY17RcjmWfJsXGBn",
userAgent: "claude-cli/2.0.67 (external, cli)",
expectedClient: aibridge.ClientClaudeCode,
},
{
name: config.ProviderOpenAI,
fixture: fixtures.OaiChatSimple,
basePath: "",
expectedPath: "/chat/completions",
getResponseIDFunc: getOpenAIResponseID,
path: pathOpenAIChatCompletions,
expectedMsgID: "chatcmpl-BwoiPTGRbKkY5rncfaM0s9KtWrq5N",
userAgent: "codex_cli_rs/0.87.0 (Mac OS 26.2.0; arm64)",
expectedClient: aibridge.ClientCodex,
},
{
name: config.ProviderOpenAI + "_opencode",
fixture: fixtures.OaiChatSimple,
basePath: "",
expectedPath: "/chat/completions",
getResponseIDFunc: getOpenAIResponseID,
path: pathOpenAIChatCompletions,
expectedMsgID: "chatcmpl-BwoiPTGRbKkY5rncfaM0s9KtWrq5N",
userAgent: "opencode/1.16.0 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.14",
expectedClient: aibridge.ClientOpenCode,
},
{
name: config.ProviderAnthropic + "_baseURL_path",
fixture: fixtures.AntSimple,
basePath: "/api",
expectedPath: "/api/v1/messages",
getResponseIDFunc: getAnthropicResponseID,
path: pathAnthropicMessages,
expectedMsgID: "msg_01Pvyf26bY17RcjmWfJsXGBn",
userAgent: "GitHubCopilotChat/0.37.2026011603",
expectedClient: aibridge.ClientCopilotVSC,
},
{
name: config.ProviderOpenAI + "_baseURL_path",
fixture: fixtures.OaiChatSimple,
basePath: "/api",
expectedPath: "/api/chat/completions",
getResponseIDFunc: getOpenAIResponseID,
path: pathOpenAIChatCompletions,
expectedMsgID: "chatcmpl-BwoiPTGRbKkY5rncfaM0s9KtWrq5N",
userAgent: "Zed/0.219.4+stable.119.abc123 (macos; aarch64)",
expectedClient: aibridge.ClientZed,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("streaming=%v", streaming), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL+tc.basePath)
// When: calling the "API server" with the fixture's request body.
reqBody, err := sjson.SetBytes(fix.Request(), "stream", streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, reqBody, http.Header{"User-Agent": {tc.userAgent}})
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Then: I expect the upstream request to have the correct path.
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
require.Equal(t, tc.expectedPath, received[0].Path)
// Then: I expect a non-empty response.
bodyBytes, err := io.ReadAll(resp.Body)
require.NoError(t, err)
assert.NotEmpty(t, bodyBytes, "should have received response body")
// Reset the body after being read.
resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
// Then: I expect the prompt to have been tracked.
promptUsages := bridgeServer.Recorder.RecordedPromptUsages()
require.NotEmpty(t, promptUsages, "no prompts tracked")
assert.Contains(t, promptUsages[0].Prompt, "how many angels can dance on the head of a pin")
// Validate that responses have their IDs overridden with a interception ID rather than the original ID from the upstream provider.
// The reason for this is that Bridge may make multiple upstream requests (i.e. to invoke injected tools), and clients will not be expecting
// multiple messages in response to a single request.
id, err := tc.getResponseIDFunc(streaming, resp)
require.NoError(t, err, "failed to retrieve response ID")
require.Nilf(t, uuid.Validate(id), "%s is not a valid UUID", id)
tokenUsages := bridgeServer.Recorder.RecordedTokenUsages()
require.GreaterOrEqual(t, len(tokenUsages), 1)
require.Equal(t, tokenUsages[0].MsgID, tc.expectedMsgID)
// Validate user agent and client have been recorded.
interceptions := bridgeServer.Recorder.RecordedInterceptions()
require.Len(t, interceptions, 1, "expected exactly one interception, got: %v", interceptions)
assert.Equal(t, id, interceptions[0].ID)
assert.Equal(t, tc.userAgent, interceptions[0].UserAgent)
assert.Equal(t, string(tc.expectedClient), interceptions[0].Client)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
}
}
func TestSessionIDTracking(t *testing.T) {
t.Parallel()
testCases := []struct {
name string
fixture []byte
header http.Header
metadataSessionID string
expectedClient aibridge.Client
expectSessionID string
}{
// Session in header.
{
name: "mux",
fixture: fixtures.AntSimple,
expectedClient: aibridge.ClientMux,
expectSessionID: "mux-workspace-321",
header: http.Header{
"User-Agent": []string{"mux/1.0.0"},
"X-Mux-Workspace-Id": []string{"mux-workspace-321"},
},
},
// Session in body.
{
name: "claude_code",
fixture: fixtures.AntSimple,
expectedClient: aibridge.ClientClaudeCode,
expectSessionID: "f47ac10b-58cc-4372-a567-0e02b2c3d479",
header: http.Header{
"User-Agent": []string{"claude-cli/2.0.67 (external, cli)"},
},
metadataSessionID: "user_abc123_account_456_session_f47ac10b-58cc-4372-a567-0e02b2c3d479",
},
// No session.
{
name: "zed",
fixture: fixtures.AntSimple,
expectedClient: aibridge.ClientZed,
header: http.Header{
"User-Agent": []string{"Zed/0.219.4+stable.119.abc123 (macos; aarch64)"},
},
},
{
name: "opencode",
fixture: fixtures.AntSimple,
expectedClient: aibridge.ClientOpenCode,
expectSessionID: "ses_15a48edefffe7oY0YcIHRv29dD",
header: http.Header{
"User-Agent": []string{"opencode/1.16.0 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.14"},
"X-OpenCode-Session": []string{"ses_15a48edefffe7oY0YcIHRv29dD"},
},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withProvider(config.ProviderAnthropic))
reqBody := fix.Request()
if tc.metadataSessionID != "" {
var err error
reqBody, err = sjson.SetBytes(reqBody, "metadata.user_id", tc.metadataSessionID)
require.NoError(t, err)
}
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody, tc.header)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Drain the body to let the stream complete.
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
interceptions := bridgeServer.Recorder.RecordedInterceptions()
require.Len(t, interceptions, 1, "expected exactly one interception")
assert.Equal(t, string(tc.expectedClient), interceptions[0].Client)
if tc.expectSessionID == "" {
assert.Nil(t, interceptions[0].ClientSessionID, "expected nil session ID for %s", tc.name)
} else {
require.NotNil(t, interceptions[0].ClientSessionID, "expected non-nil session ID for %s", tc.name)
assert.Equal(t, tc.expectSessionID, *interceptions[0].ClientSessionID)
}
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
}
func TestFallthrough(t *testing.T) {
t.Parallel()
testCases := []struct {
name string
fixture []byte
basePath string
requestPath string
expectedUpstreamPath string
expectAuthHeader string
}{
{
name: "ant_empty_base_url_path",
fixture: fixtures.AntFallthrough,
basePath: "",
requestPath: "/anthropic/v1/models",
expectedUpstreamPath: "/v1/models",
expectAuthHeader: "X-Api-Key",
},
{
name: "oai_empty_base_url_path",
fixture: fixtures.OaiChatFallthrough,
basePath: "",
requestPath: "/openai/v1/models",
expectedUpstreamPath: "/models",
expectAuthHeader: "Authorization",
},
{
name: "ant_some_base_url_path",
fixture: fixtures.AntFallthrough,
basePath: "/api",
requestPath: "/anthropic/v1/models",
expectedUpstreamPath: "/api/v1/models",
expectAuthHeader: "X-Api-Key",
},
{
name: "oai_some_base_url_path",
fixture: fixtures.OaiChatFallthrough,
basePath: "/api",
requestPath: "/openai/v1/models",
expectedUpstreamPath: "/api/models",
expectAuthHeader: "Authorization",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(t.Context(), t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL+tc.basePath)
resp, err := bridgeServer.makeRequest(t, http.MethodGet, tc.requestPath, nil)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Verify upstream received the request at the expected path
// with the API key header.
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
require.Equal(t, tc.expectedUpstreamPath, received[0].Path)
require.Contains(t, received[0].Header.Get(tc.expectAuthHeader), apiKey)
gotBytes, err := io.ReadAll(resp.Body)
require.NoError(t, err)
// Compare JSON bodies for semantic equality.
var got any
var exp any
require.NoError(t, json.Unmarshal(gotBytes, &got))
require.NoError(t, json.Unmarshal(fix.NonStreaming(), &exp))
require.EqualValues(t, exp, got)
})
}
}
func TestAnthropicInjectedTools(t *testing.T) {
t.Parallel()
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("streaming=%v", streaming), func(t *testing.T) {
t.Parallel()
// Build the requirements & make the assertions which are common to all providers.
bridgeServer, mockMCP, resp := setupInjectedToolTest(t, fixtures.AntSingleInjectedTool, streaming, defaultTracer, pathAnthropicMessages, anthropicToolResultValidator(t))
defer resp.Body.Close()
// Ensure expected tool was invoked with expected input.
toolUsages := bridgeServer.Recorder.RecordedToolUsages()
require.Len(t, toolUsages, 1)
require.Equal(t, mockToolName, toolUsages[0].Tool)
expected, err := json.Marshal(map[string]any{"owner": "admin"})
require.NoError(t, err)
actual, err := json.Marshal(toolUsages[0].Args)
require.NoError(t, err)
require.EqualValues(t, expected, actual)
invocations := mockMCP.getCallsByTool(mockToolName)
require.Len(t, invocations, 1)
actual, err = json.Marshal(invocations[0])
require.NoError(t, err)
require.EqualValues(t, expected, actual)
var (
content *anthropic.ContentBlockUnion
message anthropic.Message
)
if streaming {
// Parse the response stream.
decoder := ssestream.NewDecoder(resp)
stream := ssestream.NewStream[anthropic.MessageStreamEventUnion](decoder, nil)
for stream.Next() {
event := stream.Current()
require.NoError(t, message.Accumulate(event), "accumulate event")
}
require.NoError(t, stream.Err(), "stream error")
require.Len(t, message.Content, 2)
content = &message.Content[1]
} else {
// Parse & unmarshal the response.
body, err := io.ReadAll(resp.Body)
require.NoError(t, err, "read response body")
require.NoError(t, json.Unmarshal(body, &message), "unmarshal response")
require.GreaterOrEqual(t, len(message.Content), 1)
content = &message.Content[0]
}
// Ensure tool returned expected value.
require.NotNil(t, content)
require.Contains(t, content.Text, "dd711d5c-83c6-4c08-a0af-b73055906e8c") // The ID of the workspace to be returned.
// Check the token usage from the client's perspective.
//
// We overwrite the final message_delta which is relayed to the client to include the
// accumulated tokens but currently the SDK only supports accumulating output tokens
// for message_delta events.
//
// For non-streaming requests the token usage is also overwritten and should be faithfully
// represented in the response.
//
// See https://github.com/anthropics/anthropic-sdk-go/blob/v1.12.0/message.go#L2619-L2622
if !streaming {
assert.EqualValues(t, 15308, message.Usage.InputTokens)
}
assert.EqualValues(t, 204, message.Usage.OutputTokens)
// Ensure tokens used during injected tool invocation are accounted for.
assert.EqualValues(t, 15308, bridgeServer.Recorder.TotalInputTokens())
assert.EqualValues(t, 204, bridgeServer.Recorder.TotalOutputTokens())
// Ensure we received exactly one prompt.
promptUsages := bridgeServer.Recorder.RecordedPromptUsages()
require.Len(t, promptUsages, 1)
})
}
}
func TestOpenAIInjectedTools(t *testing.T) {
t.Parallel()
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("streaming=%v", streaming), func(t *testing.T) {
t.Parallel()
// Build the requirements & make the assertions which are common to all providers.
bridgeServer, mockMCP, resp := setupInjectedToolTest(t, fixtures.OaiChatSingleInjectedTool, streaming, defaultTracer, pathOpenAIChatCompletions, openaiChatToolResultValidator(t))
defer resp.Body.Close()
// Ensure expected tool was invoked with expected input.
toolUsages := bridgeServer.Recorder.RecordedToolUsages()
require.Len(t, toolUsages, 1)
require.Equal(t, mockToolName, toolUsages[0].Tool)
expected, err := json.Marshal(map[string]any{"owner": "admin"})
require.NoError(t, err)
actual, err := json.Marshal(toolUsages[0].Args)
require.NoError(t, err)
require.EqualValues(t, expected, actual)
invocations := mockMCP.getCallsByTool(mockToolName)
require.Len(t, invocations, 1)
actual, err = json.Marshal(invocations[0])
require.NoError(t, err)
require.EqualValues(t, expected, actual)
var (
content *openai.ChatCompletionChoice
message openai.ChatCompletion
)
if streaming {
// Parse the response stream.
decoder := oaissestream.NewDecoder(resp)
stream := oaissestream.NewStream[openai.ChatCompletionChunk](decoder, nil)
var acc openai.ChatCompletionAccumulator
detectedToolCalls := make(map[string]struct{})
for stream.Next() {
chunk := stream.Current()
acc.AddChunk(chunk)
if len(chunk.Choices) == 0 {
continue
}
for _, c := range chunk.Choices {
if len(c.Delta.ToolCalls) == 0 {
continue
}
for _, t := range c.Delta.ToolCalls {
if t.Function.Name == "" {
continue
}
detectedToolCalls[t.Function.Name] = struct{}{}
}
}
}
// Verify that no injected tool call events (or partials thereof) were sent to the client.
require.Len(t, detectedToolCalls, 0)
message = acc.ChatCompletion
require.NoError(t, stream.Err(), "stream error")
} else {
// Parse & unmarshal the response.
body, err := io.ReadAll(resp.Body)
require.NoError(t, err, "read response body")
require.NoError(t, json.Unmarshal(body, &message), "unmarshal response")
// Verify that no injected tools were sent to the client.
require.GreaterOrEqual(t, len(message.Choices), 1)
require.Len(t, message.Choices[0].Message.ToolCalls, 0)
}
require.GreaterOrEqual(t, len(message.Choices), 1)
content = &message.Choices[0]
// Ensure tool returned expected value.
require.NotNil(t, content)
require.Contains(t, content.Message.Content, "dd711d5c-83c6-4c08-a0af-b73055906e8c") // The ID of the workspace to be returned.
// Check the token usage from the client's perspective.
// This *should* work but the openai SDK doesn't accumulate the prompt token details :(.
// See https://github.com/openai/openai-go/blob/v2.7.0/streamaccumulator.go#L145-L147.
// assert.EqualValues(t, 5047, message.Usage.PromptTokens-message.Usage.PromptTokensDetails.CachedTokens)
assert.EqualValues(t, 105, message.Usage.CompletionTokens)
// Ensure tokens used during injected tool invocation are accounted for.
require.EqualValues(t, 5047, bridgeServer.Recorder.TotalInputTokens())
require.EqualValues(t, 105, bridgeServer.Recorder.TotalOutputTokens())
// Ensure we received exactly one prompt.
promptUsages := bridgeServer.Recorder.RecordedPromptUsages()
require.Len(t, promptUsages, 1)
})
}
}
// anthropicToolResultValidator returns a request validator that asserts the second
// upstream request contains the assistant's tool_use and user's tool_result messages
// appended by the inner agentic loop. If the raw payload is not kept in sync with
// the structured messages, the second request will be identical to the first.
func anthropicToolResultValidator(t *testing.T) func(*http.Request, []byte) {
t.Helper()
return func(_ *http.Request, raw []byte) {
messages := gjson.GetBytes(raw, "messages").Array()
// After the agentic loop the messages must contain at minimum:
// [0] original user message
// [N-2] assistant message with tool_use content block
// [N-1] user message with tool_result content block
require.GreaterOrEqual(t, len(messages), 3,
"second upstream request must contain the original message, assistant tool_use, and user tool_result")
assistantMsg := messages[len(messages)-2]
require.Equal(t, "assistant", assistantMsg.Get("role").Str,
"penultimate message must be from the assistant")
var hasToolUse bool
for _, block := range assistantMsg.Get("content").Array() {
if block.Get("type").Str == "tool_use" {
hasToolUse = true
break
}
}
require.True(t, hasToolUse, "assistant message must contain a tool_use content block")
toolResultMsg := messages[len(messages)-1]
require.Equal(t, "user", toolResultMsg.Get("role").Str,
"last message must be a user message carrying the tool_result")
var hasToolResult bool
for _, block := range toolResultMsg.Get("content").Array() {
if block.Get("type").Str == "tool_result" {
hasToolResult = true
break
}
}
require.True(t, hasToolResult, "user message must contain a tool_result content block")
}
}
// openaiChatToolResultValidator returns a request validator that asserts the second
// upstream request contains the assistant's tool_calls and a role=tool result message
// appended by the inner agentic loop.
func openaiChatToolResultValidator(t *testing.T) func(*http.Request, []byte) {
t.Helper()
return func(_ *http.Request, raw []byte) {
messages := gjson.GetBytes(raw, "messages").Array()
// After the agentic loop the messages must contain at minimum:
// [0] original user message
// [N-2] assistant message with tool_calls array
// [N-1] message with role=tool
require.GreaterOrEqual(t, len(messages), 3,
"second upstream request must contain the original message, assistant tool_calls, and tool result")
assistantMsg := messages[len(messages)-2]
require.Equal(t, "assistant", assistantMsg.Get("role").Str,
"penultimate message must be from the assistant")
require.NotEmpty(t, len(assistantMsg.Get("tool_calls").Array()),
"assistant message must contain a tool_calls array")
toolResultMsg := messages[len(messages)-1]
require.Equal(t, "tool", toolResultMsg.Get("role").Str,
"last message must have role=tool")
require.NotEmpty(t, toolResultMsg.Get("tool_call_id").Str,
"tool result message must have a tool_call_id")
}
}
func TestErrorHandling(t *testing.T) {
t.Parallel()
// Tests that errors which occur *before* a streaming response begins, or in non-streaming requests, are handled as expected.
t.Run("non-stream error", func(t *testing.T) {
t.Parallel()
cases := []struct {
name string
fixture []byte
path string
responseHandlerFn func(resp *http.Response)
}{
{
name: config.ProviderAnthropic,
fixture: fixtures.AntNonStreamError,
path: pathAnthropicMessages,
responseHandlerFn: func(resp *http.Response) {
require.Equal(t, http.StatusBadRequest, resp.StatusCode)
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
require.Equal(t, "error", gjson.GetBytes(body, "type").Str)
require.Equal(t, "invalid_request_error", gjson.GetBytes(body, "error.type").Str)
require.Contains(t, gjson.GetBytes(body, "error.message").Str, "prompt is too long")
},
},
{
name: config.ProviderOpenAI,
fixture: fixtures.OaiChatNonStreamError,
path: pathOpenAIChatCompletions,
responseHandlerFn: func(resp *http.Response) {
require.Equal(t, http.StatusBadRequest, resp.StatusCode)
body, err := io.ReadAll(resp.Body)
require.NoError(t, err)
require.Equal(t, "context_length_exceeded", gjson.GetBytes(body, "error.code").Str)
require.Equal(t, "invalid_request_error", gjson.GetBytes(body, "error.type").Str)
require.Contains(t, gjson.GetBytes(body, "error.message").Str, "Input tokens exceed the configured limit")
},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("streaming=%v", streaming), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
// Setup mock server. Error fixtures contain raw HTTP
// responses that may cause the bridge to retry.
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
// Add the stream param to the request.
reqBody, err := sjson.SetBytes(fix.Request(), "stream", streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
tc.responseHandlerFn(resp)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
}
})
// Tests that errors which occur *during* a streaming response are handled as expected.
t.Run("mid-stream error", func(t *testing.T) {
cases := []struct {
name string
fixture []byte
path string
responseHandlerFn func(resp *http.Response)
}{
{
name: config.ProviderAnthropic,
fixture: fixtures.AntMidStreamError,
path: pathAnthropicMessages,
responseHandlerFn: func(resp *http.Response) {
// Server responds first with 200 OK then starts streaming.
require.Equal(t, http.StatusOK, resp.StatusCode)
sp := aibridge.NewSSEParser()
require.NoError(t, sp.Parse(resp.Body))
require.Len(t, sp.EventsByType("error"), 1)
require.Contains(t, sp.EventsByType("error")[0].Data, "Overloaded")
},
},
{
name: config.ProviderOpenAI,
fixture: fixtures.OaiChatMidStreamError,
path: pathOpenAIChatCompletions,
responseHandlerFn: func(resp *http.Response) {
// Server responds first with 200 OK then starts streaming.
require.Equal(t, http.StatusOK, resp.StatusCode)
sp := aibridge.NewSSEParser()
require.NoError(t, sp.Parse(resp.Body))
// OpenAI sends all events under the same type.
messageEvents := sp.MessageEvents()
require.NotEmpty(t, messageEvents)
errEvent := sp.MessageEvents()[len(sp.MessageEvents())-2] // Last event is termination marker ("[DONE]").
require.NotEmpty(t, errEvent)
require.Contains(t, errEvent.Data, "The server had an error while processing your request. Sorry about that!")
},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
// Setup mock server.
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
upstream.StatusCode = http.StatusInternalServerError
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, fix.Request())
require.NoError(t, err)
defer resp.Body.Close()
tc.responseHandlerFn(resp)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
}
// TestStableRequestEncoding validates that a given intercepted request and a
// given set of injected tools should result identical payloads.
//
// Should the payload vary, it may subvert any caching mechanisms the provider may have.
func TestStableRequestEncoding(t *testing.T) {
t.Parallel()
cases := []struct {
name string
fixture []byte
path string
}{
{
name: config.ProviderAnthropic,
fixture: fixtures.AntSimple,
path: pathAnthropicMessages,
},
{
name: config.ProviderOpenAI,
fixture: fixtures.OaiChatSimple,
path: pathOpenAIChatCompletions,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
// Setup MCP tools.
mockMCP := setupMCPForTest(t, defaultTracer)
fix := fixtures.Parse(t, tc.fixture)
// Create a mock upstream that serves the same blocking response for each request.
count := 10
responses := make([]testutil.UpstreamResponse, count)
for i := range count {
responses[i] = testutil.NewFixtureResponse(fix)
}
upstream := testutil.NewMockUpstream(ctx, t, responses...)
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withMCP(mockMCP),
)
// Make multiple requests and verify they all have identical payloads.
for range count {
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, fix.Request())
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
}
// All upstream request bodies should be identical.
received := upstream.ReceivedRequests()
require.Len(t, received, count)
reference := string(received[0].Body)
for _, r := range received[1:] {
assert.JSONEq(t, reference, string(r.Body))
}
})
}
}
// TestAnthropicToolChoiceParallelDisabled verifies that parallel tool use is
// correctly disabled based on the tool_choice parameter in the request.
// See https://github.com/coder/aibridge/issues/2
func TestAnthropicToolChoiceParallelDisabled(t *testing.T) {
t.Parallel()
var (
toolChoiceAuto = string(constant.ValueOf[constant.Auto]())
toolChoiceAny = string(constant.ValueOf[constant.Any]())
toolChoiceNone = string(constant.ValueOf[constant.None]())
toolChoiceTool = string(constant.ValueOf[constant.Tool]())
)
cases := []struct {
name string
fixture []byte
toolChoice any // nil, or map with "type" key.
withInjectedTools bool
expectDisableParallel *bool // nil = field should not be present, non-nil = expected value.
expectToolChoiceTypeInRequest string
}{
// With injected tools - disable_parallel_tool_use should be set to true.
{
name: "with injected tools: no tool_choice defined defaults to auto",
fixture: fixtures.AntSimple,
toolChoice: nil,
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with injected tools: tool_choice auto",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with injected tools: tool_choice any",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAny},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAny,
},
{
name: "with injected tools: tool_choice tool",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceTool, "name": "some_tool"},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceTool,
},
{
name: "with injected tools: tool_choice none",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceNone},
withInjectedTools: true,
expectDisableParallel: nil,
expectToolChoiceTypeInRequest: toolChoiceNone,
},
// With injected tools and builtin tools - disable_parallel_tool_use should be set to true.
{
name: "with injected and builtin tools: no tool_choice defined defaults to auto",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: nil,
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with injected and builtin tools: tool_choice auto",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAuto},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with injected and builtin tools: tool_choice any",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAny},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAny,
},
{
name: "with injected and builtin tools: tool_choice tool",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceTool, "name": "some_tool"},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceTool,
},
{
name: "with injected and builtin tools: tool_choice none",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceNone},
withInjectedTools: true,
expectDisableParallel: nil,
expectToolChoiceTypeInRequest: toolChoiceNone,
},
{
name: "with injected and builtin tools: request already disables parallel",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": true},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with injected and builtin tools: request explicitly enables parallel",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": false},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
// Without injected or builtin tools - disable_parallel_tool_use should NOT be set.
{
name: "without injected tools or builtin tools: tool_choice auto",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto},
withInjectedTools: false,
expectDisableParallel: nil,
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "without injected tools or builtin tools: tool_choice any",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAny},
withInjectedTools: false,
expectDisableParallel: nil,
expectToolChoiceTypeInRequest: toolChoiceAny,
},
// With builtin tools but without injected tools - disable_parallel_tool_use should NOT be set.
{
name: "with builtin tools only: tool_choice auto",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAuto},
withInjectedTools: false,
expectDisableParallel: nil,
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with builtin tools only: tool_choice any",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAny},
withInjectedTools: false,
expectDisableParallel: nil,
expectToolChoiceTypeInRequest: toolChoiceAny,
},
{
name: "with builtin tools only: request explicitly disables parallel",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": true},
withInjectedTools: false,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with builtin tools only: request explicitly enables parallel",
fixture: fixtures.AntSingleBuiltinTool,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": false},
withInjectedTools: false,
expectDisableParallel: utils.PtrTo(false),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
// Without injected or builtin tools - disable_parallel_tool_use should be preserved if set.
{
name: "no tools: request explicitly disables parallel",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": true},
withInjectedTools: false,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "no tools: request explicitly enables parallel",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": false},
withInjectedTools: false,
expectDisableParallel: utils.PtrTo(false),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
// Request already has disable_parallel_tool_use set - with injected tools it should be set to true.
{
name: "with injected tools: request already disables parallel",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": true},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "with injected tools: request explicitly enables parallel",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": false},
withInjectedTools: true,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
// Request already has disable_parallel_tool_use set - without injected tools it should be preserved.
{
name: "without injected tools: request already disables parallel",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": true},
withInjectedTools: false,
expectDisableParallel: utils.PtrTo(true),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
{
name: "without injected tools: request explicitly enables parallel",
fixture: fixtures.AntSimple,
toolChoice: map[string]any{"type": toolChoiceAuto, "disable_parallel_tool_use": false},
withInjectedTools: false,
expectDisableParallel: utils.PtrTo(false),
expectToolChoiceTypeInRequest: toolChoiceAuto,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
// Setup MCP tools conditionally.
var mockMCP mcp.ServerProxier
if tc.withInjectedTools {
mockMCP = setupMCPForTest(t, defaultTracer)
} else {
mockMCP = newNoopMCPManager()
}
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withMCP(mockMCP),
)
// Prepare request body with tool_choice set.
reqBody, err := sjson.SetBytes(fix.Request(), "tool_choice", tc.toolChoice)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Verify tool_choice in the upstream request.
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
var receivedRequest map[string]any
require.NoError(t, json.Unmarshal(received[0].Body, &receivedRequest))
toolChoice, ok := receivedRequest["tool_choice"].(map[string]any)
require.True(t, ok, "expected tool_choice in upstream request")
// Verify the type matches expectation.
assert.Equal(t, tc.expectToolChoiceTypeInRequest, toolChoice["type"])
// Verify name is preserved for tool_choice=tool.
if tc.expectToolChoiceTypeInRequest == toolChoiceTool {
assert.Equal(t, "some_tool", toolChoice["name"])
}
// Verify disable_parallel_tool_use based on expectations.
// See https://platform.claude.com/docs/en/agents-and-tools/tool-use/implement-tool-use#parallel-tool-use
disableParallel, hasDisableParallel := toolChoice["disable_parallel_tool_use"].(bool)
require.Equal(t, tc.expectDisableParallel != nil, hasDisableParallel,
"disable_parallel_tool_use presence mismatch")
if tc.expectDisableParallel != nil {
assert.Equal(t, *tc.expectDisableParallel, disableParallel)
}
})
}
}
// TestChatCompletionsParallelToolCallsDisabled verifies that parallel_tool_calls
// is set to false only when injectable MCP tools are present and the request
// includes tools.
func TestChatCompletionsParallelToolCallsDisabled(t *testing.T) {
t.Parallel()
cases := []struct {
name string
fixture []byte
withInjectedTools bool
initialSetting *bool
expectedSetting *bool
}{
// With injected tools and builtin tools: parallel_tool_calls should be forced false.
{
name: "with injected and builtin tools: parallel_tool_calls true",
fixture: fixtures.OaiChatSingleBuiltinTool,
withInjectedTools: true,
initialSetting: utils.PtrTo(true),
expectedSetting: utils.PtrTo(false),
},
{
name: "with injected and builtin tools: parallel_tool_calls false",
fixture: fixtures.OaiChatSingleBuiltinTool,
withInjectedTools: true,
initialSetting: utils.PtrTo(false),
expectedSetting: utils.PtrTo(false),
},
{
name: "with injected and builtin tools: parallel_tool_calls unset",
fixture: fixtures.OaiChatSingleBuiltinTool,
withInjectedTools: true,
initialSetting: nil,
expectedSetting: utils.PtrTo(false),
},
// With injected tools but without builtin tools: parallel_tool_calls should be forced false.
{
name: "with injected tools only: parallel_tool_calls true",
fixture: fixtures.OaiChatSimple,
withInjectedTools: true,
initialSetting: utils.PtrTo(true),
expectedSetting: utils.PtrTo(false),
},
{
name: "with injected tools only: parallel_tool_calls false",
fixture: fixtures.OaiChatSimple,
withInjectedTools: true,
initialSetting: utils.PtrTo(false),
expectedSetting: utils.PtrTo(false),
},
{
name: "with injected tools only: parallel_tool_calls unset",
fixture: fixtures.OaiChatSimple,
withInjectedTools: true,
initialSetting: nil,
expectedSetting: utils.PtrTo(false),
},
// With builtin tools but without injected tools: parallel_tool_calls should be preserved.
{
name: "with builtin tools only: parallel_tool_calls true",
fixture: fixtures.OaiChatSingleBuiltinTool,
withInjectedTools: false,
initialSetting: utils.PtrTo(true),
expectedSetting: utils.PtrTo(true),
},
{
name: "with builtin tools only: parallel_tool_calls false",
fixture: fixtures.OaiChatSingleBuiltinTool,
withInjectedTools: false,
initialSetting: utils.PtrTo(false),
expectedSetting: utils.PtrTo(false),
},
{
name: "with builtin tools only: parallel_tool_calls unset",
fixture: fixtures.OaiChatSingleBuiltinTool,
withInjectedTools: false,
initialSetting: nil,
expectedSetting: nil,
},
// Without any tools: nothing is modified.
{
name: "no tools: parallel_tool_calls true",
fixture: fixtures.OaiChatSimple,
withInjectedTools: false,
initialSetting: utils.PtrTo(true),
expectedSetting: utils.PtrTo(true),
},
{
name: "no tools: parallel_tool_calls false",
fixture: fixtures.OaiChatSimple,
withInjectedTools: false,
initialSetting: utils.PtrTo(false),
expectedSetting: utils.PtrTo(false),
},
{
name: "no tools: parallel_tool_calls unset",
fixture: fixtures.OaiChatSimple,
withInjectedTools: false,
initialSetting: nil,
expectedSetting: nil,
},
}
for _, tc := range cases {
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("%s/streaming=%v", tc.name, streaming), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
var opts []bridgeOption
if tc.withInjectedTools {
opts = append(opts, withMCP(setupMCPForTest(t, defaultTracer)))
}
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, opts...)
var (
reqBody = fix.Request()
err error
)
if tc.initialSetting != nil {
reqBody, err = sjson.SetBytes(reqBody, "parallel_tool_calls", *tc.initialSetting)
require.NoError(t, err)
}
reqBody, err = sjson.SetBytes(reqBody, "stream", streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIChatCompletions, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
var upstreamReq map[string]any
require.NoError(t, json.Unmarshal(received[0].Body, &upstreamReq))
ptc, ok := upstreamReq["parallel_tool_calls"].(bool)
require.Equal(t, tc.expectedSetting != nil, ok,
"parallel_tool_calls presence mismatch")
if tc.expectedSetting != nil {
assert.Equal(t, *tc.expectedSetting, ptc)
}
})
}
}
}
func TestThinkingAdaptiveIsPreserved(t *testing.T) {
t.Parallel()
fix := fixtures.Parse(t, fixtures.AntSimple)
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("streaming=%v", streaming), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
// Create a mock server that captures the request body sent upstream.
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
// Inject adaptive thinking into the fixture request.
reqBody, err := sjson.SetBytes(fix.Request(), "thinking", map[string]string{"type": "adaptive"})
require.NoError(t, err)
reqBody, err = sjson.SetBytes(reqBody, "stream", streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
// Verify the thinking field was preserved in the upstream request.
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
assert.Equal(t, "adaptive", gjson.GetBytes(received[0].Body, "thinking.type").Str)
})
}
}
func TestEnvironmentDoNotLeak(t *testing.T) {
// NOTE: Cannot use t.Parallel() here because subtests use t.Setenv which requires sequential execution.
// Test that environment variables containing API keys/tokens are not leaked to upstream requests.
// See https://github.com/coder/aibridge/issues/60.
testCases := []struct {
name string
fixture []byte
path string
envVars map[string]string
headerName string
}{
{
name: config.ProviderAnthropic,
fixture: fixtures.AntSimple,
path: pathAnthropicMessages,
envVars: map[string]string{
"ANTHROPIC_AUTH_TOKEN": "should-not-leak",
},
headerName: "Authorization", // We only send through the X-Api-Key, so this one should not be present.
},
{
name: config.ProviderOpenAI,
fixture: fixtures.OaiChatSimple,
path: pathOpenAIChatCompletions,
envVars: map[string]string{
"OPENAI_ORG_ID": "should-not-leak",
},
headerName: "OpenAI-Organization",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// NOTE: Cannot use t.Parallel() here because t.Setenv requires sequential execution.
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
// Set environment variables that the SDK would automatically read.
// These should NOT leak into upstream requests.
for key, val := range tc.envVars {
t.Setenv(key, val)
}
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, fix.Request())
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// Verify that environment values did not leak.
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
require.Empty(t, received[0].Header.Get(tc.headerName))
})
}
}
func TestActorHeaders(t *testing.T) {
t.Parallel()
actorUsername := "bob"
cases := []struct {
name string
path string
createProviderFn func(url, key string, sendHeaders bool) aibridge.Provider
fixture []byte
streaming bool
}{
{
name: "openai/v1/chat/completions",
path: pathOpenAIChatCompletions,
createProviderFn: func(url, key string, sendHeaders bool) aibridge.Provider {
cfg := openAICfg(url, key)
cfg.SendActorHeaders = sendHeaders
return provider.NewOpenAI(cfg)
},
fixture: fixtures.OaiChatSimple,
streaming: true,
},
{
name: "openai/v1/chat/completions",
path: pathOpenAIChatCompletions,
createProviderFn: func(url, key string, sendHeaders bool) aibridge.Provider {
cfg := openAICfg(url, key)
cfg.SendActorHeaders = sendHeaders
return provider.NewOpenAI(cfg)
},
fixture: fixtures.OaiChatSimple,
streaming: false,
},
{
name: "openai/v1/responses",
path: pathOpenAIResponses,
createProviderFn: func(url, key string, sendHeaders bool) aibridge.Provider {
cfg := openAICfg(url, key)
cfg.SendActorHeaders = sendHeaders
return provider.NewOpenAI(cfg)
},
fixture: fixtures.OaiResponsesStreamingSimple,
streaming: true,
},
{
name: "openai/v1/responses",
path: pathOpenAIResponses,
createProviderFn: func(url, key string, sendHeaders bool) aibridge.Provider {
cfg := openAICfg(url, key)
cfg.SendActorHeaders = sendHeaders
return provider.NewOpenAI(cfg)
},
fixture: fixtures.OaiResponsesBlockingSimple,
streaming: false,
},
{
name: "anthropic/v1/messages",
path: pathAnthropicMessages,
createProviderFn: func(url, key string, sendHeaders bool) aibridge.Provider {
cfg := anthropicCfg(url, key)
cfg.SendActorHeaders = sendHeaders
return provider.NewAnthropic(cfg, nil)
},
fixture: fixtures.AntSimple,
streaming: true,
},
{
name: "anthropic/v1/messages",
path: pathAnthropicMessages,
createProviderFn: func(url, key string, sendHeaders bool) aibridge.Provider {
cfg := anthropicCfg(url, key)
cfg.SendActorHeaders = sendHeaders
return provider.NewAnthropic(cfg, nil)
},
fixture: fixtures.AntSimple,
streaming: false,
},
}
for _, tc := range cases {
for _, send := range []bool{true, false} {
t.Run(fmt.Sprintf("%s/streaming=%v/send-headers=%v", tc.name, tc.streaming, send), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
metadataKey := "Username"
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withCustomProvider(tc.createProviderFn(upstream.URL, apiKey, send)),
withActor(defaultActorID, recorder.Metadata{
metadataKey: actorUsername,
}),
)
// Add the stream param to the request.
reqBody, err := sjson.SetBytes(fix.Request(), "stream", tc.streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, reqBody)
require.NoError(t, err)
defer resp.Body.Close()
// Drain the body so streaming responses complete without
// a "connection reset" error in the mock upstream.
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
received := upstream.ReceivedRequests()
require.NotEmpty(t, received)
receivedHeaders := received[0].Header
// Verify that the actor headers were only received if intended.
found := make(map[string][]string)
for k, v := range receivedHeaders {
k = strings.ToLower(k)
if intercept.IsActorHeader(k) {
found[k] = v
}
}
if send {
require.Equal(t, found[strings.ToLower(intercept.ActorIDHeader())], []string{defaultActorID})
require.Equal(t, found[strings.ToLower(intercept.ActorMetadataHeader(metadataKey))], []string{actorUsername})
} else {
require.Empty(t, found)
}
})
}
}
}
// extractSigV4Field extracts a named field from an AWS SigV4
// Authorization header value.
func extractSigV4Field(authHeader, prefix string) string {
idx := strings.Index(authHeader, prefix)
if idx == -1 {
return ""
}
val := authHeader[idx+len(prefix):]
if end := strings.IndexByte(val, ','); end != -1 {
val = val[:end]
}
return strings.TrimSpace(val)
}
Reference in New Issue View Git Blame Copy Permalink