OmniRoute/open-sse/services/comboAgentMiddleware.ts

/**
 * comboAgentMiddleware.ts — Combo Agent Features
 *
 * Implements the "combo as agent" features from issues #399 and #401:
 *
 * 1. **System Message Override** (#399): If the combo defines a `system_message`,
 *    it is injected as the first system message, replacing any existing system message.
 *
 * 2. **Tool Filter Regex** (#399): If the combo defines a `tool_filter_regex`,
 *    only tools whose name matches the pattern are forwarded to the provider.
 *
 * 3. **Context Caching Protection** (#401): If the combo enables
 *    `context_cache_protection`, the proxy:
 *    a. On response: injects `<omniModel>provider/model</omniModel>` tag into
 *       the first assistant message content string.
 *    b. On request: scans the message history for the tag, and if found,
 *       overrides the requested model with the pinned one.
 *
 * All features are opt-in per combo and backward compatible with existing setups.
 */

import { isFingerprintProvider } from "./combo/fingerprintExpansion.ts";

interface ComboConfig {
  system_message?: string | null;
  tool_filter_regex?: string | null;
  context_cache_protection?: number | boolean;
  [key: string]: unknown;
}

interface Message {
  role?: string;
  content?: unknown;
  [key: string]: unknown;
}

// ── Context Caching Tag ─────────────────────────────────────────────────────

// Detection / extraction pattern. The newline runs combo.ts streaming wraps the
// tag with (#531) are irrelevant to *finding* the tag or capturing the model id,
// so they are intentionally NOT matched here: an unbounded `(?:\\n|\n|\r)*` prefix
// on this unanchored regex made `.test()` / `.exec()` run in O(n²) on inputs with
// many newlines (polynomial ReDoS — CodeQL js/polynomial-redos, #3870). Non-global
// so `.exec()` / `.test()` stay stateless (a global regex carries lastIndex between
// calls and would skip matches).
const CACHE_TAG_PATTERN = /<omniModel>([^<]+)<\/omniModel>/;

// Global variant for `.replace()` callers that must strip EVERY tag (a non-global
// regex only removes the first match, so a message carrying more than one
// <omniModel> tag — e.g. an Open WebUI follow-up/title request that inlines the
// whole chat history — leaked the remaining tags to the provider, defeating the
// cache-session protection stripModelTags enforces, #454). This variant still
// consumes the newline run wrapping the tag (combo.ts streaming, #531) so removal
// leaves no blank line, but the runs are BOUNDED ({0,16}) to keep the regex linear
// (no polynomial backtracking, #3870); 16 is far beyond any real streaming wrap.
const CACHE_TAG_PATTERN_GLOBAL =
  /(?:\\n|\n|\r){0,16}<omniModel>([^<]+)<\/omniModel>(?:\\n|\n|\r){0,16}/g;

/**
 * Inject the model tag into the last assistant message (or append a new one).
 * Only modifies string content — does not touch array content to avoid breaking
 * Claude/Gemini multi-part message formats.
 */
export function injectModelTag(messages: Message[], providerModel: string): Message[] {
  // Remove any existing tags first to avoid duplication on context compaction
  const cleaned = messages.map((msg) => {
    if (msg.role === "assistant" && typeof msg.content === "string") {
      return { ...msg, content: msg.content.replace(CACHE_TAG_PATTERN_GLOBAL, "").trimEnd() };
    }
    return msg;
  });

  // Find last assistant message with string content
  const lastAssistantIdx = cleaned.map((m) => m.role).lastIndexOf("assistant");

  // #474: If no assistant message exists yet (first turn), append a synthetic one
  // so the tag is present when the client sends the next request with the response.
  if (lastAssistantIdx === -1) {
    return [...cleaned, { role: "assistant", content: `<omniModel>${providerModel}</omniModel>` }];
  }

  const msg = cleaned[lastAssistantIdx];
  // Fix #721: Handle messages where content is not a string (tool_calls responses).
  // In this case, append a synthetic assistant message with the tag so the pin
  // roundtrips through the conversation history.
  if (typeof msg.content !== "string") {
    // If the message has tool_calls but no string content, append a new assistant
    // message with the tag rather than silently failing.
    return [...cleaned, { role: "assistant", content: `<omniModel>${providerModel}</omniModel>` }];
  }

  const tagged = [...cleaned];
  tagged[lastAssistantIdx] = {
    ...msg,
    content: `${msg.content}<omniModel>${providerModel}</omniModel>`,
  };
  return tagged;
}

