diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index b8cb36de..d8d329dc 100644 --- a/components/chat-panel.tsx +++ b/components/chat-panel.tsx @@ -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() diff --git a/components/chat/ToolCallCard.tsx b/components/chat/ToolCallCard.tsx index f4c7c258..8ada80e2 100644 --- a/components/chat/ToolCallCard.tsx +++ b/components/chat/ToolCallCard.tsx @@ -30,7 +30,7 @@ function OperationsDisplay({ operations }: { operations: DiagramOperation[] }) {
{shown.map((op, index) => (
diff --git a/components/model-config-dialog.tsx b/components/model-config-dialog.tsx index 10b522ec..52dbea78 100644 --- a/components/model-config-dialog.tsx +++ b/components/model-config-dialog.tsx @@ -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") diff --git a/electron/main/app-menu.ts b/electron/main/app-menu.ts index ae129559..bb38ab9d 100644 --- a/electron/main/app-menu.ts +++ b/electron/main/app-menu.ts @@ -44,12 +44,12 @@ let lastSwitch = 0 export async function switchPreset( id: string, ): Promise> { - 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}`, ) diff --git a/electron/main/env-loader.ts b/electron/main/env-loader.ts index bb3ff9da..563b3729 100644 --- a/electron/main/env-loader.ts +++ b/electron/main/env-loader.ts @@ -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 } /** diff --git a/hooks/use-diagram-tool-handlers.ts b/hooks/use-diagram-tool-handlers.ts index 55313b85..c5cbdbe7 100644 --- a/hooks/use-diagram-tool-handlers.ts +++ b/hooks/use-diagram-tool-handlers.ts @@ -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) { diff --git a/hooks/use-session-manager.ts b/hooks/use-session-manager.ts index b0cfe308..e950644a 100644 --- a/hooks/use-session-manager.ts +++ b/hooks/use-session-manager.ts @@ -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 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 refreshSessions: () => Promise 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.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 => { - // The data is of the chat on screen when the save was asked for - const generation = chatGeneration ?? chatGenerationRef.current + (data: SessionData, ticket?: SaveTicket): Promise => { + // 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 => { // 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, } } diff --git a/lib/ai-providers.ts b/lib/ai-providers.ts index 25529f51..b2ef3390 100644 --- a/lib/ai-providers.ts +++ b/lib/ai-providers.ts @@ -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

_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" diff --git a/lib/ssrf-protection.ts b/lib/ssrf-protection.ts index 80357cf3..5e682893 100644 --- a/lib/ssrf-protection.ts +++ b/lib/ssrf-protection.ts @@ -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 } } } diff --git a/packages/mcp-server/src/history.ts b/packages/mcp-server/src/history.ts index 253a223e..896ff5f4 100644 --- a/packages/mcp-server/src/history.ts +++ b/packages/mcp-server/src/history.ts @@ -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() +/** 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 } diff --git a/packages/mcp-server/src/preview/preview.js b/packages/mcp-server/src/preview/preview.js index 10b2ec67..2ba68ec2 100644 --- a/packages/mcp-server/src/preview/preview.js +++ b/packages/mcp-server/src/preview/preview.js @@ -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 diff --git a/packages/mcp-server/tests/http-server.test.ts b/packages/mcp-server/tests/http-server.test.ts index 60c09011..7168911a 100644 --- a/packages/mcp-server/tests/http-server.test.ts +++ b/packages/mcp-server/tests/http-server.test.ts @@ -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( + "", + '', + ), + ) + expect(getHistory(id)).toHaveLength(1) }) it("keeps a version that changed only the background", () => { diff --git a/tests/unit/admin-test-model.test.ts b/tests/unit/admin-test-model.test.ts index a9254b0f..11801fae 100644 --- a/tests/unit/admin-test-model.test.ts +++ b/tests/unit/admin-test-model.test.ts @@ -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 })) +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 = {} 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") + }) }) diff --git a/tests/unit/ai-providers-credentials.test.ts b/tests/unit/ai-providers-credentials.test.ts index d74ed915..4f2738ca 100644 --- a/tests/unit/ai-providers-credentials.test.ts +++ b/tests/unit/ai-providers-credentials.test.ts @@ -11,6 +11,7 @@ const settings = vi.hoisted(() => ({ values: {} as Record })) vi.mock("@/lib/admin/settings", () => ({ loadSettings: () => settings.values, + getEnvFallback: (key: string) => process.env[key] ?? null, })) vi.mock("@ai-sdk/google-vertex", () => { diff --git a/tests/unit/app-menu.test.ts b/tests/unit/app-menu.test.ts index 8ab3caac..a3493318 100644 --- a/tests/unit/app-menu.test.ts +++ b/tests/unit/app-menu.test.ts @@ -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") + }) }) diff --git a/tests/unit/chat-route-quota.test.ts b/tests/unit/chat-route-quota.test.ts index 74933b8f..28a3fdac 100644 --- a/tests/unit/chat-route-quota.test.ts +++ b/tests/unit/chat-route-quota.test.ts @@ -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", diff --git a/tests/unit/env-loader.test.ts b/tests/unit/env-loader.test.ts index 2cffdeed..3bdf06e7 100644 --- a/tests/unit/env-loader.test.ts +++ b/tests/unit/env-loader.test.ts @@ -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\\") + }) }) diff --git a/tests/unit/mcp-preview-recovery.test.ts b/tests/unit/mcp-preview-recovery.test.ts index 379c0090..9dd504dd 100644 --- a/tests/unit/mcp-preview-recovery.test.ts +++ b/tests/unit/mcp-preview-recovery.test.ts @@ -238,6 +238,27 @@ describe("MCP preview after the server recreated its session", () => { expect(t.next("POST").body.xml).toBe("A") }) + it("ignores an edit's answer that comes after a newer AI write loaded", async () => { + const t = await inStep() + t.fromDrawio({ event: "autosave", xml: "B" }) + t.fromDrawio({ event: "export", data: "" }) + 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, "X")) + 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: "X", + }) + // 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, "C")) + 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: "", - message: { thumbExport: n - 1 }, + data: "B", + message: { thumbExport: n }, }) await t.settle() expect(thumbnailPosts(t)).toHaveLength(0) + t.fromDrawio({ + event: "export", + data: "C", + 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: "B edited" }) t.fromDrawio({ event: "export", - data: "", + data: "thumbnail", 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: "edit" }) + await t.settle() + const push = t.next("POST") + expect(push.body.xml).toBe("B edited") + expect(atob(push.body.svg.split(",")[1])).toBe("edit") }) it("downloads the canvas with an edit the server did not get", async () => { diff --git a/tests/unit/ssrf-protection.test.ts b/tests/unit/ssrf-protection.test.ts index c2e54cca..0ff98d11 100644 --- a/tests/unit/ssrf-protection.test.ts +++ b/tests/unit/ssrf-protection.test.ts @@ -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() diff --git a/tests/unit/tool-call-card.test.tsx b/tests/unit/tool-call-card.test.tsx index 1979fc6c..4225f533 100644 --- a/tests/unit/tool-call-card.test.tsx +++ b/tests/unit/tool-call-card.test.tsx @@ -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: '' }, ] diff --git a/tests/unit/use-diagram-tool-handlers.test.tsx b/tests/unit/use-diagram-tool-handlers.test.tsx index 7d3b6be8..e247c8ff 100644 --- a/tests/unit/use-diagram-tool-handlers.test.tsx +++ b/tests/unit/use-diagram-tool-handlers.test.tsx @@ -47,8 +47,9 @@ function setup(partialXml: string) { describe("the screenshot check and Stop", () => { const draw = async (opts: { - isStopped: () => boolean + watchStop: () => () => boolean validateDiagram: () => Promise + captureValidationPng?: () => Promise }) => { 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", () => { diff --git a/tests/unit/use-session-manager.test.tsx b/tests/unit/use-session-manager.test.tsx index fe753a30..c697db6f 100644 --- a/tests/unit/use-session-manager.test.tsx +++ b/tests/unit/use-session-manager.test.tsx @@ -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 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 let autoSave!: Promise 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 + 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 + 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 + act(() => { + open = result.current.switchSession("other") + }) + await finishWrites() + await open + let late!: Promise + 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) + }) +}) diff --git a/tests/unit/validate-model-route.test.ts b/tests/unit/validate-model-route.test.ts index db623b9a..d11cdeca 100644 --- a/tests/unit/validate-model-route.test.ts +++ b/tests/unit/validate-model-route.test.ts @@ -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 }))