iop/apps/edge/internal/openai/stream_gate_runtime.go
toki 7634ca8962 feat(stream-evidence-gate-core): complete stream evidence gate core implementation
- OpenAI request rebuilder with tool validation and provider tunnel
- Edge config runtime refresh for stream evidence gate
- Filter observation contract and runtime with sink/correlation
- Stream gate dispatcher, release sink, and vertical slice
- Recovery coordinator for evidence tail
- Parallel evaluation and commit boundary
- E2E test script for OpenAI vLLM
- Archive completed task groups to archive/2026/07
2026-07-28 04:08:12 +09:00

1292 lines
49 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/streamgate"
iop "iop/proto/gen/iop"
)
const (
streamGateChannelDefault = "default"
streamGateEnvironment = "edge"
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 last-attempt usage observation out
// of the event source for post-run usage metric emission. It is reset by
// every new attempt so an aborted attempt's partial usage never survives a
// recovery replace.
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
mu sync.Mutex
startSent bool
}
func newOpenAIRunEventSource(stream edgeservice.RunStream, waitTimeout time.Duration, usage *openAIStreamGateUsageHolder) *openAIRunEventSource {
return &openAIRunEventSource{stream: stream, waitTimeout: waitTimeout, usage: usage}
}
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
}
return streamgate.NewReasoningDeltaEvent(streamGateChannelDefault, event.GetDelta(), time.Now())
case "complete":
if s.usage != nil {
s.usage.set(runEventUsageObservation(event))
}
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
mu sync.Mutex
stage int
}
func newOpenAIBufferedChatEventSource(dc *chatDispatchContext, handle edgeservice.RunResult, holder *openAIBufferedResultHolder, usage *openAIStreamGateUsageHolder) *openAIBufferedChatEventSource {
return &openAIBufferedChatEventSource{
req: dc.req,
outputPolicy: dc.outputPolicy,
validation: dc.validation,
handle: handle,
holder: holder,
usage: usage,
}
}
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})
if s.usage != nil {
s.usage.set(usageObservationFromOpenAIUsage(result.usage, result.reasoningLen))
}
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.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
// streamgate.NormalizedEventSource. Raw provider body bytes travel through
// Core as opaque text_delta payloads (Go string<->[]byte round trips are
// byte-exact) so pure passthrough is preserved end-to-end. The caller-facing
// model echo rewrite runs before bytes enter Core so the transform is
// evaluated exactly once, in original provider byte order.
type openAITunnelEventSource struct {
frames <-chan *iop.ProviderTunnelFrame
waitTimeout time.Duration
rewriter *providerModelRewriter
assembler *providerChatAssembler
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 (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
}
ev, err := streamgate.NewResponseStartEvent(streamGateChannelDefault, status, sanitizedTunnelResponseHeaders(frame.GetHeaders()), 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.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)
}
if len(rewritten) == 0 {
return events, 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)
}
if s.rewriter != nil {
if flushed := s.rewriter.FlushStream(); 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()
}
// 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,
},
}, 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
// 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) {
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 {
return newOpenAIBufferedChatEventSource(dc, transport.run, cfg.holder, usage), nil
}
return newOpenAIRunEventSource(transport.run.Stream(), transport.run.WaitTimeout(), usage), nil
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)
src := newOpenAITunnelEventSource(transport.tunnel.Stream(), transport.tunnel.WaitTimeout(), rewriter, assembler)
return &openAIStreamGateUsageTrackingTunnelSource{openAITunnelEventSource: src, usage: usage}, nil
default:
return nil, fmt.Errorf("openai stream gate: unsupported attempt transport path %q for chat completions", transport.path)
}
}
}
// 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.
func (s *Server) buildOpenAIChatStreamGateRuntimeFor(dc *chatDispatchContext, cfg openAIChatStreamGateConfig) (*streamgate.RequestRuntime, *openAIStreamGateUsageHolder, error) {
usage := &openAIStreamGateUsageHolder{}
rebuilder, err := newOpenAIRequestRebuilder(dc.ingress, openAIRebuildEndpointChat)
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))
if err != nil {
return nil, nil, err
}
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
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}
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, streamGateEnvironment, 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, usage, err := s.buildOpenAIChatStreamGateRuntimeFor(dc, cfg)
if err != nil {
cfg.closeAll()
s.logger.Warn("openai stream gate chat runtime build failed", zap.Error(err))
writeBuildError()
emitUsageMetrics(dc.usageLabels(s, openAIStreamGateResponseMode(cfg.selector.get())), usageStatusError, usageObservation{})
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()
}
metricLabels := dc.usageLabels(s, openAIStreamGateResponseMode(codec))
status := streamGateUsageStatus(runErr, terminalCommitted, terminalSuccess)
if status == usageStatusSuccess {
emitUsageMetrics(metricLabels, status, usage.get())
return
}
emitUsageMetrics(metricLabels, status, usageObservation{})
}
// 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)
registry, err := openAIStreamGateRegistrySnapshot()
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")
emitUsageMetrics(dc.usageLabels(s, responseModeNormalized), usageStatusError, usageObservation{})
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()
emitUsageMetrics(dc.usageLabels(s, responseModeNormalized), usageStatusError, usageObservation{})
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...)
return openAIStreamGateRegistrySnapshotWith(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())
emitUsageMetrics(dc.usageLabels(s, responseModePassthrough), usageStatusError, usageObservation{})
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
modelGroupKey string
metadata map[string]string
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)
// 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,
},
}, 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: true,
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{}
rebuilder, err := newOpenAIRequestRebuilder(req.ingress, req.endpoint)
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: true}
rewriter := newProviderModelRewriter(true, req.requestModel)
src := newOpenAITunnelEventSource(transport.tunnel.Stream(), transport.tunnel.WaitTimeout(), rewriter, assembler)
return &openAIStreamGateUsageTrackingTunnelSource{openAITunnelEventSource: src, usage: usage}, nil
}
dispatcher, err := newOpenAIAttemptDispatcher(s.service, rebuilder.RebuiltStore(), build, eventSourceFactory)
if err != nil {
return nil, nil, err
}
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()
initialAssembler := &providerChatAssembler{streaming: true}
initialRewriter := newProviderModelRewriter(true, req.requestModel)
initialSource := &openAIStreamGateUsageTrackingTunnelSource{
openAITunnelEventSource: newOpenAITunnelEventSource(handle.Stream(), handle.WaitTimeout(), initialRewriter, initialAssembler),
usage: usage,
}
initialController := &openAIAttemptController{service: s.service, dispatch: dispatch, closeTransport: handle.Close}
initialBinding, err := streamgate.NewAttemptBinding(
openAIStreamGateSafeToken("attempt", dispatch.RunID),
actualOpenAIModel(dispatch),
actualOpenAIProvider(dispatch),
actualOpenAIExecutionPath(dispatch, openAIAdmissionTunnel),
initialSource,
initialController,
)
if err != nil {
return nil, nil, err
}
requestID := openAIStreamGateSafeToken("req", dispatch.RunID)
snapshot, err := streamgate.NewRequestRuntimeSnapshot(
requestID, streamGateConfigGeneration, streamGateEnvironment, 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 attempt's assembled usage observation into the shared holder
// once the terminal event is read, so the top-level runner can report final
// usage after the last (successful) attempt without holding a direct
// reference to the per-attempt assembler.
type openAIStreamGateUsageTrackingTunnelSource struct {
*openAITunnelEventSource
usage *openAIStreamGateUsageHolder
}
func (s *openAIStreamGateUsageTrackingTunnelSource) NextEvent(ctx context.Context) (streamgate.NormalizedEvent, error) {
ev, err := s.openAITunnelEventSource.NextEvent(ctx)
if err == nil && ev.Kind() == streamgate.EventKindTerminal && s.assembler != nil {
s.usage.set(s.assembler.usageObservation())
}
return ev, err
}
var _ streamgate.NormalizedEventSource = (*openAIStreamGateUsageTrackingTunnelSource)(nil)
// runOpenAITunnelStreamGate drives streaming provider tunnel passthrough
// through the Core request runtime.
func (s *Server) runOpenAITunnelStreamGate(w http.ResponseWriter, r *http.Request, req openAITunnelStreamGateRequest, handle edgeservice.ProviderTunnelResult, metricLabels usageLabels) {
flusher, _ := w.(http.Flusher)
sink := newOpenAITunnelReleaseSink(w, flusher)
registry, err := openAIStreamGateRegistrySnapshot()
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")
emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
return
}
rt, usage, 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")
emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
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)
if status == usageStatusSuccess {
emitUsageMetrics(metricLabels, status, usage.get())
return
}
emitUsageMetrics(metricLabels, status, usageObservation{})
}