/**
 * Scan message history for the model tag injected by a previous response.
 * Returns the pinned "provider/model" string, or null if not found.
 */
export function extractPinnedModel(messages: Message[]): string | null {
  // Scan from newest to oldest for efficiency
  for (let i = messages.length - 1; i >= 0; i--) {
    const msg = messages[i];
    if (msg.role === "assistant" && typeof msg.content === "string") {
      const match = CACHE_TAG_PATTERN.exec(msg.content);
      if (match) return match[1];
    }
  }
  return null;
}

// ── System Message Override ──────────────────────────────────────────────────

/**
 * Replace or inject a system message at the beginning of the messages array.
 * Existing system messages are removed if a combo override is set.
 */
export function applySystemMessageOverride(messages: Message[], systemMessage: string): Message[] {
  // Remove all existing system messages
  const filtered = messages.filter((m) => m.role !== "system");
  // Inject combo system message at start
  return [{ role: "system", content: systemMessage }, ...filtered];
}

// ── Tool Filter Regex ────────────────────────────────────────────────────────

/**
 * Filter the tools array, keeping only tools whose name matches the regex.
 * Returns the original array unchanged if pattern is null/empty.
 */
export function applyToolFilter(
  tools: unknown[] | undefined,
  pattern: string | null | undefined
): unknown[] | undefined {
  if (!tools || !pattern) return tools;

  let regex: RegExp;
  try {
    regex = new RegExp(pattern);
  } catch {
    // Invalid regex — return tools unchanged rather than crashing
    console.warn(`[ComboAgent] Invalid tool_filter_regex: "${pattern}"`);
    return tools;
  }

  return tools.filter((tool) => {
    const t = tool as Record<string, unknown>;
    // Support both OpenAI format ({ function: { name } }) and Anthropic ({ name })
    const name = (t.function as Record<string, unknown> | undefined)?.name ?? t.name ?? "";
    return regex.test(String(name));
  });
}

/**
 * Strip all <omniModel> tags from message content before forwarding to the provider.
 * The tag is an internal OmniRoute marker; providers must never see it or their
 * cache will treat every tagged request as a new session (#454).
 */
export function stripModelTags(messages: Message[]): Message[] {
  return messages.map((msg) => {
    if (typeof msg.content === "string" && CACHE_TAG_PATTERN.test(msg.content)) {
      return { ...msg, content: msg.content.replace(CACHE_TAG_PATTERN_GLOBAL, "").trimEnd() };
    }
    return msg;
  });
}

// ── Main Middleware ──────────────────────────────────────────────────────────

/**
 * Apply all combo agent features to the request body.
 * Safe to call with null/undefined comboConfig — returns body unchanged.
 */
export function applyComboAgentMiddleware(
  body: Record<string, unknown>,
  comboConfig: ComboConfig | null | undefined,
  providerModel: string // "provider/model" string for context caching
): { body: Record<string, unknown>; pinnedModel: string | null } {
  if (!comboConfig) return { body, pinnedModel: null };

  const hasMessages = Array.isArray(body.messages);
  const isResponsesRequest =
    Object.prototype.hasOwnProperty.call(body, "input") ||
    Object.prototype.hasOwnProperty.call(body, "instructions");
  const systemMessage =
    typeof comboConfig.system_message === "string" && comboConfig.system_message.trim()
      ? comboConfig.system_message
      : null;
  let messages: Message[] = hasMessages ? [...(body.messages as Message[])] : [];
  let pinnedModel: string | null = null;

  // Context cache pinning is handled server-side in combo.ts via
  // session_model_history. No client-side <omniModel> tag extraction needed.
  pinnedModel = null;

  // 2. System message override. Responses API uses top-level instructions instead of messages.
  if (systemMessage && !isResponsesRequest) {
    messages = applySystemMessageOverride(messages, systemMessage);
  }

  // 3. Tool filter
  const filteredTools = applyToolFilter(
    body.tools as unknown[] | undefined,
    comboConfig.tool_filter_regex
  );

  // 4. Strip internal <omniModel> tags before forwarding to provider (#454)
  //    These tags are OmniRoute-internal markers and must never reach the provider
  //    since providers would treat each tagged request as a new cache session.
  messages = stripModelTags(messages);

  return {
    body: {
      ...body,
      ...(isResponsesRequest && systemMessage ? { instructions: systemMessage } : {}),
      ...(hasMessages ? { messages } : {}),
      ...(filteredTools !== body.tools && { tools: filteredTools }),
    },
    pinnedModel,
  };
}

