OmniRoute/src/shared/validation/schemas/combo.ts

import { z } from "zod";
import {
  ACCOUNT_FALLBACK_STRATEGY_VALUES,
  ROUTING_STRATEGY_VALUES,
} from "@/shared/constants/routingStrategies";
import { SUPPORTED_BATCH_ENDPOINTS } from "@/shared/constants/batchEndpoints";
import { MAX_REQUEST_BODY_LIMIT_MB, MIN_REQUEST_BODY_LIMIT_MB } from "@/shared/constants/bodySize";
import { COMBO_CONFIG_MODES } from "@/shared/constants/comboConfigMode";
import { providerAllowsOptionalApiKey } from "@/shared/constants/providers";
import { HIDEABLE_SIDEBAR_ITEM_IDS } from "@/shared/constants/sidebarVisibility";
import {
  isForbiddenUpstreamHeaderName,
  isForbiddenCustomHeaderName,
} from "@/shared/constants/upstreamHeaders";
import { MAX_TIMER_TIMEOUT_MS } from "@/shared/utils/runtimeTimeouts";

// ──── Combo Schemas ────

export const comboStepMetaSchema = {
  id: z.string().trim().min(1).max(200).optional(),
  weight: z.number().min(0).max(100).optional().default(0),
  label: z.string().trim().min(1).max(200).optional(),
  fallbackOnlyOnQuotaExhaustion: z.boolean().optional(),
};

export const comboModelStepInputSchema = z.object({
  kind: z.literal("model").optional(),
  provider: z.string().trim().min(1).max(120).optional(),
  providerId: z.string().trim().min(1).max(120).optional(),
  model: z.string().trim().min(1).max(300),
  connectionId: z.string().trim().min(1).max(200).nullable().optional(),
  allowedConnectionIds: z.array(z.string().trim().min(1).max(200)).max(50).optional(),
  tags: z.array(z.string().trim().min(1).max(100)).max(20).optional(),
  // Pipeline strategy (open-sse/services/pipeline.ts): an optional per-step
  // instruction. Steps run in `models` order — each step's output feeds the next
  // step's input, and this `prompt` is injected as that step's system instruction.
  // Ignored by every other strategy, so it is fully backward-compatible.
  prompt: z.string().trim().min(1).max(20000).optional(),
  ...comboStepMetaSchema,
});

export const comboRefStepInputSchema = z.object({
  kind: z.literal("combo-ref"),
  comboName: z.string().trim().min(1).max(100),
  ...comboStepMetaSchema,
});

// A combo entry can be a plain string (legacy), a legacy object, or a structured ComboStep.
export const comboModelEntry = z.union([
  z.string().trim().min(1).max(300),
  comboModelStepInputSchema,
  comboRefStepInputSchema,
]);

export const shadowRoutingSchema = z
  .object({
    enabled: z.boolean().optional(),
    targets: z.array(comboModelEntry).max(20).optional(),
    sampleRate: z.coerce.number().min(0).max(1).optional(),
    maxTargets: z.coerce.number().int().min(1).max(10).optional(),
    timeoutMs: z.coerce.number().int().min(1000).max(120000).optional(),
  })
  .strict();

export const evalRoutingSchema = z
  .object({
    enabled: z.boolean().optional(),
    suiteIds: z.array(z.string().trim().min(1).max(200)).max(50).optional(),
    maxAgeHours: z.coerce.number().min(1).max(8760).optional(),
    minCases: z.coerce.number().int().min(1).max(100000).optional(),
    qualityWeight: z.coerce.number().min(0).max(1).optional(),
    latencyWeight: z.coerce.number().min(0).max(1).optional(),
    cacheTtlMs: z.coerce.number().int().min(1000).max(300000).optional(),
  })
  .strict();

export const comboStrategySchema = z.enum(ROUTING_STRATEGY_VALUES);

