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

package openai

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"go.uber.org/zap"
	edgeservice "iop/apps/edge/internal/service"
	iop "iop/proto/gen/iop"
)

// chatStreamFixed holds the values a live SSE chat stream resolves once at
// dispatch time and never changes while consuming events.
type chatStreamFixed struct {
	req             chatCompletionRequest
	dispatch        edgeservice.RunDispatch
	id              string
	created         int64
	model           string
	runID           string
	outputPolicy    strictOutputPolicy
	metricLabels    usageLabels
	traceStream     bool
	exposeReasoning bool
}

// chatStreamSession owns the accumulated state and the terminal emission of a
// single live SSE chat completion. Fixed request values stay in fixed; every
// mutable accumulator lives on the session, and exactly one terminal event
// (finish, fail, or cancel) may be emitted per session.
type chatStreamSession struct {
	server  *Server
	w       http.ResponseWriter
	flusher http.Flusher
	fixed   chatStreamFixed

	content   strings.Builder
	reasoning strings.Builder
	emitted   strings.Builder

	contentSentinel   *streamSentinelFilter
	reasoningSentinel *streamSentinelFilter
	toolTextFilter    *streamToolTextFilter

	traceSeq   int
	terminated bool
}

func (s *Server) newChatStreamSession(
	w http.ResponseWriter,
	flusher http.Flusher,
	req chatCompletionRequest,
	submitReq edgeservice.SubmitRunRequest,
	handle edgeservice.RunResult,
	outputPolicy strictOutputPolicy,
	metricLabels usageLabels,
) *chatStreamSession {
	sess := &chatStreamSession{
		server:  s,
		w:       w,
		flusher: flusher,
		fixed: chatStreamFixed{
			req:          req,
			dispatch:     handle.Dispatch(),
			id:           "chatcmpl-" + handle.Dispatch().RunID,
			created:      time.Now().Unix(),
			model:        responseModel(req.Model, handle.Dispatch().Target),
			runID:        handle.Dispatch().RunID,
			outputPolicy: outputPolicy,
			metricLabels: metricLabels,
			traceStream:  openAICompatTraceStreamEnabled(submitReq.Metadata),
			// Reasoning stays hidden on a strict stream unless the caller asked
			// for it explicitly.
			exposeReasoning: req.includeReasoning() && (!outputPolicy.Strict || req.explicitlyIncludesReasoning()),
		},
		contentSentinel:   &streamSentinelFilter{},
		reasoningSentinel: &streamSentinelFilter{},
	}
	if len(req.Tools) > 0 {
		sess.toolTextFilter = &streamToolTextFilter{}
	}
	return sess
}

// observation reports the usage counters that every non-complete terminal path
// emits.
func (sess *chatStreamSession) observation() usageObservation {
	return usageObservation{reasoningChars: sess.reasoning.Len()}
}

// terminate claims the session's single terminal emission. It returns false if
// a terminal event was already emitted, so no caller can write twice.
func (sess *chatStreamSession) terminate() bool {
	if sess.terminated {
		return false
	}
	sess.terminated = true
	return true
}

// fail ends the session with an SSE error and error usage metrics.
func (sess *chatStreamSession) fail(message string) {
	sess.failWithType("run_error", message)
}

func (sess *chatStreamSession) failWithType(errType, message string) {
	if !sess.terminate() {
		return
	}
	emitUsageMetrics(sess.fixed.metricLabels, usageStatusError, sess.observation())
	writeSSEErrorWithType(sess.w, sess.flusher, errType, message)
}

// failStream ends the session on a broken run stream. The stream never
// produced a usable run, so no usage sample is recorded for it.
func (sess *chatStreamSession) failStream(message string) {
	if !sess.terminate() {
		return
	}
	writeSSEError(sess.w, sess.flusher, message)
}

// cancel ends the session when the caller goes away or the run times out: the
// run is cancelled node-side and no SSE body is written for a caller that is no
// longer reading, except for the timeout notice the caller asked to wait for.
func (sess *chatStreamSession) cancel(err error, status string, notify bool) {
	if !sess.terminate() {
		return
	}
	sess.server.cancelRunOnHTTPGiveUp(sess.fixed.dispatch, err)
	emitUsageMetrics(sess.fixed.metricLabels, status, sess.observation())
	if notify {
		writeSSEError(sess.w, sess.flusher, err.Error())
	}
}

// writeRole emits the opening assistant role chunk.
func (sess *chatStreamSession) writeRole() {
	writeSSE(sess.w, sess.flusher, chatCompletionChunk{
		ID:      sess.fixed.id,
		Object:  "chat.completion.chunk",
		Created: sess.fixed.created,
		Model:   sess.fixed.model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})
}

// writeContent emits a content delta chunk and its trace record.
func (sess *chatStreamSession) writeContent(content, source string, filtered bool) {
	if content == "" {
		return
	}
	sess.traceSeq++
	if sess.fixed.traceStream {
		logOpenAICompatStreamOutput(sess.server.logger, sess.fixed.runID, sess.traceSeq, "content", content,
			zap.Int("source_len", len(source)),
			zap.Bool("filtered", filtered),
			zap.Bool("source_has_open_think", hasOpenThinkTag(source)),
			zap.Bool("source_has_close_think", hasCloseThinkTag(source)),
		)
	}
	writeContentDeltaSSE(sess.w, sess.flusher, sess.fixed.id, sess.fixed.created, sess.fixed.model, content)
}

