added mcp support

Author: viveknotfound

src/app.ts

@@ -6,6 +6,9 @@ import { invokeRouter } from "./routes/invoke.js";
 import { httpLoggerOptions } from "./utils/logger.js";
 import { register, httpRequestDuration, httpRequestsTotal } from "./utils/metrics.js";
 import { deployQueue } from "./deploy/queue.js";
+import { execRouter } from "./routes/exec.js";
+import { mcpRouter } from "./mcp/routes.js";
+import { startSessionReaper } from "./exec/session.js";
 
 export const app = express();
 
@@ -37,6 +40,8 @@ app.use((req, res, next) => {
 app.use(express.json());
 app.use("/deploy", deployRouter);
 app.use("/f", invokeRouter);
+app.use("/exec", execRouter)
+app.use('/mcp', mcpRouter);
 
 app.get("/metrics", async (_req, res) => {
   res.set("Content-Type", register.contentType);
@@ -57,3 +62,5 @@ app.get("/ready", (_req, res) => {
     .status(healthy ? 200 : 503)
     .json({ status: healthy ? "ready" : "not_ready", checks });
 });
+
+startSessionReaper()

src/exec/gateway.ts

@@ -3,6 +3,12 @@ import { readVsockResponse } from "../runtime/protocol.js";
 import { getVmSocket } from "../runtime/transport.js";
 import { getSession, createSession, touchSession } from "./session.js";
 import { gatewayLogger } from "../utils/logger.js";
+import { 
+  execMessageTotal, 
+  execMessageDurationSeconds, 
+  execProcessExitCode, 
+  execWorkspaceBytesWritten 
+} from "../utils/metrics.js";
 import crypto from "crypto";
 import type { Vm } from "../types/types.js";
 
@@ -61,6 +67,26 @@ export async function sendSessionMessage(
     "message sent to VM"
   );
 
-  const result = await readVsockResponse(socket, timeout, onStream);
-  return { ...result, messageId: id };
+  const startTime = process.hrtime.bigint();
+  let status = "success";
+  let result;
+
+  try {
+    result = await readVsockResponse(socket, timeout, onStream);
+    
+    if (message.type === "execute" && result.data?.exitCode !== undefined) {
+      execProcessExitCode.inc({ command: message.command, exit_code: result.data.exitCode.toString() });
+    } else if (message.type === "write_file" && result.data?.bytesWritten) {
+      execWorkspaceBytesWritten.inc(result.data.bytesWritten);
+    }
+
+    return { ...result, messageId: id };
+  } catch (err) {
+    status = "error";
+    throw err;
+  } finally {
+    const duration = Number(process.hrtime.bigint() - startTime) / 1_000_000_000;
+    execMessageDurationSeconds.observe({ type: message.type }, duration);
+    execMessageTotal.inc({ type: message.type, status });
+  }
 }

src/exec/session.ts

@@ -1,6 +1,7 @@
 import { runtimeStore } from "../runtime/store.js";
 import { cleanupVm } from "../runtime/cleanup.js";
 import { sessionLogger } from "../utils/logger.js";
+import { execSessionsActive, execSessionDurationSeconds } from "../utils/metrics.js";
 
 export interface Session {
   sessionId: string;
@@ -23,6 +24,7 @@ export function createSession(sessionId: string): Session {
     state: "creating",
   };
   sessions.set(sessionId, session);
+  execSessionsActive.inc();
   return session;
 }
 
@@ -46,6 +48,8 @@ export async function destroySession(sessionId: string): Promise<boolean> {
   }
 
   sessions.delete(sessionId);
+  execSessionsActive.dec();
+  execSessionDurationSeconds.observe((Date.now() - session.createdAt) / 1000);
   sessionLogger.info({ sessionId }, "session destroyed");
   return true;
 }

src/mcp/routes.ts

@@ -0,0 +1,51 @@
+import { Router } from "express";
+import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
+import { createMcpServer } from "./server.js";
+
+export const mcpRouter = Router();
+
+mcpRouter.use((req, res, next) => {
+  const authToken = process.env.MCP_AUTH_TOKEN;
+
+  if (!authToken) {
+    res.status(503).json({ error: "MCP_AUTH_TOKEN not configured" });
+    return;
+  }
+
+  const authHeader = req.headers.authorization;
+  if (!authHeader || authHeader !== `Bearer ${authToken}`) {
+    res.status(401).json({ error: "Unauthorized" });
+    return;
+  }
+
+  next();
+});
+
+const transports = new Map<string, SSEServerTransport>();
+
+mcpRouter.get("/", async (req, res) => {
+  const mcpSessionId = req.id as string;
+
+  const transport = new SSEServerTransport(`/mcp/messages?mcpSessionId=${mcpSessionId}`, res);
+  transports.set(mcpSessionId, transport);
+
+  const server = createMcpServer();
+  await server.connect(transport);
+
+  req.on("close", () => {
+    transports.delete(mcpSessionId);
+    server.close().catch(console.error);
+  });
+});
+
+mcpRouter.post("/messages", async (req, res) => {
+  const mcpSessionId = req.query.mcpSessionId as string;
+  const transport = transports.get(mcpSessionId);
+
+  if (!transport) {
+    res.status(404).json({ error: "Session not found or disconnected" });
+    return;
+  }
+
+  await transport.handlePostMessage(req, res);
+});

src/mcp/server.ts

@@ -0,0 +1,119 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import { sendSessionMessage } from "../exec/gateway.js";
+import { destroySession } from "../exec/session.js";
+
+
+export function createMcpServer(): McpServer {
+  const server = new McpServer({
+    name: "firecracker-sandbox",
+    version: "1.0.0",
+  });
+
+  server.tool(
+    "execute",
+    "Execute a command inside an isolated Firecracker microVM. " +
+    "The workspace persists across calls within the same sessionId.",
+    {
+      sessionId: z.string().describe("Session identifier for workspace persistence"),
+      command: z.string().describe("Command to run: node, python3, bash, etc."),
+      args: z.array(z.string()).optional().describe("Command arguments"),
+      cwd: z.string().optional().describe("Working directory relative to /workspace"),
+      timeout: z.number().optional().describe("Timeout in milliseconds (default 30000)"),
+    },
+    async ({ sessionId, command, args, cwd, timeout }) => {
+      const parts: string[] = [];
+
+      const result = await sendSessionMessage(
+        sessionId,
+        { type: "execute", command, args, cwd, timeout },
+        (chunk) => {
+          parts.push(`[${chunk.stream}] ${chunk.data}`);
+        },
+      );
+
+      const exitCode = result.data?.exitCode ?? -1;
+      parts.push(`\n--- exit code: ${exitCode} ---`);
+
+      return {
+        content: [{ type: "text", text: parts.join("") }],
+        isError: exitCode !== 0,
+      };
+    }
+  );
+
+  server.tool(
+    "write_file",
+    "Write a file to the session workspace.",
+    {
+      sessionId: z.string(),
+      path: z.string().describe("File path relative to /workspace"),
+      content: z.string().describe("File content (will be base64-encoded automatically)"),
+    },
+    async ({ sessionId, path, content }) => {
+      const encoded = Buffer.from(content).toString("base64");
+      const result = await sendSessionMessage(sessionId, {
+        type: "write_file", path, content: encoded,
+      });
+      return {
+        content: [{ type: "text", text: `Wrote ${result.data?.bytesWritten} bytes to ${path}` }],
+      };
+    }
+  );
+
+  server.tool(
+    "read_file",
+    "Read a file from the session workspace.",
+    {
+      sessionId: z.string(),
+      path: z.string().describe("File path relative to /workspace"),
+    },
+    async ({ sessionId, path }) => {
+      const result = await sendSessionMessage(sessionId, {
+        type: "read_file", path,
+      });
+      const content = Buffer.from(result.data?.content || "", "base64").toString("utf-8");
+      return {
+        content: [{ type: "text", text: content }],
+      };
+    }
+  );
+
+  server.tool(
+    "list_files",
+    "List files in the session workspace.",
+    {
+      sessionId: z.string(),
+      path: z.string().optional().describe("Directory path relative to /workspace"),
+      recursive: z.boolean().optional().describe("List recursively"),
+    },
+    async ({ sessionId, path, recursive }) => {
+      const result = await sendSessionMessage(sessionId, {
+        type: "list_files", path, recursive,
+      });
+      const listing = (result.data?.files || [])
+        .map((f: any) => `${f.type === "dir" ? "📁" : "📄"} ${f.path} (${f.size}b)`)
+        .join("\n");
+      return {
+        content: [{ type: "text", text: listing || "(empty)" }],
+      };
+    }
+  );
+
+  server.tool(
+    "reset_session",
+    "Destroy a session and its VM. The workspace is lost.",
+    { sessionId: z.string() },
+    async ({ sessionId }) => {
+      const destroyed = await destroySession(sessionId);
+      return {
+        content: [{
+          type: "text",
+          text: destroyed ? "Session destroyed." : "No active session found.",
+        }],
+      };
+    }
+  );
+
+  return server;
+}

src/utils/metrics.ts

@@ -165,3 +165,44 @@ export const schedulerQueueWaitTime = new Histogram({
   buckets: [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10],
   registers: [register],
 });
+
+export const execSessionsActive = new Gauge({
+  name: "exec_sessions_active",
+  help: "Active sessions",
+  registers: [register],
+});
+
+export const execSessionDurationSeconds = new Histogram({
+  name: "exec_session_duration_seconds",
+  help: "Session lifetime",
+  buckets: [1, 5, 30, 60, 300, 900, 1800, 3600],
+  registers: [register],
+});
+
+export const execMessageTotal = new Counter({
+  name: "exec_message_total",
+  help: "Messages by type + status",
+  labelNames: ["type", "status"],
+  registers: [register],
+});
+
+export const execMessageDurationSeconds = new Histogram({
+  name: "exec_message_duration_seconds",
+  help: "Message round-trip time",
+  labelNames: ["type"],
+  buckets: [0.01, 0.05, 0.1, 0.5, 1, 5, 10, 30, 60],
+  registers: [register],
+});
+
+export const execProcessExitCode = new Counter({
+  name: "exec_process_exit_code",
+  help: "Exit codes by command",
+  labelNames: ["command", "exit_code"],
+  registers: [register],
+});
+
+export const execWorkspaceBytesWritten = new Counter({
+  name: "exec_workspace_bytes_written",
+  help: "Bytes written to workspaces",
+  registers: [register],
+});