export const scoringWeightsSchema = z
  .object({
    quota: z.number().min(0).max(1),
    health: z.number().min(0).max(1),
    costInv: z.number().min(0).max(1),
    latencyInv: z.number().min(0).max(1),
    taskFit: z.number().min(0).max(1),
    stability: z.number().min(0).max(1),
    tierPriority: z.number().min(0).max(1).optional().default(0.05),
    tierAffinity: z.number().min(0).max(1).optional().default(0.05),
    specificityMatch: z.number().min(0).max(1).optional().default(0.05),
    contextAffinity: z.number().min(0).max(1).optional().default(0.08),
    cacheAffinity: z.number().min(0).max(1).optional().default(0),
    sessionAvailability: z.number().min(0).max(1).optional().default(0.05),
    resetWindowAffinity: z.number().min(0).max(1).optional().default(0),
    // The scorer weighs these two as well (`DEFAULT_WEIGHTS`); leaving them out
    // meant zod dropped them from a saved config, and `normalizeScoringWeights`
    // then read the gap as a deliberate zero — silently disabling
    // anti-concentration and the quality signal, and inflating every other
    // weight to make the distribution sum to 1 again.
    //
    // The defaults below therefore DO change the effective weights of a stored
    // config that omitted these keys: the other thirteen stop being renormalized
    // upward (quota 0.1549 -> 0.1429, health 0.1740 -> 0.1605, and so on). That is
    // the correction, not a side effect — but it is a behaviour change, and the
    // PR says so rather than claiming the routing is untouched.
    connectionDensity: z.number().min(0).max(1).optional().default(0.0476),
    quality: z.number().min(0).max(1).optional().default(0.03),
    reliability: z.number().min(0).max(1).optional().default(0),
  })
  .optional();

export const compositeTierEntrySchema = z
  .object({
    stepId: z.string().trim().min(1).max(200),
    fallbackTier: z.string().trim().min(1).max(100).optional(),
    label: z.string().trim().min(1).max(200).optional(),
    description: z.string().trim().min(1).max(500).optional(),
  })
  .strict();

export const compositeTiersSchema = z
  .object({
    defaultTier: z.string().trim().min(1).max(100),
    tiers: z.record(z.string().trim().min(1).max(100), compositeTierEntrySchema),
  })
  .strict();

export const compressionModeSchema = z.enum([
  "off",
  "lite",
  "standard",
  "aggressive",
  "ultra",
  "rtk",
  "stacked",
  "codex-responses",
  "omniglyph",
]);

export const comboCompressionOverrideSchema = z.union([z.literal(""), compressionModeSchema]);

export const slaRoutingPolicySchema = z
  .object({
    targetP95Ms: z.coerce.number().int().positive().max(300000).optional(),
    maxErrorRate: z.coerce.number().min(0).max(1).optional(),
    maxCostPer1MTokens: z.coerce.number().positive().max(1000000).optional(),
    hardConstraints: z.boolean().optional(),
  })
  .strict();

// Feature 4985 — configurable response-body validation for combo routing. A 200 OK whose
// body fails this predicate fails over to the next target (same path as an HTTP error).
export const responseValidationSchema = z
  .object({
    forbiddenSubstrings: z.array(z.string().min(1).max(500)).max(50).optional(),
    requiredSubstrings: z.array(z.string().min(1).max(500)).max(50).optional(),
    minContentLength: z.coerce.number().int().min(0).max(1_000_000).optional(),
    jsonPathPredicates: z
      .array(
        z.object({
          path: z.string().trim().min(1).max(300),
          condition: z.enum(["exists", "nonEmpty", "equals", "notEquals"]),
          value: z.union([z.string().max(1000), z.number(), z.boolean()]).optional(),
        })
      )
      .max(20)
      .optional(),
  })
  .strict();

