OmniRoute/open-sse/utils/cacheControlPolicy.ts

/**
 * Cache Control Policy
 *
 * Determines when to preserve client-side prompt caching headers (cache_control)
 * vs. applying OmniRoute's own caching strategy.
 *
 * Client-side caching (e.g., Claude Code) is preserved when:
 * 1. Client is Claude Code or similar caching-aware client
 * 2. Provider supports prompt caching (Anthropic, etc.)
 *
 * Combo membership / routing strategy no longer gates preservation: rewriting a
 * caching-aware client's markers produced per-request breakpoint positions that
 * thrashed the provider prompt cache, while preserving them is never worse
 * (see shouldPreserveCacheControl).
 */

import type { RoutingStrategyValue } from "../../src/shared/constants/routingStrategies";

/**
 * Cache control preservation modes
 */
export type CacheControlMode = "auto" | "always" | "never";

/**
 * Cache control settings from the database
 */
export interface CacheControlSettings {
  alwaysPreserveClientCache?: CacheControlMode;
}

/**
 * Cache metrics for tracking effectiveness
 */
export interface CacheControlMetrics {
  // Totals
  totalRequests: number;
  requestsWithCacheControl: number;

  // Token counts
  totalInputTokens: number;
  totalCachedTokens: number;
  totalCacheCreationTokens: number;

  // Savings
  tokensSaved: number;
  estimatedCostSaved: number;

  // Breakdowns
  byProvider: Record<
    string,
    {
      requests: number;
      inputTokens: number;
      cachedTokens: number;
      cacheCreationTokens: number;
    }
  >;
  byStrategy: Record<
    string,
    {
      requests: number;
      inputTokens: number;
      cachedTokens: number;
      cacheCreationTokens: number;
    }
  >;

  lastUpdated: string;
}

/**
 * Routing strategies that are deterministic (same request → same provider)
 */
const DETERMINISTIC_STRATEGIES: Set<RoutingStrategyValue> = new Set(["priority", "cost-optimized"]);

/**
 * Providers that support prompt caching
 */
const CACHING_PROVIDERS = new Set([
  "claude",
  "anthropic",
  "zai",
  "deepseek",
  // Kimi Code's OpenAI protocol requires prompt_cache_key for Coding Plan
  // cache affinity. The OAuth card and hidden API-key compatibility ID share
  // the same upstream API.
  "kimi-coding",
  "kimi-coding-apikey",
  // #3088 — Xiaomi MiMo honors OpenAI-format cache_control breakpoints. Without
  // this entry, shouldPreserveCacheControl() returns false for Claude Code
  // clients and filterToOpenAIFormat() strips cache_control, so Xiaomi never
  // sees the cache hints and every request is a cache miss.
  "xiaomi-mimo",
  // #3955 — OpenAI / Codex / Azure-OpenAI use AUTOMATIC prefix caching: the longest
  // matching prefix of a request is cached upstream WITHOUT any explicit cache_control
  // markers. They must count as caching providers so the cache-aware compression guard
  // preserves the cacheable prefix (system prompt / earliest messages) instead of
  // rewriting it and forcing a cache miss. This also activates the intended
  // `prompt_cache_key` cache-routing hint for OpenAI in chatCore.
  "openai",
  "codex",
  "azure",
  // #2069 — DashScope's OpenAI-compatible endpoints (Alibaba Model Studio and
  // Qwen Cloud pay-as-you-go, upstream "alicode"/"alicode-intl") natively honor
  // `cache_control: {type:"ephemeral"}` breakpoints. Without these entries
  // shouldPreserveCacheControl() returns false for Claude Code clients and the
  // OpenAI-format translator strips cache_control, so DashScope never sees the
  // hints and every request is a cache miss.
  "alibaba",
  "alibaba-cn",
  "qwen-cloud",
]);

/**
 * Providers that honor EXPLICIT `cache_control` breakpoints carried inside an
 * OpenAI-format request body (i.e. the markers must be passed THROUGH the
 * Claude→OpenAI translation instead of stripped).
 *
 * This is a strict subset of CACHING_PROVIDERS and deliberately excludes
 * `openai` / `codex` / `azure`: those use AUTOMATIC prefix caching (#3955) and
 * do NOT accept explicit `cache_control` fields in the request — forwarding the
 * markers there is meaningless at best and a 400 "unknown field" at worst, and
 * it broke the chatCore "strips cache markers for non-Claude providers" test.
 * Claude-format providers re-inject markers via prepareClaudeRequest, so they
 * are not listed here either.
 */
const OPENAI_FORMAT_CACHE_CONTROL_PROVIDERS = new Set([
  // #2069 — DashScope OpenAI-compatible endpoints accept ephemeral breakpoints.
  "alibaba",
  "alibaba-cn",
  "qwen-cloud",
  // #3088 — Xiaomi MiMo honors OpenAI-format cache_control breakpoints.
  "xiaomi-mimo",
]);

