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

1694 lines
58 KiB
Go

package openai
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"net/http"
"os"
"strconv"
"strings"
"time"
"unicode"
"go.uber.org/zap"
edgeservice "iop/apps/edge/internal/service"
iop "iop/proto/gen/iop"
)
const (
streamTraceMetadataKey = "iop_trace_stream"
streamTraceEnvKey = "IOP_OPENAI_COMPAT_TRACE_STREAM"
)
func (s *Server) streamChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, submitReq edgeservice.SubmitRunRequest, handle edgeservice.RunResult, outputPolicy strictOutputPolicy, validation toolValidationContract) {
flusher, ok := w.(http.Flusher)
if !ok {
handle.Close()
writeError(w, http.StatusInternalServerError, "streaming_not_supported", "response writer does not support streaming")
return
}
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
// Buffered streams collect and validate the full response before emitting
// user-visible chunks. Strict buffered output opts in explicitly, and
// tool-bearing streams use the same path so malformed tool calls can be
// retried or rejected before reaching a client tool runner.
if (outputPolicy.Strict && outputPolicy.StreamBuffer) || len(req.Tools) > 0 {
s.streamBufferedChatCompletion(w, r, req, submitReq, handle, flusher, outputPolicy, validation)
return
}
// Live SSE may emit content deltas before the terminal event, so runtime
// tool validation is excluded upstream; write the role chunk immediately.
defer handle.Close()
created := time.Now().Unix()
id := "chatcmpl-" + handle.Dispatch().RunID
model := responseModel(req.Model, handle.Dispatch().Target)
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, chatResponseModeLabel(req))
traceStream := openAICompatTraceStreamEnabled(submitReq.Metadata)
traceSeq := 0
shouldExposeReasoning := req.includeReasoning() && (!outputPolicy.Strict || req.explicitlyIncludesReasoning())
writeTracedContentDelta := func(content, source string, filtered bool) {
if content == "" {
return
}
traceSeq++
if traceStream {
logOpenAICompatStreamOutput(s.logger, handle.Dispatch().RunID, traceSeq, "content", content,
zap.Int("source_len", len(source)),
zap.Bool("filtered", filtered),
zap.Bool("source_has_open_think", hasOpenThinkTag(source)),
zap.Bool("source_has_close_think", hasCloseThinkTag(source)),
)
}
writeContentDeltaSSE(w, flusher, id, created, model, content)
}
writeTracedReasoningDelta := func(reasoning string) {
if reasoning == "" {
return
}
traceSeq++
if traceStream {
logOpenAICompatStreamOutput(s.logger, handle.Dispatch().RunID, traceSeq, "reasoning", reasoning)
}
writeSSE(w, flusher, chatCompletionChunk{
ID: id,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []chatCompletionChunkChoice{{
Index: 0,
Delta: chatDelta{ReasoningContent: reasoning},
}},
})
}
writeSSE(w, flusher, chatCompletionChunk{
ID: id,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []chatCompletionChunkChoice{{
Index: 0,
Delta: chatDelta{Role: "assistant"},
}},
})
var contentBuilder strings.Builder
var reasoningBuilder strings.Builder
var emittedContent strings.Builder
var toolTextFilter *streamToolTextFilter
contentSentinelFilter := &streamSentinelFilter{}
reasoningSentinelFilter := &streamSentinelFilter{}
if len(req.Tools) > 0 {
toolTextFilter = &streamToolTextFilter{}
}
defer func() {
s.logger.Info("openai chat completion stream closed",
zap.String("run_id", handle.Dispatch().RunID),
zap.Bool("strict_output", outputPolicy.Strict),
zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
zap.Int("content_len", contentBuilder.Len()),
zap.Int("reasoning_len", reasoningBuilder.Len()),
)
}()
stream := handle.Stream()
if stream.Events == nil {
writeSSEError(w, flusher, "run stream unavailable")
return
}
processContentDelta := func(sourceDelta string) {
if sourceDelta == "" {
return
}
delta := contentSentinelFilter.Append(sourceDelta)
if delta == "" {
return
}
contentBuilder.WriteString(delta)
filteredDelta := delta
if toolTextFilter != nil {
filteredDelta = toolTextFilter.Append(filteredDelta)
}
if filteredDelta == "" {
if traceStream {
s.logger.Info("openai chat completion stream output suppressed",
zap.String("run_id", handle.Dispatch().RunID),
zap.String("type", "content"),
zap.Int("source_len", len(sourceDelta)),
zap.Bool("source_has_open_think", hasOpenThinkTag(sourceDelta)),
zap.Bool("source_has_close_think", hasCloseThinkTag(sourceDelta)),
)
}
return
}
emittedContent.WriteString(filteredDelta)
writeTracedContentDelta(filteredDelta, sourceDelta, filteredDelta != sourceDelta)
}
for {
select {
case <-r.Context().Done():
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
emitUsageMetrics(metricLabels, usageStatusForError(r.Context().Err()), usageObservation{reasoningChars: reasoningBuilder.Len()})
return
case nodeEvent, ok := <-stream.NodeEvents:
if !ok {
stream.NodeEvents = nil
continue
}
if edgeservice.IsNodeDisconnected(nodeEvent) {
emitUsageMetrics(metricLabels, usageStatusError, usageObservation{reasoningChars: reasoningBuilder.Len()})
writeSSEError(w, flusher, "node disconnected")
return
}
case event, ok := <-stream.Events:
if !ok {
writeSSEError(w, flusher, "run stream closed")
return
}
if event == nil {
continue
}
switch event.GetType() {
case "delta":
processContentDelta(event.GetDelta())
case "reasoning_delta":
if event.GetDelta() == "" {
continue
}
reasoning := reasoningSentinelFilter.Append(event.GetDelta())
if reasoning == "" {
continue
}
reasoningBuilder.WriteString(reasoning)
if !shouldExposeReasoning {
continue
}
writeTracedReasoningDelta(reasoning)
case "complete":
processContentDelta(contentSentinelFilter.Flush())
if reasoning := reasoningSentinelFilter.Flush(); reasoning != "" {
reasoningBuilder.WriteString(reasoning)
if shouldExposeReasoning {
writeTracedReasoningDelta(reasoning)
}
}
metadata := event.GetMetadata()
streamUsage := usageObservation{reasoningChars: reasoningBuilder.Len()}
if u := event.GetUsage(); u != nil {
streamUsage.inputTokens = int(u.GetInputTokens())
streamUsage.outputTokens = int(u.GetOutputTokens())
streamUsage.reasoningTokens = int(u.GetReasoningTokens())
streamUsage.cachedInputTokens = int(u.GetCachedInputTokens())
}
finishReason := metadata["finish_reason"]
if finishReason == "" {
finishReason = "stop"
}
nativeToolCalls := toolCallsFromRunMetadata(metadata)
if len(nativeToolCalls) > 0 {
if toolTextFilter != nil {
pending := toolTextFilter.FlushNonCandidate()
if pending != "" {
emittedContent.WriteString(pending)
writeTracedContentDelta(pending, pending, false)
}
}
nativeToolCalls = normalizeNativeToolCallsForResponse(nativeToolCalls, req.Tools)
if traceStream {
s.logger.Info("openai chat completion stream output chunk",
zap.String("run_id", handle.Dispatch().RunID),
zap.Int("seq", traceSeq+1),
zap.String("type", "tool_calls"),
zap.Int("tool_call_count", len(nativeToolCalls)),
)
}
writeToolCallsDeltaSSE(w, flusher, id, created, model, nativeToolCalls)
finishReason = "tool_calls"
} else if len(req.Tools) > 0 {
res := synthesizeToolCallsFromTextResult(contentBuilder.String(), req.Tools, handle.Dispatch().RunID)
if res.validationErr != nil {
emitUsageMetrics(metricLabels, usageStatusError, usageObservation{reasoningChars: reasoningBuilder.Len()})
writeSSEErrorWithType(w, flusher, "tool_validation_error", res.validationErr.Error())
return
}
if len(res.toolCalls) > 0 {
if remaining := unstreamedCleanedText(res.cleaned, emittedContent.String()); remaining != "" {
emittedContent.WriteString(remaining)
writeTracedContentDelta(remaining, remaining, false)
}
if traceStream {
s.logger.Info("openai chat completion stream output chunk",
zap.String("run_id", handle.Dispatch().RunID),
zap.Int("seq", traceSeq+1),
zap.String("type", "tool_calls"),
zap.Int("tool_call_count", len(res.toolCalls)),
)
}
writeToolCallsDeltaSSE(w, flusher, id, created, model, res.toolCalls)
finishReason = "tool_calls"
} else if !res.candidateFound {
if toolTextFilter != nil {
if pending := toolTextFilter.Flush(); pending != "" {
emittedContent.WriteString(pending)
writeTracedContentDelta(pending, pending, false)
}
}
}
} else if toolTextFilter != nil {
if pending := toolTextFilter.Flush(); pending != "" {
emittedContent.WriteString(pending)
writeTracedContentDelta(pending, pending, false)
}
}
if !outputPolicy.Strict &&
strings.TrimSpace(contentBuilder.String()) == "" &&
strings.TrimSpace(emittedContent.String()) == "" &&
strings.TrimSpace(reasoningBuilder.String()) != "" &&
len(nativeToolCalls) == 0 {
fallback := hiddenReasoningFallbackContent(finishReason)
if req.includeReasoning() {
fallback = reasoningOnlyFallbackContent(reasoningBuilder.String(), finishReason)
}
emittedContent.WriteString(fallback)
writeTracedContentDelta(fallback, reasoningBuilder.String(), false)
}
if traceStream {
s.logger.Info("openai chat completion stream output done",
zap.String("run_id", handle.Dispatch().RunID),
zap.Int("seq", traceSeq+1),
zap.String("finish_reason", finishReason),
zap.Int("content_len", contentBuilder.Len()),
zap.Int("reasoning_len", reasoningBuilder.Len()),
)
}
writeSSE(w, flusher, chatCompletionChunk{
ID: id,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []chatCompletionChunkChoice{{
Index: 0,
Delta: chatDelta{},
FinishReason: finishReason,
}},
})
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
emitUsageMetrics(metricLabels, usageStatusSuccess, streamUsage)
return
case "error", "cancelled":
msg := event.GetError()
if msg == "" {
msg = event.GetMessage()
}
if msg == "" {
msg = "run failed"
}
emitUsageMetrics(metricLabels, usageStatusError, usageObservation{reasoningChars: reasoningBuilder.Len()})
writeSSEError(w, flusher, msg)
return
}
case <-time.After(handle.WaitTimeout()):
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
emitUsageMetrics(metricLabels, usageStatusCancel, usageObservation{reasoningChars: reasoningBuilder.Len()})
writeSSEError(w, flusher, errRunTimedOut.Error())
return
}
}
}
// tunnelChatCompletionPassthrough serves a Chat Completions request over the
// provider tunnel (SDD S04/S06): provider status, headers, and content are
// relayed without IOP sideband fields or events. Chat Completions model echoes
// are normalized back to the caller-facing model alias.
func (s *Server) tunnelChatCompletionPassthrough(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, dispatch routeDispatch, runMeta map[string]string, rawBody []byte, estimate int, contextClass string) {
handle, ok := s.submitChatCompletionTunnel(w, r, req, dispatch, runMeta, rawBody, estimate, contextClass, responseModePassthrough)
if !ok {
return
}
defer handle.Close()
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModePassthrough)
s.writeProviderTunnelResponse(w, r, handle, req.Stream, req.Model, metricLabels)
}
// tunnelChatCompletionPassthroughSideband serves a Chat Completions request
// over the raw provider tunnel with the explicit IOP sideband extension
// surface (SDD S05): provider content is preserved, but the response also
// carries IOP route/usage/assembled observations and is never claimed as
// provider-original byte identity.
func (s *Server) tunnelChatCompletionPassthroughSideband(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, dispatch routeDispatch, runMeta map[string]string, rawBody []byte, estimate int, contextClass string) {
handle, ok := s.submitChatCompletionTunnel(w, r, req, dispatch, runMeta, rawBody, estimate, contextClass, responseModePassthroughSideband)
if !ok {
return
}
defer handle.Close()
if req.Stream {
s.writeProviderTunnelSidebandStream(w, r, handle)
return
}
s.writeProviderTunnelSidebandResponse(w, r, handle)
}
// errProviderAuthRequired signals that provider auth forwarding is configured
// as required but the caller did not supply the configured provider token
// header. The raw token is never part of this error.
var errProviderAuthRequired = errors.New("provider auth token is required")
// providerTunnelAuthHeaders builds the provider tunnel request headers that
// forward a request-time raw user token to the provider (SDD S03). It reads the
// token only from the configured `openai.provider_auth.from_header` and never
// reuses the inbound IOP `Authorization` header. Returns nil when provider auth
// is disabled or the (non-required) token header is absent, and
// errProviderAuthRequired when a required token header is missing.
func (s *Server) providerTunnelAuthHeaders(r *http.Request) (map[string]string, error) {
auth := s.cfg.ProviderAuth
if !auth.Enabled {
return nil, nil
}
raw := strings.TrimSpace(r.Header.Get(auth.FromHeader))
if raw == "" {
if auth.Required {
return nil, errProviderAuthRequired
}
return nil, nil
}
value := raw
if scheme := strings.TrimSpace(auth.Scheme); scheme != "" &&
!strings.HasPrefix(strings.ToLower(raw), strings.ToLower(scheme)+" ") {
value = scheme + " " + raw
}
return map[string]string{auth.TargetHeader: value}, nil
}
// submitChatCompletionTunnel dispatches a Chat Completions provider tunnel
// request with the given response mode. On failure it writes the dispatch
// error to the caller and returns ok=false.
func (s *Server) submitChatCompletionTunnel(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, dispatch routeDispatch, runMeta map[string]string, rawBody []byte, estimate int, contextClass, responseMode string) (edgeservice.ProviderTunnelResult, bool) {
providerAuthHeaders, err := s.providerTunnelAuthHeaders(r)
if err != nil {
// Missing required provider auth is rejected before dispatch; the raw
// token is never echoed into the error surface.
writeError(w, http.StatusBadRequest, "invalid_request_error", "provider auth token is required")
return nil, false
}
metadata := make(map[string]string, len(runMeta)+5)
for k, v := range runMeta {
metadata[k] = v
}
metadata["openai_model"] = req.Model
metadata["openai_stream"] = strconv.FormatBool(req.Stream)
metadata[responseModeMetadataKey] = responseMode
metadata["estimated_input_tokens"] = strconv.Itoa(estimate)
metadata["context_class"] = contextClass
tunnelReq := edgeservice.SubmitProviderTunnelRequest{
NodeRef: dispatch.NodeRef,
ModelGroupKey: strings.TrimSpace(req.Model),
Adapter: dispatch.Adapter,
Target: dispatch.Target,
SessionID: dispatch.SessionID,
Method: http.MethodPost,
Path: "/v1/chat/completions",
Headers: providerAuthHeaders,
BuildBody: func(target string) ([]byte, error) {
return rewriteChatCompletionModel(rawBody, target, req)
},
Stream: req.Stream,
TimeoutSec: dispatch.TimeoutSec,
MaxQueue: dispatch.MaxQueue,
QueueTimeoutMS: dispatch.QueueTimeoutMS,
Metadata: metadata,
EstimatedInputTokens: estimate,
ContextClass: contextClass,
ProviderPool: dispatch.ProviderPool,
}
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseMode)
handle, err := s.service.SubmitProviderTunnel(r.Context(), tunnelReq)
if err != nil {
emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
return nil, false
}
s.logger.Info("openai chat completion passthrough dispatch",
zap.String("run_id", handle.Dispatch().RunID),
zap.String("node_id", handle.Dispatch().NodeID),
zap.String("model_group", handle.Dispatch().ModelGroupKey),
zap.String("adapter", handle.Dispatch().Adapter),
zap.String("target", handle.Dispatch().Target),
zap.Bool("stream", req.Stream),
zap.String("response_mode", responseMode),
zap.Int("estimated_input_tokens", handle.Dispatch().EstimatedInputTokens),
zap.String("context_class", handle.Dispatch().ContextClass),
zap.String("queue_reason", handle.Dispatch().QueueReason),
)
return handle, true
}
// writeProviderTunnelResponse relays ordered tunnel frames to the HTTP caller.
// The response-start frame sets status/headers, body frames are written and
// flushed in order, and END terminates the response. Caller disconnect and
// wait timeout propagate cancellation to the Node cancel path. The caller
// supplies metricLabels so success/error/cancel usage is attributed to the
// right endpoint (Chat vs Responses) and model_group; the writer never
// recomputes labels from requestModel, which is reserved for model-echo
// rewriting only.
func (s *Server) writeProviderTunnelResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult, reqStream bool, requestModel string, metricLabels usageLabels) {
frames := handle.Stream().Frames
if frames == nil {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream unavailable")
return
}
flusher, _ := w.(http.Flusher)
timer := time.NewTimer(handle.WaitTimeout())
defer timer.Stop()
assembler := &providerChatAssembler{streaming: reqStream}
modelRewriter := newProviderModelRewriter(reqStream, requestModel)
wroteHeader := false
bodyBytes := 0
var bufferedBody []byte
rewriteModelEcho := modelRewriter != nil
// metricStatus is the terminal status reported to usage metrics. It defaults
// to error and is upgraded to success on END or downgraded to cancel on
// caller give-up; provider usage observed from the passthrough body is
// reported without ever mutating that body (SDD S05/D04).
metricStatus := usageStatusError
defer func() {
obs := assembler.observation()
s.logger.Info("openai chat completion passthrough closed",
zap.String("run_id", handle.Dispatch().RunID),
zap.Bool("wrote_header", wroteHeader),
zap.Int("body_bytes", bodyBytes),
zap.String("assembled_content", obs.Content),
zap.String("assembled_reasoning", obs.Reasoning),
zap.Strings("assembled_tool_calls", obs.ToolCallNames),
zap.Int("assembled_tool_call_count", len(obs.ToolCallNames)),
)
emitUsageMetrics(metricLabels, metricStatus, assembler.usageObservation())
}()
for {
select {
case <-r.Context().Done():
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
metricStatus = usageStatusForError(r.Context().Err())
return
case <-timer.C:
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
metricStatus = usageStatusCancel
if !wroteHeader {
writeError(w, http.StatusBadGateway, "run_error", errRunTimedOut.Error())
}
return
case frame, ok := <-frames:
if !ok {
if !wroteHeader {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream closed before provider response")
}
return
}
switch frame.GetKind() {
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
if wroteHeader {
continue
}
copyProviderResponseHeaders(w.Header(), frame.GetHeaders())
status := int(frame.GetStatusCode())
if status == 0 {
status = http.StatusOK
}
if status >= http.StatusBadRequest {
rewriteModelEcho = false
}
if rewriteModelEcho {
// The response body may change length when provider-served
// model names are rewritten to caller-facing model aliases.
w.Header().Del("Content-Length")
}
w.WriteHeader(status)
wroteHeader = true
if flusher != nil {
flusher.Flush()
}
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
body := frame.GetBody()
if len(body) == 0 {
continue
}
if !wroteHeader {
// Defensive: a body frame before response-start still reaches the
// caller instead of being dropped.
w.WriteHeader(http.StatusOK)
wroteHeader = true
}
bodyBytes += len(body)
assembler.Write(body)
if rewriteModelEcho {
if reqStream {
body = modelRewriter.AppendStream(body)
} else {
bufferedBody = append(bufferedBody, body...)
continue
}
}
if len(body) == 0 {
continue
}
if _, err := w.Write(body); err != nil {
// The caller is gone mid-stream; propagate cancel to the Node.
s.sendCancelRun(handle.Dispatch())
metricStatus = usageStatusCancel
return
}
if flusher != nil {
flusher.Flush()
}
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
msg := frame.GetError()
if msg == "" {
msg = "provider tunnel failed"
}
if !wroteHeader {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", msg)
return
}
// Status/headers are already committed; the response is truncated
// and the caller observes the broken stream.
s.logger.Warn("openai passthrough tunnel error after response start",
zap.String("run_id", handle.Dispatch().RunID),
zap.String("error", msg),
)
if rewriteModelEcho && !reqStream && len(bufferedBody) > 0 {
body := modelRewriter.RewriteComplete(bufferedBody)
_, _ = w.Write(body)
if flusher != nil {
flusher.Flush()
}
}
return
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
if !wroteHeader {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel ended before provider response")
return
}
if rewriteModelEcho {
var body []byte
if reqStream {
body = modelRewriter.FlushStream()
} else if len(bufferedBody) > 0 {
body = modelRewriter.RewriteComplete(bufferedBody)
}
if len(body) > 0 {
if _, err := w.Write(body); err != nil {
s.sendCancelRun(handle.Dispatch())
metricStatus = usageStatusCancel
return
}
if flusher != nil {
flusher.Flush()
}
}
}
metricStatus = usageStatusSuccess
return
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
// Usage frames are Edge-internal observation candidates; pure
// passthrough records them for metrics but never merges them into
// the response body (SDD D04).
assembler.recordProtoUsage(frame.GetUsage())
continue
}
}
}
}
type providerModelRewriter struct {
streaming bool
model string
pending []byte
}
func newProviderModelRewriter(streaming bool, model string) *providerModelRewriter {
model = strings.TrimSpace(model)
if model == "" {
return nil
}
return &providerModelRewriter{streaming: streaming, model: model}
}
func (r *providerModelRewriter) AppendStream(chunk []byte) []byte {
if r == nil || !r.streaming || len(chunk) == 0 {
return chunk
}
r.pending = append(r.pending, chunk...)
var out bytes.Buffer
for {
idx := bytes.IndexByte(r.pending, '\n')
if idx < 0 {
break
}
line := r.pending[:idx+1]
out.Write(rewriteProviderSSEModelLine(line, r.model))
r.pending = r.pending[idx+1:]
}
return out.Bytes()
}
func (r *providerModelRewriter) FlushStream() []byte {
if r == nil || len(r.pending) == 0 {
return nil
}
pending := r.pending
r.pending = nil
return rewriteProviderSSEModelLine(pending, r.model)
}
func (r *providerModelRewriter) RewriteComplete(body []byte) []byte {
if r == nil || len(body) == 0 {
return body
}
return rewriteProviderJSONModel(body, r.model)
}
func rewriteProviderSSEModelLine(line []byte, model string) []byte {
body, ending := splitLineEnding(line)
prefix, payload, ok := bytes.Cut(body, []byte(":"))
if !ok || strings.TrimSpace(string(prefix)) != "data" {
return line
}
payload = bytes.TrimSpace(payload)
if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
return line
}
rewritten := rewriteProviderJSONModel(payload, model)
if bytes.Equal(rewritten, payload) {
return line
}
out := make([]byte, 0, len("data: ")+len(rewritten)+len(ending))
out = append(out, "data: "...)
out = append(out, rewritten...)
out = append(out, ending...)
return out
}
func splitLineEnding(line []byte) ([]byte, []byte) {
if len(line) == 0 || line[len(line)-1] != '\n' {
return line, nil
}
if len(line) >= 2 && line[len(line)-2] == '\r' {
return line[:len(line)-2], []byte("\r\n")
}
return line[:len(line)-1], []byte("\n")
}
func rewriteProviderJSONModel(body []byte, model string) []byte {
var raw map[string]json.RawMessage
if err := json.Unmarshal(body, &raw); err != nil {
return body
}
if _, ok := raw["model"]; !ok {
return body
}
modelJSON, err := json.Marshal(model)
if err != nil {
return body
}
raw["model"] = modelJSON
rewritten, err := json.Marshal(raw)
if err != nil {
return body
}
return rewritten
}
// Sideband extension surface (SDD S05): explicit passthrough+sideband
// responses expose IOP observations alongside provider-original content.
// Streaming responses interleave `event: iop.sideband` SSE events at provider
// event boundaries; non-streaming responses wrap the provider body in the
// iop.chat.passthrough_sideband envelope. Both are IOP extension outputs and
// are never claimed as provider-original byte identity.
const (
sidebandSSEEventName = "iop.sideband"
sidebandEnvelopeObject = "iop.chat.passthrough_sideband"
)
type sidebandRouteObservation struct {
RunID string `json:"run_id"`
NodeID string `json:"node_id,omitempty"`
Adapter string `json:"adapter,omitempty"`
Target string `json:"target,omitempty"`
ModelGroup string `json:"model_group,omitempty"`
ResponseMode string `json:"response_mode"`
ProviderStatusCode int `json:"provider_status_code,omitempty"`
}
type sidebandUsageObservation struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
}
type sidebandAssembledObservation struct {
Content string `json:"content,omitempty"`
Reasoning string `json:"reasoning,omitempty"`
ToolCallNames []string `json:"tool_call_names,omitempty"`
BodyBytes int `json:"body_bytes"`
}
// sidebandObservation is one `event: iop.sideband` SSE payload on the
// streaming sideband surface.
type sidebandObservation struct {
Object string `json:"object"`
Kind string `json:"kind"`
Route *sidebandRouteObservation `json:"route,omitempty"`
Usage *sidebandUsageObservation `json:"usage,omitempty"`
Assembled *sidebandAssembledObservation `json:"assembled,omitempty"`
}
// sidebandEnvelope is the non-streaming sideband response schema: the provider
// body is carried verbatim next to the IOP observations.
type sidebandEnvelope struct {
Object string `json:"object"`
Sideband sidebandObservations `json:"iop_sideband"`
ProviderStatusCode int `json:"provider_status_code"`
ProviderResponse json.RawMessage `json:"provider_response,omitempty"`
ProviderBody string `json:"provider_body,omitempty"`
}
type sidebandObservations struct {
Route sidebandRouteObservation `json:"route"`
Usage *sidebandUsageObservation `json:"usage,omitempty"`
Assembled *sidebandAssembledObservation `json:"assembled,omitempty"`
}
func sidebandRouteFromDispatch(dispatch edgeservice.RunDispatch, providerStatus int) sidebandRouteObservation {
return sidebandRouteObservation{
RunID: dispatch.RunID,
NodeID: dispatch.NodeID,
Adapter: dispatch.Adapter,
Target: dispatch.Target,
ModelGroup: dispatch.ModelGroupKey,
ResponseMode: responseModePassthroughSideband,
ProviderStatusCode: providerStatus,
}
}
func sidebandUsageFromFrame(frame *iop.ProviderTunnelFrame) *sidebandUsageObservation {
usage := frame.GetUsage()
if usage == nil {
return nil
}
return &sidebandUsageObservation{
InputTokens: int(usage.GetInputTokens()),
OutputTokens: int(usage.GetOutputTokens()),
}
}
// usageObservationFromProtoUsage converts a proto iop.Usage into the internal
// usageObservation used for metrics. Returns zero values when u is nil.
func usageObservationFromProtoUsage(u *iop.Usage) usageObservation {
if u == nil {
return usageObservation{}
}
return usageObservation{
inputTokens: int(u.GetInputTokens()),
outputTokens: int(u.GetOutputTokens()),
reasoningTokens: int(u.GetReasoningTokens()),
cachedInputTokens: int(u.GetCachedInputTokens()),
}
}
// mergeUsageObservation applies the sideband merge rule: use the body value
// when body observed a token type; use the proto value only when body is zero.
// This prevents double-counting input/output when both body and proto report
// the same provider usage, while preserving proto-only fields like
// reasoning/cached_input (SDD S05).
func mergeUsageObservation(body, proto usageObservation) usageObservation {
return usageObservation{
inputTokens: selectFirstNonZero(body.inputTokens, proto.inputTokens),
outputTokens: selectFirstNonZero(body.outputTokens, proto.outputTokens),
reasoningTokens: selectFirstNonZero(body.reasoningTokens, proto.reasoningTokens),
cachedInputTokens: selectFirstNonZero(body.cachedInputTokens, proto.cachedInputTokens),
reasoningChars: body.reasoningChars,
}
}
func selectFirstNonZero(a, b int) int {
if a != 0 {
return a
}
return b
}
// providerChatAssembler accumulates a human-readable view of the provider
// Chat Completions response for the sideband assembled observation. It
// tolerates non-JSON and partial payloads: whatever cannot be parsed simply
// yields an empty summary.
type providerChatAssembler struct {
streaming bool
nonStreamingParsed bool
bodyBytes int
pending []byte
content strings.Builder
reasoning strings.Builder
toolCallNames []string
// usage holds provider-reported token counts observed in the passthrough
// body. It feeds Edge-internal metrics only.
usage usageObservation
}
// providerChatDeltaEnvelope matches the provider fields the assembler
// observes; streaming uses delta, non-streaming uses message.
type providerChatDeltaEnvelope struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
Reasoning string `json:"reasoning"`
ToolCalls []struct {
Function struct {
Name string `json:"name"`
} `json:"function"`
} `json:"tool_calls"`
}
// providerUsageEnvelope matches the OpenAI-compatible usage object from
// both Chat Completions and Responses. Chat Completions uses prompt_tokens/
// completion_tokens; Responses uses input_tokens/output_tokens. The optional
// detail objects carry reasoning and cached-input tokens for both formats.
type providerUsageEnvelope struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
PromptTokensDetails *struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
CompletionTokensDetails *struct {
ReasoningTokens int `json:"reasoning_tokens"`
} `json:"completion_tokens_details"`
InputTokensDetails *struct {
CachedTokens int `json:"cached_tokens"`
} `json:"input_tokens_details"`
OutputTokensDetails *struct {
ReasoningTokens int `json:"reasoning_tokens"`
} `json:"output_tokens_details"`
}
// recordUsage stores the latest provider-reported usage. The provider tunnel
// body is never mutated; this observation feeds Edge-internal metrics only
// (SDD S05/D04). Chat Completions keys (prompt_tokens/completion_tokens) take
// precedence; if absent, Responses keys (input_tokens/output_tokens) are used.
func (a *providerChatAssembler) recordUsage(u *providerUsageEnvelope) {
if u == nil {
return
}
// Prefer Chat Completions keys; fall back to Responses keys.
if u.PromptTokens != 0 {
a.usage.inputTokens = u.PromptTokens
a.usage.outputTokens = u.CompletionTokens
if d := u.PromptTokensDetails; d != nil {
a.usage.cachedInputTokens = d.CachedTokens
}
if d := u.CompletionTokensDetails; d != nil {
a.usage.reasoningTokens = d.ReasoningTokens
}
return
}
if u.InputTokens != 0 {
a.usage.inputTokens = u.InputTokens
a.usage.outputTokens = u.OutputTokens
if d := u.InputTokensDetails; d != nil {
a.usage.cachedInputTokens = d.CachedTokens
}
if d := u.OutputTokensDetails; d != nil {
a.usage.reasoningTokens = d.ReasoningTokens
}
}
}
// recordProtoUsage stores usage carried on a provider tunnel USAGE frame. Like
// recordUsage it observes for metrics only and never alters the response body.
func (a *providerChatAssembler) recordProtoUsage(u *iop.Usage) {
if u == nil {
return
}
a.usage.inputTokens = int(u.GetInputTokens())
a.usage.outputTokens = int(u.GetOutputTokens())
a.usage.reasoningTokens = int(u.GetReasoningTokens())
a.usage.cachedInputTokens = int(u.GetCachedInputTokens())
}
// usageObservation returns the observed token usage plus the observed reasoning
// character count. Reasoning chars are never converted into a token estimate.
func (a *providerChatAssembler) usageObservation() usageObservation {
obs := a.usage
obs.reasoningChars = a.reasoning.Len()
return obs
}
func (a *providerChatAssembler) Write(chunk []byte) {
a.bodyBytes += len(chunk)
a.pending = append(a.pending, chunk...)
if !a.streaming {
return
}
for {
idx := bytes.IndexByte(a.pending, '\n')
if idx < 0 {
return
}
line := strings.TrimRight(string(a.pending[:idx]), "\r")
a.pending = a.pending[idx+1:]
a.consumeSSELine(line)
}
}
func (a *providerChatAssembler) consumeSSELine(line string) {
payload, ok := strings.CutPrefix(line, "data:")
if !ok {
return
}
payload = strings.TrimSpace(payload)
if payload == "" || payload == "[DONE]" {
return
}
var chunk struct {
Choices []struct {
Delta providerChatDeltaEnvelope `json:"delta"`
} `json:"choices"`
Usage *providerUsageEnvelope `json:"usage"`
}
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
return
}
for _, choice := range chunk.Choices {
a.consumeDelta(choice.Delta)
}
a.recordUsage(chunk.Usage)
// Responses streaming: nested response.usage from events like
// response.completed. This captures provider-reported token usage from
// the Responses API SSE payload (SDD S05/D04).
if a.usage.inputTokens == 0 && a.usage.outputTokens == 0 {
var respEvent struct {
Response struct {
Usage *providerUsageEnvelope `json:"usage"`
} `json:"response"`
}
if err := json.Unmarshal([]byte(payload), &respEvent); err == nil {
if respEvent.Response.Usage != nil {
a.recordUsage(respEvent.Response.Usage)
}
}
}
}
func (a *providerChatAssembler) consumeDelta(delta providerChatDeltaEnvelope) {
a.content.WriteString(delta.Content)
a.reasoning.WriteString(delta.ReasoningContent)
a.reasoning.WriteString(delta.Reasoning)
for _, call := range delta.ToolCalls {
if name := strings.TrimSpace(call.Function.Name); name != "" {
a.toolCallNames = append(a.toolCallNames, name)
}
}
}
func (a *providerChatAssembler) observation() *sidebandAssembledObservation {
if !a.streaming && !a.nonStreamingParsed {
a.nonStreamingParsed = true
var resp struct {
Choices []struct {
Message providerChatDeltaEnvelope `json:"message"`
} `json:"choices"`
Usage *providerUsageEnvelope `json:"usage"`
}
if err := json.Unmarshal(a.pending, &resp); err == nil {
for _, choice := range resp.Choices {
a.consumeDelta(choice.Message)
}
a.recordUsage(resp.Usage)
}
}
return &sidebandAssembledObservation{
Content: a.content.String(),
Reasoning: a.reasoning.String(),
ToolCallNames: a.toolCallNames,
BodyBytes: a.bodyBytes,
}
}
// writeProviderTunnelSidebandStream relays provider SSE bytes unchanged and
// interleaves explicit `event: iop.sideband` extension events. IOP events are
// written only at provider event boundaries (before the first body byte,
// after a "\n\n"-terminated provider event, or at end of stream) so a
// provider event is never split by sideband injection.
func (s *Server) writeProviderTunnelSidebandStream(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult) {
frames := handle.Stream().Frames
if frames == nil {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream unavailable")
return
}
flusher, _ := w.(http.Flusher)
timer := time.NewTimer(handle.WaitTimeout())
defer timer.Stop()
assembler := &providerChatAssembler{streaming: true}
wroteHeader := false
atEventBoundary := true
tail := ""
var pendingObservations []sidebandObservation
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(handle.Dispatch().ModelGroupKey), usageEndpointChatCompletions, responseModePassthroughSideband)
// metricStatus is the terminal status reported to usage metrics. It defaults
// to error and is upgraded to success on END.
metricStatus := usageStatusError
// pendingMerged accumulates usage from all USAGE frames so that emitUsageMetrics
// sees the full usage even after pendingObservations is flushed.
var pendingMerged usageObservation
writeObservation := func(obs sidebandObservation) {
obs.Object = sidebandSSEEventName
payload, err := json.Marshal(obs)
if err != nil {
return
}
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", sidebandSSEEventName, payload)
if flusher != nil {
flusher.Flush()
}
}
flushObservations := func() {
for _, obs := range pendingObservations {
writeObservation(obs)
}
pendingObservations = nil
}
defer func() {
obs := assembler.observation()
s.logger.Info("openai chat completion sideband stream closed",
zap.String("run_id", handle.Dispatch().RunID),
zap.Bool("wrote_header", wroteHeader),
zap.Int("body_bytes", assembler.bodyBytes),
zap.String("assembled_content", obs.Content),
zap.String("assembled_reasoning", obs.Reasoning),
zap.Strings("assembled_tool_calls", obs.ToolCallNames),
zap.Int("assembled_tool_call_count", len(obs.ToolCallNames)),
)
// Merge body-parsed usage (primary) with USAGE frame data
// (auxiliary reasoning/cached_input). Per the merge rule, body values
// are used when observed; proto values fill in when body is zero.
// This avoids double-counting input/output when both sources report
// the same provider usage, while preserving proto-only fields like
// reasoning/cached_input (SDD S05).
merged := mergeUsageObservation(assembler.usageObservation(), pendingMerged)
emitUsageMetrics(metricLabels, metricStatus, merged)
}()
for {
select {
case <-r.Context().Done():
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
metricStatus = usageStatusCancel
return
case <-timer.C:
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
metricStatus = usageStatusCancel
if !wroteHeader {
writeError(w, http.StatusBadGateway, "run_error", errRunTimedOut.Error())
}
return
case frame, ok := <-frames:
if !ok {
if !wroteHeader {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream closed before provider response")
}
return
}
switch frame.GetKind() {
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
if wroteHeader {
continue
}
copyProviderResponseHeaders(w.Header(), frame.GetHeaders())
// The sideband stream is an IOP extension surface: label it
// and drop the provider Content-Length, which no longer
// matches the extended body.
w.Header().Del("Content-Length")
w.Header().Set(responseModeHeaderName, responseModePassthroughSideband)
status := int(frame.GetStatusCode())
if status == 0 {
status = http.StatusOK
}
w.WriteHeader(status)
wroteHeader = true
if flusher != nil {
flusher.Flush()
}
route := sidebandRouteFromDispatch(handle.Dispatch(), status)
writeObservation(sidebandObservation{Kind: "route", Route: &route})
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
body := frame.GetBody()
if len(body) == 0 {
continue
}
if !wroteHeader {
// Defensive: a body frame before response-start still
// reaches the caller instead of being dropped.
w.Header().Set(responseModeHeaderName, responseModePassthroughSideband)
w.WriteHeader(http.StatusOK)
wroteHeader = true
}
if _, err := w.Write(body); err != nil {
// The caller is gone mid-stream; propagate cancel to the Node.
s.sendCancelRun(handle.Dispatch())
metricStatus = usageStatusCancel
return
}
assembler.Write(body)
if flusher != nil {
flusher.Flush()
}
tail = sseTail(tail, body)
atEventBoundary = strings.HasSuffix(tail, "\n\n")
if atEventBoundary {
flushObservations()
}
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
msg := frame.GetError()
if msg == "" {
msg = "provider tunnel failed"
}
if !wroteHeader {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", msg)
return
}
// Status/headers are already committed; the response is
// truncated and the caller observes the broken stream.
s.logger.Warn("openai sideband tunnel error after response start",
zap.String("run_id", handle.Dispatch().RunID),
zap.String("error", msg),
)
return
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
usage := sidebandUsageFromFrame(frame)
if usage == nil {
continue
}
// Accumulate proto usage from all USAGE frames so emitUsageMetrics
// sees the full usage even after pendingObservations is flushed.
pendingMerged = mergeUsageObservation(pendingMerged,
usageObservationFromProtoUsage(frame.GetUsage()))
pendingObservations = append(pendingObservations, sidebandObservation{Kind: "usage", Usage: usage})
if wroteHeader && atEventBoundary {
flushObservations()
}
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
if !wroteHeader {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel ended before provider response")
return
}
metricStatus = usageStatusSuccess
if !atEventBoundary {
// The provider stream ended mid-event; terminate it so the
// trailing IOP events stay well-formed SSE.
fmt.Fprint(w, "\n\n")
}
flushObservations()
writeObservation(sidebandObservation{Kind: "assembled", Assembled: assembler.observation()})
return
}
}
}
}
// sseTail keeps the last two bytes of the relayed provider stream so the
// sideband writer can detect "\n\n" event boundaries across chunk splits.
func sseTail(tail string, chunk []byte) string {
combined := tail + string(chunk)
if len(combined) > 2 {
return combined[len(combined)-2:]
}
return combined
}
// writeProviderTunnelSidebandResponse buffers the provider response and
// answers with the iop.chat.passthrough_sideband envelope: provider status is
// preserved, the provider body is carried verbatim inside the envelope, and
// IOP route/usage/assembled observations sit alongside it.
func (s *Server) writeProviderTunnelSidebandResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult) {
frames := handle.Stream().Frames
if frames == nil {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream unavailable")
return
}
timer := time.NewTimer(handle.WaitTimeout())
defer timer.Stop()
assembler := &providerChatAssembler{}
var body bytes.Buffer
providerStatus := 0
var protoUsage *sidebandUsageObservation
var protoObs usageObservation
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(handle.Dispatch().ModelGroupKey), usageEndpointChatCompletions, responseModePassthroughSideband)
// metricStatus is the terminal status reported to usage metrics. It defaults
// to error and is upgraded to success on END.
metricStatus := usageStatusError
defer func() {
obs := assembler.observation()
s.logger.Info("openai chat completion sideband response",
zap.String("run_id", handle.Dispatch().RunID),
zap.Int("provider_status", providerStatus),
zap.Int("body_bytes", body.Len()),
zap.String("assembled_content", obs.Content),
zap.String("assembled_reasoning", obs.Reasoning),
zap.Strings("assembled_tool_calls", obs.ToolCallNames),
zap.Int("assembled_tool_call_count", len(obs.ToolCallNames)),
)
// Use body-parsed usage as primary source; layer in USAGE frame
// data per the merge rule. sidebandUsageObservation in the response
// schema carries only input/output, while the metric path uses the
// full breakdown including reasoning/cached_input from the provider
// body or proto (SDD S05).
merged := mergeUsageObservation(assembler.usageObservation(), protoObs)
emitUsageMetrics(metricLabels, metricStatus, merged)
}()
for {
select {
case <-r.Context().Done():
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
metricStatus = usageStatusCancel
return
case <-timer.C:
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
metricStatus = usageStatusCancel
writeError(w, http.StatusBadGateway, "run_error", errRunTimedOut.Error())
return
case frame, ok := <-frames:
if !ok {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel stream closed before provider response")
return
}
switch frame.GetKind() {
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
if providerStatus != 0 {
continue
}
providerStatus = int(frame.GetStatusCode())
if providerStatus == 0 {
providerStatus = http.StatusOK
}
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
body.Write(frame.GetBody())
assembler.Write(frame.GetBody())
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
msg := frame.GetError()
if msg == "" {
msg = "provider tunnel failed"
}
writeError(w, http.StatusBadGateway, "provider_tunnel_error", msg)
return
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
if u := sidebandUsageFromFrame(frame); u != nil {
protoUsage = u
}
protoObs = mergeUsageObservation(protoObs, usageObservationFromProtoUsage(frame.GetUsage()))
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
if providerStatus == 0 && body.Len() == 0 {
writeError(w, http.StatusBadGateway, "provider_tunnel_error", "tunnel ended before provider response")
return
}
metricStatus = usageStatusSuccess
if providerStatus == 0 {
providerStatus = http.StatusOK
}
envelope := sidebandEnvelope{
Object: sidebandEnvelopeObject,
Sideband: sidebandObservations{
Route: sidebandRouteFromDispatch(handle.Dispatch(), providerStatus),
Usage: protoUsage,
Assembled: assembler.observation(),
},
ProviderStatusCode: providerStatus,
}
raw := body.Bytes()
if json.Valid(raw) {
envelope.ProviderResponse = json.RawMessage(raw)
} else if len(raw) > 0 {
envelope.ProviderBody = string(raw)
}
w.Header().Set(responseModeHeaderName, responseModePassthroughSideband)
writeJSON(w, providerStatus, envelope)
return
}
}
}
}
// hopByHopResponseHeaders are transport-level headers owned by each hop; they
// are not copied from the provider response to the caller response.
var hopByHopResponseHeaders = map[string]struct{}{
"Connection": {},
"Keep-Alive": {},
"Proxy-Authenticate": {},
"Proxy-Authorization": {},
"Te": {},
"Trailer": {},
"Transfer-Encoding": {},
"Upgrade": {},
}
func copyProviderResponseHeaders(dst http.Header, headers map[string]string) {
for k, v := range headers {
if _, hop := hopByHopResponseHeaders[http.CanonicalHeaderKey(k)]; hop {
continue
}
dst.Set(k, v)
}
}
type streamToolTextFilter struct {
pending string
}
func (f *streamToolTextFilter) Append(delta string) string {
f.pending += delta
if idx := firstStreamToolTextCandidateIndex(f.pending); idx >= 0 {
out := strings.TrimRightFunc(f.pending[:idx], unicode.IsSpace)
f.pending = f.pending[idx:]
return out
}
flushLen := streamToolTextSafeFlushLen(f.pending)
out := strings.TrimRightFunc(f.pending[:flushLen], unicode.IsSpace)
f.pending = f.pending[len(out):]
return out
}
func (f *streamToolTextFilter) Flush() string {
out := f.pending
f.pending = ""
return out
}
func (f *streamToolTextFilter) FlushNonCandidate() string {
if firstStreamToolTextCandidateIndex(f.pending) >= 0 {
f.pending = ""
return ""
}
// Even without a full candidate, chunk-boundary protection can leave a
// partial candidate suffix (e.g. "<tool_ca", trailing "{") in pending.
// Return only the safe non-candidate prefix and drop the partial suffix so
// a native tool_calls completion never flushes a raw candidate fragment.
flushLen := streamToolTextSafeFlushLen(f.pending)
out := strings.TrimRightFunc(f.pending[:flushLen], unicode.IsSpace)
f.pending = ""
return out
}
func firstStreamToolTextCandidateIndex(s string) int {
xml := strings.Index(strings.ToLower(s), "<tool_call")
mustache := strings.Index(s, textMustacheToolCallOpenHint)
switch {
case xml < 0:
return mustache
case mustache < 0:
return xml
case xml < mustache:
return xml
default:
return mustache
}
}
func streamToolTextSafeFlushLen(s string) int {
if s == "" {
return 0
}
const xmlHint = "<tool_call"
lower := strings.ToLower(s)
start := len(s) - len(xmlHint) + 1
if start < 0 {
start = 0
}
for i := start; i < len(s); i++ {
if strings.HasPrefix(xmlHint, lower[i:]) ||
strings.HasPrefix(textMustacheToolCallOpenHint, s[i:]) {
return i
}
}
return len(s)
}
func unstreamedCleanedText(cleaned, emitted string) string {
if cleaned == "" {
return ""
}
if emitted == "" {
return cleaned
}
if strings.HasPrefix(cleaned, emitted) {
return cleaned[len(emitted):]
}
cleaned = strings.TrimSpace(cleaned)
emitted = strings.TrimSpace(emitted)
if cleaned == "" || cleaned == emitted {
return ""
}
if emitted != "" && strings.HasPrefix(cleaned, emitted) {
return strings.TrimSpace(strings.TrimPrefix(cleaned, emitted))
}
return ""
}
func writeContentDeltaSSE(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string, content string) {
if content == "" {
return
}
writeSSE(w, flusher, chatCompletionChunk{
ID: id,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []chatCompletionChunkChoice{{
Index: 0,
Delta: chatDelta{Content: content},
}},
})
}
func writeToolCallsDeltaSSE(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string, toolCalls []any) {
if len(toolCalls) == 0 {
return
}
writeSSE(w, flusher, chatCompletionChunk{
ID: id,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []chatCompletionChunkChoice{{
Index: 0,
Delta: chatDelta{ToolCalls: toolCallsForStreamDelta(toolCalls)},
}},
})
}
// streamBufferedChatCompletion collects the full run output, validates tool
// calls against the request schema before emitting anything, and — when
// validation fails on a run that has not yet written a user-visible chunk —
// resubmits the same request as a bounded exact replay. Only a validated (or
// validation-disabled) result reaches the SSE writer; a malformed tool call
// that survives the retry budget is surfaced as an SSE error instead of a
// successful tool_calls chunk.
func (s *Server) streamBufferedChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, submitReq edgeservice.SubmitRunRequest, handle edgeservice.RunResult, flusher http.Flusher, outputPolicy strictOutputPolicy, validation toolValidationContract) {
attempt := 1
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, chatResponseModeLabel(req))
for {
result, err := collectChatCompletionOutput(r.Context(), req, handle, outputPolicy)
if err != nil {
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), err)
handle.Close()
emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
writeSSEError(w, flusher, err.Error())
return
}
verr := result.toolValidationErr
if verr == nil {
verr = validateToolCallResponse(validation, result.toolCalls, result.toolCallOrigin)
}
if verr != nil {
failedRunID := handle.Dispatch().RunID
if attempt < maxToolValidationAttempts {
handle.Close()
attempt++
retryReq := submitReq
retryReq.Metadata = toolValidationAttemptMetadata(submitReq.Metadata, attempt, failedRunID, verr.Error())
s.logger.Warn("openai chat completion stream tool validation retry",
zap.String("run_id", failedRunID),
zap.Int("attempt", attempt),
zap.String("reason", verr.Error()),
)
next, submitErr := s.service.SubmitRun(r.Context(), retryReq)
if submitErr != nil {
emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
writeSSEErrorWithType(w, flusher, "tool_validation_retry_error", submitErr.Error())
return
}
handle = next
continue
}
handle.Close()
s.logger.Warn("openai chat completion stream tool validation failed",
zap.String("run_id", failedRunID),
zap.Int("attempt", attempt),
zap.String("reason", verr.Error()),
)
emitUsageMetrics(metricLabels, usageStatusError, usageObservation{})
writeSSEErrorWithType(w, flusher, "tool_validation_error", verr.Error())
return
}
handle.Close()
s.writeBufferedStreamOutput(w, flusher, req, handle.Dispatch(), result, outputPolicy, openAICompatTraceStreamEnabled(submitReq.Metadata))
emitUsageMetrics(metricLabels, usageStatusSuccess, usageObservationFromOpenAIUsage(result.usage, result.reasoningLen))
return
}
}
// writeBufferedStreamOutput emits a validated buffered result as SSE chunks:
// the role chunk, an optional content chunk, an optional tool_calls chunk, and
// the terminal finish chunk.
func (s *Server) writeBufferedStreamOutput(w http.ResponseWriter, flusher http.Flusher, req chatCompletionRequest, dispatch edgeservice.RunDispatch, result chatCompletionOutput, outputPolicy strictOutputPolicy, traceStream bool) {
created := time.Now().Unix()
id := "chatcmpl-" + dispatch.RunID
model := responseModel(req.Model, dispatch.Target)
content := result.message.Content
toolCallNames := extractToolCallNames(result.toolCalls)
s.logger.Info("openai chat completion stream closed",
zap.String("run_id", dispatch.RunID),
zap.Bool("strict_output", outputPolicy.Strict),
zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
zap.Bool("normalized", result.normalized),
zap.String("response_mode", result.responseMode),
zap.Int("content_len", len(content)),
zap.Int("reasoning_len", result.reasoningLen),
zap.String("assembled_content", content),
zap.String("assembled_reasoning", result.message.ReasoningContent),
zap.Strings("assembled_tool_calls", toolCallNames),
zap.Int("assembled_tool_call_count", len(toolCallNames)),
)
writeSSE(w, flusher, chatCompletionChunk{
ID: id,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []chatCompletionChunkChoice{{
Index: 0,
Delta: chatDelta{Role: "assistant"},
}},
})
if traceStream {
logOpenAICompatStreamOutput(s.logger, dispatch.RunID, 1, "content", content)
}
writeContentDeltaSSE(w, flusher, id, created, model, content)
if traceStream && len(result.toolCalls) > 0 {
s.logger.Info("openai chat completion stream output chunk",
zap.String("run_id", dispatch.RunID),
zap.Int("seq", 2),
zap.String("type", "tool_calls"),
zap.Int("tool_call_count", len(result.toolCalls)),
)
}
writeToolCallsDeltaSSE(w, flusher, id, created, model, result.toolCalls)
if traceStream {
s.logger.Info("openai chat completion stream output done",
zap.String("run_id", dispatch.RunID),
zap.Int("seq", 3),
zap.String("finish_reason", result.finishReason),
zap.Int("content_len", len(content)),
zap.Int("reasoning_len", result.reasoningLen),
)
}
writeSSE(w, flusher, chatCompletionChunk{
ID: id,
Object: "chat.completion.chunk",
Created: created,
Model: model,
Choices: []chatCompletionChunkChoice{{
Index: 0,
Delta: chatDelta{},
FinishReason: result.finishReason,
}},
})
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
}
func writeSSE(w http.ResponseWriter, flusher http.Flusher, v any) {
b, err := json.Marshal(v)
if err != nil {
return
}
fmt.Fprintf(w, "data: %s\n\n", b)
flusher.Flush()
}
func writeSSEError(w http.ResponseWriter, flusher http.Flusher, message string) {
writeSSEErrorWithType(w, flusher, "run_error", message)
}
func writeSSEErrorWithType(w http.ResponseWriter, flusher http.Flusher, errType, message string) {
writeSSE(w, flusher, errorResponse{Error: errorBody{Type: errType, Message: message}})
fmt.Fprint(w, "data: [DONE]\n\n")
flusher.Flush()
}
func openAICompatTraceStreamEnabled(metadata map[string]string) bool {
if traceBool(os.Getenv(streamTraceEnvKey)) {
return true
}
if metadata == nil {
return false
}
return traceBool(metadata[streamTraceMetadataKey])
}
func traceBool(v string) bool {
switch strings.ToLower(strings.TrimSpace(v)) {
case "1", "true", "yes", "on":
return true
default:
return false
}
}
func logOpenAICompatStreamOutput(logger *zap.Logger, runID string, seq int, typ string, delta string, extra ...zap.Field) {
fields := []zap.Field{
zap.String("run_id", runID),
zap.Int("seq", seq),
zap.String("type", typ),
zap.Int("delta_len", len(delta)),
zap.Bool("delta_has_open_think", hasOpenThinkTag(delta)),
zap.Bool("delta_has_close_think", hasCloseThinkTag(delta)),
}
fields = append(fields, extra...)
logger.Info("openai chat completion stream output chunk", fields...)
}
func hasOpenThinkTag(s string) bool {
return strings.Contains(strings.ToLower(s), "<think")
}
func hasCloseThinkTag(s string) bool {
return strings.Contains(strings.ToLower(s), "</think")
}