export const comboRuntimeConfigSchema = z
  .object({
    responseValidation: responseValidationSchema.optional(),
    strategy: comboStrategySchema.optional(),
    maxRetries: z.coerce.number().int().min(0).max(10).optional(),
    retryDelayMs: z.coerce.number().int().min(0).max(60000).optional(),
    fallbackDelayMs: z.coerce.number().int().min(0).max(60000).optional(),
    timeoutMs: z.coerce.number().int().min(1000).optional(),
    targetTimeoutMs: z.coerce.number().int().min(0).max(MAX_TIMER_TIMEOUT_MS).optional(),
    // Whole-combo wall-clock budget. 0 (default) means unlimited iteration;
    // the 10-minute COMBO_LOOP_SAFETY_TIMEOUT_MS hang-stop still applies.
    // A positive value replaces that safety net for this combo.
    comboTimeoutMs: z.coerce.number().int().min(0).max(MAX_TIMER_TIMEOUT_MS).optional(),
    concurrencyPerModel: z.coerce.number().int().min(1).max(20).optional(),
    queueTimeoutMs: z.coerce.number().int().min(1000).max(120000).optional(),
    // #3872: pre-cascade semaphore queue depth (round-robin). 0 = fail over immediately.
    queueDepth: z.coerce.number().int().min(0).max(100).optional(),
    // Per-combo sticky round-robin batch size. When unset, handleRoundRobinCombo
    // falls back to the global `settings.stickyRoundRobinLimit` so the existing
    // knob still controls the default. 0 clamps to 1 (no batching) upstream.
    stickyRoundRobinLimit: z.coerce.number().int().min(0).max(1000).optional(),
    // #6168: opt-out for per-conversation session stickiness. When true, round-robin
    // and random/weighted/priority combos rotate freely instead of pinning a whole
    // conversation to one connection by the first-message hash. Per-combo `config`
    // wins over the global `settings.disableSessionStickiness` fallback. Default false
    // preserves the #3825 prompt-cache/504 fix.
    disableSessionStickiness: z.boolean().optional(),
    stickyWeightedLimit: z.coerce.number().int().min(0).max(1000).optional(),
    healthCheckEnabled: z.boolean().optional(),
    healthCheckTimeoutMs: z.coerce.number().int().min(100).max(30000).optional(),
    handoffThreshold: z.coerce.number().min(0.5).max(0.94).optional(),
    handoffModel: z.string().trim().max(200).optional(),
    handoffProviders: z.array(z.string().trim().min(1).max(100)).max(10).optional(),
    maxMessagesForSummary: z.coerce.number().int().min(5).max(100).optional(),
    maxComboDepth: z.coerce.number().int().min(1).max(10).optional(),
    // #11134: shared per-request attempt budget. Bounds mirror
    // MAX_GLOBAL_ATTEMPTS_HARD_CAP (200) in comboPredicates.ts.
    maxGlobalAttempts: z.coerce.number().int().min(1).max(200).optional(),
    nestedComboMode: z.enum(["flatten", "execute"]).optional(),
    trackMetrics: z.boolean().optional(),
    reasoningTokenBufferEnabled: z.boolean().optional(),
    reasoningTransportFallback: z.enum(["skip", "drop"]).optional(),
    compressionMode: compressionModeSchema.optional(),
    failoverBeforeRetry: z.boolean().optional(),
    maxSetRetries: z.coerce.number().int().min(0).max(10).optional(),
    setRetryDelayMs: z.coerce.number().int().min(0).max(60000).optional(),
    zeroLatencyOptimizationsEnabled: z.boolean().optional(),
    hedging: z.boolean().optional(),
    hedgeDelayMs: z.coerce.number().int().min(0).max(60000).optional(),
    fallbackCompressionMode: compressionModeSchema.optional(),
    fallbackCompressionThreshold: z.coerce.number().int().min(0).max(2_000_000).optional(),
    predictiveTtftMs: z.coerce.number().int().min(0).max(300000).optional(),
    relayMode: z.enum(["schema-locked", "standard"]).optional(),
    // Auto-Combo / LKGP Extensions
    candidatePool: z.array(z.string().min(1)).optional(),
    weights: scoringWeightsSchema.optional(),
    modePack: z.string().max(100).optional(),
    budgetCap: z.number().positive().optional(),
    explorationRate: z.number().min(0).max(1).optional(),
    routerStrategy: z.string().optional(),
    slaTargetP95Ms: z.coerce.number().int().positive().max(300000).optional(),
    slaMaxErrorRate: z.coerce.number().min(0).max(1).optional(),
    slaMaxCostPer1MTokens: z.coerce.number().positive().max(1000000).optional(),
    slaHardConstraints: z.boolean().optional(),
    sla: slaRoutingPolicySchema.optional(),
    compositeTiers: compositeTiersSchema.optional(),
    resetAwareSessionWeight: z.coerce.number().min(0).max(100).optional(),
    resetAwareWeeklyWeight: z.coerce.number().min(0).max(100).optional(),
    resetAwareTieBandPercent: z.coerce.number().min(0).max(100).optional(),
    resetAwareExhaustionGuardPercent: z.coerce.number().min(0).max(100).optional(),
    quotaWeightedFloorPercent: z.coerce.number().min(0).max(100).optional(),
    resetAwareQuotaCacheTtlMs: z.coerce.number().int().min(0).max(300_000).optional(),
    resetAwareQuotaCacheMaxStaleMs: z.coerce.number().int().min(0).max(3_600_000).optional(),
    resetWindowWindows: z.array(z.enum(["weekly", "session", "monthly"])).optional(),
    resetWindowIncludeSession: z.boolean().optional(),
    resetWindowTieBandMs: z.coerce.number().int().min(0).max(86_400_000).optional(),
    resetWindowQuotaCacheTtlMs: z.coerce.number().int().min(0).max(300_000).optional(),
    resetWindowQuotaCacheMaxStaleMs: z.coerce.number().int().min(0).max(3_600_000).optional(),
    // Connection-aware expansion for group-B combo strategies is opt-in.
    connectionAwareExpansion: z.boolean().optional(),
    connectionAwareExpansionMaxPerTarget: z.coerce.number().int().min(1).max(64).optional(),
    shadowRouting: shadowRoutingSchema.optional(),
    evalRouting: evalRoutingSchema.optional(),
    // Fusion strategy (open-sse/services/fusion.ts): the panel is the combo's
    // targets; `judgeModel` synthesizes the final answer (defaults to the first
    // panel model when unset); `fusionTuning` controls quorum-grace collection.
    judgeModel: z.string().trim().max(200).optional(),
    fusionTuning: z
      .object({
        minPanel: z.coerce.number().int().min(1).max(50).optional(),
        stragglerGraceMs: z.coerce.number().int().min(0).max(120_000).optional(),
        panelHardTimeoutMs: z.coerce.number().int().min(1000).max(600_000).optional(),
        // Hard cap on panel size (issue #1905) — see FUSION_DEFAULTS.maxPanel in
        // open-sse/services/fusion.ts. Bounds how many models can be fanned out
        // and buffered in memory concurrently before the container's heap ceiling
        // is at risk.
        maxPanel: z.coerce.number().int().min(1).max(200).optional(),
      })
      .strict()
      .optional(),
    // Context window requirements for combo target filtering and sorting.
    // minContextWindow: filters out models with context windows below this threshold.
    // maxContextWindow: filters out models with context windows above this threshold.
    // preferLargeContext: sorts remaining targets by context size (descending).
    // contextFilterMode: "strict" excludes unknown-context models, "lenient" includes them.
    contextRequirements: z
      .object({
        minContextWindow: z.coerce.number().int().min(0).max(10_000_000).optional(),
        maxContextWindow: z.coerce.number().int().min(0).max(10_000_000).optional(),
        preferLargeContext: z.boolean().optional(),
        contextFilterMode: z.enum(["strict", "lenient"]).optional(),
      })
      .strict()
      .optional(),
    // Optional client-side sort hint for combo models.
    // Honored in the dashboard builder; reserved for future server-side use. Inert on execution.
    modelSort: z
      .object({ method: z.enum(["manual", "provider", "score", "name"]) })
      .passthrough()
      .optional(),
  })
  .passthrough()
  .transform((config) => {
    // Backward-compat shim: combos stored prior to v3.8.33 may carry zero-latency
    // feature flags (fallbackCompressionMode !== "off", hedging === true, or
    // predictiveTtftMs > 0) without the accompanying zeroLatencyOptimizationsEnabled
    // gate that the new schema requires. Auto-promote the flag when any such feature
    // is enabled but the gate is unset/false, so stored combos continue to round-trip
    // through PUT /api/combos/{id} without returning 400. This replaces the prior
    // superRefine that hard-rejected these payloads (see issue #4382).
    if (config.zeroLatencyOptimizationsEnabled === true) return config;

    const hasZeroLatencyFeature =
      config.hedging === true ||
      (typeof config.predictiveTtftMs === "number" && config.predictiveTtftMs > 0) ||
      (!!config.fallbackCompressionMode && config.fallbackCompressionMode !== "off");

    if (hasZeroLatencyFeature) {
      return { ...config, zeroLatencyOptimizationsEnabled: true };
    }
    return config;
  });

