Code-MCP

import pg from "pg"; import crypto from "crypto"; const { Pool } = pg; // Lazy pool - only created when getPool() is first called let _pool: pg.Pool | null = null; function getPool(): pg.Pool { if (!_pool) { if (!process.env.DATABASE_URL) { throw new Error("DATABASE_URL environment variable is required"); } _pool = new Pool({ connectionString: process.env.DATABASE_URL, ssl: process.env.NODE_ENV === "production" ? { rejectUnauthorized: false } : undefined, }); } return _pool; } export type UserRole = "admin" | "user"; export interface ApiKey { id: number; key_hash: string; // Now stores hash, not plaintext key_prefix: string; // Store prefix for display (e.g., "sk-abc12345") role: UserRole; description?: string; created_at: Date; } /** * Hash an API key using SHA-256 */ function hashKey(key: string): string { return crypto.createHash("sha256").update(key).digest("hex"); } /** * Initialize the database: create tables and bootstrap admin key */ let dbInitStarted = false; export async function initDatabase() { // Prevent multiple initializations in same process if (dbInitStarted) { console.log("Database initialization already in progress, skipping..."); return; } dbInitStarted = true; const client = await getPool().connect(); try { // Acquire advisory lock to prevent race conditions across instances await client.query("SELECT pg_advisory_lock(12345)"); console.log("Initializing database..."); // Create role enum type if not exists await client.query(` DO $$ BEGIN CREATE TYPE user_role AS ENUM ('admin', 'user'); EXCEPTION WHEN duplicate_object THEN null; END $$; `); // Create api_keys table with hash column await client.query(` CREATE TABLE IF NOT EXISTS api_keys ( id SERIAL PRIMARY KEY, key_hash TEXT NOT NULL UNIQUE, key_prefix TEXT NOT NULL, role user_role NOT NULL DEFAULT 'user', description TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); `); // Create memories table (per-user persistent storage) await client.query(` CREATE TABLE IF NOT EXISTS memories ( id SERIAL PRIMARY KEY, user_id INT REFERENCES api_keys(id) ON DELETE CASCADE, key TEXT NOT NULL, value JSONB NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, UNIQUE(user_id, key) ); `); // Create usage_logs table (analytics) await client.query(` CREATE TABLE IF NOT EXISTS usage_logs ( id SERIAL PRIMARY KEY, api_key_id INT REFERENCES api_keys(id) ON DELETE SET NULL, tool_name TEXT NOT NULL, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP, response_time_ms INT ); `); // Bootstrap Admin Key let adminKey = process.env.ADMIN_KEY; const resetAdminKey = process.env.RESET_ADMIN_KEY === "true"; // Reset admin keys if requested if (resetAdminKey) { console.log( "⚠️ RESET_ADMIN_KEY=true detected - deleting all admin keys...", ); await client.query("DELETE FROM api_keys WHERE role = 'admin'"); console.log("✅ All admin keys deleted."); } // Check if any admin key exists const existingAdmin = await client.query( "SELECT id FROM api_keys WHERE role = 'admin' LIMIT 1", ); if (existingAdmin.rows.length === 0) { // No admin key exists - generate one if not provided if (!adminKey) { adminKey = "sk-" + crypto.randomBytes(32).toString("hex"); console.log(""); console.log("═══════════════════════════════════════════════════════"); console.log(" 🔐 AUTO-GENERATED ADMIN KEY (save this!)"); console.log("═══════════════════════════════════════════════════════"); console.log(` ${adminKey}`); console.log("═══════════════════════════════════════════════════════"); console.log(" Add to Railway environment variables as ADMIN_KEY"); console.log(" This key will NOT be shown again!"); console.log("═══════════════════════════════════════════════════════"); console.log(""); } else { // First time setup with provided ADMIN_KEY console.log(""); console.log("═══════════════════════════════════════════════════════"); console.log(" 🔐 ADMIN KEY (from env var)"); console.log("═══════════════════════════════════════════════════════"); console.log(` ${adminKey}`); console.log("═══════════════════════════════════════════════════════"); console.log(""); } const keyHash = hashKey(adminKey); const keyPrefix = adminKey.substring(0, 12) + "..."; await client.query( `INSERT INTO api_keys (key_hash, key_prefix, role, description) VALUES ($1, $2, 'admin', 'Root Admin Key')`, [keyHash, keyPrefix], ); console.log("✅ Admin key created and hashed."); } else if (adminKey) { // Admin exists but ADMIN_KEY env is provided - ensure it's registered console.log(""); console.log("═══════════════════════════════════════════════════════"); console.log(" 🔐 ADMIN KEY (from env var)"); console.log("═══════════════════════════════════════════════════════"); console.log(` ${adminKey}`); console.log("═══════════════════════════════════════════════════════"); console.log(""); const keyHash = hashKey(adminKey); await client.query( `INSERT INTO api_keys (key_hash, key_prefix, role, description) VALUES ($1, $2, 'admin', 'Root Admin Key') ON CONFLICT (key_hash) DO UPDATE SET role = 'admin'`, [keyHash, adminKey.substring(0, 12) + "..."], ); console.log("✅ Admin key ensured."); } else { console.log("✅ Existing admin key found (no ADMIN_KEY env var set)."); } console.log( "✅ Database initialized: api_keys, memories, usage_logs tables ready.", ); } catch (err) { console.error("❌ Error initializing database:", err); throw err; } finally { // Release advisory lock await client.query("SELECT pg_advisory_unlock(12345)"); client.release(); } } /** * Validate an API key - hash incoming key and compare */ export async function validateApiKey(key: string): Promise<ApiKey | null> { try { const keyHash = hashKey(key); const res = await getPool().query<ApiKey>( "SELECT * FROM api_keys WHERE key_hash = $1", [keyHash], ); return res.rows[0] || null; } catch (err) { console.error("Error validating API key:", err); return null; } } /** * Create a new API key (Admin only) * Returns the plaintext key ONCE - it's never stored */ export async function createApiKey( plaintextKey: string, role: UserRole, description?: string, ): Promise<{ key: string; record: ApiKey }> { const keyHash = hashKey(plaintextKey); const keyPrefix = plaintextKey.substring(0, 10) + "..."; const res = await getPool().query<ApiKey>( `INSERT INTO api_keys (key_hash, key_prefix, role, description) VALUES ($1, $2, $3, $4) RETURNING *`, [keyHash, keyPrefix, role, description], ); // Return plaintext for display to admin, but it's not stored return { key: plaintextKey, record: res.rows[0] }; } /** * Revoke/Delete an API key by hash */ export async function revokeApiKey(plaintextKey: string): Promise<boolean> { const keyHash = hashKey(plaintextKey); const res = await getPool().query( "DELETE FROM api_keys WHERE key_hash = $1", [keyHash], ); return (res.rowCount || 0) > 0; } /** * List all keys (Admin only) - returns only prefixes, never full keys */ export async function listApiKeys(): Promise<ApiKey[]> { const res = await getPool().query<ApiKey>( "SELECT id, key_prefix, role, description, created_at FROM api_keys ORDER BY created_at DESC", ); return res.rows; } // ============================================ // Memory Persistence (Per-User) // ============================================ export interface Memory { id: number; user_id: number; key: string; value: any; created_at: Date; updated_at: Date; } export async function saveMemory( userId: number, key: string, value: any, ): Promise<Memory> { const res = await getPool().query<Memory>( `INSERT INTO memories (user_id, key, value) VALUES ($1, $2, $3) ON CONFLICT (user_id, key) DO UPDATE SET value = $3, updated_at = CURRENT_TIMESTAMP RETURNING *`, [userId, key, JSON.stringify(value)], ); return res.rows[0]; } export async function readMemory( userId: number, key: string, ): Promise<Memory | null> { const res = await getPool().query<Memory>( "SELECT * FROM memories WHERE user_id = $1 AND key = $2", [userId, key], ); return res.rows[0] || null; } export async function listMemories(userId: number): Promise<Memory[]> { const res = await getPool().query<Memory>( "SELECT * FROM memories WHERE user_id = $1 ORDER BY updated_at DESC", [userId], ); return res.rows; } export async function clearMemory( userId: number, key?: string, ): Promise<number> { if (key) { const res = await getPool().query( "DELETE FROM memories WHERE user_id = $1 AND key = $2", [userId, key], ); return res.rowCount || 0; } const res = await getPool().query("DELETE FROM memories WHERE user_id = $1", [ userId, ]); return res.rowCount || 0; } // ============================================ // Usage Logging (Analytics) // ============================================ export async function logUsage( apiKeyId: number, toolName: string, responseTimeMs?: number, ): Promise<void> { try { await getPool().query( "INSERT INTO usage_logs (api_key_id, tool_name, response_time_ms) VALUES ($1, $2, $3)", [apiKeyId, toolName, responseTimeMs], ); } catch (err) { console.error("Error logging usage:", err); } } export async function getUsageStats( apiKeyId?: number, ): Promise<{ tool_name: string; count: number }[]> { const query = apiKeyId ? "SELECT tool_name, COUNT(*) as count FROM usage_logs WHERE api_key_id = $1 GROUP BY tool_name ORDER BY count DESC" : "SELECT tool_name, COUNT(*) as count FROM usage_logs GROUP BY tool_name ORDER BY count DESC"; const res = await getPool().query<{ tool_name: string; count: string }>( query, apiKeyId ? [apiKeyId] : [], ); return res.rows.map((r) => ({ tool_name: r.tool_name, count: parseInt(r.count), })); }