tryaii-mcp-server
Version:
TryAII MCP Server - 15+ AI models with comparison, cost tracking, and collective intelligence
740 lines • 33.6 kB
JavaScript
import { logger } from '../utils/logger.js';
import { crossDatabaseService } from '../services/CrossDatabaseService.js';
import { balanceService } from '../services/BalanceService.js';
export class McpProtocolHandler {
sessionManager;
userApiKeyService;
connectionSessions = new Map();
constructor(sessionManager, userApiKeyService) {
this.sessionManager = sessionManager;
this.userApiKeyService = userApiKeyService;
}
async handleRequest(connectionId, request) {
try {
switch (request.method) {
case 'initialize':
return await this.handleInitialize(connectionId, request);
case 'tools/list':
return await this.handleToolsList(connectionId, request);
case 'tools/call':
return await this.handleToolCall(connectionId, request);
case 'resources/list':
return await this.handleResourcesList(connectionId, request);
case 'prompts/list':
return await this.handlePromptsList(connectionId, request);
// Handle notifications (they don't require responses)
case 'notifications/initialized':
case 'notifications/cancelled':
logger.debug('MCP notification received', { connectionId, method: request.method });
return null; // No response needed for notifications
default:
logger.warn('Unknown MCP method', { connectionId, method: request.method });
throw new Error(`Unknown method: ${request.method}`);
}
}
catch (error) {
logger.error('Error handling MCP request', {
connectionId,
method: request.method,
id: request.id,
error: error instanceof Error ? error.message : error
});
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32603,
message: error instanceof Error ? error.message : 'Internal error'
}
};
}
}
async handleInitialize(connectionId, request) {
const { clientInfo } = request.params || {};
logger.info('MCP initialize request', {
connectionId,
clientInfo,
id: request.id
});
// 🔍 DEBUG: Log the complete request to see what we're receiving
logger.debug('🔍 DEBUG: Complete MCP initialize request', {
connectionId,
request: JSON.stringify(request, null, 2),
clientInfo: JSON.stringify(clientInfo, null, 2),
requestParams: JSON.stringify(request.params, null, 2)
});
try {
// 🔥 MANDATORY: API key authentication required
const providedApiKey = this.extractApiKey(request, clientInfo);
// 🔍 DEBUG: Log API key extraction attempt
logger.debug('🔍 DEBUG: API key extraction attempt', {
connectionId,
providedApiKey: providedApiKey ? `${providedApiKey.substring(0, 8)}...` : 'NULL',
apiKeyFound: !!providedApiKey,
clientInfoMeta: clientInfo?.meta,
clientInfoApiKey: clientInfo?.apiKey,
requestParamsApiKey: request.params?.apiKey,
hasClientInfo: !!clientInfo,
hasRequestParams: !!request.params
});
if (!providedApiKey) {
logger.error('MCP connection rejected - no API key provided', {
connectionId,
clientName: clientInfo?.name,
clientVersion: clientInfo?.version,
message: 'Authentication required: API key must be provided'
});
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32600,
message: 'Authentication required: API key must be provided in clientInfo.meta.apiKey, clientInfo.apiKey, or request.params.apiKey'
}
};
}
// 🔥 MANDATORY: API key authentication
const authResult = await crossDatabaseService.validateApiKey(providedApiKey);
if (!authResult || !authResult.isActive) {
logger.warn('MCP connection rejected - invalid API key', {
connectionId,
keyPrefix: providedApiKey.substring(0, 8)
});
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32600,
message: 'Invalid or inactive API key'
}
};
}
// 🔥 Create session with authenticated user data
const session = await this.sessionManager.getOrCreateSession(authResult.userId, authResult.keyId);
const connectionSession = {
connectionId,
sessionId: session.sessionId,
auth: authResult,
initialized: true
};
this.connectionSessions.set(connectionId, connectionSession);
// 🔥 Update usage stats
await crossDatabaseService.updateUsageStats(authResult.keyId, {
totalRequests: 1,
lastUsed: new Date()
});
logger.info('MCP connection initialized with authenticated user', {
connectionId,
sessionId: session.sessionId,
userId: authResult.userId,
keyId: authResult.keyId,
keyPrefix: authResult.keyPrefix,
permissions: authResult.permissions
});
return {
jsonrpc: '2.0',
id: request.id,
result: {
protocolVersion: '2024-11-05',
capabilities: {
tools: {},
resources: {},
prompts: {}
},
serverInfo: {
name: 'TryAII-MCP-Server',
version: '1.0.0'
}
}
};
}
catch (error) {
logger.error('Failed to initialize MCP connection', {
connectionId,
error: error instanceof Error ? error.message : error
});
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32603,
message: 'Failed to initialize connection'
}
};
}
}
/**
* Multi-source API key extraction
*/
extractApiKey(request, clientInfo) {
logger.debug('🔍 DEBUG: extractApiKey - checking all sources', {
clientInfo_meta_apiKey: clientInfo?.meta?.apiKey,
clientInfo_apiKey: clientInfo?.apiKey,
clientInfo_name: clientInfo?.name,
clientInfo_version: clientInfo?.version,
request_params_apiKey: request.params?.apiKey,
request_params_keys: request.params ? Object.keys(request.params) : []
});
// Try multiple sources for API key (no environment fallback)
const apiKey = clientInfo?.meta?.apiKey || // Primary: clientInfo.meta.apiKey
clientInfo?.apiKey || // Secondary: clientInfo.apiKey
request.params?.apiKey || // Tertiary: request.params.apiKey
this.extractFromHeaders(request) || // Quaternary: headers
null; // No fallback - authentication required
logger.debug('🔍 DEBUG: extractApiKey - result', {
foundApiKey: apiKey ? `${apiKey.substring(0, 8)}...` : 'NULL',
apiKeyValid: apiKey && typeof apiKey === 'string' && apiKey.startsWith('tai_'),
apiKeyLength: apiKey ? apiKey.length : 0,
apiKeyPrefix: apiKey ? apiKey.substring(0, 4) : 'N/A'
});
if (apiKey && typeof apiKey === 'string' && apiKey.startsWith('tai_')) {
return apiKey;
}
return null;
}
/**
* Extract API key from request headers (if available)
*/
extractFromHeaders(request) {
const headers = request.headers || {};
const authHeader = headers.authorization || headers.Authorization;
if (authHeader) {
if (authHeader.startsWith('Bearer ')) {
return authHeader.substring(7);
}
if (authHeader.startsWith('ApiKey ')) {
return authHeader.substring(7);
}
return authHeader;
}
return null;
}
async handleToolsList(connectionId, request) {
const connectionSession = this.connectionSessions.get(connectionId);
if (!connectionSession || !connectionSession.initialized) {
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32002,
message: 'Connection not initialized'
}
};
}
logger.debug('MCP tools/list request', {
connectionId,
sessionId: connectionSession.sessionId
});
return {
jsonrpc: '2.0',
id: request.id,
result: {
tools: [
{
name: 'list_available_models',
description: 'List all available AI models with their capabilities and pricing',
inputSchema: {
type: 'object',
properties: {
provider: {
type: 'string',
description: 'Filter by provider (openai, anthropic, google, deepseek, xai, mistral)'
}
}
}
},
{
name: 'chat_with_model',
description: 'Chat with a specific AI model',
inputSchema: {
type: 'object',
properties: {
modelId: { type: 'string', description: 'The ID of the model to use' },
message: { type: 'string', description: 'The message to send to the model' },
enableWebSearch: { type: 'boolean', description: 'Enable web search (if supported)' },
temperature: { type: 'number', minimum: 0, maximum: 2, description: 'Sampling temperature' },
maxTokens: { type: 'number', minimum: 1, maximum: 12000, description: 'Maximum tokens to generate' },
conversationHistory: {
type: 'array',
items: {
type: 'object',
properties: {
role: { type: 'string', enum: ['user', 'assistant', 'system'] },
content: { type: 'string' }
},
required: ['role', 'content']
},
description: 'Previous conversation messages'
}
},
required: ['modelId', 'message']
}
},
{
name: 'compare_models',
description: 'Compare responses from multiple AI models',
inputSchema: {
type: 'object',
properties: {
modelIds: {
type: 'array',
items: { type: 'string' },
description: 'Array of model IDs to compare'
},
message: { type: 'string', description: 'The message to send to all models' },
enableWebSearch: { type: 'boolean', description: 'Enable web search for supported models' },
temperature: { type: 'number', minimum: 0, maximum: 2, description: 'Sampling temperature' },
maxTokens: { type: 'number', minimum: 1, maximum: 12000, description: 'Maximum tokens per model' }
},
required: ['modelIds', 'message']
}
},
{
name: 'get_model_info',
description: 'Get detailed information about a specific model',
inputSchema: {
type: 'object',
properties: {
modelId: { type: 'string', description: 'The ID of the model to get information about' }
},
required: ['modelId']
}
},
{
name: 'brains',
description: 'Collective intelligence from 5 top AI models: o3, Claude Opus 4, DeepSeek Chat, Gemini 2.5 Pro, and Grok 3. Creates a beautiful HTML report saved to a file with a clickable URL for instant browser viewing - always return it to the user. You should sum the results of the 5 models and return the final answer.',
inputSchema: {
type: 'object',
properties: {
question: { type: 'string', description: 'Question to ask the 5 smartest AI models' },
enableWebSearch: { type: 'boolean', description: 'Enable web search for supported models' },
temperature: { type: 'number', minimum: 0, maximum: 2, description: 'Sampling temperature' },
maxTokens: { type: 'number', minimum: 1, maximum: 12000, description: 'Maximum tokens per model' }
},
required: ['question']
}
}
]
}
};
}
async handleToolCall(connectionId, request) {
const connectionSession = this.connectionSessions.get(connectionId);
if (!connectionSession || !connectionSession.initialized) {
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32002,
message: 'Connection not initialized'
}
};
}
const { name, arguments: args } = request.params || {};
if (!name) {
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32602,
message: 'Tool name is required'
}
};
}
// Permission checking for authenticated connections
const hasPermission = this.checkToolPermission(name, connectionSession.auth.permissions);
if (!hasPermission) {
logger.warn('MCP tool call denied - insufficient permissions', {
connectionId,
toolName: name,
userId: connectionSession.auth.userId,
userPermissions: connectionSession.auth.permissions
});
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32601,
message: `Access denied: Tool '${name}' requires permission not granted to your API key`
}
};
}
// Calculate tool cost and check balance BEFORE execution
const toolCost = this.calculateToolCost(name, args);
if (toolCost > 0) {
const hasSufficientBalance = await balanceService.checkSufficientBalance(connectionSession.auth.userId, toolCost);
if (!hasSufficientBalance) {
const currentBalance = await balanceService.getCurrentBalance(connectionSession.auth.userId);
logger.warn('MCP tool call denied - insufficient balance', {
connectionId,
toolName: name,
userId: connectionSession.auth.userId,
requiredCost: toolCost,
currentBalance
});
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32602,
message: `Insufficient balance. Required: $${toolCost}, Available: $${currentBalance}`
}
};
}
}
logger.info('MCP tool call', {
connectionId,
sessionId: connectionSession.sessionId,
toolName: name,
hasArgs: !!args
});
const startTime = Date.now();
try {
// Check if session still exists, if not recreate it
let sessionId = connectionSession.sessionId;
const session = this.sessionManager.getSession(sessionId);
if (!session) {
logger.info('Session not found, recreating session for MCP connection', {
connectionId,
oldSessionId: sessionId,
userId: connectionSession.auth.userId,
keyId: connectionSession.auth.keyId
});
// Create session with authenticated user data
const newSession = await this.sessionManager.getOrCreateSession(connectionSession.auth.userId, connectionSession.auth.keyId);
connectionSession.sessionId = newSession.sessionId;
sessionId = newSession.sessionId;
logger.info('New session created for MCP connection', {
connectionId,
newSessionId: sessionId
});
}
const result = await this.sessionManager.queueRequest(sessionId, name, args || {});
const executionTime = Date.now() - startTime;
// 🚀 ENHANCED: Extract actual cost immediately using our improved method
const actualCost = this.extractActualCostFromResult(result, name);
logger.info('McpProtocolHandler: Cost extraction complete', {
toolName: name,
userId: connectionSession.auth.userId,
sessionId,
actualCost,
executionTime: `${executionTime}ms`,
hasCalculatedCost: !!result?._mcpCalculatedCost,
calculatedCostValue: result?._mcpCalculatedCost
});
// Deduct actual balance AFTER execution
if (actualCost > 0) {
const balanceResult = await balanceService.deductBalance(connectionSession.auth.userId, actualCost, `MCP ${name} tool execution`, {
ipAddress: 'mcp-connection',
userAgent: 'MCP Protocol',
requestId: `mcp_${Date.now()}`
});
logger.info('Balance deducted for MCP tool call', {
toolName: name,
userId: connectionSession.auth.userId,
cost: actualCost,
success: balanceResult.success,
balanceAfter: balanceResult.balanceAfter,
transactionId: balanceResult.transactionId
});
// Store balance information in session for database save
const currentSession = this.sessionManager.getSession(sessionId);
if (currentSession) {
currentSession.balanceInfo = {
totalCost: actualCost,
balanceResult: balanceResult
};
logger.debug('Updated session with balance info after execution - cost consistency check', {
sessionId,
mcpCalculatedCost: actualCost,
balanceSuccess: balanceResult.success,
transactionId: balanceResult.transactionId
});
}
}
else {
logger.warn('McpProtocolHandler: No cost to deduct - tool may be free or cost extraction failed', {
toolName: name,
actualCost,
resultKeys: Object.keys(result || {}),
userId: connectionSession.auth.userId
});
}
// Update usage stats
await crossDatabaseService.updateUsageStats(connectionSession.auth.keyId, {
totalRequests: 1,
totalCost: actualCost,
lastUsed: new Date()
});
return {
jsonrpc: '2.0',
id: request.id,
result: {
content: [
{
type: 'text',
text: JSON.stringify(result, null, 2)
}
]
}
};
}
catch (error) {
logger.error('MCP tool call failed', {
connectionId,
sessionId: connectionSession.sessionId,
toolName: name,
error: error instanceof Error ? error.message : error
});
return {
jsonrpc: '2.0',
id: request.id,
error: {
code: -32603,
message: error instanceof Error ? error.message : 'Tool execution failed'
}
};
}
}
/**
* Calculate estimated cost for a tool before execution
*/
calculateToolCost(toolName, args) {
const costs = {
'brains': 0.10, // 5 top models
'compare_models': 0.06, // Variable based on model count
'chat_with_model': 0.02, // Single model
'list_available_models': 0, // Free
'get_model_info': 0 // Free
};
let baseCost = costs[toolName] || 0;
// Adjust costs based on parameters
if (toolName === 'compare_models' && args?.modelIds) {
baseCost = Math.min(args.modelIds.length * 0.02, 0.20); // Cap at $0.20
}
return baseCost;
}
/**
* Extract actual cost from result content
* 🚀 ENHANCED: Fixed to properly extract costs from the session data structure
*/
extractActualCostFromResult(result, toolName) {
try {
let totalCost = 0;
let modelCount = 0;
// 🚀 PRIORITY FIX: Check for calculated cost attached by SessionManager
if (result?._mcpCalculatedCost && typeof result._mcpCalculatedCost === 'number' && result._mcpCalculatedCost > 0) {
logger.info('McpProtocolHandler: Successfully extracted _mcpCalculatedCost from SessionManager', {
toolName,
totalCost: result._mcpCalculatedCost,
source: '_mcpCalculatedCost'
});
return result._mcpCalculatedCost;
}
// 🚀 ENHANCED: First check for direct cost fields in the result
if (result?.totalCost && typeof result.totalCost === 'number' && result.totalCost > 0) {
logger.info('McpProtocolHandler: Successfully extracted direct totalCost', {
toolName,
totalCost: result.totalCost,
source: 'directCostField'
});
return result.totalCost;
}
// 🚀 ENHANCED: Check for cost in the result metadata or content
if (result?.content?.[0]?.text) {
const resultText = result.content[0].text;
// Try to parse as JSON to extract cost information
try {
const parsed = JSON.parse(resultText);
// 🚀 ENHANCED: Check for totalCost field in parsed content
if (parsed?.totalCost && typeof parsed.totalCost === 'number' && parsed.totalCost > 0) {
logger.info('McpProtocolHandler: Successfully extracted totalCost from content', {
toolName,
totalCost: parsed.totalCost,
source: 'parsedContent'
});
return parsed.totalCost;
}
// 🚀 ENHANCED: Look for model interactions with enhanced structure
if (parsed?.modelInteractions && Array.isArray(parsed.modelInteractions)) {
for (const interaction of parsed.modelInteractions) {
if (interaction.cost && typeof interaction.cost === 'number') {
totalCost += interaction.cost;
modelCount++;
}
}
}
// 🚀 ENHANCED: Look for responses array (from brains/compare tools)
if (parsed?.responses && Array.isArray(parsed.responses)) {
for (const response of parsed.responses) {
if (response?.cost && typeof response.cost === 'number') {
totalCost += response.cost;
modelCount++;
}
}
}
// 🚀 ENHANCED: Look for results array (from compare_models)
if (parsed?.results && Array.isArray(parsed.results)) {
for (const result of parsed.results) {
if (result?.cost && typeof result.cost === 'number') {
totalCost += result.cost;
modelCount++;
}
}
}
}
catch (parseError) {
// 🚀 ENHANCED: Use more sophisticated pattern matching for cost extraction
const patterns = [
/["']?totalCost["']?\s*:\s*([\d.]+)/,
/["']?cost["']?\s*:\s*([\d.]+)/g
];
for (const pattern of patterns) {
const matches = resultText.match(pattern);
if (matches) {
if (pattern.global) {
// For global patterns, extract all costs
for (const match of matches) {
const cost = parseFloat(match.replace(/["']?\w+["']?\s*:\s*/, ''));
if (!isNaN(cost)) {
totalCost += cost;
modelCount++;
}
}
}
else {
// For single patterns like totalCost, use directly
const cost = parseFloat(matches[1]);
if (!isNaN(cost)) {
logger.info('McpProtocolHandler: Successfully extracted cost via pattern matching', {
toolName,
totalCost: cost,
source: 'patternMatching'
});
return cost;
}
}
}
}
}
}
// 🚀 ENHANCED: Look directly in responses array from mcp_tryaii
if (result?.responses && Array.isArray(result.responses)) {
for (const response of result.responses) {
if (response?.cost && typeof response.cost === 'number') {
totalCost += response.cost;
modelCount++;
}
}
}
// 🚀 ENHANCED: Look directly in results array from compare_models
if (result?.results && Array.isArray(result.results)) {
for (const resultItem of result.results) {
if (resultItem?.cost && typeof resultItem.cost === 'number') {
totalCost += resultItem.cost;
modelCount++;
}
}
}
// Validate the extracted cost
if (totalCost > 0 && modelCount > 0) {
logger.info('McpProtocolHandler: Successfully extracted actual cost from model interactions', {
toolName,
totalCost,
modelCount,
source: 'modelInteractions'
});
return totalCost;
}
// If no cost found, use conservative fallback
logger.warn('McpProtocolHandler: No cost found in result, using conservative fallback', {
toolName,
resultKeys: Object.keys(result || {}),
hasCalculatedCost: !!result?._mcpCalculatedCost,
hasTotalCost: !!result?.totalCost,
hasResponses: !!result?.responses,
hasResults: !!result?.results,
contentPreview: result?.content?.[0]?.text?.substring(0, 200) || 'no content'
});
return this.getConservativeFallbackCost(toolName);
}
catch (error) {
logger.error('McpProtocolHandler: Error extracting cost from result', {
toolName,
error: error instanceof Error ? error.message : error
});
return this.getConservativeFallbackCost(toolName);
}
}
/**
* Get conservative fallback cost to prevent overcharging
*/
getConservativeFallbackCost(toolName) {
const conservativeCosts = {
'brains': 0.02, // Instead of 0.15, use 0.02
'compare_models': 0.025, // Instead of 0.08, use 0.025
'chat_with_model': 0.01, // Instead of 0.03, use 0.01
'list_available_models': 0,
'get_model_info': 0
};
const cost = conservativeCosts[toolName] || 0.01;
logger.info('McpProtocolHandler: Using conservative fallback cost', {
toolName,
cost,
reason: 'Could not extract actual cost from result'
});
return cost;
}
/**
* Check if user has permission to use a specific tool
*/
checkToolPermission(toolName, userPermissions) {
// Map tools to required permissions
const toolPermissions = {
'list_available_models': 'models',
'get_model_info': 'models',
'chat_with_model': 'chat',
'compare_models': 'compare',
'brains': 'chat', // Brains requires chat permission
};
const requiredPermission = toolPermissions[toolName];
// If tool doesn't have a specific permission requirement, allow it
if (!requiredPermission) {
return true;
}
// Check if user has the required permission
return userPermissions.includes(requiredPermission);
}
async handleResourcesList(connectionId, request) {
return {
jsonrpc: '2.0',
id: request.id,
result: {
resources: []
}
};
}
async handlePromptsList(connectionId, request) {
return {
jsonrpc: '2.0',
id: request.id,
result: {
prompts: []
}
};
}
removeConnection(connectionId) {
const connectionSession = this.connectionSessions.get(connectionId);
if (connectionSession) {
logger.info('Removing MCP connection session', {
connectionId,
sessionId: connectionSession.sessionId
});
this.connectionSessions.delete(connectionId);
}
}
getConnectionCount() {
return this.connectionSessions.size;
}
}
//# sourceMappingURL=McpProtocolHandler.js.map