diff --git a/config/quality/file-size-baseline.json b/config/quality/file-size-baseline.json index 374cef8c1a..1ca6cccb9f 100644 --- a/config/quality/file-size-baseline.json +++ b/config/quality/file-size-baseline.json @@ -142,14 +142,15 @@ "_rebaseline_pr4592_exclude_exhausted_auto": "Reconcile #4592 already-merged growth: combo.ts 2991->3036 (+45, terminal-status quota-cutoff exclusion in buildAutoCandidates + opt-in gate). Fast-gate PR->release does not run check:file-size.", "_rebaseline_2026_06_24_headroom_strategy": "Headroom-aware connection selection (dario technique): combo.ts 3168->3180 (+12 = a new `else if (strategy === \"headroom\")` dispatch branch in handleComboChat that delegates to orderTargetsByHeadroom + its log line, plus the import). The actual logic lives OUT of the god-file: the pure ranker rankByHeadroom/computeHeadroom is the new leaf open-sse/services/combo/headroomRanking.ts (91 LOC, 3190 (+10 = one new `else if (strategy === \"quota-share\")` dispatch branch in handleComboChat that delegates 100% to selectQuotaShareTarget + its log line, plus the import). All the new logic lives OUT of the god-file in two new leaves under open-sse/services/combo/: quotaShareInflight.ts (in-flight counter with TTL/lease, ~150 LOC 3225 (+35) = one new `else if (strategy === \"task-aware\")` dispatch branch delegating 100% to selectTaskAwareTarget + its imports/log lines. All scoring/classification logic lives OUT of the god-file in the new leaf open-sse/services/taskAwareRouting.ts (553 LOC Boolean(m)); const cfg = config as Record; - const judgeModel = - typeof cfg.judgeModel === "string" ? cfg.judgeModel : undefined; + const judgeModel = typeof cfg.judgeModel === "string" ? cfg.judgeModel : undefined; const tuning = cfg.fusionTuning && typeof cfg.fusionTuning === "object" ? (cfg.fusionTuning as FusionTuning) @@ -1020,7 +1024,8 @@ export async function handleComboChat({ const isTargetSelectableForWeighted = async (target: ResolvedComboTarget): Promise => { const rawModel = parseModel(target.modelStr).model || target.modelStr; - if (target.provider && getCircuitBreaker(target.provider).getStatus().state === "OPEN") return false; + if (target.provider && getCircuitBreaker(target.provider).getStatus().state === "OPEN") + return false; if ( resilienceSettings.providerCooldown.enabled && Boolean(target.provider && target.provider !== "unknown") && @@ -1098,16 +1103,21 @@ export async function handleComboChat({ : null; const getWeightedStepKeyForTarget = (target: ResolvedComboTarget): string | null => { if (!weightedResolution?.orderedSteps) return null; - const step = weightedResolution.orderedSteps.find((entry) => - target.executionKey === entry.executionKey || - target.executionKey.startsWith(entry.executionKey + ">") + const step = weightedResolution.orderedSteps.find( + (entry) => + target.executionKey === entry.executionKey || + target.executionKey.startsWith(entry.executionKey + ">") ); return step?.executionKey || null; }; let orderedTargets = strategy === "weighted" ? weightedResolution?.orderedTargets || [] - : resolveComboTargets(expandedCombo, expandedAllCombos, clampComboDepth(config.maxComboDepth)); + : resolveComboTargets( + expandedCombo, + expandedAllCombos, + clampComboDepth(config.maxComboDepth) + ); orderedTargets = await applyRequestTagRouting(orderedTargets, body, log); @@ -1566,10 +1576,33 @@ export async function handleComboChat({ `Quota-share ordering: ${orderedTargets[0]?.modelStr}${orderedTargets[0]?.connectionId ? ` (${orderedTargets[0].connectionId})` : ""} selected (DRR+P2C)` ); } - const _sticky = await applySessionStickiness(orderedTargets, body.messages as Array<{ role?: string; content?: unknown }>); orderedTargets = _sticky.targets; + const _sticky = await applySessionStickiness( + orderedTargets, + body.messages as Array<{ role?: string; content?: unknown }> + ); + orderedTargets = _sticky.targets; orderedTargets = orderTargetsByEvalScores(orderedTargets, config.evalRouting, log); orderedTargets = filterTargetsByRequestCompatibility(orderedTargets, body, log); + // Task-aware reordering: only active for strategies ["smart","task","task-aware","task_aware","auto"]. + // Additive — does not affect any of the other 15 strategies. + if (isTaskRoutingStrategy(strategy)) { + const task = classifyTask(body); + const conversationCacheKey = getConversationCacheKey(body); + const taskReordered = reorderByTaskWeight(orderedTargets, task); + if (taskReordered[0]?.modelStr !== orderedTargets[0]?.modelStr) { + const reasons = + Array.isArray(task.reasons) && task.reasons.length > 0 + ? ` (${task.reasons.join(",")})` + : ""; + log.info( + "COMBO", + `task-route task=${task.level}${reasons} cacheKey=${conversationCacheKey ?? "none"} → ${taskReordered[0]?.modelStr}` + ); + } + orderedTargets = taskReordered; + } + // Parallel pre-screen: check provider profiles and model availability for all targets // Only runs for priority strategy where sequential checking causes latency const preScreenMap = @@ -2101,7 +2134,8 @@ export async function handleComboChat({ } } } - if (_sticky.messageHash && target.connectionId) recordStickyBinding(_sticky.messageHash, target.connectionId); // LKGP (#919): + if (_sticky.messageHash && target.connectionId) + recordStickyBinding(_sticky.messageHash, target.connectionId); // LKGP (#919): if (provider) { const connId = effectiveConnectionId || undefined; void (async () => { @@ -2648,7 +2682,8 @@ async function handleRoundRobinCombo({ if (stickyTarget) { const rawModel = parseModel(stickyTarget.modelStr).model || stickyTarget.modelStr; const stickyAvailable = - (!stickyTarget.provider || getCircuitBreaker(stickyTarget.provider).getStatus().state !== "OPEN") && + (!stickyTarget.provider || + getCircuitBreaker(stickyTarget.provider).getStatus().state !== "OPEN") && !( resilienceSettings.providerCooldown.enabled && Boolean(stickyTarget.provider && stickyTarget.provider !== "unknown") && diff --git a/open-sse/services/taskAwareRouting.ts b/open-sse/services/taskAwareRouting.ts new file mode 100644 index 0000000000..3cc2da730c --- /dev/null +++ b/open-sse/services/taskAwareRouting.ts @@ -0,0 +1,553 @@ +/** + * Task-aware routing layer for combo routing. + * + * Derives request difficulty (light / standard / heavy / critical) from cheap, + * local structural signals — no LLM call. Routes heavier tasks toward higher-power + * models using a continuous `modelPowerScore`. Also provides conversation-affinity + * sticky round-robin for prompt-cache efficiency. + * + * Ported from upstream PR #2045 (decolua/9router) by @nguyenxvotanminh3. + * Adaptation: operates on ResolvedComboTarget[] (TS target objects) rather than + * plain string arrays, uses getResolvedModelCapabilities (OmniRoute TS API) instead + * of getCapabilitiesForModel, and is wired additively — only applies when + * isTaskRoutingStrategy() returns true. All other strategies are unaffected. + * + * ReDoS safety: all regexes use word-boundary anchors with alternation of fixed + * literals, no variable-length quantifiers on overlapping groups. Safe per + * CLAUDE.md PII/Regex learnings. + */ + +import { createHash } from "node:crypto"; +import { getResolvedModelCapabilities } from "./modelCapabilities.ts"; +import type { ResolvedComboTarget } from "./combo/types.ts"; + +// ── Task level constants ────────────────────────────────────────────────────── + +export const TASK_LEVEL_WEIGHT = { + light: 1, + standard: 2, + heavy: 3, + critical: 4, +} as const; + +export type TaskLevel = keyof typeof TASK_LEVEL_WEIGHT; + +export const TASK_TARGET_POWER: Record = { + light: 35, + standard: 65, + heavy: 95, + critical: 120, +}; + +// ReDoS-safe: all alternations are fixed-length word literals under \b anchors. +// No overlapping or nested quantifiers. +export const LIGHT_TASK_RE = + /\b(hi|hello|thanks|thank you|ping|format|rewrite|grammar|translate|summari[sz]e|short|quick|one[- ]?liner|explain briefly)\b/i; + +export const HEAVY_TASK_RE = + /\b(debug|root cause|architecture|architectural|refactor|migrate|implementation|implement|design|analy[sz]e|investigate|compare|benchmark|whitebox|codebase|end[- ]?to[- ]?end|e2e)\b/i; + +export const CRITICAL_TASK_RE = + /\b(critical|security|vulnerability|exploit|rce|remote code execution|supply chain|account takeover|auth bypass|privilege escalation|tenant|cross[- ]tenant|sandbox escape|ssrf|deserialization|prod incident|data exfiltration|bug bounty)\b/i; + +// ── Conversation affinity state ─────────────────────────────────────────────── + +/** @internal exported only for testing */ +export const comboConversationAffinity = new Map(); +const CONVERSATION_AFFINITY_TTL_MS = 60 * 60 * 1000; // 1 hour +const MAX_CONVERSATION_AFFINITY_ENTRIES = 1000; + +// ── Strategy gate ───────────────────────────────────────────────────────────── + +/** + * Returns true when the combo strategy is one of the task-aware strategies. + * Task routing is additive: other strategies are wholly unaffected. + */ +export function isTaskRoutingStrategy(strategy: unknown): boolean { + return ["smart", "task", "task-aware", "task_aware", "auto"].includes( + String(strategy ?? "").toLowerCase() + ); +} + +// ── Internal helpers ────────────────────────────────────────────────────────── + +function taskWeight(level: TaskLevel): number { + return TASK_LEVEL_WEIGHT[level]; +} + +function collectText(value: unknown, out: string[] = []): string[] { + if (value == null) return out; + if (typeof value === "string" || typeof value === "number" || typeof value === "boolean") { + out.push(String(value)); + return out; + } + if (Array.isArray(value)) { + for (const item of value) collectText(item, out); + return out; + } + if (typeof value !== "object") return out; + + const rec = value as Record; + if (typeof rec["text"] === "string") out.push(rec["text"]); + if (typeof rec["input_text"] === "string") out.push(rec["input_text"]); + if (typeof rec["output_text"] === "string") out.push(rec["output_text"]); + if (typeof rec["content"] === "string") out.push(rec["content"]); + else if (Array.isArray(rec["content"])) collectText(rec["content"], out); + if (Array.isArray(rec["parts"])) collectText(rec["parts"], out); + if (typeof rec["query"] === "string") out.push(rec["query"]); + if (typeof rec["url"] === "string") out.push(rec["url"]); + + return out; +} + +function estimatePromptChars(body: Record): number { + const contents = + (body["contents"] as unknown) ?? + ((body["request"] as Record | undefined)?.["contents"] as unknown); + const parts = [ + body["system"], + body["instructions"], + body["messages"], + body["input"], + contents, + body["query"], + body["url"], + ]; + return collectText(parts).join("\n").length; +} + +function countMessages(body: Record): number { + const contents = + (body["contents"] as unknown[]) ?? + ((body["request"] as Record | undefined)?.["contents"] as unknown[]); + return ( + (Array.isArray(body["messages"]) ? body["messages"].length : 0) + + (Array.isArray(body["input"]) ? body["input"].length : 0) + + (Array.isArray(body["contents"]) ? (body["contents"] as unknown[]).length : 0) + + (Array.isArray(contents) ? contents.length : 0) + ); +} + +function maxRequestedOutput(body: Record): number { + const genConf = body["generationConfig"] as Record | undefined; + const candidates = [ + body["max_tokens"], + body["max_output_tokens"], + body["max_completion_tokens"], + genConf?.["maxOutputTokens"], + ] + .map((v) => Number.parseInt(String(v ?? ""), 10)) + .filter((v) => Number.isFinite(v)); + return candidates.length > 0 ? Math.max(...candidates) : 0; +} + +function getTaskText(body: Record): string { + const contents = + (body?.["contents"] as unknown) ?? + ((body?.["request"] as Record | undefined)?.["contents"] as unknown); + return collectText([ + body?.["system"], + body?.["instructions"], + body?.["messages"], + body?.["input"], + contents, + body?.["query"], + body?.["url"], + ]).join("\n"); +} + +function normalizeEffort(body: Record): string { + const reasoning = body?.["reasoning"] as Record | undefined; + return String(body?.["reasoning_effort"] ?? reasoning?.["effort"] ?? "").toLowerCase(); +} + +// ── Task signals ────────────────────────────────────────────────────────────── + +export interface TaskSignals { + promptChars: number; + messageCount: number; + toolCount: number; + outputTokens: number; + effort: string; + hasExplicitReasoning: boolean; + lightKeyword: boolean; + heavyKeyword: boolean; + criticalKeyword: boolean; +} + +export function getTaskSignals(body: Record): TaskSignals { + const promptChars = estimatePromptChars(body); + const messageCount = countMessages(body); + const toolCount = Array.isArray(body?.["tools"]) ? (body["tools"] as unknown[]).length : 0; + const outputTokens = maxRequestedOutput(body); + const effort = normalizeEffort(body); + const text = getTaskText(body); + + return { + promptChars, + messageCount, + toolCount, + outputTokens, + effort, + hasExplicitReasoning: Boolean( + effort && effort !== "none" && effort !== "off" && effort !== "disabled" + ), + lightKeyword: LIGHT_TASK_RE.test(text), + heavyKeyword: HEAVY_TASK_RE.test(text), + criticalKeyword: CRITICAL_TASK_RE.test(text), + }; +} + +// ── Task classification ─────────────────────────────────────────────────────── + +export interface TaskClassification extends TaskSignals { + level: TaskLevel; + weight: number; + reasons: string[]; +} + +/** + * Classify request difficulty for smart combo routing. + * + * Deliberately uses cheap, local signals only — no LLM call. It is a routing + * hint: light requests stay on fast/cheap models while large, tool-heavy, + * security-sensitive, or reasoning-heavy requests try stronger models first. + * Fallback still tries every model. + */ +export function classifyTask(body: Record): TaskClassification { + const s = getTaskSignals(body ?? {}); + const reasons: string[] = []; + const add = (condition: boolean, reason: string): boolean => { + if (condition) reasons.push(reason); + return condition; + }; + + const effortIsHigh = /^(high|xhigh|max|maximum|hard|deep)$/.test(s.effort); + const effortIsLight = + !s.hasExplicitReasoning || /^(low|minimal|none|off|disabled)$/.test(s.effort); + + const critical = + add(s.promptChars >= 100000, "huge-context") || + add(s.outputTokens >= 32768, "huge-output") || + add(s.toolCount >= 8 && s.promptChars >= 16000, "many-tools-large-context") || + add( + s.criticalKeyword && (effortIsHigh || s.toolCount >= 3 || s.promptChars >= 8000), + "critical-domain" + ); + + if (critical) { + return { level: "critical", weight: taskWeight("critical"), ...s, reasons }; + } + + const heavySignalCount = [ + add(s.promptChars >= 50000, "large-context"), + add(s.promptChars >= 24000, "medium-large-context"), + add(s.messageCount >= 16, "long-conversation"), + add(s.toolCount >= 4, "many-tools"), + add(s.outputTokens >= 8192, "large-output"), + add(effortIsHigh, "high-reasoning-effort"), + add(s.criticalKeyword, "security-sensitive"), + add(s.heavyKeyword && s.promptChars >= 4000, "complex-task"), + ].filter(Boolean).length; + + if (heavySignalCount >= 2 || s.promptChars >= 50000 || effortIsHigh) { + return { level: "heavy", weight: taskWeight("heavy"), ...s, reasons }; + } + + const light = + s.promptChars <= 2000 && + s.messageCount <= 3 && + s.toolCount === 0 && + s.outputTokens <= 1500 && + effortIsLight && + !s.criticalKeyword && + !s.heavyKeyword; + + if ( + light || + (s.lightKeyword && + s.promptChars <= 4000 && + s.toolCount === 0 && + effortIsLight && + !s.criticalKeyword) + ) { + return { + level: "light", + weight: taskWeight("light"), + ...s, + reasons: reasons.length > 0 ? reasons : ["small-simple-request"], + }; + } + + return { + level: "standard", + weight: taskWeight("standard"), + ...s, + reasons: reasons.length > 0 ? reasons : ["default"], + }; +} + +// ── Model power scoring ─────────────────────────────────────────────────────── + +/** + * Estimate how capable a model is on a continuous 0–150 scale, using both + * registry-derived capabilities and heuristic name matching. + */ +export function modelPowerScore(modelStr: string): number { + const id = `${modelStr ?? ""}`.toLowerCase(); + const caps = getResolvedModelCapabilities(modelStr); + + let score = 35; + if (caps.reasoning) score += 18; + if (caps.supportsVision === true) score += 3; + if (caps.toolCalling) score += 3; + + const ctx = caps.contextWindow ?? 0; + if (ctx >= 1_000_000) score += 22; + else if (ctx >= 400_000) score += 15; + else if (ctx >= 200_000) score += 9; + else if (ctx > 0 && ctx <= 32_000) score -= 10; + + const maxOut = caps.maxOutputTokens ?? 0; + if (maxOut >= 128_000) score += 12; + else if (maxOut >= 64_000) score += 8; + else if (maxOut > 0 && maxOut <= 8_192) score -= 8; + + if ( + /\b(opus|mythos|gpt-5|o3|o4|pro|max|ultra|deepseek-v4-pro|sonnet-4|glm-5|kimi-k2\.7|minimax-m3|reasoner)\b/i.test( + id + ) + ) + score += 28; + if (/\b(coder|code|coding)\b/i.test(id)) score += 8; + if (/\b(haiku|flash|mini|lite|small|nano|instant|fast|turbo|3\.5|8b|7b)\b/i.test(id)) score -= 24; + + return Math.max(0, Math.min(150, score)); +} + +// Hard capabilities: missing one drops request data (images, PDFs). Maps TS +// ResolvedModelCapabilities fields that correspond to hard-cap modalities. +const HARD_CAP_CHECKS = new Set(["vision"]); + +/** + * Score a single model for a given task + capability requirements. + * Higher score = better fit. Negative score = capability hard-miss. + */ +export function scoreModelForTask( + modelStr: string, + task: TaskClassification = classifyTask({}), + required: Set = new Set() +): number { + const caps = getResolvedModelCapabilities(modelStr); + const target = TASK_TARGET_POWER[task.level]; + const power = modelPowerScore(modelStr); + let score = 100 - Math.abs(power - target); + + // Hard capability misses: drop score heavily so the model sorts to the back + // but stays in the list (fallback still tries every model). + for (const cap of required) { + if (!HARD_CAP_CHECKS.has(cap)) continue; + if (cap === "vision" && caps.supportsVision !== true) score -= 10000; + } + + if ((required.has("reasoning") || task.weight >= TASK_LEVEL_WEIGHT.heavy) && !caps.reasoning) + score -= 120; + if (required.has("search") && !caps.toolCalling) score -= 30; + + const estimatedPromptTokens = Math.ceil((task.promptChars ?? 0) / 4); + const ctxLimit = caps.contextWindow ?? 0; + if (ctxLimit > 0 && estimatedPromptTokens > ctxLimit * 0.85) score -= 200; + + const maxOut = caps.maxOutputTokens ?? 0; + if (maxOut > 0 && task.outputTokens > 0 && task.outputTokens > maxOut) score -= 80; + + if (task.level === "light" && power > 95) score -= 35; + if (task.level === "standard" && power > 125) score -= 10; + if (task.level === "heavy" && power < 65) score -= 60; + if (task.level === "critical" && power < 85) score -= 100; + + return score; +} + +/** + * Reorder ResolvedComboTarget[] so the best-fit model for the task comes first. + * Stable: ties keep original order. Identity-returns when no reordering needed + * (avoids allocations on the common path). Never removes targets. + */ +export function reorderByTaskWeight( + targets: ResolvedComboTarget[], + task: TaskClassification = classifyTask({}), + required: Set = new Set() +): ResolvedComboTarget[] { + if (!Array.isArray(targets) || targets.length <= 1) return targets; + + const reordered = targets + .map((t, i) => ({ t, i, score: scoreModelForTask(t.modelStr, task, required) })) + .sort((a, b) => b.score - a.score || a.i - b.i) + .map((x) => x.t); + + return reordered.every((t, i) => t === targets[i]) ? targets : reordered; +} + +// ── Conversation affinity (cache-key derivation) ────────────────────────────── + +function normalizeFingerprintText(value: unknown): string { + return String(value ?? "") + .replace(/\s+/g, " ") + .trim() + .slice(0, 12000); +} + +function firstRoleText( + items: unknown[], + roles: Set, + contentKey: "content" | "parts" = "content" +): string { + if (!Array.isArray(items)) return ""; + for (const item of items) { + if (!item || typeof item !== "object") continue; + const rec = item as Record; + if (!roles.has(String(rec["role"] ?? ""))) continue; + const raw = contentKey === "parts" ? rec["parts"] : rec["content"]; + const text = normalizeFingerprintText(collectText(raw).join("\n")); + if (text) return text; + } + return ""; +} + +function allRoleText( + items: unknown[], + roles: Set, + contentKey: "content" | "parts" = "content" +): string { + if (!Array.isArray(items)) return ""; + return normalizeFingerprintText( + items + .filter( + (item): item is Record => + !!item && + typeof item === "object" && + roles.has(String((item as Record)["role"] ?? "")) + ) + .map((item) => + collectText(contentKey === "parts" ? item["parts"] : item["content"]).join("\n") + ) + .filter(Boolean) + .join("\n") + ); +} + +function hashConversationSeed(seed: string): string | null { + const normalized = normalizeFingerprintText(seed); + if (!normalized) return null; + return createHash("sha1").update(normalized).digest("hex").slice(0, 24); +} + +/** + * Derive a stable cache-affinity key from explicit thread metadata when present, + * otherwise from the immutable start of the prompt (system + first user turn). + * Appended turns should not move an existing conversation to another model. + */ +export function getConversationCacheKey(body: Record): string | null { + if (!body || typeof body !== "object") return null; + + const meta = body["metadata"] as Record | undefined; + const explicitCandidates = [ + body["conversation_id"], + body["conversationId"], + body["thread_id"], + body["threadId"], + body["session_id"], + body["sessionId"], + meta?.["conversation_id"], + meta?.["conversationId"], + meta?.["thread_id"], + meta?.["threadId"], + meta?.["session_id"], + meta?.["sessionId"], + ]; + const explicit = explicitCandidates.find((v) => v != null && String(v).trim()); + if (explicit != null) return hashConversationSeed(`explicit:${String(explicit).trim()}`); + + const systemRoles = new Set(["system", "developer"]); + const userRoles = new Set(["user"]); + const contents = + (body["contents"] as unknown[]) ?? + ((body["request"] as Record | undefined)?.["contents"] as unknown[]); + + const seedParts = [ + collectText(body["system"]).join("\n"), + collectText(body["instructions"]).join("\n"), + allRoleText((body["messages"] as unknown[]) ?? [], systemRoles), + allRoleText((body["input"] as unknown[]) ?? [], systemRoles), + allRoleText(contents ?? [], systemRoles, "parts"), + firstRoleText((body["messages"] as unknown[]) ?? [], userRoles), + typeof body["input"] === "string" + ? body["input"] + : firstRoleText((body["input"] as unknown[]) ?? [], userRoles), + firstRoleText(contents ?? [], userRoles, "parts"), + body["query"], + body["url"], + ].filter(Boolean); + + return hashConversationSeed(seedParts.join("\n")); +} + +// ── Affinity management (for round-robin integration) ──────────────────────── + +/** @internal exported for testing */ +export function pruneConversationAffinity(now = Date.now()): void { + for (const [key, value] of comboConversationAffinity) { + if (!value || now - value.lastUsed > CONVERSATION_AFFINITY_TTL_MS) { + comboConversationAffinity.delete(key); + } + } + while (comboConversationAffinity.size > MAX_CONVERSATION_AFFINITY_ENTRIES) { + const oldestKey = comboConversationAffinity.keys().next().value; + if (oldestKey === undefined) break; + comboConversationAffinity.delete(oldestKey); + } +} + +/** + * Returns the pinned target index for a conversation, or null if none. + * Creates and stores an affinity entry on first call for a new conversation. + * + * Used by getRotatedModels (round-robin) when stickyLimit > 1. + */ +export function getOrSetConversationAffinityIndex( + rotationKey: string, + conversationCacheKey: string, + currentIndex: number +): number { + const now = Date.now(); + pruneConversationAffinity(now); + + const affinityKey = `${rotationKey}:${conversationCacheKey}`; + const existing = comboConversationAffinity.get(affinityKey); + if (existing) { + const pinnedIndex = existing.index; + // Refresh TTL (move to end of Map iteration order) + comboConversationAffinity.delete(affinityKey); + comboConversationAffinity.set(affinityKey, { index: pinnedIndex, lastUsed: now }); + return pinnedIndex; + } + + comboConversationAffinity.set(affinityKey, { index: currentIndex, lastUsed: now }); + return currentIndex; +} + +/** + * Clear affinity entries for a specific combo (or all if no name given). + * Called by resetComboRotation. + */ +export function clearConversationAffinity(comboName?: string): void { + if (comboName) { + const prefix = `${comboName}:`; + for (const key of comboConversationAffinity.keys()) { + if (key.startsWith(prefix)) comboConversationAffinity.delete(key); + } + } else { + comboConversationAffinity.clear(); + } +} diff --git a/tests/unit/combo-task-aware.test.ts b/tests/unit/combo-task-aware.test.ts new file mode 100644 index 0000000000..b8a9887c7b --- /dev/null +++ b/tests/unit/combo-task-aware.test.ts @@ -0,0 +1,429 @@ +/** + * Tests for task-aware combo routing (port of upstream PR #2045). + * + * Coverage: + * - classifyTask: light / standard / heavy / critical levels + * - reorderByTaskWeight: floats higher-power models for heavy tasks + * - getConversationCacheKey: stable per conversation, distinct per thread + * - getOrSetConversationAffinityIndex: affinity pin + TTL behaviour + * - isTaskRoutingStrategy: gate for task-aware strategies + * - Guard: non-task-aware strategies are unaffected (identity path) + */ + +import { describe, it, before, after } from "node:test"; +import assert from "node:assert/strict"; +import { + classifyTask, + getConversationCacheKey, + getOrSetConversationAffinityIndex, + isTaskRoutingStrategy, + reorderByTaskWeight, + scoreModelForTask, + TASK_LEVEL_WEIGHT, + comboConversationAffinity, + clearConversationAffinity, +} from "../../open-sse/services/taskAwareRouting.ts"; +import type { ResolvedComboTarget } from "../../open-sse/services/combo/types.ts"; + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +function makeTarget(modelStr: string, idx = 0): ResolvedComboTarget { + const [provider = "", model = modelStr] = modelStr.includes("/") + ? modelStr.split("/") + : ["", modelStr]; + return { + kind: "model" as const, + stepId: `step-${idx}`, + executionKey: `exec-${idx}`, + modelStr, + provider, + providerId: null, + connectionId: null, + weight: 1, + label: null, + }; +} + +function makeBody(overrides: Record = {}): Record { + return { messages: [{ role: "user", content: "Hello" }], ...overrides }; +} + +// ── classifyTask ────────────────────────────────────────────────────────────── + +describe("classifyTask", () => { + it("classifies empty body as light (all signals tiny → default-light path)", () => { + const t = classifyTask({}); + // Empty body: 0 chars, 0 messages, 0 tools, 0 output tokens → light + assert.equal(t.level, "light"); + assert.equal(t.weight, TASK_LEVEL_WEIGHT.light); + }); + + it("classifies a simple short greeting as light", () => { + const t = classifyTask({ messages: [{ role: "user", content: "hi there" }], max_tokens: 100 }); + assert.equal(t.level, "light"); + assert.equal(t.weight, TASK_LEVEL_WEIGHT.light); + }); + + it("classifies a translate request as light", () => { + const t = classifyTask({ + messages: [{ role: "user", content: "translate this short sentence" }], + max_tokens: 256, + }); + assert.equal(t.level, "light"); + }); + + it("classifies a quick rewrite request as light", () => { + const t = classifyTask({ + messages: [{ role: "user", content: "quick rewrite this sentence to be more concise" }], + max_tokens: 300, + }); + assert.equal(t.level, "light"); + }); + + it("classifies long conversation + big prompt as heavy (two signals)", () => { + // 20 messages (1 signal) + prompt >= 24000 chars (1 signal) = 2 → heavy + const bigContent = "context ".repeat(3500); // ~28000 chars + const messages = [ + { role: "user", content: bigContent }, + ...Array.from({ length: 20 }, (_, i) => ({ + role: i % 2 === 0 ? "assistant" : "user", + content: "turn " + i, + })), + ]; + const t = classifyTask({ messages }); + assert.equal(t.level, "heavy"); + assert.equal(t.weight, TASK_LEVEL_WEIGHT.heavy); + }); + + it("classifies high reasoning effort as heavy", () => { + const t = classifyTask({ + messages: [{ role: "user", content: "analyze this code" }], + reasoning_effort: "high", + }); + assert.equal(t.level, "heavy"); + }); + + it("classifies huge context (>= 100k chars) as critical", () => { + const t = classifyTask({ + messages: [{ role: "user", content: "x".repeat(120_000) }], + }); + assert.equal(t.level, "critical"); + assert.equal(t.weight, TASK_LEVEL_WEIGHT.critical); + }); + + it("classifies huge output request (>= 32768 tokens) as critical", () => { + const t = classifyTask({ + messages: [{ role: "user", content: "write a comprehensive report" }], + max_tokens: 40_000, + }); + assert.equal(t.level, "critical"); + }); + + it("classifies security+tools+effort as critical", () => { + const t = classifyTask({ + messages: [ + { + role: "user", + content: `Find a critical bug bounty attack chain for RCE or supply chain impact.\n${"context ".repeat(2000)}`, + }, + ], + tools: [{ name: "read" }, { name: "grep" }, { name: "web" }, { name: "bash" }], + reasoning_effort: "high", + max_tokens: 12_000, + }); + assert.equal(t.level, "critical"); + }); + + it("classifies moderate prompt without signals as standard (must exceed light thresholds)", () => { + // To be classified standard (not light), prompt must either: be > 4000 chars, + // have > 3 messages, tools, or not match a light keyword — "explain briefly" would be light. + // A prompt with "analyze" keyword + 5000 chars crosses into standard. + const t = classifyTask({ + messages: [ + { + role: "user", + content: "Please analyze how TCP works in detail.\n" + "context ".repeat(500), + }, + { role: "assistant", content: "TCP uses..." }, + { role: "user", content: "And what about UDP?" }, + { role: "assistant", content: "UDP is..." }, + ], + }); + assert.equal(t.level, "standard"); + }); + + it("returns reasons array for each classification", () => { + const t = classifyTask({ + messages: [{ role: "user", content: "x".repeat(120_000) }], + }); + assert.ok(Array.isArray(t.reasons), "reasons should be an array"); + assert.ok(t.reasons.length > 0, "critical task should have at least one reason"); + }); +}); + +// ── reorderByTaskWeight ─────────────────────────────────────────────────────── + +describe("reorderByTaskWeight", () => { + it("returns same reference for single-target list", () => { + const targets = [makeTarget("anthropic/claude-haiku-4.5", 0)]; + const task = classifyTask({ messages: [{ role: "user", content: "hello" }] }); + const out = reorderByTaskWeight(targets, task, new Set()); + assert.strictEqual(out, targets); + }); + + it("returns same reference for empty list", () => { + const targets: ResolvedComboTarget[] = []; + const task = classifyTask({}); + const out = reorderByTaskWeight(targets, task, new Set()); + assert.strictEqual(out, targets); + }); + + it("routes light tasks to lighter models (haiku before opus)", () => { + const targets = [ + makeTarget("anthropic/claude-opus-4.6", 0), + makeTarget("anthropic/claude-haiku-4.5", 1), + ]; + const task = classifyTask({ + messages: [{ role: "user", content: "quick rewrite this sentence" }], + max_tokens: 300, + }); + const out = reorderByTaskWeight(targets, task, new Set()); + assert.equal(out[0].modelStr, "anthropic/claude-haiku-4.5"); + }); + + it("routes critical tasks to stronger models (opus before haiku)", () => { + const targets = [ + makeTarget("anthropic/claude-haiku-4.5", 0), + makeTarget("anthropic/claude-opus-4.6", 1), + ]; + const task = classifyTask({ + messages: [ + { + role: "user", + content: `Security review for account takeover and cross-tenant RCE impact.\n${"code ".repeat(3000)}`, + }, + ], + tools: [{ name: "read" }, { name: "grep" }, { name: "bash" }, { name: "web" }], + reasoning_effort: "high", + }); + const required = new Set(["reasoning"]); + const out = reorderByTaskWeight(targets, task, required); + assert.equal(out[0].modelStr, "anthropic/claude-opus-4.6"); + assert.ok( + scoreModelForTask(out[0].modelStr, task, required) >= + scoreModelForTask(out[1].modelStr, task, required), + "first model should score >= second" + ); + }); + + it("keeps hard-cap vision model first even for light tasks", () => { + // anthropic/claude-haiku-4.5 has vision; deepseek-chat does not + const targets = [ + makeTarget("deepseek/deepseek-chat", 0), + makeTarget("anthropic/claude-haiku-4.5", 1), + ]; + const task = classifyTask({ messages: [{ role: "user", content: "what is in this image?" }] }); + const out = reorderByTaskWeight(targets, task, new Set(["vision"])); + assert.equal(out[0].modelStr, "anthropic/claude-haiku-4.5"); + }); + + it("never drops any targets", () => { + const targets = [ + makeTarget("anthropic/claude-opus-4.6", 0), + makeTarget("anthropic/claude-haiku-4.5", 1), + makeTarget("deepseek/deepseek-chat", 2), + ]; + const task = classifyTask({ messages: [{ role: "user", content: "x".repeat(120_000) }] }); + const out = reorderByTaskWeight(targets, task, new Set()); + assert.equal(out.length, 3); + }); + + it("is stable: ties preserve original order", () => { + // Two identical models should not swap positions + const targets = [ + makeTarget("anthropic/claude-haiku-4.5", 0), + makeTarget("anthropic/claude-haiku-4.5", 1), + ]; + const task = classifyTask({}); + const out = reorderByTaskWeight(targets, task, new Set()); + assert.equal(out[0].stepId, "step-0"); + assert.equal(out[1].stepId, "step-1"); + }); +}); + +// ── isTaskRoutingStrategy ───────────────────────────────────────────────────── + +describe("isTaskRoutingStrategy", () => { + it("returns true for task-aware strategy names", () => { + for (const name of ["smart", "task", "task-aware", "task_aware", "auto"]) { + assert.ok(isTaskRoutingStrategy(name), `Expected ${name} to be task-routing`); + } + }); + + it("is case-insensitive", () => { + assert.ok(isTaskRoutingStrategy("SMART")); + assert.ok(isTaskRoutingStrategy("Task-Aware")); + }); + + it("returns false for standard strategy names", () => { + for (const name of [ + "priority", + "fallback", + "round-robin", + "weighted", + "random", + "fill-first", + "p2c", + "least-used", + "cost-optimized", + "reset-aware", + "reset-window", + "context-optimized", + "headroom", + "lkgp", + "context-relay", + ]) { + assert.ok(!isTaskRoutingStrategy(name), `Expected ${name} to NOT be task-routing`); + } + }); + + it("returns false for null/undefined/empty", () => { + assert.ok(!isTaskRoutingStrategy(null)); + assert.ok(!isTaskRoutingStrategy(undefined)); + assert.ok(!isTaskRoutingStrategy("")); + }); +}); + +// ── getConversationCacheKey ─────────────────────────────────────────────────── + +describe("getConversationCacheKey", () => { + it("returns null for null input", () => { + assert.equal(getConversationCacheKey(null as unknown as Record), null); + }); + + it("returns a string for a body with messages", () => { + const key = getConversationCacheKey({ + messages: [ + { role: "system", content: "You are helpful" }, + { role: "user", content: "hi" }, + ], + }); + assert.equal(typeof key, "string"); + assert.ok((key as string).length > 0); + }); + + it("is stable across calls for the same conversation start", () => { + const body = { + messages: [ + { role: "system", content: "You are a helpful assistant" }, + { role: "user", content: "Hello, how are you?" }, + ], + }; + const k1 = getConversationCacheKey(body); + const k2 = getConversationCacheKey(body); + assert.equal(k1, k2); + }); + + it("differs for different first messages", () => { + const k1 = getConversationCacheKey({ + messages: [{ role: "user", content: "Explain TCP" }], + }); + const k2 = getConversationCacheKey({ + messages: [{ role: "user", content: "Explain UDP" }], + }); + assert.notEqual(k1, k2); + }); + + it("same key even when later messages are appended (conversation affinity)", () => { + const systemMsg = { role: "system", content: "You are a coding assistant" }; + const firstUser = { role: "user", content: "What is a pointer?" }; + const k1 = getConversationCacheKey({ messages: [systemMsg, firstUser] }); + const k2 = getConversationCacheKey({ + messages: [ + systemMsg, + firstUser, + { role: "assistant", content: "A pointer stores an address..." }, + { role: "user", content: "Can you give an example?" }, + ], + }); + assert.equal(k1, k2); + }); + + it("prefers explicit conversation_id when present", () => { + const k1 = getConversationCacheKey({ + conversation_id: "thread-abc-123", + messages: [{ role: "user", content: "hello" }], + }); + const k2 = getConversationCacheKey({ conversation_id: "thread-abc-123" }); + assert.equal(k1, k2); + assert.notEqual(k1, null); + }); + + it("different conversation_ids produce different keys", () => { + const k1 = getConversationCacheKey({ conversation_id: "thread-1" }); + const k2 = getConversationCacheKey({ conversation_id: "thread-2" }); + assert.notEqual(k1, k2); + }); +}); + +// ── Conversation affinity (getOrSetConversationAffinityIndex) ───────────────── + +describe("conversation affinity", () => { + before(() => clearConversationAffinity()); + after(() => clearConversationAffinity()); + + it("returns currentIndex for a new conversation key", () => { + const idx = getOrSetConversationAffinityIndex("combo-x", "conv-aaa", 2); + assert.equal(idx, 2); + }); + + it("returns the pinned index for the same conversation key", () => { + clearConversationAffinity(); + getOrSetConversationAffinityIndex("combo-y", "conv-bbb", 0); + // Second call should return the pinned index (0), not a new currentIndex (1) + const idx = getOrSetConversationAffinityIndex("combo-y", "conv-bbb", 1); + assert.equal(idx, 0); + }); + + it("different conversation keys get independent pins", () => { + clearConversationAffinity(); + const i1 = getOrSetConversationAffinityIndex("combo-z", "conv-c1", 0); + const i2 = getOrSetConversationAffinityIndex("combo-z", "conv-c2", 1); + assert.equal(i1, 0); + assert.equal(i2, 1); + // Both pins should stay + assert.equal(getOrSetConversationAffinityIndex("combo-z", "conv-c1", 99), 0); + assert.equal(getOrSetConversationAffinityIndex("combo-z", "conv-c2", 99), 1); + }); + + it("clearConversationAffinity clears only the named combo prefix", () => { + clearConversationAffinity(); + getOrSetConversationAffinityIndex("combo-a", "conv-1", 0); + getOrSetConversationAffinityIndex("combo-b", "conv-1", 1); + clearConversationAffinity("combo-a"); + // combo-a entry should be gone + const sizeAfter = comboConversationAffinity.size; + assert.equal(sizeAfter, 1); + }); +}); + +// ── Guard: non-task-aware strategy leaves targets unchanged ─────────────────── + +describe("non-task-aware strategy guard", () => { + it("reorderByTaskWeight is still safe to call externally but only wired for task strategies", () => { + // The function itself is pure — what matters is that combo.ts only calls it + // when isTaskRoutingStrategy() is true. Verify that explicitly. + const nonTaskStrategies = ["priority", "fallback", "round-robin", "cost-optimized"]; + for (const strategy of nonTaskStrategies) { + assert.ok(!isTaskRoutingStrategy(strategy), `${strategy} should not trigger task routing`); + } + }); + + it("reorderByTaskWeight on standard task with equal-power models returns same reference", () => { + // If all scores are equal, no reordering → returns same array reference + const targets = [makeTarget("anthropic/claude-haiku-4.5", 0)]; + const task = classifyTask({}); + const out = reorderByTaskWeight(targets, task, new Set()); + assert.strictEqual(out, targets); + }); +});