/**
 * Per-connection override for cache behavior, resolved from the connection's
 * `provider_specific_data.cache` JSON sub-object (see `resolveConnectionCacheOverride`).
 * Lets an operator opt a custom/openai-compatible connection into prompt-cache
 * behavior that the hardcoded provider-name sets above can never match (#6880).
 */
export interface ConnectionCacheOverride {
  supportsPromptCaching?: boolean;
  cacheControlPassthrough?: "strip" | "openai-format" | "claude-format";
}

/**
 * Extract and validate a `ConnectionCacheOverride` from a connection's
 * `providerSpecificData` bag. Returns `null` when absent/malformed so every
 * call site can safely pass the result straight through.
 */
export function resolveConnectionCacheOverride(
  providerSpecificData: unknown
): ConnectionCacheOverride | null {
  if (!providerSpecificData || typeof providerSpecificData !== "object") return null;
  const cache = (providerSpecificData as Record<string, unknown>).cache;
  if (!cache || typeof cache !== "object" || Array.isArray(cache)) return null;
  const record = cache as Record<string, unknown>;
  const result: ConnectionCacheOverride = {};
  if (typeof record.supportsPromptCaching === "boolean") {
    result.supportsPromptCaching = record.supportsPromptCaching;
  }
  if (
    record.cacheControlPassthrough === "strip" ||
    record.cacheControlPassthrough === "openai-format" ||
    record.cacheControlPassthrough === "claude-format"
  ) {
    result.cacheControlPassthrough = record.cacheControlPassthrough;
  }
  return Object.keys(result).length > 0 ? result : null;
}

/**
 * Whether `cache_control` markers should be PASSED THROUGH the OpenAI-format
 * translation for this provider (vs. stripped). Used to gate the request-side
 * passthrough so generic / implicit-cache OpenAI providers keep getting cleaned.
 */
export function providerHonorsOpenAIFormatCacheControl(
  provider: string | null | undefined,
  connectionCacheOverride?: ConnectionCacheOverride | null
): boolean {
  if (connectionCacheOverride?.cacheControlPassthrough === "openai-format") return true;
  if (connectionCacheOverride?.cacheControlPassthrough === "strip") return false;
  if (!provider) return false;
  return OPENAI_FORMAT_CACHE_CONTROL_PROVIDERS.has(provider.toLowerCase());
}

/**
 * Detect if the client is Claude Code or another caching-aware client
 */
export function isClaudeCodeClient(userAgent: string | null | undefined): boolean {
  if (!userAgent) return false;
  const ua = userAgent.toLowerCase();

  // Claude Code user agents
  if (ua.includes("claude-code") || ua.includes("claude_code")) return true;
  if (ua.includes("claude-cli/")) return true;
  if (ua.includes("sdk-cli")) return true;
  if (ua.includes("anthropic") && ua.includes("cli")) return true;

  return false;
}

/**
 * Check if a provider supports prompt caching
 * Supports caching if:
 * 1. Provider is in the known caching providers list, OR
 * 2. Provider uses Claude protocol (detected via targetFormat)
 */
export function providerSupportsCaching(
  provider: string | null | undefined,
  targetFormat?: string | null,
  connectionCacheOverride?: ConnectionCacheOverride | null
): boolean {
  if (typeof connectionCacheOverride?.supportsPromptCaching === "boolean") {
    return connectionCacheOverride.supportsPromptCaching;
  }
  if (!provider) return false;
  if (CACHING_PROVIDERS.has(provider.toLowerCase())) return true;
  // All Claude-protocol providers support prompt caching
  if (targetFormat === "claude") return true;
  return false;
}

/**
 * Check if a routing strategy is deterministic
 */
export function isDeterministicStrategy(
  strategy: RoutingStrategyValue | null | undefined
): boolean {
  if (!strategy) return false;
  return DETERMINISTIC_STRATEGIES.has(strategy);
}

/**
 * Determine if client-side cache_control headers should be preserved
 *
 * Auto mode preserves for every caching-aware client talking to a
 * caching-capable provider — regardless of combo membership or routing
 * strategy. The old gate (combos only preserved on "deterministic"
 * strategies) forced OmniRoute to strip the client's markers and re-derive
 * breakpoints per request; the re-derived positions are not stable
 * turn-over-turn, which thrashed the provider prompt cache (observed in
 * production as ~200k cache_write tokens per turn on quota-share combos).
 * Preserving the client's markers is never worse than rewriting them: on a
 * stable target the client's breakpoints advance deterministically, and on a
 * target switch both approaches miss equally.
 *
 * @param userAgent - User-Agent header from the request
 * @param isCombo - Whether this is a combo model (kept for callers/telemetry;
 *   no longer gates preservation)
 * @param comboStrategy - The combo's routing strategy (kept for
 *   callers/telemetry; no longer gates preservation)
 * @param targetProvider - The target provider for the request
 * @param settings - Cache control settings from database (optional)
 * @returns true if cache_control should be preserved, false if OmniRoute should manage it
 */
