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

package openai

import (
	"context"
	"encoding/json"
	"net/http"
	"strings"

	edgeservice "iop/apps/edge/internal/service"
)

// openAIRequestContext carries the routing, identity, and metrics values that
// every OpenAI-compatible surface resolves once at ingress, before any
// path-specific decode or dispatch. Chat and Responses share it so both
// surfaces classify and meter a request the same way.
//
// It is built at a single ingress point per surface and read-only afterwards.
type openAIRequestContext struct {
	r     *http.Request
	route routeDispatch
	// rawBody is the caller's original body, preserved verbatim for the tunnel
	// passthrough path. The rewritten provider body is produced separately by
	// the tunnel BuildBody hook and never replaces rawBody.
	rawBody []byte
	// callerMetadata is the parsed request metadata with the authenticated
	// principal applied over any caller-supplied identity. Context constructors
	// copy it before adding path-specific fields, so this ingress snapshot stays
	// immutable while the request moves through tunnel and normalized branches.
	callerMetadata map[string]string
	workspace      string
	estimate       int
	contextClass   string
	endpoint       string
}

// cloneMetadata derives a request-local metadata map. OpenAI ingress metadata
// is shared by multiple possible dispatch paths, while service requests need
// to add execution fields such as openai_stream and context_class. Copying at
// that boundary prevents one derived path from changing another's view.
func cloneMetadata(metadata map[string]string) map[string]string {
	cloned := make(map[string]string, len(metadata))
	for key, value := range metadata {
		cloned[key] = value
	}
	return cloned
}

// resolveCallerIdentity parses the caller metadata, applies the authenticated
// principal over it, and validates the workspace against the resolved route.
// Every OpenAI-compatible surface resolves identity through this one path so a
// caller-supplied iop_principal_* value can never win on any of them.
func resolveCallerIdentity(r *http.Request, route routeDispatch, rawMetadata json.RawMessage) (map[string]string, string, error) {
	runMeta, workspace, err := parseOpenAIMetadata(rawMetadata)
	if err != nil {
		return nil, "", err
	}
	// Overwrite (not merge-if-absent): the authenticated caller identity must
	// win over any caller-supplied metadata.iop_principal_* spoof attempt.
	for k, v := range principalMetadata(r.Context()) {
		runMeta[k] = v
	}
	if err := validateWorkspaceForRoute(route, workspace); err != nil {
		return nil, "", err
	}
	return runMeta, workspace, nil
}

// chatDispatchContext is the immutable per-request context of a dispatched chat
// completion. Everything downstream of ingress — provider-pool dispatch, live
// SSE, buffered SSE, and the non-stream completion — reads this context instead
// of taking the same values as a long parameter list.
//
// req is the effective request: the model-catalog generation policy is applied
// once at ingress, before the context is built, so no consumer sees a
// pre-policy request or re-applies the policy.
type chatDispatchContext struct {
	openAIRequestContext

	req          chatCompletionRequest
	basePrompt   string
	prompt       string
	input        map[string]any
	runMetadata  map[string]string
	outputPolicy strictOutputPolicy
	validation   toolValidationContract

	// submitReq is the run request for the resolved path: provider-pool shaped
	// (adapter/target resolved per candidate at admission) or direct.
	submitReq edgeservice.SubmitRunRequest
	// retrySubmit resubmits the run for a bounded tool-validation replay,
	// preserving the path the request was dispatched on.
	retrySubmit func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error)
}

// responsesRequestContext is the immutable ingress context for a
// /v1/responses request. It keeps the lenient routing envelope next to the
// common request context so raw provider tunnel paths never need to pass a
// parallel list of route, body, metadata, and estimate arguments.
type responsesRequestContext struct {
	openAIRequestContext
	envelope responsesEnvelope
}

// responsesDispatchContext is the normalized /v1/responses dispatch context.
// It is derived only after strict decode and policy application. The raw body
// remains on responsesRequestContext for the tunnel path; the normalized run
// request is represented separately by submitReq.
type responsesDispatchContext struct {
	*responsesRequestContext

	req          responsesRequest
	prompt       string
	input        map[string]any
	outputPolicy strictOutputPolicy
	runMetadata  map[string]string
	submitReq    edgeservice.SubmitRunRequest
}

func (dc *responsesDispatchContext) usageLabels(s *Server, responseMode string) usageLabels {
	return s.usageLabelsFor(dc.r.Context(), strings.TrimSpace(dc.req.Model), dc.endpoint, responseMode)
}

// withRetrySubmit derives a context bound to a retry function. The
// provider-pool retry closure can only be built once the pool request exists,
// so the pool path derives a new context rather than mutating this one.
func (dc *chatDispatchContext) withRetrySubmit(fn func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error)) *chatDispatchContext {
	derived := *dc
	derived.retrySubmit = fn
	return &derived
}

// usageLabels resolves the metric labels for this request's response mode.
func (dc *chatDispatchContext) usageLabels(s *Server, responseMode string) usageLabels {
	return s.usageLabelsFor(dc.r.Context(), strings.TrimSpace(dc.req.Model), dc.endpoint, responseMode)
}