cprime-supergateway 3.4.3

package/src/lib/getLogger.ts

@@ -0,0 +1,83 @@
+import util from 'node:util'
+import { Logger } from '../types.js'
+
+const defaultFormatArgs = (args: any[]) => args
+
+const log =
+  (
+    {
+      formatArgs = defaultFormatArgs,
+    }: {
+      formatArgs?: typeof defaultFormatArgs
+    } = { formatArgs: defaultFormatArgs },
+  ) =>
+  (...args: any[]) =>
+    console.log('[supergateway]', ...formatArgs(args))
+
+const logStderr =
+  (
+    {
+      formatArgs = defaultFormatArgs,
+    }: {
+      formatArgs?: typeof defaultFormatArgs
+    } = { formatArgs: defaultFormatArgs },
+  ) =>
+  (...args: any[]) =>
+    console.error('[supergateway]', ...formatArgs(args))
+
+const noneLogger: Logger = {
+  info: () => {},
+  error: () => {},
+}
+
+const infoLogger: Logger = {
+  info: log(),
+  error: logStderr(),
+}
+
+const infoLoggerStdio: Logger = {
+  info: logStderr(),
+  error: logStderr(),
+}
+
+const debugFormatArgs = (args: any[]) =>
+  args.map((arg) => {
+    if (typeof arg === 'object') {
+      return util.inspect(arg, {
+        depth: null,
+        colors: process.stderr.isTTY,
+        compact: false,
+      })
+    }
+
+    return arg
+  })
+
+const debugLogger: Logger = {
+  info: log({ formatArgs: debugFormatArgs }),
+  error: logStderr({ formatArgs: debugFormatArgs }),
+}
+
+const debugLoggerStdio: Logger = {
+  info: logStderr({ formatArgs: debugFormatArgs }),
+  error: logStderr({ formatArgs: debugFormatArgs }),
+}
+
+export const getLogger = ({
+  logLevel,
+  outputTransport,
+}: {
+  logLevel: string
+  outputTransport: string
+}): Logger => {
+  if (logLevel === 'none') {
+    return noneLogger
+  }
+
+  if (logLevel === 'debug') {
+    return outputTransport === 'stdio' ? debugLoggerStdio : debugLogger
+  }
+
+  // info logLevel
+  return outputTransport === 'stdio' ? infoLoggerStdio : infoLogger
+}

package/src/lib/getVersion.ts

@@ -0,0 +1,17 @@
+import { fileURLToPath } from 'url'
+import { join, dirname } from 'path'
+import { readFileSync } from 'fs'
+
+const __filename = fileURLToPath(import.meta.url)
+const __dirname = dirname(__filename)
+
+export function getVersion(): string {
+  try {
+    const packageJsonPath = join(__dirname, '../../package.json')
+    const packageJson = JSON.parse(readFileSync(packageJsonPath, 'utf-8'))
+    return packageJson.version || '1.0.0'
+  } catch (err) {
+    console.error('[supergateway]', 'Unable to retrieve version:', err)
+    return 'unknown'
+  }
+}

package/src/lib/headers.ts

@@ -0,0 +1,55 @@
+import { Logger } from '../types.js'
+
+const parseHeaders = ({
+  argvHeader,
+  logger,
+}: {
+  argvHeader: (string | number)[]
+  logger: Logger
+}): Record<string, string> => {
+  return argvHeader.reduce<Record<string, string>>((acc, rawHeader) => {
+    const header = `${rawHeader}`
+
+    const colonIndex = header.indexOf(':')
+    if (colonIndex === -1) {
+      logger.error(`Invalid header format: ${header}, ignoring`)
+      return acc
+    }
+
+    const key = header.slice(0, colonIndex).trim()
+    const value = header.slice(colonIndex + 1).trim()
+
+    if (!key || !value) {
+      logger.error(`Invalid header format: ${header}, ignoring`)
+      return acc
+    }
+
+    acc[key] = value
+    return acc
+  }, {})
+}
+
+export const headers = ({
+  argv,
+  logger,
+}: {
+  argv: {
+    header: (string | number)[]
+    oauth2Bearer: string | undefined
+  }
+  logger: Logger
+}): Record<string, string> => {
+  const headers = parseHeaders({
+    argvHeader: argv.header,
+    logger,
+  })
+
+  if ('oauth2Bearer' in argv) {
+    return {
+      ...headers,
+      Authorization: `Bearer ${argv.oauth2Bearer}`,
+    }
+  }
+
+  return headers
+}

package/src/lib/initMongoClient.ts

@@ -0,0 +1,10 @@
+import { MongoClient } from 'mongodb'
+
+export function initMongoClient(uri = process.env.MONGO_URI): MongoClient {
+  if (!uri) {
+    throw new Error('MONGO_URI is not set')
+  }
+
+  const client = new MongoClient(uri)
+  return client
+}

package/src/lib/mcpServerLogRepository.ts

