main.ts

// deno-lint-ignore-file no-explicit-any
/// <reference types="npm:@types/node@22.12.0" />

import './polyfill.ts'
import http from 'node:http'
import { randomUUID } from 'node:crypto'
import { parseArgs } from '@std/cli/parse-args'
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js'
import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/sdk/types.js'
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { z } from 'zod'

import { asJson, asXml, RunCode } from './runCode.ts'
import { Buffer } from 'node:buffer'

const VERSION = '0.0.13'

// Tool injection support
const EMPTY_TOOLS_RESULT = { toolNames: [], toolSchemas: {} }

const ELICIT_RESPONSE_SCHEMA = z.object({
  action: z.enum(['accept', 'decline', 'cancel']),
  content: z.optional(z.record(z.string(), z.unknown())),
})

// Schema for tool discovery responses
const TOOL_DISCOVERY_RESPONSE_SCHEMA = {
  type: 'object',
  properties: {
    action: {
      type: 'string',
      enum: ['accept', 'decline', 'cancel'],
      description: 'Action to take with the tool discovery request',
    },
    content: {
      type: 'object',
      properties: {
        data: {
          type: 'string',
          description: 'JSON string containing tool_names and tool_schemas',
        },
      },
      description: 'Tool discovery result data',
    },
  },
  required: ['action'],
}

// Schema for tool execution responses
const TOOL_EXECUTION_RESPONSE_SCHEMA = {
  type: 'object',
  properties: {
    action: {
      type: 'string',
      enum: ['accept', 'decline', 'cancel'],
      description: 'Action to take with the tool execution request',
    },
    content: {
      type: 'object',
      properties: {
        result: {
          type: 'string',
          description: 'JSON string containing tool execution result',
        },
        error: {
          type: 'string',
          description: 'Error message if execution failed',
        },
        retry: {
          type: 'string',
          description: 'JSON string containing retry information if tool needs retry',
        },
      },
      description: 'Tool execution result or error information',
    },
  },
  required: ['action'],
}

function makeElicitationRequest(
  server: McpServer,
  message: string,
  requestedSchema: Record<string, unknown>,
): Promise<z.infer<typeof ELICIT_RESPONSE_SCHEMA>> {
  return server.server.request(
    {
      method: 'elicitation/create',
      params: { message, requestedSchema },
    },
    ELICIT_RESPONSE_SCHEMA,
  )
}

async function discoverAvailableTools(
  server: McpServer,
): Promise<{ toolNames: string[]; toolSchemas: Record<string, Record<string, unknown>> }> {
  try {
    const result = await makeElicitationRequest(
      server,
      JSON.stringify({ action: 'discover_tools' }),
      TOOL_DISCOVERY_RESPONSE_SCHEMA,
    )

    if (result.action === 'accept' && result.content?.data) {
      const data = JSON.parse(result.content.data as string)
      return {
        toolNames: data.tool_names || [],
        toolSchemas: data.tool_schemas || {},
      }
    }
  } catch (error) {
    console.warn('Tool discovery failed:', error)
  }

  return EMPTY_TOOLS_RESULT
}

function createToolExecutionCallback(
  server: McpServer,
): (message: string) => Promise<z.infer<typeof ELICIT_RESPONSE_SCHEMA>> {
  return async (message: string) => {
    try {
      return await makeElicitationRequest(
        server,
        message,
        TOOL_EXECUTION_RESPONSE_SCHEMA,
      )
    } catch (error) {
      console.error('Tool execution failed:', error)
      return {
        action: 'decline',
        content: { error: `Tool execution failed: ${error}` },
      }
    }
  }
}

export async function main() {
  const { args } = Deno
  const flags = parseArgs(Deno.args, {
    string: ['deps', 'return-mode', 'port'],
    default: { port: '3001', 'return-mode': 'xml' },
  })
  const deps = flags.deps?.split(',') ?? []
  if (args.length >= 1) {
    if (args[0] === 'stdio') {
      await runStdio(deps, flags['return-mode'])
      return
    } else if (args[0] === 'streamable_http') {
      const port = parseInt(flags.port)
      runStreamableHttp(port, deps, flags['return-mode'])
      return
    } else if (args[0] === 'example') {
      await example(deps)
      return
    } else if (args[0] === 'noop') {
      await installDeps(deps)
      return
    }
  }
  console.error(
    `\
Invalid arguments: ${args.join(' ')}

Usage: deno ... deno/main.ts [stdio|streamable_http|install_deps|noop]

options:
--port <port>             Port to run the HTTP server on (default: 3001)
--deps <deps>             Comma separated list of dependencies to install
--return-mode <xml/json>  Return mode for output data (default: xml)`,
  )
  Deno.exit(1)
}

