AIRouter/internal/handler/sse.go

package handler

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/edgeai/gateway/internal/adapter"
	"github.com/edgeai/gateway/pkg/api"
)

// SSEWriter writes Server-Sent Events to an HTTP response.
type SSEWriter struct {
	w       http.ResponseWriter
	flusher http.Flusher
}

// NewSSEWriter creates a new SSEWriter. Returns nil if streaming is not supported.
func NewSSEWriter(w http.ResponseWriter) *SSEWriter {
	flusher, ok := w.(http.Flusher)
	if !ok {
		return nil
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	return &SSEWriter{w: w, flusher: flusher}
}

// WriteChunk writes a single SSE data event.
func (s *SSEWriter) WriteChunk(data any) error {
	jsonData, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("marshal sse data: %w", err)
	}
	fmt.Fprintf(s.w, "data: %s\n\n", jsonData)
	s.flusher.Flush()
	return nil
}

// WriteDone writes the [DONE] marker.
func (s *SSEWriter) WriteDone() {
	fmt.Fprintf(s.w, "data: [DONE]\n\n")
	s.flusher.Flush()
}

// WritePing 发送 SSE 注释行作为心跳，防止代理/负载均衡器因空闲超时断开连接。
// 注释行以冒号开头，客户端会忽略，不影响事件流。
func (s *SSEWriter) WritePing() {
	fmt.Fprintf(s.w, ": keepalive\n\n")
	s.flusher.Flush()
}

// StreamChatCompletion streams chunks from an adapter to the client in OpenAI SSE format.
// Returns input tokens, output tokens, full content string, and error.
// onFirstToken is called when the first content chunk is received (may be nil).
// 每 15 秒发送一次 keepalive ping，防止连接被代理断开。
func StreamChatCompletion(sse *SSEWriter, ch <-chan adapter.StreamChunk, requestID, taskID, model string, onFirstToken func()) (int, int, string, error) {
	inputTokens := 0
	outputTokens := 0
	var contentBuilder strings.Builder
	firstTokenSent := false

	// 启动 keepalive 心跳定时器
	keepalive := time.NewTicker(15 * time.Second)
	defer keepalive.Stop()

	for {
		select {
		case chunk, ok := <-ch:
			if !ok {
				return inputTokens, outputTokens, contentBuilder.String(), nil
			}
			if chunk.Error != nil {
				return inputTokens, outputTokens, contentBuilder.String(), chunk.Error
			}

			if chunk.Done {
				if chunk.InputTokens > 0 {
					inputTokens = chunk.InputTokens
				}
				if chunk.OutputTokens > 0 {
					outputTokens = chunk.OutputTokens
				}

				// Write final chunk with finish_reason
				sseChunk := map[string]any{
					"id":     requestID,
					"object": "chat.completion.chunk",
					"model":  model,
					"choices": []map[string]any{
						{
							"index":         0,
							"delta":         map[string]any{},
							"finish_reason": chunk.FinishReason,
						},
					},
				}
				if inputTokens > 0 || outputTokens > 0 {
					sseChunk["usage"] = map[string]int{
						"input_tokens":  inputTokens,
						"output_tokens": outputTokens,
						"total_tokens":  inputTokens + outputTokens,
					}
				}
				sse.WriteChunk(sseChunk)
				sse.WriteDone()
				return inputTokens, outputTokens, contentBuilder.String(), nil
			}

			// 首 token 延迟回调
			if !firstTokenSent && chunk.Delta != "" {
				firstTokenSent = true
				if onFirstToken != nil {
					onFirstToken()
				}
			}

			// Write content delta
			sseChunk := map[string]any{
				"id":     requestID,
				"object": "chat.completion.chunk",
				"model":  model,
				"choices": []map[string]any{
					{
						"index": 0,
						"delta": map[string]any{
							"content": chunk.Delta,
						},
						"finish_reason": nil,
					},
				},
			}
			sse.WriteChunk(sseChunk)
			contentBuilder.WriteString(chunk.Delta)

		case <-keepalive.C:
			// 发送心跳，保持连接活跃
			sse.WritePing()
		}
	}
}

// BuildChatResponse creates a non-streaming ChatResponse from adapter result.
func BuildChatResponse(requestID, taskID, logicalModel string, resp *adapter.ChatResponse) api.ChatResponse {
	return api.ChatResponse{
		RequestID: requestID,
		TaskID:    taskID,
		Status:    "completed",
		Model:     logicalModel,
		Choices: []api.Choice{
			{
				Index: 0,
				Message: &api.Message{
					Role:    "assistant",
					Content: resp.Content,
				},
				FinishReason: resp.FinishReason,
			},
		},
		LogicalModel: logicalModel,
		ActualModel:  resp.ActualModel,
		Usage: &api.Usage{
			InputTokens:  resp.InputTokens,
			OutputTokens: resp.OutputTokens,
			TotalTokens:  resp.InputTokens + resp.OutputTokens,
		},
	}
}