export const comboNameSchema = z
  .string()
  .trim()
  .min(1, "Name is required")
  .max(100)
  .regex(
    /^[a-zA-Z0-9_/.\-\[\] ]+$/,
    "Name can only contain letters, numbers, spaces, -, _, /, ., [ and ]."
  );

type QuotaOnlyComboRefState = {
  models?: Array<z.infer<typeof comboModelEntry>>;
  strategy?: string;
  config?: z.infer<typeof comboRuntimeConfigSchema>;
};

export function requiresQuotaOnlyComboRefExecute(value: QuotaOnlyComboRefState): boolean {
  const hasProtectedComboRef = value.models?.some(
    (step) =>
      typeof step === "object" &&
      step.kind === "combo-ref" &&
      step.fallbackOnlyOnQuotaExhaustion === true
  );
  return (
    (value.strategy === undefined || value.strategy === "priority") &&
    hasProtectedComboRef === true &&
    value.config?.nestedComboMode !== "execute"
  );
}

function validateQuotaOnlyComboRefs(value: QuotaOnlyComboRefState, ctx: z.RefinementCtx): void {
  if (requiresQuotaOnlyComboRefExecute(value)) {
    ctx.addIssue({
      code: z.ZodIssueCode.custom,
      message: "Quota-only combo references require nestedComboMode execute",
      path: ["config", "nestedComboMode"],
    });
  }
}

