mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-16 15:43:13 +08:00
399 lines
11 KiB
Go
399 lines
11 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/url"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
openaiwsv2 "github.com/Wei-Shaw/sub2api/internal/service/openai_ws_v2"
|
||
coderws "github.com/coder/websocket"
|
||
"github.com/coder/websocket/wsjson"
|
||
)
|
||
|
||
const openAIWSMessageReadLimitBytes int64 = 16 * 1024 * 1024
|
||
const (
|
||
openAIWSProxyTransportMaxIdleConns = 128
|
||
openAIWSProxyTransportMaxIdleConnsPerHost = 64
|
||
openAIWSProxyTransportIdleConnTimeout = 90 * time.Second
|
||
openAIWSProxyClientCacheMaxEntries = 256
|
||
openAIWSProxyClientCacheIdleTTL = 15 * time.Minute
|
||
)
|
||
|
||
type OpenAIWSTransportMetricsSnapshot struct {
|
||
ProxyClientCacheHits int64 `json:"proxy_client_cache_hits"`
|
||
ProxyClientCacheMisses int64 `json:"proxy_client_cache_misses"`
|
||
TransportReuseRatio float64 `json:"transport_reuse_ratio"`
|
||
}
|
||
|
||
// openAIWSClientConn 抽象 WS 客户端连接，便于替换底层实现。
|
||
type openAIWSClientConn interface {
|
||
WriteJSON(ctx context.Context, value any) error
|
||
ReadMessage(ctx context.Context) ([]byte, error)
|
||
Ping(ctx context.Context) error
|
||
Close() error
|
||
}
|
||
|
||
// openAIWSIdlePingCapable is intentionally separate from openAIWSClientConn.
|
||
// A pool probe happens while no goroutine is reading an idle connection, which
|
||
// is not safe for every WebSocket implementation.
|
||
type openAIWSIdlePingCapable interface {
|
||
SupportsIdlePingWithoutReader() bool
|
||
}
|
||
|
||
// openAIWSReaderLoopCapable 声明该实现的控制帧只在阻塞读期间被消费，
|
||
// 连接池需为其常驻一个读循环，否则空闲连接无法应答上游 ping。
|
||
type openAIWSReaderLoopCapable interface {
|
||
RequiresReaderLoop() bool
|
||
}
|
||
|
||
// openAIWSUpstreamPingCounter 报告连接收到过多少个上游 ping 帧，用于核对读循环是否在应答保活。
|
||
type openAIWSUpstreamPingCounter interface {
|
||
UpstreamPingCount() int64
|
||
}
|
||
|
||
// openAIWSForceCloser 不做关闭握手直接切断连接，用于对端已不响应的场景。
|
||
type openAIWSForceCloser interface {
|
||
CloseNow() error
|
||
}
|
||
|
||
// openAIWSClientDialer 抽象 WS 建连器。
|
||
type openAIWSClientDialer interface {
|
||
Dial(ctx context.Context, wsURL string, headers http.Header, proxyURL string) (openAIWSClientConn, int, http.Header, error)
|
||
}
|
||
|
||
type openAIWSTransportMetricsDialer interface {
|
||
SnapshotTransportMetrics() OpenAIWSTransportMetricsSnapshot
|
||
}
|
||
|
||
func newDefaultOpenAIWSClientDialer() openAIWSClientDialer {
|
||
return &coderOpenAIWSClientDialer{
|
||
proxyClients: make(map[string]*openAIWSProxyClientEntry),
|
||
}
|
||
}
|
||
|
||
type coderOpenAIWSClientDialer struct {
|
||
proxyMu sync.Mutex
|
||
proxyClients map[string]*openAIWSProxyClientEntry
|
||
proxyHits atomic.Int64
|
||
proxyMisses atomic.Int64
|
||
}
|
||
|
||
// openAIWSHandshakeError keeps a bounded, non-logged HTTP error body so the
|
||
// Agent Identity recovery path can distinguish an invalid task from other
|
||
// 401 handshake failures.
|
||
type openAIWSHandshakeError struct {
|
||
Body []byte
|
||
Err error
|
||
}
|
||
|
||
func (e *openAIWSHandshakeError) Error() string {
|
||
if e == nil || e.Err == nil {
|
||
return "openai ws handshake failed"
|
||
}
|
||
return e.Err.Error()
|
||
}
|
||
|
||
func (e *openAIWSHandshakeError) Unwrap() error {
|
||
if e == nil {
|
||
return nil
|
||
}
|
||
return e.Err
|
||
}
|
||
|
||
type openAIWSProxyClientEntry struct {
|
||
client *http.Client
|
||
lastUsedUnixNano int64
|
||
}
|
||
|
||
func (d *coderOpenAIWSClientDialer) Dial(
|
||
ctx context.Context,
|
||
wsURL string,
|
||
headers http.Header,
|
||
proxyURL string,
|
||
) (openAIWSClientConn, int, http.Header, error) {
|
||
targetURL := strings.TrimSpace(wsURL)
|
||
if targetURL == "" {
|
||
return nil, 0, nil, errors.New("ws url is empty")
|
||
}
|
||
|
||
wrapped := &coderOpenAIWSClientConn{}
|
||
opts := &coderws.DialOptions{
|
||
HTTPHeader: cloneHeader(headers),
|
||
CompressionMode: coderws.CompressionContextTakeover,
|
||
OnPingReceived: func(context.Context, []byte) bool {
|
||
wrapped.upstreamPings.Add(1)
|
||
return true
|
||
},
|
||
}
|
||
if proxy := strings.TrimSpace(proxyURL); proxy != "" {
|
||
proxyClient, err := d.proxyHTTPClient(proxy)
|
||
if err != nil {
|
||
return nil, 0, nil, err
|
||
}
|
||
opts.HTTPClient = proxyClient
|
||
}
|
||
|
||
conn, resp, err := coderws.Dial(ctx, targetURL, opts)
|
||
if err != nil {
|
||
status := 0
|
||
respHeaders := http.Header(nil)
|
||
if resp != nil {
|
||
status = resp.StatusCode
|
||
respHeaders = cloneHeader(resp.Header)
|
||
}
|
||
var body []byte
|
||
if resp != nil && resp.Body != nil {
|
||
body, _ = io.ReadAll(io.LimitReader(resp.Body, 8<<10))
|
||
_ = resp.Body.Close()
|
||
}
|
||
return nil, status, respHeaders, &openAIWSHandshakeError{Body: body, Err: err}
|
||
}
|
||
// coder/websocket 默认单消息读取上限为 32KB，Codex WS 事件（如 rate_limits/大 delta）
|
||
// 可能超过该阈值，需显式提高上限，避免本地 read_fail(message too big)。
|
||
conn.SetReadLimit(openAIWSMessageReadLimitBytes)
|
||
respHeaders := http.Header(nil)
|
||
if resp != nil {
|
||
respHeaders = cloneHeader(resp.Header)
|
||
}
|
||
wrapped.conn = conn
|
||
return wrapped, 0, respHeaders, nil
|
||
}
|
||
|
||
func (d *coderOpenAIWSClientDialer) proxyHTTPClient(proxy string) (*http.Client, error) {
|
||
if d == nil {
|
||
return nil, errors.New("openai ws dialer is nil")
|
||
}
|
||
normalizedProxy := strings.TrimSpace(proxy)
|
||
if normalizedProxy == "" {
|
||
return nil, errors.New("proxy url is empty")
|
||
}
|
||
parsedProxyURL, err := url.Parse(normalizedProxy)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("invalid proxy url: %w", err)
|
||
}
|
||
now := time.Now().UnixNano()
|
||
|
||
d.proxyMu.Lock()
|
||
defer d.proxyMu.Unlock()
|
||
if entry, ok := d.proxyClients[normalizedProxy]; ok && entry != nil && entry.client != nil {
|
||
entry.lastUsedUnixNano = now
|
||
d.proxyHits.Add(1)
|
||
return entry.client, nil
|
||
}
|
||
d.cleanupProxyClientsLocked(now)
|
||
transport := &http.Transport{
|
||
Proxy: http.ProxyURL(parsedProxyURL),
|
||
MaxIdleConns: openAIWSProxyTransportMaxIdleConns,
|
||
MaxIdleConnsPerHost: openAIWSProxyTransportMaxIdleConnsPerHost,
|
||
IdleConnTimeout: openAIWSProxyTransportIdleConnTimeout,
|
||
TLSHandshakeTimeout: 10 * time.Second,
|
||
ForceAttemptHTTP2: true,
|
||
}
|
||
client := &http.Client{Transport: transport}
|
||
d.proxyClients[normalizedProxy] = &openAIWSProxyClientEntry{
|
||
client: client,
|
||
lastUsedUnixNano: now,
|
||
}
|
||
d.ensureProxyClientCapacityLocked()
|
||
d.proxyMisses.Add(1)
|
||
return client, nil
|
||
}
|
||
|
||
func (d *coderOpenAIWSClientDialer) cleanupProxyClientsLocked(nowUnixNano int64) {
|
||
if d == nil || len(d.proxyClients) == 0 {
|
||
return
|
||
}
|
||
idleTTL := openAIWSProxyClientCacheIdleTTL
|
||
if idleTTL <= 0 {
|
||
return
|
||
}
|
||
now := time.Unix(0, nowUnixNano)
|
||
for key, entry := range d.proxyClients {
|
||
if entry == nil || entry.client == nil {
|
||
delete(d.proxyClients, key)
|
||
continue
|
||
}
|
||
lastUsed := time.Unix(0, entry.lastUsedUnixNano)
|
||
if now.Sub(lastUsed) > idleTTL {
|
||
closeOpenAIWSProxyClient(entry.client)
|
||
delete(d.proxyClients, key)
|
||
}
|
||
}
|
||
}
|
||
|
||
func (d *coderOpenAIWSClientDialer) ensureProxyClientCapacityLocked() {
|
||
if d == nil {
|
||
return
|
||
}
|
||
maxEntries := openAIWSProxyClientCacheMaxEntries
|
||
if maxEntries <= 0 {
|
||
return
|
||
}
|
||
for len(d.proxyClients) > maxEntries {
|
||
var oldestKey string
|
||
var oldestLastUsed int64
|
||
hasOldest := false
|
||
for key, entry := range d.proxyClients {
|
||
lastUsed := int64(0)
|
||
if entry != nil {
|
||
lastUsed = entry.lastUsedUnixNano
|
||
}
|
||
if !hasOldest || lastUsed < oldestLastUsed {
|
||
hasOldest = true
|
||
oldestKey = key
|
||
oldestLastUsed = lastUsed
|
||
}
|
||
}
|
||
if !hasOldest {
|
||
return
|
||
}
|
||
if entry := d.proxyClients[oldestKey]; entry != nil {
|
||
closeOpenAIWSProxyClient(entry.client)
|
||
}
|
||
delete(d.proxyClients, oldestKey)
|
||
}
|
||
}
|
||
|
||
func closeOpenAIWSProxyClient(client *http.Client) {
|
||
if client == nil || client.Transport == nil {
|
||
return
|
||
}
|
||
if transport, ok := client.Transport.(*http.Transport); ok && transport != nil {
|
||
transport.CloseIdleConnections()
|
||
}
|
||
}
|
||
|
||
func (d *coderOpenAIWSClientDialer) SnapshotTransportMetrics() OpenAIWSTransportMetricsSnapshot {
|
||
if d == nil {
|
||
return OpenAIWSTransportMetricsSnapshot{}
|
||
}
|
||
hits := d.proxyHits.Load()
|
||
misses := d.proxyMisses.Load()
|
||
total := hits + misses
|
||
reuseRatio := 0.0
|
||
if total > 0 {
|
||
reuseRatio = float64(hits) / float64(total)
|
||
}
|
||
return OpenAIWSTransportMetricsSnapshot{
|
||
ProxyClientCacheHits: hits,
|
||
ProxyClientCacheMisses: misses,
|
||
TransportReuseRatio: reuseRatio,
|
||
}
|
||
}
|
||
|
||
type coderOpenAIWSClientConn struct {
|
||
conn *coderws.Conn
|
||
upstreamPings atomic.Int64
|
||
}
|
||
|
||
func (c *coderOpenAIWSClientConn) UpstreamPingCount() int64 {
|
||
if c == nil {
|
||
return 0
|
||
}
|
||
return c.upstreamPings.Load()
|
||
}
|
||
|
||
var _ openaiwsv2.FrameConn = (*coderOpenAIWSClientConn)(nil)
|
||
|
||
func (c *coderOpenAIWSClientConn) WriteJSON(ctx context.Context, value any) error {
|
||
if c == nil || c.conn == nil {
|
||
return errOpenAIWSConnClosed
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
return wsjson.Write(ctx, c.conn, value)
|
||
}
|
||
|
||
func (c *coderOpenAIWSClientConn) ReadMessage(ctx context.Context) ([]byte, error) {
|
||
if c == nil || c.conn == nil {
|
||
return nil, errOpenAIWSConnClosed
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
|
||
msgType, payload, err := c.conn.Read(ctx)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
switch msgType {
|
||
case coderws.MessageText, coderws.MessageBinary:
|
||
return payload, nil
|
||
default:
|
||
return nil, errOpenAIWSConnClosed
|
||
}
|
||
}
|
||
|
||
func (c *coderOpenAIWSClientConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
|
||
if c == nil || c.conn == nil {
|
||
return coderws.MessageText, nil, errOpenAIWSConnClosed
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
msgType, payload, err := c.conn.Read(ctx)
|
||
if err != nil {
|
||
return coderws.MessageText, nil, err
|
||
}
|
||
return msgType, payload, nil
|
||
}
|
||
|
||
func (c *coderOpenAIWSClientConn) WriteFrame(ctx context.Context, msgType coderws.MessageType, payload []byte) error {
|
||
if c == nil || c.conn == nil {
|
||
return errOpenAIWSConnClosed
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
return c.conn.Write(ctx, msgType, payload)
|
||
}
|
||
|
||
func (c *coderOpenAIWSClientConn) Ping(ctx context.Context) error {
|
||
if c == nil || c.conn == nil {
|
||
return errOpenAIWSConnClosed
|
||
}
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
return c.conn.Ping(ctx)
|
||
}
|
||
|
||
// SupportsIdlePingWithoutReader reports the actual coder/websocket contract.
|
||
// Conn.Ping waits for a pong, while control frames are only consumed by Read.
|
||
// Without a reader, using Ping as a health probe would deterministically time
|
||
// out a healthy socket; the pool compensates with a resident reader loop.
|
||
func (*coderOpenAIWSClientConn) SupportsIdlePingWithoutReader() bool {
|
||
return false
|
||
}
|
||
|
||
// RequiresReaderLoop 让池为 coder/websocket 连接常驻读循环，空闲期也能应答 ping。
|
||
func (*coderOpenAIWSClientConn) RequiresReaderLoop() bool {
|
||
return true
|
||
}
|
||
|
||
func (c *coderOpenAIWSClientConn) Close() error {
|
||
if c == nil || c.conn == nil {
|
||
return nil
|
||
}
|
||
// Close 为幂等，忽略重复关闭错误。
|
||
_ = c.conn.Close(coderws.StatusNormalClosure, "")
|
||
_ = c.conn.CloseNow()
|
||
return nil
|
||
}
|
||
|
||
func (c *coderOpenAIWSClientConn) CloseNow() error {
|
||
if c == nil || c.conn == nil {
|
||
return nil
|
||
}
|
||
_ = c.conn.CloseNow()
|
||
return nil
|
||
}
|