i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-08 10:47:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: what the batch C review found

Browse Source 
- A redirect followed for a custom base URL (quota on) no longer carries
  the user's key or cookies to another origin, as fetch itself does, and
  a private address may redirect to another private one (already counted).
- The admin Test of an Ollama or Vertex AI entry without a URL goes where
  chat sends that entry's key: the environment's own URL variable, for
  Ollama else the local default. The Test of an Ollama Cloud key without
  a URL went to the cloud while chat went to local Ollama.
- Chat saves: each save notes the chat on screen and the order of the
  reads before reading its data. A save read before switching chats no
  longer writes into the chat switched to, and a copy that waited for its
  thumbnail no longer replaces a newer one.
- "Continue without saving" keeps its button when a later auto-save fails,
  and goes away when a new message is sent.
- A screenshot check that was waiting for its image when the user pressed
  Stop stays skipped after the next message.
- MCP History: draw.io's own copy of a diagram (after get_diagram) no
  longer adds an entry without a picture; a change of background is still
  its own version. The tab ignores an edit's answer that arrives after a
  newer AI write loaded.
- Desktop: a deleted preset is not brought back by a failed switch, and a
  request naming no preset does not stop a rollback. An origin keeping
  an access code counts as having settings.
- .env: a quoted value ending in a backslash ("C:\dir\") is read as dotenv
  reads it.
- The tool card no longer crashes on an id that does not turn into text;
  an older Test's success timer no longer ends a newer Test's spinner.
- Tests that passed without their fix now check it.
This commit is contained in:
dayuan.jiang
2026-10-05 20:09:42 +09:00
parent c0fa997186
commit 69c7896002
23 changed files with 461 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
components/chat-panel.tsx
+45 -10
View File
@@ -277,6 +277,9 @@ export default function ChatPanel({
// Set by Stop until the user sends the next message
const stoppedRef = 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)
// Store original XML for display_diagram and edit_diagram streaming -
// shared between streaming preview and tool handler
@@ -347,7 +350,11 @@ export default function ChatPanel({
validateDiagram: validateWithFallback,
enableVlmValidation: vlmValidationEnabled,
sessionId,
isStopped: () => stoppedRef.current,
watchStop: () => {
const stopsBefore = stopCountRef.current
return () =>
stoppedRef.current || stopCountRef.current !== stopsBefore
},
onValidationStateChange: handleValidationStateChange,
})
@@ -672,6 +679,7 @@ export default function ChatPanel({
currentSessionId,
saveCurrentSession,
getChatGeneration,
getSaveTicket,
} = sessionManager
// Use ref for saveCurrentSession to avoid infinite loop
@@ -709,14 +717,17 @@ export default function ChatPanel({
localStorageDebounceRef.current = setTimeout(async () => {
try {
if (messages.length > 0 || hasDiagramNow) {
// Taken before the data is read, for the chat it was
// scheduled for
const ticket = {
...getSaveTicket(),
generation: scheduledForChat,
}
const sessionData = await buildSessionData({
// Only capture thumbnail if there was a diagram AND this isn't a no-diagram session
withThumbnail: hasDiagramNow && !isNodiagramSession,
})
await saveCurrentSessionRef.current(
sessionData,
scheduledForChat,
)
await saveCurrentSessionRef.current(sessionData, ticket)
}
} catch (error) {
console.error("Failed to save session:", error)
@@ -736,6 +747,7 @@ export default function ChatPanel({
sessionIsAvailable,
currentSessionId,
getChatGeneration,
getSaveTicket,
buildSessionData,
])
@@ -768,10 +780,11 @@ export default function ChatPanel({
try {
// Attempt to save session - browser may not wait for completion
// Skip thumbnail capture as it may not complete in time
const ticket = sessionManager.getSaveTicket()
const sessionData = await buildSessionData({
withThumbnail: false,
})
await sessionManager.saveCurrentSession(sessionData)
await sessionManager.saveCurrentSession(sessionData, ticket)
} catch (error) {
console.error(
"Failed to save session on visibility change:",
@@ -922,11 +935,13 @@ export default function ChatPanel({
// The current chat could not be saved (storage full). The list where
// old chats can be deleted shows only in an empty chat, so let the user
// go on without saving (same toast id: it replaces the plain message)
// go on without saving. It replaces the plain message, and has its own
// id so a later failed auto-save does not take its button away.
const offerToContinueUnsaved = useCallback(
(proceed: () => void) => {
toast.dismiss("session-save-failed")
toast.error(dict.errors.sessionSaveFailedLeave, {
id: "session-save-failed",
id: "session-save-leave",
duration: 15000,
action: {
label: dict.errors.continueWithoutSaving,
@@ -937,6 +952,14 @@ export default function ChatPanel({
[dict],
)
// A new turn makes the offer stale: going on would clear the chat while
// the answer streams in
useEffect(() => {
if (status === "submitted" || status === "streaming") {
toast.dismiss("session-save-leave")
}
}, [status])
// Handle session switching from history dropdown
const handleSelectSession = useCallback(
async (sessionId: string) => {
@@ -969,10 +992,18 @@ export default function ChatPanel({
// without messages); if that failed (storage full), stay on it
// unless the user goes on without saving it
if (messages.length > 0 || isRealDiagram(chartXMLRef.current)) {
// Of the chat on screen now, also if another one comes on
// screen while the thumbnail is taken
const ticket = sessionManager.getSaveTicket()
const sessionData = await buildSessionData({
withThumbnail: true,
})
if (!(await sessionManager.saveCurrentSession(sessionData))) {
if (
!(await sessionManager.saveCurrentSession(
sessionData,
ticket,
))
) {
offerToContinueUnsaved(open)
return
}
@@ -1049,10 +1080,13 @@ export default function ChatPanel({
sessionManager.isAvailable &&
(messages.length > 0 || isRealDiagram(chartXMLRef.current))
) {
const ticket = sessionManager.getSaveTicket()
const sessionData = await buildSessionData({ withThumbnail: true })
// Not saved (storage full): keep the chat on screen, unless the
// user goes on without saving it
if (!(await sessionManager.saveCurrentSession(sessionData))) {
if (
!(await sessionManager.saveCurrentSession(sessionData, ticket))
) {
offerToContinueUnsaved(startNewChat)
return
}
@@ -1128,6 +1162,7 @@ export default function ChatPanel({
// Handle stop button click
const handleStop = useCallback(() => {
stoppedRef.current = true
stopCountRef.current++
// A running screenshot check holds up the chat (the SDK waits for
// the tool handler): end it, so the call gets its result now
cancelValidation()
components/chat/ToolCallCard.tsx
+1 -1
View File
@@ -30,7 +30,7 @@ function OperationsDisplay({ operations }: { operations: DiagramOperation[] }) {
<div className="space-y-3">
{shown.map((op, index) => (
<div
key={`${op.operation}-${op.cell_id}-${index}`}
key={`${op.operation}-${text(op.cell_id)}-${index}`}
className="rounded-lg border border-border/50 overflow-hidden bg-background/50"
>
<div className="px-3 py-1.5 bg-muted/40 border-b border-border/30 flex items-center gap-2">
components/model-config-dialog.tsx
+2 -1
View File
@@ -559,8 +559,9 @@ export function ModelConfigDialog({
clearTimeout(validationResetTimeoutRef.current)
}
validationResetTimeoutRef.current = setTimeout(() => {
setValidationStatus("idle")
validationResetTimeoutRef.current = null
if (run !== validationRunRef.current) return
setValidationStatus("idle")
}, 1500)
} else {
setValidationStatus("error")
electron/main/app-menu.ts
+5 -4
View File
@@ -44,12 +44,12 @@ let lastSwitch = 0
export async function switchPreset(
id: string,
): Promise<Record<string, string>> {
const switchNumber = ++lastSwitch
const previousPresetId = getCurrentPresetId()
const env = applyPresetToEnv(id)
if (!env) {
throw new Error("Preset not found")
}
const switchNumber = ++lastSwitch
rebuildAppMenu()
// In development, scripts/electron-dev.mjs restarts the Next.js dev server
@@ -64,9 +64,10 @@ export async function switchPreset(
console.error("Failed to restart server:", error)
const reason = error instanceof Error ? error.message : String(error)
// A newer switch started meanwhile (also of this same preset): its
// own restart follows, and undoing would lose that choice
if (switchNumber !== lastSwitch) {
// A newer switch started meanwhile (also of this same preset), or
// the preset was deleted: its own restart follows, and undoing
// would lose that choice
if (switchNumber !== lastSwitch || getCurrentPresetId() !== id) {
throw new Error(
`The server could not be restarted.\n\nError: ${reason}`,
)
electron/main/env-loader.ts
+9 -4
View File
@@ -31,15 +31,20 @@ export function loadEnvFile(): void {
/**
* Index of the quote that closes a value starting with a quote, or -1. A
* backslash before the quote character escapes it, as in dotenv; the
* backslash stays in the value.
* backslash stays in the value. As in dotenv, an escaped quote with only a
* comment or nothing after it still closes the value when no other quote
* does ("C:\dir\" keeps its trailing backslash).
*/
function findClosingQuote(value: string): number {
const quote = value[0]
let lastEscaped = -1
for (let i = 1; i < value.length; i++) {
if (value[i] === "\\" && value[i + 1] === quote) i++
else if (value[i] === quote) return i
if (value[i] === "\\" && value[i + 1] === quote) {
i++
if (/^\s*(#.*)?$/.test(value.slice(i + 1))) lastEscaped = i
} else if (value[i] === quote) return i
}
return -1
return lastEscaped
}
/**
hooks/use-diagram-tool-handlers.ts
+8 -5
View File
@@ -64,9 +64,11 @@ interface UseDiagramToolHandlersParams {
validateDiagram?: ValidateDiagramFn
enableVlmValidation?: boolean
sessionId?: string
// The user pressed Stop: a screenshot check that has not started is
// skipped (one already running is cancelled by the caller)
isStopped?: () => boolean
// Called when a screenshot check begins; the function it returns
// tells whether the user pressed Stop in this turn, also after the next
// message was sent. A check that has not started then is skipped (one
// already running is cancelled by the caller).
watchStop?: () => () => boolean
onValidationStateChange?: (
toolCallId: string,
state: ValidationState,
@@ -93,7 +95,7 @@ export function useDiagramToolHandlers({
validateDiagram,
enableVlmValidation = true,
sessionId,
isStopped,
watchStop,
onValidationStateChange,
}: UseDiagramToolHandlersParams) {
// Helper to update validation state
@@ -252,6 +254,7 @@ ${finalXml}
validateDiagram
) {
let capturedPngData: string | null = null
const stopped = watchStop?.()
try {
// Notify UI that we're starting capture
updateValidationState(toolCall.toolCallId, "capturing")
@@ -263,7 +266,7 @@ ${finalXml}
capturedPngData = await captureValidationPng()
// Stopped while the screenshot was taken: no check. The
// chat waits for this handler, so it must end now.
if (isStopped?.()) {
if (stopped?.()) {
updateValidationState(toolCall.toolCallId, "skipped")
} else if (capturedPngData) {
if (DEBUG) {
hooks/use-session-manager.ts
+40 -9
View File
@@ -28,6 +28,13 @@ export interface SessionData {
diagramHistory?: { svg: string; xml: string }[]
}
// Taken right before a save's data is read: the chat on screen then, and
// the order of the reads
export interface SaveTicket {
generation: number
seq: number
}
export interface UseSessionManagerReturn {
// State
sessions: SessionMetadata[]
@@ -39,17 +46,18 @@ export interface UseSessionManagerReturn {
// Actions
switchSession: (id: string) => Promise<SessionData | null>
deleteSession: (id: string) => Promise<{ wasCurrentSession: boolean }>
// chatGeneration: getChatGeneration() when the save was scheduled (by
// default, now); the save is dropped if another chat is on screen when
// its turn comes
// ticket: getSaveTicket() before the data was read (by default, now).
// The save is dropped if another chat is on screen when its turn comes,
// or if a copy of this chat read later was saved already.
// Resolves to false when the save failed (the user was told)
saveCurrentSession: (
data: SessionData,
chatGeneration?: number,
ticket?: SaveTicket,
) => Promise<boolean>
refreshSessions: () => Promise<void>
clearCurrentSession: () => void
getChatGeneration: () => number
getSaveTicket: () => SaveTicket
}
// Reading the session list loads every stored session in full, and window
@@ -91,6 +99,9 @@ export function useSessionManager(
const chatGenerationRef = useRef(0)
// Saves run one at a time, so two saves of a new chat create it once
const saveQueueRef = useRef<Promise<unknown>>(Promise.resolve())
// The last ticket number, and that of the newest data saved
const saveSeqRef = useRef(0)
const savedSeqRef = useRef(0)
const changeChat = useCallback((session: ChatSession | null) => {
chatGenerationRef.current++
@@ -140,14 +151,17 @@ export function useSessionManager(
if (window.electronAPI?.chatsLoaded) {
const count = await readSessionCount()
// The app saves an empty config on its first load; the
// providers are what holds the keys
// providers are what holds the keys, besides an access
// code
let hasSettings = true
try {
const config = JSON.parse(
localStorage.getItem(STORAGE_KEYS.modelConfigs) ??
"{}",
)
hasSettings = (config.providers?.length ?? 0) > 0
hasSettings =
(config.providers?.length ?? 0) > 0 ||
!!localStorage.getItem(STORAGE_KEYS.accessCode)
} catch {
// Unreadable: treat as settings, and stay
}
@@ -277,12 +291,18 @@ export function useSessionManager(
// Save current session data (debounced externally by caller)
const saveCurrentSession = useCallback(
(data: SessionData, chatGeneration?: number): Promise<boolean> => {
// The data is of the chat on screen when the save was asked for
const generation = chatGeneration ?? chatGenerationRef.current
(data: SessionData, ticket?: SaveTicket): Promise<boolean> => {
// The data is of the chat on screen when it was read
const { generation, seq } = ticket ?? {
generation: chatGenerationRef.current,
seq: ++saveSeqRef.current,
}
const run = async (): Promise<boolean> => {
// That chat is no longer on screen (leaving it saved it)
if (generation !== chatGenerationRef.current) return true
// A copy read later was saved already (one that waited for
// its thumbnail must not undo it)
if (seq < savedSeqRef.current) return true
// Nothing can be stored without IndexedDB
if (!isIndexedDBAvailable()) return true
// The user may put another chat on screen while this one is
@@ -308,6 +328,7 @@ export function useSessionManager(
notifySaveFailed(dict.errors.sessionSaveFailed)
return false
}
savedSeqRef.current = seq
await enforceSessionLimit()
if (stillOnScreen()) {
currentSessionRef.current = newSession
@@ -342,6 +363,7 @@ export function useSessionManager(
notifySaveFailed(dict.errors.sessionSaveFailed)
return false
}
savedSeqRef.current = seq
if (stillOnScreen()) {
currentSessionRef.current = updatedSession
setCurrentSession(updatedSession)
@@ -382,6 +404,14 @@ export function useSessionManager(
const getChatGeneration = useCallback(() => chatGenerationRef.current, [])
const getSaveTicket = useCallback(
(): SaveTicket => ({
generation: chatGenerationRef.current,
seq: ++saveSeqRef.current,
}),
[],
)
return {
sessions,
currentSessionId,
@@ -394,5 +424,6 @@ export function useSessionManager(
refreshSessions,
clearCurrentSession,
getChatGeneration,
getSaveTicket,
}
}
lib/ai-providers.ts
+15 -5
View File
@@ -21,6 +21,7 @@ import {
adminProvidersToConfig,
loadAdminProviders,
} from "@/lib/admin/providers"
import { getEnvFallback } from "@/lib/admin/settings"
import { getApiEndpoint } from "@/lib/base-path"
import { redirectGuardedFetch } from "@/lib/ssrf-protection"
import {
@@ -803,8 +804,11 @@ export function getAIModel(clientOverrides?: ClientOverrides): ModelConfig {
// Exception: EdgeOne doesn't require API keys.
// Ollama is exempt only when no server OLLAMA_API_KEY is configured;
// when it IS configured, the outer guard also enforces client apiKey for custom baseUrls.
// A trusted URL is the server's own (the admin Test of an entry without
// one), not a user's
if (
overrides?.baseUrl &&
!overrides?.trustedBaseUrl &&
!overrides?.apiKey &&
!(overrides?.provider === "vertexai" && overrides?.vertexApiKey) &&
overrides?.provider !== "edgeone" &&
@@ -1118,14 +1122,20 @@ export function edgeOneEndpoint(req: Request): string {
/**
* The server's <P>_BASE_URL for a provider, which getAIModel uses for a
* server model without a URL variable of its own (an admin panel entry
* without a URL). None for Bedrock and EdgeOne, and none for Ollama and
* Vertex AI, whose variables the panel writes itself (before a save they
* still hold the entry's previous URL).
* without a URL). None for Bedrock and EdgeOne. Ollama and Vertex AI share
* one variable with the panel, which writes an entry's URL into it: an
* entry without a URL gets the environment's value once saved (before a
* save the variable may still hold the entry's previous URL), and Ollama
* without one goes to the SDK's local default.
*/
export function globalBaseUrl(provider: ProviderName): string | undefined {
if (["bedrock", "edgeone", "ollama", "vertexai"].includes(provider)) {
return undefined
if (provider === "ollama") {
return getEnvFallback("OLLAMA_BASE_URL") || "http://127.0.0.1:11434/api"
}
if (provider === "vertexai") {
return getEnvFallback("GOOGLE_VERTEX_BASE_URL") || undefined
}
if (provider === "bedrock" || provider === "edgeone") return undefined
const name =
provider === "gateway"
? "AI_GATEWAY_BASE_URL"
lib/ssrf-protection.ts
+28 -6
View File
@@ -131,8 +131,8 @@ const MAX_REDIRECTS = 5
* blocked, a public URL could still redirect the request to an internal
* host, so redirects are refused. With private URLs allowed but the quota
* on (DYNAMODB_QUOTA_TABLE), a request to a private address counts as the
* server's: redirects are followed only to public addresses, or a public
* URL could reach the server's own network uncounted. Undefined otherwise.
* server's: a public URL's redirects are followed only to public addresses,
* or it could reach the server's own network uncounted. Undefined otherwise.
*/
export function redirectGuardedFetch(): typeof fetch | undefined {
const blockAll = !allowPrivateUrls()
@@ -140,6 +140,8 @@ export function redirectGuardedFetch(): typeof fetch | undefined {
return async (input, init) => {
let url = input instanceof Request ? input.url : String(input)
let next = init
// A request to a private address already counts as the server's
let startsPrivate: boolean | undefined
for (let hop = 0; hop <= MAX_REDIRECTS; hop++) {
const response = await fetch(url, { ...next, redirect: "manual" })
const location = response.headers.get("location")
@@ -147,20 +149,40 @@ export function redirectGuardedFetch(): typeof fetch | undefined {
return response
}
if (blockAll) throw new RedirectRefusedError()
startsPrivate ??= await isPrivateUrl(url)
const from = new URL(url)
url = new URL(location, url).toString()
if (await isPrivateUrl(url)) {
if (!startsPrivate && (await isPrivateUrl(url))) {
throw new RedirectRefusedError(
"Redirects to private addresses are not allowed",
)
}
// As fetch itself does: 303, and 301 or 302 after a POST, go on
// as a GET without the body
const method = (next?.method ?? "GET").toUpperCase()
// The rest as fetch itself does it. Another origin gets no
// 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")
}
next = { ...next, headers }
// 303, and 301 or 302 after a POST, go on as a GET without the
// body
const method = (next.method ?? "GET").toUpperCase()
if (
response.status === 303 ||
((response.status === 301 || response.status === 302) &&
method === "POST")
) {
for (const name of [
"content-type",
"content-length",
"content-encoding",
"content-language",
"content-location",
]) {
headers.delete(name)
}
next = { ...next, method: "GET", body: undefined }
}
}
packages/mcp-server/src/history.ts
+29 -3
View File
@@ -3,7 +3,9 @@
* 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
@@ -16,6 +18,30 @@ 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) {
@@ -23,10 +49,10 @@ export function addHistory(sessionId: string, xml: string, svg = ""): number {
historyStore.set(sessionId, history)
}
// Dedupe: skip if same as last entry (a change of page settings or
// background only is a new version)
// Dedupe: skip if same as last entry, also when only re-serialized
// (a change of background only is a new version)
const last = history[history.length - 1]
if (last && last.xml === xml) {
if (last && sameDiagram(last.xml, xml)) {
if (svg && !last.svg) last.svg = svg
return history.length - 1
}
packages/mcp-server/src/preview/preview.js
+3 -2
View File
@@ -190,8 +190,9 @@ async function pushState(xml, svg = '', baseVersion = currentVersion, source = '
pushFailing = false;
if (r.ok) {
const d = await r.json();
// An answer about a state this tab has left since
if (sid !== stateId) return;
// An answer about a state this tab has left since, or one that
// comes after a newer version was loaded or saved
if (sid !== stateId || d.version < currentVersion) return;
currentVersion = d.version;
lastXml = xml;
// The canvas changed while this edit was on its way, to
packages/mcp-server/tests/http-server.test.ts
+9
View File
@@ -582,6 +582,15 @@ 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">',
),
)
expect(getHistory(id)).toHaveLength(1)
})
it("keeps a version that changed only the background", () => {
tests/unit/admin-test-model.test.ts
+24 -1
View File
@@ -11,13 +11,19 @@ vi.mock("@/app/api/validate-model/route", () => ({
},
}))
vi.mock("@/lib/admin/auth", () => ({ checkAdminAuth: () => null }))
vi.mock("@/lib/admin/settings", () => ({ loadSettings: () => ({}) }))
// The environment's own values, under the panel's settings
const envFallback = vi.hoisted(() => ({ values: {} as Record<string, string> }))
vi.mock("@/lib/admin/settings", () => ({
loadSettings: () => ({}),
getEnvFallback: (key: string) => envFallback.values[key] ?? null,
}))
import { POST as testModel } from "@/app/api/admin/test-model/route"
const ENV = ["OPENAI_BASE_URL", "SGLANG_BASE_URL", "AI_GATEWAY_BASE_URL"]
const saved: Record<string, string | undefined> = {}
beforeEach(() => {
envFallback.values = {}
for (const k of ENV) {
saved[k] = process.env[k]
delete process.env[k]
@@ -74,8 +80,25 @@ describe("admin Test of an entry without a URL", () => {
try {
await test({ provider: "vertexai", vertexApiKey: "new-key" })
expect(sent.body.baseUrl).toBeUndefined()
// The environment's own URL, which chat uses once it is saved
envFallback.values.GOOGLE_VERTEX_BASE_URL =
"https://vertex-proxy.example.com"
await test({ provider: "vertexai", vertexApiKey: "new-key" })
expect(sent.body.baseUrl).toBe("https://vertex-proxy.example.com")
expect(sent.body.serverBaseUrl).toBe(true)
} finally {
delete process.env.GOOGLE_VERTEX_BASE_URL
}
})
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)
await test({ provider: "ollama", apiKey: "k" })
expect(sent.body.baseUrl).toBe("http://127.0.0.1:11434/api")
expect(sent.body.serverBaseUrl).toBe(true)
envFallback.values.OLLAMA_BASE_URL = "http://gpu:11434/api"
await test({ provider: "ollama", apiKey: "k" })
expect(sent.body.baseUrl).toBe("http://gpu:11434/api")
})
})
tests/unit/ai-providers-credentials.test.ts
+1
View File
@@ -11,6 +11,7 @@ const settings = vi.hoisted(() => ({ values: {} as Record<string, string> }))
vi.mock("@/lib/admin/settings", () => ({
loadSettings: () => settings.values,
getEnvFallback: (key: string) => process.env[key] ?? null,
}))
vi.mock("@ai-sdk/google-vertex", () => {
tests/unit/app-menu.test.ts
+21
View File
@@ -21,6 +21,7 @@ const state = vi.hoisted(() => ({
}))
vi.mock("@/electron/main/config-manager", () => ({
applyPresetToEnv: (id: string) => {
if (id === "missing") return null
state.current = id
return { AI_PROVIDER: id }
},
@@ -80,4 +81,24 @@ describe("switchPreset", () => {
await third
expect(state.current).toBe("B")
})
it("does not bring back the old preset over a deletion", async () => {
const toB = switchPreset("B").catch(() => {})
// B is deleted while its restart is pending
state.current = null
state.restarts[0].reject(new Error("timed out"))
await toB
expect(state.current).toBeNull()
expect(state.restarts).toHaveLength(1)
})
it("still rolls back when a later request named no preset", async () => {
const toB = switchPreset("B").catch(() => {})
await expect(switchPreset("missing")).rejects.toThrow("not found")
state.restarts[0].reject(new Error("timed out"))
await new Promise((r) => setTimeout(r, 0))
state.restarts[1]?.resolve()
await toB
expect(state.current).toBe("A")
})
})
tests/unit/chat-route-quota.test.ts
+3 -1
View File
@@ -163,11 +163,13 @@ describe("chat quota", () => {
it("does not count a provider that never uses the base URL header", async () => {
// Bedrock on the user's own AWS keys goes to AWS, whatever the
// leftover base URL says
// leftover base URL says (with a key header too, so the request gets
// past the custom URL check to the quota decision)
const res = await send({
"x-ai-provider": "bedrock",
"x-ai-model": "amazon.nova-lite-v1:0",
"x-ai-base-url": "http://127.0.0.1:8080",
"x-ai-api-key": "leftover",
"x-aws-access-key-id": "id",
"x-aws-secret-access-key": "secret",
"x-aws-region": "us-east-1",
tests/unit/env-loader.test.ts
+16
View File
@@ -28,6 +28,8 @@ const KEYS = [
"T_ESC_HASH",
"T_ESC_INNER",
"T_ESC_COMMENT",
"T_DIR",
"T_DIR_COMMENT",
]
afterEach(() => {
for (const k of KEYS) delete process.env[k]
@@ -92,4 +94,18 @@ describe("loadEnvFile", () => {
expect(process.env.T_ESC_INNER).toBe('a # \\"b\\"')
expect(process.env.T_ESC_COMMENT).toBe('x\\"y')
})
it("keeps a backslash before the closing quote, like dotenv", () => {
dir.path = mkdtempSync(join(tmpdir(), "env-loader-"))
writeFileSync(
join(dir.path, ".env"),
['T_DIR="C:\\dir\\"', 'T_DIR_COMMENT="C:\\data\\" # dir'].join(
"\n",
),
)
loadEnvFile()
// Windows folders; dotenv 16.6.1 reads them the same
expect(process.env.T_DIR).toBe("C:\\dir\\")
expect(process.env.T_DIR_COMMENT).toBe("C:\\data\\")
})
})
tests/unit/mcp-preview-recovery.test.ts
+44 -3
View File
@@ -238,6 +238,27 @@ describe("MCP preview after the server recreated its session", () => {
expect(t.next("POST").body.xml).toBe("<mxfile>A</mxfile>")
})
it("ignores an edit's answer that comes after a newer AI write loaded", 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")
// The AI wrote X after B; the poll's answer comes first
const poll = t.page.poll()
t.next("GET").answer(state("S1", 4, "<mxfile>X</mxfile>"))
await poll
pushB.answer({ status: 200, body: { success: true, version: 3 } })
await t.settle()
await t.settle()
expect(t.page.read()).toMatchObject({
currentVersion: 4,
lastXml: "<mxfile>X</mxfile>",
})
// No push of the AI's diagram as the user's edit
expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0)
})
it("sends nothing more after a sync reply", async () => {
const t = await inStep()
const poll = t.page.poll()
@@ -298,13 +319,27 @@ describe("MCP preview thumbnails and downloads", () => {
it("drops the reply to an older thumbnail export", async () => {
const { t, n } = await loadedB()
// The next AI write loads before draw.io answered the first export
const poll = t.page.poll()
t.next("GET").answer(state("S1", 4, "<mxfile>C</mxfile>"))
await poll
await new Promise((r) => setTimeout(r, 600))
const newer = t.toDrawio.at(-1).thumbExport
expect(newer).toBeGreaterThan(n)
t.fromDrawio({
event: "export",
data: "<svg/>",
message: { thumbExport: n - 1 },
data: "<svg>B</svg>",
message: { thumbExport: n },
})
await t.settle()
expect(thumbnailPosts(t)).toHaveLength(0)
t.fromDrawio({
event: "export",
data: "<svg>C</svg>",
message: { thumbExport: newer },
})
await t.settle()
expect(thumbnailPosts(t).map((c) => c.body.version)).toEqual([4])
})
it("drops the image when the user changed the canvas since the load", async () => {
@@ -312,11 +347,17 @@ describe("MCP preview thumbnails and downloads", () => {
t.fromDrawio({ event: "autosave", xml: "<mxfile>B edited</mxfile>" })
t.fromDrawio({
event: "export",
data: "<svg/>",
data: "<svg>thumbnail</svg>",
message: { thumbExport: n },
})
await t.settle()
expect(thumbnailPosts(t)).toHaveLength(0)
// The edit is saved with the image of its own export
t.fromDrawio({ event: "export", data: "<svg>edit</svg>" })
await t.settle()
const push = t.next("POST")
expect(push.body.xml).toBe("<mxfile>B edited</mxfile>")
expect(atob(push.body.svg.split(",")[1])).toBe("<svg>edit</svg>")
})
it("downloads the canvas with an edit the server did not get", async () => {
tests/unit/ssrf-protection.test.ts
+52
View File
@@ -143,6 +143,58 @@ describe("redirectGuardedFetch with the quota on", () => {
expect(fetch).toHaveBeenCalledTimes(1)
})
it("follows a private address's redirect to another one", async () => {
// Counted as the server's from the start
vi.stubGlobal(
"fetch",
answers({
"http://10.0.0.5:4000/v1/chat": new Response(null, {
status: 307,
headers: { location: "http://10.0.0.6:4000/v1/chat" },
}),
"http://10.0.0.6:4000/v1/chat": new Response("ok"),
}),
)
const res = await redirectGuardedFetch()?.(
"http://10.0.0.5:4000/v1/chat",
{ method: "POST", body: "{}" },
)
expect(await res?.text()).toBe("ok")
})
it("sends no credentials to another origin", async () => {
const fetchMock = answers({
"https://proxy.example/v1/chat": new Response(null, {
status: 307,
headers: { location: "https://other.example/v1/chat" },
}),
"https://other.example/v1/chat": new Response("ok"),
})
vi.stubGlobal("fetch", fetchMock)
await redirectGuardedFetch()?.("https://proxy.example/v1/chat", {
method: "POST",
body: "{}",
headers: {
Authorization: "Bearer user-key",
Cookie: "eo_token=1",
"Content-Type": "application/json",
},
})
const sent = (call: number) =>
new Headers(
(
fetchMock.mock.calls[call] as unknown as [
string,
RequestInit,
]
)[1].headers,
)
expect(sent(0).get("authorization")).toBe("Bearer user-key")
expect(sent(1).get("authorization")).toBeNull()
expect(sent(1).get("cookie")).toBeNull()
expect(sent(1).get("content-type")).toBe("application/json")
})
it("is not used without the quota", () => {
delete process.env.DYNAMODB_QUOTA_TABLE
expect(redirectGuardedFetch()).toBeUndefined()
tests/unit/tool-call-card.test.tsx
+2
View File
@@ -16,6 +16,8 @@ describe("ToolCallCard", () => {
null,
{ operation: {} },
{ operation: "add", cell_id: {} },
// JSON can hold an object that does not turn into text
JSON.parse('{"operation":"add","cell_id":{"toString":null}}'),
{ operation: "add", cell_id: "2", new_xml: {} },
{ operation: "update", cell_id: "3", new_xml: '<mxCell id="3"/>' },
]
tests/unit/use-diagram-tool-handlers.test.tsx
+36 -5
View File
@@ -47,8 +47,9 @@ function setup(partialXml: string) {
describe("the screenshot check and Stop", () => {
const draw = async (opts: {
isStopped: () => boolean
watchStop: () => () => boolean
validateDiagram: () => Promise<any>
captureValidationPng?: () => Promise<string>
}) => {
const onValidationStateChange = vi.fn()
const { result } = renderHook(() =>
@@ -62,9 +63,11 @@ describe("the screenshot check and Stop", () => {
onFetchChart: async () => "",
onExport: () => {},
enableVlmValidation: true,
captureValidationPng: async () => "data:image/png;base64,AA",
captureValidationPng:
opts.captureValidationPng ??
(async () => "data:image/png;base64,AA"),
validateDiagram: opts.validateDiagram,
isStopped: opts.isStopped,
watchStop: opts.watchStop,
onValidationStateChange,
}),
)
@@ -89,7 +92,7 @@ describe("the screenshot check and Stop", () => {
suggestions: [],
}))
const { addToolOutput, onValidationStateChange } = await draw({
isStopped: () => true,
watchStop: () => () => true,
validateDiagram,
})
expect(validateDiagram).not.toHaveBeenCalled()
@@ -103,7 +106,7 @@ describe("the screenshot check and Stop", () => {
it("ends with the diagram's result when Stop cancels a running check", async () => {
const { addToolOutput, onValidationStateChange } = await draw({
isStopped: () => false,
watchStop: () => () => false,
validateDiagram: async () => {
throw new DOMException("Validation cancelled", "AbortError")
},
@@ -114,6 +117,34 @@ describe("the screenshot check and Stop", () => {
expect(addToolOutput).toHaveBeenCalledTimes(1)
expect(addToolOutput.mock.lastCall?.[0].state).toBeUndefined()
})
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
let stops = 0
let stoppedNow = false
const validateDiagram = vi.fn(async () => ({
valid: true,
issues: [],
suggestions: [],
}))
const { onValidationStateChange } = await draw({
watchStop: () => {
const before = stops
return () => stoppedNow || stops !== before
},
captureValidationPng: async () => {
stops++ // Stop
stoppedNow = false // the next message
return "data:image/png;base64,AA"
},
validateDiagram,
})
expect(validateDiagram).not.toHaveBeenCalled()
expect(onValidationStateChange.mock.lastCall?.[1].status).toBe(
"skipped",
)
})
})
describe("append_diagram and the stored previews", () => {
tests/unit/use-session-manager.test.tsx
+64 -2
View File
@@ -96,7 +96,7 @@ describe("saving the chat on screen", () => {
it("drops a save scheduled before New Chat", async () => {
const { result } = await setup()
const scheduled = result.current.getChatGeneration()
const scheduled = result.current.getSaveTicket()
act(() => result.current.clearCurrentSession())
let save!: Promise<boolean>
act(() => {
@@ -110,7 +110,7 @@ describe("saving the chat on screen", () => {
it("drops a save of the old chat waiting behind New Chat's save", async () => {
const { result } = await setup()
// The auto-save is scheduled, then New Chat saves and clears
const scheduled = result.current.getChatGeneration()
const scheduled = result.current.getSaveTicket()
let newChatSave!: Promise<boolean>
let autoSave!: Promise<boolean>
act(() => {
@@ -168,3 +168,65 @@ describe("saving the chat on screen", () => {
expect(hook.result.current.currentSessionId).toBeNull()
})
})
describe("save tickets", () => {
beforeEach(() => {
stored.clear()
pendingWrites = []
})
const textOf = (session: any) => session?.messages[0].parts[0].text
const said = (text: string) => ({
...data,
messages: [{ ...data.messages[0], parts: [{ type: "text", text }] }],
})
it("never put an older copy of a chat over a newer one", async () => {
const { result } = await setup()
let first!: Promise<boolean>
act(() => {
first = result.current.saveCurrentSession(said("first"))
})
await finishWrites()
await first
// An auto-save read its data, then waits for its thumbnail; a save
// without a thumbnail reads newer data and is done first
const older = result.current.getSaveTicket()
const newer = result.current.getSaveTicket()
let saves!: Promise<boolean[]>
act(() => {
saves = Promise.all([
result.current.saveCurrentSession(said("newer"), newer),
result.current.saveCurrentSession(said("older"), older),
])
})
await finishWrites()
await saves
expect(textOf([...stored.values()][0])).toBe("newer")
})
it("keep a chat read before a switch out of the chat switched to", async () => {
stored.set("other", {
...said("other chat"),
id: "other",
title: "Other",
})
const { result } = await setup()
// New Chat reads this chat, then waits for its thumbnail
const ticket = result.current.getSaveTicket()
// Meanwhile the user opens the other chat
let open!: Promise<unknown>
act(() => {
open = result.current.switchSession("other")
})
await finishWrites()
await open
let late!: Promise<boolean>
act(() => {
late = result.current.saveCurrentSession(said("this chat"), ticket)
})
await finishWrites()
await late
expect(textOf(stored.get("other"))).toBe("other chat")
expect(stored.size).toBe(1)
})
})
tests/unit/validate-model-route.test.ts
+4 -1
View File
@@ -6,7 +6,10 @@ import { POST as validateModel } from "@/app/api/validate-model/route"
import { getAIModel } from "@/lib/ai-providers"
// No saved admin providers
vi.mock("@/lib/admin/settings", () => ({ loadSettings: () => ({}) }))
vi.mock("@/lib/admin/settings", () => ({
loadSettings: () => ({}),
getEnvFallback: (key: string) => process.env[key] ?? null,
}))
// Every URL is public (no DNS in tests), unless a test says otherwise
const privateUrls = vi.hoisted(() => ({ all: false }))