// Package proxy provides high-performance OpenAI multi-key proxy server package proxy import ( "bytes" "context" "encoding/json" "errors" "fmt" "io" "net/http" "strconv" "time" "gpt-load/internal/channel" "gpt-load/internal/config" app_errors "gpt-load/internal/errors" "gpt-load/internal/keypool" "gpt-load/internal/models" "gpt-load/internal/response" "gpt-load/internal/services" "gpt-load/internal/types" "gpt-load/internal/utils" "github.com/gin-gonic/gin" "github.com/sirupsen/logrus" ) // ProxyServer represents the proxy server type ProxyServer struct { keyProvider *keypool.KeyProvider groupManager *services.GroupManager settingsManager *config.SystemSettingsManager channelFactory *channel.Factory requestLogService *services.RequestLogService } // NewProxyServer creates a new proxy server func NewProxyServer( keyProvider *keypool.KeyProvider, groupManager *services.GroupManager, settingsManager *config.SystemSettingsManager, channelFactory *channel.Factory, requestLogService *services.RequestLogService, ) (*ProxyServer, error) { return &ProxyServer{ keyProvider: keyProvider, groupManager: groupManager, settingsManager: settingsManager, channelFactory: channelFactory, requestLogService: requestLogService, }, nil } // HandleProxy is the main entry point for proxy requests, refactored based on the stable .bak logic. func (ps *ProxyServer) HandleProxy(c *gin.Context) { startTime := time.Now() groupName := c.Param("group_name") group, err := ps.groupManager.GetGroupByName(groupName) if err != nil { response.Error(c, app_errors.ParseDBError(err)) return } channelHandler, err := ps.channelFactory.GetChannel(group) if err != nil { response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, fmt.Sprintf("Failed to get channel for group '%s': %v", groupName, err))) return } bodyBytes, err := io.ReadAll(c.Request.Body) if err != nil { logrus.Errorf("Failed to read request body: %v", err) response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Failed to read request body")) return } c.Request.Body.Close() finalBodyBytes, err := ps.applyParamOverrides(bodyBytes, group) if err != nil { response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, fmt.Sprintf("Failed to apply parameter overrides: %v", err))) return } isStream := channelHandler.IsStreamRequest(c, bodyBytes) ps.executeRequestWithRetry(c, channelHandler, group, finalBodyBytes, isStream, startTime, 0, nil) } // executeRequestWithRetry is the core recursive function for handling requests and retries. func (ps *ProxyServer) executeRequestWithRetry( c *gin.Context, channelHandler channel.ChannelProxy, group *models.Group, bodyBytes []byte, isStream bool, startTime time.Time, retryCount int, retryErrors []types.RetryError, ) { cfg := group.EffectiveConfig if retryCount > cfg.MaxRetries { if len(retryErrors) > 0 { lastError := retryErrors[len(retryErrors)-1] var errorJSON map[string]any if err := json.Unmarshal([]byte(lastError.ErrorMessage), &errorJSON); err == nil { c.JSON(lastError.StatusCode, errorJSON) } else { response.Error(c, app_errors.NewAPIErrorWithUpstream(lastError.StatusCode, "UPSTREAM_ERROR", lastError.ErrorMessage)) } logrus.Debugf("Max retries exceeded for group %s after %d attempts. Parsed Error: %s", group.Name, retryCount, lastError.ErrorMessage) keyID, _ := strconv.ParseUint(lastError.KeyID, 10, 64) ps.logRequest(c, group, uint(keyID), startTime, lastError.StatusCode, retryCount, errors.New(lastError.ErrorMessage)) } else { response.Error(c, app_errors.ErrMaxRetriesExceeded) logrus.Debugf("Max retries exceeded for group %s after %d attempts.", group.Name, retryCount) ps.logRequest(c, group, 0, startTime, http.StatusServiceUnavailable, retryCount, app_errors.ErrMaxRetriesExceeded) } return } apiKey, err := ps.keyProvider.SelectKey(group.ID) if err != nil { logrus.Errorf("Failed to select a key for group %s on attempt %d: %v", group.Name, retryCount+1, err) response.Error(c, app_errors.NewAPIError(app_errors.ErrNoKeysAvailable, err.Error())) ps.logRequest(c, group, 0, startTime, http.StatusServiceUnavailable, retryCount, err) return } upstreamURL, err := channelHandler.BuildUpstreamURL(c.Request.URL, group) if err != nil { response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, fmt.Sprintf("Failed to build upstream URL: %v", err))) return } var ctx context.Context var cancel context.CancelFunc if isStream { ctx, cancel = context.WithCancel(c.Request.Context()) } else { timeout := time.Duration(cfg.RequestTimeout) * time.Second ctx, cancel = context.WithTimeout(c.Request.Context(), timeout) } defer cancel() req, err := http.NewRequestWithContext(ctx, c.Request.Method, upstreamURL, bytes.NewReader(bodyBytes)) if err != nil { logrus.Errorf("Failed to create upstream request: %v", err) response.Error(c, app_errors.ErrInternalServer) return } req.ContentLength = int64(len(bodyBytes)) req.Header = c.Request.Header.Clone() channelHandler.ModifyRequest(req, apiKey, group) var client *http.Client if isStream { client = channelHandler.GetStreamClient() req.Header.Set("X-Accel-Buffering", "no") } else { client = channelHandler.GetHTTPClient() } resp, err := client.Do(req) if resp != nil { defer resp.Body.Close() } // Unified error handling for retries. if err != nil || (resp != nil && resp.StatusCode >= 400) { if err != nil && app_errors.IsIgnorableError(err) { logrus.Debugf("Client-side ignorable error for key %s, aborting retries: %v", utils.MaskAPIKey(apiKey.KeyValue), err) ps.logRequest(c, group, apiKey.ID, startTime, 499, retryCount+1, err) return } ps.keyProvider.UpdateStatus(apiKey, group, false) var statusCode int var errorMessage string var parsedError string if err != nil { statusCode = 0 errorMessage = err.Error() logrus.Debugf("Request failed (attempt %d/%d) for key %s: %v", retryCount+1, cfg.MaxRetries, utils.MaskAPIKey(apiKey.KeyValue), err) } else { // HTTP-level error (status >= 400) statusCode = resp.StatusCode errorBody, readErr := io.ReadAll(resp.Body) if readErr != nil { logrus.Errorf("Failed to read error body: %v", readErr) errorBody = []byte("Failed to read error body") } errorBody = handleGzipCompression(resp, errorBody) errorMessage = string(errorBody) parsedError = app_errors.ParseUpstreamError(errorBody) logrus.Debugf("Request failed with status %d (attempt %d/%d) for key %s. Parsed Error: %s", statusCode, retryCount+1, cfg.MaxRetries, utils.MaskAPIKey(apiKey.KeyValue), parsedError) } newRetryErrors := append(retryErrors, types.RetryError{ StatusCode: statusCode, ErrorMessage: errorMessage, KeyID: fmt.Sprintf("%d", apiKey.ID), Attempt: retryCount + 1, }) ps.executeRequestWithRetry(c, channelHandler, group, bodyBytes, isStream, startTime, retryCount+1, newRetryErrors) return } ps.keyProvider.UpdateStatus(apiKey, group, true) logrus.Debugf("Request for group %s succeeded on attempt %d with key %s", group.Name, retryCount+1, utils.MaskAPIKey(apiKey.KeyValue)) ps.logRequest(c, group, apiKey.ID, startTime, resp.StatusCode, retryCount+1, nil) for key, values := range resp.Header { for _, value := range values { c.Header(key, value) } } c.Status(resp.StatusCode) if isStream { ps.handleStreamingResponse(c, resp) } else { ps.handleNormalResponse(c, resp) } } // logRequest is a helper function to create and record a request log. func (ps *ProxyServer) logRequest( c *gin.Context, group *models.Group, keyID uint, startTime time.Time, statusCode int, retries int, finalError error, ) { if ps.requestLogService == nil { return } duration := time.Since(startTime).Milliseconds() logEntry := &models.RequestLog{ GroupID: group.ID, KeyID: keyID, IsSuccess: finalError == nil && statusCode < 400, SourceIP: c.ClientIP(), StatusCode: statusCode, RequestPath: c.Request.URL.String(), Duration: duration, UserAgent: c.Request.UserAgent(), Retries: retries, } if finalError != nil { logEntry.ErrorMessage = finalError.Error() } if err := ps.requestLogService.Record(logEntry); err != nil { logrus.Errorf("Failed to record request log: %v", err) } }