// ── System Prompt Template Expansion (#5501) ─────────────────────────────────

export interface ComboSystemPromptTemplateContext {
  modelId: string;
  providerId: string;
  account: string;
  fingerprint: string;
}

/**
 * Replace allowlisted `{{TOKEN}}` placeholders in a single left-to-right scan.
 * No regex (ReDoS-averse, cf. #3870) and no recursion: an expanded value is
 * appended to the output and never re-scanned. Unknown tokens ({{FOO}}) and
 * dangling "{{" stay literal.
 */
function expandStringTemplates(value: string, values: Record<string, string>): string {
  let out = "";
  let rest = value;
  while (rest.length > 0) {
    const start = rest.indexOf("{{");
    if (start === -1) {
      out += rest;
      break;
    }
    const end = rest.indexOf("}}", start + 2);
    if (end === -1) {
      out += rest;
      break;
    }
    const token = rest.slice(start, end + 2);
    out += rest.slice(0, start);
    out += token in values ? values[token] : token;
    rest = rest.slice(end + 2);
  }
  return out;
}

/**
 * Expand allowlisted placeholders in the combo-injected system prompt (#5501).
 *
 * Strictly scoped to the content the combo override produced — never
 * client-owned system content:
 *   - Responses API body (has `instructions`) → expand `body.instructions`.
 *   - messages body → expand `body.messages[0]` when it is the injected combo
 *     system message (the override filters all system messages and injects its
 *     own at index 0 with string content).
 *   - otherwise → body unchanged.
 */
export function expandComboSystemPromptTemplates(
  body: Record<string, unknown>,
  ctx: ComboSystemPromptTemplateContext
): Record<string, unknown> {
  const values: Record<string, string> = {
    "{{MODEL_ID}}": ctx.modelId,
    "{{PROVIDER_ID}}": ctx.providerId,
    "{{ACCOUNT}}": ctx.account,
    "{{FINGERPRINT}}": ctx.fingerprint,
  };
  const result = { ...body };
  if (typeof result.instructions === "string") {
    result.instructions = expandStringTemplates(result.instructions, values);
    return result;
  }
  const messages = result.messages;
  if (Array.isArray(messages)) {
    const first = messages[0] as Record<string, unknown> | undefined;
    if (
      first &&
      (first.role === "system" || first.role === "developer") &&
      typeof first.content === "string"
    ) {
      const next = [...messages];
      next[0] = { ...first, content: expandStringTemplates(first.content, values) };
      result.messages = next;
    }
  }
  return result;
}

/**
 * Gate + expand: expand the combo `system_message` template placeholders only
 * when the combo actually defines a non-empty `system_message`. Client-owned
 * content passes through untouched (single gate shared by every dispatch path).
 */
export function expandComboSystemPromptIfPresent(
  body: Record<string, unknown>,
  combo: { system_message?: string | null },
  ctx: ComboSystemPromptTemplateContext
): Record<string, unknown> {
  if (typeof combo.system_message === "string" && combo.system_message.trim()) {
    return expandComboSystemPromptTemplates(body, ctx);
  }
  return body;
}

/**
 * Resolve the device fingerprint for a combo target (#5501, #6087).
 * Only fingerprint-based providers carry fingerprints (see isFingerprintProvider).
 * Priority: explicit pin (`pinnedFingerprint`, combo builder) → the `@fp:`
 * suffix in `executionKey` (auto-rotation).
 * Returns null when none is knowable (the first fingerprint of an auto-rotated
 * set keeps the bare execution key — documented limitation).
 */
export function resolveTargetFingerprint(target: {
  provider: string;
  pinnedFingerprint?: string;
  executionKey?: string;
}): string | null {
  if (!isFingerprintProvider(target.provider)) return null;
  if (target.pinnedFingerprint) return target.pinnedFingerprint;
  const key = target.executionKey;
  if (key) {
    const marker = "@fp:";
    const idx = key.lastIndexOf(marker);
    if (idx !== -1) return key.slice(idx + marker.length);
  }
  return null;
}