CrossRelay/internal/responses/stream.go

package responses

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

	"gitea.rutua.cn/kate522/CrossRelay/internal/chatcmpl"
)

// StreamConverter synthesizes Responses SSE events from Chat Completions
// stream chunks.
//
// Output items are tracked by their upstream origin: one message item
// (opened on the first text delta) and one function_call item per
// upstream tool-call index, in order of first appearance. Finish closes
// every open item and emits response.completed with usage.
type StreamConverter struct {
	w       http.ResponseWriter
	flusher http.Flusher
	model   string
	respID  string

	started bool
	seq     int
	nextOut int

	msgOpen  bool
	msgIndex int
	msgID    string
	msgText  strings.Builder

	tools     map[int]*toolItem
	toolOrder []int

	inputEstimate int
	outputChars   int
	usage         *chatcmpl.Usage
	finishReason  string
}

type toolItem struct {
	outputIndex int
	id          string
	callID      string
	name        string
	started     bool
	pending     string
	args        strings.Builder
}

// NewStreamConverter builds a converter writing to w. model is echoed
// in every event's response object; inputTokens seeds the usage shown
// when the upstream stream carries none.
func NewStreamConverter(w http.ResponseWriter, model string, inputTokens int) *StreamConverter {
	return &StreamConverter{
		w:             w,
		flusher:       flusherOf(w),
		model:         model,
		respID:        "resp_" + randHex(16),
		tools:         map[int]*toolItem{},
		inputEstimate: inputTokens,
	}
}

// WriteStart sends the SSE headers and the opening response events.
func (s *StreamConverter) WriteStart() error {
	h := s.w.Header()
	h.Set("Content-Type", "text/event-stream; charset=utf-8")
	h.Set("Cache-Control", "no-cache")
	h.Set("Connection", "keep-alive")
	h.Set("X-Accel-Buffering", "no")
	s.w.WriteHeader(http.StatusOK)
	s.started = true
	initial := map[string]any{
		"id":         s.respID,
		"object":     "response",
		"created_at": time.Now().Unix(),
		"status":     "in_progress",
		"model":      s.model,
		"output":     []any{},
	}
	if err := s.emit("response.created", map[string]any{"response": initial}); err != nil {
		return err
	}
	return s.emit("response.in_progress", map[string]any{"response": initial})
}

// OnChunk folds one upstream chunk into the Responses event stream.
func (s *StreamConverter) OnChunk(chunk *chatcmpl.StreamChunk) error {
	if chunk == nil {
		return nil
	}
	if chunk.Usage != nil {
		u := *chunk.Usage
		s.usage = &u
	}
	for _, choice := range chunk.Choices {
		if choice.Delta.Content != nil && *choice.Delta.Content != "" {
			if err := s.writeText(*choice.Delta.Content); err != nil {
				return err
			}
		}
		for _, tc := range choice.Delta.ToolCalls {
			if err := s.toolDelta(tc); err != nil {
				return err
			}
		}
		if choice.FinishReason != nil && *choice.FinishReason != "" {
			s.finishReason = *choice.FinishReason
		}
	}
	return nil
}

func (s *StreamConverter) writeText(text string) error {
	if !s.msgOpen {
		s.msgIndex = s.nextOut
		s.nextOut++
		s.msgOpen = true
		s.msgID = "msg_" + randHex(12)
		if err := s.emit("response.output_item.added", map[string]any{
			"output_index": s.msgIndex,
			"item": map[string]any{
				"id":      s.msgID,
				"type":    "message",
				"role":    "assistant",
				"status":  "in_progress",
				"content": []any{},
			},
		}); err != nil {
			return err
		}
		if err := s.emit("response.content_part.added", map[string]any{
			"item_id":       s.msgID,
			"output_index":  s.msgIndex,
			"content_index": 0,
			"part":          map[string]any{"type": "output_text", "text": ""},
		}); err != nil {
			return err
		}
	}
	s.msgText.WriteString(text)
	s.outputChars += len(text)
	return s.emit("response.output_text.delta", map[string]any{
		"item_id":       s.msgID,
		"output_index":  s.msgIndex,
		"content_index": 0,
		"delta":         text,
	})
}

func (s *StreamConverter) toolDelta(tc chatcmpl.ToolCall) error {
	st := s.tools[tc.Index]
	if st == nil {
		st = &toolItem{}
		s.tools[tc.Index] = st
		s.toolOrder = append(s.toolOrder, tc.Index)
	}
	if tc.ID != "" && st.callID == "" {
		st.callID = tc.ID
	}
	if tc.Function.Name != "" && st.name == "" {
		st.name = tc.Function.Name
	}
	if tc.Function.Arguments != "" {
		st.pending += tc.Function.Arguments
	}
	// Open the item once we know a name or have arguments to flush, so
	// the client never sees an unnamed function_call.
	if !st.started && (st.name != "" || st.pending != "") {
		st.outputIndex = s.nextOut
		s.nextOut++
		st.started = true
		st.id = "fc_" + randHex(12)
		if st.callID == "" {
			st.callID = fmt.Sprintf("call_%s_%d", s.respID[len("resp_"):], tc.Index)
		}
		if err := s.emit("response.output_item.added", map[string]any{
			"output_index": st.outputIndex,
			"item": map[string]any{
				"id":        st.id,
				"type":      "function_call",
				"status":    "in_progress",
				"call_id":   st.callID,
				"name":      st.name,
				"arguments": "",
			},
		}); err != nil {
			return err
		}
		s.outputChars += len(st.name)
	}
	if st.started && st.pending != "" {
		s.outputChars += len(st.pending)
		st.args.WriteString(st.pending)
		pending := st.pending
		st.pending = ""
		return s.emit("response.function_call_arguments.delta", map[string]any{
			"item_id":      st.id,
			"output_index": st.outputIndex,
			"delta":        pending,
		})
	}
	return nil
}

