feat: implement MCP output compression mechanism

- Add compressionService.ts with AI-powered compression using OpenAI
- Add compression configuration to SystemConfig type
- Update SystemConfig database entity with compression field
- Integrate compression with mcpService tool call responses
- Add migration support for compression settings
- Add comprehensive unit tests (23 new tests)
- Support graceful fallback when compression is unavailable

Co-authored-by: samanhappy <2755122+samanhappy@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
2025-12-31 15:11:59 +00:00
parent 1b752827a5
commit d25a85a1bf
7 changed files with 725 additions and 4 deletions

View File

@@ -33,6 +33,9 @@ export class SystemConfig {
@Column({ type: 'boolean', nullable: true }) @Column({ type: 'boolean', nullable: true })
enableSessionRebuild?: boolean; enableSessionRebuild?: boolean;
@Column({ type: 'simple-json', nullable: true })
compression?: Record<string, any>;
@CreateDateColumn({ name: 'created_at', type: 'timestamp' }) @CreateDateColumn({ name: 'created_at', type: 'timestamp' })
createdAt: Date; createdAt: Date;

View File

@@ -32,6 +32,7 @@ export class SystemConfigRepository {
oauth: {}, oauth: {},
oauthServer: {}, oauthServer: {},
enableSessionRebuild: false, enableSessionRebuild: false,
compression: {},
}); });
config = await this.repository.save(config); config = await this.repository.save(config);
} }

View File

