From 7fb78c2de6f9378b3acd6ab98ac4e354c08497a0 Mon Sep 17 00:00:00 2001
From: "dayuan.jiang" <jdy.toh@gmail.com>
Date: Sun, 11 Oct 2026 12:06:12 +0900
Subject: [PATCH] fix(mcp-server): review fixes for the preview token and the
 draw.io static files

- The preview page fetches a fresh copy of itself and retries once when an
  API request is refused with 403: another MCP process, with its own token,
  now answers on this port, and the recovery logic (recoverState) needs its
  polls to go through. The page is sent with Cache-Control: no-store.
- draw.io files are served with an ETag and Cache-Control: no-cache instead
  of a 24 hour max-age: their names do not change between versions, so a
  package upgrade must reach the browser on the next preview. HEAD and
  If-None-Match (304) are answered.
- The file read stream goes through stream.pipeline, so a read error no
  longer ends the MCP process and a client that leaves mid-download no
  longer leaks the file handle.
---
 packages/mcp-server/src/http-server.ts        | 32 ++++++++--
 packages/mcp-server/src/preview/index.html    |  2 +-
 packages/mcp-server/src/preview/preview.js    | 23 +++++++-
 packages/mcp-server/tests/http-server.test.ts | 48 ++++++++++++++-
 tests/unit/mcp-preview-recovery.test.ts       | 59 ++++++++++++++++++-
 5 files changed, 152 insertions(+), 12 deletions(-)

diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts
index f77e9d6c..f24ef57f 100644
--- a/packages/mcp-server/src/http-server.ts
+++ b/packages/mcp-server/src/http-server.ts
@@ -7,6 +7,7 @@ import { randomBytes, randomUUID } from "node:crypto"
 import { createReadStream, existsSync, readFileSync, statSync } from "node:fs"
 import http from "node:http"
 import { dirname, extname, join, posix, resolve, sep } from "node:path"
+import { pipeline } from "node:stream"
 import { fileURLToPath } from "node:url"
 
 const MAX_BODY_BYTES = 10 * 1024 * 1024 // 10 MiB
