import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import type { McpDependencies } from '../server/mcp-server.js';
import { contextCompressInput } from './schemas.js';
import { chat } from '../services/ollama-client.js';
import { SemanticCache } from '../services/semantic-cache.js';
import { estimateTokens } from '../utils/tokens.js';
import { auditToolCall } from '../middleware/audit.js';
import { logger } from '../utils/logger.js';
const FORMAT_PROMPTS: Record<string, string> = {
bullets: 'Compress the following text into concise bullet points. Preserve all key information.',
json: 'Compress the following text into a structured JSON object with relevant keys. Output ONLY valid JSON.',
steps: 'Compress the following text into numbered step-by-step instructions.',
summary: 'Write a concise summary of the following text, preserving key facts and decisions.',
};
export function registerContextCompress(server: McpServer, deps: McpDependencies): void {
server.registerTool('context_compress', {
description:
'Compress text to reduce token usage while preserving key information. Supports bullets, json, steps, and summary formats.',
inputSchema: contextCompressInput,
}, async (args, extra) => {
const start = Date.now();
const { text, format = 'bullets', max_tokens = 500 } = args;
try {
// Check cache
const cacheKey = SemanticCache.keyFromString(`compress:${format}:${text.slice(0, 200)}`);
const cached = deps.cache.get(cacheKey);
if (cached) {
logger.debug('context.compress cache hit');
return { content: [{ type: 'text' as const, text: cached }] };
}
const inputTokens = estimateTokens(text);
const compressed = await chat([
{
role: 'system',
content: `${FORMAT_PROMPTS[format]}\n\nKeep output under ${max_tokens} tokens. Be extremely concise.`,
},
{
role: 'user',
content: text,
},
]);
const outputTokens = estimateTokens(compressed);
const ratio = inputTokens > 0 ? Math.round((1 - outputTokens / inputTokens) * 100) : 0;
const result = JSON.stringify({
compressed,
stats: {
input_tokens: inputTokens,
output_tokens: outputTokens,
reduction_percent: ratio,
format,
},
});
deps.cache.set(cacheKey, result);
auditToolCall(deps.db, 'context_compress', extra.sessionId, `${inputTokens} tokens`, `${outputTokens} tokens (${ratio}% reduction)`, Date.now() - start);
return { content: [{ type: 'text' as const, text: result }] };
} catch (err) {
const msg = err instanceof Error ? err.message : String(err);
logger.error({ err }, 'context.compress failed');
auditToolCall(deps.db, 'context_compress', extra.sessionId, `${text.length} chars`, `error: ${msg}`, Date.now() - start);
return { content: [{ type: 'text' as const, text: `Error: ${msg}` }], isError: true };
}
});
}
