世界
608b7e7fa2
When the access token expires and refreshToken() gets 429, getAccessToken() returned the error but left credentials unchanged with no cooldown. Every subsequent request re-attempted the refresh, creating a burst that overwhelmed the token endpoint. - refreshToken() now returns Retry-After duration from 429 response headers (-1 when no header present, meaning permanently blocked) - getAccessToken() caches the 429 and blocks further refresh attempts until Retry-After expires (or permanently if no header) - reloadCredentials() clears the block when new credentials are loaded from file - Remove go pollUsage() on upstream errors (unrelated to usage state)
513 lines
16 KiB
Go
513 lines
16 KiB
Go
package ocm
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"mime"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/sagernet/sing-box/log"
|
|
"github.com/sagernet/sing-box/option"
|
|
"github.com/sagernet/sing/common/buf"
|
|
E "github.com/sagernet/sing/common/exceptions"
|
|
|
|
"github.com/openai/openai-go/v3"
|
|
"github.com/openai/openai-go/v3/responses"
|
|
)
|
|
|
|
func weeklyCycleHintForLimit(headers http.Header, limitIdentifier string) *WeeklyCycleHint {
|
|
normalizedLimitIdentifier := normalizeRateLimitIdentifier(limitIdentifier)
|
|
if normalizedLimitIdentifier == "" {
|
|
return nil
|
|
}
|
|
|
|
windowHeader := "x-" + normalizedLimitIdentifier + "-secondary-window-minutes"
|
|
resetHeader := "x-" + normalizedLimitIdentifier + "-secondary-reset-at"
|
|
|
|
windowMinutes, hasWindowMinutes := parseInt64Header(headers, windowHeader)
|
|
resetAtUnix, hasResetAt := parseInt64Header(headers, resetHeader)
|
|
if !hasWindowMinutes || !hasResetAt || windowMinutes <= 0 || resetAtUnix <= 0 {
|
|
return nil
|
|
}
|
|
|
|
return &WeeklyCycleHint{
|
|
WindowMinutes: windowMinutes,
|
|
ResetAt: time.Unix(resetAtUnix, 0).UTC(),
|
|
}
|
|
}
|
|
|
|
func extractWeeklyCycleHint(headers http.Header) *WeeklyCycleHint {
|
|
activeLimitIdentifier := normalizeRateLimitIdentifier(headers.Get("x-codex-active-limit"))
|
|
if activeLimitIdentifier != "" {
|
|
if activeHint := weeklyCycleHintForLimit(headers, activeLimitIdentifier); activeHint != nil {
|
|
return activeHint
|
|
}
|
|
}
|
|
return weeklyCycleHintForLimit(headers, "codex")
|
|
}
|
|
|
|
func (s *Service) resolveCredentialProvider(username string) (credentialProvider, error) {
|
|
if len(s.options.Users) > 0 {
|
|
return credentialForUser(s.userConfigMap, s.providers, username)
|
|
}
|
|
provider := s.providers[s.options.Credentials[0].Tag]
|
|
if provider == nil {
|
|
return nil, E.New("no credential available")
|
|
}
|
|
return provider, nil
|
|
}
|
|
|
|
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
ctx := log.ContextWithNewID(r.Context())
|
|
if r.URL.Path == "/ocm/v1/status" {
|
|
s.handleStatusEndpoint(w, r)
|
|
return
|
|
}
|
|
|
|
if r.URL.Path == "/ocm/v1/reverse" {
|
|
s.handleReverseConnect(ctx, w, r)
|
|
return
|
|
}
|
|
|
|
path := r.URL.Path
|
|
if !strings.HasPrefix(path, "/v1/") {
|
|
writeJSONError(w, r, http.StatusNotFound, "invalid_request_error", "path must start with /v1/")
|
|
return
|
|
}
|
|
|
|
if r.Header.Get("X-Api-Key") != "" || r.Header.Get("Api-Key") != "" {
|
|
writeJSONError(w, r, http.StatusBadRequest, "invalid_request_error",
|
|
"API key authentication is not supported; use Authorization: Bearer with an OCM user token")
|
|
return
|
|
}
|
|
|
|
var username string
|
|
if len(s.options.Users) > 0 {
|
|
authHeader := r.Header.Get("Authorization")
|
|
if authHeader == "" {
|
|
s.logger.WarnContext(ctx, "authentication failed for request from ", r.RemoteAddr, ": missing Authorization header")
|
|
writeJSONError(w, r, http.StatusUnauthorized, "authentication_error", "missing api key")
|
|
return
|
|
}
|
|
clientToken := strings.TrimPrefix(authHeader, "Bearer ")
|
|
if clientToken == authHeader {
|
|
s.logger.WarnContext(ctx, "authentication failed for request from ", r.RemoteAddr, ": invalid Authorization format")
|
|
writeJSONError(w, r, http.StatusUnauthorized, "authentication_error", "invalid api key format")
|
|
return
|
|
}
|
|
var ok bool
|
|
username, ok = s.userManager.Authenticate(clientToken)
|
|
if !ok {
|
|
s.logger.WarnContext(ctx, "authentication failed for request from ", r.RemoteAddr, ": unknown key: ", clientToken)
|
|
writeJSONError(w, r, http.StatusUnauthorized, "authentication_error", "invalid api key")
|
|
return
|
|
}
|
|
}
|
|
|
|
sessionID := r.Header.Get("session_id")
|
|
|
|
// Resolve credential provider and user config
|
|
var provider credentialProvider
|
|
var userConfig *option.OCMUser
|
|
if len(s.options.Users) > 0 {
|
|
userConfig = s.userConfigMap[username]
|
|
var err error
|
|
provider, err = credentialForUser(s.userConfigMap, s.providers, username)
|
|
if err != nil {
|
|
s.logger.ErrorContext(ctx, "resolve credential: ", err)
|
|
writeJSONError(w, r, http.StatusInternalServerError, "api_error", err.Error())
|
|
return
|
|
}
|
|
} else {
|
|
provider = s.providers[s.options.Credentials[0].Tag]
|
|
}
|
|
if provider == nil {
|
|
writeJSONError(w, r, http.StatusInternalServerError, "api_error", "no credential available")
|
|
return
|
|
}
|
|
|
|
provider.pollIfStale()
|
|
|
|
selection := credentialSelectionForUser(userConfig)
|
|
|
|
selectedCredential, isNew, err := provider.selectCredential(sessionID, selection)
|
|
if err != nil {
|
|
writeNonRetryableCredentialError(w, unavailableCredentialMessage(provider, err.Error()))
|
|
return
|
|
}
|
|
|
|
if strings.EqualFold(r.Header.Get("Upgrade"), "websocket") && strings.HasPrefix(path, "/v1/responses") {
|
|
s.handleWebSocket(ctx, w, r, path, username, sessionID, userConfig, provider, selectedCredential, selection, isNew)
|
|
return
|
|
}
|
|
|
|
if !selectedCredential.isExternal() && selectedCredential.ocmIsAPIKeyMode() {
|
|
// API key mode path handling
|
|
} else if !selectedCredential.isExternal() {
|
|
if path == "/v1/chat/completions" {
|
|
writeJSONError(w, r, http.StatusBadRequest, "invalid_request_error",
|
|
"chat completions endpoint is only available in API key mode")
|
|
return
|
|
}
|
|
}
|
|
|
|
shouldTrackUsage := selectedCredential.usageTrackerOrNil() != nil &&
|
|
(path == "/v1/chat/completions" || strings.HasPrefix(path, "/v1/responses"))
|
|
canRetryRequest := len(provider.allCredentials()) > 1
|
|
|
|
// Read body for model extraction and retry buffer when JSON replay is useful.
|
|
var bodyBytes []byte
|
|
var requestModel string
|
|
var requestServiceTier string
|
|
if r.Body != nil && (shouldTrackUsage || canRetryRequest) {
|
|
mediaType, _, parseErr := mime.ParseMediaType(r.Header.Get("Content-Type"))
|
|
isJSONRequest := parseErr == nil && (mediaType == "application/json" || strings.HasSuffix(mediaType, "+json"))
|
|
if isJSONRequest {
|
|
bodyBytes, err = io.ReadAll(r.Body)
|
|
if err != nil {
|
|
s.logger.ErrorContext(ctx, "read request body: ", err)
|
|
writeJSONError(w, r, http.StatusInternalServerError, "api_error", "failed to read request body")
|
|
return
|
|
}
|
|
var request struct {
|
|
Model string `json:"model"`
|
|
ServiceTier string `json:"service_tier"`
|
|
}
|
|
if json.Unmarshal(bodyBytes, &request) == nil {
|
|
requestModel = request.Model
|
|
requestServiceTier = request.ServiceTier
|
|
}
|
|
r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
|
}
|
|
}
|
|
|
|
if isNew {
|
|
logParts := []any{"assigned credential ", selectedCredential.tagName()}
|
|
if sessionID != "" {
|
|
logParts = append(logParts, " for session ", sessionID)
|
|
}
|
|
if username != "" {
|
|
logParts = append(logParts, " by user ", username)
|
|
}
|
|
if requestModel != "" {
|
|
logParts = append(logParts, ", model=", requestModel)
|
|
}
|
|
if requestServiceTier == "priority" {
|
|
logParts = append(logParts, ", fast")
|
|
}
|
|
s.logger.DebugContext(ctx, logParts...)
|
|
}
|
|
|
|
requestContext := selectedCredential.wrapRequestContext(ctx)
|
|
{
|
|
currentRequestContext := requestContext
|
|
requestContext.addInterruptLink(provider.linkProviderInterrupt(selectedCredential, selection, func() {
|
|
currentRequestContext.cancelOnce.Do(currentRequestContext.cancelFunc)
|
|
}))
|
|
}
|
|
defer func() {
|
|
requestContext.cancelRequest()
|
|
}()
|
|
proxyRequest, err := selectedCredential.buildProxyRequest(requestContext, r, bodyBytes, s.httpHeaders)
|
|
if err != nil {
|
|
s.logger.ErrorContext(ctx, "create proxy request: ", err)
|
|
writeJSONError(w, r, http.StatusInternalServerError, "api_error", "Internal server error")
|
|
return
|
|
}
|
|
|
|
response, err := selectedCredential.httpClient().Do(proxyRequest)
|
|
if err != nil {
|
|
if r.Context().Err() != nil {
|
|
return
|
|
}
|
|
if requestContext.Err() != nil {
|
|
writeCredentialUnavailableError(w, r, provider, selectedCredential, selection, "credential became unavailable while processing the request")
|
|
return
|
|
}
|
|
writeJSONError(w, r, http.StatusBadGateway, "api_error", err.Error())
|
|
return
|
|
}
|
|
requestContext.releaseCredentialInterrupt()
|
|
|
|
// Transparent 429 retry
|
|
for response.StatusCode == http.StatusTooManyRequests {
|
|
resetAt := parseOCMRateLimitResetFromHeaders(response.Header)
|
|
nextCredential := provider.onRateLimited(sessionID, selectedCredential, resetAt, selection)
|
|
needsBodyReplay := r.Method != http.MethodGet && r.Method != http.MethodHead && r.Method != http.MethodDelete
|
|
selectedCredential.updateStateFromHeaders(response.Header)
|
|
if (needsBodyReplay && bodyBytes == nil) || nextCredential == nil {
|
|
response.Body.Close()
|
|
writeCredentialUnavailableError(w, r, provider, selectedCredential, selection, "all credentials rate-limited")
|
|
return
|
|
}
|
|
response.Body.Close()
|
|
s.logger.InfoContext(ctx, "retrying with credential ", nextCredential.tagName(), " after 429 from ", selectedCredential.tagName())
|
|
requestContext.cancelRequest()
|
|
requestContext = nextCredential.wrapRequestContext(ctx)
|
|
{
|
|
currentRequestContext := requestContext
|
|
requestContext.addInterruptLink(provider.linkProviderInterrupt(nextCredential, selection, func() {
|
|
currentRequestContext.cancelOnce.Do(currentRequestContext.cancelFunc)
|
|
}))
|
|
}
|
|
retryRequest, buildErr := nextCredential.buildProxyRequest(requestContext, r, bodyBytes, s.httpHeaders)
|
|
if buildErr != nil {
|
|
s.logger.ErrorContext(ctx, "retry request: ", buildErr)
|
|
writeJSONError(w, r, http.StatusBadGateway, "api_error", buildErr.Error())
|
|
return
|
|
}
|
|
retryResponse, retryErr := nextCredential.httpClient().Do(retryRequest)
|
|
if retryErr != nil {
|
|
if r.Context().Err() != nil {
|
|
return
|
|
}
|
|
if requestContext.Err() != nil {
|
|
writeCredentialUnavailableError(w, r, provider, nextCredential, selection, "credential became unavailable while retrying the request")
|
|
return
|
|
}
|
|
s.logger.ErrorContext(ctx, "retry request: ", retryErr)
|
|
writeJSONError(w, r, http.StatusBadGateway, "api_error", retryErr.Error())
|
|
return
|
|
}
|
|
requestContext.releaseCredentialInterrupt()
|
|
response = retryResponse
|
|
selectedCredential = nextCredential
|
|
}
|
|
defer response.Body.Close()
|
|
|
|
selectedCredential.updateStateFromHeaders(response.Header)
|
|
|
|
if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusTooManyRequests {
|
|
body, _ := io.ReadAll(response.Body)
|
|
s.logger.ErrorContext(ctx, "upstream error from ", selectedCredential.tagName(), ": status ", response.StatusCode, " ", string(body))
|
|
writeJSONError(w, r, http.StatusInternalServerError, "api_error",
|
|
"proxy request (status "+strconv.Itoa(response.StatusCode)+"): "+string(body))
|
|
return
|
|
}
|
|
|
|
s.rewriteResponseHeaders(response.Header, provider, userConfig)
|
|
|
|
for key, values := range response.Header {
|
|
if !isHopByHopHeader(key) && !isReverseProxyHeader(key) {
|
|
w.Header()[key] = values
|
|
}
|
|
}
|
|
w.WriteHeader(response.StatusCode)
|
|
|
|
usageTracker := selectedCredential.usageTrackerOrNil()
|
|
if usageTracker != nil && response.StatusCode == http.StatusOK &&
|
|
(path == "/v1/chat/completions" || strings.HasPrefix(path, "/v1/responses")) {
|
|
s.handleResponseWithTracking(ctx, w, response, usageTracker, path, requestModel, username)
|
|
} else {
|
|
mediaType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type"))
|
|
if err == nil && mediaType != "text/event-stream" {
|
|
_, _ = io.Copy(w, response.Body)
|
|
return
|
|
}
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok {
|
|
s.logger.ErrorContext(ctx, "streaming not supported")
|
|
return
|
|
}
|
|
buffer := make([]byte, buf.BufferSize)
|
|
for {
|
|
n, err := response.Body.Read(buffer)
|
|
if n > 0 {
|
|
_, writeError := w.Write(buffer[:n])
|
|
if writeError != nil {
|
|
if E.IsClosedOrCanceled(writeError) {
|
|
return
|
|
}
|
|
s.logger.ErrorContext(ctx, "write streaming response: ", writeError)
|
|
return
|
|
}
|
|
flusher.Flush()
|
|
}
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Service) handleResponseWithTracking(ctx context.Context, writer http.ResponseWriter, response *http.Response, usageTracker *AggregatedUsage, path string, requestModel string, username string) {
|
|
isChatCompletions := path == "/v1/chat/completions"
|
|
weeklyCycleHint := extractWeeklyCycleHint(response.Header)
|
|
mediaType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type"))
|
|
isStreaming := err == nil && mediaType == "text/event-stream"
|
|
if !isStreaming && !isChatCompletions && response.Header.Get("Content-Type") == "" {
|
|
isStreaming = true
|
|
}
|
|
if !isStreaming {
|
|
bodyBytes, err := io.ReadAll(response.Body)
|
|
if err != nil {
|
|
s.logger.ErrorContext(ctx, "read response body: ", err)
|
|
return
|
|
}
|
|
|
|
var responseModel, serviceTier string
|
|
var inputTokens, outputTokens, cachedTokens int64
|
|
|
|
if isChatCompletions {
|
|
var chatCompletion openai.ChatCompletion
|
|
if json.Unmarshal(bodyBytes, &chatCompletion) == nil {
|
|
responseModel = chatCompletion.Model
|
|
serviceTier = string(chatCompletion.ServiceTier)
|
|
inputTokens = chatCompletion.Usage.PromptTokens
|
|
outputTokens = chatCompletion.Usage.CompletionTokens
|
|
cachedTokens = chatCompletion.Usage.PromptTokensDetails.CachedTokens
|
|
}
|
|
} else {
|
|
var responsesResponse responses.Response
|
|
if json.Unmarshal(bodyBytes, &responsesResponse) == nil {
|
|
responseModel = string(responsesResponse.Model)
|
|
serviceTier = string(responsesResponse.ServiceTier)
|
|
inputTokens = responsesResponse.Usage.InputTokens
|
|
outputTokens = responsesResponse.Usage.OutputTokens
|
|
cachedTokens = responsesResponse.Usage.InputTokensDetails.CachedTokens
|
|
}
|
|
}
|
|
|
|
if inputTokens > 0 || outputTokens > 0 {
|
|
if responseModel == "" {
|
|
responseModel = requestModel
|
|
}
|
|
if responseModel != "" {
|
|
contextWindow := detectContextWindow(responseModel, serviceTier, inputTokens)
|
|
usageTracker.AddUsageWithCycleHint(
|
|
responseModel,
|
|
contextWindow,
|
|
inputTokens,
|
|
outputTokens,
|
|
cachedTokens,
|
|
serviceTier,
|
|
username,
|
|
time.Now(),
|
|
weeklyCycleHint,
|
|
)
|
|
}
|
|
}
|
|
|
|
_, _ = writer.Write(bodyBytes)
|
|
return
|
|
}
|
|
|
|
flusher, ok := writer.(http.Flusher)
|
|
if !ok {
|
|
s.logger.ErrorContext(ctx, "streaming not supported")
|
|
return
|
|
}
|
|
|
|
var inputTokens, outputTokens, cachedTokens int64
|
|
var responseModel, serviceTier string
|
|
buffer := make([]byte, buf.BufferSize)
|
|
var leftover []byte
|
|
|
|
for {
|
|
n, err := response.Body.Read(buffer)
|
|
if n > 0 {
|
|
data := append(leftover, buffer[:n]...)
|
|
lines := bytes.Split(data, []byte("\n"))
|
|
|
|
if err == nil {
|
|
leftover = lines[len(lines)-1]
|
|
lines = lines[:len(lines)-1]
|
|
} else {
|
|
leftover = nil
|
|
}
|
|
|
|
for _, line := range lines {
|
|
line = bytes.TrimSpace(line)
|
|
if len(line) == 0 {
|
|
continue
|
|
}
|
|
|
|
if bytes.HasPrefix(line, []byte("data: ")) {
|
|
eventData := bytes.TrimPrefix(line, []byte("data: "))
|
|
if bytes.Equal(eventData, []byte("[DONE]")) {
|
|
continue
|
|
}
|
|
|
|
if isChatCompletions {
|
|
var chatChunk openai.ChatCompletionChunk
|
|
if json.Unmarshal(eventData, &chatChunk) == nil {
|
|
if chatChunk.Model != "" {
|
|
responseModel = chatChunk.Model
|
|
}
|
|
if chatChunk.ServiceTier != "" {
|
|
serviceTier = string(chatChunk.ServiceTier)
|
|
}
|
|
if chatChunk.Usage.PromptTokens > 0 {
|
|
inputTokens = chatChunk.Usage.PromptTokens
|
|
cachedTokens = chatChunk.Usage.PromptTokensDetails.CachedTokens
|
|
}
|
|
if chatChunk.Usage.CompletionTokens > 0 {
|
|
outputTokens = chatChunk.Usage.CompletionTokens
|
|
}
|
|
}
|
|
} else {
|
|
var streamEvent responses.ResponseStreamEventUnion
|
|
if json.Unmarshal(eventData, &streamEvent) == nil {
|
|
if streamEvent.Type == "response.completed" {
|
|
completedEvent := streamEvent.AsResponseCompleted()
|
|
if string(completedEvent.Response.Model) != "" {
|
|
responseModel = string(completedEvent.Response.Model)
|
|
}
|
|
if completedEvent.Response.ServiceTier != "" {
|
|
serviceTier = string(completedEvent.Response.ServiceTier)
|
|
}
|
|
if completedEvent.Response.Usage.InputTokens > 0 {
|
|
inputTokens = completedEvent.Response.Usage.InputTokens
|
|
cachedTokens = completedEvent.Response.Usage.InputTokensDetails.CachedTokens
|
|
}
|
|
if completedEvent.Response.Usage.OutputTokens > 0 {
|
|
outputTokens = completedEvent.Response.Usage.OutputTokens
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
_, writeError := writer.Write(buffer[:n])
|
|
if writeError != nil {
|
|
if E.IsClosedOrCanceled(writeError) {
|
|
return
|
|
}
|
|
s.logger.ErrorContext(ctx, "write streaming response: ", writeError)
|
|
return
|
|
}
|
|
flusher.Flush()
|
|
}
|
|
|
|
if err != nil {
|
|
if responseModel == "" {
|
|
responseModel = requestModel
|
|
}
|
|
|
|
if inputTokens > 0 || outputTokens > 0 {
|
|
if responseModel != "" {
|
|
contextWindow := detectContextWindow(responseModel, serviceTier, inputTokens)
|
|
usageTracker.AddUsageWithCycleHint(
|
|
responseModel,
|
|
contextWindow,
|
|
inputTokens,
|
|
outputTokens,
|
|
cachedTokens,
|
|
serviceTier,
|
|
username,
|
|
time.Now(),
|
|
weeklyCycleHint,
|
|
)
|
|
}
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|