@@ -0,0 +1,48 @@
+import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js'
+import { Collection, MongoClient } from 'mongodb'
+
+export class McpServerLogRepository {
+  private readonly collection: Collection
+
+  constructor(
+    private readonly client: MongoClient,
+    dbName = 'local',
+    collectionName = 'mcp_server_logs',
+  ) {
+    this.collection = this.client.db(dbName).collection(collectionName)
+  }
+
+  insert(data: McpServerLogDto) {
+    return this.collection.insertOne(data)
+  }
+
+  update(
+    sessionId: string,
+    update: {
+      result?: unknown
+      id: number
+      error?: { message: string; code: number }
+    },
+  ) {
+    return this.collection.updateOne(
+      { sessionId, 'data.id': update.id },
+      {
+        $set: {
+          'data.result': update.result,
+          'data.error': update.error,
+          updatedAt: new Date(),
+        },
+      },
+    )
+  }
+}
+
+export type McpServerLogDto = {
+  ip: string
+  type: 'rpc' | 'error' | 'system'
+  userId: string
+  sessionId: string
+  data: JSONRPCMessage | string
+  createdAt: Date
+  updatedAt: Date
+}

package/src/lib/onSignals.ts

@@ -0,0 +1,39 @@
+import { Logger } from '../types.js'
+
+export interface OnSignalsOptions {
+  logger: Logger
+  cleanup?: () => void
+}
+
+/**
+ * Sets up signal handlers for graceful shutdown.
+ *
+ * @param options Configuration options
+ * @param options.logger Logger instance
+ * @param options.cleanup Optional cleanup function to be called before exit
+ */
+export function onSignals(options: OnSignalsOptions): void {
+  const { logger, cleanup } = options
+
+  const handleSignal = (signal: string) => {
+    logger.info(`Caught ${signal}. Exiting...`)
+    if (cleanup) {
+      cleanup()
+    }
+    process.exit(0)
+  }
+
+  process.on('SIGINT', () => handleSignal('SIGINT'))
+
+  process.on('SIGTERM', () => handleSignal('SIGTERM'))
+
+  process.on('SIGHUP', () => handleSignal('SIGHUP'))
+
+  process.stdin.on('close', () => {
+    logger.info('stdin closed. Exiting...')
+    if (cleanup) {
+      cleanup()
+    }
+    process.exit(0)
+  })
+}

package/src/lib/serializeCorsOrigin.ts

@@ -0,0 +1,14 @@
+import type { CorsOptions } from 'cors'
+
+export const serializeCorsOrigin = ({
+  corsOrigin,
+}: {
+  corsOrigin: CorsOptions['origin']
+}) =>
+  JSON.stringify(corsOrigin, (_key, value) => {
+    if (value instanceof RegExp) {
+      return value.toString()
+    }
+
+    return value
+  })

package/src/lib/sessionAccessCounter.ts

@@ -0,0 +1,118 @@
+import { Logger } from '../types.js'
+
+export class SessionAccessCounter {
+  private sessions: Map<
+    string,
+    { accessCount: number } | { timeout: NodeJS.Timeout }
+  > = new Map()
+
+  constructor(
+    public timeout: number,
+    public cleanup: (sessionId: string) => unknown,
+    public logger: Logger,
+  ) {}
+
+  inc(sessionId: string, reason: string) {
+    this.logger.info(
+      `SessionAccessCounter.inc() ${sessionId}, caused by ${reason}`,
+    )
+
+    const session = this.sessions.get(sessionId)
+
+    if (!session) {
+      // New session
+      this.logger.info(
+        `Session access count 0 -> 1 for ${sessionId} (new session)`,
+      )
+      this.sessions.set(sessionId, { accessCount: 1 })
+      return
+    }
+
+    if ('timeout' in session) {
+      // Clear pending cleanup and reactivate
+      this.logger.info(
+        `Session access count 0 -> 1, clearing cleanup timeout for ${sessionId}`,
+      )
+      clearTimeout(session.timeout)
+      this.sessions.set(sessionId, { accessCount: 1 })
+    } else {
+      // Increment active session
+      this.logger.info(
+        `Session access count ${session.accessCount} -> ${session.accessCount + 1} for ${sessionId}`,
+      )
+      session.accessCount++
+    }
+  }
+
+  dec(sessionId: string, reason: string) {
+    this.logger.info(
+      `SessionAccessCounter.dec() ${sessionId}, caused by ${reason}`,
+    )
+
+    const session = this.sessions.get(sessionId)
+
+    if (!session) {
+      this.logger.error(
+        `Called dec() on non-existent session ${sessionId}, ignoring`,
+      )
+      return
+    }
+
+    if ('timeout' in session) {
+      this.logger.error(
+        `Called dec() on session ${sessionId} that is already pending cleanup, ignoring`,
+      )
+      return
+    }
+
+    if (session.accessCount <= 0) {
+      throw new Error(
+        `Invalid access count ${session.accessCount} for session ${sessionId}`,
+      )
+    }
+
+    session.accessCount--
+    this.logger.info(
+      `Session access count ${session.accessCount + 1} -> ${session.accessCount} for ${sessionId}`,
+    )
+
+    if (session.accessCount === 0) {
+      this.logger.info(
+        `Session access count reached 0, setting cleanup timeout for ${sessionId}`,
+      )
+
+      this.sessions.set(sessionId, {
+        timeout: setTimeout(() => {
+          this.logger.info(`Session ${sessionId} timed out, cleaning up`)
+          this.sessions.delete(sessionId)
+          this.cleanup(sessionId)
+        }, this.timeout),
+      })
+    }
+  }
+
+  clear(sessionId: string, runCleanup: boolean, reason: string) {
+    this.logger.info(
+      `SessionAccessCounter.clear() ${sessionId}, caused by ${reason}`,
+    )
+
+    const session = this.sessions.get(sessionId)
+    if (!session) {
+      this.logger.info(`Attempted to clear non-existent session ${sessionId}`)
+      return
+    }
+
+    // Clear any pending timeout
+    if ('timeout' in session) {
+      clearTimeout(session.timeout)
+    }
+
+    // Remove from tracking
+    this.sessions.delete(sessionId)
+
+    // Run cleanup if requested
+    if (runCleanup) {
+      this.cleanup(sessionId)
+    }
+  }
+}