export const createComboSchema = z
  .object({
    name: comboNameSchema,
    description: z.string().max(2000).optional(),
    // Optional label advertised as `display_name` in /v1/models. Lets a combo
    // carry a machine-oriented name while clients show something readable.
    displayName: z.string().trim().max(200).optional(),
    models: z.array(comboModelEntry).min(1, "a combo requires at least one model"),
    strategy: comboStrategySchema.optional().default("priority"),
    config: comboRuntimeConfigSchema.optional(),
    allowedProviders: z.array(z.string().trim().min(1).max(200)).max(100).optional(),
    allowedModelFamilies: z.array(z.string().trim().min(1).max(100)).max(100).optional(),
    system_message: z.string().max(50000).optional(),
    tool_filter_regex: z.string().max(1000).optional(),
    context_cache_protection: z.boolean().optional(),
    context_length: z.number().int().min(1000).max(2000000).optional(),
    // Optional embedding dimensions override for embedding combos.
    // When set, the value is injected into every upstream embedding request as
    // the `dimensions` field (and translated to `outputDimensionality` for Gemini).
    // Stored as a string to match the OpenAI API convention; coerced to number
    // by the embedding handler. Leave unset to use each model's default.
    dimensions: z
      .string()
      .regex(/^\d+$/, "dimensions must be a positive integer string")
      .optional()
      .nullable(),
  })
  .superRefine(validateQuotaOnlyComboRefs);