// emitPending records and emits text held back by the tool-text filter.
func (sess *chatStreamSession) emitPending(pending string) {
	if pending == "" {
		return
	}
	sess.emitted.WriteString(pending)
	sess.writeContent(pending, pending, false)
}

func (sess *chatStreamSession) writeReasoning(reasoning string) {
	if reasoning == "" {
		return
	}
	sess.traceSeq++
	if sess.fixed.traceStream {
		logOpenAICompatStreamOutput(sess.server.logger, sess.fixed.runID, sess.traceSeq, "reasoning", reasoning)
	}
	writeSSE(sess.w, sess.flusher, chatCompletionChunk{
		ID:      sess.fixed.id,
		Object:  "chat.completion.chunk",
		Created: sess.fixed.created,
		Model:   sess.fixed.model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{ReasoningContent: reasoning},
		}},
	})
}

// emitDelta accumulates a source content delta and emits the part that
// survives sentinel and tool-text filtering.
func (sess *chatStreamSession) emitDelta(sourceDelta string) {
	if sourceDelta == "" {
		return
	}
	delta := sess.contentSentinel.Append(sourceDelta)
	if delta == "" {
		return
	}
	sess.content.WriteString(delta)
	filteredDelta := delta
	if sess.toolTextFilter != nil {
		filteredDelta = sess.toolTextFilter.Append(filteredDelta)
	}
	if filteredDelta == "" {
		if sess.fixed.traceStream {
			sess.server.logger.Info("openai chat completion stream output suppressed",
				zap.String("run_id", sess.fixed.runID),
				zap.String("type", "content"),
				zap.Int("source_len", len(sourceDelta)),
				zap.Bool("source_has_open_think", hasOpenThinkTag(sourceDelta)),
				zap.Bool("source_has_close_think", hasCloseThinkTag(sourceDelta)),
			)
		}
		return
	}
	sess.emitted.WriteString(filteredDelta)
	sess.writeContent(filteredDelta, sourceDelta, filteredDelta != sourceDelta)
}

// emitReasoningDelta accumulates a reasoning delta and emits it when the
// caller is allowed to see reasoning.
func (sess *chatStreamSession) emitReasoningDelta(sourceDelta string) {
	if sourceDelta == "" {
		return
	}
	reasoning := sess.reasoningSentinel.Append(sourceDelta)
	if reasoning == "" {
		return
	}
	sess.reasoning.WriteString(reasoning)
	if !sess.fixed.exposeReasoning {
		return
	}
	sess.writeReasoning(reasoning)
}

// consumeEvent applies one run event to the session and reports whether the
// session reached its terminal state.
func (sess *chatStreamSession) consumeEvent(event *iop.RunEvent) bool {
	if event == nil {
		return false
	}
	switch event.GetType() {
	case "delta":
		sess.emitDelta(event.GetDelta())
	case "reasoning_delta":
		sess.emitReasoningDelta(event.GetDelta())
	case "complete":
		sess.finish(event)
		return true
	case "error", "cancelled":
		msg := event.GetError()
		if msg == "" {
			msg = event.GetMessage()
		}
		if msg == "" {
			msg = "run failed"
		}
		sess.fail(msg)
		return true
	}
	return false
}

// streamUsage reads the provider-reported usage off a complete event.
func (sess *chatStreamSession) streamUsage(event *iop.RunEvent) usageObservation {
	obs := sess.observation()
	if u := event.GetUsage(); u != nil {
		obs.inputTokens = int(u.GetInputTokens())
		obs.outputTokens = int(u.GetOutputTokens())
		obs.reasoningTokens = int(u.GetReasoningTokens())
		obs.cachedInputTokens = int(u.GetCachedInputTokens())
	}
	return obs
}

// chatStreamToolOutcome reports how the completion's tool calls resolved.
type chatStreamToolOutcome struct {
	finishReason string
	nativeCount  int
	failed       bool
}

// finish emits the tool calls, reasoning fallback, and terminal chunk of a
// completed run.
func (sess *chatStreamSession) finish(event *iop.RunEvent) {
	sess.emitDelta(sess.contentSentinel.Flush())
	if reasoning := sess.reasoningSentinel.Flush(); reasoning != "" {
		sess.reasoning.WriteString(reasoning)
		if sess.fixed.exposeReasoning {
			sess.writeReasoning(reasoning)
		}
	}

	metadata := event.GetMetadata()
	usage := sess.streamUsage(event)

	outcome := sess.emitToolCalls(metadata)
	if outcome.failed {
		return
	}
	sess.emitReasoningFallback(&outcome)
	sess.writeTerminal(outcome.finishReason, usage)
}

