mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-11 20:19:51 +08:00
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.
This commit is contained in:
@@ -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) */
|
||||
|
||||
@@ -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}}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user