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

package openai

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

func (s *Server) streamBufferedChatCompletion(w http.ResponseWriter, dc *chatDispatchContext, handle edgeservice.RunResult, flusher http.Flusher) {
	r := dc.r
	req := dc.req
	outputPolicy := dc.outputPolicy
	attempt := 1
	metricLabels := dc.usageLabels(s, responseModeNormalized)
	for {
		result, err := collectChatCompletionOutput(r.Context(), req, handle, outputPolicy)
		if err != nil {
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), err)
			handle.Close()
			emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
			writeSSEError(w, flusher, err.Error())
			return
		}
		verr := result.toolValidationErr
		if verr == nil {
			verr = validateToolCallResponse(dc.validation, result.toolCalls, result.toolCallOrigin)
		}
		if verr != nil {
			failedRunID := handle.Dispatch().RunID
			if attempt < maxToolValidationAttempts {
				handle.Close()
				attempt++
				retryReq := dc.submitReq
				retryReq.Metadata = toolValidationAttemptMetadata(dc.submitReq.Metadata, attempt, failedRunID, verr.Error())
				s.logger.Warn("openai chat completion stream tool validation retry",
					zap.String("run_id", failedRunID),
					zap.Int("attempt", attempt),
					zap.String("reason", verr.Error()),
				)
				next, submitErr := dc.retrySubmit(r.Context(), retryReq)
				if submitErr != nil {
					emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
					writeSSEErrorWithType(w, flusher, "tool_validation_retry_error", submitErr.Error())
					return
				}
				// Extract RunResult from either a direct RunResult or a
				// ProviderPoolDispatchResult (pool retry). For provider-pool
				// retries, only normalized results with a non-nil Run are
				// accepted; tunnel or nil Run results are rejected to prevent
				// nil handle panic downstream.
				if rr, ok := next.(edgeservice.RunResult); ok {
					handle = rr
				} else if ppd, ok := next.(*edgeservice.ProviderPoolDispatchResult); ok {
					if ppd.Path != edgeservice.ProviderPoolPathNormalized || ppd.Run == nil {
						if ppd.Tunnel != nil {
							ppd.Tunnel.Close()
						}
						writeSSEErrorWithType(w, flusher, "tool_validation_retry_error", "provider-pool retry selected tunnel path")
						return
					}
					handle = ppd.Run
				} else {
					handle.Close()
					writeSSEErrorWithType(w, flusher, "run_error", "unexpected retry result type")
					return
				}
				continue
			}
			handle.Close()
			s.logger.Warn("openai chat completion stream tool validation failed",
				zap.String("run_id", failedRunID),
				zap.Int("attempt", attempt),
				zap.String("reason", verr.Error()),
			)
			emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
			writeSSEErrorWithType(w, flusher, "tool_validation_error", verr.Error())
			return
		}
		handle.Close()
		s.writeBufferedStreamOutput(w, flusher, req, handle.Dispatch(), result, outputPolicy, openAICompatTraceStreamEnabled(dc.submitReq.Metadata))
		emitUsageMetrics(metricLabels, usageStatusSuccess, usageObservationFromOpenAIUsage(result.usage, result.reasoningLen))
		return
	}
}

// writeBufferedStreamOutput emits a validated buffered result as SSE chunks:
// the role chunk, an optional content chunk, an optional tool_calls chunk, and
// the terminal finish chunk.
func (s *Server) writeBufferedStreamOutput(w http.ResponseWriter, flusher http.Flusher, req chatCompletionRequest, dispatch edgeservice.RunDispatch, result chatCompletionOutput, outputPolicy strictOutputPolicy, traceStream bool) {
	created := time.Now().Unix()
	id := "chatcmpl-" + dispatch.RunID
	model := responseModel(req.Model, dispatch.Target)
	content := result.message.Content
	toolCallNames := extractToolCallNames(result.toolCalls)
	s.logger.Info("openai chat completion stream closed",
		zap.String("run_id", 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", result.normalized),
		zap.String("response_mode", result.responseMode),
		zap.Int("content_len", len(content)),
		zap.Int("reasoning_len", result.reasoningLen),
		zap.String("assembled_content", content),
		zap.String("assembled_reasoning", result.message.ReasoningContent),
		zap.Strings("assembled_tool_calls", toolCallNames),
		zap.Int("assembled_tool_call_count", len(toolCallNames)),
	)
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})
	if traceStream {
		logOpenAICompatStreamOutput(s.logger, dispatch.RunID, 1, "content", content)
	}
	writeContentDeltaSSE(w, flusher, id, created, model, content)
	if traceStream && len(result.toolCalls) > 0 {
		s.logger.Info("openai chat completion stream output chunk",
			zap.String("run_id", dispatch.RunID),
			zap.Int("seq", 2),
			zap.String("type", "tool_calls"),
			zap.Int("tool_call_count", len(result.toolCalls)),
		)
	}
	writeToolCallsDeltaSSE(w, flusher, id, created, model, result.toolCalls)
	if traceStream {
		s.logger.Info("openai chat completion stream output done",
			zap.String("run_id", dispatch.RunID),
			zap.Int("seq", 3),
			zap.String("finish_reason", result.finishReason),
			zap.Int("content_len", len(content)),
			zap.Int("reasoning_len", result.reasoningLen),
		)
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index:        0,
			Delta:        chatDelta{},
			FinishReason: result.finishReason,
		}},
	})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}