test: add global setup and simplify web server

Author: zhensherlock

src/services/web.ts

@@ -1,92 +1,41 @@
-import { nanoid } from 'nanoid'
 import express from 'express'
-import cors from 'cors'
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
-import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js'
 import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
-import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'
 import type { OptionsType } from '@/types'
 
 export async function webServer(server: McpServer, options: OptionsType) {
   const app = express()
-  app.use(cors())
   app.use(express.json())
 
-  const transports = {
-    streamable: {} as Record<string, StreamableHTTPServerTransport>,
-    sse: {} as Record<string, SSEServerTransport>,
-  }
   app.post('/mcp', async (req, res) => {
-    const sessionId = req.headers['mcp-session-id'] as string | undefined
-    let transport: StreamableHTTPServerTransport
-
-    if (sessionId && transports.streamable[sessionId]) {
-      transport = transports.streamable[sessionId]
-    } else if (!sessionId && isInitializeRequest(req.body)) {
-      transport = new StreamableHTTPServerTransport({
-        sessionIdGenerator: () => nanoid(),
-        onsessioninitialized: sessionId => {
-          transports.streamable[sessionId] = transport
-        },
-      })
+    const transport: StreamableHTTPServerTransport = new StreamableHTTPServerTransport({
+      sessionIdGenerator: undefined,
+    })
+    res.on('close', () => {
+      transport.close()
+      server.close()
+    })
+    await server.connect(transport)
+    await transport.handleRequest(req, res, req.body)
+  })
 
-      transport.onclose = () => {
-        if (transport.sessionId) {
-          delete transports.streamable[transport.sessionId]
-        }
-      }
-      await server.connect(transport)
-    } else {
-      res.status(400).json({
+  const handleRequest = async (req: express.Request, res: express.Response) => {
+    res.writeHead(405).end(
+      JSON.stringify({
         jsonrpc: '2.0',
         error: {
           code: -32000,
-          message: 'Bad Request: No valid session ID provided',
+          message: 'Method not allowed.',
         },
         id: null,
-      })
-      return
-    }
-
-    await transport.handleRequest(req, res, req.body)
-  })
-
-  const handleSessionRequest = async (req: express.Request, res: express.Response) => {
-    const sessionId = req.headers['mcp-session-id'] as string | undefined
-    if (!sessionId || !transports.streamable[sessionId]) {
-      res.status(400).send('Invalid or missing session ID')
-      return
-    }
-
-    const transport = transports.streamable[sessionId]
-    await transport.handleRequest(req, res)
+      }),
+    )
   }
 
-  app.get('/mcp', handleSessionRequest)
-
-  app.delete('/mcp', handleSessionRequest)
+  app.get('/mcp', handleRequest)
 
-  app.get('/sse', async (req, res) => {
-    const transport = new SSEServerTransport('/messages', res)
-    transports.sse[transport.sessionId] = transport
-
-    res.on('close', () => {
-      delete transports.sse[transport.sessionId]
-    })
-
-    await server.connect(transport)
-  })
-
-  app.post('/messages', async (req, res) => {
-    const sessionId = req.query.sessionId as string
-    const transport = transports.sse[sessionId]
-    if (transport) {
-      await transport.handlePostMessage(req, res, req.body)
-    } else {
-      res.status(400).send('No transport found for sessionId')
-    }
-  })
+  app.delete('/mcp', handleRequest)
 
   app.listen(options.port)
-  console.log(`MCP server started on port ${options.port}. SSE endpoint: /sse, streamable endpoint: /mcp`)
+  console.log(`MCP server started on port ${options.port}, streamable endpoint: /mcp`)
 }

tests/tools/watermarkJsPlus.test.ts

@@ -10,7 +10,7 @@ describe('watermarkJsPlusDocTool', () => {
     })) as { content: unknown[] }
 
     expect(res.content.length).toBeGreaterThan(0)
-  }, 20000)
+  })
 
   test('returns a "not found" response for an unrecognized input', async () => {
     const res = (await global.client.callTool({
@@ -20,5 +20,5 @@ describe('watermarkJsPlusDocTool', () => {
       },
     })) as { content: unknown[] }
     expect(res.content.length).toEqual(0)
-  }, 20000)
+  })
 })

tests/utils.ts

@@ -0,0 +1,20 @@
+export function waitForValue<T>(
+  getterFn: () => T | undefined | null,
+  checkInterval = 100,
+  timeout = 10000,
+): Promise<T> {
+  return new Promise((resolve, reject) => {
+    const start = Date.now()
+
+    const intervalId = setInterval(() => {
+      const value = getterFn()
+      if (value) {
+        clearInterval(intervalId)
+        resolve(value)
+      } else if (Date.now() - start > timeout) {
+        clearInterval(intervalId)
+        reject(new Error('Timeout waiting for value'))
+      }
+    }, checkInterval)
+  })
+}

tsconfig.json

@@ -12,6 +12,6 @@
       "@/*": ["./src/*"]
     }
   },
-  "include": ["src", "tests", "vitest.setup.ts", "vitest.config.ts"],
+  "include": ["src", "tests", "vitest.setup.ts", "vitest.config.ts", "vitest.global.ts"],
   "exclude": ["node_modules"]
 }

vitest.config.ts

@@ -6,6 +6,7 @@ export default defineConfig({
       '@': '/src',
     },
     setupFiles: ['./vitest.setup.ts'],
+    globalSetup: ['./vitest.global.ts'],
     coverage: {
       include: ['src/**/*.ts'],
     },

vitest.global.ts

@@ -0,0 +1,24 @@
+import { spawn } from 'child_process';
+import { waitForValue } from './tests/utils'
+
+export default async function setup() {
+  const webProcess = spawn('c8', ['--reporter=lcov', '--reporter=text', 'tsx', './src/index.ts', 'web'], {
+    stdio: 'pipe',
+    env: {
+      ...process.env,
+      NODE_V8_COVERAGE: './coverage/tmp',
+    },
+  })
+  let webStarted = false
+  webProcess.stdout?.on('data', async (data) => {
+    const output = data.toString()
+    if (output.includes('MCP server started')) {
+      webStarted = true
+    }
+  });
+  await waitForValue(() => webStarted)
+  return () => {
+    webProcess.kill('SIGINT')
+  }
+}
+

vitest.setup.ts

@@ -1,19 +1,26 @@
 import 'dotenv/config'
 import { Client } from '@modelcontextprotocol/sdk/client/index.js'
+import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
 
-const serverParams = new StdioClientTransport({
+const stdioClient = new StdioClientTransport({
   command: 'c8',
   args: ['--reporter=lcov', '--reporter=text', 'tsx', './src/index.ts'],
   env: {
     ...process.env,
     NODE_V8_COVERAGE: './coverage/tmp',
   },
 })
+
+const baseUrl = new URL('http://localhost:8401/mcp')
+const streamableClient = new StreamableHTTPClientTransport(
+  new URL(baseUrl)
+)
 const client = new Client({
   name: 'doc-mcp-client',
   version: '1.0.0',
 })
-await client.connect(serverParams)
+await client.connect(stdioClient)
+await client.connect(streamableClient)
 
 global.client = client