tryaii-mcp-server
Version:
TryAII MCP Server - 15+ AI models with comparison, cost tracking, and collective intelligence
228 lines • 10.2 kB
JavaScript
import { authenticationService } from './services/AuthenticationService.js';
import { rateLimitService } from './services/RateLimitService.js';
import { metricsService } from './services/MetricsService.js';
import { logger } from './utils/logger.js';
export class AuthMiddleware {
static async authenticate(options = {}) {
return async (req, res, next) => {
const startTime = Date.now();
req.startTime = startTime;
req.requestId = `req_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
try {
// Extract API key from header
const authHeader = req.headers.authorization;
const apiKey = authHeader?.startsWith('Bearer ') ? authHeader.substring(7) : null;
// Track request start
if (!options.skipMetrics) {
metricsService.recordCounter('requests.total');
logger.debug('Request started', {
requestId: req.requestId,
method: req.method,
path: req.path,
userAgent: req.headers['user-agent']
});
}
// Handle unauthenticated requests
if (!apiKey) {
if (options.requireAuth) {
return AuthMiddleware.sendErrorResponse(res, 401, 'Missing API key', req.requestId);
}
// Allow unauthenticated requests for public endpoints
return next();
}
// Authenticate the API key
const user = await authenticationService.authenticateApiKey(apiKey);
if (!user) {
if (!options.skipMetrics) {
metricsService.recordCounter('auth.failures');
}
logger.warn('Authentication failed', {
requestId: req.requestId,
keyPrefix: apiKey.substring(0, 8)
});
return AuthMiddleware.sendErrorResponse(res, 401, 'Invalid API key', req.requestId);
}
// Attach user to request
req.user = user;
// Check permissions if required
if (options.requiredPermissions?.length) {
const hasPermissions = await authenticationService.validatePermissions(user, options.requiredPermissions);
if (!hasPermissions) {
if (!options.skipMetrics) {
metricsService.recordCounter('auth.permission_denied');
}
logger.warn('Insufficient permissions', {
requestId: req.requestId,
userId: user.userId,
required: options.requiredPermissions,
user: user.permissions
});
return AuthMiddleware.sendErrorResponse(res, 403, 'Insufficient permissions', req.requestId);
}
}
// Check rate limits
if (!options.skipRateLimit) {
const userTier = options.rateLimitTier || user.user?.balance?.currency || 'free';
const estimatedCost = AuthMiddleware.estimateRequestCost(req);
const rateLimitResult = await rateLimitService.checkRateLimit(user.userId, userTier, estimatedCost);
if (!rateLimitResult.allowed) {
if (!options.skipMetrics) {
metricsService.recordCounter('rate_limit.violations');
}
logger.warn('Rate limit exceeded', {
requestId: req.requestId,
userId: user.userId,
tier: userTier,
reason: rateLimitResult.reason
});
return AuthMiddleware.sendRateLimitResponse(res, rateLimitResult, req.requestId);
}
// Add rate limit headers
res.set({
'X-RateLimit-Remaining': rateLimitResult.remaining.toString(),
'X-RateLimit-Reset': Math.ceil(rateLimitResult.resetTime / 1000).toString()
});
}
// Record successful authentication
if (!options.skipMetrics) {
metricsService.recordCounter('auth.successes');
metricsService.recordCounter('requests.authenticated');
}
logger.debug('Authentication successful', {
requestId: req.requestId,
userId: user.userId,
permissions: user.permissions.length
});
next();
}
catch (error) {
logger.error('Authentication middleware error', {
requestId: req.requestId,
error: error instanceof Error ? error.message : String(error),
stack: error instanceof Error ? error.stack : undefined
});
if (!options.skipMetrics) {
metricsService.recordCounter('auth.errors');
}
return AuthMiddleware.sendErrorResponse(res, 500, 'Authentication service error', req.requestId);
}
};
}
// Middleware to track request completion
static requestTracker() {
return (req, res, next) => {
const originalSend = res.send;
const originalJson = res.json;
// Override send method to capture response
res.send = function (body) {
AuthMiddleware.trackRequestCompletion(req, res, body);
return originalSend.call(this, body);
};
// Override json method to capture response
res.json = function (body) {
AuthMiddleware.trackRequestCompletion(req, res, body);
return originalJson.call(this, body);
};
next();
};
}
static trackRequestCompletion(req, res, body) {
if (!req.startTime || !req.requestId)
return;
const responseTime = Date.now() - req.startTime;
const success = res.statusCode < 400;
// Record metrics
metricsService.recordRequest(success, responseTime, req.route?.path || req.path, req.user?.userId);
// Log request completion
logger.info('Request completed', {
requestId: req.requestId,
method: req.method,
path: req.path,
statusCode: res.statusCode,
responseTime,
userId: req.user?.userId,
success
});
// Track AI usage if applicable
if (body && typeof body === 'object' && body.usage) {
metricsService.recordAIUsage(body.usage.modelId || 'unknown', body.usage.provider || 'unknown', body.usage.cost || 0, body.usage.tokensUsed || 0, req.user?.userId);
}
}
static estimateRequestCost(req) {
// Simple cost estimation based on endpoint and body size
const bodySize = req.body ? JSON.stringify(req.body).length : 0;
const baseCost = 0.001; // Base cost per request
const sizeCost = bodySize * 0.000001; // Cost per byte
// Higher cost for AI endpoints
if (req.path.includes('/chat') || req.path.includes('/compare') || req.path.includes('/brains')) {
return baseCost * 10 + sizeCost;
}
return baseCost + sizeCost;
}
static sendErrorResponse(res, statusCode, message, requestId) {
return res.status(statusCode).json({
success: false,
error: {
code: statusCode,
message,
requestId,
timestamp: new Date().toISOString()
}
});
}
static sendRateLimitResponse(res, rateLimitResult, requestId) {
const retryAfter = rateLimitResult.retryAfter || 60;
res.set({
'Retry-After': retryAfter.toString(),
'X-RateLimit-Remaining': '0',
'X-RateLimit-Reset': Math.ceil(rateLimitResult.resetTime / 1000).toString()
});
return res.status(429).json({
success: false,
error: {
code: 429,
message: 'Rate limit exceeded',
reason: rateLimitResult.reason,
retryAfter,
requestId,
timestamp: new Date().toISOString()
}
});
}
// Utility method for checking user permissions inline
static requirePermission(permission) {
return AuthMiddleware.authenticate({
requireAuth: true,
requiredPermissions: [permission]
});
}
// Utility method for different rate limit tiers
static withRateLimit(tier) {
return AuthMiddleware.authenticate({
requireAuth: true,
rateLimitTier: tier
});
}
// Utility method for public endpoints with optional auth
static optionalAuth() {
return AuthMiddleware.authenticate({
requireAuth: false,
skipRateLimit: true
});
}
// Admin only middleware
static adminOnly() {
return AuthMiddleware.authenticate({
requireAuth: true,
requiredPermissions: ['admin'],
rateLimitTier: 'admin'
});
}
}
export const authenticate = AuthMiddleware.authenticate;
export const requestTracker = AuthMiddleware.requestTracker;
export const requirePermission = AuthMiddleware.requirePermission;
export const withRateLimit = AuthMiddleware.withRateLimit;
export const optionalAuth = AuthMiddleware.optionalAuth;
export const adminOnly = AuthMiddleware.adminOnly;
//# sourceMappingURL=AuthMiddleware.js.map