383 lines
12 KiB
Go
383 lines
12 KiB
Go
|
package service
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"bufio"
|
|||
|
|
"bytes"
|
|||
|
|
"context"
|
|||
|
|
"encoding/json"
|
|||
|
|
"fmt"
|
|||
|
|
"io"
|
|||
|
|
"net/http"
|
|||
|
|
"strings"
|
|||
|
|
"sync/atomic"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
|||
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
|||
|
|
"github.com/gin-gonic/gin"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// ForwardUpstream 使用 base_url + /v1/messages + 双 header 认证透传上游 Claude 请求
|
|||
|
|
func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
|||
|
|
beginUpstreamResponseModelObservation(c)
|
|||
|
|
startTime := time.Now()
|
|||
|
|
sessionID := getSessionID(c)
|
|||
|
|
prefix := logPrefix(sessionID, account.Name)
|
|||
|
|
|
|||
|
|
// 获取上游配置
|
|||
|
|
baseURL := strings.TrimSpace(account.GetCredential("base_url"))
|
|||
|
|
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
|
|||
|
|
if baseURL == "" || apiKey == "" {
|
|||
|
|
return nil, fmt.Errorf("upstream account missing base_url or api_key")
|
|||
|
|
}
|
|||
|
|
baseURL = strings.TrimSuffix(baseURL, "/")
|
|||
|
|
|
|||
|
|
// 解析请求获取模型信息
|
|||
|
|
var claudeReq antigravity.ClaudeRequest
|
|||
|
|
if err := json.Unmarshal(body, &claudeReq); err != nil {
|
|||
|
|
return nil, fmt.Errorf("parse claude request: %w", err)
|
|||
|
|
}
|
|||
|
|
if strings.TrimSpace(claudeReq.Model) == "" {
|
|||
|
|
return nil, fmt.Errorf("missing model")
|
|||
|
|
}
|
|||
|
|
originalModel := claudeReq.Model
|
|||
|
|
|
|||
|
|
// 构建上游请求 URL
|
|||
|
|
upstreamURL := baseURL + "/v1/messages"
|
|||
|
|
|
|||
|
|
// 能力维度 sanitize：Anthropic-compatible 上游透传路径也需要保证 body↔beta header
|
|||
|
|
// 对称。客户端 anthropic-beta header 不含 context-management-2025-06-27 但 body 带
|
|||
|
|
// context_management 时 strip，与 Anthropic 直连 / Bedrock / Vertex 路径保持一致。
|
|||
|
|
clientBeta := c.GetHeader("anthropic-beta")
|
|||
|
|
if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed {
|
|||
|
|
body = sanitized
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 创建请求
|
|||
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body))
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, fmt.Errorf("create upstream request: %w", err)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 设置请求头
|
|||
|
|
req.Header.Set("Content-Type", "application/json")
|
|||
|
|
req.Header.Set("Authorization", "Bearer "+apiKey)
|
|||
|
|
req.Header.Set("x-api-key", apiKey) // Claude API 兼容
|
|||
|
|
|
|||
|
|
// 透传 Claude 相关 headers
|
|||
|
|
if v := c.GetHeader("anthropic-version"); v != "" {
|
|||
|
|
req.Header.Set("anthropic-version", v)
|
|||
|
|
}
|
|||
|
|
if v := clientBeta; v != "" {
|
|||
|
|
req.Header.Set("anthropic-beta", v)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 代理 URL
|
|||
|
|
proxyURL := ""
|
|||
|
|
if account.ProxyID != nil && account.Proxy != nil {
|
|||
|
|
proxyURL = account.Proxy.URL()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 发送请求
|
|||
|
|
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
|
|||
|
|
if err != nil {
|
|||
|
|
logger.LegacyPrintf("service.antigravity_gateway", "%s upstream request failed: %v", prefix, err)
|
|||
|
|
return nil, fmt.Errorf("upstream request failed: %w", err)
|
|||
|
|
}
|
|||
|
|
defer func() { _ = resp.Body.Close() }()
|
|||
|
|
|
|||
|
|
// 处理错误响应
|
|||
|
|
if resp.StatusCode >= 400 {
|
|||
|
|
respBody := s.readUpstreamErrorBody(resp)
|
|||
|
|
|
|||
|
|
// 429 错误时标记账号限流
|
|||
|
|
if resp.StatusCode == http.StatusTooManyRequests {
|
|||
|
|
s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", false)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 透传上游错误
|
|||
|
|
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
|||
|
|
c.Status(resp.StatusCode)
|
|||
|
|
_, _ = c.Writer.Write(respBody)
|
|||
|
|
|
|||
|
|
return &ForwardResult{
|
|||
|
|
Model: originalModel,
|
|||
|
|
}, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 处理成功响应（流式/非流式）
|
|||
|
|
var usage *ClaudeUsage
|
|||
|
|
var firstTokenMs *int
|
|||
|
|
var clientDisconnect bool
|
|||
|
|
|
|||
|
|
if claudeReq.Stream {
|
|||
|
|
// 流式响应：透传
|
|||
|
|
c.Header("Content-Type", "text/event-stream")
|
|||
|
|
c.Header("Cache-Control", "no-cache")
|
|||
|
|
c.Header("Connection", "keep-alive")
|
|||
|
|
c.Header("X-Accel-Buffering", "no")
|
|||
|
|
c.Status(http.StatusOK)
|
|||
|
|
|
|||
|
|
streamRes := s.streamUpstreamResponse(c, resp, startTime)
|
|||
|
|
usage = streamRes.usage
|
|||
|
|
firstTokenMs = streamRes.firstTokenMs
|
|||
|
|
clientDisconnect = streamRes.clientDisconnect
|
|||
|
|
} else {
|
|||
|
|
// 非流式响应：直接透传
|
|||
|
|
respBody, err := io.ReadAll(resp.Body)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, fmt.Errorf("read upstream response: %w", err)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 提取 usage
|
|||
|
|
upstreamResponseModelObserverFromContext(c).ObserveAnthropic(respBody)
|
|||
|
|
usage = s.extractClaudeUsage(respBody)
|
|||
|
|
|
|||
|
|
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
|||
|
|
c.Status(http.StatusOK)
|
|||
|
|
_, _ = c.Writer.Write(respBody)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 构建计费结果
|
|||
|
|
duration := time.Since(startTime)
|
|||
|
|
logger.LegacyPrintf("service.antigravity_gateway", "%s status=success duration_ms=%d", prefix, duration.Milliseconds())
|
|||
|
|
|
|||
|
|
return &ForwardResult{
|
|||
|
|
Model: originalModel,
|
|||
|
|
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
|||
|
|
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
|||
|
|
Stream: claudeReq.Stream,
|
|||
|
|
Duration: duration,
|
|||
|
|
FirstTokenMs: firstTokenMs,
|
|||
|
|
ClientDisconnect: clientDisconnect,
|
|||
|
|
Usage: ClaudeUsage{
|
|||
|
|
InputTokens: usage.InputTokens,
|
|||
|
|
OutputTokens: usage.OutputTokens,
|
|||
|
|
CacheReadInputTokens: usage.CacheReadInputTokens,
|
|||
|
|
CacheCreationInputTokens: usage.CacheCreationInputTokens,
|
|||
|
|
},
|
|||
|
|
}, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// streamUpstreamResponse 透传上游 SSE 流并提取 Claude usage
|
|||
|
|
func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp *http.Response, startTime time.Time) *antigravityStreamResult {
|
|||
|
|
usage := &ClaudeUsage{}
|
|||
|
|
var firstTokenMs *int
|
|||
|
|
|
|||
|
|
scanner := bufio.NewScanner(resp.Body)
|
|||
|
|
maxLineSize := defaultMaxLineSize
|
|||
|
|
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
|
|||
|
|
maxLineSize = s.settingService.cfg.Gateway.MaxLineSize
|
|||
|
|
}
|
|||
|
|
scanner.Buffer(make([]byte, 64*1024), maxLineSize)
|
|||
|
|
|
|||
|
|
type scanEvent struct {
|
|||
|
|
line string
|
|||
|
|
err error
|
|||
|
|
}
|
|||
|
|
events := make(chan scanEvent, 16)
|
|||
|
|
done := make(chan struct{})
|
|||
|
|
sendEvent := func(ev scanEvent) bool {
|
|||
|
|
select {
|
|||
|
|
case events <- ev:
|
|||
|
|
return true
|
|||
|
|
case <-done:
|
|||
|
|
return false
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
var lastReadAt int64
|
|||
|
|
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
|
|||
|
|
go func() {
|
|||
|
|
defer close(events)
|
|||
|
|
for scanner.Scan() {
|
|||
|
|
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
|
|||
|
|
if !sendEvent(scanEvent{line: scanner.Text()}) {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
if err := scanner.Err(); err != nil {
|
|||
|
|
_ = sendEvent(scanEvent{err: err})
|
|||
|
|
}
|
|||
|
|
}()
|
|||
|
|
defer close(done)
|
|||
|
|
|
|||
|
|
streamInterval := time.Duration(0)
|
|||
|
|
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 {
|
|||
|
|
streamInterval = time.Duration(s.settingService.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
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 下游 keepalive：防止代理/Cloudflare Tunnel 因连接空闲而断开
|
|||
|
|
keepaliveInterval := time.Duration(0)
|
|||
|
|
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 {
|
|||
|
|
keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second
|
|||
|
|
}
|
|||
|
|
var keepaliveTicker *time.Ticker
|
|||
|
|
if keepaliveInterval > 0 {
|
|||
|
|
keepaliveTicker = time.NewTicker(keepaliveInterval)
|
|||
|
|
defer keepaliveTicker.Stop()
|
|||
|
|
}
|
|||
|
|
var keepaliveCh <-chan time.Time
|
|||
|
|
if keepaliveTicker != nil {
|
|||
|
|
keepaliveCh = keepaliveTicker.C
|
|||
|
|
}
|
|||
|
|
lastDataAt := time.Now()
|
|||
|
|
|
|||
|
|
flusher, _ := c.Writer.(http.Flusher)
|
|||
|
|
cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity upstream")
|
|||
|
|
|
|||
|
|
for {
|
|||
|
|
select {
|
|||
|
|
case ev, ok := <-events:
|
|||
|
|
if !ok {
|
|||
|
|
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}
|
|||
|
|
}
|
|||
|
|
if ev.err != nil {
|
|||
|
|
if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity upstream"); handled {
|
|||
|
|
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect}
|
|||
|
|
}
|
|||
|
|
logger.LegacyPrintf("service.antigravity_gateway", "Stream read error (antigravity upstream): %v", ev.err)
|
|||
|
|
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
lastDataAt = time.Now()
|
|||
|
|
|
|||
|
|
line := ev.line
|
|||
|
|
if data, ok := extractAnthropicSSEDataLine(line); ok {
|
|||
|
|
upstreamResponseModelObserverFromContext(c).ObserveAnthropic([]byte(strings.TrimSpace(data)))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 记录首 token 时间
|
|||
|
|
if firstTokenMs == nil && len(line) > 0 {
|
|||
|
|
ms := int(time.Since(startTime).Milliseconds())
|
|||
|
|
firstTokenMs = &ms
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 尝试从 message_delta 或 message_stop 事件提取 usage
|
|||
|
|
s.extractSSEUsage(line, usage)
|
|||
|
|
|
|||
|
|
// 透传行
|
|||
|
|
cw.Fprintf("%s\n", line)
|
|||
|
|
|
|||
|
|
case <-intervalCh:
|
|||
|
|
lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
|
|||
|
|
if time.Since(lastRead) < streamInterval {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if cw.Disconnected() {
|
|||
|
|
logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity upstream), returning collected usage")
|
|||
|
|
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}
|
|||
|
|
}
|
|||
|
|
logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity upstream)")
|
|||
|
|
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}
|
|||
|
|
|
|||
|
|
case <-keepaliveCh:
|
|||
|
|
if cw.Disconnected() {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
if time.Since(lastDataAt) < keepaliveInterval {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
// SSE ping 事件：Anthropic 原生格式，客户端会正确处理，
|
|||
|
|
// 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开
|
|||
|
|
if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") {
|
|||
|
|
logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity upstream), continuing to drain upstream for billing")
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// extractSSEUsage 从 SSE data 行中提取 Claude usage（用于流式透传场景）
|
|||
|
|
//
|
|||
|
|
// Anthropic streaming 的 usage 字段分布在两类事件中：
|
|||
|
|
// - message_start：嵌套在 event.message.usage（input_tokens、cache_creation_input_tokens、
|
|||
|
|
// cache_read_input_tokens 等输入侧字段）
|
|||
|
|
// - message_delta：位于顶层 event.usage（流结束时的最终 output_tokens）
|
|||
|
|
//
|
|||
|
|
// 仅读取顶层 event.usage 会漏掉 message_start 的输入侧字段，导致流式透传请求落库的
|
|||
|
|
// usage_logs 记录 input_tokens=0。
|
|||
|
|
func (s *AntigravityGatewayService) extractSSEUsage(line string, usage *ClaudeUsage) {
|
|||
|
|
if !strings.HasPrefix(line, "data: ") {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
dataStr := strings.TrimPrefix(line, "data: ")
|
|||
|
|
var event map[string]any
|
|||
|
|
if json.Unmarshal([]byte(dataStr), &event) != nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
var u map[string]any
|
|||
|
|
if eventType, _ := event["type"].(string); eventType == "message_start" {
|
|||
|
|
if msg, ok := event["message"].(map[string]any); ok {
|
|||
|
|
u, _ = msg["usage"].(map[string]any)
|
|||
|
|
}
|
|||
|
|
} else {
|
|||
|
|
u, _ = event["usage"].(map[string]any)
|
|||
|
|
}
|
|||
|
|
if u == nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
if v, ok := u["input_tokens"].(float64); ok && int(v) > 0 {
|
|||
|
|
usage.InputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := u["output_tokens"].(float64); ok && int(v) > 0 {
|
|||
|
|
usage.OutputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := u["cache_read_input_tokens"].(float64); ok && int(v) > 0 {
|
|||
|
|
usage.CacheReadInputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := u["cache_creation_input_tokens"].(float64); ok && int(v) > 0 {
|
|||
|
|
usage.CacheCreationInputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
// 解析嵌套的 cache_creation 对象中的 5m/1h 明细
|
|||
|
|
if cc, ok := u["cache_creation"].(map[string]any); ok {
|
|||
|
|
if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok {
|
|||
|
|
usage.CacheCreation5mTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok {
|
|||
|
|
usage.CacheCreation1hTokens = int(v)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// extractClaudeUsage 从非流式 Claude 响应提取 usage
|
|||
|
|
func (s *AntigravityGatewayService) extractClaudeUsage(body []byte) *ClaudeUsage {
|
|||
|
|
usage := &ClaudeUsage{}
|
|||
|
|
var resp map[string]any
|
|||
|
|
if json.Unmarshal(body, &resp) != nil {
|
|||
|
|
return usage
|
|||
|
|
}
|
|||
|
|
if u, ok := resp["usage"].(map[string]any); ok {
|
|||
|
|
if v, ok := u["input_tokens"].(float64); ok {
|
|||
|
|
usage.InputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := u["output_tokens"].(float64); ok {
|
|||
|
|
usage.OutputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := u["cache_read_input_tokens"].(float64); ok {
|
|||
|
|
usage.CacheReadInputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := u["cache_creation_input_tokens"].(float64); ok {
|
|||
|
|
usage.CacheCreationInputTokens = int(v)
|
|||
|
|
}
|
|||
|
|
// 解析嵌套的 cache_creation 对象中的 5m/1h 明细
|
|||
|
|
if cc, ok := u["cache_creation"].(map[string]any); ok {
|
|||
|
|
if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok {
|
|||
|
|
usage.CacheCreation5mTokens = int(v)
|
|||
|
|
}
|
|||
|
|
if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok {
|
|||
|
|
usage.CacheCreation1hTokens = int(v)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return usage
|
|||
|
|
}
|