token-server.js•16.4 kB
/**
* Custom Atrax server launcher with token authentication
* This file directly executes the CLI with our token middleware added
*/
import { spawn } from 'child_process';
import createTokenMiddleware from './token-middleware.js';
import express from 'express';
import http from 'http';
import fs from 'fs';
// Get token from environment
const TOKEN = process.env.MCP_TOKEN;
if (!TOKEN) {
console.error('Error: No token provided. Set MCP_TOKEN environment variable.');
process.exit(1);
}
// Start the server with CLI
const args = ['./src/cli.ts', 'serve', '-f', './examples/inspector-config.json'];
console.log(`Starting Atrax CLI with args: ${args.join(' ')}`);
// Create a proxy server that adds token authentication
const app = express();
app.use(express.json());
// Add CORS middleware
app.use((req, res, next) => {
res.setHeader('Access-Control-Allow-Origin', '*');
res.setHeader('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept, Authorization');
res.setHeader('Access-Control-Allow-Methods', 'POST, GET, OPTIONS');
if (req.method === 'OPTIONS') {
res.status(200).end();
return;
}
// Add logging
console.log(`${req.method} ${req.path} request`);
next();
});
// Add token middleware
const tokenMiddleware = createTokenMiddleware(TOKEN);
app.use(tokenMiddleware);
// Start the proxy server
const proxyServer = http.createServer(app);
proxyServer.listen(4000, () => {
console.log('=== Token Authentication Proxy Running ===');
console.log('Proxy listening on port 4000');
console.log(`Token: ${TOKEN}`);
console.log('\nFor MCP Inspector:');
console.log('1. Set transport type to SSE');
console.log('2. Set SSE URL to: http://localhost:4000/sse');
console.log('3. Add Authorization header: Bearer ' + TOKEN);
console.log(`Alternatively, use: http://localhost:4000/sse?token=${TOKEN}`);
console.log('===========================================');
// Forward SSE requests to Atrax
app.all('*', (req, res) => {
// Server is authenticated, forward the request to Atrax
console.log(`Forwarding ${req.method} ${req.url}`);
// TODO: Implement actual proxy forwarding to Atrax
if (req.path === '/sse' || req.path === '/') {
// Both root and /sse should work for SSE connections
console.log('SSE connection established');
// Generate a session ID
const sessionId = req.query.sessionId || `session_${Date.now()}_${Math.random().toString(36).substring(2, 10)}`;
console.log(`Assigned session ID: ${sessionId}`);
res.setHeader('Connection', 'keep-alive');
res.setHeader('Cache-Control', 'no-cache');
res.setHeader('Content-Type', 'text/event-stream');
res.setHeader('Access-Control-Allow-Origin', '*');
res.setHeader('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept, Authorization');
// Send the message endpoint with session ID - format must match exactly what MCP Inspector expects
console.log(`Sending endpoint event with session ID: ${sessionId}`);
res.write(`event: endpoint\ndata: /message?sessionId=${sessionId}\n\n`);
// Store this session for message handling
app.locals.sessions = app.locals.sessions || new Map();
app.locals.sessions.set(sessionId, {
res,
messageHandler: (message) => {
console.log(`Sending message to session ${sessionId}:`, message);
res.write(`data: ${JSON.stringify(message)}\n\n`);
}
});
// Send a test event to verify SSE is working
setTimeout(() => {
console.log('Sending test event');
res.write(`event: test\ndata: {"message":"Test SSE connection is working"}\n\n`);
}, 2000);
// Send a prompt after 4 seconds to encourage next steps
setTimeout(() => {
console.log('Sending tools notification');
res.write(`data: {"jsonrpc":"2.0","method":"notifications/toolListChanged"}\n\n`);
}, 4000);
// Keep the connection open
const interval = setInterval(() => {
res.write(': ping\n\n');
}, 30000);
req.on('close', () => {
clearInterval(interval);
// Clean up the session
if (app.locals.sessions && app.locals.sessions.has(sessionId)) {
app.locals.sessions.delete(sessionId);
console.log(`Session ${sessionId} closed and removed`);
}
console.log('SSE connection closed');
});
} else if (req.path === '/message') {
// Get session ID from query
const sessionId = req.query.sessionId;
console.log(`Received message for session ${sessionId}:`, req.body);
// Set CORS headers
res.setHeader('Access-Control-Allow-Origin', '*');
res.setHeader('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept, Authorization');
// Handle OPTIONS request (CORS preflight)
if (req.method === 'OPTIONS') {
res.setHeader('Access-Control-Allow-Methods', 'POST, GET, OPTIONS');
res.status(200).end();
return;
}
// Verify the session ID
if (!sessionId) {
console.error('No sessionId provided');
return res.status(400).json({
jsonrpc: '2.0',
error: {
code: -32000,
message: 'Session ID required',
data: { details: 'No sessionId provided in request' }
},
id: req.body?.id || null
});
}
// Make sure sessions are initialized
app.locals.sessions = app.locals.sessions || new Map();
// For debugging purposes, log the current sessions
const sessionIds = Array.from(app.locals.sessions.keys());
console.log(`Current sessions: ${sessionIds.join(', ') || 'none'}`);
// Check if the session exists
if (!app.locals.sessions.has(sessionId)) {
console.error(`Session not found: ${sessionId}`);
// Create a new session for this connection
console.log(`Auto-creating session ${sessionId}`);
app.locals.sessions.set(sessionId, {
created: Date.now(),
messageHandler: (message) => {
console.log(`Cannot send message to client ${sessionId}: No active SSE connection`);
}
});
}
const session = app.locals.sessions.get(sessionId);
if (!req.body || !req.body.jsonrpc || req.body.jsonrpc !== '2.0') {
console.error('Invalid JSON-RPC request:', req.body);
res.status(400).json({
jsonrpc: '2.0',
error: {
code: -32600,
message: 'Invalid Request',
data: { details: 'Not a valid JSON-RPC 2.0 request' }
},
id: req.body?.id || null
});
return;
}
// Process the request and send response
try {
// Check method
if (req.body.method === 'initialize') {
console.log('Processing initialize request with params:', JSON.stringify(req.body.params));
// Send response to client
const response = {
jsonrpc: '2.0',
result: {
protocolVersion: req.body.params?.protocolVersion || '2024-11-05',
serverInfo: {
name: 'atrax-echo',
version: '0.1.0'
},
capabilities: {
tools: {}
}
},
id: req.body.id
};
console.log('Sending initialize response:', JSON.stringify(response));
res.json(response);
} else if (req.body.method === 'mcp.listTools') {
console.log('Processing listTools request');
const response = {
jsonrpc: '2.0',
result: {
tools: [
{
name: 'echo',
description: 'Echoes back the input message',
inputSchema: {
type: 'object',
properties: {
message: {
type: 'string',
description: 'The message to echo back',
},
},
required: ['message'],
}
}
]
},
id: req.body.id
};
res.json(response);
} else if (req.body.method === 'mcp.callTool') {
console.log('Processing callTool request');
const toolName = req.body.params?.name;
const args = req.body.params?.arguments || {};
if (toolName === 'echo') {
const response = {
jsonrpc: '2.0',
result: {
content: [
{
type: 'text',
text: args.message || 'No message provided'
}
]
},
id: req.body.id
};
res.json(response);
} else {
res.json({
jsonrpc: '2.0',
error: {
code: -32601,
message: 'Method not found',
data: { details: `Tool ${toolName} not found` }
},
id: req.body.id
});
}
} else {
console.log(`Unknown method: ${req.body.method}`);
res.json({
jsonrpc: '2.0',
error: {
code: -32601,
message: 'Method not found',
data: { details: `Method ${req.body.method} not found` }
},
id: req.body.id
});
}
} catch (error) {
console.error('Error processing message:', error);
res.status(500).json({
jsonrpc: '2.0',
error: {
code: -32603,
message: 'Internal error',
data: { details: error instanceof Error ? error.message : String(error) }
},
id: req.body?.id || null
});
}
} else {
// Handle CORS preflight for other routes
if (req.method === 'OPTIONS') {
res.setHeader('Access-Control-Allow-Origin', '*');
res.setHeader('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept, Authorization');
res.setHeader('Access-Control-Allow-Methods', 'POST, GET, OPTIONS');
res.status(200).end();
return;
}
res.status(404).json({
error: 'Not found',
path: req.path
});
}
});
});
// Add debug endpoints
app.get('/debug/sessions', (req, res) => {
if (!app.locals.sessions) {
return res.json({ sessions: "No active sessions" });
}
const sessions = Array.from(app.locals.sessions.keys()).map(id => ({
id,
created: id.split('_')[1],
}));
res.json({ sessions });
});
app.get('/debug/status', (req, res) => {
// Return connection status and information
res.json({
serverStarted: new Date(proxyServer.startTime || Date.now()).toISOString(),
sessionCount: app.locals.sessions ? app.locals.sessions.size : 0,
token: TOKEN ? TOKEN.substring(0, 8) + '...' : 'not set',
endpoints: {
sse: 'http://localhost:4000/',
message: 'http://localhost:4000/message?sessionId=[SESSION_ID]',
health: 'http://localhost:4000/health',
debug: 'http://localhost:4000/debug/sessions',
test: 'http://localhost:4000/debug/test'
},
connectionInstructions: {
headers: {
'Authorization': `Bearer ${TOKEN}`
},
queryString: `token=${TOKEN}`
}
});
});
// Add a test endpoint that runs a full sequence
app.get('/debug/test', (req, res) => {
const sessionId = `test_${Date.now()}_${Math.random().toString(36).substring(2, 10)}`;
console.log(`Running test sequence with session ID: ${sessionId}`);
// Create a full HTML response with JavaScript to test the connection
res.send(`
<!DOCTYPE html>
<html>
<head>
<title>MCP Connection Test</title>
<style>
body { font-family: sans-serif; margin: 20px; }
pre { background: #f5f5f5; padding: 10px; border-radius: 5px; }
.success { color: green; }
.error { color: red; }
.log { margin-bottom: 5px; }
</style>
</head>
<body>
<h1>MCP Connection Test</h1>
<p>Session ID: <strong>${sessionId}</strong></p>
<p>Testing connection to server with token: <strong>${TOKEN ? TOKEN.substring(0, 8) + '...' : 'not set'}</strong></p>
<h2>Results:</h2>
<div id="results"></div>
<h2>Event Stream:</h2>
<pre id="eventStream"></pre>
<script>
const results = document.getElementById('results');
const eventStream = document.getElementById('eventStream');
function log(message, isError = false) {
const div = document.createElement('div');
div.className = isError ? 'log error' : 'log success';
div.textContent = message;
results.appendChild(div);
console.log(message);
}
async function runTest() {
try {
// Step 1: Connect to SSE
log('1. Opening SSE connection...');
const eventSource = new EventSource('/sse?token=${TOKEN}');
let messageEndpoint = '';
eventSource.addEventListener('endpoint', function(e) {
messageEndpoint = e.data;
log('2. Received endpoint: ' + messageEndpoint);
// Step 3: Send initialize request
sendInitializeRequest();
});
eventSource.addEventListener('message', function(e) {
const message = e.data;
const line = document.createElement('div');
line.textContent = message;
eventStream.appendChild(line);
});
eventSource.addEventListener('error', function(e) {
log('SSE Error: ' + JSON.stringify(e), true);
});
// Step 3: Send initialize request
async function sendInitializeRequest() {
try {
log('3. Sending initialize request...');
const initResponse = await fetch(messageEndpoint, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'Authorization': 'Bearer ${TOKEN}'
},
body: JSON.stringify({
jsonrpc: '2.0',
id: 1,
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: { tools: {} },
clientInfo: { name: 'test-client', version: '1.0.0' }
}
})
});
const initResult = await initResponse.json();
log('4. Initialize response: ' + JSON.stringify(initResult));
// Step 4: Send listTools request
sendListToolsRequest();
} catch (error) {
log('Initialize request error: ' + error.message, true);
}
}
// Step 4: Send listTools request
async function sendListToolsRequest() {
try {
log('5. Sending listTools request...');
const toolsResponse = await fetch(messageEndpoint, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'Authorization': 'Bearer ${TOKEN}'
},
body: JSON.stringify({
jsonrpc: '2.0',
id: 2,
method: 'mcp.listTools'
})
});
const toolsResult = await toolsResponse.json();
log('6. Tools response: ' + JSON.stringify(toolsResult));
// Test complete
log('Test sequence completed successfully!');
} catch (error) {
log('List tools request error: ' + error.message, true);
}
}
} catch (error) {
log('Test error: ' + error.message, true);
}
}
// Start the test
runTest();
</script>
</body>
</html>
`);
});
// Add a health check endpoint
app.get('/health', (req, res) => {
res.json({ status: 'ok' });
});
// Handle shutdown
process.on('SIGINT', () => {
console.log('Shutting down...');
proxyServer.close();
process.exit(0);
});