复用 Grok Realtime 预握手连接

backend/internal/handler/grok_audio.go

@@ -56,6 +56,7 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
 	var selection *service.AccountSelectionResult
 	var release func()
 	var token string
+	var upstream *service.GrokRealtimeUpstream
 	for attempts := 0; attempts < 4; attempts++ {
 		candidate, _, selectErr := h.gatewayService.SelectAccountWithSchedulerForCapability(
 			c.Request.Context(), apiKey.GroupID, "", "", "grok-4.6", failed,
@@ -86,23 +87,24 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
 			continue
 		}
 		probeCtx, cancelProbe := context.WithTimeout(c.Request.Context(), 15*time.Second)
-		probeErr := h.gatewayService.ProbeGrokRealtime(probeCtx, account, token, model)
+		candidateUpstream, openErr := h.gatewayService.OpenGrokRealtime(probeCtx, account, token, model)
 		cancelProbe()
-		if probeErr != nil {
-			reqLog.Warn("grok_realtime.pre_accept_failed", zap.Int64("account_id", account.ID), zap.Error(probeErr))
+		if openErr != nil {
+			reqLog.Warn("grok_realtime.pre_accept_failed", zap.Int64("account_id", account.ID), zap.Error(openErr))
 			release()
 			release = nil
 			failed[account.ID] = struct{}{}
 			continue
 		}
-		selection = candidate
+		selection, upstream = candidate, candidateUpstream
 		break
 	}
-	if selection == nil || selection.Account == nil || release == nil {
+	if selection == nil || selection.Account == nil || release == nil || upstream == nil {
 		h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Grok realtime upstream unavailable")
 		return
 	}
 	defer release()
+	defer upstream.Close()
 
 	conn, err := coderws.Accept(c.Writer, c.Request, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
 	if err != nil {
@@ -111,7 +113,7 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
 	defer func() { _ = conn.CloseNow() }()
 
 	started := time.Now()
-	audioObserved, proxyErr := h.gatewayService.ProxyGrokRealtime(c.Request.Context(), c, conn, selection.Account, token, model)
+	audioObserved, proxyErr := h.gatewayService.ProxyGrokRealtimeConn(c.Request.Context(), c, conn, upstream)
 	elapsed := time.Since(started)
 	if proxyErr != nil {
 		reqLog.Info("grok_realtime.proxy_failed", zap.Error(proxyErr))

backend/internal/service/grok_audio.go

@@ -125,35 +125,60 @@ func (s *OpenAIGatewayService) ProxyGrokRealtime(ctx context.Context, c *gin.Con
 	if account.Platform != PlatformGrok {
 		return false, fmt.Errorf("account platform %s is not supported for grok realtime", account.Platform)
 	}
-	base, err := buildGrokVoiceURL(account, s.cfg, "realtime")
+	upstream, err := s.OpenGrokRealtime(ctx, account, token, model)
 	if err != nil {
 		return false, err
 	}
+	defer upstream.Close()
+	return s.ProxyGrokRealtimeConn(ctx, c, client, upstream)
+}
+
+type GrokRealtimeUpstream struct{ conn openAIWSClientConn }
+
+func (u *GrokRealtimeUpstream) Close() error {
+	if u == nil || u.conn == nil {
+		return nil
+	}
+	return u.conn.Close()
+}
+
+func (s *OpenAIGatewayService) OpenGrokRealtime(ctx context.Context, account *Account, token, model string) (*GrokRealtimeUpstream, error) {
+	if s == nil || account == nil || account.Platform != PlatformGrok {
+		return nil, fmt.Errorf("grok realtime account is required")
+	}
+	base, err := buildGrokVoiceURL(account, s.cfg, "realtime")
+	if err != nil {
+		return nil, err
+	}
 	u, err := url.Parse(base)
 	if err != nil {
-		return false, err
+		return nil, err
 	}
 	u.Scheme = "wss"
-	u.RawQuery = "model=" + url.QueryEscape(firstNonEmpty(model, "grok-voice-latest"))
+	q := u.Query()
+	q.Set("model", firstNonEmpty(model, "grok-voice-latest"))
+	u.RawQuery = q.Encode()
 	headers := http.Header{"Authorization": []string{"Bearer " + token}}
-	// Match media/voice HTTP: CLI headers only on CLI proxy hosts.
 	if account.IsGrokOAuth() && isGrokCLIProxyTarget(u.String()) {
 		applyGrokCLIHeaders(headers)
 	}
-	if account != nil {
-		account.ApplyHeaderOverrides(headers)
-	}
-
-	dialer := s.getOpenAIWSPassthroughDialer()
+	account.ApplyHeaderOverrides(headers)
 	proxyURL := ""
 	if account.ProxyID != nil && account.Proxy != nil {
 		proxyURL = account.Proxy.URL()
 	}
-	upstream, _, _, err := dialer.Dial(ctx, u.String(), headers, proxyURL)
+	conn, _, _, err := s.getOpenAIWSPassthroughDialer().Dial(ctx, u.String(), headers, proxyURL)
 	if err != nil {
-		return false, err
+		return nil, err
 	}
-	defer func() { _ = upstream.Close() }()
+	return &GrokRealtimeUpstream{conn: conn}, nil
+}
+
+func (s *OpenAIGatewayService) ProxyGrokRealtimeConn(ctx context.Context, c *gin.Context, client *coderws.Conn, upstream *GrokRealtimeUpstream) (bool, error) {
+	if s == nil || client == nil || upstream == nil || upstream.conn == nil {
+		return false, fmt.Errorf("realtime connection is required")
+	}
+	conn := upstream.conn
 
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
@@ -163,7 +188,7 @@ func (s *OpenAIGatewayService) ProxyGrokRealtime(ctx context.Context, c *gin.Con
 	// Upstream → client
 	go func() {
 		for {
-			msg, readErr := upstream.ReadMessage(ctx)
+			msg, readErr := conn.ReadMessage(ctx)
 			if readErr != nil {
 				errCh <- readErr
 				return
@@ -197,7 +222,7 @@ func (s *OpenAIGatewayService) ProxyGrokRealtime(ctx context.Context, c *gin.Con
 				errCh <- fmt.Errorf("invalid realtime event: %w", unmarshalErr)
 				return
 			}
-			if writeErr := upstream.WriteJSON(ctx, raw); writeErr != nil {
+			if writeErr := conn.WriteJSON(ctx, raw); writeErr != nil {
 				errCh <- writeErr
 				return
 			}