iop/apps/edge/internal/openai/normalized_sse.go

package openai

import (
	"go.uber.org/zap"
	edgeservice "iop/apps/edge/internal/service"
	"net/http"
	"os"
	"strings"
	"time"
	"unicode"
)

const (
	streamTraceMetadataKey = "iop_trace_stream"
	streamTraceEnvKey      = "IOP_OPENAI_COMPAT_TRACE_STREAM"
)

func (s *Server) streamChatCompletion(w http.ResponseWriter, dc *chatDispatchContext, handle edgeservice.RunResult) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		handle.Close()
		writeError(w, http.StatusInternalServerError, "streaming_not_supported", "response writer does not support streaming")
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	// Buffered streams collect and validate the full response before emitting
	// user-visible chunks. Strict buffered output opts in explicitly, and
	// tool-bearing streams use the same path so malformed tool calls can be
	// retried or rejected before reaching a client tool runner.
	if (dc.outputPolicy.Strict && dc.outputPolicy.StreamBuffer) || len(dc.req.Tools) > 0 {
		s.streamBufferedChatCompletion(w, dc, handle, flusher)
		return
	}

	// Live SSE may emit content deltas before the terminal event, so runtime
	// tool validation is excluded upstream; write the role chunk immediately.
	defer handle.Close()
	sess := s.newChatStreamSession(w, flusher, dc.req, dc.submitReq, handle, dc.outputPolicy, dc.usageLabels(s, responseModeNormalized))
	sess.writeRole()
	defer sess.logClosed()

	stream := handle.Stream()
	if stream.Events == nil {
		sess.failStream("run stream unavailable")
		return
	}
	sess.consume(dc.r, stream, handle.WaitTimeout())
}

// consume drives the session until it reaches a terminal state: a terminal run
// event, a node disconnect, a caller cancel, or the run wait timeout.
func (sess *chatStreamSession) consume(r *http.Request, stream edgeservice.RunStream, waitTimeout time.Duration) {
	for {
		select {
		case <-r.Context().Done():
			sess.cancel(r.Context().Err(), usageStatusForError(r.Context().Err()), false)
			return
		case nodeEvent, ok := <-stream.NodeEvents:
			if !ok {
				stream.NodeEvents = nil
				continue
			}
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				sess.fail("node disconnected")
				return
			}
		case event, ok := <-stream.Events:
			if !ok {
				sess.failStream("run stream closed")
				return
			}
			if sess.consumeEvent(event) {
				return
			}
		case <-time.After(waitTimeout):
			sess.cancel(errRunTimedOut, usageStatusCancel, true)
			return
		}
	}
}

type streamToolTextFilter struct {
	pending string
}

func (f *streamToolTextFilter) Append(delta string) string {
	f.pending += delta
	if idx := firstStreamToolTextCandidateIndex(f.pending); idx >= 0 {
		out := strings.TrimRightFunc(f.pending[:idx], unicode.IsSpace)
		f.pending = f.pending[idx:]
		return out
	}
	flushLen := streamToolTextSafeFlushLen(f.pending)
	out := strings.TrimRightFunc(f.pending[:flushLen], unicode.IsSpace)
	f.pending = f.pending[len(out):]
	return out
}

func (f *streamToolTextFilter) Flush() string {
	out := f.pending
	f.pending = ""
	return out
}

func (f *streamToolTextFilter) FlushNonCandidate() string {
	if firstStreamToolTextCandidateIndex(f.pending) >= 0 {
		f.pending = ""
		return ""
	}
	// Even without a full candidate, chunk-boundary protection can leave a
	// partial candidate suffix (e.g. "<tool_ca", trailing "{") in pending.
	// Return only the safe non-candidate prefix and drop the partial suffix so
	// a native tool_calls completion never flushes a raw candidate fragment.
	flushLen := streamToolTextSafeFlushLen(f.pending)
	out := strings.TrimRightFunc(f.pending[:flushLen], unicode.IsSpace)
	f.pending = ""
	return out
}

func firstStreamToolTextCandidateIndex(s string) int {
	xml := strings.Index(strings.ToLower(s), "<tool_call")
	mustache := strings.Index(s, textMustacheToolCallOpenHint)
	switch {
	case xml < 0:
		return mustache
	case mustache < 0:
		return xml
	case xml < mustache:
		return xml
	default:
		return mustache
	}
}

func streamToolTextSafeFlushLen(s string) int {
	if s == "" {
		return 0
	}
	const xmlHint = "<tool_call"
	lower := strings.ToLower(s)
	start := len(s) - len(xmlHint) + 1
	if start < 0 {
		start = 0
	}
	for i := start; i < len(s); i++ {
		if strings.HasPrefix(xmlHint, lower[i:]) ||
			strings.HasPrefix(textMustacheToolCallOpenHint, s[i:]) {
			return i
		}
	}
	return len(s)
}

func unstreamedCleanedText(cleaned, emitted string) string {
	if cleaned == "" {
		return ""
	}
	if emitted == "" {
		return cleaned
	}
	if strings.HasPrefix(cleaned, emitted) {
		return cleaned[len(emitted):]
	}
	cleaned = strings.TrimSpace(cleaned)
	emitted = strings.TrimSpace(emitted)
	if cleaned == "" || cleaned == emitted {
		return ""
	}
	if emitted != "" && strings.HasPrefix(cleaned, emitted) {
		return strings.TrimSpace(strings.TrimPrefix(cleaned, emitted))
	}
	return ""
}

func openAICompatTraceStreamEnabled(metadata map[string]string) bool {
	if traceBool(os.Getenv(streamTraceEnvKey)) {
		return true
	}
	if metadata == nil {
		return false
	}
	return traceBool(metadata[streamTraceMetadataKey])
}

func traceBool(v string) bool {
	switch strings.ToLower(strings.TrimSpace(v)) {
	case "1", "true", "yes", "on":
		return true
	default:
		return false
	}
}

func logOpenAICompatStreamOutput(logger *zap.Logger, runID string, seq int, typ string, delta string, extra ...zap.Field) {
	fields := []zap.Field{
		zap.String("run_id", runID),
		zap.Int("seq", seq),
		zap.String("type", typ),
		zap.Int("delta_len", len(delta)),
		zap.Bool("delta_has_open_think", hasOpenThinkTag(delta)),
		zap.Bool("delta_has_close_think", hasCloseThinkTag(delta)),
	}
	fields = append(fields, extra...)
	logger.Info("openai chat completion stream output chunk", fields...)
}

func hasOpenThinkTag(s string) bool {
	return strings.Contains(strings.ToLower(s), "<think")
}

func hasCloseThinkTag(s string) bool {
	return strings.Contains(strings.ToLower(s), "</think")
}