Skip to content

Commit 3113a3e

Browse files
committed
adding MCP
1 parent 697796b commit 3113a3e

File tree

5 files changed

+926
-2
lines changed

5 files changed

+926
-2
lines changed

proxy/env.example

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,6 @@ DATABASE_URL=postgresql://user:password@host:port/dbname
1111
JWT_SECRET=
1212
# Google OAuth Client ID (from Google Cloud Console)
1313
GOOGLE_CLIENT_ID=
14+
# Frontend URL used by MCP to generate diagram links
15+
FRONTEND_BASE_URL=https://kennethkutyn.github.io/AmpliStack
1416

proxy/index.js

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import cors from 'cors';
66
import express from 'express';
77
import OpenAI from 'openai';
88
import pg from 'pg';
9+
import { mountMcp } from './mcp.js';
910

1011
const __dirname = path.dirname(fileURLToPath(import.meta.url));
1112
loadEnv({ path: path.join(__dirname, '.env') });
@@ -477,6 +478,10 @@ app.post('/api/diagrams/:shortCode/fork', authMiddleware(true), async (req, res)
477478
}
478479
});
479480

481+
// --- MCP server ---
482+
const FRONTEND_BASE_URL = process.env.FRONTEND_BASE_URL || 'https://kennethkutyn.github.io/AmpliStack';
483+
mountMcp(app, { openai, openaiModel: OPENAI_MODEL, baseUrl: FRONTEND_BASE_URL });
484+
480485
// --- Start server ---
481486
(async () => {
482487
await initDb();

proxy/mcp.js

Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
2+
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
3+
import { z } from 'zod';
4+
import { promisify } from 'util';
5+
import { gzip } from 'zlib';
6+
7+
const gzipAsync = promisify(gzip);
8+
9+
// Mirror of the known item catalog from the frontend config.js
10+
const KNOWN_ITEMS = {
11+
'paid-ads': 'marketing', 'email': 'marketing', 'sms': 'marketing',
12+
'push-notifications': 'marketing', 'social-media': 'marketing',
13+
'search': 'marketing', 'referral': 'marketing',
14+
'website': 'experiences', 'web-app': 'experiences', 'mobile-app': 'experiences',
15+
'ott': 'experiences', 'call-center': 'experiences', 'pos': 'experiences',
16+
'amplitude-sdk': 'sources', 'segment': 'sources', 'tealium': 'sources',
17+
'api': 'sources', 'cdp': 'sources', 'etl': 'sources', 'crm': 'sources',
18+
'amplitude-analytics': 'analysis', 'snowflake': 'analysis', 'bigquery': 'analysis',
19+
'databricks': 'analysis', 'bi': 'analysis', 's3': 'analysis', 'llm': 'analysis',
20+
'amp-gs': 'activation', 'amp-webexp': 'activation', 'amp-feaexp': 'activation',
21+
'amp-assistant': 'activation', 'braze': 'activation', 'iterable': 'activation',
22+
'salesforce': 'activation', 'hubspot': 'activation', 'marketo': 'activation',
23+
'intercom': 'activation'
24+
};
25+
26+
const VALID_LAYERS = new Set(['marketing', 'experiences', 'sources', 'analysis', 'activation']);
27+
28+
function slugify(value) {
29+
return (value || '').toString().trim().toLowerCase()
30+
.replace(/[^a-z0-9]+/g, '-').replace(/^-+|-+$/g, '') || 'node';
31+
}
32+
33+
function guessLayerFromKind(kind) {
34+
const lower = (kind || '').toString().trim().toLowerCase();
35+
if (!lower) return null;
36+
if (lower === 'activation') return 'activation';
37+
if (lower === 'warehouse' || lower === 'analysis') return 'analysis';
38+
if (lower === 'amplitude') return 'analysis';
39+
if (lower === 'datasource' || lower === 'source') return 'sources';
40+
return null;
41+
}
42+
43+
function buildConnectionKey(sourceId, targetId) {
44+
return [sourceId, targetId].sort().join('::');
45+
}
46+
47+
/**
48+
* Build a serialized diagram state from the AI response (mirrors frontend applyDiagramFromAi).
49+
*/
50+
function buildDiagramState(aiResult, title) {
51+
const nodes = Array.isArray(aiResult.diagramNodes) ? aiResult.diagramNodes : [];
52+
const edges = Array.isArray(aiResult.diagramEdges) ? aiResult.diagramEdges : [];
53+
54+
const addedItems = { marketing: [], experiences: [], sources: [], analysis: [], activation: [] };
55+
const customEntries = { marketing: [], experiences: [], sources: [], analysis: [], activation: [] };
56+
const layerOrder = { marketing: [], experiences: [], sources: [], analysis: [], activation: [] };
57+
const customConnections = [];
58+
const itemCategoryIndex = { ...KNOWN_ITEMS };
59+
60+
for (const node of nodes) {
61+
const id = node?.id || slugify(node?.label);
62+
const label = node?.label || id;
63+
const rawLayer = (node?.layer || '').toString().trim().toLowerCase();
64+
const normalizedLayer = VALID_LAYERS.has(rawLayer) ? rawLayer : null;
65+
const kindLayer = guessLayerFromKind(node?.kind);
66+
67+
// Known catalog item
68+
if (KNOWN_ITEMS[id]) {
69+
const category = KNOWN_ITEMS[id];
70+
if (!addedItems[category].includes(id)) {
71+
addedItems[category].push(id);
72+
layerOrder[category].push(id);
73+
}
74+
continue;
75+
}
76+
77+
// Custom item — resolve layer
78+
const category = normalizedLayer || kindLayer || 'analysis';
79+
itemCategoryIndex[id] = category;
80+
81+
if (!customEntries[category].some(e => e.id === id)) {
82+
customEntries[category].push({ id, name: label, icon: 'custom', isCustom: true });
83+
}
84+
if (!addedItems[category].includes(id)) {
85+
addedItems[category].push(id);
86+
layerOrder[category].push(id);
87+
}
88+
}
89+
90+
for (const edge of edges) {
91+
const sourceId = edge?.sourceId;
92+
const targetId = edge?.targetId;
93+
if (!sourceId || !targetId) continue;
94+
if (!itemCategoryIndex[sourceId] || !itemCategoryIndex[targetId]) continue;
95+
const key = buildConnectionKey(sourceId, targetId);
96+
if (!customConnections.includes(key)) {
97+
customConnections.push(key);
98+
}
99+
}
100+
101+
return {
102+
version: 1,
103+
activeCategory: 'marketing',
104+
activeModel: null,
105+
diagramTitle: title || 'MCP-generated Diagram',
106+
lastEditedAt: new Date().toISOString(),
107+
addedItems,
108+
customEntries,
109+
layerOrder,
110+
customConnections,
111+
dismissedConnections: [],
112+
dottedConnections: [],
113+
connectionAnnotations: {},
114+
amplitudeSdkSelectedBadges: [],
115+
nodeNotes: {}
116+
};
117+
}
118+
119+
async function encodeStateForUrl(state) {
120+
const json = JSON.stringify(state);
121+
const compressed = await gzipAsync(Buffer.from(json, 'utf8'));
122+
// base64url encode
123+
return compressed.toString('base64url');
124+
}
125+
126+
const SYSTEM_PROMPT = `
127+
You are an Amplitude analytics solutions architect. Given a call transcript, extract architecture-ready details for building a data/activation diagram.
128+
129+
Return ONLY valid JSON with this shape:
130+
{
131+
"architecture": { "goal": string, "scope": string },
132+
"events": [
133+
{ "name": string, "properties": [string], "notes": string }
134+
],
135+
"risks": [string],
136+
"assumptions": [string],
137+
"diagramNodes": [
138+
{ "id": string, "label": string, "layer": "marketing|experiences|sources|analysis|activation", "kind": "amplitude|warehouse|activation|custom", "notes": string }
139+
],
140+
"diagramEdges": [
141+
{ "sourceId": string, "targetId": string, "label": string }
142+
]
143+
}
144+
145+
Rules:
146+
- JSON only; no prose.
147+
- Use best-effort extraction even if partial.
148+
- If there is mention of web site, web app, or mobile app, make sure to add an appropriate node to the "owned experiences" layer.
149+
- if they mention a service or vendor that doesn't exist, you can suggest a new node and guess the most appropriate layer.
150+
- Pay special attention to mention of Amplitude SDK. If they mention Mobile or Web app, assume an AMplitude SDK will be present unless the specifically say otherwise or mention a CDP.
151+
- Prefer concise labels; derive stable ids from names (lowercase, dashes).
152+
- Map AmpliStack layers: marketing, experiences (owned surfaces/apps), sources (ingest), analysis (warehouse/BI/Amplitude), activation (destinations/engagement).
153+
- For diagramEdges, keep labels descriptive (e.g., "track events", "sync audiences").
154+
- If there is a mention of push, email, ads or similar, make sure to add the appropriate node to the "marketing channesl" layer
155+
- If something is unknown, use an empty array or empty string rather than guessing.
156+
`;
157+
158+
/**
159+
* Mount MCP Streamable HTTP endpoints on an existing Express app.
160+
*/
161+
export function mountMcp(app, { openai, openaiModel, baseUrl }) {
162+
// Stateless: create a fresh server+transport per request
163+
function createServer() {
164+
const server = new McpServer({
165+
name: 'amplistack',
166+
version: '1.0.0'
167+
});
168+
169+
server.registerTool(
170+
'create_diagram',
171+
{
172+
title: 'Create Diagram',
173+
description: 'Generate an Amplistack architecture diagram from a text description. Returns a link to the interactive diagram.',
174+
inputSchema: { description: z.string().describe('A text description of the architecture, customer setup, or call transcript to turn into a diagram') }
175+
},
176+
async ({ description }) => {
177+
if (!openai) {
178+
return { content: [{ type: 'text', text: 'Error: OpenAI API key not configured on server.' }] };
179+
}
180+
181+
// 1. Call OpenAI
182+
const completion = await openai.chat.completions.create({
183+
model: openaiModel,
184+
messages: [
185+
{ role: 'system', content: SYSTEM_PROMPT },
186+
{ role: 'user', content: `Transcript:\n${description}` }
187+
],
188+
response_format: { type: 'json_object' },
189+
temperature: 0.2
190+
});
191+
192+
const content = completion.choices?.[0]?.message?.content;
193+
if (!content) {
194+
return { content: [{ type: 'text', text: 'Error: Empty response from AI model.' }] };
195+
}
196+
197+
let aiResult;
198+
try { aiResult = JSON.parse(content); }
199+
catch { return { content: [{ type: 'text', text: 'Error: Failed to parse AI response as JSON.' }] }; }
200+
201+
// 2. Build diagram state
202+
const state = buildDiagramState(aiResult, 'AI-generated Diagram');
203+
204+
// 3. Encode as URL
205+
const encoded = await encodeStateForUrl(state);
206+
const diagramUrl = `${baseUrl}/?state=${encoded}`;
207+
208+
// 4. Build summary
209+
const nodeCount = (aiResult.diagramNodes || []).length;
210+
const edgeCount = (aiResult.diagramEdges || []).length;
211+
const goal = aiResult.architecture?.goal || '';
212+
const layers = [...new Set((aiResult.diagramNodes || []).map(n => n.layer).filter(Boolean))];
213+
214+
const summary = [
215+
`Diagram created with ${nodeCount} nodes and ${edgeCount} connections.`,
216+
goal ? `\nGoal: ${goal}` : '',
217+
layers.length ? `\nLayers used: ${layers.join(', ')}` : '',
218+
`\nView diagram: ${diagramUrl}`
219+
].join('');
220+
221+
return {
222+
content: [
223+
{ type: 'text', text: summary }
224+
]
225+
};
226+
}
227+
);
228+
229+
return server;
230+
}
231+
232+
// POST /mcp — main MCP endpoint (stateless)
233+
app.post('/mcp', async (req, res) => {
234+
try {
235+
const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined });
236+
const server = createServer();
237+
await server.connect(transport);
238+
await transport.handleRequest(req, res, req.body);
239+
} catch (error) {
240+
console.error('MCP request error:', error);
241+
if (!res.headersSent) {
242+
res.status(500).json({ jsonrpc: '2.0', error: { code: -32603, message: 'Internal server error' }, id: null });
243+
}
244+
}
245+
});
246+
247+
// GET /mcp and DELETE /mcp — not supported in stateless mode
248+
app.get('/mcp', (req, res) => {
249+
res.status(405).json({ jsonrpc: '2.0', error: { code: -32000, message: 'SSE not supported in stateless mode. Use POST.' }, id: null });
250+
});
251+
252+
app.delete('/mcp', (req, res) => {
253+
res.status(405).json({ jsonrpc: '2.0', error: { code: -32000, message: 'Session management not supported in stateless mode.' }, id: null });
254+
});
255+
256+
console.log('MCP server mounted at /mcp (Streamable HTTP, stateless)');
257+
}

0 commit comments

Comments
 (0)