Compare commits

..

2 Commits

Author SHA1 Message Date
dayuan.jiang
f3edceb9c0 fix: align get_shape_library guidance for non-cloud icon libraries 2026-02-07 13:55:39 +09:00
dayuan.jiang
5f87ffae1e feat: add Material Design Icons shape library (#685)
Add Google Material Design Icons as a new shape library using Google's
CDN. Includes top 300 most popular icons by usage, and updates system
prompts to guide the AI to call get_shape_library before using any icon
library.
2026-02-07 13:37:20 +09:00
15 changed files with 575 additions and 1030 deletions

View File

@@ -5,7 +5,6 @@ import { allowPrivateUrls, isPrivateUrl } from "@/lib/ssrf-protection"
const MAX_CONTENT_LENGTH = 150000 // Match PDF limit
const EXTRACT_TIMEOUT_MS = 15000
const USER_AGENT = "Mozilla/5.0 (compatible; NextAIDrawio/1.0)"
export async function POST(req: Request) {
try {
@@ -35,31 +34,6 @@ export async function POST(req: Request) {
{ status: 400 },
)
}
const headController = new AbortController()
const headTimeout = setTimeout(() => headController.abort(), 3000)
try {
const headResponse = await fetch(url, {
method: "HEAD",
headers: { "User-Agent": USER_AGENT },
signal: headController.signal,
})
const contentType = headResponse.headers.get("content-type")
if (contentType?.includes("application/pdf")) {
return NextResponse.json(
{
error: "PDF URLs are not supported. Please download and upload the PDF file directly",
},
{ status: 422 },
)
}
} catch (err) {
console.warn(
"HEAD pre-check failed, proceeding with extraction:",
err,
)
} finally {
clearTimeout(headTimeout)
}
// Extract article content with timeout to avoid tying up server resources
const controller = new AbortController()
@@ -70,7 +44,9 @@ export async function POST(req: Request) {
let article
try {
article = await extract(url, undefined, {
headers: { "User-Agent": USER_AGENT },
headers: {
"User-Agent": "Mozilla/5.0 (compatible; NextAIDrawio/1.0)",
},
signal: controller.signal,
})
} catch (err: any) {

View File

@@ -174,21 +174,10 @@ export async function POST(req: Request) {
}
case "ollama": {
// SECURITY: Mirror ai-providers.ts guard — only use server
// OLLAMA_API_KEY when the URL is also from server config.
const ollamaApiKey = baseUrl
? apiKey || undefined
: apiKey || process.env.OLLAMA_API_KEY || undefined
const ollamaProvider = createOllama({
baseURL:
baseUrl ||
process.env.OLLAMA_BASE_URL ||
"https://ollama.com/api",
...(ollamaApiKey && {
headers: { Authorization: `Bearer ${ollamaApiKey}` },
}),
const ollama = createOllama({
baseURL: baseUrl || "http://localhost:11434",
})
model = ollamaProvider(modelId)
model = ollama(modelId)
break
}

View File

@@ -282,7 +282,6 @@ export function ModelConfigDialog({
// Check credentials based on provider type
const isBedrock = selectedProvider.provider === "bedrock"
const isEdgeOne = selectedProvider.provider === "edgeone"
const isOllama = selectedProvider.provider === "ollama"
const isVertexAI = selectedProvider.provider === "vertexai"
if (isBedrock) {
if (
@@ -297,7 +296,7 @@ export function ModelConfigDialog({
if (!selectedProvider.vertexApiKey) {
return
}
} else if (!isEdgeOne && !isOllama && !selectedProvider.apiKey) {
} else if (!isEdgeOne && !selectedProvider.apiKey) {
return
}
@@ -1031,7 +1030,9 @@ export function ModelConfigDialog({
</div>
</>
) : selectedProvider.provider ===
"edgeone" ? (
"ollama" ||
selectedProvider.provider ===
"edgeone" ? (
<div className="space-y-3">
<div className="flex items-center gap-2">
<Button
@@ -1099,9 +1100,6 @@ export function ModelConfigDialog({
dict.modelConfig
.apiKey
}
{selectedProvider.provider ===
"ollama" &&
` ${dict.modelConfig.optional}`}
</Label>
<div className="flex gap-2">
<div className="relative flex-1">
@@ -1165,9 +1163,7 @@ export function ModelConfigDialog({
handleValidate
}
disabled={
(selectedProvider.provider !==
"ollama" &&
!selectedProvider.apiKey) ||
!selectedProvider.apiKey ||
validationStatus ===
"validating"
}

View File

@@ -359,9 +359,9 @@ const PROVIDER_ENV_MAP: Record<string, { apiKey: string; baseUrl: string }> = {
baseUrl: "MODELSCOPE_BASE_URL",
},
gateway: { apiKey: "AI_GATEWAY_API_KEY", baseUrl: "AI_GATEWAY_BASE_URL" },
// bedrock doesn't use API keys in the same way
// bedrock and ollama don't use API keys in the same way
bedrock: { apiKey: "", baseUrl: "" },
ollama: { apiKey: "OLLAMA_API_KEY", baseUrl: "OLLAMA_BASE_URL" },
ollama: { apiKey: "", baseUrl: "OLLAMA_BASE_URL" },
}
/**

View File

@@ -94,8 +94,7 @@ if (!gotTheLock) {
if (
url.includes("diagrams.net") ||
url.includes("draw.io") ||
url.startsWith("http://localhost") ||
url.startsWith("http://127.0.0.1")
url.startsWith("http://localhost")
) {
return { action: "allow" }
}

View File

@@ -68,7 +68,7 @@ export async function startNextServer(): Promise<string> {
const env: Record<string, string> = {
NODE_ENV: "production",
PORT: String(port),
HOSTNAME: "127.0.0.1",
HOSTNAME: "localhost",
// Enable Node.js built-in proxy support for fetch (Node.js 24+)
NODE_USE_ENV_PROXY: "1",
}

View File

@@ -9,11 +9,9 @@ import { app } from "electron"
const PORT_CONFIG = {
// Development mode uses fixed port for hot reload compatibility
development: 6002,
// Legacy production port — tried first to preserve localStorage for existing users
legacyProduction: 61337,
// New production port below the ephemeral range (49152-65535)
// to avoid conflicts with Windows Hyper-V / ephemeral port reservations
production: 13370,
// Production mode uses fixed port (61337) to preserve localStorage
// Falls back to sequential ports if unavailable
production: 61337,
// Maximum attempts to find an available port (fallback)
maxAttempts: 100,
}
@@ -29,10 +27,7 @@ let allocatedPort: number | null = null
export function isPortAvailable(port: number): Promise<boolean> {
return new Promise((resolve) => {
const server = net.createServer()
server.once("error", (err: NodeJS.ErrnoException) => {
console.warn(`Port ${port} unavailable: ${err.code}`)
resolve(false)
})
server.once("error", () => resolve(false))
server.once("listening", () => {
server.close()
resolve(true)
@@ -44,12 +39,12 @@ export function isPortAvailable(port: number): Promise<boolean> {
/**
* Find an available port
* - In development: uses fixed port (6002)
* - In production: uses fixed port (13370) to preserve localStorage
* - In production: uses fixed port (61337) to preserve localStorage
* - Falls back to sequential ports if preferred port is unavailable
* - Last resort: lets the OS assign a port (port 0)
*
* @param reuseExisting If true, try to reuse the previously allocated port
* @returns Promise<number> The available port
* @throws Error if no available port found after max attempts
*/
export async function findAvailablePort(reuseExisting = true): Promise<number> {
const isDev = !app.isPackaged
@@ -69,16 +64,7 @@ export async function findAvailablePort(reuseExisting = true): Promise<number> {
allocatedPort = null
}
// In production, try legacy port first to preserve existing users' localStorage
if (!isDev) {
const legacyPort = PORT_CONFIG.legacyProduction
if (await isPortAvailable(legacyPort)) {
allocatedPort = legacyPort
return legacyPort
}
}
// Try preferred port
// Try preferred port first
if (await isPortAvailable(preferredPort)) {
allocatedPort = preferredPort
return preferredPort
@@ -98,23 +84,9 @@ export async function findAvailablePort(reuseExisting = true): Promise<number> {
}
}
// Last resort: let the OS pick an available port
console.warn(
"All sequential ports failed. Requesting OS-assigned port (localStorage may not persist across restarts).",
throw new Error(
`Failed to find available port after ${PORT_CONFIG.maxAttempts} attempts`,
)
const osPort = await new Promise<number>((resolve, reject) => {
const server = net.createServer()
server.once("error", reject)
server.once("listening", () => {
const addr = server.address()
const port = (addr as net.AddressInfo).port
server.close(() => resolve(port))
})
server.listen(0, "127.0.0.1")
})
allocatedPort = osPort
console.log(`OS assigned port: ${osPort}`)
return osPort
}
/**
@@ -141,5 +113,5 @@ export function getServerUrl(): string {
"No port allocated yet. Call findAvailablePort() first.",
)
}
return `http://127.0.0.1:${allocatedPort}`
return `http://localhost:${allocatedPort}`
}

View File

@@ -66,11 +66,7 @@ export function createWindow(serverUrl: string): BrowserWindow {
// Handle page title updates
mainWindow.webContents.on("page-title-updated", (event, title) => {
if (
title &&
!title.includes("localhost") &&
!title.includes("127.0.0.1")
) {
if (title && !title.includes("localhost")) {
mainWindow?.setTitle(title)
} else {
event.preventDefault()

View File

@@ -59,9 +59,8 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0
# AZURE_REASONING_EFFORT=low # Optional: Azure reasoning effort (low, medium, high)
# AZURE_REASONING_SUMMARY=detailed
# Ollama Configuration (Local or Cloud)
# OLLAMA_BASE_URL=https://ollama.com/api # Optional, defaults to Ollama Cloud
# OLLAMA_API_KEY=your-ollama-cloud-api-key # Optional: For Ollama Cloud or authenticated remote instances
# Ollama (Local) Configuration
# OLLAMA_BASE_URL=http://localhost:11434/api # Optional, defaults to localhost
# OLLAMA_ENABLE_THINKING=true # Optional: Enable thinking for models that support it (e.g., qwen3)
# OpenRouter Configuration

View File

@@ -596,7 +596,7 @@ function validateProviderCredentials(
* - GOOGLE_GENERATIVE_AI_API_KEY: Google API key
* - AZURE_RESOURCE_NAME, AZURE_API_KEY: Azure OpenAI credentials
* - AWS_REGION, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY: AWS Bedrock credentials
* - OLLAMA_BASE_URL: Ollama server URL (optional, defaults to https://ollama.com/api)
* - OLLAMA_BASE_URL: Ollama server URL (optional, defaults to http://localhost:11434)
* - OPENROUTER_API_KEY: OpenRouter API key
* - DEEPSEEK_API_KEY: DeepSeek API key
* - DEEPSEEK_BASE_URL: DeepSeek endpoint (optional)
@@ -611,15 +611,13 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
// SECURITY: Prevent SSRF attacks (GHSA-9qf7-mprq-9qgm)
// If a custom baseUrl is provided, an API key MUST also be provided.
// This prevents attackers from redirecting server API keys to malicious endpoints.
// Exception: EdgeOne doesn't require API keys.
// Ollama is exempt only when no server OLLAMA_API_KEY is configured;
// when it IS configured, the outer guard also enforces client apiKey for custom baseUrls.
// Exception: EdgeOne and Ollama providers don't require API keys
if (
overrides?.baseUrl &&
!overrides?.apiKey &&
!(overrides?.provider === "vertexai" && overrides?.vertexApiKey) &&
overrides?.provider !== "edgeone" &&
!(overrides?.provider === "ollama" && !process.env.OLLAMA_API_KEY)
overrides?.provider !== "ollama"
) {
throw new Error(
`API key is required when using a custom base URL. ` +
@@ -880,19 +878,8 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
case "ollama": {
const baseURL = overrides?.baseUrl || process.env.OLLAMA_BASE_URL
// SECURITY: When client provides a custom base URL, only use
// client-provided API key. Never fall back to server OLLAMA_API_KEY
// to prevent leaking server credentials to user-controlled endpoints.
const apiKey = overrides?.baseUrl
? overrides?.apiKey || undefined
: resolveApiKey(overrides, "OLLAMA_API_KEY")
if (baseURL || apiKey) {
const customOllama = createOllama({
...(baseURL && { baseURL }),
...(apiKey && {
headers: { Authorization: `Bearer ${apiKey}` },
}),
})
if (baseURL) {
const customOllama = createOllama({ baseURL })
model = customOllama(modelId)
} else {
model = ollama(modelId)
@@ -1230,12 +1217,7 @@ export function supportsImageInput(modelId: string): boolean {
}
// Qwen text models (not vision variants like qwen-vl)
// qwen3.5-plus is a vision model
if (
lowerModelId.includes("qwen") &&
!hasVisionIndicator &&
!lowerModelId.includes("qwen3.5-plus")
) {
if (lowerModelId.includes("qwen") && !hasVisionIndicator) {
return false
}

View File

@@ -104,7 +104,7 @@ export const PROVIDER_INFO: Record<
bedrock: { label: "Amazon Bedrock" },
ollama: {
label: "Ollama",
defaultBaseUrl: "https://ollama.com/api",
defaultBaseUrl: "http://localhost:11434",
},
openrouter: {
label: "OpenRouter",
@@ -264,7 +264,6 @@ export const SUGGESTED_MODELS: Partial<Record<ProviderName, string[]>> = {
"Qwen/Qwen2.5-Coder-32B-Instruct",
"Qwen/Qwen2.5-7B-Instruct",
"Qwen/Qwen2-VL-72B-Instruct",
"qwen3.5-plus",
],
sglang: [
// SGLang is OpenAI-compatible, models depend on deployment
@@ -294,7 +293,6 @@ export const SUGGESTED_MODELS: Partial<Record<ProviderName, string[]>> = {
"Qwen/Qwen3-235B-A22B-Instruct-2507",
"Qwen/Qwen3-VL-235B-A22B-Instruct",
"Qwen/Qwen3-32B",
"qwen3.5-plus",
// DeepSeek
"deepseek-ai/DeepSeek-R1-0528",
"deepseek-ai/DeepSeek-V3.2",

1247
package-lock.json generated

File diff suppressed because it is too large Load Diff

View File

@@ -43,7 +43,7 @@
"@aws-sdk/client-dynamodb": "^3.957.0",
"@aws-sdk/credential-providers": "^3.943.0",
"@extractus/article-extractor": "^8.0.18",
"@formatjs/intl-localematcher": "^0.8.0",
"@formatjs/intl-localematcher": "^0.7.2",
"@langfuse/client": "^4.4.9",
"@langfuse/otel": "^4.4.4",
"@langfuse/tracing": "^4.4.9",
@@ -51,7 +51,7 @@
"@opennextjs/cloudflare": "1.16.1",
"@openrouter/ai-sdk-provider": "^1.5.4",
"@opentelemetry/api": "^1.9.0",
"@opentelemetry/exporter-trace-otlp-http": "^0.211.0",
"@opentelemetry/exporter-trace-otlp-http": "^0.209.0",
"@opentelemetry/sdk-trace-node": "^2.2.0",
"@radix-ui/react-alert-dialog": "^1.1.15",
"@radix-ui/react-collapsible": "^1.1.12",
@@ -72,7 +72,7 @@
"cmdk": "^1.1.1",
"idb": "^8.0.3",
"jsonrepair": "^3.13.1",
"lucide-react": "^0.563.0",
"lucide-react": "^0.562.0",
"motion": "^12.23.25",
"nanoid": "^5.0.0",
"negotiator": "^1.0.0",

View File

@@ -1,12 +1,12 @@
{
"name": "@next-ai-drawio/mcp-server",
"version": "0.1.16",
"version": "0.1.12",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "@next-ai-drawio/mcp-server",
"version": "0.1.16",
"version": "0.1.12",
"license": "Apache-2.0",
"dependencies": {
"@modelcontextprotocol/sdk": "^1.0.4",
@@ -469,9 +469,9 @@
}
},
"node_modules/@hono/node-server": {
"version": "1.19.9",
"resolved": "https://registry.npmjs.org/@hono/node-server/-/node-server-1.19.9.tgz",
"integrity": "sha512-vHL6w3ecZsky+8P5MD+eFfaGTyCeOHUIFYMGpQGbrBTSmNNoxv0if69rEZ5giu36weC5saFuznL411gRX7bJDw==",
"version": "1.19.7",
"resolved": "https://registry.npmjs.org/@hono/node-server/-/node-server-1.19.7.tgz",
"integrity": "sha512-vUcD0uauS7EU2caukW8z5lJKtoGMokxNbJtBiwHgpqxEXokaHCBkQUmCHhjFB1VUTWdqj25QoMkMKzgjq+uhrw==",
"license": "MIT",
"engines": {
"node": ">=18.14.1"
@@ -481,12 +481,12 @@
}
},
"node_modules/@modelcontextprotocol/sdk": {
"version": "1.26.0",
"resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.26.0.tgz",
"integrity": "sha512-Y5RmPncpiDtTXDbLKswIJzTqu2hyBKxTNsgKqKclDbhIgg1wgtf1fRuvxgTnRfcnxtvvgbIEcqUOzZrJ6iSReg==",
"version": "1.25.2",
"resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.25.2.tgz",
"integrity": "sha512-LZFeo4F9M5qOhC/Uc1aQSrBHxMrvxett+9KLHt7OhcExtoiRN9DKgbZffMP/nxjutWDQpfMDfP3nkHI4X9ijww==",
"license": "MIT",
"dependencies": {
"@hono/node-server": "^1.19.9",
"@hono/node-server": "^1.19.7",
"ajv": "^8.17.1",
"ajv-formats": "^3.0.1",
"content-type": "^1.0.5",
@@ -494,15 +494,14 @@
"cross-spawn": "^7.0.5",
"eventsource": "^3.0.2",
"eventsource-parser": "^3.0.0",
"express": "^5.2.1",
"express-rate-limit": "^8.2.1",
"hono": "^4.11.4",
"jose": "^6.1.3",
"express": "^5.0.1",
"express-rate-limit": "^7.5.0",
"jose": "^6.1.1",
"json-schema-typed": "^8.0.2",
"pkce-challenge": "^5.0.0",
"raw-body": "^3.0.0",
"zod": "^3.25 || ^4.0",
"zod-to-json-schema": "^3.25.1"
"zod-to-json-schema": "^3.25.0"
},
"engines": {
"node": ">=18"
@@ -521,9 +520,9 @@
}
},
"node_modules/@types/node": {
"version": "24.10.12",
"resolved": "https://registry.npmjs.org/@types/node/-/node-24.10.12.tgz",
"integrity": "sha512-68e+T28EbdmLSTkPgs3+UacC6rzmqrcWFPQs1C8mwJhI/r5Uxr0yEuQotczNRROd1gq30NGxee+fo0rSIxpyAw==",
"version": "24.10.6",
"resolved": "https://registry.npmjs.org/@types/node/-/node-24.10.6.tgz",
"integrity": "sha512-B8h60xgJMR/xmgyX9fncRzEW9gCxoJjdenUhke2v1JGOd/V66KopmWrLPXi5oUI4VuiGK+d+HlXJjDRZMj21EQ==",
"dev": true,
"license": "MIT",
"dependencies": {
@@ -1075,13 +1074,10 @@
}
},
"node_modules/express-rate-limit": {
"version": "8.2.1",
"resolved": "https://registry.npmjs.org/express-rate-limit/-/express-rate-limit-8.2.1.tgz",
"integrity": "sha512-PCZEIEIxqwhzw4KF0n7QF4QqruVTcF73O5kFKUnGOyjbCCgizBBiFaYpd/fnBLUMPw/BWw9OsiN7GgrNYr7j6g==",
"version": "7.5.1",
"resolved": "https://registry.npmjs.org/express-rate-limit/-/express-rate-limit-7.5.1.tgz",
"integrity": "sha512-7iN8iPMDzOMHPUYllBEsQdWVB6fPDMPqwjBaFrgr4Jgr/+okjvzAy+UHlYYL/Vs0OsOrMkwS6PJDkFlJwoxUnw==",
"license": "MIT",
"dependencies": {
"ip-address": "10.0.1"
},
"engines": {
"node": ">= 16"
},
@@ -1264,9 +1260,9 @@
}
},
"node_modules/hono": {
"version": "4.11.9",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.11.9.tgz",
"integrity": "sha512-Eaw2YTGM6WOxA6CXbckaEvslr2Ne4NFsKrvc0v97JD5awbmeBLO5w9Ho9L9kmKonrwF9RJlW6BxT1PVv/agBHQ==",
"version": "4.11.1",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.11.1.tgz",
"integrity": "sha512-KsFcH0xxHes0J4zaQgWbYwmz3UPOOskdqZmItstUG93+Wk1ePBLkLGwbP9zlmh1BFUiL8Qp+Xfu9P7feJWpGNg==",
"license": "MIT",
"peer": true,
"engines": {
@@ -1352,15 +1348,6 @@
"integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==",
"license": "ISC"
},
"node_modules/ip-address": {
"version": "10.0.1",
"resolved": "https://registry.npmjs.org/ip-address/-/ip-address-10.0.1.tgz",
"integrity": "sha512-NWv9YLW4PoW2B7xtzaS3NCot75m6nK7Icdv0o3lfMceJVRfSoQwqD4wEH5rLwoKJwUiZ/rfpiVBhnaF0FK4HoA==",
"license": "MIT",
"engines": {
"node": ">= 12"
}
},
"node_modules/ipaddr.js": {
"version": "1.9.1",
"resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-1.9.1.tgz",
@@ -2064,9 +2051,9 @@
}
},
"node_modules/zod": {
"version": "4.3.6",
"resolved": "https://registry.npmjs.org/zod/-/zod-4.3.6.tgz",
"integrity": "sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==",
"version": "4.3.5",
"resolved": "https://registry.npmjs.org/zod/-/zod-4.3.5.tgz",
"integrity": "sha512-k7Nwx6vuWx1IJ9Bjuf4Zt1PEllcwe7cls3VNzm4CQ1/hgtFUK2bRNG3rvnpPUhFjmqJKAKtjV576KnUkHocg/g==",
"license": "MIT",
"peer": true,
"funding": {
@@ -2074,9 +2061,9 @@
}
},
"node_modules/zod-to-json-schema": {
"version": "3.25.1",
"resolved": "https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.25.1.tgz",
"integrity": "sha512-pM/SU9d3YAggzi6MtR4h7ruuQlqKtad8e9S0fmxcMi+ueAK5Korys/aWcV9LIIHTVbj01NdzxcnXSN+O74ZIVA==",
"version": "3.25.0",
"resolved": "https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.25.0.tgz",
"integrity": "sha512-HvWtU2UG41LALjajJrML6uQejQhNJx+JBO9IflpSja4R03iNWfKXrj6W2h7ljuLyc1nKS+9yDyL/9tD1U/yBnQ==",
"license": "ISC",
"peerDependencies": {
"zod": "^3.25 || ^4"

View File

@@ -1,6 +1,5 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { describe, expect, it } from "vitest"
import {
getAIModel,
resolveBaseURL,
supportsImageInput,
supportsPromptCaching,
@@ -190,120 +189,3 @@ describe("supportsImageInput", () => {
expect(supportsImageInput("gemini-pro")).toBe(true)
})
})
vi.mock("ollama-ai-provider-v2", () => {
const mockModel = { modelId: "test-model" }
const mockProviderFn = vi.fn(() => mockModel)
const mockCreateOllama = vi.fn(() => mockProviderFn)
const mockOllama = vi.fn(() => mockModel)
return { createOllama: mockCreateOllama, ollama: mockOllama }
})
describe("Ollama API key security", () => {
let createOllamaMock: ReturnType<typeof vi.fn>
const savedEnv: Record<string, string | undefined> = {}
beforeEach(async () => {
savedEnv.OLLAMA_API_KEY = process.env.OLLAMA_API_KEY
savedEnv.OLLAMA_BASE_URL = process.env.OLLAMA_BASE_URL
delete process.env.OLLAMA_BASE_URL
const mod = await import("ollama-ai-provider-v2")
createOllamaMock = mod.createOllama as ReturnType<typeof vi.fn>
createOllamaMock.mockClear()
})
afterEach(() => {
process.env.OLLAMA_API_KEY = savedEnv.OLLAMA_API_KEY
process.env.OLLAMA_BASE_URL = savedEnv.OLLAMA_BASE_URL
})
it("applies server OLLAMA_API_KEY when no client baseUrl is provided", () => {
process.env.OLLAMA_API_KEY = "server-secret-key"
getAIModel({ provider: "ollama", modelId: "llama2" })
expect(createOllamaMock).toHaveBeenCalledWith(
expect.objectContaining({
headers: { Authorization: "Bearer server-secret-key" },
}),
)
})
it("does NOT leak server OLLAMA_API_KEY when client provides a custom baseUrl", () => {
process.env.OLLAMA_API_KEY = "server-secret-key"
// When server has OLLAMA_API_KEY, the SSRF guard rejects
// client-provided baseUrl without an apiKey outright
expect(() =>
getAIModel({
provider: "ollama",
baseUrl: "https://evil-server.com",
modelId: "llama2",
}),
).toThrow("API key is required")
})
it("uses client API key when client provides both baseUrl and apiKey", () => {
process.env.OLLAMA_API_KEY = "server-secret-key"
getAIModel({
provider: "ollama",
baseUrl: "https://my-ollama.com",
apiKey: "client-key",
modelId: "llama2",
})
expect(createOllamaMock).toHaveBeenCalledWith(
expect.objectContaining({
baseURL: "https://my-ollama.com",
headers: { Authorization: "Bearer client-key" },
}),
)
})
it("applies both server OLLAMA_BASE_URL and OLLAMA_API_KEY when no client overrides", () => {
process.env.OLLAMA_BASE_URL = "https://cloud.ollama.com"
process.env.OLLAMA_API_KEY = "server-key"
getAIModel({ provider: "ollama", modelId: "llama2" })
expect(createOllamaMock).toHaveBeenCalledWith(
expect.objectContaining({
baseURL: "https://cloud.ollama.com",
headers: { Authorization: "Bearer server-key" },
}),
)
})
it("works when OLLAMA_API_KEY is set but OLLAMA_BASE_URL is not", () => {
process.env.OLLAMA_API_KEY = "server-key"
delete process.env.OLLAMA_BASE_URL
getAIModel({ provider: "ollama", modelId: "llama2" })
expect(createOllamaMock).toHaveBeenCalledTimes(1)
const callArgs = createOllamaMock.mock.calls[0][0]
expect(callArgs).not.toHaveProperty("baseURL")
expect(callArgs).toEqual(
expect.objectContaining({
headers: { Authorization: "Bearer server-key" },
}),
)
})
it("allows client custom baseUrl without apiKey when no server OLLAMA_API_KEY", () => {
delete process.env.OLLAMA_API_KEY
getAIModel({
provider: "ollama",
baseUrl: "https://my-ollama.com",
modelId: "llama2",
})
expect(createOllamaMock).toHaveBeenCalledTimes(1)
const callArgs = createOllamaMock.mock.calls[0][0]
expect(callArgs.baseURL).toBe("https://my-ollama.com")
expect(callArgs).not.toHaveProperty("headers")
})
})