feat: MCP feature parity, the web app on the MCP core, and provider fixes (#951)

* fix(chat): close credential leaks and harden the chat route

- Vertex: a client-supplied base URL only works with the client's own Vertex key
- Accept only data: URLs for file parts in every message, so the server never downloads them
- Output budget retry accounts for the thinking budget Bedrock/Anthropic add, and reads
  Volcengine, DashScope, SGLang and vLLM rejections; falls back to 16000 once
- x-max-output-tokens can only lower the budget on server credentials
- On server credentials only server models or AI_MODEL entries can be used
- Drop tool results together with the invalid tool calls they belong to
- Count quota tokens as input + output (cached tokens were counted twice)
- Private-URL check for custom base URLs, end Langfuse traces on error/abort/early return
- Fix repairToolCall ordering and placeholder, align edit_diagram prompt with operations
- Panel Bedrock keys are read from ADMIN_AWS_*; forward the access code to EdgeOne
- isMinimalDiagram only treats root cells as an empty canvas

* fix(api): require access codes and limit sizes on helper routes

- Shared checkAccessCode for validate-diagram, validate-model, parse-url, verify-access-code
- parse-url: 5 MB streamed body limit; validate-diagram: 5 MB image limit
- validate-model refuses redirects when private URLs are blocked
- Admin settings state shared across module instances via globalThis
- Server model ids: unique slugs (non-ASCII names encoded), duplicates rejected
- Panel Bedrock credentials stored as ADMIN_AWS_* so the DynamoDB client keeps its own
- Locale redirect keeps basePath and query; EdgeOne function drops open CORS and checks the access code
- Providers payload reports whether .env sets a default model

* fix(chat): keep saved diagrams and pages when restoring, editing and retrying

- Restored sessions no longer replay the last display_diagram over the saved diagram
- Failed or stopped edit_diagram restores the canvas
- Message snapshots keep the full multi-page document
- "Improve with suggestions" uses the normal send path (headers, xml, retry counters)
- Editing a message keeps its file/URL sections; cached example edits work
- New chat's first autosave no longer resets the UI
- Validation retries counted per user turn; validate-diagram sends the access code
- Cached examples only match the example files on an empty canvas
- Template sends keep attachments and wait for extraction

* fix(diagram): fix autosave staleness and XML repair corrupting valid diagrams

- Autosave guard reads refs, so edits after a theme or dark mode switch are kept
- Duplicate-id check and rename run per page; repair loop no longer quadratic
- autoFixXml no longer breaks style values, rich text " or single-line cells
- extractCompleteMxCells keeps the cell after a self-closing cell
- Better truncation detection; object/UserObject wrapped cells are editable
- Exports for thumbnail, PNG and save are routed by tag instead of a shared resolver
- History stores the full document; storage errors are reported, no auto-deletion of chats
- IndexedDB connection reopens after errors; focus refresh throttled
- Keep ?session= on locale redirect, map zh-Hant to zh-tw for draw.io

* fix(chat-input): stop template dialogs from sending and fix attachment races

- Template dialogs no longer submit the outer chat form
- Sending is blocked while files or URLs are still extracting
- File and URL extraction no longer drop or resurrect entries
- IME composition Enter no longer sends
- Tool call cards show the error text; keyboard handling on cards fixed
- Template import available when empty, edit dialog resets, saved templates refresh
- Only png/jpeg/gif/webp images accepted, SVG sent as text; PDF objects released
- parse-url request sends the access code

* fix(model-config): keep model selection valid and fix admin panel edge cases

- Fall back to the default server model when a saved one disappears
- Sync model config across tabs
- Validation uses the base path and sends the access code
- Model ids edited as drafts (no empty, duplicate or padded ids)
- Credential changes reset validation; stale validation results are dropped
- Admin: generateId over HTTP, env-locked group switches, discard and toggle fixes,
  clearing a secret field keeps the saved key, first provider not auto-default when .env sets AI_MODEL
- Model selector items use unique values

* fix(electron): decrypt keys after ready and harden navigation and IPC

- Apply preset env after app ready, so Windows/Linux get decrypted keys
- Never re-encrypt ciphertext; restore env when switching or removing presets
- Block navigation away from the app, open external links in the browser, check IPC senders
- Keep inherited proxy settings, default NO_PROXY for localhost
- Serialize server start/restart, kill stuck processes, follow port changes
- Atomic config writes, keep corrupt files as backups, remember the server port
- Menu and settings window stay in sync; dev script gets the decrypted preset env
- Use app.isPackaged, parse inline .env comments, drop .env files from the bundle

* fix(mcp-server): fix XSS and crashes, make XML validation strict

- Validate and escape the mcp session id; only serve localhost Host/Origin
- Malformed URLs and session ids return errors instead of crashing the process
- Strict XML syntax check with saxes (linkedom never reports parse errors)
- autoFixXml no longer corrupts valid XML; attribute newlines serialized as entities
- Sessions stay alive while polled; browser pushes carry a base version (409 on conflict)
- Page tools respect the edit gate; UTF-8 bodies decoded correctly
- Export replies matched to requests and serialized; xml sync export handled
- UserObject/object cells addressable by id; history restored by stable id; logs off stdout

* fix(mcp-server): make edit_diagram all-or-nothing and fix preview sync races

- edit_diagram applies nothing when any operation fails, rejects invalid or
  multi-cell new_xml, validates only the target page, and returns the
  current page XML on every rejection (including stale edits)
- Fix get_diagram reading the old diagram right after an AI write: the
  preview pushed its sync reply with a newer version than it was taken at
- Keep a user edit that loses the race with an AI write in history and
  tell the user in the preview
- Autofix removes only exact foreign tags (a stray <mxGraph/> deleted
  <mxGraphModel>), fixes tag case, drops orphan <mxPoint>s, and rejects
  unknown element names in model XML
- Edit empty and compressed pages; PNG exports use the page on screen;
  tag download exports; reload from the server after a page export
- Expand ~ in paths, tell the model when the browser sync timed out,
  use registerPrompt, require SDK ^1.31.0

* feat(mcp-server): bring the web app's drawing knowledge to MCP

- Add a drawing guide adapted from the web system prompt (layout, edge
  routing, styles, minimal style, editing rules), returned by
  start_session, a new get_drawing_guide tool and the diagram-workflow prompt
- Add get_shape_library with the 30 icon libraries; the build copies
  docs/shape-libraries into dist and CI checks the packed files
- Accept bare mxCell lists in create_new_diagram and add_page; the server
  adds the wrapper and root cells
- Send server instructions, shorten create_new_diagram's description to
  fit Claude Code's 2,048 character limit, and annotate every tool
- Fix dead links and the totals in docs/shape-libraries/README.md

* feat(mcp-server): add screenshot_diagram so the model can check its render

- New read-only screenshot_diagram tool returns the rendered page as a PNG
  plus the web app's visual checklist (overlaps, edges crossing shapes,
  readability, layout, rendering errors), replacing the web app's
  separate vision model with the host model's own vision
- PNG exports use draw.io's width and pageId options: screenshots stay
  under ~140,000 base64 characters and page exports no longer swap the
  page on screen
- Fail fast with a clear message when the preview tab stopped polling
  (browsers throttle background tabs)
- Mention the screenshot step in the drawing guide and instructions

* feat(mcp-server): auto-save each session's diagram to a .drawio file

- Save the latest diagram of every session 1 second after each change
  (AI write, browser edit, history restore) to ~/.next-ai-drawio/<id>.drawio,
  keep the newest 50, flush on shutdown; DRAWIO_DATA_DIR changes the folder
  and "off" disables it, like the web app's IndexedDB sessions
- start_session names the file, so a resumed conversation can reopen the
  diagram with load_diagram after the MCP process restarted
- Fix PNG/SVG exports randomly timing out: a previous export's 10 second
  timer cleared the export in progress, and a late reply could be taken
  for the current one; exports are now numbered

* refactor(mcp-server): move the preview page into src/preview

The 580-line page template in http-server.ts becomes index.html,
preview.css and preview.js, copied to dist/preview by the build and
filled at request time. The rendered page is unchanged apart from the
session id and draw.io origin now coming from a small config script.
Biome skips the folder because of the {{placeholders}}, as it never
linted the old template string either.

* feat(mcp-server): add theme menu, dark mode and editable SVG to the preview

- Pass themes=1 and dark=auto to the draw.io iframe, so the Extras menu
  offers the Theme submenu and draw.io follows the system dark mode; the
  header and dialogs follow it too through prefers-color-scheme
- Download dialog: new Editable SVG (.drawio.svg) format through draw.io's
  xmlsvg export, default name diagram-YYYY-MM-DD, Enter saves and Escape
  closes it and the history dialog
- export_diagram accepts format "drawio.svg" and detects the .drawio.svg
  extension
- Use system fonts instead of Google Fonts, so the page works offline with
  a self-hosted draw.io (DRAWIO_BASE_URL)

* docs(mcp-server): describe the new MCP features in all READMEs

- MCP server README: 13 tools, drawing rules and shape libraries,
  screenshots, all-or-nothing edits, auto-save and how to continue a
  diagram later, DRAWIO_DATA_DIR and DEBUG, offline use with a local
  draw.io, and what to do when a background tab makes exports time out
- Fix two errors: History is a button at the top right of the preview
  page, and exports are not limited to .drawio
- Claude Code plugin README: same tool list, formats and settings
- Root READMEs (English, Chinese, Japanese): short list of what the MCP
  server can do

* fix(mcp-server): fix duplicate page exports and auto-save deleting user files

- Preview page: keep an MCP export open until the server has its result.
  A poll answered before that still saw the request and started the same
  export again, so a parallel page export could write the previous
  page's image into its file
- Auto-save only removes its own mcp-*.drawio files, so a DRAWIO_DATA_DIR
  that also holds the user's diagrams keeps them
- screenshot_diagram captures a page that has no id attribute by loading
  just that page, like export_diagram
- An empty <Array as="points"/> no longer hides orphan mxPoints that
  come after it
- POST /api/state refuses a push without xml, which used to wipe the
  stored diagram
- Clear exportOptions when an export ends, reuse hasCells for the empty
  diagram check, and reword two log lines

* fix: log Editable SVG saves, use local dates in file names, drop unused Electron dialogs

- /api/log-save accepted only drawio, png and svg, so saving as Editable
  SVG (xmlsvg) got a 400 and was never recorded in Langfuse
- The default download name and the template export name used the UTC
  date, which is the previous day on mornings in East Asia
- Remove the Electron openFile/saveFile IPC handlers, their preload
  bindings and types; nothing in the app calls them

* refactor(mcp-server): make the diagram modules usable from the web app

The web app will reuse the MCP server's XML engine instead of its own
copy in lib/utils.ts, so these modules now run in the browser too.

- Relative imports end in .ts, rewritten to .js by tsc
  (rewriteRelativeImportExtensions); Next.js resolves them directly
- Every module uses the global DOMParser/XMLSerializer: native in the
  browser, linkedom in Node via installDomPolyfill. pages.ts parsed with
  linkedom but serialized with the global serializer, which throws in
  the browser
- The saxes syntax check moves to xml-syntax.ts, so the browser does not
  pull in linkedom; it now also rejects undeclared prefixes such as
  xlink:, as the browser does
- Page decompression uses pako and atob instead of node:zlib and Buffer
- hasCells moves to pages.ts, away from the file system code
- The duplicate cell id check counts UserObject/object ids
- wrapCellsInModel drops comments and text before the first cell, which
  the web app accepts today
- validateAndFixXml takes { strict: false } for diagrams with user content
- Web tests run these modules with a browser DOM (jsdom)
- saxes becomes a direct dependency of the web app

* refactor(web): validate and repair diagram XML with the MCP server's engine

- Delete the web app's own copy of the XML checks and repairs from
  lib/utils.ts (1,074 lines). loadDiagram now uses the MCP server's
  validateAndFixXml without the strict checks, because the XML may hold
  the user's own diagram
- display_diagram and append_diagram prepare the model's XML with the new
  shared prepareNewDiagram, also used by the MCP create_new_diagram: wrap,
  validate strictly and auto-fix while it is still a bare model (where
  duplicate ids are renamed), then turn it into an mxfile
- The streaming preview of display_diagram no longer redraws the model's
  raw cells after the tool handler loaded the checked diagram, and drops
  a queued preview once the input is complete. That redraw lost
  auto-fixes and UserObject/object wrappers, so a linked cell lost its
  label; it also showed a second error toast
- The web repair regression tests now run against the MCP functions
- New e2e test checks the canvas content after display_diagram
- Fix the e2e upload tests, whose file input locator also matched the
  template import input

* refactor(web): edit and wrap diagrams with the MCP server's code

- edit_diagram runs the MCP server's editDiagram: every new_xml is checked
  first, one cell per operation, and after the edit only the target page
  is checked, rejecting only errors this edit introduced. An unrelated
  problem elsewhere in the document no longer blocks every edit. The
  error lists each failed operation
- The streaming edit preview uses the MCP applyDiagramOperations
- Delete applyDiagramOperations (292 lines) and wrapWithMxFile from
  lib/utils.ts, and the unused hand-copied scripts/test-diagram-operations.mjs
- One blank document (BLANK_MXFILE) for the web app and the MCP preview,
  replacing four copies
- Saving a .drawio wraps a bare model with normalizeToMxfile
- The empty-diagram check uses hasCells, which also counts cells wrapped
  in a UserObject/object
- DiagramOperation is the MCP type
- The wrapped-cell and empty-diagram tests now run against the MCP code
- New e2e test: edit_diagram changes the canvas, and a failing edit
  leaves it as it was

* refactor(web): share prompt examples and the shape library with the MCP server

- The three XML examples (swimlanes, two edges, waypoints) that the web
  system prompt, the display_diagram description and the MCP drawing
  guide each had a copy of now live in packages/mcp-server/src/xml-examples.ts
- The shape library group list and reader come from the MCP
  shape-library.ts; getShapeLibrary takes the folder, and the web route
  passes docs/shape-libraries under the working directory as before.
  Only the 30 known library names are read. Error texts differ slightly
