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) + }) +})