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

package openai

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

	"go.uber.org/zap"

	edgeservice "iop/apps/edge/internal/service"
)

func (s *Server) streamChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, handle edgeservice.RunResult, outputPolicy strictOutputPolicy) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		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")

	created := time.Now().Unix()
	id := "chatcmpl-" + handle.Dispatch().RunID
	model := responseModel(req.Model, handle.Dispatch().Target)
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})

	if outputPolicy.Strict && outputPolicy.StreamBuffer {
		s.streamBufferedChatCompletion(w, r, handle, id, created, model, flusher, outputPolicy, req.Tools)
		return
	}

	var contentBuilder strings.Builder
	var reasoningBuilder strings.Builder
	var emittedContent strings.Builder
	var toolTextFilter *streamToolTextFilter
	if len(req.Tools) > 0 {
		toolTextFilter = &streamToolTextFilter{}
	}
	defer func() {
		s.logger.Info("openai chat completion stream closed",
			zap.String("run_id", handle.Dispatch().RunID),
			zap.Bool("strict_output", outputPolicy.Strict),
			zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
			zap.Int("content_len", contentBuilder.Len()),
			zap.Int("reasoning_len", reasoningBuilder.Len()),
			zap.String("content_preview", previewString(contentBuilder.String(), 1000)),
			zap.String("reasoning_preview", previewString(reasoningBuilder.String(), 1000)),
		)
	}()

	stream := handle.Stream()
	if stream.Events == nil {
		writeSSEError(w, flusher, "run stream unavailable")
		return
	}
	for {
		select {
		case <-r.Context().Done():
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
			return
		case nodeEvent, ok := <-stream.NodeEvents:
			if !ok {
				stream.NodeEvents = nil
				continue
			}
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				writeSSEError(w, flusher, "node disconnected")
				return
			}
		case event, ok := <-stream.Events:
			if !ok {
				writeSSEError(w, flusher, "run stream closed")
				return
			}
			if event == nil {
				continue
			}
			switch event.GetType() {
			case "delta":
				if event.GetDelta() == "" {
					continue
				}
				delta := event.GetDelta()
				contentBuilder.WriteString(delta)
				if toolTextFilter != nil {
					delta = toolTextFilter.Append(delta)
				}
				if delta == "" {
					continue
				}
				emittedContent.WriteString(delta)
				writeContentDeltaSSE(w, flusher, id, created, model, delta)
			case "reasoning_delta":
				if event.GetDelta() == "" {
					continue
				}
				reasoningBuilder.WriteString(event.GetDelta())
				if outputPolicy.Strict {
					continue
				}
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index: 0,
						Delta: chatDelta{ReasoningContent: event.GetDelta()},
					}},
				})
			case "complete":
				metadata := event.GetMetadata()
				finishReason := metadata["finish_reason"]
				if finishReason == "" {
					finishReason = "stop"
				}
				nativeToolCalls := toolCallsFromRunMetadata(metadata)
				if len(nativeToolCalls) > 0 {
					if toolTextFilter != nil {
						pending := toolTextFilter.Flush()
						if pending != "" {
							emittedContent.WriteString(pending)
							writeContentDeltaSSE(w, flusher, id, created, model, pending)
						}
					}
					nativeToolCalls = normalizeNativeToolCallsForResponse(nativeToolCalls, req.Tools)
					writeToolCallsDeltaSSE(w, flusher, id, created, model, nativeToolCalls)
					finishReason = "tool_calls"
				} else if toolTextFilter != nil && shouldSynthesizeTextToolCalls(handle.Dispatch(), isTextToolFallback(metadata)) {
					cleaned, toolCalls := synthesizeToolCallsFromText(contentBuilder.String(), req.Tools, handle.Dispatch().RunID)
					if len(toolCalls) > 0 {
						if remaining := unstreamedCleanedText(cleaned, emittedContent.String()); remaining != "" {
							emittedContent.WriteString(remaining)
							writeContentDeltaSSE(w, flusher, id, created, model, remaining)
						}
						writeToolCallsDeltaSSE(w, flusher, id, created, model, toolCalls)
						finishReason = "tool_calls"
					} else if pending := toolTextFilter.Flush(); pending != "" {
						emittedContent.WriteString(pending)
						writeContentDeltaSSE(w, flusher, id, created, model, pending)
					}
				} else if toolTextFilter != nil {
					if pending := toolTextFilter.Flush(); pending != "" {
						emittedContent.WriteString(pending)
						writeContentDeltaSSE(w, flusher, id, created, model, pending)
					}
				}
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index:        0,
						Delta:        chatDelta{},
						FinishReason: finishReason,
					}},
				})
				fmt.Fprint(w, "data: [DONE]\n\n")
				flusher.Flush()
				return
			case "error", "cancelled":
				msg := event.GetError()
				if msg == "" {
					msg = event.GetMessage()
				}
				if msg == "" {
					msg = "run failed"
				}
				writeSSEError(w, flusher, msg)
				return
			}
		case <-time.After(handle.WaitTimeout()):
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
			writeSSEError(w, flusher, errRunTimedOut.Error())
			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 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 writeContentDeltaSSE(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string, content string) {
	if content == "" {
		return
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Content: content},
		}},
	})
}

