|
| 1 | +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; |
| 2 | +import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; |
| 3 | +import type { ReadResourceResult } from "@modelcontextprotocol/sdk/types.js"; |
| 4 | +import cors from "cors"; |
| 5 | +import express, { type Request, type Response } from "express"; |
| 6 | +import fs from "node:fs/promises"; |
| 7 | +import path from "node:path"; |
| 8 | +import { z } from "zod"; |
| 9 | +import { RESOURCE_URI_META_KEY } from "../../dist/src/app"; |
| 10 | + |
| 11 | +const PORT = process.env.PORT ? parseInt(process.env.PORT, 10) : 3001; |
| 12 | +const DIST_DIR = path.join(import.meta.dirname, "dist"); |
| 13 | + |
| 14 | +// Schemas - types are derived from these using z.infer |
| 15 | +const GetCohortDataInputSchema = z.object({ |
| 16 | + metric: z |
| 17 | + .enum(["retention", "revenue", "active"]) |
| 18 | + .optional() |
| 19 | + .default("retention"), |
| 20 | + periodType: z.enum(["monthly", "weekly"]).optional().default("monthly"), |
| 21 | + cohortCount: z.number().min(3).max(24).optional().default(12), |
| 22 | + maxPeriods: z.number().min(3).max(24).optional().default(12), |
| 23 | +}); |
| 24 | + |
| 25 | +const CohortCellSchema = z.object({ |
| 26 | + cohortIndex: z.number(), |
| 27 | + periodIndex: z.number(), |
| 28 | + retention: z.number(), |
| 29 | + usersRetained: z.number(), |
| 30 | + usersOriginal: z.number(), |
| 31 | +}); |
| 32 | + |
| 33 | +const CohortRowSchema = z.object({ |
| 34 | + cohortId: z.string(), |
| 35 | + cohortLabel: z.string(), |
| 36 | + originalUsers: z.number(), |
| 37 | + cells: z.array(CohortCellSchema), |
| 38 | +}); |
| 39 | + |
| 40 | +const CohortDataSchema = z.object({ |
| 41 | + cohorts: z.array(CohortRowSchema), |
| 42 | + periods: z.array(z.string()), |
| 43 | + periodLabels: z.array(z.string()), |
| 44 | + metric: z.string(), |
| 45 | + periodType: z.string(), |
| 46 | + generatedAt: z.string(), |
| 47 | +}); |
| 48 | + |
| 49 | +// Types derived from schemas |
| 50 | +type CohortCell = z.infer<typeof CohortCellSchema>; |
| 51 | +type CohortRow = z.infer<typeof CohortRowSchema>; |
| 52 | +type CohortData = z.infer<typeof CohortDataSchema>; |
| 53 | + |
| 54 | +// Internal types (not part of API schema) |
| 55 | +interface RetentionParams { |
| 56 | + baseRetention: number; |
| 57 | + decayRate: number; |
| 58 | + floor: number; |
| 59 | + noise: number; |
| 60 | +} |
| 61 | + |
| 62 | +// Retention curve generator using exponential decay |
| 63 | +function generateRetention(period: number, params: RetentionParams): number { |
| 64 | + if (period === 0) return 1.0; |
| 65 | + |
| 66 | + const { baseRetention, decayRate, floor, noise } = params; |
| 67 | + const base = baseRetention * Math.exp(-decayRate * (period - 1)) + floor; |
| 68 | + const variation = (Math.random() - 0.5) * 2 * noise; |
| 69 | + |
| 70 | + return Math.max(0, Math.min(1, base + variation)); |
| 71 | +} |
| 72 | + |
| 73 | +// Generate cohort data |
| 74 | +function generateCohortData( |
| 75 | + metric: string, |
| 76 | + periodType: string, |
| 77 | + cohortCount: number, |
| 78 | + maxPeriods: number, |
| 79 | +): CohortData { |
| 80 | + const now = new Date(); |
| 81 | + const cohorts: CohortRow[] = []; |
| 82 | + const periods: string[] = []; |
| 83 | + const periodLabels: string[] = []; |
| 84 | + |
| 85 | + // Generate period headers |
| 86 | + for (let i = 0; i < maxPeriods; i++) { |
| 87 | + periods.push(`M${i}`); |
| 88 | + periodLabels.push(i === 0 ? "Month 0" : `Month ${i}`); |
| 89 | + } |
| 90 | + |
| 91 | + // Retention parameters vary by metric type |
| 92 | + const paramsMap: Record<string, RetentionParams> = { |
| 93 | + retention: { |
| 94 | + baseRetention: 0.75, |
| 95 | + decayRate: 0.12, |
| 96 | + floor: 0.08, |
| 97 | + noise: 0.04, |
| 98 | + }, |
| 99 | + revenue: { baseRetention: 0.7, decayRate: 0.1, floor: 0.15, noise: 0.06 }, |
| 100 | + active: { baseRetention: 0.6, decayRate: 0.18, floor: 0.05, noise: 0.05 }, |
| 101 | + }; |
| 102 | + const params = paramsMap[metric] ?? paramsMap.retention; |
| 103 | + |
| 104 | + // Generate cohorts (oldest first) |
| 105 | + for (let c = 0; c < cohortCount; c++) { |
| 106 | + const cohortDate = new Date(now); |
| 107 | + cohortDate.setMonth(cohortDate.getMonth() - (cohortCount - 1 - c)); |
| 108 | + |
| 109 | + const cohortId = `${cohortDate.getFullYear()}-${String(cohortDate.getMonth() + 1).padStart(2, "0")}`; |
| 110 | + const cohortLabel = cohortDate.toLocaleDateString("en-US", { |
| 111 | + month: "short", |
| 112 | + year: "numeric", |
| 113 | + }); |
| 114 | + |
| 115 | + // Random cohort size: 1000-5000 users |
| 116 | + const originalUsers = Math.floor(1000 + Math.random() * 4000); |
| 117 | + |
| 118 | + // Number of periods this cohort has data for (newer cohorts have fewer periods) |
| 119 | + const periodsAvailable = cohortCount - c; |
| 120 | + |
| 121 | + const cells: CohortCell[] = []; |
| 122 | + let previousRetention = 1.0; |
| 123 | + |
| 124 | + for (let p = 0; p < Math.min(periodsAvailable, maxPeriods); p++) { |
| 125 | + // Retention must decrease or stay same (with small exceptions for noise) |
| 126 | + let retention = generateRetention(p, params); |
| 127 | + retention = Math.min(retention, previousRetention + 0.02); |
| 128 | + previousRetention = retention; |
| 129 | + |
| 130 | + cells.push({ |
| 131 | + cohortIndex: c, |
| 132 | + periodIndex: p, |
| 133 | + retention, |
| 134 | + usersRetained: Math.round(originalUsers * retention), |
| 135 | + usersOriginal: originalUsers, |
| 136 | + }); |
| 137 | + } |
| 138 | + |
| 139 | + cohorts.push({ cohortId, cohortLabel, originalUsers, cells }); |
| 140 | + } |
| 141 | + |
| 142 | + return { |
| 143 | + cohorts, |
| 144 | + periods, |
| 145 | + periodLabels, |
| 146 | + metric, |
| 147 | + periodType, |
| 148 | + generatedAt: new Date().toISOString(), |
| 149 | + }; |
| 150 | +} |
| 151 | + |
| 152 | +function formatCohortSummary(data: CohortData): string { |
| 153 | + const avgRetention = data.cohorts |
| 154 | + .flatMap((c) => c.cells) |
| 155 | + .filter((cell) => cell.periodIndex > 0) |
| 156 | + .reduce((sum, cell, _, arr) => sum + cell.retention / arr.length, 0); |
| 157 | + |
| 158 | + return `Cohort Analysis: ${data.cohorts.length} cohorts, ${data.periods.length} periods |
| 159 | +Average retention: ${(avgRetention * 100).toFixed(1)}% |
| 160 | +Metric: ${data.metric}, Period: ${data.periodType}`; |
| 161 | +} |
| 162 | + |
| 163 | +const server = new McpServer({ |
| 164 | + name: "Cohort Heatmap Server", |
| 165 | + version: "1.0.0", |
| 166 | +}); |
| 167 | + |
| 168 | +// Register tool and resource |
| 169 | +{ |
| 170 | + const resourceUri = "ui://get-cohort-data/mcp-app.html"; |
| 171 | + |
| 172 | + server.registerTool( |
| 173 | + "get-cohort-data", |
| 174 | + { |
| 175 | + title: "Get Cohort Retention Data", |
| 176 | + description: |
| 177 | + "Returns cohort retention heatmap data showing customer retention over time by signup month", |
| 178 | + inputSchema: GetCohortDataInputSchema.shape, |
| 179 | + outputSchema: CohortDataSchema.shape, |
| 180 | + _meta: { [RESOURCE_URI_META_KEY]: resourceUri }, |
| 181 | + }, |
| 182 | + async ({ metric, periodType, cohortCount, maxPeriods }) => { |
| 183 | + const data = generateCohortData( |
| 184 | + metric, |
| 185 | + periodType, |
| 186 | + cohortCount, |
| 187 | + maxPeriods, |
| 188 | + ); |
| 189 | + |
| 190 | + return { |
| 191 | + content: [{ type: "text", text: formatCohortSummary(data) }], |
| 192 | + structuredContent: data, |
| 193 | + }; |
| 194 | + }, |
| 195 | + ); |
| 196 | + |
| 197 | + server.registerResource( |
| 198 | + resourceUri, |
| 199 | + resourceUri, |
| 200 | + {}, |
| 201 | + async (): Promise<ReadResourceResult> => { |
| 202 | + const html = await fs.readFile( |
| 203 | + path.join(DIST_DIR, "mcp-app.html"), |
| 204 | + "utf-8", |
| 205 | + ); |
| 206 | + |
| 207 | + return { |
| 208 | + contents: [{ uri: resourceUri, mimeType: "text/html+mcp", text: html }], |
| 209 | + }; |
| 210 | + }, |
| 211 | + ); |
| 212 | +} |
| 213 | + |
| 214 | +const app = express(); |
| 215 | +app.use(cors()); |
| 216 | +app.use(express.json()); |
| 217 | + |
| 218 | +app.post("/mcp", async (req: Request, res: Response) => { |
| 219 | + try { |
| 220 | + const transport = new StreamableHTTPServerTransport({ |
| 221 | + sessionIdGenerator: undefined, |
| 222 | + enableJsonResponse: true, |
| 223 | + }); |
| 224 | + res.on("close", () => { |
| 225 | + transport.close(); |
| 226 | + }); |
| 227 | + |
| 228 | + await server.connect(transport); |
| 229 | + |
| 230 | + await transport.handleRequest(req, res, req.body); |
| 231 | + } catch (error) { |
| 232 | + console.error("Error handling MCP request:", error); |
| 233 | + if (!res.headersSent) { |
| 234 | + res.status(500).json({ |
| 235 | + jsonrpc: "2.0", |
| 236 | + error: { code: -32603, message: "Internal server error" }, |
| 237 | + id: null, |
| 238 | + }); |
| 239 | + } |
| 240 | + } |
| 241 | +}); |
| 242 | + |
| 243 | +const httpServer = app.listen(PORT, () => { |
| 244 | + console.log(`Server listening on http://localhost:${PORT}/mcp`); |
| 245 | +}); |
| 246 | + |
| 247 | +function shutdown() { |
| 248 | + console.log("\nShutting down..."); |
| 249 | + httpServer.close(() => { |
| 250 | + console.log("Server closed"); |
| 251 | + process.exit(0); |
| 252 | + }); |
| 253 | +} |
| 254 | + |
| 255 | +process.on("SIGINT", shutdown); |
| 256 | +process.on("SIGTERM", shutdown); |
0 commit comments