openteam/server/internal/proxy/stream.go

package proxy

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"

	"openteam/server/internal/proxy/claude"
	"openteam/server/internal/proxy/convert"
	"openteam/server/internal/proxy/openai"
	"openteam/server/internal/proxy/stream"
	"openteam/server/internal/store"
)

// streamOpenAI forwards an SSE stream to an OpenAI-protocol client.
func (g *Gateway) streamOpenAI(c *gin.Context, route Route, resp *http.Response,
	start time.Time, reqID string, auth *authContext, model *store.Model, ch *store.Channel) {

	w := c.Writer
	tr := convert.NewTranslator(clientProto(route), upstreamProvider(route, ch))
	done := make(chan struct{})

	if tr == nil {
		// Passthrough: copy raw frames, sniffing usage from OpenAI chunks
		// and Responses-API stream events.
		var usage *tokenUsage
		onEvent := func(ev stream.SSEEvent) {
			if ev.Done || ev.Data == "" {
				return
			}
			var chunk openai.ChatChunk
			if json.Unmarshal([]byte(ev.Data), &chunk) == nil && chunk.Usage != nil {
				usage = &tokenUsage{
					input:  chunk.Usage.PromptTokens,
					output: chunk.Usage.CompletionTokens,
				}
				return
			}
			var resp respStreamEvent
			if json.Unmarshal([]byte(ev.Data), &resp) == nil && resp.Response != nil && resp.Response.Usage != nil {
				usage = &tokenUsage{
					input:  resp.Response.Usage.InputTokens,
					output: resp.Response.Usage.OutputTokens,
				}
			}
		}
		go func() {
			defer close(done)
			copyRawSSE(w, resp.Body, onEvent)
		}()
		<-done
		g.afterComplete(start, reqID, auth, model, ch, usage, "success", "")
		return
	}

	// Converted: read upstream events, emit translated frames.
	var finalUsage *tokenUsage
	go func() {
		defer close(done)
		err := stream.ReadSSE(resp.Body, func(ev stream.SSEEvent) error {
			frames, ferr := tr.Feed(ev)
			if ferr != nil {
				return ferr
			}
			for _, f := range frames {
				if err := stream.Write(w, f); err != nil {
					return err
				}
			}
			return nil
		})
		if err != nil {
			g.log.Debug("upstream stream read ended", zap.Error(err))
		}
		fin, _ := tr.Finish()
		for _, f := range fin {
			_ = stream.Write(w, f)
		}
		if u := tr.Usage(); u != nil {
			finalUsage = &tokenUsage{input: u.Input, output: u.Output, cacheRead: u.CacheRead, cacheCreation: u.CacheCreation}
		}
	}()
	<-done
	g.afterComplete(start, reqID, auth, model, ch, finalUsage, "success", "")
}

// streamAnthropic forwards an SSE stream to a Claude-protocol client.
func (g *Gateway) streamAnthropic(c *gin.Context, route Route, resp *http.Response,
	start time.Time, reqID string, auth *authContext, model *store.Model, ch *store.Channel) {

	w := c.Writer
	tr := convert.NewTranslator(clientProto(route), upstreamProvider(route, ch))
	done := make(chan struct{})

	if tr == nil {
		// Passthrough: copy raw frames, sniffing usage from message_delta.
		var usage *tokenUsage
		onEvent := func(ev stream.SSEEvent) {
			if ev.Done || ev.Data == "" {
				return
			}
			var e claude.StreamEvent
			if json.Unmarshal([]byte(ev.Data), &e) != nil || e.Type != "message_delta" {
				return
			}
			var u claude.Usage
			if json.Unmarshal(e.Usage, &u) == nil {
				usage = &tokenUsage{
					input: u.InputTokens, output: u.OutputTokens,
					cacheRead: u.CacheReadInputTokens, cacheCreation: u.CacheCreationInputTokens,
				}
			}
		}
		go func() {
			defer close(done)
			copyRawSSE(w, resp.Body, onEvent)
		}()
		<-done
		g.afterComplete(start, reqID, auth, model, ch, usage, "success", "")
		return
	}

	// Converted: read upstream events, emit translated frames.
	var finalUsage *tokenUsage
	go func() {
		defer close(done)
		err := stream.ReadSSE(resp.Body, func(ev stream.SSEEvent) error {
			frames, ferr := tr.Feed(ev)
			if ferr != nil {
				return ferr
			}
			for _, f := range frames {
				if err := stream.Write(w, f); err != nil {
					return err
				}
			}
			return nil
		})
		if err != nil {
			g.log.Debug("upstream stream read ended", zap.Error(err))
		}
		fin, _ := tr.Finish()
		for _, f := range fin {
			_ = stream.Write(w, f)
		}
		if u := tr.Usage(); u != nil {
			finalUsage = &tokenUsage{input: u.Input, output: u.Output, cacheRead: u.CacheRead, cacheCreation: u.CacheCreation}
		}
	}()
	<-done
	g.afterComplete(start, reqID, auth, model, ch, finalUsage, "success", "")
}

// copyRawSSE copies an upstream SSE stream verbatim, invoking onEvent for
// each data frame (used for passthrough + usage sniffing).
func copyRawSSE(w http.ResponseWriter, r io.Reader, onEvent func(stream.SSEEvent)) {
	br := bufio.NewReader(r)
	for {
		line, err := br.ReadBytes('\n')
		if len(line) > 0 {
			if _, werr := w.Write(line); werr != nil {
				return
			}
			trimmed := bytes.TrimSpace(line)
			if bytes.HasPrefix(trimmed, []byte("data:")) {
				data := bytes.TrimSpace(trimmed[len("data:"):])
				ev := stream.SSEEvent{
					Data: string(data),
					Done: bytes.Equal(data, []byte("[DONE]")),
				}
				if ev.Data != "" {
					onEvent(ev)
				}
			}
		}
		if err != nil {
			break
		}
	}
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
}

// clientProto maps a proxy Protocol to the convert package's protocol type.
func clientProto(route Route) convert.ClientProtocol {
	switch route.Protocol {
	case ProtocolOpenAIResponses:
		return convert.ClientOpenAIResponses
	case ProtocolAnthropic:
		return convert.ClientAnthropic
	default:
		return convert.ClientOpenAIChat
	}
}