diff --git a/AGENTS.md b/AGENTS.md index c2ba57d..9f6f270 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,7 +19,8 @@ enter this repository, its history, or the packaged VSIX. manager, a second lockfile, or edit the lockfile by hand. - Node.js 22 or newer for development and CI. The bundle targets `node20` because that is the runtime shipped by the supported VS Code extension host. -- `engines.vscode` is `^1.104.0`, the release that introduced the language model chat provider API. +- `engines.vscode` is `^1.106.0`, the minimum version that provides the image/data chat-part API used by the + extension. ## Architecture @@ -49,6 +50,7 @@ Every command below exists in `package.json`; do not invent others. | Type-check | `npm run check-types` | | Lint (code only) | `npm run lint:code` | | Lint (types + code + format) | `npm run lint` | +| Dependency audit | `npm run audit` | | Format / verify | `npm run format` / `npm run format:check` | | Offline tests | `npm test` | | Live tests | `npm run test:live` (requires `CHUTES_KEY`) | @@ -57,8 +59,8 @@ Every command below exists in `package.json`; do not invent others. | Package a VSIX | `npm run vsix` | | Run in the editor | `F5` in VS Code (Extension Development Host) | -`npm run check` is the gate: strict TypeScript, Oxlint with `--deny-warnings`, Prettier verification, the offline -test suite, a production bundle, and `vsce ls` to confirm the VSIX contents. +`npm run check` is the gate: strict TypeScript, Oxlint with `--deny-warnings`, Prettier verification, a high-severity +dependency audit, the offline test suite, a production bundle, and `vsce ls` to confirm the VSIX contents. ## Release and publishing diff --git a/CHANGELOG.md b/CHANGELOG.md index aa3cc6c..c4cb5a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,30 @@ ## Unreleased +### Fixed + +- Propagated model-list cancellation and prevented invalidated in-flight requests from restoring stale cache data. +- Reserved `model-router` for the virtual **Auto (router)** entry so a catalogue collision cannot create duplicate + models or route a normal model through the router endpoint. +- Corrected daily quota normalization for mixed unlimited/finite entries, negative API values, and account payloads + that provide quota usage without a separate quota list. +- Raised the minimum supported VS Code version to 1.106 because image attachments use `LanguageModelDataPart`, which + is not present in the previously declared 1.104 API. + +### Changed + +- Bounded model, account, error, streaming-event, tool-call, and quota-fallback payload processing; quota fallback + requests now use limited concurrency. +- Potentially expensive model-filter regular expressions now fall back to literal substring matching. +- Pinned the VS Code API types to the declared minimum and added a high-severity dependency audit to the quality gate. + +### Security + +- Tool calls now fail closed unless their streamed id, type, name, arguments, and advertised availability are valid; + missing ids are no longer synthesized. +- Updated vulnerable transitive development dependencies used by VSIX packaging (`brace-expansion`, `fast-uri`, + `js-yaml`, and `undici`). + ## 0.4.4 - 2026-08-01 ### Fixed diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 23a4d19..06327d9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -6,7 +6,7 @@ Thanks for your interest in improving Chutes AI — Chat Model Provider. - Node.js 22 or newer - `npm` -- Visual Studio Code 1.104 or newer +- Visual Studio Code 1.106 or newer ## Development setup diff --git a/README.md b/README.md index 340a900..2038d03 100644 --- a/README.md +++ b/README.md @@ -41,18 +41,18 @@ You can also set the key anytime via **`Chutes AI: Manage API Key`** in the Comm ## Requirements -- **VS Code 1.104.0 or newer** (the language model provider API). VS Code **1.125+** also lets you discover this extension from the _Language Models_ editor via **Install Model Providers**. +- **VS Code 1.106.0 or newer** (required for image/data chat parts). VS Code **1.125+** also lets you discover this extension from the _Language Models_ editor via **Install Model Providers**. - A **Chutes API key** (starts with `cpk_`). Create one at [chutes.ai](https://chutes.ai). ## Settings -| Setting | Default | Description | -| -------------------------- | ---------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `chutes.endpoint` | `https://llm.chutes.ai/v1` | OpenAI-compatible API base URL. Change only for self-hosted or proxy endpoints. | -| `chutes.modelFilter` | _(empty)_ | Restrict which models appear. Comma-separated terms matched against the model id as a case-insensitive substring or regex (e.g. `deepseek, qwen` or `Qwen3.*TEE`). Empty shows all chat models. | -| `chutes.requestTimeoutMs` | `15000` | Timeout (ms) for fetching the model list. Does not limit streaming responses. | -| `chutes.autoRouterEnabled` | `true` | Show the **Auto (router)** model that delegates selection and automatic cold/unavailable fallback to Chutes' native router. | -| `chutes.routerEndpoint` | `https://model-router-ten.vercel.app/v1` | Base URL of Chutes' native router, used by the **Auto (router)** model. Change only for a self-hosted router. | +| Setting | Default | Description | +| -------------------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `chutes.endpoint` | `https://llm.chutes.ai/v1` | OpenAI-compatible API base URL. Change only for self-hosted or proxy endpoints. | +| `chutes.modelFilter` | _(empty)_ | Restrict which models appear. Comma-separated terms use a case-insensitive substring or safe regex (e.g. `deepseek, qwen` or `Qwen3.*TEE`); unsafe or invalid regexes are treated literally. Empty shows all chat models. | +| `chutes.requestTimeoutMs` | `15000` | Timeout (ms) for fetching the model list. Does not limit streaming responses. | +| `chutes.autoRouterEnabled` | `true` | Show the **Auto (router)** model that delegates selection and automatic cold/unavailable fallback to Chutes' native router. | +| `chutes.routerEndpoint` | `https://model-router-ten.vercel.app/v1` | Base URL of Chutes' native router, used by the **Auto (router)** model. Change only for a self-hosted router. | Changes to any `chutes.*` setting invalidate the model cache immediately; no window reload is required. diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index e44c716..6ecb3d2 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -5,7 +5,7 @@ - Make sure an API key is set: run **`Chutes AI: Manage API Key`**. - Run **`Chutes AI: Refresh Models`** to re-fetch the list. - Check `chutes.modelFilter` — an over-strict filter can hide everything. Clear it to show all chat models. -- Confirm you are on **VS Code 1.104+**. +- Confirm you are on **VS Code 1.106+**. ## "Could not load models" error diff --git a/docs/user-guide.md b/docs/user-guide.md index eca2f04..3962505 100644 --- a/docs/user-guide.md +++ b/docs/user-guide.md @@ -25,13 +25,13 @@ Models that accept image input (the picker marks them via their capabilities) ca ## Settings -| Setting | Default | Description | -| -------------------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------- | -| `chutes.endpoint` | `https://llm.chutes.ai/v1` | OpenAI-compatible API base URL. | -| `chutes.modelFilter` | _(empty)_ | Comma-separated terms (substring or regex) matched against model ids to narrow the picker. Example: `deepseek, qwen` or `Qwen3.*TEE`. | -| `chutes.requestTimeoutMs` | `15000` | Timeout for fetching the model list. | -| `chutes.autoRouterEnabled` | `true` | Show the **Auto (router)** model with automatic model selection and fallback. | -| `chutes.routerEndpoint` | `https://model-router-ten.vercel.app/v1` | OpenAI-compatible endpoint used only by **Auto (router)**. | +| Setting | Default | Description | +| -------------------------- | ---------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `chutes.endpoint` | `https://llm.chutes.ai/v1` | OpenAI-compatible API base URL. | +| `chutes.modelFilter` | _(empty)_ | Comma-separated terms (substring or safe regex) matched against model ids. Unsafe or invalid regexes are treated literally. Example: `deepseek, qwen` or `Qwen3.*TEE`. | +| `chutes.requestTimeoutMs` | `15000` | Timeout for fetching the model list. | +| `chutes.autoRouterEnabled` | `true` | Show the **Auto (router)** model with automatic model selection and fallback. | +| `chutes.routerEndpoint` | `https://model-router-ten.vercel.app/v1` | OpenAI-compatible endpoint used only by **Auto (router)**. | Setting changes refresh the model list immediately. A custom model or router endpoint receives your API key and request content; configure only services you trust. diff --git a/package-lock.json b/package-lock.json index a1f7d8e..48c1683 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10,7 +10,7 @@ "license": "MIT", "devDependencies": { "@types/node": "^22.20.1", - "@types/vscode": "^1.104.0", + "@types/vscode": "1.106.0", "@vscode/vsce": "^3.2.0", "esbuild": "^0.28.1", "oxlint": "^1.76.0", @@ -19,7 +19,7 @@ }, "engines": { "node": ">=22", - "vscode": "^1.104.0" + "vscode": "^1.106.0" } }, "node_modules/@azu/format-text": { @@ -839,9 +839,6 @@ "arm64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -859,9 +856,6 @@ "arm64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -879,9 +873,6 @@ "ppc64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -899,9 +890,6 @@ "riscv64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -919,9 +907,6 @@ "riscv64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -939,9 +924,6 @@ "s390x" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -959,9 +941,6 @@ "x64" ], "dev": true, - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -979,9 +958,6 @@ "x64" ], "dev": true, - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -1355,9 +1331,9 @@ "license": "MIT" }, "node_modules/@types/vscode": { - "version": "1.125.0", - "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.125.0.tgz", - "integrity": "sha512-0icm/ZQAaism87P0ekHqi4/Ju9du+Tm0RUW+y7vqRsxY2cY0FNRX1nAnaW7nT6npPt2tfHiheZ55Zm9UhqonFA==", + "version": "1.106.0", + "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.106.0.tgz", + "integrity": "sha512-88oUcEl9Wmlyt64pbvjLzyyFuFuPotdjwy+P+5ggg3DyTJSMWJD3ShX2ppya5mqrAYTKEhcaJBerdc5JTeb32w==", "dev": true, "license": "MIT" }, @@ -2091,9 +2067,9 @@ "license": "BSD-2-Clause" }, "node_modules/brace-expansion": { - "version": "5.0.8", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.8.tgz", - "integrity": "sha512-JZyDyq3D4AUifKTPOB7DELf6XsB3WdPuNxCtob1vFXPsSXhdAiHBWJ/tJ8HAc9aH84BK+5JFZLNkJKx3G9kzQg==", + "version": "5.0.9", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.9.tgz", + "integrity": "sha512-ScQ4IuvIEF1TMlP7Zt+vjJ//9zlPb2SDcxWxM3bk8s6t6GGdJ7KO1dCcTidOPJKePW30LE/2cT7wCyPho9/Wxg==", "dev": true, "license": "MIT", "dependencies": { @@ -2754,9 +2730,9 @@ } }, "node_modules/fast-uri": { - "version": "3.1.4", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.4.tgz", - "integrity": "sha512-8JnbkQ4juDyvYs4mgFGQqg4yCYtFDtUtmp2QIQq11ZZe5CFQ5wcqm1rqDgAh/QdMySuBnPzMUiJUNZG5N/AiQw==", + "version": "3.1.5", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", + "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", "dev": true, "funding": [ { @@ -3282,9 +3258,9 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.0.tgz", - "integrity": "sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==", + "version": "4.3.1", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.1.tgz", + "integrity": "sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==", "dev": true, "funding": [ { @@ -4928,9 +4904,9 @@ "license": "MIT" }, "node_modules/undici": { - "version": "7.28.0", - "resolved": "https://registry.npmjs.org/undici/-/undici-7.28.0.tgz", - "integrity": "sha512-cRZYrTDwWznlnRiPjggAGxZXanty6M8RV1ff8Wm4LWXBp7/IG8v5DnOm74DtUBp9OONpK75YlPnIjQqX0dBDtA==", + "version": "7.29.0", + "resolved": "https://registry.npmjs.org/undici/-/undici-7.29.0.tgz", + "integrity": "sha512-IDxfleLmmbSskfWSUATiN1nfn2rDuvnMOqb5CWR92iIfojA0Ud+ulOAAEQ57LPr9rWmsreUyf5lwyao+7GNNVw==", "dev": true, "license": "MIT", "engines": { diff --git a/package.json b/package.json index d92dd63..a8c92ec 100644 --- a/package.json +++ b/package.json @@ -30,7 +30,7 @@ }, "engines": { "node": ">=22", - "vscode": "^1.104.0" + "vscode": "^1.106.0" }, "categories": [ "AI", @@ -105,7 +105,7 @@ "chutes.modelFilter": { "type": "string", "default": "", - "markdownDescription": "Restrict which models appear in the picker. Comma-separated terms; each is matched against the model id as a case-insensitive substring or a regular expression (e.g. `deepseek, qwen` or `Qwen3.*TEE`). Leave empty to show all chat models." + "markdownDescription": "Restrict which models appear in the picker. Comma-separated terms; each is matched against the model id as a case-insensitive substring or a safe regular expression (e.g. `deepseek, qwen` or `Qwen3.*TEE`). Unsafe or invalid expressions are treated as literal substrings. Leave empty to show all chat models." }, "chutes.requestTimeoutMs": { "type": "number", @@ -133,18 +133,19 @@ "format:check": "prettier --check .", "lint:code": "oxlint --deny-warnings src test esbuild.js prettier.config.cjs", "lint": "npm run check-types && npm run lint:code && npm run format:check", + "audit": "npm audit --audit-level=high", "compile": "npm run check-types && node esbuild.js", "watch": "node esbuild.js --watch", "package": "npm run check-types && node esbuild.js --production", "test:build": "node test/build.js", "test": "npm run test:build && node --test test/unit.test.cjs", "test:live": "esbuild test/harness.ts --bundle --platform=node --format=cjs --outfile=test/harness.cjs --alias:vscode=./test/vscode-stub.cjs && node test/harness.cjs", - "check": "npm run lint && npm test && node esbuild.js --production && vsce ls", + "check": "npm run lint && npm run audit && npm test && node esbuild.js --production && vsce ls", "vsix": "vsce package" }, "devDependencies": { "@types/node": "^22.20.1", - "@types/vscode": "^1.104.0", + "@types/vscode": "1.106.0", "@vscode/vsce": "^3.2.0", "esbuild": "^0.28.1", "oxlint": "^1.76.0", diff --git a/src/chutesClient.ts b/src/chutesClient.ts index dd9acf0..5e69dd2 100644 --- a/src/chutesClient.ts +++ b/src/chutesClient.ts @@ -1,4 +1,10 @@ import { ChutesConfig } from './config'; +import { readResponseTextLimited } from './http'; + +const MAX_MODEL_RESPONSE_BYTES = 5 * 1024 * 1024; +const MAX_ERROR_RESPONSE_BYTES = 8 * 1024; +const MAX_SSE_LINE_CHARS = 1024 * 1024; +const MAX_MODEL_ID_CHARS = 512; /** Shape of a single entry returned by `GET /v1/models` (only fields we use). */ export interface ChutesRawModel { @@ -18,6 +24,7 @@ export interface ChutesRawModel { /** A streamed delta from `POST /v1/chat/completions` (OpenAI-compatible shape). */ export interface ChatCompletionDelta { content?: string; + toolCallError?: string; tool_calls?: Array<{ index: number; id?: string; @@ -43,9 +50,15 @@ export class ChutesClient { ) {} /** Fetches the full model catalogue. Aborts after the configured timeout. */ - async listModels(apiKey: string): Promise { + async listModels(apiKey: string, signal?: AbortSignal): Promise { const { endpoint, requestTimeoutMs } = this.config(); const controller = new AbortController(); + const abort = () => controller.abort(signal?.reason); + if (signal?.aborted) { + abort(); + } else { + signal?.addEventListener('abort', abort, { once: true }); + } const timer = setTimeout(() => controller.abort(), requestTimeoutMs); try { const res = await this.request(`${endpoint}/models`, { @@ -55,10 +68,20 @@ export class ChutesClient { if (!res.ok) { throw new ChutesApiError(await describeError(res, 'GET /models'), res.status); } - const json = (await res.json()) as { data?: unknown }; - return Array.isArray(json?.data) ? json.data.filter(isRawModel) : []; + const payload = await readResponseTextLimited(res, MAX_MODEL_RESPONSE_BYTES); + if (payload.truncated) { + throw new ChutesApiError(`Chutes: GET /models response exceeded ${MAX_MODEL_RESPONSE_BYTES} bytes`); + } + let json: unknown; + try { + json = JSON.parse(payload.text) as unknown; + } catch { + throw new ChutesApiError('Chutes: GET /models returned invalid JSON'); + } + return isRecord(json) && Array.isArray(json.data) ? json.data.filter(isRawModel) : []; } finally { clearTimeout(timer); + signal?.removeEventListener('abort', abort); } } @@ -103,9 +126,20 @@ export class ChutesClient { const lines = buffer.split('\n'); // Keep the last (possibly partial) line in the buffer. buffer = lines.pop() ?? ''; + if (buffer.length > MAX_SSE_LINE_CHARS) { + throw new ChutesApiError(`Chutes: SSE event exceeded ${MAX_SSE_LINE_CHARS} characters`); + } for (const rawLine of lines) { + if (rawLine.length > MAX_SSE_LINE_CHARS) { + throw new ChutesApiError(`Chutes: SSE event exceeded ${MAX_SSE_LINE_CHARS} characters`); + } const event = parseStreamLine(rawLine); if (event?.done) { + try { + await reader.cancel(); + } catch { + /* the response may already be closed */ + } return; } if (event?.delta) { @@ -117,6 +151,9 @@ export class ChutesClient { // Flush the decoder and process a final event even when the server closes // the stream without a trailing newline. buffer += decoder.decode(); + if (buffer.length > MAX_SSE_LINE_CHARS) { + throw new ChutesApiError(`Chutes: SSE event exceeded ${MAX_SSE_LINE_CHARS} characters`); + } const finalEvent = parseStreamLine(buffer); if (finalEvent?.delta) { yield finalEvent.delta; @@ -128,7 +165,12 @@ export class ChutesClient { } function isRawModel(value: unknown): value is ChutesRawModel { - return isRecord(value) && typeof value.id === 'string' && value.id.trim().length > 0; + return ( + isRecord(value) && + typeof value.id === 'string' && + value.id.trim().length > 0 && + value.id.length <= MAX_MODEL_ID_CHARS + ); } function parseStreamLine( @@ -165,13 +207,36 @@ function normalizeDelta(value: unknown): ChatCompletionDelta | undefined { if (typeof value.content === 'string') { delta.content = value.content; } - if (Array.isArray(value.tool_calls)) { + if (value.tool_calls !== undefined && !Array.isArray(value.tool_calls)) { + delta.toolCallError = 'malformed tool calls'; + } else if (Array.isArray(value.tool_calls)) { const toolCalls: NonNullable = []; for (const item of value.tool_calls) { if (!isRecord(item) || typeof item.index !== 'number' || !Number.isInteger(item.index) || item.index < 0) { + delta.toolCallError = 'malformed tool call metadata'; + continue; + } + if (item.id !== undefined && typeof item.id !== 'string') { + delta.toolCallError = 'malformed tool call id'; + continue; + } + if (item.type !== undefined && typeof item.type !== 'string') { + delta.toolCallError = 'malformed tool call type'; + continue; + } + if (item.function !== undefined && !isRecord(item.function)) { + delta.toolCallError = 'malformed tool call function'; continue; } const fn = isRecord(item.function) ? item.function : undefined; + if (fn?.name !== undefined && typeof fn.name !== 'string') { + delta.toolCallError = 'malformed tool call name'; + continue; + } + if (fn?.arguments !== undefined && typeof fn.arguments !== 'string') { + delta.toolCallError = 'malformed tool call arguments'; + continue; + } toolCalls.push({ index: item.index, id: typeof item.id === 'string' ? item.id : undefined, @@ -189,7 +254,9 @@ function normalizeDelta(value: unknown): ChatCompletionDelta | undefined { } } - return delta.content !== undefined || delta.tool_calls !== undefined ? delta : undefined; + return delta.content !== undefined || delta.tool_calls !== undefined || delta.toolCallError !== undefined + ? delta + : undefined; } function isRecord(value: unknown): value is Record { @@ -199,7 +266,11 @@ function isRecord(value: unknown): value is Record { async function describeError(res: Response, op: string): Promise { let detail = ''; try { - detail = (await res.text()).slice(0, 500); + const payload = await readResponseTextLimited(res, MAX_ERROR_RESPONSE_BYTES); + detail = payload.text.slice(0, 500); + if (payload.truncated) { + detail += '…'; + } } catch { /* ignore */ } diff --git a/src/http.ts b/src/http.ts new file mode 100644 index 0000000..4cc7c0d --- /dev/null +++ b/src/http.ts @@ -0,0 +1,48 @@ +export interface LimitedResponseText { + text: string; + truncated: boolean; +} + +/** Reads at most `maxBytes` from a response body and cancels the remainder. */ +export async function readResponseTextLimited(response: Response, maxBytes: number): Promise { + if (!response.body) { + return { text: '', truncated: false }; + } + + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let text = ''; + let bytesRead = 0; + let truncated = false; + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) { + text += decoder.decode(); + break; + } + + const remaining = maxBytes - bytesRead; + if (value.byteLength > remaining) { + if (remaining > 0) { + text += decoder.decode(value.subarray(0, remaining), { stream: true }); + } + truncated = true; + try { + await reader.cancel(); + } catch { + /* the response may already be closed */ + } + break; + } + + bytesRead += value.byteLength; + text += decoder.decode(value, { stream: true }); + } + } finally { + reader.releaseLock(); + } + + return { text, truncated }; +} diff --git a/src/modelMapping.ts b/src/modelMapping.ts index 63b2842..3181369 100644 --- a/src/modelMapping.ts +++ b/src/modelMapping.ts @@ -2,6 +2,7 @@ import * as vscode from 'vscode'; import { ChutesRawModel } from './chutesClient'; const FALLBACK_CONTEXT = 32768; +const MAX_FILTER_TERM_CHARS = 128; /** Reserved model id of Chutes' native router; selecting it delegates routing + fallback. */ export const AUTO_MODEL_ID = 'model-router'; @@ -54,17 +55,71 @@ export function applyUserFilter(models: ChutesRawModel[], filter: string): Chute return models; } const matchers = terms.map((term) => { - try { - const re = new RegExp(term, 'i'); - return (id: string) => re.test(id); - } catch { - const lower = term.toLowerCase(); - return (id: string) => id.toLowerCase().includes(lower); + if (isSafeRegex(term)) { + try { + const re = new RegExp(term, 'i'); + return (id: string) => re.test(id); + } catch { + // Invalid expressions retain the documented substring fallback. + } } + const lower = term.toLowerCase(); + return (id: string) => id.toLowerCase().includes(lower); }); return models.filter((m) => matchers.some((match) => match(m.id))); } +/** Conservatively excludes regex constructs that can cause excessive backtracking. */ +function isSafeRegex(pattern: string): boolean { + if (pattern.length > MAX_FILTER_TERM_CHARS || /\\[1-9]/.test(pattern) || pattern.includes('(?')) { + return false; + } + + const groups: Array<{ riskyContent: boolean }> = []; + let escaped = false; + let inCharacterClass = false; + for (let index = 0; index < pattern.length; index++) { + const char = pattern[index]; + if (escaped) { + escaped = false; + continue; + } + if (char === '\\') { + escaped = true; + continue; + } + if (char === '[') { + inCharacterClass = true; + continue; + } + if (char === ']' && inCharacterClass) { + inCharacterClass = false; + continue; + } + if (inCharacterClass) { + continue; + } + if (char === '(') { + groups.push({ riskyContent: false }); + continue; + } + if (char === '|' || char === '*' || char === '+' || char === '{') { + for (const group of groups) { + group.riskyContent = true; + } + continue; + } + if (char === ')' && groups.length > 0) { + const group = groups.pop(); + const next = pattern[index + 1]; + if (group?.riskyContent && (next === '*' || next === '+' || next === '?' || next === '{')) { + return false; + } + } + } + return true; +} + /** Maps a Chutes model to the VS Code chat model descriptor. */ export function toChatInformation(m: ChutesRawModel): vscode.LanguageModelChatInformation { const features = stringArrayOrDefault(m.supported_features, []); diff --git a/src/provider.ts b/src/provider.ts index 3446198..e921132 100644 --- a/src/provider.ts +++ b/src/provider.ts @@ -6,6 +6,10 @@ import { isChatModel, applyUserFilter, toChatInformation, autoRouterInfo, AUTO_M import { convertMessages, convertTools, convertToolMode, messageToText } from './messageConverter'; const CACHE_TTL_MS = 5 * 60 * 1000; +const MAX_TOOL_CALLS = 64; +const MAX_TOOL_CALL_ID_CHARS = 512; +const MAX_TOOL_NAME_CHARS = 256; +const MAX_TOOL_ARGUMENT_CHARS = 1024 * 1024; /** Sampling parameters we forward from VS Code's modelOptions to the Chutes API. */ const PASSTHROUGH_OPTIONS = [ @@ -25,6 +29,7 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider readonly onDidChangeLanguageModelChatInformation = this.changed.event; private cache?: { at: number; models: vscode.LanguageModelChatInformation[] }; + private cacheGeneration = 0; private pendingKeyPrompt?: Thenable; constructor( @@ -34,14 +39,18 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider /** Drops the cached model list and asks VS Code to re-query. */ invalidate(): void { + this.cacheGeneration++; this.cache = undefined; this.changed.fire(); } async provideLanguageModelChatInformation( options: { silent: boolean }, - _token: vscode.CancellationToken + token: vscode.CancellationToken ): Promise { + if (token.isCancellationRequested) { + return []; + } let apiKey = await this.secrets.get(); if (!apiKey) { if (options.silent) { @@ -55,14 +64,27 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider } } + if (token.isCancellationRequested) { + return []; + } + if (this.cache && Date.now() - this.cache.at < CACHE_TTL_MS) { return this.cache.models; } + const generation = this.cacheGeneration; + const controller = new AbortController(); + const cancellation = token.onCancellationRequested(() => controller.abort()); try { const cfg = getConfig(); - const raw = await this.client.listModels(apiKey); - const models = applyUserFilter(raw.filter(isChatModel), cfg.modelFilter) + const raw = await this.client.listModels(apiKey, controller.signal); + if (token.isCancellationRequested || generation !== this.cacheGeneration) { + return []; + } + const models = applyUserFilter( + raw.filter((candidate) => candidate.id !== AUTO_MODEL_ID).filter(isChatModel), + cfg.modelFilter + ) .map(toChatInformation) .sort((a, b) => a.id.localeCompare(b.id)); // Pin the virtual "Auto" entry at the top so it is easy to find in the picker. @@ -70,11 +92,16 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider this.cache = { at: Date.now(), models: withAuto }; return withAuto; } catch (err) { + if (controller.signal.aborted || token.isCancellationRequested || generation !== this.cacheGeneration) { + return []; + } if (!options.silent) { const msg = err instanceof ChutesApiError ? err.message : String(err); void vscode.window.showErrorMessage(`Chutes AI: could not load models. ${msg}`); } return []; + } finally { + cancellation.dispose(); } } @@ -85,10 +112,16 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider progress: vscode.Progress, token: vscode.CancellationToken ): Promise { + if (token.isCancellationRequested) { + return; + } const apiKey = await this.secrets.get(); if (!apiKey) { throw new Error('Chutes AI: no API key configured. Run "Chutes AI: Manage API Key".'); } + if (token.isCancellationRequested) { + return; + } // The virtual "Auto" model routes to Chutes' native router endpoint instead of // the configured one; everything else (OpenAI-compatible body, SSE) is identical. @@ -115,29 +148,58 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider } const controller = new AbortController(); + if (token.isCancellationRequested) { + controller.abort(); + } const cancel = token.onCancellationRequested(() => controller.abort()); // Tool-call fragments arrive split across deltas, keyed by index; assemble then emit. - const toolCalls = new Map(); + const toolCalls = new Map(); + const availableToolNames = new Set(options.tools?.map((tool) => tool.name) ?? []); try { for await (const delta of this.client.streamChatCompletion(apiKey, body, controller.signal, endpointOverride)) { if (token.isCancellationRequested) { break; } + if (delta.toolCallError) { + throw new Error(`Chutes AI: ${delta.toolCallError} returned by the model.`); + } if (delta.content) { progress.report(new vscode.LanguageModelTextPart(delta.content)); } if (delta.tool_calls) { for (const tc of delta.tool_calls) { - const current = toolCalls.get(tc.index) ?? { id: '', name: '', args: '' }; - if (tc.id) { + if (tc.index >= MAX_TOOL_CALLS || (!toolCalls.has(tc.index) && toolCalls.size >= MAX_TOOL_CALLS)) { + throw new Error(`Chutes AI: response exceeded the ${MAX_TOOL_CALLS}-tool-call limit.`); + } + const current = toolCalls.get(tc.index) ?? { args: '' }; + if (tc.id !== undefined) { + if (!tc.id || tc.id.length > MAX_TOOL_CALL_ID_CHARS || (current.id && current.id !== tc.id)) { + throw new Error('Chutes AI: malformed tool call id returned by the model.'); + } current.id = tc.id; } - if (tc.function?.name) { + if (tc.type !== undefined) { + if (tc.type !== 'function' || (current.type && current.type !== tc.type)) { + throw new Error('Chutes AI: unsupported tool call type returned by the model.'); + } + current.type = tc.type; + } + if (tc.function?.name !== undefined) { + if ( + !tc.function.name || + tc.function.name.length > MAX_TOOL_NAME_CHARS || + (current.name && current.name !== tc.function.name) + ) { + throw new Error('Chutes AI: malformed tool name returned by the model.'); + } current.name = tc.function.name; } - if (tc.function?.arguments) { + if (tc.function?.arguments !== undefined) { + if (current.args.length + tc.function.arguments.length > MAX_TOOL_ARGUMENT_CHARS) { + throw new Error('Chutes AI: tool arguments exceeded the supported size limit.'); + } current.args += tc.function.arguments; } toolCalls.set(tc.index, current); @@ -160,8 +222,11 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider } for (const [index, call] of Array.from(toolCalls.entries()).sort(([a], [b]) => a - b)) { - if (!call.name) { - continue; + if (!call.id || call.type !== 'function' || !call.name) { + throw new Error(`Chutes AI: incomplete tool call returned at index ${index}.`); + } + if (!availableToolNames.has(call.name)) { + throw new Error(`Chutes AI: unavailable tool "${call.name}" requested by the model.`); } let input: unknown = {}; try { @@ -172,7 +237,7 @@ export class ChutesChatModelProvider implements vscode.LanguageModelChatProvider if (typeof input !== 'object' || input === null || Array.isArray(input)) { throw new Error(`Chutes AI: non-object arguments returned for tool "${call.name}".`); } - progress.report(new vscode.LanguageModelToolCallPart(call.id || `chutes-tool-${index}`, call.name, input)); + progress.report(new vscode.LanguageModelToolCallPart(call.id, call.name, input)); } } diff --git a/src/usage/accountClient.ts b/src/usage/accountClient.ts index 66042c2..9c1c448 100644 --- a/src/usage/accountClient.ts +++ b/src/usage/accountClient.ts @@ -2,9 +2,14 @@ // chutes-usage project. The Chutes `cpk_` key authenticates here too (verified). // No `vscode` dependency. import type { JsonContainer, JsonObject } from './types'; +import { readResponseTextLimited } from '../http'; const API_BASE_URL = 'https://api.chutes.ai'; const REQUEST_TIMEOUT_MS = 15000; +const MAX_ACCOUNT_RESPONSE_BYTES = 5 * 1024 * 1024; +const MAX_QUOTA_USAGE_REQUESTS = 50; +const QUOTA_USAGE_CONCURRENCY = 4; +const MAX_CHUTE_ID_CHARS = 512; export interface DashboardPayload { subscriptionUsage: JsonObject; @@ -15,7 +20,10 @@ export interface DashboardPayload { } export class ChutesAccountClient { - constructor(private readonly apiKey: string) {} + constructor( + private readonly apiKey: string, + private readonly request: typeof fetch = globalThis.fetch + ) {} /** Fetches the account/usage endpoints needed to summarize spend and quotas. */ async getDashboardPayload(signal?: AbortSignal): Promise { @@ -44,7 +52,7 @@ export class ChutesAccountClient { } const timeout = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS); try { - const response = await fetch(`${API_BASE_URL}${path}`, { + const response = await this.request(`${API_BASE_URL}${path}`, { method: 'GET', headers: { Authorization: this.apiKey, Accept: 'application/json' }, signal: controller.signal @@ -52,7 +60,16 @@ export class ChutesAccountClient { if (!response.ok) { throw new Error(`Chutes account API: ${path} failed (HTTP ${response.status})`); } - const json = (await response.json()) as unknown; + const payload = await readResponseTextLimited(response, MAX_ACCOUNT_RESPONSE_BYTES); + if (payload.truncated) { + throw new Error(`Chutes account API: ${path} response exceeded ${MAX_ACCOUNT_RESPONSE_BYTES} bytes`); + } + let json: unknown; + try { + json = JSON.parse(payload.text) as unknown; + } catch { + throw new Error(`Chutes account API: ${path} returned invalid JSON`); + } if (!isJsonContainer(json)) { throw new Error(`Unexpected API response shape for ${path}`); } @@ -65,16 +82,29 @@ export class ChutesAccountClient { private async getQuotaUsagePayload(quotas: JsonContainer, signal?: AbortSignal): Promise { const chuteIds = getQuotaUsageChuteIds(quotas); - if (chuteIds.length === 0) { + if (chuteIds.length === 0 || chuteIds.length > MAX_QUOTA_USAGE_REQUESTS) { return null; } - const entries = await Promise.all( - chuteIds.map(async (chuteId) => { + const entries: Array = Array.from({ length: chuteIds.length }, () => null); + let nextIndex = 0; + const workers = Array.from({ length: Math.min(QUOTA_USAGE_CONCURRENCY, chuteIds.length) }, async () => { + while (nextIndex < chuteIds.length) { + if (signal?.aborted) { + throw signal.reason instanceof Error ? signal.reason : new Error('Chutes account API request cancelled'); + } + const index = nextIndex++; + const chuteId = chuteIds[index]; const path = `/users/me/quota_usage/${encodePathSegment(chuteId)}`; - const payload = await this.getJsonContainer(path, signal).catch(() => null); - return payload === null ? null : ([chuteId, payload] as const); - }) - ); + const payload = await this.getJsonContainer(path, signal).catch((error: unknown) => { + if (signal?.aborted) { + throw error; + } + return null; + }); + entries[index] = payload === null ? null : ([chuteId, payload] as const); + } + }); + await Promise.all(workers); const valid = entries.filter((entry): entry is readonly [string, JsonContainer] => entry !== null); return valid.length === 0 ? null : Object.fromEntries(valid); } @@ -92,21 +122,22 @@ function hasQuotaUsageData(payload: JsonContainer | null): boolean { if (payload === null || Array.isArray(payload)) { return false; } - if (isFiniteNumberLike(payload.used) || isFiniteNumberLike(payload.quota)) { + if (isNonNegativeNumberLike(payload.used) || isNonNegativeNumberLike(payload.quota)) { return true; } return Object.values(payload).some((value) => { const object = isJsonObject(value) ? value : null; - return isFiniteNumberLike(object?.used) || isFiniteNumberLike(object?.quota); + return isNonNegativeNumberLike(object?.used) || isNonNegativeNumberLike(object?.quota); }); } -function isFiniteNumberLike(value: unknown): boolean { +function isNonNegativeNumberLike(value: unknown): boolean { if (typeof value === 'number') { - return Number.isFinite(value); + return Number.isFinite(value) && value >= 0; } if (typeof value === 'string' && value.trim().length > 0) { - return Number.isFinite(Number(value)); + const parsed = Number(value); + return Number.isFinite(parsed) && parsed >= 0; } return false; } @@ -122,9 +153,15 @@ function getQuotaUsageChuteIds(payload: JsonContainer): string[] { const chuteIds = new Set(); for (const item of items) { const object = isJsonObject(item) ? item : null; - const chuteId = typeof object?.chute_id === 'string' && object.chute_id.length > 0 ? object.chute_id : null; + const chuteId = + typeof object?.chute_id === 'string' && object.chute_id.length > 0 && object.chute_id.length <= MAX_CHUTE_ID_CHARS + ? object.chute_id + : null; if (chuteId) { chuteIds.add(chuteId); + if (chuteIds.size > MAX_QUOTA_USAGE_REQUESTS) { + break; + } } } return Array.from(chuteIds); diff --git a/src/usage/normalize.ts b/src/usage/normalize.ts index fc2daa6..01b3ab1 100644 --- a/src/usage/normalize.ts +++ b/src/usage/normalize.ts @@ -182,6 +182,7 @@ export function normalizeQuotaUsage(payload: JsonContainer | null): QuotaUsageSu } let usedTotal: number | null = null; let quotaTotal: number | null = null; + let hasUnlimitedQuota = false; for (const value of Object.values(payload)) { const item = asObject(value); if (!item) { @@ -193,9 +194,16 @@ export function normalizeQuotaUsage(payload: JsonContainer | null): QuotaUsageSu usedTotal = (usedTotal ?? 0) + used; } if (quota !== null) { - quotaTotal = (quotaTotal ?? 0) + quota; + if (quota === 0) { + hasUnlimitedQuota = true; + } else if (!hasUnlimitedQuota) { + quotaTotal = (quotaTotal ?? 0) + quota; + } } } + if (hasUnlimitedQuota) { + quotaTotal = 0; + } if (usedTotal !== null || quotaTotal !== null) { return { used: usedTotal, quota: quotaTotal, trusted: true }; } @@ -280,16 +288,18 @@ function buildDailyQuotaWindow( if (entry.quota === null) { return sum; } + if (entry.quota === 0 || sum === 0) { + return 0; + } return (sum ?? 0) + entry.quota; }, null); - if (totalQuota === null) { - return []; - } - const preferredQuotaUsage = quotaUsageMe ?? quotaUsageFallback; const liveTotalRequests = invocationStats?.totalRequests ?? 0; const limit = preferredQuotaUsage?.quota ?? totalQuota; + if (limit === null) { + return []; + } const isUnlimited = limit === 0; const isStale = preferredQuotaUsage?.used === 0 && liveTotalRequests > 0; const used = isStale ? null : preferredQuotaUsage?.trusted ? preferredQuotaUsage.used : null; @@ -381,9 +391,10 @@ function countUsageSignals(payload: JsonObject): number { function findBestUsageObject( payload: JsonObject, - seen = new Set() + seen = new Set(), + depth = 0 ): { object: JsonObject; score: number } | null { - if (seen.has(payload)) { + if (seen.has(payload) || depth >= 50) { return null; } seen.add(payload); @@ -397,7 +408,7 @@ function findBestUsageObject( if (!child) { continue; } - const nestedBest = findBestUsageObject(child, seen); + const nestedBest = findBestUsageObject(child, seen, depth + 1); if (nestedBest && (best === null || nestedBest.score > best.score)) { best = nestedBest; } @@ -416,12 +427,12 @@ function pickObject(payload: JsonObject, aliases: string[]): { key: string; valu } function asNumber(value: unknown): number | null { - if (typeof value === 'number' && Number.isFinite(value)) { + if (typeof value === 'number' && Number.isFinite(value) && value >= 0) { return value; } if (typeof value === 'string' && value.trim().length > 0) { const parsed = Number(value); - return Number.isFinite(parsed) ? parsed : null; + return Number.isFinite(parsed) && parsed >= 0 ? parsed : null; } return null; } diff --git a/test/unit.test.ts b/test/unit.test.ts index b7a08ce..3a900c3 100644 --- a/test/unit.test.ts +++ b/test/unit.test.ts @@ -12,10 +12,13 @@ import { ChutesChatModelProvider } from '../src/provider'; import { SecretStore } from '../src/secrets'; import { formatUsageMarkdown, formatQuotasMarkdown } from '../src/chatParticipant'; import { normalizeDashboardData } from '../src/usage/normalize'; +import { normalizeQuotaUsage } from '../src/usage/normalize'; +import { ChutesAccountClient } from '../src/usage/accountClient'; import type { DashboardData } from '../src/usage/types'; import type { ChutesRawModel } from '../src/chutesClient'; import { DEFAULT_ROUTER_ENDPOINT } from '../src/config'; import { ChutesClient } from '../src/chutesClient'; +import { readResponseTextLimited } from '../src/http'; function model(partial: Partial & { id: string }): ChutesRawModel { return { input_modalities: ['text'], output_modalities: ['text'], ...partial }; @@ -57,6 +60,17 @@ test('applyUserFilter matches substrings, regex, and comma lists', () => { ); }); +test('applyUserFilter treats potentially catastrophic regexes as literal text', () => { + // Build the hostile pattern as test data so static analyzers do not mistake it + // for an expression this test executes directly. + const unsafePattern = String.fromCharCode(40, 97, 43, 41, 43, 36); + const models = [model({ id: `literal/${unsafePattern}` }), model({ id: `x/${'a'.repeat(64)}!` })]; + assert.deepEqual( + applyUserFilter(models, unsafePattern).map((entry) => entry.id), + [`literal/${unsafePattern}`] + ); +}); + test('toChatInformation maps fields and capabilities', () => { const info = toChatInformation( model({ @@ -113,6 +127,13 @@ test('ChutesClient drops malformed model rows', async () => { assert.deepEqual(await client.listModels('cpk_test'), [{ id: 'valid/model' }]); }); +test('response reader caps payloads and cancels the remainder', async () => { + assert.deepEqual(await readResponseTextLimited(new Response('123456'), 4), { + text: '1234', + truncated: true + }); +}); + test('ChutesClient parses the final SSE event without a trailing newline', async () => { const request = (async () => new Response('data: {"choices":[{"delta":{"content":"done"}}]}', { @@ -136,6 +157,102 @@ test('ChutesClient parses the final SSE event without a trailing newline', async assert.deepEqual(deltas, [{ content: 'done' }]); }); +test('ChutesClient propagates model-list cancellation to fetch', async () => { + let requestSignal: AbortSignal | undefined; + const request = ((_input: string | URL | Request, init?: RequestInit) => { + requestSignal = init?.signal ?? undefined; + return new Promise((_resolve, reject) => { + requestSignal?.addEventListener( + 'abort', + () => reject(requestSignal?.reason ?? new DOMException('The operation was aborted', 'AbortError')), + { once: true } + ); + }); + }) as typeof fetch; + const client = new ChutesClient( + () => ({ + endpoint: 'https://example.test/v1', + modelFilter: '', + requestTimeoutMs: 1000, + autoRouterEnabled: true, + routerEndpoint: DEFAULT_ROUTER_ENDPOINT + }), + request + ); + const controller = new AbortController(); + const pending = client.listModels('cpk_test', controller.signal); + controller.abort(new DOMException('The operation was aborted', 'AbortError')); + await assert.rejects(pending, { name: 'AbortError' }); + assert.equal(requestSignal?.aborted, true); +}); + +test('ChutesClient rejects oversized SSE events', async () => { + const request = (async () => + new Response(`data: ${'x'.repeat(1024 * 1024)}\n`, { + status: 200, + headers: { 'Content-Type': 'text/event-stream' } + })) as typeof fetch; + const client = new ChutesClient( + () => ({ + endpoint: 'https://example.test/v1', + modelFilter: '', + requestTimeoutMs: 1000, + autoRouterEnabled: true, + routerEndpoint: DEFAULT_ROUTER_ENDPOINT + }), + request + ); + await assert.rejects(async () => { + for await (const _delta of client.streamChatCompletion('cpk_test', {}, new AbortController().signal)) { + // No valid delta should be emitted from an oversized event. + } + }, /SSE event exceeded/); +}); + +test('ChutesClient marks malformed streamed tool-call fields', async () => { + const request = (async () => + new Response(`data: ${JSON.stringify({ choices: [{ delta: { tool_calls: [{ index: 0, id: 42 }] } }] })}\n`, { + headers: { 'Content-Type': 'text/event-stream' } + })) as typeof fetch; + const client = new ChutesClient( + () => ({ + endpoint: 'https://example.test/v1', + modelFilter: '', + requestTimeoutMs: 1000, + autoRouterEnabled: true, + routerEndpoint: DEFAULT_ROUTER_ENDPOINT + }), + request + ); + const deltas = []; + for await (const delta of client.streamChatCompletion('cpk_test', {}, new AbortController().signal)) { + deltas.push(delta); + } + assert.deepEqual(deltas, [{ toolCallError: 'malformed tool call id' }]); +}); + +test('ChutesClient rejects a non-array streamed tool_calls field', async () => { + const request = (async () => + new Response(`data: ${JSON.stringify({ choices: [{ delta: { tool_calls: {} } }] })}\n`, { + headers: { 'Content-Type': 'text/event-stream' } + })) as typeof fetch; + const client = new ChutesClient( + () => ({ + endpoint: 'https://example.test/v1', + modelFilter: '', + requestTimeoutMs: 1000, + autoRouterEnabled: true, + routerEndpoint: DEFAULT_ROUTER_ENDPOINT + }), + request + ); + const deltas = []; + for await (const delta of client.streamChatCompletion('cpk_test', {}, new AbortController().signal)) { + deltas.push(delta); + } + assert.deepEqual(deltas, [{ toolCallError: 'malformed tool calls' }]); +}); + test('autoRouterInfo describes the virtual router model', () => { const info = autoRouterInfo(); assert.equal(info.id, AUTO_MODEL_ID); @@ -220,7 +337,10 @@ function memSecrets(initial?: string): SecretStore { } const RAW: ChutesRawModel[] = [model({ id: 'a/Chat-One', supported_features: ['tools'], context_length: 8000 })]; -const noToken = {} as never; +const noToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose() {} }) +} as never; test('provider: silent + no key returns [] and never prompts', async () => { let prompts = 0; @@ -317,6 +437,52 @@ test('provider: omits the Auto model when autoRouterEnabled is false', async () } }); +test('provider: reserves model-router for the virtual Auto model', async () => { + const raw = [model({ id: AUTO_MODEL_ID }), ...RAW]; + const provider = new ChutesChatModelProvider(memSecrets('cpk_test'), fakeClient(raw)); + const info = await provider.provideLanguageModelChatInformation({ silent: false }, noToken); + assert.equal(info.filter((entry) => entry.id === AUTO_MODEL_ID).length, 1); + assert.equal(info.length, RAW.length + 1); +}); + +test('provider: cancellation prevents model-list requests', async () => { + let requests = 0; + const provider = new ChutesChatModelProvider(memSecrets('cpk_test'), { + listModels: async () => (requests++, RAW) + } as never); + const cancelledToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose() {} }) + } as never; + assert.deepEqual(await provider.provideLanguageModelChatInformation({ silent: false }, cancelledToken), []); + assert.equal(requests, 0); +}); + +test('provider: invalidation prevents an in-flight model request from restoring stale cache', async () => { + let resolveFirst!: (models: ChutesRawModel[]) => void; + const firstResponse = new Promise((resolve) => { + resolveFirst = resolve; + }); + let requests = 0; + const client = { + listModels: async () => { + requests++; + if (requests === 1) { + return firstResponse; + } + return RAW; + } + } as never; + const provider = new ChutesChatModelProvider(memSecrets('cpk_test'), client); + const pending = provider.provideLanguageModelChatInformation({ silent: false }, noToken); + await new Promise((resolve) => setImmediate(resolve)); + provider.invalidate(); + resolveFirst(RAW); + assert.deepEqual(await pending, []); + assert.ok((await provider.provideLanguageModelChatInformation({ silent: false }, noToken)).length > 0); + assert.equal(requests, 2); +}); + test('provider: Auto model streams via the router endpoint; normal models do not', async () => { const captured: Array = []; const client = { @@ -349,7 +515,7 @@ test('provider: malformed tool arguments fail closed', async () => { listModels: async () => RAW, async *streamChatCompletion() { yield { - tool_calls: [{ index: 0, id: 'call_1', function: { name: 'dangerous_tool', arguments: '{' } }] + tool_calls: [{ index: 0, id: 'call_1', type: 'function', function: { name: 'dangerous_tool', arguments: '{' } }] }; } } as never; @@ -358,11 +524,53 @@ test('provider: malformed tool arguments fail closed', async () => { const token = { isCancellationRequested: false, onCancellationRequested: () => ({ dispose() {} }) } as never; await assert.rejects( - provider.provideLanguageModelChatResponse(autoRouterInfo(), [], {} as never, progress, token), + provider.provideLanguageModelChatResponse( + autoRouterInfo(), + [], + { tools: [{ name: 'dangerous_tool', description: '', inputSchema: {} }] } as never, + progress, + token + ), /invalid arguments returned for tool/ ); }); +test('provider: rejects unsupported, incomplete, and unavailable tool calls', async () => { + const cases = [ + { + call: { index: 0, id: 'call_1', type: 'not-a-function', function: { name: 'safe_tool', arguments: '{}' } }, + error: /unsupported tool call type/ + }, + { + call: { index: 0, type: 'function', function: { name: 'safe_tool', arguments: '{}' } }, + error: /incomplete tool call/ + }, + { + call: { index: 0, id: 'call_1', type: 'function', function: { name: 'other_tool', arguments: '{}' } }, + error: /unavailable tool/ + } + ]; + for (const { call, error } of cases) { + const client = { + listModels: async () => RAW, + async *streamChatCompletion() { + yield { tool_calls: [call] }; + } + } as never; + const provider = new ChutesChatModelProvider(memSecrets('cpk_test'), client); + await assert.rejects( + provider.provideLanguageModelChatResponse( + autoRouterInfo(), + [], + { tools: [{ name: 'safe_tool', description: '', inputSchema: {} }] } as never, + { report() {} } as never, + noToken + ), + error + ); + } +}); + test('provider: cancelling mid-stream drops half-assembled tool calls', async () => { let cancelled = false; const listeners: Array<() => void> = []; @@ -371,7 +579,14 @@ test('provider: cancelling mid-stream drops half-assembled tool calls', async () async *streamChatCompletion() { // Truncated JSON, exactly what a stream cut short by the user looks like. yield { - tool_calls: [{ index: 0, id: 'call_1', function: { name: 'dangerous_tool', arguments: '{"path": "sr' } }] + tool_calls: [ + { + index: 0, + id: 'call_1', + type: 'function', + function: { name: 'dangerous_tool', arguments: '{"path": "sr' } + } + ] }; cancelled = true; for (const listener of listeners) { @@ -394,7 +609,13 @@ test('provider: cancelling mid-stream drops half-assembled tool calls', async () } as never; // Must resolve (no "invalid arguments" error) and must not run the tool. - await provider.provideLanguageModelChatResponse(autoRouterInfo(), [], {} as never, progress, token); + await provider.provideLanguageModelChatResponse( + autoRouterInfo(), + [], + { tools: [{ name: 'dangerous_tool', description: '', inputSchema: {} }] } as never, + progress, + token + ); assert.ok(reported.every((part) => !(part instanceof vscode.LanguageModelToolCallPart))); }); @@ -515,3 +736,112 @@ test('normalizeDashboardData parses spend windows and derives the plan', () => { assert.equal(billing?.limit, 100); assert.ok(data.windows.some((w) => w.kind === 'daily-requests')); }); + +test('quota normalization preserves unlimited and rejects negative API values', () => { + assert.deepEqual(normalizeQuotaUsage({ a: { used: 3, quota: 100 }, b: { used: 2, quota: 0 } }), { + used: 5, + quota: 0, + trusted: true + }); + assert.deepEqual(normalizeQuotaUsage({ used: -1, quota: -10 }), null); + + const data = normalizeDashboardData( + {}, + [ + { model: 'Unlimited', quota: 0 }, + { model: 'Finite', quota: 100 }, + { model: 'Invalid', quota: -1 } + ] as never, + null, + null, + null + ); + const daily = data.windows.find((window) => window.kind === 'daily-requests'); + assert.equal(daily?.limit, 0); + assert.equal(data.quotas[2].quota, null); +}); + +test('preferred quota usage can provide a daily window when the quota list is empty', () => { + const data = normalizeDashboardData({}, [], null, { used: 2, quota: 100 }, null); + const daily = data.windows.find((window) => window.kind === 'daily-requests'); + assert.equal(daily?.used, 2); + assert.equal(daily?.limit, 100); +}); + +test('account client caps quota fallback fan-out', async () => { + let perChuteRequests = 0; + const quotas = Array.from({ length: 51 }, (_, index) => ({ chute_id: `chute-${index}`, quota: 10 })); + const request = (async (input: string | URL | Request) => { + const path = new URL(String(input)).pathname; + if (path === '/users/me/subscription_usage') { + return Response.json({}); + } + if (path === '/users/me/quotas') { + return Response.json(quotas); + } + if (path === '/users/me/quota_usage/me') { + return Response.json({}); + } + if (path === '/invocations/stats/llm') { + return Response.json([]); + } + perChuteRequests++; + return Response.json({ used: 0, quota: 10 }); + }) as typeof fetch; + const payload = await new ChutesAccountClient('cpk_test', request).getDashboardPayload(); + assert.equal(payload.quotaUsageFallback, null); + assert.equal(perChuteRequests, 0); +}); + +test('account client limits concurrent quota fallback requests', async () => { + let active = 0; + let maxActive = 0; + const quotas = Array.from({ length: 12 }, (_, index) => ({ chute_id: `chute-${index}`, quota: 10 })); + const request = (async (input: string | URL | Request) => { + const path = new URL(String(input)).pathname; + if (path === '/users/me/subscription_usage') { + return Response.json({}); + } + if (path === '/users/me/quotas') { + return Response.json(quotas); + } + if (path === '/users/me/quota_usage/me') { + return Response.json({}); + } + if (path === '/invocations/stats/llm') { + return Response.json([]); + } + active++; + maxActive = Math.max(maxActive, active); + await new Promise((resolve) => setTimeout(resolve, 2)); + active--; + return Response.json({ used: 0, quota: 10 }); + }) as typeof fetch; + const payload = await new ChutesAccountClient('cpk_test', request).getDashboardPayload(); + assert.equal(Object.keys(payload.quotaUsageFallback ?? {}).length, quotas.length); + assert.ok(maxActive <= 4); +}); + +test('account client ignores negative quota-usage sentinels and uses fallback data', async () => { + let fallbackRequests = 0; + const request = (async (input: string | URL | Request) => { + const path = new URL(String(input)).pathname; + if (path === '/users/me/subscription_usage') { + return Response.json({}); + } + if (path === '/users/me/quotas') { + return Response.json([{ chute_id: 'chute-1', quota: 10 }]); + } + if (path === '/users/me/quota_usage/me') { + return Response.json({ used: -1, quota: -1 }); + } + if (path === '/invocations/stats/llm') { + return Response.json([]); + } + fallbackRequests++; + return Response.json({ used: 2, quota: 10 }); + }) as typeof fetch; + const payload = await new ChutesAccountClient('cpk_test', request).getDashboardPayload(); + assert.equal(fallbackRequests, 1); + assert.deepEqual(payload.quotaUsageFallback, { 'chute-1': { used: 2, quota: 10 } }); +});