911 lines
35 KiB
Go
911 lines
35 KiB
Go
package aws
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"maps"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
relaytypes "github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/aws/aws-sdk-go-v2/aws"
|
|
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
|
|
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi"
|
|
"github.com/aws/aws-sdk-go-v2/aws/signer/v4"
|
|
"github.com/aws/aws-sdk-go-v2/credentials"
|
|
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
const awsTestModel = "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
|
|
|
type awsHTTPClientFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f awsHTTPClientFunc) Do(request *http.Request) (*http.Response, error) {
|
|
return f(request)
|
|
}
|
|
|
|
type awsNotifyingResponseWriter struct {
|
|
*httptest.ResponseRecorder
|
|
notifyOn []byte
|
|
notified chan int
|
|
once sync.Once
|
|
}
|
|
|
|
func newAwsNotifyingResponseWriter(notifyOn string) *awsNotifyingResponseWriter {
|
|
return &awsNotifyingResponseWriter{
|
|
ResponseRecorder: httptest.NewRecorder(),
|
|
notifyOn: []byte(notifyOn),
|
|
notified: make(chan int, 1),
|
|
}
|
|
}
|
|
|
|
func (w *awsNotifyingResponseWriter) Write(data []byte) (int, error) {
|
|
return w.ResponseRecorder.Write(data)
|
|
}
|
|
|
|
func (w *awsNotifyingResponseWriter) Flush() {
|
|
w.ResponseRecorder.Flush()
|
|
if bytes.Contains(w.Body.Bytes(), w.notifyOn) {
|
|
w.once.Do(func() {
|
|
w.notified <- w.Body.Len()
|
|
})
|
|
}
|
|
}
|
|
|
|
func newAwsTestClient(httpClient bedrockruntime.HTTPClient) *bedrockruntime.Client {
|
|
return bedrockruntime.New(bedrockruntime.Options{
|
|
Region: "us-east-1",
|
|
BaseEndpoint: aws.String("https://bedrock.test"),
|
|
Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(
|
|
"access-key", "secret-key", "",
|
|
)),
|
|
HTTPClient: httpClient,
|
|
Retryer: aws.NopRetryer{},
|
|
})
|
|
}
|
|
|
|
func newAwsTestContext(writer http.ResponseWriter, requestContext context.Context) *gin.Context {
|
|
c, _ := gin.CreateTestContext(writer)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(requestContext)
|
|
return c
|
|
}
|
|
|
|
func newAwsTestRelayInfo() *relaycommon.RelayInfo {
|
|
return &relaycommon.RelayInfo{
|
|
StartTime: time.Now(),
|
|
IsStream: true,
|
|
OriginModelName: awsTestModel,
|
|
RelayFormat: relaytypes.RelayFormatOpenAI,
|
|
ShouldIncludeUsage: true,
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
UpstreamModelName: awsTestModel,
|
|
},
|
|
}
|
|
}
|
|
|
|
func newAwsInvokeModelInput() *bedrockruntime.InvokeModelInput {
|
|
return &bedrockruntime.InvokeModelInput{
|
|
ModelId: aws.String(awsTestModel),
|
|
Body: []byte(`{}`),
|
|
Accept: aws.String("application/json"),
|
|
ContentType: aws.String("application/json"),
|
|
}
|
|
}
|
|
|
|
func newAwsStreamInput() *bedrockruntime.InvokeModelWithResponseStreamInput {
|
|
return &bedrockruntime.InvokeModelWithResponseStreamInput{
|
|
ModelId: aws.String(awsTestModel),
|
|
Body: []byte(`{}`),
|
|
Accept: aws.String("application/json"),
|
|
ContentType: aws.String("application/json"),
|
|
}
|
|
}
|
|
|
|
func writeAwsStreamEvent(writer io.Writer, data string) error {
|
|
payload, err := common.Marshal(struct {
|
|
Bytes []byte `json:"bytes"`
|
|
}{Bytes: []byte(data)})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return eventstream.NewEncoder().Encode(writer, eventstream.Message{
|
|
Headers: eventstream.Headers{
|
|
{Name: eventstreamapi.MessageTypeHeader, Value: eventstream.StringValue(eventstreamapi.EventMessageType)},
|
|
{Name: eventstreamapi.EventTypeHeader, Value: eventstream.StringValue("chunk")},
|
|
{Name: eventstreamapi.ContentTypeHeader, Value: eventstream.StringValue("application/json")},
|
|
},
|
|
Payload: payload,
|
|
})
|
|
}
|
|
|
|
func newAwsStreamResponse(request *http.Request, body io.ReadCloser) *http.Response {
|
|
return &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Status: "200 OK",
|
|
Header: http.Header{
|
|
"Content-Type": []string{"application/vnd.amazon.eventstream"},
|
|
"X-Amzn-Bedrock-Content-Type": []string{"application/json"},
|
|
},
|
|
Body: body,
|
|
Request: request,
|
|
}
|
|
}
|
|
|
|
func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
recorder := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(recorder)
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
|
|
|
info := &relaycommon.RelayInfo{
|
|
OriginModelName: "claude-3-5-sonnet-20240620",
|
|
IsStream: false,
|
|
UseRuntimeHeadersOverride: true,
|
|
RuntimeHeadersOverride: map[string]any{
|
|
"anthropic-beta": "computer-use-2025-01-24",
|
|
},
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ApiKey: "access-key|secret-key|us-east-1",
|
|
UpstreamModelName: "claude-3-5-sonnet-20240620",
|
|
},
|
|
}
|
|
|
|
requestBody := bytes.NewBufferString(`{"messages":[{"role":"user","content":"hello"}],"max_tokens":128}`)
|
|
adaptor := &Adaptor{}
|
|
|
|
_, err := doAwsClientRequest(ctx, info, adaptor, requestBody)
|
|
require.NoError(t, err)
|
|
|
|
awsReq, ok := adaptor.AwsReq.(*bedrockruntime.InvokeModelInput)
|
|
require.True(t, ok)
|
|
|
|
var payload map[string]any
|
|
require.NoError(t, common.Unmarshal(awsReq.Body, &payload))
|
|
|
|
anthropicBeta, exists := payload["anthropic_beta"]
|
|
require.True(t, exists)
|
|
|
|
values, ok := anthropicBeta.([]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, []any{"computer-use-2025-01-24"}, values)
|
|
}
|
|
|
|
func TestDoAwsClientRequest_ForwardsThinkingBlockBindingAndBetaList(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
|
ctx.Request.Header.Set("anthropic-beta", "fast-mode-2026-02-01, thinking-binding-controls-2026-08-01")
|
|
|
|
info := &relaycommon.RelayInfo{
|
|
OriginModelName: "claude-opus-5-5",
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ApiKey: "access-key|secret-key|us-east-1",
|
|
UpstreamModelName: "claude-opus-5-5",
|
|
},
|
|
}
|
|
requestBody := bytes.NewBufferString(`{"max_tokens":128,"messages":[{"role":"user","content":"hello"}],` +
|
|
`"thinking":{"type":"adaptive","display":"summarized","block_binding":{"prefix_mismatch_behavior":"drop_block"}}}`)
|
|
adaptor := &Adaptor{}
|
|
|
|
_, err := doAwsClientRequest(ctx, info, adaptor, requestBody)
|
|
require.NoError(t, err)
|
|
awsReq, ok := adaptor.AwsReq.(*bedrockruntime.InvokeModelInput)
|
|
require.True(t, ok)
|
|
|
|
var payload struct {
|
|
AnthropicBeta []string `json:"anthropic_beta"`
|
|
Thinking json.RawMessage `json:"thinking"`
|
|
}
|
|
require.NoError(t, common.Unmarshal(awsReq.Body, &payload))
|
|
assert.Equal(t, []string{"fast-mode-2026-02-01", "thinking-binding-controls-2026-08-01"}, payload.AnthropicBeta)
|
|
assert.JSONEq(t, `{"type":"adaptive","display":"summarized","block_binding":{"prefix_mismatch_behavior":"drop_block"}}`, string(payload.Thinking))
|
|
}
|
|
|
|
func TestNewAwsInvokeContextInheritsParent(t *testing.T) {
|
|
originalRelayTimeout := common.RelayTimeout
|
|
t.Cleanup(func() {
|
|
common.RelayTimeout = originalRelayTimeout
|
|
})
|
|
|
|
tests := []struct {
|
|
name string
|
|
relayTimeout int
|
|
wantDeadline bool
|
|
}{
|
|
{name: "without relay timeout", relayTimeout: 0, wantDeadline: false},
|
|
{name: "with relay timeout", relayTimeout: 30, wantDeadline: true},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
common.RelayTimeout = test.relayTimeout
|
|
parent, cancelParent := context.WithCancel(context.Background())
|
|
invokeContext, cancelInvoke := newAwsInvokeContext(parent)
|
|
defer cancelInvoke()
|
|
|
|
_, hasDeadline := invokeContext.Deadline()
|
|
assert.Equal(t, test.wantDeadline, hasDeadline)
|
|
|
|
cancelParent()
|
|
require.ErrorIs(t, invokeContext.Err(), context.Canceled)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNewAwsInvokeErrorSkipsRetryOnlyForClientCancellation(t *testing.T) {
|
|
canceledContext, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
requestContext context.Context
|
|
err error
|
|
wantSkipRetry bool
|
|
}{
|
|
{
|
|
name: "client context canceled",
|
|
requestContext: canceledContext,
|
|
err: context.Canceled,
|
|
wantSkipRetry: true,
|
|
},
|
|
{
|
|
name: "relay timeout with live client context",
|
|
requestContext: context.Background(),
|
|
err: context.DeadlineExceeded,
|
|
wantSkipRetry: false,
|
|
},
|
|
{
|
|
name: "upstream error with live client context",
|
|
requestContext: context.Background(),
|
|
err: errors.New("upstream failed"),
|
|
wantSkipRetry: false,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
err := newAwsInvokeError(test.requestContext, test.err, "InvokeModel")
|
|
assert.Equal(t, test.wantSkipRetry, relaytypes.IsSkipRetryError(err))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAwsHandlersCancelSdkRequestAndSkipRetry(t *testing.T) {
|
|
originalRelayTimeout := common.RelayTimeout
|
|
common.RelayTimeout = 0
|
|
t.Cleanup(func() {
|
|
common.RelayTimeout = originalRelayTimeout
|
|
})
|
|
|
|
tests := []struct {
|
|
name string
|
|
request any
|
|
handle func(*gin.Context, *relaycommon.RelayInfo, *Adaptor) (*relaytypes.NewAPIError, *dto.Usage)
|
|
}{
|
|
{name: "non-stream", request: newAwsInvokeModelInput(), handle: awsHandler},
|
|
{name: "stream", request: newAwsStreamInput(), handle: awsStreamHandler},
|
|
{name: "nova", request: newAwsInvokeModelInput(), handle: handleNovaRequest},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
requestContext, cancelRequest := context.WithCancel(context.Background())
|
|
t.Cleanup(cancelRequest)
|
|
|
|
upstreamContexts := make(chan context.Context, 1)
|
|
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
|
|
upstreamContexts <- request.Context()
|
|
<-request.Context().Done()
|
|
return nil, request.Context().Err()
|
|
}))
|
|
adaptor := &Adaptor{AwsClient: client, AwsReq: test.request}
|
|
c := newAwsTestContext(httptest.NewRecorder(), requestContext)
|
|
info := newAwsTestRelayInfo()
|
|
|
|
type handlerResult struct {
|
|
err *relaytypes.NewAPIError
|
|
usage *dto.Usage
|
|
}
|
|
results := make(chan handlerResult, 1)
|
|
go func() {
|
|
err, usage := test.handle(c, info, adaptor)
|
|
results <- handlerResult{err: err, usage: usage}
|
|
}()
|
|
|
|
var upstreamContext context.Context
|
|
select {
|
|
case upstreamContext = <-upstreamContexts:
|
|
case result := <-results:
|
|
t.Fatalf("handler returned before issuing AWS request: %v", result.err)
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("AWS request did not start")
|
|
}
|
|
|
|
cancelRequest()
|
|
|
|
var result handlerResult
|
|
select {
|
|
case result = <-results:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("handler did not stop after client cancellation")
|
|
}
|
|
|
|
require.ErrorIs(t, upstreamContext.Err(), context.Canceled)
|
|
require.NotNil(t, result.err)
|
|
assert.True(t, relaytypes.IsSkipRetryError(result.err))
|
|
assert.Nil(t, result.usage)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAwsStreamHandlerUsesFinalUpstreamUsage(t *testing.T) {
|
|
originalRelayTimeout := common.RelayTimeout
|
|
common.RelayTimeout = 0
|
|
t.Cleanup(func() {
|
|
common.RelayTimeout = originalRelayTimeout
|
|
})
|
|
|
|
events := []string{
|
|
`{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`,
|
|
`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
|
`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`,
|
|
`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`,
|
|
`{"type":"message_stop"}`,
|
|
}
|
|
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
|
|
var body bytes.Buffer
|
|
for _, event := range events {
|
|
if err := writeAwsStreamEvent(&body, event); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return newAwsStreamResponse(request, io.NopCloser(bytes.NewReader(body.Bytes()))), nil
|
|
}))
|
|
adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()}
|
|
recorder := httptest.NewRecorder()
|
|
c := newAwsTestContext(recorder, context.Background())
|
|
|
|
handlerErr, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor)
|
|
|
|
require.Nil(t, handlerErr)
|
|
require.NotNil(t, usage)
|
|
require.NotNil(t, usage.BillingUsage)
|
|
require.NotNil(t, usage.BillingUsage.ClaudeUsage)
|
|
assert.Equal(t, 100, usage.BillingUsage.ClaudeUsage.InputTokens)
|
|
assert.Equal(t, 423, usage.BillingUsage.ClaudeUsage.OutputTokens)
|
|
assert.Contains(t, recorder.Body.String(), "[DONE]")
|
|
}
|
|
|
|
func TestAwsStreamHandlerStopsAtClientCancellation(t *testing.T) {
|
|
originalRelayTimeout := common.RelayTimeout
|
|
common.RelayTimeout = 0
|
|
t.Cleanup(func() {
|
|
common.RelayTimeout = originalRelayTimeout
|
|
})
|
|
|
|
requestContext, cancelRequest := context.WithCancel(context.Background())
|
|
t.Cleanup(cancelRequest)
|
|
releaseFinal := make(chan struct{})
|
|
var releaseFinalOnce sync.Once
|
|
release := func() {
|
|
releaseFinalOnce.Do(func() {
|
|
close(releaseFinal)
|
|
})
|
|
}
|
|
t.Cleanup(release)
|
|
|
|
producerResults := make(chan error, 1)
|
|
upstreamContexts := make(chan context.Context, 1)
|
|
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
|
|
upstreamContexts <- request.Context()
|
|
reader, writer := io.Pipe()
|
|
go func() {
|
|
defer writer.Close()
|
|
initialEvents := []string{
|
|
`{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`,
|
|
`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
|
`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`,
|
|
}
|
|
for _, event := range initialEvents {
|
|
if err := writeAwsStreamEvent(writer, event); err != nil {
|
|
producerResults <- err
|
|
return
|
|
}
|
|
}
|
|
|
|
<-releaseFinal
|
|
producerResults <- writeAwsStreamEvent(writer, `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`)
|
|
}()
|
|
return newAwsStreamResponse(request, reader), nil
|
|
}))
|
|
|
|
responseWriter := newAwsNotifyingResponseWriter("partial")
|
|
c := newAwsTestContext(responseWriter, requestContext)
|
|
adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()}
|
|
|
|
type handlerResult struct {
|
|
err *relaytypes.NewAPIError
|
|
usage *dto.Usage
|
|
}
|
|
results := make(chan handlerResult, 1)
|
|
go func() {
|
|
err, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor)
|
|
results <- handlerResult{err: err, usage: usage}
|
|
}()
|
|
|
|
var upstreamContext context.Context
|
|
select {
|
|
case upstreamContext = <-upstreamContexts:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("AWS stream request did not start")
|
|
}
|
|
|
|
var bodyLengthBeforeCancel int
|
|
select {
|
|
case bodyLengthBeforeCancel = <-responseWriter.notified:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("partial response was not written")
|
|
}
|
|
cancelRequest()
|
|
|
|
var result handlerResult
|
|
select {
|
|
case result = <-results:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("stream handler did not stop after client cancellation")
|
|
}
|
|
|
|
require.ErrorIs(t, upstreamContext.Err(), context.Canceled)
|
|
require.Nil(t, result.err)
|
|
require.NotNil(t, result.usage)
|
|
assert.Equal(t, bodyLengthBeforeCancel, responseWriter.Body.Len())
|
|
assert.NotContains(t, responseWriter.Body.String(), "[DONE]")
|
|
|
|
release()
|
|
select {
|
|
case producerErr := <-producerResults:
|
|
require.Error(t, producerErr)
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("upstream producer did not observe the closed stream")
|
|
}
|
|
}
|
|
|
|
func newMantleTestInfo() *relaycommon.RelayInfo {
|
|
info := newAwsTestRelayInfo()
|
|
info.IsStream = false
|
|
info.RelayFormat = relaytypes.RelayFormatClaude
|
|
info.UpstreamModelName = "claude-opus-4-7"
|
|
info.ApiKey = "test-api-key|us-east-1"
|
|
info.ChannelOtherSettings.AwsApiType = dto.AwsApiTypeMantle
|
|
return info
|
|
}
|
|
|
|
func TestMantleRequestPreservesNativeFieldsAndSignsFinalPayload(t *testing.T) {
|
|
for _, key := range []string{"test-api-key|us-east-1", "access-key|secret-key|eu-west-1"} {
|
|
t.Run(strings.Split(key, "|")[0], func(t *testing.T) {
|
|
info := newMantleTestInfo()
|
|
info.ApiKey = key
|
|
info.UseRuntimeHeadersOverride = true
|
|
info.RuntimeHeadersOverride = map[string]any{"anthropic-beta": "test-beta", "X-Test": "override"}
|
|
info.ChannelSetting.PassThroughBodyEnabled = true
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
c := newAwsTestContext(httptest.NewRecorder(), ctx)
|
|
c.Request.Header.Set("Authorization", "Bearer client-secret")
|
|
c.Request.Header.Set("x-api-key", "client-secret")
|
|
a := &Adaptor{}
|
|
body := `{"model":"client-alias","stream":true,"max_tokens":128,"temperature":0,"cache_control":{"type":"ephemeral","ttl":"1h"},"messages":[{"role":"user","content":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}]}],"future_option":{"enabled":false},"anthropic_version":"bedrock-2023-05-31"}`
|
|
req, err := a.newMantleRequest(c, info, strings.NewReader(body))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "/anthropic/v1/messages", req.URL.Path)
|
|
assert.Equal(t, http.MethodPost, req.Method)
|
|
assert.Equal(t, "2023-06-01", req.Header.Get("anthropic-version"))
|
|
assert.Equal(t, "test-beta", req.Header.Get("anthropic-beta"))
|
|
assert.Equal(t, "override", req.Header.Get("X-Test"))
|
|
payload, err := io.ReadAll(req.Body)
|
|
require.NoError(t, err)
|
|
var data map[string]json.RawMessage
|
|
require.NoError(t, common.Unmarshal(payload, &data))
|
|
assert.JSONEq(t, `"anthropic.claude-opus-4-7"`, string(data["model"]))
|
|
assert.JSONEq(t, `{"type":"ephemeral","ttl":"1h"}`, string(data["cache_control"]))
|
|
assert.JSONEq(t, `[{"role":"user","content":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}]}]`, string(data["messages"]))
|
|
assert.JSONEq(t, `{"enabled":false}`, string(data["future_option"]))
|
|
assert.Equal(t, "0", string(data["temperature"]))
|
|
assert.Equal(t, "true", string(data["stream"]))
|
|
assert.NotContains(t, data, "anthropic_version")
|
|
replay, err := req.GetBody()
|
|
require.NoError(t, err)
|
|
defer replay.Close()
|
|
replayed, err := io.ReadAll(replay)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, payload, replayed)
|
|
if strings.HasPrefix(key, "access-key|") {
|
|
assert.Empty(t, req.Header.Get("x-api-key"))
|
|
auth := req.Header.Get("Authorization")
|
|
assert.Contains(t, auth, "/eu-west-1/bedrock-mantle/aws4_request")
|
|
assert.NotContains(t, auth, "secret-key")
|
|
signedAt, err := time.Parse("20060102T150405Z", req.Header.Get("X-Amz-Date"))
|
|
require.NoError(t, err)
|
|
verification := req.Clone(ctx)
|
|
verification.Header.Del("Authorization")
|
|
hash := sha256.Sum256(payload)
|
|
err = v4.NewSigner().SignHTTP(ctx, aws.Credentials{AccessKeyID: "access-key", SecretAccessKey: "secret-key"}, verification, hex.EncodeToString(hash[:]), "bedrock-mantle", "eu-west-1", signedAt)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, auth, verification.Header.Get("Authorization"))
|
|
} else {
|
|
assert.Equal(t, "test-api-key", req.Header.Get("x-api-key"))
|
|
assert.Empty(t, req.Header.Get("Authorization"))
|
|
}
|
|
cancel()
|
|
assert.ErrorIs(t, req.Context().Err(), context.Canceled)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMantleEndpointAndCredentialValidation(t *testing.T) {
|
|
for _, test := range []struct {
|
|
name, key, base, want string
|
|
invalid bool
|
|
}{
|
|
{name: "regional default", key: "key|us-west-2", want: "https://bedrock-mantle.us-west-2.api.aws/anthropic/v1/messages"},
|
|
{name: "proxy root", key: "ak|sk|us-east-1", base: "https://proxy.example/gateway/", want: "https://proxy.example/gateway/anthropic/v1/messages"},
|
|
{name: "anthropic root", key: "key|us-east-1", base: "https://proxy.example/anthropic/", want: "https://proxy.example/anthropic/v1/messages"},
|
|
{name: "version root", key: "key|us-east-1", base: "https://proxy.example/anthropic/v1", want: "https://proxy.example/anthropic/v1/messages"},
|
|
{name: "full endpoint", key: "key|us-east-1", base: "https://proxy.example/anthropic/v1/messages", want: "https://proxy.example/anthropic/v1/messages"},
|
|
{name: "missing region", key: "secret", invalid: true},
|
|
{name: "empty key", key: "|us-east-1", invalid: true},
|
|
{name: "empty secret", key: "ak||us-east-1", invalid: true},
|
|
{name: "injected host", key: "key|us-east-1.evil.example/", invalid: true},
|
|
{name: "plaintext URL", key: "key|us-east-1", base: "http://proxy.example", invalid: true},
|
|
{name: "URL credentials", key: "key|us-east-1", base: "https://user:secret@proxy.example", invalid: true},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
info := newMantleTestInfo()
|
|
info.ApiKey, info.ChannelBaseUrl = test.key, test.base
|
|
endpoint, err := (&Adaptor{}).GetRequestURL(info)
|
|
if test.invalid {
|
|
require.Error(t, err)
|
|
assert.NotContains(t, err.Error(), "secret")
|
|
assert.Empty(t, endpoint)
|
|
return
|
|
}
|
|
require.NoError(t, err)
|
|
assert.Equal(t, test.want, endpoint)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMantleHTTPResponsesPreserveCacheUsage(t *testing.T) {
|
|
originalTimeout := constant.StreamingTimeout
|
|
constant.StreamingTimeout = 30
|
|
t.Cleanup(func() { constant.StreamingTimeout = originalTimeout })
|
|
const usageJSON = `{"input_tokens":100,"output_tokens":23,"cache_read_input_tokens":1000,"cache_creation_input_tokens":50,"cache_creation":{"ephemeral_5m_input_tokens":20,"ephemeral_1h_input_tokens":30}}`
|
|
for _, stream := range []bool{false, true} {
|
|
for _, format := range []relaytypes.RelayFormat{relaytypes.RelayFormatClaude, relaytypes.RelayFormatOpenAI} {
|
|
t.Run(fmt.Sprintf("stream=%t/format=%s", stream, format), func(t *testing.T) {
|
|
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if stream {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
fmt.Fprintf(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_mantle\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"anthropic.claude-opus-4-7\",\"content\":[],\"usage\":%s}}\n\n", usageJSON)
|
|
fmt.Fprint(w, "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n")
|
|
fmt.Fprint(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":23}}\n\n")
|
|
fmt.Fprint(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
fmt.Fprintf(w, `{"id":"msg_mantle","type":"message","role":"assistant","model":"anthropic.claude-opus-4-7","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","usage":%s}`, usageJSON)
|
|
}))
|
|
defer server.Close()
|
|
original := http.DefaultClient
|
|
http.DefaultClient = server.Client()
|
|
defer func() { http.DefaultClient = original }()
|
|
info := newMantleTestInfo()
|
|
info.ChannelBaseUrl, info.IsStream, info.RelayFormat = server.URL, stream, format
|
|
recorder := httptest.NewRecorder()
|
|
c := newAwsTestContext(recorder, context.Background())
|
|
a := &Adaptor{}
|
|
a.Init(info)
|
|
var request dto.ClaudeRequest
|
|
require.NoError(t, common.UnmarshalJsonStr(`{"model":"claude-opus-4-7","max_tokens":128,"messages":[{"role":"user","content":"hello"}],"cache_control":{"type":"ephemeral"}}`, &request))
|
|
request.Stream = &stream
|
|
converted, err := a.ConvertClaudeRequest(c, info, &request)
|
|
require.NoError(t, err)
|
|
body, err := common.Marshal(converted)
|
|
require.NoError(t, err)
|
|
response, err := a.DoRequest(c, info, bytes.NewReader(body))
|
|
require.NoError(t, err)
|
|
httpResponse, ok := response.(*http.Response)
|
|
require.True(t, ok)
|
|
result, apiErr := a.DoResponse(c, httpResponse, info)
|
|
require.Nil(t, apiErr)
|
|
usage, ok := result.(*dto.Usage)
|
|
require.True(t, ok)
|
|
require.NotNil(t, usage.BillingUsage)
|
|
require.NotNil(t, usage.BillingUsage.ClaudeUsage)
|
|
claudeUsage := usage.BillingUsage.ClaudeUsage
|
|
assert.Equal(t, 100, claudeUsage.InputTokens)
|
|
assert.Equal(t, 23, claudeUsage.OutputTokens)
|
|
assert.Equal(t, 1000, claudeUsage.CacheReadInputTokens)
|
|
assert.Equal(t, 50, claudeUsage.CacheCreationInputTokens)
|
|
assert.Equal(t, 20, claudeUsage.GetCacheCreation5mTokens())
|
|
assert.Equal(t, 30, claudeUsage.GetCacheCreation1hTokens())
|
|
assert.Contains(t, recorder.Body.String(), "hello")
|
|
if stream && format == relaytypes.RelayFormatOpenAI {
|
|
assert.Contains(t, recorder.Body.String(), "[DONE]")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestMantleDoesNotFollowRedirectsOrFallbackOnAuthFailure(t *testing.T) {
|
|
for _, status := range []int{http.StatusTemporaryRedirect, http.StatusUnauthorized, http.StatusTooManyRequests} {
|
|
t.Run(fmt.Sprint(status), func(t *testing.T) {
|
|
calls := 0
|
|
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls++
|
|
w.Header().Set("Location", "/credential-leak")
|
|
w.WriteHeader(status)
|
|
}))
|
|
defer server.Close()
|
|
original := http.DefaultClient
|
|
http.DefaultClient = server.Client()
|
|
defer func() { http.DefaultClient = original }()
|
|
info := newMantleTestInfo()
|
|
info.ChannelBaseUrl = server.URL
|
|
c := newAwsTestContext(httptest.NewRecorder(), context.Background())
|
|
response, err := (&Adaptor{}).DoRequest(c, info, strings.NewReader(`{"messages":[]}`))
|
|
require.NoError(t, err)
|
|
httpResponse := response.(*http.Response)
|
|
defer httpResponse.Body.Close()
|
|
assert.Equal(t, status, httpResponse.StatusCode)
|
|
assert.Equal(t, 1, calls)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAWSLegacyInterfaceRemainsDefault(t *testing.T) {
|
|
for _, mode := range []dto.AwsApiType{"", dto.AwsApiTypeInvokeModel} {
|
|
for _, stream := range []bool{false, true} {
|
|
t.Run(fmt.Sprintf("%s/stream=%t", mode, stream), func(t *testing.T) {
|
|
info := newMantleTestInfo()
|
|
info.ChannelOtherSettings.AwsApiType = mode
|
|
info.IsStream = stream
|
|
info.UpstreamModelName = "claude-3-5-sonnet-20240620"
|
|
c := newAwsTestContext(httptest.NewRecorder(), context.Background())
|
|
a := &Adaptor{}
|
|
a.Init(info)
|
|
response, err := a.DoRequest(c, info, strings.NewReader(`{"max_tokens":128,"messages":[{"role":"user","content":[{"type":"text","text":"hi","cache_control":{"type":"ephemeral"}}]}],"cache_control":{"type":"ephemeral"}}`))
|
|
require.NoError(t, err)
|
|
assert.Nil(t, response)
|
|
var body []byte
|
|
if stream {
|
|
input, ok := a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput)
|
|
require.True(t, ok)
|
|
body = input.Body
|
|
assert.Equal(t, "us."+awsTestModel, *input.ModelId)
|
|
} else {
|
|
input, ok := a.AwsReq.(*bedrockruntime.InvokeModelInput)
|
|
require.True(t, ok)
|
|
body = input.Body
|
|
assert.Equal(t, "us."+awsTestModel, *input.ModelId)
|
|
}
|
|
var payload map[string]json.RawMessage
|
|
require.NoError(t, common.Unmarshal(body, &payload))
|
|
assert.Equal(t, `"bedrock-2023-05-31"`, string(payload["anthropic_version"]))
|
|
assert.NotContains(t, payload, "cache_control")
|
|
assert.Contains(t, string(payload["messages"]), "cache_control")
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestMantleRejectsInvalidRequestsAndPreservesExplicitModelIDs(t *testing.T) {
|
|
for _, test := range []struct {
|
|
name, model, body, wantModel string
|
|
invalid bool
|
|
}{
|
|
{name: "legacy Haiku alias", model: "claude-haiku-4-5-20251001", body: `{"messages":[]}`, wantModel: "anthropic.claude-haiku-4-5"},
|
|
{name: "short Haiku alias", model: "claude-haiku-4-5", body: `{"messages":[]}`, wantModel: "anthropic.claude-haiku-4-5"},
|
|
{name: "explicit Haiku model", model: "anthropic.claude-haiku-4-5", body: `{"messages":[]}`},
|
|
{name: "explicit Haiku inference profile", model: "us.anthropic.claude-haiku-4-5", body: `{"messages":[]}`},
|
|
{name: "explicit regional model", model: "us.anthropic.claude-opus-4-7", body: `{"messages":[]}`},
|
|
{name: "explicit model", model: "anthropic.claude-opus-4-7", body: `{"messages":[]}`},
|
|
{name: "Nova requires legacy", model: "nova-pro-v1:0", body: `{"messages":[]}`, invalid: true},
|
|
{name: "null body", model: "claude-opus-4-7", body: `null`, invalid: true},
|
|
{name: "malformed JSON", model: "claude-opus-4-7", body: `{`, invalid: true},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
info := newMantleTestInfo()
|
|
info.UpstreamModelName = test.model
|
|
c := newAwsTestContext(httptest.NewRecorder(), context.Background())
|
|
req, err := (&Adaptor{}).newMantleRequest(c, info, strings.NewReader(test.body))
|
|
if test.invalid {
|
|
require.Error(t, err)
|
|
assert.Nil(t, req)
|
|
return
|
|
}
|
|
require.NoError(t, err)
|
|
defer req.Body.Close()
|
|
var payload map[string]any
|
|
require.NoError(t, common.DecodeJson(req.Body, &payload))
|
|
wantModel := test.wantModel
|
|
if wantModel == "" {
|
|
wantModel = test.model
|
|
}
|
|
assert.Equal(t, wantModel, payload["model"])
|
|
assert.NotContains(t, payload, "cache_control")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMantleSignatureMatchesWireHeaders(t *testing.T) {
|
|
for _, test := range []struct {
|
|
name string
|
|
http2 bool
|
|
allHopHeaders bool
|
|
}{
|
|
{name: "HTTP2 keep-alive", http2: true},
|
|
{name: "HTTP1 hop-by-hop headers", allHopHeaders: true},
|
|
{name: "HTTP2 hop-by-hop headers", http2: true, allHopHeaders: true},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
type receivedRequest struct {
|
|
request *http.Request
|
|
body []byte
|
|
err error
|
|
}
|
|
received := make(chan receivedRequest, 1)
|
|
server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
body, err := io.ReadAll(r.Body)
|
|
received <- receivedRequest{request: r.Clone(context.Background()), body: body, err: err}
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
server.EnableHTTP2 = test.http2
|
|
server.StartTLS()
|
|
defer server.Close()
|
|
originalClient := http.DefaultClient
|
|
http.DefaultClient = server.Client()
|
|
defer func() { http.DefaultClient = originalClient }()
|
|
|
|
info := newMantleTestInfo()
|
|
info.ChannelBaseUrl = server.URL
|
|
info.ApiKey = "access-key|secret-key|us-east-1"
|
|
info.UseRuntimeHeadersOverride = true
|
|
info.RuntimeHeadersOverride = map[string]any{
|
|
"Connection": "keep-alive",
|
|
"anthropic-beta": "test-beta",
|
|
"X-Test": "end-to-end",
|
|
}
|
|
if test.allHopHeaders {
|
|
maps.Copy(info.RuntimeHeadersOverride, map[string]any{
|
|
"Connection": "keep-alive, x-hop-token",
|
|
"X-Hop-Token": "hop-only",
|
|
"Keep-Alive": "timeout=5",
|
|
"Proxy-Connection": "keep-alive",
|
|
"Proxy-Authenticate": "Basic",
|
|
"Proxy-Authorization": "Basic test-proxy-token",
|
|
"Te": "trailers",
|
|
"Trailer": "X-Checksum",
|
|
"Transfer-Encoding": "chunked",
|
|
"Upgrade": "websocket",
|
|
})
|
|
}
|
|
c := newAwsTestContext(httptest.NewRecorder(), context.Background())
|
|
response, err := (&Adaptor{}).DoRequest(c, info, strings.NewReader(`{"messages":[],"max_tokens":128,"cache_control":{"type":"ephemeral","ttl":"1h"}}`))
|
|
require.NoError(t, err)
|
|
httpResponse, ok := response.(*http.Response)
|
|
require.True(t, ok)
|
|
defer httpResponse.Body.Close()
|
|
wire := <-received
|
|
require.NoError(t, wire.err)
|
|
wantProtocol := 1
|
|
if test.http2 {
|
|
wantProtocol = 2
|
|
}
|
|
require.Equal(t, wantProtocol, wire.request.ProtoMajor)
|
|
for _, name := range []string{"Connection", "X-Hop-Token", "Keep-Alive", "Proxy-Connection", "Proxy-Authenticate", "Proxy-Authorization", "Te", "Trailer", "Transfer-Encoding", "Upgrade"} {
|
|
assert.Empty(t, wire.request.Header.Values(name), "%s must not reach the upstream", name)
|
|
}
|
|
assert.Equal(t, "test-beta", wire.request.Header.Get("anthropic-beta"))
|
|
assert.Equal(t, "end-to-end", wire.request.Header.Get("X-Test"))
|
|
|
|
// Recompute the signature from only the bytes and signed headers that
|
|
// reached the server, not from the pre-transport request object.
|
|
authorization := wire.request.Header.Get("Authorization")
|
|
_, signedHeaders, found := strings.Cut(authorization, "SignedHeaders=")
|
|
require.True(t, found)
|
|
signedHeaders, _, found = strings.Cut(signedHeaders, ",")
|
|
require.True(t, found)
|
|
verification := wire.request.Clone(context.Background())
|
|
verification.URL.Scheme = "https"
|
|
verification.URL.Host = wire.request.Host
|
|
verification.Header = make(http.Header)
|
|
for name := range strings.SplitSeq(signedHeaders, ";") {
|
|
if values, ok := wire.request.Header[http.CanonicalHeaderKey(name)]; ok {
|
|
verification.Header[http.CanonicalHeaderKey(name)] = values
|
|
}
|
|
}
|
|
signedAt, err := time.Parse("20060102T150405Z", wire.request.Header.Get("X-Amz-Date"))
|
|
require.NoError(t, err)
|
|
hash := sha256.Sum256(wire.body)
|
|
err = v4.NewSigner().SignHTTP(context.Background(), aws.Credentials{AccessKeyID: "access-key", SecretAccessKey: "secret-key"}, verification, hex.EncodeToString(hash[:]), "bedrock-mantle", "us-east-1", signedAt)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, authorization, verification.Header.Get("Authorization"))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAWSLegacyHaikuMappingPreservesVersionAndRegion(t *testing.T) {
|
|
for _, test := range []struct {
|
|
name string
|
|
mode dto.AwsApiType
|
|
stream bool
|
|
key string
|
|
}{
|
|
{name: "default API key", key: "test-api-key|us-east-1"},
|
|
{name: "default API key stream", stream: true, key: "test-api-key|us-east-1"},
|
|
{name: "default AK/SK", key: "access-key|secret-key|us-east-1"},
|
|
{name: "default AK/SK stream", stream: true, key: "access-key|secret-key|us-east-1"},
|
|
{name: "explicit API key", mode: dto.AwsApiTypeInvokeModel, key: "test-api-key|us-east-1"},
|
|
{name: "explicit API key stream", mode: dto.AwsApiTypeInvokeModel, stream: true, key: "test-api-key|us-east-1"},
|
|
{name: "explicit AK/SK", mode: dto.AwsApiTypeInvokeModel, key: "access-key|secret-key|us-east-1"},
|
|
{name: "explicit AK/SK stream", mode: dto.AwsApiTypeInvokeModel, stream: true, key: "access-key|secret-key|us-east-1"},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
info := newMantleTestInfo()
|
|
info.ChannelOtherSettings.AwsApiType = test.mode
|
|
info.ApiKey, info.IsStream = test.key, test.stream
|
|
info.UpstreamModelName = "claude-haiku-4-5-20251001"
|
|
c := newAwsTestContext(httptest.NewRecorder(), context.Background())
|
|
a := &Adaptor{}
|
|
a.Init(info)
|
|
response, err := a.DoRequest(c, info, strings.NewReader(`{"messages":[{"role":"user","content":"hello"}],"max_tokens":128}`))
|
|
require.NoError(t, err)
|
|
assert.Nil(t, response)
|
|
const wantModel = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
if test.stream {
|
|
request, ok := a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput)
|
|
require.True(t, ok)
|
|
assert.Equal(t, wantModel, *request.ModelId)
|
|
} else {
|
|
request, ok := a.AwsReq.(*bedrockruntime.InvokeModelInput)
|
|
require.True(t, ok)
|
|
assert.Equal(t, wantModel, *request.ModelId)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMantleCanceledRequestSkipsRetry(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
c := newAwsTestContext(httptest.NewRecorder(), ctx)
|
|
response, err := (&Adaptor{}).DoRequest(c, newMantleTestInfo(), strings.NewReader(`{"messages":[]}`))
|
|
require.Error(t, err)
|
|
assert.Nil(t, response)
|
|
var apiError *relaytypes.NewAPIError
|
|
require.ErrorAs(t, err, &apiError)
|
|
assert.True(t, relaytypes.IsSkipRetryError(apiError))
|
|
}
|