diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 7c4154129a..0ec3a8485c 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -48,7 +48,7 @@ import * as semaphore from "./rateLimitSemaphore.ts"; import { getCircuitBreaker } from "../../src/shared/utils/circuitBreaker"; import { fisherYatesShuffle, getNextFromDeck } from "../../src/shared/utils/shuffleDeck"; import { parseModel } from "./model.ts"; -import { applyComboAgentMiddleware, injectModelTag } from "./comboAgentMiddleware.ts"; +import { applyComboAgentMiddleware } from "./comboAgentMiddleware.ts"; import { checkCredentialGate, logCredentialSkip } from "./credentialGate.ts"; import { emit } from "../../src/lib/events/eventBus"; import { notifyWebhookEvent } from "../../src/lib/webhookDispatcher"; @@ -2682,161 +2682,34 @@ export async function handleComboChat({ | undefined, relayOptions?.universalHandoffConfig as Record | null | undefined ); + // ── Server-side context cache pinning (replaces tag roundtrip) ─ + // Uses session_model_history — no client-side tag injection, no visible output pollution. + let pinnedModel: string | null = null; + if ( + combo.context_cache_protection && + relayOptions?.sessionId && + !(body as Record)?.[SKIP_UNIVERSAL_HANDOFF_FLAG] + ) { + const pinned = getLastSessionModel(relayOptions.sessionId, combo.name); + if (pinned) { + body = { ...body, model: pinned }; + pinnedModel = pinned; + log.info("COMBO", `[#401] Context cache: pinned model=${pinned} (server-side)`); + } + } + // ── Combo Agent Middleware (#399 + #401) ──────────────────────────────── - // Apply system_message override, tool_filter_regex, and extract pinned model - // from context caching tag. These are all opt-in per combo config. - const { body: agentBody, pinnedModel } = applyComboAgentMiddleware( + // Apply system_message override, tool_filter_regex. + // Context cache pinning is handled above via session_model_history. + const { body: agentBody } = applyComboAgentMiddleware( body, combo, "" // provider/model not yet known — resolved per-model in loop ); body = agentBody; - if (pinnedModel) { - log.info("COMBO", `[#401] Context caching: pinned model=${pinnedModel}`); - } const clientRequestedStream = body?.stream === true; - // Wrap handleSingleModel to inject context caching tag on response (#401) - const handleSingleModelWrapped = combo.context_cache_protection - ? async (b: Record, modelStr: string, target?: SingleModelTarget) => { - const res = await handleSingleModel(b, modelStr, target); - if (!res.ok) return res; - - // Non-streaming: inject tag into JSON response - // Fix #721: Use OpenAI choices format (json.choices[0].message) not json.messages - if (!b.stream) { - try { - const json = await res.clone().json(); - const choice = json?.choices?.[0]; - if (choice?.message) { - // Wrap single message in array for injectModelTag, then unwrap - const tagged = injectModelTag([choice.message], modelStr); - // If the message had tool_calls but no string content, injectModelTag - // appends a synthetic assistant message — use the last one - const taggedMsg = tagged.at(-1); - const updatedJson = { - ...json, - choices: [{ ...choice, message: taggedMsg }, ...(json.choices?.slice(1) || [])], - }; - return new Response(JSON.stringify(updatedJson), { - status: res.status, - headers: res.headers, - }); - } - } catch { - /* non-JSON — skip tagging */ - } - return res; - } - - // Streaming (Fix #490 + #511): prepend omniModel tag into the first - // non-empty content chunk so it arrives BEFORE finish_reason:stop. - // SDKs close the connection on finish_reason, so anything sent after - // that marker is silently dropped. - if (!res.body) return res; - const tagContent = `${modelStr}`; - const encoder = new TextEncoder(); - const decoder = new TextDecoder(); - let tagInjected = false; - - const transform = new TransformStream( - { - transform(chunk, controller) { - if (tagInjected) { - // Already injected — passthrough - controller.enqueue(chunk); - return; - } - - const text = decoder.decode(chunk, { stream: true }); - - // Fix #721: Look for either non-empty content OR tool_calls in the - // SSE data. Tool-call-only responses have content:null, so we inject - // the tag when we see a finish_reason approaching, or on first content. - const contentMatch = RegExp(/"content":"([^"]+)/).exec(text); - if (contentMatch) { - // Inject tag at the beginning of the first content value - const injected = text.replace( - /"content":"([^"]+)/, - `"content":"${tagContent.replaceAll("\\", "\\\\").replaceAll('"', String.raw`\"`)}$1` - ); - tagInjected = true; - controller.enqueue(encoder.encode(injected)); - return; - } - - // Fix #721: For tool-call-only streams, inject the tag when we see - // the finish_reason chunk (before it reaches the client SDK which - // would close the connection). This ensures the tag roundtrips - // through the conversation history even when there's no text content. - if (text.includes('"finish_reason"') && !text.includes('"finish_reason":null')) { - // Inject a content chunk with the tag just before this finish chunk - const tagChunk = `data: ${JSON.stringify({ - choices: [ - { - delta: { content: tagContent }, - index: 0, - finish_reason: null, - }, - ], - })}\n\n`; - tagInjected = true; - controller.enqueue(encoder.encode(tagChunk)); - controller.enqueue(chunk); - return; - } - - // No content yet — passthrough - controller.enqueue(chunk); - }, - flush(controller) { - // If stream ends without ever finding content (edge case), - // inject tag as a standalone chunk before the stream closes - if (!tagInjected) { - const tagChunk = `data: ${JSON.stringify({ - choices: [ - { - delta: { content: tagContent }, - index: 0, - finish_reason: null, - }, - ], - })}\n\n`; - controller.enqueue(encoder.encode(tagChunk)); - } - }, - }, - { highWaterMark: 16384 }, - { highWaterMark: 16384 } - ); - - const transformedStream = res.body.pipeThrough(transform); - const headers = new Headers(); - if (res.headers) { - try { - res.headers.forEach((v, k) => { - headers.set(k, v); - }); - } catch { - try { - for (const [k, v] of res.headers as unknown as Iterable<[string, string]>) { - headers.set(k, v); - } - } catch { - try { - for (const [k, v] of Object.entries(res.headers)) { - headers.set(k, v == null ? "" : String(v)); - } - } catch {} - } - } - } - headers.set("X-OmniRoute-Model", modelStr); - return new Response(transformedStream, { - status: res.status, - headers, - }); - } - : handleSingleModel; + // Context cache pinning is handled above via server-side session_model_history. + // No tag injection on response — use handleSingleModel directly. // ───────────────────────────────────────────────────────────────────────── // Use config cascade before dispatch so all strategies, pinned context routes, @@ -2862,7 +2735,7 @@ export async function handleComboChat({ target?: SingleModelTarget ): Promise => { if (comboTargetTimeoutMs <= 0) { - return handleSingleModelWrapped(b, modelStr, target).catch((err) => + return handleSingleModel(b, modelStr, target).catch((err) => errorResponse(502, err?.message ?? "Upstream model error") ); } @@ -2904,7 +2777,7 @@ export async function handleComboChat({ } try { return await Promise.race([ - handleSingleModelWrapped(b, modelStr, targetWithSignal).catch((err) => { + handleSingleModel(b, modelStr, targetWithSignal).catch((err) => { if (timedOut) { // Inner call rejected because we aborted it. The synthetic 524 from // timeoutPromise already wins the race; return an empty response so @@ -3353,7 +3226,7 @@ export async function handleComboChat({ config, body, resolveShadowTargets(combo, config, allCombos), - handleSingleModelWrapped, + handleSingleModel, isModelAvailable, strategy, log @@ -3643,6 +3516,22 @@ export async function handleComboChat({ fallbackCount, }); + // Context cache pinning: record model usage for session-based pinning + // (independent of universal handoff — always fires when context_cache_protection is on) + if ( + combo.context_cache_protection && + relayOptions?.sessionId && + !(body as Record)?.[SKIP_UNIVERSAL_HANDOFF_FLAG] + ) { + recordSessionModelUsage( + relayOptions.sessionId, + combo.name, + modelStr, + provider, + target.connectionId ?? undefined + ); + } + // Universal handoff: record model usage for session if ( universalHandoffConfig.enabled && diff --git a/open-sse/services/comboAgentMiddleware.ts b/open-sse/services/comboAgentMiddleware.ts index 1b86b1cffd..28b5de5f4f 100644 --- a/open-sse/services/comboAgentMiddleware.ts +++ b/open-sse/services/comboAgentMiddleware.ts @@ -176,17 +176,9 @@ export function applyComboAgentMiddleware( let messages: Message[] = Array.isArray(body.messages) ? [...body.messages] : []; let pinnedModel: string | null = null; - // 1. Context caching: check for pinned model in history - if (comboConfig.context_cache_protection) { - pinnedModel = extractPinnedModel(messages); - if (pinnedModel) { - // (#535) Model is pinned via tag — override body.model so the combo - // router uses exactly this model instead of picking a different one. Without this, - // the extracted pinnedModel is returned but body.model is unchanged, breaking - // context cache sessions by sending subsequent turns to a different model. - body = { ...body, model: pinnedModel }; - } - } + // Context cache pinning is handled server-side in combo.ts via + // session_model_history. No client-side tag extraction needed. + pinnedModel = null; // 2. System message override if (comboConfig.system_message && comboConfig.system_message.trim()) { diff --git a/open-sse/utils/stream.ts b/open-sse/utils/stream.ts index 59dfa5eb62..a4e60c2650 100644 --- a/open-sse/utils/stream.ts +++ b/open-sse/utils/stream.ts @@ -421,8 +421,7 @@ type ClaudeEmptyResponseLifecycle = { warningLogged: boolean; }; -const SYNTHETIC_CLAUDE_EMPTY_RESPONSE_TEXT = - "[Proxy Error] The upstream API returned an empty response. Please retry the request."; +const SYNTHETIC_CLAUDE_EMPTY_RESPONSE_TEXT = ""; function createClaudeEmptyResponseLifecycle(): ClaudeEmptyResponseLifecycle { return { diff --git a/src/lib/db/combos.ts b/src/lib/db/combos.ts index 2695756e0d..51cb0892b2 100644 --- a/src/lib/db/combos.ts +++ b/src/lib/db/combos.ts @@ -68,7 +68,19 @@ function normalizeStoredCombo( function parseComboRow(row: unknown): JsonRecord | null { const payload = getSerializedData(row); if (!payload) return null; - return withSortOrder(payload, getSortOrder(row)); + const parsed = withSortOrder(payload, getSortOrder(row)); + // Merge deduplicated column values back into the record + const record = asRecord(row); + if (record.context_cache_protection !== undefined && record.context_cache_protection !== null) { + // Column is authoritative when explicitly enabled (1). + // When column is 0 (unset default) preserve the JSON blob value + // to avoid silently disabling the feature on pre-migration rows. + if (record.context_cache_protection === 1) { + parsed.context_cache_protection = true; + } + // Column is 0 — keep existing JSON blob value + } + return parsed; } function getNextSortOrder() { @@ -81,7 +93,7 @@ function getNextSortOrder() { export async function getCombos() { const db = getDbInstance(); const rawCombos = db - .prepare("SELECT data, sort_order FROM combos ORDER BY sort_order ASC, name COLLATE NOCASE ASC") + .prepare("SELECT data, sort_order, context_cache_protection FROM combos ORDER BY sort_order ASC, name COLLATE NOCASE ASC") .all() .map((row) => parseComboRow(row)) .filter((row): row is JsonRecord => row !== null); @@ -99,7 +111,7 @@ export async function getCombos() { export async function getComboById(id: string) { const db = getDbInstance(); - const row = db.prepare("SELECT data, sort_order FROM combos WHERE id = ?").get(id); + const row = db.prepare("SELECT data, sort_order, context_cache_protection FROM combos WHERE id = ?").get(id); const combo = parseComboRow(row); if (!combo) return null; return normalizeStoredCombo(combo, db, typeof combo.name === "string" ? [combo.name] : []); @@ -107,7 +119,7 @@ export async function getComboById(id: string) { export async function getComboByName(name: string) { const db = getDbInstance(); - const row = db.prepare("SELECT data, sort_order FROM combos WHERE name = ?").get(name); + const row = db.prepare("SELECT data, sort_order, context_cache_protection FROM combos WHERE name = ?").get(name); const combo = parseComboRow(row); if (!combo) return null; return normalizeStoredCombo(combo, db, [name]); @@ -135,9 +147,10 @@ export async function createCombo(data: JsonRecord) { typeof data.name === "string" ? [data.name] : [] ); + const contextCache = data.context_cache_protection ? 1 : 0; db.prepare( - "INSERT INTO combos (id, name, data, sort_order, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)" - ).run(combo.id, combo.name, JSON.stringify(combo), sortOrder, now, now); + "INSERT INTO combos (id, name, data, sort_order, created_at, updated_at, context_cache_protection) VALUES (?, ?, ?, ?, ?, ?, ?)" + ).run(combo.id, combo.name, JSON.stringify(combo), sortOrder, now, now, contextCache); invalidateDbCache("combos"); backupDbFile("pre-write"); @@ -146,7 +159,7 @@ export async function createCombo(data: JsonRecord) { export async function updateCombo(id: string, data: JsonRecord) { const db = getDbInstance(); - const existing = db.prepare("SELECT data, sort_order FROM combos WHERE id = ?").get(id); + const existing = db.prepare("SELECT data, sort_order, context_cache_protection FROM combos WHERE id = ?").get(id); if (!existing) return null; const current = parseComboRow(existing); @@ -175,10 +188,11 @@ export async function updateCombo(id: string, data: JsonRecord) { ? merged["name"] : currentName; const normalizedMerged = normalizeStoredCombo({ ...merged, name: nextName }, db, [nextName]); + const contextCacheProtection = normalizedMerged.context_cache_protection ? 1 : 0; db.prepare( - "UPDATE combos SET name = ?, data = ?, sort_order = ?, updated_at = ? WHERE id = ?" - ).run(nextName, JSON.stringify(normalizedMerged), sortOrder, normalizedMerged.updatedAt, id); + "UPDATE combos SET name = ?, data = ?, sort_order = ?, updated_at = ?, context_cache_protection = ? WHERE id = ?" + ).run(nextName, JSON.stringify(normalizedMerged), sortOrder, normalizedMerged.updatedAt, contextCacheProtection, id); invalidateDbCache("combos"); backupDbFile("pre-write"); diff --git a/src/lib/db/migrations/096_sync_context_cache_protection.sql b/src/lib/db/migrations/096_sync_context_cache_protection.sql new file mode 100644 index 0000000000..542adcf2ae --- /dev/null +++ b/src/lib/db/migrations/096_sync_context_cache_protection.sql @@ -0,0 +1,9 @@ +-- 096_sync_context_cache_protection.sql +-- Sync the context_cache_protection column with the JSON blob for existing combos. +-- Before this migration, the column was never written by the API, so existing rows +-- have 0 (from ADD COLUMN DEFAULT 0) even when the JSON blob has it set to true. +-- This migration brings them in sync so the column becomes the authoritative source. +UPDATE combos +SET context_cache_protection = 1 +WHERE json_extract(data, '$.context_cache_protection') = 1 + AND (context_cache_protection IS NULL OR context_cache_protection = 0); diff --git a/tests/unit/combo-context-relay.test.ts b/tests/unit/combo-context-relay.test.ts index f25d69bb8d..66ef46fef5 100644 --- a/tests/unit/combo-context-relay.test.ts +++ b/tests/unit/combo-context-relay.test.ts @@ -531,3 +531,83 @@ test("handleComboChat universal handoff detects model switch before recording cu assert.ok(saved); assert.equal(saved.lastModel, "openai/previous"); }); + +// ── Rule #18 gate — PR #3399: server-side context cache pinning ───────────── +// Proves that when context_cache_protection=true and session_model_history has +// a prior model, handleComboChat overrides body.model with the pinned model +// (no client-side tag injection required). + +test("context_cache_protection: pins body.model to last session model when history exists", async () => { + const sessionId = "sess-cache-pin-active"; + const comboName = "cache-pin-combo"; + + // Pre-record a prior model usage for this session/combo + handoffDb.recordSessionModelUsage(sessionId, comboName, "anthropic/claude-3-5-sonnet", "anthropic"); + + const capturedModels: string[] = []; + + const result = await handleComboChat({ + body: { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "Continue the task" }], + }, + combo: { + name: comboName, + strategy: "priority", + models: ["openai/gpt-4o", "anthropic/claude-3-5-sonnet"], + config: { maxRetries: 0 }, + context_cache_protection: true, + }, + handleSingleModel: async (body, modelStr) => { + capturedModels.push(modelStr); + capturedModels.push(body?.model as string); + return okResponse(); + }, + isModelAvailable: async () => true, + log: createLog(), + settings: null, + allCombos: null, + relayOptions: { sessionId }, + }); + + assert.equal(result.ok, true); + // The first model tried must be the pinned one, not the combo's first model + assert.equal(capturedModels[0], "anthropic/claude-3-5-sonnet", "modelStr must be pinned model"); + // body.model must also reflect the pinned model + assert.equal(capturedModels[1], "anthropic/claude-3-5-sonnet", "body.model must be pinned model"); +}); + +test("context_cache_protection: does NOT pin when no session history exists (first request)", async () => { + const sessionId = "sess-cache-pin-first"; + const comboName = "cache-pin-first-combo"; + // No prior recordSessionModelUsage call — fresh session + + const capturedModels: string[] = []; + + const result = await handleComboChat({ + body: { + model: "openai/gpt-4o", + messages: [{ role: "user", content: "First message" }], + }, + combo: { + name: comboName, + strategy: "priority", + models: ["openai/gpt-4o"], + config: { maxRetries: 0 }, + context_cache_protection: true, + }, + handleSingleModel: async (body, modelStr) => { + capturedModels.push(modelStr); + return okResponse(); + }, + isModelAvailable: async () => true, + log: createLog(), + settings: null, + allCombos: null, + relayOptions: { sessionId }, + }); + + assert.equal(result.ok, true); + // No pinning on first request — should use the combo's first model + assert.equal(capturedModels[0], "openai/gpt-4o", "first request must use combo model (no pinning)"); +});