/*
 * Create an MCP server with the `run_python_code` tool registered.
 */
function createServer(deps: string[], returnMode: string): McpServer {
  const runCode = new RunCode()
  const server = new McpServer(
    {
      name: 'MCP Run Python',
      version: VERSION,
    },
    {
      instructions: 'Call the "run_python_code" tool with the Python code to run.',
      capabilities: {
        logging: {},
      },
    },
  )

  const toolDescription = `Tool to execute Python code and return stdout, stderr, and return value.

The code may be async, and the value on the last line will be returned as the return value.

The code will be executed with Python 3.13.
`

  let setLogLevel: LoggingLevel = 'emergency'

  server.server.setRequestHandler(SetLevelRequestSchema, (request) => {
    setLogLevel = request.params.level
    return {}
  })

  server.registerTool(
    'discover_available_tools',
    {
      title: 'Discover available tools',
      description:
        'Discover what tools are available for injection into Python code execution environment. Call this before writing Python code to know which tools you can use.',
      inputSchema: {},
    },
    async () => {
      // Discover available tools via elicitation
      const { toolNames: availableTools, toolSchemas } = await discoverAvailableTools(server)

      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify(
              {
                available_tools: availableTools,
                tool_schemas: toolSchemas,
                usage_note:
                  "These tools will be available as Python functions when you run Python code. Tool names with hyphens become underscores (e.g., 'web-search' -> 'web_search').",
              },
              null,
              2,
            ),
          },
        ],
      }
    },
  )

  server.registerTool(
    'run_python_code',
    {
      title: 'Run Python code',
      description: toolDescription,
      inputSchema: {
        python_code: z.string().describe('Python code to run'),
        global_variables: z.record(z.string(), z.any()).default({}).describe(
          'Map of global variables in context when the code is executed',
        ),
      },
    },
    async ({ python_code, global_variables }: { python_code: string; global_variables: Record<string, any> }) => {
      const logPromises: Promise<void>[] = []

      // Discover available tools via elicitation
      const { toolNames: availableTools, toolSchemas } = await discoverAvailableTools(server)

      // Create elicitation callback for tool execution if tools are available
      let toolInjectionOptions
      if (availableTools.length > 0) {
        toolInjectionOptions = {
          elicitationCallback: createToolExecutionCallback(server),
          availableTools: availableTools,
          toolSchemas: toolSchemas,
        }
      }

      const result = await runCode.run(
        deps,
        (level, data) => {
          if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
            logPromises.push(server.server.sendLoggingMessage({ level, data }))
          }
        },
        { name: 'main.py', content: python_code },
        global_variables,
        returnMode !== 'xml',
        toolInjectionOptions,
      )
      await Promise.all(logPromises)
      return {
        content: [{ type: 'text', text: returnMode === 'xml' ? asXml(result) : asJson(result) }],
      }
    },
  )
  return server
}

/*
 * Define some QOL functions for both the Streamable HTTP server implementation
 */
function httpGetUrl(req: http.IncomingMessage): URL {
  return new URL(
    req.url ?? '',
    `http://${req.headers.host ?? 'unknown'}`,
  )
}

function httpGetBody(req: http.IncomingMessage): Promise<JSON> {
  // https://nodejs.org/en/learn/modules/anatomy-of-an-http-transaction#request-body
  return new Promise((resolve) => {
    const bodyParts: any[] = []
    let body
    req.on('data', (chunk) => {
      bodyParts.push(chunk)
    }).on('end', () => {
      body = Buffer.concat(bodyParts).toString()
      resolve(JSON.parse(body))
    })
  })
}

function httpSetTextResponse(res: http.ServerResponse, status: number, text: string) {
  res.setHeader('Content-Type', 'text/plain')
  res.statusCode = status
  res.end(`${text}\n`)
}

