@solarisdk/mcp 0.2.1 → 0.3.0

package/dist/http.js

@@ -11,41 +11,76 @@
 //   SOLARI_BASE_URL      sandbox/desktop gateway (default https://api.getsolari.com)
 //   SOLARI_BROWSER_URL   browser gateway (default = SOLARI_BASE_URL)
 //   SESSION_IDLE_MS      evict MCP sessions idle longer than this (default 30 min)
+//   MAX_SESSIONS         global cap on live MCP sessions (default 500)
+//   MAX_SESSIONS_PER_KEY per-API-key cap (default 20)
 import { createServer } from "node:http";
 import { createHash, randomUUID } from "node:crypto";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
-import { buildServerParts } from "./server.js";
+import { buildServerParts, closeAllVmSessions } from "./server.js";
+import { releaseAllBrowserSessions } from "./browser.js";
 const PORT = Number(process.env.PORT ?? 8080);
 const IDLE_MS = Number(process.env.SESSION_IDLE_MS ?? 30 * 60 * 1000);
-const BROWSER_URL = process.env.SOLARI_BROWSER_URL ?? process.env.SOLARI_BASE_URL ?? "https://api.getsolari.com";
+const MAX_SESSIONS = Number(process.env.MAX_SESSIONS ?? 500);
+const MAX_PER_KEY = Number(process.env.MAX_SESSIONS_PER_KEY ?? 20);
+const MAX_BODY_BYTES = 4 * 1024 * 1024;
+const DESKTOP_URL = process.env.SOLARI_BASE_URL ?? "https://api.getsolari.com";
+const BROWSER_URL = process.env.SOLARI_BROWSER_URL ?? DESKTOP_URL;
 const sessions = new Map();
 const hash = (key) => createHash("sha256").update(key).digest("hex");
+const countForKey = (h) => [...sessions.values()].filter((s) => s.keyHash === h).length;
 // Cheap upfront key validation (the gateways are the real authority; this just
 // gives a clean 401 at connect time instead of failures on the first tool
