mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-07 02:07:47 +08:00
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
+6
-11
@@ -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<Response> {
|
||||
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) &&
|
||||
|
||||
@@ -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({
|
||||
|
||||
@@ -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<string | null>(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<unknown[] | null>(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<object>())
|
||||
|
||||
// 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<HTMLFormElement>) => {
|
||||
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<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.
|
||||
// onSent runs right after sending, so the input empties as the message
|
||||
// shows in the chat
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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!",
|
||||
|
||||
+34
-4
@@ -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<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
|
||||
* IAM role) rather than credentials sent with the request. Mirrors which key
|
||||
|
||||
+13
-3
@@ -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
|
||||
|
||||
+4
-1
@@ -59,7 +59,10 @@ export function isMxCellXmlComplete(xml: string | undefined | null): boolean {
|
||||
const lastCellStart = trimmed.lastIndexOf("<mxCell")
|
||||
if (
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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<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 {
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 =
|
||||
`<mxCell id="s" value="Hello
|
||||
world" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell>` +
|
||||
`<mxCell id="m" value="Line 1
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
Line 2"`)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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(
|
||||
"<mxGraphModel>",
|
||||
'<mxGraphModel dx="1244" dy="534" grid="1" pageWidth="850">',
|
||||
),
|
||||
})
|
||||
|
||||
// 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(
|
||||
"<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)
|
||||
})
|
||||
|
||||
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"))
|
||||
|
||||
@@ -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 = `<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 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) =>
|
||||
`<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 () => {
|
||||
const dir = tempDir()
|
||||
const saver = new Autosaver(dir, 10)
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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<string, string | undefined> = {}
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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<typeof import("node:fs")>()
|
||||
return {
|
||||
@@ -32,6 +35,17 @@ vi.mock("node:fs", async (importOriginal) => {
|
||||
}
|
||||
return fs.readFileSync(...args)
|
||||
}) 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(() => {
|
||||
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([])
|
||||
|
||||
@@ -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: "<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 () => {
|
||||
const t = await inStep()
|
||||
const poll = t.page.poll()
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
|
||||
|
||||
@@ -50,6 +50,8 @@ describe("the screenshot check and Stop", () => {
|
||||
watchStop: () => () => boolean
|
||||
validateDiagram: () => Promise<any>
|
||||
captureValidationPng?: () => Promise<string>
|
||||
// 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
|
||||
|
||||
@@ -44,6 +44,12 @@ describe("isMxCellXmlComplete", () => {
|
||||
const xml =
|
||||
'<mxCell id="2" value="Hello" style="rounded=1;" vertex="1" parent="1"/>'
|
||||
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", () => {
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user