- next.config.ts traces docs/shape-libraries/*.md for /api/chat, since
  the read now happens in another module
- Every prompt is byte-for-byte unchanged: the web system prompt for five
  model ids with and without minimal style, and the MCP drawing guide,
  compared before and after; the library list and swimlane example match
  the old tool description text exactly

* fix(providers): update the v6 SDK packages and fix Claude and Gemini settings

- Update ai to 6.0.300 and the @ai-sdk providers to their latest v6-line
  versions. @ai-sdk/anthropic 3.0.47 did not know claude-opus-4-7/4-8
  and capped their output at 32000 tokens; 3.0.127 allows 128000
- Drop the fine-grained-tool-streaming beta header for the Anthropic API:
  the provider now streams tool input per tool (eager_input_streaming)
- Claude 4.7 and later reject a non-default temperature/top_p/top_k and
  the extended thinking budget with a 400. A middleware retries once
  without them, so TEMPERATURE and *_THINKING_BUDGET_TOKENS no longer
  break those models
- Prompt caching also reaches Claude on the Anthropic API and OpenRouter;
  before, only Bedrock got a cache marker
- GOOGLE_TOP_K and GOOGLE_TOP_P never reached Gemini: they were sent as
  Google provider options, which drops them. They are call settings now.
  GOOGLE_CANDIDATE_COUNT and GOOGLE_REASONING_EFFORT, which the provider
  does not support, are removed
- Add @ai-sdk/openai-compatible as a direct dependency

* refactor: simpler streaming preview and small AI SDK cleanups

- useChat throttles streamed message updates (experimental_throttle,
  150 ms), replacing the two hand-written 150 ms timers of the
  display_diagram and edit_diagram previews (94 lines less). The preview
  now only runs while the input streams; once it is complete the tool
  handler takes over, so a queued preview can no longer redraw an edit
  the handler rejected and rolled back. Measured on a streamed 60-cell
  diagram: 41 redraws at least 97 ms apart, before 37 with gaps down to
  48 ms
- The diagram check endpoint uses streamText with Output.object instead
  of the deprecated streamObject, and returns its fixed result as a plain
  text response; new route test
- Import createGateway/gateway from ai and drop the direct
  @ai-sdk/gateway dependency
- The per-request message structure logs only print with
  DEBUG_LLM_PAYLOAD=true
- Remove an empty onFinish callback

* refactor(chat): check the last tool part with the SDK's isToolUIPart

Drop the hand-written MessagePart and ChatMessage types that only served
this check.

* refactor(providers): one model factory for chat and the settings Test button

- getAIModel resolves credentials (client key, server env vars, the
  existing SSRF rules) and createModel builds the model by SDK. The
  provider-by-provider switch shrinks from 24 cases to the few that
  differ (lib/ai-providers.ts 1531 -> 1106 lines)
- /api/validate-model calls getAIModel instead of its own 24-case switch
  (503 -> 175 lines), which had drifted from the chat: it built Azure
  with createOpenAI, Kimi and MiMo with createOpenAI instead of
  createDeepSeek, and the official OpenAI endpoint with Chat Completions.
  A passing test now means the chat works
- Plain OpenAI-compatible providers (SiliconFlow, SGLang, ModelScope,
  GLM, Qwen, Qiniu, Novita, Atlas Cloud, EdgeOne, Doubao, MiniMax in
  OpenAI mode, AIHubMix on a custom URL) use @ai-sdk/openai-compatible,
  which reads reasoning_content, so their reasoning shows, and accepts
  SGLang's stream as is (its 95-line stream rewrite is gone).
  includeUsage keeps token usage for quotas. <think> tags in their text
  become reasoning (extractReasoningMiddleware)
- SGLang without a base URL used OpenAI's endpoint; it now defaults to
  http://127.0.0.1:8000/v1 like the Test button did
- Chat requests to a client base URL refuse redirects, as the Test
  button already did (redirectGuardedFetch moves to lib/ssrf-protection)
- The Test button streams like the chat (the ModelScope special case is
  gone), times out after 15 s, does not retry, asks the model to call a
  ping tool and warns when it answers without one, and tests all models
  at once. The time each test took shows on its check mark
- Unknown provider names are rejected with Object.hasOwn, and the error
  texts list providers from PROVIDER_INFO instead of hand-kept lists

* feat(settings): link to each provider's key page and clean up base URLs

- A "Get API key" link next to the API Key field for the 19 providers
  that have a key page (from env.example and the providers' docs). 17
  answered 200 to curl; OpenAI's is behind a Cloudflare challenge and
  DeepSeek's behind a regional block, both checked in Chrome
- Base URLs drop spaces, trailing slashes and a pasted endpoint path
  (/chat/completions, /completions, /messages, /responses), which the
  SDK would otherwise append a second time and get a 404. getAIModel does
  this for the chat and the Test button; the field does it on blur and
  shows the URL requests go to

* feat(errors): classify provider errors and show a hint the user can act on

- lib/llm-errors.ts sorts an error into about a dozen kinds (key
  rejected, no access, unknown model, no credit, rate limited, context
  too long, no image input, no tool calls, output cut off, provider down,
  cannot connect, timeout): first texts that name the cause precisely,
  then the HTTP status code, then general texts. It unwraps RetryError and
  hides keys and Bearer tokens in the provider's message
- The chat route uses it for errors before the stream and, through
  toUIMessageStreamResponse's onError, for errors in the stream. Errors
  of the model's own tool call stay as they are: the same text goes back
  to the model so it can fix the call
- The chat shows the hint in the user's language, then the provider's
  message; a rejected key, missing access or unknown model adds an "Open
  model settings" button. The Test button shows the same hints
- Fixes: our message "API key is required when using a custom base URL"
  was replaced by "Authentication failed" because it contains "key"; a
  provider's "Rate limit exceeded" opened this site's quota toast; an
  error body like {"error": ...} was shown as raw JSON; the Test button
  matched "401" in the message, where providers rarely put it
- Remove the string matching fallbacks in the chat panel

* refactor(settings): drop an unused index parameter

* feat(settings): fetch the model list from any provider and flag models that cannot draw

The "Fetch models" button asks the provider for its models (OpenAI-style
/models, Anthropic, Google, Ollama, OpenRouter, Vercel Gateway, AIHubMix)
and shows them in a searchable picker. This replaces the route that only
worked for AIHubMix.

A snapshot of models.dev (MIT) says which models support tool calls.
Models without them get a "no tool calls" badge in the picker and a hint
in the model list, since drawing needs tool calls. Refresh the snapshot
with scripts/update-model-catalog.mjs.

* fix(mcp-server): reject text between tags, which draw.io cannot open

draw.io reads any text inside a page as compressed page data, so a stray
text node makes the whole page fail with an atob error. gpt-5-mini sends
new cells with a literal "\n" between the tags; the edit card said
Complete while draw.io showed the error and kept the old diagram.

Validation now reports text between tags, and auto-fix turns a literal
\n, \t or \r between tags into whitespace. Other text goes back to the
model as an error. The compressed data directly under <diagram> is fine.

* fix(chat): clearer provider errors and no empty bubble, found with real models

- An error object sent inside the stream (OpenRouter's { code, message })
  showed as "[object Object]"; its message and status code are read now.
- A problem+json "detail" is added to the message: NVIDIA only said
  "Gone" for a retired model. 410 counts as model not found.
- "Cannot connect to API" from the SDK gets the connection hint.
- Text that is only whitespace (Kimi K2.6 sends a space before a tool
  call) no longer shows an empty bubble.
- allowSystemInMessages stops the warning on every request. Our system
  messages carry cache points; a client's own system messages are already
  dropped by the empty-content filter.

* feat(providers): suggest the Claude 5 and GPT 6 models, show GPT 6 reasoning

The suggested models stopped at Claude Opus 4.8 and GPT 5.5. They now
start with Claude Opus 5.5, Sonnet 5.5 and Fable 5.1, and gpt-6.1-sol,
gpt-6-sol, gpt-6-luna and gpt-6-astra (ids checked against the provider
lists). The Bedrock list is unchanged until its ids are checked.

The reasoning summary was only turned on for model ids containing o1,
o3, o4 or gpt-5, so GPT 6 models showed no thinking. It now matches the
o-series and gpt-5 or later by version, like the OpenAI SDK.

* fix(providers): keep thinking on the newest Claude models, fix the Bedrock ids

Measured on Bedrock: Claude Opus 4.7, 4.8 and every Claude 5 model reject
a thinking budget ("thinking.type.enabled") and a temperature, and all of
them accept adaptive thinking. The retry used to drop thinking, so with
a thinking budget set these models did not think at all. It now switches
to adaptive thinking with display "summarized"; without that setting the
models think but send no thinking text to show.

The suggested Bedrock ids had no region prefix, and the newer models only
answer through an inference profile id: "anthropic.claude-sonnet-5-5"
fails with "on-demand throughput isn't supported". Each suggested id was
called once; the Claude ones now start with "global.", Llama, Pixtral and
Nova 2 Lite with "us.". Nova Premier (end of life) and the ids without a
working profile are gone. The docs example had the same problem and a
wrong date.

* fix(chat): an edit after a broken edit call no longer fails, found with Opus 5.5

- Claude Opus 5.5 sent an edit with invalid JSON, then the same edit
  again. The first call's streamed preview was never undone: its input
  has no operations, and the undo sat behind that check. The second edit
  then started from the preview, failed on a duplicate id, and the model
  had to try a third time. The undo now runs first, and an edit that
  starts in the same render uses the undone diagram.
- The SDK passes an invalid tool call's error as a string, which was
  wrapped as a provider error. streamErrorText keeps it as the text the
  model reads.
- Bedrock's "on-demand throughput isn't supported" gets the model id hint.
- The thinking header uses the page language ("Thought for 1 second" in
  English), from the dictionary entries that were already there.

* chore(mcp-server): release the new features as 0.3.0

npm has 0.2.3. The earlier commits raised the version step by step to
0.6.0, but they ship together as one release.

* style: auto-format with Biome

* fix(chat): draw the built-in examples again and undo edit previews on errors

Found by the PR review:
- The built-in examples showed a finished card and an empty canvas. They
  are answered in the browser, never reach the tool handler, and relied
  on the final redraw that an earlier commit removed. The example branch
  now loads its diagram itself.
- When the request failed while an edit was streaming (a provider error,
  a lost connection), its preview stayed on the canvas. The error handler
  now restores the diagram from before the preview.
- The model picker could not scroll with the wheel or touch: the settings
  dialog blocks those events outside itself, and the picker is rendered
  outside it. The popover is modal now.
- A fetch error and the open picker stayed when switching providers.
- Editing a model id kept the old test warning and response time, which
  also hid the "may not be able to draw" hint for the new id.

* fix(mcp-server): keep both pages when get_diagram meets a page export, and more review fixes

Found by the PR review, each with a test that failed first:
- get_diagram during a page export returned the one-page projection on
  screen as the whole document (6 of 6 times when timed so). The preview
  page no longer answers a sync while a projection shows, and syncs after
  reloading, so the poll that restores the real document exports it.
- Exports are numbered on the server too: a late result of an export that
  timed out was saved as the next export's file.
- In Chrome, a new_xml with a syntax error counted the <parsererror>
  element as a second cell, so the web app rejected edits that auto-fix
  repairs ("must contain exactly one cell").
- hasCells missed single-quoted ids, so screenshot_diagram called such a
  diagram empty and auto-save never created its file.
- A literal \n directly under a <diagram> that has a model passed
  validation; only text-only pages are compressed data.
- A wrapped mxCell repeating its UserObject's id took the wrapper's place
  in edits, so delete and update left an empty or nested wrapper.
- Bare cells with a shape or edge id of "0" or "1" are rejected with a
  clear message instead of being renamed, which broke their edges.
- DRAWIO_DATA_DIR expands ~, which JSON configs pass on as it is.

* fix(server): count quota by the key actually used, and more review fixes

Found by the PR review, each with a test that failed first:
- Quota: any key header skipped it, even one the provider never reads
  (x-aws-access-key-id with OpenAI), so a request ran on the server's
  key without being counted. The check now runs after the model is
  resolved and uses usesServerCredentials. On main already.
- usesServerCredentials read the raw base URL; "/" cleans up to none, so
  an Ollama request ran on the server's key past the server-model check.
- SGLang's default 127.0.0.1:8000 only fills the settings form. Chat and
  the model list used it as a real address, so the server called its own
  machine even with private URLs blocked. Now a base URL is required.
- With a user's OpenAI key and no base URL, the SDK read the server's
  OPENAI_BASE_URL. The official endpoint is now passed. On main already.
- The Test button refused nothing on the server's keys (Ollama Cloud),
  and a 15 s timeout reported "connected, no tool call".
- The model list for Ollama without a base URL came from ollama.com while
  chat went to the server's Ollama.
- Bedrock's "Too many tokens, please wait" counted as context too long.
- On the server's keys the provider's error text stays in the server log;
  it can name the server's AWS account, role or internal hosts.
- Desktop app: the preset keys are the user's own (NEXT_AI_DRAWIO_DESKTOP),
  so Max Output Tokens can be raised and keyless models in settings work
  again. A launch that found the remembered port taken no longer replaces
  it, which hid the user's chats and settings for good.

* ci: run auto-format with the Biome version in package.json again

package.json moved to Biome 2.5.7 but the auto-format job stayed on
2.4.13. The two format some files differently, so on this PR the bot
reformatted tests/unit/log-save-route.test.ts with 2.4.13 and the lint
job, on 2.5.7, then failed on it. The pin now matches package.json, as
its comment asks, and the file is back in 2.5.7's format.

* fix(chat): keep the canvas after an unrelated error, and more review fixes

Found by the second PR review:
- After a streamed edit, an older render of the stream stored the edit's
  original diagram again, and the next failed request (no quota, a lost
  connection) put that old diagram back on the canvas. The tool handler
  now marks its call as handled, so the preview code leaves it alone.
- An edit applied before the UI showed an earlier broken edit's error was
  erased when that error undid its preview, or was built on that preview.
  The handler now starts from the diagram before all unhandled previews,
  and reads the diagram state that updates at once.
- A failed or stopped display_diagram left its half drawn diagram on the
  canvas. Its preview is undone now, like an edit's.
- "New chat" cleared a chat that could not be saved (storage full).
- The settings dialog showed a model list, a fetch error or a test result
  on the provider that was opened after the request started, and marked a
  model id changed during the test as tested.
- A tool call with broken JSON was shown as cut off by the output limit.
- History entries and session thumbnails could pair with a later diagram
  when draw.io answered an export late.
- A server model saved before non-ASCII provider names got into the id
  was reset to the default model.

* fix(mcp-server): count a one-page view only for that page, and more review fixes

Found by the second PR review:
- get_diagram with a page selector, or a rejected edit's error, counted
  the whole document as seen, so an edit on another page could overwrite
  the user's change there. A one-page view now counts for all pages only
  if the others are unchanged; otherwise the reply says to get them.
- add_page accepted shapes with the root cell ids "0" and "1" and renamed
  them, breaking their edges. The check also missed ids on UserObject
  wrappers and ids written with spaces around the "=".
- Root cells written over two lines were kept as an extra layer, cells
  with id = "a" did not count as cells, and CDATA text before a page's
  model passed the check although draw.io cannot open the page.
- Auto-save cleanup deleted the user's own files that start with mcp-.
  Only names in the session id format are removed now.
- Restoring a history entry dropped edits made in the browser since the
  last entry. They are added to history first.
- A session whose state expired showed a blank page, and the next change
  overwrote its auto-save file. The saved file is loaded instead.
- An edit on a page export's one-page projection, made before the real
  document was back, replaced the whole document.
- A late sync reply could overwrite a newer edit: each sync export is
  numbered, and the server ignores replies older than the current state.
- screenshot_diagram could return another session's image after
  start_session ran during its retries.

* fix(server): use the keys the user sent, and more review fixes

Found by the second PR review:
- With AWS_BEARER_TOKEN_BEDROCK set on the server, a request with the
  user's AWS keys ran on the server's token: the Bedrock SDK prefers it.
  Checked with Bedrock: invalid user keys used to get an answer.
- An OpenAI key with the official URL filled in (the settings form does
  that) went to the Responses API. Back to main's rule: a configured base
  URL uses Chat Completions.
- A user's Ollama key went to the server's OLLAMA_BASE_URL, for chat and
  for the model list. Like every other provider, it goes to the user's
  base URL or Ollama Cloud.
- The server's keyless Ollama and EdgeOne were not counted in the quota.
- AI_MODEL models on the server's keys ran on any provider with a server
  key, not only on AI_PROVIDER.
- A user's Azure key without a base URL used the server's resource name.
- The admin panel's Test button failed whenever access codes were set.
- DeepSeek's errors in the stream (plain text) were shown as they were,
  without a hint and also on the server's keys. Bedrock's throttling in
  the stream was not recognised as a rate limit.
- The EdgeOne function accepted text/plain; x=application/json, which
  other sites can send without a CORS preflight.
- Desktop app: a launch that found the old port taken for a moment (the
  previous version still quitting after an update) remembered the new
  port for good. The new port is kept only when Windows reserves the old
  one. A failed read of the presets file moved it aside as corrupt, and a
  save could then replace the presets. Switching presets on the same port
  now reloads the page. The dev launcher no longer misses a preset change
  made before or during a restart.

* fix(chat): undo rejected or stopped previews, and more fixes from the third review

- A call the server runs (get_shape_library) still reaches the browser's
  tool handler, and it dropped the stored diagram of an earlier broken
  edit, whose preview then stayed. Only the tools that draw take it now.
- A display_diagram whose final XML fails the checks loads the diagram
  from before its preview again, as a failed edit does.
- After Stop, a tool result that arrives later (a screenshot check still
  running) no longer sends a new request; Stop also skips calls the tool
  handler already took.
- New chat and opening another chat kept nothing of a diagram drawn
  without messages when it could not be saved; now they stay on it.
- The settings dialog drops a model list or test result whose provider
  credentials changed meanwhile, also in another tab.
- A saved provider this version does not know crashed the whole page on
  load; it is skipped.
- The input emptied a moment after the message showed in the chat, so it
  briefly appeared twice (seen as a flaky e2e test).
- The desktop app's preset switch on the same port refetches the server
  models instead of reloading the page, which lost unsent attachments.

* fix(mcp-server): recover sessions in one place, and more fixes from the third review

- A session whose state expired was recovered from its auto-save file only
  for the preview page; the tools built on their older copy and then
  overwrote the file. They now recover it first (restoreSavedSession).
- A preview tab that missed the last AI write pushed its older diagram
  over the recovered one after a restart. It now shows the recovered
  diagram and keeps its own copy in History.
- An empty record of what the model has seen (after load_diagram or a
  page tool on unseen changes) no longer lets one page of a multi-page
  document count for all, and get_diagram counts a page only once found.
- History: a thumbnail goes only to the entry it shows, the cached image
  never belongs to an older diagram, a re-serialized copy adds no entry,
  and a cleared document with its own pages is kept before a restore.
- The root cell id check reads attributes one by one: rack-id="1" or id
  text inside a label no longer counts.
- A compressed page counts as having cells; a saved file that could not
  be read is never written over.

* fix(server): keep users' keys at their own endpoints, and more fixes from the third review

- Bedrock: a user's AWS keys no longer go to an endpoint the server sets
  in AWS_ENDPOINT_URL_BEDROCK_RUNTIME / AWS_ENDPOINT_URL (read by the
  upgraded SDK), and admin panel keys win over AWS_BEARER_TOKEN_BEDROCK,
  as the Test button checks them. Checked with Bedrock.
- Ollama: a server key without a base URL (admin panel, OLLAMA_API_KEY)
  goes to Ollama Cloud, as env.example says, instead of 127.0.0.1.
- Quota: EdgeOne counts whatever key header comes along, keyless Ollama at
  a private address counts, and their provider texts stay in the log.
- An EdgeOne server model (admin panel, ai-models.json) works: the route
  checked the raw provider header, which holds the name's slug.
- parse-url ends downloads it does not read (too large, PDF, errors).
- Desktop app: the port follows where the chats are (IndexedDB per
  origin) instead of a remembered port, which could hide them for good;
  a same-port restart tells the page to refetch the server models; a
  failed preset switch no longer undoes a newer choice; a presets file
  removed after a failed read can be saved again; .env values quoted from
  start to end keep their inner quotes, as dotenv reads them.

* fix(security): check request sources, regions and endpoints

- Bedrock: a request's AWS region must be a region name. It becomes part
  of the endpoint's host name, so a value such as
  "us-east-1.attacker.example/" sent the server's bearer token or signed
  request to another host.
- MCP preview server: only the preview page itself (Origin equal to the
  Host) or a non-browser client may call it; a page on another localhost
  port could replace the diagram with a plain text POST. History builds
  its thumbnails element by element and shows only SVG data images, so a
  stored value can no longer run script in the preview.
- chat, validate-model, validate-diagram, provider-models and parse-url
  take JSON bodies only, so another website cannot make the user's own
  server (the desktop app, a local install) run models with their keys;
  the desktop app also refuses a foreign Host (DNS rebinding).
- The model list reads at most 2 MB, also through the Gateway SDK, and
  answers only with its own error texts: the URL is the caller's and may
  be an internal address.
- An admin panel provider with its own key and no URL no longer inherits
  the global <P>_BASE_URL, which may be a proxy for another key; OpenAI
  then gets the official endpoint, as its Test. Azure keeps the server's
  resource.

* fix: what the third round broke, and the first batch's review

MCP preview after the server lost a session (it expired, or the MCP
process restarted):
- Every server state has an id, made when the state is created. The tab
  notices a new id even when the version numbers happen to match, and
  every push names the state it was based on, so one based on a lost state
  is refused, also when it comes before the tab's first poll (the server
  recovers the saved file first).
- The tab keeps the newest canvas XML, saved or not. When the server knows
  nothing (no file) or exactly what the tab last saved, the canvas wins and
  is saved, so edits made while the server was down are kept. Otherwise
  the server's diagram (an AI write the tab missed, a cleared document
  that was saved) is shown and the tab's copy goes to History.
- Late answers to an old state's push or poll are dropped; a failed push
  says the server is unreachable; Download as .drawio saves the canvas.

Settings and server:
- Saved providers this version does not know stay in storage with their
  keys, and sending no longer trips over them.
- The desktop "Ollama (Local)" preset with a key goes to local Ollama
  again; a server model's Ollama URL variable is read; the admin panel
  writes Ollama Cloud's URL for a key without one.
- Provider error texts show again in the desktop app and for EdgeOne.
- .env: a quoted value followed by a comment ending in a quote is read as
  dotenv reads it; unquoted values are unchanged.
- Desktop app: the next launch opens the port where a chat was last
  saved; a launch elsewhere that saves nothing does not move it, and a
  page with no chats lets the next launch try the other port once.
- The Test button no longer stays busy after another tab changed the key.
- A completed append_diagram is no longer undone by an earlier failed
  edit's preview; a file read once in vain is saved again once it is read
  or gone.

From the first batch's review:
- The admin panel's Test of an entry without a URL now tests the server's
  <P>_BASE_URL, where chat sends the entry's key; chat is unchanged (the
  first fix rerouted working setups).
- The model list ends downloads that are too large, accepts answers
  without a body, and keeps the "redirects are not allowed" explanation.
- A test covers the preview's History rendering.

* fix: older defects (batch C) and the second batch's review

Chats:
- New Chat right after an answer saves that chat once. Saves run one at
  a time and read the chat on screen when their turn comes; a save
  scheduled for a chat that is no longer on screen is dropped. A chat
  whose id was still on its way to the URL no longer comes back after
  New Chat (the next answer went into it).
- Crossing the 768 px breakpoint keeps the chat panel: a streaming answer,
  unsaved messages and attachments stay. The panel gets the sizes of each
  side, and a panel collapsed on desktop opens on mobile.
- The chat's export waits for its own reply: an edit's history export
  still on its way no longer answers it with the older diagram, and two
  file saves at once no longer swap results.
- A second edit in one answer is previewed on the first edit's result.
- Stop also ends a running screenshot check; a chat that cannot be saved
  (storage full) can be left with "Continue without saving".
- Small diagrams with shapes count as diagrams; the tool card no longer
  crashes on malformed operations.

Quota and providers:
- Requests that reach the server's own endpoints count toward the quota:
  EdgeOne (always its own endpoint now), a private base URL whatever key
  header is sent, keyless Ollama without a URL. With the quota on, a
  redirect is followed only to a public address. The output cap applies
  to these requests too.
- Stop records the tokens of the steps that finished; the screenshot check
  counts its tokens without counting a request.
- EdgeOne configured only by AI_PROVIDER works, also in the admin Test,
  which forwards the access code. Azure set up only in the admin panel
  works in chat. The Test sends a Bedrock session token.
- The admin panel's Test of an entry without a URL uses the server's URL
  as the server does (no private address check for it); the admin panel
  no longer writes an Ollama URL.

MCP server:
- Write tools and start_session run one at a time, so two at once never
  drop each other's change; a cancelled call waiting its turn is skipped.
  get_diagram and export_diagram keep the session they started with.
- Export to .drawio first gets the user's latest edits from the browser.
- History thumbnails: one that arrives after the next AI write is
  dropped; a sync reply keeps the image; a version that changed only page
  settings is its own entry.
- A diagram over the 10 MB limit is saved without its image, or the user
  is told to download it (the server now answers 413 instead of cutting
  the connection).
- Labels holding text like id='1' or parent='1' are no longer read as
  attributes (a layer or a parent was deleted). A broken bare
  <mxGraphModel> file is refused.
- After a sync reply the tab no longer sends its autosave copy again.

Desktop and files:
- A newer switch of the same preset is not rolled back by an older one
  that failed. .env values with escaped quotes are read whole.
- MCP saved files: a file that could not be read stays protected while a
  folder without permission hides it, and is saved again once deleted.
- The desktop app reports "no chats" only when the count was read and no
  model settings are stored.

* fix: what the batch C review found

- 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.

* 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.

* fix(chat): stop saving and exporting an idle chat every second

Each auto-save takes a thumbnail, and draw.io's SVG gets a new random id
every time, so latestSvg changed and the diagram context rendered again.
getThumbnailSvg was a new function on every render, the auto-save depends
on it, so the next save was scheduled a second later: a chat with some
diagrams was saved and exported once a second for as long as it was open.
getThumbnailSvg (and the export helper it uses) now keep one identity;
they read refs only.

Found by hand in the production build; the new e2e test fails without the
fix on a production build (as CI runs it).

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Dayuan Jiang
2026-10-05 22:39:32 +09:00
committed by GitHub
co-authored by github-actions[bot]
parent a45e5b6796
commit 0b63e28e5a
192 changed files with 20781 additions and 6839 deletions
+71
View File
@@ -0,0 +1,71 @@
// @vitest-environment node
import fs from "fs"
import os from "os"
import path from "path"
import { afterEach, beforeEach, describe, expect, it } from "vitest"
import { GET, PUT } from "@/app/api/admin/providers/route"
import { _resetForTests } from "@/lib/admin/settings"
let tmpDir: string
beforeEach(() => {
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "admin-providers-route-"))
process.env.SETTINGS_FILE = path.join(tmpDir, "settings.json")
process.env.ADMIN_PASSWORD = "pw"
process.env.AI_MODELS_CONFIG_PATH = path.join(tmpDir, "none.json")
_resetForTests()
})
afterEach(() => {
_resetForTests()
delete process.env.SETTINGS_FILE
delete process.env.ADMIN_PASSWORD
delete process.env.AI_MODELS_CONFIG_PATH
delete process.env.AI_MODEL
fs.rmSync(tmpDir, { recursive: true, force: true })
})
const headers = { "x-admin-password": "pw" }
async function saveDefaultPanelProvider() {
const res = await PUT(
new Request("http://localhost/api/admin/providers", {
method: "PUT",
headers: { ...headers, "Content-Type": "application/json" },
body: JSON.stringify({
providers: [
{
id: "p1",
provider: "openai",
apiKey: "sk-test",
models: ["gpt-panel"],
isDefault: true,
},
],
}),
}),
)
expect(res.status).toBe(200)
}
async function envHasDefaultModel(): Promise<boolean> {
const res = await GET(
new Request("http://localhost/api/admin/providers", { headers }),
)
return (await res.json()).envHasDefaultModel
}
describe("envHasDefaultModel", () => {
it("is true when .env sets AI_MODEL, even after a panel default", async () => {
process.env.AI_MODEL = "gpt-env"
expect(await envHasDefaultModel()).toBe(true)
await saveDefaultPanelProvider()
expect(await envHasDefaultModel()).toBe(true)
})
it("ignores the AI_MODEL the panel default writes", async () => {
await saveDefaultPanelProvider()
expect(process.env.AI_MODEL).toBe("gpt-panel")
expect(await envHasDefaultModel()).toBe(false)
})
})
+75 -4
View File
@@ -69,7 +69,35 @@ describe("deriveEnvUpdates", () => {
expect(updates.ADMIN_OPENAI_API_KEY_2).toBe("sk-second")
})
it("maps bedrock credentials to AWS env vars", () => {
it("writes an Ollama URL only when the entry has one", () => {
// Without one, the operator's own OLLAMA_BASE_URL (or local Ollama)
// stays, also for the AI_PROVIDER=ollama default model
const keyOnly = deriveEnvUpdates(
[provider({ provider: "ollama", apiKey: "ollama-key" })],
[],
)
expect(keyOnly.OLLAMA_API_KEY).toBe("ollama-key")
expect(keyOnly.OLLAMA_BASE_URL ?? null).toBeNull()
const own = deriveEnvUpdates(
[
provider({
provider: "ollama",
apiKey: "k",
baseUrl: "https://ollama.internal/api",
}),
],
[],
)
expect(own.OLLAMA_BASE_URL).toBe("https://ollama.internal/api")
// No key: local Ollama, nothing to write
const local = deriveEnvUpdates(
[provider({ provider: "ollama", apiKey: undefined })],
[],
)
expect(local.OLLAMA_BASE_URL ?? null).toBeNull()
})
it("maps bedrock credentials to ADMIN_AWS_* env vars", () => {
const updates = deriveEnvUpdates(
[
provider({
@@ -83,9 +111,26 @@ describe("deriveEnvUpdates", () => {
],
[],
)
expect(updates.AWS_ACCESS_KEY_ID).toBe("AKIA123")
expect(updates.AWS_SECRET_ACCESS_KEY).toBe("secret")
expect(updates.AWS_REGION).toBe("us-west-2")
expect(updates.ADMIN_AWS_ACCESS_KEY_ID).toBe("AKIA123")
expect(updates.ADMIN_AWS_SECRET_ACCESS_KEY).toBe("secret")
expect(updates.ADMIN_AWS_REGION).toBe("us-west-2")
// Standard AWS vars are left to the environment
expect(updates.AWS_ACCESS_KEY_ID).toBeUndefined()
})
it("clears AWS_* bedrock keys written by older versions", () => {
const bedrock = provider({
provider: "bedrock",
apiKey: undefined,
awsAccessKeyId: "AKIA123",
awsSecretAccessKey: "secret",
models: ["claude-x"],
})
const updates = deriveEnvUpdates([bedrock], [bedrock])
expect(updates.AWS_ACCESS_KEY_ID).toBeNull()
expect(updates.AWS_SECRET_ACCESS_KEY).toBeNull()
expect(updates.AWS_REGION).toBeNull()
expect(updates.ADMIN_AWS_ACCESS_KEY_ID).toBe("AKIA123")
})
it("clears keys owned by the previous list when providers are removed", () => {
@@ -315,6 +360,32 @@ describe("validateAdminProviders", () => {
expect(validateAdminProviders(list)).toMatch(/unique/)
})
it("rejects names that differ only in case or punctuation", () => {
const list = [
provider({ id: "p1", name: "Open AI" }),
provider({ id: "p2", name: "open-ai" }),
]
expect(validateAdminProviders(list)).toMatch(/unique/)
})
it("rejects a case-only clash with an env-configured name", () => {
expect(
validateAdminProviders([provider({ name: "openai" })], {
providers: [
{ name: "OpenAI", provider: "openai", models: ["gpt-x"] },
],
}),
).toMatch(/already defined/)
})
it("accepts distinct CJK names", () => {
const list = [
provider({ id: "p1", provider: "deepseek", name: "主力" }),
provider({ id: "p2", provider: "deepseek", name: "备用" }),
]
expect(validateAdminProviders(list)).toBeNull()
})
it("rejects multiple defaults", () => {
const list = [
provider({ id: "p1", isDefault: true }),
+87
View File
@@ -0,0 +1,87 @@
import { cleanup, fireEvent, render, screen } from "@testing-library/react"
import type { ReactNode } from "react"
import { afterEach, describe, expect, it, vi } from "vitest"
import { SecretInput, SettingField } from "@/app/[lang]/admin/setting-field"
import { DictionaryProvider } from "@/hooks/use-dictionary"
import type { SettingDef } from "@/lib/admin/settings-registry"
import type { Dictionary } from "@/lib/i18n/dictionaries"
import en from "@/lib/i18n/dictionaries/en.json"
const STORED = { isSet: true as const, hint: "…abcd" }
function withDict(node: ReactNode) {
return (
<DictionaryProvider dictionary={en as unknown as Dictionary}>
{node}
</DictionaryProvider>
)
}
function typeInto(label: string, text: string) {
fireEvent.change(screen.getByLabelText(label, { selector: "input" }), {
target: { value: text },
})
}
afterEach(cleanup)
describe("SecretInput", () => {
it("reverts to a key saved after mount instead of deleting it", () => {
const onChange = vi.fn()
const props = { id: "secret", keepOnEmpty: true, onChange }
// New provider: nothing stored at mount, then saved
const { rerender } = render(
withDict(
<>
<label htmlFor="secret">secret</label>
<SecretInput {...props} value={undefined} />
</>,
),
)
rerender(
withDict(
<>
<label htmlFor="secret">secret</label>
<SecretInput {...props} value={STORED} />
</>,
),
)
rerender(
withDict(
<>
<label htmlFor="secret">secret</label>
<SecretInput {...props} value="abc" />
</>,
),
)
typeInto("secret", "")
expect(onChange).toHaveBeenLastCalledWith(STORED)
})
})
describe("SettingField secret", () => {
const def: SettingDef = {
key: "LANGFUSE_SECRET_KEY",
group: "observability",
type: "secret",
label: "Langfuse Secret Key",
}
it("drops the pending change when a saved secret is typed over and cleared", () => {
const onChange = vi.fn()
const props = {
def,
state: { key: def.key, source: "file" as const, value: STORED },
disabled: false,
onChange,
}
const { rerender } = render(
withDict(<SettingField {...props} pendingValue={undefined} />),
)
rerender(withDict(<SettingField {...props} pendingValue="a" />))
typeInto(en.admin.settings.LANGFUSE_SECRET_KEY.label, "")
expect(onChange).toHaveBeenLastCalledWith(undefined)
})
})
+19 -1
View File
@@ -1,7 +1,7 @@
import fs from "fs"
import os from "os"
import path from "path"
import { afterEach, beforeEach, describe, expect, it } from "vitest"
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import {
_resetForTests,
applyToEnv,
@@ -100,6 +100,24 @@ describe("applyToEnv / saveSettings", () => {
expect(process.env.TEST_ADMIN_VAR).toBeUndefined()
})
it("a second module instance can remove a key the first one overlaid", async () => {
// instrumentation.ts and API routes load separate copies in a build
process.env.TEST_ADMIN_VAR = "from-env"
fs.writeFileSync(
process.env.SETTINGS_FILE!,
JSON.stringify({ version: 1, values: { TEST_ADMIN_VAR: "abc" } }),
)
applyToEnv()
expect(process.env.TEST_ADMIN_VAR).toBe("abc")
vi.resetModules()
const second = await import("@/lib/admin/settings")
expect(second.getValueSource("TEST_ADMIN_VAR")).toBe("file")
second.saveSettings({ TEST_ADMIN_VAR: null })
expect(process.env.TEST_ADMIN_VAR).toBe("from-env")
expect(second.getEnvFallback("TEST_ADMIN_VAR")).toBe("from-env")
})
it("persists across cache reset (file round-trip)", () => {
saveSettings({ TEST_ADMIN_VAR: "persisted" })
_resetForTests()
+119
View File
@@ -0,0 +1,119 @@
// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
// The request the admin Test hands to validate-model
const sent = vi.hoisted(() => ({ body: null as any, headers: null as any }))
vi.mock("@/app/api/validate-model/route", () => ({
POST: async (req: Request) => {
sent.body = await req.json()
sent.headers = Object.fromEntries(req.headers)
return Response.json({ valid: true })
},
}))
vi.mock("@/lib/admin/auth", () => ({ checkAdminAuth: () => null }))
// 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",
"AZURE_BASE_URL",
"AZURE_RESOURCE_NAME",
]
const saved: Record<string, string | undefined> = {}
beforeEach(() => {
envFallback.values = {}
for (const k of ENV) {
saved[k] = process.env[k]
delete process.env[k]
}
})
afterEach(() => {
for (const k of ENV) {
if (saved[k] === undefined) delete process.env[k]
else process.env[k] = saved[k]
}
})
const test = (provider: Record<string, unknown>) =>
testModel(
new Request("http://localhost/api/admin/test-model", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
provider: { id: "p1", models: ["m"], ...provider },
modelId: "m",
}),
}),
)
describe("admin Test of an entry without a URL", () => {
it("tests the server's <P>_BASE_URL, where chat sends the entry's key", async () => {
// A server model without baseUrlEnv reads the global variable
process.env.OPENAI_BASE_URL = "https://operator-proxy.example.com/v1"
await test({ provider: "openai", apiKey: "panel-key" })
expect(sent.body.baseUrl).toBe("https://operator-proxy.example.com/v1")
// The server's own URL, tested without the rules for typed URLs
expect(sent.body.serverBaseUrl).toBe(true)
process.env.AI_GATEWAY_BASE_URL = "https://gateway.example.com/v3/ai"
await test({ provider: "gateway", apiKey: "k" })
expect(sent.body.baseUrl).toBe("https://gateway.example.com/v3/ai")
})
it("keeps the entry's own URL, and none when the server has none", async () => {
process.env.SGLANG_BASE_URL = "http://gpu-box:8000/v1"
await test({
provider: "sglang",
apiKey: "k",
baseUrl: "http://other:8000/v1",
})
expect(sent.body.baseUrl).toBe("http://other:8000/v1")
await test({ provider: "deepseek", apiKey: "k" })
expect(sent.body.baseUrl).toBeUndefined()
})
it("does not use Vertex's variable, which the panel writes itself", async () => {
// Before a save it still holds the entry's previous URL
process.env.GOOGLE_VERTEX_BASE_URL = "https://old-proxy.example.com"
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 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)
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")
})
})
+85
View File
@@ -0,0 +1,85 @@
// @vitest-environment node
import { generateText, type ModelMessage } from "ai"
import { afterEach, describe, expect, it, vi } from "vitest"
import { CACHE_POINT, getAIModel } from "@/lib/ai-providers"
const ENV = ["ANTHROPIC_API_KEY", "OPENROUTER_API_KEY"]
const saved = Object.fromEntries(ENV.map((k) => [k, process.env[k]]))
afterEach(() => {
for (const k of ENV) {
if (saved[k] === undefined) delete process.env[k]
else process.env[k] = saved[k]
}
vi.unstubAllGlobals()
})
// The chat route marks its system messages like this
const messages: ModelMessage[] = [
{ role: "system", content: "Instructions", providerOptions: CACHE_POINT },
{ role: "user", content: "hi" },
]
/** Send the messages through a real provider and return the request body */
async function requestBody(
provider: "anthropic" | "openrouter",
reply: unknown,
) {
let body: any
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init: RequestInit) => {
body = JSON.parse(init.body as string)
return new Response(JSON.stringify(reply), {
headers: { "content-type": "application/json" },
})
}),
)
const { model } = getAIModel({
provider,
modelId:
provider === "anthropic"
? "claude-sonnet-4-5"
: "anthropic/claude-sonnet-4.5",
})
await generateText({ model, messages, maxRetries: 0 }).catch(() => {})
return body
}
describe("prompt cache breakpoints", () => {
it("reach the Anthropic API", async () => {
process.env.ANTHROPIC_API_KEY = "test-key"
const body = await requestBody("anthropic", {
id: "msg_1",
type: "message",
role: "assistant",
content: [{ type: "text", text: "ok" }],
stop_reason: "end_turn",
usage: { input_tokens: 1, output_tokens: 1 },
})
expect(body.system).toEqual([
{
type: "text",
text: "Instructions",
cache_control: { type: "ephemeral" },
},
])
})
it("reach OpenRouter", async () => {
process.env.OPENROUTER_API_KEY = "test-key"
const body = await requestBody("openrouter", {
id: "gen-1",
choices: [
{
index: 0,
message: { role: "assistant", content: "ok" },
finish_reason: "stop",
},
],
})
expect(JSON.stringify(body.messages[0])).toContain(
'"cache_control":{"type":"ephemeral"}',
)
})
})
+508
View File
@@ -0,0 +1,508 @@
import { createOpenAI } from "@ai-sdk/openai"
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { AWS_REGIONS } from "@/components/provider-credentials-fields"
import {
getAIModel,
getValidationModel,
usesServerCredentials,
} from "@/lib/ai-providers"
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", () => {
const mockProviderFn = vi.fn(() => ({ modelId: "test-model" }))
return { createVertex: vi.fn(() => mockProviderFn) }
})
vi.mock("@ai-sdk/openai", () => {
const mockModel = { modelId: "test-model" }
const mockProviderFn = vi.fn(() => mockModel) as any
mockProviderFn.chat = vi.fn(() => mockModel)
return {
createOpenAI: vi.fn(() => mockProviderFn),
openai: vi.fn(() => mockModel),
}
})
vi.mock("@ai-sdk/azure", () => {
const mockModel = { modelId: "test-model" }
const mockProviderFn = vi.fn(() => mockModel) as any
mockProviderFn.chat = vi.fn(() => mockModel)
mockProviderFn.responses = vi.fn(() => mockModel)
return { createAzure: vi.fn(() => mockProviderFn) }
})
vi.mock("@ai-sdk/amazon-bedrock", () => {
const mockProviderFn = vi.fn(() => ({ modelId: "test-model" }))
return { createAmazonBedrock: vi.fn(() => mockProviderFn) }
})
vi.mock("@aws-sdk/credential-providers", () => ({
fromNodeProviderChain: vi.fn(() => "node-chain"),
}))
vi.mock("ollama-ai-provider-v2", () => {
const mockProviderFn = vi.fn(() => ({ modelId: "test-model" }))
return { createOllama: vi.fn(() => mockProviderFn) }
})
vi.mock("@openrouter/ai-sdk-provider", () => {
const mockProviderFn = vi.fn(() => ({ modelId: "test-model" }))
return { createOpenRouter: vi.fn(() => mockProviderFn) }
})
const ENV_KEYS = [
"GOOGLE_VERTEX_API_KEY",
"GOOGLE_VERTEX_BASE_URL",
"OPENAI_API_KEY",
"OPENAI_BASE_URL",
"OPENROUTER_API_KEY",
"ADMIN_OPENAI_API_KEY",
"ADMIN_OPENROUTER_API_KEY",
"OLLAMA_API_KEY",
"ADMIN_AWS_ACCESS_KEY_ID",
"ADMIN_AWS_SECRET_ACCESS_KEY",
"ADMIN_AWS_REGION",
"AWS_REGION",
"AI_PROVIDER",
"AI_MODEL",
"VALIDATION_MODEL",
"NEXT_AI_DRAWIO_DESKTOP",
"SGLANG_API_KEY",
"SGLANG_BASE_URL",
"AZURE_RESOURCE_NAME",
"OLLAMA_BASE_URL",
]
const savedEnv: Record<string, string | undefined> = {}
beforeEach(() => {
for (const key of ENV_KEYS) {
savedEnv[key] = process.env[key]
delete process.env[key]
}
settings.values = {}
vi.clearAllMocks()
})
afterEach(() => {
for (const key of ENV_KEYS) {
if (savedEnv[key] === undefined) delete process.env[key]
else process.env[key] = savedEnv[key]
}
})
describe("Vertex AI key security", () => {
it("never sends the server key to a client base URL", () => {
process.env.GOOGLE_VERTEX_API_KEY = "server-vertex-key"
// Any x-ai-api-key passes the outer guard; the branch must still refuse
expect(() =>
getAIModel({
provider: "vertexai",
apiKey: "x",
baseUrl: "https://attacker.example",
modelId: "gemini-2.5-flash",
}),
).toThrow("Vertex AI requires an API key")
})
it("sends the client key to the client base URL", async () => {
process.env.GOOGLE_VERTEX_API_KEY = "server-vertex-key"
const { createVertex } = await import("@ai-sdk/google-vertex")
getAIModel({
provider: "vertexai",
vertexApiKey: "client-key",
baseUrl: "https://my-proxy.example",
modelId: "gemini-2.5-flash",
})
expect(createVertex).toHaveBeenCalledWith({
apiKey: "client-key",
baseURL: "https://my-proxy.example",
})
})
it("does not send the client key to the server's base URL", async () => {
process.env.GOOGLE_VERTEX_BASE_URL = "https://server-proxy.internal"
const { createVertex } = await import("@ai-sdk/google-vertex")
getAIModel({
provider: "vertexai",
vertexApiKey: "client-key",
modelId: "gemini-2.5-flash",
})
expect(createVertex).toHaveBeenCalledWith({ apiKey: "client-key" })
})
it("still uses the server key and base URL without client overrides", async () => {
process.env.GOOGLE_VERTEX_API_KEY = "server-vertex-key"
process.env.GOOGLE_VERTEX_BASE_URL = "https://server-proxy.internal"
const { createVertex } = await import("@ai-sdk/google-vertex")
getAIModel({ provider: "vertexai", modelId: "gemini-2.5-flash" })
expect(createVertex).toHaveBeenCalledWith({
apiKey: "server-vertex-key",
baseURL: "https://server-proxy.internal",
})
})
})
describe("Bedrock admin panel credentials", () => {
it("uses the ADMIN_AWS_* keys when the client sends none", async () => {
process.env.ADMIN_AWS_ACCESS_KEY_ID = "panel-id"
process.env.ADMIN_AWS_SECRET_ACCESS_KEY = "panel-secret"
process.env.ADMIN_AWS_REGION = "eu-west-1"
process.env.AWS_REGION = "us-east-1"
const { createAmazonBedrock } = await import("@ai-sdk/amazon-bedrock")
getAIModel({ provider: "bedrock", modelId: "amazon.nova-lite-v1:0" })
expect(createAmazonBedrock).toHaveBeenCalledWith({
region: "eu-west-1",
accessKeyId: "panel-id",
secretAccessKey: "panel-secret",
// The keys the Test button checked, not AWS_BEARER_TOKEN_BEDROCK
apiKey: "",
})
})
it("prefers the client's keys and region", async () => {
process.env.ADMIN_AWS_ACCESS_KEY_ID = "panel-id"
process.env.ADMIN_AWS_SECRET_ACCESS_KEY = "panel-secret"
process.env.ADMIN_AWS_REGION = "eu-west-1"
const { createAmazonBedrock } = await import("@ai-sdk/amazon-bedrock")
getAIModel({
provider: "bedrock",
modelId: "amazon.nova-lite-v1:0",
awsAccessKeyId: "client-id",
awsSecretAccessKey: "client-secret",
awsRegion: "ap-northeast-1",
})
expect(createAmazonBedrock).toHaveBeenCalledWith({
region: "ap-northeast-1",
accessKeyId: "client-id",
secretAccessKey: "client-secret",
// The SDK would otherwise use the server's AWS_BEARER_TOKEN_BEDROCK
apiKey: "",
// and the server's AWS_ENDPOINT_URL_BEDROCK_RUNTIME
baseURL: "https://bedrock-runtime.ap-northeast-1.amazonaws.com",
})
})
it("refuses a region that is not a region name", async () => {
// It becomes part of the endpoint's host name, with the server's
// credentials too
process.env.ADMIN_AWS_ACCESS_KEY_ID = "panel-id"
process.env.ADMIN_AWS_SECRET_ACCESS_KEY = "panel-secret"
const { createAmazonBedrock } = await import("@ai-sdk/amazon-bedrock")
for (const awsRegion of [
"us-east-1.attacker.example/",
"x/#",
"US-EAST-1",
"us-east-1 ",
]) {
expect(() =>
getAIModel({
provider: "bedrock",
modelId: "amazon.nova-lite-v1:0",
awsRegion,
}),
).toThrow(/Invalid AWS region/)
expect(() =>
getAIModel({
provider: "bedrock",
modelId: "amazon.nova-lite-v1:0",
awsAccessKeyId: "client-id",
awsSecretAccessKey: "client-secret",
awsRegion,
}),
).toThrow(/Invalid AWS region/)
}
expect(createAmazonBedrock).not.toHaveBeenCalled()
})
it("accepts every region the settings offer, and other partitions", () => {
for (const awsRegion of [
...AWS_REGIONS.map(([region]) => region),
"us-gov-west-1",
"cn-northwest-1",
"us-iso-east-1",
"eusc-de-east-1",
]) {
expect(() =>
getAIModel({
provider: "bedrock",
modelId: "amazon.nova-lite-v1:0",
awsAccessKeyId: "client-id",
awsSecretAccessKey: "client-secret",
awsRegion,
}),
).not.toThrow()
}
})
it("falls back to the default AWS credential chain", async () => {
process.env.AWS_REGION = "us-east-1"
const { createAmazonBedrock } = await import("@ai-sdk/amazon-bedrock")
getAIModel({ provider: "bedrock", modelId: "amazon.nova-lite-v1:0" })
expect(createAmazonBedrock).toHaveBeenCalledWith({
region: "us-east-1",
credentialProvider: "node-chain",
})
})
})
describe("usesServerCredentials", () => {
it("is true when no key comes with the request", () => {
expect(usesServerCredentials("openai", {})).toBe(true)
expect(usesServerCredentials("openai", { apiKey: "k" })).toBe(false)
})
it("looks at the credential each provider actually uses", () => {
// A stray x-ai-api-key does not replace the IAM role or Vertex key
expect(usesServerCredentials("bedrock", { apiKey: "x" })).toBe(true)
expect(
usesServerCredentials("bedrock", {
awsAccessKeyId: "id",
awsSecretAccessKey: "secret",
}),
).toBe(false)
expect(usesServerCredentials("vertexai", { apiKey: "x" })).toBe(true)
expect(usesServerCredentials("vertexai", { vertexApiKey: "k" })).toBe(
false,
)
})
it("treats keyless EdgeOne and local Ollama as free", () => {
expect(usesServerCredentials("edgeone", {})).toBe(false)
expect(usesServerCredentials("ollama", {})).toBe(false)
expect(
usesServerCredentials("ollama", {
baseUrl: "http://localhost:11434",
}),
).toBe(false)
process.env.OLLAMA_API_KEY = "server-ollama-key"
expect(usesServerCredentials("ollama", {})).toBe(true)
})
})
describe("server model apiKeyEnv", () => {
it("uses the custom env var on the official OpenAI endpoint", async () => {
process.env.ADMIN_OPENAI_API_KEY = "panel-key"
const { createOpenAI, openai } = await import("@ai-sdk/openai")
getAIModel({
provider: "openai",
modelId: "gpt-4o",
apiKeyEnv: "ADMIN_OPENAI_API_KEY",
})
// The default instance would read OPENAI_API_KEY instead
expect(openai).not.toHaveBeenCalled()
expect(createOpenAI).toHaveBeenCalledWith({ apiKey: "panel-key" })
})
})
describe("getValidationModel", () => {
it("uses the admin panel default's ADMIN_ key", async () => {
settings.values = {
ADMIN_PROVIDERS: JSON.stringify([
{
id: "p1",
provider: "openrouter",
name: "My OpenRouter",
apiKey: "panel-key",
models: ["openai/gpt-4o"],
isDefault: true,
},
]),
}
// What deriveEnvUpdates writes for that panel config
process.env.AI_PROVIDER = "openrouter"
process.env.AI_MODEL = "openai/gpt-4o"
process.env.ADMIN_OPENROUTER_API_KEY = "panel-key"
const { createOpenRouter } = await import("@openrouter/ai-sdk-provider")
expect(() => getValidationModel()).not.toThrow()
expect(createOpenRouter).toHaveBeenCalledWith({ apiKey: "panel-key" })
})
it("uses the standard env vars without a panel default", async () => {
process.env.AI_PROVIDER = "openrouter"
process.env.AI_MODEL = "openai/gpt-4o"
process.env.OPENROUTER_API_KEY = "env-key"
const { createOpenRouter } = await import("@openrouter/ai-sdk-provider")
getValidationModel()
expect(createOpenRouter).toHaveBeenCalledWith({ apiKey: "env-key" })
})
})
describe("whose keys a request uses", () => {
it("cleans the base URL like the request does", () => {
// "/" and a pasted path clean up to no base URL: the server's Ollama
process.env.OLLAMA_API_KEY = "server-ollama-key"
expect(usesServerCredentials("ollama", { baseUrl: "/" })).toBe(true)
expect(
usesServerCredentials("ollama", { baseUrl: "/chat/completions" }),
).toBe(true)
})
it("counts the desktop app's keys as the user's own", () => {
// Electron passes the user's preset keys as server env vars
process.env.NEXT_AI_DRAWIO_DESKTOP = "1"
expect(usesServerCredentials("openai", {})).toBe(false)
})
it("sends a user's OpenAI key to the official endpoint", () => {
// The SDK would otherwise read the server's OPENAI_BASE_URL
process.env.OPENAI_BASE_URL = "https://operator-proxy.example.com/v1"
getAIModel({
provider: "openai",
apiKey: "user-key",
modelId: "gpt-5.5",
})
expect(createOpenAI).toHaveBeenLastCalledWith(
expect.objectContaining({
apiKey: "user-key",
baseURL: "https://api.openai.com/v1",
}),
)
// Still the Responses API, like without a base URL
const provider = vi.mocked(createOpenAI).mock.results.at(-1)?.value
expect(provider.chat).not.toHaveBeenCalled()
})
it("uses Chat Completions for any configured base URL", () => {
// The settings form fills in the official URL for a new provider
getAIModel({
provider: "openai",
apiKey: "user-key",
baseUrl: "https://api.openai.com/v1",
modelId: "gpt-5.5",
})
const provider = vi.mocked(createOpenAI).mock.results.at(-1)?.value
expect(provider.chat).toHaveBeenCalledWith("gpt-5.5")
})
it("sends a user's Ollama key to Ollama Cloud, not the server's Ollama", async () => {
process.env.OLLAMA_BASE_URL = "http://ollama.internal:11434/api"
const { createOllama } = await import("ollama-ai-provider-v2")
getAIModel({ provider: "ollama", apiKey: "user-key", modelId: "m" })
expect(createOllama).toHaveBeenLastCalledWith(
expect.objectContaining({ baseURL: "https://ollama.com/api" }),
)
// Without a key: the server's Ollama
getAIModel({ provider: "ollama", modelId: "m" })
expect(createOllama).toHaveBeenLastCalledWith(
expect.objectContaining({
baseURL: "http://ollama.internal:11434/api",
}),
)
})
it("sends the server's Ollama key where OLLAMA_BASE_URL says, or to local Ollama", async () => {
// The desktop app's "Ollama (Local)" preset puts its API Key field
// into OLLAMA_API_KEY; with no base URL that is the local Ollama
process.env.OLLAMA_API_KEY = "server-key"
const { createOllama } = await import("ollama-ai-provider-v2")
getAIModel({ provider: "ollama", modelId: "m" })
expect(vi.mocked(createOllama).mock.lastCall?.[0]).not.toHaveProperty(
"baseURL",
)
process.env.OLLAMA_BASE_URL = "https://ollama.com/api"
getAIModel({ provider: "ollama", modelId: "m" })
expect(createOllama).toHaveBeenLastCalledWith(
expect.objectContaining({ baseURL: "https://ollama.com/api" }),
)
})
it("uses a server model's own Ollama URL variable", async () => {
process.env.OLLAMA_BASE_URL = "http://other.internal:11434/api"
process.env.MY_OLLAMA_URL = "https://ollama.proxy.example/api"
process.env.MY_OLLAMA_KEY = "proxy-key"
try {
const { createOllama } = await import("ollama-ai-provider-v2")
getAIModel({
provider: "ollama",
modelId: "m",
apiKeyEnv: "MY_OLLAMA_KEY",
baseUrlEnv: "MY_OLLAMA_URL",
})
expect(createOllama).toHaveBeenLastCalledWith(
expect.objectContaining({
baseURL: "https://ollama.proxy.example/api",
headers: { Authorization: "Bearer proxy-key" },
}),
)
} finally {
delete process.env.MY_OLLAMA_URL
delete process.env.MY_OLLAMA_KEY
}
})
it("runs an Azure entry set up only in the admin panel", async () => {
// No AZURE_BASE_URL or AZURE_RESOURCE_NAME: the entry's own
// variables hold the key and the resource URL
process.env.ADMIN_AZURE_API_KEY = "panel-key"
process.env.ADMIN_AZURE_BASE_URL = "https://res.openai.azure.com/openai"
try {
const { createAzure } = await import("@ai-sdk/azure")
expect(() =>
getAIModel({
provider: "azure",
modelId: "gpt-4o",
apiKeyEnv: "ADMIN_AZURE_API_KEY",
baseUrlEnv: "ADMIN_AZURE_BASE_URL",
}),
).not.toThrow()
expect(createAzure).toHaveBeenLastCalledWith(
expect.objectContaining({
apiKey: "panel-key",
baseURL: "https://res.openai.azure.com/openai",
}),
)
// Without any URL it still says what is missing
delete process.env.ADMIN_AZURE_BASE_URL
expect(() =>
getAIModel({
provider: "azure",
modelId: "gpt-4o",
apiKeyEnv: "ADMIN_AZURE_API_KEY",
baseUrlEnv: "ADMIN_AZURE_BASE_URL",
}),
).toThrow(/AZURE_BASE_URL/)
} finally {
delete process.env.ADMIN_AZURE_API_KEY
delete process.env.ADMIN_AZURE_BASE_URL
}
})
it("needs a base URL with a user's Azure key", () => {
// The SDK would otherwise read the server's AZURE_RESOURCE_NAME
process.env.AZURE_RESOURCE_NAME = "operator-resource"
expect(() =>
getAIModel({ provider: "azure", apiKey: "k", modelId: "gpt-4o" }),
).toThrow(/base URL/)
})
it("needs a base URL for SGLang instead of using 127.0.0.1", () => {
expect(() =>
getAIModel({ provider: "sglang", apiKey: "k", modelId: "m" }),
).toThrow(/base URL/)
})
})
+62
View File
@@ -0,0 +1,62 @@
// @vitest-environment node
import { generateText } from "ai"
import { afterEach, describe, expect, it, vi } from "vitest"
import { getAIModel } from "@/lib/ai-providers"
const ENV = ["GOOGLE_GENERATIVE_AI_API_KEY", "GOOGLE_TOP_K", "GOOGLE_TOP_P"]
const saved = Object.fromEntries(ENV.map((k) => [k, process.env[k]]))
afterEach(() => {
for (const k of ENV) {
if (saved[k] === undefined) delete process.env[k]
else process.env[k] = saved[k]
}
vi.unstubAllGlobals()
})
/** Send one request through the real Google provider and return its body */
async function googleRequestBody() {
let body: any
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init: RequestInit) => {
body = JSON.parse(init.body as string)
return new Response(
JSON.stringify({
candidates: [
{
content: { role: "model", parts: [{ text: "ok" }] },
finishReason: "STOP",
},
],
}),
{ headers: { "content-type": "application/json" } },
)
}),
)
const { model } = getAIModel({
provider: "google",
modelId: "gemini-2.5-flash",
})
await generateText({ model, prompt: "hi", maxRetries: 0 })
return body
}
describe("Google sampling settings", () => {
it("sends GOOGLE_TOP_K and GOOGLE_TOP_P in the generation config", async () => {
process.env.GOOGLE_GENERATIVE_AI_API_KEY = "test-key"
process.env.GOOGLE_TOP_K = "40"
process.env.GOOGLE_TOP_P = "0.9"
const body = await googleRequestBody()
expect(body.generationConfig).toMatchObject({ topK: 40, topP: 0.9 })
})
it("sends neither when they are not set", async () => {
process.env.GOOGLE_GENERATIVE_AI_API_KEY = "test-key"
delete process.env.GOOGLE_TOP_K
delete process.env.GOOGLE_TOP_P
const body = await googleRequestBody()
expect(body.generationConfig?.topK).toBeUndefined()
expect(body.generationConfig?.topP).toBeUndefined()
})
})
+27
View File
@@ -0,0 +1,27 @@
// @vitest-environment node
import { describe, expect, it } from "vitest"
import { getAIModel } from "@/lib/ai-providers"
const summaryFor = (modelId: string) =>
getAIModel({ provider: "openai", apiKey: "test-key", modelId })
.providerOptions?.openai?.reasoningSummary
describe("OpenAI reasoning summary", () => {
it("is on for the o-series and gpt-5 or later", () => {
for (const id of [
"o3",
"o4-mini",
"gpt-5.5",
"gpt-6-luna",
"gpt-6.1-sol",
]) {
expect(summaryFor(id), id).toBe("auto")
}
})
it("is off for older chat models", () => {
for (const id of ["gpt-4.1", "gpt-4o", "gpt-4o-mini"]) {
expect(summaryFor(id), id).toBeUndefined()
}
})
})
+37 -31
View File
@@ -5,31 +5,6 @@ import {
resolveBaseURL,
supportsPromptCaching,
} from "@/lib/ai-providers"
import { extractAihubmixModelIds } from "@/lib/aihubmix-models"
describe("extractAihubmixModelIds", () => {
it("extracts unique chat model IDs from the AIHubMix model list payload", () => {
const models = extractAihubmixModelIds({
data: [
{ model_id: "claude-sonnet-4-5-20250929", types: "llm" },
{ model_id: "gpt-5.1", types: "llm" },
{ model_id: "gpt-5.1", types: "llm" },
{ model_id: "gpt-image-2", types: "image_generation,llm" },
{ model_id: "cohere-rerank-v4.0", types: "rerank" },
{ model_id: "", types: "llm" },
{ types: "llm" },
],
})
expect(models).toEqual(["claude-sonnet-4-5-20250929", "gpt-5.1"])
})
it("returns an empty list for malformed payloads", () => {
expect(extractAihubmixModelIds({ data: null })).toEqual([])
expect(extractAihubmixModelIds({})).toEqual([])
expect(extractAihubmixModelIds(null)).toEqual([])
})
})
describe("resolveBaseURL", () => {
const SERVER_BASE_URL = "https://server-proxy.example.com"
@@ -272,8 +247,15 @@ describe("AIHubMix provider", () => {
})
})
vi.mock("@ai-sdk/openai-compatible", () => {
const mockModel = { specificationVersion: "v3", modelId: "test-model" }
const mockProviderFn = vi.fn(() => mockModel)
const mockCreate = vi.fn(() => mockProviderFn)
return { createOpenAICompatible: mockCreate }
})
describe("Atlas Cloud provider", () => {
let createOpenAIMock: ReturnType<typeof vi.fn>
let createCompatibleMock: ReturnType<typeof vi.fn>
const savedEnv: Record<string, string | undefined> = {}
beforeEach(async () => {
@@ -281,9 +263,11 @@ describe("Atlas Cloud provider", () => {
savedEnv.ATLASCLOUD_BASE_URL = process.env.ATLASCLOUD_BASE_URL
delete process.env.ATLASCLOUD_BASE_URL
const mod = await import("@ai-sdk/openai")
createOpenAIMock = mod.createOpenAI as ReturnType<typeof vi.fn>
createOpenAIMock.mockClear()
const mod = await import("@ai-sdk/openai-compatible")
createCompatibleMock = mod.createOpenAICompatible as ReturnType<
typeof vi.fn
>
createCompatibleMock.mockClear()
})
afterEach(() => {
@@ -299,9 +283,12 @@ describe("Atlas Cloud provider", () => {
modelId: "qwen/qwen3.5-flash",
})
expect(createOpenAIMock).toHaveBeenCalledWith({
// An OpenAI-compatible API; includeUsage keeps quota tracking working
expect(createCompatibleMock).toHaveBeenCalledWith({
name: "atlascloud",
apiKey: "server-atlas-key",
baseURL: "https://api.atlascloud.ai/v1",
includeUsage: true,
})
})
@@ -313,11 +300,28 @@ describe("Atlas Cloud provider", () => {
modelId: "deepseek-ai/deepseek-v4-pro",
})
expect(createOpenAIMock).toHaveBeenCalledWith({
expect(createCompatibleMock).toHaveBeenCalledWith({
name: "atlascloud",
apiKey: "client-atlas-key",
baseURL: "https://proxy.example.com/v1",
includeUsage: true,
})
})
it("drops an endpoint path pasted along with the base URL", () => {
getAIModel({
provider: "atlascloud",
apiKey: "client-atlas-key",
baseUrl: "https://proxy.example.com/v1/chat/completions/",
modelId: "deepseek-ai/deepseek-v4-pro",
})
expect(createCompatibleMock).toHaveBeenCalledWith(
expect.objectContaining({
baseURL: "https://proxy.example.com/v1",
}),
)
})
})
describe("Kimi provider uses createDeepSeek for reasoning_content support", () => {
@@ -458,6 +462,8 @@ describe("Ollama API key security", () => {
expect(createOllamaMock).toHaveBeenCalledTimes(1)
const callArgs = createOllamaMock.mock.calls[0][0]
// The SDK's local default: the desktop app's "Ollama (Local)"
// preset puts its API Key field into OLLAMA_API_KEY
expect(callArgs).not.toHaveProperty("baseURL")
expect(callArgs).toEqual(
expect.objectContaining({
+156
View File
@@ -0,0 +1,156 @@
// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { POST as parseUrl } from "@/app/api/parse-url/route"
import { POST as validateDiagram } from "@/app/api/validate-diagram/route"
import { POST as validateModel } from "@/app/api/validate-model/route"
import { POST as verifyAccessCode } from "@/app/api/verify-access-code/route"
import { checkAccessCode } from "@/lib/access-code"
// Treat every URL as public so no test hits DNS
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
isPrivateUrl: async () => false,
}))
function post(path: string, body: unknown, accessCode?: string): Request {
return new Request(`http://localhost${path}`, {
method: "POST",
headers: {
"Content-Type": "application/json",
...(accessCode ? { "x-access-code": accessCode } : {}),
},
body: JSON.stringify(body),
})
}
beforeEach(() => {
process.env.ACCESS_CODE_LIST = "secret, other"
})
afterEach(() => {
delete process.env.ACCESS_CODE_LIST
delete process.env.ALLOW_PRIVATE_URLS
vi.unstubAllGlobals()
})
describe("checkAccessCode", () => {
it("passes when no access codes are configured", () => {
delete process.env.ACCESS_CODE_LIST
expect(checkAccessCode(post("/x", {}))).toBeNull()
})
it("rejects a missing or wrong code and accepts a listed one", () => {
expect(checkAccessCode(post("/x", {}))?.status).toBe(401)
expect(checkAccessCode(post("/x", {}, "nope"))?.status).toBe(401)
expect(checkAccessCode(post("/x", {}, "other"))).toBeNull()
})
})
describe("routes that spend server resources require the access code", () => {
it("parse-url", async () => {
const res = await parseUrl(
post("/api/parse-url", { url: "https://example.com" }),
)
expect(res.status).toBe(401)
})
it("validate-diagram", async () => {
const res = await validateDiagram(
post("/api/validate-diagram", {
imageData: "data:image/png;base64,AAAA",
}),
)
expect(res.status).toBe(401)
})
it("validate-model", async () => {
const res = await validateModel(
post("/api/validate-model", {
provider: "openai",
apiKey: "sk",
modelId: "m",
}),
)
expect(res.status).toBe(401)
})
it("verify-access-code", async () => {
const bad = await verifyAccessCode(post("/api/verify-access-code", {}))
expect(bad.status).toBe(401)
expect((await bad.json()).valid).toBe(false)
const good = await verifyAccessCode(
post("/api/verify-access-code", {}, "secret"),
)
expect((await good.json()).valid).toBe(true)
})
})
describe("size limits", () => {
it("parse-url stops reading a body over the download limit", async () => {
const chunk = new Uint8Array(1024 * 1024)
let sent = 0
const body = new ReadableStream<Uint8Array>({
pull(controller) {
sent++
controller.enqueue(chunk)
},
})
vi.stubGlobal(
"fetch",
vi.fn(
async () =>
new Response(body, {
headers: { "content-type": "text/html" },
}),
),
)
const res = await parseUrl(
post("/api/parse-url", { url: "https://example.com" }, "secret"),
)
expect(res.status).toBe(413)
expect(sent).toBeLessThan(10)
})
it("validate-diagram rejects oversized image data", async () => {
const imageData = `data:image/png;base64,${"A".repeat(6 * 1024 * 1024)}`
const res = await validateDiagram(
post("/api/validate-diagram", { imageData }, "secret"),
)
expect(res.status).toBe(413)
})
})
describe("validate-model redirects", () => {
it("refuses redirects when private URLs are blocked", async () => {
process.env.ALLOW_PRIVATE_URLS = "false"
vi.stubGlobal(
"fetch",
vi.fn(
async () =>
new Response(null, {
status: 302,
headers: { location: "http://169.254.169.254/" },
}),
),
)
const res = await validateModel(
post(
"/api/validate-model",
{
provider: "openai",
apiKey: "sk",
modelId: "m",
baseUrl: "https://attacker.example/v1",
},
"secret",
),
)
const data = await res.json()
expect(data.valid).toBe(false)
expect(data.error).toMatch(/Redirects are not allowed/)
expect(fetch).toHaveBeenCalledTimes(1)
})
})
+104
View File
@@ -0,0 +1,104 @@
// @vitest-environment node
import { beforeEach, describe, expect, it, vi } from "vitest"
vi.mock("electron", () => ({
app: {
isPackaged: true,
getName: () => "app",
getVersion: () => "1",
getLocale: () => "en",
},
BrowserWindow: { getFocusedWindow: () => null },
dialog: {},
Menu: { buildFromTemplate: () => ({}), setApplicationMenu: () => {} },
shell: {},
}))
// The saved current preset, and restarts that wait until the test ends them
const state = vi.hoisted(() => ({
current: "A" as string | null,
restarts: [] as Array<{ resolve: () => void; reject: (e: Error) => void }>,
}))
vi.mock("@/electron/main/config-manager", () => ({
applyPresetToEnv: (id: string) => {
if (id === "missing") return null
state.current = id
return { AI_PROVIDER: id }
},
getAllPresets: () => [],
getCurrentPresetId: () => state.current,
setCurrentPreset: (id: string | null) => {
state.current = id
return true
},
}))
vi.mock("@/electron/main/next-server", () => ({
restartNextServer: () =>
new Promise<void>((resolve, reject) =>
state.restarts.push({ resolve, reject }),
),
}))
vi.mock("@/electron/main/menu-i18n", () => ({
getMenuTranslations: () => new Proxy({}, { get: () => "x" }),
getPreferredLocale: () => "en",
}))
vi.mock("@/electron/main/settings-window", () => ({
showSettingsWindow: () => {},
}))
import { switchPreset } from "@/electron/main/app-menu"
beforeEach(() => {
state.current = "A"
state.restarts = []
})
describe("switchPreset", () => {
it("keeps a preset chosen while a failed switch was restarting", async () => {
const toB = switchPreset("B").catch(() => {})
const toC = switchPreset("C")
// B's restart fails after the user already picked C
state.restarts[0].reject(new Error("timed out"))
await new Promise((r) => setTimeout(r, 0))
for (const r of state.restarts.slice(1)) r.resolve()
await toB
await toC
expect(state.current).toBe("C")
})
it("keeps a newer choice of the same preset", async () => {
// A, then B, C, and B again while the first restart is pending
const first = switchPreset("B").catch(() => {})
const second = switchPreset("C").catch(() => {})
const third = switchPreset("B")
// The first restart fails: the current preset is B again, but it is
// the third switch's, which must not be undone
state.restarts[0].reject(new Error("timed out"))
await new Promise((r) => setTimeout(r, 0))
for (const r of state.restarts.slice(1)) r.resolve()
await first
await second
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")
})
})
+42
View File
@@ -0,0 +1,42 @@
import { describe, expect, it } from "vitest"
import { chatRequestUrl, normalizeBaseUrl } from "@/lib/types/model-config"
describe("normalizeBaseUrl", () => {
it("drops spaces, trailing slashes and a pasted endpoint path", () => {
expect(normalizeBaseUrl(" https://api.x.com/v1/ ")).toBe(
"https://api.x.com/v1",
)
expect(normalizeBaseUrl("https://api.x.com/v1/chat/completions")).toBe(
"https://api.x.com/v1",
)
expect(
normalizeBaseUrl("https://api.x.com/anthropic/v1/messages/"),
).toBe("https://api.x.com/anthropic/v1")
})
it("keeps provider paths such as /api/paas/v4 and /api/v3", () => {
expect(normalizeBaseUrl("https://open.bigmodel.cn/api/paas/v4")).toBe(
"https://open.bigmodel.cn/api/paas/v4",
)
expect(
normalizeBaseUrl("https://ark.cn-beijing.volces.com/api/v3"),
).toBe("https://ark.cn-beijing.volces.com/api/v3")
})
})
describe("chatRequestUrl", () => {
it("shows the endpoint the SDK will call", () => {
expect(chatRequestUrl("glm", "https://api.x.com/v1/")).toBe(
"https://api.x.com/v1/chat/completions",
)
expect(
chatRequestUrl("anthropic", "https://proxy.example.com/v1"),
).toBe("https://proxy.example.com/v1/messages")
})
it("stays out of the way for SDKs that build their own paths", () => {
expect(chatRequestUrl("google", "https://x.example.com")).toBeNull()
expect(chatRequestUrl("minimax", "https://x.example.com")).toBeNull()
expect(chatRequestUrl("glm", " ")).toBeNull()
})
})
+25 -2
View File
@@ -14,12 +14,35 @@ describe("findCachedResponse", () => {
expect(result?.xml).toContain("Transformer Architecture")
})
it("returns cached response for exact match with image", () => {
const result = findCachedResponse("Replicate this in aws style", true)
it("returns cached response for exact match with the example file", () => {
const result = findCachedResponse(
"Replicate this in aws style",
true,
"architecture.png",
)
expect(result).toBeDefined()
expect(result?.xml).toContain("AWS")
})
it("returns undefined when the user attached their own file", () => {
expect(
findCachedResponse("Replicate this flowchart.", true, "mine.png"),
).toBeUndefined()
expect(
findCachedResponse(
"Summarize this paper as a diagram",
true,
"thesis.pdf",
),
).toBeUndefined()
})
it("returns undefined for file examples when the file name is unknown", () => {
// The server only knows whether a file is attached, not which one
const result = findCachedResponse("Replicate this in aws style", true)
expect(result).toBeUndefined()
})
it("returns undefined for non-matching prompt", () => {
const result = findCachedResponse(
"random prompt that doesn't exist",
+146 -17
View File
@@ -1,7 +1,11 @@
// @vitest-environment node
import { convertToModelMessages } from "ai"
import { jsonrepair } from "jsonrepair"
import { describe, expect, it } from "vitest"
import {
isMinimalDiagram,
dropInvalidToolCalls,
fixToolInputJson,
replaceHistoricalToolInputs,
validateFileParts,
} from "@/lib/chat-helpers"
@@ -65,23 +69,28 @@ describe("validateFileParts", () => {
expect(result.valid).toBe(false)
expect(result.error).toContain("exceeds")
})
})
describe("isMinimalDiagram", () => {
it("returns true for empty diagram", () => {
const xml = '<mxCell id="0"/><mxCell id="1" parent="0"/>'
expect(isMinimalDiagram(xml)).toBe(true)
it("rejects file URLs the server would have to download", () => {
for (const url of [
"http://10.0.0.5/secret.png",
"https://example.com/a.png",
undefined,
]) {
const messages = [{ role: "user", parts: [{ type: "file", url }] }]
expect(validateFileParts(messages).valid).toBe(false)
}
})
it("returns false for diagram with content", () => {
const xml =
'<mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="Hello"/>'
expect(isMinimalDiagram(xml)).toBe(false)
})
it("handles whitespace correctly", () => {
const xml = ' <mxCell id="0"/> <mxCell id="1" parent="0"/> '
expect(isMinimalDiagram(xml)).toBe(true)
it("checks files in earlier messages too", () => {
const messages = [
{
role: "user",
parts: [{ type: "file", url: "http://169.254.169.254/x" }],
},
{ role: "assistant", parts: [{ type: "text", text: "ok" }] },
{ role: "user", parts: [{ type: "text", text: "hello" }] },
]
expect(validateFileParts(messages).valid).toBe(false)
})
})
@@ -124,7 +133,7 @@ describe("replaceHistoricalToolInputs", () => {
)
})
it("removes tool calls with invalid inputs", () => {
it("leaves tool calls with invalid inputs for dropInvalidToolCalls", () => {
const messages = [
{
role: "assistant",
@@ -143,7 +152,7 @@ describe("replaceHistoricalToolInputs", () => {
},
]
const result = replaceHistoricalToolInputs(messages)
expect(result[0].content).toHaveLength(0)
expect(result[0].content).toEqual(messages[0].content)
})
it("preserves non-assistant messages", () => {
@@ -169,3 +178,123 @@ describe("replaceHistoricalToolInputs", () => {
expect(result[0].content[0].input).toEqual({ foo: "bar" })
})
})
describe("dropInvalidToolCalls", () => {
it("drops an invalid tool-call together with its tool-result", () => {
const messages = [
{ role: "user", content: [{ type: "text", text: "draw" }] },
{
role: "assistant",
content: [
{
type: "tool-call",
toolCallId: "call-1",
toolName: "display_diagram",
input: undefined,
},
],
},
{
role: "tool",
content: [
{
type: "tool-result",
toolCallId: "call-1",
toolName: "display_diagram",
output: { type: "error-text", value: "Stopped" },
},
],
},
{ role: "user", content: [{ type: "text", text: "again" }] },
]
const result = dropInvalidToolCalls(messages)
expect(result.map((m) => m.role)).toEqual(["user", "user"])
})
it("keeps valid calls and results in the same messages", () => {
const messages = [
{
role: "assistant",
content: [
{ type: "text", text: "Here you go" },
{
type: "tool-call",
toolCallId: "bad",
toolName: "edit_diagram",
input: "{broken",
},
{
type: "tool-call",
toolCallId: "good",
toolName: "display_diagram",
input: { xml: "<mxCell/>" },
},
],
},
{
role: "tool",
content: [
{ type: "tool-result", toolCallId: "bad", output: {} },
{ type: "tool-result", toolCallId: "good", output: {} },
],
},
]
const result = dropInvalidToolCalls(messages)
expect(result[0].content.map((p: any) => p.toolCallId)).toEqual([
undefined,
"good",
])
expect(result[1].content.map((p: any) => p.toolCallId)).toEqual([
"good",
])
})
it("cleans up a tool call the user stopped before its input arrived", async () => {
// handleStop turns a still-streaming call into output-error with no input
const modelMessages = await convertToModelMessages([
{ role: "user", parts: [{ type: "text", text: "draw" }] },
{
role: "assistant",
parts: [
{
type: "tool-display_diagram",
toolCallId: "call-1",
state: "output-error",
input: undefined,
errorText: "Stopped by user",
} as any,
],
},
{ role: "user", parts: [{ type: "text", text: "again" }] },
])
const result = dropInvalidToolCalls(modelMessages)
expect(result.map((m) => m.role)).toEqual(["user", "user"])
})
it("leaves messages with string content alone", () => {
const messages = [{ role: "system", content: "You are..." }]
expect(dropInvalidToolCalls(messages)).toEqual(messages)
})
})
describe("fixToolInputJson", () => {
it("fixes an attribute whose closing quote alone is escaped", () => {
const input =
'{"xml": "<mxCell id=\\"2\\" vertex=\\"1\\"><mxGeometry x=\\"10\\" y="-20\\" as=\\"geometry\\"/></mxCell>"}'
const parsed = JSON.parse(jsonrepair(fixToolInputJson(input)))
expect(parsed.xml).toContain('y="-20"')
expect(parsed.xml).toContain('id="2"')
})
it("fixes = used instead of : after a JSON key", () => {
const input = '{"xml"= "<mxCell id=\\"2\\"/>"}'
const parsed = JSON.parse(jsonrepair(fixToolInputJson(input)))
expect(parsed.xml).toBe('<mxCell id="2"/>')
})
it("leaves well-formed input unchanged", () => {
const input =
'{"operations": [{"operation": "add", "cell_id": "a", "new_xml": "<mxCell id=\\"a\\" value=\\"x=1\\"/>"}]}'
expect(fixToolInputJson(input)).toBe(input)
})
})
@@ -0,0 +1,68 @@
import { render } from "@testing-library/react"
import { describe, expect, it, vi } from "vitest"
import en from "@/lib/i18n/dictionaries/en.json"
const page = (cells: string) =>
`<mxfile><diagram id="p" name="Page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel></diagram></mxfile>`
const box = (id: string) =>
`<mxCell id="${id}" value="${id}" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
// The first edit's result is loaded (the ref has it); the chartXML state
// has not caught up yet
const BEFORE_FIRST_EDIT = page(box("a"))
const AFTER_FIRST_EDIT = page(box("a") + box("b"))
vi.mock("@/contexts/diagram-context", () => ({
useDiagram: () => ({
chartXML: BEFORE_FIRST_EDIT,
chartXMLRef: { current: AFTER_FIRST_EDIT },
loadDiagram: vi.fn(() => null),
}),
}))
vi.mock("@/hooks/use-dictionary", () => ({ useDictionary: () => en }))
import { ChatMessageDisplay } from "@/components/chat-message-display"
// jsdom has no layout
Element.prototype.scrollIntoView = () => {}
describe("the streaming preview of a second edit", () => {
it("starts from the first edit's result", () => {
const editDiagramOriginalXmlRef = { current: new Map<string, string>() }
const messages = [
{
id: "m1",
role: "assistant",
parts: [
{
type: "tool-edit_diagram",
toolCallId: "edit-2",
state: "input-streaming",
input: {
operations: [
{
operation: "add",
cell_id: "c",
new_xml: box("c"),
},
],
},
},
],
},
] as any
render(
<ChatMessageDisplay
messages={messages}
setInput={() => {}}
setFiles={() => {}}
processedToolCallsRef={{ current: new Set() }}
editDiagramOriginalXmlRef={editDiagramOriginalXmlRef}
status="streaming"
/>,
)
expect(editDiagramOriginalXmlRef.current.get("edit-2")).toBe(
AFTER_FIRST_EDIT,
)
})
})
+139
View File
@@ -0,0 +1,139 @@
// @vitest-environment node
import { afterEach, describe, expect, it, vi } from "vitest"
const quota = vi.hoisted(() => ({ recorded: [] as number[] }))
vi.mock("@/lib/dynamo-quota-manager", () => ({
isQuotaEnabled: () => true,
checkAndIncrementRequest: async () => ({ allowed: true }),
recordTokenUsage: async (_ip: string, tokens: number) => {
quota.recorded.push(tokens)
},
}))
// No DNS in tests: only loopback addresses are private
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
isPrivateUrl: async (url: string) =>
/^https?:\/\/(127\.0\.0\.1|localhost)\b/.test(url),
}))
import { POST as chat } from "@/app/api/chat/route"
afterEach(() => {
quota.recorded = []
vi.unstubAllGlobals()
})
const sse = (chunks: object[], end = true) =>
chunks.map((c) => `data: ${JSON.stringify(c)}\n\n`).join("") +
(end ? "data: [DONE]\n\n" : "")
describe("a request stopped after a finished step", () => {
it("counts that step's tokens", async () => {
// Step 1 asks for a shape library (run on the server) and reports
// its usage; step 2 never ends, and the user stops
let call = 0
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init?: RequestInit) => {
call++
if (call === 1) {
return new Response(
sse([
{
id: "c1",
choices: [
{
index: 0,
delta: {
role: "assistant",
tool_calls: [
{
index: 0,
id: "call_1",
type: "function",
function: {
name: "get_shape_library",
arguments:
'{"library":"aws4"}',
},
},
],
},
finish_reason: null,
},
],
},
{
id: "c1",
choices: [
{
index: 0,
delta: {},
finish_reason: "tool_calls",
},
],
usage: {
prompt_tokens: 1200,
completion_tokens: 30,
},
},
]),
{ headers: { "content-type": "text/event-stream" } },
)
}
// Never ends, until the request is aborted (as fetch does)
const body = new ReadableStream({
start(controller) {
init?.signal?.addEventListener("abort", () =>
controller.error(
new DOMException("aborted", "AbortError"),
),
)
},
})
return new Response(body, {
headers: { "content-type": "text/event-stream" },
})
}),
)
const stop = new AbortController()
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
signal: stop.signal,
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
// The server's own network: counted
"x-ai-provider": "glm",
"x-ai-base-url": "http://127.0.0.1:9000/v1",
"x-ai-api-key": "dummy",
"x-ai-model": "glm-5",
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw AWS" }],
},
],
xml: "",
}),
}),
)
const reader = res.body?.getReader()
// Read until the second step has started
await vi.waitFor(() => expect(call).toBe(2), { timeout: 3000 })
stop.abort()
// The answer stream ends; the SDK handles the stop as it is read
while (
reader &&
!(await reader.read().catch(() => ({ done: true }))).done
) {
// drain
}
await vi.waitFor(() => expect(quota.recorded).toEqual([1230]))
})
})
+139
View File
@@ -0,0 +1,139 @@
// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
vi.mock("@/lib/dynamo-quota-manager", () => ({
isQuotaEnabled: () => false,
checkAndIncrementRequest: async () => ({ allowed: true }),
recordTokenUsage: async () => {},
}))
import { POST as chat } from "@/app/api/chat/route"
const ENV = ["AI_MODELS_CONFIG", "AI_PROVIDER", "AI_MODEL"]
const saved: Record<string, string | undefined> = {}
const calls: Array<{ url: string; headers: Headers }> = []
beforeEach(() => {
for (const k of ENV) saved[k] = process.env[k]
delete process.env.AI_PROVIDER
delete process.env.AI_MODEL
calls.length = 0
vi.stubGlobal(
"fetch",
vi.fn(async (url: string, init: RequestInit) => {
calls.push({ url: String(url), headers: new Headers(init.headers) })
throw new Error("no network in tests")
}),
)
})
afterEach(() => {
for (const k of ENV) {
if (saved[k] === undefined) delete process.env[k]
else process.env[k] = saved[k]
}
vi.unstubAllGlobals()
})
describe("EdgeOne as a server model", () => {
it("calls the site's Edge AI function with the cookies", async () => {
// An admin panel or ai-models.json provider: the client sends the
// provider name's slug, not "edgeone"
process.env.AI_MODELS_CONFIG = JSON.stringify({
providers: [
{
name: "Edge Pages",
provider: "edgeone",
models: ["@tx/deepseek-ai/deepseek-v3-0324"],
},
],
})
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: {
"Content-Type": "application/json",
"x-ai-provider": "edge-pages",
"x-selected-model-id":
"server:edge-pages:@tx/deepseek-ai/deepseek-v3-0324",
cookie: "eo_token=t; eo_time=1",
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw two boxes" }],
},
],
xml: "",
}),
}),
)
await res.text()
expect(calls[0]?.url).toBe(
"http://localhost/api/edgeai/chat/completions",
)
expect(calls[0]?.headers.get("cookie")).toBe("eo_token=t; eo_time=1")
})
})
const send = (headers: Record<string, string>) =>
chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: { "Content-Type": "application/json", ...headers },
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw two boxes" }],
},
],
xml: "",
}),
}),
).then((r) => r.text())
describe("EdgeOne endpoints", () => {
it("works when the deployment names EdgeOne only in AI_PROVIDER", async () => {
process.env.AI_PROVIDER = "edgeone"
process.env.AI_MODEL = "@tx/deepseek-ai/deepseek-v3-0324"
await send({})
expect(calls[0]?.url).toBe(
"http://localhost/api/edgeai/chat/completions",
)
})
it("always calls the site's own function, whatever URL the request names", async () => {
// Another host would get the user's EdgeOne cookies
await send({
"x-ai-provider": "edgeone",
"x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
"x-ai-base-url": "https://elsewhere.example/api/edgeai",
cookie: "eo_token=t",
})
expect(calls[0]?.url).toBe(
"http://localhost/api/edgeai/chat/completions",
)
})
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 {
await send({
"x-ai-provider": "edgeone",
"x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
})
expect(calls[0]?.url).toBe(
"http://localhost/api/edgeai/chat/completions",
)
} finally {
if (savedPath === undefined)
delete process.env.NEXT_PUBLIC_BASE_PATH
else process.env.NEXT_PUBLIC_BASE_PATH = savedPath
}
})
})
+248
View File
@@ -0,0 +1,248 @@
// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
// No DNS in tests: only loopback addresses are private
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
isPrivateUrl: async (url: string) =>
/^https?:\/\/(127\.0\.0\.1|localhost)\b/.test(url),
}))
import { POST as chat } from "@/app/api/chat/route"
const ENV = [
"AI_PROVIDER",
"AI_MODEL",
"OPENAI_API_KEY",
"OLLAMA_BASE_URL",
"OLLAMA_API_KEY",
"NEXT_AI_DRAWIO_DESKTOP",
]
const saved: Record<string, string | undefined> = {}
beforeEach(() => {
for (const k of ENV) saved[k] = process.env[k]
for (const k of ENV) delete process.env[k]
})
afterEach(() => {
for (const k of ENV) {
if (saved[k] === undefined) delete process.env[k]
else process.env[k] = saved[k]
}
vi.unstubAllGlobals()
})
/** Every provider request answers with this status and text */
const providerAnswers = (status: number, body: string) =>
vi.stubGlobal(
"fetch",
vi.fn(
async () =>
new Response(body, {
status,
headers: { "Content-Type": "application/json" },
}),
),
)
/** The error text the chat panel gets from the stream */
async function streamedError(headers: Record<string, string>) {
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: { "Content-Type": "application/json", ...headers },
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw two boxes" }],
},
],
xml: "",
}),
}),
)
const text = await res.text()
const line = text.split("\n").find((l) => l.includes('"type":"error"'))
return line ? JSON.parse(JSON.parse(line.slice(6)).errorText).message : ""
}
describe("provider error texts in the stream", () => {
it("shows the user's own local Ollama error in the desktop app", async () => {
process.env.NEXT_AI_DRAWIO_DESKTOP = "1"
process.env.AI_PROVIDER = "ollama"
process.env.AI_MODEL = "llama3"
// Ollama is not running
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw Object.assign(new TypeError("fetch failed"), {
cause: new Error("connect ECONNREFUSED 127.0.0.1:11434"),
})
}),
)
expect(await streamedError({})).toMatch(
/127\.0\.0\.1:11434|fetch failed/,
)
// The SDK retries a refused connection twice, waiting between
}, 20_000)
it("shows the server's keyless Ollama error on the web too", async () => {
// No key, no money involved; round three hid this text
process.env.AI_PROVIDER = "ollama"
process.env.AI_MODEL = "llama3"
process.env.OLLAMA_BASE_URL = "http://ollama.internal:11434/api"
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw Object.assign(new TypeError("fetch failed"), {
cause: new Error("connect ECONNREFUSED 10.0.0.9:11434"),
})
}),
)
expect(await streamedError({})).not.toBe(
"The provider returned an error.",
)
}, 20_000)
it("shows EdgeOne's own daily quota explanation", async () => {
// The function answers 429, which the SDK retries with a wait;
// the status does not decide whether the text is shown
providerAnswers(
400,
JSON.stringify({
error: {
message:
"The daily public quota has been exhausted. After deployment, you can enjoy a personal daily exclusive quota.",
},
}),
)
expect(
await streamedError({
"x-ai-provider": "edgeone",
"x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
}),
).toMatch(/daily public quota/)
})
it("hides the provider's text on the server's own key", async () => {
process.env.AI_PROVIDER = "openai"
process.env.AI_MODEL = "gpt-5.5"
process.env.OPENAI_API_KEY = "server-key"
providerAnswers(
403,
JSON.stringify({
error: { message: "Organization org-operator is suspended" },
}),
)
const message = await streamedError({})
expect(message).not.toMatch(/org-operator/)
expect(message).toBe("The provider returned an error.")
})
})
describe("the output cap", () => {
it("holds for the server's own keyless endpoints too", async () => {
process.env.MAX_OUTPUT_TOKENS = "8000"
const sent: string[] = []
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init?: RequestInit) => {
sent.push(String(init?.body ?? ""))
return new Response("{}", { status: 400 })
}),
)
try {
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: {
"Content-Type": "application/json",
"x-ai-provider": "ollama",
"x-ai-base-url": "http://127.0.0.1:11434/api",
"x-ai-model": "llama3",
"x-max-output-tokens": "200000",
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw" }],
},
],
xml: "",
}),
}),
)
await res.text()
expect(sent[0]).toContain('"max_output_tokens":8000')
} finally {
delete process.env.MAX_OUTPUT_TOKENS
}
})
})
describe("a tool call that never got its result", () => {
it("is left out of the prompt instead of failing every later message", async () => {
// Stop while the screenshot check ran left display_diagram without
// a result, and the chat was saved like that
process.env.AI_PROVIDER = "openai"
process.env.AI_MODEL = "gpt-5.5"
process.env.OPENAI_API_KEY = "server-key"
const sent: string[] = []
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init?: RequestInit) => {
sent.push(String(init?.body ?? ""))
return new Response(
JSON.stringify({ error: { message: "x" } }),
{
status: 400,
headers: { "Content-Type": "application/json" },
},
)
}),
)
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw a box" }],
},
{
id: "a1",
role: "assistant",
parts: [
{
type: "tool-display_diagram",
toolCallId: "call-without-result",
state: "input-available",
input: { xml: "<mxCell/>" },
},
],
},
{
id: "u2",
role: "user",
parts: [{ type: "text", text: "Make it red" }],
},
],
xml: "",
}),
}),
)
await res.text()
// The request reached the model, without the unanswered call
expect(sent).toHaveLength(1)
expect(sent[0]).toContain("Make it red")
expect(sent[0]).not.toContain("call-without-result")
})
})
+242
View File
@@ -0,0 +1,242 @@
// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
// Quota on, and every check answers that the daily limit is used up
const quota = vi.hoisted(() => ({ checks: 0 }))
vi.mock("@/lib/dynamo-quota-manager", () => ({
isQuotaEnabled: () => true,
checkAndIncrementRequest: async () => {
quota.checks++
return {
allowed: false,
error: "Daily limit reached",
type: "request",
used: 10,
limit: 10,
}
},
recordTokenUsage: async () => {},
}))
// No DNS in tests: only loopback addresses are private
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
isPrivateUrl: async (url: string) =>
/^https?:\/\/(127\.0\.0\.1|localhost)\b/.test(url),
}))
import { POST as chat } from "@/app/api/chat/route"
const ENV = [
"AI_PROVIDER",
"AI_MODEL",
"OPENAI_API_KEY",
"OLLAMA_BASE_URL",
"OLLAMA_API_KEY",
"AI_GATEWAY_API_KEY",
"ALLOW_PRIVATE_URLS",
]
const saved: Record<string, string | undefined> = {}
beforeEach(() => {
for (const k of ENV) saved[k] = process.env[k]
process.env.AI_PROVIDER = "openai"
process.env.AI_MODEL = "gpt-5.5"
process.env.OPENAI_API_KEY = "server-key"
quota.checks = 0
// No request may reach a provider
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw new Error("no network in tests")
}),
)
})
afterEach(() => {
for (const k of ENV) {
if (saved[k] === undefined) delete process.env[k]
else process.env[k] = saved[k]
}
vi.unstubAllGlobals()
})
const send = (headers: Record<string, string>) =>
chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
...headers,
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw two boxes" }],
},
],
xml: "",
}),
}),
)
describe("chat quota", () => {
it("counts a request whose key header the provider never reads", async () => {
// OpenAI ignores the AWS key, so this runs on the server's key
const res = await send({
"x-ai-provider": "openai",
"x-aws-access-key-id": "x",
})
expect(res.status).toBe(429)
expect(quota.checks).toBe(1)
})
it("does not count a request on the user's own key", async () => {
const res = await send({
"x-ai-provider": "openai",
"x-ai-api-key": "user-key",
"x-ai-model": "gpt-5.5",
})
expect(res.status).not.toBe(429)
expect(quota.checks).toBe(0)
})
it("counts the server's keyless Ollama and EdgeOne", async () => {
process.env.AI_PROVIDER = "ollama"
process.env.AI_MODEL = "llama3.2"
process.env.OLLAMA_BASE_URL = "http://ollama.internal:11434/api"
expect((await send({})).status).toBe(429)
expect(
(
await send({
"x-ai-provider": "edgeone",
"x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
})
).status,
).toBe(429)
expect(quota.checks).toBe(2)
})
it("counts EdgeOne with a key header it never reads", async () => {
const res = await send({
"x-ai-provider": "edgeone",
"x-ai-api-key": "ignored",
"x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
})
expect(res.status).toBe(429)
expect(quota.checks).toBe(1)
})
it("counts Ollama at a private address, the server's network", async () => {
const res = await send({
"x-ai-provider": "ollama",
"x-ai-base-url": "http://127.0.0.1:11434/api",
"x-ai-model": "llama3.2",
})
expect(res.status).toBe(429)
expect(quota.checks).toBe(1)
})
it("counts the server's network whatever key header comes along", async () => {
// A keyless Ollama or a local SGLang ignores a dummy key
for (const headers of [
{
"x-ai-provider": "ollama",
"x-ai-base-url": "http://127.0.0.1:11434/api",
"x-ai-api-key": "dummy",
"x-ai-model": "llama3.2",
},
{
"x-ai-provider": "openai",
"x-ai-base-url": "http://127.0.0.1:30000/v1",
"x-ai-api-key": "dummy",
"x-ai-model": "m",
},
]) {
expect((await send(headers)).status).toBe(429)
}
expect(quota.checks).toBe(2)
})
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 (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",
})
expect(res.status).not.toBe(429)
expect(quota.checks).toBe(0)
})
it("counts EdgeOne even with a base URL header", async () => {
const res = await send({
"x-ai-provider": "edgeone",
"x-ai-model": "@tx/deepseek-ai/deepseek-v3-0324",
"x-ai-base-url": "https://this-site.example/api/edgeai",
})
expect(res.status).toBe(429)
expect(quota.checks).toBe(1)
})
it("never counts in the desktop app, where every endpoint is the user's", async () => {
process.env.NEXT_AI_DRAWIO_DESKTOP = "1"
try {
const res = await send({
"x-ai-provider": "ollama",
"x-ai-base-url": "http://127.0.0.1:11434/api",
"x-ai-model": "llama3.2",
})
expect(res.status).not.toBe(429)
expect(quota.checks).toBe(0)
} finally {
delete process.env.NEXT_AI_DRAWIO_DESKTOP
}
})
it("does not count Ollama on the user's own server", async () => {
const res = await send({
"x-ai-provider": "ollama",
"x-ai-base-url": "https://ollama.example.com/api",
"x-ai-model": "llama3.2",
})
expect(res.status).not.toBe(429)
expect(quota.checks).toBe(0)
})
})
describe("request checks", () => {
it("refuses an AWS region that is not a region name", async () => {
process.env.AI_PROVIDER = "bedrock"
process.env.AI_MODEL = "amazon.nova-lite-v1:0"
const res = await send({
"x-aws-region": "us-east-1.attacker.example/",
})
expect(res.status).toBe(400)
expect(await res.text()).toMatch(/Invalid AWS region/)
expect(fetch).not.toHaveBeenCalled()
})
})
describe("server model allowlist", () => {
it("runs AI_MODEL only on the server's AI_PROVIDER", async () => {
// Another provider's server key must not run it
process.env.AI_GATEWAY_API_KEY = "server-gateway-key"
const res = await send({
"x-ai-provider": "gateway",
"x-ai-model": "gpt-5.5",
})
expect(res.status).toBe(400)
expect(await res.text()).toMatch(/not available on this server/)
expect(quota.checks).toBe(0)
})
})
+109
View File
@@ -0,0 +1,109 @@
// @vitest-environment node
import {
existsSync,
mkdtempSync,
readdirSync,
readFileSync,
rmSync,
writeFileSync,
} from "node:fs"
import { tmpdir } from "node:os"
import { join } from "node:path"
import { beforeEach, describe, expect, it, vi } from "vitest"
const userData = vi.hoisted(() => ({ dir: "" }))
vi.mock("electron", () => ({
app: { getPath: () => userData.dir },
safeStorage: { isEncryptionAvailable: () => false },
}))
// 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 {
...fs,
readFileSync: ((...args: Parameters<typeof fs.readFileSync>) => {
if (readFails.next && String(args[0]).endsWith(".json")) {
readFails.next = false
throw Object.assign(new Error("EBUSY: resource busy"), {
code: "EBUSY",
})
}
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,
}
})
import { createPreset, loadPresets } from "@/electron/main/config-manager"
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", () => {
it("keeps a file it could not read for now", () => {
createPreset({ name: "Mine", config: { AI_PROVIDER: "openai" } })
readFails.next = true
expect(loadPresets().presets).toEqual([])
expect(existsSync(presetsFile())).toBe(true)
// A save based on that empty read must not replace the presets
readFails.next = true
expect(() =>
createPreset({ name: "New", config: { AI_PROVIDER: "openai" } }),
).toThrow()
expect(loadPresets().presets.map((p) => p.name)).toEqual(["Mine"])
})
it("saves again once a file it could not read is gone", () => {
createPreset({ name: "Mine", config: { AI_PROVIDER: "openai" } })
readFails.next = true
loadPresets()
// The user removes the file to start over
rmSync(presetsFile())
createPreset({ name: "New", config: { AI_PROVIDER: "openai" } })
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([])
expect(existsSync(presetsFile())).toBe(false)
expect(
readdirSync(userData.dir).some((f) =>
f.startsWith("config-presets.json.corrupt-"),
),
).toBe(true)
})
})
+95
View File
@@ -0,0 +1,95 @@
// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { POST as chat } from "@/app/api/chat/route"
import { POST as parseUrl } from "@/app/api/parse-url/route"
import { POST as providerModels } from "@/app/api/provider-models/route"
import { POST as validateDiagram } from "@/app/api/validate-diagram/route"
import { POST as validateModel } from "@/app/api/validate-model/route"
// The routes that run models or fetch URLs, which a page on another site
// could otherwise make the user's own server do
const ROUTES = {
chat,
"parse-url": parseUrl,
"provider-models": providerModels,
"validate-diagram": validateDiagram,
"validate-model": validateModel,
}
const body = JSON.stringify({
messages: [
{ id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
],
url: "https://example.com",
provider: "openai",
apiKey: "k",
modelId: "gpt-5.5",
imageData: "data:image/png;base64,AAAA",
})
const saved = process.env.NEXT_AI_DRAWIO_DESKTOP
beforeEach(() => {
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw new Error("no network in tests")
}),
)
})
afterEach(() => {
if (saved === undefined) delete process.env.NEXT_AI_DRAWIO_DESKTOP
else process.env.NEXT_AI_DRAWIO_DESKTOP = saved
vi.unstubAllGlobals()
})
describe("requests another website could send", () => {
for (const [name, post] of Object.entries(ROUTES)) {
it(`${name}: refuses a text body, which needs no CORS preflight`, async () => {
// fetch(..., { mode: "no-cors", body: JSON.stringify(...) })
// from another site arrives as text/plain
const res = await post(
new Request(`http://127.0.0.1:61337/api/${name}`, {
method: "POST",
body,
}),
)
expect(res.status).toBe(415)
expect(fetch).not.toHaveBeenCalled()
})
it(`${name}: desktop app refuses a foreign Host (DNS rebinding)`, async () => {
process.env.NEXT_AI_DRAWIO_DESKTOP = "1"
const res = await post(
new Request(`http://127.0.0.1:61337/api/${name}`, {
method: "POST",
headers: {
"Content-Type": "application/json",
host: "rebind.attacker.example:61337",
},
body,
}),
)
expect(res.status).toBe(403)
expect(fetch).not.toHaveBeenCalled()
})
}
it("lets the desktop window's own requests through", async () => {
process.env.NEXT_AI_DRAWIO_DESKTOP = "1"
const res = await validateModel(
new Request("http://127.0.0.1:61337/api/validate-model", {
method: "POST",
headers: {
"Content-Type": "application/json; charset=utf-8",
host: "127.0.0.1:61337",
},
body: JSON.stringify({ provider: "openai" }),
}),
)
// Past the check: the route's own validation answers
expect(res.status).toBe(400)
expect(await res.text()).toMatch(/required/)
})
})
+119
View File
@@ -0,0 +1,119 @@
// @vitest-environment node
import { simulateReadableStream, streamText } from "ai"
import { MockLanguageModelV3 } from "ai/test"
import { describe, expect, it } from "vitest"
import {
withDeprecatedParamsFallback,
withoutDeprecatedParams,
} from "@/lib/deprecated-params"
// The error texts Claude 4.7 and later return (the Anthropic API and Bedrock)
const rejection = (message: string) => ({
statusCode: 400,
message,
responseBody: JSON.stringify({ error: { message } }),
})
const TEMPERATURE = rejection("`temperature` is deprecated for this model.")
const THINKING = rejection(
'"thinking.type.enabled" is not supported for this model. Use "thinking.type.adaptive" and "output_config.effort" to control thinking behavior.',
)
describe("withoutDeprecatedParams", () => {
const params = {
temperature: 0.2,
topP: 0.9,
maxOutputTokens: 1000,
providerOptions: {
anthropic: {
thinking: { type: "enabled", budgetTokens: 4000 },
cacheControl: { type: "ephemeral" },
},
},
}
// Measured on Bedrock: Opus 4.7, 4.8 and every Claude 5 model accept
// adaptive thinking, and return its text only with display "summarized"
const ADAPTIVE = { type: "adaptive", display: "summarized" }
it("drops sampling settings and switches to adaptive thinking", () => {
for (const error of [TEMPERATURE, THINKING]) {
expect(withoutDeprecatedParams(error, params)).toEqual({
maxOutputTokens: 1000,
providerOptions: {
anthropic: {
thinking: ADAPTIVE,
cacheControl: { type: "ephemeral" },
},
},
})
}
})
it("switches a Bedrock thinking budget to adaptive thinking", () => {
const bedrock = {
providerOptions: {
bedrock: {
reasoningConfig: { type: "enabled", budgetTokens: 4000 },
},
},
}
expect(withoutDeprecatedParams(THINKING, bedrock)).toEqual({
providerOptions: { bedrock: { reasoningConfig: ADAPTIVE } },
})
})
it("leaves other errors and requests with nothing to drop alone", () => {
expect(withoutDeprecatedParams(rejection("bad key"), params)).toBeNull()
expect(
withoutDeprecatedParams(
{ ...TEMPERATURE, statusCode: 401 },
params,
),
).toBeNull()
const nothingToDrop = {
providerOptions: {
anthropic: { cacheControl: { type: "ephemeral" } },
},
}
expect(withoutDeprecatedParams(TEMPERATURE, nothingToDrop)).toBeNull()
})
})
describe("withDeprecatedParamsFallback", () => {
it("retries the stream once without the rejected settings", async () => {
const calls: any[] = []
const model = new MockLanguageModelV3({
// The test stream only has the parts this check needs
doStream: (async (options: any) => {
calls.push(options)
if (options.temperature !== undefined) throw TEMPERATURE
return {
stream: simulateReadableStream({
chunks: [
{ type: "text-start", id: "t" },
{ type: "text-delta", id: "t", delta: "ok" },
{ type: "text-end", id: "t" },
{
type: "finish",
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
},
],
}),
}
}) as any,
})
const result = streamText({
model: withDeprecatedParamsFallback(model as any),
prompt: "hi",
temperature: 0.2,
maxRetries: 0,
})
expect(await result.text).toBe("ok")
expect(calls).toHaveLength(2)
expect(calls[1].temperature).toBeUndefined()
})
})
+119
View File
@@ -0,0 +1,119 @@
import { deflateRawSync } from "node:zlib"
import { act, renderHook } from "@testing-library/react"
import type React from "react"
import { afterEach, describe, expect, it, vi } from "vitest"
import { DiagramProvider, useDiagram } from "@/contexts/diagram-context"
vi.mock("sonner", () => ({ toast: { success: vi.fn() } }))
// The provider with a stand-in draw.io that records each export request
function setup() {
const { result } = renderHook(() => useDiagram(), {
wrapper: ({ children }: { children: React.ReactNode }) => (
<DiagramProvider>{children}</DiagramProvider>
),
})
const requests: { format: string; message: string }[] = []
result.current.drawioRef.current = {
exportDiagram: (r: any) => requests.push(r),
load: vi.fn(),
} as any
// draw.io's reply to a request: it echoes the request in `message`
const reply = (request: { message: string }, data: string, xml = "") =>
act(() =>
result.current.handleDiagramExport({
event: "export",
data,
xml,
format: "xmlsvg",
message: request,
} as any),
)
return { result, requests, reply }
}
// An editable SVG as draw.io exports it: the diagram, compressed, in its
// content attribute
const svgOf = (label: string) => {
const model = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="${label}" vertex="1" parent="1"/></root></mxGraphModel>`
const packed = deflateRawSync(
Buffer.from(encodeURIComponent(model)),
).toString("base64")
const content = `<mxfile><diagram id="p">${packed}</diagram></mxfile>`
.replaceAll("&", "&amp;")
.replaceAll("<", "&lt;")
.replaceAll(">", "&gt;")
.replaceAll('"', "&quot;")
const svg = `<svg xmlns="http://www.w3.org/2000/svg" content="${content}"></svg>`
return `data:image/svg+xml;base64,${btoa(svg)}`
}
afterEach(() => {
vi.restoreAllMocks()
})
describe("exports in flight at the same time", () => {
it("give the chat's export only its own reply", () => {
const { result, requests, reply } = setup()
// An edit's history export is still on its way when the chat exports
act(() => {
result.current.handleExport()
})
let tag = ""
const got: string[] = []
act(() => {
tag = result.current.handleExportWithoutHistory()
result.current.exportResolversRef.current[tag] = (xml) =>
got.push(xml)
})
reply(requests[0], svgOf("older"))
expect(got).toEqual([])
reply(requests[1], svgOf("current"))
expect(got).toHaveLength(1)
expect(got[0]).toContain('value="current"')
expect(result.current.exportResolversRef.current[tag]).toBeUndefined()
})
it("save each file with its own result", async () => {
const { result, requests, reply } = setup()
const saved: { name: string; href: string }[] = []
vi.spyOn(HTMLAnchorElement.prototype, "click").mockImplementation(
function (this: HTMLAnchorElement) {
saved.push({ name: this.download, href: this.href })
},
)
const blobs = new Map<string, Blob>()
URL.createObjectURL = vi.fn((blob: Blob) => {
const url = `blob:test-${blobs.size}`
blobs.set(url, blob)
return url
})
URL.revokeObjectURL = vi.fn()
vi.stubGlobal(
"fetch",
vi.fn(async () => new Response("{}")),
)
const twoPages =
'<mxfile><diagram id="a" name="A"><mxGraphModel><root/></mxGraphModel></diagram><diagram id="b" name="B"><mxGraphModel><root/></mxGraphModel></diagram></mxfile>'
act(() => {
result.current.saveDiagramToFile("doc", "drawio")
result.current.saveDiagramToFile("pic", "png")
})
// The PNG answers first
reply(requests[1], "data:image/png;base64,iVBORw0KGgo=")
reply(requests[0], svgOf("doc"), twoPages)
expect(saved.map((s) => s.name)).toEqual(["pic.png", "doc.drawio"])
expect(saved[0].href).toMatch(/^data:image\/png/)
const file = blobs.get(saved[1].href)
const text = await new Promise<string>((resolve) => {
const reader = new FileReader()
reader.onload = () => resolve(String(reader.result))
reader.readAsText(file as Blob)
})
expect(text).toContain('name="A"')
expect(text).toContain('name="B"')
vi.unstubAllGlobals()
})
})
+62
View File
@@ -0,0 +1,62 @@
// @vitest-environment node
import { describe, expect, it } from "vitest"
import { onRequest } from "@/edge-functions/api/edgeai/chat/completions"
function request(headers: Record<string, string>): Request {
return new Request("http://localhost/api/edgeai/chat/completions", {
method: "POST",
headers,
// Non-streaming requests return a mock reply without calling AI
body: JSON.stringify({ messages: [{ role: "user", content: "hi" }] }),
})
}
const json = { "Content-Type": "application/json" }
describe("EdgeOne chat completions function", () => {
it("sends no CORS headers", async () => {
const res = await onRequest({ request: request(json), env: {} })
expect(res.status).toBe(200)
expect(res.headers.get("access-control-allow-origin")).toBeNull()
})
it("rejects non-JSON requests", async () => {
const res = await onRequest({
request: request({ "Content-Type": "text/plain" }),
env: {},
})
expect(res.status).toBe(400)
// A plain-text type that only mentions JSON needs no CORS preflight
const disguised = await onRequest({
request: request({
"Content-Type": "text/plain; x=application/json",
}),
env: {},
})
expect(disguised.status).toBe(400)
const withCharset = await onRequest({
request: request({
"Content-Type": "application/json; charset=utf-8",
}),
env: {},
})
expect(withCharset.status).toBe(200)
})
it("checks the access code when ACCESS_CODE_LIST is set", async () => {
const env = { ACCESS_CODE_LIST: "secret" }
const missing = await onRequest({ request: request(json), env })
expect(missing.status).toBe(401)
const ok = await onRequest({
request: request({ ...json, "x-access-code": "secret" }),
env,
})
expect(ok.status).toBe(200)
})
it("lets requests through when env is unavailable", async () => {
const res = await onRequest({ request: request(json) })
expect(res.status).toBe(200)
})
})
+111
View File
@@ -0,0 +1,111 @@
// @vitest-environment node
import { mkdtempSync, writeFileSync } from "node:fs"
import { tmpdir } from "node:os"
import { join } from "node:path"
import { afterEach, describe, expect, it, vi } from "vitest"
const dir = vi.hoisted(() => ({ path: "" }))
vi.mock("electron", () => ({
app: {
getPath: (name: string) =>
name === "exe" ? `${dir.path}/app/exe` : dir.path,
getAppPath: () => `${dir.path}/app`,
},
}))
import { loadEnvFile } from "@/electron/main/env-loader"
const KEYS = [
"T_JSON",
"T_COMMENT",
"T_PLAIN",
"T_DOUBLE",
"T_QUOTED_COMMENT",
"T_KEY_COMMENT",
"T_HASH",
"T_AFTER",
"T_JOINED",
"T_ESC_HASH",
"T_ESC_INNER",
"T_ESC_COMMENT",
"T_DIR",
"T_DIR_COMMENT",
]
afterEach(() => {
for (const k of KEYS) delete process.env[k]
})
describe("loadEnvFile", () => {
it("reads quoted values like dotenv", () => {
dir.path = mkdtempSync(join(tmpdir(), "env-loader-"))
writeFileSync(
join(dir.path, ".env"),
[
// An apostrophe inside a single-quoted JSON value
`T_JSON='{"name":"Team's models"}'`,
`T_COMMENT="value" # a comment`,
"T_PLAIN=plain # a comment",
`T_DOUBLE="say "hi""`,
].join("\n"),
)
loadEnvFile()
expect(process.env.T_JSON).toBe(`{"name":"Team's models"}`)
expect(process.env.T_COMMENT).toBe("value")
expect(process.env.T_PLAIN).toBe("plain")
expect(process.env.T_DOUBLE).toBe(`say "hi"`)
})
it("drops a comment that ends with a quote, like dotenv", () => {
// Expected values checked against dotenv 16.6.1's parse
dir.path = mkdtempSync(join(tmpdir(), "env-loader-"))
writeFileSync(
join(dir.path, ".env"),
[
`T_QUOTED_COMMENT="gpt-5" # pick "fast"`,
`T_KEY_COMMENT="sk-abc" # from "Team A"`,
`T_AFTER='a' b`,
`T_JOINED="a"b`,
// Unquoted: a # without a space before it stays in the value
// (dotenv would cut it; this loader never did)
"T_HASH=http://host/#/x",
].join("\n"),
)
loadEnvFile()
expect(process.env.T_QUOTED_COMMENT).toBe("gpt-5")
expect(process.env.T_KEY_COMMENT).toBe("sk-abc")
expect(process.env.T_AFTER).toBe(`'a' b`)
expect(process.env.T_JOINED).toBe(`"a"b`)
expect(process.env.T_HASH).toBe("http://host/#/x")
})
it("does not end a quoted value at an escaped quote, like dotenv", () => {
dir.path = mkdtempSync(join(tmpdir(), "env-loader-"))
writeFileSync(
join(dir.path, ".env"),
[
'T_ESC_HASH="abc\\" #def"',
'T_ESC_INNER="a # \\"b\\""',
'T_ESC_COMMENT="x\\"y" # c',
].join("\n"),
)
loadEnvFile()
// Expected values from dotenv 16.6.1, which keeps the backslashes
expect(process.env.T_ESC_HASH).toBe('abc\\" #def')
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\\")
})
})
+265
View File
@@ -0,0 +1,265 @@
// @vitest-environment node
import {
APICallError,
InvalidToolInputError,
RetryError,
simulateReadableStream,
streamText,
tool,
} from "ai"
import { MockLanguageModelV3 } from "ai/test"
import { describe, expect, it } from "vitest"
import { z } from "zod"
import {
classifyLLMError,
isToolCallError,
streamErrorText,
} from "@/lib/llm-errors"
const apiError = (statusCode: number, message: string, responseBody = "") =>
new APICallError({
message,
url: "https://api.example.com/v1/chat/completions",
requestBodyValues: {},
statusCode,
responseBody,
})
describe("classifyLLMError", () => {
it("reads the status code, not the message", () => {
// Providers rarely put the number in their message
expect(
classifyLLMError(apiError(401, "Authentication Fails")).code,
).toBe("invalid_api_key")
expect(classifyLLMError(apiError(404, "Unknown")).code).toBe(
"model_not_found",
)
expect(classifyLLMError(apiError(503, "busy")).code).toBe(
"provider_unavailable",
)
})
it("lets a specific text win over the status code", () => {
expect(
classifyLLMError(
apiError(429, "You exceeded your current quota, check billing"),
).code,
).toBe("insufficient_quota")
expect(
classifyLLMError(
apiError(400, "This model's maximum context length is 128000"),
).code,
).toBe("context_too_long")
expect(
classifyLLMError(
apiError(400, "bad", '{"message":"toolUse.input is invalid"}'),
).code,
).toBe("output_truncated")
})
it("does not call a 403 an invalid key", () => {
expect(classifyLLMError(apiError(403, "Forbidden")).code).toBe(
"forbidden",
)
})
it("uses the last attempt after retries", () => {
const retry = new RetryError({
message: "Failed after 3 attempts",
reason: "maxRetriesExceeded",
errors: [apiError(500, "x"), apiError(429, "slow down")],
})
expect(classifyLLMError(retry).code).toBe("rate_limited")
})
it("keeps the message but hides secrets in it", () => {
const { code, message } = classifyLLMError(
apiError(
401,
"Incorrect API key provided: sk-proj-abcdefghijklmnop. Header Bearer abc.def",
),
)
expect(code).toBe("invalid_api_key")
expect(message).toContain("Incorrect API key provided")
expect(message).not.toContain("abcdefghijklmnop")
expect(message).not.toContain("abc.def")
})
it("leaves our own messages readable", () => {
// This one used to be replaced by "Authentication failed" for
// containing the word key
const { message } = classifyLLMError(
new Error(
"API key is required when using a custom base URL. Please provide your own API key in Settings.",
),
)
expect(message).toContain("API key is required when using a custom")
})
it("names a timeout", () => {
const timeout = new Error("The operation was aborted due to timeout")
timeout.name = "TimeoutError"
expect(classifyLLMError(timeout).code).toBe("timeout")
})
it("points to the model id when Bedrock wants an inference profile", () => {
const error = apiError(
400,
"Invocation of model ID anthropic.claude-sonnet-5-5 with on-demand throughput isn’t supported. Retry your request with the ID or ARN of an inference profile that contains this model.",
)
expect(classifyLLMError(error).code).toBe("model_not_found")
})
it("reads Bedrock's token throttling as a rate limit", () => {
const error = apiError(
429,
"Too many tokens, please wait before trying again.",
)
expect(classifyLLMError(error).code).toBe("rate_limited")
})
it("names a network error the SDK wrapped", () => {
const error = new APICallError({
message:
"Cannot connect to API: Connect Timeout Error (attempted address: api.example.com:443, timeout: 10000ms)",
url: "https://api.example.com/v1/chat/completions",
requestBodyValues: {},
})
expect(classifyLLMError(error).code).toBe("cannot_connect")
})
it("reads an error object sent in the stream", () => {
// OpenRouter, when the upstream provider is overloaded
const error = {
code: 503,
message:
"Upstream error from Nvidia: Service temporarily overloaded",
metadata: { error_type: "provider_overloaded" },
}
expect(classifyLLMError(error)).toEqual({
type: "provider",
code: "provider_unavailable",
message:
"Upstream error from Nvidia: Service temporarily overloaded",
})
})
it("adds the reason from a problem+json body", () => {
// NVIDIA, for a retired model; the SDK's message is only "Gone"
const body = JSON.stringify({
title: "Gone",
status: 410,
detail: "The model 'deepseek-v4-flash' has reached its end of life",
})
expect(classifyLLMError(apiError(410, "Gone", body))).toEqual({
type: "provider",
code: "model_not_found",
message:
"Gone: The model 'deepseek-v4-flash' has reached its end of life",
})
})
})
describe("streamErrorText", () => {
it("keeps the text of a tool call the model got wrong", async () => {
// Seen with Claude Opus 5.5: a quote left unescaped in the input
const model = new MockLanguageModelV3({
doStream: (async () => ({
stream: simulateReadableStream({
chunks: [
{
type: "tool-call",
toolCallId: "c1",
toolName: "edit_diagram",
input: '{"operations": [{"new_xml": "as="x""}]}',
},
{
type: "finish",
finishReason: {
unified: "tool-calls",
raw: "tool_use",
},
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
},
],
}),
})) as any,
})
const result = streamText({
model: model as any,
prompt: "edit",
tools: {
edit_diagram: tool({
inputSchema: z.object({ operations: z.array(z.any()) }),
}),
},
})
const errors: string[] = []
for await (const chunk of result.toUIMessageStream({
onError: streamErrorText,
})) {
if ("errorText" in chunk) errors.push(chunk.errorText)
}
expect(errors.length).toBeGreaterThan(0)
for (const text of errors) {
expect(text).toMatch(/^Invalid input for tool edit_diagram/)
}
})
it("hides the provider's text on the server's keys", () => {
const error = apiError(
403,
"User: arn:aws:sts::123456789012:assumed-role/app/s is not authorized to perform: bedrock:InvokeModel",
)
const hidden = JSON.parse(streamErrorText(error, true))
expect(hidden.code).toBe("forbidden")
expect(hidden.message).not.toMatch(/arn:aws|123456789012/)
expect(JSON.parse(streamErrorText(error)).message).toMatch(
/not authorized/,
)
})
it("classifies a provider error", () => {
expect(JSON.parse(streamErrorText(apiError(401, "bad key")))).toEqual({
type: "provider",
code: "invalid_api_key",
message: "bad key",
})
})
it("classifies a provider error sent as plain text", () => {
// DeepSeek's SDK sends errors in the stream as a string
const text = "Insufficient Balance for account 42"
expect(JSON.parse(streamErrorText(text))).toEqual({
type: "provider",
code: "insufficient_quota",
message: text,
})
expect(JSON.parse(streamErrorText(text, true)).message).not.toMatch(
/account 42/,
)
})
it("classifies Bedrock's throttling sent in the stream", () => {
// Bedrock's ThrottlingException as a plain object, not an API error
const throttled = {
message: "Too many tokens, please wait before trying again.",
}
expect(JSON.parse(streamErrorText(throttled)).code).toBe("rate_limited")
})
})
describe("isToolCallError", () => {
it("spots errors the model must see unchanged", () => {
const invalid = new InvalidToolInputError({
toolName: "display_diagram",
toolInput: "{",
cause: new Error("bad JSON"),
})
expect(isToolCallError(invalid)).toBe(true)
expect(isToolCallError(apiError(500, "x"))).toBe(false)
})
})
+42
View File
@@ -0,0 +1,42 @@
// @vitest-environment node
import { describe, expect, it, vi } from "vitest"
import { POST as logSave } from "@/app/api/log-save/route"
// A Langfuse client with a trace to attach the save score to
const batch = vi.fn(async () => ({}))
vi.mock("@/lib/langfuse", () => ({
getLangfuseClient: () => ({
api: {
trace: { list: async () => ({ data: [{ id: "trace-1" }] }) },
ingestion: { batch },
},
}),
}))
function post(body: unknown): Request {
return new Request("http://localhost/api/log-save", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(body),
})
}
describe("POST /api/log-save", () => {
it.each(["drawio", "png", "svg", "xmlsvg"])(
"logs a save in %s format",
async (format) => {
const res = await logSave(
post({ filename: "diagram", format, sessionId: "s1" }),
)
expect(res.status).toBe(200)
expect(await res.json()).toEqual({ success: true, logged: true })
},
)
it("rejects an unknown format", async () => {
const res = await logSave(
post({ filename: "diagram", format: "pdf", sessionId: "s1" }),
)
expect(res.status).toBe(400)
})
})
+265
View File
@@ -0,0 +1,265 @@
/**
* The web app runs the MCP server's diagram modules in the browser, where
* DOMParser and XMLSerializer are the native ones (jsdom here), not the
* linkedom polyfill the MCP tests use.
*/
import { deflateRaw } from "pako"
import { describe, expect, it } from "vitest"
import { applyDiagramOperations } from "@/packages/mcp-server/src/diagram-operations.ts"
import { editDiagram } from "@/packages/mcp-server/src/edit-diagram.ts"
import { decompressPageContent } from "@/packages/mcp-server/src/load-diagram.ts"
import {
hasCells,
normalizeToMxfile,
wrapCellsInModel,
} from "@/packages/mcp-server/src/pages.ts"
import { getXmlSyntaxError } from "@/packages/mcp-server/src/xml-syntax.ts"
import {
autoFixXml,
validateAndFixXml,
validateMxCellStructure,
} from "@/packages/mcp-server/src/xml-validation.ts"
const box = (id: string, parent = "1") =>
`<mxCell id="${id}" value="${id}" vertex="1" parent="${parent}"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
const edge = (id: string, source: string, target: string) =>
`<mxCell id="${id}" edge="1" parent="1" source="${source}" target="${target}"><mxGeometry relative="1" as="geometry"/></mxCell>`
const file = (cells: string) =>
normalizeToMxfile(wrapCellsInModel(cells), {
pageId: "p1",
pageName: "Page-1",
}) as string
describe("MCP diagram modules with a browser DOM", () => {
it("wraps bare cells into a valid file", () => {
const xml = file(box("a") + box("b"))
expect(xml).toContain('<diagram id="p1" name="Page-1">')
expect(validateAndFixXml(xml).valid).toBe(true)
expect(hasCells(xml)).toBe(true)
expect(hasCells(file(""))).toBe(false)
})
it("fixes the case of a misspelled tag in model XML", () => {
const result = validateAndFixXml(
file(box("a"))
.replace('<mxCell id="a"', '<mxcell id="a"')
.replace("</mxCell></root>", "</mxcell></root>"),
)
expect(result.valid).toBe(true)
expect(result.fixes.join(" ")).toMatch(/tag case/)
})
it("reports syntax errors with line and column", () => {
expect(getXmlSyntaxError("<mxfile><diagram></mxfile>")).toMatch(
/^1:\d+/,
)
expect(getXmlSyntaxError(file(box("a")))).toBeNull()
})
it("deletes a cell with its edges", () => {
const xml = file(box("a") + box("b") + edge("e", "a", "b"))
const { result, errors } = applyDiagramOperations(xml, [
{ operation: "delete", cell_id: "a" },
])
expect(errors).toEqual([])
expect(result).not.toContain('id="a"')
expect(result).not.toContain('id="e"')
expect(result).toContain('id="b"')
})
it("runs a whole edit and serializes the target page", () => {
const outcome = editDiagram(
file(box("a")),
[{ operation: "add", cell_id: "b", new_xml: box("b") }],
{},
)
expect(outcome.ok).toBe(true)
if (outcome.ok) expect(outcome.xml).toContain('id="b"')
const failed = editDiagram(
file(box("a")),
[{ operation: "add", cell_id: "b", new_xml: box("b") + box("c") }],
{},
)
expect(failed.ok).toBe(false)
})
it("decompresses a draw.io compressed page", () => {
const model = wrapCellsInModel(box("a"))
const deflated = deflateRaw(encodeURIComponent(model))
const base64 = btoa(String.fromCharCode(...deflated))
expect(decompressPageContent(base64)).toBe(model)
expect(decompressPageContent("not compressed")).toBeNull()
})
})
// Repair cases fixed in the web app's own copy before it moved here
const page = (id: string, cells: string) =>
`<diagram name="${id}" id="${id}"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel></diagram>`
describe("duplicate ids in multi-page documents", () => {
const shape = (id: string, value = "Box") =>
`<mxCell id="${id}" value="${value}" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
it("accepts the same ids on different pages", () => {
const xml = `<mxfile>${page("p1", shape("2"))}${page("p2", shape("2"))}</mxfile>`
expect(validateMxCellStructure(xml)).toBeNull()
})
it("still reports duplicate ids within one page", () => {
const xml = `<mxfile>${page("p1", shape("2") + shape("2"))}${page("p2", "")}</mxfile>`
expect(validateMxCellStructure(xml)).toMatch(/duplicate cell ID/i)
})
it("does not rename the root cells of other pages when fixing", () => {
const xml = `<mxfile>${page("p1", shape("2", "R&D"))}${page("p2", shape("3"))}</mxfile>`
const result = validateAndFixXml(xml)
expect(result.valid).toBe(true)
expect(result.fixed).not.toContain("_dup")
expect(result.fixed).toContain("R&amp;D")
})
it("renames a duplicate id in a bare model, as display_diagram has", () => {
// In an <mxfile> the duplicate is reported instead (above)
const xml = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${shape("d") + shape("d")}</root></mxGraphModel>`
const { fixed } = autoFixXml(xml)
expect(fixed).toContain('<mxCell id="d" ')
expect(fixed).toContain('<mxCell id="d_dup1" ')
})
})
describe("autoFixXml", () => {
it("does not insert a space at the start of style values", () => {
const xml = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="R&D" style="rounded=1;whiteSpace=wrap;" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></root></mxGraphModel>`
const { fixed } = autoFixXml(xml)
expect(fixed).toContain('style="rounded=1;whiteSpace=wrap;"')
})
it("adds a missing space between attributes", () => {
const xml = `<mxCell id="2" vertex="1"parent="1"/>`
expect(autoFixXml(xml).fixed).toContain('vertex="1" parent="1"')
})
it("keeps &quot; inside rich text labels", () => {
const label = "&lt;font color=&quot;#ff0000&quot;&gt;Hello&lt;/font&gt;"
const xml = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="${label}" style="html=1;" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell><mxCell id="3" value="Q&A" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></root></mxGraphModel>`
const result = validateAndFixXml(xml)
expect(result.valid).toBe(true)
expect(result.fixed).toContain(`value="${label}"`)
})
it("fixes an attribute delimited by &quot;", () => {
const xml = `<mxCell id="2" dashPattern=&quot;1 1;&quot; vertex="1" parent="1"/>`
expect(autoFixXml(xml).fixed).toContain('dashPattern="1 1;"')
})
it("keeps cells written on one line next to multi-line cells", () => {
const xml = `<mxGraphModel><root>
<mxCell id="0"/>
<mxCell id="1" parent="0"/>
<mxCell id="2" value="Q&A" vertex="1" parent="1">
<mxGeometry x="0" y="0" width="80" height="40" as="geometry"/>
</mxCell>
<mxCell id="e1" edge="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"/></mxCell>
<mxCell id="3" value="B" vertex="1" parent="1">
<mxGeometry x="200" y="0" width="80" height="40" as="geometry"/>
</mxCell>
</root></mxGraphModel>`
const result = validateAndFixXml(xml)
expect(result.valid).toBe(true)
for (const id of ["2", "e1", "3"]) {
expect(result.fixed).toContain(`<mxCell id="${id}"`)
}
})
it("keeps object and UserObject wrappers", () => {
const xml = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><UserObject id="2" label="Docs" link="https://example.com"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject><object id="3" label="A&B" owner="me"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></object></root></mxGraphModel>`
const result = validateAndFixXml(xml)
expect(result.valid).toBe(true)
expect(result.fixed).toContain('<UserObject id="2"')
expect(result.fixed).toContain('<object id="3"')
})
})
describe("hasCells (was isMinimalDiagram in the web app)", () => {
it("returns true for empty diagram", () => {
const xml = '<mxCell id="0"/><mxCell id="1" parent="0"/>'
expect(hasCells(xml)).toBe(false)
})
it("returns false for diagram with content", () => {
const xml =
'<mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="Hello"/>'
expect(hasCells(xml)).toBe(true)
})
it("handles whitespace correctly", () => {
const xml = ' <mxCell id="0"/> <mxCell id="1" parent="0"/> '
expect(hasCells(xml)).toBe(false)
})
it("returns false for a shape drawn in draw.io with a random id", () => {
const xml =
'<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="xY3kQ9-1" value="" style="rounded=0;" vertex="1" parent="1"><mxGeometry x="10" y="10" width="120" height="60" as="geometry"/></mxCell></root></mxGraphModel>'
expect(hasCells(xml)).toBe(true)
})
it("does not mistake ids that start with 0 or 1 for root cells", () => {
const xml =
'<mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="10"/>'
expect(hasCells(xml)).toBe(true)
})
it("counts a cell wrapped in a UserObject", () => {
const xml =
'<mxCell id="0"/><mxCell id="1" parent="0"/><UserObject id="u" link="x"><mxCell vertex="1" parent="1"/></UserObject>'
expect(hasCells(xml)).toBe(true)
})
})
describe("applyDiagramOperations with wrapped cells", () => {
const xml = `<mxfile><diagram id="p1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><UserObject id="5" label="Docs" link="https://example.com"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject><mxCell id="6" value="B" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell><mxCell id="e1" edge="1" parent="1" source="5" target="6"><mxGeometry relative="1" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>`
it("deletes a wrapped cell and its edges", () => {
const { result, errors } = applyDiagramOperations(xml, [
{ operation: "delete", cell_id: "5" },
{ operation: "delete", cell_id: "e1" },
])
expect(errors).toEqual([])
expect(result).not.toContain("UserObject")
expect(result).not.toContain('id="e1"')
expect(result).toContain('id="6"')
})
it("rejects adding a cell with the id of a wrapped cell", () => {
const { errors } = applyDiagramOperations(xml, [
{
operation: "add",
cell_id: "5",
new_xml: '<mxCell id="5" vertex="1" parent="1"/>',
},
])
expect(errors[0]?.message).toContain("already exists")
})
it("updates a wrapped cell", () => {
const { result, errors } = applyDiagramOperations(xml, [
{
operation: "update",
cell_id: "5",
new_xml:
'<UserObject id="5" label="New" link="https://example.org"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject>',
},
])
expect(errors).toEqual([])
expect(result).toContain('label="New"')
expect(result).not.toContain('label="Docs"')
})
it("reports deleting a cell that does not exist", () => {
const { errors } = applyDiagramOperations(xml, [
{ operation: "delete", cell_id: "missing" },
])
expect(errors[0]?.message).toContain("not found")
})
})
+46
View File
@@ -0,0 +1,46 @@
import { readFileSync } from "node:fs"
import { join } from "node:path"
import { describe, expect, it } from "vitest"
// The MCP preview page, as the server fills it in, with its script run in
// this document (no session id, so it does not poll)
const dir = join(process.cwd(), "packages/mcp-server/src/preview")
const html = readFileSync(join(dir, "index.html"), "utf8")
.replace("{{CSS}}", "")
.replace("{{SESSION_BADGE}}", "")
.replaceAll("{{DISABLED}}", "")
.replace("{{DRAWIO_URL}}", "about:blank")
.replace("{{SESSION_JSON}}", '""')
.replace("{{ORIGIN_JSON}}", '"https://embed.diagrams.net"')
const scripts = [...html.matchAll(/<script>([\s\S]*?)<\/script>/g)].map((m) =>
m[1].replace("{{SCRIPT}}", ""),
)
const preview = readFileSync(join(dir, "preview.js"), "utf8")
function renderHistory(entries: unknown[]): HTMLElement {
document.body.innerHTML = html.replace(/<script>[\s\S]*?<\/script>/g, "")
// One scope, as the page's scripts share one; returns its renderHistory
const run = new Function(
`${scripts.join("\n")}\n${preview}\nreturn (d) => { historyData = d; renderHistory(); }`,
)
run()(entries)
return document.getElementById("history-grid") as HTMLElement
}
describe("MCP preview History", () => {
it("never reads a stored thumbnail as HTML", () => {
const grid = renderHistory([
{ id: 1, index: 0, svg: 'x" onerror="window.__xss=1' },
{ id: 2, index: 1, svg: "javascript:window.__xss=2" },
{ id: 3, index: 2, svg: "data:image/svg+xml;base64,PHN2Zy8+" },
])
const images = [...grid.querySelectorAll("img")]
expect(images.map((i) => i.getAttribute("src"))).toEqual([
"data:image/svg+xml;base64,PHN2Zy8+",
])
expect(grid.querySelector("[onerror]")).toBeNull()
// Entries without a usable picture show their number
expect(grid.textContent).toContain("#0")
expect(grid.textContent).toContain("#1")
})
})
+436
View File
@@ -0,0 +1,436 @@
import { readFileSync } from "node:fs"
import { join } from "node:path"
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
// The MCP preview page's script, run in this document with a stubbed
// server (fetch) and draw.io iframe (its postMessage), so the tab's side of
// a recreated session can be driven step by step
const dir = join(process.cwd(), "packages/mcp-server/src/preview")
const DRAWIO = "https://embed.diagrams.net"
const html = readFileSync(join(dir, "index.html"), "utf8")
.replace("{{CSS}}", "")
.replace("{{SESSION_BADGE}}", "")
.replaceAll("{{DISABLED}}", "")
.replace("{{DRAWIO_URL}}", "about:blank")
.replace("{{SESSION_JSON}}", '"mcp-test"')
.replace("{{ORIGIN_JSON}}", JSON.stringify(DRAWIO))
const scripts = [...html.matchAll(/<script>([\s\S]*?)<\/script>/g)].map((m) =>
m[1].replace("{{SCRIPT}}", ""),
)
const preview = readFileSync(join(dir, "preview.js"), "utf8")
const pageListeners: Array<[string, EventListener]> = []
type Answer = { status: number; body: unknown }
interface Call {
url: string
method: string
body: any
answer: (a: Answer) => void
fail: () => void
}
function openPage() {
document.body.innerHTML = html.replace(/<script>[\s\S]*?<\/script>/g, "")
const toDrawio: any[] = []
const iframe = document.getElementById("drawio") as HTMLIFrameElement
Object.defineProperty(iframe, "contentWindow", {
value: { postMessage: (m: string) => toDrawio.push(JSON.parse(m)) },
})
// Every request waits until the test answers it
const calls: Call[] = []
vi.stubGlobal(
"fetch",
vi.fn(
(url: string, init?: RequestInit) =>
new Promise((resolve, reject) => {
calls.push({
url,
method: init?.method ?? "GET",
body: init?.body ? JSON.parse(String(init.body)) : null,
answer: ({ status, body }) =>
resolve(
new Response(JSON.stringify(body), { status }),
),
fail: () => reject(new TypeError("Failed to fetch")),
})
}),
),
)
// The page's window listeners, removed after the test
const addListener = window.addEventListener
window.addEventListener = ((type: string, listener: any, options?: any) => {
pageListeners.push([type, listener])
addListener.call(window, type, listener, options)
}) as typeof window.addEventListener
const run = new Function(
`${scripts.join("\n")}\n${preview}\nreturn { poll, read: () => ({ stateId, currentVersion, lastXml, latestXml }) }`,
)
let created: unknown
try {
created = run()
} finally {
window.addEventListener = addListener
}
const page = created as {
poll: () => Promise<void>
read: () => {
stateId: string | null
currentVersion: number
lastXml: string | null
latestXml: string | null
}
}
const fromDrawio = (msg: object) =>
window.dispatchEvent(
new MessageEvent("message", {
data: JSON.stringify(msg),
origin: DRAWIO,
}),
)
const settle = () => new Promise((r) => setTimeout(r, 0))
const next = (method: string) => {
const call = calls.find((c) => c.method === method)
if (!call) throw new Error(`no pending ${method}`)
calls.splice(calls.indexOf(call), 1)
return call
}
return { page, toDrawio, calls, fromDrawio, settle, next }
}
const state = (
stateId: string,
version: number,
xml: string,
blank = false,
) => ({
status: 200,
body: { stateId, version, xml, blank, syncRequested: false },
})
/** A tab in step with state S1 at version 2, showing diagram A */
async function inStep() {
const t = openPage()
t.next("GET").answer(state("S1", 2, "<mxfile>A</mxfile>"))
await t.settle()
t.fromDrawio({ event: "init" })
await t.settle()
expect(t.page.read().lastXml).toBe("<mxfile>A</mxfile>")
return t
}
beforeEach(() => {
vi.useFakeTimers({ toFake: ["setInterval"] })
})
afterEach(() => {
for (const [type, listener] of pageListeners.splice(0)) {
window.removeEventListener(type, listener)
}
vi.useRealTimers()
vi.unstubAllGlobals()
})
describe("MCP preview after the server recreated its session", () => {
it("keeps edits made while the server was down", async () => {
const t = await inStep()
// The user edits; the server is gone
t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
// draw.io's SVG export answers, then the push fails
t.fromDrawio({ event: "export", data: "<svg/>" })
await t.settle()
t.next("POST").fail()
await t.settle()
// A new process recovered the file, which holds A
const poll = t.page.poll()
t.next("GET").answer(state("S2", 1, "<mxfile>A</mxfile>"))
await poll
const push = t.next("POST")
expect(push.body).toMatchObject({
xml: "<mxfile>B</mxfile>",
stateId: "S2",
baseVersion: 1,
})
})
it("shows the server's diagram and keeps the tab's in History", async () => {
const t = await inStep()
const poll = t.page.poll()
t.next("GET").answer(state("S2", 1, "<mxfile>C</mxfile>"))
await poll
expect(t.toDrawio.at(-1)).toMatchObject({
action: "load",
xml: "<mxfile>C</mxfile>",
})
expect(t.next("POST").body).toMatchObject({
xml: "<mxfile>A</mxfile>",
source: "recover",
stateId: "S2",
})
})
it("sends an edit of the replaced canvas to History, until draw.io loaded", async () => {
const t = await inStep()
const poll = t.page.poll()
t.next("GET").answer(state("S2", 1, "<mxfile>C</mxfile>"))
await poll
t.next("POST") // the tab's copy, to History
// An autosave the old canvas sent before the load
t.fromDrawio({ event: "autosave", xml: "<mxfile>A edited</mxfile>" })
await t.settle()
expect(t.next("POST").body).toMatchObject({
xml: "<mxfile>A edited</mxfile>",
source: "recover",
})
// After the load, edits are edits again
t.fromDrawio({ event: "load" })
t.fromDrawio({ event: "autosave", xml: "<mxfile>C edited</mxfile>" })
t.fromDrawio({ event: "export", data: "<svg/>" })
await t.settle()
expect(t.next("POST").body).toMatchObject({
xml: "<mxfile>C edited</mxfile>",
source: "edit",
})
})
it("ignores the late answer to an old state's push", async () => {
const t = await inStep()
t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
t.fromDrawio({ event: "export", data: "<svg/>" })
await t.settle()
const oldPush = t.next("POST")
const poll = t.page.poll()
t.next("GET").answer(state("S2", 1, "<mxfile>B</mxfile>"))
await poll
oldPush.answer({ status: 200, body: { success: true, version: 21 } })
await t.settle()
expect(t.page.read()).toMatchObject({
stateId: "S2",
currentVersion: 1,
})
})
it("drops a poll answer older than one already handled", async () => {
const t = await inStep()
const first = t.page.poll()
const firstGet = t.next("GET")
const second = t.page.poll()
t.next("GET").answer(state("S2", 1, "<mxfile>A</mxfile>"))
await second
// The answer from before the restart comes last
firstGet.answer(state("S1", 3, "<mxfile>old</mxfile>"))
await first
expect(t.page.read().stateId).toBe("S2")
})
it("saves an undo made while a push was on its way", 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: equal to the last saved diagram, so no push
t.fromDrawio({ event: "autosave", xml: "<mxfile>A</mxfile>" })
await t.settle()
expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0)
pushB.answer({ status: 200, body: { success: true, version: 3 } })
await t.settle()
await t.settle()
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("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()
t.next("GET").answer({
status: 200,
body: {
...state("S1", 2, "<mxfile>A</mxfile>").body,
syncRequested: true,
},
})
await poll
const request = t.toDrawio.at(-1)
expect(request).toMatchObject({ action: "export", format: "xml" })
// draw.io's export of the canvas, formatted unlike its autosave
t.fromDrawio({
event: "export",
format: "xml",
xml: '<mxfile host="drawio">A</mxfile>',
message: request,
})
await t.settle()
const sync = t.next("POST")
expect(sync.body.source).toBe("sync")
sync.answer({ status: 200, body: { success: true, version: 3 } })
await t.settle()
await t.settle()
expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0)
})
})
describe("MCP preview thumbnails and downloads", () => {
/** The tab loads the server write B at version 3 and asks for its image */
async function loadedB() {
const t = await inStep()
const poll = t.page.poll()
t.next("GET").answer(state("S1", 3, "<mxfile>B</mxfile>"))
await poll
await new Promise((r) => setTimeout(r, 600))
const request = t.toDrawio.at(-1)
expect(request).toMatchObject({ action: "export", format: "svg" })
return { t, n: request.thumbExport as number }
}
const thumbnailPosts = (t: ReturnType<typeof openPage>) =>
t.calls.filter((c) => c.url === "/api/history-svg")
it("sends the image with the state and version it shows", async () => {
const { t, n } = await loadedB()
t.fromDrawio({
event: "export",
data: "<svg/>",
message: { thumbExport: n },
})
await t.settle()
expect(thumbnailPosts(t).map((c) => c.body)).toEqual([
expect.objectContaining({ stateId: "S1", version: 3 }),
])
})
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>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 () => {
const { t, n } = await loadedB()
t.fromDrawio({ event: "autosave", xml: "<mxfile>B edited</mxfile>" })
t.fromDrawio({
event: "export",
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 () => {
const t = await inStep()
t.fromDrawio({ event: "autosave", xml: "<mxfile>B</mxfile>" })
t.fromDrawio({ event: "export", data: "<svg/>" })
await t.settle()
t.next("POST").fail()
await t.settle()
let saved: Blob | undefined
URL.createObjectURL = vi.fn((blob: Blob) => {
saved = blob
return "blob:test"
})
URL.revokeObjectURL = vi.fn()
;(document.getElementById("save-format") as HTMLSelectElement).value =
"drawio"
document.getElementById("save-confirm-btn")?.click()
const text = await new Promise<string>((resolve) => {
const reader = new FileReader()
reader.onload = () => resolve(String(reader.result))
reader.readAsText(saved as Blob)
})
expect(text).toBe("<mxfile>B</mxfile>")
})
})
describe("MCP preview with a diagram over the size limit", () => {
it("retries without the image, then tells the user", async () => {
const t = await inStep()
t.fromDrawio({ event: "autosave", xml: "<mxfile>huge</mxfile>" })
t.fromDrawio({ event: "export", data: "<svg/>" })
await t.settle()
const first = t.next("POST")
expect(first.body.svg).not.toBe("")
first.answer({ status: 413, body: { error: "Payload too large" } })
await t.settle()
const retry = t.next("POST")
expect(retry.body).toMatchObject({
xml: "<mxfile>huge</mxfile>",
svg: "",
})
expect(
document.getElementById("notice")?.classList.contains("open"),
).toBe(false)
retry.answer({ status: 413, body: { error: "Payload too large" } })
await t.settle()
expect(t.calls.filter((c) => c.method === "POST")).toHaveLength(0)
expect(document.getElementById("notice")?.textContent).toContain(
"too large",
)
})
})
+32
View File
@@ -0,0 +1,32 @@
import { describe, expect, it } from "vitest"
import { getModelInfo } from "@/lib/model-catalog"
import catalog from "@/lib/model-catalog.json"
describe("getModelInfo", () => {
it("finds a model by its exact id, ignoring case", () => {
const id = Object.keys(catalog.openai).find((m) => m === "gpt-4.1")
expect(id).toBe("gpt-4.1")
expect(getModelInfo("openai", "GPT-4.1")).toEqual(
catalog.openai["gpt-4.1"],
)
})
it("finds a dated or tagged variant by the longest id it starts with", () => {
expect(getModelInfo("openai", "gpt-4.1-2025-04-14")).toEqual(
catalog.openai["gpt-4.1"],
)
})
it("does not match a different model that shares a prefix", () => {
// gpt-4.1-mini is its own entry, not gpt-4.1
expect(getModelInfo("openai", "gpt-4.1-mini")).toEqual(
catalog.openai["gpt-4.1-mini"],
)
expect(getModelInfo("openai", "gpt-4.1x")).toBeUndefined()
})
it("knows nothing about providers models.dev does not list", () => {
expect(getModelInfo("sglang", "anything")).toBeUndefined()
expect(getModelInfo("openai", "not-a-model")).toBeUndefined()
})
})
+212 -13
View File
@@ -3,6 +3,7 @@ import {
DEFAULT_MAX_OUTPUT_TOKENS,
parseOutputTokenLimit,
resolveMaxOutputTokens,
retryOutputTokens,
withOutputTokenLimitFallback,
} from "@/lib/output-token-limit"
@@ -98,34 +99,204 @@ describe("parseOutputTokenLimit", () => {
}
expect(parseOutputTokenLimit(error)).toBeNull()
})
it("reads the ceiling from a Volcengine Ark rejection", () => {
const error = {
message:
"The parameter `max_tokens` specified in the request are not valid: integer above maximum value, expected a value <= 32768, but got 64000 instead.",
statusCode: 400,
}
expect(parseOutputTokenLimit(error)).toBe(32768)
// Same message JSON-escaped in the response body
expect(
parseOutputTokenLimit({
message: "Bad request",
responseBody:
'{"error":{"message":"The parameter `max_tokens` specified in the request are not valid: integer above maximum value, expected a value \\u003c= 16384, but got 64000 instead."}}',
}),
).toBe(16384)
})
it("reads the ceiling from a DashScope rejection", () => {
const error = {
message:
"<400> InternalError.Algo.InvalidParameter: Range of max_tokens should be [1, 8192]",
}
expect(parseOutputTokenLimit(error)).toBe(8192)
})
it("subtracts the input in SGLang and vLLM context rejections", () => {
// SGLang
expect(
parseOutputTokenLimit({
message:
"Requested token count exceeds the model's maximum context length of 32768 tokens. You requested a total of 70000 tokens: 6000 tokens from the input messages and 64000 tokens for the completion.",
}),
).toBe(32768 - 6000 - 1024)
// vLLM, older wording
expect(
parseOutputTokenLimit({
message:
"This model's maximum context length is 32768 tokens. However, you requested 70000 tokens (6000 in the messages, 64000 in the completion).",
}),
).toBe(32768 - 6000 - 1024)
// vLLM, newer wording
expect(
parseOutputTokenLimit({
message:
"This model's maximum context length is 32768 tokens and your request has 6000 input tokens (64000 > 32768 - 6000).",
}),
).toBe(32768 - 6000 - 1024)
})
})
describe("retryOutputTokens", () => {
const bedrockLimit = Object.assign(
new Error(
"The maximum tokens you requested exceeds the model limit of 64000.",
),
{ statusCode: 400 },
)
it("leaves room for the Bedrock thinking budget the provider adds", () => {
// 64000 + 12000 thinking was sent, so the ceiling of 64000 is below it
expect(
retryOutputTokens(bedrockLimit, {
maxOutputTokens: 64000,
providerOptions: {
bedrock: {
reasoningConfig: {
type: "enabled",
budgetTokens: 12000,
},
},
},
}),
).toBe(52000)
})
it("leaves room for the Anthropic thinking budget the provider adds", () => {
const error = Object.assign(
new Error(
"max_tokens: 76000 > 64000, which is the maximum allowed number of output tokens",
),
{ statusCode: 400 },
)
expect(
retryOutputTokens(error, {
maxOutputTokens: 64000,
providerOptions: {
anthropic: {
thinking: { type: "enabled", budgetTokens: 12000 },
},
},
}),
).toBe(52000)
})
it("does not retry when the ceiling covers what was sent", () => {
expect(
retryOutputTokens(bedrockLimit, { maxOutputTokens: 64000 }),
).toBeNull()
})
it("does not retry when the thinking budget leaves no usable room", () => {
const error = Object.assign(new Error("model limit of 16000"), {
statusCode: 400,
})
expect(
retryOutputTokens(error, {
maxOutputTokens: 64000,
providerOptions: {
bedrock: {
reasoningConfig: {
type: "enabled",
budgetTokens: 15500,
},
},
},
}),
).toBeNull()
})
it("falls back to 16000 when the budget is named but no number can be read", () => {
const error = Object.assign(
new Error("max_tokens (64000) exceeds the limit for this model"),
{ statusCode: 400 },
)
expect(retryOutputTokens(error, { maxOutputTokens: 64000 })).toBe(16000)
// Nothing to gain when the request was already that small
expect(retryOutputTokens(error, { maxOutputTokens: 16000 })).toBeNull()
})
it("does not fall back for errors that do not name the budget", () => {
const error = Object.assign(new Error("temperature must be <= 2"), {
statusCode: 400,
})
expect(retryOutputTokens(error, { maxOutputTokens: 64000 })).toBeNull()
// Not a bad request, even though it names the budget
const auth = Object.assign(new Error("max_tokens: invalid API key"), {
statusCode: 401,
})
expect(retryOutputTokens(auth, { maxOutputTokens: 64000 })).toBeNull()
})
it("does not fall back when a ceiling was found but is too small", () => {
const error = Object.assign(
new Error(
"max_tokens: 64000 > 512, which is the maximum allowed number of output tokens",
),
{ statusCode: 400 },
)
expect(retryOutputTokens(error, { maxOutputTokens: 64000 })).toBeNull()
})
})
describe("resolveMaxOutputTokens", () => {
it("uses a valid header value", () => {
expect(resolveMaxOutputTokens("32000")).toBe(32000)
expect(resolveMaxOutputTokens("32000", false)).toBe(32000)
expect(resolveMaxOutputTokens("32000", true)).toBe(32000)
})
it("falls back to the default for missing or bogus values", () => {
expect(resolveMaxOutputTokens(null)).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
expect(resolveMaxOutputTokens("")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
expect(resolveMaxOutputTokens("abc")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
expect(resolveMaxOutputTokens("0")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
expect(resolveMaxOutputTokens("-5")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
expect(resolveMaxOutputTokens("1.5")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
// Above the sanity ceiling, e.g. an extra zero
expect(resolveMaxOutputTokens("640000")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
for (const value of [null, "", "abc", "0", "-5", "1.5", "640000"]) {
// "640000" is above the sanity ceiling, e.g. an extra zero
expect(resolveMaxOutputTokens(value, false)).toBe(
DEFAULT_MAX_OUTPUT_TOKENS,
)
}
})
it("uses the env value when no header is sent, and validates it too", () => {
const original = process.env.MAX_OUTPUT_TOKENS
try {
process.env.MAX_OUTPUT_TOKENS = "24000"
expect(resolveMaxOutputTokens(null)).toBe(24000)
// Header still wins
expect(resolveMaxOutputTokens("8000")).toBe(8000)
expect(resolveMaxOutputTokens(null, true)).toBe(24000)
// A lower header still wins
expect(resolveMaxOutputTokens("8000", true)).toBe(8000)
process.env.MAX_OUTPUT_TOKENS = "-1"
expect(resolveMaxOutputTokens(null)).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
expect(resolveMaxOutputTokens(null, true)).toBe(
DEFAULT_MAX_OUTPUT_TOKENS,
)
} finally {
if (original === undefined) delete process.env.MAX_OUTPUT_TOKENS
else process.env.MAX_OUTPUT_TOKENS = original
}
})
it("lets the header raise the budget only on the user's own credentials", () => {
const original = process.env.MAX_OUTPUT_TOKENS
try {
process.env.MAX_OUTPUT_TOKENS = "16000"
expect(resolveMaxOutputTokens("200000", true)).toBe(16000)
expect(resolveMaxOutputTokens("200000", false)).toBe(200000)
// Without MAX_OUTPUT_TOKENS the default is the cap
delete process.env.MAX_OUTPUT_TOKENS
expect(resolveMaxOutputTokens("100000", true)).toBe(
DEFAULT_MAX_OUTPUT_TOKENS,
)
} finally {
if (original === undefined) delete process.env.MAX_OUTPUT_TOKENS
else process.env.MAX_OUTPUT_TOKENS = original
@@ -178,6 +349,34 @@ describe("withOutputTokenLimitFallback", () => {
expect(calls.map((c) => c.maxOutputTokens)).toEqual([64000, 4096])
})
it("subtracts the thinking budget from the retry", async () => {
const [model, calls] = fakeModel([
() =>
Promise.reject(
Object.assign(
new Error(
"The maximum tokens you requested exceeds the model limit of 64000.",
),
{ statusCode: 400 },
),
),
() => Promise.resolve(STREAM_OK),
])
const wrapped = withOutputTokenLimitFallback(model)
await wrapped.doStream({
prompt: [],
maxOutputTokens: 64000,
providerOptions: {
bedrock: {
reasoningConfig: { type: "enabled", budgetTokens: 12000 },
},
},
} as any)
expect(calls.map((c) => c.maxOutputTokens)).toEqual([64000, 52000])
})
it("does not retry an error it cannot attribute to the budget", async () => {
const [model, calls] = fakeModel([
() =>
+41
View File
@@ -0,0 +1,41 @@
// @vitest-environment node
import { afterEach, describe, expect, it, vi } from "vitest"
import { POST as parseUrl } from "@/app/api/parse-url/route"
// Treat every URL as public so no test hits DNS
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
isPrivateUrl: async () => false,
}))
afterEach(() => {
vi.unstubAllGlobals()
})
describe("POST /api/parse-url", () => {
it("stops the download of a page announced as too large", async () => {
let signal: AbortSignal | undefined
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init: RequestInit) => {
signal = init.signal ?? undefined
// A body that never ends unless the request is aborted
return new Response(new ReadableStream(), {
headers: {
"content-type": "text/html",
"content-length": String(50 * 1024 * 1024),
},
})
}),
)
const res = await parseUrl(
new Request("http://localhost/api/parse-url", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ url: "https://example.com/huge" }),
}),
)
expect(res.status).toBe(413)
expect(signal?.aborted).toBe(true)
})
})
+16
View File
@@ -0,0 +1,16 @@
import { describe, expect, it } from "vitest"
import { isTextFile } from "@/lib/pdf-utils"
describe("isTextFile", () => {
it("treats SVG files as text so their markup is sent to the model", () => {
const svg = new File(["<svg/>"], "diagram.svg", {
type: "image/svg+xml",
})
expect(isTextFile(svg)).toBe(true)
})
it("does not treat raster images as text", () => {
const png = new File(["x"], "photo.png", { type: "image/png" })
expect(isTextFile(png)).toBe(false)
})
})
+153
View File
@@ -0,0 +1,153 @@
// @vitest-environment node
import { mkdirSync, mkdtempSync } from "node:fs"
import { tmpdir } from "node:os"
import { join } from "node:path"
import { beforeEach, describe, expect, it, vi } from "vitest"
const userData = vi.hoisted(() => ({ dir: "" }))
vi.mock("electron", () => ({
app: { isPackaged: true, getPath: () => userData.dir },
}))
// Ports that fail to listen, with the error code
const busy = vi.hoisted(() => ({ ports: {} as Record<number, string> }))
vi.mock("node:net", () => ({
default: {
createServer: () => {
const handlers: Record<string, (arg?: unknown) => void> = {}
const server = {
once: (event: string, cb: (arg?: unknown) => void) => {
handlers[event] = cb
return server
},
listen: (port: number) => {
const code = busy.ports[port]
if (code) handlers.error?.({ code })
else handlers.listening?.()
},
close: () => {},
}
return server
},
},
}))
import {
findAvailablePort,
noteNoChats,
rememberChatPort,
resetAllocatedPort,
} from "@/electron/main/port-manager"
/** Chats saved under http://127.0.0.1:<port>, as Electron stores them */
const storeData = (port: number) =>
mkdirSync(
join(
userData.dir,
"IndexedDB",
`http_127.0.0.1_${port}.indexeddb.leveldb`,
),
{ recursive: true },
)
const launch = () => findAvailablePort(false)
beforeEach(() => {
userData.dir = mkdtempSync(join(tmpdir(), "port-manager-"))
busy.ports = {}
resetAllocatedPort()
})
describe("findAvailablePort", () => {
it("uses the legacy port first, as main does", async () => {
expect(await launch()).toBe(61337)
storeData(61337)
storeData(13370)
expect(await launch()).toBe(61337)
})
it("uses 13370 when only it has the user's chats", async () => {
// Windows reserved 61337 when they started using the app
storeData(13370)
expect(await launch()).toBe(13370)
})
it("goes back to the port with the chats once it is free", async () => {
// The previous version still quitting after an update
storeData(61337)
busy.ports[61337] = "EADDRINUSE"
expect(await launch()).toBe(13370)
storeData(13370)
busy.ports = {}
expect(await launch()).toBe(61337)
})
it("does not hide the chats for good after one reserved launch", async () => {
// Windows reserves port ranges per boot
storeData(61337)
busy.ports[61337] = "EACCES"
expect(await launch()).toBe(13370)
storeData(13370)
busy.ports = {}
expect(await launch()).toBe(61337)
})
it("falls back to the next ports", async () => {
storeData(13370)
busy.ports[13370] = "EADDRINUSE"
expect(await launch()).toBe(61337)
busy.ports[61337] = "EACCES"
expect(await launch()).toBe(13371)
})
})
describe("the port where chats were last saved", () => {
it("opens there first", async () => {
// Windows reserved 61337 for a while, and the user kept working
storeData(61337)
busy.ports[61337] = "EACCES"
expect(await launch()).toBe(13370)
storeData(13370)
rememberChatPort()
busy.ports = {}
expect(await launch()).toBe(13370)
})
it("does not move after a launch elsewhere that saved nothing", async () => {
storeData(61337)
expect(await launch()).toBe(61337)
rememberChatPort()
busy.ports[61337] = "EADDRINUSE"
expect(await launch()).toBe(13370)
storeData(13370)
busy.ports = {}
expect(await launch()).toBe(61337)
})
it("never stores a last-resort port, which changes between launches", async () => {
busy.ports[61337] = "EACCES"
busy.ports[13370] = "EADDRINUSE"
expect(await launch()).toBe(13371)
rememberChatPort()
busy.ports = {}
expect(await launch()).toBe(61337)
})
it("tries the other port after opening on one without chats", async () => {
// Split before this version: chats only on 13370, and a launch on
// 61337 created that origin's folder
storeData(13370)
storeData(61337)
expect(await launch()).toBe(61337)
noteNoChats()
expect(await launch()).toBe(13370)
// Once a choice is stored, an empty page changes nothing
noteNoChats()
expect(await launch()).toBe(13370)
})
it("stays put for a new user", async () => {
expect(await launch()).toBe(61337)
noteNoChats()
expect(await launch()).toBe(61337)
})
})
+314
View File
@@ -0,0 +1,314 @@
// @vitest-environment node
import { afterEach, describe, expect, it, vi } from "vitest"
// No DNS in tests: only loopback addresses are private
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
isPrivateUrl: async (url: string) =>
/^https?:\/\/(127\.0\.0\.1|localhost)\b/.test(url),
}))
import { POST as providerModels } from "@/app/api/provider-models/route"
import {
canListModels,
extractAihubmixModelIds,
listProviderModels,
} from "@/lib/provider-models"
afterEach(() => {
vi.unstubAllGlobals()
})
/** A fetch that answers with this JSON and records the request */
function answer(json: unknown, status = 200) {
const calls: Array<{ url: string; headers: Record<string, string> }> = []
const fn = vi.fn(async (url: string, init?: RequestInit) => {
calls.push({ url, headers: (init?.headers ?? {}) as any })
return new Response(JSON.stringify(json), { status })
}) as unknown as typeof fetch
return { fn, calls }
}
describe("listProviderModels", () => {
it("reads an OpenAI-style list and drops models that are not for chat", async () => {
const { fn, calls } = answer({
data: [
{ id: "gpt-4.1" },
{ id: "text-embedding-3-small" },
{ id: "whisper-1" },
{ id: "gpt-image-1" },
],
})
const models = await listProviderModels("openai", { apiKey: "k" }, fn)
expect(models.map((m) => m.id)).toEqual(["gpt-4.1"])
// Tool support comes from models.dev when the list has none
expect(models[0].tools).toBe(true)
expect(calls[0].url).toBe("https://api.openai.com/v1/models")
expect(calls[0].headers.Authorization).toBe("Bearer k")
})
it("uses the base URL the user gave, without a pasted path", async () => {
const { fn, calls } = answer({ data: [{ id: "m" }] })
await listProviderModels(
"glm",
{
apiKey: "k",
baseUrl: "https://proxy.example.com/v4/chat/completions",
},
fn,
)
expect(calls[0].url).toBe("https://proxy.example.com/v4/models")
})
it("asks Anthropic with its own headers", async () => {
const { fn, calls } = answer({ data: [{ id: "claude-sonnet-4-5" }] })
await listProviderModels("anthropic", { apiKey: "k" }, fn)
expect(calls[0].url).toBe(
"https://api.anthropic.com/v1/models?limit=1000",
)
expect(calls[0].headers["x-api-key"]).toBe("k")
})
it("keeps Gemini models that generate content, without models/", async () => {
const { fn, calls } = answer({
models: [
{
name: "models/gemini-2.5-flash",
supportedGenerationMethods: ["generateContent"],
},
{
name: "models/text-embedding-004",
supportedGenerationMethods: ["embedContent"],
},
],
})
const models = await listProviderModels("google", { apiKey: "k" }, fn)
expect(models.map((m) => m.id)).toEqual(["gemini-2.5-flash"])
// The key is a header, not part of the URL
expect(calls[0].url).not.toContain("k&")
expect(calls[0].headers["x-goog-api-key"]).toBe("k")
})
it("reads Ollama's tags and OpenRouter's tool support", async () => {
const ollama = answer({ models: [{ name: "llama3.2" }] })
await listProviderModels(
"ollama",
{ baseUrl: "http://localhost:11434" },
ollama.fn,
)
expect(ollama.calls[0].url).toBe("http://localhost:11434/api/tags")
const openrouter = answer({
data: [
{ id: "a/with-tools", supported_parameters: ["tools"] },
{ id: "b/no-tools", supported_parameters: ["temperature"] },
],
})
const models = await listProviderModels("openrouter", {}, openrouter.fn)
expect(models).toEqual([
{ id: "a/with-tools", tools: true },
{ id: "b/no-tools", tools: false },
])
})
it("lists Ollama from where chat goes without a base URL", async () => {
const { fn, calls } = answer({ models: [{ name: "llama3.2" }] })
process.env.OLLAMA_BASE_URL = "http://ollama.internal:11434"
try {
await listProviderModels("ollama", {}, fn)
} finally {
delete process.env.OLLAMA_BASE_URL
}
await listProviderModels("ollama", {}, fn)
expect(calls.map((c) => c.url)).toEqual([
"http://ollama.internal:11434/api/tags",
"http://127.0.0.1:11434/api/tags",
])
})
it("lists Ollama Cloud with a user's key, like chat", async () => {
// The user's key must not go to the server's Ollama
const { fn, calls } = answer({ models: [{ name: "gpt-oss:120b" }] })
process.env.OLLAMA_BASE_URL = "http://ollama.internal:11434"
try {
await listProviderModels("ollama", { apiKey: "user-key" }, fn)
} finally {
delete process.env.OLLAMA_BASE_URL
}
expect(calls[0].url).toBe("https://ollama.com/api/tags")
})
it("does not use SGLang's local address as a default", async () => {
const { fn, calls } = answer({ data: [] })
await expect(
listProviderModels("sglang", { apiKey: "k" }, fn),
).rejects.toThrow(/base URL/)
expect(calls).toHaveLength(0)
})
it("turns a failed request into an error with its status", async () => {
const { fn } = answer({ error: "bad key" }, 401)
await expect(
listProviderModels("deepseek", { apiKey: "k" }, fn),
).rejects.toMatchObject({ statusCode: 401 })
})
})
describe("extractAihubmixModelIds", () => {
it("keeps unique chat models", () => {
expect(
extractAihubmixModelIds({
data: [
{ model_id: "claude-sonnet-4-5", types: "llm" },
{ model_id: "gpt-5.1", types: "llm" },
{ model_id: "gpt-5.1", types: "llm" },
{ model_id: "gpt-image-2", types: "image_generation,llm" },
{ model_id: "", types: "llm" },
],
}),
).toEqual(["claude-sonnet-4-5", "gpt-5.1"])
expect(extractAihubmixModelIds({ data: null })).toEqual([])
})
})
describe("POST /api/provider-models", () => {
const post = (body: unknown) =>
providerModels(
new Request("http://localhost/api/provider-models", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(body),
}),
)
it("answers null for providers that cannot list models", async () => {
expect(canListModels("bedrock")).toBe(false)
expect(canListModels("toString" as never)).toBe(false)
const res = await post({ provider: "bedrock" })
expect(await res.json()).toEqual({ models: null })
})
it("needs the user's key where the list is not public", async () => {
const res = await post({ provider: "deepseek" })
expect(res.status).toBe(400)
})
it("explains a failure with the error hints", async () => {
vi.stubGlobal("fetch", answer({}, 401).fn)
const res = await post({ provider: "deepseek", apiKey: "k" })
expect(await res.json()).toMatchObject({ code: "invalid_api_key" })
})
// The base URL is the caller's, and private addresses are allowed by
// default (local Ollama), so the answer must not reveal what an
// internal address sent back
const text = (body: string) =>
vi.fn(
async () => new Response(body, { status: 200 }),
) as unknown as typeof fetch
it("does not repeat a body that is not JSON", async () => {
vi.stubGlobal("fetch", text("ROLE-NAME-OF-THE-SERVER"))
const res = await post({
provider: "ollama",
baseUrl: "http://169.254.169.254/latest/meta-data/x?",
})
const data = await res.json()
expect(data.error).toBe("The model list was not valid JSON.")
expect(JSON.stringify(data)).not.toContain("ROLE")
})
it("stops reading a list over 2 MB, also through the Gateway SDK", async () => {
const huge = JSON.stringify({ data: [{ id: "x".repeat(3_000_000) }] })
for (const body of [
{ provider: "ollama", baseUrl: "https://big.example.com" },
{
provider: "gateway",
apiKey: "k",
baseUrl: "https://big.example.com/v3/ai",
},
]) {
vi.stubGlobal("fetch", text(huge))
const data = await (await post(body)).json()
expect(data.error).toBe("The model list is too large.")
expect(data.models).toBeUndefined()
}
})
it("ends the download of a list that is too large", async () => {
// The answer announces 4 MB and never finishes
let signal: AbortSignal | undefined
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init?: RequestInit) => {
signal = init?.signal ?? undefined
const body = new ReadableStream({ start() {} })
return new Response(body, {
headers: { "content-length": String(4 * 1024 * 1024) },
})
}),
)
const data = await (
await post({
provider: "ollama",
baseUrl: "https://big.example.com",
})
).json()
expect(data.error).toBe("The model list is too large.")
expect(signal?.aborted).toBe(true)
})
it("handles answers without a body", async () => {
for (const status of [204, 304]) {
vi.stubGlobal(
"fetch",
vi.fn(async () => new Response(null, { status })),
)
const data = await (
await post({ provider: "ollama", baseUrl: "https://x.example" })
).json()
expect(data.error).toMatch(/not valid JSON|failed \(304\)/)
}
})
it("explains a refused redirect", async () => {
process.env.ALLOW_PRIVATE_URLS = "false"
try {
vi.stubGlobal(
"fetch",
vi.fn(
async () =>
new Response(null, {
status: 301,
headers: { location: "https://elsewhere.example" },
}),
),
)
const data = await (
await post({ provider: "ollama", baseUrl: "https://x.example" })
).json()
expect(data.error).toMatch(/Redirects are not allowed/)
} finally {
delete process.env.ALLOW_PRIVATE_URLS
}
})
it("keeps its own explanations and hides other error texts", async () => {
// Our own: no base URL for SGLang
const own = await (
await post({ provider: "sglang", apiKey: "k" })
).json()
expect(own.error).toMatch(/needs a base URL/)
// Not ours: an exception text from the network layer
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw new Error("connect ECONNREFUSED 10.1.2.3:8080")
}),
)
const other = await (
await post({ provider: "ollama", baseUrl: "http://10.1.2.3:8080" })
).json()
expect(other.error).toBe("The model list request failed.")
})
})
+48
View File
@@ -4,6 +4,7 @@ import {
loadFlattenedServerModels,
type ServerModelsConfig,
ServerModelsConfigSchema,
slugify,
} from "@/lib/server-model-config"
const ORIGINAL_ENV = { ...process.env }
@@ -233,3 +234,50 @@ describe("loadFlattenedServerModels", () => {
expect(models[0].apiKeyEnv).toEqual(["OPENAI_KEY_1", "OPENAI_KEY_2"])
})
})
describe("slugify", () => {
it("keeps ASCII names readable", () => {
expect(slugify("OpenAI Production")).toBe("openai-production")
})
it("gives distinct ASCII slugs to distinct CJK names", () => {
const slugs = ["主力", "备用", "DeepSeek 官方", "DeepSeek 备用"].map(
slugify,
)
expect(new Set(slugs).size).toBe(4)
for (const slug of slugs) expect(slug).toMatch(/^[a-z0-9-]+$/)
})
})
describe("loadFlattenedServerModels id collisions", () => {
it("drops a model whose id repeats an earlier provider's", async () => {
const config: ServerModelsConfig = {
providers: [
{ name: "OpenAI", provider: "openai", models: ["gpt-4o"] },
{
name: "openai",
provider: "openai",
models: ["gpt-4o"],
apiKeyEnv: "OTHER_KEY",
},
{
name: "主力",
provider: "deepseek",
models: ["deepseek-chat"],
},
{
name: "备用",
provider: "deepseek",
models: ["deepseek-chat"],
},
],
}
process.env.AI_MODELS_CONFIG = JSON.stringify(config)
const models = await loadFlattenedServerModels()
const ids = models.map((m) => m.id)
expect(new Set(ids).size).toBe(ids.length)
expect(ids).toHaveLength(3)
expect(models[0].apiKeyEnv).toBeUndefined()
})
})
+131 -2
View File
@@ -1,5 +1,5 @@
import { beforeEach, describe, expect, it, vi } from "vitest"
import { isPrivateUrl } from "@/lib/ssrf-protection"
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { isPrivateUrl, redirectGuardedFetch } from "@/lib/ssrf-protection"
// Mock DNS so tests are deterministic and never hit the network.
const lookupMock = vi.hoisted(() => vi.fn())
@@ -77,3 +77,132 @@ describe("isPrivateUrl", () => {
)
})
})
describe("redirectGuardedFetch with the quota on", () => {
const answers = (map: Record<string, Response>) =>
vi.fn(
async (url: string) =>
map[String(url)] ?? new Response("?", { status: 404 }),
)
beforeEach(() => {
lookupMock.mockReset()
// Hosts ending in .example are public
lookupMock.mockImplementation(async (host: string) =>
host.endsWith(".example")
? [{ address: "93.184.216.34", family: 4 }]
: [],
)
process.env.DYNAMODB_QUOTA_TABLE = "quota"
delete process.env.ALLOW_PRIVATE_URLS
})
afterEach(() => {
delete process.env.DYNAMODB_QUOTA_TABLE
vi.unstubAllGlobals()
})
it("follows a redirect to a public address", async () => {
// A user's own proxy that moves http to https
vi.stubGlobal(
"fetch",
answers({
"http://proxy.example/v1/chat": new Response(null, {
status: 308,
headers: { location: "https://proxy.example/v1/chat" },
}),
"https://proxy.example/v1/chat": new Response("ok"),
}),
)
const guarded = redirectGuardedFetch()
expect(guarded).toBeDefined()
const res = await guarded?.("http://proxy.example/v1/chat", {
method: "POST",
body: "{}",
})
expect(await res?.text()).toBe("ok")
})
it("refuses a redirect to the server's own network", async () => {
// It would be counted as a public endpoint while using the server's
vi.stubGlobal(
"fetch",
answers({
"https://public.example/api/chat": new Response(null, {
status: 307,
headers: { location: "http://127.0.0.1:11434/api/chat" },
}),
}),
)
await expect(
redirectGuardedFetch()?.("https://public.example/api/chat", {
method: "POST",
body: "{}",
}),
).rejects.toThrow(/private addresses/)
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",
"x-api-key": "anthropic-key",
"x-goog-api-key": "google-key",
"api-key": "azure-key",
"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()
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")
})
it("is not used without the quota", () => {
delete process.env.DYNAMODB_QUOTA_TABLE
expect(redirectGuardedFetch()).toBeUndefined()
})
})
+45
View File
@@ -0,0 +1,45 @@
import { cleanup, render } from "@testing-library/react"
import { afterEach, describe, expect, it } from "vitest"
import { ToolCallCard } from "@/components/chat/ToolCallCard"
afterEach(cleanup)
const dict = {
tools: { complete: "Complete" },
chat: { copied: "Copied", failedToCopy: "Failed", copyResponse: "Copy" },
}
describe("ToolCallCard", () => {
it("shows streamed operations without crashing on broken entries", () => {
// A partly streamed or invalid edit_diagram input
const operations = [
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"/>' },
]
const { container } = render(
<ToolCallCard
part={
{
type: "tool-edit_diagram",
toolCallId: "t1",
state: "input-streaming",
input: { operations },
} as any
}
expandedTools={{ t1: true }}
setExpandedTools={() => {}}
onCopy={() => {}}
copiedToolCallId={null}
copyFailedToolCallId={null}
dict={dict}
/>,
)
expect(container.textContent).toContain("cell_id: 3")
expect(container.textContent).toContain('<mxCell id="3"/>')
})
})
+25
View File
@@ -0,0 +1,25 @@
import { afterEach, describe, expect, it, vi } from "vitest"
import { STORAGE_KEYS } from "@/lib/storage"
import { extractUrlContent } from "@/lib/url-utils"
describe("extractUrlContent", () => {
afterEach(() => {
vi.unstubAllGlobals()
localStorage.clear()
})
it("sends the saved access code with the request", async () => {
localStorage.setItem(STORAGE_KEYS.accessCode, "secret")
const body = { title: "T", content: "body", charCount: 4 }
const fetchMock = vi
.fn()
.mockResolvedValue(new Response(JSON.stringify(body)))
vi.stubGlobal("fetch", fetchMock)
const data = await extractUrlContent("https://example.com")
expect(data.content).toBe("body")
const headers = fetchMock.mock.calls[0][1].headers
expect(headers["x-access-code"]).toBe("secret")
})
})
@@ -0,0 +1,204 @@
import { renderHook } from "@testing-library/react"
import { describe, expect, it, vi } from "vitest"
import { useDiagramToolHandlers } from "@/hooks/use-diagram-tool-handlers"
const geometry =
'<mxGeometry x="0" y="0" width="80" height="40" as="geometry"/>'
const box = (id: string) =>
`<mxCell id="${id}" value="${id}" vertex="1" parent="1">${geometry}</mxCell>`
function setup(partialXml: string) {
const refs = {
partialXmlRef: { current: partialXml },
// A failed edit's preview is still on the canvas, its original kept
editDiagramOriginalXmlRef: {
current: new Map([["edit-1", "<mxfile>original</mxfile>"]]),
},
processedToolCallsRef: { current: new Set<string>() },
validationRetryCountRef: { current: 0 },
chartXMLRef: { current: "" },
}
const onDisplayChart = vi.fn(
(_xml: string, _skipValidation?: boolean): string | null => null,
)
const { result } = renderHook(() =>
useDiagramToolHandlers({
...refs,
onDisplayChart,
onFetchChart: async () => "",
onExport: () => {},
enableVlmValidation: false,
}),
)
const addToolOutput = vi.fn()
const append = (xml: string) =>
result.current.handleToolCall(
{
toolCall: {
toolCallId: "append-1",
toolName: "append_diagram",
input: { xml },
},
},
addToolOutput,
)
return { refs, onDisplayChart, addToolOutput, append }
}
describe("the screenshot check and Stop", () => {
const draw = async (opts: {
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(() =>
useDiagramToolHandlers({
partialXmlRef: { current: "" },
editDiagramOriginalXmlRef: { current: new Map() },
processedToolCallsRef: { current: new Set() },
validationRetryCountRef: opts.retryCount ?? { current: 0 },
chartXMLRef: { current: "" },
onDisplayChart: () => null,
onFetchChart: async () => "",
onExport: () => {},
enableVlmValidation: true,
captureValidationPng:
opts.captureValidationPng ??
(async () => "data:image/png;base64,AA"),
validateDiagram: opts.validateDiagram,
watchStop: opts.watchStop,
onValidationStateChange,
}),
)
const addToolOutput = vi.fn()
await result.current.handleToolCall(
{
toolCall: {
toolCallId: "d1",
toolName: "display_diagram",
input: { xml: box("2") },
},
},
addToolOutput,
)
return { addToolOutput, onValidationStateChange }
}
it("skips a check that had not started when the user stopped", async () => {
const validateDiagram = vi.fn(async () => ({
valid: true,
issues: [],
suggestions: [],
}))
const { addToolOutput, onValidationStateChange } = await draw({
watchStop: () => () => true,
validateDiagram,
})
expect(validateDiagram).not.toHaveBeenCalled()
expect(onValidationStateChange.mock.lastCall?.[1].status).toBe(
"skipped",
)
expect(addToolOutput.mock.lastCall?.[0].output).toMatch(
/Successfully displayed/,
)
})
it("ends with the diagram's result when Stop cancels a running check", async () => {
const { addToolOutput, onValidationStateChange } = await draw({
watchStop: () => () => false,
validateDiagram: async () => {
throw new DOMException("Validation cancelled", "AbortError")
},
})
expect(onValidationStateChange.mock.lastCall?.[1].status).toBe(
"skipped",
)
expect(addToolOutput).toHaveBeenCalledTimes(1)
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
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", () => {
it("takes the stored originals when it draws the completed diagram", async () => {
// Otherwise the preview code later loads the failed edit's original
// over the completed diagram
const { refs, onDisplayChart, append } = setup(
`${box("2")}<mxCell id="3" value="3" vertex="1" parent="1"><mxGeometry x="0" y="0" width="8`,
)
await append('0" height="40" as="geometry"/></mxCell>')
expect(onDisplayChart).toHaveBeenCalledTimes(1)
expect(onDisplayChart.mock.calls[0][0]).toContain('id="3"')
expect(refs.editDiagramOriginalXmlRef.current.size).toBe(0)
expect(refs.processedToolCallsRef.current.has("edit-1")).toBe(true)
})
it("leaves them while the diagram is still incomplete", async () => {
// Nothing is drawn, so the failed edit's preview must still be undone
const { refs, onDisplayChart, append } = setup(
`${box("2")}<mxCell id="3" value="3" vertex="1" parent="1"><mxGeometry x="0" y="0" width="8`,
)
await append('0" height="40"')
expect(onDisplayChart).not.toHaveBeenCalled()
expect(refs.editDiagramOriginalXmlRef.current.size).toBe(1)
expect(refs.processedToolCallsRef.current.has("edit-1")).toBe(false)
})
it("leaves them when the assembled diagram is invalid", async () => {
const { refs, onDisplayChart, addToolOutput, append } = setup(
`<mxCell id="1" value="root id" vertex="1" parent="1">${geometry}</mxCell><mxCell id="3" value="3" vertex="1" parent="1"><mxGeometry x="0" y="0" width="8`,
)
await append('0" height="40" as="geometry"/></mxCell>')
expect(onDisplayChart).not.toHaveBeenCalled()
expect(addToolOutput.mock.calls[0][0].state).toBe("output-error")
expect(refs.editDiagramOriginalXmlRef.current.size).toBe(1)
})
})
+102
View File
@@ -0,0 +1,102 @@
import { act, renderHook } from "@testing-library/react"
import { beforeEach, describe, expect, it, vi } from "vitest"
import { extractPdfText, extractTextFileContent } from "@/lib/pdf-utils"
import { useFileProcessor } from "@/lib/use-file-processor"
vi.mock("sonner", () => ({ toast: { error: vi.fn() } }))
vi.mock("@/lib/pdf-utils", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/pdf-utils")>()),
extractPdfText: vi.fn(),
extractTextFileContent: vi.fn(),
}))
// A promise we can resolve from the test, to control extraction timing
function deferred<T>() {
let resolve!: (value: T) => void
const promise = new Promise<T>((r) => {
resolve = r
})
return { promise, resolve }
}
const pdfFile = () =>
new File(["%PDF"], "slow.pdf", { type: "application/pdf" })
const textFile = () => new File(["notes"], "notes.txt", { type: "text/plain" })
describe("useFileProcessor", () => {
beforeEach(() => {
vi.mocked(extractPdfText).mockReset()
vi.mocked(extractTextFileContent).mockReset()
})
it("marks queued files as extracting before the first one finishes", async () => {
const pdf = deferred<string>()
vi.mocked(extractPdfText).mockReturnValue(pdf.promise)
vi.mocked(extractTextFileContent).mockResolvedValue("notes")
const a = pdfFile()
const b = textFile()
const { result } = renderHook(() => useFileProcessor())
let done!: Promise<void>
act(() => {
done = result.current.handleFileChange([a, b])
})
expect(result.current.pdfData.get(a)?.isExtracting).toBe(true)
expect(result.current.pdfData.get(b)?.isExtracting).toBe(true)
await act(async () => {
pdf.resolve("pdf text")
await done
})
expect(result.current.pdfData.get(a)?.text).toBe("pdf text")
expect(result.current.pdfData.get(b)?.text).toBe("notes")
})
it("keeps text of a file added while an earlier file is extracting", async () => {
const pdf = deferred<string>()
vi.mocked(extractPdfText).mockReturnValue(pdf.promise)
vi.mocked(extractTextFileContent).mockResolvedValue("notes")
const a = pdfFile()
const b = textFile()
const { result } = renderHook(() => useFileProcessor())
let first!: Promise<void>
act(() => {
first = result.current.handleFileChange([a])
})
await act(async () => {
await result.current.handleFileChange([a, b])
})
expect(result.current.pdfData.get(b)?.text).toBe("notes")
await act(async () => {
pdf.resolve("pdf text")
await first
})
expect(result.current.pdfData.get(a)?.text).toBe("pdf text")
expect(result.current.pdfData.get(b)?.text).toBe("notes")
})
it("does not bring back a file removed while extracting", async () => {
const pdf = deferred<string>()
vi.mocked(extractPdfText).mockReturnValue(pdf.promise)
const a = pdfFile()
const { result } = renderHook(() => useFileProcessor())
let first!: Promise<void>
act(() => {
first = result.current.handleFileChange([a])
})
await act(async () => {
await result.current.handleFileChange([])
})
await act(async () => {
pdf.resolve("pdf text")
await first
})
expect(result.current.pdfData.has(a)).toBe(false)
expect(result.current.files).toEqual([])
})
})
+262
View File
@@ -0,0 +1,262 @@
import { act, cleanup, renderHook, waitFor } from "@testing-library/react"
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { getSelectedAIConfig, useModelConfig } from "@/hooks/use-model-config"
import type { FlattenedServerModel } from "@/lib/server-model-config"
import { STORAGE_KEYS } from "@/lib/storage"
import type { MultiModelConfig } from "@/lib/types/model-config"
const SERVER_MODELS: FlattenedServerModel[] = [
{
id: "server:openai-main:gpt-4o-mini",
modelId: "gpt-4o-mini",
provider: "openai",
providerLabel: "OpenAI Main",
isDefault: false,
},
{
id: "server:openai-main:gpt-4o",
modelId: "gpt-4o",
provider: "openai",
providerLabel: "OpenAI Main",
isDefault: true,
},
]
const USER_CONFIG: MultiModelConfig = {
version: 1,
providers: [
{
id: "p1",
provider: "openai",
apiKey: "sk-test",
models: [{ id: "m1", modelId: "gpt-4o" }],
},
],
}
function storeConfig(config: MultiModelConfig) {
localStorage.setItem(STORAGE_KEYS.modelConfigs, JSON.stringify(config))
}
async function renderLoaded() {
const hook = renderHook(() => useModelConfig())
await waitFor(() => expect(hook.result.current.isLoaded).toBe(true))
return hook
}
beforeEach(() => {
localStorage.clear()
vi.stubGlobal(
"fetch",
vi.fn(async () => ({
ok: true,
json: async () => ({ models: SERVER_MODELS }),
})),
)
})
afterEach(() => {
cleanup()
vi.unstubAllGlobals()
})
describe("useModelConfig server model selection", () => {
it("replaces a saved server model that no longer exists", async () => {
storeConfig({
...USER_CONFIG,
selectedModelId: "server:openai-production:gpt-4o",
})
const { result } = await renderLoaded()
expect(result.current.selectedModelId).toBe("server:openai-main:gpt-4o")
})
it("keeps a server model saved under its old id", async () => {
// Before non-ASCII provider names got their own slug, "主力 OpenAI"
// became "openai"
const renamed: FlattenedServerModel = {
id: "server:4e3b-529b-openai:gpt-4o-mini",
modelId: "gpt-4o-mini",
provider: "openai",
providerLabel: "主力 OpenAI",
isDefault: false,
}
vi.stubGlobal(
"fetch",
vi.fn(async () => ({
ok: true,
json: async () => ({ models: [...SERVER_MODELS, renamed] }),
})),
)
storeConfig({
...USER_CONFIG,
selectedModelId: "server:openai:gpt-4o-mini",
})
const { result } = await renderLoaded()
expect(result.current.selectedModelId).toBe(renamed.id)
})
it("keeps a saved server model that still exists", async () => {
storeConfig({
...USER_CONFIG,
selectedModelId: "server:openai-main:gpt-4o-mini",
})
const { result } = await renderLoaded()
expect(result.current.selectedModelId).toBe(
"server:openai-main:gpt-4o-mini",
)
})
it("skips a saved provider this version does not know", async () => {
// Saved by another version, or edited by hand: it used to crash the
// whole page on load
storeConfig({
...USER_CONFIG,
providers: [
...USER_CONFIG.providers,
{
id: "p9",
provider: "not-a-provider" as any,
apiKey: "k",
models: [{ id: "m9", modelId: "x" }],
},
],
})
const { result } = await renderLoaded()
expect(result.current.config.providers.map((p) => p.id)).toEqual(["p1"])
expect(result.current.models.map((m) => m.id)).toContain("m1")
})
it("keeps an unknown provider and its key in storage", async () => {
// The version that saved it may be opened again (an older desktop
// build, another tab): the provider must still be there
storeConfig({
...USER_CONFIG,
providers: [
...USER_CONFIG.providers,
{
id: "p9",
provider: "not-a-provider" as any,
apiKey: "k9",
models: [{ id: "m9", modelId: "x" }],
},
],
selectedModelId: "m1",
})
const { result } = await renderLoaded()
act(() => result.current.setSelectedModelId(undefined))
await waitFor(() => {
const stored = JSON.parse(
localStorage.getItem(STORAGE_KEYS.modelConfigs) ?? "{}",
)
expect(stored.selectedModelId).toBeUndefined()
expect(stored.providers.map((p: { id: string }) => p.id)).toEqual([
"p1",
"p9",
])
expect(stored.providers[1].apiKey).toBe("k9")
})
// Sending reads the stored config too, and must not trip over it
act(() => result.current.setSelectedModelId("m1"))
expect(getSelectedAIConfig()).toMatchObject({
aiProvider: "openai",
aiModel: "gpt-4o",
})
})
it("keeps a selected user model", async () => {
storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
const { result } = await renderLoaded()
expect(result.current.selectedModelId).toBe("m1")
})
it("falls back to the default server model when the selected model is deleted", async () => {
storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
const { result } = await renderLoaded()
act(() => result.current.deleteModel("p1", "m1"))
expect(result.current.selectedModelId).toBe("server:openai-main:gpt-4o")
})
it("falls back to the default server model when the selected provider is deleted", async () => {
storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
const { result } = await renderLoaded()
act(() => result.current.deleteProvider("p1"))
expect(result.current.selectedModelId).toBe("server:openai-main:gpt-4o")
})
})
describe("useModelConfig in the desktop app", () => {
it("reloads the server models after a preset switch restarts the server", async () => {
// The new preset offers other models; the saved one is gone
let restarted: (() => void) | undefined
;(window as any).electronAPI = {
onServerRestarted: (callback: () => void) => {
restarted = callback
return () => {
restarted = undefined
}
},
}
try {
storeConfig({
...USER_CONFIG,
selectedModelId: "server:openai-main:gpt-4o-mini",
})
const { result } = await renderLoaded()
await waitFor(() => expect(restarted).toBeDefined())
const nextModels: FlattenedServerModel[] = [
{
id: "server:claude:claude-sonnet-5-5",
modelId: "claude-sonnet-5-5",
provider: "anthropic",
providerLabel: "Claude",
isDefault: true,
},
]
vi.stubGlobal(
"fetch",
vi.fn(async () => ({
ok: true,
json: async () => ({ models: nextModels }),
})),
)
act(() => restarted?.())
await waitFor(() =>
expect(result.current.selectedModelId).toBe(
"server:claude:claude-sonnet-5-5",
),
)
} finally {
delete (window as any).electronAPI
}
})
})
describe("useModelConfig across tabs", () => {
it("reloads the config when another tab saves it", async () => {
storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
const { result } = await renderLoaded()
const fromOtherTab: MultiModelConfig = {
...USER_CONFIG,
providers: [
...USER_CONFIG.providers,
{
id: "p2",
provider: "anthropic",
apiKey: "sk-ant",
models: [{ id: "m2", modelId: "claude-sonnet-4-5" }],
},
],
selectedModelId: "m2",
}
act(() => {
storeConfig(fromOtherTab)
window.dispatchEvent(
new StorageEvent("storage", { key: STORAGE_KEYS.modelConfigs }),
)
})
expect(result.current.selectedModelId).toBe("m2")
expect(result.current.config.providers).toHaveLength(2)
})
})
+232
View File
@@ -0,0 +1,232 @@
import { act, renderHook, waitFor } from "@testing-library/react"
import { beforeEach, describe, expect, it, vi } from "vitest"
const stored = new Map<string, any>()
// Each write waits for the test to let it finish
let pendingWrites: (() => void)[] = []
// Reads wait for this
let readGate: Promise<void> = Promise.resolve()
vi.mock("@/hooks/use-dictionary", () => ({
useDictionary: () => ({ errors: { sessionSaveFailed: "Not saved" } }),
}))
vi.mock("@/lib/session-storage", async (importActual) => {
const actual = await importActual<typeof import("@/lib/session-storage")>()
return {
createEmptySession: actual.createEmptySession,
extractTitle: actual.extractTitle,
isIndexedDBAvailable: () => true,
migrateFromLocalStorage: async () => null,
readSessionCount: async () => stored.size,
enforceSessionLimit: async () => {},
getSession: async (id: string) => {
await readGate
return stored.get(id) ?? null
},
deleteSession: async (id: string) => {
stored.delete(id)
},
getAllSessionMetadata: async () =>
[...stored.values()].map((s) => ({ id: s.id, title: s.title })),
saveSession: vi.fn(
(session: any) =>
new Promise<boolean>((resolve) => {
pendingWrites.push(() => {
stored.set(session.id, session)
resolve(true)
})
}),
),
}
})
import { useSessionManager } from "@/hooks/use-session-manager"
const data = {
messages: [
{
id: "m1",
role: "user" as const,
parts: [{ type: "text", text: "Draw a cat" }],
},
],
xmlSnapshots: [] as [number, string][],
diagramXml: "",
}
// Let every write waiting now (and those it leads to) finish
async function finishWrites() {
for (let i = 0; i < 10; i++) {
await act(async () => {
const writes = pendingWrites
pendingWrites = []
for (const finish of writes) finish()
await new Promise((r) => setTimeout(r, 0))
})
}
}
async function setup() {
const hook = renderHook(() => useSessionManager())
await waitFor(() => expect(hook.result.current.isLoading).toBe(false))
return hook
}
describe("saving the chat on screen", () => {
beforeEach(() => {
stored.clear()
pendingWrites = []
})
it("creates one session when two saves of a new chat overlap", async () => {
const { result } = await setup()
let saves!: Promise<boolean[]>
act(() => {
saves = Promise.all([
result.current.saveCurrentSession(data),
result.current.saveCurrentSession(data),
])
})
await finishWrites()
expect(await saves).toEqual([true, true])
expect(stored.size).toBe(1)
expect(result.current.currentSessionId).toBe([...stored.keys()][0])
})
it("drops a save scheduled before New Chat", async () => {
const { result } = await setup()
const scheduled = result.current.getSaveTicket()
act(() => result.current.clearCurrentSession())
let save!: Promise<boolean>
act(() => {
save = result.current.saveCurrentSession(data, scheduled)
})
await finishWrites()
expect(await save).toBe(true)
expect(stored.size).toBe(0)
})
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.getSaveTicket()
let newChatSave!: Promise<boolean>
let autoSave!: Promise<boolean>
act(() => {
newChatSave = result.current.saveCurrentSession(data)
autoSave = result.current.saveCurrentSession(data, scheduled)
})
await act(async () => {
await waitFor(() => expect(pendingWrites).toHaveLength(1))
pendingWrites.shift()?.()
await newChatSave
result.current.clearCurrentSession()
})
await finishWrites()
await autoSave
expect(stored.size).toBe(1)
expect(result.current.currentSessionId).toBeNull()
})
it("keeps the new blank chat when a save of the old one ends later", async () => {
const { result } = await setup()
let save!: Promise<boolean>
act(() => {
save = result.current.saveCurrentSession(data)
})
await waitFor(() => expect(pendingWrites).toHaveLength(1))
// New Chat while the write runs
act(() => result.current.clearCurrentSession())
await finishWrites()
await save
expect(stored.size).toBe(1)
expect(result.current.currentSessionId).toBeNull()
expect(result.current.currentSession).toBeNull()
})
it("keeps New Chat when the URL's chat finishes loading after it", async () => {
stored.set("s1", { id: "s1", title: "Old", messages: [] })
const hook = renderHook(
({ id }: { id: string | null }) =>
useSessionManager({ initialSessionId: id }),
{ initialProps: { id: null as string | null } },
)
await waitFor(() => expect(hook.result.current.isLoading).toBe(false))
// The new chat's id reaches the URL; reading it takes a moment
let release!: () => void
readGate = new Promise((r) => {
release = r
})
hook.rerender({ id: "s1" })
act(() => hook.result.current.clearCurrentSession())
await act(async () => {
release()
await new Promise((r) => setTimeout(r, 0))
})
readGate = Promise.resolve()
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)
})
})
+81 -31
View File
@@ -1,5 +1,37 @@
import { describe, expect, it } from "vitest"
import { cn, isMxCellXmlComplete, wrapWithMxFile } from "@/lib/utils"
import {
cn,
extractCompleteMxCells,
isMxCellXmlComplete,
isRealDiagram,
} from "@/lib/utils"
import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts"
describe("isRealDiagram", () => {
it("counts a small diagram with one shape", () => {
// 234 characters: valid, shown, and saved with its chat
const xml =
'<mxfile><diagram id="p"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="A" vertex="1" parent="1"><mxGeometry width="80" height="30" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>'
expect(xml.length).toBeLessThan(300)
expect(isRealDiagram(xml)).toBe(true)
})
it("does not count a blank page or nothing", () => {
expect(isRealDiagram(BLANK_MXFILE)).toBe(false)
expect(isRealDiagram("")).toBe(false)
expect(isRealDiagram(null)).toBe(false)
})
it("still counts a longer document of empty named pages", () => {
// Pages and page settings are worth keeping, as before
const pages = Array.from(
{ length: 3 },
(_, i) =>
`<diagram id="p${i}" name="Page ${i}"><mxGraphModel pageWidth="1600" pageHeight="900"><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram>`,
).join("")
expect(isRealDiagram(`<mxfile>${pages}</mxfile>`)).toBe(true)
})
})
describe("isMxCellXmlComplete", () => {
it("returns false for empty/null input", () => {
@@ -12,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", () => {
@@ -33,6 +71,31 @@ describe("isMxCellXmlComplete", () => {
expect(isMxCellXmlComplete(xml)).toBe(false)
})
it("returns false when output stops after a child of an open mxCell", () => {
const xml = `<mxCell id="2" value="A" vertex="1" parent="1">
<mxGeometry x="0" y="0" width="80" height="40" as="geometry"/>
</mxCell>
<mxCell id="3" value="B" vertex="1" parent="1">
<mxGeometry x="100" y="0" width="80" height="40" as="geometry"/>`
expect(isMxCellXmlComplete(xml)).toBe(false)
})
it("returns false when output stops after </mxGeometry> of an open mxCell", () => {
const xml = `<mxCell id="e1" edge="1" parent="1" source="2" target="3">
<mxGeometry relative="1" as="geometry">
<mxPoint x="10" y="10" as="sourcePoint"/>
</mxGeometry>`
expect(isMxCellXmlComplete(xml)).toBe(false)
})
it("returns true for a self-closing last mxCell with > in its value", () => {
const xml = `<mxCell id="2" value="A" vertex="1" parent="1">
<mxGeometry as="geometry"/>
</mxCell>
<mxCell id="3" value="A -> B" vertex="1" parent="1"/></root>`
expect(isMxCellXmlComplete(xml)).toBe(true)
})
it("returns true for multiple complete mxCells", () => {
const xml = `<mxCell id="2" value="A" vertex="1" parent="1"/>
<mxCell id="3" value="B" vertex="1" parent="1"/>`
@@ -40,36 +103,6 @@ describe("isMxCellXmlComplete", () => {
})
})
describe("wrapWithMxFile", () => {
it("wraps empty string with default structure", () => {
const result = wrapWithMxFile("")
expect(result).toContain("<mxfile>")
expect(result).toContain("<mxGraphModel>")
expect(result).toContain('<mxCell id="0"/>')
expect(result).toContain('<mxCell id="1" parent="0"/>')
})
it("wraps raw mxCell content", () => {
const xml = '<mxCell id="2" value="Hello"/>'
const result = wrapWithMxFile(xml)
expect(result).toContain("<mxfile>")
expect(result).toContain(xml)
expect(result).toContain("</mxfile>")
})
it("returns full mxfile unchanged", () => {
const fullXml =
'<mxfile><diagram name="Page-1"><mxGraphModel></mxGraphModel></diagram></mxfile>'
const result = wrapWithMxFile(fullXml)
expect(result).toBe(fullXml)
})
it("handles whitespace in input", () => {
const result = wrapWithMxFile(" ")
expect(result).toContain("<mxfile>")
})
})
describe("cn (class name utility)", () => {
it("merges class names", () => {
expect(cn("foo", "bar")).toBe("foo bar")
@@ -84,3 +117,20 @@ describe("cn (class name utility)", () => {
expect(cn("text-red-500", "text-blue-500")).toBe("text-blue-500")
})
})
describe("extractCompleteMxCells", () => {
it("keeps the cell right after self-closing root cells", () => {
const xml = `<mxfile><diagram id="p1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="A" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell><mxCell id="3" value="B" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>`
const ids = [
...extractCompleteMxCells(xml).matchAll(/<mxCell id="([^"]+)"/g),
].map((m) => m[1])
expect(ids).toEqual(["0", "1", "2", "3"])
})
it("drops an incomplete trailing cell", () => {
const xml = `<mxCell id="2" vertex="1" parent="1"/><mxCell id="3" vertex="1" parent="1"><mxGeometry as="geometry"/>`
expect(extractCompleteMxCells(xml)).toBe(
'<mxCell id="2" vertex="1" parent="1"/>',
)
})
})
+136
View File
@@ -0,0 +1,136 @@
// @vitest-environment node
import { simulateReadableStream } from "ai"
import { MockLanguageModelV3 } from "ai/test"
import { afterEach, describe, expect, it, vi } from "vitest"
import { POST as validateDiagram } from "@/app/api/validate-diagram/route"
const RESULT = {
valid: false,
issues: [
{
type: "overlap",
severity: "critical",
description: "Box A covers box B",
},
],
suggestions: ["Move box B to the right"],
}
// A vision model that answers with the JSON in a few text chunks
vi.mock("@/lib/ai-providers", () => ({
getValidationModel: () =>
new MockLanguageModelV3({
doStream: (async () => {
const json = JSON.stringify(RESULT)
return {
stream: simulateReadableStream({
chunks: [
{ type: "text-start", id: "t" },
...[json.slice(0, 20), json.slice(20)].map(
(delta) => ({
type: "text-delta",
id: "t",
delta,
}),
),
{ type: "text-end", id: "t" },
{
type: "finish",
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
},
],
}),
}
}) as any,
}),
}))
// The quota, off unless a test turns it on
const quota = vi.hoisted(() => ({
enabled: false,
allowed: true,
checks: [] as Array<{ limits: any; increment?: number }>,
recorded: [] as number[],
}))
vi.mock("@/lib/dynamo-quota-manager", () => ({
isQuotaEnabled: () => quota.enabled,
checkAndIncrementRequest: async (
_ip: string,
limits: unknown,
increment?: number,
) => {
quota.checks.push({ limits, increment })
return quota.allowed
? { allowed: true }
: {
allowed: false,
type: "token",
error: "Daily token limit exceeded",
used: 10,
limit: 10,
}
},
recordTokenUsage: async (_ip: string, tokens: number) => {
quota.recorded.push(tokens)
},
}))
const post = () =>
new Request("http://localhost/api/validate-diagram", {
method: "POST",
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
},
body: JSON.stringify({ imageData: "data:image/png;base64,AAAA" }),
})
afterEach(() => {
delete process.env.ENABLE_VLM_VALIDATION
quota.enabled = false
quota.allowed = true
quota.checks = []
quota.recorded = []
})
describe("the quota", () => {
it("refuses a check once the daily tokens are used up", async () => {
quota.enabled = true
quota.allowed = false
const res = await validateDiagram(post())
expect(res.status).toBe(429)
})
it("applies the token limits only, and records the tokens", async () => {
// The request limit is for chats: the day's last chat must still
// get its check, and the check does not count as a chat
quota.enabled = true
const res = await validateDiagram(post())
expect(JSON.parse(await res.text())).toEqual(RESULT)
expect(quota.checks).toHaveLength(1)
expect(quota.checks[0].increment).toBe(0)
expect(quota.checks[0].limits.requests).toBe(0)
await vi.waitFor(() => expect(quota.recorded).toEqual([2]))
})
})
describe("POST /api/validate-diagram", () => {
it("streams the model's result as JSON text for useObject", async () => {
const res = await validateDiagram(post())
expect(JSON.parse(await res.text())).toEqual(RESULT)
})
it("answers valid when the check is turned off", async () => {
process.env.ENABLE_VLM_VALIDATION = "false"
const res = await validateDiagram(post())
expect(JSON.parse(await res.text())).toEqual({
valid: true,
issues: [],
suggestions: [],
})
})
})
+34
View File
@@ -0,0 +1,34 @@
// @vitest-environment node
import { describe, expect, it, vi } from "vitest"
// What the Test hands to the provider factory
const factory = vi.hoisted(() => ({ overrides: null as any }))
vi.mock("@/lib/ai-providers", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ai-providers")>()),
getAIModel: (overrides: unknown) => {
factory.overrides = overrides
throw new Error("stop here")
},
}))
import { POST as validateModel } from "@/app/api/validate-model/route"
describe("testing Bedrock", () => {
it("passes temporary credentials' session token on, as the chat does", async () => {
await validateModel(
new Request("http://localhost/api/validate-model", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
provider: "bedrock",
modelId: "amazon.nova-lite-v1:0",
awsAccessKeyId: "ASIA-temporary",
awsSecretAccessKey: "secret",
awsRegion: "us-east-1",
awsSessionToken: "session-token",
}),
}),
)
expect(factory.overrides.awsSessionToken).toBe("session-token")
})
})
+368
View File
@@ -0,0 +1,368 @@
// @vitest-environment node
import { streamText } from "ai"
import { afterEach, describe, expect, it, vi } from "vitest"
import { POST as testModel } from "@/app/api/admin/test-model/route"
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: () => ({}),
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 }))
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
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()
})
/** An OpenAI-compatible streaming reply made of the given deltas */
function streamReply(...deltas: object[]) {
const chunk = (delta: object, finish: string | null) =>
`data: ${JSON.stringify({
id: "c1",
object: "chat.completion.chunk",
created: 1,
model: "m",
choices: [{ index: 0, delta, finish_reason: finish }],
})}\n\n`
const body =
deltas.map((d) => chunk(d, null)).join("") +
chunk({}, "stop") +
"data: [DONE]\n\n"
vi.stubGlobal(
"fetch",
vi.fn(
async () =>
new Response(body, {
headers: { "content-type": "text/event-stream" },
}),
),
)
}
const testGlm = async () => {
const res = await validateModel(
new Request("http://localhost/api/validate-model", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
provider: "glm",
apiKey: "key",
modelId: "glm-5",
}),
}),
)
return res.json()
}
describe("POST /api/validate-model", () => {
it("passes when the model calls the test tool", async () => {
streamReply({
role: "assistant",
tool_calls: [
{
index: 0,
id: "call_1",
type: "function",
function: { name: "ping", arguments: "{}" },
},
],
})
const data = await testGlm()
expect(data.valid).toBe(true)
expect(data.warning).toBeUndefined()
expect(typeof data.responseTime).toBe("number")
})
it("reports a model that did not answer in time", async () => {
// The 15 s timeout has fired: the SDK ends the stream with an
// abort part instead of throwing
const timedOut = AbortSignal.abort(
new DOMException("The operation timed out.", "TimeoutError"),
)
const timeout = vi
.spyOn(AbortSignal, "timeout")
.mockReturnValue(timedOut)
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw timedOut.reason
}),
)
try {
const data = await testGlm()
expect(data.valid).toBe(false)
expect(data.code).toBe("timeout")
} finally {
timeout.mockRestore()
}
})
it("does not run on the server's keys", async () => {
process.env.OLLAMA_API_KEY = "server-ollama-key"
try {
const res = await validateModel(
new Request("http://localhost/api/validate-model", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
provider: "ollama",
modelId: "any-cloud-model",
}),
}),
)
expect(res.status).toBe(400)
expect((await res.json()).error).toMatch(/API key/)
} finally {
delete process.env.OLLAMA_API_KEY
}
})
it("warns when the model answers without a tool call", async () => {
streamReply({ role: "assistant", content: "OK" })
const data = await testGlm()
expect(data.valid).toBe(true)
expect(data.warning).toMatch(/without calling a tool/)
})
})
describe("chat requests to a client base URL", () => {
it("refuse redirects when private URLs are blocked", async () => {
process.env.ALLOW_PRIVATE_URLS = "false"
vi.stubGlobal(
"fetch",
vi.fn(
async () =>
new Response(null, {
status: 302,
headers: { location: "http://169.254.169.254/" },
}),
),
)
const { model } = getAIModel({
provider: "glm",
apiKey: "key",
baseUrl: "https://attacker.example/v1",
modelId: "glm-5",
})
let error: unknown
const result = streamText({
model,
prompt: "hi",
maxRetries: 0,
onError: ({ error: e }) => {
error = e
},
})
await result.consumeStream()
expect(String(error)).toMatch(/Redirects are not allowed/)
})
})
describe("testing EdgeOne", () => {
// The request validate-model sends to the EdgeOne function
const capture = () => {
const calls: Array<{ url: string; headers: Headers }> = []
vi.stubGlobal(
"fetch",
vi.fn(async (url: string, init?: RequestInit) => {
calls.push({
url: String(url),
headers: new Headers(init?.headers),
})
throw new Error("no network in tests")
}),
)
return calls
}
it("calls the site's own function, also without a base URL", async () => {
// The admin panel's Test sends none; a relative one cannot be fetched
const calls = capture()
await validateModel(
new Request("http://localhost/api/validate-model", {
method: "POST",
headers: {
"Content-Type": "application/json",
origin: "https://draw.example",
},
body: JSON.stringify({
provider: "edgeone",
modelId: "@tx/deepseek-ai/deepseek-v3-0324",
baseUrl: "https://elsewhere.example/api/edgeai",
}),
}),
)
expect(calls[0]?.url).toBe(
"https://draw.example/api/edgeai/chat/completions",
)
})
it("passes the admin's access code and cookies on", async () => {
// The EdgeOne function checks the access code too
process.env.ADMIN_PASSWORD = "admin-pw"
try {
const calls = capture()
await testModel(
new Request("http://localhost/api/admin/test-model", {
method: "POST",
headers: {
"Content-Type": "application/json",
"x-admin-password": "admin-pw",
"x-access-code": "visitor-code",
cookie: "eo_token=t",
origin: "https://draw.example",
},
body: JSON.stringify({
provider: {
id: "p1",
provider: "edgeone",
models: ["@tx/deepseek-ai/deepseek-v3-0324"],
},
modelId: "@tx/deepseek-ai/deepseek-v3-0324",
}),
}),
)
expect(calls[0]?.url).toBe(
"https://draw.example/api/edgeai/chat/completions",
)
expect(calls[0]?.headers.get("x-access-code")).toBe("visitor-code")
expect(calls[0]?.headers.get("cookie")).toBe("eo_token=t")
} finally {
delete process.env.ADMIN_PASSWORD
}
})
})
describe("the admin Test of the server's own base URL", () => {
const test = (headers: Record<string, string>) =>
validateModel(
new Request("http://localhost/api/validate-model", {
method: "POST",
headers: { "Content-Type": "application/json", ...headers },
body: JSON.stringify({
provider: "openai",
apiKey: "panel-key",
modelId: "gpt-5.5",
baseUrl: "http://10.0.0.5:8000/v1",
serverBaseUrl: true,
}),
}),
)
it("tests it as chat uses it: an internal address is allowed", async () => {
// ALLOW_PRIVATE_URLS=false guards URLs users type, not the server's
process.env.ALLOW_PRIVATE_URLS = "false"
process.env.OPENAI_BASE_URL = "http://10.0.0.5:8000/v1"
process.env.ADMIN_PASSWORD = "admin-pw"
privateUrls.all = true
try {
streamReply({ role: "assistant", content: "OK" })
const admin = await (
await test({ "x-admin-password": "admin-pw" })
).json()
expect(admin.valid).toBe(true)
// Anyone else claiming it is still refused
const other = await test({})
expect(other.status).toBe(400)
} finally {
privateUrls.all = false
delete process.env.OPENAI_BASE_URL
delete process.env.ADMIN_PASSWORD
}
})
})
describe("the admin panel's Test button", () => {
it("works when access codes are set", async () => {
// The admin password stands in for the visitor access code
process.env.ACCESS_CODE_LIST = "visitor-code"
process.env.ADMIN_PASSWORD = "admin-pw"
try {
streamReply({ role: "assistant", content: "OK" })
const res = await testModel(
new Request("http://localhost/api/admin/test-model", {
method: "POST",
headers: {
"Content-Type": "application/json",
"x-admin-password": "admin-pw",
},
body: JSON.stringify({
provider: {
id: "p1",
provider: "glm",
apiKey: "key",
models: ["glm-5"],
},
modelId: "glm-5",
}),
}),
)
expect(res.status).toBe(200)
expect((await res.json()).valid).toBe(true)
} finally {
delete process.env.ACCESS_CODE_LIST
delete process.env.ADMIN_PASSWORD
}
})
})
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)
})
})