From b8b851fc5bffca04db8e334fde16005d79a869a9 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Mon, 5 Oct 2026 21:19:32 +0900 Subject: [PATCH] fix: what the whole-PR review and Copilot found - A redirect followed for a custom base URL also drops the key headers of providers that do not use Authorization (x-api-key, x-goog-api-key, api-key) when it goes to another origin. - A second Enter or click while a message is being prepared (attachments read, diagram exported) no longer sends it twice. - The admin Test on the deployment's own endpoints (EdgeOne, the server's keyless Ollama, an address on the server's network) counts toward the quota like a chat; the chat and the Test share one rule for it. The Test of an Azure entry set up by AZURE_RESOURCE_NAME only goes where chat goes. - EdgeOne's function is called at the site root again, as on main: EdgeOne serves edge functions there, outside Next's base path. - MCP History: the state before a write is kept unless the browser saved no change of the user's since the last server write (draw.io's sync copy of it adds no entry), and the dedupe compares the exact text again, so a change of page size or other settings only is its own version. - MCP: an edit keeps untouched labels as draw.io shows them (a literal line break in an attribute is a space); a new document of empty pages the user named is auto-saved; load_diagram reads only regular files, so a pipe cannot hold up the other write tools; the preview does not load back its own push still on its way (an undo made meanwhile is saved). - Two overlapping saves of a new chat no longer reload the canvas from the older copy. - At most three screenshot checks per user turn, passed or failed, as documented. - Desktop: the main window navigates only within the app (draw.io stays in its frame); a presets file that is not JSON and cannot be moved aside is not overwritten. - A last self-closing cell with a raw "<" in a value is not taken for cut off output. - README: Material Design shapes load their icons from fonts.gstatic.com. --- app/api/chat/route.ts | 17 +++--- app/api/validate-model/route.ts | 35 +++++++++++- components/chat-panel.tsx | 29 +++++++--- components/model-config-dialog.tsx | 2 +- electron/main/config-manager.ts | 2 + electron/main/index.ts | 6 ++- hooks/use-diagram-tool-handlers.ts | 14 +++-- lib/ai-providers.ts | 38 +++++++++++-- lib/ssrf-protection.ts | 16 ++++-- lib/utils.ts | 5 +- packages/mcp-server/README.md | 2 +- packages/mcp-server/src/dom.ts | 27 +++++++++- packages/mcp-server/src/history.ts | 33 ++---------- packages/mcp-server/src/http-server.ts | 23 +++++++- packages/mcp-server/src/index.ts | 18 +++++-- packages/mcp-server/src/persistence.ts | 11 +++- packages/mcp-server/src/preview/preview.js | 9 +++- .../mcp-server/tests/edit-diagram.test.ts | 20 +++++++ packages/mcp-server/tests/http-server.test.ts | 53 ++++++++++++++++--- packages/mcp-server/tests/persistence.test.ts | 22 +++++++- tests/e2e/chat.spec.ts | 28 ++++++++++ tests/unit/admin-test-model.test.ts | 17 +++++- tests/unit/chat-route-edgeone.test.ts | 4 +- tests/unit/config-manager.test.ts | 26 +++++++++ tests/unit/mcp-preview-recovery.test.ts | 23 ++++++++ tests/unit/ssrf-protection.test.ts | 6 +++ tests/unit/use-diagram-tool-handlers.test.tsx | 22 +++++++- tests/unit/utils.test.ts | 6 +++ tests/unit/validate-model-route.test.ts | 52 ++++++++++++++++++ 29 files changed, 476 insertions(+), 90 deletions(-) diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 3c66756b..c5ea47a2 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -19,6 +19,7 @@ import { SINGLE_SYSTEM_PROVIDERS, supportsPromptCaching, usesServerCredentials, + usesServerEndpoint, } from "@/lib/ai-providers" import { findCachedResponse } from "@/lib/cached-responses" import { @@ -287,17 +288,11 @@ async function handleChatRequest(req: Request): Promise { const clientBaseUrl = normalizeBaseUrl( req.headers.get("x-ai-base-url") ?? "", ) - const usesClientBaseUrl = - resolvedProvider !== "bedrock" && resolvedProvider !== "edgeone" - const onServerEndpoint = - process.env.NEXT_AI_DRAWIO_DESKTOP !== "1" && - (resolvedProvider === "edgeone" || - (resolvedProvider === "ollama" && - !clientBaseUrl && - !clientOverrides.apiKey) || - (usesClientBaseUrl && - !!clientBaseUrl && - (await isPrivateUrl(clientBaseUrl)))) + const onServerEndpoint = await usesServerEndpoint( + resolvedProvider, + clientBaseUrl, + clientOverrides.apiKey, + ) const countsQuota = isQuotaEnabled() && (onServerCredentials || onServerEndpoint) && diff --git a/app/api/validate-model/route.ts b/app/api/validate-model/route.ts index 3a396abe..1109a976 100644 --- a/app/api/validate-model/route.ts +++ b/app/api/validate-model/route.ts @@ -8,10 +8,16 @@ import { getAIModel, globalBaseUrl, usesServerCredentials, + usesServerEndpoint, } from "@/lib/ai-providers" +import { + checkAndIncrementRequest, + isQuotaEnabled, +} from "@/lib/dynamo-quota-manager" import { classifyLLMError } from "@/lib/llm-errors" import { allowPrivateUrls, isPrivateUrl } from "@/lib/ssrf-protection" -import type { ProviderName } from "@/lib/types/model-config" +import { normalizeBaseUrl, type ProviderName } from "@/lib/types/model-config" +import { getUserIdFromRequest } from "@/lib/user-id" export const runtime = "nodejs" @@ -140,6 +146,33 @@ export async function POST(req: Request) { ) } + // On the deployment's own endpoints a Test runs a model as a chat + // does, so with the quota on it counts as a chat request (an + // admin's Test of the server's URL does not) + const userId = getUserIdFromRequest(req) + if ( + isQuotaEnabled() && + !serverUrl && + userId !== "anonymous" && + (await usesServerEndpoint( + provider, + normalizeBaseUrl(body.baseUrl ?? ""), + apiKey, + )) + ) { + const quotaCheck = await checkAndIncrementRequest(userId, { + requests: Number(process.env.DAILY_REQUEST_LIMIT) || 10, + tokens: Number(process.env.DAILY_TOKEN_LIMIT) || 200000, + tpm: Number(process.env.TPM_LIMIT) || 20000, + }) + if (!quotaCheck.allowed) { + return NextResponse.json( + { valid: false, error: quotaCheck.error }, + { status: 429 }, + ) + } + } + // The same model the chat would use. A client base URL makes it // refuse redirects to internal hosts. const { model } = getAIModel({ diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index d8d329dc..c378a3e1 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -277,6 +277,7 @@ export default function ChatPanel({ // Set by Stop until the user sends the next message const stoppedRef = useRef(false) + const preparingSendRef = useRef(false) // Presses of Stop: a check that began before one still knows of it after // the next message clears stoppedRef const stopCountRef = useRef(0) @@ -524,9 +525,10 @@ export default function ChatPanel({ // Track last synced session ID to detect external changes (e.g., URL back/forward) const lastSyncedSessionIdRef = useRef(null) - // Messages array from our latest save. A session holding this exact array was - // created by our own save, so it must not be treated as an external switch. - const lastSavedMessagesRef = useRef(null) + // Message arrays of our own saves. A session holding one of them was + // created by our own save, so it must not be treated as an external + // switch (with two saves of a new chat at once, the first creates it). + const savedMessagesRef = useRef(new WeakSet()) // Helper: Sync UI state with session data (eliminates duplication) // Track message IDs that are being loaded from session (to skip animations/scroll) @@ -594,7 +596,7 @@ export default function ChatPanel({ } } const messages = sanitizeMessages(messagesRef.current) - lastSavedMessagesRef.current = messages + savedMessagesRef.current.add(messages) return { messages, xmlSnapshots: Array.from(xmlSnapshotsRef.current.entries()), @@ -651,7 +653,7 @@ export default function ChatPanel({ // Our own save created this session; the UI already shows its content const isOwnNewSession = - newSession?.messages === lastSavedMessagesRef.current + !!newSession && savedMessagesRef.current.has(newSession.messages) // Update last synced ID lastSyncedSessionIdRef.current = newSessionId @@ -802,8 +804,7 @@ export default function ChatPanel({ ) }, [sessionManager, buildSessionData]) - const onFormSubmit = async (e: React.FormEvent) => { - e.preventDefault() + const submitInput = async () => { const isProcessing = status === "streaming" || status === "submitted" // Attachments still extracting have no text yet. Template sends call // requestSubmit() and skip the disabled send button, so check here too. @@ -901,6 +902,20 @@ export default function ChatPanel({ } } + const onFormSubmit = async (e: React.FormEvent) => { + e.preventDefault() + // While a send is prepared (attachments read, diagram exported) the + // status is still "ready": a second Enter or click would send the + // message again + if (preparingSendRef.current) return + preparingSendRef.current = true + try { + await submitInput() + } finally { + preparingSendRef.current = false + } + } + // Export the current diagram, snapshot it for this message, and send. // onSent runs right after sending, so the input empties as the message // shows in the chat diff --git a/components/model-config-dialog.tsx b/components/model-config-dialog.tsx index 52dbea78..dfe242e7 100644 --- a/components/model-config-dialog.tsx +++ b/components/model-config-dialog.tsx @@ -422,7 +422,7 @@ export function ModelConfigDialog({ // For EdgeOne, construct baseUrl from current origin const baseUrl = isEdgeOne - ? `${window.location.origin}${getApiEndpoint("/api/edgeai")}` + ? `${window.location.origin}/api/edgeai` : selectedProvider.baseUrl // Test every model at once; each row updates when its answer arrives diff --git a/electron/main/config-manager.ts b/electron/main/config-manager.ts index 7b4fc5fb..1c3499f7 100644 --- a/electron/main/config-manager.ts +++ b/electron/main/config-manager.ts @@ -217,7 +217,9 @@ export function loadPresets(): ConfigPresetsFile { renameSync(configPath, backupPath) console.error(`Unreadable config presets moved to ${backupPath}`) } catch (renameError) { + // Still there: refuse saves that would overwrite it console.error("Failed to back up config presets:", renameError) + presetsUnreadable = true } return { version: 1, diff --git a/electron/main/index.ts b/electron/main/index.ts index e17ddb02..08313b1a 100644 --- a/electron/main/index.ts +++ b/electron/main/index.ts @@ -124,9 +124,11 @@ if (!gotTheLock) { }) // Clicking a plain link would otherwise replace the app page with - // an external site that keeps the preload API + // an external site that keeps the preload API. Only the page + // itself may navigate there; draw.io stays in its frame (this event + // is for the main frame only) contents.on("will-navigate", (event) => { - if (isInAppUrl(event.url)) { + if (isAppUrl(event.url)) { return } event.preventDefault() diff --git a/hooks/use-diagram-tool-handlers.ts b/hooks/use-diagram-tool-handlers.ts index c5cbdbe7..ad67ec28 100644 --- a/hooks/use-diagram-tool-handlers.ts +++ b/hooks/use-diagram-tool-handlers.ts @@ -251,7 +251,9 @@ ${finalXml} if ( enableVlmValidation && captureValidationPng && - validateDiagram + validateDiagram && + // At most this many checks per user turn, passed or not + validationRetryCountRef.current < MAX_VALIDATION_RETRIES ) { let capturedPngData: string | null = null const stopped = watchStop?.() @@ -275,8 +277,10 @@ ${finalXml} ) } - // Each retry is a new tool call, so count attempts per user turn + // Each retry is a new tool call, so count attempts + // per user turn (the chat resets it when the user sends) const attempt = validationRetryCountRef.current + 1 + validationRetryCountRef.current = attempt // Notify UI that we're validating (include the image) updateValidationState( @@ -296,8 +300,6 @@ ${finalXml} if (!result.valid) { if (attempt < MAX_VALIDATION_RETRIES) { - validationRetryCountRef.current = attempt - const feedback = formatValidationFeedback(result) if (DEBUG) { @@ -333,8 +335,6 @@ ${finalXml} "[display_diagram] Max validation retries reached, accepting diagram", ) } - validationRetryCountRef.current = 0 - // Notify UI that we're accepting with issues (include the image) updateValidationState( toolCall.toolCallId, @@ -350,8 +350,6 @@ ${finalXml} return } } else { - // Validation passed - reset retry count - validationRetryCountRef.current = 0 if (DEBUG) { console.log( "[display_diagram] Validation passed!", diff --git a/lib/ai-providers.ts b/lib/ai-providers.ts index b2ef3390..bf040440 100644 --- a/lib/ai-providers.ts +++ b/lib/ai-providers.ts @@ -22,8 +22,7 @@ import { loadAdminProviders, } from "@/lib/admin/providers" import { getEnvFallback } from "@/lib/admin/settings" -import { getApiEndpoint } from "@/lib/base-path" -import { redirectGuardedFetch } from "@/lib/ssrf-protection" +import { isPrivateUrl, redirectGuardedFetch } from "@/lib/ssrf-protection" import { normalizeBaseUrl, PROVIDER_INFO, @@ -1112,11 +1111,12 @@ export function getAIModel(clientOverrides?: ClientOverrides): ModelConfig { /** * The deployment's EdgeOne Pages function, as an absolute URL (the SDK - * needs one), under the deployment's base path + * needs one). EdgeOne serves functions by their folder from the site root, + * so Next's base path does not apply. */ export function edgeOneEndpoint(req: Request): string { const origin = req.headers.get("origin") || new URL(req.url).origin - return `${origin}${getApiEndpoint("/api/edgeai")}` + return `${origin}/api/edgeai` } /** @@ -1136,6 +1136,14 @@ export function globalBaseUrl(provider: ProviderName): string | undefined { return getEnvFallback("GOOGLE_VERTEX_BASE_URL") || undefined } if (provider === "bedrock" || provider === "edgeone") return undefined + // Azure set up by resource name only: the URL the SDK builds from it + if ( + provider === "azure" && + !process.env.AZURE_BASE_URL && + process.env.AZURE_RESOURCE_NAME + ) { + return `https://${process.env.AZURE_RESOURCE_NAME}.openai.azure.com/openai` + } const name = provider === "gateway" ? "AI_GATEWAY_BASE_URL" @@ -1148,6 +1156,28 @@ export function getServerProvider(): ProviderName | null { return (process.env.AI_PROVIDER as ProviderName) || detectProvider() } +/** + * Whether a call made with the caller's own settings runs on an endpoint of + * the deployment: its EdgeOne function, the server's keyless Ollama, or an + * address on the server's network (which ignores a dummy key header). + * Bedrock and EdgeOne never use a client base URL. Never in the desktop + * app, where every endpoint is the user's. clientBaseUrl: normalized. + */ +export async function usesServerEndpoint( + provider: ProviderName | null | undefined, + clientBaseUrl: string, + apiKey: string | null | undefined, +): Promise { + if (process.env.NEXT_AI_DRAWIO_DESKTOP === "1") return false + if (provider === "edgeone") return true + if (provider === "ollama" && !clientBaseUrl && !apiKey) return true + return ( + provider !== "bedrock" && + !!clientBaseUrl && + (await isPrivateUrl(clientBaseUrl)) + ) +} + /** * Whether the call is paid for by the server's own credentials (env keys or * IAM role) rather than credentials sent with the request. Mirrors which key diff --git a/lib/ssrf-protection.ts b/lib/ssrf-protection.ts index 5e682893..20f2e191 100644 --- a/lib/ssrf-protection.ts +++ b/lib/ssrf-protection.ts @@ -126,6 +126,18 @@ export class RedirectRefusedError extends Error { const MAX_REDIRECTS = 5 +// Dropped when a redirect goes to another origin: those fetch drops, and +// the key headers of providers that do not use Authorization (Anthropic, +// Google, Azure) +const CREDENTIAL_HEADERS = [ + "authorization", + "proxy-authorization", + "cookie", + "x-api-key", + "x-goog-api-key", + "api-key", +] + /** * A fetch for requests to a base URL the client chose. With private URLs * blocked, a public URL could still redirect the request to an internal @@ -161,9 +173,7 @@ export function redirectGuardedFetch(): typeof fetch | undefined { // credentials (the user's key, EdgeOne's cookies) const headers = new Headers(next?.headers) if (new URL(url).origin !== from.origin) { - headers.delete("authorization") - headers.delete("proxy-authorization") - headers.delete("cookie") + for (const name of CREDENTIAL_HEADERS) headers.delete(name) } next = { ...next, headers } // 303, and 301 or 302 after a POST, go on as a GET without the diff --git a/lib/utils.ts b/lib/utils.ts index 3f648a1b..a4749b92 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -59,7 +59,10 @@ export function isMxCellXmlComplete(xml: string | undefined | null): boolean { const lastCellStart = trimmed.lastIndexOf(" lastMxCellClose && - !/^/.test(trimmed.slice(lastCellStart)) + // (quoted values may hold a raw "<", which the auto-fix escapes) + !/^"']|"[^"]*"|'[^']*')*\/>/.test( + trimmed.slice(lastCellStart), + ) ) { return false } diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index 9dcfa195..968688a2 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -195,7 +195,7 @@ You can deploy your own draw.io instance using the official Docker image: docker run -d -p 8080:8080 jgraph/drawio ``` -Then set `DRAWIO_BASE_URL=http://localhost:8080` (or your server's URL). The preview page loads nothing else from the internet, so with a local draw.io it works fully offline. +Then set `DRAWIO_BASE_URL=http://localhost:8080` (or your server's URL). The preview page loads nothing else from the internet, so with a local draw.io it works offline. One exception: shapes from the Material Design library show icons from `fonts.gstatic.com`. ## Troubleshooting diff --git a/packages/mcp-server/src/dom.ts b/packages/mcp-server/src/dom.ts index c777c03d..13972b3a 100644 --- a/packages/mcp-server/src/dom.ts +++ b/packages/mcp-server/src/dom.ts @@ -60,6 +60,31 @@ export function serializeXml(node: Node): string { } } +/** + * XML parsers read a literal tab or line break inside an attribute value as + * a space (a line break written as stays one). linkedom keeps it, and + * serializeXml would then write it as a real line break, so an edit would + * change labels it never touched. Applied to the text before linkedom. + */ +function normalizeAttributeWhitespace(xml: string): string { + return xml.replace( + /<[A-Za-z][^"'<>]*(?:(?:"[^"]*"|'[^']*')[^"'<>]*)*>/g, + (tag) => + tag.replace(/"[^"]*"|'[^']*'/g, (value) => + value.replace(/\r\n|[\t\n\r]/g, " "), + ), + ) +} + +class XmlDomParser extends DOMParser { + parseFromString(text: string, type: string) { + return super.parseFromString( + type.includes("xml") ? normalizeAttributeWhitespace(text) : text, + type as any, + ) + } +} + class XMLSerializerPolyfill { serializeToString(node: Node): string { return serializeXml(node) @@ -68,6 +93,6 @@ class XMLSerializerPolyfill { /** Install the DOMParser and XMLSerializer globals the XML helpers use. */ export function installDomPolyfill(): void { - ;(globalThis as any).DOMParser = DOMParser + ;(globalThis as any).DOMParser = XmlDomParser ;(globalThis as any).XMLSerializer = XMLSerializerPolyfill } diff --git a/packages/mcp-server/src/history.ts b/packages/mcp-server/src/history.ts index 896ff5f4..7fff09a8 100644 --- a/packages/mcp-server/src/history.ts +++ b/packages/mcp-server/src/history.ts @@ -3,9 +3,7 @@ * Stores {xml, svg} entries in a circular buffer */ -import { contentFingerprint } from "./edit-gate.ts" import { log } from "./logger.ts" -import { normalizeToMxfile, parseMxfile } from "./pages.ts" const MAX_HISTORY = 20 @@ -18,30 +16,6 @@ interface HistoryEntry { let nextEntryId = 0 const historyStore = new Map() -/** Each page's background colour */ -function backgrounds(xml: string): string { - const doc = parseMxfile(normalizeToMxfile(xml) ?? xml) - if (!doc) return "" - return Array.from(doc.querySelectorAll("mxGraphModel")) - .map((m) => m.getAttribute("background") || "none") - .join(",") -} - -// The same pages, cells and backgrounds. draw.io's own copy of a diagram -// (a sync reply) adds view and page attributes such as dx, grid and the -// page size, which the model's XML leaves out, so those are not compared. -// A document without pages has an empty fingerprint and is compared as -// text only. -function sameDiagram(a: string, b: string): boolean { - if (a === b) return true - const fingerprint = contentFingerprint(a) - return ( - fingerprint !== "" && - fingerprint === contentFingerprint(b) && - backgrounds(a) === backgrounds(b) - ) -} - export function addHistory(sessionId: string, xml: string, svg = ""): number { let history = historyStore.get(sessionId) if (!history) { @@ -49,10 +23,11 @@ export function addHistory(sessionId: string, xml: string, svg = ""): number { historyStore.set(sessionId, history) } - // Dedupe: skip if same as last entry, also when only re-serialized - // (a change of background only is a new version) + // Dedupe: skip if same as last entry (any other change, also of page + // settings only, is a new version; keepInHistory leaves out draw.io's + // copy of a server write) const last = history[history.length - 1] - if (last && sameDiagram(last.xml, xml)) { + if (last && last.xml === xml) { if (svg && !last.svg) last.svg = svg return history.length - 1 } diff --git a/packages/mcp-server/src/http-server.ts b/packages/mcp-server/src/http-server.ts index 3b40cf70..26170052 100644 --- a/packages/mcp-server/src/http-server.ts +++ b/packages/mcp-server/src/http-server.ts @@ -138,6 +138,9 @@ interface SessionState { serverVersion?: number // The XML of that write: what a thumbnail taken after loading it shows serverXml?: string + // The browser saved a change of the user's since that write (a sync + // reply is no change) + userEdited?: boolean lastUpdated: Date lastPolled?: number // Last browser poll; an open tab keeps the session alive svg?: string // Cached SVG from last browser save @@ -206,6 +209,7 @@ export function setState( stateId: existing?.stateId ?? randomUUID(), serverVersion: fromBrowser ? existing?.serverVersion : newVersion, serverXml: fromBrowser ? existing?.serverXml : xml, + userEdited: fromBrowser ? existing?.userEdited : false, lastUpdated: new Date(), lastPolled: existing?.lastPolled, // The image of this XML, never an older one's: a write without an @@ -226,6 +230,21 @@ export function setState( return newVersion } +/** + * Keep the session's diagram in History before a write replaces it. + * Nothing to keep when the browser saved no change of the user's since the + * last server write and History ends with that write: the state is that + * write, or draw.io's own copy of it from a sync (other text, same diagram). + */ +export function keepInHistory(sessionId: string, xml: string, svg = ""): void { + const state = stateStore.get(sessionId) + const last = getHistory(sessionId).at(-1) + if (state && !state.userEdited && last && last.xml === state.serverXml) { + return + } + addHistory(sessionId, xml, svg) +} + /** * Ask the browser bridge to export the current diagram as png/svg. * @@ -604,6 +623,8 @@ function handleStateApi( return } const version = setState(sessionId, data.xml, data.svg, true) + const saved = stateStore.get(sessionId) + if (saved && data.source !== "sync") saved.userEdited = true res.writeHead(200, { "Content-Type": "application/json" }) res.end(JSON.stringify({ success: true, version })) } catch { @@ -685,7 +706,7 @@ function handleRestoreApi( contentFingerprint(current.xml) !== contentFingerprint(BLANK_MXFILE) ) { - addHistory(sessionId, current.xml, current.svg) + keepInHistory(sessionId, current.xml, current.svg) } const newVersion = setState(sessionId, entry.xml) addHistory(sessionId, entry.xml, entry.svg) diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index 35d5e26e..fb8939b9 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -35,6 +35,7 @@ import { type ExportOptions, getServerPort, getState, + keepInHistory, onSessionRecreate, onStateChange, requestExport, @@ -389,7 +390,7 @@ Rules: cells are siblings (never nested), ids are unique per page and start from // Save user's state before AI overwrites (with cached SVG) if (currentSession.xml) { - addHistory( + keepInHistory( currentSession.id, currentSession.xml, browserState?.svg || "", @@ -481,6 +482,11 @@ registerWriteTool( let content: string try { + // A pipe or device could be read forever, and the other + // write tools wait for this one + if (!(await fs.stat(absolutePath)).isFile()) { + throw new Error("not a regular file") + } content = await fs.readFile(absolutePath, "utf-8") } catch (e) { const msg = e instanceof Error ? e.message : String(e) @@ -515,7 +521,7 @@ registerWriteTool( currentSession.xml = browserState.xml } if (currentSession.xml) { - addHistory( + keepInHistory( currentSession.id, currentSession.xml, browserState?.svg || "", @@ -728,7 +734,7 @@ registerWriteTool( // Save the pre-edit state for undo (with cached SVG from browser). // Done only once the edit applied: a rejected edit returns above // without leaving a phantom history entry. - addHistory( + keepInHistory( currentSession.id, currentSession.xml, browserState?.svg || "", @@ -1456,7 +1462,11 @@ async function loadMxfileForMutation(): Promise< browserState?.xml ?? "", ).ok // Save history before overwriting so the user can undo. - addHistory(sessionRef.id, sessionRef.xml, browserState?.svg || "") + keepInHistory( + sessionRef.id, + sessionRef.xml, + browserState?.svg || "", + ) sessionRef.xml = newXml sessionRef.version++ setState(sessionRef.id, newXml) diff --git a/packages/mcp-server/src/persistence.ts b/packages/mcp-server/src/persistence.ts index 1d35330c..2ba44b44 100644 --- a/packages/mcp-server/src/persistence.ts +++ b/packages/mcp-server/src/persistence.ts @@ -18,8 +18,15 @@ import { } from "node:fs" import { homedir } from "node:os" import { join } from "node:path" +import { contentFingerprint } from "./edit-gate.ts" import { log } from "./logger.ts" -import { hasCells } from "./pages.ts" +import { BLANK_MXFILE, hasCells } from "./pages.ts" + +// The blank page the browser shows before any drawing (page names count: +// empty pages the user named or added are kept) +const isBlank = (xml: string) => + !hasCells(xml) && + contentFingerprint(xml) === contentFingerprint(BLANK_MXFILE) const DELAY_MS = 1000 const MAX_FILES = 50 @@ -126,7 +133,7 @@ export class Autosaver { try { const isNew = !existsSync(path) // A blank page the browser shows before any drawing: nothing to keep - if (isNew && !hasCells(entry.xml)) return + if (isNew && isBlank(entry.xml)) return mkdirSync(this.dir, { recursive: true }) // Write to a temporary file first so a crash never leaves half a file writeFileSync(`${path}.tmp`, entry.xml, "utf-8") diff --git a/packages/mcp-server/src/preview/preview.js b/packages/mcp-server/src/preview/preview.js index 2ba68ec2..0a88fb05 100644 --- a/packages/mcp-server/src/preview/preview.js +++ b/packages/mcp-server/src/preview/preview.js @@ -7,6 +7,7 @@ let stateId = null; // the last one the server has let latestXml = null; let pushFailing = false; // the last push could not reach the server +const pushesInFlight = []; // XML of pushes not answered yet // After recovery replaced the canvas, until draw.io reports the load: an // autosave still on its way belongs to the canvas being replaced let awaitingLoad = false; @@ -181,6 +182,7 @@ function showNotice(text) { // server state the push is based on. async function pushState(xml, svg = '', baseVersion = currentVersion, source = 'edit', sid = stateId) { if (!sessionId) return; + pushesInFlight.push(xml); try { const r = await fetch('/api/state', { method: 'POST', @@ -226,6 +228,8 @@ async function pushState(xml, svg = '', baseVersion = currentVersion, source = ' pushFailing = true; showNotice("Can't reach the MCP server. Your changes are only in this tab for now; use Download to keep a copy."); } + } finally { + pushesInFlight.splice(pushesInFlight.indexOf(xml), 1); } } @@ -291,7 +295,10 @@ async function poll() { // (forceReload) replaces it, so a new version doesn't fight the // projection; currentVersion stays unadvanced until then, so the // bump is applied with the real document. - if ((forceReload || (s.version > currentVersion && !projectionExportActive)) && s.xml) { + // The tab's own push still on its way is not loaded back: the + // canvas may have moved on since (an undo), and its answer follows + const ownPush = pushesInFlight.includes(s.xml); + if ((forceReload || (s.version > currentVersion && !projectionExportActive && !ownPush)) && s.xml) { forceReload = false; projectionExportActive = false; currentVersion = s.version; diff --git a/packages/mcp-server/tests/edit-diagram.test.ts b/packages/mcp-server/tests/edit-diagram.test.ts index 7354c99c..34368ef0 100644 --- a/packages/mcp-server/tests/edit-diagram.test.ts +++ b/packages/mcp-server/tests/edit-diagram.test.ts @@ -141,3 +141,23 @@ describe("targetPageXml", () => { expect(xml).not.toContain('id="a"') }) }) + +describe("labels an edit does not touch", () => { + it("keep their line breaks and spaces as draw.io reads them", () => { + // A literal line break in an attribute reads as a space; is a + // real line break + const labels = + `` + + `` + const out = editDiagram( + `${page("p1", labels)}`, + [{ operation: "add", cell_id: "c", new_xml: cell("c") }], + {}, + ) + expect(out.ok).toBe(true) + if (!out.ok) return + expect(out.xml).toContain(`value="Hello world"`) + expect(out.xml).toContain(`value="Line 1 Line 2"`) + }) +}) diff --git a/packages/mcp-server/tests/http-server.test.ts b/packages/mcp-server/tests/http-server.test.ts index 7168911a..1f4f2bca 100644 --- a/packages/mcp-server/tests/http-server.test.ts +++ b/packages/mcp-server/tests/http-server.test.ts @@ -12,6 +12,7 @@ import { installDomPolyfill } from "../src/dom.ts" import { addHistory, getHistory } from "../src/history.ts" import { getState, + keepInHistory, onSessionRecreate, requestExport, requestSync, @@ -582,17 +583,55 @@ describe("history restore", () => { expect(getHistory(id)).toHaveLength(1) // The missing image is filled in expect(getHistory(id)[0].svg).toBe("SVG") - // draw.io's copy of it, as a sync reply brings it back - addHistory( - id, - page("same").replace( - "", - '', - ), + }) + + // draw.io's own copy of a diagram, as a sync reply or an edit brings it + const drawioCopy = (xml: string, attrs = 'pageWidth="850"') => + xml.replace( + "", + ``, ) + + it("keeps no copy of a server write that a sync brought back", async () => { + const id = "mcp-history-keep-sync" + const version = setState(id, page("ai")) + addHistory(id, page("ai")) + await postJson("/api/state", { + sessionId: id, + xml: drawioCopy(page("ai")), + baseVersion: version, + source: "sync", + stateId: getState(id)?.stateId, + }) + // The next AI write keeps the state it replaces + keepInHistory(id, getState(id)?.xml ?? "") expect(getHistory(id)).toHaveLength(1) }) + it("keeps a change of page settings only before a write", async () => { + const id = "mcp-history-keep-settings" + const version = setState(id, page("ai")) + addHistory(id, page("ai")) + // The user turns the page to A3 in the preview + await postJson("/api/state", { + sessionId: id, + xml: drawioCopy(page("ai"), 'pageWidth="1169" pageHeight="1654"'), + baseVersion: version, + stateId: getState(id)?.stateId, + }) + keepInHistory(id, getState(id)?.xml ?? "") + expect(getHistory(id)).toHaveLength(2) + expect(getHistory(id).at(-1)?.xml).toContain('pageWidth="1169"') + }) + + it("keeps a diagram restored from its file before a write", () => { + const id = "mcp-history-keep-restored" + // As restoreSavedSession puts it back after a restart + setState(id, page("from file"), undefined, false, false) + keepInHistory(id, page("from file")) + expect(getHistory(id).map((e) => e.xml)).toEqual([page("from file")]) + }) + it("keeps a version that changed only the background", () => { const id = "mcp-history-background" addHistory(id, page("same")) diff --git a/packages/mcp-server/tests/persistence.test.ts b/packages/mcp-server/tests/persistence.test.ts index 85e63199..8e4ddd66 100644 --- a/packages/mcp-server/tests/persistence.test.ts +++ b/packages/mcp-server/tests/persistence.test.ts @@ -14,11 +14,16 @@ import { } from "node:fs" import { homedir, tmpdir } from "node:os" import { join } from "node:path" -import { afterEach, describe, expect, it } from "vitest" +import { afterEach, beforeAll, describe, expect, it } from "vitest" +import { installDomPolyfill } from "../src/dom.ts" import { Autosaver, defaultDataDir } from "../src/persistence.ts" +beforeAll(() => { + installDomPolyfill() +}) + const DIAGRAM = `` -const BLANK = `` +const BLANK = `` const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)) const tempDir = () => mkdtempSync(join(tmpdir(), "mcp-autosave-")) @@ -35,6 +40,19 @@ describe("Autosaver", () => { expect(readdirSync(dir)).toEqual(["mcp-a.drawio"]) }) + it("saves a new document of empty pages the user named", async () => { + const dir = tempDir() + const saver = new Autosaver(dir, 10) + const emptyPage = (name: string) => + `` + const pages = `${emptyPage("Planning")}${emptyPage("Notes")}` + saver.schedule("mcp-pages", pages) + saver.flush() + expect( + readFileSync(saver.pathFor("mcp-pages") as string, "utf-8"), + ).toBe(pages) + }) + it("skips a blank page that was never saved, but saves clearing a diagram", async () => { const dir = tempDir() const saver = new Autosaver(dir, 10) diff --git a/tests/e2e/chat.spec.ts b/tests/e2e/chat.spec.ts index 554a8efb..13f410d1 100644 --- a/tests/e2e/chat.spec.ts +++ b/tests/e2e/chat.spec.ts @@ -5,6 +5,7 @@ import { sendMessage, test, } from "./lib/fixtures" +import { createTextOnlyResponse } from "./lib/helpers" test.describe("Chat Panel", () => { test.beforeEach(async ({ page }) => { @@ -105,3 +106,30 @@ test.describe("Crossing the mobile breakpoint", () => { await expect(getChatInput(page)).toBeVisible() }) }) + +test.describe("Sending", () => { + test("a double Enter sends the message once", async ({ page }) => { + let requests = 0 + await page.route("**/api/chat", async (route) => { + requests++ + await route.fulfill({ + status: 200, + contentType: "text/event-stream", + body: createTextOnlyResponse("Hello there."), + }) + }) + await page.goto("/", { waitUntil: "networkidle" }) + await getIframe(page).waitFor({ state: "visible", timeout: 30000 }) + const input = getChatInput(page) + await input.fill("Hi") + // The second press comes while the diagram is being exported + await input.press("ControlOrMeta+Enter") + await input.press("ControlOrMeta+Enter") + await expect(page.getByText("Hello there.")).toBeVisible({ + timeout: 10000, + }) + await page.waitForTimeout(1500) + expect(requests).toBe(1) + await expect(page.getByText("Hello there.")).toHaveCount(1) + }) +}) diff --git a/tests/unit/admin-test-model.test.ts b/tests/unit/admin-test-model.test.ts index 11801fae..f2440cc6 100644 --- a/tests/unit/admin-test-model.test.ts +++ b/tests/unit/admin-test-model.test.ts @@ -20,7 +20,13 @@ vi.mock("@/lib/admin/settings", () => ({ import { POST as testModel } from "@/app/api/admin/test-model/route" -const ENV = ["OPENAI_BASE_URL", "SGLANG_BASE_URL", "AI_GATEWAY_BASE_URL"] +const ENV = [ + "OPENAI_BASE_URL", + "SGLANG_BASE_URL", + "AI_GATEWAY_BASE_URL", + "AZURE_BASE_URL", + "AZURE_RESOURCE_NAME", +] const saved: Record = {} beforeEach(() => { envFallback.values = {} @@ -91,6 +97,15 @@ describe("admin Test of an entry without a URL", () => { } }) + it("tests Azure set up by resource name where chat goes", async () => { + process.env.AZURE_RESOURCE_NAME = "team-openai" + await test({ provider: "azure", apiKey: "k" }) + expect(sent.body.baseUrl).toBe( + "https://team-openai.openai.azure.com/openai", + ) + expect(sent.body.serverBaseUrl).toBe(true) + }) + it("tests Ollama where chat sends the entry's key", async () => { // Chat on the saved entry: OLLAMA_BASE_URL of the environment, else // the SDK's local default (the Test used to go to Ollama Cloud) diff --git a/tests/unit/chat-route-edgeone.test.ts b/tests/unit/chat-route-edgeone.test.ts index d77f2440..20991bad 100644 --- a/tests/unit/chat-route-edgeone.test.ts +++ b/tests/unit/chat-route-edgeone.test.ts @@ -119,7 +119,7 @@ describe("EdgeOne endpoints", () => { ) }) - it("keeps the deployment's base path", async () => { + it("calls the function at the site root, also with a base path", async () => { const savedPath = process.env.NEXT_PUBLIC_BASE_PATH process.env.NEXT_PUBLIC_BASE_PATH = "/draw" try { @@ -128,7 +128,7 @@ describe("EdgeOne endpoints", () => { "x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324", }) expect(calls[0]?.url).toBe( - "http://localhost/draw/api/edgeai/chat/completions", + "http://localhost/api/edgeai/chat/completions", ) } finally { if (savedPath === undefined) diff --git a/tests/unit/config-manager.test.ts b/tests/unit/config-manager.test.ts index 0ca38a3a..eaff86bf 100644 --- a/tests/unit/config-manager.test.ts +++ b/tests/unit/config-manager.test.ts @@ -3,6 +3,7 @@ import { existsSync, mkdtempSync, readdirSync, + readFileSync, rmSync, writeFileSync, } from "node:fs" @@ -19,6 +20,8 @@ vi.mock("electron", () => ({ // Make the next read of the presets file fail, like a file an antivirus // scanner holds on Windows const readFails = vi.hoisted(() => ({ next: false })) +// Make renaming the presets file fail, as when a sync tool holds it +const renameFails = vi.hoisted(() => ({ on: false })) vi.mock("node:fs", async (importOriginal) => { const fs = await importOriginal() return { @@ -32,6 +35,17 @@ vi.mock("node:fs", async (importOriginal) => { } return fs.readFileSync(...args) }) as typeof fs.readFileSync, + renameSync: ((...args: Parameters) => { + if (renameFails.on && String(args[0]).endsWith(".json")) { + throw Object.assign( + new Error("EPERM: operation not permitted"), + { + code: "EPERM", + }, + ) + } + return fs.renameSync(...args) + }) as typeof fs.renameSync, } }) @@ -42,6 +56,7 @@ const presetsFile = () => join(userData.dir, "config-presets.json") beforeEach(() => { userData.dir = mkdtempSync(join(tmpdir(), "config-manager-")) readFails.next = false + renameFails.on = false }) describe("config presets file", () => { @@ -70,6 +85,17 @@ describe("config presets file", () => { expect(loadPresets().presets.map((p) => p.name)).toEqual(["New"]) }) + it("keeps a file that is not JSON when it cannot be moved aside", () => { + writeFileSync(presetsFile(), "{not json") + renameFails.on = true + expect(loadPresets().presets).toEqual([]) + // A save based on that empty read must not replace it + expect(() => + createPreset({ name: "New", config: { AI_PROVIDER: "openai" } }), + ).toThrow() + expect(readFileSync(presetsFile(), "utf-8")).toBe("{not json") + }) + it("moves a file that is not JSON aside", () => { writeFileSync(presetsFile(), "{not json") expect(loadPresets().presets).toEqual([]) diff --git a/tests/unit/mcp-preview-recovery.test.ts b/tests/unit/mcp-preview-recovery.test.ts index 9dd504dd..cf7b4286 100644 --- a/tests/unit/mcp-preview-recovery.test.ts +++ b/tests/unit/mcp-preview-recovery.test.ts @@ -259,6 +259,29 @@ describe("MCP preview after the server recreated its session", () => { expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0) }) + it("keeps an undo when a poll sees the tab's own push first", async () => { + const t = await inStep() + t.fromDrawio({ event: "autosave", xml: "B" }) + t.fromDrawio({ event: "export", data: "" }) + await t.settle() + const pushB = t.next("POST") + // Undo back to A while B is on its way (equal to the saved A: not sent) + t.fromDrawio({ event: "autosave", xml: "A" }) + // The server already has B, and the poll's answer comes first + const loadsBefore = t.toDrawio.filter((m) => m.action === "load").length + const poll = t.page.poll() + t.next("GET").answer(state("S1", 3, "B")) + await poll + expect(t.toDrawio.filter((m) => m.action === "load")).toHaveLength( + loadsBefore, + ) + pushB.answer({ status: 200, body: { success: true, version: 3 } }) + await t.settle() + await t.settle() + // The undo is saved + expect(t.next("POST").body.xml).toBe("A") + }) + it("sends nothing more after a sync reply", async () => { const t = await inStep() const poll = t.page.poll() diff --git a/tests/unit/ssrf-protection.test.ts b/tests/unit/ssrf-protection.test.ts index 0ff98d11..80b3c1ae 100644 --- a/tests/unit/ssrf-protection.test.ts +++ b/tests/unit/ssrf-protection.test.ts @@ -177,6 +177,9 @@ describe("redirectGuardedFetch with the quota on", () => { headers: { Authorization: "Bearer user-key", Cookie: "eo_token=1", + "x-api-key": "anthropic-key", + "x-goog-api-key": "google-key", + "api-key": "azure-key", "Content-Type": "application/json", }, }) @@ -192,6 +195,9 @@ describe("redirectGuardedFetch with the quota on", () => { expect(sent(0).get("authorization")).toBe("Bearer user-key") expect(sent(1).get("authorization")).toBeNull() expect(sent(1).get("cookie")).toBeNull() + for (const name of ["x-api-key", "x-goog-api-key", "api-key"]) { + expect(sent(1).get(name)).toBeNull() + } expect(sent(1).get("content-type")).toBe("application/json") }) diff --git a/tests/unit/use-diagram-tool-handlers.test.tsx b/tests/unit/use-diagram-tool-handlers.test.tsx index e247c8ff..56e64a45 100644 --- a/tests/unit/use-diagram-tool-handlers.test.tsx +++ b/tests/unit/use-diagram-tool-handlers.test.tsx @@ -50,6 +50,8 @@ describe("the screenshot check and Stop", () => { watchStop: () => () => boolean validateDiagram: () => Promise captureValidationPng?: () => Promise + // Checks already made in this user turn + retryCount?: { current: number } }) => { const onValidationStateChange = vi.fn() const { result } = renderHook(() => @@ -57,7 +59,7 @@ describe("the screenshot check and Stop", () => { partialXmlRef: { current: "" }, editDiagramOriginalXmlRef: { current: new Map() }, processedToolCallsRef: { current: new Set() }, - validationRetryCountRef: { current: 0 }, + validationRetryCountRef: opts.retryCount ?? { current: 0 }, chartXMLRef: { current: "" }, onDisplayChart: () => null, onFetchChart: async () => "", @@ -118,6 +120,24 @@ describe("the screenshot check and Stop", () => { expect(addToolOutput.mock.lastCall?.[0].state).toBeUndefined() }) + it("checks at most three diagrams in one user turn", async () => { + const validateDiagram = vi.fn(async () => ({ + valid: true, + issues: [], + suggestions: [], + })) + const retryCount = { current: 0 } + for (let i = 0; i < 4; i++) { + await draw({ + watchStop: () => () => false, + validateDiagram, + retryCount, + }) + } + // Passed checks count too + expect(validateDiagram).toHaveBeenCalledTimes(3) + }) + it("skips the check when Stop came during the screenshot", async () => { // As the chat panel counts it: the next message already cleared // the stop flag when the screenshot arrives diff --git a/tests/unit/utils.test.ts b/tests/unit/utils.test.ts index 2560274c..381b7e94 100644 --- a/tests/unit/utils.test.ts +++ b/tests/unit/utils.test.ts @@ -44,6 +44,12 @@ describe("isMxCellXmlComplete", () => { const xml = '' expect(isMxCellXmlComplete(xml)).toBe(true) + // A raw "<" in a value (escaped later by the auto-fix) + expect( + isMxCellXmlComplete( + '', + ), + ).toBe(true) }) it("returns true for mxCell with closing tag", () => { diff --git a/tests/unit/validate-model-route.test.ts b/tests/unit/validate-model-route.test.ts index d11cdeca..5a224e07 100644 --- a/tests/unit/validate-model-route.test.ts +++ b/tests/unit/validate-model-route.test.ts @@ -18,8 +18,22 @@ vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({ isPrivateUrl: async () => privateUrls.all, })) +// The quota, off unless a test turns it on; every request is refused +const quota = vi.hoisted(() => ({ enabled: false, checks: 0 })) +vi.mock("@/lib/dynamo-quota-manager", () => ({ + isQuotaEnabled: () => quota.enabled, + checkAndIncrementRequest: async () => { + quota.checks++ + return { allowed: false, error: "Daily limit reached" } + }, +})) +vi.mock("@/lib/user-id", () => ({ getUserIdFromRequest: () => "user-1" })) + afterEach(() => { delete process.env.ALLOW_PRIVATE_URLS + quota.enabled = false + quota.checks = 0 + privateUrls.all = false vi.unstubAllGlobals() }) @@ -314,3 +328,41 @@ describe("the admin panel's Test button", () => { } }) }) + +describe("the Test on the deployment's own endpoints", () => { + const test = (body: object) => + validateModel( + new Request("http://localhost/api/validate-model", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ modelId: "m", ...body }), + }), + ) + + it("counts as a chat request with the quota on", async () => { + quota.enabled = true + // EdgeOne, and a model server on the server's network with a dummy key + const edgeone = await test({ provider: "edgeone" }) + expect(edgeone.status).toBe(429) + privateUrls.all = true + const internal = await test({ + provider: "openai", + apiKey: "x", + baseUrl: "http://10.0.0.5:8000/v1", + }) + expect(internal.status).toBe(429) + expect(quota.checks).toBe(2) + }) + + it("does not count a user's own endpoint", async () => { + quota.enabled = true + streamReply({ role: "assistant", content: "OK" }) + const res = await test({ + provider: "openai", + apiKey: "user-key", + baseUrl: "https://api.example.com/v1", + }) + expect(res.status).toBe(200) + expect(quota.checks).toBe(0) + }) +})