feat: 代理调试版本

internal/channel/base_channel.go

@@ -3,16 +3,11 @@ package channel
 import (
 	"bytes"
 	"encoding/json"
-	"fmt"
 	"gpt-load/internal/models"
-	"io"
 	"net/http"
-	"net/http/httputil"
 	"net/url"
-	"strings"
 	"sync"
 
-	"github.com/gin-gonic/gin"
 	"github.com/sirupsen/logrus"
 	"gorm.io/datatypes"
 )
@@ -29,15 +24,13 @@ type BaseChannel struct {
 	Name            string
 	Upstreams       []UpstreamInfo
 	HTTPClient      *http.Client
+	StreamClient    *http.Client
 	TestModel       string
 	upstreamLock    sync.Mutex
 	groupUpstreams  datatypes.JSON
 	groupConfig     datatypes.JSONMap
 }
 
-// RequestModifier is a function that can modify the request before it's sent.
-type RequestModifier func(req *http.Request, key *models.APIKey)
-
 // getUpstreamURL selects an upstream URL using a smooth weighted round-robin algorithm.
 func (b *BaseChannel) getUpstreamURL() *url.URL {
 	b.upstreamLock.Lock()
@@ -99,100 +92,12 @@ func (b *BaseChannel) IsConfigStale(group *models.Group) bool {
 	return false
 }
 
-// ProcessRequest handles the common logic of processing and forwarding a request.
-func (b *BaseChannel) ProcessRequest(c *gin.Context, apiKey *models.APIKey, modifier RequestModifier, ch ChannelProxy) error {
-	upstreamURL := b.getUpstreamURL()
-	if upstreamURL == nil {
-		return fmt.Errorf("no upstream URL configured for channel %s", b.Name)
-	}
-
-	director := func(req *http.Request) {
-		req.URL.Scheme = upstreamURL.Scheme
-		req.URL.Host = upstreamURL.Host
-		req.URL.Path = singleJoiningSlash(upstreamURL.Path, req.URL.Path)
-		req.Host = upstreamURL.Host
-
-		// Apply the channel-specific modifications
-		if modifier != nil {
-			modifier(req, apiKey)
-		}
-
-		// Remove headers that should not be forwarded
-		req.Header.Del("Cookie")
-		req.Header.Del("X-Real-Ip")
-		req.Header.Del("X-Forwarded-For")
-	}
-
-	errorHandler := func(rw http.ResponseWriter, req *http.Request, err error) {
-		logrus.WithFields(logrus.Fields{
-			"channel": b.Name,
-			"key_id":  apiKey.ID,
-			"error":   err,
-		}).Error("HTTP proxy error")
-		rw.WriteHeader(http.StatusBadGateway)
-	}
-
-	proxy := &httputil.ReverseProxy{
-		Director:     director,
-		ErrorHandler: errorHandler,
-		Transport:    b.HTTPClient.Transport,
-	}
-
-	// Check if the client request is for a streaming endpoint
-	if ch.IsStreamingRequest(c) {
-		return b.handleStreaming(c, proxy)
-	}
-
-	proxy.ServeHTTP(c.Writer, c.Request)
-	return nil
+// GetHTTPClient returns the client for standard requests.
+func (b *BaseChannel) GetHTTPClient() *http.Client {
+	return b.HTTPClient
 }
 
-func (b *BaseChannel) handleStreaming(c *gin.Context, proxy *httputil.ReverseProxy) error {
-	var wg sync.WaitGroup
-	wg.Add(1)
-
-	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("Access-Control-Allow-Origin", "*")
-
-	// Use a pipe to avoid buffering the entire response
-	pr, pw := io.Pipe()
-	defer pr.Close()
-
-	req := c.Request.Clone(c.Request.Context())
-	req.Body = pr
-
-	// Start the proxy in a goroutine
-	go func() {
-		defer wg.Done()
-		defer pw.Close()
-		proxy.ServeHTTP(c.Writer, req)
-	}()
-
-	// Copy the original request body to the pipe writer
-	_, err := io.Copy(pw, c.Request.Body)
-	if err != nil {
-		logrus.Errorf("Error copying request body to pipe: %v", err)
-		wg.Wait() // Wait for the goroutine to finish even if copy fails
-		return err
-	}
-
-	// Wait for the proxy to finish
-	wg.Wait()
-
-	return nil
-}
-
-// singleJoiningSlash joins two URL paths with a single slash.
-func singleJoiningSlash(a, b string) string {
-	aslash := strings.HasSuffix(a, "/")
-	bslash := strings.HasPrefix(b, "/")
-	switch {
-	case aslash && bslash:
-		return a + b[1:]
-	case !aslash && !bslash:
-		return a + "/" + b
-	}
-	return a + b
+// GetStreamClient returns the client for streaming requests.
+func (b *BaseChannel) GetStreamClient() *http.Client {
+	return b.StreamClient
 }

internal/channel/channel.go

@@ -3,22 +3,35 @@ package channel
 import (
 	"context"
 	"gpt-load/internal/models"
+	"net/http"
+	"net/url"
 
 	"github.com/gin-gonic/gin"
 )
 
 // ChannelProxy defines the interface for different API channel proxies.
+// It's responsible for channel-specific logic like URL building and request modification.
 type ChannelProxy interface {
-	// Handle takes a context, an API key, and the original request,
-	// then forwards the request to the upstream service.
-	Handle(c *gin.Context, apiKey *models.APIKey, group *models.Group) error
+	// BuildUpstreamURL constructs the target URL for the upstream service.
+	BuildUpstreamURL(originalURL *url.URL, group *models.Group) (string, error)
+
+	// ModifyRequest allows the channel to add specific headers or modify the request
+	// before it's sent to the upstream service.
+	ModifyRequest(req *http.Request, apiKey *models.APIKey, group *models.Group)
+
+	// IsStreamRequest checks if the request is for a streaming response,
+	// now using the cached request body to avoid re-reading the stream.
+	IsStreamRequest(c *gin.Context, bodyBytes []byte) bool
 
 	// ValidateKey checks if the given API key is valid.
 	ValidateKey(ctx context.Context, key string) (bool, error)
 
-	// IsStreamingRequest checks if the request is for a streaming response.
-	IsStreamingRequest(c *gin.Context) bool
-
 	// IsConfigStale checks if the channel's configuration is stale compared to the provided group.
 	IsConfigStale(group *models.Group) bool
+
+	// GetHTTPClient returns the client for standard requests.
+	GetHTTPClient() *http.Client
+
+	// GetStreamClient returns the client for streaming requests.
+	GetStreamClient() *http.Client
 }

internal/channel/factory.go

@@ -4,8 +4,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"gpt-load/internal/config"
+	"gpt-load/internal/httpclient"
 	"gpt-load/internal/models"
-	"net/http"
 	"net/url"
 	"sync"
 	"time"
@@ -42,14 +42,16 @@ func GetChannels() []string {
 // Factory is responsible for creating channel proxies.
 type Factory struct {
 	settingsManager *config.SystemSettingsManager
+	clientManager   *httpclient.HTTPClientManager
 	channelCache    map[uint]ChannelProxy
 	cacheLock       sync.Mutex
 }
 
 // NewFactory creates a new channel factory.
-func NewFactory(settingsManager *config.SystemSettingsManager) *Factory {
+func NewFactory(settingsManager *config.SystemSettingsManager, clientManager *httpclient.HTTPClientManager) *Factory {
 	return &Factory{
 		settingsManager: settingsManager,
+		clientManager:   clientManager,
 		channelCache:    make(map[uint]ChannelProxy),
 	}
 }
@@ -109,21 +111,39 @@ func (f *Factory) newBaseChannel(name string, group *models.Group) (*BaseChannel
 		upstreamInfos = append(upstreamInfos, UpstreamInfo{URL: u, Weight: weight})
 	}
 
-	// Get effective settings by merging system and group configs
-	effectiveSettings := f.settingsManager.GetEffectiveConfig(group.Config)
-
-	// Configure the HTTP client with the effective timeouts
-	httpClient := &http.Client{
-		Transport: &http.Transport{
-			IdleConnTimeout: time.Duration(effectiveSettings.IdleConnTimeout) * time.Second,
-		},
-		Timeout: time.Duration(effectiveSettings.RequestTimeout) * time.Second,
+	// Base configuration for regular requests, derived from the group's effective settings.
+	clientConfig := &httpclient.Config{
+		ConnectTimeout:        time.Duration(group.EffectiveConfig.ConnectTimeout) * time.Second,
+		RequestTimeout:        time.Duration(group.EffectiveConfig.RequestTimeout) * time.Second,
+		IdleConnTimeout:       time.Duration(group.EffectiveConfig.IdleConnTimeout) * time.Second,
+		MaxIdleConns:          group.EffectiveConfig.MaxIdleConns,
+		MaxIdleConnsPerHost:   group.EffectiveConfig.MaxIdleConnsPerHost,
+		ResponseHeaderTimeout: time.Duration(group.EffectiveConfig.ResponseHeaderTimeout) * time.Second,
+		DisableCompression:    group.EffectiveConfig.DisableCompression,
+		WriteBufferSize:       32 * 1024, // Use a reasonable default buffer size for regular requests
+		ReadBufferSize:        32 * 1024,
 	}
 
+	// Create a dedicated configuration for streaming requests.
+	// This configuration is optimized for low-latency, long-running connections.
+	streamConfig := *clientConfig
+	streamConfig.RequestTimeout = 0      // No overall timeout for the entire request.
+	streamConfig.DisableCompression = true // Always disable compression for streaming to reduce latency.
+	streamConfig.WriteBufferSize = 0     // Disable buffering for real-time data transfer.
+	streamConfig.ReadBufferSize = 0
+	// For stream-specific connection pool, we can use a simple heuristic like doubling the regular one.
+	streamConfig.MaxIdleConns = group.EffectiveConfig.MaxIdleConns * 2
+	streamConfig.MaxIdleConnsPerHost = group.EffectiveConfig.MaxIdleConnsPerHost * 2
+
+	// Get both clients from the manager using their respective configurations.
+	httpClient := f.clientManager.GetClient(clientConfig)
+	streamClient := f.clientManager.GetClient(&streamConfig)
+
 	return &BaseChannel{
 		Name:           name,
 		Upstreams:      upstreamInfos,
 		HTTPClient:     httpClient,
+		StreamClient:   streamClient,
 		TestModel:      group.TestModel,
 		groupUpstreams: group.Upstreams,
 		groupConfig:    group.Config,

internal/channel/gemini_channel.go

@@ -9,6 +9,7 @@ import (
 	"gpt-load/internal/models"
 	"io"
 	"net/http"
+	"net/url"
 	"strings"
 
 	"github.com/gin-gonic/gin"
@@ -33,13 +34,36 @@ func newGeminiChannel(f *Factory, group *models.Group) (ChannelProxy, error) {
 	}, nil
 }
 
-func (ch *GeminiChannel) Handle(c *gin.Context, apiKey *models.APIKey, group *models.Group) error {
-	modifier := func(req *http.Request, key *models.APIKey) {
-		q := req.URL.Query()
-		q.Set("key", key.KeyValue)
-		req.URL.RawQuery = q.Encode()
+// BuildUpstreamURL constructs the target URL for the Gemini service.
+func (ch *GeminiChannel) BuildUpstreamURL(originalURL *url.URL, group *models.Group) (string, error) {
+	base := ch.getUpstreamURL()
+	if base == nil {
+		// Fallback to default Gemini URL
+		base, _ = url.Parse("https://generativelanguage.googleapis.com")
 	}
-	return ch.ProcessRequest(c, apiKey, modifier, ch)
+
+	finalURL := *base
+	// The originalURL.Path contains the full path, e.g., "/proxy/gemini/v1beta/models/gemini-pro:generateContent".
+	// We need to strip the proxy prefix to get the correct upstream path.
+	proxyPrefix := "/proxy/" + group.Name
+	if strings.HasPrefix(originalURL.Path, proxyPrefix) {
+		finalURL.Path = strings.TrimPrefix(originalURL.Path, proxyPrefix)
+	} else {
+		// Fallback for safety.
+		finalURL.Path = originalURL.Path
+	}
+
+	// The API key will be added to RawQuery in ModifyRequest.
+	finalURL.RawQuery = originalURL.RawQuery
+
+	return finalURL.String(), nil
+}
+
+// ModifyRequest adds the API key as a query parameter for Gemini requests.
+func (ch *GeminiChannel) ModifyRequest(req *http.Request, apiKey *models.APIKey, group *models.Group) {
+	q := req.URL.Query()
+	q.Set("key", apiKey.KeyValue)
+	req.URL.RawQuery = q.Encode()
 }
 
 // ValidateKey checks if the given API key is valid by making a generateContent request.
@@ -95,12 +119,21 @@ func (ch *GeminiChannel) ValidateKey(ctx context.Context, key string) (bool, err
 	return false, fmt.Errorf("[status %d] %s", resp.StatusCode, parsedError)
 }
 
-// IsStreamingRequest checks if the request is for a streaming response.
-func (ch *GeminiChannel) IsStreamingRequest(c *gin.Context) bool {
-	// For Gemini, streaming is indicated by the path containing streaming keywords
+// IsStreamRequest checks if the request is for a streaming response.
+// For Gemini, this is primarily determined by the URL path.
+func (ch *GeminiChannel) IsStreamRequest(c *gin.Context, bodyBytes []byte) bool {
 	path := c.Request.URL.Path
-	return strings.Contains(path, ":streamGenerateContent") ||
-		strings.Contains(path, "streamGenerateContent") ||
-		strings.Contains(path, ":stream") ||
-		strings.Contains(path, "/stream")
+	if strings.HasSuffix(path, ":streamGenerateContent") {
+		return true
+	}
+
+	// Also check for standard streaming indicators as a fallback.
+	if strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
+		return true
+	}
+	if c.Query("stream") == "true" {
+		return true
+	}
+
+	return false
 }

internal/channel/openai_channel.go

@@ -9,9 +9,10 @@ import (
 	"gpt-load/internal/models"
 	"io"
 	"net/http"
+	"net/url"
+	"strings"
 
 	"github.com/gin-gonic/gin"
-	"github.com/gin-gonic/gin/binding"
 )
 
 func init() {
@@ -33,11 +34,38 @@ func newOpenAIChannel(f *Factory, group *models.Group) (ChannelProxy, error) {
 	}, nil
 }
 
-func (ch *OpenAIChannel) Handle(c *gin.Context, apiKey *models.APIKey, group *models.Group) error {
-	modifier := func(req *http.Request, key *models.APIKey) {
-		req.Header.Set("Authorization", "Bearer "+key.KeyValue)
+// BuildUpstreamURL constructs the target URL for the OpenAI service.
+func (ch *OpenAIChannel) BuildUpstreamURL(originalURL *url.URL, group *models.Group) (string, error) {
+	// Use the weighted round-robin selection from the base channel.
+	// This method already handles parsing the group's Upstreams JSON.
+	base := ch.getUpstreamURL()
+	if base == nil {
+		// If no upstreams are configured in the group, fallback to a default.
+		// This can be considered an error or a feature depending on requirements.
+		// For now, we'll use the official OpenAI URL as a last resort.
+		base, _ = url.Parse("https://api.openai.com")
 	}
-	return ch.ProcessRequest(c, apiKey, modifier, ch)
+
+	// It's crucial to create a copy to avoid modifying the cached URL object in BaseChannel.
+	finalURL := *base
+	// The originalURL.Path contains the full path, e.g., "/proxy/openai/v1/chat/completions".
+	// We need to strip the proxy prefix to get the correct upstream path.
+	proxyPrefix := "/proxy/" + group.Name
+	if strings.HasPrefix(originalURL.Path, proxyPrefix) {
+		finalURL.Path = strings.TrimPrefix(originalURL.Path, proxyPrefix)
+	} else {
+		// Fallback for safety, though this case should ideally not be hit.
+		finalURL.Path = originalURL.Path
+	}
+
+	finalURL.RawQuery = originalURL.RawQuery
+
+	return finalURL.String(), nil
+}
+
+// ModifyRequest sets the Authorization header for the OpenAI service.
+func (ch *OpenAIChannel) ModifyRequest(req *http.Request, apiKey *models.APIKey, group *models.Group) {
+	req.Header.Set("Authorization", "Bearer "+apiKey.KeyValue)
 }
 
 // ValidateKey checks if the given API key is valid by making a chat completion request.
@@ -92,16 +120,23 @@ func (ch *OpenAIChannel) ValidateKey(ctx context.Context, key string) (bool, err
 	return false, fmt.Errorf("[status %d] %s", resp.StatusCode, parsedError)
 }
 
-// IsStreamingRequest checks if the request is for a streaming response.
-func (ch *OpenAIChannel) IsStreamingRequest(c *gin.Context) bool {
-	// For OpenAI, streaming is indicated by a "stream": true field in the JSON body.
-	// We use ShouldBindBodyWith to check the body without consuming it, so it can be read again by the proxy.
+// IsStreamRequest checks if the request is for a streaming response using the pre-read body.
+func (ch *OpenAIChannel) IsStreamRequest(c *gin.Context, bodyBytes []byte) bool {
+	if strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
+		return true
+	}
+
+	if c.Query("stream") == "true" {
+		return true
+	}
+
 	type streamPayload struct {
 		Stream bool `json:"stream"`
 	}
 	var p streamPayload
-	if err := c.ShouldBindBodyWith(&p, binding.JSON); err == nil {
+	if err := json.Unmarshal(bodyBytes, &p); err == nil {
 		return p.Stream
 	}
+
 	return false
 }