package/src/server/websocket.ts

@@ -0,0 +1,121 @@
+import {
+  Transport,
+  TransportSendOptions,
+} from '@modelcontextprotocol/sdk/shared/transport.js'
+import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js'
+import { v4 as uuidv4 } from 'uuid'
+import { WebSocket, WebSocketServer } from 'ws'
+import { Server } from 'http'
+
+export class WebSocketServerTransport implements Transport {
+  private wss!: WebSocketServer
+  private clients: Map<string, WebSocket> = new Map()
+
+  onclose?: () => void
+  onerror?: (err: Error) => void
+  private messageHandler?: (msg: JSONRPCMessage, clientId: string) => void
+  onconnection?: (clientId: string) => void
+  ondisconnection?: (clientId: string) => void
+
+  set onmessage(handler: ((message: JSONRPCMessage) => void) | undefined) {
+    this.messageHandler = handler
+      ? (msg, clientId) => {
+          // @ts-ignore
+          if (msg.id === undefined) {
+            console.log('Broadcast message:', msg)
+            return handler(msg)
+          }
+          // @ts-ignore
+          return handler({
+            ...msg,
+            // @ts-ignore
+            id: clientId + ':' + msg.id,
+          })
+        }
+      : undefined
+  }
+
+  constructor({ path, server }: { path: string; server: Server }) {
+    this.wss = new WebSocketServer({
+      path,
+      server,
+    })
+  }
+
+  async start(): Promise<void> {
+    this.wss.on('connection', (ws: WebSocket) => {
+      const clientId = uuidv4()
+      this.clients.set(clientId, ws)
+      this.onconnection?.(clientId)
+
+      ws.on('message', (data: Buffer) => {
+        try {
+          const msg = JSON.parse(data.toString())
+          this.messageHandler?.(msg, clientId)
+        } catch (err) {
+          this.onerror?.(new Error(`Failed to parse message: ${err}`))
+        }
+      })
+
+      ws.on('close', () => {
+        this.clients.delete(clientId)
+        this.ondisconnection?.(clientId)
+      })
+
+      ws.on('error', (err: Error) => {
+        this.onerror?.(err)
+      })
+    })
+  }
+
+  async send(
+    msg: JSONRPCMessage,
+    options?: TransportSendOptions | string,
+  ): Promise<void> {
+    // decide if they passed a raw clientId (legacy) or options object
+    const clientId = typeof options === 'string' ? options : undefined
+
+    // if your protocol mangles IDs to include clientId, strip it off
+    const [cId, rawId] = clientId?.split(':') ?? []
+    if (rawId !== undefined) {
+      // @ts-ignore
+      msg.id = parseInt(rawId, 10)
+    }
+
+    const payload = JSON.stringify(msg)
+
+    if (cId) {
+      // send only to the one client
+      const ws = this.clients.get(cId)
+      if (ws?.readyState === WebSocket.OPEN) {
+        ws.send(payload)
+      } else {
+        this.clients.delete(cId)
+        this.ondisconnection?.(cId)
+      }
+    } else {
+      // broadcast to everyone
+      for (const [id, ws] of this.clients) {
+        if (ws.readyState === WebSocket.OPEN) {
+          ws.send(payload)
+        } else {
+          this.clients.delete(id)
+          this.ondisconnection?.(id)
+        }
+      }
+    }
+  }
+
+  async broadcast(msg: JSONRPCMessage): Promise<void> {
+    return this.send(msg)
+  }
+
+  async close(): Promise<void> {
+    return new Promise((resolve) => {
+      this.wss.close(() => {
+        this.clients.clear()
+        resolve()
+      })
+    })
+  }
+}