function httpSetJsonResponse(res: http.ServerResponse, status: number, text: string, code: number) {
  res.setHeader('Content-Type', 'application/json')
  res.statusCode = status
  res.write(JSON.stringify({
    jsonrpc: '2.0',
    error: {
      code: code,
      message: text,
    },
    id: null,
  }))
  res.end()
}

/*
 * Run the MCP server using the Streamable HTTP transport
 */
function runStreamableHttp(port: number, deps: string[], returnMode: string) {
  // https://github.com/modelcontextprotocol/typescript-sdk?tab=readme-ov-file#with-session-management
  const mcpServer = createServer(deps, returnMode)
  const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}

  const server = http.createServer(async (req, res) => {
    const url = httpGetUrl(req)
    let pathMatch = false
    function match(method: string, path: string): boolean {
      if (url.pathname === path) {
        pathMatch = true
        return req.method === method
      }
      return false
    }

    // Reusable handler for GET and DELETE requests
    async function handleSessionRequest() {
      const sessionId = req.headers['mcp-session-id'] as string | undefined
      if (!sessionId || !transports[sessionId]) {
        httpSetTextResponse(res, 400, 'Invalid or missing session ID')
        return
      }

      const transport = transports[sessionId]
      await transport.handleRequest(req, res)
    }

    // Handle different request methods and paths
    if (match('POST', '/mcp')) {
      // Check for existing session ID
      const sessionId = req.headers['mcp-session-id'] as string | undefined
      let transport: StreamableHTTPServerTransport

      const body = await httpGetBody(req)

      if (sessionId && transports[sessionId]) {
        // Reuse existing transport
        transport = transports[sessionId]
      } else if (!sessionId && isInitializeRequest(body)) {
        // New initialization request
        transport = new StreamableHTTPServerTransport({
          sessionIdGenerator: () => randomUUID(),
          onsessioninitialized: (sessionId) => {
            // Store the transport by session ID
            transports[sessionId] = transport
          },
        })

        // Clean up transport when closed
        transport.onclose = () => {
          if (transport.sessionId) {
            delete transports[transport.sessionId]
          }
        }

        await mcpServer.connect(transport)
      } else {
        httpSetJsonResponse(res, 400, 'Bad Request: No valid session ID provided', -32000)
        return
      }

      // Handle the request
      await transport.handleRequest(req, res, body)
    } else if (match('GET', '/mcp')) {
      // Handle server-to-client notifications
      await handleSessionRequest()
    } else if (match('DELETE', '/mcp')) {
      // Handle requests for session termination
      await handleSessionRequest()
    } else if (pathMatch) {
      httpSetTextResponse(res, 405, 'Method not allowed')
    } else {
      httpSetTextResponse(res, 404, 'Page not found')
    }
  })

  server.listen(port, () => {
    console.log(`Listening on port ${port}`)
  })
}

/*
 * Run the MCP server using the Stdio transport.
 */
async function runStdio(deps: string[], returnMode: string) {
  const mcpServer = createServer(deps, returnMode)
  const transport = new StdioServerTransport()
  await mcpServer.connect(transport)
}

/*
 * Run pyodide to download and install dependencies.
 */
async function installDeps(deps: string[]) {
  const runCode = new RunCode()
  const result = await runCode.run(
    deps,
    (level, data) => console.error(`${level}|${data}`),
  )
  if (result.status !== 'success') {
    console.error('error|Failed to install dependencies')
    Deno.exit(1)
  }
}

/*
 * Run a short example script that requires numpy.
 */
async function example(deps: string[]) {
  console.error(
    `Running example script for MCP Run Python version ${VERSION}...`,
  )
  const code = `
import numpy
a = numpy.array([1, 2, 3])
print('numpy array:', a)
a
`
  const runCode = new RunCode()
  const result = await runCode.run(
    deps,
    // use warn to avoid recursion since console.log is patched in runCode
    (level, data) => console.warn(`${level}: ${data}`),
    { name: 'example.py', content: code },
  )
  console.log('Tool return value:')
  console.log(asXml(result))
  if (result.status !== 'success') {
    Deno.exit(1)
  }
}

// list of log levels to use for level comparison
const LogLevels: LoggingLevel[] = [
  'debug',
  'info',
  'notice',
  'warning',
  'error',
  'critical',
  'alert',
  'emergency',
]

await main()