func writeToolCallsDeltaSSE(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string, toolCalls []any) {
	if len(toolCalls) == 0 {
		return
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{ToolCalls: toolCallsForStreamDelta(toolCalls)},
		}},
	})
}

func (s *Server) streamBufferedChatCompletion(w http.ResponseWriter, r *http.Request, handle edgeservice.RunResult, id string, created int64, model string, flusher http.Flusher, outputPolicy strictOutputPolicy, tools []any) {
	text, reasoning, finishReason, nativeToolCalls, _, textToolFallback, err := collectRunResult(r.Context(), handle.Stream(), handle.WaitTimeout())
	if err != nil {
		s.cancelRunOnHTTPGiveUp(handle.Dispatch(), err)
		writeSSEError(w, flusher, err.Error())
		return
	}
	text, reasoning, normalized := normalizeCompletionOutput(outputPolicy, text, reasoning)
	s.logger.Info("openai chat completion stream closed",
		zap.String("run_id", handle.Dispatch().RunID),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("normalized", normalized),
		zap.Int("content_len", len(text)),
		zap.Int("reasoning_len", len(reasoning)),
		zap.String("content_preview", previewString(text, 1000)),
		zap.String("reasoning_preview", previewString(reasoning, 1000)),
	)
	if len(nativeToolCalls) > 0 {
		nativeToolCalls = normalizeNativeToolCallsForResponse(nativeToolCalls, tools)
		if text != "" {
			writeSSE(w, flusher, chatCompletionChunk{
				ID:      id,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   model,
				Choices: []chatCompletionChunkChoice{{
					Index: 0,
					Delta: chatDelta{Content: text},
				}},
			})
		}
		writeSSE(w, flusher, chatCompletionChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []chatCompletionChunkChoice{{
				Index: 0,
				Delta: chatDelta{ToolCalls: toolCallsForStreamDelta(nativeToolCalls)},
			}},
		})
		finishReason = "tool_calls"
	} else if shouldSynthesizeTextToolCalls(handle.Dispatch(), textToolFallback) {
		cleaned, toolCalls := synthesizeToolCallsFromText(text, tools, handle.Dispatch().RunID)
		if len(toolCalls) > 0 {
			if cleaned != "" {
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index: 0,
						Delta: chatDelta{Content: cleaned},
					}},
				})
			}
			writeSSE(w, flusher, chatCompletionChunk{
				ID:      id,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   model,
				Choices: []chatCompletionChunkChoice{{
					Index: 0,
					Delta: chatDelta{ToolCalls: toolCallsForStreamDelta(toolCalls)},
				}},
			})
			finishReason = "tool_calls"
		} else if text != "" {
			writeSSE(w, flusher, chatCompletionChunk{
				ID:      id,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   model,
				Choices: []chatCompletionChunkChoice{{
					Index: 0,
					Delta: chatDelta{Content: text},
				}},
			})
		}
	} else if text != "" {
		writeSSE(w, flusher, chatCompletionChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []chatCompletionChunkChoice{{
				Index: 0,
				Delta: chatDelta{Content: text},
			}},
		})
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index:        0,
			Delta:        chatDelta{},
			FinishReason: finishReason,
		}},
	})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}

func writeSSE(w http.ResponseWriter, flusher http.Flusher, v any) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	fmt.Fprintf(w, "data: %s\n\n", b)
	flusher.Flush()
}

func writeSSEError(w http.ResponseWriter, flusher http.Flusher, message string) {
	writeSSE(w, flusher, errorResponse{Error: errorBody{Type: "run_error", Message: message}})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}