// emitToolCalls resolves native or text-synthesized tool calls and emits the
// tool_calls delta. A tool validation failure terminates the session.
func (sess *chatStreamSession) emitToolCalls(metadata map[string]string) chatStreamToolOutcome {
	outcome := chatStreamToolOutcome{finishReason: metadata["finish_reason"]}
	if outcome.finishReason == "" {
		outcome.finishReason = "stop"
	}

	nativeToolCalls := toolCallsFromRunMetadata(metadata)
	switch {
	case len(nativeToolCalls) > 0:
		// A native tool call wins: drop any partial text candidate the filter
		// is holding and emit only the safe non-candidate prefix.
		if sess.toolTextFilter != nil {
			sess.emitPending(sess.toolTextFilter.FlushNonCandidate())
		}
		nativeToolCalls = normalizeNativeToolCallsForResponse(nativeToolCalls, sess.fixed.req.Tools)
		sess.traceToolCalls(len(nativeToolCalls))
		writeToolCallsDeltaSSE(sess.w, sess.flusher, sess.fixed.id, sess.fixed.created, sess.fixed.model, nativeToolCalls)
		outcome.finishReason = "tool_calls"
		outcome.nativeCount = len(nativeToolCalls)

	case len(sess.fixed.req.Tools) > 0:
		res := synthesizeToolCallsFromTextResult(sess.content.String(), sess.fixed.req.Tools, sess.fixed.runID)
		if res.validationErr != nil {
			sess.failWithType("tool_validation_error", res.validationErr.Error())
			outcome.failed = true
			return outcome
		}
		if len(res.toolCalls) > 0 {
			sess.emitPending(unstreamedCleanedText(res.cleaned, sess.emitted.String()))
			sess.traceToolCalls(len(res.toolCalls))
			writeToolCallsDeltaSSE(sess.w, sess.flusher, sess.fixed.id, sess.fixed.created, sess.fixed.model, res.toolCalls)
			outcome.finishReason = "tool_calls"
		} else if !res.candidateFound && sess.toolTextFilter != nil {
			sess.emitPending(sess.toolTextFilter.Flush())
		}

	case sess.toolTextFilter != nil:
		sess.emitPending(sess.toolTextFilter.Flush())
	}
	return outcome
}

func (sess *chatStreamSession) traceToolCalls(count int) {
	if !sess.fixed.traceStream {
		return
	}
	sess.server.logger.Info("openai chat completion stream output chunk",
		zap.String("run_id", sess.fixed.runID),
		zap.Int("seq", sess.traceSeq+1),
		zap.String("type", "tool_calls"),
		zap.Int("tool_call_count", count),
	)
}

// emitReasoningFallback emits a placeholder when a non-strict run produced only
// reasoning and no visible content or native tool call, so the caller never
// receives an empty completion.
func (sess *chatStreamSession) emitReasoningFallback(outcome *chatStreamToolOutcome) {
	if sess.fixed.outputPolicy.Strict ||
		outcome.nativeCount > 0 ||
		strings.TrimSpace(sess.content.String()) != "" ||
		strings.TrimSpace(sess.emitted.String()) != "" ||
		strings.TrimSpace(sess.reasoning.String()) == "" {
		return
	}
	fallback := hiddenReasoningFallbackContent(outcome.finishReason)
	if sess.fixed.req.includeReasoning() {
		fallback = reasoningOnlyFallbackContent(sess.reasoning.String(), outcome.finishReason)
	}
	sess.emitted.WriteString(fallback)
	sess.writeContent(fallback, sess.reasoning.String(), false)
}

// writeTerminal emits the finish chunk, the [DONE] sentinel, and success usage.
func (sess *chatStreamSession) writeTerminal(finishReason string, usage usageObservation) {
	if !sess.terminate() {
		return
	}
	if sess.fixed.traceStream {
		sess.server.logger.Info("openai chat completion stream output done",
			zap.String("run_id", sess.fixed.runID),
			zap.Int("seq", sess.traceSeq+1),
			zap.String("finish_reason", finishReason),
			zap.Int("content_len", sess.content.Len()),
			zap.Int("reasoning_len", sess.reasoning.Len()),
		)
	}
	writeSSE(sess.w, sess.flusher, chatCompletionChunk{
		ID:      sess.fixed.id,
		Object:  "chat.completion.chunk",
		Created: sess.fixed.created,
		Model:   sess.fixed.model,
		Choices: []chatCompletionChunkChoice{{
			Index:        0,
			Delta:        chatDelta{},
			FinishReason: finishReason,
		}},
	})
	fmt.Fprint(sess.w, "data: [DONE]\n\n")
	sess.flusher.Flush()
	emitUsageMetrics(sess.fixed.metricLabels, usageStatusSuccess, usage)
}

// logClosed records the final accumulated state of the session.
func (sess *chatStreamSession) logClosed() {
	sess.server.logger.Info("openai chat completion stream closed",
		zap.String("run_id", sess.fixed.runID),
		zap.Bool("strict_output", sess.fixed.outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", sess.fixed.outputPolicy.StreamBuffer),
		zap.Int("content_len", sess.content.Len()),
		zap.Int("reasoning_len", sess.reasoning.Len()),
	)
}