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

package openai

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"go.uber.org/zap"
	"iop/packages/go/streamgate"
)

// openAIStreamGateErrorMessage derives a caller-facing message from a Core
// TerminalResult. Core external descriptors carry only sanitized stable
// tokens (no raw provider text), so the code is the most specific safe value
// available.
func openAIStreamGateErrorMessage(tr streamgate.TerminalResult) string {
	desc := tr.ExternalDesc()
	if desc == nil {
		return "stream gate terminal error"
	}
	if code := desc.Code(); code != "" {
		return code
	}
	return desc.Type()
}

// openAIChatSSEReleaseSink implements streamgate.ReleaseSink for the
// normalized live-SSE chat completion path. It stages the HTTP status,
// SSE headers, and the opening assistant role chunk behind the Core's first
// safe release: CommitResponseStart is the only place status/header/role are
// written, and it is only called by CommitBoundary once a release or a
// success terminal is ready to commit.
type openAIChatSSEReleaseSink struct {
	w       http.ResponseWriter
	flusher http.Flusher
	id      string
	created int64
	model   string

	mu                sync.Mutex
	wroteHeader       bool
	terminalCommitted bool
	terminalSuccess   bool
}

func newOpenAIChatSSEReleaseSink(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string) *openAIChatSSEReleaseSink {
	return &openAIChatSSEReleaseSink{w: w, flusher: flusher, id: id, created: created, model: model}
}

// terminalStatus reports whether the Core committed a terminal through this
// sink and, if so, whether it was a success terminal. The host runner uses it
// to derive the terminal usage-metric status as the single source of truth for
// the response outcome.
func (s *openAIChatSSEReleaseSink) terminalStatus() (committed bool, success bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.terminalCommitted, s.terminalSuccess
}

func (s *openAIChatSSEReleaseSink) commitHeaderLocked(status int) {
	if s.wroteHeader {
		return
	}
	s.w.Header().Set("Content-Type", "text/event-stream")
	s.w.Header().Set("Cache-Control", "no-cache")
	s.w.Header().Set("Connection", "keep-alive")
	if status == 0 {
		status = http.StatusOK
	}
	s.w.WriteHeader(status)
	s.wroteHeader = true
	writeSSE(s.w, s.flusher, chatCompletionChunk{
		ID:      s.id,
		Object:  "chat.completion.chunk",
		Created: s.created,
		Model:   s.model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})
}

func (s *openAIChatSSEReleaseSink) CommitResponseStart(ctx context.Context, rs streamgate.ResponseStart) (streamgate.CommitState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.commitHeaderLocked(rs.Status())
	return streamgate.CommitStateStreamOpen, nil
}

func (s *openAIChatSSEReleaseSink) Release(ctx context.Context, ev streamgate.ReleaseEvent) (streamgate.CommitState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	// CommitBoundary always commits a staged response start before the first
	// Release call, so the header/role chunk is already flushed here.
	switch ev.Kind() {
	case streamgate.EventKindTextDelta:
		text, err := ev.AsTextDelta()
		if err != nil {
			return streamgate.CommitStateStreamOpen, err
		}
		writeContentDeltaSSE(s.w, s.flusher, s.id, s.created, s.model, text)
	case streamgate.EventKindReasoningDelta:
		reasoning, err := ev.AsReasoningDelta()
		if err != nil {
			return streamgate.CommitStateStreamOpen, err
		}
		writeSSE(s.w, s.flusher, chatCompletionChunk{
			ID:      s.id,
			Object:  "chat.completion.chunk",
			Created: s.created,
			Model:   s.model,
			Choices: []chatCompletionChunkChoice{{
				Index: 0,
				Delta: chatDelta{ReasoningContent: reasoning},
			}},
		})
	default:
		return streamgate.CommitStateStreamOpen, fmt.Errorf("openai stream gate: chat SSE sink does not support release event kind %q", ev.Kind())
	}
	return streamgate.CommitStateStreamOpen, nil
}

func (s *openAIChatSSEReleaseSink) CommitTerminal(ctx context.Context, tr streamgate.TerminalResult) (streamgate.CommitState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.terminalCommitted = true
	s.terminalSuccess = tr.Success()
	if tr.Success() {
		s.commitHeaderLocked(http.StatusOK)
		writeSSE(s.w, s.flusher, chatCompletionChunk{
			ID:      s.id,
			Object:  "chat.completion.chunk",
			Created: s.created,
			Model:   s.model,
			Choices: []chatCompletionChunkChoice{{
				Index:        0,
				Delta:        chatDelta{},
				FinishReason: "stop",
			}},
		})
		fmt.Fprint(s.w, "data: [DONE]\n\n")
		if s.flusher != nil {
			s.flusher.Flush()
		}
		return streamgate.CommitStateTerminalCommitted, nil
	}

	message := openAIStreamGateErrorMessage(tr)
	if !s.wroteHeader {
		writeError(s.w, http.StatusBadGateway, "run_error", message)
		s.wroteHeader = true
		return streamgate.CommitStateTerminalCommitted, nil
	}
	writeSSEErrorWithType(s.w, s.flusher, "run_error", message)
	return streamgate.CommitStateTerminalCommitted, nil
}

var _ streamgate.ReleaseSink = (*openAIChatSSEReleaseSink)(nil)

// openAITunnelReleaseSink implements streamgate.ReleaseSink for the raw
// provider tunnel passthrough path. It relays the provider's own
// status/headers verbatim (already sanitized of hop-by-hop/content-length
// fields by the event source) and writes release payloads as raw bytes with
// no IOP framing, preserving pure passthrough semantics.
type openAITunnelReleaseSink struct {
	w       http.ResponseWriter
	flusher http.Flusher

	// buffered is set for a non-streaming passthrough attempt: the provider
	// body only becomes rewritable once complete, so released bytes accumulate
	// and the caller-facing model echo rewrite runs exactly once at terminal.
	// A streaming attempt rewrites in provider byte order inside the event
	// source instead and writes each release straight through.
	buffered bool
	rewriter *providerModelRewriter

	mu                sync.Mutex
	wroteHeader       bool
	body              []byte
	terminalCommitted bool
	terminalSuccess   bool
}

func newOpenAITunnelReleaseSink(w http.ResponseWriter, flusher http.Flusher) *openAITunnelReleaseSink {
	return &openAITunnelReleaseSink{w: w, flusher: flusher}
}

// newOpenAIBufferedTunnelReleaseSink builds the non-streaming passthrough sink.
// requestModel is the caller-facing alias used to rewrite the provider model
// echo; an empty value keeps provider-original bytes.
func newOpenAIBufferedTunnelReleaseSink(w http.ResponseWriter, flusher http.Flusher, requestModel string) *openAITunnelReleaseSink {
	return &openAITunnelReleaseSink{
		w:        w,
		flusher:  flusher,
		buffered: true,
		rewriter: newProviderModelRewriter(false, requestModel),
	}
}

// terminalStatus reports whether the Core committed a terminal through this
// sink and, if so, whether it was a success terminal.
func (s *openAITunnelReleaseSink) terminalStatus() (committed bool, success bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.terminalCommitted, s.terminalSuccess
}

func (s *openAITunnelReleaseSink) CommitResponseStart(ctx context.Context, rs streamgate.ResponseStart) (streamgate.CommitState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.wroteHeader {
		return streamgate.CommitStateStreamOpen, nil
	}
	for k, v := range rs.Headers() {
		s.w.Header().Set(k, v)
	}
	status := rs.Status()
	if status == 0 {
		status = http.StatusOK
	}
	s.w.WriteHeader(status)
	s.wroteHeader = true
	if s.flusher != nil {
		s.flusher.Flush()
	}
	return streamgate.CommitStateStreamOpen, nil
}

func (s *openAITunnelReleaseSink) Release(ctx context.Context, ev streamgate.ReleaseEvent) (streamgate.CommitState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if ev.Kind() != streamgate.EventKindTextDelta {
		return streamgate.CommitStateStreamOpen, fmt.Errorf("openai stream gate: tunnel sink does not support release event kind %q", ev.Kind())
	}
	// Raw provider bytes travel through Core as an opaque text_delta payload;
	// Go strings are byte-exact so the round trip is lossless pure passthrough.
	text, err := ev.AsTextDelta()
	if err != nil {
		return streamgate.CommitStateStreamOpen, err
	}
	if s.buffered {
		s.body = append(s.body, text...)
		return streamgate.CommitStateStreamOpen, nil
	}
	if _, err := s.w.Write([]byte(text)); err != nil {
		return streamgate.CommitStateStreamOpen, err
	}
	if s.flusher != nil {
		s.flusher.Flush()
	}
	return streamgate.CommitStateStreamOpen, nil
}

func (s *openAITunnelReleaseSink) CommitTerminal(ctx context.Context, tr streamgate.TerminalResult) (streamgate.CommitState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.terminalCommitted = true
	s.terminalSuccess = tr.Success()
	if tr.Success() {
		if s.buffered {
			body := s.body
			s.body = nil
			if s.rewriter != nil {
				body = s.rewriter.RewriteComplete(body)
			}
			if len(body) > 0 {
				if _, err := s.w.Write(body); err != nil {
					return streamgate.CommitStateTerminalCommitted, err
				}
				if s.flusher != nil {
					s.flusher.Flush()
				}
			}
		}
		return streamgate.CommitStateTerminalCommitted, nil
	}
	s.body = nil
	if !s.wroteHeader {
		writeError(s.w, http.StatusBadGateway, "provider_tunnel_error", openAIStreamGateErrorMessage(tr))
		s.wroteHeader = true
		return streamgate.CommitStateTerminalCommitted, nil
	}
	// Status/headers are already committed; the stream ends truncated, matching
	// the legacy tunnel writer's post-commit error behavior.
	return streamgate.CommitStateTerminalCommitted, nil
}

var _ streamgate.ReleaseSink = (*openAITunnelReleaseSink)(nil)

// --- attempt codec selection (provider-pool path switch) ----------------------

// openAIStreamGateCodec names the response framing an attempt binding produces.
// It is derived from the actual admission result, never from the caller request
// or the rebuilder.
type openAIStreamGateCodec string

const (
	openAIStreamGateCodecNormalized openAIStreamGateCodec = "normalized"
	openAIStreamGateCodecTunnel     openAIStreamGateCodec = "tunnel"
)

func openAIStreamGateCodecForPath(path openAIAdmissionKind) openAIStreamGateCodec {
	if path == openAIAdmissionTunnel {
		return openAIStreamGateCodecTunnel
	}
	return openAIStreamGateCodecNormalized
}

// openAIStreamGateCodecSelector is the request-local handoff between the
// attempt event-source factory (which knows the actual admission path) and the
// composite release sink (which must expose exactly one framing). The factory
// updates it before the new attempt produces its first event; the sink freezes
// it at its first commit, so a provider-pool recovery may switch framing only
// while the transport is still uncommitted.
type openAIStreamGateCodecSelector struct {
	mu    sync.Mutex
	codec openAIStreamGateCodec
}

func newOpenAIStreamGateCodecSelector(initial openAIStreamGateCodec) *openAIStreamGateCodecSelector {
	return &openAIStreamGateCodecSelector{codec: initial}
}

func (s *openAIStreamGateCodecSelector) set(codec openAIStreamGateCodec) {
	s.mu.Lock()
	s.codec = codec
	s.mu.Unlock()
}

func (s *openAIStreamGateCodecSelector) get() openAIStreamGateCodec {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.codec
}

// openAIStreamGateSink is a release sink that also reports the terminal
// disposition the Core committed through it.
type openAIStreamGateSink interface {
	streamgate.ReleaseSink
	terminalStatus() (committed bool, success bool)
}

// openAICompositeReleaseSink delegates to the normalized or the raw tunnel sink
// for a provider-pool request whose actual execution path is only known after
// admission and may still change across a pre-commit recovery. The delegate is
// resolved from the codec selector at the first commit call and frozen from
// then on: the caller therefore observes exactly one framing, and a post-commit
// path change can never re-frame an already-open response.
type openAICompositeReleaseSink struct {
	selector   *openAIStreamGateCodecSelector
	normalized openAIStreamGateSink
	tunnel     openAIStreamGateSink

	mu     sync.Mutex
	active openAIStreamGateSink
	frozen openAIStreamGateCodec
}

func newOpenAICompositeReleaseSink(selector *openAIStreamGateCodecSelector, normalized, tunnel openAIStreamGateSink) *openAICompositeReleaseSink {
	return &openAICompositeReleaseSink{selector: selector, normalized: normalized, tunnel: tunnel}
}

func (s *openAICompositeReleaseSink) resolve() openAIStreamGateSink {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active != nil {
		return s.active
	}
	codec := s.selector.get()
	if codec == openAIStreamGateCodecTunnel {
		s.active = s.tunnel
	} else {
		s.active = s.normalized
	}
	s.frozen = codec
	return s.active
}

// resolvedCodec reports the framing actually exposed to the caller, or the
// currently selected one when nothing has been committed yet. The host uses it
// to attribute usage metrics to the right response mode.
func (s *openAICompositeReleaseSink) resolvedCodec() openAIStreamGateCodec {
	s.mu.Lock()
	frozen := s.frozen
	s.mu.Unlock()
	if frozen != "" {
		return frozen
	}
	return s.selector.get()
}

func (s *openAICompositeReleaseSink) CommitResponseStart(ctx context.Context, rs streamgate.ResponseStart) (streamgate.CommitState, error) {
	return s.resolve().CommitResponseStart(ctx, rs)
}

func (s *openAICompositeReleaseSink) Release(ctx context.Context, ev streamgate.ReleaseEvent) (streamgate.CommitState, error) {
	return s.resolve().Release(ctx, ev)
}

func (s *openAICompositeReleaseSink) CommitTerminal(ctx context.Context, tr streamgate.TerminalResult) (streamgate.CommitState, error) {
	return s.resolve().CommitTerminal(ctx, tr)
}

func (s *openAICompositeReleaseSink) terminalStatus() (bool, bool) {
	s.mu.Lock()
	active := s.active
	s.mu.Unlock()
	if active == nil {
		return false, false
	}
	return active.terminalStatus()
}

var (
	_ streamgate.ReleaseSink = (*openAICompositeReleaseSink)(nil)
	_ openAIStreamGateSink   = (*openAICompositeReleaseSink)(nil)
	_ openAIStreamGateSink   = (*openAIChatSSEReleaseSink)(nil)
	_ openAIStreamGateSink   = (*openAITunnelReleaseSink)(nil)
)

// --- buffered chat completion sink (buffered SSE and non-stream JSON) ---------

// openAIBufferedChatReleaseSink renders a fully assembled chat completion once
// the Core commits a terminal. The buffered paths publish no content events —
// the assembled output lives in the request-local result holder — so Release is
// evidence-only here and the authoritative render happens exactly once at
// terminal, reproducing the legacy writeBufferedStreamOutput / chat completion
// JSON bytes including structured tool_calls, content cleanup, finish reason,
// and usage.
type openAIBufferedChatReleaseSink struct {
	server       *Server
	w            http.ResponseWriter
	flusher      http.Flusher
	req          chatCompletionRequest
	outputPolicy strictOutputPolicy
	traceStream  bool
	// stream selects the buffered SSE framing; false renders the non-stream
	// chat.completion JSON object.
	stream bool
	holder *openAIBufferedResultHolder

	mu                sync.Mutex
	wroteHeader       bool
	terminalCommitted bool
	terminalSuccess   bool
}

func newOpenAIBufferedChatReleaseSink(
	server *Server,
	w http.ResponseWriter,
	flusher http.Flusher,
	dc *chatDispatchContext,
	stream bool,
	holder *openAIBufferedResultHolder,
) *openAIBufferedChatReleaseSink {
	return &openAIBufferedChatReleaseSink{
		server:       server,
		w:            w,
		flusher:      flusher,
		req:          dc.req,
		outputPolicy: dc.outputPolicy,
		traceStream:  openAICompatTraceStreamEnabled(dc.submitReq.Metadata),
		stream:       stream,
		holder:       holder,
	}
}

func (s *openAIBufferedChatReleaseSink) terminalStatus() (bool, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.terminalCommitted, s.terminalSuccess
}

// CommitResponseStart is a no-op for both buffered framings: the SSE headers
// are written together with the first rendered chunk at terminal, and the
// non-stream JSON response writes its status and body in one call. Nothing
// reaches the caller before the terminal verdict.
func (s *openAIBufferedChatReleaseSink) CommitResponseStart(ctx context.Context, rs streamgate.ResponseStart) (streamgate.CommitState, error) {
	return streamgate.CommitStateStreamOpen, nil
}

// Release is evidence-only: buffered attempts carry their payload in the result
// holder, so a released event never produces caller-visible bytes here.
func (s *openAIBufferedChatReleaseSink) Release(ctx context.Context, ev streamgate.ReleaseEvent) (streamgate.CommitState, error) {
	switch ev.Kind() {
	case streamgate.EventKindTextDelta, streamgate.EventKindReasoningDelta:
		return streamgate.CommitStateStreamOpen, nil
	default:
		return streamgate.CommitStateStreamOpen, fmt.Errorf("openai stream gate: buffered chat sink does not support release event kind %q", ev.Kind())
	}
}

func (s *openAIBufferedChatReleaseSink) CommitTerminal(ctx context.Context, tr streamgate.TerminalResult) (streamgate.CommitState, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.terminalCommitted = true
	s.terminalSuccess = tr.Success()

	result, ok := s.holder.get()
	if tr.Success() && ok && result.collectErr == nil && result.validErr == nil {
		s.renderSuccessLocked(result)
		return streamgate.CommitStateTerminalCommitted, nil
	}
	s.renderErrorLocked(tr, result, ok)
	return streamgate.CommitStateTerminalCommitted, nil
}

func (s *openAIBufferedChatReleaseSink) renderSuccessLocked(result openAIBufferedAttemptResult) {
	if s.stream {
		s.commitSSEHeaderLocked(http.StatusOK)
		s.server.writeBufferedStreamOutput(s.w, s.flusher, s.req, result.dispatch, result.output, s.outputPolicy, s.traceStream)
		return
	}
	s.server.logChatCompletionOutput(result.dispatch, result.output, s.outputPolicy)
	writeJSON(s.w, http.StatusOK, chatCompletionResponse{
		ID:      "chatcmpl-" + result.dispatch.RunID,
		Object:  "chat.completion",
		Created: time.Now().Unix(),
		Model:   responseModel(s.req.Model, result.dispatch.Target),
		Choices: []chatCompletionChoice{{
			Index:        0,
			Message:      result.output.message,
			FinishReason: result.output.finishReason,
		}},
		Usage: result.output.usage,
	})
	s.wroteHeader = true
}

// renderErrorLocked preserves the legacy OpenAI-compatible error envelope for
// the buffered paths: an exhausted tool-validation recovery still surfaces as
// tool_validation_error with the last validation reason, and any other terminal
// error keeps the run_error envelope.
func (s *openAIBufferedChatReleaseSink) renderErrorLocked(tr streamgate.TerminalResult, result openAIBufferedAttemptResult, ok bool) {
	errType := "run_error"
	message := openAIStreamGateErrorMessage(tr)
	switch {
	case ok && result.validErr != nil:
		errType = "tool_validation_error"
		message = result.validErr.Error()
		s.server.logger.Warn("openai chat completion tool validation failed",
			zap.String("run_id", result.dispatch.RunID),
			zap.String("reason", message),
		)
	case ok && result.collectErr != nil:
		message = result.collectErr.Error()
	}

	if s.stream {
		// The buffered SSE contract keeps a 200 event stream and reports the
		// failure as an SSE error event, matching the legacy writer.
		s.commitSSEHeaderLocked(http.StatusOK)
		writeSSEErrorWithType(s.w, s.flusher, errType, message)
		return
	}
	if s.wroteHeader {
		return
	}
	status := http.StatusBadGateway
	if ok && result.collectErr != nil {
		status = httpStatusForRunError(result.collectErr)
	}
	writeError(s.w, status, errType, message)
	s.wroteHeader = true
}

func (s *openAIBufferedChatReleaseSink) commitSSEHeaderLocked(status int) {
	if s.wroteHeader {
		return
	}
	s.w.Header().Set("Content-Type", "text/event-stream")
	s.w.Header().Set("Cache-Control", "no-cache")
	s.w.Header().Set("Connection", "keep-alive")
	if status == 0 {
		status = http.StatusOK
	}
	s.w.WriteHeader(status)
	s.wroteHeader = true
}

var (
	_ streamgate.ReleaseSink = (*openAIBufferedChatReleaseSink)(nil)
	_ openAIStreamGateSink   = (*openAIBufferedChatReleaseSink)(nil)
)

// openAIStreamGateResponseMode maps a resolved attempt codec to the usage
// metric response mode label.
func openAIStreamGateResponseMode(codec openAIStreamGateCodec) string {
	if codec == openAIStreamGateCodecTunnel {
		return responseModePassthrough
	}
	return responseModeNormalized
}