@@ -0,0 +1,266 @@
import OpenAI from 'openai';
import { getSmartRoutingConfig, SmartRoutingConfig } from '../utils/smartRouting.js';
import { getSystemConfigDao } from '../dao/index.js';
/**
* Compression configuration interface
*/
export interface CompressionConfig {
enabled: boolean;
model?: string;
maxInputTokens?: number;
targetReductionRatio?: number;
}
/**
* Default compression configuration
*/
const DEFAULT_COMPRESSION_CONFIG: CompressionConfig = {
enabled: false,
model: 'gpt-4o-mini',
maxInputTokens: 100000,
targetReductionRatio: 0.5,
};
/**
* Get compression configuration from system settings
*/
export async function getCompressionConfig(): Promise<CompressionConfig> {
try {
const systemConfigDao = getSystemConfigDao();
const systemConfig = await systemConfigDao.get();
const compressionSettings = systemConfig?.compression || {};
return {
enabled: compressionSettings.enabled ?? DEFAULT_COMPRESSION_CONFIG.enabled,
model: compressionSettings.model ?? DEFAULT_COMPRESSION_CONFIG.model,
maxInputTokens: compressionSettings.maxInputTokens ?? DEFAULT_COMPRESSION_CONFIG.maxInputTokens,
targetReductionRatio:
compressionSettings.targetReductionRatio ?? DEFAULT_COMPRESSION_CONFIG.targetReductionRatio,
};
} catch (error) {
console.warn('Failed to get compression config, using defaults:', error);
return DEFAULT_COMPRESSION_CONFIG;
}
}
/**
* Check if compression is available and enabled
*/
export async function isCompressionEnabled(): Promise<boolean> {
const config = await getCompressionConfig();
if (!config.enabled) {
return false;
}
// Check if we have OpenAI API key configured (via smart routing config)
const smartRoutingConfig = await getSmartRoutingConfig();
return !!smartRoutingConfig.openaiApiKey;
}
/**
* Get OpenAI client for compression
*/
async function getOpenAIClient(smartRoutingConfig: SmartRoutingConfig): Promise<OpenAI | null> {
if (!smartRoutingConfig.openaiApiKey) {
return null;
}
return new OpenAI({
apiKey: smartRoutingConfig.openaiApiKey,
baseURL: smartRoutingConfig.openaiApiBaseUrl || 'https://api.openai.com/v1',
});
}
/**
* Estimate token count for a string (rough approximation)
* Uses ~4 characters per token as a rough estimate
*/
export function estimateTokenCount(text: string): number {
return Math.ceil(text.length / 4);
}
/**
* Check if content should be compressed based on token count
*/
export function shouldCompress(content: string, maxInputTokens: number): boolean {
const estimatedTokens = estimateTokenCount(content);
// Only compress if content is larger than a reasonable threshold
const compressionThreshold = Math.min(maxInputTokens * 0.1, 1000);
return estimatedTokens > compressionThreshold;
}
/**
* Compress MCP tool output using AI
*
* @param content The MCP tool output content to compress
* @param context Optional context about the tool that generated this output
* @returns Compressed content or original content if compression fails/is disabled
*/
export async function compressOutput(
content: string,
context?: {
toolName?: string;
serverName?: string;
},
): Promise<{ compressed: string; originalLength: number; compressedLength: number; wasCompressed: boolean }> {
const originalLength = content.length;
// Check if compression is enabled
const compressionConfig = await getCompressionConfig();
if (!compressionConfig.enabled) {
return {
compressed: content,
originalLength,
compressedLength: originalLength,
wasCompressed: false,
};
}
// Check if content should be compressed
if (!shouldCompress(content, compressionConfig.maxInputTokens || 100000)) {
return {
compressed: content,
originalLength,
compressedLength: originalLength,
wasCompressed: false,
};
}
try {
const smartRoutingConfig = await getSmartRoutingConfig();
const openai = await getOpenAIClient(smartRoutingConfig);
if (!openai) {
console.warn('Compression enabled but OpenAI API key not configured');
return {
compressed: content,
originalLength,
compressedLength: originalLength,
wasCompressed: false,
};
}
const targetRatio = compressionConfig.targetReductionRatio || 0.5;
const toolContext = context?.toolName ? `from tool "${context.toolName}"` : '';
const serverContext = context?.serverName ? `on server "${context.serverName}"` : '';
const systemPrompt = `You are a data compression assistant. Your task is to compress MCP (Model Context Protocol) tool outputs while preserving all essential information.
Guidelines:
- Remove redundant information, formatting, and verbose descriptions
- Preserve all data values, identifiers, and critical information
- Keep error messages and status information intact
- Maintain structured data (JSON, arrays) in a compact but readable format
- Target approximately ${Math.round(targetRatio * 100)}% reduction in size
- If the content cannot be meaningfully compressed, return it as-is
The output is ${toolContext} ${serverContext}.`;
const userPrompt = `Compress the following MCP tool output while preserving all essential information:
${content}`;
const response = await openai.chat.completions.create({
model: compressionConfig.model || 'gpt-4o-mini',
messages: [
{ role: 'system', content: systemPrompt },
{ role: 'user', content: userPrompt },
],
temperature: 0.1,
max_tokens: Math.ceil(estimateTokenCount(content) * targetRatio * 1.5),
});
const compressedContent = response.choices[0]?.message?.content;
if (!compressedContent) {
console.warn('Compression returned empty result, using original content');
return {
compressed: content,
originalLength,
compressedLength: originalLength,
wasCompressed: false,
};
}
const compressedLength = compressedContent.length;
// Only use compressed version if it's actually smaller
if (compressedLength >= originalLength) {
console.log('Compression did not reduce size, using original content');
return {
compressed: content,
originalLength,
compressedLength: originalLength,
wasCompressed: false,
};
}
const reductionPercent = ((originalLength - compressedLength) / originalLength * 100).toFixed(1);
console.log(`Compressed output: ${originalLength} -> ${compressedLength} chars (${reductionPercent}% reduction)`);
return {
compressed: compressedContent,
originalLength,
compressedLength,
wasCompressed: true,
};
} catch (error) {
console.error('Compression failed, using original content:', error);
return {
compressed: content,
originalLength,
compressedLength: originalLength,
wasCompressed: false,
};
}
}
/**
* Compress tool call result content
* This handles the MCP tool result format with content array
*/
export async function compressToolResult(
result: any,
context?: {
toolName?: string;
serverName?: string;
},
): Promise<any> {
// Check if compression is enabled first
const compressionEnabled = await isCompressionEnabled();
if (!compressionEnabled) {
return result;
}
// Handle error results - don't compress error messages
if (result?.isError) {
return result;
}
// Handle content array format
if (!result?.content || !Array.isArray(result.content)) {
return result;
}
const compressedContent = await Promise.all(
result.content.map(async (item: any) => {
// Only compress text content
if (item?.type !== 'text' || !item?.text) {
return item;
}
const compressionResult = await compressOutput(item.text, context);
return {
...item,
text: compressionResult.compressed,
};
}),
);
return {
...result,
content: compressedContent,
};
}

View File

