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

package openai

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

	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
	// ingress is the sole request-local owner of the canonical caller body and
	// bounded typed ledger. Auth headers are intentionally outside this owner.
	ingress *openAIIngressSnapshot
	// 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
	usage          *openAIUsageRecorder
}

func (c openAIRequestContext) canonicalBody() ([]byte, error) {
	if c.ingress == nil {
		return nil, fmt.Errorf("OpenAI ingress snapshot is unavailable")
	}
	return c.ingress.canonicalBody()
}

// 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. It is only used by the
	// runtime-disabled compatibility path.
	retrySubmit func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error)
	// poolDispatch is the provider-pool admission template this request was
	// dispatched with, or nil for a direct model_routes request. The stream
	// evidence gate runtime re-enters SubmitProviderPool with it for every
	// recovery attempt, so the actual provider/model/execution path is
	// re-selected per attempt instead of being pinned to the initial candidate.
	poolDispatch *edgeservice.ProviderPoolDispatchRequest
}

// 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
	poolDispatch *edgeservice.ProviderPoolDispatchRequest
}

func (c openAIRequestContext) recordUsageAttempt(dispatch edgeservice.RunDispatch, responseMode string, obs usageObservation) {
	c.usage.RecordAttempt(newUsageDispatchBinding(dispatch, responseMode), obs)
}

func (c openAIRequestContext) finishUsageRequest(status, responseMode string) {
	c.usage.FinishRequest(status, responseMode)
}

func (dc *responsesDispatchContext) withPoolDispatch(pool edgeservice.ProviderPoolDispatchRequest) *responsesDispatchContext {
	derived := *dc
	derived.poolDispatch = &pool
	return &derived
}

// 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
}

// withPoolDispatch derives a context bound to the provider-pool admission
// template the request was dispatched with. Like withRetrySubmit it derives
// rather than mutates: the template can only be built once the pool request
// exists, and the base context stays immutable for the other branches.
func (dc *chatDispatchContext) withPoolDispatch(poolReq edgeservice.ProviderPoolDispatchRequest) *chatDispatchContext {
	derived := *dc
	derived.poolDispatch = &poolReq
	return &derived
}