export function shouldPreserveCacheControl({
  userAgent,
  targetProvider,
  targetFormat,
  settings,
  connectionCacheOverride,
}: {
  userAgent: string | null | undefined;
  isCombo: boolean;
  comboStrategy?: RoutingStrategyValue | null;
  targetProvider: string | null | undefined;
  targetFormat?: string | null;
  settings?: CacheControlSettings;
  connectionCacheOverride?: ConnectionCacheOverride | null;
}): boolean {
  // User override takes precedence
  if (settings?.alwaysPreserveClientCache === "always") {
    return true;
  }
  if (settings?.alwaysPreserveClientCache === "never") {
    return false;
  }

  // Auto mode: must be a caching-aware client…
  if (!isClaudeCodeClient(userAgent)) {
    return false;
  }

  // …talking to a provider that supports prompt caching.
  return providerSupportsCaching(targetProvider, targetFormat, connectionCacheOverride);
}

/**
 * Track cache control metrics for a request
 */
export function trackCacheMetrics({
  preserved,
  provider,
  strategy,
  metrics,
  inputTokens,
  cachedTokens,
  cacheCreationTokens,
}: {
  preserved: boolean;
  provider: string;
  strategy: string | null | undefined;
  metrics: CacheControlMetrics;
  inputTokens?: number;
  cachedTokens?: number;
  cacheCreationTokens?: number;
}): CacheControlMetrics {
  const now = new Date().toISOString();

  // Initialize metrics if empty
  if (!metrics) {
    metrics = {
      totalRequests: 0,
      requestsWithCacheControl: 0,
      totalInputTokens: 0,
      totalCachedTokens: 0,
      totalCacheCreationTokens: 0,
      tokensSaved: 0,
      estimatedCostSaved: 0,
      byProvider: {},
      byStrategy: {},
      lastUpdated: now,
    };
  }

  // Increment total requests
  metrics.totalRequests++;

  // Track token counts
  const input = inputTokens || 0;
  const cached = cachedTokens || 0;
  const creation = cacheCreationTokens || 0;

  metrics.totalInputTokens += input;
  metrics.totalCachedTokens += cached;
  metrics.totalCacheCreationTokens += creation;

  // Calculate tokens saved (cached tokens are reused, not charged)
  if (cached > 0) {
    metrics.tokensSaved += cached;
  }

  // Only track requests where cache_control was preserved
  if (preserved) {
    metrics.requestsWithCacheControl++;

    // Initialize provider tracking
    if (!metrics.byProvider[provider]) {
      metrics.byProvider[provider] = {
        requests: 0,
        inputTokens: 0,
        cachedTokens: 0,
        cacheCreationTokens: 0,
      };
    }
    metrics.byProvider[provider].requests++;
    metrics.byProvider[provider].inputTokens += input;
    metrics.byProvider[provider].cachedTokens += cached;
    metrics.byProvider[provider].cacheCreationTokens += creation;

    // Initialize strategy tracking
    if (strategy && !metrics.byStrategy[strategy]) {
      metrics.byStrategy[strategy] = {
        requests: 0,
        inputTokens: 0,
        cachedTokens: 0,
        cacheCreationTokens: 0,
      };
    }
    if (strategy) {
      metrics.byStrategy[strategy].requests++;
      metrics.byStrategy[strategy].inputTokens += input;
      metrics.byStrategy[strategy].cachedTokens += cached;
      metrics.byStrategy[strategy].cacheCreationTokens += creation;
    }
  }

  metrics.lastUpdated = now;
  return metrics;
}

/**
 * Record cache token usage and update metrics
 */
export function updateCacheTokenMetrics({
  metrics,
  provider,
  strategy,
  inputTokens,
  cachedTokens,
  cacheCreationTokens,
  costSaved,
}: {
  metrics: CacheControlMetrics;
  provider: string;
  strategy: string | null | undefined;
  inputTokens: number;
  cachedTokens: number;
  cacheCreationTokens: number;
  costSaved?: number;
}): CacheControlMetrics {
  metrics.totalCachedTokens += cachedTokens;
  metrics.totalCacheCreationTokens += cacheCreationTokens;
  metrics.totalInputTokens += inputTokens;

  // Cached tokens are reused (saved), creation tokens are new cache writes
  metrics.tokensSaved += cachedTokens;
  if (costSaved !== undefined) {
    metrics.estimatedCostSaved += costSaved;
  }

  // Update provider tracking
  if (metrics.byProvider[provider]) {
    metrics.byProvider[provider].cachedTokens += cachedTokens;
    metrics.byProvider[provider].cacheCreationTokens += cacheCreationTokens;
    metrics.byProvider[provider].inputTokens += inputTokens;
  }

  // Update strategy tracking
  if (strategy && metrics.byStrategy[strategy]) {
    metrics.byStrategy[strategy].cachedTokens += cachedTokens;
    metrics.byStrategy[strategy].cacheCreationTokens += cacheCreationTokens;
    metrics.byStrategy[strategy].inputTokens += inputTokens;
  }

  metrics.lastUpdated = new Date().toISOString();
  return metrics;
}