export const updateComboDefaultsSchema = z
  .object({
    comboDefaults: comboRuntimeConfigSchema.optional(),
    providerOverrides: z.record(z.string().trim().min(1), comboRuntimeConfigSchema).optional(),
  })
  .superRefine((value, ctx) => {
    if (!value.comboDefaults && !value.providerOverrides) {
      ctx.addIssue({
        code: z.ZodIssueCode.custom,
        message: "Nothing to update",
        path: [],
      });
    }

    if (value.comboDefaults?.compositeTiers) {
      ctx.addIssue({
        code: z.ZodIssueCode.custom,
        message: "compositeTiers is only supported on concrete combos",
        path: ["comboDefaults", "compositeTiers"],
      });
    }

    for (const [providerId, config] of Object.entries(value.providerOverrides || {})) {
      if (config?.compositeTiers) {
        ctx.addIssue({
          code: z.ZodIssueCode.custom,
          message: "compositeTiers is only supported on concrete combos",
          path: ["providerOverrides", providerId, "compositeTiers"],
        });
      }
    }
  });

export const updateComboSchema = z
  .object({
    name: comboNameSchema.optional(),
    description: z.string().max(2000).optional().nullable(),
    displayName: z.string().trim().max(200).optional().nullable(),
    // An update may not remove every model from a combo, or a working combo
    // loses every target. Creation refuses an empty list too: since the CLI
    // gained --models (#10954), an empty draft has no remaining legitimate path.
    models: z
      .array(comboModelEntry)
      .min(1, "an update cannot remove every model from a combo")
      .optional(),
    strategy: comboStrategySchema.optional(),
    config: comboRuntimeConfigSchema.optional(),
    isActive: z.boolean().optional(),
    // Stored on the combo record and honoured by the readers — the builder's
    // option list and the dashboard grid both filter on it — but omitted here,
    // so the one endpoint a client can flip it through stripped the field and
    // a visibility-only update was rejected as empty. #12836
    isHidden: z.boolean().optional(),
    allowedProviders: z.array(z.string().trim().min(1).max(200)).max(100).optional().nullable(),
    allowedModelFamilies: z.array(z.string().trim().min(1).max(100)).max(100).optional().nullable(),
    overrideAllowedProviders: z.boolean().optional(),
    // Nullable like `description` and `context_length` above: an absent field means
    // "leave unchanged" because updateCombo merges over the stored record, so clearing
    // one needs an explicit null for updateCombo's null-means-delete pass (#12158).
    system_message: z.string().max(50000).optional().nullable(),
    tool_filter_regex: z.string().max(1000).optional().nullable(),
    context_cache_protection: z.boolean().optional().nullable(),
    context_length: z.number().int().min(1000).max(2000000).optional().nullable(),
    compressionOverride: comboCompressionOverrideSchema.optional(),
    dimensions: z
      .string()
      .regex(/^\d+$/, "dimensions must be a positive integer string")
      .optional()
      .nullable(),
  })
  .superRefine((value, ctx) => {
    if (
      value.name === undefined &&
      value.description === undefined &&
      value.displayName === undefined &&
      value.models === undefined &&
      value.strategy === undefined &&
      value.config === undefined &&
      value.isActive === undefined &&
      value.isHidden === undefined &&
      value.allowedProviders === undefined &&
      value.allowedModelFamilies === undefined &&
      value.system_message === undefined &&
      value.tool_filter_regex === undefined &&
      value.context_cache_protection === undefined &&
      value.context_length === undefined &&
      value.compressionOverride === undefined &&
      value.dimensions === undefined
    ) {
      ctx.addIssue({
        code: z.ZodIssueCode.custom,
        message: "No valid fields to update",
        path: [],
      });
    }
  });

export const reorderCombosSchema = z
  .object({
    comboIds: z.array(z.string().trim().min(1).max(200)).min(1).max(1000),
  })
  .superRefine((value, ctx) => {
    if (new Set(value.comboIds).size !== value.comboIds.length) {
      ctx.addIssue({
        code: z.ZodIssueCode.custom,
        message: "comboIds must be unique",
        path: ["comboIds"],
      });
    }
  });

export const testComboSchema = z.object({
  comboName: z.string().trim().min(1, "comboName is required"),
});

// POST /api/combos/duplicate - Resolve an auto-combo template (e.g. "auto/best-coding")
// into a static, editable combo snapshot.
export const duplicateAutoComboSchema = z.object({
  name: z.string().trim().min(1, 'Missing required field: "name" (e.g. auto/best-coding)'),
  strategy: comboStrategySchema.optional(),
});