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.
This commit is contained in:
dayuan.jiang
2026-10-05 21:19:32 +09:00
parent 69c7896002
commit b8b851fc5b
29 changed files with 476 additions and 90 deletions
+6 -11
View File
@@ -19,6 +19,7 @@ import {
SINGLE_SYSTEM_PROVIDERS, SINGLE_SYSTEM_PROVIDERS,
supportsPromptCaching, supportsPromptCaching,
usesServerCredentials, usesServerCredentials,
usesServerEndpoint,
} from "@/lib/ai-providers" } from "@/lib/ai-providers"
import { findCachedResponse } from "@/lib/cached-responses" import { findCachedResponse } from "@/lib/cached-responses"
import { import {
@@ -287,17 +288,11 @@ async function handleChatRequest(req: Request): Promise<Response> {
const clientBaseUrl = normalizeBaseUrl( const clientBaseUrl = normalizeBaseUrl(
req.headers.get("x-ai-base-url") ?? "", req.headers.get("x-ai-base-url") ?? "",
) )
const usesClientBaseUrl = const onServerEndpoint = await usesServerEndpoint(
resolvedProvider !== "bedrock" && resolvedProvider !== "edgeone" resolvedProvider,
const onServerEndpoint = clientBaseUrl,
process.env.NEXT_AI_DRAWIO_DESKTOP !== "1" && clientOverrides.apiKey,
(resolvedProvider === "edgeone" || )
(resolvedProvider === "ollama" &&
!clientBaseUrl &&
!clientOverrides.apiKey) ||
(usesClientBaseUrl &&
!!clientBaseUrl &&
(await isPrivateUrl(clientBaseUrl))))
const countsQuota = const countsQuota =
isQuotaEnabled() && isQuotaEnabled() &&
(onServerCredentials || onServerEndpoint) && (onServerCredentials || onServerEndpoint) &&
+34 -1
View File
@@ -8,10 +8,16 @@ import {
getAIModel, getAIModel,
globalBaseUrl, globalBaseUrl,
usesServerCredentials, usesServerCredentials,
usesServerEndpoint,
} from "@/lib/ai-providers" } from "@/lib/ai-providers"
import {
checkAndIncrementRequest,
isQuotaEnabled,
} from "@/lib/dynamo-quota-manager"
import { classifyLLMError } from "@/lib/llm-errors" import { classifyLLMError } from "@/lib/llm-errors"
import { allowPrivateUrls, isPrivateUrl } from "@/lib/ssrf-protection" 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" 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 // The same model the chat would use. A client base URL makes it
// refuse redirects to internal hosts. // refuse redirects to internal hosts.
const { model } = getAIModel({ const { model } = getAIModel({
+22 -7
View File
@@ -277,6 +277,7 @@ export default function ChatPanel({
// Set by Stop until the user sends the next message // Set by Stop until the user sends the next message
const stoppedRef = useRef(false) const stoppedRef = useRef(false)
const preparingSendRef = useRef(false)
// Presses of Stop: a check that began before one still knows of it after // Presses of Stop: a check that began before one still knows of it after
// the next message clears stoppedRef // the next message clears stoppedRef
const stopCountRef = useRef(0) 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) // Track last synced session ID to detect external changes (e.g., URL back/forward)
const lastSyncedSessionIdRef = useRef<string | null>(null) const lastSyncedSessionIdRef = useRef<string | null>(null)
// Messages array from our latest save. A session holding this exact array was // 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. // created by our own save, so it must not be treated as an external
const lastSavedMessagesRef = useRef<unknown[] | null>(null) // switch (with two saves of a new chat at once, the first creates it).
const savedMessagesRef = useRef(new WeakSet<object>())
// Helper: Sync UI state with session data (eliminates duplication) // Helper: Sync UI state with session data (eliminates duplication)
// Track message IDs that are being loaded from session (to skip animations/scroll) // 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) const messages = sanitizeMessages(messagesRef.current)
lastSavedMessagesRef.current = messages savedMessagesRef.current.add(messages)
return { return {
messages, messages,
xmlSnapshots: Array.from(xmlSnapshotsRef.current.entries()), 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 // Our own save created this session; the UI already shows its content
const isOwnNewSession = const isOwnNewSession =
newSession?.messages === lastSavedMessagesRef.current !!newSession && savedMessagesRef.current.has(newSession.messages)
// Update last synced ID // Update last synced ID
lastSyncedSessionIdRef.current = newSessionId lastSyncedSessionIdRef.current = newSessionId
@@ -802,8 +804,7 @@ export default function ChatPanel({
) )
}, [sessionManager, buildSessionData]) }, [sessionManager, buildSessionData])
const onFormSubmit = async (e: React.FormEvent<HTMLFormElement>) => { const submitInput = async () => {
e.preventDefault()
const isProcessing = status === "streaming" || status === "submitted" const isProcessing = status === "streaming" || status === "submitted"
// Attachments still extracting have no text yet. Template sends call // Attachments still extracting have no text yet. Template sends call
// requestSubmit() and skip the disabled send button, so check here too. // 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<HTMLFormElement>) => {
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. // Export the current diagram, snapshot it for this message, and send.
// onSent runs right after sending, so the input empties as the message // onSent runs right after sending, so the input empties as the message
// shows in the chat // shows in the chat
+1 -1
View File
@@ -422,7 +422,7 @@ export function ModelConfigDialog({
// For EdgeOne, construct baseUrl from current origin // For EdgeOne, construct baseUrl from current origin
const baseUrl = isEdgeOne const baseUrl = isEdgeOne
? `${window.location.origin}${getApiEndpoint("/api/edgeai")}` ? `${window.location.origin}/api/edgeai`
: selectedProvider.baseUrl : selectedProvider.baseUrl
// Test every model at once; each row updates when its answer arrives // Test every model at once; each row updates when its answer arrives
+2
View File
@@ -217,7 +217,9 @@ export function loadPresets(): ConfigPresetsFile {
renameSync(configPath, backupPath) renameSync(configPath, backupPath)
console.error(`Unreadable config presets moved to ${backupPath}`) console.error(`Unreadable config presets moved to ${backupPath}`)
} catch (renameError) { } catch (renameError) {
// Still there: refuse saves that would overwrite it
console.error("Failed to back up config presets:", renameError) console.error("Failed to back up config presets:", renameError)
presetsUnreadable = true
} }
return { return {
version: 1, version: 1,
+4 -2
View File
@@ -124,9 +124,11 @@ if (!gotTheLock) {
}) })
// Clicking a plain link would otherwise replace the app page with // 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) => { contents.on("will-navigate", (event) => {
if (isInAppUrl(event.url)) { if (isAppUrl(event.url)) {
return return
} }
event.preventDefault() event.preventDefault()
+6 -8
View File
@@ -251,7 +251,9 @@ ${finalXml}
if ( if (
enableVlmValidation && enableVlmValidation &&
captureValidationPng && captureValidationPng &&
validateDiagram validateDiagram &&
// At most this many checks per user turn, passed or not
validationRetryCountRef.current < MAX_VALIDATION_RETRIES
) { ) {
let capturedPngData: string | null = null let capturedPngData: string | null = null
const stopped = watchStop?.() 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 const attempt = validationRetryCountRef.current + 1
validationRetryCountRef.current = attempt
// Notify UI that we're validating (include the image) // Notify UI that we're validating (include the image)
updateValidationState( updateValidationState(
@@ -296,8 +300,6 @@ ${finalXml}
if (!result.valid) { if (!result.valid) {
if (attempt < MAX_VALIDATION_RETRIES) { if (attempt < MAX_VALIDATION_RETRIES) {
validationRetryCountRef.current = attempt
const feedback = const feedback =
formatValidationFeedback(result) formatValidationFeedback(result)
if (DEBUG) { if (DEBUG) {
@@ -333,8 +335,6 @@ ${finalXml}
"[display_diagram] Max validation retries reached, accepting diagram", "[display_diagram] Max validation retries reached, accepting diagram",
) )
} }
validationRetryCountRef.current = 0
// Notify UI that we're accepting with issues (include the image) // Notify UI that we're accepting with issues (include the image)
updateValidationState( updateValidationState(
toolCall.toolCallId, toolCall.toolCallId,
@@ -350,8 +350,6 @@ ${finalXml}
return return
} }
} else { } else {
// Validation passed - reset retry count
validationRetryCountRef.current = 0
if (DEBUG) { if (DEBUG) {
console.log( console.log(
"[display_diagram] Validation passed!", "[display_diagram] Validation passed!",
+34 -4
View File
@@ -22,8 +22,7 @@ import {
loadAdminProviders, loadAdminProviders,
} from "@/lib/admin/providers" } from "@/lib/admin/providers"
import { getEnvFallback } from "@/lib/admin/settings" import { getEnvFallback } from "@/lib/admin/settings"
import { getApiEndpoint } from "@/lib/base-path" import { isPrivateUrl, redirectGuardedFetch } from "@/lib/ssrf-protection"
import { redirectGuardedFetch } from "@/lib/ssrf-protection"
import { import {
normalizeBaseUrl, normalizeBaseUrl,
PROVIDER_INFO, PROVIDER_INFO,
@@ -1112,11 +1111,12 @@ export function getAIModel(clientOverrides?: ClientOverrides): ModelConfig {
/** /**
* The deployment's EdgeOne Pages function, as an absolute URL (the SDK * 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 { export function edgeOneEndpoint(req: Request): string {
const origin = req.headers.get("origin") || new URL(req.url).origin 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 return getEnvFallback("GOOGLE_VERTEX_BASE_URL") || undefined
} }
if (provider === "bedrock" || provider === "edgeone") return 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 = const name =
provider === "gateway" provider === "gateway"
? "AI_GATEWAY_BASE_URL" ? "AI_GATEWAY_BASE_URL"
@@ -1148,6 +1156,28 @@ export function getServerProvider(): ProviderName | null {
return (process.env.AI_PROVIDER as ProviderName) || detectProvider() 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<boolean> {
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 * 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 * IAM role) rather than credentials sent with the request. Mirrors which key
+13 -3
View File
@@ -126,6 +126,18 @@ export class RedirectRefusedError extends Error {
const MAX_REDIRECTS = 5 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 * 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 * 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) // credentials (the user's key, EdgeOne's cookies)
const headers = new Headers(next?.headers) const headers = new Headers(next?.headers)
if (new URL(url).origin !== from.origin) { if (new URL(url).origin !== from.origin) {
headers.delete("authorization") for (const name of CREDENTIAL_HEADERS) headers.delete(name)
headers.delete("proxy-authorization")
headers.delete("cookie")
} }
next = { ...next, headers } next = { ...next, headers }
// 303, and 301 or 302 after a POST, go on as a GET without the // 303, and 301 or 302 after a POST, go on as a GET without the
+4 -1
View File
@@ -59,7 +59,10 @@ export function isMxCellXmlComplete(xml: string | undefined | null): boolean {
const lastCellStart = trimmed.lastIndexOf("<mxCell") const lastCellStart = trimmed.lastIndexOf("<mxCell")
if ( if (
lastCellStart > lastMxCellClose && lastCellStart > lastMxCellClose &&
!/^<mxCell\b[^<]*\/>/.test(trimmed.slice(lastCellStart)) // (quoted values may hold a raw "<", which the auto-fix escapes)
!/^<mxCell\b(?:[^<>"']|"[^"]*"|'[^']*')*\/>/.test(
trimmed.slice(lastCellStart),
)
) { ) {
return false return false
} }
+1 -1
View File
@@ -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 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 ## Troubleshooting
+26 -1
View File
@@ -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 &#xa; 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 { class XMLSerializerPolyfill {
serializeToString(node: Node): string { serializeToString(node: Node): string {
return serializeXml(node) return serializeXml(node)
@@ -68,6 +93,6 @@ class XMLSerializerPolyfill {
/** Install the DOMParser and XMLSerializer globals the XML helpers use. */ /** Install the DOMParser and XMLSerializer globals the XML helpers use. */
export function installDomPolyfill(): void { export function installDomPolyfill(): void {
;(globalThis as any).DOMParser = DOMParser ;(globalThis as any).DOMParser = XmlDomParser
;(globalThis as any).XMLSerializer = XMLSerializerPolyfill ;(globalThis as any).XMLSerializer = XMLSerializerPolyfill
} }
+4 -29
View File
@@ -3,9 +3,7 @@
* Stores {xml, svg} entries in a circular buffer * Stores {xml, svg} entries in a circular buffer
*/ */
import { contentFingerprint } from "./edit-gate.ts"
import { log } from "./logger.ts" import { log } from "./logger.ts"
import { normalizeToMxfile, parseMxfile } from "./pages.ts"
const MAX_HISTORY = 20 const MAX_HISTORY = 20
@@ -18,30 +16,6 @@ interface HistoryEntry {
let nextEntryId = 0 let nextEntryId = 0
const historyStore = new Map<string, HistoryEntry[]>() const historyStore = new Map<string, HistoryEntry[]>()
/** 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 { export function addHistory(sessionId: string, xml: string, svg = ""): number {
let history = historyStore.get(sessionId) let history = historyStore.get(sessionId)
if (!history) { if (!history) {
@@ -49,10 +23,11 @@ export function addHistory(sessionId: string, xml: string, svg = ""): number {
historyStore.set(sessionId, history) historyStore.set(sessionId, history)
} }
// Dedupe: skip if same as last entry, also when only re-serialized // Dedupe: skip if same as last entry (any other change, also of page
// (a change of background only is a new version) // settings only, is a new version; keepInHistory leaves out draw.io's
// copy of a server write)
const last = history[history.length - 1] const last = history[history.length - 1]
if (last && sameDiagram(last.xml, xml)) { if (last && last.xml === xml) {
if (svg && !last.svg) last.svg = svg if (svg && !last.svg) last.svg = svg
return history.length - 1 return history.length - 1
} }
+22 -1
View File
@@ -138,6 +138,9 @@ interface SessionState {
serverVersion?: number serverVersion?: number
// The XML of that write: what a thumbnail taken after loading it shows // The XML of that write: what a thumbnail taken after loading it shows
serverXml?: string serverXml?: string
// The browser saved a change of the user's since that write (a sync
// reply is no change)
userEdited?: boolean
lastUpdated: Date lastUpdated: Date
lastPolled?: number // Last browser poll; an open tab keeps the session alive lastPolled?: number // Last browser poll; an open tab keeps the session alive
svg?: string // Cached SVG from last browser save svg?: string // Cached SVG from last browser save
@@ -206,6 +209,7 @@ export function setState(
stateId: existing?.stateId ?? randomUUID(), stateId: existing?.stateId ?? randomUUID(),
serverVersion: fromBrowser ? existing?.serverVersion : newVersion, serverVersion: fromBrowser ? existing?.serverVersion : newVersion,
serverXml: fromBrowser ? existing?.serverXml : xml, serverXml: fromBrowser ? existing?.serverXml : xml,
userEdited: fromBrowser ? existing?.userEdited : false,
lastUpdated: new Date(), lastUpdated: new Date(),
lastPolled: existing?.lastPolled, lastPolled: existing?.lastPolled,
// The image of this XML, never an older one's: a write without an // The image of this XML, never an older one's: a write without an
@@ -226,6 +230,21 @@ export function setState(
return newVersion 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. * Ask the browser bridge to export the current diagram as png/svg.
* *
@@ -604,6 +623,8 @@ function handleStateApi(
return return
} }
const version = setState(sessionId, data.xml, data.svg, true) 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.writeHead(200, { "Content-Type": "application/json" })
res.end(JSON.stringify({ success: true, version })) res.end(JSON.stringify({ success: true, version }))
} catch { } catch {
@@ -685,7 +706,7 @@ function handleRestoreApi(
contentFingerprint(current.xml) !== contentFingerprint(current.xml) !==
contentFingerprint(BLANK_MXFILE) contentFingerprint(BLANK_MXFILE)
) { ) {
addHistory(sessionId, current.xml, current.svg) keepInHistory(sessionId, current.xml, current.svg)
} }
const newVersion = setState(sessionId, entry.xml) const newVersion = setState(sessionId, entry.xml)
addHistory(sessionId, entry.xml, entry.svg) addHistory(sessionId, entry.xml, entry.svg)
+14 -4
View File
@@ -35,6 +35,7 @@ import {
type ExportOptions, type ExportOptions,
getServerPort, getServerPort,
getState, getState,
keepInHistory,
onSessionRecreate, onSessionRecreate,
onStateChange, onStateChange,
requestExport, 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) // Save user's state before AI overwrites (with cached SVG)
if (currentSession.xml) { if (currentSession.xml) {
addHistory( keepInHistory(
currentSession.id, currentSession.id,
currentSession.xml, currentSession.xml,
browserState?.svg || "", browserState?.svg || "",
@@ -481,6 +482,11 @@ registerWriteTool(
let content: string let content: string
try { 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") content = await fs.readFile(absolutePath, "utf-8")
} catch (e) { } catch (e) {
const msg = e instanceof Error ? e.message : String(e) const msg = e instanceof Error ? e.message : String(e)
@@ -515,7 +521,7 @@ registerWriteTool(
currentSession.xml = browserState.xml currentSession.xml = browserState.xml
} }
if (currentSession.xml) { if (currentSession.xml) {
addHistory( keepInHistory(
currentSession.id, currentSession.id,
currentSession.xml, currentSession.xml,
browserState?.svg || "", browserState?.svg || "",
@@ -728,7 +734,7 @@ registerWriteTool(
// Save the pre-edit state for undo (with cached SVG from browser). // Save the pre-edit state for undo (with cached SVG from browser).
// Done only once the edit applied: a rejected edit returns above // Done only once the edit applied: a rejected edit returns above
// without leaving a phantom history entry. // without leaving a phantom history entry.
addHistory( keepInHistory(
currentSession.id, currentSession.id,
currentSession.xml, currentSession.xml,
browserState?.svg || "", browserState?.svg || "",
@@ -1456,7 +1462,11 @@ async function loadMxfileForMutation(): Promise<
browserState?.xml ?? "", browserState?.xml ?? "",
).ok ).ok
// Save history before overwriting so the user can undo. // 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.xml = newXml
sessionRef.version++ sessionRef.version++
setState(sessionRef.id, newXml) setState(sessionRef.id, newXml)
+9 -2
View File
@@ -18,8 +18,15 @@ import {
} from "node:fs" } from "node:fs"
import { homedir } from "node:os" import { homedir } from "node:os"
import { join } from "node:path" import { join } from "node:path"
import { contentFingerprint } from "./edit-gate.ts"
import { log } from "./logger.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 DELAY_MS = 1000
const MAX_FILES = 50 const MAX_FILES = 50
@@ -126,7 +133,7 @@ export class Autosaver {
try { try {
const isNew = !existsSync(path) const isNew = !existsSync(path)
// A blank page the browser shows before any drawing: nothing to keep // 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 }) mkdirSync(this.dir, { recursive: true })
// Write to a temporary file first so a crash never leaves half a file // Write to a temporary file first so a crash never leaves half a file
writeFileSync(`${path}.tmp`, entry.xml, "utf-8") writeFileSync(`${path}.tmp`, entry.xml, "utf-8")
+8 -1
View File
@@ -7,6 +7,7 @@ let stateId = null;
// the last one the server has // the last one the server has
let latestXml = null; let latestXml = null;
let pushFailing = false; // the last push could not reach the server 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 // After recovery replaced the canvas, until draw.io reports the load: an
// autosave still on its way belongs to the canvas being replaced // autosave still on its way belongs to the canvas being replaced
let awaitingLoad = false; let awaitingLoad = false;
@@ -181,6 +182,7 @@ function showNotice(text) {
// server state the push is based on. // server state the push is based on.
async function pushState(xml, svg = '', baseVersion = currentVersion, source = 'edit', sid = stateId) { async function pushState(xml, svg = '', baseVersion = currentVersion, source = 'edit', sid = stateId) {
if (!sessionId) return; if (!sessionId) return;
pushesInFlight.push(xml);
try { try {
const r = await fetch('/api/state', { const r = await fetch('/api/state', {
method: 'POST', method: 'POST',
@@ -226,6 +228,8 @@ async function pushState(xml, svg = '', baseVersion = currentVersion, source = '
pushFailing = true; pushFailing = true;
showNotice("Can't reach the MCP server. Your changes are only in this tab for now; use Download to keep a copy."); 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 // (forceReload) replaces it, so a new version doesn't fight the
// projection; currentVersion stays unadvanced until then, so the // projection; currentVersion stays unadvanced until then, so the
// bump is applied with the real document. // 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; forceReload = false;
projectionExportActive = false; projectionExportActive = false;
currentVersion = s.version; currentVersion = s.version;
@@ -141,3 +141,23 @@ describe("targetPageXml", () => {
expect(xml).not.toContain('id="a"') 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; &#xa; is a
// real line break
const labels =
`<mxCell id="s" value="Hello
world" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell>` +
`<mxCell id="m" value="Line 1&#xa;Line 2" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell>`
const out = editDiagram(
`<mxfile>${page("p1", labels)}</mxfile>`,
[{ 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&#xa;Line 2"`)
})
})
+46 -7
View File
@@ -12,6 +12,7 @@ import { installDomPolyfill } from "../src/dom.ts"
import { addHistory, getHistory } from "../src/history.ts" import { addHistory, getHistory } from "../src/history.ts"
import { import {
getState, getState,
keepInHistory,
onSessionRecreate, onSessionRecreate,
requestExport, requestExport,
requestSync, requestSync,
@@ -582,17 +583,55 @@ describe("history restore", () => {
expect(getHistory(id)).toHaveLength(1) expect(getHistory(id)).toHaveLength(1)
// The missing image is filled in // The missing image is filled in
expect(getHistory(id)[0].svg).toBe("SVG") expect(getHistory(id)[0].svg).toBe("SVG")
// draw.io's copy of it, as a sync reply brings it back })
addHistory(
id, // draw.io's own copy of a diagram, as a sync reply or an edit brings it
page("same").replace( const drawioCopy = (xml: string, attrs = 'pageWidth="850"') =>
"<mxGraphModel>", xml.replace(
'<mxGraphModel dx="1244" dy="534" grid="1" pageWidth="850">', "<mxGraphModel>",
), `<mxGraphModel dx="1244" dy="534" grid="1" ${attrs}>`,
) )
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) 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", () => { it("keeps a version that changed only the background", () => {
const id = "mcp-history-background" const id = "mcp-history-background"
addHistory(id, page("same")) addHistory(id, page("same"))
+20 -2
View File
@@ -14,11 +14,16 @@ import {
} from "node:fs" } from "node:fs"
import { homedir, tmpdir } from "node:os" import { homedir, tmpdir } from "node:os"
import { join } from "node:path" 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" import { Autosaver, defaultDataDir } from "../src/persistence.ts"
beforeAll(() => {
installDomPolyfill()
})
const DIAGRAM = `<mxfile><diagram id="p" name="P"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="a" vertex="1" parent="1"/></root></mxGraphModel></diagram></mxfile>` const DIAGRAM = `<mxfile><diagram id="p" name="P"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="a" vertex="1" parent="1"/></root></mxGraphModel></diagram></mxfile>`
const BLANK = `<mxfile><diagram id="p" name="P"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram></mxfile>` const BLANK = `<mxfile><diagram id="page-1" name="Page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram></mxfile>`
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)) const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms))
const tempDir = () => mkdtempSync(join(tmpdir(), "mcp-autosave-")) const tempDir = () => mkdtempSync(join(tmpdir(), "mcp-autosave-"))
@@ -35,6 +40,19 @@ describe("Autosaver", () => {
expect(readdirSync(dir)).toEqual(["mcp-a.drawio"]) 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) =>
`<diagram id="${name}" name="${name}"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram>`
const pages = `<mxfile>${emptyPage("Planning")}${emptyPage("Notes")}</mxfile>`
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 () => { it("skips a blank page that was never saved, but saves clearing a diagram", async () => {
const dir = tempDir() const dir = tempDir()
const saver = new Autosaver(dir, 10) const saver = new Autosaver(dir, 10)
+28
View File
@@ -5,6 +5,7 @@ import {
sendMessage, sendMessage,
test, test,
} from "./lib/fixtures" } from "./lib/fixtures"
import { createTextOnlyResponse } from "./lib/helpers"
test.describe("Chat Panel", () => { test.describe("Chat Panel", () => {
test.beforeEach(async ({ page }) => { test.beforeEach(async ({ page }) => {
@@ -105,3 +106,30 @@ test.describe("Crossing the mobile breakpoint", () => {
await expect(getChatInput(page)).toBeVisible() 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)
})
})
+16 -1
View File
@@ -20,7 +20,13 @@ vi.mock("@/lib/admin/settings", () => ({
import { POST as testModel } from "@/app/api/admin/test-model/route" 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<string, string | undefined> = {} const saved: Record<string, string | undefined> = {}
beforeEach(() => { beforeEach(() => {
envFallback.values = {} 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 () => { it("tests Ollama where chat sends the entry's key", async () => {
// Chat on the saved entry: OLLAMA_BASE_URL of the environment, else // 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) // the SDK's local default (the Test used to go to Ollama Cloud)
+2 -2
View File
@@ -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 const savedPath = process.env.NEXT_PUBLIC_BASE_PATH
process.env.NEXT_PUBLIC_BASE_PATH = "/draw" process.env.NEXT_PUBLIC_BASE_PATH = "/draw"
try { try {
@@ -128,7 +128,7 @@ describe("EdgeOne endpoints", () => {
"x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324", "x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
}) })
expect(calls[0]?.url).toBe( expect(calls[0]?.url).toBe(
"http://localhost/draw/api/edgeai/chat/completions", "http://localhost/api/edgeai/chat/completions",
) )
} finally { } finally {
if (savedPath === undefined) if (savedPath === undefined)
+26
View File
@@ -3,6 +3,7 @@ import {
existsSync, existsSync,
mkdtempSync, mkdtempSync,
readdirSync, readdirSync,
readFileSync,
rmSync, rmSync,
writeFileSync, writeFileSync,
} from "node:fs" } from "node:fs"
@@ -19,6 +20,8 @@ vi.mock("electron", () => ({
// Make the next read of the presets file fail, like a file an antivirus // Make the next read of the presets file fail, like a file an antivirus
// scanner holds on Windows // scanner holds on Windows
const readFails = vi.hoisted(() => ({ next: false })) 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) => { vi.mock("node:fs", async (importOriginal) => {
const fs = await importOriginal<typeof import("node:fs")>() const fs = await importOriginal<typeof import("node:fs")>()
return { return {
@@ -32,6 +35,17 @@ vi.mock("node:fs", async (importOriginal) => {
} }
return fs.readFileSync(...args) return fs.readFileSync(...args)
}) as typeof fs.readFileSync, }) as typeof fs.readFileSync,
renameSync: ((...args: Parameters<typeof fs.renameSync>) => {
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(() => { beforeEach(() => {
userData.dir = mkdtempSync(join(tmpdir(), "config-manager-")) userData.dir = mkdtempSync(join(tmpdir(), "config-manager-"))
readFails.next = false readFails.next = false
renameFails.on = false
}) })
describe("config presets file", () => { describe("config presets file", () => {
@@ -70,6 +85,17 @@ describe("config presets file", () => {
expect(loadPresets().presets.map((p) => p.name)).toEqual(["New"]) 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", () => { it("moves a file that is not JSON aside", () => {
writeFileSync(presetsFile(), "{not json") writeFileSync(presetsFile(), "{not json")
expect(loadPresets().presets).toEqual([]) expect(loadPresets().presets).toEqual([])
+23
View File
@@ -259,6 +259,29 @@ describe("MCP preview after the server recreated its session", () => {
expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0) 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: "<mxfile>B</mxfile>" })
t.fromDrawio({ event: "export", data: "<svg/>" })
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: "<mxfile>A</mxfile>" })
// 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, "<mxfile>B</mxfile>"))
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("<mxfile>A</mxfile>")
})
it("sends nothing more after a sync reply", async () => { it("sends nothing more after a sync reply", async () => {
const t = await inStep() const t = await inStep()
const poll = t.page.poll() const poll = t.page.poll()
+6
View File
@@ -177,6 +177,9 @@ describe("redirectGuardedFetch with the quota on", () => {
headers: { headers: {
Authorization: "Bearer user-key", Authorization: "Bearer user-key",
Cookie: "eo_token=1", Cookie: "eo_token=1",
"x-api-key": "anthropic-key",
"x-goog-api-key": "google-key",
"api-key": "azure-key",
"Content-Type": "application/json", "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(0).get("authorization")).toBe("Bearer user-key")
expect(sent(1).get("authorization")).toBeNull() expect(sent(1).get("authorization")).toBeNull()
expect(sent(1).get("cookie")).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") expect(sent(1).get("content-type")).toBe("application/json")
}) })
+21 -1
View File
@@ -50,6 +50,8 @@ describe("the screenshot check and Stop", () => {
watchStop: () => () => boolean watchStop: () => () => boolean
validateDiagram: () => Promise<any> validateDiagram: () => Promise<any>
captureValidationPng?: () => Promise<string> captureValidationPng?: () => Promise<string>
// Checks already made in this user turn
retryCount?: { current: number }
}) => { }) => {
const onValidationStateChange = vi.fn() const onValidationStateChange = vi.fn()
const { result } = renderHook(() => const { result } = renderHook(() =>
@@ -57,7 +59,7 @@ describe("the screenshot check and Stop", () => {
partialXmlRef: { current: "" }, partialXmlRef: { current: "" },
editDiagramOriginalXmlRef: { current: new Map() }, editDiagramOriginalXmlRef: { current: new Map() },
processedToolCallsRef: { current: new Set() }, processedToolCallsRef: { current: new Set() },
validationRetryCountRef: { current: 0 }, validationRetryCountRef: opts.retryCount ?? { current: 0 },
chartXMLRef: { current: "" }, chartXMLRef: { current: "" },
onDisplayChart: () => null, onDisplayChart: () => null,
onFetchChart: async () => "", onFetchChart: async () => "",
@@ -118,6 +120,24 @@ describe("the screenshot check and Stop", () => {
expect(addToolOutput.mock.lastCall?.[0].state).toBeUndefined() 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 () => { it("skips the check when Stop came during the screenshot", async () => {
// As the chat panel counts it: the next message already cleared // As the chat panel counts it: the next message already cleared
// the stop flag when the screenshot arrives // the stop flag when the screenshot arrives
+6
View File
@@ -44,6 +44,12 @@ describe("isMxCellXmlComplete", () => {
const xml = const xml =
'<mxCell id="2" value="Hello" style="rounded=1;" vertex="1" parent="1"/>' '<mxCell id="2" value="Hello" style="rounded=1;" vertex="1" parent="1"/>'
expect(isMxCellXmlComplete(xml)).toBe(true) expect(isMxCellXmlComplete(xml)).toBe(true)
// A raw "<" in a value (escaped later by the auto-fix)
expect(
isMxCellXmlComplete(
'<mxCell id="3" value="<b>Title</b>" style="text;html=1;" vertex="1" parent="1"/>',
),
).toBe(true)
}) })
it("returns true for mxCell with closing tag", () => { it("returns true for mxCell with closing tag", () => {
+52
View File
@@ -18,8 +18,22 @@ vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
isPrivateUrl: async () => privateUrls.all, 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(() => { afterEach(() => {
delete process.env.ALLOW_PRIVATE_URLS delete process.env.ALLOW_PRIVATE_URLS
quota.enabled = false
quota.checks = 0
privateUrls.all = false
vi.unstubAllGlobals() 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)
})
})