diff --git a/server.mjs b/server.mjs index 929ff2b..a0f1a27 100644 --- a/server.mjs +++ b/server.mjs @@ -350,18 +350,15 @@ async function saveImage({ dataUrl, model, prefix }) { return { filepath, filename, bytes, size: bytes.length }; } -// ---------- MCP server ---------- -const server = new Server( - { name: 'mcp-img-gen', version: '1.2.0' }, - { capabilities: { tools: {} } }, -); +// ---------- Server instances (one per transport) ---------- -server.setRequestHandler(ListToolsRequestSchema, async () => { - const models = await getImageModels(); - return { - tools: [ - { - name: 'generate_image', +// Both instances get the same handlers — each serves one transport type. +function registerHandlers(srv) { + srv.setRequestHandler(ListToolsRequestSchema, async () => { + const models = await getImageModels(); + return { + tools: [ + { name: 'generate_image', description: 'Generate an opaque (background filled) image via OpenRouter. ' + 'Aspect: square|landscape|portrait. ' + @@ -462,133 +459,82 @@ server.setRequestHandler(ListToolsRequestSchema, async () => { }, ], }; -}); + }); // end ListToolsRequestSchema -server.setRequestHandler(CallToolRequestSchema, async (req) => { - const { name, arguments: args } = req.params; - try { - if (name === 'list_image_models') { - const list = (await getImageModels()).join('\n'); - return { - content: [ - { - type: 'text', - text: `# OpenRouter image models (${MODEL_CACHE.length})\n\n${list}`, - }, - ], - }; - } - if (name === 'generate_image' || name === 'generate_image_transparent') { - const prompt = String(args?.prompt || '').trim(); - if (!prompt) return errResp('prompt required'); - const aspect = args?.aspect || 'square'; - const imageUrls = normalizeImageUrls(args?.image_urls || []); - - let model; - if (name === 'generate_image_transparent') { - model = args?.model || 'recraft/recraft-v3'; - if (!isTransparentCapable(model)) { - return errResp( - `Model ${model} is not in the transparent-capable list. ` + - `Use one of: ${TRANSPARENT_MODELS.join(', ')}`, - ); + // ---- CallTool handler ---- + srv.setRequestHandler(CallToolRequestSchema, async (req) => { + const { name, arguments: args } = req.params; + try { + if (name === 'list_image_models') { + const list = (await getImageModels()).join('\n'); + return { content: [{ type: 'text', text: `# OpenRouter image models (${MODEL_CACHE.length})\n\n${list}` }] }; + } + if (name === 'generate_image' || name === 'generate_image_transparent') { + const prompt = String(args?.prompt || '').trim(); + if (!prompt) return errResp('prompt required'); + const aspect = args?.aspect || 'square'; + const imageUrls = normalizeImageUrls(args?.image_urls || []); + let model; + if (name === 'generate_image_transparent') { + model = args?.model || 'recraft/recraft-v3'; + if (!isTransparentCapable(model)) + return errResp(`Model ${model} not transparent-capable. Use: ${TRANSPARENT_MODELS.join(', ')}`); + } else { + model = args?.model || 'google/gemini-2.5-flash-image'; } - } else { - model = args?.model || 'google/gemini-2.5-flash-image'; - } - const transparent = name === 'generate_image_transparent'; + const transparent = name === 'generate_image_transparent'; + if (imageUrls.length > 0 && !IMAGE_REF_MODELS.has(model)) + process.stderr.write(`[gen] WARN: ${imageUrls.length} refs on model that may not support them\n`); - if (imageUrls.length > 0 && !IMAGE_REF_MODELS.has(model)) { - process.stderr.write( - `[gen] WARN: ${imageUrls.length} image_ref(s) passed but model ${model} may not support refs. Will be forwarded anyway as chat-completions image_url parts or top-level image[] field.\n`, - ); - } + process.stderr.write(`[gen] req tool=${name} model=${model} aspect=${aspect} transparent=${transparent} refs=${imageUrls.length} prompt="${prompt.slice(0, 80)}"\n`); + const t0 = Date.now(); + const { dataUrl, provider, mime } = await callOpenRouter({ prompt, model, aspect, transparent, imageUrls }); + const { filepath, filename, bytes, size } = await saveImage({ dataUrl, model, prefix: transparent ? 'transparent' : 'opaque' }); + process.stderr.write(`[gen] OK provider=${provider} mime=${mime} ${Date.now() - t0}ms ${filepath} ${size}b\n`); - process.stderr.write( - `[gen] req tool=${name} model=${model} aspect=${aspect} transparent=${transparent} refs=${imageUrls.length} prompt="${prompt.slice(0, 80)}"\n`, - ); - const t0 = Date.now(); - const { dataUrl, provider, mime } = await callOpenRouter({ - prompt, model, aspect, transparent, imageUrls, - }); - const { filepath, filename, bytes, size } = await saveImage({ - dataUrl, model, prefix: transparent ? 'transparent' : 'opaque', - }); - process.stderr.write( - `[gen] OK provider=${provider} mime=${mime} ${Date.now() - t0}ms ${filepath} ${size}b\n`, - ); - // Return multipart content: image inline (base64) + text metadata. - // The image block lets MCP clients (Cursor, Claude Desktop, Figma) render - // the picture directly in the chat without an extra tool call. - // The text block exposes structured metadata for the agent / user. - const meta = { - tool: name, - model, - provider, - mime, - aspect, - transparent, - refs_used: imageUrls.length, - filepath, - filename, - bytes: size, - latency_ms: Date.now() - t0, - }; - return { - content: [ - { - type: 'image', - mimeType: mime, - data: bytes.toString('base64'), - }, - { - type: 'text', - text: JSON.stringify(meta, null, 2), - }, - ], - }; - } - if (name === 'get_image') { - const filename = String(args?.filename || '').trim(); - if (!filename) return errResp('filename required'); - // Security: prevent path escape — only flat filenames in OUTPUT_DIR. - if (filename.includes('/') || filename.includes('\\') || filename.includes('..')) { - return errResp('filename must be a flat name (no path separators)'); + return { + content: [ + { type: 'image', mimeType: mime, data: bytes.toString('base64') }, + { type: 'text', text: JSON.stringify({ + tool: name, model, provider, mime, aspect, transparent, + refs_used: imageUrls.length, filepath, filename, bytes: size, + latency_ms: Date.now() - t0, + }, null, 2) }, + ], + }; } - const filepath = path.join(OUTPUT_DIR, filename); - let buf; - try { - buf = await fs.readFile(filepath); - } catch (e) { - return errResp(`Image not found: ${filename} (container may have restarted)`); + if (name === 'get_image') { + const filename = String(args?.filename || '').trim(); + if (!filename) return errResp('filename required'); + if (!/^[\w.-]+$/.test(filename)) return errResp('filename must be alphanumeric with dots/hyphens only'); + const filepath = path.join(OUTPUT_DIR, filename); + let buf; + try { buf = await fs.readFile(filepath); } catch (e) { return errResp(`Image not found: ${filename}`); } + const mime = filename.endsWith('.png') ? 'image/png' + : filename.endsWith('.jpg') || filename.endsWith('.jpeg') ? 'image/jpeg' + : filename.endsWith('.webp') ? 'image/webp' : 'application/octet-stream'; + process.stderr.write(`[get] OK ${filename} ${buf.length}b\n`); + return { + content: [ + { type: 'image', mimeType: mime, data: buf.toString('base64') }, + { type: 'text', text: JSON.stringify({ filename, filepath, bytes: buf.length, mime }, null, 2) }, + ], + }; } - const mime = filename.endsWith('.png') - ? 'image/png' - : filename.endsWith('.jpg') || filename.endsWith('.jpeg') - ? 'image/jpeg' - : filename.endsWith('.webp') - ? 'image/webp' - : 'application/octet-stream'; - process.stderr.write(`[get] OK ${filename} ${buf.length}b\n`); - return { - content: [ - { type: 'image', mimeType: mime, data: buf.toString('base64') }, - { - type: 'text', - text: JSON.stringify({ filename, filepath, bytes: buf.length, mime }, null, 2), - }, - ], - }; + return errResp(`Unknown tool: ${name}`); + } catch (e) { + process.stderr.write(`[gen-ERR] ${e.stack || e}\n`); + return { content: [{ type: 'text', text: `Error: ${e.message}` }], isError: true }; } - return errResp(`Unknown tool: ${name}`); - } catch (e) { - process.stderr.write(`[gen-ERR] ${e.stack || e}\n`); - return { - content: [{ type: 'text', text: `Error: ${e.message}` }], - isError: true, - }; - } -}); + }); // end CallToolRequestSchema +} // end registerHandlers + +// ---------- Create MCP server instance ---------- +const mcpServer = new Server( + { name: 'mcp-img-gen', version: '1.2.0' }, + { capabilities: { tools: {} } }, +); +registerHandlers(mcpServer); function errResp(text) { return { content: [{ type: 'text', text }], isError: true }; @@ -603,7 +549,7 @@ const sseSessions = new Map(); const mcpTransport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => crypto.randomUUID(), }); -await server.connect(mcpTransport); +await mcpServer.connect(mcpTransport); const httpServer = http.createServer(async (req, res) => { const url = new URL(req.url, `http://${req.headers.host || 'localhost'}`); @@ -638,12 +584,17 @@ const httpServer = http.createServer(async (req, res) => { const sseTransport = new SSEServerTransport('/messages', res); const sessionId = sseTransport.sessionId; sseSessions.set(sessionId, sseTransport); - // Clean up on disconnect res.on('close', () => { sseSessions.delete(sessionId); process.stderr.write(`[sse] session ${sessionId} closed\n`); }); - await server.connect(sseTransport); + // Each SSE session gets its own server instance (can't reuse mcpServer) + const sseServer = new Server( + { name: 'mcp-img-gen', version: '1.2.0' }, + { capabilities: { tools: {} } }, + ); + registerHandlers(sseServer); + await sseServer.connect(sseTransport); process.stderr.write(`[sse] session ${sessionId} established\n`); return; }