From 6b8cedf15eec662e73fd050fde1ea767771ab8dd Mon Sep 17 00:00:00 2001 From: NotXf1le <89696340+NotXf1le@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:33:32 +0200 Subject: [PATCH] Add Ollama and vision to MCP server --- README.md | 6 +- packages/choosekit-mcp/README.md | 39 +++++++++-- packages/choosekit-mcp/package-lock.json | 12 ++-- packages/choosekit-mcp/package.json | 7 +- packages/choosekit-mcp/src/cli.ts | 9 +++ packages/choosekit-mcp/src/config.ts | 34 ++++++++- packages/choosekit-mcp/src/images.ts | 74 ++++++++++++++++++++ packages/choosekit-mcp/src/server.ts | 35 +++++++-- packages/choosekit-mcp/tests/cli.test.mjs | 68 +++++++++++++++++- packages/choosekit-mcp/tests/images.test.mjs | 51 ++++++++++++++ packages/choosekit-mcp/tests/server.test.mjs | 74 ++++++++++++++++++++ 11 files changed, 382 insertions(+), 27 deletions(-) create mode 100644 packages/choosekit-mcp/src/images.ts create mode 100644 packages/choosekit-mcp/tests/images.test.mjs diff --git a/README.md b/README.md index 3d0a666..b854407 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,7 @@ The project was inspired by [Jev and the System One model interface](https://typ ## MCP server -[`choosekit-mcp`](packages/choosekit-mcp/README.md) exposes choosekit through llama.cpp or OpenRouter as a read-only stdio tool for Claude Code, Codex, and OpenCode. Select the backend and configure it with environment variables when starting the MCP server. +[`choosekit-mcp`](packages/choosekit-mcp/README.md) exposes choosekit through llama.cpp, Ollama, or OpenRouter as a read-only stdio tool for Claude Code, Codex, and OpenCode. Select the backend and configure it with environment variables when starting the MCP server. ## llama.cpp @@ -98,7 +98,7 @@ import { fromOpenRouter } from "choosekit/openrouter"; const choose = fromOpenRouter({ apiKey: process.env.OPENROUTER_API_KEY!, - model: "qwen/qwen3.8-27b", + model: "your-model", }); ``` @@ -131,6 +131,8 @@ const decision = await choose({ Supported media types are PNG, JPEG, and WebP. llama.cpp image inputs currently support `labels` mode only. +The MCP server accepts image file paths through `imagePaths` in `labels` mode. Paths are resolved from the server's working directory by default; set `CHOOSEKIT_IMAGE_ROOT` to use another root. Every image must remain within that root and be a PNG, JPEG, or WebP file. With OpenRouter, the image contents are sent to the remote service. + ## Scoring modes | Mode | Candidate representation | Use when | diff --git a/packages/choosekit-mcp/README.md b/packages/choosekit-mcp/README.md index 0a4f301..34cda70 100644 --- a/packages/choosekit-mcp/README.md +++ b/packages/choosekit-mcp/README.md @@ -1,6 +1,6 @@ # choosekit-mcp -`choosekit-mcp` exposes [choosekit](https://github.com/NotXf1le/choosekit) as a read-only MCP tool backed by llama.cpp or OpenRouter. +`choosekit-mcp` exposes [choosekit](https://github.com/NotXf1le/choosekit) as a read-only MCP tool backed by llama.cpp, Ollama, or OpenRouter. The server uses stdio. It returns a decision and probability distribution. Execution of the selected action remains with the MCP client. @@ -10,16 +10,18 @@ Select and configure the backend when the MCP process starts. `CHOOSEKIT_BACKEND | Backend | Environment variable | Default | Description | |---|---|---|---| -| Both | `CHOOSEKIT_BACKEND` | `llama-cpp` | `llama-cpp` or `openrouter` | -| Both | `CHOOSEKIT_MODEL` | Not set | Optional for llama.cpp; required for OpenRouter | -| Both | `CHOOSEKIT_MODE` | `labels` | `labels` for either backend; `minimal-prefix` is llama.cpp only | +| All | `CHOOSEKIT_BACKEND` | `llama-cpp` | `llama-cpp`, `ollama`, or `openrouter` | +| All | `CHOOSEKIT_MODEL` | Not set | Optional for llama.cpp; required for Ollama and OpenRouter | +| All | `CHOOSEKIT_MODE` | `labels` | `labels` for every backend; `minimal-prefix` is llama.cpp only | +| All | `CHOOSEKIT_IMAGE_ROOT` | Current working directory | Root directory for files supplied through `imagePaths` | | llama.cpp | `CHOOSEKIT_BASE_URL` | Required | Base URL of the llama.cpp server, for example `http://127.0.0.1:8080` | +| Ollama | `CHOOSEKIT_BASE_URL` | `http://127.0.0.1:11434` | Base URL of the Ollama server | | OpenRouter | `OPENROUTER_API_KEY` | Required | OpenRouter API key | | OpenRouter | `OPENROUTER_PROVIDER` | Not set | Pins one provider and disables fallback | Every tool call uses the configuration set when the MCP process starts. Backend settings are not accepted as tool arguments. The llama.cpp server must provide its native `/tokenize` and `/completion` endpoints. OpenRouter receives the context, question, and choice descriptions. -OpenRouter supports `labels` mode with up to 20 choices. llama.cpp supports up to 26 choices in `labels` mode and has no additional MCP choice limit in `minimal-prefix` mode. +Ollama and OpenRouter support `labels` mode with up to 20 choices. llama.cpp supports up to 26 choices in `labels` mode and has no additional MCP choice limit in `minimal-prefix` mode. ### Claude Code @@ -39,12 +41,20 @@ codex mcp add choosekit --env CHOOSEKIT_BASE_URL=http://127.0.0.1:8080 -- npx -y opencode mcp add choosekit --env CHOOSEKIT_BASE_URL=http://127.0.0.1:8080 -- npx -y choosekit-mcp ``` +### Ollama + +For example, with Codex: + +```sh +codex mcp add choosekit --env CHOOSEKIT_BACKEND=ollama --env CHOOSEKIT_MODEL=your-model -- npx -y choosekit-mcp +``` + ### OpenRouter For example, with Codex: ```sh -codex mcp add choosekit --env CHOOSEKIT_BACKEND=openrouter --env OPENROUTER_API_KEY=... --env CHOOSEKIT_MODEL=... -- npx -y choosekit-mcp +codex mcp add choosekit --env CHOOSEKIT_BACKEND=openrouter --env OPENROUTER_API_KEY=... --env CHOOSEKIT_MODEL=your-model -- npx -y choosekit-mcp ``` On native Windows, launch `npx` through `cmd` in any example above: replace `-- npx -y choosekit-mcp` with `-- cmd /c npx -y choosekit-mcp`. @@ -65,6 +75,22 @@ The server exposes one tool named `choose`: } ``` +To score images, add `imagePaths` in `labels` mode: + +```json +{ + "context": "Inspect the attached screenshot.", + "question": "Which state is the interface in?", + "choices": { + "ready": "The interface is ready for input.", + "loading": "The interface is still loading." + }, + "imagePaths": ["screenshots/status.png"] +} +``` + +Relative paths are resolved from `CHOOSEKIT_IMAGE_ROOT`, which defaults to the MCP process working directory. Every image path must remain within that root. PNG, JPEG, and WebP files are supported, and the selected model must support vision. With OpenRouter, the image contents are sent to the remote service. + With `labels`, choosekit maps the supplied choice keys to A/B/C labels for scoring, so each description must contain the option's full meaning. With `minimal-prefix`, it scores the shortest token prefixes that distinguish the original keys. When no supplied option may apply, add an explicit choice such as `insufficient_information`. The result contains the selected key, the complete normalized distribution, raw scores, margin, entropy, token-boundary rollback, and backend usage when available. A score is `null` when the upstream backend did not return a log probability for that choice; its probability in the distribution is `0`. Probabilities represent relative preference among the supplied choices. Estimating correctness requires separate calibration. @@ -73,6 +99,7 @@ The result contains the selected key, the complete normalized distribution, raw - Node.js 20 or newer. - For llama.cpp, a reachable server with a compatible model already loaded. +- For Ollama, version 0.12.11 or newer and a compatible model. - For OpenRouter, an API key and a model that returns the required log probabilities. [Apache-2.0](LICENSE). Copyright 2026 NotXf1le. diff --git a/packages/choosekit-mcp/package-lock.json b/packages/choosekit-mcp/package-lock.json index 1bd08ec..832a955 100644 --- a/packages/choosekit-mcp/package-lock.json +++ b/packages/choosekit-mcp/package-lock.json @@ -1,16 +1,16 @@ { "name": "choosekit-mcp", - "version": "0.2.0", + "version": "0.3.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "choosekit-mcp", - "version": "0.2.0", + "version": "0.3.0", "license": "Apache-2.0", "dependencies": { "@modelcontextprotocol/server": "2.0.0", - "choosekit": "^0.5.0", + "choosekit": "^0.6.0", "zod": "^4.6.4" }, "bin": { @@ -60,9 +60,9 @@ } }, "node_modules/choosekit": { - "version": "0.5.0", - "resolved": "https://registry.npmjs.org/choosekit/-/choosekit-0.5.0.tgz", - "integrity": "sha512-eTnaA5zOr1W5sN62K3ou+/GFZoRl6mhpr7m63AvgR5aQa1sljEiglfigHGVHC/BR3XLVocdq9LyYn8iqQ4XoyA==", + "version": "0.6.0", + "resolved": "https://registry.npmjs.org/choosekit/-/choosekit-0.6.0.tgz", + "integrity": "sha512-dhTRxTbS5uCQapuxNLL1csglaKzsc87u7KaKqyIEuE5vfunOH13+k7k6JoJ2CRzBiwaH2IC/rHItc3Jfgr9txA==", "license": "Apache-2.0", "engines": { "node": ">=20" diff --git a/packages/choosekit-mcp/package.json b/packages/choosekit-mcp/package.json index 27da1dc..ab1f955 100644 --- a/packages/choosekit-mcp/package.json +++ b/packages/choosekit-mcp/package.json @@ -1,7 +1,7 @@ { "name": "choosekit-mcp", - "version": "0.2.0", - "description": "Use choosekit as a local MCP server.", + "version": "0.3.0", + "description": "Use choosekit as an MCP server.", "license": "Apache-2.0", "author": { "name": "NotXf1le", @@ -37,7 +37,7 @@ }, "dependencies": { "@modelcontextprotocol/server": "2.0.0", - "choosekit": "^0.5.0", + "choosekit": "^0.6.0", "zod": "^4.6.4" }, "devDependencies": { @@ -51,6 +51,7 @@ "mcp", "llm", "llama.cpp", + "ollama", "choosekit" ] } diff --git a/packages/choosekit-mcp/src/cli.ts b/packages/choosekit-mcp/src/cli.ts index e0224bb..38566fe 100644 --- a/packages/choosekit-mcp/src/cli.ts +++ b/packages/choosekit-mcp/src/cli.ts @@ -3,8 +3,10 @@ import { serveStdio } from "@modelcontextprotocol/server/stdio"; import type { Chooser } from "choosekit"; import { fromLlamaCpp } from "choosekit/llama-cpp"; +import { fromOllama } from "choosekit/ollama"; import { fromOpenRouter } from "choosekit/openrouter"; import { loadConfig } from "./config.js"; +import { createImageLoader } from "./images.js"; import { buildServer } from "./server.js"; try { @@ -18,6 +20,12 @@ try { mode: config.mode, }); break; + case "ollama": + chooser = fromOllama({ + model: config.model, + ...(config.baseURL === undefined ? {} : { baseURL: config.baseURL }), + }); + break; case "openrouter": chooser = fromOpenRouter({ apiKey: config.apiKey, @@ -28,6 +36,7 @@ try { } serveStdio(() => buildServer(chooser, { backend: config.backend, + imageLoader: createImageLoader(config.imageRoot), mode: config.mode, })); } catch (error) { diff --git a/packages/choosekit-mcp/src/config.ts b/packages/choosekit-mcp/src/config.ts index 8aa1649..c25be6d 100644 --- a/packages/choosekit-mcp/src/config.ts +++ b/packages/choosekit-mcp/src/config.ts @@ -6,15 +6,24 @@ export type Config = | { readonly backend: "llama-cpp"; readonly baseURL: string; + readonly imageRoot?: string; readonly model?: string; readonly mode: "labels" | "minimal-prefix"; } | { readonly backend: "openrouter"; readonly apiKey: string; + readonly imageRoot?: string; readonly model: string; readonly provider?: string; readonly mode: "labels"; + } + | { + readonly backend: "ollama"; + readonly baseURL?: string; + readonly imageRoot?: string; + readonly model: string; + readonly mode: "labels"; }; function requiredText(value: string | undefined, name: string): string { @@ -32,8 +41,27 @@ function optionalText(value: string | undefined, name: string): string | undefin export function loadConfig(env: NodeJS.ProcessEnv = process.env): Config { const backend = env.CHOOSEKIT_BACKEND?.trim() ?? "llama-cpp"; - if (backend !== "llama-cpp" && backend !== "openrouter") { - throw new TypeError("CHOOSEKIT_BACKEND must be llama-cpp or openrouter."); + const imageRoot = optionalText(env.CHOOSEKIT_IMAGE_ROOT, "CHOOSEKIT_IMAGE_ROOT"); + if (backend !== "llama-cpp" && backend !== "ollama" && backend !== "openrouter") { + throw new TypeError("CHOOSEKIT_BACKEND must be llama-cpp, ollama, or openrouter."); + } + + if (backend === "ollama") { + const mode = env.CHOOSEKIT_MODE?.trim() ?? "labels"; + if (mode !== "labels") { + throw new TypeError("CHOOSEKIT_MODE must be labels when using Ollama."); + } + const baseURL = optionalText(env.CHOOSEKIT_BASE_URL, "CHOOSEKIT_BASE_URL"); + if (baseURL !== undefined && !urlSchema.safeParse(baseURL).success) { + throw new TypeError("CHOOSEKIT_BASE_URL must be a valid URL."); + } + return { + backend, + ...(baseURL === undefined ? {} : { baseURL }), + ...(imageRoot === undefined ? {} : { imageRoot }), + model: requiredText(env.CHOOSEKIT_MODEL, "CHOOSEKIT_MODEL"), + mode, + }; } if (backend === "openrouter") { @@ -45,6 +73,7 @@ export function loadConfig(env: NodeJS.ProcessEnv = process.env): Config { return { backend, apiKey: requiredText(env.OPENROUTER_API_KEY, "OPENROUTER_API_KEY"), + ...(imageRoot === undefined ? {} : { imageRoot }), model: requiredText(env.CHOOSEKIT_MODEL, "CHOOSEKIT_MODEL"), ...(provider === undefined ? {} : { provider }), mode, @@ -63,6 +92,7 @@ export function loadConfig(env: NodeJS.ProcessEnv = process.env): Config { return { backend, baseURL, + ...(imageRoot === undefined ? {} : { imageRoot }), ...(model === undefined ? {} : { model }), mode, }; diff --git a/packages/choosekit-mcp/src/images.ts b/packages/choosekit-mcp/src/images.ts new file mode 100644 index 0000000..4fe2234 --- /dev/null +++ b/packages/choosekit-mcp/src/images.ts @@ -0,0 +1,74 @@ +import type { ImageInput, ImageMediaType } from "choosekit"; +import { readFile, realpath, stat } from "node:fs/promises"; +import { isAbsolute, relative, resolve, sep } from "node:path"; + +export type ImageLoader = ( + paths: readonly string[], + signal?: AbortSignal, +) => Promise; + +export class ImageLoadError extends Error { + override readonly name = "ImageLoadError"; +} + +function isInside(root: string, path: string): boolean { + const fromRoot = relative(root, path); + return fromRoot !== ".." && !fromRoot.startsWith(`..${sep}`) && !isAbsolute(fromRoot); +} + +function mediaType(data: Uint8Array): ImageMediaType | undefined { + if (data.length >= 8 + && data[0] === 0x89 && data[1] === 0x50 && data[2] === 0x4e && data[3] === 0x47 + && data[4] === 0x0d && data[5] === 0x0a && data[6] === 0x1a && data[7] === 0x0a) { + return "image/png"; + } + if (data.length >= 3 && data[0] === 0xff && data[1] === 0xd8 && data[2] === 0xff) { + return "image/jpeg"; + } + if (data.length >= 12 + && data[0] === 0x52 && data[1] === 0x49 && data[2] === 0x46 && data[3] === 0x46 + && data[8] === 0x57 && data[9] === 0x45 && data[10] === 0x42 && data[11] === 0x50) { + return "image/webp"; + } + return undefined; +} + +export function createImageLoader(root = process.cwd()): ImageLoader { + let resolvedRoot: Promise | undefined; + return async (paths, signal) => { + try { + resolvedRoot ??= realpath(resolve(root)).then(async (path) => { + if (!(await stat(path)).isDirectory()) { + throw new ImageLoadError("image root must be a directory."); + } + return path; + }); + const rootPath = await resolvedRoot; + const images: ImageInput[] = []; + for (const input of paths) { + signal?.throwIfAborted(); + if (typeof input !== "string" || input.trim().length === 0) { + throw new ImageLoadError("image paths must be non-empty strings."); + } + const path = await realpath(resolve(rootPath, input)); + if (!isInside(rootPath, path)) { + throw new ImageLoadError("image path is outside the configured root."); + } + if (!(await stat(path)).isFile()) { + throw new ImageLoadError("image path must identify a regular file."); + } + const data = await readFile(path, signal ? { signal } : undefined); + const detected = mediaType(data); + if (detected === undefined) { + throw new ImageLoadError("image must be PNG, JPEG, or WebP."); + } + images.push(Object.freeze({ mediaType: detected, base64: data.toString("base64") })); + } + return Object.freeze(images); + } catch (error) { + if (signal?.aborted || (error instanceof Error && error.name === "AbortError")) throw error; + if (error instanceof ImageLoadError) throw error; + throw new ImageLoadError("image file could not be read.", { cause: error }); + } + }; +} diff --git a/packages/choosekit-mcp/src/server.ts b/packages/choosekit-mcp/src/server.ts index 101b32d..9f46dc1 100644 --- a/packages/choosekit-mcp/src/server.ts +++ b/packages/choosekit-mcp/src/server.ts @@ -2,12 +2,14 @@ import { McpServer } from "@modelcontextprotocol/server"; import { ScoringError, type Chooser, type Decision, type Usage } from "choosekit"; import { createRequire } from "node:module"; import { z } from "zod"; +import { ImageLoadError, type ImageLoader } from "./images.js"; export type ChoiceMode = "labels" | "minimal-prefix"; export interface ServerOptions { - readonly backend?: "llama-cpp" | "openrouter"; + readonly backend?: "llama-cpp" | "ollama" | "openrouter"; readonly mode?: ChoiceMode; + readonly imageLoader?: ImageLoader; } const packageVersion = (createRequire(import.meta.url)("../package.json") as { version: string }).version; @@ -29,7 +31,8 @@ const decisionSchema = z.object({ usage: usageSchema.optional(), }).strict(); -function inputSchema(backend: "llama-cpp" | "openrouter", mode: ChoiceMode) { +function inputSchema(backend: "llama-cpp" | "ollama" | "openrouter", mode: ChoiceMode, + imagesEnabled: boolean) { const choicesSchema = z.fromJSONSchema({ type: "object", propertyNames: { @@ -41,7 +44,7 @@ function inputSchema(backend: "llama-cpp" | "openrouter", mode: ChoiceMode) { pattern: "\\S", }, minProperties: 2, - ...(backend === "openrouter" + ...(backend === "openrouter" || backend === "ollama" ? { maxProperties: 20 } : mode === "labels" ? { maxProperties: 26 } : {}), }) as z.ZodType>; @@ -51,6 +54,11 @@ function inputSchema(backend: "llama-cpp" | "openrouter", mode: ChoiceMode) { question: z.string() .refine((question) => question.trim().length > 0, "The question must not be empty."), choices: choicesSchema, + ...(mode === "labels" && imagesEnabled ? { + imagePaths: z.array(z.string().min(1)).min(1) + .describe("Image file paths inside the configured image root, in display order.") + .optional(), + } : {}), }).strict(); } @@ -71,6 +79,15 @@ function errorResult(error: unknown, signal: AbortSignal) { content: [{ type: "text" as const, text: "The model could not score the supplied choices." }], }; } + if (error instanceof ImageLoadError) { + return { + isError: true as const, + content: [{ + type: "text" as const, + text: "The image could not be loaded. Check imagePaths and CHOOSEKIT_IMAGE_ROOT.", + }], + }; + } return { isError: true as const, content: [{ type: "text" as const, text: "The choice request failed unexpectedly." }], @@ -86,8 +103,8 @@ Readonly> { export function buildServer(chooser: Chooser, options: ServerOptions = {}): McpServer { if (typeof chooser !== "function") throw new TypeError("chooser must be a function."); const backend = options.backend ?? "llama-cpp"; - if (backend !== "llama-cpp" && backend !== "openrouter") { - throw new TypeError("backend must be llama-cpp or openrouter."); + if (backend !== "llama-cpp" && backend !== "ollama" && backend !== "openrouter") { + throw new TypeError("backend must be llama-cpp, ollama, or openrouter."); } const mode = options.mode ?? "labels"; if (mode !== "labels" && mode !== "minimal-prefix") { @@ -99,19 +116,23 @@ export function buildServer(chooser: Chooser, options: ServerOptions = {}): McpS "choose", { description: toolDescription(mode), - inputSchema: inputSchema(backend, mode), + inputSchema: inputSchema(backend, mode, options.imageLoader !== undefined), outputSchema: decisionSchema, annotations: { readOnlyHint: true, openWorldHint: backend === "openrouter", }, }, - async ({ context, question, choices }, ctx) => { + async ({ context, question, choices, imagePaths }, ctx) => { try { + const images = imagePaths === undefined + ? undefined + : await options.imageLoader!(imagePaths as readonly string[], ctx.mcpReq.signal); const decision = await chooser({ context, question, choices, + ...(images === undefined ? {} : { images }), signal: ctx.mcpReq.signal, }) as Decision; const structuredContent: { diff --git a/packages/choosekit-mcp/tests/cli.test.mjs b/packages/choosekit-mcp/tests/cli.test.mjs index 9af1e14..69b6a9e 100644 --- a/packages/choosekit-mcp/tests/cli.test.mjs +++ b/packages/choosekit-mcp/tests/cli.test.mjs @@ -1,6 +1,9 @@ import assert from "node:assert/strict"; import { spawn, spawnSync } from "node:child_process"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; import { createServer } from "node:http"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import { test } from "node:test"; import { fileURLToPath } from "node:url"; @@ -19,6 +22,7 @@ const request = Object.freeze({ const configurationKeys = [ "CHOOSEKIT_BACKEND", "CHOOSEKIT_BASE_URL", + "CHOOSEKIT_IMAGE_ROOT", "CHOOSEKIT_MODEL", "CHOOSEKIT_MODE", "OPENROUTER_API_KEY", @@ -131,7 +135,15 @@ test("reports invalid environment configuration without writing to stdout", () = ["invalid mode", { CHOOSEKIT_MODE: "keys" }, "CHOOSEKIT_MODE must be labels or minimal-prefix."], ["unknown backend", { CHOOSEKIT_BACKEND: "other" }, - "CHOOSEKIT_BACKEND must be llama-cpp or openrouter."], + "CHOOSEKIT_BACKEND must be llama-cpp, ollama, or openrouter."], + ["Ollama without model", { + CHOOSEKIT_BACKEND: "ollama", + }, "CHOOSEKIT_MODEL is required."], + ["Ollama with minimal-prefix", { + CHOOSEKIT_BACKEND: "ollama", + CHOOSEKIT_MODEL: "test-model", + CHOOSEKIT_MODE: "minimal-prefix", + }, "CHOOSEKIT_MODE must be labels when using Ollama."], ["OpenRouter without API key", { CHOOSEKIT_BACKEND: "openrouter", CHOOSEKIT_MODEL: "test/model", }, "OPENROUTER_API_KEY is required."], @@ -195,6 +207,60 @@ test("serves an OpenRouter choice with key, model, and provider kept process-loc assert.doesNotMatch(stderr, new RegExp(secret)); }); +test("serves an Ollama image choice through its configured chat endpoint", async (t) => { + const imageRoot = await mkdtemp(join(tmpdir(), "choosekit-mcp-images-")); + const image = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); + await writeFile(join(imageRoot, "example.png"), image); + t.after(() => rm(imageRoot, { recursive: true, force: true })); + const requests = []; + const ollama = createServer(async (request, response) => { + let body = ""; + for await (const chunk of request) body += chunk; + requests.push({ url: request.url, body: JSON.parse(body) }); + response.setHeader("content-type", "application/json"); + response.end(JSON.stringify({ + model: "fixture-model", + message: { role: "assistant", content: "B" }, + done: true, + done_reason: "length", + prompt_eval_count: 42, + eval_count: 1, + logprobs: [{ + token: "B", + bytes: [66], + logprob: -0.1, + top_logprobs: [ + { token: "A", bytes: [65], logprob: -1.1 }, + { token: "B", bytes: [66], logprob: -0.1 }, + ], + }], + })); + }); + await new Promise((resolve) => ollama.listen(0, "127.0.0.1", resolve)); + t.after(() => new Promise((resolve) => ollama.close(resolve))); + const { port } = ollama.address(); + const process_ = startCli(cliEnvironment({ + CHOOSEKIT_BACKEND: "ollama", + CHOOSEKIT_BASE_URL: `http://127.0.0.1:${port}`, + CHOOSEKIT_IMAGE_ROOT: imageRoot, + CHOOSEKIT_MODEL: "fixture-model", + })); + t.after(() => process_.close()); + await initialize(process_); + + const response = await process_.rpc("tools/call", { + name: "choose", + arguments: { ...request, imagePaths: ["example.png"] }, + }); + + assert.equal(response.result.structuredContent.choice, "technical"); + assert.deepEqual(JSON.parse(response.result.content[0].text), response.result.structuredContent); + assert.equal(requests.length, 1); + assert.equal(requests[0].url, "/api/chat"); + assert.equal(requests[0].body.model, "fixture-model"); + assert.deepEqual(requests[0].body.messages[0].images, [image.toString("base64")]); +}); + test("serves a minimal-prefix choice using the configured llama.cpp endpoint", async (t) => { const requests = []; const llama = createServer(async (request, response) => { diff --git a/packages/choosekit-mcp/tests/images.test.mjs b/packages/choosekit-mcp/tests/images.test.mjs new file mode 100644 index 0000000..f855941 --- /dev/null +++ b/packages/choosekit-mcp/tests/images.test.mjs @@ -0,0 +1,51 @@ +import assert from "node:assert/strict"; +import { mkdir, mkdtemp, rm, symlink, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { test } from "node:test"; +import { createImageLoader } from "../dist/images.js"; + +const PNG = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); +const JPEG = Buffer.from([0xff, 0xd8, 0xff, 0x00]); +const WEBP = Buffer.from([0x52, 0x49, 0x46, 0x46, 0, 0, 0, 0, 0x57, 0x45, 0x42, 0x50]); + +test("loads supported images in path order", async (t) => { + const root = await mkdtemp(join(tmpdir(), "choosekit-images-")); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, "first.bin"), WEBP); + await writeFile(join(root, "second.bin"), PNG); + await writeFile(join(root, "third.bin"), JPEG); + + const images = await createImageLoader(root)(["first.bin", "second.bin", "third.bin"]); + + assert.deepEqual(images.map(({ mediaType }) => mediaType), + ["image/webp", "image/png", "image/jpeg"]); + assert.deepEqual(images.map(({ base64 }) => base64), + [WEBP, PNG, JPEG].map((data) => data.toString("base64"))); +}); + +test("rejects paths outside the image root, including symlink escapes", async (t) => { + const parent = await mkdtemp(join(tmpdir(), "choosekit-images-")); + t.after(() => rm(parent, { recursive: true, force: true })); + const root = join(parent, "root"); + const outside = join(parent, "outside"); + await mkdir(root); + await mkdir(outside); + await writeFile(join(outside, "image.png"), PNG); + await symlink(outside, join(root, "escape"), process.platform === "win32" ? "junction" : "dir"); + const load = createImageLoader(root); + + await assert.rejects(load([join("..", "outside", "image.png")]), /outside/i); + await assert.rejects(load([join("escape", "image.png")]), /outside/i); +}); + +test("rejects directories and unsupported file contents", async (t) => { + const root = await mkdtemp(join(tmpdir(), "choosekit-images-")); + t.after(() => rm(root, { recursive: true, force: true })); + await mkdir(join(root, "directory")); + await writeFile(join(root, "not-an-image.png"), "plain text"); + const load = createImageLoader(root); + + await assert.rejects(load(["directory"]), /regular file/i); + await assert.rejects(load(["not-an-image.png"]), /PNG, JPEG, or WebP/i); +}); diff --git a/packages/choosekit-mcp/tests/server.test.mjs b/packages/choosekit-mcp/tests/server.test.mjs index d840dba..fcacb6d 100644 --- a/packages/choosekit-mcp/tests/server.test.mjs +++ b/packages/choosekit-mcp/tests/server.test.mjs @@ -2,6 +2,7 @@ import assert from "node:assert/strict"; import { test } from "node:test"; import { InMemoryTransport } from "@modelcontextprotocol/server"; import { createChooser, ScoringError } from "choosekit"; +import { ImageLoadError } from "../dist/images.js"; import { buildServer } from "../dist/server.js"; const request = Object.freeze({ @@ -114,6 +115,79 @@ test("marks the OpenRouter tool as open-world", async (t) => { }); }); +test("exposes Ollama as a closed-world labels backend with at most 20 choices", async (t) => { + const connection = await connect(async () => assert.fail("chooser should not run"), { + backend: "ollama", + }); + t.after(() => connection.close()); + + const listed = await connection.rpc("tools/list"); + const [tool] = listed.result.tools; + assert.deepEqual(tool.annotations, { readOnlyHint: true, openWorldHint: false }); + assert.equal(tool.inputSchema.properties.choices.maxProperties, 20); +}); + +test("passes loaded images to the chooser in labels mode", async (t) => { + const image = Object.freeze({ mediaType: "image/png", base64: "aW1hZ2U=" }); + const seenPaths = []; + let receivedImages; + const imageLoader = async (paths) => { + seenPaths.push(...paths); + return [image]; + }; + const chooser = async ({ choices, images }) => { + receivedImages = images; + return { + choice: "billing", + distribution: { billing: 1, technical: 0 }, + scores: { billing: -0.1, technical: -2 }, + margin: 1, + entropy: 0, + boundaryTokens: 1, + }; + }; + const connection = await connect(chooser, { backend: "ollama", imageLoader }); + t.after(() => connection.close()); + + const listed = await connection.rpc("tools/list"); + assert.ok("imagePaths" in listed.result.tools[0].inputSchema.properties); + const response = await callChoose(connection, { ...request, imagePaths: ["screen.png"] }); + + assert.notEqual(response.result.isError, true); + assert.deepEqual(seenPaths, ["screen.png"]); + assert.deepEqual(receivedImages, [image]); +}); + +test("reports image loading failures without exposing local paths", async (t) => { + const secretPath = "C:\\Users\\example\\private.png"; + const connection = await connect(async () => assert.fail("chooser should not run"), { + imageLoader: async () => { + throw new ImageLoadError(`could not read ${secretPath}`); + }, + }); + t.after(() => connection.close()); + + const response = await callChoose(connection, { ...request, imagePaths: ["missing.png"] }); + + assert.equal(response.result.isError, true); + assert.equal(response.result.content[0].text, + "The image could not be loaded. Check imagePaths and CHOOSEKIT_IMAGE_ROOT."); + assert.doesNotMatch(JSON.stringify(response), /private\.png/); +}); + +test("does not expose imagePaths without an image loader or in minimal-prefix mode", async (t) => { + for (const options of [{}, { mode: "minimal-prefix", imageLoader: async () => [] }]) { + await t.test(JSON.stringify(options), async (t) => { + const connection = await connect(async () => assert.fail("chooser should not run"), options); + t.after(() => connection.close()); + const listed = await connection.rpc("tools/list"); + assert.ok(!("imagePaths" in listed.result.tools[0].inputSchema.properties)); + const response = await callChoose(connection, { ...request, imagePaths: ["screen.png"] }); + assert.equal(response.result.isError, true); + }); + } +}); + test("OpenRouter accepts 20 choices and preserves the selected choice key", async (t) => { const choices = Object.fromEntries(Array.from({ length: 20 }, (_, index) => [`choice_${index}`, `Choice ${index}`]));