// Finish closes every open item and ends the response.
func (s *StreamConverter) Finish() error {
	if s.msgOpen {
		text := s.msgText.String()
		if err := s.emit("response.output_text.done", map[string]any{
			"item_id":       s.msgID,
			"output_index":  s.msgIndex,
			"content_index": 0,
			"text":          text,
		}); err != nil {
			return err
		}
		if err := s.emit("response.content_part.done", map[string]any{
			"item_id":       s.msgID,
			"output_index":  s.msgIndex,
			"content_index": 0,
			"part":          map[string]any{"type": "output_text", "text": text},
		}); err != nil {
			return err
		}
		if err := s.emit("response.output_item.done", map[string]any{
			"output_index": s.msgIndex,
			"item": map[string]any{
				"id":      s.msgID,
				"type":    "message",
				"role":    "assistant",
				"status":  "completed",
				"content": []any{map[string]any{"type": "output_text", "text": text}},
			},
		}); err != nil {
			return err
		}
		s.msgOpen = false
	}
	for _, idx := range s.toolOrder {
		st := s.tools[idx]
		if st == nil || !st.started {
			continue
		}
		if err := s.emit("response.function_call_arguments.done", map[string]any{
			"item_id":      st.id,
			"output_index": st.outputIndex,
			"arguments":    st.args.String(),
		}); err != nil {
			return err
		}
		if err := s.emit("response.output_item.done", map[string]any{
			"output_index": st.outputIndex,
			"item": map[string]any{
				"id":        st.id,
				"type":      "function_call",
				"status":    "completed",
				"call_id":   st.callID,
				"name":      st.name,
				"arguments": st.args.String(),
			},
		}); err != nil {
			return err
		}
	}
	return s.emit("response.completed", map[string]any{"response": s.finalResponse()})
}

// WriteError terminates the stream with a response.failed event.
func (s *StreamConverter) WriteError(code, message string) {
	_ = s.emit("response.failed", map[string]any{
		"response": map[string]any{
			"id":         s.respID,
			"object":     "response",
			"created_at": time.Now().Unix(),
			"status":     "failed",
			"model":      s.model,
			"output":     []any{},
			"error":      map[string]any{"code": code, "message": message},
		},
	})
}

func (s *StreamConverter) finalResponse() map[string]any {
	status, incomplete := MapStatus(s.finishReason)
	output := []any{}
	text := s.msgText.String()
	if text != "" {
		output = append(output, map[string]any{
			"id":      s.msgID,
			"type":    "message",
			"role":    "assistant",
			"status":  "completed",
			"content": []any{map[string]any{"type": "output_text", "text": text}},
		})
	}
	for _, idx := range s.toolOrder {
		st := s.tools[idx]
		if st == nil || !st.started {
			continue
		}
		output = append(output, map[string]any{
			"id":        st.id,
			"type":      "function_call",
			"status":    "completed",
			"call_id":   st.callID,
			"name":      st.name,
			"arguments": st.args.String(),
		})
	}
	resp := map[string]any{
		"id":         s.respID,
		"object":     "response",
		"created_at": time.Now().Unix(),
		"status":     status,
		"model":      s.model,
		"output":     output,
	}
	if text != "" {
		resp["output_text"] = text
	}
	if incomplete != nil {
		resp["incomplete_details"] = incomplete
	}
	resp["usage"] = s.usageMap()
	return resp
}

func (s *StreamConverter) usageMap() map[string]any {
	if s.usage != nil {
		total := s.usage.TotalTokens
		if total == 0 {
			total = s.usage.PromptTokens + s.usage.CompletionTokens
		}
		return map[string]any{
			"input_tokens":  s.usage.PromptTokens,
			"output_tokens": s.usage.CompletionTokens,
			"total_tokens":  total,
		}
	}
	out := (s.outputChars + 3) / 4
	return map[string]any{
		"input_tokens":  s.inputEstimate,
		"output_tokens": out,
		"total_tokens":  s.inputEstimate + out,
	}
}

func (s *StreamConverter) emit(event string, payload map[string]any) error {
	s.seq++
	payload["type"] = event
	payload["sequence_number"] = s.seq
	return writeEvent(s.w, s.flusher, event, payload)
}

func writeEvent(w http.ResponseWriter, flusher http.Flusher, event string, payload any) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	var sb strings.Builder
	sb.WriteString("event: ")
	sb.WriteString(event)
	sb.WriteString("\ndata: ")
	sb.Write(data)
	sb.WriteString("\n\n")
	if _, err := io.WriteString(w, sb.String()); err != nil {
		return err
	}
	if flusher != nil {
		flusher.Flush()
	}
	return nil
}

func flusherOf(w http.ResponseWriter) http.Flusher {
	if f, ok := w.(http.Flusher); ok {
		return f
	}
	return nil
}