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

400 lines
15 KiB
Go

package openai
import (
"errors"
"go.uber.org/zap"
edgeservice "iop/apps/edge/internal/service"
iop "iop/proto/gen/iop"
"net/http"
"strconv"
"strings"
"time"
)
// 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, dc *chatDispatchContext) {
handle, ok := s.submitChatCompletionTunnel(w, dc)
if !ok {
return
}
defer handle.Close()
metricLabels := dc.usageLabels(s, responseModePassthrough)
s.writeProviderTunnelResponse(w, dc.r, handle, dc.req.Stream, dc.req.Model, metricLabels)
}
// 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 for pure passthrough. On failure it writes the dispatch error to the
// caller and returns ok=false.
func (s *Server) submitChatCompletionTunnel(w http.ResponseWriter, dc *chatDispatchContext) (edgeservice.ProviderTunnelResult, bool) {
providerAuthHeaders, err := s.providerTunnelAuthHeaders(dc.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 := cloneMetadata(dc.callerMetadata)
metadata["openai_model"] = dc.req.Model
metadata["openai_stream"] = strconv.FormatBool(dc.req.Stream)
metadata["estimated_input_tokens"] = strconv.Itoa(dc.estimate)
metadata["context_class"] = dc.contextClass
tunnelReq := edgeservice.SubmitProviderTunnelRequest{
NodeRef: dc.route.NodeRef,
ModelGroupKey: strings.TrimSpace(dc.req.Model),
Adapter: dc.route.Adapter,
Target: dc.route.Target,
SessionID: dc.route.SessionID,
Method: http.MethodPost,
Path: "/v1/chat/completions",
Headers: providerAuthHeaders,
BuildBody: func(target string) ([]byte, error) {
return rewriteChatCompletionModel(dc.rawBody, target, dc.req)
},
Stream: dc.req.Stream,
TimeoutSec: dc.route.TimeoutSec,
MaxQueue: dc.route.MaxQueue,
QueueTimeoutMS: dc.route.QueueTimeoutMS,
Metadata: metadata,
EstimatedInputTokens: dc.estimate,
ContextClass: dc.contextClass,
ProviderPool: dc.route.ProviderPool,
}
metricLabels := dc.usageLabels(s, responseModePassthrough)
handle, err := s.service.SubmitProviderTunnel(dc.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("provider_id", handle.Dispatch().ProviderID),
zap.String("provider_type", handle.Dispatch().ProviderType),
zap.String("execution_path", handle.Dispatch().ExecutionPath),
zap.String("model_group", handle.Dispatch().ModelGroupKey),
zap.String("adapter", handle.Dispatch().Adapter),
zap.String("target", handle.Dispatch().Target),
zap.Bool("stream", dc.req.Stream),
zap.String("response_mode", responseModePassthrough),
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
}
}
}
}
// 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)
}
}
// tunnelResponsesPassthrough serves a /v1/responses request over the raw
// provider tunnel (SDD S04): the caller's original body is forwarded with only
// the model field rewritten to the served target, provider auth is injected
// from the configured request header, and provider status/headers/body bytes
// are relayed to the caller unmodified. Streaming is honored when the caller
// requested it. No IOP sideband fields, model-echo rewrite, or output-token
// normalization are applied.
func (s *Server) tunnelResponsesPassthrough(w http.ResponseWriter, requestCtx *responsesRequestContext) {
providerAuthHeaders, err := s.providerTunnelAuthHeaders(requestCtx.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
}
metadata := cloneMetadata(requestCtx.callerMetadata)
metadata["openai_model"] = requestCtx.envelope.Model
metadata["openai_stream"] = strconv.FormatBool(requestCtx.envelope.Stream)
metadata["estimated_input_tokens"] = strconv.Itoa(requestCtx.estimate)
metadata["context_class"] = requestCtx.contextClass
tunnelReq := edgeservice.SubmitProviderTunnelRequest{
NodeRef: requestCtx.route.NodeRef,
ModelGroupKey: strings.TrimSpace(requestCtx.envelope.Model),
Adapter: requestCtx.route.Adapter,
Target: requestCtx.route.Target,
SessionID: requestCtx.route.SessionID,
Method: http.MethodPost,
Path: "/v1/responses",
Headers: providerAuthHeaders,
BuildBody: func(target string) ([]byte, error) {
return rewriteResponsesModel(requestCtx.rawBody, target)
},
Stream: requestCtx.envelope.Stream,
TimeoutSec: requestCtx.route.TimeoutSec,
MaxQueue: requestCtx.route.MaxQueue,
QueueTimeoutMS: requestCtx.route.QueueTimeoutMS,
Metadata: metadata,
EstimatedInputTokens: requestCtx.estimate,
ContextClass: requestCtx.contextClass,
ProviderPool: requestCtx.route.ProviderPool,
}
metricLabels := s.usageLabelsFor(requestCtx.r.Context(), strings.TrimSpace(requestCtx.envelope.Model), requestCtx.endpoint, responseModePassthrough)
handle, err := s.service.SubmitProviderTunnel(requestCtx.r.Context(), tunnelReq)
if err != nil {
emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
return
}
defer handle.Close()
s.logger.Info("openai responses passthrough dispatch",
zap.String("run_id", handle.Dispatch().RunID),
zap.String("node_id", handle.Dispatch().NodeID),
zap.String("provider_id", handle.Dispatch().ProviderID),
zap.String("provider_type", handle.Dispatch().ProviderType),
zap.String("execution_path", handle.Dispatch().ExecutionPath),
zap.String("model_group", handle.Dispatch().ModelGroupKey),
zap.String("adapter", handle.Dispatch().Adapter),
zap.String("target", handle.Dispatch().Target),
zap.Bool("stream", requestCtx.envelope.Stream),
zap.Int("estimated_input_tokens", handle.Dispatch().EstimatedInputTokens),
zap.String("context_class", handle.Dispatch().ContextClass),
zap.String("queue_reason", handle.Dispatch().QueueReason),
)
// requestModel is left empty so the shared tunnel writer relays provider
// bytes verbatim without rewriting the provider-echoed model back to a
// caller alias: Responses passthrough prefers provider-original bytes.
// metricLabels carries endpoint=responses, response_mode=passthrough, and
// the request model alias so success/error usage is attributed correctly.
s.writeProviderTunnelResponse(w, requestCtx.r, handle, requestCtx.envelope.Stream, "", metricLabels)
}