Major refactor: Multi-user Chrome MCP extension with remote server architecture
This commit is contained in:
nasir@endelospay.com
487
app/remote-server/src/index.ts
Normal file
487
app/remote-server/src/index.ts
Normal file
@@ -0,0 +1,487 @@
|
||||
import Fastify from 'fastify';
|
||||
import cors from '@fastify/cors';
|
||||
import websocket from '@fastify/websocket';
|
||||
import { pino } from 'pino';
|
||||
import chalk from 'chalk';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
|
||||
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
|
||||
import { MCPRemoteServer } from './server/mcp-remote-server.js';
|
||||
|
||||
const logger = pino({
|
||||
level: 'info',
|
||||
});
|
||||
|
||||
async function startServer() {
|
||||
const fastify = Fastify({
|
||||
logger: true,
|
||||
});
|
||||
|
||||
// Register CORS
|
||||
await fastify.register(cors, {
|
||||
origin: true,
|
||||
credentials: true,
|
||||
});
|
||||
|
||||
// Register WebSocket support
|
||||
await fastify.register(websocket);
|
||||
|
||||
// Create MCP Remote Server instance
|
||||
const mcpServer = new MCPRemoteServer(logger);
|
||||
|
||||
// Transport mapping for streaming connections
|
||||
const transportsMap: Map<string, StreamableHTTPServerTransport | SSEServerTransport> = new Map();
|
||||
|
||||
// Health check endpoint
|
||||
fastify.get('/health', async (request, reply) => {
|
||||
return { status: 'ok', timestamp: new Date().toISOString() };
|
||||
});
|
||||
|
||||
// SSE endpoint for streaming MCP communication
|
||||
fastify.get('/sse', async (request, reply) => {
|
||||
try {
|
||||
// Set SSE headers
|
||||
reply.raw.writeHead(200, {
|
||||
'Content-Type': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache',
|
||||
Connection: 'keep-alive',
|
||||
'Access-Control-Allow-Origin': '*',
|
||||
'Access-Control-Allow-Headers': 'Cache-Control',
|
||||
});
|
||||
|
||||
// Create SSE transport
|
||||
const transport = new SSEServerTransport('/messages', reply.raw);
|
||||
transportsMap.set(transport.sessionId, transport);
|
||||
|
||||
reply.raw.on('close', () => {
|
||||
transportsMap.delete(transport.sessionId);
|
||||
logger.info(`SSE connection closed for session: ${transport.sessionId}`);
|
||||
});
|
||||
|
||||
// Start the transport first
|
||||
await transport.start();
|
||||
|
||||
// Connect the MCP server to this transport
|
||||
await mcpServer.connectTransport(transport);
|
||||
|
||||
// Hijack the reply to prevent Fastify from sending additional headers
|
||||
reply.hijack();
|
||||
} catch (error) {
|
||||
logger.error('Error setting up SSE connection:', error);
|
||||
if (!reply.sent && !reply.raw.headersSent) {
|
||||
reply.code(500).send({ error: 'Internal server error' });
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// POST /messages: Handle SSE POST messages
|
||||
fastify.post('/messages', async (request, reply) => {
|
||||
const sessionId = request.headers['x-session-id'] as string | undefined;
|
||||
const transport = sessionId ? (transportsMap.get(sessionId) as SSEServerTransport) : undefined;
|
||||
|
||||
if (!transport) {
|
||||
reply.code(400).send({ error: 'Invalid session ID for SSE' });
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
await transport.handlePostMessage(request.raw, reply.raw, request.body);
|
||||
} catch (error) {
|
||||
logger.error('Error handling SSE POST message:', error);
|
||||
if (!reply.sent) {
|
||||
reply.code(500).send({ error: 'Internal server error' });
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// POST /mcp: Handle client-to-server messages for streamable HTTP
|
||||
fastify.post('/mcp', async (request, reply) => {
|
||||
// Extract session ID and user ID from headers for routing
|
||||
const sessionId = request.headers['mcp-session-id'] as string | undefined;
|
||||
const userId = request.headers['chrome-user-id'] as string | undefined;
|
||||
let transport: StreamableHTTPServerTransport | undefined = transportsMap.get(
|
||||
sessionId || '',
|
||||
) as StreamableHTTPServerTransport;
|
||||
|
||||
if (transport) {
|
||||
// Transport found, use existing one
|
||||
} else if (!sessionId && isInitializeRequest(request.body)) {
|
||||
// Create new session for initialization request
|
||||
const newSessionId = randomUUID();
|
||||
transport = new StreamableHTTPServerTransport({
|
||||
sessionIdGenerator: () => newSessionId,
|
||||
onsessioninitialized: (initializedSessionId) => {
|
||||
if (transport && initializedSessionId === newSessionId) {
|
||||
transportsMap.set(initializedSessionId, transport);
|
||||
logger.info(`New streamable HTTP session initialized: ${initializedSessionId}`);
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
// Connect the MCP server to this transport
|
||||
await mcpServer.connectTransport(transport);
|
||||
} else {
|
||||
reply.code(400).send({ error: 'Invalid session or missing initialization' });
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
// Set user context for routing if user ID is provided
|
||||
if (userId) {
|
||||
mcpServer.setUserContext(userId, sessionId);
|
||||
logger.info(
|
||||
`🎯 [MCP] User context set for request - User: ${userId}, Session: ${sessionId}`,
|
||||
);
|
||||
}
|
||||
|
||||
await transport.handleRequest(request.raw, reply.raw, request.body);
|
||||
if (!reply.sent) {
|
||||
reply.hijack(); // Prevent Fastify from automatically sending response
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error handling streamable HTTP POST request:', error);
|
||||
if (!reply.sent) {
|
||||
reply.code(500).send({ error: 'Internal server error' });
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// GET /mcp: Handle SSE stream for streamable HTTP
|
||||
fastify.get('/mcp', async (request, reply) => {
|
||||
const sessionId = request.headers['mcp-session-id'] as string | undefined;
|
||||
const transport = sessionId
|
||||
? (transportsMap.get(sessionId) as StreamableHTTPServerTransport)
|
||||
: undefined;
|
||||
|
||||
if (!transport) {
|
||||
reply.code(400).send({ error: 'Invalid session ID' });
|
||||
return;
|
||||
}
|
||||
|
||||
reply.raw.setHeader('Content-Type', 'text/event-stream');
|
||||
reply.raw.setHeader('Cache-Control', 'no-cache');
|
||||
reply.raw.setHeader('Connection', 'keep-alive');
|
||||
reply.raw.setHeader('Access-Control-Allow-Origin', '*');
|
||||
reply.raw.flushHeaders();
|
||||
|
||||
try {
|
||||
await transport.handleRequest(request.raw, reply.raw);
|
||||
if (!reply.sent) {
|
||||
reply.hijack(); // Prevent Fastify from automatically sending response
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error handling streamable HTTP GET request:', error);
|
||||
if (!reply.raw.writableEnded) {
|
||||
reply.raw.end();
|
||||
}
|
||||
}
|
||||
|
||||
request.socket.on('close', () => {
|
||||
logger.info(`Streamable HTTP client disconnected for session: ${sessionId}`);
|
||||
});
|
||||
});
|
||||
|
||||
// DELETE /mcp: Handle session termination for streamable HTTP
|
||||
fastify.delete('/mcp', async (request, reply) => {
|
||||
const sessionId = request.headers['mcp-session-id'] as string | undefined;
|
||||
const transport = sessionId
|
||||
? (transportsMap.get(sessionId) as StreamableHTTPServerTransport)
|
||||
: undefined;
|
||||
|
||||
if (!transport) {
|
||||
reply.code(400).send({ error: 'Invalid session ID' });
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
await transport.handleRequest(request.raw, reply.raw);
|
||||
if (sessionId) {
|
||||
transportsMap.delete(sessionId);
|
||||
logger.info(`Streamable HTTP session terminated: ${sessionId}`);
|
||||
}
|
||||
|
||||
if (!reply.sent) {
|
||||
reply.code(204).send();
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error handling streamable HTTP DELETE request:', error);
|
||||
if (!reply.sent) {
|
||||
reply.code(500).send({ error: 'Internal server error' });
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// WebSocket endpoint for MCP communication
|
||||
fastify.register(async function (fastify) {
|
||||
fastify.get('/ws/mcp', { websocket: true }, (connection: any, req) => {
|
||||
logger.info('New MCP WebSocket connection established');
|
||||
|
||||
// Set up ping/pong to keep connection alive
|
||||
const pingInterval = setInterval(() => {
|
||||
if (connection.readyState === connection.OPEN) {
|
||||
connection.ping();
|
||||
}
|
||||
}, 30000); // Ping every 30 seconds
|
||||
|
||||
connection.on('pong', () => {
|
||||
logger.debug('Received pong from MCP client');
|
||||
});
|
||||
|
||||
connection.on('message', async (message: any) => {
|
||||
try {
|
||||
const data = JSON.parse(message.toString());
|
||||
console.log('🔵 [MCP WebSocket] Received message:', JSON.stringify(data, null, 2));
|
||||
logger.info('🔵 [MCP WebSocket] Received message:', {
|
||||
method: data.method,
|
||||
id: data.id,
|
||||
params: data.params,
|
||||
fullMessage: data,
|
||||
});
|
||||
|
||||
// Handle MCP protocol messages with proper ID preservation
|
||||
let response;
|
||||
|
||||
if (data.method === 'tools/list') {
|
||||
try {
|
||||
const toolsResponse = await mcpServer.handleMessage(data);
|
||||
response = {
|
||||
jsonrpc: '2.0',
|
||||
id: data.id,
|
||||
result: toolsResponse,
|
||||
};
|
||||
} catch (error) {
|
||||
response = {
|
||||
jsonrpc: '2.0',
|
||||
id: data.id,
|
||||
error: {
|
||||
code: -32603,
|
||||
message: error instanceof Error ? error.message : 'Internal error',
|
||||
},
|
||||
};
|
||||
}
|
||||
} else if (data.method === 'tools/call') {
|
||||
try {
|
||||
const toolResponse = await mcpServer.handleMessage(data);
|
||||
response = {
|
||||
jsonrpc: '2.0',
|
||||
id: data.id,
|
||||
result: toolResponse,
|
||||
};
|
||||
} catch (error) {
|
||||
response = {
|
||||
jsonrpc: '2.0',
|
||||
id: data.id,
|
||||
error: {
|
||||
code: -32603,
|
||||
message: error instanceof Error ? error.message : 'Tool execution failed',
|
||||
},
|
||||
};
|
||||
}
|
||||
} else {
|
||||
// Unknown method
|
||||
response = {
|
||||
jsonrpc: '2.0',
|
||||
id: data.id,
|
||||
error: {
|
||||
code: -32601,
|
||||
message: 'Method not found',
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
if (response) {
|
||||
logger.info('🟢 [MCP WebSocket] Sending response:', {
|
||||
id: response.id,
|
||||
hasResult: !!response.result,
|
||||
hasError: !!response.error,
|
||||
fullResponse: response,
|
||||
});
|
||||
connection.send(JSON.stringify(response));
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('🚨 [MCP WebSocket] Error handling MCP message:');
|
||||
console.error('Error message:', error instanceof Error ? error.message : String(error));
|
||||
console.error('Error stack:', error instanceof Error ? error.stack : undefined);
|
||||
console.error('Full error:', error);
|
||||
logger.error('🚨 [MCP WebSocket] Error handling MCP message:', {
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
stack: error instanceof Error ? error.stack : undefined,
|
||||
errorType: typeof error,
|
||||
fullError: error,
|
||||
});
|
||||
// Send error response with proper MCP format
|
||||
const errorResponse = {
|
||||
jsonrpc: '2.0',
|
||||
id: null, // Use null if we can't parse the original ID
|
||||
error: {
|
||||
code: -32700,
|
||||
message: 'Parse error',
|
||||
},
|
||||
};
|
||||
connection.send(JSON.stringify(errorResponse));
|
||||
}
|
||||
});
|
||||
|
||||
connection.on('close', () => {
|
||||
logger.info('MCP WebSocket connection closed');
|
||||
clearInterval(pingInterval);
|
||||
});
|
||||
|
||||
connection.on('error', (error: any) => {
|
||||
logger.error('WebSocket error:', error);
|
||||
clearInterval(pingInterval);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
// WebSocket endpoint for Chrome extension connections
|
||||
fastify.register(async function (fastify) {
|
||||
fastify.get('/chrome', { websocket: true }, (connection: any, req) => {
|
||||
logger.info('New Chrome extension WebSocket connection established');
|
||||
|
||||
// Set up ping/pong to keep Chrome extension connection alive
|
||||
const chromeExtensionPingInterval = setInterval(() => {
|
||||
if (connection.readyState === connection.OPEN) {
|
||||
connection.ping();
|
||||
}
|
||||
}, 30000); // Ping every 30 seconds
|
||||
|
||||
// Create a connection wrapper for the Chrome tools
|
||||
const connectionWrapper = {
|
||||
socket: connection,
|
||||
send: (data: string) => connection.send(data),
|
||||
on: (event: string, handler: Function) => connection.on(event, handler),
|
||||
off: (event: string, handler: Function) => connection.off(event, handler),
|
||||
get readyState() {
|
||||
// WebSocket states: 0=CONNECTING, 1=OPEN, 2=CLOSING, 3=CLOSED
|
||||
return connection.readyState || 1; // Default to OPEN if not available
|
||||
},
|
||||
};
|
||||
|
||||
// Extract user information from connection headers or query params
|
||||
const userAgent = req.headers['user-agent'] || 'Unknown';
|
||||
const ipAddress = req.headers['x-forwarded-for'] || req.socket?.remoteAddress || 'Unknown';
|
||||
|
||||
// Initialize with temporary user ID (will be updated when Chrome extension sends connection_info)
|
||||
let currentUserId = `temp_user_${Date.now()}_${Math.random().toString(36).substring(2, 8)}`;
|
||||
|
||||
// Register this connection with the Chrome tools with session management
|
||||
const sessionInfo = mcpServer.registerChromeExtension(connectionWrapper, currentUserId, {
|
||||
userAgent,
|
||||
ipAddress,
|
||||
connectedAt: new Date().toISOString(),
|
||||
connectionType: 'anonymous',
|
||||
});
|
||||
|
||||
logger.info('🟢 [Chrome Extension] Connection registered:', sessionInfo);
|
||||
|
||||
connection.on('message', async (message: any) => {
|
||||
try {
|
||||
const data = JSON.parse(message.toString());
|
||||
|
||||
// Handle connection info message
|
||||
if (data.type === 'connection_info') {
|
||||
logger.info('🔗 [Chrome Extension] Received connection info:', data);
|
||||
|
||||
// Update user ID if provided by Chrome extension
|
||||
if (data.userId && data.userId !== sessionInfo.userId) {
|
||||
logger.info(
|
||||
`🔄 [Chrome Extension] Updating user ID from ${sessionInfo.userId} to ${data.userId}`,
|
||||
);
|
||||
|
||||
// Update the session with the Chrome extension's user ID
|
||||
const updatedSessionInfo = mcpServer.updateChromeExtensionUserId(
|
||||
connectionWrapper,
|
||||
data.userId,
|
||||
);
|
||||
if (updatedSessionInfo) {
|
||||
// Update our local reference
|
||||
Object.assign(sessionInfo, updatedSessionInfo);
|
||||
logger.info(
|
||||
`✅ [Chrome Extension] User ID updated successfully: ${sessionInfo.userId}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Send session info back to extension
|
||||
const sessionResponse = {
|
||||
type: 'session_info',
|
||||
sessionInfo: {
|
||||
userId: sessionInfo.userId,
|
||||
sessionId: sessionInfo.sessionId,
|
||||
connectionId: sessionInfo.connectionId,
|
||||
},
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
connection.send(JSON.stringify(sessionResponse));
|
||||
return;
|
||||
}
|
||||
|
||||
logger.info('🟡 [Chrome Extension] Received message:', {
|
||||
action: data.action,
|
||||
id: data.id,
|
||||
type: data.type,
|
||||
sessionId: sessionInfo.sessionId,
|
||||
userId: sessionInfo.userId,
|
||||
fullMessage: data,
|
||||
});
|
||||
|
||||
// Handle responses from Chrome extension
|
||||
mcpServer.handleChromeResponse(data);
|
||||
} catch (error) {
|
||||
logger.error('Error handling Chrome extension message:', error);
|
||||
}
|
||||
});
|
||||
|
||||
connection.on('close', () => {
|
||||
logger.info('Chrome extension WebSocket connection closed');
|
||||
clearInterval(chromeExtensionPingInterval);
|
||||
mcpServer.unregisterChromeExtension(connectionWrapper);
|
||||
});
|
||||
|
||||
connection.on('error', (error: any) => {
|
||||
logger.error('Chrome extension WebSocket error:', error);
|
||||
clearInterval(chromeExtensionPingInterval);
|
||||
mcpServer.unregisterChromeExtension(connectionWrapper);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
// Start the server
|
||||
const port = process.env.PORT ? parseInt(process.env.PORT) : 3001;
|
||||
const host = process.env.HOST || '0.0.0.0';
|
||||
|
||||
try {
|
||||
await fastify.listen({ port, host });
|
||||
console.log(chalk.green(`🚀 MCP Remote Server started successfully!`));
|
||||
console.log(chalk.blue(`📡 Server running at: http://${host}:${port}`));
|
||||
console.log(chalk.blue(`🌊 Streaming HTTP endpoint: http://${host}:${port}/mcp`));
|
||||
console.log(chalk.blue(`📡 SSE endpoint: http://${host}:${port}/sse`));
|
||||
console.log(chalk.blue(`🔌 WebSocket endpoint: ws://${host}:${port}/ws/mcp`));
|
||||
console.log(chalk.blue(`🔌 Chrome extension endpoint: ws://${host}:${port}/chrome`));
|
||||
console.log(chalk.yellow(`💡 Use 'npm run start:server' to start the server`));
|
||||
} catch (err) {
|
||||
console.error('Error starting server:', err);
|
||||
logger.error('Error starting server:', err);
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
// Handle graceful shutdown
|
||||
process.on('SIGINT', () => {
|
||||
console.log(chalk.yellow('\n🛑 Shutting down server...'));
|
||||
process.exit(0);
|
||||
});
|
||||
|
||||
process.on('SIGTERM', () => {
|
||||
console.log(chalk.yellow('\n🛑 Shutting down server...'));
|
||||
process.exit(0);
|
||||
});
|
||||
|
||||
startServer().catch((error) => {
|
||||
console.error('Failed to start server:', error);
|
||||
logger.error('Failed to start server:', error);
|
||||
process.exit(1);
|
||||
});
|
||||
647
app/remote-server/src/server/chrome-tools.ts
Normal file
647
app/remote-server/src/server/chrome-tools.ts
Normal file
@@ -0,0 +1,647 @@
|
||||
import { Logger } from 'pino';
|
||||
import { TOOL_NAMES } from 'chrome-mcp-shared';
|
||||
import { SessionManager, ExtensionConnection } from './session-manager.js';
|
||||
import { ConnectionRouter, RouteResult } from './connection-router.js';
|
||||
import { LiveKitAgentManager } from './livekit-agent-manager.js';
|
||||
|
||||
export class ChromeTools {
|
||||
private logger: Logger;
|
||||
private sessionManager: SessionManager;
|
||||
private connectionRouter: ConnectionRouter;
|
||||
private liveKitAgentManager: LiveKitAgentManager;
|
||||
private currentUserId?: string;
|
||||
private currentSessionId?: string;
|
||||
|
||||
// Common URL mappings for natural language requests
|
||||
private urlMappings: Map<string, string> = new Map([
|
||||
['google', 'https://www.google.com'],
|
||||
['google.com', 'https://www.google.com'],
|
||||
['youtube', 'https://www.youtube.com'],
|
||||
['youtube.com', 'https://www.youtube.com'],
|
||||
['facebook', 'https://www.facebook.com'],
|
||||
['facebook.com', 'https://www.facebook.com'],
|
||||
['twitter', 'https://www.twitter.com'],
|
||||
['twitter.com', 'https://www.twitter.com'],
|
||||
['x.com', 'https://www.x.com'],
|
||||
['github', 'https://www.github.com'],
|
||||
['github.com', 'https://www.github.com'],
|
||||
['stackoverflow', 'https://www.stackoverflow.com'],
|
||||
['stackoverflow.com', 'https://www.stackoverflow.com'],
|
||||
['reddit', 'https://www.reddit.com'],
|
||||
['reddit.com', 'https://www.reddit.com'],
|
||||
['amazon', 'https://www.amazon.com'],
|
||||
['amazon.com', 'https://www.amazon.com'],
|
||||
['netflix', 'https://www.netflix.com'],
|
||||
['netflix.com', 'https://www.netflix.com'],
|
||||
['linkedin', 'https://www.linkedin.com'],
|
||||
['linkedin.com', 'https://www.linkedin.com'],
|
||||
['instagram', 'https://www.instagram.com'],
|
||||
['instagram.com', 'https://www.instagram.com'],
|
||||
]);
|
||||
|
||||
constructor(logger: Logger) {
|
||||
this.logger = logger;
|
||||
this.sessionManager = new SessionManager(logger);
|
||||
this.connectionRouter = new ConnectionRouter(logger, this.sessionManager);
|
||||
this.liveKitAgentManager = new LiveKitAgentManager(logger, this.sessionManager);
|
||||
}
|
||||
|
||||
// Register a Chrome extension connection with session management
|
||||
registerExtension(
|
||||
connection: any,
|
||||
userId?: string,
|
||||
metadata?: any,
|
||||
): { userId: string; sessionId: string; connectionId: string } {
|
||||
const result = this.sessionManager.registerExtensionConnection(connection, userId, metadata);
|
||||
this.logger.info(
|
||||
`🔗 Chrome extension connected - User: ${result.userId}, Session: ${result.sessionId}`,
|
||||
);
|
||||
|
||||
// Note: LiveKit agent is no longer started automatically on connection
|
||||
// Agents should be started manually when needed
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// Unregister a Chrome extension connection
|
||||
unregisterExtension(connection: any): boolean {
|
||||
const result = this.sessionManager.unregisterExtensionConnection(connection);
|
||||
if (result) {
|
||||
this.logger.info('🔌 Chrome extension disconnected');
|
||||
// Note: LiveKit agent is no longer stopped automatically on disconnection
|
||||
// Agents should be managed manually when needed
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// Update Chrome extension user ID
|
||||
updateExtensionUserId(connection: any, newUserId: string): any {
|
||||
const result = this.sessionManager.updateExtensionUserId(connection, newUserId);
|
||||
if (result) {
|
||||
this.logger.info(`🔄 Chrome extension user ID updated to: ${newUserId}`);
|
||||
// Note: LiveKit agent is no longer restarted automatically on user ID update
|
||||
// Agents should be managed manually when needed
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// Set user context for routing
|
||||
setUserContext(userId: string, sessionId?: string) {
|
||||
this.currentUserId = userId;
|
||||
this.currentSessionId = sessionId;
|
||||
this.logger.info(`🎯 [Chrome Tools] User context set - User: ${userId}, Session: ${sessionId}`);
|
||||
}
|
||||
|
||||
// Handle responses from Chrome extension
|
||||
handleResponse(data: any) {
|
||||
const stats = this.sessionManager.getStats();
|
||||
this.logger.info(`📨 [Chrome Tools] Received response from Chrome extension:`, {
|
||||
messageId: data.id,
|
||||
hasResult: !!data.result,
|
||||
hasError: !!data.error,
|
||||
pendingRequestsCount: stats.pendingRequests,
|
||||
fullData: data,
|
||||
});
|
||||
|
||||
if (data.id) {
|
||||
if (data.error) {
|
||||
this.logger.error(`📨 [Chrome Tools] Chrome extension returned error: ${data.error}`);
|
||||
this.sessionManager.rejectPendingRequest(data.id, new Error(data.error));
|
||||
} else {
|
||||
this.logger.info(
|
||||
`📨 [Chrome Tools] Chrome extension returned success result:`,
|
||||
data.result,
|
||||
);
|
||||
this.sessionManager.resolvePendingRequest(data.id, data.result);
|
||||
}
|
||||
} else {
|
||||
// Filter out ping/heartbeat messages and other non-request messages to reduce noise
|
||||
const isPingMessage =
|
||||
data.type === 'ping' || (data.id && data.id.toString().startsWith('ping_'));
|
||||
const isHeartbeatMessage = !data.id || data.id === undefined;
|
||||
|
||||
if (!isPingMessage && !isHeartbeatMessage) {
|
||||
this.logger.warn(
|
||||
`📨 [Chrome Tools] Received response for unknown or expired request ID: ${data.id}`,
|
||||
);
|
||||
} else {
|
||||
// Log ping/heartbeat messages at debug level to reduce noise
|
||||
this.logger.debug(
|
||||
`📨 [Chrome Tools] Received ping/heartbeat message (ID: ${data.id}, type: ${data.type})`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process natural language navigation requests
|
||||
private processNavigationRequest(args: any): any {
|
||||
if (!args || !args.url) {
|
||||
return args;
|
||||
}
|
||||
|
||||
const url = args.url.toLowerCase().trim();
|
||||
|
||||
// Check if it's a natural language request like "google", "open google", etc.
|
||||
const patterns = [/^(?:open\s+|go\s+to\s+|navigate\s+to\s+)?(.+?)(?:\.com)?$/i, /^(.+?)$/i];
|
||||
|
||||
for (const pattern of patterns) {
|
||||
const match = url.match(pattern);
|
||||
if (match) {
|
||||
const site = match[1].toLowerCase().trim();
|
||||
const mappedUrl = this.urlMappings.get(site);
|
||||
if (mappedUrl) {
|
||||
this.logger.info(`Mapped natural language request "${url}" to "${mappedUrl}"`);
|
||||
return { ...args, url: mappedUrl };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If no mapping found, check if it's already a valid URL
|
||||
if (!url.startsWith('http://') && !url.startsWith('https://')) {
|
||||
// Try to make it a valid URL
|
||||
const processedUrl = url.includes('.')
|
||||
? `https://${url}`
|
||||
: `https://www.google.com/search?q=${encodeURIComponent(url)}`;
|
||||
this.logger.info(`Processed URL "${url}" to "${processedUrl}"`);
|
||||
return { ...args, url: processedUrl };
|
||||
}
|
||||
|
||||
return args;
|
||||
}
|
||||
|
||||
// Send a general tool call to Chrome extension with routing
|
||||
async callTool(name: string, args: any, sessionId?: string, userId?: string): Promise<any> {
|
||||
// Use current user context if not provided
|
||||
const effectiveUserId = userId || this.currentUserId;
|
||||
const effectiveSessionId = sessionId || this.currentSessionId;
|
||||
|
||||
this.logger.info(`🔧 [Chrome Tools] Calling tool: ${name} with routing context:`, {
|
||||
args,
|
||||
sessionId: effectiveSessionId,
|
||||
userId: effectiveUserId,
|
||||
usingCurrentContext: !userId && !sessionId,
|
||||
});
|
||||
|
||||
const message = {
|
||||
action: 'callTool',
|
||||
params: { name, arguments: args },
|
||||
};
|
||||
|
||||
this.logger.info(`🔧 [Chrome Tools] Sending routed message to extensions:`, message);
|
||||
|
||||
const result = await this.sendToExtensions(message, effectiveSessionId, effectiveUserId);
|
||||
|
||||
this.logger.info(`🔧 [Chrome Tools] Received result from extensions:`, result);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// Get session statistics
|
||||
getSessionStats(): any {
|
||||
return this.sessionManager.getStats();
|
||||
}
|
||||
|
||||
// Get routing statistics
|
||||
getRoutingStats(): any {
|
||||
return this.connectionRouter.getRoutingStats();
|
||||
}
|
||||
|
||||
// Get connection by session ID
|
||||
getConnectionBySessionId(sessionId: string): ExtensionConnection | null {
|
||||
return this.sessionManager.getConnectionBySessionId(sessionId);
|
||||
}
|
||||
|
||||
// Get connection by user ID
|
||||
getConnectionByUserId(userId: string): ExtensionConnection | null {
|
||||
return this.sessionManager.getConnectionByUserId(userId);
|
||||
}
|
||||
|
||||
// Route message to specific connection type
|
||||
async callToolWithConnectionType(
|
||||
name: string,
|
||||
args: any,
|
||||
connectionType: 'newest' | 'oldest' | 'most_active',
|
||||
): Promise<any> {
|
||||
this.logger.info(
|
||||
`🔧 [Chrome Tools] Calling tool: ${name} with connection type: ${connectionType}`,
|
||||
);
|
||||
|
||||
const message = {
|
||||
action: 'callTool',
|
||||
params: { name, arguments: args },
|
||||
};
|
||||
|
||||
const routeResult = this.connectionRouter.routeToConnectionType(message, connectionType);
|
||||
const result = await this.sendToExtensions(message, routeResult.sessionId);
|
||||
|
||||
this.logger.info(`🔧 [Chrome Tools] Tool result from ${connectionType} connection:`, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Check if session can handle message
|
||||
canSessionHandleMessage(sessionId: string, messageType: string): boolean {
|
||||
return this.connectionRouter.canSessionHandleMessage(sessionId, messageType);
|
||||
}
|
||||
|
||||
// Get recommended session for user
|
||||
getRecommendedSessionForUser(userId: string): string | null {
|
||||
return this.connectionRouter.getRecommendedSessionForUser(userId);
|
||||
}
|
||||
|
||||
// Get LiveKit agent for user
|
||||
getLiveKitAgentForUser(userId: string): any {
|
||||
return this.liveKitAgentManager.getAgentForUser(userId);
|
||||
}
|
||||
|
||||
// Get LiveKit agent statistics
|
||||
getLiveKitAgentStats(): any {
|
||||
return this.liveKitAgentManager.getAgentStats();
|
||||
}
|
||||
|
||||
// Get all active LiveKit agents
|
||||
getAllActiveLiveKitAgents(): any[] {
|
||||
return this.liveKitAgentManager.getAllActiveAgents();
|
||||
}
|
||||
|
||||
// Cleanup resources
|
||||
destroy(): void {
|
||||
this.connectionRouter.cleanupRoutingRules();
|
||||
this.liveKitAgentManager.shutdownAllAgents();
|
||||
this.sessionManager.destroy();
|
||||
}
|
||||
|
||||
// Send a message to Chrome extensions with intelligent routing
|
||||
private async sendToExtensions(message: any, sessionId?: string, userId?: string): Promise<any> {
|
||||
const stats = this.sessionManager.getStats();
|
||||
this.logger.info(`📤 [Chrome Tools] Routing message to Chrome extensions:`, {
|
||||
action: message.action,
|
||||
connectionsCount: stats.activeConnections,
|
||||
sessionId,
|
||||
userId,
|
||||
fullMessage: message,
|
||||
});
|
||||
|
||||
if (stats.activeConnections === 0) {
|
||||
this.logger.error('🚫 [Chrome Tools] No Chrome extensions connected');
|
||||
throw new Error('No Chrome extensions connected');
|
||||
}
|
||||
|
||||
// Use connection router to find the best connection
|
||||
let routeResult: RouteResult;
|
||||
try {
|
||||
routeResult = this.connectionRouter.routeMessage(message, sessionId, userId);
|
||||
} catch (error) {
|
||||
this.logger.error('Failed to route message:', error);
|
||||
throw error;
|
||||
}
|
||||
|
||||
const { connection: extensionConnection, routingReason } = routeResult;
|
||||
const connection = extensionConnection.connection;
|
||||
const readyState = (connection as any).readyState;
|
||||
|
||||
this.logger.info(
|
||||
`📤 [Chrome Tools] Routed to connection - Session: ${extensionConnection.sessionId}, User: ${extensionConnection.userId}, Reason: ${routingReason}, ReadyState: ${readyState}`,
|
||||
);
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
const messageId = Date.now().toString() + Math.random().toString(36).substring(2, 11);
|
||||
const messageWithId = { ...message, id: messageId };
|
||||
|
||||
// Store the request with session context
|
||||
this.sessionManager.storePendingRequest(
|
||||
messageId,
|
||||
resolve,
|
||||
reject,
|
||||
extensionConnection.sessionId,
|
||||
60000, // 60 second timeout
|
||||
);
|
||||
|
||||
try {
|
||||
// Check if connection is still open before sending
|
||||
if (readyState === 1) {
|
||||
// WebSocket.OPEN
|
||||
this.logger.info(
|
||||
`📤 [Chrome Tools] Sending message with ID ${messageId} to Chrome extension (Session: ${extensionConnection.sessionId}, Routing: ${routingReason}):`,
|
||||
messageWithId,
|
||||
);
|
||||
(connection as any).send(JSON.stringify(messageWithId));
|
||||
} else {
|
||||
this.sessionManager.rejectPendingRequest(
|
||||
messageId,
|
||||
new Error(`Chrome extension connection is not open (readyState: ${readyState})`),
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
this.sessionManager.rejectPendingRequest(messageId, error);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async navigateToUrl(url: string): Promise<any> {
|
||||
this.logger.info(`Navigating to URL: ${url}`);
|
||||
|
||||
// Process natural language navigation requests
|
||||
const processedArgs = this.processNavigationRequest({ url });
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'navigate',
|
||||
params: processedArgs,
|
||||
});
|
||||
}
|
||||
|
||||
async getPageContent(selector?: string): Promise<any> {
|
||||
this.logger.info(`Getting page content${selector ? ` with selector: ${selector}` : ''}`);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'getContent',
|
||||
params: { selector },
|
||||
});
|
||||
}
|
||||
|
||||
async clickElement(selector: string): Promise<any> {
|
||||
this.logger.info(`Clicking element: ${selector}`);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'click',
|
||||
params: { selector },
|
||||
});
|
||||
}
|
||||
|
||||
async fillInput(selector: string, value: string): Promise<any> {
|
||||
this.logger.info(`Filling input ${selector} with value: ${value}`);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'fillInput',
|
||||
params: { selector, value },
|
||||
});
|
||||
}
|
||||
|
||||
async takeScreenshot(fullPage: boolean = false): Promise<any> {
|
||||
this.logger.info(`Taking screenshot (fullPage: ${fullPage})`);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'screenshot',
|
||||
params: { fullPage },
|
||||
});
|
||||
}
|
||||
|
||||
async executeScript(script: string): Promise<any> {
|
||||
this.logger.info('Executing script');
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'executeScript',
|
||||
params: { script },
|
||||
});
|
||||
}
|
||||
|
||||
async getCurrentTab(): Promise<any> {
|
||||
this.logger.info('Getting current tab info');
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'getCurrentTab',
|
||||
params: {},
|
||||
});
|
||||
}
|
||||
|
||||
async getAllTabs(): Promise<any> {
|
||||
this.logger.info('Getting all tabs');
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'getAllTabs',
|
||||
params: {},
|
||||
});
|
||||
}
|
||||
|
||||
async switchToTab(tabId: number): Promise<any> {
|
||||
this.logger.info(`Switching to tab: ${tabId}`);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'switchTab',
|
||||
params: { tabId },
|
||||
});
|
||||
}
|
||||
|
||||
async createNewTab(url?: string): Promise<any> {
|
||||
this.logger.info(`Creating new tab${url ? ` with URL: ${url}` : ''}`);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'createTab',
|
||||
params: { url },
|
||||
});
|
||||
}
|
||||
|
||||
async closeTab(tabId?: number): Promise<any> {
|
||||
this.logger.info(`Closing tab${tabId ? `: ${tabId}` : ' (current)'}`);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: 'closeTab',
|
||||
params: { tabId },
|
||||
});
|
||||
}
|
||||
|
||||
// Browser automation tools matching the native server functionality
|
||||
|
||||
async getWindowsAndTabs(): Promise<any> {
|
||||
this.logger.info('Getting all windows and tabs');
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.GET_WINDOWS_AND_TABS,
|
||||
params: {},
|
||||
});
|
||||
}
|
||||
|
||||
async searchTabsContent(query: string): Promise<any> {
|
||||
this.logger.info(`Searching tabs content for: ${query}`);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.SEARCH_TABS_CONTENT,
|
||||
params: { query },
|
||||
});
|
||||
}
|
||||
|
||||
async chromeNavigate(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome navigate with args:`, args);
|
||||
|
||||
// Process natural language navigation requests
|
||||
const processedArgs = this.processNavigationRequest(args);
|
||||
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.NAVIGATE,
|
||||
params: processedArgs,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeScreenshot(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome screenshot with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.SCREENSHOT,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeCloseTabs(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome close tabs with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.CLOSE_TABS,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeGoBackOrForward(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome go back/forward with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.GO_BACK_OR_FORWARD,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeGetWebContent(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome get web content with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.WEB_FETCHER,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeClickElement(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome click element with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.CLICK,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeFillOrSelect(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome fill or select with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.FILL,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeGetInteractiveElements(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome get interactive elements with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.GET_INTERACTIVE_ELEMENTS,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeNetworkCaptureStart(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome network capture start with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.NETWORK_CAPTURE_START,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeNetworkCaptureStop(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome network capture stop with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.NETWORK_CAPTURE_STOP,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeNetworkRequest(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome network request with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.NETWORK_REQUEST,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeNetworkDebuggerStart(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome network debugger start with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.NETWORK_DEBUGGER_START,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeNetworkDebuggerStop(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome network debugger stop with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.NETWORK_DEBUGGER_STOP,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeKeyboard(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome keyboard with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.KEYBOARD,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeHistory(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome history with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.HISTORY,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeBookmarkSearch(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome bookmark search with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.BOOKMARK_SEARCH,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeBookmarkAdd(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome bookmark add with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.BOOKMARK_ADD,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeBookmarkDelete(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome bookmark delete with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.BOOKMARK_DELETE,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeInjectScript(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome inject script with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.INJECT_SCRIPT,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeSendCommandToInjectScript(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome send command to inject script with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.SEND_COMMAND_TO_INJECT_SCRIPT,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeConsole(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome console with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.CONSOLE,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeSearchGoogle(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome search Google with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.SEARCH_GOOGLE,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
|
||||
async chromeSubmitForm(args: any): Promise<any> {
|
||||
this.logger.info(`Chrome submit form with args:`, args);
|
||||
return await this.sendToExtensions({
|
||||
action: TOOL_NAMES.BROWSER.SUBMIT_FORM,
|
||||
params: args,
|
||||
});
|
||||
}
|
||||
}
|
||||
287
app/remote-server/src/server/connection-router.ts
Normal file
287
app/remote-server/src/server/connection-router.ts
Normal file
@@ -0,0 +1,287 @@
|
||||
import { Logger } from 'pino';
|
||||
import { SessionManager, ExtensionConnection } from './session-manager.js';
|
||||
|
||||
export interface RoutingRule {
|
||||
sessionId?: string;
|
||||
userId?: string;
|
||||
priority: number;
|
||||
condition?: (connection: ExtensionConnection) => boolean;
|
||||
}
|
||||
|
||||
export interface RouteResult {
|
||||
connection: ExtensionConnection;
|
||||
sessionId: string;
|
||||
userId: string;
|
||||
routingReason: string;
|
||||
}
|
||||
|
||||
export class ConnectionRouter {
|
||||
private logger: Logger;
|
||||
private sessionManager: SessionManager;
|
||||
private routingRules: RoutingRule[] = [];
|
||||
|
||||
constructor(logger: Logger, sessionManager: SessionManager) {
|
||||
this.logger = logger;
|
||||
this.sessionManager = sessionManager;
|
||||
|
||||
// Set up default routing rules
|
||||
this.setupDefaultRoutingRules();
|
||||
}
|
||||
|
||||
/**
|
||||
* Set up default routing rules
|
||||
*/
|
||||
private setupDefaultRoutingRules(): void {
|
||||
// Rule 1: Route by exact session ID match (highest priority)
|
||||
this.addRoutingRule({
|
||||
priority: 100,
|
||||
condition: (connection: ExtensionConnection) => true, // Will be filtered by sessionId parameter
|
||||
});
|
||||
|
||||
// Rule 2: Route by user ID (medium priority)
|
||||
this.addRoutingRule({
|
||||
priority: 50,
|
||||
condition: (connection: ExtensionConnection) => connection.isActive,
|
||||
});
|
||||
|
||||
// Rule 3: Route to any active connection (lowest priority)
|
||||
this.addRoutingRule({
|
||||
priority: 10,
|
||||
condition: (connection: ExtensionConnection) => connection.isActive,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a custom routing rule
|
||||
*/
|
||||
addRoutingRule(rule: RoutingRule): void {
|
||||
this.routingRules.push(rule);
|
||||
// Sort by priority (highest first)
|
||||
this.routingRules.sort((a, b) => b.priority - a.priority);
|
||||
}
|
||||
|
||||
/**
|
||||
* Route a message to the appropriate Chrome extension connection
|
||||
*/
|
||||
routeMessage(message: any, sessionId?: string, userId?: string): RouteResult {
|
||||
this.logger.info('Routing message:', {
|
||||
action: message.action,
|
||||
sessionId,
|
||||
userId,
|
||||
messageId: message.id,
|
||||
});
|
||||
|
||||
// Try to route by session ID first
|
||||
if (sessionId) {
|
||||
const connection = this.sessionManager.getConnectionBySessionId(sessionId);
|
||||
if (connection && connection.isActive) {
|
||||
return {
|
||||
connection,
|
||||
sessionId: connection.sessionId,
|
||||
userId: connection.userId,
|
||||
routingReason: 'exact_session_match',
|
||||
};
|
||||
} else {
|
||||
this.logger.warn(`No active connection found for session: ${sessionId}`);
|
||||
}
|
||||
}
|
||||
|
||||
// Try to route by user ID
|
||||
if (userId) {
|
||||
const connection = this.sessionManager.getConnectionByUserId(userId);
|
||||
if (connection && connection.isActive) {
|
||||
return {
|
||||
connection,
|
||||
sessionId: connection.sessionId,
|
||||
userId: connection.userId,
|
||||
routingReason: 'user_id_match',
|
||||
};
|
||||
} else {
|
||||
this.logger.warn(`No active connection found for user: ${userId}`);
|
||||
}
|
||||
}
|
||||
|
||||
// Apply routing rules to find best connection
|
||||
const activeConnections = this.sessionManager.getAllActiveConnections();
|
||||
|
||||
if (activeConnections.length === 0) {
|
||||
throw new Error('No active Chrome extension connections available');
|
||||
}
|
||||
|
||||
// Apply routing rules in priority order
|
||||
for (const rule of this.routingRules) {
|
||||
const candidates = activeConnections.filter((conn) => {
|
||||
// Apply session/user filters if specified in rule
|
||||
if (rule.sessionId && conn.sessionId !== rule.sessionId) return false;
|
||||
if (rule.userId && conn.userId !== rule.userId) return false;
|
||||
|
||||
// Apply custom condition
|
||||
if (rule.condition && !rule.condition(conn)) return false;
|
||||
|
||||
return true;
|
||||
});
|
||||
|
||||
if (candidates.length > 0) {
|
||||
// Use the first candidate (could implement load balancing here)
|
||||
const selectedConnection = candidates[0];
|
||||
|
||||
return {
|
||||
connection: selectedConnection,
|
||||
sessionId: selectedConnection.sessionId,
|
||||
userId: selectedConnection.userId,
|
||||
routingReason: `rule_priority_${rule.priority}`,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: use first available active connection
|
||||
const fallbackConnection = activeConnections[0];
|
||||
return {
|
||||
connection: fallbackConnection,
|
||||
sessionId: fallbackConnection.sessionId,
|
||||
userId: fallbackConnection.userId,
|
||||
routingReason: 'fallback_first_available',
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Route a message with load balancing
|
||||
*/
|
||||
routeMessageWithLoadBalancing(message: any, sessionId?: string, userId?: string): RouteResult {
|
||||
// For session-specific requests, use exact routing
|
||||
if (sessionId || userId) {
|
||||
return this.routeMessage(message, sessionId, userId);
|
||||
}
|
||||
|
||||
// For general requests, implement round-robin load balancing
|
||||
const activeConnections = this.sessionManager.getAllActiveConnections();
|
||||
|
||||
if (activeConnections.length === 0) {
|
||||
throw new Error('No active Chrome extension connections available');
|
||||
}
|
||||
|
||||
// Simple round-robin based on message timestamp
|
||||
const index = Date.now() % activeConnections.length;
|
||||
const selectedConnection = activeConnections[index];
|
||||
|
||||
return {
|
||||
connection: selectedConnection,
|
||||
sessionId: selectedConnection.sessionId,
|
||||
userId: selectedConnection.userId,
|
||||
routingReason: 'load_balanced_round_robin',
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get routing statistics
|
||||
*/
|
||||
getRoutingStats(): any {
|
||||
const stats = this.sessionManager.getStats();
|
||||
return {
|
||||
...stats,
|
||||
routingRules: this.routingRules.length,
|
||||
routingRulesPriorities: this.routingRules.map((rule) => rule.priority),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Route message to specific connection type
|
||||
*/
|
||||
routeToConnectionType(
|
||||
message: any,
|
||||
connectionType: 'newest' | 'oldest' | 'most_active',
|
||||
): RouteResult {
|
||||
const activeConnections = this.sessionManager.getAllActiveConnections();
|
||||
|
||||
if (activeConnections.length === 0) {
|
||||
throw new Error('No active Chrome extension connections available');
|
||||
}
|
||||
|
||||
let selectedConnection: ExtensionConnection;
|
||||
|
||||
switch (connectionType) {
|
||||
case 'newest':
|
||||
selectedConnection = activeConnections.reduce((newest, current) =>
|
||||
current.connectedAt > newest.connectedAt ? current : newest,
|
||||
);
|
||||
break;
|
||||
|
||||
case 'oldest':
|
||||
selectedConnection = activeConnections.reduce((oldest, current) =>
|
||||
current.connectedAt < oldest.connectedAt ? current : oldest,
|
||||
);
|
||||
break;
|
||||
|
||||
case 'most_active':
|
||||
selectedConnection = activeConnections.reduce((mostActive, current) =>
|
||||
current.lastActivity > mostActive.lastActivity ? current : mostActive,
|
||||
);
|
||||
break;
|
||||
|
||||
default:
|
||||
selectedConnection = activeConnections[0];
|
||||
}
|
||||
|
||||
return {
|
||||
connection: selectedConnection,
|
||||
sessionId: selectedConnection.sessionId,
|
||||
userId: selectedConnection.userId,
|
||||
routingReason: `connection_type_${connectionType}`,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a specific session can handle a message type
|
||||
*/
|
||||
canSessionHandleMessage(sessionId: string, messageType: string): boolean {
|
||||
const connection = this.sessionManager.getConnectionBySessionId(sessionId);
|
||||
|
||||
if (!connection || !connection.isActive) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check if connection has been active recently
|
||||
const timeSinceActivity = Date.now() - connection.lastActivity;
|
||||
const maxInactiveTime = 5 * 60 * 1000; // 5 minutes
|
||||
|
||||
if (timeSinceActivity > maxInactiveTime) {
|
||||
this.logger.warn(`Session ${sessionId} has been inactive for ${timeSinceActivity}ms`);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Add message type specific checks here if needed
|
||||
// For now, assume all active connections can handle all message types
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get recommended session for a user
|
||||
*/
|
||||
getRecommendedSessionForUser(userId: string): string | null {
|
||||
const connection = this.sessionManager.getConnectionByUserId(userId);
|
||||
return connection ? connection.sessionId : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Cleanup inactive routing rules
|
||||
*/
|
||||
cleanupRoutingRules(): void {
|
||||
// Remove rules that reference non-existent sessions
|
||||
const validSessionIds = new Set(
|
||||
this.sessionManager.getAllActiveConnections().map((conn) => conn.sessionId),
|
||||
);
|
||||
|
||||
const initialRuleCount = this.routingRules.length;
|
||||
this.routingRules = this.routingRules.filter((rule) => {
|
||||
if (rule.sessionId && !validSessionIds.has(rule.sessionId)) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
});
|
||||
|
||||
const removedRules = initialRuleCount - this.routingRules.length;
|
||||
if (removedRules > 0) {
|
||||
this.logger.info(`Cleaned up ${removedRules} invalid routing rules`);
|
||||
}
|
||||
}
|
||||
}
|
||||
317
app/remote-server/src/server/livekit-agent-manager.ts
Normal file
317
app/remote-server/src/server/livekit-agent-manager.ts
Normal file
@@ -0,0 +1,317 @@
|
||||
import { Logger } from 'pino';
|
||||
import { spawn, ChildProcess } from 'child_process';
|
||||
import { SessionManager, ExtensionConnection } from './session-manager.js';
|
||||
import path from 'path';
|
||||
|
||||
export interface LiveKitAgentInstance {
|
||||
userId: string;
|
||||
sessionId: string;
|
||||
process: ChildProcess;
|
||||
roomName: string;
|
||||
startedAt: number;
|
||||
status: 'starting' | 'running' | 'stopping' | 'stopped' | 'error';
|
||||
pid?: number;
|
||||
}
|
||||
|
||||
export class LiveKitAgentManager {
|
||||
private logger: Logger;
|
||||
private sessionManager: SessionManager;
|
||||
private agentInstances: Map<string, LiveKitAgentInstance> = new Map(); // sessionId -> agent
|
||||
private userToAgent: Map<string, string> = new Map(); // userId -> sessionId
|
||||
private agentPath: string;
|
||||
private liveKitConfig: any;
|
||||
|
||||
constructor(logger: Logger, sessionManager: SessionManager, agentPath?: string) {
|
||||
this.logger = logger;
|
||||
this.sessionManager = sessionManager;
|
||||
this.agentPath = agentPath || path.join(process.cwd(), '../../agent-livekit');
|
||||
this.liveKitConfig = this.loadLiveKitConfig();
|
||||
}
|
||||
|
||||
private loadLiveKitConfig(): any {
|
||||
// Default LiveKit configuration
|
||||
return {
|
||||
livekit_url: process.env.LIVEKIT_URL || 'ws://localhost:7880',
|
||||
api_key: process.env.LIVEKIT_API_KEY || 'devkey',
|
||||
api_secret: process.env.LIVEKIT_API_SECRET || 'secret',
|
||||
room_prefix: 'mcp-chrome-user-',
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Start a LiveKit agent for a Chrome extension connection
|
||||
*/
|
||||
async startAgentForConnection(connection: ExtensionConnection): Promise<LiveKitAgentInstance> {
|
||||
const { userId, sessionId } = connection;
|
||||
|
||||
// Check if agent already exists for this user
|
||||
const existingSessionId = this.userToAgent.get(userId);
|
||||
if (existingSessionId && this.agentInstances.has(existingSessionId)) {
|
||||
const existingAgent = this.agentInstances.get(existingSessionId)!;
|
||||
if (existingAgent.status === 'running' || existingAgent.status === 'starting') {
|
||||
this.logger.info(`Agent already running for user ${userId}, reusing existing agent`);
|
||||
return existingAgent;
|
||||
}
|
||||
}
|
||||
|
||||
// Create room name based on user ID
|
||||
const roomName = `${this.liveKitConfig.room_prefix}${userId}`;
|
||||
|
||||
this.logger.info(
|
||||
`Starting LiveKit agent for user ${userId}, session ${sessionId}, room ${roomName}`,
|
||||
);
|
||||
|
||||
// Create agent instance record
|
||||
const agentInstance: LiveKitAgentInstance = {
|
||||
userId,
|
||||
sessionId,
|
||||
process: null as any, // Will be set below
|
||||
roomName,
|
||||
startedAt: Date.now(),
|
||||
status: 'starting',
|
||||
};
|
||||
|
||||
try {
|
||||
// Spawn the full LiveKit agent process directly
|
||||
const agentProcess = spawn(
|
||||
'python',
|
||||
[
|
||||
'livekit_agent.py',
|
||||
'start',
|
||||
'--url',
|
||||
this.liveKitConfig.livekit_url,
|
||||
'--api-key',
|
||||
this.liveKitConfig.api_key,
|
||||
'--api-secret',
|
||||
this.liveKitConfig.api_secret,
|
||||
],
|
||||
{
|
||||
cwd: this.agentPath,
|
||||
env: {
|
||||
...process.env,
|
||||
LIVEKIT_URL: this.liveKitConfig.livekit_url,
|
||||
LIVEKIT_API_KEY: this.liveKitConfig.api_key,
|
||||
LIVEKIT_API_SECRET: this.liveKitConfig.api_secret,
|
||||
MCP_SERVER_URL: 'http://localhost:3001/mcp',
|
||||
CHROME_USER_ID: userId, // Pass the user ID as environment variable
|
||||
// Voice processing optimization
|
||||
LIVEKIT_ROOM_NAME: roomName,
|
||||
OPENAI_API_KEY: process.env.OPENAI_API_KEY || '',
|
||||
DEEPGRAM_API_KEY: process.env.DEEPGRAM_API_KEY || '',
|
||||
},
|
||||
stdio: ['pipe', 'pipe', 'pipe'],
|
||||
},
|
||||
);
|
||||
|
||||
agentInstance.process = agentProcess;
|
||||
agentInstance.pid = agentProcess.pid;
|
||||
|
||||
// Set up process event handlers
|
||||
agentProcess.stdout?.on('data', (data) => {
|
||||
const output = data.toString();
|
||||
this.logger.info(`[Agent ${userId}] ${output.trim()}`);
|
||||
|
||||
// Check for successful startup
|
||||
if (
|
||||
output.includes('Agent initialized successfully') ||
|
||||
output.includes('LiveKit agent started')
|
||||
) {
|
||||
agentInstance.status = 'running';
|
||||
this.logger.info(`LiveKit agent for user ${userId} is now running`);
|
||||
}
|
||||
});
|
||||
|
||||
agentProcess.stderr?.on('data', (data) => {
|
||||
const error = data.toString();
|
||||
this.logger.error(`[Agent ${userId}] ERROR: ${error.trim()}`);
|
||||
});
|
||||
|
||||
agentProcess.on('close', (code) => {
|
||||
this.logger.info(`LiveKit agent for user ${userId} exited with code ${code}`);
|
||||
agentInstance.status = code === 0 ? 'stopped' : 'error';
|
||||
|
||||
// Clean up mappings
|
||||
this.agentInstances.delete(sessionId);
|
||||
this.userToAgent.delete(userId);
|
||||
});
|
||||
|
||||
agentProcess.on('error', (error) => {
|
||||
this.logger.error(`Failed to start LiveKit agent for user ${userId}:`, error);
|
||||
agentInstance.status = 'error';
|
||||
});
|
||||
|
||||
// Store the agent instance
|
||||
this.agentInstances.set(sessionId, agentInstance);
|
||||
this.userToAgent.set(userId, sessionId);
|
||||
|
||||
this.logger.info(
|
||||
`LiveKit agent process started for user ${userId} with PID ${agentProcess.pid}`,
|
||||
);
|
||||
|
||||
return agentInstance;
|
||||
} catch (error) {
|
||||
this.logger.error(`Error starting LiveKit agent for user ${userId}:`, error);
|
||||
agentInstance.status = 'error';
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Stop a LiveKit agent for a user
|
||||
*/
|
||||
async stopAgentForUser(userId: string): Promise<boolean> {
|
||||
const sessionId = this.userToAgent.get(userId);
|
||||
if (!sessionId) {
|
||||
this.logger.warn(`No agent found for user ${userId}`);
|
||||
return false;
|
||||
}
|
||||
|
||||
return this.stopAgentForSession(sessionId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Stop a LiveKit agent for a session
|
||||
*/
|
||||
async stopAgentForSession(sessionId: string): Promise<boolean> {
|
||||
const agentInstance = this.agentInstances.get(sessionId);
|
||||
if (!agentInstance) {
|
||||
this.logger.warn(`No agent found for session ${sessionId}`);
|
||||
return false;
|
||||
}
|
||||
|
||||
this.logger.info(
|
||||
`Stopping LiveKit agent for user ${agentInstance.userId}, session ${sessionId}`,
|
||||
);
|
||||
|
||||
agentInstance.status = 'stopping';
|
||||
|
||||
try {
|
||||
if (agentInstance.process && !agentInstance.process.killed) {
|
||||
// Try graceful shutdown first
|
||||
agentInstance.process.kill('SIGTERM');
|
||||
|
||||
// Force kill after 5 seconds if still running
|
||||
setTimeout(() => {
|
||||
if (agentInstance.process && !agentInstance.process.killed) {
|
||||
this.logger.warn(`Force killing LiveKit agent for user ${agentInstance.userId}`);
|
||||
agentInstance.process.kill('SIGKILL');
|
||||
}
|
||||
}, 5000);
|
||||
}
|
||||
|
||||
return true;
|
||||
} catch (error) {
|
||||
this.logger.error(`Error stopping LiveKit agent for user ${agentInstance.userId}:`, error);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle Chrome extension connection
|
||||
*/
|
||||
async onChromeExtensionConnected(connection: ExtensionConnection): Promise<void> {
|
||||
this.logger.info(
|
||||
`Chrome extension connected, starting LiveKit agent for user ${connection.userId}`,
|
||||
);
|
||||
|
||||
try {
|
||||
await this.startAgentForConnection(connection);
|
||||
} catch (error) {
|
||||
this.logger.error(`Failed to start LiveKit agent for Chrome connection:`, error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle Chrome extension disconnection
|
||||
*/
|
||||
async onChromeExtensionDisconnected(connection: ExtensionConnection): Promise<void> {
|
||||
this.logger.info(
|
||||
`Chrome extension disconnected, stopping LiveKit agent for user ${connection.userId}`,
|
||||
);
|
||||
|
||||
try {
|
||||
await this.stopAgentForUser(connection.userId);
|
||||
} catch (error) {
|
||||
this.logger.error(`Failed to stop LiveKit agent for Chrome disconnection:`, error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get agent instance for a user
|
||||
*/
|
||||
getAgentForUser(userId: string): LiveKitAgentInstance | null {
|
||||
const sessionId = this.userToAgent.get(userId);
|
||||
return sessionId ? this.agentInstances.get(sessionId) || null : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get agent instance for a session
|
||||
*/
|
||||
getAgentForSession(sessionId: string): LiveKitAgentInstance | null {
|
||||
return this.agentInstances.get(sessionId) || null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all active agents
|
||||
*/
|
||||
getAllActiveAgents(): LiveKitAgentInstance[] {
|
||||
return Array.from(this.agentInstances.values()).filter(
|
||||
(agent) => agent.status === 'running' || agent.status === 'starting',
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get agent statistics
|
||||
*/
|
||||
getAgentStats(): any {
|
||||
const agents = Array.from(this.agentInstances.values());
|
||||
return {
|
||||
totalAgents: agents.length,
|
||||
runningAgents: agents.filter((a) => a.status === 'running').length,
|
||||
startingAgents: agents.filter((a) => a.status === 'starting').length,
|
||||
stoppedAgents: agents.filter((a) => a.status === 'stopped').length,
|
||||
errorAgents: agents.filter((a) => a.status === 'error').length,
|
||||
agentsByUser: Object.fromEntries(this.userToAgent.entries()),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Cleanup stopped agents
|
||||
*/
|
||||
cleanupStoppedAgents(): void {
|
||||
const stoppedAgents: string[] = [];
|
||||
|
||||
for (const [sessionId, agent] of this.agentInstances.entries()) {
|
||||
if (agent.status === 'stopped' || agent.status === 'error') {
|
||||
stoppedAgents.push(sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
for (const sessionId of stoppedAgents) {
|
||||
const agent = this.agentInstances.get(sessionId);
|
||||
if (agent) {
|
||||
this.agentInstances.delete(sessionId);
|
||||
this.userToAgent.delete(agent.userId);
|
||||
this.logger.info(`Cleaned up stopped agent for user ${agent.userId}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Shutdown all agents
|
||||
*/
|
||||
async shutdownAllAgents(): Promise<void> {
|
||||
this.logger.info('Shutting down all LiveKit agents...');
|
||||
|
||||
const shutdownPromises = Array.from(this.agentInstances.keys()).map((sessionId) =>
|
||||
this.stopAgentForSession(sessionId),
|
||||
);
|
||||
|
||||
await Promise.all(shutdownPromises);
|
||||
|
||||
this.agentInstances.clear();
|
||||
this.userToAgent.clear();
|
||||
|
||||
this.logger.info('All LiveKit agents shut down');
|
||||
}
|
||||
}
|
||||
256
app/remote-server/src/server/mcp-remote-server.ts
Normal file
256
app/remote-server/src/server/mcp-remote-server.ts
Normal file
@@ -0,0 +1,256 @@
|
||||
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
|
||||
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
|
||||
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
|
||||
import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { Logger } from 'pino';
|
||||
import { ChromeTools } from './chrome-tools.js';
|
||||
import { TOOL_SCHEMAS, TOOL_NAMES } from 'chrome-mcp-shared';
|
||||
|
||||
export class MCPRemoteServer {
|
||||
private server: Server;
|
||||
private chromeTools: ChromeTools;
|
||||
private logger: Logger;
|
||||
|
||||
constructor(logger: Logger) {
|
||||
this.logger = logger;
|
||||
this.server = new Server(
|
||||
{
|
||||
name: 'mcp-chrome-remote-server',
|
||||
version: '1.0.0',
|
||||
},
|
||||
{
|
||||
capabilities: {
|
||||
tools: {},
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
this.chromeTools = new ChromeTools(logger);
|
||||
this.setupHandlers();
|
||||
}
|
||||
|
||||
// Register Chrome extension connection with session management
|
||||
registerChromeExtension(
|
||||
connection: any,
|
||||
userId?: string,
|
||||
metadata?: any,
|
||||
): { userId: string; sessionId: string; connectionId: string } {
|
||||
return this.chromeTools.registerExtension(connection, userId, metadata);
|
||||
}
|
||||
|
||||
// Unregister Chrome extension connection
|
||||
unregisterChromeExtension(connection: any): boolean {
|
||||
return this.chromeTools.unregisterExtension(connection);
|
||||
}
|
||||
|
||||
// Get session statistics
|
||||
getSessionStats(): any {
|
||||
return this.chromeTools.getSessionStats();
|
||||
}
|
||||
|
||||
// Handle responses from Chrome extension
|
||||
handleChromeResponse(data: any) {
|
||||
this.chromeTools.handleResponse(data);
|
||||
}
|
||||
|
||||
// Update Chrome extension user ID
|
||||
updateChromeExtensionUserId(connection: any, newUserId: string): any {
|
||||
return this.chromeTools.updateExtensionUserId(connection, newUserId);
|
||||
}
|
||||
|
||||
// Set user context for routing
|
||||
setUserContext(userId: string, sessionId?: string) {
|
||||
this.chromeTools.setUserContext(userId, sessionId);
|
||||
}
|
||||
|
||||
// Connect a streaming transport to the MCP server
|
||||
async connectTransport(transport: SSEServerTransport | StreamableHTTPServerTransport) {
|
||||
try {
|
||||
await this.server.connect(transport);
|
||||
this.logger.info('MCP server connected to streaming transport');
|
||||
} catch (error) {
|
||||
this.logger.error('Error connecting MCP server to transport:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private setupHandlers() {
|
||||
// List available tools
|
||||
this.server.setRequestHandler(ListToolsRequestSchema, async () => {
|
||||
return { tools: TOOL_SCHEMAS };
|
||||
});
|
||||
|
||||
// Handle tool calls
|
||||
this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
|
||||
const { name, arguments: args } = request.params;
|
||||
|
||||
this.logger.info('🔧 [MCP Server] Handling tool call:', {
|
||||
toolName: name,
|
||||
hasArgs: !!args,
|
||||
args,
|
||||
});
|
||||
|
||||
try {
|
||||
let result;
|
||||
|
||||
switch (name) {
|
||||
// Legacy tool names for backward compatibility
|
||||
case 'navigate_to_url':
|
||||
result = await this.chromeTools.navigateToUrl((args as any)?.url);
|
||||
break;
|
||||
case 'get_page_content':
|
||||
result = await this.chromeTools.getPageContent((args as any)?.selector);
|
||||
break;
|
||||
case 'click_element':
|
||||
result = await this.chromeTools.clickElement((args as any)?.selector);
|
||||
break;
|
||||
case 'fill_input':
|
||||
result = await this.chromeTools.fillInput(
|
||||
(args as any)?.selector,
|
||||
(args as any)?.value,
|
||||
);
|
||||
break;
|
||||
case 'take_screenshot':
|
||||
result = await this.chromeTools.takeScreenshot((args as any)?.fullPage);
|
||||
break;
|
||||
|
||||
// Browser automation tools matching native server
|
||||
case TOOL_NAMES.BROWSER.GET_WINDOWS_AND_TABS:
|
||||
result = await this.chromeTools.getWindowsAndTabs();
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.SEARCH_TABS_CONTENT:
|
||||
result = await this.chromeTools.searchTabsContent((args as any)?.query);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.NAVIGATE:
|
||||
result = await this.chromeTools.chromeNavigate(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.SCREENSHOT:
|
||||
result = await this.chromeTools.chromeScreenshot(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.CLOSE_TABS:
|
||||
result = await this.chromeTools.chromeCloseTabs(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.GO_BACK_OR_FORWARD:
|
||||
result = await this.chromeTools.chromeGoBackOrForward(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.WEB_FETCHER:
|
||||
result = await this.chromeTools.chromeGetWebContent(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.CLICK:
|
||||
result = await this.chromeTools.chromeClickElement(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.FILL:
|
||||
result = await this.chromeTools.chromeFillOrSelect(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.GET_INTERACTIVE_ELEMENTS:
|
||||
result = await this.chromeTools.chromeGetInteractiveElements(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.NETWORK_CAPTURE_START:
|
||||
result = await this.chromeTools.chromeNetworkCaptureStart(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.NETWORK_CAPTURE_STOP:
|
||||
result = await this.chromeTools.chromeNetworkCaptureStop(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.NETWORK_REQUEST:
|
||||
result = await this.chromeTools.chromeNetworkRequest(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.NETWORK_DEBUGGER_START:
|
||||
result = await this.chromeTools.chromeNetworkDebuggerStart(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.NETWORK_DEBUGGER_STOP:
|
||||
result = await this.chromeTools.chromeNetworkDebuggerStop(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.KEYBOARD:
|
||||
result = await this.chromeTools.chromeKeyboard(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.HISTORY:
|
||||
result = await this.chromeTools.chromeHistory(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.BOOKMARK_SEARCH:
|
||||
result = await this.chromeTools.chromeBookmarkSearch(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.BOOKMARK_ADD:
|
||||
result = await this.chromeTools.chromeBookmarkAdd(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.BOOKMARK_DELETE:
|
||||
result = await this.chromeTools.chromeBookmarkDelete(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.INJECT_SCRIPT:
|
||||
result = await this.chromeTools.chromeInjectScript(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.SEND_COMMAND_TO_INJECT_SCRIPT:
|
||||
result = await this.chromeTools.chromeSendCommandToInjectScript(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.CONSOLE:
|
||||
result = await this.chromeTools.chromeConsole(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.SEARCH_GOOGLE:
|
||||
result = await this.chromeTools.chromeSearchGoogle(args);
|
||||
break;
|
||||
case TOOL_NAMES.BROWSER.SUBMIT_FORM:
|
||||
result = await this.chromeTools.chromeSubmitForm(args);
|
||||
break;
|
||||
default:
|
||||
// Use the general tool call method for any tools not explicitly mapped
|
||||
result = await this.chromeTools.callTool(name, args);
|
||||
}
|
||||
|
||||
this.logger.info('🔧 [MCP Server] Tool call completed:', {
|
||||
toolName: name,
|
||||
hasResult: !!result,
|
||||
result,
|
||||
});
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
text: JSON.stringify(result, null, 2),
|
||||
},
|
||||
],
|
||||
};
|
||||
} catch (error) {
|
||||
this.logger.error(`🔧 [MCP Server] Error executing tool ${name}:`, error);
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
text: `Error: ${error instanceof Error ? error.message : 'Unknown error'}`,
|
||||
},
|
||||
],
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async handleMessage(message: any): Promise<any> {
|
||||
// This method will handle incoming WebSocket messages
|
||||
// and route them to the appropriate MCP server handlers
|
||||
try {
|
||||
// For now, we'll implement a simple message routing
|
||||
// In a full implementation, you'd want to properly handle the MCP protocol
|
||||
|
||||
if (message.method === 'tools/list') {
|
||||
const response = await this.server.request(
|
||||
{ method: 'tools/list', params: {} },
|
||||
ListToolsRequestSchema,
|
||||
);
|
||||
return response;
|
||||
}
|
||||
|
||||
if (message.method === 'tools/call') {
|
||||
const response = await this.server.request(
|
||||
{ method: 'tools/call', params: message.params },
|
||||
CallToolRequestSchema,
|
||||
);
|
||||
return response;
|
||||
}
|
||||
|
||||
return { error: 'Unknown method' };
|
||||
} catch (error) {
|
||||
this.logger.error('Error handling message:', error);
|
||||
return { error: error instanceof Error ? error.message : 'Unknown error' };
|
||||
}
|
||||
}
|
||||
}
|
||||
476
app/remote-server/src/server/session-manager.ts
Normal file
476
app/remote-server/src/server/session-manager.ts
Normal file
@@ -0,0 +1,476 @@
|
||||
import { Logger } from 'pino';
|
||||
import { randomUUID } from 'crypto';
|
||||
|
||||
export interface UserSession {
|
||||
userId: string;
|
||||
sessionId: string;
|
||||
connectionId: string;
|
||||
createdAt: number;
|
||||
lastActivity: number;
|
||||
metadata: {
|
||||
userAgent?: string;
|
||||
ipAddress?: string;
|
||||
extensionVersion?: string;
|
||||
[key: string]: any;
|
||||
};
|
||||
}
|
||||
|
||||
export interface ExtensionConnection {
|
||||
connection: any;
|
||||
userId: string;
|
||||
sessionId: string;
|
||||
connectionId: string;
|
||||
connectedAt: number;
|
||||
lastActivity: number;
|
||||
isActive: boolean;
|
||||
metadata: any;
|
||||
}
|
||||
|
||||
export interface PendingRequest {
|
||||
resolve: Function;
|
||||
reject: Function;
|
||||
userId: string;
|
||||
sessionId: string;
|
||||
createdAt: number;
|
||||
timeout: NodeJS.Timeout;
|
||||
}
|
||||
|
||||
export class SessionManager {
|
||||
private logger: Logger;
|
||||
private userSessions: Map<string, UserSession> = new Map();
|
||||
private extensionConnections: Map<string, ExtensionConnection> = new Map();
|
||||
private sessionToConnection: Map<string, string> = new Map();
|
||||
private userToSessions: Map<string, Set<string>> = new Map();
|
||||
private pendingRequests: Map<string, PendingRequest> = new Map();
|
||||
private cleanupInterval: NodeJS.Timeout;
|
||||
|
||||
constructor(logger: Logger) {
|
||||
this.logger = logger;
|
||||
|
||||
// Start cleanup interval for stale sessions and connections
|
||||
this.cleanupInterval = setInterval(() => {
|
||||
this.cleanupStaleConnections();
|
||||
this.cleanupExpiredRequests();
|
||||
}, 30000); // Check every 30 seconds
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a unique user ID
|
||||
*/
|
||||
generateUserId(): string {
|
||||
return `user_${randomUUID()}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a unique session ID
|
||||
*/
|
||||
generateSessionId(): string {
|
||||
return `session_${randomUUID()}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a unique connection ID
|
||||
*/
|
||||
generateConnectionId(): string {
|
||||
return `conn_${randomUUID()}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Register a new Chrome extension connection
|
||||
*/
|
||||
registerExtensionConnection(
|
||||
connection: any,
|
||||
userId?: string,
|
||||
metadata: any = {},
|
||||
): { userId: string; sessionId: string; connectionId: string } {
|
||||
const actualUserId = userId || this.generateUserId();
|
||||
const sessionId = this.generateSessionId();
|
||||
const connectionId = this.generateConnectionId();
|
||||
|
||||
// Create user session
|
||||
const userSession: UserSession = {
|
||||
userId: actualUserId,
|
||||
sessionId,
|
||||
connectionId,
|
||||
createdAt: Date.now(),
|
||||
lastActivity: Date.now(),
|
||||
metadata: {
|
||||
userAgent: metadata.userAgent,
|
||||
ipAddress: metadata.ipAddress,
|
||||
extensionVersion: metadata.extensionVersion,
|
||||
...metadata,
|
||||
},
|
||||
};
|
||||
|
||||
// Create extension connection
|
||||
const extensionConnection: ExtensionConnection = {
|
||||
connection,
|
||||
userId: actualUserId,
|
||||
sessionId,
|
||||
connectionId,
|
||||
connectedAt: Date.now(),
|
||||
lastActivity: Date.now(),
|
||||
isActive: true,
|
||||
metadata,
|
||||
};
|
||||
|
||||
// Store mappings
|
||||
this.userSessions.set(sessionId, userSession);
|
||||
this.extensionConnections.set(connectionId, extensionConnection);
|
||||
this.sessionToConnection.set(sessionId, connectionId);
|
||||
|
||||
// Track user sessions
|
||||
if (!this.userToSessions.has(actualUserId)) {
|
||||
this.userToSessions.set(actualUserId, new Set());
|
||||
}
|
||||
this.userToSessions.get(actualUserId)!.add(sessionId);
|
||||
|
||||
this.logger.info(
|
||||
`Extension registered - User: ${actualUserId}, Session: ${sessionId}, Connection: ${connectionId}`,
|
||||
);
|
||||
this.logConnectionStats();
|
||||
|
||||
return { userId: actualUserId, sessionId, connectionId };
|
||||
}
|
||||
|
||||
/**
|
||||
* Unregister a Chrome extension connection
|
||||
*/
|
||||
unregisterExtensionConnection(connection: any): boolean {
|
||||
// Find connection by reference
|
||||
let connectionToRemove: ExtensionConnection | null = null;
|
||||
let connectionId: string | null = null;
|
||||
|
||||
for (const [id, extConnection] of this.extensionConnections.entries()) {
|
||||
if (extConnection.connection === connection) {
|
||||
connectionToRemove = extConnection;
|
||||
connectionId = id;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!connectionToRemove || !connectionId) {
|
||||
this.logger.warn('Attempted to unregister unknown connection');
|
||||
return false;
|
||||
}
|
||||
|
||||
const { userId, sessionId } = connectionToRemove;
|
||||
|
||||
// Remove from all mappings
|
||||
this.extensionConnections.delete(connectionId);
|
||||
this.sessionToConnection.delete(sessionId);
|
||||
this.userSessions.delete(sessionId);
|
||||
|
||||
// Update user sessions
|
||||
const userSessions = this.userToSessions.get(userId);
|
||||
if (userSessions) {
|
||||
userSessions.delete(sessionId);
|
||||
if (userSessions.size === 0) {
|
||||
this.userToSessions.delete(userId);
|
||||
}
|
||||
}
|
||||
|
||||
// Cancel any pending requests for this session
|
||||
this.cancelPendingRequestsForSession(sessionId);
|
||||
|
||||
this.logger.info(
|
||||
`Extension unregistered - User: ${userId}, Session: ${sessionId}, Connection: ${connectionId}`,
|
||||
);
|
||||
this.logConnectionStats();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get extension connection by session ID
|
||||
*/
|
||||
getConnectionBySessionId(sessionId: string): ExtensionConnection | null {
|
||||
const connectionId = this.sessionToConnection.get(sessionId);
|
||||
if (!connectionId) {
|
||||
return null;
|
||||
}
|
||||
return this.extensionConnections.get(connectionId) || null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get extension connection by user ID (returns first active connection)
|
||||
*/
|
||||
getConnectionByUserId(userId: string): ExtensionConnection | null {
|
||||
const userSessions = this.userToSessions.get(userId);
|
||||
if (!userSessions || userSessions.size === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Find first active connection
|
||||
for (const sessionId of userSessions) {
|
||||
const connection = this.getConnectionBySessionId(sessionId);
|
||||
if (connection && connection.isActive) {
|
||||
return connection;
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all active connections
|
||||
*/
|
||||
getAllActiveConnections(): ExtensionConnection[] {
|
||||
return Array.from(this.extensionConnections.values()).filter((conn) => conn.isActive);
|
||||
}
|
||||
|
||||
/**
|
||||
* Update last activity for a session
|
||||
*/
|
||||
updateSessionActivity(sessionId: string): void {
|
||||
const session = this.userSessions.get(sessionId);
|
||||
if (session) {
|
||||
session.lastActivity = Date.now();
|
||||
}
|
||||
|
||||
const connectionId = this.sessionToConnection.get(sessionId);
|
||||
if (connectionId) {
|
||||
const connection = this.extensionConnections.get(connectionId);
|
||||
if (connection) {
|
||||
connection.lastActivity = Date.now();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Update user ID for an existing extension connection
|
||||
*/
|
||||
updateExtensionUserId(connection: any, newUserId: string): any {
|
||||
// Find the extension connection
|
||||
let targetConnection: ExtensionConnection | null = null;
|
||||
let targetConnectionId: string | null = null;
|
||||
|
||||
for (const [connectionId, extConnection] of this.extensionConnections.entries()) {
|
||||
if (extConnection.connection === connection) {
|
||||
targetConnection = extConnection;
|
||||
targetConnectionId = connectionId;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!targetConnection || !targetConnectionId) {
|
||||
this.logger.warn('Extension connection not found for user ID update');
|
||||
return null;
|
||||
}
|
||||
|
||||
const oldUserId = targetConnection.userId;
|
||||
const sessionId = targetConnection.sessionId;
|
||||
|
||||
// Update the extension connection
|
||||
targetConnection.userId = newUserId;
|
||||
targetConnection.lastActivity = Date.now();
|
||||
|
||||
// Update the user session
|
||||
const userSession = this.userSessions.get(sessionId);
|
||||
if (userSession) {
|
||||
userSession.userId = newUserId;
|
||||
userSession.lastActivity = Date.now();
|
||||
}
|
||||
|
||||
// Update user to sessions mapping
|
||||
const oldUserSessions = this.userToSessions.get(oldUserId);
|
||||
if (oldUserSessions) {
|
||||
oldUserSessions.delete(sessionId);
|
||||
if (oldUserSessions.size === 0) {
|
||||
this.userToSessions.delete(oldUserId);
|
||||
}
|
||||
}
|
||||
|
||||
if (!this.userToSessions.has(newUserId)) {
|
||||
this.userToSessions.set(newUserId, new Set());
|
||||
}
|
||||
this.userToSessions.get(newUserId)!.add(sessionId);
|
||||
|
||||
this.logger.info(`Updated extension user ID from ${oldUserId} to ${newUserId}`);
|
||||
|
||||
return {
|
||||
userId: newUserId,
|
||||
oldUserId: oldUserId,
|
||||
sessionId: sessionId,
|
||||
connectionId: targetConnectionId,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Store a pending request with session context
|
||||
*/
|
||||
storePendingRequest(
|
||||
requestId: string,
|
||||
resolve: Function,
|
||||
reject: Function,
|
||||
sessionId: string,
|
||||
timeoutMs: number = 60000,
|
||||
): void {
|
||||
const session = this.userSessions.get(sessionId);
|
||||
if (!session) {
|
||||
reject(new Error(`Session ${sessionId} not found`));
|
||||
return;
|
||||
}
|
||||
|
||||
const timeout = setTimeout(() => {
|
||||
this.pendingRequests.delete(requestId);
|
||||
reject(new Error(`Request ${requestId} timed out after ${timeoutMs}ms`));
|
||||
}, timeoutMs);
|
||||
|
||||
const pendingRequest: PendingRequest = {
|
||||
resolve,
|
||||
reject,
|
||||
userId: session.userId,
|
||||
sessionId,
|
||||
createdAt: Date.now(),
|
||||
timeout,
|
||||
};
|
||||
|
||||
this.pendingRequests.set(requestId, pendingRequest);
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve a pending request
|
||||
*/
|
||||
resolvePendingRequest(requestId: string, result: any): boolean {
|
||||
const request = this.pendingRequests.get(requestId);
|
||||
if (!request) {
|
||||
return false;
|
||||
}
|
||||
|
||||
clearTimeout(request.timeout);
|
||||
this.pendingRequests.delete(requestId);
|
||||
request.resolve(result);
|
||||
|
||||
// Update session activity
|
||||
this.updateSessionActivity(request.sessionId);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reject a pending request
|
||||
*/
|
||||
rejectPendingRequest(requestId: string, error: any): boolean {
|
||||
const request = this.pendingRequests.get(requestId);
|
||||
if (!request) {
|
||||
return false;
|
||||
}
|
||||
|
||||
clearTimeout(request.timeout);
|
||||
this.pendingRequests.delete(requestId);
|
||||
request.reject(error);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Cancel all pending requests for a session
|
||||
*/
|
||||
private cancelPendingRequestsForSession(sessionId: string): void {
|
||||
const requestsToCancel: string[] = [];
|
||||
|
||||
for (const [requestId, request] of this.pendingRequests.entries()) {
|
||||
if (request.sessionId === sessionId) {
|
||||
requestsToCancel.push(requestId);
|
||||
}
|
||||
}
|
||||
|
||||
for (const requestId of requestsToCancel) {
|
||||
this.rejectPendingRequest(requestId, new Error(`Session ${sessionId} disconnected`));
|
||||
}
|
||||
|
||||
this.logger.info(
|
||||
`Cancelled ${requestsToCancel.length} pending requests for session ${sessionId}`,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Clean up stale connections and sessions
|
||||
*/
|
||||
private cleanupStaleConnections(): void {
|
||||
const now = Date.now();
|
||||
const staleThreshold = 5 * 60 * 1000; // 5 minutes
|
||||
const connectionsToRemove: string[] = [];
|
||||
|
||||
for (const [connectionId, connection] of this.extensionConnections.entries()) {
|
||||
if (now - connection.lastActivity > staleThreshold) {
|
||||
connectionsToRemove.push(connectionId);
|
||||
}
|
||||
}
|
||||
|
||||
for (const connectionId of connectionsToRemove) {
|
||||
const connection = this.extensionConnections.get(connectionId);
|
||||
if (connection) {
|
||||
this.logger.info(`Cleaning up stale connection: ${connectionId}`);
|
||||
this.unregisterExtensionConnection(connection.connection);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Clean up expired requests
|
||||
*/
|
||||
private cleanupExpiredRequests(): void {
|
||||
const now = Date.now();
|
||||
const expiredThreshold = 2 * 60 * 1000; // 2 minutes
|
||||
const requestsToRemove: string[] = [];
|
||||
|
||||
for (const [requestId, request] of this.pendingRequests.entries()) {
|
||||
if (now - request.createdAt > expiredThreshold) {
|
||||
requestsToRemove.push(requestId);
|
||||
}
|
||||
}
|
||||
|
||||
for (const requestId of requestsToRemove) {
|
||||
this.rejectPendingRequest(requestId, new Error('Request expired'));
|
||||
}
|
||||
|
||||
if (requestsToRemove.length > 0) {
|
||||
this.logger.info(`Cleaned up ${requestsToRemove.length} expired requests`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Log connection statistics
|
||||
*/
|
||||
private logConnectionStats(): void {
|
||||
this.logger.info(
|
||||
`Connection Stats - Users: ${this.userToSessions.size}, Sessions: ${this.userSessions.size}, Connections: ${this.extensionConnections.size}, Pending Requests: ${this.pendingRequests.size}`,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get session statistics
|
||||
*/
|
||||
getStats(): any {
|
||||
return {
|
||||
totalUsers: this.userToSessions.size,
|
||||
totalSessions: this.userSessions.size,
|
||||
totalConnections: this.extensionConnections.size,
|
||||
activeConnections: this.getAllActiveConnections().length,
|
||||
pendingRequests: this.pendingRequests.size,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Cleanup resources
|
||||
*/
|
||||
destroy(): void {
|
||||
if (this.cleanupInterval) {
|
||||
clearInterval(this.cleanupInterval);
|
||||
}
|
||||
|
||||
// Cancel all pending requests
|
||||
for (const [requestId, request] of this.pendingRequests.entries()) {
|
||||
clearTimeout(request.timeout);
|
||||
request.reject(new Error('Session manager destroyed'));
|
||||
}
|
||||
|
||||
this.pendingRequests.clear();
|
||||
this.extensionConnections.clear();
|
||||
this.userSessions.clear();
|
||||
this.sessionToConnection.clear();
|
||||
this.userToSessions.clear();
|
||||
}
|
||||
}
|
||||
304
app/remote-server/src/server/user-auth.ts
Normal file
304
app/remote-server/src/server/user-auth.ts
Normal file
@@ -0,0 +1,304 @@
|
||||
import { Logger } from 'pino';
|
||||
import { randomUUID } from 'crypto';
|
||||
|
||||
export interface UserToken {
|
||||
userId: string;
|
||||
tokenId: string;
|
||||
createdAt: number;
|
||||
expiresAt: number;
|
||||
metadata: {
|
||||
userAgent?: string;
|
||||
ipAddress?: string;
|
||||
[key: string]: any;
|
||||
};
|
||||
}
|
||||
|
||||
export interface AuthResult {
|
||||
success: boolean;
|
||||
userId?: string;
|
||||
sessionId?: string;
|
||||
token?: string;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export class UserAuthManager {
|
||||
private logger: Logger;
|
||||
private userTokens: Map<string, UserToken> = new Map(); // tokenId -> UserToken
|
||||
private userSessions: Map<string, Set<string>> = new Map(); // userId -> Set<tokenId>
|
||||
private tokenCleanupInterval: NodeJS.Timeout;
|
||||
|
||||
constructor(logger: Logger) {
|
||||
this.logger = logger;
|
||||
|
||||
// Start token cleanup interval
|
||||
this.tokenCleanupInterval = setInterval(() => {
|
||||
this.cleanupExpiredTokens();
|
||||
}, 60000); // Check every minute
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a new user authentication token
|
||||
*/
|
||||
generateUserToken(metadata: any = {}): AuthResult {
|
||||
const userId = `user_${randomUUID()}`;
|
||||
const tokenId = `token_${randomUUID()}`;
|
||||
const now = Date.now();
|
||||
const expiresAt = now + (24 * 60 * 60 * 1000); // 24 hours
|
||||
|
||||
const userToken: UserToken = {
|
||||
userId,
|
||||
tokenId,
|
||||
createdAt: now,
|
||||
expiresAt,
|
||||
metadata: {
|
||||
userAgent: metadata.userAgent,
|
||||
ipAddress: metadata.ipAddress,
|
||||
...metadata
|
||||
}
|
||||
};
|
||||
|
||||
// Store token
|
||||
this.userTokens.set(tokenId, userToken);
|
||||
|
||||
// Track user sessions
|
||||
if (!this.userSessions.has(userId)) {
|
||||
this.userSessions.set(userId, new Set());
|
||||
}
|
||||
this.userSessions.get(userId)!.add(tokenId);
|
||||
|
||||
this.logger.info(`Generated user token - User: ${userId}, Token: ${tokenId}`);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
userId,
|
||||
token: tokenId,
|
||||
sessionId: `session_${userId}_${Date.now()}`
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate a user token
|
||||
*/
|
||||
validateToken(tokenId: string): AuthResult {
|
||||
const userToken = this.userTokens.get(tokenId);
|
||||
|
||||
if (!userToken) {
|
||||
return {
|
||||
success: false,
|
||||
error: 'Invalid token'
|
||||
};
|
||||
}
|
||||
|
||||
// Check if token is expired
|
||||
if (Date.now() > userToken.expiresAt) {
|
||||
this.revokeToken(tokenId);
|
||||
return {
|
||||
success: false,
|
||||
error: 'Token expired'
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
userId: userToken.userId,
|
||||
sessionId: `session_${userToken.userId}_${userToken.createdAt}`
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh a user token (extend expiration)
|
||||
*/
|
||||
refreshToken(tokenId: string): AuthResult {
|
||||
const userToken = this.userTokens.get(tokenId);
|
||||
|
||||
if (!userToken) {
|
||||
return {
|
||||
success: false,
|
||||
error: 'Invalid token'
|
||||
};
|
||||
}
|
||||
|
||||
// Extend expiration by 24 hours
|
||||
userToken.expiresAt = Date.now() + (24 * 60 * 60 * 1000);
|
||||
|
||||
this.logger.info(`Refreshed token: ${tokenId} for user: ${userToken.userId}`);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
userId: userToken.userId,
|
||||
token: tokenId,
|
||||
sessionId: `session_${userToken.userId}_${userToken.createdAt}`
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Revoke a user token
|
||||
*/
|
||||
revokeToken(tokenId: string): boolean {
|
||||
const userToken = this.userTokens.get(tokenId);
|
||||
|
||||
if (!userToken) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Remove from user sessions
|
||||
const userSessions = this.userSessions.get(userToken.userId);
|
||||
if (userSessions) {
|
||||
userSessions.delete(tokenId);
|
||||
if (userSessions.size === 0) {
|
||||
this.userSessions.delete(userToken.userId);
|
||||
}
|
||||
}
|
||||
|
||||
// Remove token
|
||||
this.userTokens.delete(tokenId);
|
||||
|
||||
this.logger.info(`Revoked token: ${tokenId} for user: ${userToken.userId}`);
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Revoke all tokens for a user
|
||||
*/
|
||||
revokeUserTokens(userId: string): number {
|
||||
const userSessions = this.userSessions.get(userId);
|
||||
|
||||
if (!userSessions) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let revokedCount = 0;
|
||||
for (const tokenId of userSessions) {
|
||||
if (this.userTokens.delete(tokenId)) {
|
||||
revokedCount++;
|
||||
}
|
||||
}
|
||||
|
||||
this.userSessions.delete(userId);
|
||||
|
||||
this.logger.info(`Revoked ${revokedCount} tokens for user: ${userId}`);
|
||||
return revokedCount;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get user information by token
|
||||
*/
|
||||
getUserInfo(tokenId: string): UserToken | null {
|
||||
return this.userTokens.get(tokenId) || null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all active tokens for a user
|
||||
*/
|
||||
getUserTokens(userId: string): UserToken[] {
|
||||
const userSessions = this.userSessions.get(userId);
|
||||
|
||||
if (!userSessions) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const tokens: UserToken[] = [];
|
||||
for (const tokenId of userSessions) {
|
||||
const token = this.userTokens.get(tokenId);
|
||||
if (token) {
|
||||
tokens.push(token);
|
||||
}
|
||||
}
|
||||
|
||||
return tokens;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract user ID from session ID
|
||||
*/
|
||||
extractUserIdFromSession(sessionId: string): string | null {
|
||||
// Session format: session_{userId}_{timestamp}
|
||||
const match = sessionId.match(/^session_(.+?)_\d+$/);
|
||||
return match ? match[1] : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create anonymous user session (no token required)
|
||||
*/
|
||||
createAnonymousSession(metadata: any = {}): AuthResult {
|
||||
const userId = `anon_${randomUUID()}`;
|
||||
const sessionId = `session_${userId}_${Date.now()}`;
|
||||
|
||||
this.logger.info(`Created anonymous session - User: ${userId}, Session: ${sessionId}`);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
userId,
|
||||
sessionId
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Clean up expired tokens
|
||||
*/
|
||||
private cleanupExpiredTokens(): void {
|
||||
const now = Date.now();
|
||||
const tokensToRemove: string[] = [];
|
||||
|
||||
for (const [tokenId, userToken] of this.userTokens.entries()) {
|
||||
if (now > userToken.expiresAt) {
|
||||
tokensToRemove.push(tokenId);
|
||||
}
|
||||
}
|
||||
|
||||
for (const tokenId of tokensToRemove) {
|
||||
this.revokeToken(tokenId);
|
||||
}
|
||||
|
||||
if (tokensToRemove.length > 0) {
|
||||
this.logger.info(`Cleaned up ${tokensToRemove.length} expired tokens`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication statistics
|
||||
*/
|
||||
getAuthStats(): any {
|
||||
return {
|
||||
totalTokens: this.userTokens.size,
|
||||
totalUsers: this.userSessions.size,
|
||||
activeTokens: Array.from(this.userTokens.values()).filter(token => Date.now() < token.expiresAt).length
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Authenticate request from headers
|
||||
*/
|
||||
authenticateRequest(headers: any): AuthResult {
|
||||
// Try to get token from Authorization header
|
||||
const authHeader = headers.authorization || headers.Authorization;
|
||||
if (authHeader && authHeader.startsWith('Bearer ')) {
|
||||
const token = authHeader.substring(7);
|
||||
return this.validateToken(token);
|
||||
}
|
||||
|
||||
// Try to get token from custom header
|
||||
const tokenHeader = headers['x-auth-token'] || headers['X-Auth-Token'];
|
||||
if (tokenHeader) {
|
||||
return this.validateToken(tokenHeader);
|
||||
}
|
||||
|
||||
// Create anonymous session if no token provided
|
||||
return this.createAnonymousSession({
|
||||
userAgent: headers['user-agent'],
|
||||
ipAddress: headers['x-forwarded-for'] || 'unknown'
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Cleanup resources
|
||||
*/
|
||||
destroy(): void {
|
||||
if (this.tokenCleanupInterval) {
|
||||
clearInterval(this.tokenCleanupInterval);
|
||||
}
|
||||
|
||||
this.userTokens.clear();
|
||||
this.userSessions.clear();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user