Skip to content

GitLab

Commit 11f7b835 authored by Ylarod's avatar Ylarod
Browse files

sub2api: add bedrock support

parent 1ee98447
Expand all Hide whitespace changes
Inline Side-by-side
backend/go.mod
View file @ 11f7b835
......@@ -7,7 +7,7 @@ require (
github.com/DATA-DOG/go-sqlmock v1.5.2
github.com/DouDOU-start/go-sora2api v1.1.0
github.com/alitto/pond/v2 v2.6.2
github.com/aws/aws-sdk-go-v2 v1.41.2
github.com/aws/aws-sdk-go-v2 v1.41.3
github.com/aws/aws-sdk-go-v2/config v1.32.10
github.com/aws/aws-sdk-go-v2/credentials v1.19.10
github.com/aws/aws-sdk-go-v2/service/s3 v1.96.2
......@@ -66,7 +66,7 @@ require (
github.com/aws/aws-sdk-go-v2/service/sso v1.30.11 // indirect
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.15 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.41.7 // indirect
github.com/aws/smithy-go v1.24.1 // indirect
github.com/aws/smithy-go v1.24.2 // indirect
github.com/bdandy/go-errors v1.2.2 // indirect
github.com/bdandy/go-socks4 v1.2.3 // indirect
github.com/bmatcuk/doublestar v1.3.4 // indirect
......
backend/go.sum
View file @ 11f7b835
......@@ -24,6 +24,8 @@ github.com/apparentlymart/go-textseg/v15 v15.0.0 h1:uYvfpb3DyLSCGWnctWKGj857c6ew
github.com/apparentlymart/go-textseg/v15 v15.0.0/go.mod h1:K8XmNZdhEBkdlyDdvbmmsvpAG721bKi0joRfFdHIWJ4=
github.com/aws/aws-sdk-go-v2 v1.41.2 h1:LuT2rzqNQsauaGkPK/7813XxcZ3o3yePY0Iy891T2ls=
github.com/aws/aws-sdk-go-v2 v1.41.2/go.mod h1:IvvlAZQXvTXznUPfRVfryiG1fbzE2NGK6m9u39YQ+S4=
github.com/aws/aws-sdk-go-v2 v1.41.3 h1:4kQ/fa22KjDt13QCy1+bYADvdgcxpfH18f0zP542kZA=
github.com/aws/aws-sdk-go-v2 v1.41.3/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.5 h1:zWFmPmgw4sveAYi1mRqG+E/g0461cJ5M4bJ8/nc6d3Q=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.5/go.mod h1:nVUlMLVV8ycXSb7mSkcNu9e3v/1TJq2RTlrPwhYWr5c=
github.com/aws/aws-sdk-go-v2/config v1.32.10 h1:9DMthfO6XWZYLfzZglAgW5Fyou2nRI5CuV44sTedKBI=
......@@ -60,6 +62,8 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.41.7 h1:NITQpgo9A5NrDZ57uOWj+abvXSb8
github.com/aws/aws-sdk-go-v2/service/sts v1.41.7/go.mod h1:sks5UWBhEuWYDPdwlnRFn1w7xWdH29Jcpe+/PJQefEs=
github.com/aws/smithy-go v1.24.1 h1:VbyeNfmYkWoxMVpGUAbQumkODcYmfMRfZ8yQiH30SK0=
github.com/aws/smithy-go v1.24.1/go.mod h1:LEj2LM3rBRQJxPZTB4KuzZkaZYnZPnvgIhb4pu07mx0=
github.com/aws/smithy-go v1.24.2 h1:FzA3bu/nt/vDvmnkg+R8Xl46gmzEDam6mZ1hzmwXFng=
github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc=
github.com/bdandy/go-errors v1.2.2 h1:WdFv/oukjTJCLa79UfkGmwX7ZxONAihKu4V0mLIs11Q=
github.com/bdandy/go-errors v1.2.2/go.mod h1:NkYHl4Fey9oRRdbB1CoC6e84tuqQHiqrOcZpqFEkBxM=
github.com/bdandy/go-socks4 v1.2.3 h1:Q6Y2heY1GRjCtHbmlKfnwrKVU/k81LS8mRGLRlmDlic=
......
backend/internal/domain/constants.go
View file @ 11f7b835
......@@ -31,6 +31,8 @@ const (
AccountTypeSetupToken = "setup-token" // Setup Token类型账号（inference only scope）
AccountTypeAPIKey = "apikey" // API Key类型账号
AccountTypeUpstream = "upstream" // 上游透传类型账号（通过 Base URL + API Key 连接上游）
AccountTypeBedrock = "bedrock" // AWS Bedrock 类型账号（通过 SigV4 签名连接 Bedrock）
AccountTypeBedrockAPIKey = "bedrock-apikey" // AWS Bedrock API Key 类型账号（通过 Bearer Token 连接 Bedrock）
)
// Redeem type constants
......@@ -113,3 +115,27 @@ var DefaultAntigravityModelMapping = map[string]string{
"gpt-oss-120b-medium": "gpt-oss-120b-medium",
"tab_flash_lite_preview": "tab_flash_lite_preview",
}
// DefaultBedrockModelMapping 是 AWS Bedrock 平台的默认模型映射
// 将 Anthropic 标准模型名映射到 Bedrock 模型 ID
// 注意：此处的 "us." 前缀仅为默认值，ResolveBedrockModelID 会根据账号配置的
// aws_region 自动调整为匹配的区域前缀（如 eu.、apac.、jp. 等）
var DefaultBedrockModelMapping = map[string]string{
// Claude Opus
"claude-opus-4-6-thinking": "us.anthropic.claude-opus-4-6-v1",
"claude-opus-4-6": "us.anthropic.claude-opus-4-6-v1",
"claude-opus-4-5-thinking": "us.anthropic.claude-opus-4-5-20251101-v1:0",
"claude-opus-4-5-20251101": "us.anthropic.claude-opus-4-5-20251101-v1:0",
"claude-opus-4-1": "us.anthropic.claude-opus-4-1-20250805-v1:0",
"claude-opus-4-20250514": "us.anthropic.claude-opus-4-20250514-v1:0",
// Claude Sonnet
"claude-sonnet-4-6-thinking": "us.anthropic.claude-sonnet-4-6",
"claude-sonnet-4-6": "us.anthropic.claude-sonnet-4-6",
"claude-sonnet-4-5": "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"claude-sonnet-4-5-thinking": "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"claude-sonnet-4-5-20250929": "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"claude-sonnet-4-20250514": "us.anthropic.claude-sonnet-4-20250514-v1:0",
// Claude Haiku
"claude-haiku-4-5": "us.anthropic.claude-haiku-4-5-20251001-v1:0",
"claude-haiku-4-5-20251001": "us.anthropic.claude-haiku-4-5-20251001-v1:0",
}
backend/internal/handler/admin/account_handler.go
View file @ 11f7b835
......@@ -97,7 +97,7 @@ type CreateAccountRequest struct {
Name string `json:"name" binding:"required"`
Notes *string `json:"notes"`
Platform string `json:"platform" binding:"required"`
Type string `json:"type" binding:"required,oneof=oauth setup-token apikey upstream"`
Type string `json:"type" binding:"required,oneof=oauth setup-token apikey upstream bedrock bedrock-apikey"`
Credentials map[string]any `json:"credentials" binding:"required"`
Extra map[string]any `json:"extra"`
ProxyID *int64 `json:"proxy_id"`
......@@ -116,7 +116,7 @@ type CreateAccountRequest struct {
type UpdateAccountRequest struct {
Name string `json:"name"`
Notes *string `json:"notes"`
Type string `json:"type" binding:"omitempty,oneof=oauth setup-token apikey upstream"`
Type string `json:"type" binding:"omitempty,oneof=oauth setup-token apikey upstream bedrock bedrock-apikey"`
Credentials map[string]any `json:"credentials"`
Extra map[string]any `json:"extra"`
ProxyID *int64 `json:"proxy_id"`
......
backend/internal/service/account.go
View file @ 11f7b835
......@@ -412,6 +412,7 @@ func (a *Account) resolveModelMapping(rawMapping map[string]any) map[string]stri
if a.Platform == domain.PlatformAntigravity {
return domain.DefaultAntigravityModelMapping
}
// Bedrock 默认映射由 forwardBedrock 统一处理（需配合 region prefix 调整）
return nil
}
if len(rawMapping) == 0 {
......@@ -764,6 +765,14 @@ func (a *Account) IsInterceptWarmupEnabled() bool {
return false
}
func (a *Account) IsBedrock() bool {
return a.Platform == PlatformAnthropic && (a.Type == AccountTypeBedrock || a.Type == AccountTypeBedrockAPIKey)
}
func (a *Account) IsBedrockAPIKey() bool {
return a.Platform == PlatformAnthropic && a.Type == AccountTypeBedrockAPIKey
}
func (a *Account) IsOpenAI() bool {
return a.Platform == PlatformOpenAI
}
......
backend/internal/service/account_test_service.go
View file @ 11f7b835
......@@ -207,14 +207,14 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
testModelID = claude.DefaultTestModel
}
// For API Key accounts with model mapping, map the model
// API Key 账号测试连接时也需要应用通配符模型映射。
if account.Type == "apikey" {
mapping := account.GetModelMapping()
if len(mapping) > 0 {
if mappedModel, exists := mapping[testModelID]; exists {
testModelID = mappedModel
}
}
testModelID = account.GetMappedModel(testModelID)
}
// Bedrock accounts use a separate test path
if account.IsBedrock() {
return s.testBedrockAccountConnection(c, ctx, account, testModelID)
}
// Determine authentication method and API URL
......@@ -312,6 +312,109 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
return s.processClaudeStream(c, resp.Body)
}
// testBedrockAccountConnection tests a Bedrock (SigV4 or API Key) account using non-streaming invoke
func (s *AccountTestService) testBedrockAccountConnection(c *gin.Context, ctx context.Context, account *Account, testModelID string) error {
region := bedrockRuntimeRegion(account)
resolvedModelID, ok := ResolveBedrockModelID(account, testModelID)
if !ok {
return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Bedrock model: %s", testModelID))
}
testModelID = resolvedModelID
// Set SSE headers (test UI expects SSE)
c.Writer.Header().Set("Content-Type", "text/event-stream")
c.Writer.Header().Set("Cache-Control", "no-cache")
c.Writer.Header().Set("Connection", "keep-alive")
c.Writer.Header().Set("X-Accel-Buffering", "no")
c.Writer.Flush()
// Create a minimal Bedrock-compatible payload (no stream, no cache_control)
bedrockPayload := map[string]any{
"anthropic_version": "bedrock-2023-05-31",
"messages": []map[string]any{
{
"role": "user",
"content": []map[string]any{
{
"type": "text",
"text": "hi",
},
},
},
},
"max_tokens": 256,
"temperature": 1,
}
bedrockBody, _ := json.Marshal(bedrockPayload)
// Use non-streaming endpoint (response is standard Claude JSON)
apiURL := BuildBedrockURL(region, testModelID, false)
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(bedrockBody))
if err != nil {
return s.sendErrorAndEnd(c, "Failed to create request")
}
req.Header.Set("Content-Type", "application/json")
// Sign or set auth based on account type
if account.IsBedrockAPIKey() {
apiKey := account.GetCredential("api_key")
if apiKey == "" {
return s.sendErrorAndEnd(c, "No API key available")
}
req.Header.Set("Authorization", "Bearer "+apiKey)
} else {
signer, err := NewBedrockSignerFromAccount(account)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to create Bedrock signer: %s", err.Error()))
}
if err := signer.SignRequest(ctx, req, bedrockBody); err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to sign request: %s", err.Error()))
}
}
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
resp, err := s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, false)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Request failed: %s", err.Error()))
}
defer func() { _ = resp.Body.Close() }()
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
}
// Bedrock non-streaming response is standard Claude JSON, extract the text
var result struct {
Content []struct {
Text string `json:"text"`
} `json:"content"`
}
if err := json.Unmarshal(body, &result); err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to parse response: %s", err.Error()))
}
text := ""
if len(result.Content) > 0 {
text = result.Content[0].Text
}
if text == "" {
text = "(empty response)"
}
s.sendEvent(c, TestEvent{Type: "content", Text: text})
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
return nil
}
// testOpenAIAccountConnection tests an OpenAI account's connection
func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account *Account, modelID string) error {
ctx := c.Request.Context()
......
backend/internal/service/bedrock_request.go 0 → 100644
View file @ 11f7b835
This diff is collapsed.
backend/internal/service/bedrock_request_test.go 0 → 100644
View file @ 11f7b835
This diff is collapsed.
backend/internal/service/bedrock_signer.go 0 → 100644
View file @ 11f7b835
package service
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"net/http"
"time"
"github.com/aws/aws-sdk-go-v2/aws"
v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
)
// BedrockSigner 使用 AWS SigV4 对 Bedrock 请求签名
type BedrockSigner struct {
credentials aws.Credentials
region string
signer *v4.Signer
}
// NewBedrockSigner 创建 BedrockSigner
func NewBedrockSigner(accessKeyID, secretAccessKey, sessionToken, region string) *BedrockSigner {
return &BedrockSigner{
credentials: aws.Credentials{
AccessKeyID: accessKeyID,
SecretAccessKey: secretAccessKey,
SessionToken: sessionToken,
},
region: region,
signer: v4.NewSigner(),
}
}
// NewBedrockSignerFromAccount 从 Account 凭证创建 BedrockSigner
func NewBedrockSignerFromAccount(account *Account) (*BedrockSigner, error) {
accessKeyID := account.GetCredential("aws_access_key_id")
if accessKeyID == "" {
return nil, fmt.Errorf("aws_access_key_id not found in credentials")
}
secretAccessKey := account.GetCredential("aws_secret_access_key")
if secretAccessKey == "" {
return nil, fmt.Errorf("aws_secret_access_key not found in credentials")
}
region := account.GetCredential("aws_region")
if region == "" {
region = defaultBedrockRegion
}
sessionToken := account.GetCredential("aws_session_token") // 可选
return NewBedrockSigner(accessKeyID, secretAccessKey, sessionToken, region), nil
}
// SignRequest 对 HTTP 请求进行 SigV4 签名
// 重要约束：调用此方法前，req 应只包含 AWS 相关的 header（如 Content-Type、Accept）。
// 非 AWS header（如 anthropic-beta）会参与签名计算，如果 Bedrock 服务端不识别这些 header，
// 签名验证可能失败。litellm 通过 _filter_headers_for_aws_signature 实现头过滤，
// 当前实现中 buildUpstreamRequestBedrock 仅设置了 Content-Type 和 Accept，因此是安全的。
func (s *BedrockSigner) SignRequest(ctx context.Context, req *http.Request, body []byte) error {
payloadHash := sha256Hash(body)
return s.signer.SignHTTP(ctx, s.credentials, req, payloadHash, "bedrock", s.region, time.Now())
}
func sha256Hash(data []byte) string {
h := sha256.Sum256(data)
return hex.EncodeToString(h[:])
}
backend/internal/service/bedrock_signer_test.go 0 → 100644
View file @ 11f7b835
package service
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestNewBedrockSignerFromAccount_DefaultRegion(t *testing.T) {
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeBedrock,
Credentials: map[string]any{
"aws_access_key_id": "test-akid",
"aws_secret_access_key": "test-secret",
},
}
signer, err := NewBedrockSignerFromAccount(account)
require.NoError(t, err)
require.NotNil(t, signer)
assert.Equal(t, defaultBedrockRegion, signer.region)
}
func TestFilterBetaTokens(t *testing.T) {
tokens := []string{"interleaved-thinking-2025-05-14", "tool-search-tool-2025-10-19"}
filterSet := map[string]struct{}{
"tool-search-tool-2025-10-19": {},
}
assert.Equal(t, []string{"interleaved-thinking-2025-05-14"}, filterBetaTokens(tokens, filterSet))
assert.Equal(t, tokens, filterBetaTokens(tokens, nil))
assert.Nil(t, filterBetaTokens(nil, filterSet))
}
backend/internal/service/bedrock_stream.go 0 → 100644
View file @ 11f7b835
package service
import (
"bufio"
"context"
"encoding/base64"
"errors"
"fmt"
"hash/crc32"
"io"
"net/http"
"sync/atomic"
"time"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
)
// handleBedrockStreamingResponse 处理 Bedrock InvokeModelWithResponseStream 的 EventStream 响应
// Bedrock 返回 AWS EventStream 二进制格式，每个事件的 payload 中 chunk.bytes 是 base64 编码的
// Claude SSE 事件 JSON。本方法解码后转换为标准 SSE 格式写入客户端。
func (s *GatewayService) handleBedrockStreamingResponse(
ctx context.Context,
resp *http.Response,
c *gin.Context,
account *Account,
startTime time.Time,
model string,
) (*streamingResult, error) {
w := c.Writer
flusher, ok := w.(http.Flusher)
if !ok {
return nil, errors.New("streaming not supported")
}
c.Header("Content-Type", "text/event-stream")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")
c.Header("X-Accel-Buffering", "no")
if v := resp.Header.Get("x-amzn-requestid"); v != "" {
c.Header("x-request-id", v)
}
usage := &ClaudeUsage{}
var firstTokenMs *int
clientDisconnected := false
// Bedrock EventStream 使用 application/vnd.amazon.eventstream 二进制格式。
// 每个帧结构：total_length(4) + headers_length(4) + prelude_crc(4) + headers + payload + message_crc(4)
// 但更实用的方式是使用行扫描找 JSON chunks，因为 Bedrock 的响应在二进制帧中。
// 我们使用 EventStream decoder 来正确解析。
decoder := newBedrockEventStreamDecoder(resp.Body)
type decodeEvent struct {
payload []byte
err error
}
events := make(chan decodeEvent, 16)
done := make(chan struct{})
sendEvent := func(ev decodeEvent) bool {
select {
case events <- ev:
return true
case <-done:
return false
}
}
var lastReadAt atomic.Int64
lastReadAt.Store(time.Now().UnixNano())
go func() {
defer close(events)
for {
payload, err := decoder.Decode()
if err != nil {
if err == io.EOF {
return
}
_ = sendEvent(decodeEvent{err: err})
return
}
lastReadAt.Store(time.Now().UnixNano())
if !sendEvent(decodeEvent{payload: payload}) {
return
}
}
}()
defer close(done)
streamInterval := time.Duration(0)
if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 {
streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
}
var intervalTicker *time.Ticker
if streamInterval > 0 {
intervalTicker = time.NewTicker(streamInterval)
defer intervalTicker.Stop()
}
var intervalCh <-chan time.Time
if intervalTicker != nil {
intervalCh = intervalTicker.C
}
for {
select {
case ev, ok := <-events:
if !ok {
if !clientDisconnected {
flusher.Flush()
}
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil
}
if ev.err != nil {
if clientDisconnected {
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
}
if errors.Is(ev.err, context.Canceled) || errors.Is(ev.err, context.DeadlineExceeded) {
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
}
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("bedrock stream read error: %w", ev.err)
}
// payload 是 JSON，提取 chunk.bytes（base64 编码的 Claude SSE 事件数据）
sseData := extractBedrockChunkData(ev.payload)
if sseData == nil {
continue
}
if firstTokenMs == nil {
ms := int(time.Since(startTime).Milliseconds())
firstTokenMs = &ms
}
// 转换 Bedrock 特有的 amazon-bedrock-invocationMetrics 为标准 Anthropic usage 格式
// 同时移除该字段避免透传给客户端
sseData = transformBedrockInvocationMetrics(sseData)
// 解析 SSE 事件数据提取 usage
s.parseSSEUsagePassthrough(string(sseData), usage)
// 确定 SSE event type
eventType := gjson.GetBytes(sseData, "type").String()
// 写入标准 SSE 格式
if !clientDisconnected {
var writeErr error
if eventType != "" {
_, writeErr = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", eventType, sseData)
} else {
_, writeErr = fmt.Fprintf(w, "data: %s\n\n", sseData)
}
if writeErr != nil {
clientDisconnected = true
logger.LegacyPrintf("service.gateway", "[Bedrock] Client disconnected during streaming, continue draining for usage: account=%d", account.ID)
} else {
flusher.Flush()
}
}
case <-intervalCh:
lastRead := time.Unix(0, lastReadAt.Load())
if time.Since(lastRead) < streamInterval {
continue
}
if clientDisconnected {
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil
}
logger.LegacyPrintf("service.gateway", "[Bedrock] Stream data interval timeout: account=%d model=%s interval=%s", account.ID, model, streamInterval)
if s.rateLimitService != nil {
s.rateLimitService.HandleStreamTimeout(ctx, account, model)
}
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")
}
}
}
// extractBedrockChunkData 从 Bedrock EventStream payload 中提取 Claude SSE 事件数据
// Bedrock payload 格式：{"bytes":"<base64-encoded-json>"}
func extractBedrockChunkData(payload []byte) []byte {
b64 := gjson.GetBytes(payload, "bytes").String()
if b64 == "" {
return nil
}
decoded, err := base64.StdEncoding.DecodeString(b64)
if err != nil {
return nil
}
return decoded
}
// transformBedrockInvocationMetrics 将 Bedrock 特有的 amazon-bedrock-invocationMetrics
// 转换为标准 Anthropic usage 格式，并从 SSE 数据中移除该字段。
//
// Bedrock Invoke 返回的 message_delta 事件可能包含：
//
// {"type":"message_delta","delta":{...},"amazon-bedrock-invocationMetrics":{"inputTokenCount":150,"outputTokenCount":42}}
//
// 转换为：
//
// {"type":"message_delta","delta":{...},"usage":{"input_tokens":150,"output_tokens":42}}
func transformBedrockInvocationMetrics(data []byte) []byte {
metrics := gjson.GetBytes(data, "amazon-bedrock-invocationMetrics")
if !metrics.Exists() || !metrics.IsObject() {
return data
}
// 移除 Bedrock 特有字段
data, _ = sjson.DeleteBytes(data, "amazon-bedrock-invocationMetrics")
// 如果已有标准 usage 字段，不覆盖
if gjson.GetBytes(data, "usage").Exists() {
return data
}
// 转换 camelCase → snake_case 写入 usage
inputTokens := metrics.Get("inputTokenCount")
outputTokens := metrics.Get("outputTokenCount")
if inputTokens.Exists() {
data, _ = sjson.SetBytes(data, "usage.input_tokens", inputTokens.Int())
}
if outputTokens.Exists() {
data, _ = sjson.SetBytes(data, "usage.output_tokens", outputTokens.Int())
}
return data
}
// bedrockEventStreamDecoder 解码 AWS EventStream 二进制帧
// EventStream 帧格式：
//
// [total_byte_length: 4 bytes]
// [headers_byte_length: 4 bytes]
// [prelude_crc: 4 bytes]
// [headers: variable]
// [payload: variable]
// [message_crc: 4 bytes]
type bedrockEventStreamDecoder struct {
reader *bufio.Reader
}
func newBedrockEventStreamDecoder(r io.Reader) *bedrockEventStreamDecoder {
return &bedrockEventStreamDecoder{
reader: bufio.NewReaderSize(r, 64*1024),
}
}
// Decode 读取下一个 EventStream 帧并返回 chunk 类型事件的 payload
func (d *bedrockEventStreamDecoder) Decode() ([]byte, error) {
for {
// 读取 prelude: total_length(4) + headers_length(4) + prelude_crc(4) = 12 bytes
prelude := make([]byte, 12)
if _, err := io.ReadFull(d.reader, prelude); err != nil {
return nil, err
}
// 验证 prelude CRC（AWS EventStream 使用标准 CRC32 / IEEE）
preludeCRC := bedrockReadUint32(prelude[8:12])
if crc32.Checksum(prelude[0:8], crc32IEEETable) != preludeCRC {
return nil, fmt.Errorf("eventstream prelude CRC mismatch")
}
totalLength := bedrockReadUint32(prelude[0:4])
headersLength := bedrockReadUint32(prelude[4:8])
if totalLength < 16 { // minimum: 12 prelude + 4 message_crc
return nil, fmt.Errorf("invalid eventstream frame: total_length=%d", totalLength)
}
// 读取 headers + payload + message_crc
remaining := int(totalLength) - 12
if remaining <= 0 {
continue
}
data := make([]byte, remaining)
if _, err := io.ReadFull(d.reader, data); err != nil {
return nil, err
}
// 验证 message CRC（覆盖 prelude + headers + payload）
messageCRC := bedrockReadUint32(data[len(data)-4:])
h := crc32.New(crc32IEEETable)
h.Write(prelude)
h.Write(data[:len(data)-4])
if h.Sum32() != messageCRC {
return nil, fmt.Errorf("eventstream message CRC mismatch")
}
// 解析 headers
headers := data[:headersLength]
payload := data[headersLength : len(data)-4] // 去掉 message_crc
// 从 headers 中提取 :event-type
eventType := extractEventStreamHeaderValue(headers, ":event-type")
// 只处理 chunk 事件
if eventType == "chunk" {
// payload 是完整的 JSON，包含 bytes 字段
return payload, nil
}
// 检查异常事件
exceptionType := extractEventStreamHeaderValue(headers, ":exception-type")
if exceptionType != "" {
return nil, fmt.Errorf("bedrock exception: %s: %s", exceptionType, string(payload))
}
messageType := extractEventStreamHeaderValue(headers, ":message-type")
if messageType == "exception" || messageType == "error" {
return nil, fmt.Errorf("bedrock error: %s", string(payload))
}
// 跳过其他事件类型（如 initial-response）
}
}
// extractEventStreamHeaderValue 从 EventStream headers 二进制数据中提取指定 header 的字符串值
// EventStream header 格式：
//
// [name_length: 1 byte][name: variable][value_type: 1 byte][value: variable]
//
// value_type = 7 表示 string 类型，前 2 bytes 为长度
func extractEventStreamHeaderValue(headers []byte, targetName string) string {
pos := 0
for pos < len(headers) {
if pos >= len(headers) {
break
}
nameLen := int(headers[pos])
pos++
if pos+nameLen > len(headers) {
break
}
name := string(headers[pos : pos+nameLen])
pos += nameLen
if pos >= len(headers) {
break
}
valueType := headers[pos]
pos++
switch valueType {
case 7: // string
if pos+2 > len(headers) {
return ""
}
valueLen := int(bedrockReadUint16(headers[pos : pos+2]))
pos += 2
if pos+valueLen > len(headers) {
return ""
}
value := string(headers[pos : pos+valueLen])
pos += valueLen
if name == targetName {
return value
}
case 0: // bool true
if name == targetName {
return "true"
}
case 1: // bool false
if name == targetName {
return "false"
}
case 2: // byte
pos++
if name == targetName {
return ""
}
case 3: // short
pos += 2
if name == targetName {
return ""
}
case 4: // int
pos += 4
if name == targetName {
return ""
}
case 5: // long
pos += 8
if name == targetName {
return ""
}
case 6: // bytes
if pos+2 > len(headers) {
return ""
}
valueLen := int(bedrockReadUint16(headers[pos : pos+2]))
pos += 2 + valueLen
case 8: // timestamp
pos += 8
case 9: // uuid
pos += 16
default:
return "" // 未知类型，无法继续解析
}
}
return ""
}
// crc32IEEETable is the CRC32 / IEEE table used by AWS EventStream.
var crc32IEEETable = crc32.MakeTable(crc32.IEEE)
func bedrockReadUint32(b []byte) uint32 {
return uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
}
func bedrockReadUint16(b []byte) uint16 {
return uint16(b[0])<<8 | uint16(b[1])
}
backend/internal/service/bedrock_stream_test.go 0 → 100644
View file @ 11f7b835
package service
import (
"bytes"
"encoding/base64"
"encoding/binary"
"hash/crc32"
"io"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestExtractBedrockChunkData(t *testing.T) {
t.Run("valid base64 payload", func(t *testing.T) {
original := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}`
b64 := base64.StdEncoding.EncodeToString([]byte(original))
payload := []byte(`{"bytes":"` + b64 + `"}`)
result := extractBedrockChunkData(payload)
require.NotNil(t, result)
assert.JSONEq(t, original, string(result))
})
t.Run("empty bytes field", func(t *testing.T) {
result := extractBedrockChunkData([]byte(`{"bytes":""}`))
assert.Nil(t, result)
})
t.Run("no bytes field", func(t *testing.T) {
result := extractBedrockChunkData([]byte(`{"other":"value"}`))
assert.Nil(t, result)
})
t.Run("invalid base64", func(t *testing.T) {
result := extractBedrockChunkData([]byte(`{"bytes":"not-valid-base64!!!"}`))
assert.Nil(t, result)
})
}
func TestTransformBedrockInvocationMetrics(t *testing.T) {
t.Run("converts metrics to usage", func(t *testing.T) {
input := `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"amazon-bedrock-invocationMetrics":{"inputTokenCount":150,"outputTokenCount":42}}`
result := transformBedrockInvocationMetrics([]byte(input))
// amazon-bedrock-invocationMetrics should be removed
assert.False(t, gjson.GetBytes(result, "amazon-bedrock-invocationMetrics").Exists())
// usage should be set
assert.Equal(t, int64(150), gjson.GetBytes(result, "usage.input_tokens").Int())
assert.Equal(t, int64(42), gjson.GetBytes(result, "usage.output_tokens").Int())
// original fields preserved
assert.Equal(t, "message_delta", gjson.GetBytes(result, "type").String())
assert.Equal(t, "end_turn", gjson.GetBytes(result, "delta.stop_reason").String())
})
t.Run("no metrics present", func(t *testing.T) {
input := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}`
result := transformBedrockInvocationMetrics([]byte(input))
assert.JSONEq(t, input, string(result))
})
t.Run("does not overwrite existing usage", func(t *testing.T) {
input := `{"type":"message_delta","usage":{"output_tokens":100},"amazon-bedrock-invocationMetrics":{"inputTokenCount":150,"outputTokenCount":42}}`
result := transformBedrockInvocationMetrics([]byte(input))
// metrics removed but existing usage preserved
assert.False(t, gjson.GetBytes(result, "amazon-bedrock-invocationMetrics").Exists())
assert.Equal(t, int64(100), gjson.GetBytes(result, "usage.output_tokens").Int())
})
}
func TestExtractEventStreamHeaderValue(t *testing.T) {
// Build a header with :event-type = "chunk" (string type = 7)
buildStringHeader := func(name, value string) []byte {
var buf bytes.Buffer
// name length (1 byte)
buf.WriteByte(byte(len(name)))
// name
buf.WriteString(name)
// value type (7 = string)
buf.WriteByte(7)
// value length (2 bytes, big-endian)
_ = binary.Write(&buf, binary.BigEndian, uint16(len(value)))
// value
buf.WriteString(value)
return buf.Bytes()
}
t.Run("find string header", func(t *testing.T) {
headers := buildStringHeader(":event-type", "chunk")
assert.Equal(t, "chunk", extractEventStreamHeaderValue(headers, ":event-type"))
})
t.Run("header not found", func(t *testing.T) {
headers := buildStringHeader(":event-type", "chunk")
assert.Equal(t, "", extractEventStreamHeaderValue(headers, ":message-type"))
})
t.Run("multiple headers", func(t *testing.T) {
var buf bytes.Buffer
buf.Write(buildStringHeader(":content-type", "application/json"))
buf.Write(buildStringHeader(":event-type", "chunk"))
buf.Write(buildStringHeader(":message-type", "event"))
headers := buf.Bytes()
assert.Equal(t, "chunk", extractEventStreamHeaderValue(headers, ":event-type"))
assert.Equal(t, "application/json", extractEventStreamHeaderValue(headers, ":content-type"))
assert.Equal(t, "event", extractEventStreamHeaderValue(headers, ":message-type"))
})
t.Run("empty headers", func(t *testing.T) {
assert.Equal(t, "", extractEventStreamHeaderValue([]byte{}, ":event-type"))
})
}
func TestBedrockEventStreamDecoder(t *testing.T) {
crc32IeeeTab := crc32.MakeTable(crc32.IEEE)
// Build a valid EventStream frame with correct CRC32/IEEE checksums.
buildFrame := func(eventType string, payload []byte) []byte {
// Build headers
var headersBuf bytes.Buffer
// :event-type header
headersBuf.WriteByte(byte(len(":event-type")))
headersBuf.WriteString(":event-type")
headersBuf.WriteByte(7) // string type
_ = binary.Write(&headersBuf, binary.BigEndian, uint16(len(eventType)))
headersBuf.WriteString(eventType)
// :message-type header
headersBuf.WriteByte(byte(len(":message-type")))
headersBuf.WriteString(":message-type")
headersBuf.WriteByte(7)
_ = binary.Write(&headersBuf, binary.BigEndian, uint16(len("event")))
headersBuf.WriteString("event")
headers := headersBuf.Bytes()
headersLen := uint32(len(headers))
// total = 12 (prelude) + headers + payload + 4 (message_crc)
totalLen := uint32(12 + len(headers) + len(payload) + 4)
// Prelude: total_length(4) + headers_length(4)
var preludeBuf bytes.Buffer
_ = binary.Write(&preludeBuf, binary.BigEndian, totalLen)
_ = binary.Write(&preludeBuf, binary.BigEndian, headersLen)
preludeBytes := preludeBuf.Bytes()
preludeCRC := crc32.Checksum(preludeBytes, crc32IeeeTab)
// Build frame: prelude + prelude_crc + headers + payload
var frame bytes.Buffer
frame.Write(preludeBytes)
_ = binary.Write(&frame, binary.BigEndian, preludeCRC)
frame.Write(headers)
frame.Write(payload)
// Message CRC covers everything before itself
messageCRC := crc32.Checksum(frame.Bytes(), crc32IeeeTab)
_ = binary.Write(&frame, binary.BigEndian, messageCRC)
return frame.Bytes()
}
t.Run("decode chunk event", func(t *testing.T) {
payload := []byte(`{"bytes":"dGVzdA=="}`) // base64("test")
frame := buildFrame("chunk", payload)
decoder := newBedrockEventStreamDecoder(bytes.NewReader(frame))
result, err := decoder.Decode()
require.NoError(t, err)
assert.Equal(t, payload, result)
})
t.Run("skip non-chunk events", func(t *testing.T) {
// Write initial-response followed by chunk
var buf bytes.Buffer
buf.Write(buildFrame("initial-response", []byte(`{}`)))
chunkPayload := []byte(`{"bytes":"aGVsbG8="}`)
buf.Write(buildFrame("chunk", chunkPayload))
decoder := newBedrockEventStreamDecoder(&buf)
result, err := decoder.Decode()
require.NoError(t, err)
assert.Equal(t, chunkPayload, result)
})
t.Run("EOF on empty input", func(t *testing.T) {
decoder := newBedrockEventStreamDecoder(bytes.NewReader(nil))
_, err := decoder.Decode()
assert.Equal(t, io.EOF, err)
})
t.Run("corrupted prelude CRC", func(t *testing.T) {
frame := buildFrame("chunk", []byte(`{"bytes":"dGVzdA=="}`))
// Corrupt the prelude CRC (bytes 8-11)
frame[8] ^= 0xFF
decoder := newBedrockEventStreamDecoder(bytes.NewReader(frame))
_, err := decoder.Decode()
require.Error(t, err)
assert.Contains(t, err.Error(), "prelude CRC mismatch")
})
t.Run("corrupted message CRC", func(t *testing.T) {
frame := buildFrame("chunk", []byte(`{"bytes":"dGVzdA=="}`))
// Corrupt the message CRC (last 4 bytes)
frame[len(frame)-1] ^= 0xFF
decoder := newBedrockEventStreamDecoder(bytes.NewReader(frame))
_, err := decoder.Decode()
require.Error(t, err)
assert.Contains(t, err.Error(), "message CRC mismatch")
})
t.Run("castagnoli encoded frame is rejected", func(t *testing.T) {
castagnoliTab := crc32.MakeTable(crc32.Castagnoli)
payload := []byte(`{"bytes":"dGVzdA=="}`)
var headersBuf bytes.Buffer
headersBuf.WriteByte(byte(len(":event-type")))
headersBuf.WriteString(":event-type")
headersBuf.WriteByte(7)
_ = binary.Write(&headersBuf, binary.BigEndian, uint16(len("chunk")))
headersBuf.WriteString("chunk")
headers := headersBuf.Bytes()
headersLen := uint32(len(headers))
totalLen := uint32(12 + len(headers) + len(payload) + 4)
var preludeBuf bytes.Buffer
_ = binary.Write(&preludeBuf, binary.BigEndian, totalLen)
_ = binary.Write(&preludeBuf, binary.BigEndian, headersLen)
preludeBytes := preludeBuf.Bytes()
var frame bytes.Buffer
frame.Write(preludeBytes)
_ = binary.Write(&frame, binary.BigEndian, crc32.Checksum(preludeBytes, castagnoliTab))
frame.Write(headers)
frame.Write(payload)
_ = binary.Write(&frame, binary.BigEndian, crc32.Checksum(frame.Bytes(), castagnoliTab))
decoder := newBedrockEventStreamDecoder(bytes.NewReader(frame.Bytes()))
_, err := decoder.Decode()
require.Error(t, err)
assert.Contains(t, err.Error(), "prelude CRC mismatch")
})
}
func TestBuildBedrockURL(t *testing.T) {
t.Run("stream URL with colon in model ID", func(t *testing.T) {
url := BuildBedrockURL("us-east-1", "us.anthropic.claude-opus-4-5-20251101-v1:0", true)
assert.Equal(t, "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-opus-4-5-20251101-v1%3A0/invoke-with-response-stream", url)
})
t.Run("non-stream URL with colon in model ID", func(t *testing.T) {
url := BuildBedrockURL("eu-west-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", false)
assert.Equal(t, "https://bedrock-runtime.eu-west-1.amazonaws.com/model/eu.anthropic.claude-sonnet-4-5-20250929-v1%3A0/invoke", url)
})
t.Run("model ID without colon", func(t *testing.T) {
url := BuildBedrockURL("us-east-1", "us.anthropic.claude-sonnet-4-6", true)
assert.Equal(t, "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-sonnet-4-6/invoke-with-response-stream", url)
})
}
backend/internal/service/domain_constants.go
View file @ 11f7b835
......@@ -33,6 +33,8 @@ const (
AccountTypeSetupToken = domain.AccountTypeSetupToken // Setup Token类型账号（inference only scope）
AccountTypeAPIKey = domain.AccountTypeAPIKey // API Key类型账号
AccountTypeUpstream = domain.AccountTypeUpstream // 上游透传类型账号（通过 Base URL + API Key 连接上游）
AccountTypeBedrock = domain.AccountTypeBedrock // AWS Bedrock 类型账号（通过 SigV4 签名连接 Bedrock）
AccountTypeBedrockAPIKey = domain.AccountTypeBedrockAPIKey // AWS Bedrock API Key 类型账号（通过 Bearer Token 连接 Bedrock）
)
// Redeem type constants
......
backend/internal/service/gateway_service.go
View file @ 11f7b835
......@@ -3336,6 +3336,10 @@ func (s *GatewayService) isModelSupportedByAccount(account *Account, requestedMo
if account.Platform == PlatformSora {
return s.isSoraModelSupportedByAccount(account, requestedModel)
}
if account.IsBedrock() {
_, ok := ResolveBedrockModelID(account, requestedModel)
return ok
}
// OAuth/SetupToken 账号使用 Anthropic 标准映射（短ID → 长ID）
if account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey {
requestedModel = claude.NormalizeModelID(requestedModel)
......@@ -3493,6 +3497,10 @@ func (s *GatewayService) GetAccessToken(ctx context.Context, account *Account) (
return "", "", errors.New("api_key not found in credentials")
}
return apiKey, "apikey", nil
case AccountTypeBedrock:
return "", "bedrock", nil // Bedrock 使用 SigV4 签名，不需要 token
case AccountTypeBedrockAPIKey:
return "", "bedrock-apikey", nil // Bedrock API Key 使用 Bearer Token，由 forwardBedrock 处理
default:
return "", "", fmt.Errorf("unsupported account type: %s", account.Type)
}
......@@ -3948,6 +3956,10 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
return s.forwardAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody, passthroughModel, parsed.Stream, startTime)
}
if account != nil && account.IsBedrock() {
return s.forwardBedrock(ctx, c, account, parsed, startTime)
}
// Beta policy: evaluate once; block check + cache filter set for buildUpstreamRequest.
// Always overwrite the cache to prevent stale values from a previous retry with a different account.
if account.Platform == PlatformAnthropic && c != nil {
......@@ -5068,6 +5080,366 @@ func writeAnthropicPassthroughResponseHeaders(dst http.Header, src http.Header,
}
}
// forwardBedrock 转发请求到 AWS Bedrock
func (s *GatewayService) forwardBedrock(
ctx context.Context,
c *gin.Context,
account *Account,
parsed *ParsedRequest,
startTime time.Time,
) (*ForwardResult, error) {
reqModel := parsed.Model
reqStream := parsed.Stream
body := parsed.Body
region := bedrockRuntimeRegion(account)
mappedModel, ok := ResolveBedrockModelID(account, reqModel)
if !ok {
return nil, fmt.Errorf("unsupported bedrock model: %s", reqModel)
}
if mappedModel != reqModel {
logger.LegacyPrintf("service.gateway", "[Bedrock] Model mapping: %s -> %s (account: %s)", reqModel, mappedModel, account.Name)
}
betaHeader := ""
if c != nil && c.Request != nil {
betaHeader = c.GetHeader("anthropic-beta")
}
// 准备请求体（注入 anthropic_version/anthropic_beta，移除 Bedrock 不支持的字段，清理 cache_control）
betaTokens, err := s.resolveBedrockBetaTokensForRequest(ctx, account, betaHeader, body, mappedModel)
if err != nil {
return nil, err
}
bedrockBody, err := PrepareBedrockRequestBodyWithTokens(body, mappedModel, betaTokens)
if err != nil {
return nil, fmt.Errorf("prepare bedrock request body: %w", err)
}
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
logger.LegacyPrintf("service.gateway", "[Bedrock] 命中 Bedrock 分支: account=%d name=%s model=%s->%s stream=%v",
account.ID, account.Name, reqModel, mappedModel, reqStream)
// 根据账号类型选择认证方式
var signer *BedrockSigner
var bedrockAPIKey string
if account.IsBedrockAPIKey() {
bedrockAPIKey = account.GetCredential("api_key")
if bedrockAPIKey == "" {
return nil, fmt.Errorf("api_key not found in bedrock-apikey credentials")
}
} else {
signer, err = NewBedrockSignerFromAccount(account)
if err != nil {
return nil, fmt.Errorf("create bedrock signer: %w", err)
}
}
// 执行上游请求（含重试）
resp, err := s.executeBedrockUpstream(ctx, c, account, bedrockBody, mappedModel, region, reqStream, signer, bedrockAPIKey, proxyURL)
if err != nil {
return nil, err
}
defer func() { _ = resp.Body.Close() }()
// 将 Bedrock 的 x-amzn-requestid 映射到 x-request-id，
// 使通用错误处理函数（handleErrorResponse、handleRetryExhaustedError）能正确提取 AWS request ID。
if awsReqID := resp.Header.Get("x-amzn-requestid"); awsReqID != "" && resp.Header.Get("x-request-id") == "" {
resp.Header.Set("x-request-id", awsReqID)
}
// 错误/failover 处理
if resp.StatusCode >= 400 {
return s.handleBedrockUpstreamErrors(ctx, resp, c, account)
}
// 响应处理
var usage *ClaudeUsage
var firstTokenMs *int
var clientDisconnect bool
if reqStream {
streamResult, err := s.handleBedrockStreamingResponse(ctx, resp, c, account, startTime, reqModel)
if err != nil {
return nil, err
}
usage = streamResult.usage
firstTokenMs = streamResult.firstTokenMs
clientDisconnect = streamResult.clientDisconnect
} else {
usage, err = s.handleBedrockNonStreamingResponse(ctx, resp, c, account)
if err != nil {
return nil, err
}
}
if usage == nil {
usage = &ClaudeUsage{}
}
return &ForwardResult{
RequestID: resp.Header.Get("x-amzn-requestid"),
Usage: *usage,
Model: reqModel,
Stream: reqStream,
Duration: time.Since(startTime),
FirstTokenMs: firstTokenMs,
ClientDisconnect: clientDisconnect,
}, nil
}
// executeBedrockUpstream 执行 Bedrock 上游请求（含重试逻辑）
func (s *GatewayService) executeBedrockUpstream(
ctx context.Context,
c *gin.Context,
account *Account,
body []byte,
modelID string,
region string,
stream bool,
signer *BedrockSigner,
apiKey string,
proxyURL string,
) (*http.Response, error) {
var resp *http.Response
var err error
retryStart := time.Now()
for attempt := 1; attempt <= maxRetryAttempts; attempt++ {
var upstreamReq *http.Request
if account.IsBedrockAPIKey() {
upstreamReq, err = s.buildUpstreamRequestBedrockAPIKey(ctx, body, modelID, region, stream, apiKey)
} else {
upstreamReq, err = s.buildUpstreamRequestBedrock(ctx, body, modelID, region, stream, signer)
}
if err != nil {
return nil, err
}
resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, false)
if err != nil {
if resp != nil && resp.Body != nil {
_ = resp.Body.Close()
}
safeErr := sanitizeUpstreamErrorMessage(err.Error())
setOpsUpstreamError(c, 0, safeErr, "")
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: 0,
Kind: "request_error",
Message: safeErr,
})
c.JSON(http.StatusBadGateway, gin.H{
"type": "error",
"error": gin.H{
"type": "upstream_error",
"message": "Upstream request failed",
},
})
return nil, fmt.Errorf("upstream request failed: %s", safeErr)
}
if resp.StatusCode >= 400 && resp.StatusCode != 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) {
if attempt < maxRetryAttempts {
elapsed := time.Since(retryStart)
if elapsed >= maxRetryElapsed {
break
}
delay := retryBackoffDelay(attempt)
remaining := maxRetryElapsed - elapsed
if delay > remaining {
delay = remaining
}
if delay <= 0 {
break
}
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = resp.Body.Close()
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
Kind: "retry",
Message: extractUpstreamErrorMessage(respBody),
Detail: func() string {
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes)
}
return ""
}(),
})
logger.LegacyPrintf("service.gateway", "[Bedrock] account %d: upstream error %d, retry %d/%d after %v",
account.ID, resp.StatusCode, attempt, maxRetryAttempts, delay)
if err := sleepWithContext(ctx, delay); err != nil {
return nil, err
}
continue
}
break
}
break
}
if resp == nil || resp.Body == nil {
return nil, errors.New("upstream request failed: empty response")
}
return resp, nil
}
// handleBedrockUpstreamErrors 处理 Bedrock 上游 4xx/5xx 错误（failover + 错误响应）
func (s *GatewayService) handleBedrockUpstreamErrors(
ctx context.Context,
resp *http.Response,
c *gin.Context,
account *Account,
) (*ForwardResult, error) {
// retry exhausted + failover
if s.shouldRetryUpstreamError(account, resp.StatusCode) {
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
logger.LegacyPrintf("service.gateway", "[Bedrock] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d Body=%s",
account.ID, account.Name, resp.StatusCode, truncateString(string(respBody), 1000))
s.handleRetryExhaustedSideEffects(ctx, resp, account)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
Kind: "retry_exhausted_failover",
Message: extractUpstreamErrorMessage(respBody),
})
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
}
}
return s.handleRetryExhaustedError(ctx, resp, c, account)
}
// non-retryable failover
if s.shouldFailoverUpstreamError(resp.StatusCode) {
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = resp.Body.Close()
resp.Body = io.NopCloser(bytes.NewReader(respBody))
s.handleFailoverSideEffects(ctx, resp, account)
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
Platform: account.Platform,
AccountID: account.ID,
AccountName: account.Name,
UpstreamStatusCode: resp.StatusCode,
Kind: "failover",
Message: extractUpstreamErrorMessage(respBody),
})
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
}
}
// other errors
return s.handleErrorResponse(ctx, resp, c, account)
}
// buildUpstreamRequestBedrock 构建 Bedrock 上游请求
func (s *GatewayService) buildUpstreamRequestBedrock(
ctx context.Context,
body []byte,
modelID string,
region string,
stream bool,
signer *BedrockSigner,
) (*http.Request, error) {
targetURL := BuildBedrockURL(region, modelID, stream)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
// SigV4 签名
if err := signer.SignRequest(ctx, req, body); err != nil {
return nil, fmt.Errorf("sign bedrock request: %w", err)
}
return req, nil
}
// buildUpstreamRequestBedrockAPIKey 构建 Bedrock API Key (Bearer Token) 上游请求
func (s *GatewayService) buildUpstreamRequestBedrockAPIKey(
ctx context.Context,
body []byte,
modelID string,
region string,
stream bool,
apiKey string,
) (*http.Request, error) {
targetURL := BuildBedrockURL(region, modelID, stream)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
return req, nil
}
// handleBedrockNonStreamingResponse 处理 Bedrock 非流式响应
// Bedrock InvokeModel 非流式响应的 body 格式与 Claude API 兼容
func (s *GatewayService) handleBedrockNonStreamingResponse(
ctx context.Context,
resp *http.Response,
c *gin.Context,
account *Account,
) (*ClaudeUsage, error) {
maxBytes := resolveUpstreamResponseReadLimit(s.cfg)
body, err := readUpstreamResponseBodyLimited(resp.Body, maxBytes)
if err != nil {
if errors.Is(err, ErrUpstreamResponseBodyTooLarge) {
setOpsUpstreamError(c, http.StatusBadGateway, "upstream response too large", "")
c.JSON(http.StatusBadGateway, gin.H{
"type": "error",
"error": gin.H{
"type": "upstream_error",
"message": "Upstream response too large",
},
})
}
return nil, err
}
// 转换 Bedrock 特有的 amazon-bedrock-invocationMetrics 为标准 Anthropic usage 格式
// 并移除该字段避免透传给客户端
body = transformBedrockInvocationMetrics(body)
usage := parseClaudeUsageFromResponseBody(body)
c.Header("Content-Type", "application/json")
if v := resp.Header.Get("x-amzn-requestid"); v != "" {
c.Header("x-request-id", v)
}
c.Data(resp.StatusCode, "application/json", body)
return usage, nil
}
func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, error) {
// 确定目标URL
targetURL := claudeAPIURL
......@@ -5481,6 +5853,95 @@ func containsBetaToken(header, token string) bool {
return false
}
// filterBetaTokensFromHeader removes tokens present in filterSet from a comma-separated header value.
// Returns the filtered header string, or "" if all tokens were removed.
func filterBetaTokensFromHeader(header string, filterSet map[string]struct{}) string {
if header == "" || len(filterSet) == 0 {
return header
}
var kept []string
for _, p := range strings.Split(header, ",") {
t := strings.TrimSpace(p)
if t == "" {
continue
}
if _, filtered := filterSet[t]; !filtered {
kept = append(kept, t)
}
}
return strings.Join(kept, ", ")
}
func filterBetaTokens(tokens []string, filterSet map[string]struct{}) []string {
if len(tokens) == 0 || len(filterSet) == 0 {
return tokens
}
kept := make([]string, 0, len(tokens))
for _, token := range tokens {
if _, filtered := filterSet[token]; !filtered {
kept = append(kept, token)
}
}
return kept
}
func (s *GatewayService) resolveBedrockBetaTokensForRequest(
ctx context.Context,
account *Account,
betaHeader string,
body []byte,
modelID string,
) ([]string, error) {
// 1. 对原始 header 中的 beta token 做 block 检查（快速失败）
policy := s.evaluateBetaPolicy(ctx, betaHeader, account)
if policy.blockErr != nil {
return nil, policy.blockErr
}
// 2. 解析 header + body 自动注入 + Bedrock 转换/过滤
betaTokens := ResolveBedrockBetaTokens(betaHeader, body, modelID)
// 3. 对最终 token 列表再做 block 检查，捕获通过 body 自动注入绕过 header block 的情况。
// 例如：管理员 block 了 interleaved-thinking，客户端不在 header 中带该 token，
// 但请求体中包含 thinking 字段 → autoInjectBedrockBetaTokens 会自动补齐 →
// 如果不做此检查，block 规则会被绕过。
if blockErr := s.checkBetaPolicyBlockForTokens(ctx, betaTokens, account); blockErr != nil {
return nil, blockErr
}
return filterBetaTokens(betaTokens, policy.filterSet), nil
}
// checkBetaPolicyBlockForTokens 检查 token 列表中是否有被管理员 block 规则命中的 token。
// 用于补充 evaluateBetaPolicy 对 header 的检查，覆盖 body 自动注入的 token。
func (s *GatewayService) checkBetaPolicyBlockForTokens(ctx context.Context, tokens []string, account *Account) *BetaBlockedError {
if s.settingService == nil || len(tokens) == 0 {
return nil
}
settings, err := s.settingService.GetBetaPolicySettings(ctx)
if err != nil || settings == nil {
return nil
}
isOAuth := account.IsOAuth()
tokenSet := buildBetaTokenSet(tokens)
for _, rule := range settings.Rules {
if rule.Action != BetaPolicyActionBlock {
continue
}
if !betaPolicyScopeMatches(rule.Scope, isOAuth) {
continue
}
if _, present := tokenSet[rule.BetaToken]; present {
msg := rule.ErrorMessage
if msg == "" {
msg = "beta feature " + rule.BetaToken + " is not allowed"
}
return &BetaBlockedError{Message: msg}
}
}
return nil
}
func buildBetaTokenSet(tokens []string) map[string]struct{} {
m := make(map[string]struct{}, len(tokens))
for _, t := range tokens {
......@@ -7064,6 +7525,12 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
return s.forwardCountTokensAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody)
}
// Bedrock 不支持 count_tokens 端点
if account != nil && account.IsBedrock() {
s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported for Bedrock")
return nil
}
body := parsed.Body
reqModel := parsed.Model
......
backend/internal/service/gateway_service_bedrock_beta_test.go 0 → 100644
View file @ 11f7b835
package service
import (
"context"
"encoding/json"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
)
type betaPolicySettingRepoStub struct {
values map[string]string
}
func (s *betaPolicySettingRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
panic("unexpected Get call")
}
func (s *betaPolicySettingRepoStub) GetValue(ctx context.Context, key string) (string, error) {
if v, ok := s.values[key]; ok {
return v, nil
}
return "", ErrSettingNotFound
}
func (s *betaPolicySettingRepoStub) Set(ctx context.Context, key, value string) error {
panic("unexpected Set call")
}
func (s *betaPolicySettingRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
panic("unexpected GetMultiple call")
}
func (s *betaPolicySettingRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
panic("unexpected SetMultiple call")
}
func (s *betaPolicySettingRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
panic("unexpected GetAll call")
}
func (s *betaPolicySettingRepoStub) Delete(ctx context.Context, key string) error {
panic("unexpected Delete call")
}
func TestResolveBedrockBetaTokensForRequest_BlocksOnOriginalAnthropicToken(t *testing.T) {
settings := &BetaPolicySettings{
Rules: []BetaPolicyRule{
{
BetaToken: "advanced-tool-use-2025-11-20",
Action: BetaPolicyActionBlock,
Scope: BetaPolicyScopeAll,
ErrorMessage: "advanced tool use is blocked",
},
},
}
raw, err := json.Marshal(settings)
if err != nil {
t.Fatalf("marshal settings: %v", err)
}
svc := &GatewayService{
settingService: NewSettingService(
&betaPolicySettingRepoStub{values: map[string]string{
SettingKeyBetaPolicySettings: string(raw),
}},
&config.Config{},
),
}
account := &Account{Platform: PlatformAnthropic, Type: AccountTypeBedrock}
_, err = svc.resolveBedrockBetaTokensForRequest(
context.Background(),
account,
"advanced-tool-use-2025-11-20",
[]byte(`{"messages":[{"role":"user","content":"hi"}]}`),
"us.anthropic.claude-opus-4-6-v1",
)
if err == nil {
t.Fatal("expected raw advanced-tool-use token to be blocked before Bedrock transform")
}
if err.Error() != "advanced tool use is blocked" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestResolveBedrockBetaTokensForRequest_FiltersAfterBedrockTransform(t *testing.T) {
settings := &BetaPolicySettings{
Rules: []BetaPolicyRule{
{
BetaToken: "tool-search-tool-2025-10-19",
Action: BetaPolicyActionFilter,
Scope: BetaPolicyScopeAll,
},
},
}
raw, err := json.Marshal(settings)
if err != nil {
t.Fatalf("marshal settings: %v", err)
}
svc := &GatewayService{
settingService: NewSettingService(
&betaPolicySettingRepoStub{values: map[string]string{
SettingKeyBetaPolicySettings: string(raw),
}},
&config.Config{},
),
}
account := &Account{Platform: PlatformAnthropic, Type: AccountTypeBedrock}
betaTokens, err := svc.resolveBedrockBetaTokensForRequest(
context.Background(),
account,
"advanced-tool-use-2025-11-20",
[]byte(`{"messages":[{"role":"user","content":"hi"}]}`),
"us.anthropic.claude-opus-4-6-v1",
)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
for _, token := range betaTokens {
if token == "tool-search-tool-2025-10-19" {
t.Fatalf("expected transformed Bedrock token to be filtered")
}
}
}
// TestResolveBedrockBetaTokensForRequest_BlocksBodyAutoInjectedThinking 验证：
// 管理员 block 了 interleaved-thinking，客户端不在 header 中带该 token，
// 但请求体包含 thinking 字段 → 自动注入后应被 block。
func TestResolveBedrockBetaTokensForRequest_BlocksBodyAutoInjectedThinking(t *testing.T) {
settings := &BetaPolicySettings{
Rules: []BetaPolicyRule{
{
BetaToken: "interleaved-thinking-2025-05-14",
Action: BetaPolicyActionBlock,
Scope: BetaPolicyScopeAll,
ErrorMessage: "thinking is blocked",
},
},
}
raw, err := json.Marshal(settings)
if err != nil {
t.Fatalf("marshal settings: %v", err)
}
svc := &GatewayService{
settingService: NewSettingService(
&betaPolicySettingRepoStub{values: map[string]string{
SettingKeyBetaPolicySettings: string(raw),
}},
&config.Config{},
),
}
account := &Account{Platform: PlatformAnthropic, Type: AccountTypeBedrock}
// header 中不带 beta token，但 body 中有 thinking 字段
_, err = svc.resolveBedrockBetaTokensForRequest(
context.Background(),
account,
"", // 空 header
[]byte(`{"thinking":{"type":"enabled","budget_tokens":10000},"messages":[{"role":"user","content":"hi"}]}`),
"us.anthropic.claude-opus-4-6-v1",
)
if err == nil {
t.Fatal("expected body-injected interleaved-thinking to be blocked")
}
if err.Error() != "thinking is blocked" {
t.Fatalf("unexpected error: %v", err)
}
}
// TestResolveBedrockBetaTokensForRequest_BlocksBodyAutoInjectedToolSearch 验证：
// 管理员 block 了 tool-search-tool，客户端不在 header 中带 beta token，
// 但请求体包含 tool search 工具 → 自动注入后应被 block。
func TestResolveBedrockBetaTokensForRequest_BlocksBodyAutoInjectedToolSearch(t *testing.T) {
settings := &BetaPolicySettings{
Rules: []BetaPolicyRule{
{
BetaToken: "tool-search-tool-2025-10-19",
Action: BetaPolicyActionBlock,
Scope: BetaPolicyScopeAll,
ErrorMessage: "tool search is blocked",
},
},
}
raw, err := json.Marshal(settings)
if err != nil {
t.Fatalf("marshal settings: %v", err)
}
svc := &GatewayService{
settingService: NewSettingService(
&betaPolicySettingRepoStub{values: map[string]string{
SettingKeyBetaPolicySettings: string(raw),
}},
&config.Config{},
),
}
account := &Account{Platform: PlatformAnthropic, Type: AccountTypeBedrock}
// header 中不带 beta token，但 body 中有 tool_search_tool 工具
_, err = svc.resolveBedrockBetaTokensForRequest(
context.Background(),
account,
"",
[]byte(`{"tools":[{"type":"tool_search_tool_regex_20251119","name":"search"}],"messages":[{"role":"user","content":"hi"}]}`),
"us.anthropic.claude-sonnet-4-6",
)
if err == nil {
t.Fatal("expected body-injected tool-search-tool to be blocked")
}
if err.Error() != "tool search is blocked" {
t.Fatalf("unexpected error: %v", err)
}
}
// TestResolveBedrockBetaTokensForRequest_PassesWhenNoBlockRuleMatches 验证：
// body 自动注入的 token 如果没有对应的 block 规则，应正常通过。
func TestResolveBedrockBetaTokensForRequest_PassesWhenNoBlockRuleMatches(t *testing.T) {
settings := &BetaPolicySettings{
Rules: []BetaPolicyRule{
{
BetaToken: "computer-use-2025-11-24",
Action: BetaPolicyActionBlock,
Scope: BetaPolicyScopeAll,
ErrorMessage: "computer use is blocked",
},
},
}
raw, err := json.Marshal(settings)
if err != nil {
t.Fatalf("marshal settings: %v", err)
}
svc := &GatewayService{
settingService: NewSettingService(
&betaPolicySettingRepoStub{values: map[string]string{
SettingKeyBetaPolicySettings: string(raw),
}},
&config.Config{},
),
}
account := &Account{Platform: PlatformAnthropic, Type: AccountTypeBedrock}
// body 中有 thinking（会注入 interleaved-thinking），但 block 规则只针对 computer-use
tokens, err := svc.resolveBedrockBetaTokensForRequest(
context.Background(),
account,
"",
[]byte(`{"thinking":{"type":"enabled","budget_tokens":10000},"messages":[{"role":"user","content":"hi"}]}`),
"us.anthropic.claude-opus-4-6-v1",
)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
found := false
for _, token := range tokens {
if token == "interleaved-thinking-2025-05-14" {
found = true
}
}
if !found {
t.Fatal("expected interleaved-thinking token to be present")
}
}
backend/internal/service/gateway_service_bedrock_model_support_test.go 0 → 100644
View file @ 11f7b835
package service
import "testing"
func TestGatewayServiceIsModelSupportedByAccount_BedrockDefaultMappingRestrictsModels(t *testing.T) {
svc := &GatewayService{}
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeBedrock,
Credentials: map[string]any{
"aws_region": "us-east-1",
},
}
if !svc.isModelSupportedByAccount(account, "claude-sonnet-4-5") {
t.Fatalf("expected default Bedrock alias to be supported")
}
if svc.isModelSupportedByAccount(account, "claude-3-5-sonnet-20241022") {
t.Fatalf("expected unsupported alias to be rejected for Bedrock account")
}
}
func TestGatewayServiceIsModelSupportedByAccount_BedrockCustomMappingStillActsAsAllowlist(t *testing.T) {
svc := &GatewayService{}
account := &Account{
Platform: PlatformAnthropic,
Type: AccountTypeBedrock,
Credentials: map[string]any{
"aws_region": "eu-west-1",
"model_mapping": map[string]any{
"claude-sonnet-*": "claude-sonnet-4-6",
},
},
}
if !svc.isModelSupportedByAccount(account, "claude-sonnet-4-6") {
t.Fatalf("expected matched custom mapping to be supported")
}
if !svc.isModelSupportedByAccount(account, "claude-opus-4-6") {
t.Fatalf("expected default Bedrock alias fallback to remain supported")
}
if svc.isModelSupportedByAccount(account, "claude-3-5-sonnet-20241022") {
t.Fatalf("expected unsupported model to still be rejected")
}
}
frontend/src/components/account/EditAccountModal.vue
View file @ 11f7b835
......@@ -563,6 +563,233 @@
</div>
</div>
<!-- Bedrock fields (only for bedrock type) -->
<div v-if="account.type === 'bedrock'" class="space-y-4">
<div>
<label class="input-label">{{ t('admin.accounts.bedrockAccessKeyId') }}</label>
<input
v-model="editBedrockAccessKeyId"
type="text"
class="input font-mono"
placeholder="AKIA..."
/>
</div>
<div>
<label class="input-label">{{ t('admin.accounts.bedrockSecretAccessKey') }}</label>
<input
v-model="editBedrockSecretAccessKey"
type="password"
class="input font-mono"
:placeholder="t('admin.accounts.bedrockSecretKeyLeaveEmpty')"
/>
<p class="input-hint">{{ t('admin.accounts.bedrockSecretKeyLeaveEmpty') }}</p>
</div>
<div>
<label class="input-label">{{ t('admin.accounts.bedrockSessionToken') }}</label>
<input
v-model="editBedrockSessionToken"
type="password"
class="input font-mono"
:placeholder="t('admin.accounts.bedrockSecretKeyLeaveEmpty')"
/>
<p class="input-hint">{{ t('admin.accounts.bedrockSessionTokenHint') }}</p>
</div>
<div>
<label class="input-label">{{ t('admin.accounts.bedrockRegion') }}</label>
<input
v-model="editBedrockRegion"
type="text"
class="input"
placeholder="us-east-1"
/>
<p class="input-hint">{{ t('admin.accounts.bedrockRegionHint') }}</p>
</div>
<div>
<label class="flex items-center gap-2 cursor-pointer">
<input
v-model="editBedrockForceGlobal"
type="checkbox"
class="rounded border-gray-300 text-primary-600 focus:ring-primary-500 dark:border-dark-500"
/>
<span class="text-sm text-gray-700 dark:text-gray-300">{{ t('admin.accounts.bedrockForceGlobal') }}</span>
</label>
<p class="input-hint mt-1">{{ t('admin.accounts.bedrockForceGlobalHint') }}</p>
</div>
<!-- Model Restriction for Bedrock -->
<div class="border-t border-gray-200 pt-4 dark:border-dark-600">
<label class="input-label">{{ t('admin.accounts.modelRestriction') }}</label>
<!-- Mode Toggle -->
<div class="mb-4 flex gap-2">
<button
type="button"
@click="modelRestrictionMode = 'whitelist'"
:class="[
'flex-1 rounded-lg px-4 py-2 text-sm font-medium transition-all',
modelRestrictionMode === 'whitelist'
? 'bg-primary-100 text-primary-700 dark:bg-primary-900/30 dark:text-primary-400'
: 'bg-gray-100 text-gray-600 hover:bg-gray-200 dark:bg-dark-600 dark:text-gray-400 dark:hover:bg-dark-500'
]"
>
{{ t('admin.accounts.modelWhitelist') }}
</button>
<button
type="button"
@click="modelRestrictionMode = 'mapping'"
:class="[
'flex-1 rounded-lg px-4 py-2 text-sm font-medium transition-all',
modelRestrictionMode === 'mapping'
? 'bg-purple-100 text-purple-700 dark:bg-purple-900/30 dark:text-purple-400'
: 'bg-gray-100 text-gray-600 hover:bg-gray-200 dark:bg-dark-600 dark:text-gray-400 dark:hover:bg-dark-500'
]"
>
{{ t('admin.accounts.modelMapping') }}
</button>
</div>
<!-- Whitelist Mode -->
<div v-if="modelRestrictionMode === 'whitelist'">
<ModelWhitelistSelector v-model="allowedModels" platform="anthropic" />
<p class="text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.selectedModels', { count: allowedModels.length }) }}
<span v-if="allowedModels.length === 0">{{ t('admin.accounts.supportsAllModels') }}</span>
</p>
</div>
<!-- Mapping Mode -->
<div v-else class="space-y-3">
<div v-for="(mapping, index) in modelMappings" :key="getModelMappingKey(mapping)" class="flex items-center gap-2">
<input v-model="mapping.from" type="text" class="input flex-1" :placeholder="t('admin.accounts.fromModel')" />
<span class="text-gray-400">→</span>
<input v-model="mapping.to" type="text" class="input flex-1" :placeholder="t('admin.accounts.toModel')" />
<button type="button" @click="modelMappings.splice(index, 1)" class="text-red-500 hover:text-red-700">
<Icon name="trash" size="sm" />
</button>
</div>
<button type="button" @click="modelMappings.push({ from: '', to: '' })" class="btn btn-secondary text-sm">
+ {{ t('admin.accounts.addMapping') }}
</button>
<!-- Bedrock Preset Mappings -->
<div class="flex flex-wrap gap-2">
<button
v-for="preset in bedrockPresets"
:key="preset.from"
type="button"
@click="modelMappings.push({ from: preset.from, to: preset.to })"
:class="['rounded-lg px-3 py-1 text-xs transition-colors', preset.color]"
>
+ {{ preset.label }}
</button>
</div>
</div>
</div>
</div>
<!-- Bedrock API Key fields (only for bedrock-apikey type) -->
<div v-if="account.type === 'bedrock-apikey'" class="space-y-4">
<div>
<label class="input-label">{{ t('admin.accounts.bedrockApiKeyInput') }}</label>
<input
v-model="editBedrockApiKeyValue"
type="password"
class="input font-mono"
:placeholder="t('admin.accounts.bedrockApiKeyLeaveEmpty')"
/>
<p class="input-hint">{{ t('admin.accounts.bedrockApiKeyLeaveEmpty') }}</p>
</div>
<div>
<label class="input-label">{{ t('admin.accounts.bedrockRegion') }}</label>
<input
v-model="editBedrockApiKeyRegion"
type="text"
class="input"
placeholder="us-east-1"
/>
<p class="input-hint">{{ t('admin.accounts.bedrockRegionHint') }}</p>
</div>
<div>
<label class="flex items-center gap-2 cursor-pointer">
<input
v-model="editBedrockApiKeyForceGlobal"
type="checkbox"
class="rounded border-gray-300 text-primary-600 focus:ring-primary-500 dark:border-dark-500"
/>
<span class="text-sm text-gray-700 dark:text-gray-300">{{ t('admin.accounts.bedrockForceGlobal') }}</span>
</label>
<p class="input-hint mt-1">{{ t('admin.accounts.bedrockForceGlobalHint') }}</p>
</div>
<!-- Model Restriction for Bedrock API Key -->
<div class="border-t border-gray-200 pt-4 dark:border-dark-600">
<label class="input-label">{{ t('admin.accounts.modelRestriction') }}</label>
<!-- Mode Toggle -->
<div class="mb-4 flex gap-2">
<button
type="button"
@click="modelRestrictionMode = 'whitelist'"
:class="[
'flex-1 rounded-lg px-4 py-2 text-sm font-medium transition-all',
modelRestrictionMode === 'whitelist'
? 'bg-primary-100 text-primary-700 dark:bg-primary-900/30 dark:text-primary-400'
: 'bg-gray-100 text-gray-600 hover:bg-gray-200 dark:bg-dark-600 dark:text-gray-400 dark:hover:bg-dark-500'
]"
>
{{ t('admin.accounts.modelWhitelist') }}
</button>
<button
type="button"
@click="modelRestrictionMode = 'mapping'"
:class="[
'flex-1 rounded-lg px-4 py-2 text-sm font-medium transition-all',
modelRestrictionMode === 'mapping'
? 'bg-purple-100 text-purple-700 dark:bg-purple-900/30 dark:text-purple-400'
: 'bg-gray-100 text-gray-600 hover:bg-gray-200 dark:bg-dark-600 dark:text-gray-400 dark:hover:bg-dark-500'
]"
>
{{ t('admin.accounts.modelMapping') }}
</button>
</div>
<!-- Whitelist Mode -->
<div v-if="modelRestrictionMode === 'whitelist'">
<ModelWhitelistSelector v-model="allowedModels" platform="anthropic" />
<p class="text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.selectedModels', { count: allowedModels.length }) }}
<span v-if="allowedModels.length === 0">{{ t('admin.accounts.supportsAllModels') }}</span>
</p>
</div>
<!-- Mapping Mode -->
<div v-else class="space-y-3">
<div v-for="(mapping, index) in modelMappings" :key="getModelMappingKey(mapping)" class="flex items-center gap-2">
<input v-model="mapping.from" type="text" class="input flex-1" :placeholder="t('admin.accounts.fromModel')" />
<span class="text-gray-400">→</span>
<input v-model="mapping.to" type="text" class="input flex-1" :placeholder="t('admin.accounts.toModel')" />
<button type="button" @click="modelMappings.splice(index, 1)" class="text-red-500 hover:text-red-700">
<Icon name="trash" size="sm" />
</button>
</div>
<button type="button" @click="modelMappings.push({ from: '', to: '' })" class="btn btn-secondary text-sm">
+ {{ t('admin.accounts.addMapping') }}
</button>
<!-- Bedrock Preset Mappings -->
<div class="flex flex-wrap gap-2">
<button
v-for="preset in bedrockPresets"
:key="preset.from"
type="button"
@click="modelMappings.push({ from: preset.from, to: preset.to })"
:class="['rounded-lg px-3 py-1 text-xs transition-colors', preset.color]"
>
+ {{ preset.label }}
</button>
</div>
</div>
</div>
</div>
<!-- Antigravity model restriction (applies to all antigravity types) -->
<!-- Antigravity 只支持模型映射模式，不支持白名单模式 -->
<div v-if="account.platform === 'antigravity'" class="border-t border-gray-200 pt-4 dark:border-dark-600">
......@@ -1529,6 +1756,7 @@ const baseUrlHint = computed(() => {
})
const antigravityPresetMappings = computed(() => getPresetMappingsByPlatform('antigravity'))
const bedrockPresets = computed(() => getPresetMappingsByPlatform('bedrock'))
// Model mapping type
interface ModelMapping {
......@@ -1547,6 +1775,17 @@ interface TempUnschedRuleForm {
const submitting = ref(false)
const editBaseUrl = ref('https://api.anthropic.com')
const editApiKey = ref('')
// Bedrock credentials
const editBedrockAccessKeyId = ref('')
const editBedrockSecretAccessKey = ref('')
const editBedrockSessionToken = ref('')
const editBedrockRegion = ref('')
const editBedrockForceGlobal = ref(false)
// Bedrock API Key credentials
const editBedrockApiKeyValue = ref('')
const editBedrockApiKeyRegion = ref('')
const editBedrockApiKeyForceGlobal = ref(false)
const modelMappings = ref<ModelMapping[]>([])
const modelRestrictionMode = ref<'whitelist' | 'mapping'>('whitelist')
const allowedModels = ref<string[]>([])
......@@ -1889,6 +2128,58 @@ watch(
} else {
selectedErrorCodes.value = []
}
} else if (newAccount.type === 'bedrock' && newAccount.credentials) {
const bedrockCreds = newAccount.credentials as Record<string, unknown>
editBedrockAccessKeyId.value = (bedrockCreds.aws_access_key_id as string) || ''
editBedrockRegion.value = (bedrockCreds.aws_region as string) || ''
editBedrockForceGlobal.value = (bedrockCreds.aws_force_global as string) === 'true'
editBedrockSecretAccessKey.value = ''
editBedrockSessionToken.value = ''
// Load model mappings for bedrock
const existingMappings = bedrockCreds.model_mapping as Record<string, string> | undefined
if (existingMappings && typeof existingMappings === 'object') {
const entries = Object.entries(existingMappings)
const isWhitelistMode = entries.length > 0 && entries.every(([from, to]) => from === to)
if (isWhitelistMode) {
modelRestrictionMode.value = 'whitelist'
allowedModels.value = entries.map(([from]) => from)
modelMappings.value = []
} else {
modelRestrictionMode.value = 'mapping'
modelMappings.value = entries.map(([from, to]) => ({ from, to }))
allowedModels.value = []
}
} else {
modelRestrictionMode.value = 'whitelist'
modelMappings.value = []
allowedModels.value = []
}
} else if (newAccount.type === 'bedrock-apikey' && newAccount.credentials) {
const bedrockApiKeyCreds = newAccount.credentials as Record<string, unknown>
editBedrockApiKeyRegion.value = (bedrockApiKeyCreds.aws_region as string) || 'us-east-1'
editBedrockApiKeyForceGlobal.value = (bedrockApiKeyCreds.aws_force_global as string) === 'true'
editBedrockApiKeyValue.value = ''
// Load model mappings for bedrock-apikey
const existingMappings = bedrockApiKeyCreds.model_mapping as Record<string, string> | undefined
if (existingMappings && typeof existingMappings === 'object') {
const entries = Object.entries(existingMappings)
const isWhitelistMode = entries.length > 0 && entries.every(([from, to]) => from === to)
if (isWhitelistMode) {
modelRestrictionMode.value = 'whitelist'
allowedModels.value = entries.map(([from]) => from)
modelMappings.value = []
} else {
modelRestrictionMode.value = 'mapping'
modelMappings.value = entries.map(([from, to]) => ({ from, to }))
allowedModels.value = []
}
} else {
modelRestrictionMode.value = 'whitelist'
modelMappings.value = []
allowedModels.value = []
}
} else if (newAccount.type === 'upstream' && newAccount.credentials) {
const credentials = newAccount.credentials as Record<string, unknown>
editBaseUrl.value = (credentials.base_url as string) || ''
......@@ -2431,6 +2722,70 @@ const handleSubmit = async () => {
return
}
updatePayload.credentials = newCredentials
} else if (props.account.type === 'bedrock') {
const currentCredentials = (props.account.credentials as Record<string, unknown>) || {}
const newCredentials: Record<string, unknown> = { ...currentCredentials }
newCredentials.aws_access_key_id = editBedrockAccessKeyId.value.trim()
newCredentials.aws_region = editBedrockRegion.value.trim()
if (editBedrockForceGlobal.value) {
newCredentials.aws_force_global = 'true'
} else {
delete newCredentials.aws_force_global
}
// Only update secrets if user provided new values
if (editBedrockSecretAccessKey.value.trim()) {
newCredentials.aws_secret_access_key = editBedrockSecretAccessKey.value.trim()
}
if (editBedrockSessionToken.value.trim()) {
newCredentials.aws_session_token = editBedrockSessionToken.value.trim()
}
// Model mapping
const modelMapping = buildModelMappingObject(modelRestrictionMode.value, allowedModels.value, modelMappings.value)
if (modelMapping) {
newCredentials.model_mapping = modelMapping
} else {
delete newCredentials.model_mapping
}
applyInterceptWarmup(newCredentials, interceptWarmupRequests.value, 'edit')
if (!applyTempUnschedConfig(newCredentials)) {
return
}
updatePayload.credentials = newCredentials
} else if (props.account.type === 'bedrock-apikey') {
const currentCredentials = (props.account.credentials as Record<string, unknown>) || {}
const newCredentials: Record<string, unknown> = { ...currentCredentials }
newCredentials.aws_region = editBedrockApiKeyRegion.value.trim() || 'us-east-1'
if (editBedrockApiKeyForceGlobal.value) {
newCredentials.aws_force_global = 'true'
} else {
delete newCredentials.aws_force_global
}
// Only update API key if user provided new value
if (editBedrockApiKeyValue.value.trim()) {
newCredentials.api_key = editBedrockApiKeyValue.value.trim()
}
// Model mapping
const modelMapping = buildModelMappingObject(modelRestrictionMode.value, allowedModels.value, modelMappings.value)
if (modelMapping) {
newCredentials.model_mapping = modelMapping
} else {
delete newCredentials.model_mapping
}
applyInterceptWarmup(newCredentials, interceptWarmupRequests.value, 'edit')
if (!applyTempUnschedConfig(newCredentials)) {
return
}
updatePayload.credentials = newCredentials
} else {
// For oauth/setup-token types, only update intercept_warmup_requests if changed
......
frontend/src/components/common/PlatformTypeBadge.vue
View file @ 11f7b835
......@@ -82,6 +82,8 @@ const typeLabel = computed(() => {
return 'Token'
case 'apikey':
return 'Key'
case 'bedrock':
return 'Bedrock'
default:
return props.type
}
......
frontend/src/composables/useModelWhitelist.ts
View file @ 11f7b835
......@@ -331,6 +331,15 @@ const antigravityPresetMappings = [
{ label: 'Opus 4.6-thinking', from: 'claude-opus-4-6-thinking', to: 'claude-opus-4-6-thinking', color: 'bg-pink-100 text-pink-700 hover:bg-pink-200 dark:bg-pink-900/30 dark:text-pink-400' }
]
// Bedrock 预设映射（与后端 DefaultBedrockModelMapping 保持一致）
const bedrockPresetMappings = [
{ label: 'Opus 4.6', from: 'claude-opus-4-6', to: 'us.anthropic.claude-opus-4-6-v1', color: 'bg-pink-100 text-pink-700 hover:bg-pink-200 dark:bg-pink-900/30 dark:text-pink-400' },
{ label: 'Sonnet 4.6', from: 'claude-sonnet-4-6', to: 'us.anthropic.claude-sonnet-4-6', color: 'bg-cyan-100 text-cyan-700 hover:bg-cyan-200 dark:bg-cyan-900/30 dark:text-cyan-400' },
{ label: 'Opus 4.5', from: 'claude-opus-4-5-thinking', to: 'us.anthropic.claude-opus-4-5-20251101-v1:0', color: 'bg-pink-100 text-pink-700 hover:bg-pink-200 dark:bg-pink-900/30 dark:text-pink-400' },
{ label: 'Sonnet 4.5', from: 'claude-sonnet-4-5', to: 'us.anthropic.claude-sonnet-4-5-20250929-v1:0', color: 'bg-cyan-100 text-cyan-700 hover:bg-cyan-200 dark:bg-cyan-900/30 dark:text-cyan-400' },
{ label: 'Haiku 4.5', from: 'claude-haiku-4-5', to: 'us.anthropic.claude-haiku-4-5-20251001-v1:0', color: 'bg-green-100 text-green-700 hover:bg-green-200 dark:bg-green-900/30 dark:text-green-400' },
]
// Antigravity 默认映射（从后端 API 获取，与 constants.go 保持一致）
// 使用 fetchAntigravityDefaultMappings() 异步获取
import { getAntigravityDefaultModelMapping } from '@/api/admin/accounts'
......@@ -403,6 +412,7 @@ export function getPresetMappingsByPlatform(platform: string) {
if (platform === 'gemini') return geminiPresetMappings
if (platform === 'sora') return soraPresetMappings
if (platform === 'antigravity') return antigravityPresetMappings
if (platform === 'bedrock' || platform === 'bedrock-apikey') return bedrockPresetMappings
return anthropicPresetMappings
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment