iop/apps/edge/internal/openai/stream_gate_runtime.go
toki 04879f2b43 feat(openai): 실제 provider별 사용량 귀속을 기록한다
요청 종료 계수와 실제 provider 시도 사용량을 분리하고, 직접·pool·retry 경로의 attribution을 보존한다. 관련 계약·스펙과 완료된 task archive 정리도 함께 반영한다.
2026-07-31 20:22:23 +09:00

1529 lines
59 KiB
Go

package openai
import (
"bytes"
"context"
"crypto/sha256"
"encoding/json"
"fmt"
"net/http"
"strings"
"sync"
"time"
"go.uber.org/zap"
edgeservice "iop/apps/edge/internal/service"
"iop/packages/go/config"
"iop/packages/go/streamgate"
iop "iop/proto/gen/iop"
)
const (
streamGateChannelDefault = "default"
streamGateEnvironment = config.StreamGateEnvironmentDev
streamGateConfigGeneration = "edge.vertical-slice.1"
streamGateNoopCapability = "core.always"
streamGateFilterTimeout = 5 * time.Second
streamGateErrorNodeDisconnect = "node_disconnected"
streamGateErrorStreamClosed = "stream_closed"
streamGateErrorRunFailed = "run_failed"
streamGateErrorTunnelClosed = "provider_tunnel_closed"
streamGateErrorTunnelFailed = "provider_tunnel_error"
)
// openAIStreamGateSafeToken derives a streamgate.StableToken-safe identifier
// ([a-z0-9][a-z0-9._:-]*) from a host-controlled prefix and an arbitrary raw
// value such as a run id. Any character outside the grammar is replaced with
// '-' so an unexpected run id shape never fails Core validation.
func openAIStreamGateSafeToken(prefix, raw string) string {
var b strings.Builder
b.WriteString(prefix)
b.WriteByte('.')
lower := strings.ToLower(strings.TrimSpace(raw))
if lower == "" {
b.WriteString("unknown")
return b.String()
}
for _, r := range lower {
switch {
case r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '.', r == '_', r == ':', r == '-':
b.WriteRune(r)
default:
b.WriteByte('-')
}
}
safe := b.String()
if len(safe) <= 128 {
return safe
}
digest := sha256.Sum256([]byte(raw))
suffix := fmt.Sprintf("%x", digest[:8])
return safe[:128-1-len(suffix)] + "-" + suffix
}
// newOpenAIProviderErrorEvent builds a raw-free provider-error normalized
// event from a stable, sanitized code. Raw provider/error text is
// intentionally never carried into the Core event contract.
func newOpenAIProviderErrorEvent(code string) (streamgate.NormalizedEvent, error) {
desc, err := streamgate.NewExternalDescriptor("provider_error", code, code, "")
if err != nil {
return streamgate.NormalizedEvent{}, err
}
cause, err := streamgate.NewFailureCause("openai_event_source", code, "", "", "")
if err != nil {
return streamgate.NormalizedEvent{}, err
}
causes, err := streamgate.NewFailureCauseChain([]streamgate.FailureCause{cause})
if err != nil {
return streamgate.NormalizedEvent{}, err
}
return streamgate.NewProviderErrorEvent(streamGateChannelDefault, desc, causes, time.Now())
}
// openAIStreamGateUsageHolder carries the final attempt observation used by
// response renderers. Provider metrics use the separate per-attempt owner and
// never discard an aborted attempt when recovery replaces it.
type openAIStreamGateUsageHolder struct {
mu sync.Mutex
obs usageObservation
}
func (h *openAIStreamGateUsageHolder) set(obs usageObservation) {
h.mu.Lock()
h.obs = obs
h.mu.Unlock()
}
func (h *openAIStreamGateUsageHolder) get() usageObservation {
h.mu.Lock()
defer h.mu.Unlock()
return h.obs
}
// --- Normalized RunEvent -> NormalizedEvent source -------------------------
// openAIRunEventSource adapts an edgeservice.RunStream to
// streamgate.NormalizedEventSource. Its first call synthesizes a
// response-start event so Core can stage it before any real content; the
// synthesized start carries no side effects until the Core commits it
// through ReleaseSink.
type openAIRunEventSource struct {
stream edgeservice.RunStream
waitTimeout time.Duration
usage *openAIStreamGateUsageHolder
attempt *openAIAttemptUsage
mu sync.Mutex
startSent bool
}
func newOpenAIRunEventSource(stream edgeservice.RunStream, waitTimeout time.Duration, usage *openAIStreamGateUsageHolder, attempts ...*openAIAttemptUsage) *openAIRunEventSource {
source := &openAIRunEventSource{stream: stream, waitTimeout: waitTimeout, usage: usage}
if len(attempts) > 0 {
source.attempt = attempts[0]
}
return source
}
func (s *openAIRunEventSource) NextEvent(ctx context.Context) (streamgate.NormalizedEvent, error) {
s.mu.Lock()
sendStart := !s.startSent
s.startSent = true
s.mu.Unlock()
if sendStart {
return streamgate.NewResponseStartEvent(streamGateChannelDefault, http.StatusOK, openAIStreamGateChatSSEHeaders(), time.Now())
}
if s.stream.Events == nil {
return newOpenAIProviderErrorEvent(streamGateErrorStreamClosed)
}
timer := time.NewTimer(s.waitTimeout)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return streamgate.NormalizedEvent{}, ctx.Err()
case <-timer.C:
return streamgate.NormalizedEvent{}, errRunTimedOut
case nodeEvent, ok := <-s.stream.NodeEvents:
if !ok {
s.stream.NodeEvents = nil
continue
}
if edgeservice.IsNodeDisconnected(nodeEvent) {
return newOpenAIProviderErrorEvent(streamGateErrorNodeDisconnect)
}
case event, ok := <-s.stream.Events:
if !ok {
return newOpenAIProviderErrorEvent(streamGateErrorStreamClosed)
}
if event == nil {
continue
}
switch event.GetType() {
case "delta":
if event.GetDelta() == "" {
continue
}
return streamgate.NewTextDeltaEvent(streamGateChannelDefault, event.GetDelta(), time.Now())
case "reasoning_delta":
if event.GetDelta() == "" {
continue
}
s.attempt.addReasoningChars(len(event.GetDelta()))
return streamgate.NewReasoningDeltaEvent(streamGateChannelDefault, event.GetDelta(), time.Now())
case "complete":
obs := runEventUsageObservation(event)
s.attempt.observe(obs)
if s.usage != nil {
s.usage.set(obs)
}
return streamgate.NewTerminalEvent(streamGateChannelDefault, time.Now())
case "error", "cancelled":
return newOpenAIProviderErrorEvent(streamGateErrorRunFailed)
default:
continue
}
}
}
}
// --- Buffered chat completion -> NormalizedEvent source ---------------------
// openAIBufferedChatEventSource adapts one buffered chat attempt to
// streamgate.NormalizedEventSource. The buffered surfaces (strict/tool-bearing
// SSE and non-stream JSON) assemble their whole response before anything may be
// shown, so this source publishes no content events at all: it stages a
// response start, runs the existing collectChatCompletionOutput exactly once,
// publishes the assembled result plus the existing validateToolCallResponse
// verdict into the request-local holder, and then emits the terminal event that
// triggers the terminal-gate consumer. Holding nothing keeps the whole
// (arbitrarily large) assembled response out of the Core evidence buffer while
// still giving the consumer a single decision point before any commit.
type openAIBufferedChatEventSource struct {
req chatCompletionRequest
outputPolicy strictOutputPolicy
validation toolValidationContract
handle edgeservice.RunResult
holder *openAIBufferedResultHolder
usage *openAIStreamGateUsageHolder
attempt *openAIAttemptUsage
mu sync.Mutex
stage int
}
func newOpenAIBufferedChatEventSource(dc *chatDispatchContext, handle edgeservice.RunResult, holder *openAIBufferedResultHolder, usage *openAIStreamGateUsageHolder, attempts ...*openAIAttemptUsage) *openAIBufferedChatEventSource {
source := &openAIBufferedChatEventSource{
req: dc.req,
outputPolicy: dc.outputPolicy,
validation: dc.validation,
handle: handle,
holder: holder,
usage: usage,
}
if len(attempts) > 0 {
source.attempt = attempts[0]
}
return source
}
func (s *openAIBufferedChatEventSource) NextEvent(ctx context.Context) (streamgate.NormalizedEvent, error) {
s.mu.Lock()
stage := s.stage
s.stage++
s.mu.Unlock()
switch stage {
case 0:
s.holder.beginAttempt()
return streamgate.NewResponseStartEvent(streamGateChannelDefault, http.StatusOK, openAIStreamGateChatSSEHeaders(), time.Now())
case 1:
result, err := collectChatCompletionOutput(ctx, s.req, s.handle, s.outputPolicy)
if err != nil {
if ctx.Err() != nil {
return streamgate.NormalizedEvent{}, ctx.Err()
}
s.holder.set(openAIBufferedAttemptResult{dispatch: s.handle.Dispatch(), collectErr: err})
return newOpenAIProviderErrorEvent(streamGateErrorRunFailed)
}
verr := result.toolValidationErr
if verr == nil {
verr = validateToolCallResponse(s.validation, result.toolCalls, result.toolCallOrigin)
}
s.holder.set(openAIBufferedAttemptResult{output: result, dispatch: s.handle.Dispatch(), validErr: verr})
obs := usageObservationFromOpenAIUsage(result.usage, result.reasoningLen)
s.attempt.observe(obs)
if s.usage != nil {
s.usage.set(obs)
}
return streamgate.NewTerminalEvent(streamGateChannelDefault, time.Now())
default:
return newOpenAIProviderErrorEvent(streamGateErrorStreamClosed)
}
}
var _ streamgate.NormalizedEventSource = (*openAIBufferedChatEventSource)(nil)
func runEventUsageObservation(event *iop.RunEvent) usageObservation {
var obs usageObservation
if u := event.GetUsage(); u != nil {
obs.providerReported = true
obs.inputTokens = int(u.GetInputTokens())
obs.outputTokens = int(u.GetOutputTokens())
obs.reasoningTokens = int(u.GetReasoningTokens())
obs.cachedInputTokens = int(u.GetCachedInputTokens())
}
return obs
}
func openAIStreamGateChatSSEHeaders() map[string]string {
return map[string]string{"Content-Type": "text/event-stream", "Cache-Control": "no-cache"}
}
var _ streamgate.NormalizedEventSource = (*openAIRunEventSource)(nil)
// --- ProviderTunnelFrame -> NormalizedEvent source --------------------------
// streamGateForbiddenResponseHeaders mirrors the Core's forbidden
// response-start header set (hop-by-hop and transport metadata headers must
// never enter the event contract).
var streamGateForbiddenResponseHeaders = map[string]struct{}{
"connection": {},
"keep-alive": {},
"proxy-authenticate": {},
"proxy-authorization": {},
"proxy-connection": {},
"te": {},
"trailer": {},
"transfer-encoding": {},
"upgrade": {},
"content-length": {},
}
func sanitizedTunnelResponseHeaders(headers map[string]string) map[string]string {
if len(headers) == 0 {
return nil
}
out := make(map[string]string, len(headers))
for k, v := range headers {
if _, forbidden := streamGateForbiddenResponseHeaders[strings.ToLower(k)]; forbidden {
continue
}
out[k] = v
}
if len(out) == 0 {
return nil
}
return out
}
// openAITunnelEventSource adapts a ProviderTunnelStream to endpoint-semantic
// Core events. The endpoint codec retains caller-facing provider frames in a
// request-local release queue so Chat/Responses parsing and lossless passthrough
// coexist. The model echo rewrite still runs exactly once in provider order.
type openAITunnelEventSource struct {
frames <-chan *iop.ProviderTunnelFrame
waitTimeout time.Duration
rewriter *providerModelRewriter
assembler *providerChatAssembler
codec *openAITunnelEndpointCodec
responseStatus int
mu sync.Mutex
started bool
pending []streamgate.NormalizedEvent
}
func newOpenAITunnelEventSource(stream edgeservice.ProviderTunnelStream, waitTimeout time.Duration, rewriter *providerModelRewriter, assembler *providerChatAssembler) *openAITunnelEventSource {
return &openAITunnelEventSource{frames: stream.Frames, waitTimeout: waitTimeout, rewriter: rewriter, assembler: assembler}
}
func newOpenAITunnelEndpointEventSource(stream edgeservice.ProviderTunnelStream, waitTimeout time.Duration, rewriter *providerModelRewriter, assembler *providerChatAssembler, endpoint string, state *openAITunnelCodecState) *openAITunnelEventSource {
source := newOpenAITunnelEventSource(stream, waitTimeout, rewriter, assembler)
source.codec = newOpenAITunnelEndpointCodec(endpoint, state)
return source
}
func (s *openAITunnelEventSource) NextEvent(ctx context.Context) (streamgate.NormalizedEvent, error) {
s.mu.Lock()
if len(s.pending) > 0 {
ev := s.pending[0]
s.pending = s.pending[1:]
s.mu.Unlock()
return ev, nil
}
s.mu.Unlock()
if s.frames == nil {
return newOpenAIProviderErrorEvent(streamGateErrorTunnelClosed)
}
timer := time.NewTimer(s.waitTimeout)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return streamgate.NormalizedEvent{}, ctx.Err()
case <-timer.C:
return streamgate.NormalizedEvent{}, errRunTimedOut
case frame, ok := <-s.frames:
if !ok {
return newOpenAIProviderErrorEvent(streamGateErrorTunnelClosed)
}
events, err := s.translateFrame(frame)
if err != nil {
return streamgate.NormalizedEvent{}, err
}
if len(events) == 0 {
continue
}
first := events[0]
if len(events) > 1 {
s.mu.Lock()
s.pending = append(s.pending, events[1:]...)
s.mu.Unlock()
}
return first, nil
}
}
}
func (s *openAITunnelEventSource) markStarted() bool {
s.mu.Lock()
defer s.mu.Unlock()
if s.started {
return true
}
s.started = true
return false
}
func (s *openAITunnelEventSource) translateFrame(frame *iop.ProviderTunnelFrame) ([]streamgate.NormalizedEvent, error) {
switch frame.GetKind() {
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
if s.markStarted() {
return nil, nil
}
status := int(frame.GetStatusCode())
if status == 0 {
status = http.StatusOK
}
s.responseStatus = status
headers := sanitizedTunnelResponseHeaders(frame.GetHeaders())
if s.codec != nil && status >= http.StatusBadRequest {
s.codec.state.stageErrorResponseStart(status, headers)
}
ev, err := streamgate.NewResponseStartEvent(streamGateChannelDefault, status, headers, time.Now())
if err != nil {
return nil, err
}
return []streamgate.NormalizedEvent{ev}, nil
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
body := frame.GetBody()
if len(body) == 0 {
return nil, nil
}
if s.codec != nil && s.responseStatus >= http.StatusBadRequest {
// A non-2xx body is opaque provider wire even when it resembles a
// successful Chat/Responses payload. It is committed only if this
// attempt's provider-error terminal wins final arbitration.
s.codec.state.appendErrorResponseWire(body)
return nil, nil
}
if s.assembler != nil {
s.assembler.Write(body)
}
rewritten := body
if s.rewriter != nil {
rewritten = s.rewriter.AppendStream(body)
}
var events []streamgate.NormalizedEvent
if !s.markStarted() {
ev, err := streamgate.NewResponseStartEvent(streamGateChannelDefault, http.StatusOK, nil, time.Now())
if err != nil {
return nil, err
}
events = append(events, ev)
s.responseStatus = http.StatusOK
}
if len(rewritten) == 0 {
return events, nil
}
if s.codec != nil {
decoded, err := s.codec.decode(rewritten, false)
if err != nil {
return nil, err
}
return append(events, decoded...), nil
}
ev, err := streamgate.NewTextDeltaEvent(streamGateChannelDefault, string(rewritten), time.Now())
if err != nil {
return nil, err
}
return append(events, ev), nil
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
if s.assembler != nil {
s.assembler.recordProtoUsage(frame.GetUsage())
}
return nil, nil
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
ev, err := newOpenAIProviderErrorEvent(streamGateErrorTunnelFailed)
if err != nil {
return nil, err
}
return []streamgate.NormalizedEvent{ev}, nil
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
var events []streamgate.NormalizedEvent
if !s.markStarted() {
ev, err := streamgate.NewResponseStartEvent(streamGateChannelDefault, http.StatusOK, nil, time.Now())
if err != nil {
return nil, err
}
events = append(events, ev)
s.responseStatus = http.StatusOK
}
var flushed []byte
if s.rewriter != nil {
flushed = s.rewriter.FlushStream()
}
if s.codec != nil {
decoded, err := s.codec.finishTransport(flushed, s.responseStatus >= http.StatusBadRequest)
if err != nil {
return nil, err
}
events = append(events, decoded...)
if len(decoded) > 0 {
return events, nil
}
} else if len(flushed) > 0 {
ev, err := streamgate.NewTextDeltaEvent(streamGateChannelDefault, string(flushed), time.Now())
if err != nil {
return nil, err
}
events = append(events, ev)
}
term, err := streamgate.NewTerminalEvent(streamGateChannelDefault, time.Now())
if err != nil {
return nil, err
}
return append(events, term), nil
default:
return nil, nil
}
}
var _ streamgate.NormalizedEventSource = (*openAITunnelEventSource)(nil)
// --- Filter registry (production: Noop only) --------------------------------
// openAIStreamGateNoopRegistrations returns the production filter set for the
// vertical slice: a single always-applicable, non-holding Noop filter. Real
// semantic filters (repeat/malformed/schema/tool-syntax) belong to the
// consumer Milestones and are intentionally absent here.
func openAIStreamGateNoopRegistrations() ([]streamgate.FilterRegistration, error) {
noop, err := streamgate.NewNoopFilter()
if err != nil {
return nil, err
}
reg, err := streamgate.NewFilterRegistration(noop, streamGateNoopCapability, true, streamgate.FilterEnforcementBlocking, streamGateFilterTimeout, 0)
if err != nil {
return nil, err
}
return []streamgate.FilterRegistration{reg}, nil
}
func openAIStreamGateRegistrySnapshot() (streamgate.FilterRegistrySnapshot, error) {
return openAIStreamGateRegistrySnapshotWith()
}
// openAIStreamGateRegistrySnapshotFor builds the production registry snapshot for
// one request: the always-applicable Noop mechanics filter, the configured
// semantic output filters (repeat/schema/provider-error) translated from the
// stream_evidence_gate policy, plus any request-local extra registrations (e.g.
// the tool-validation terminal gate). An empty Filters policy reduces to the
// legacy Noop+extra set exactly, so the default production behavior is unchanged.
func openAIStreamGateRegistrySnapshotFor(gateCfg config.StreamEvidenceGateConf, fctx openAIOutputFilterContext, extra ...streamgate.FilterRegistration) (streamgate.FilterRegistrySnapshot, error) {
regs, err := openAIStreamGateNoopRegistrations()
if err != nil {
return streamgate.FilterRegistrySnapshot{}, err
}
outputRegs, policies, err := openAIOutputFilterRegistrations(gateCfg, fctx)
if err != nil {
return streamgate.FilterRegistrySnapshot{}, err
}
regs = append(regs, outputRegs...)
regs = append(regs, extra...)
return streamgate.NewFilterRegistrySnapshot(streamGateConfigGeneration, regs, policies)
}
// streamGateConfig returns a copy of the request-stable stream-gate config the
// request runtime pins at request start (generation isolation).
func (s *Server) streamGateConfig() config.StreamEvidenceGateConf {
s.mu.RLock()
defer s.mu.RUnlock()
return s.cfg.StreamEvidenceGate
}
// openAIChatOutputFilterContext resolves the immutable request-start selector,
// model-group, endpoint, and schema facts for a chat request.
func (s *Server) openAIChatOutputFilterContext(dc *chatDispatchContext) (openAIOutputFilterContext, error) {
snapRef, err := dc.ingress.recoveryRef()
if err != nil {
return openAIOutputFilterContext{}, err
}
body, err := dc.ingress.canonicalBody()
if err != nil {
return openAIOutputFilterContext{}, err
}
history, err := decodeOpenAIChatRepeatHistory(body)
if err != nil {
return openAIOutputFilterContext{}, err
}
return openAIOutputFilterContext{
environment: s.streamGateConfig().EffectiveEnvironment(),
endpoint: openAIRebuildEndpointChat,
modelGroup: strings.TrimSpace(dc.req.Model),
hasScheme: chatRequestHasSchemeMetadata(dc.req.Metadata),
requestRef: snapRef.SnapshotRef(),
history: history,
}, nil
}
// openAIResponsesOutputFilterContext resolves the same request-local policy
// facts for a Responses request. The endpoint-specific lossless Rebuilder is
// still the canonical snapshot owner; only the semantic filter selection is
// shared with Chat.
func (s *Server) openAIResponsesOutputFilterContext(requestCtx *responsesRequestContext) (openAIOutputFilterContext, error) {
snapRef, err := requestCtx.ingress.recoveryRef()
if err != nil {
return openAIOutputFilterContext{}, err
}
body, err := requestCtx.ingress.canonicalBody()
if err != nil {
return openAIOutputFilterContext{}, err
}
history, err := decodeOpenAIResponsesRepeatHistory(body)
if err != nil {
return openAIOutputFilterContext{}, err
}
return openAIOutputFilterContext{
environment: s.streamGateConfig().EffectiveEnvironment(),
endpoint: openAIRebuildEndpointResponses,
modelGroup: strings.TrimSpace(requestCtx.envelope.Model),
hasScheme: chatRequestHasSchemeMetadata(requestCtx.envelope.Metadata),
requestRef: snapRef.SnapshotRef(),
history: history,
}, nil
}
// openAITunnelOutputFilterContext preserves the same selector, model-group,
// endpoint, and schema facts when the actual execution path is a tunnel.
func (s *Server) openAITunnelOutputFilterContext(req openAITunnelStreamGateRequest) (openAIOutputFilterContext, error) {
snapRef, err := req.ingress.recoveryRef()
if err != nil {
return openAIOutputFilterContext{}, err
}
body, err := req.ingress.canonicalBody()
if err != nil {
return openAIOutputFilterContext{}, err
}
var history openAIRepeatHistorySnapshot
switch req.endpoint {
case openAIRebuildEndpointResponses:
history, err = decodeOpenAIResponsesRepeatHistory(body)
default:
history, err = decodeOpenAIChatRepeatHistory(body)
}
if err != nil {
return openAIOutputFilterContext{}, err
}
return openAIOutputFilterContext{
environment: s.streamGateConfig().EffectiveEnvironment(),
modelGroup: req.modelGroupKey,
endpoint: req.endpoint,
hasScheme: req.hasScheme,
requestRef: snapRef.SnapshotRef(),
history: history,
}, nil
}
// openAIStreamGateRegistrySnapshotWith builds the production registry snapshot
// plus request-local registrations. The only production request-local filter is
// the tool-validation terminal gate, which needs this request's result holder;
// it is registered only when the request actually carries a tool-validation
// contract.
func openAIStreamGateRegistrySnapshotWith(extra ...streamgate.FilterRegistration) (streamgate.FilterRegistrySnapshot, error) {
regs, err := openAIStreamGateNoopRegistrations()
if err != nil {
return streamgate.FilterRegistrySnapshot{}, err
}
regs = append(regs, extra...)
return streamgate.NewFilterRegistrySnapshot(streamGateConfigGeneration, regs, nil)
}
// openAIToolValidationRegistrations returns the request-local tool-validation
// registration for a buffered/non-stream chat request, or nil when the request
// has no tool-validation contract.
func openAIToolValidationRegistrations(dc *chatDispatchContext, holder *openAIBufferedResultHolder) ([]streamgate.FilterRegistration, error) {
if !dc.validation.enabled || holder == nil {
return nil, nil
}
snapRef, err := dc.ingress.recoveryRef()
if err != nil {
return nil, err
}
filter, err := newOpenAIToolValidationFilter(holder, snapRef.SnapshotRef())
if err != nil {
return nil, err
}
reg, err := streamgate.NewFilterRegistration(
filter, streamGateNoopCapability, true,
streamgate.FilterEnforcementBlocking, streamGateFilterTimeout, openAIToolValidationPriority,
)
if err != nil {
return nil, err
}
return []streamgate.FilterRegistration{reg}, nil
}
// --- Runtime options ---------------------------------------------------------
func (s *Server) streamGateRuntimeOptions() (streamgate.RuntimeOptions, error) {
s.mu.RLock()
gateCfg := s.cfg.StreamEvidenceGate
s.mu.RUnlock()
opts, err := streamgate.NewRuntimeOptions(
streamgate.DefaultMaxEvidenceRunes,
streamgate.DefaultMaxBufferRunes,
int64(gateCfg.EffectiveMaxIngressSnapshotBytes()),
gateCfg.EffectiveMaxRequestFaultRecovery(),
streamgate.GateCoordinatorOptions{},
streamgate.RecoveryCoordinatorOptions{},
)
if err != nil {
return streamgate.RuntimeOptions{}, err
}
// An omitted strategy cap inherits the request-total cap for every fault
// strategy; an explicit value (including 0) applies uniformly to each fault
// strategy, so an explicit 0 forbids strategy-level fault recovery entirely.
if gateCfg.MaxStrategyFaultRecovery != nil {
strategyCap := gateCfg.EffectiveMaxStrategyFaultRecovery()
opts, err = opts.WithRecoveryStrategyLimits(map[streamgate.RecoveryStrategy]int{
streamgate.RecoveryStrategyExactReplay: strategyCap,
streamgate.RecoveryStrategyContinuationRepair: strategyCap,
streamgate.RecoveryStrategySchemaRepair: strategyCap,
})
if err != nil {
return streamgate.RuntimeOptions{}, err
}
}
return opts, nil
}
// streamGateEnabled reports whether the request runtime should own this
// request's response lifecycle. Disabled (default) always uses the legacy
// eager-write path unchanged.
func (s *Server) streamGateEnabled() bool {
s.mu.RLock()
defer s.mu.RUnlock()
return s.cfg.StreamEvidenceGate.Enabled
}
// streamGateUsageStatus derives the terminal usage-metric status from the run
// outcome. A non-nil Run error is a caller give-up (cancel/timeout) or an
// internal error; otherwise the committed terminal disposition captured by the
// release sink is authoritative, so a Core-committed error terminal is never
// mislabeled as success.
func streamGateUsageStatus(runErr error, terminalCommitted, terminalSuccess bool) string {
if runErr != nil {
return usageStatusForError(runErr)
}
if terminalCommitted && terminalSuccess {
return usageStatusSuccess
}
return usageStatusError
}
// --- Chat (normalized) runtime wiring ---------------------------------------
// newOpenAIChatRecoveryAdmissionBuilder rebuilds a chat completion admission
// from the Core-rebuilt canonical body. It replays the same catalog generation
// policy and dispatch-context construction ingress used, so an exact-replay
// recovery attempt is admitted identically to the original request.
//
// The admission shape follows the route the request was dispatched on: a
// direct model_routes request re-enters SubmitRun, and a provider-pool request
// re-enters SubmitProviderPool with the same tunnel template, auth preparation,
// and normalized preparation the initial admission used. The pool therefore
// re-selects a candidate for every recovery attempt, which is what lets the
// actual execution path change inside one request runtime.
func newOpenAIChatRecoveryAdmissionBuilder(s *Server, dc *chatDispatchContext, holder *openAIBufferedResultHolder) openAIAttemptAdmissionBuilder {
requestCtx := dc.openAIRequestContext
poolTemplate := dc.poolDispatch
return func(ctx context.Context, request streamgate.RebuiltRequest, body []byte) (openAIAttemptAdmission, error) {
var req chatCompletionRequest
decoder := json.NewDecoder(bytes.NewReader(body))
var decodeErr error
if poolTemplate != nil {
// The provider-pool ingress decodes leniently so caller unknown and
// provider-specific fields survive into the tunnel body; recovery
// must decode the same way.
decodeErr = decodeChatCompletionRequestLenient(decoder, &req)
} else {
decodeErr = decodeChatCompletionRequest(decoder, &req)
}
if decodeErr != nil {
return openAIAttemptAdmission{}, fmt.Errorf("openai stream gate: recovery chat decode: %w", decodeErr)
}
basePrompt := promptFromMessages(req.Messages)
if strings.TrimSpace(basePrompt) == "" {
return openAIAttemptAdmission{}, fmt.Errorf("openai stream gate: recovery chat request has no messages")
}
outputPolicy := s.resolveOutputPolicy(basePrompt)
if catalogEntry := s.findProviderPoolEntry(req.Model); catalogEntry != nil {
providerNativeThinking := chatRequestHasProviderNativeThinking(body)
applyModelCatalogGenerationPolicyToChat(&req, *catalogEntry, outputPolicy.Strict, providerNativeThinking)
}
rebuiltDC := s.newChatDispatchContext(requestCtx, req, basePrompt, outputPolicy)
runReq := rebuiltDC.submitReq
if holder != nil {
// Preserve the legacy provider-visible retry annotation: the Core
// owns the recovery budget, this only reproduces the attempt number,
// failed run id, and reason the provider used to see.
attempt, retryOf, reason := holder.nextAttemptMetadata()
runReq.Metadata = toolValidationAttemptMetadata(runReq.Metadata, attempt, retryOf, reason)
}
if poolTemplate == nil {
return openAIAttemptAdmission{kind: openAIAdmissionRun, run: runReq}, nil
}
tunnelReq := poolTemplate.Tunnel
tunnelReq.Metadata = runReq.Metadata
tunnelReq.BuildBody = func(target string) ([]byte, error) {
var decoded chatCompletionRequest
if err := decodeChatCompletionRequestLenient(json.NewDecoder(bytes.NewReader(body)), &decoded); err != nil {
return nil, err
}
return rewriteChatCompletionModel(body, target, decoded)
}
return openAIAttemptAdmission{
kind: openAIAdmissionPool,
pool: edgeservice.ProviderPoolDispatchRequest{
Run: runReq,
Tunnel: tunnelReq,
PrepareTunnel: poolTemplate.PrepareTunnel,
PrepareRun: poolTemplate.PrepareRun,
AcceptCandidate: poolTemplate.AcceptCandidate,
},
}, nil
}
}
// openAIChatStreamGateMode selects how a runtime-enabled chat attempt on the
// normalized execution path is turned into events and rendered.
type openAIChatStreamGateMode int
const (
// openAIChatGateModeLive streams normalized run deltas as they arrive.
openAIChatGateModeLive openAIChatStreamGateMode = iota
// openAIChatGateModeBuffered assembles the whole completion before the
// terminal verdict; it renders buffered SSE or non-stream JSON.
openAIChatGateModeBuffered
)
// openAIChatStreamGateConfig is the fully resolved wiring of one runtime-enabled
// chat request: the framing of a normalized attempt, the already-dispatched
// initial transport, the release sink, and the attempt codec selector shared
// with the event-source factory.
type openAIChatStreamGateConfig struct {
mode openAIChatStreamGateMode
initial openAIAttemptTransport
dispatch edgeservice.RunDispatch
closeAll func()
sink streamgate.ReleaseSink
selector *openAIStreamGateCodecSelector
registry streamgate.FilterRegistrySnapshot
holder *openAIBufferedResultHolder
// recoverySource is the request-local model-output recorder shared by every
// attempt's event source, so a repeat-loop continuation plan can be assembled
// from the aborted attempt's own content/reasoning.
recoverySource *openAIRecoverySourceStore
// preparer and prepFactory are the optional one-shot host preparation seam
// Core calls after attempt ownership is closed and before the rebuild. The
// OpenAI surfaces have no production preparer in this slice, so both stay
// nil there; Core requires them to be set or unset together.
preparer streamgate.RecoveryPlanPreparer
prepFactory streamgate.RecoveryPreparationSnapshotFactory
obsSink streamgate.ObservationSink
}
// newOpenAIChatAttemptEventSourceFactory builds the dual event-source factory.
// It is the single place where an attempt's actual admission path is turned
// into both an event source and the response framing: the codec selector is
// updated here, before the new attempt produces any event and therefore before
// the composite sink can freeze it.
func (s *Server) newOpenAIChatAttemptEventSourceFactory(
dc *chatDispatchContext,
cfg openAIChatStreamGateConfig,
usage *openAIStreamGateUsageHolder,
) openAIAttemptEventSourceFactory {
return func(transport openAIAttemptTransport) (streamgate.NormalizedEventSource, error) {
var src streamgate.NormalizedEventSource
switch transport.path {
case openAIAdmissionRun:
if transport.run == nil {
return nil, fmt.Errorf("openai stream gate: chat normalized attempt is missing its run transport")
}
cfg.selector.set(openAIStreamGateCodecNormalized)
if cfg.mode == openAIChatGateModeBuffered {
src = newOpenAIBufferedChatEventSource(dc, transport.run, cfg.holder, usage, transport.usage)
} else {
src = newOpenAIRunEventSource(transport.run.Stream(), transport.run.WaitTimeout(), usage, transport.usage)
}
case openAIAdmissionTunnel:
if transport.tunnel == nil {
return nil, fmt.Errorf("openai stream gate: chat tunnel attempt is missing its tunnel transport")
}
cfg.selector.set(openAIStreamGateCodecTunnel)
// A fresh rewriter/assembler per attempt so an aborted attempt's
// partial rewrite or usage state never bleeds into its replacement.
assembler := &providerChatAssembler{streaming: dc.req.Stream}
rewriter := newProviderModelRewriter(dc.req.Stream, dc.req.Model)
state := openAITunnelCodecStateForSink(cfg.sink)
state.reset()
tunnelSrc := newOpenAITunnelEndpointEventSource(transport.tunnel.Stream(), transport.tunnel.WaitTimeout(), rewriter, assembler, openAIRebuildEndpointChat, state)
src = &openAIStreamGateUsageTrackingTunnelSource{openAITunnelEventSource: tunnelSrc, usage: usage, attempt: transport.usage}
default:
return nil, fmt.Errorf("openai stream gate: unsupported attempt transport path %q for chat completions", transport.path)
}
return newOpenAIRecoverySourceEventSource(src, cfg.recoverySource), nil
}
}
// buildOpenAIChatStreamGateRuntimeFor assembles the request-local Core runtime
// for a runtime-enabled chat completion. The initial attempt binding wraps the
// already-dispatched transport directly (no re-admission); the
// dispatcher/rebuilder pair serves every subsequent recovery attempt.
// openAIResumeContextWindowTokens returns the target model group's single-request
// context-window contract from the provider-pool catalog, or 0 when the model is
// unknown. The resume builder treats 0 as fail-closed.
func (s *Server) openAIResumeContextWindowTokens(model string) int {
if entry := s.findProviderPoolEntry(model); entry != nil {
return entry.ContextWindowTokens
}
return 0
}
func (s *Server) buildOpenAIChatStreamGateRuntimeFor(dc *chatDispatchContext, cfg openAIChatStreamGateConfig) (*streamgate.RequestRuntime, *openAIStreamGateUsageHolder, error) {
usage := &openAIStreamGateUsageHolder{}
// One request-local recorder is shared by the initial and every recovery
// event source so a repeat-loop continuation is assembled from the aborted
// attempt's own output. The rebuilder measures the resume body against the
// target model context window and fails closed before dispatch on overflow.
recoverySource := newOpenAIRecoverySourceStore(dc.ingress)
cfg.recoverySource = recoverySource
rebuilder, err := newOpenAIRequestRebuilder(dc.ingress, openAIRebuildEndpointChat, recoverySource, s.openAIResumeContextWindowTokens(dc.req.Model))
if err != nil {
return nil, nil, err
}
build := newOpenAIChatRecoveryAdmissionBuilder(s, dc, cfg.holder)
dispatcher, err := newOpenAIAttemptDispatcher(s.service, rebuilder.RebuiltStore(), build, s.newOpenAIChatAttemptEventSourceFactory(dc, cfg, usage), dc.usage)
if err != nil {
return nil, nil, err
}
bindOpenAIRecoveryAdmissionState(cfg.sink, dispatcher.admissionState())
opts, err := s.streamGateRuntimeOptions()
if err != nil {
return nil, nil, err
}
snapRef, err := dc.ingress.recoveryRef()
if err != nil {
return nil, nil, err
}
dispatch := cfg.dispatch
cfg.initial.bindUsage(dispatch)
initialSource, err := s.newOpenAIChatAttemptEventSourceFactory(dc, cfg, usage)(cfg.initial)
if err != nil {
return nil, nil, err
}
initialController := &openAIAttemptController{
service: s.service, dispatch: dispatch, closeTransport: cfg.closeAll,
usageRecorder: dc.usage, usageBinding: cfg.initial.usageBinding, usage: cfg.initial.usage,
}
initialBinding, err := streamgate.NewAttemptBinding(
openAIStreamGateSafeToken("attempt", dispatch.RunID),
actualOpenAIModel(dispatch),
actualOpenAIProvider(dispatch),
actualOpenAIExecutionPath(dispatch, cfg.initial.path),
initialSource,
initialController,
)
if err != nil {
return nil, nil, err
}
requestID := openAIStreamGateSafeToken("req", dispatch.RunID)
snapshot, err := streamgate.NewRequestRuntimeSnapshot(
requestID, streamGateConfigGeneration, s.streamGateConfig().EffectiveEnvironment(), openAIRebuildEndpointChat, openAIRebuildFamily,
opts, cfg.registry, nil, snapRef, dispatcher, rebuilder, cfg.preparer, cfg.prepFactory, cfg.sink,
)
if err != nil {
return nil, nil, err
}
snapshot = snapshot.WithObservationSink(cfg.obsSink)
modelGroup := strings.TrimSpace(dc.req.Model)
if modelGroup == "" {
modelGroup = actualOpenAIModel(dispatch)
}
rt, err := streamgate.NewRequestRuntime(snapshot, modelGroup, initialBinding)
if err != nil {
return nil, nil, err
}
return rt, usage, nil
}
// buildOpenAIChatStreamGateRuntime assembles a request-local Core runtime for
// the normalized live-SSE chat completion path from an already-dispatched run
// handle. The registry is caller-supplied so tests can register additional
// test-only filters (e.g. an injected-violation filter) alongside the
// production filters without changing production wiring.
func (s *Server) buildOpenAIChatStreamGateRuntime(dc *chatDispatchContext, handle edgeservice.RunResult, sink streamgate.ReleaseSink, registry streamgate.FilterRegistrySnapshot) (*streamgate.RequestRuntime, *openAIStreamGateUsageHolder, error) {
return s.buildOpenAIChatStreamGateRuntimeFor(dc, openAIChatStreamGateConfig{
mode: openAIChatGateModeLive,
initial: openAIAttemptTransport{path: openAIAdmissionRun, run: handle},
dispatch: handle.Dispatch(),
closeAll: handle.Close,
sink: sink,
selector: newOpenAIStreamGateCodecSelector(openAIStreamGateCodecNormalized),
registry: registry,
obsSink: s.observationSink(),
})
}
// runOpenAIChatStreamGateRuntime is the shared runtime-enabled chat driver:
// build, run, close request resources exactly once, and report terminal usage
// against the framing actually exposed to the caller.
func (s *Server) runOpenAIChatStreamGateRuntime(
dc *chatDispatchContext,
cfg openAIChatStreamGateConfig,
sink openAIStreamGateSink,
writeBuildError func(),
) {
rt, _, err := s.buildOpenAIChatStreamGateRuntimeFor(dc, cfg)
if err != nil {
cfg.closeAll()
s.logger.Warn("openai stream gate chat runtime build failed", zap.Error(err))
writeBuildError()
dc.finishUsageRequest(usageStatusError, openAIStreamGateResponseMode(cfg.selector.get()))
return
}
runErr := rt.Run(dc.r.Context())
terminalCommitted, terminalSuccess := sink.terminalStatus()
// The request runtime owns the current attempt binding's transport, rebuilt
// lease, and the request rebuilder. Close them exactly once across success,
// error, and caller-cancel; a graceful close (no provider cancel) is used
// only when a success terminal was committed, otherwise the latest provider
// run is canceled.
_ = rt.CloseRequestResources(context.Background(), runErr == nil && terminalCommitted && terminalSuccess)
codec := cfg.selector.get()
if composite, ok := sink.(*openAICompositeReleaseSink); ok {
codec = composite.resolvedCodec()
}
status := streamGateUsageStatus(runErr, terminalCommitted, terminalSuccess)
dc.finishUsageRequest(status, openAIStreamGateResponseMode(codec))
}
// openAIChatCompositeSink wraps the normalized sink with the raw tunnel sink
// when the request may switch execution paths (provider-pool routes only). A
// direct model_routes request can never produce a tunnel attempt, so it keeps
// its normalized sink unwrapped and byte-identical to the single-path wiring.
func (s *Server) openAIChatCompositeSink(
w http.ResponseWriter,
flusher http.Flusher,
dc *chatDispatchContext,
selector *openAIStreamGateCodecSelector,
normalized openAIStreamGateSink,
) openAIStreamGateSink {
if dc.poolDispatch == nil {
return normalized
}
var tunnel openAIStreamGateSink
if dc.req.Stream {
tunnel = newOpenAITunnelReleaseSink(w, flusher)
} else {
tunnel = newOpenAIBufferedTunnelReleaseSink(w, flusher, dc.req.Model)
}
return newOpenAICompositeReleaseSink(selector, normalized, tunnel)
}
// runOpenAIChatStreamGate drives a live-SSE chat completion through the Core
// request runtime. It owns response-start staging, content/reasoning release,
// and terminal commit for the runtime-enabled path.
func (s *Server) runOpenAIChatStreamGate(w http.ResponseWriter, flusher http.Flusher, dc *chatDispatchContext, handle edgeservice.RunResult) {
dispatch := handle.Dispatch()
selector := newOpenAIStreamGateCodecSelector(openAIStreamGateCodecNormalized)
normalized := newOpenAIChatSSEReleaseSink(w, flusher, "chatcmpl-"+dispatch.RunID, time.Now().Unix(), responseModel(dc.req.Model, dispatch.Target))
sink := s.openAIChatCompositeSink(w, flusher, dc, selector, normalized)
fctx, err := s.openAIChatOutputFilterContext(dc)
if err != nil {
handle.Close()
s.logger.Warn("openai stream gate chat output filter context build failed", zap.Error(err))
writeSSEErrorWithType(w, flusher, "run_error", "stream gate runtime unavailable")
dc.finishUsageRequest(usageStatusError, responseModeNormalized)
return
}
registry, err := openAIStreamGateRegistrySnapshotFor(s.streamGateConfig(), fctx)
if err != nil {
handle.Close()
s.logger.Warn("openai stream gate chat registry build failed", zap.Error(err))
writeSSEErrorWithType(w, flusher, "run_error", "stream gate runtime unavailable")
dc.finishUsageRequest(usageStatusError, responseModeNormalized)
return
}
s.runOpenAIChatStreamGateRuntime(dc, openAIChatStreamGateConfig{
mode: openAIChatGateModeLive,
initial: openAIAttemptTransport{path: openAIAdmissionRun, run: handle},
dispatch: dispatch,
closeAll: handle.Close,
sink: sink,
selector: selector,
registry: registry,
obsSink: s.observationSink(),
}, sink, func() {
writeSSEErrorWithType(w, flusher, "run_error", "stream gate runtime unavailable")
})
}
// runOpenAIBufferedChatStreamGate drives a buffered chat completion (strict or
// tool-bearing SSE when stream is true, non-stream JSON otherwise) through the
// Core request runtime. The Core is the single owner of hold, validate,
// rebuild, and re-admission here: the legacy retrySubmit loop is not reachable.
func (s *Server) runOpenAIBufferedChatStreamGate(w http.ResponseWriter, flusher http.Flusher, dc *chatDispatchContext, handle edgeservice.RunResult, stream bool) {
writeBuildError := func() {
if stream {
writeSSEErrorWithType(w, flusher, "run_error", "stream gate runtime unavailable")
return
}
writeError(w, http.StatusInternalServerError, "run_error", "stream gate runtime unavailable")
}
cfg, sink, err := s.newOpenAIBufferedChatStreamGateConfig(w, flusher, dc, handle, stream, nil)
if err != nil {
handle.Close()
s.logger.Warn("openai stream gate buffered chat registry build failed", zap.Error(err))
writeBuildError()
dc.finishUsageRequest(usageStatusError, responseModeNormalized)
return
}
s.runOpenAIChatStreamGateRuntime(dc, cfg, sink, writeBuildError)
}
// newOpenAIBufferedChatStreamGateConfig resolves the buffered chat wiring:
// result holder, codec selector, release sink (composite for a provider-pool
// route), and the registry snapshot including the request-local tool-validation
// terminal gate. extraFilters lets a fixture register additional filters
// alongside the production set without duplicating this wiring.
func (s *Server) newOpenAIBufferedChatStreamGateConfig(
w http.ResponseWriter,
flusher http.Flusher,
dc *chatDispatchContext,
handle edgeservice.RunResult,
stream bool,
extraFilters []streamgate.FilterRegistration,
) (openAIChatStreamGateConfig, openAIStreamGateSink, error) {
holder := newOpenAIBufferedResultHolder()
selector := newOpenAIStreamGateCodecSelector(openAIStreamGateCodecNormalized)
normalized := newOpenAIBufferedChatReleaseSink(s, w, flusher, dc, stream, holder)
sink := s.openAIChatCompositeSink(w, flusher, dc, selector, normalized)
registry, err := s.openAIChatStreamGateRegistry(dc, holder, extraFilters)
if err != nil {
return openAIChatStreamGateConfig{}, nil, err
}
return openAIChatStreamGateConfig{
mode: openAIChatGateModeBuffered,
initial: openAIAttemptTransport{path: openAIAdmissionRun, run: handle},
dispatch: handle.Dispatch(),
closeAll: handle.Close,
sink: sink,
selector: selector,
registry: registry,
holder: holder,
obsSink: s.observationSink(),
}, sink, nil
}
// openAIChatStreamGateRegistry assembles the production registry snapshot for a
// chat request: the always-applicable Noop filter, the request-local
// tool-validation terminal gate when the request carries a tool contract, and
// any caller-supplied extra registrations.
func (s *Server) openAIChatStreamGateRegistry(dc *chatDispatchContext, holder *openAIBufferedResultHolder, extraFilters []streamgate.FilterRegistration) (streamgate.FilterRegistrySnapshot, error) {
extra, err := openAIToolValidationRegistrations(dc, holder)
if err != nil {
return streamgate.FilterRegistrySnapshot{}, err
}
extra = append(extra, extraFilters...)
fctx, err := s.openAIChatOutputFilterContext(dc)
if err != nil {
return streamgate.FilterRegistrySnapshot{}, err
}
return openAIStreamGateRegistrySnapshotFor(s.streamGateConfig(), fctx, extra...)
}
// runOpenAIChatPoolStreamGate drives a provider-pool chat completion through the
// Core request runtime from the pool's initial admission result, whichever
// execution path it selected. Every recovery attempt re-enters
// SubmitProviderPool, so the actual provider, model, and execution path may
// change inside this one request runtime as long as nothing is committed yet.
func (s *Server) runOpenAIChatPoolStreamGate(w http.ResponseWriter, dc *chatDispatchContext, result *edgeservice.ProviderPoolDispatchResult) {
flusher, _ := w.(http.Flusher)
writeBuildError := func() {
if dc.req.Stream {
writeSSEErrorWithType(w, flusher, "run_error", "stream gate runtime unavailable")
return
}
writeError(w, http.StatusInternalServerError, "run_error", "stream gate runtime unavailable")
}
cfg, sink, err := s.newOpenAIChatPoolStreamGateConfig(w, flusher, dc, result, nil)
if err != nil {
if cfg.closeAll != nil {
cfg.closeAll()
}
s.logger.Warn("openai stream gate pool chat runtime wiring failed", zap.Error(err))
writeError(w, http.StatusInternalServerError, "run_error", err.Error())
dc.finishUsageRequest(usageStatusError, responseModePassthrough)
return
}
s.runOpenAIChatStreamGateRuntime(dc, cfg, sink, writeBuildError)
}
// newOpenAIChatPoolStreamGateConfig resolves the provider-pool chat wiring from
// the pool's initial admission result. The response framing of a normalized
// attempt (live SSE vs buffered SSE/JSON) follows the same rules the direct
// route uses, and the composite sink carries the raw tunnel framing for a
// tunnel attempt. extraFilters lets a fixture register additional filters
// alongside the production set without duplicating this wiring.
func (s *Server) newOpenAIChatPoolStreamGateConfig(
w http.ResponseWriter,
flusher http.Flusher,
dc *chatDispatchContext,
result *edgeservice.ProviderPoolDispatchResult,
extraFilters []streamgate.FilterRegistration,
) (openAIChatStreamGateConfig, openAIStreamGateSink, error) {
transport, closeAll, err := openAIPoolAttemptTransport(result)
if err != nil {
return openAIChatStreamGateConfig{}, nil, err
}
buffered := !dc.req.Stream || (dc.outputPolicy.Strict && dc.outputPolicy.StreamBuffer) || len(dc.req.Tools) > 0
selector := newOpenAIStreamGateCodecSelector(openAIStreamGateCodecForPath(transport.path))
var (
normalized openAIStreamGateSink
holder *openAIBufferedResultHolder
mode openAIChatStreamGateMode
)
if buffered {
holder = newOpenAIBufferedResultHolder()
normalized = newOpenAIBufferedChatReleaseSink(s, w, flusher, dc, dc.req.Stream, holder)
mode = openAIChatGateModeBuffered
} else {
normalized = newOpenAIChatSSEReleaseSink(w, flusher, "chatcmpl-"+result.DispatchInfo.RunID, time.Now().Unix(), responseModel(dc.req.Model, result.DispatchInfo.Target))
mode = openAIChatGateModeLive
}
sink := s.openAIChatCompositeSink(w, flusher, dc, selector, normalized)
registry, err := s.openAIChatStreamGateRegistry(dc, holder, extraFilters)
if err != nil {
return openAIChatStreamGateConfig{closeAll: closeAll}, nil, err
}
return openAIChatStreamGateConfig{
mode: mode,
initial: transport,
dispatch: result.DispatchInfo,
closeAll: closeAll,
sink: sink,
selector: selector,
registry: registry,
holder: holder,
obsSink: s.observationSink(),
}, sink, nil
}
// openAIPoolAttemptTransport converts a provider-pool dispatch result into the
// initial attempt transport plus its exactly-once closer.
func openAIPoolAttemptTransport(result *edgeservice.ProviderPoolDispatchResult) (openAIAttemptTransport, func(), error) {
if result == nil {
return openAIAttemptTransport{}, nil, fmt.Errorf("provider-pool selection returned no result")
}
switch result.Path {
case edgeservice.ProviderPoolPathNormalized:
if result.Run == nil {
return openAIAttemptTransport{}, nil, fmt.Errorf("provider-pool selection returned normalized path but no run result")
}
return openAIAttemptTransport{path: openAIAdmissionRun, run: result.Run}, result.Run.Close, nil
case edgeservice.ProviderPoolPathTunnel:
if result.Tunnel == nil {
return openAIAttemptTransport{}, nil, fmt.Errorf("provider-pool selection returned tunnel path but no tunnel result")
}
return openAIAttemptTransport{path: openAIAdmissionTunnel, tunnel: result.Tunnel}, result.Tunnel.Close, nil
default:
return openAIAttemptTransport{}, nil, fmt.Errorf("provider-pool selection returned an unknown path")
}
}
// --- Provider tunnel runtime wiring ------------------------------------------
// openAITunnelStreamGateRequest describes the fixed (non-recovery-varying)
// parameters of a runtime-enabled provider tunnel passthrough request.
type openAITunnelStreamGateRequest struct {
route routeDispatch
ingress *openAIIngressSnapshot
endpoint string // openAIRebuildEndpointChat or openAIRebuildEndpointResponses
method string
path string
stream bool
modelGroupKey string
metadata map[string]string
hasScheme bool
estimate int
contextClass string
requestModel string // caller-facing model alias for echo rewrite; "" disables rewrite
authorize func(context.Context) (map[string]string, error)
rewriteBody func(body []byte, target string) ([]byte, error)
usage *openAIUsageRecorder
// pool is the provider-pool admission template this tunnel request was
// dispatched with, or nil for a direct provider route. When set, every
// recovery attempt re-enters SubmitProviderPool so the pool re-selects a
// candidate instead of pinning the initial one.
pool *edgeservice.ProviderPoolDispatchRequest
}
func newOpenAITunnelRecoveryAdmissionBuilder(req openAITunnelStreamGateRequest) openAIAttemptAdmissionBuilder {
return func(ctx context.Context, request streamgate.RebuiltRequest, body []byte) (openAIAttemptAdmission, error) {
buildBody := func(target string) ([]byte, error) {
return req.rewriteBody(body, target)
}
if req.pool != nil {
tunnelReq := req.pool.Tunnel
tunnelReq.BuildBody = buildBody
return openAIAttemptAdmission{
kind: openAIAdmissionPool,
pool: edgeservice.ProviderPoolDispatchRequest{
Run: req.pool.Run,
Tunnel: tunnelReq,
PrepareTunnel: req.pool.PrepareTunnel,
PrepareRun: req.pool.PrepareRun,
AcceptCandidate: req.pool.AcceptCandidate,
},
}, nil
}
tunnelReq := edgeservice.SubmitProviderTunnelRequest{
NodeRef: req.route.NodeRef,
ModelGroupKey: req.modelGroupKey,
Adapter: req.route.Adapter,
Target: req.route.Target,
SessionID: req.route.SessionID,
Method: req.method,
Path: req.path,
Stream: req.stream,
TimeoutSec: req.route.TimeoutSec,
MaxQueue: req.route.MaxQueue,
QueueTimeoutMS: req.route.QueueTimeoutMS,
Metadata: req.metadata,
EstimatedInputTokens: req.estimate,
ContextClass: req.contextClass,
ProviderPool: req.route.ProviderPool,
}
tunnelReq.BuildBody = buildBody
return openAIAttemptAdmission{kind: openAIAdmissionTunnel, tunnel: tunnelReq, authorize: req.authorize}, nil
}
}
// buildOpenAITunnelStreamGateRuntime assembles a request-local Core runtime
// for provider tunnel passthrough streaming. A fresh model-rewriter and usage
// assembler are constructed per attempt (including the initial one) so an
// aborted attempt's partial rewrite/usage state never bleeds into a
// recovery-replaced attempt.
func (s *Server) buildOpenAITunnelStreamGateRuntime(
req openAITunnelStreamGateRequest,
handle edgeservice.ProviderTunnelResult,
sink streamgate.ReleaseSink,
registry streamgate.FilterRegistrySnapshot,
) (*streamgate.RequestRuntime, *openAIStreamGateUsageHolder, error) {
usage := &openAIStreamGateUsageHolder{}
recoverySource := newOpenAIRecoverySourceStore(req.ingress)
rebuilder, err := newOpenAIRequestRebuilder(req.ingress, req.endpoint, recoverySource, s.openAIResumeContextWindowTokens(req.modelGroupKey))
if err != nil {
return nil, nil, err
}
build := newOpenAITunnelRecoveryAdmissionBuilder(req)
// A pure passthrough surface can only render raw provider framing. A pool
// recovery that selects a normalized candidate is therefore rejected here:
// the dispatcher aborts that attempt and the Core converges on a safe error
// terminal, preserving the existing normalized-candidate validation error
// instead of re-framing an in-flight passthrough response.
eventSourceFactory := func(transport openAIAttemptTransport) (streamgate.NormalizedEventSource, error) {
if transport.path != openAIAdmissionTunnel || transport.tunnel == nil {
return nil, fmt.Errorf("openai stream gate: unsupported recovery transport path %q for provider tunnel", transport.path)
}
assembler := &providerChatAssembler{streaming: req.stream}
rewriter := newProviderModelRewriter(req.stream, req.requestModel)
state := openAITunnelCodecStateForSink(sink)
state.reset()
src := newOpenAITunnelEndpointEventSource(transport.tunnel.Stream(), transport.tunnel.WaitTimeout(), rewriter, assembler, req.endpoint, state)
tracking := &openAIStreamGateUsageTrackingTunnelSource{openAITunnelEventSource: src, usage: usage, attempt: transport.usage}
return newOpenAIRecoverySourceEventSource(tracking, recoverySource), nil
}
dispatcher, err := newOpenAIAttemptDispatcher(s.service, rebuilder.RebuiltStore(), build, eventSourceFactory, req.usage)
if err != nil {
return nil, nil, err
}
bindOpenAIRecoveryAdmissionState(sink, dispatcher.admissionState())
opts, err := s.streamGateRuntimeOptions()
if err != nil {
return nil, nil, err
}
snapRef, err := req.ingress.recoveryRef()
if err != nil {
return nil, nil, err
}
dispatch := handle.Dispatch()
initialTransport := openAIAttemptTransport{path: openAIAdmissionTunnel, tunnel: handle}
initialTransport.bindUsage(dispatch)
initialAssembler := &providerChatAssembler{streaming: req.stream}
initialRewriter := newProviderModelRewriter(req.stream, req.requestModel)
initialState := openAITunnelCodecStateForSink(sink)
initialState.reset()
initialSource := &openAIStreamGateUsageTrackingTunnelSource{
openAITunnelEventSource: newOpenAITunnelEndpointEventSource(handle.Stream(), handle.WaitTimeout(), initialRewriter, initialAssembler, req.endpoint, initialState),
usage: usage,
attempt: initialTransport.usage,
}
initialController := &openAIAttemptController{
service: s.service, dispatch: dispatch, closeTransport: handle.Close,
usageRecorder: req.usage, usageBinding: initialTransport.usageBinding, usage: initialTransport.usage,
}
initialBinding, err := streamgate.NewAttemptBinding(
openAIStreamGateSafeToken("attempt", dispatch.RunID),
actualOpenAIModel(dispatch),
actualOpenAIProvider(dispatch),
actualOpenAIExecutionPath(dispatch, openAIAdmissionTunnel),
newOpenAIRecoverySourceEventSource(initialSource, recoverySource),
initialController,
)
if err != nil {
return nil, nil, err
}
requestID := openAIStreamGateSafeToken("req", dispatch.RunID)
snapshot, err := streamgate.NewRequestRuntimeSnapshot(
requestID, streamGateConfigGeneration, s.streamGateConfig().EffectiveEnvironment(), req.endpoint, openAIRebuildFamily,
opts, registry, nil, snapRef, dispatcher, rebuilder, nil, nil, sink,
)
if err != nil {
return nil, nil, err
}
snapshot = snapshot.WithObservationSink(s.observationSink())
modelGroup := req.modelGroupKey
if modelGroup == "" {
modelGroup = actualOpenAIModel(dispatch)
}
rt, err := streamgate.NewRequestRuntime(snapshot, modelGroup, initialBinding)
if err != nil {
return nil, nil, err
}
return rt, usage, nil
}
// openAIStreamGateUsageTrackingTunnelSource wraps openAITunnelEventSource and
// publishes the terminal observation into the response-rendering holder while
// continuously updating the separately owned per-attempt metric observation.
type openAIStreamGateUsageTrackingTunnelSource struct {
*openAITunnelEventSource
usage *openAIStreamGateUsageHolder
attempt *openAIAttemptUsage
}
func (s *openAIStreamGateUsageTrackingTunnelSource) NextEvent(ctx context.Context) (streamgate.NormalizedEvent, error) {
ev, err := s.openAITunnelEventSource.NextEvent(ctx)
if s.assembler != nil {
final := err != nil || ev.Kind() == streamgate.EventKindTerminal || ev.Kind() == streamgate.EventKindProviderError
obs := s.assembler.usageObservation()
if final {
obs = s.assembler.finalizeUsageObservation()
}
s.attempt.observe(obs)
if err == nil && ev.Kind() == streamgate.EventKindTerminal && s.usage != nil {
s.usage.set(obs)
}
}
return ev, err
}
var _ streamgate.NormalizedEventSource = (*openAIStreamGateUsageTrackingTunnelSource)(nil)
// runOpenAITunnelStreamGate drives streaming or buffered provider tunnel
// passthrough through the Core request runtime.
func (s *Server) runOpenAITunnelStreamGate(w http.ResponseWriter, r *http.Request, req openAITunnelStreamGateRequest, handle edgeservice.ProviderTunnelResult, usageRecorder *openAIUsageRecorder) {
req.usage = usageRecorder
flusher, _ := w.(http.Flusher)
var sink *openAITunnelReleaseSink
if req.stream {
sink = newOpenAITunnelReleaseSink(w, flusher)
} else {
sink = newOpenAIBufferedTunnelReleaseSink(w, flusher, req.requestModel)
}
fctx, err := s.openAITunnelOutputFilterContext(req)
if err != nil {
handle.Close()
s.logger.Warn("openai stream gate tunnel output filter context build failed", zap.Error(err))
writeError(w, http.StatusInternalServerError, "provider_tunnel_error", "stream gate runtime unavailable")
usageRecorder.FinishRequest(usageStatusError, responseModePassthrough)
return
}
registry, err := openAIStreamGateRegistrySnapshotFor(s.streamGateConfig(), fctx)
if err != nil {
handle.Close()
s.logger.Warn("openai stream gate tunnel registry build failed", zap.Error(err))
writeError(w, http.StatusInternalServerError, "provider_tunnel_error", "stream gate runtime unavailable")
usageRecorder.FinishRequest(usageStatusError, responseModePassthrough)
return
}
rt, _, err := s.buildOpenAITunnelStreamGateRuntime(req, handle, sink, registry)
if err != nil {
handle.Close()
s.logger.Warn("openai stream gate tunnel runtime build failed", zap.Error(err))
writeError(w, http.StatusInternalServerError, "provider_tunnel_error", "stream gate runtime unavailable")
usageRecorder.FinishRequest(usageStatusError, responseModePassthrough)
return
}
runErr := rt.Run(r.Context())
terminalCommitted, terminalSuccess := sink.terminalStatus()
// The request runtime owns the current attempt binding's transport, rebuilt
// lease, and the request rebuilder across success, error, and caller-cancel.
_ = rt.CloseRequestResources(context.Background(), runErr == nil && terminalCommitted && terminalSuccess)
status := streamGateUsageStatus(runErr, terminalCommitted, terminalSuccess)
usageRecorder.FinishRequest(status, responseModePassthrough)
}