@@ -512,8 +513,10 @@ function routeRequest(
 
         ensureSessionStateInitialized(sessionId)
 
+        // The page holds this process's token: never served from a cache
         res.writeHead(200, {
             "Content-Type": "text/html; charset=utf-8",
+            "Cache-Control": "no-store",
             ...HTML_SECURITY_HEADERS,
         })
         res.end(getHtmlPage(sessionId))
@@ -885,7 +888,7 @@ function serveDrawioFile(
     res: http.ServerResponse,
     rawPath: string,
 ): void {
-    if (req.method !== "GET") {
+    if (req.method !== "GET" && req.method !== "HEAD") {
         res.writeHead(405)
         res.end("Method Not Allowed")
         return
@@ -914,24 +917,45 @@ function serveDrawioFile(
         return
     }
     let size: number
+    let etag: string
     try {
         const stat = statSync(file)
         if (!stat.isFile()) throw new Error("not a file")
         size = stat.size
+        etag = `"${size.toString(16)}-${Math.floor(stat.mtimeMs).toString(16)}"`
     } catch {
         res.writeHead(404)
         res.end("Not Found")
         return
     }
+    // The file names do not change between draw.io versions, so the browser
+    // may keep a copy but asks before using it (a 304 from localhost is
+    // cheap); after a package upgrade the next preview gets the new files
+    const cacheHeaders = {
+        ETag: etag,
+        "Cache-Control": "no-cache",
+        "X-Content-Type-Options": "nosniff",
+    }
+    if (req.headers["if-none-match"] === etag) {
+        res.writeHead(304, cacheHeaders)
+        res.end()
+        return
+    }
     const ext = extname(file).toLowerCase()
     res.writeHead(200, {
+        ...cacheHeaders,
         "Content-Type": MIME_TYPES[ext] || "application/octet-stream",
         "Content-Length": size,
-        "Cache-Control": "public, max-age=86400",
-        "X-Content-Type-Options": "nosniff",
         ...(ext === ".html" ? HTML_SECURITY_HEADERS : {}),
     })
-    createReadStream(file).pipe(res)
+    if (req.method === "HEAD") {
+        res.end()
+        return
+    }
+    // pipeline closes both ends when either fails or goes away (a tab closed
+    // mid-download), so no file handle leaks and no error goes unhandled; a
+    // client that left early is routine, not worth a log line
+    pipeline(createReadStream(file), res, () => {})
 }
 
 /** Where the iframe loads the editor from (without its query) */
diff --git a/packages/mcp-server/src/preview/index.html b/packages/mcp-server/src/preview/index.html
index e6b85732..5d14780a 100644
--- a/packages/mcp-server/src/preview/index.html
+++ b/packages/mcp-server/src/preview/index.html
@@ -87,7 +87,7 @@
         const sessionId = {{SESSION_JSON}};
         // Empty when draw.io is served from this origin
         const DRAWIO_ORIGIN = {{ORIGIN_JSON}} || location.origin;
-        const API_TOKEN = {{TOKEN_JSON}};
+        let API_TOKEN = {{TOKEN_JSON}};
     </script>
     <script>
 {{SCRIPT}}
diff --git a/packages/mcp-server/src/preview/preview.js b/packages/mcp-server/src/preview/preview.js
index c662f524..c173261a 100644
--- a/packages/mcp-server/src/preview/preview.js
+++ b/packages/mcp-server/src/preview/preview.js
@@ -1,6 +1,25 @@
 const iframe = document.getElementById('drawio');
-// Every API request carries the token the server wrote into this page
-const api = (url, init = {}) => fetch(url, { ...init, headers: { ...(init.headers || {}), 'X-Drawio-Token': API_TOKEN } });
+// Every API request carries the token the server wrote into this page. A 403
+// means another MCP process (with its own token) now answers on this port:
+// read its token from a fresh copy of this page and retry once, so the next
+// poll can recover the session instead of being refused forever.
+let tokenRefresh = null;
+function refreshToken() {
+    if (!tokenRefresh) {
+        tokenRefresh = fetch(location.href, { cache: 'no-store' })
+            .then((r) => r.text())
+            .then((html) => { const m = html.match(/API_TOKEN = "([0-9a-f]+)"/); if (m) API_TOKEN = m[1]; })
+            .catch(() => {})
+            .finally(() => { tokenRefresh = null; });
+    }
+    return tokenRefresh;
+}
+async function api(url, init = {}) {
+    const send = () => fetch(url, { ...init, headers: { ...(init.headers || {}), 'X-Drawio-Token': API_TOKEN } });
+    let r = await send();
+    if (r.status === 403) { await refreshToken(); r = await send(); }
+    return r;
+}
 let currentVersion = 0, isReady = false, pendingXml = null, lastXml = null;
 // The server state this tab is in step with (see stateId in http-server.ts);
 // null until the first poll
diff --git a/packages/mcp-server/tests/http-server.test.ts b/packages/mcp-server/tests/http-server.test.ts
index 1c5bdd64..fb29f883 100644
--- a/packages/mcp-server/tests/http-server.test.ts
+++ b/packages/mcp-server/tests/http-server.test.ts
@@ -487,7 +487,7 @@ describe("preview page", () => {
             .map((m) => m[1])
             .join("\n")
         expect(scripts).toContain('const sessionId = "mcp-test-script";')
-        expect(scripts).toContain(`const API_TOKEN = "${getApiToken()}";`)
+        expect(scripts).toContain(`let API_TOKEN = "${getApiToken()}";`)
         expect(() => new Function(scripts)).not.toThrow()
     })
 
@@ -515,12 +515,14 @@ describe("preview page", () => {
         }
     })
 
-    it("sends the security headers with the page", async () => {
+    it("sends the security headers with the page, which is never cached", async () => {
         const res = await request("/?mcp=mcp-test-headers")
         expect(res.headers["content-security-policy"]).toBe(
             "frame-ancestors 'self'",
         )
         expect(res.headers["x-content-type-options"]).toBe("nosniff")
+        // It holds the process's token
+        expect(res.headers["cache-control"]).toBe("no-store")
     })
 })
 
@@ -566,11 +568,51 @@ describe("bundled draw.io files", () => {
         expect(res.headers["content-type"]).toBe(
             "text/javascript; charset=utf-8",
         )
-        expect(res.headers["cache-control"]).toBe("public, max-age=86400")
+        // Revalidated on every use: the file names stay the same across
+        // draw.io versions, so a package upgrade must reach the browser
+        expect(res.headers["cache-control"]).toBe("no-cache")
+        expect(res.headers.etag).toMatch(/^"[0-9a-f]+-[0-9a-f]+"$/)
         expect(res.headers["x-content-type-options"]).toBe("nosniff")
         expect(res.headers["content-security-policy"]).toBeUndefined()
     })
 
+    it("answers 304 to a matching If-None-Match, and HEAD without a body", async () => {
+        const first = await request("/drawio/js/app.min.js")
+        const again = await request("/drawio/js/app.min.js", {
+            headers: { "if-none-match": first.headers.etag as string },
+        })
+        expect(again.status).toBe(304)
+        expect(again.body).toBe("")
+        expect(again.headers.etag).toBe(first.headers.etag)
+        const changed = await request("/drawio/js/app.min.js", {
+            headers: { "if-none-match": '"0-0"' },
+        })
+        expect(changed.status).toBe(200)
+        const head = await request("/drawio/js/app.min.js", {
+            method: "HEAD",
+        })
+        expect(head.status).toBe(200)
+        expect(head.body).toBe("")
+        expect(head.headers["content-length"]).toBe("6")
+    })
+
+    it("keeps running when a file cannot be read", async () => {
+        // Readable for stat, not for open (root reads anything: skipped)
+        if (process.getuid?.() === 0) return
+        const file = join(drawioDir, "js/locked.js")
+        writeFileSync(file, "// locked", { mode: 0o000 })
+        try {
+            // The connection is dropped (the headers may be out already)
+            const res = await request("/drawio/js/locked.js").catch(() => null)
+            expect(res?.body ?? "").toBe("")
+        } finally {
+            rmSync(file, { force: true })
+        }
+        const next = await request("/drawio/js/app.min.js")
+        expect(next.status).toBe(200)
+        expect(next.body).toBe("// app")
+    })
+
     it("frames the editor page only from this origin", async () => {
         const res = await request("/drawio/index.html")
         expect(res.status).toBe(200)
diff --git a/tests/unit/mcp-preview-recovery.test.ts b/tests/unit/mcp-preview-recovery.test.ts
index 30e252da..eedf02cb 100644
--- a/tests/unit/mcp-preview-recovery.test.ts
+++ b/tests/unit/mcp-preview-recovery.test.ts
@@ -22,10 +22,12 @@ const preview = readFileSync(join(dir, "preview.js"), "utf8")
 
 const pageListeners: Array<[string, EventListener]> = []
 
-type Answer = { status: number; body: unknown }
+// A string body is sent as it is (the page's HTML); anything else as JSON
+type Answer = { status: number; body?: unknown }
 interface Call {
     url: string
     method: string
+    headers: Record<string, string>
     body: any
     answer: (a: Answer) => void
     fail: () => void
@@ -48,10 +50,19 @@ function openPage() {
                     calls.push({
                         url,
                         method: init?.method ?? "GET",
+                        headers: (init?.headers ?? {}) as Record<
+                            string,
+                            string
+                        >,
                         body: init?.body ? JSON.parse(String(init.body)) : null,
                         answer: ({ status, body }) =>
                             resolve(
-                                new Response(JSON.stringify(body), { status }),
+                                new Response(
+                                    typeof body === "string"
+                                        ? body
+                                        : JSON.stringify(body),
+                                    { status },
+                                ),
                             ),
                         fail: () => reject(new TypeError("Failed to fetch")),
                     })
@@ -153,6 +164,50 @@ describe("MCP preview after the server recreated its session", () => {
         })
     })
 
+    it("fetches the new process's token after a 403 and then recovers", async () => {
+        const t = await inStep()
+        // The user edits; the push reaches a new process, which has its own
+        // token, and is refused, as is the next poll
+        t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
+        t.fromDrawio({ event: "export", data: "<svg/>" })
+        await t.settle()
+        const refused = t.next("POST")
+        expect(refused.headers["X-Drawio-Token"]).toBe("test-token")
+        refused.answer({ status: 403 })
+        await t.settle()
+        // The page reads its own HTML again for the token
+        const reload = t.next("GET")
+        expect(reload.url).toBe(location.href)
+        expect(reload.headers["X-Drawio-Token"]).toBeUndefined()
+        reload.answer({
+            status: 200,
+            body: '<script>let API_TOKEN = "0123abcd";</script>',
+        })
+        await t.settle()
+        // The push is retried with the new token; the new process refuses
+        // it for its state id, and the poll that follows recovers
+        const retried = t.next("POST")
+        expect(retried.headers["X-Drawio-Token"]).toBe("0123abcd")
+        retried.answer({
+            status: 409,
+            body: { error: "Session was recreated" },
+        })
+        await t.settle()
+        const poll = t.next("GET")
+        expect(poll.headers["X-Drawio-Token"]).toBe("0123abcd")
+        poll.answer(state("S2", 1, "<mxfile>A</mxfile>"))
+        await t.settle()
+        await t.settle()
+        expect(t.page.read().stateId).toBe("S2")
+        const push = t.next("POST")
+        expect(push.headers["X-Drawio-Token"]).toBe("0123abcd")
+        expect(push.body).toMatchObject({
+            xml: "<mxfile>B</mxfile>",
+            stateId: "S2",
+            baseVersion: 1,
+        })
+    })
+
     it("shows the server's diagram and keeps the tab's in History", async () => {
         const t = await inStep()
         const poll = t.page.poll()