Files
2026-07-13 13:39:12 +08:00

452 lines
17 KiB
TypeScript

/**
* OmniRoute MCP Compression Tools — Manage and monitor prompt compression.
*
* Tools:
* 1. omniroute_compression_status — Get compression config, analytics, and cache stats
* 2. omniroute_compression_configure — Update compression settings
*/
import { logToolCall } from "../audit.ts";
import {
getCompressionSettings,
updateCompressionSettings,
} from "../../../src/lib/db/compression.ts";
import { getCompressionAnalyticsSummary } from "../../../src/lib/db/compressionAnalytics.ts";
import { getCacheStatsSummary } from "../../../src/lib/db/compressionCacheStats.ts";
import { listCompressionCombos } from "../../../src/lib/db/compressionCombos.ts";
import type { McpToolExtraLike } from "../scopeEnforcement.ts";
import {
getMcpDescriptionCompressionStats,
snapshotMcpDescriptionCompressionStats,
} from "../descriptionCompressor.ts";
/**
* Handle compression_status tool: return current compression config, analytics, and cache stats
*/
export async function handleCompressionStatus(
args: Record<string, never>,
extra?: McpToolExtraLike
): Promise<{
enabled: boolean;
strategy: string;
settings: {
maxTokens: number;
autoTriggerMode: string;
targetRatio: number;
preserveSystemPrompt: boolean;
mcpDescriptionCompressionEnabled: boolean;
};
analytics: {
totalRequests: number;
compressedRequests: number;
tokensSaved: number;
avgCompressionRatio: number;
byMode: Record<string, { count: number; tokensSaved: number; avgSavingsPct: number }>;
byEngine: Record<string, { count: number; tokensSaved: number; avgSavingsPct: number }>;
byCompressionCombo: Record<string, { count: number; tokensSaved: number }>;
validationFallbacks: number;
requestsWithReceipts: number;
realUsage: {
requestsWithReceipts: number;
promptTokens: number;
completionTokens: number;
totalTokens: number;
cacheReadTokens: number;
cacheWriteTokens: number;
estimatedUsdSaved: number;
bySource: Record<string, number>;
};
mcpDescriptionCompression: {
descriptionsCompressed: number;
charsBefore: number;
charsAfter: number;
charsSaved: number;
estimatedTokensSaved: number;
persistedEstimatedTokensSaved: number;
persistedSnapshots: number;
source: "mcp_metadata_estimate";
notProviderUsage: true;
};
};
cacheStats: {
hits: number;
misses: number;
hitRate: string;
tokensSaved: number;
} | null;
}> {
const start = Date.now();
try {
const settings = await getCompressionSettings();
await snapshotMcpDescriptionCompressionStats();
const analyticsSummary = getCompressionAnalyticsSummary();
const mcpDescriptionStats = getMcpDescriptionCompressionStats();
const cacheStats = getCacheStatsSummary();
const result = {
enabled: settings.enabled,
strategy: settings.defaultMode || "standard",
settings: {
maxTokens: settings.autoTriggerTokens,
autoTriggerMode: settings.autoTriggerMode ?? "lite",
targetRatio: 0.7, // Default target ratio
preserveSystemPrompt: settings.preserveSystemPrompt,
mcpDescriptionCompressionEnabled: settings.mcpDescriptionCompressionEnabled !== false,
},
analytics: {
totalRequests: analyticsSummary.totalRequests,
compressedRequests: Object.values(analyticsSummary.byMode ?? {}).reduce(
(sum, mode) => sum + mode.count,
0
),
tokensSaved: analyticsSummary.totalTokensSaved,
avgCompressionRatio: analyticsSummary.avgSavingsPct,
byMode: analyticsSummary.byMode ?? {},
byEngine: analyticsSummary.byEngine ?? {},
byCompressionCombo: analyticsSummary.byCompressionCombo ?? {},
validationFallbacks: analyticsSummary.validationFallbacks,
requestsWithReceipts: analyticsSummary.realUsage.requestsWithReceipts,
realUsage: analyticsSummary.realUsage,
mcpDescriptionCompression: {
descriptionsCompressed: mcpDescriptionStats.descriptionsCompressed,
charsBefore: mcpDescriptionStats.charsBefore,
charsAfter: mcpDescriptionStats.charsAfter,
charsSaved: mcpDescriptionStats.charsSaved,
estimatedTokensSaved: mcpDescriptionStats.estimatedTokensSaved,
persistedEstimatedTokensSaved:
analyticsSummary.mcpDescriptionCompression.estimatedTokensSaved,
persistedSnapshots: analyticsSummary.mcpDescriptionCompression.snapshots,
source: "mcp_metadata_estimate" as const,
notProviderUsage: true as const,
},
},
cacheStats: cacheStats
? {
hits: Math.round(cacheStats.cacheHitRate * (cacheStats.totalRequests || 1)),
misses: Math.round((1 - cacheStats.cacheHitRate) * (cacheStats.totalRequests || 1)),
hitRate: `${(cacheStats.cacheHitRate * 100).toFixed(2)}%`,
tokensSaved: Math.round(cacheStats.avgNetSavings),
}
: null,
};
const duration = Date.now() - start;
await logToolCall("omniroute_compression_status", args, result, duration, true);
return result;
} catch (error) {
const duration = Date.now() - start;
const errorMessage = error instanceof Error ? error.message : String(error);
await logToolCall(
"omniroute_compression_status",
args,
{ error: errorMessage },
duration,
false,
"ERROR"
);
throw error;
}
}
/**
* Handle compression_configure tool: update compression settings
*/
export async function handleCompressionConfigure(
args: {
enabled?: boolean;
strategy?: string;
autoTriggerMode?: string;
maxTokens?: number;
targetRatio?: number;
preserveSystemPrompt?: boolean;
mcpDescriptionCompressionEnabled?: boolean;
},
extra?: McpToolExtraLike
): Promise<{
success: boolean;
updated: Record<string, unknown>;
settings: {
enabled: boolean;
strategy: string;
autoTriggerMode: string;
maxTokens: number;
targetRatio: number;
preserveSystemPrompt: boolean;
mcpDescriptionCompressionEnabled: boolean;
};
}> {
const start = Date.now();
try {
const updates: Record<string, unknown> = {};
if (args.enabled !== undefined) {
updates.enabled = args.enabled;
}
if (args.strategy !== undefined) {
updates.defaultMode = args.strategy;
}
if (args.autoTriggerMode !== undefined) {
updates.autoTriggerMode = args.autoTriggerMode;
}
if (args.maxTokens !== undefined) {
updates.autoTriggerTokens = args.maxTokens;
}
if (args.preserveSystemPrompt !== undefined) {
updates.preserveSystemPrompt = args.preserveSystemPrompt;
}
if (args.mcpDescriptionCompressionEnabled !== undefined) {
updates.mcpDescriptionCompressionEnabled = args.mcpDescriptionCompressionEnabled;
}
const settings = await updateCompressionSettings(updates);
const result = {
success: true,
updated: updates,
settings: {
enabled: settings.enabled,
strategy: settings.defaultMode || "standard",
autoTriggerMode: settings.autoTriggerMode ?? "lite",
maxTokens: settings.autoTriggerTokens,
targetRatio: 0.7, // Default target ratio
preserveSystemPrompt: settings.preserveSystemPrompt,
mcpDescriptionCompressionEnabled: settings.mcpDescriptionCompressionEnabled !== false,
},
};
const duration = Date.now() - start;
await logToolCall("omniroute_compression_configure", args, result, duration, true);
return result;
} catch (error) {
const duration = Date.now() - start;
const errorMessage = error instanceof Error ? error.message : String(error);
await logToolCall(
"omniroute_compression_configure",
args,
{ error: errorMessage },
duration,
false,
"ERROR"
);
throw error;
}
}
import { z } from "zod";
import {
compressionStatusInput,
compressionConfigureInput,
setCompressionEngineInput,
listCompressionCombosInput,
compressionComboStatsInput,
} from "../schemas/tools.ts";
import { handleCcrRetrieve } from "../../services/compression/engines/ccr/index.ts";
import {
listRtkCommandSamples,
discoverRepeatedNoise,
suggestFilter,
commandToId,
} from "../../services/compression/engines/rtk/index.ts";
import { resolveCallerScopeContext } from "../scopeEnforcement.ts";
import { resolveMcpCallerApiKeyId } from "../mcpCallerIdentity.ts";
const ccrRetrieveInput = z.object({
hash: z
.string()
.min(6)
.max(64)
.describe("24-hex content hash from a [CCR retrieve hash=<hash>] marker"),
mode: z
.enum(["full", "head", "tail", "lines", "grep", "stats"])
.optional()
.describe("Retrieval mode: full (default) | head | tail | lines | grep | stats"),
n: z.number().int().positive().max(10000).optional().describe("head/tail: number of lines"),
start: z.number().int().positive().optional().describe("lines: 1-indexed inclusive start"),
end: z.number().int().positive().optional().describe("lines: 1-indexed inclusive end"),
pattern: z.string().max(512).optional().describe("grep: regex (validated safe; ReDoS-rejected)"),
unique: z.boolean().optional().describe("grep: dedupe matching lines"),
});
export async function handleSetCompressionEngine(
args: z.infer<typeof setCompressionEngineInput>
): Promise<{ success: boolean; settings: Record<string, unknown> }> {
const updates: Record<string, unknown> = { enabled: true };
if (args.engine) {
updates.defaultMode = args.engine === "caveman" ? "standard" : args.engine;
if (args.engine === "off") updates.enabled = false;
}
if (args.cavemanIntensity) {
const current = await getCompressionSettings();
updates.cavemanConfig = {
...(current.cavemanConfig ?? {}),
intensity: args.cavemanIntensity,
};
}
if (args.rtkIntensity) {
const current = await getCompressionSettings();
updates.rtkConfig = {
...(current.rtkConfig ?? {}),
intensity: args.rtkIntensity,
};
}
if (args.outputMode !== undefined) {
const current = await getCompressionSettings();
updates.cavemanOutputMode = {
...(current.cavemanOutputMode ?? {}),
enabled: args.outputMode,
};
}
const settings = await updateCompressionSettings(updates);
return { success: true, settings: settings as unknown as Record<string, unknown> };
}
export async function handleListCompressionCombos(): Promise<{
combos: ReturnType<typeof listCompressionCombos>;
}> {
return { combos: listCompressionCombos() };
}
export async function handleCompressionComboStats(
args: z.infer<typeof compressionComboStatsInput>
): Promise<Record<string, unknown>> {
const summary = getCompressionAnalyticsSummary(args.since === "all" ? undefined : args.since);
if (!args.comboId) return summary as unknown as Record<string, unknown>;
return {
comboId: args.comboId,
summary,
combo: summary.byCompressionCombo[args.comboId] ?? { count: 0, tokensSaved: 0 },
};
}
// T07 — RTK learn/discover exposed via MCP (read-only; suggestions only). Mines the opt-in
// raw-output sample store, exactly like the /api/context/rtk/{discover,learn} routes.
const rtkDiscoverInput = z.object({
limit: z.number().int().positive().max(2000).optional().describe("Max samples to scan (default 500)"),
});
const rtkLearnInput = z.object({
command: z.string().min(1).max(500).describe("The command to learn an RTK filter draft for"),
limit: z.number().int().positive().max(2000).optional().describe("Max samples to scan (default 500)"),
});
function resolveSampleLimit(limit?: number): number {
if (!Number.isFinite(limit) || !limit || limit <= 0) return 500;
return Math.min(2000, Math.floor(limit));
}
export async function handleRtkDiscover(
args: z.infer<typeof rtkDiscoverInput>
): Promise<{ sampleCount: number; candidates: ReturnType<typeof discoverRepeatedNoise> }> {
const start = Date.now();
const samples = listRtkCommandSamples({ limit: resolveSampleLimit(args.limit) });
const candidates = discoverRepeatedNoise(samples);
const result = { sampleCount: samples.length, candidates };
await logToolCall("omniroute_rtk_discover", args, result, Date.now() - start, true);
return result;
}
export async function handleRtkLearn(
args: z.infer<typeof rtkLearnInput>
): Promise<{ command: string; sampleCount: number; filter: ReturnType<typeof suggestFilter> }> {
const start = Date.now();
const command = args.command.trim();
const targetId = commandToId(command);
const matching = listRtkCommandSamples({ limit: resolveSampleLimit(args.limit) }).filter(
(sample) => commandToId(sample.command) === targetId
);
const filter = suggestFilter(command, matching);
const result = { command, sampleCount: matching.length, filter };
await logToolCall("omniroute_rtk_learn", args, result, Date.now() - start, true);
return result;
}
export const compressionTools = {
omniroute_compression_status: {
name: "omniroute_compression_status",
description:
"Returns current compression configuration, strategy, analytics summary (requests compressed, tokens saved, avg ratio), and provider-aware cache statistics.",
scopes: ["read:compression"],
inputSchema: compressionStatusInput,
handler: (args: z.infer<typeof compressionStatusInput>) => handleCompressionStatus(args),
},
omniroute_compression_configure: {
name: "omniroute_compression_configure",
description:
"Configure compression settings at runtime. Supports enabling/disabling compression, changing strategy (off/lite/standard/aggressive/ultra/rtk/stacked), adjusting maxTokens threshold, targetRatio, auto-trigger mode, system prompt preservation, and MCP description compression.",
scopes: ["write:compression"],
inputSchema: compressionConfigureInput,
handler: (args: z.infer<typeof compressionConfigureInput>) => handleCompressionConfigure(args),
},
omniroute_set_compression_engine: {
name: "omniroute_set_compression_engine",
description: "Set the active compression engine and Caveman/RTK runtime options.",
scopes: ["write:compression"],
inputSchema: setCompressionEngineInput,
handler: (args: z.infer<typeof setCompressionEngineInput>) => handleSetCompressionEngine(args),
},
omniroute_list_compression_combos: {
name: "omniroute_list_compression_combos",
description: "List compression combos and their engine pipelines.",
scopes: ["read:compression"],
inputSchema: listCompressionCombosInput,
handler: (_args: z.infer<typeof listCompressionCombosInput>) => handleListCompressionCombos(),
},
omniroute_compression_combo_stats: {
name: "omniroute_compression_combo_stats",
description: "Get compression analytics grouped by engine and compression combo.",
scopes: ["read:compression"],
inputSchema: compressionComboStatsInput,
handler: (args: z.infer<typeof compressionComboStatsInput>) =>
handleCompressionComboStats(args),
},
omniroute_ccr_retrieve: {
name: "omniroute_ccr_retrieve",
description:
"Retrieve the verbatim content block stored by the CCR compression engine. " +
"When a large block is compressed, a marker `[CCR retrieve hash=<24hex> chars=N]` " +
"is inserted. Pass the hash from the marker to this tool to get the original text back. " +
"Optional `mode` (head/tail/lines/grep/stats) retrieves a slice or summary instead of the whole block; omit for the full block. " +
"Scope: read:compression. Always available (sticky-on).",
scopes: ["read:compression"],
inputSchema: ccrRetrieveInput,
handler: async (args: z.infer<typeof ccrRetrieveInput>, extra?: McpToolExtraLike) => {
// Retrieve must use the SAME principal the CCR store used at compression time:
// `String(apiKeyInfo.id)` (chatCore → getApiKeyMetadata(rawKey)). On MCP HTTP
// transports the raw key lives in httpAuthContext (not in extra.authInfo, since
// OmniRoute auth is API-key not OAuth-clientId) — resolve it to the same key id
// so the block is found. Without this the caller resolved to "anonymous" and the
// store-key never matched (#5649). Cross-tenant IDOR stays closed: a different
// key → different id → miss; no key → undefined → anonymous bucket only.
const apiKeyPrincipal = await resolveMcpCallerApiKeyId();
if (apiKeyPrincipal) {
return handleCcrRetrieve(args, apiKeyPrincipal);
}
// Fallback (unchanged): OAuth clientId / session scope context, then anonymous.
const { callerId } = resolveCallerScopeContext(extra, ["read:compression"]);
return handleCcrRetrieve(args, callerId === "anonymous" ? undefined : callerId);
},
},
omniroute_rtk_discover: {
name: "omniroute_rtk_discover",
description:
"Mine the opt-in RTK raw-output sample store for recurring noise lines and return them " +
"as ranked candidates the operator can turn into strip/collapse filters. Read-only; " +
"suggestions only. Scope: read:compression.",
scopes: ["read:compression"],
inputSchema: rtkDiscoverInput,
handler: (args: z.infer<typeof rtkDiscoverInput>) => handleRtkDiscover(args),
},
omniroute_rtk_learn: {
name: "omniroute_rtk_learn",
description:
"Suggest an RTK filter draft for a specific command, learned from that command's captured " +
"outputs in the opt-in raw-output sample store. Read-only; returns a draft for the operator " +
"to review and save. Scope: read:compression.",
scopes: ["read:compression"],
inputSchema: rtkLearnInput,
handler: (args: z.infer<typeof rtkLearnInput>) => handleRtkLearn(args),
},
};