-// call). Positive results cached briefly to absorb client reconnect storms.
-const keyOkUntil = new Map();
-async function validateKey(key) {
-    const h = hash(key);
-    const until = keyOkUntil.get(h);
-    if (until && until > Date.now())
-        return true;
+// call). Results are cached both ways and coalesced per key so a reconnect
+// storm can't amplify into the prod auth path.
+const keyCache = new Map();
+const inflight = new Map();
+async function probe(url, key) {
     try {
-        const res = await fetch(`${BROWSER_URL}/profiles`, {
+        const res = await fetch(url, {
             headers: { authorization: `Bearer ${key}` },
-            signal: AbortSignal.timeout(10_000),
+            signal: AbortSignal.timeout(5_000),
         });
-        if (res.ok) {
-            keyOkUntil.set(h, Date.now() + 60_000);
-            return true;
-        }
-        // Gateway auth can negative-cache a freshly minted key for a few seconds;
-        // treat only a definitive 401/403 as invalid.
-        return res.status !== 401 && res.status !== 403;
+        if (res.ok)
+            return "ok";
+        if (res.status === 401 || res.status === 403)
+            return "reject";
+        return "unknown";
     }
     catch {
-        return true; // don't lock customers out on our own transient errors
+        return "unknown";
     }
 }
+async function validateKey(key) {
+    const h = hash(key);
+    const hit = keyCache.get(h);
+    if (hit && hit.until > Date.now())
+        return hit.ok;
+    const existing = inflight.get(h);
+    if (existing)
+        return existing;
+    const p = (async () => {
+        // Prod's browser and desktop gateways have separate key stores — a key is
+        // acceptable if EITHER accepts it. Probe both concurrently.
+        const [b, d] = await Promise.all([
+            probe(`${BROWSER_URL}/profiles`, key),
+            probe(`${DESKTOP_URL}/sandboxes?limit=1`, key),
+        ]);
+        if (b === "ok" || d === "ok") {
+            keyCache.set(h, { ok: true, until: Date.now() + 60_000 });
+            return true;
+        }
+        if (b === "reject" || d === "reject") {
+            // Definitive no from a gateway that answered.
+            keyCache.set(h, { ok: false, until: Date.now() + 30_000 });
+            return false;
+        }
+        // Everything inconclusive (our blip, not the customer's fault): allow, but
+        // don't cache, so we re-check on the next connect.
+        return true;
+    })().finally(() => inflight.delete(h));
+    inflight.set(h, p);
+    return p;
+}
 function bearer(req) {
     const m = /^Bearer\s+(.+)$/i.exec(req.headers.authorization ?? "");
     return m?.[1];
@@ -54,9 +89,17 @@ function json(res, status, body) {
     res.writeHead(status, { "content-type": "application/json" }).end(JSON.stringify(body));
 }
 async function readBody(req) {
+    const declared = Number(req.headers["content-length"] ?? 0);
+    if (declared > MAX_BODY_BYTES)
+        throw Object.assign(new Error("body too large"), { status: 413 });
     const chunks = [];
-    for await (const c of req)
+    let size = 0;
+    for await (const c of req) {
+        size += c.length;
+        if (size > MAX_BODY_BYTES)
+            throw Object.assign(new Error("body too large"), { status: 413 });
         chunks.push(c);
+    }
     const raw = Buffer.concat(chunks).toString("utf8");
     if (!raw)
         return undefined;
@@ -67,16 +110,15 @@ async function evict(id, why) {
     if (!s)
         return;
     sessions.delete(id);
-    // Disconnect any live puppeteer handles (the cloud sessions themselves are
-    // governed by their own idle/pause TTLs — we only drop our sockets).
-    for (const e of s.browserReg.sessions.values()) {
-        try {
-            await e.browser.disconnect();
-        }
-        catch {
-            /* already gone */
-        }
-    }
+    // Actually RELEASE the customer's cloud resources. Browser sessions hold a
+    // concurrency-limited pool slot until their hard TTL, and an open VM control
+    // channel pins the VM "active" so it never idle-pauses — both bill until we
+    // let go, and abandoning the MCP session is the normal end-of-conversation
+    // path (models rarely call the close tools).
+    await Promise.allSettled([
+        releaseAllBrowserSessions(s.browserCfg, s.browserReg),
+        closeAllVmSessions(s.vmReg),
+    ]);
     try {
         await s.transport.close();
     }
@@ -85,12 +127,13 @@ async function evict(id, why) {
     }
     console.error(`mcp session ${id} evicted (${why})`);
 }
-setInterval(() => {
+const idleSweep = setInterval(() => {
     const now = Date.now();
     for (const [id, s] of sessions)
         if (now - s.lastSeen > IDLE_MS)
             void evict(id, "idle");
-}, 60_000).unref();
+}, 60_000);
+idleSweep.unref();
 async function handleMcp(req, res) {
     const key = bearer(req);
     if (!key || !key.startsWith("slr_live_")) {
@@ -122,20 +165,33 @@ async function handleMcp(req, res) {
         json(res, 401, { error: "invalid Solari API key" });
         return;
     }
+    const keyHash = hash(key);
+    if (sessions.size >= MAX_SESSIONS) {
+        json(res, 503, { error: "server at session capacity, retry shortly" });
+        return;
+    }
+    if (countForKey(keyHash) >= MAX_PER_KEY) {
+        json(res, 429, { error: `too many concurrent MCP sessions for this API key (max ${MAX_PER_KEY})` });
+        return;
+    }
     const body = await readBody(req);
-    const { server, browserReg } = buildServerParts(undefined, {
-        apiKey: key,
-        baseUrl: BROWSER_URL,
-    });
+    const parts = buildServerParts(undefined, { apiKey: key, baseUrl: BROWSER_URL });
     const transport = new StreamableHTTPServerTransport({
         sessionIdGenerator: randomUUID,
         onsessioninitialized: (id) => {
-            sessions.set(id, { transport, browserReg, keyHash: hash(key), lastSeen: Date.now() });
-            console.error(`mcp session ${id} started (key …${key.slice(-4)})`);
+            sessions.set(id, {
+                transport,
+                browserReg: parts.browserReg,
+                browserCfg: parts.browserCfg,
+                vmReg: parts.vmReg,
+                keyHash,
+                lastSeen: Date.now(),
+            });
+            console.error(`mcp session ${id} started (key ${keyHash.slice(0, 8)})`);
         },
         onsessionclosed: (id) => void evict(id, "client closed"),
     });
-    await server.connect(transport);
+    await parts.server.connect(transport);
     await transport.handleRequest(req, res, body);
 }
 const httpServer = createServer((req, res) => {
@@ -146,9 +202,11 @@ const httpServer = createServer((req, res) => {
     }
     if (url === "/mcp") {
         handleMcp(req, res).catch((err) => {
-            console.error("mcp request failed:", err);
+            const status = err.status ?? 500;
+            if (status !== 413)
+                console.error("mcp request failed:", err);
             if (!res.headersSent)
-                json(res, 500, { error: "internal error" });
+                json(res, status, { error: status === 413 ? "body too large" : "internal error" });
             else
                 res.end();
         });
@@ -156,7 +214,29 @@ const httpServer = createServer((req, res) => {
     }
     json(res, 404, { error: "not found (use /mcp)" });
 });
+// A stray rejection must never take down a task serving many tenants.
+process.on("unhandledRejection", (err) => {
+    console.error("solari-mcp: unhandled rejection (ignored):", err);
+});
+let shuttingDown = false;
+async function shutdown(sig) {
+    if (shuttingDown)
+        return;
+    shuttingDown = true;
+    console.error(`solari-mcp: ${sig} — draining ${sessions.size} session(s)`);
+    httpServer.close();
+    clearInterval(idleSweep);
+    // Release every tenant's cloud resources; without this an ECS rollout leaks
+    // one pool slot + one pinned VM per live conversation.
+    await Promise.race([
+        Promise.allSettled([...sessions.keys()].map((id) => evict(id, sig))),
+        new Promise((r) => setTimeout(r, 20_000)),
+    ]);
+    process.exit(0);
+}
+for (const sig of ["SIGTERM", "SIGINT"]) {
+    process.on(sig, () => void shutdown(sig));
+}
 httpServer.listen(PORT, () => {
-    console.error(`solari-mcp http listening on :${PORT} (browser gw ${BROWSER_URL}, ` +
-        `desktop gw ${process.env.SOLARI_BASE_URL ?? "https://api.getsolari.com"})`);
+    console.error(`solari-mcp http listening on :${PORT} (browser gw ${BROWSER_URL}, desktop gw ${DESKTOP_URL})`);
 });

package/dist/server.d.ts

@@ -7,6 +7,7 @@ type Kind = "sandbox" | "desktop";
 interface Entry {
     kind: Kind;
     handle: any;
+    commands: any[];
 }
 export interface Registry {
     sessions: Map<string, Entry>;
@@ -32,7 +33,19 @@ export declare function registerToolset(server: McpServer, toolset: Record<strin
 export interface ServerParts {
     server: McpServer;
     browserReg: BrowserRegistry;
+    /** Sandbox/desktop registry — the caller must close these on teardown. */
+    vmReg: Registry;
+    browserCfg: BrowserConfig;
 }
+/**
+ * Drop our control-channel sockets for every sandbox/desktop in a registry.
+ *
+ * This is required, not cosmetic: the gateway treats a session with any open
+ * control/stream connection as active and never expires it, so leaving the
+ * channel open defeats idle-pause and the VM bills indefinitely. We close
+ * (not kill) so the customer keeps their VM and its own idle policy applies.
+ */
+export declare function closeAllVmSessions(reg: Registry): Promise<void>;
 export declare function buildServerParts(client?: SolariClient, browserCfg?: BrowserConfig): ServerParts;
 export declare function buildServer(client?: SolariClient, browserCfg?: BrowserConfig): McpServer;
 export declare function main(): Promise<void>;

package/dist/server.js

@@ -12,10 +12,18 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { z } from "zod";
 import { SolariClient } from "@solarisdk/sdk";
-import { makeBrowserToolset } from "./browser.js";
-const text = (o) => ({
-    content: [{ type: "text", text: typeof o === "string" ? o : JSON.stringify(o, null, 2) }],
-});
+import { makeBrowserToolset, releaseAllBrowserSessions, } from "./browser.js";
+// Guest output (file reads, command stdout, code results) is unbounded on the
+// wire; cap it so one `cat` of a big file can't blow the MCP payload budget or
+// OOM the shared hosted task.
+const MAX_TOOL_TEXT = 30_000;
+const text = (o) => {
+    const s = typeof o === "string" ? o : JSON.stringify(o, null, 2);
+    const capped = s.length > MAX_TOOL_TEXT
+        ? `${s.slice(0, MAX_TOOL_TEXT)}\n…[truncated ${s.length - MAX_TOOL_TEXT} of ${s.length} chars]`
+        : s;
+    return { content: [{ type: "text", text: capped }] };
+};
 export function makeToolset(client, reg) {
     const need = (id) => {
         const e = reg.sessions.get(id);
@@ -23,10 +31,21 @@ export function makeToolset(client, reg) {
             throw new Error(`unknown sessionId: ${id}`);
         return e;
     };
+    // A paused session refuses the /control upgrade with 409, so resume it and
+    // retry rather than surfacing an opaque websocket error.
     const live = async (id) => {
         const e = need(id);
-        if (!e.handle.connected)
+        if (e.handle.connected)
+            return e;
+        try {
             await e.handle.connect();
+        }
+        catch (err) {
+            if (typeof e.handle.resume !== "function")
+                throw err;
+            await e.handle.resume();
+            await e.handle.connect();
+        }
         return e;
     };
     const desktop = async (id) => {
@@ -45,7 +64,7 @@ export function makeToolset(client, reg) {
                     ...(a.cpu ? { cpu: a.cpu } : {}),
                     ...(a.memMb ? { memMb: a.memMb } : {}),
                 });
-                reg.sessions.set(sbx.sandboxId, { kind: "sandbox", handle: sbx });
+                reg.sessions.set(sbx.sandboxId, { kind: "sandbox", handle: sbx, commands: [] });
                 return text({ sessionId: sbx.sandboxId });
             },
         },
@@ -57,7 +76,7 @@ export function makeToolset(client, reg) {
                     ...(a.template ? { template: a.template } : {}),
                     ...(a.resolution ? { resolution: a.resolution } : {}),
                 });
-                reg.sessions.set(d.sessionId, { kind: "desktop", handle: d });
+                reg.sessions.set(d.sessionId, { kind: "desktop", handle: d, commands: [] });
                 return text({ sessionId: d.sessionId, streamUrl: d.streamUrl });
             },
         },
@@ -92,6 +111,13 @@ export function makeToolset(client, reg) {
             handler: async (a) => {
                 const e = await live(a.sessionId);
                 const h = await e.handle.commands.start("sh", { args: ["-c", a.command] });
+                // commands.start() arms an exit promise that REJECTS when the control
+                // channel closes (kill, idle-pause, gateway redeploy). Nothing else
+                // awaits it, so without this catch the rejection is unhandled and Node
+                // terminates the process — taking every other tenant's session with it
+                // on the hosted server.
+                void h.wait?.().catch(() => { });
+                e.commands.push(h);
                 return text({ cmdId: h.cmdId });
             },
         },
@@ -104,17 +130,34 @@ export function makeToolset(client, reg) {
             },
         },
         solari_connect: {
-            description: "Re-attach to an existing session by id (across process restarts); auto-resumes it if paused. Registers it and returns { sessionId, kind }.",
+            description: "Re-attach to an existing sandbox or desktop by id (e.g. across restarts); resumes it if " +
+                "paused. The kind is read from the gateway, so you don't need to know it. Returns { sessionId, kind, state }.",
             inputSchema: { sessionId: z.string(), kind: z.enum(["sandbox", "desktop"]).optional() },
             handler: async (a) => {
                 const id = a.sessionId;
-                const existing = reg.sessions.get(id);
-                const kind = a.kind ?? existing?.kind ?? "sandbox";
+                // Trust the gateway over the caller's hint: guessing "sandbox" for a
+                // desktop permanently misclassifies it and every GUI tool then wrongly
+                // refuses the session.
+                let kind = a.kind ?? reg.sessions.get(id)?.kind ?? "sandbox";
+                let state;
+                try {
+                    const view = await client.sandboxes.get(id);
+                    if (view?.kind === "desktop" || view?.kind === "sandbox")
+                        kind = view.kind;
+                    state = view?.state;
+                }
+                catch {
+                    /* fall back to the hint */
+                }
                 const handle = kind === "desktop"
                     ? await client.desktops.connect(id)
                     : await client.sandboxes.connect(id);
-                reg.sessions.set(id, { kind, handle });
-                return text({ sessionId: id, kind });
+                // desktops.connect() auto-resumes; sandboxes.connect() does not.
+                if (kind === "sandbox" && state === "paused" && typeof handle.resume === "function") {
+                    await handle.resume();
+                }
+                reg.sessions.set(id, { kind, handle, commands: [] });
+                return text({ sessionId: id, kind, state: state ?? "unknown" });
             },
         },
         solari_run_code: {
@@ -203,6 +246,27 @@ export function registerToolset(server, toolset) {
         async (args) => t.handler(args));
     }
 }
+/**
+ * Drop our control-channel sockets for every sandbox/desktop in a registry.
+ *
+ * This is required, not cosmetic: the gateway treats a session with any open
+ * control/stream connection as active and never expires it, so leaving the
+ * channel open defeats idle-pause and the VM bills indefinitely. We close
+ * (not kill) so the customer keeps their VM and its own idle policy applies.
+ */
+export async function closeAllVmSessions(reg) {
+    const ids = [...reg.sessions.keys()];
+    await Promise.all(ids.map(async (id) => {
+        const e = reg.sessions.get(id);
+        reg.sessions.delete(id);
+        try {
+            await e?.handle.close?.();
+        }
+        catch (err) {
+            console.error(`solari-mcp: failed to close session ${id}:`, err);
+        }
+    }));
+}
 export function buildServerParts(client, browserCfg) {
     const apiKey = browserCfg?.apiKey ?? process.env.SOLARI_API_KEY ?? "";
     const c = client ??
@@ -212,19 +276,22 @@ export function buildServerParts(client, browserCfg) {
         });
     // The cloud-browser gateway is a separate service from the sandbox/desktop
     // gateway (same domain in prod, different domains on staging).
+    // Prod's browser + desktop gateways have separate key stores today; allow a
+    // browser-specific key (falls back to the shared one).
     const bCfg = browserCfg ?? {
-        apiKey,
+        apiKey: process.env.SOLARI_BROWSER_API_KEY ?? apiKey,
         baseUrl: process.env.SOLARI_BROWSER_URL ??
             process.env.SOLARI_BASE_URL ??
             "https://api.getsolari.com",
     };
-    const server = new McpServer({ name: "solari-mcp", version: "0.2.0" });
+    const server = new McpServer({ name: "solari-mcp", version: "0.3.0" });
     const browserReg = { sessions: new Map() };
+    const vmReg = { sessions: new Map() };
     registerToolset(server, {
-        ...makeToolset(c, { sessions: new Map() }),
+        ...makeToolset(c, vmReg),
         ...makeBrowserToolset(bCfg, browserReg),
     });
-    return { server, browserReg };
+    return { server, browserReg, vmReg, browserCfg: bCfg };
 }
 export function buildServer(client, browserCfg) {
     return buildServerParts(client, browserCfg).server;
@@ -234,8 +301,34 @@ export async function main() {
         console.error("solari-mcp: SOLARI_API_KEY is required");
         process.exit(1);
     }
-    const server = buildServer();
-    await server.connect(new StdioServerTransport());
+    // A stray rejection (e.g. a background command's exit promise) must never
+    // take the server down mid-session.
+    process.on("unhandledRejection", (err) => {
+        console.error("solari-mcp: unhandled rejection (ignored):", err);
+    });
+    const parts = buildServerParts();
+    // Without this, every exit (client quit, Ctrl-C, host restart) leaves browser
+    // sessions held to their TTL and VM control channels open — which pins the
+    // VMs active so they never idle-pause. Both bill until cleaned up.
+    let shuttingDown = false;
+    const shutdown = async (sig) => {
+        if (shuttingDown)
+            return;
+        shuttingDown = true;
+        console.error(`solari-mcp: ${sig} — releasing sessions`);
+        await Promise.race([
+            Promise.allSettled([
+                releaseAllBrowserSessions(parts.browserCfg, parts.browserReg),
+                closeAllVmSessions(parts.vmReg),
+            ]),
+            new Promise((r) => setTimeout(r, 10_000)),
+        ]);
+        process.exit(0);
+    };
+    for (const sig of ["SIGTERM", "SIGINT"]) {
+        process.on(sig, () => void shutdown(sig));
+    }
+    await parts.server.connect(new StdioServerTransport());
 }
 // Run only when executed as the binary (not when imported by tests).
 if (process.argv[1] && import.meta.url === `file://${process.argv[1]}`) {