pjandro/OmniRoute
1
0
Fork 0
mirror of https://github.com/diegosouzapw/OmniRoute.git synced 2026-08-03 13:52:09 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix: server-side context cache pinning, stop proxy message leaks, persist context_cache_protection toggle (#3399)

Browse Source 
Integrated into release/v3.8.16
This commit is contained in:
k0valik
2026-06-08 05:55:42 +02:00
committed by GitHub
parent 452e6cc937
commit 4c420b015d
6 changed files with 157 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files

193
open-sse/services/combo.ts
View File

@@ -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<string, unknown> | null | undefined
);
// ── Server-side context cache pinning (replaces <omniModel> 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<string, unknown>)?.[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<string, unknown>, 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 = `<omniModel>${modelStr}</omniModel>`;
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<Response> => {
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<string, unknown>)?.[SKIP_UNIVERSAL_HANDOFF_FLAG]
) {
recordSessionModelUsage(
relayOptions.sessionId,
combo.name,
modelStr,
provider,
target.connectionId ?? undefined
);
}
// Universal handoff: record model usage for session
if (
universalHandoffConfig.enabled &&

14
open-sse/services/comboAgentMiddleware.ts
View File

@@ -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 <omniModel> 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 <omniModel> tag extraction needed.
pinnedModel = null;
// 2. System message override
if (comboConfig.system_message && comboConfig.system_message.trim()) {

3
open-sse/utils/stream.ts
View File

@@ -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 {

32
src/lib/db/combos.ts
View File

@@ -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");

9
src/lib/db/migrations/096_sync_context_cache_protection.sql Normal file
View File

@@ -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);

80
tests/unit/combo-context-relay.test.ts
View File

@@ -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 <omniModel> 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)");
});