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

319 lines
13 KiB
Go

package openai
import (
"bytes"
"context"
"encoding/json"
"errors"
"go.uber.org/zap"
"io"
edgeservice "iop/apps/edge/internal/service"
"net/http"
"strconv"
"strings"
)
func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
return
}
defer r.Body.Close()
// The raw body is kept for the provider tunnel passthrough path, which
// forwards the caller's own payload (model rewritten to the served target).
rawBody, err := io.ReadAll(r.Body)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid_request_error", "failed to read request body")
return
}
// The route decision runs before strict normalization: only the
// routing-relevant envelope fields are decoded leniently so provider
// passthrough can forward Codex/provider-specific payloads verbatim.
// Strict field validation stays on the non-provider path.
env, err := decodeChatCompletionEnvelope(rawBody)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
return
}
dispatch, ok := s.resolveRouteDispatch(env.Model)
if !ok {
writeError(w, http.StatusBadRequest, "invalid_request_error", "model is required")
return
}
var req chatCompletionRequest
if dispatch.ProviderPool {
// Provider-pool path: lenient decode so caller unknown/Codex/
// provider-specific fields survive into the tunnel body.
if err := decodeChatCompletionRequestLenient(json.NewDecoder(bytes.NewReader(rawBody)), &req); err != nil {
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
return
}
} else {
// Non-provider path keeps the strict allowlist-based decode.
if err := decodeChatCompletionRequest(json.NewDecoder(bytes.NewReader(rawBody)), &req); err != nil {
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
return
}
}
runMeta, workspace, err := resolveCallerIdentity(r, dispatch, req.Metadata)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
return
}
providerTunnelRoute := routeUsesProviderTunnel(dispatch)
basePrompt := promptFromMessages(req.Messages)
if strings.TrimSpace(basePrompt) == "" {
writeError(w, http.StatusBadRequest, "invalid_request_error", "messages are required")
return
}
outputPolicy := s.resolveOutputPolicy(basePrompt)
// The mutable model-catalog generation policy is applied here, at the single
// ingress point, before req is frozen into the dispatch context. Every stage
// downstream reads the post-policy request.
if catalogEntry := s.findProviderPoolEntry(req.Model); catalogEntry != nil {
applyModelCatalogGenerationPolicyToChat(&req, *catalogEntry, outputPolicy.Strict, chatRequestHasProviderNativeThinking(rawBody))
}
// The response path is decided by the resolved route, never by caller
// metadata: provider routes relay pure passthrough over the raw tunnel;
// every other route uses the normalized RunEvent path. Caller metadata is
// arbitrary context and does not select route or response shape (SDD S01).
// Pre-compute estimate/contextClass early so provider-pool path can use them.
estimate := s.estimateChatInputTokens(basePrompt, runMeta, req.Tools, req.ToolChoice)
contextClass := classifyContext(estimate, s.longContextThreshold())
requestCtx := openAIRequestContext{
r: r,
route: dispatch,
rawBody: rawBody,
callerMetadata: runMeta,
workspace: workspace,
estimate: estimate,
contextClass: contextClass,
endpoint: usageEndpointChatCompletions,
}
// Provider-pool: use one-shot SubmitProviderPool which dispatches either
// tunnel or normalized based on the selected provider's executionPath.
dc := s.newChatDispatchContext(requestCtx, req, basePrompt, outputPolicy)
if dispatch.ProviderPool {
s.handleChatCompletionsProviderPool(w, dc)
return
}
if providerTunnelRoute {
s.tunnelChatCompletionPassthrough(w, dc)
return
}
metricLabels := dc.usageLabels(s, responseModeNormalized)
handle, err := s.service.SubmitRun(r.Context(), dc.submitReq)
if err != nil {
emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
return
}
s.logChatDispatch("openai chat completion dispatch", handle.Dispatch())
dc = dc.withRetrySubmit(func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error) {
return s.service.SubmitRun(ctx, req)
})
if req.Stream {
s.streamChatCompletion(w, dc, handle)
return
}
s.completeChatCompletion(w, dc, handle)
}
// newChatDispatchContext freezes the ingress-resolved values into the immutable
// dispatch context for the resolved route. The provider-pool route leaves
// adapter/target unset: the service layer resolves them per candidate at
// admission and rewrites the target to the winning provider's served model.
func (s *Server) newChatDispatchContext(requestCtx openAIRequestContext, req chatCompletionRequest, basePrompt string, outputPolicy strictOutputPolicy) *chatDispatchContext {
dc := &chatDispatchContext{
openAIRequestContext: requestCtx,
req: req,
basePrompt: basePrompt,
outputPolicy: outputPolicy,
validation: toolValidationContractFromRequest(req),
}
messages := req.Messages
if !requestCtx.route.ProviderPool {
// Strict output is a normalized-dispatch contract only; the pool route
// decides the path after admission, so the instruction is not baked in
// here (SDD D02).
if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" {
messages = prependSystemMessage(messages, instruction)
}
}
dc.prompt = promptFromMessages(messages)
dc.input = req.runInput(dc.prompt, messages, outputPolicy.Strict, routeSupportsNativeToolCalls(requestCtx.route))
if !requestCtx.route.ProviderPool {
// Re-estimate against the instruction-augmented prompt for the
// normalized path.
dc.estimate = s.estimateChatInputTokens(dc.prompt, requestCtx.callerMetadata, req.Tools, req.ToolChoice)
dc.contextClass = classifyContext(dc.estimate, s.longContextThreshold())
}
dc.runMetadata = chatRunMetadata(cloneMetadata(requestCtx.callerMetadata), req, outputPolicy)
dc.runMetadata["estimated_input_tokens"] = strconv.Itoa(dc.estimate)
dc.runMetadata["context_class"] = dc.contextClass
if requestCtx.route.ProviderPool {
dc.submitReq = edgeservice.SubmitRunRequest{
NodeRef: requestCtx.route.NodeRef,
ModelGroupKey: strings.TrimSpace(req.Model),
SessionID: requestCtx.route.SessionID,
Workspace: requestCtx.workspace,
Prompt: dc.prompt,
Input: dc.input,
TimeoutSec: requestCtx.route.TimeoutSec,
MaxQueue: requestCtx.route.MaxQueue,
QueueTimeoutMS: requestCtx.route.QueueTimeoutMS,
Metadata: dc.runMetadata,
EstimatedInputTokens: dc.estimate,
ContextClass: dc.contextClass,
ProviderPool: true,
}
return dc
}
if dc.validation.enabled {
dc.runMetadata = toolValidationAttemptMetadata(dc.runMetadata, 1, "", "")
}
dc.submitReq = chatSubmitRunRequest(requestCtx.route, req, requestCtx.workspace, dc.prompt, dc.input, dc.runMetadata)
dc.submitReq.EstimatedInputTokens = dc.estimate
dc.submitReq.ContextClass = dc.contextClass
return dc
}
func (s *Server) logChatDispatch(msg string, disp edgeservice.RunDispatch, extra ...zap.Field) {
fields := []zap.Field{
zap.String("run_id", disp.RunID),
zap.String("node_id", disp.NodeID),
zap.String("provider_id", disp.ProviderID),
zap.String("provider_type", disp.ProviderType),
zap.String("execution_path", disp.ExecutionPath),
zap.String("model_group", disp.ModelGroupKey),
zap.String("adapter", disp.Adapter),
zap.String("target", disp.Target),
zap.Int("estimated_input_tokens", disp.EstimatedInputTokens),
zap.String("context_class", disp.ContextClass),
zap.String("queue_reason", disp.QueueReason),
}
s.logger.Info(msg, append(fields, extra...)...)
}
// handleChatCompletionsProviderPool dispatches a provider-pool chat completion
// through the one-shot SubmitProviderPool surface. After a single queue admission,
// the selected candidate's executionPath determines whether the response is
// delivered via tunnel passthrough or the normalized RunEvent path.
//
// Provider auth is validated via providerTunnelAuthHeaders which is called for
// tunnel-path only. For normalized paths, the auth check is skipped (SDD D01).
func (s *Server) handleChatCompletionsProviderPool(w http.ResponseWriter, dc *chatDispatchContext) {
r := dc.r
req := dc.req
poolReq := edgeservice.ProviderPoolDispatchRequest{
Run: dc.submitReq,
Tunnel: edgeservice.SubmitProviderTunnelRequest{
ModelGroupKey: strings.TrimSpace(req.Model),
SessionID: dc.route.SessionID,
Method: http.MethodPost,
Path: "/v1/chat/completions",
Stream: req.Stream,
TimeoutSec: dc.route.TimeoutSec,
MaxQueue: dc.route.MaxQueue,
QueueTimeoutMS: dc.route.QueueTimeoutMS,
Metadata: dc.runMetadata,
EstimatedInputTokens: dc.estimate,
ContextClass: dc.contextClass,
ProviderPool: true,
},
}
// Pre-dispatch provider auth header injection. Runs inside SubmitProviderPool
// BEFORE buildProviderTunnelRequest and the Node Send step, so the auth
// header lands on the actual wire request. On failure the slot is released
// and no tunnel request is sent (SDD S03).
poolReq.PrepareTunnel = func(tunnelReq edgeservice.SubmitProviderTunnelRequest) (edgeservice.SubmitProviderTunnelRequest, error) {
headers, err := s.providerTunnelAuthHeaders(r)
if err != nil {
return tunnelReq, err
}
tunnelReq.Headers = headers
return tunnelReq, nil
}
// strict-output output policy only applies to normalized dispatch,
// not to raw tunnel passthrough (SDD D02). The caller's raw body is the
// passthrough source; only the model field is rewritten to the served target.
poolReq.Tunnel.BuildBody = func(target string) ([]byte, error) {
return rewriteChatCompletionModel(dc.rawBody, target, req)
}
// Normalized branch: chat completions use lenient decode so unknown
// fields survive; PrepareRun passes through the already-built Run
// without strict validation (SDD S02/S03). The tunnel branch is
// unaffected.
poolReq.PrepareRun = func(runReq edgeservice.SubmitRunRequest) (edgeservice.SubmitRunRequest, error) {
return runReq, nil
}
result, err := s.service.SubmitProviderPool(r.Context(), poolReq)
if err != nil {
emitUsageMetrics(dc.usageLabels(s, responseModePassthrough), usageStatusForError(err), usageObservation{})
// Provider auth failure is a client request error (400), not a
// backend dispatch error. The auth check now runs inside
// SubmitProviderPool (PrepareTunnel) before any tunnel request is sent.
if errors.Is(err, errProviderAuthRequired) {
writeError(w, http.StatusBadRequest, "invalid_request_error", "provider auth token is required")
return
}
writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
return
}
s.logChatDispatch("openai chat completion provider-pool dispatch", result.DispatchInfo,
zap.String("path", string(result.Path)),
)
switch result.Path {
case edgeservice.ProviderPoolPathTunnel:
// Tunnel path: provider auth was already validated and injected via
// PrepareTunnel before dispatch; on failure SubmitProviderPool returns
// an error and no tunnel handle exists. Provider bytes are relayed as
// pure passthrough; caller metadata never selects a sideband surface.
s.writeProviderTunnelResponse(w, r, result.Tunnel, req.Stream, req.Model, dc.usageLabels(s, responseModePassthrough))
case edgeservice.ProviderPoolPathNormalized:
// Normalized path: no auth required, collect from RunEvent stream.
handle := result.Run
if handle == nil {
writeError(w, http.StatusInternalServerError, "run_error", "provider-pool selection returned normalized path but no run result")
return
}
// Retry must re-enter SubmitProviderPool (not SubmitRun) so a bounded
// tool-validation replay keeps the ModelGroupKey, provider-pool
// metadata, and input of the original dispatch.
poolDC := dc.withRetrySubmit(func(ctx context.Context, retryReq edgeservice.SubmitRunRequest) (any, error) {
return s.service.SubmitProviderPool(ctx, edgeservice.ProviderPoolDispatchRequest{
Run: retryReq,
Tunnel: poolReq.Tunnel,
PrepareTunnel: poolReq.PrepareTunnel,
PrepareRun: poolReq.PrepareRun,
})
})
if req.Stream {
s.streamChatCompletion(w, poolDC, handle)
} else {
s.completeChatCompletion(w, poolDC, handle)
}
}
}