@@ -27,6 +27,7 @@ import { getDataService } from './services.js';
import { getServerDao, getSystemConfigDao, ServerConfigWithName } from '../dao/index.js'; import { getServerDao, getSystemConfigDao, ServerConfigWithName } from '../dao/index.js';
import { initializeAllOAuthClients } from './oauthService.js'; import { initializeAllOAuthClients } from './oauthService.js';
import { createOAuthProvider } from './mcpOAuthProvider.js'; import { createOAuthProvider } from './mcpOAuthProvider.js';
import { compressToolResult } from './compressionService.js';
const servers: { [sessionId: string]: Server } = {}; const servers: { [sessionId: string]: Server } = {};
@@ -1260,7 +1261,7 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
const result = await openApiClient.callTool(cleanToolName, finalArgs, passthroughHeaders); const result = await openApiClient.callTool(cleanToolName, finalArgs, passthroughHeaders);
console.log(`OpenAPI tool invocation result: ${JSON.stringify(result)}`); console.log(`OpenAPI tool invocation result: ${JSON.stringify(result)}`);
return { const openApiResult = {
content: [ content: [
{ {
type: 'text', type: 'text',
@@ -1268,6 +1269,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
}, },
], ],
}; };
return compressToolResult(openApiResult, {
toolName: cleanToolName,
serverName: targetServerInfo.name,
});
} }
// Call the tool on the target server (MCP servers) // Call the tool on the target server (MCP servers)
@@ -1297,7 +1302,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
); );
console.log(`Tool invocation result: ${JSON.stringify(result)}`); console.log(`Tool invocation result: ${JSON.stringify(result)}`);
return result; return compressToolResult(result, {
toolName,
serverName: targetServerInfo.name,
});
} }
// Regular tool handling // Regular tool handling
@@ -1356,7 +1364,7 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
); );
console.log(`OpenAPI tool invocation result: ${JSON.stringify(result)}`); console.log(`OpenAPI tool invocation result: ${JSON.stringify(result)}`);
return { const openApiResult = {
content: [ content: [
{ {
type: 'text', type: 'text',
@@ -1364,6 +1372,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
}, },
], ],
}; };
return compressToolResult(openApiResult, {
toolName: cleanToolName,
serverName: serverInfo.name,
});
} }
// Handle MCP servers // Handle MCP servers
@@ -1374,6 +1386,7 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
const separator = getNameSeparator(); const separator = getNameSeparator();
const prefix = `${serverInfo.name}${separator}`; const prefix = `${serverInfo.name}${separator}`;
const originalToolName = request.params.name;
request.params.name = request.params.name.startsWith(prefix) request.params.name = request.params.name.startsWith(prefix)
? request.params.name.substring(prefix.length) ? request.params.name.substring(prefix.length)
: request.params.name; : request.params.name;
@@ -1383,7 +1396,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
serverInfo.options || {}, serverInfo.options || {},
); );
console.log(`Tool call result: ${JSON.stringify(result)}`); console.log(`Tool call result: ${JSON.stringify(result)}`);
return result; return compressToolResult(result, {
toolName: originalToolName,
serverName: serverInfo.name,
});
} catch (error) { } catch (error) {
console.error(`Error handling CallToolRequest: ${error}`); console.error(`Error handling CallToolRequest: ${error}`);
return { return {

View File

@@ -173,6 +173,12 @@ export interface SystemConfig {
oauth?: OAuthProviderConfig; // OAuth provider configuration for upstream MCP servers oauth?: OAuthProviderConfig; // OAuth provider configuration for upstream MCP servers
oauthServer?: OAuthServerConfig; // OAuth authorization server configuration for MCPHub itself oauthServer?: OAuthServerConfig; // OAuth authorization server configuration for MCPHub itself
enableSessionRebuild?: boolean; // Controls whether server session rebuild is enabled enableSessionRebuild?: boolean; // Controls whether server session rebuild is enabled
compression?: {
enabled?: boolean; // Enable/disable AI compression of MCP tool outputs
model?: string; // AI model to use for compression (default: 'gpt-4o-mini')
maxInputTokens?: number; // Maximum input tokens for compression (default: 100000)
targetReductionRatio?: number; // Target reduction ratio, 0.0-1.0 (default: 0.5)
};
} }
export interface UserConfig { export interface UserConfig {

View File

@@ -117,6 +117,7 @@ export async function migrateToDatabase(): Promise<boolean> {
oauth: settings.systemConfig.oauth || {}, oauth: settings.systemConfig.oauth || {},
oauthServer: settings.systemConfig.oauthServer || {}, oauthServer: settings.systemConfig.oauthServer || {},
enableSessionRebuild: settings.systemConfig.enableSessionRebuild, enableSessionRebuild: settings.systemConfig.enableSessionRebuild,
compression: settings.systemConfig.compression || {},
}; };
await systemConfigRepo.update(systemConfig); await systemConfigRepo.update(systemConfig);
console.log(' - System configuration updated'); console.log(' - System configuration updated');

View File

@@ -0,0 +1,428 @@
// Mock the DAO module before imports
jest.mock('../../src/dao/index.js', () => ({
getSystemConfigDao: jest.fn(),
}));
// Mock smart routing config
jest.mock('../../src/utils/smartRouting.js', () => ({
getSmartRoutingConfig: jest.fn(),
}));
// Mock OpenAI
jest.mock('openai', () => {
return {
__esModule: true,
default: jest.fn().mockImplementation(() => ({
chat: {
completions: {
create: jest.fn(),
},
},
})),
};
});
import {
getCompressionConfig,
isCompressionEnabled,
estimateTokenCount,
shouldCompress,
compressOutput,
compressToolResult,
} from '../../src/services/compressionService.js';
import { getSystemConfigDao } from '../../src/dao/index.js';
import { getSmartRoutingConfig } from '../../src/utils/smartRouting.js';
import OpenAI from 'openai';
describe('CompressionService', () => {
const mockSystemConfigDao = {
get: jest.fn(),
getSection: jest.fn(),
update: jest.fn(),
updateSection: jest.fn(),
};
beforeEach(() => {
jest.clearAllMocks();
(getSystemConfigDao as jest.Mock).mockReturnValue(mockSystemConfigDao);
});
describe('getCompressionConfig', () => {
it('should return default config when no config is set', async () => {
mockSystemConfigDao.get.mockResolvedValue({});
const config = await getCompressionConfig();
expect(config).toEqual({
enabled: false,
model: 'gpt-4o-mini',
maxInputTokens: 100000,
targetReductionRatio: 0.5,
});
});
it('should return configured values when set', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: {
enabled: true,
model: 'gpt-4o',
maxInputTokens: 50000,
targetReductionRatio: 0.3,
},
});
const config = await getCompressionConfig();
expect(config).toEqual({
enabled: true,
model: 'gpt-4o',
maxInputTokens: 50000,
targetReductionRatio: 0.3,
});
});
it('should use defaults for missing values', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: {
enabled: true,
},
});
const config = await getCompressionConfig();
expect(config).toEqual({
enabled: true,
model: 'gpt-4o-mini',
maxInputTokens: 100000,
targetReductionRatio: 0.5,
});
});
it('should return defaults on error', async () => {
mockSystemConfigDao.get.mockRejectedValue(new Error('Test error'));
const config = await getCompressionConfig();
expect(config).toEqual({
enabled: false,
model: 'gpt-4o-mini',
maxInputTokens: 100000,
targetReductionRatio: 0.5,
});
});
});
describe('isCompressionEnabled', () => {
it('should return false when compression is disabled', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: false },
});
const enabled = await isCompressionEnabled();
expect(enabled).toBe(false);
});
it('should return false when enabled but no API key', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: '',
});
const enabled = await isCompressionEnabled();
expect(enabled).toBe(false);
});
it('should return true when enabled and API key is set', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
});
const enabled = await isCompressionEnabled();
expect(enabled).toBe(true);
});
});
describe('estimateTokenCount', () => {
it('should estimate tokens for short text', () => {
const text = 'Hello world';
const tokens = estimateTokenCount(text);
// 11 characters / 4 = 2.75, ceil = 3
expect(tokens).toBe(3);
});
it('should estimate tokens for longer text', () => {
const text = 'This is a longer piece of text that should have more tokens';
const tokens = estimateTokenCount(text);
// 59 characters / 4 = 14.75, ceil = 15
expect(tokens).toBe(15);
});
it('should handle empty string', () => {
const tokens = estimateTokenCount('');
expect(tokens).toBe(0);
});
});
describe('shouldCompress', () => {
it('should return false for small content', () => {
const content = 'Small content';
const result = shouldCompress(content, 100000);
expect(result).toBe(false);
});
it('should return true for large content', () => {
// Create content larger than the threshold (1000 tokens = ~4000 chars)
const content = 'x'.repeat(5000);
const result = shouldCompress(content, 100000);
expect(result).toBe(true);
});
it('should use 10% of maxInputTokens as threshold', () => {
// With maxInputTokens = 1000, threshold = 100 tokens = ~400 chars
const smallContent = 'x'.repeat(300);
const largeContent = 'x'.repeat(500);
expect(shouldCompress(smallContent, 1000)).toBe(false);
expect(shouldCompress(largeContent, 1000)).toBe(true);
});
});
describe('compressOutput', () => {
it('should return original content when compression is disabled', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: false },
});
const content = 'Test content';
const result = await compressOutput(content);
expect(result).toEqual({
compressed: content,
originalLength: content.length,
compressedLength: content.length,
wasCompressed: false,
});
});
it('should return original content when content is too small', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true, maxInputTokens: 100000 },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
});
const content = 'Small content';
const result = await compressOutput(content);
expect(result.wasCompressed).toBe(false);
expect(result.compressed).toBe(content);
});
it('should return original content when no API key is configured', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: '',
});
const content = 'x'.repeat(5000);
const result = await compressOutput(content);
expect(result.wasCompressed).toBe(false);
expect(result.compressed).toBe(content);
});
it('should compress content when enabled and content is large', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true, model: 'gpt-4o-mini', maxInputTokens: 100000 },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
openaiApiBaseUrl: 'https://api.openai.com/v1',
});
const originalContent = 'x'.repeat(5000);
const compressedContent = 'y'.repeat(2000);
// Mock OpenAI response
const mockCreate = jest.fn().mockResolvedValue({
choices: [{ message: { content: compressedContent } }],
});
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
chat: {
completions: {
create: mockCreate,
},
},
}));
const result = await compressOutput(originalContent, {
toolName: 'test-tool',
serverName: 'test-server',
});
expect(result.wasCompressed).toBe(true);
expect(result.compressed).toBe(compressedContent);
expect(result.originalLength).toBe(originalContent.length);
expect(result.compressedLength).toBe(compressedContent.length);
});
it('should return original content when compressed is larger', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true, model: 'gpt-4o-mini', maxInputTokens: 100000 },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
openaiApiBaseUrl: 'https://api.openai.com/v1',
});
const originalContent = 'x'.repeat(5000);
const largerContent = 'y'.repeat(6000);
const mockCreate = jest.fn().mockResolvedValue({
choices: [{ message: { content: largerContent } }],
});
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
chat: {
completions: {
create: mockCreate,
},
},
}));
const result = await compressOutput(originalContent);
expect(result.wasCompressed).toBe(false);
expect(result.compressed).toBe(originalContent);
});
it('should return original content on API error', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true, model: 'gpt-4o-mini', maxInputTokens: 100000 },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
openaiApiBaseUrl: 'https://api.openai.com/v1',
});
const mockCreate = jest.fn().mockRejectedValue(new Error('API error'));
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
chat: {
completions: {
create: mockCreate,
},
},
}));
const content = 'x'.repeat(5000);
const result = await compressOutput(content);
expect(result.wasCompressed).toBe(false);
expect(result.compressed).toBe(content);
});
});
describe('compressToolResult', () => {
it('should return original result when compression is disabled', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: false },
});
const result = {
content: [{ type: 'text', text: 'Test output' }],
};
const compressed = await compressToolResult(result);
expect(compressed).toEqual(result);
});
it('should not compress error results', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
});
const result = {
content: [{ type: 'text', text: 'Error message' }],
isError: true,
};
const compressed = await compressToolResult(result);
expect(compressed).toEqual(result);
});
it('should handle results without content array', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
});
const result = { someOtherField: 'value' };
const compressed = await compressToolResult(result);
expect(compressed).toEqual(result);
});
it('should only compress text content items', async () => {
mockSystemConfigDao.get.mockResolvedValue({
compression: { enabled: true, maxInputTokens: 100000 },
});
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
openaiApiKey: 'test-api-key',
openaiApiBaseUrl: 'https://api.openai.com/v1',
});
const largeText = 'x'.repeat(5000);
const compressedText = 'y'.repeat(2000);
const mockCreate = jest.fn().mockResolvedValue({
choices: [{ message: { content: compressedText } }],
});
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
chat: {
completions: {
create: mockCreate,
},
},
}));
const result = {
content: [
{ type: 'text', text: largeText },
{ type: 'image', data: 'base64data' },
],
};
const compressed = await compressToolResult(result);
expect(compressed.content[0].text).toBe(compressedText);
expect(compressed.content[1]).toEqual({ type: 'image', data: 'base64data' });
});
});
});