From c2a419b4f7945dcdd02938427ddfeeb115e8ab0d Mon Sep 17 00:00:00 2001 From: Michal Date: Tue, 25 Aug 2026 23:29:45 +0100 Subject: [PATCH] fix(mcplocal): one LLM budget per stage, so a loop cannot multiply it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A per-call timeout bounds one call. It does not bound a stage that calls the LLM in a loop -- and summarize-tree does worse than loop: buildTree recurses to maxDepth (3 by default) and iterates per section at every level, then groupSections iterates again. Hundreds of sequential calls are reachable, so a 10s per-call cap is a multiplier, not a bound. StageBudget is one wall-clock budget shared across everything a single stage invocation does. Once it is gone the stage takes its deterministic path -- first-line excerpts, numbered pages -- and says so. The executor builds one per stage from config.budgetMs ?? MCPCTL_STAGE_LLM_BUDGET_MS (30s) and disposes it in a finally, so a long-lived mcplocal never accumulates timers. Two details that matter more than they look: - A warm cache is never budget-gated. cachedSummarize and generatePageTitles used ctx.cache.getOrCompute, which makes the budget check impossible to place correctly; split into get/set so a cached summary -- which costs nothing -- is still returned when the budget is spent. - "No LLM configured" is NOT a degradation. It is a deliberate choice, and numbered pages are the expected output there, so it gets no warning. Only failures, timeouts and exhausted budgets do. Crying wolf on a working configuration would train people to ignore the notice. summarize-tree's single-block path previously had no try/catch at all, so a failure escaped the stage entirely and executor.ts discarded the whole stage's work. It now degrades like the others. Degradation is reported three ways, all from util/degrade.ts: the reason in the content behind the established "⚠ unavailable ()" prefix, the count of affected sections (the recursion can hit one exhausted budget dozens of times, so the user needs the reason once plus what it cost), and degraded/degradedReason on the stage_execution audit event -- previously a degradation was invisible in the trace. docs/reliability.md now names ONE helper rather than a list of compliant call sites. That list is exactly how this drifted: the doc named the gate and llm/pagination.ts, and the newer stages never joined it. Tests: the 20-section recursion completes in ~300ms against a 300ms budget; a warm cache still serves real titles with the budget exhausted. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GqMidYEGUJG5fxeoTELBu2 --- docs/reliability.md | 45 ++++++- src/mcplocal/src/proxymodel/executor.ts | 17 +++ src/mcplocal/src/proxymodel/stage-budget.ts | 50 +++++++ .../src/proxymodel/stages/paginate.ts | 96 ++++++++++---- .../src/proxymodel/stages/summarize-tree.ts | 108 +++++++++++++--- src/mcplocal/src/proxymodel/types.ts | 9 ++ src/mcplocal/tests/proxymodel-stages.test.ts | 2 + src/mcplocal/tests/proxymodel-types.test.ts | 2 + src/mcplocal/tests/stage-budget.test.ts | 122 ++++++++++++++++++ .../tests/system-prompt-fetching.test.ts | 2 + 10 files changed, 405 insertions(+), 48 deletions(-) create mode 100644 src/mcplocal/src/proxymodel/stage-budget.ts create mode 100644 src/mcplocal/tests/stage-budget.test.ts diff --git a/docs/reliability.md b/docs/reliability.md index 30423c8..0eae89f 100644 --- a/docs/reliability.md +++ b/docs/reliability.md @@ -19,10 +19,47 @@ and report the degradation — never hang and never degrade silently.** ()…` and sets `degraded: true` + `degradedReason` on the audit `gate_decision` event. -Applied in: the gate's `begin_session` prompt selection -(`proxymodel/plugins/gate.ts`, cap `MCPCTL_GATE_LLM_TIMEOUT_MS`, default 8s) and -pagination's smart index (`llm/pagination.ts`, `MCPCTL_PAGINATION_LLM_TIMEOUT_MS`, -default 10s). `read_prompts` is LLM-free by design. +**One implementation:** [`util/degrade.ts`](../src/mcplocal/src/util/degrade.ts)'s +`bounded()`. It cannot throw — the caller always gets a value or a reason — so +the deterministic fallback is unconditional rather than something a `catch` +block has to remember. `degradationNotice()` produces the `⚠ … unavailable +(reason). hint` wording, and `degradationAudit()` the `degraded`/`degradedReason` +payload, so every degradation reads the same whichever subsystem produced it. + +**Nothing optional is unbounded by construction.** The budget lives in +`LLMProviderAdapter.complete()` (`proxymodel/llm-adapter.ts`), not at each call +site, so a stage that passes no options is still bounded and a stage written +next year inherits the guarantee. One budget spans the whole failover chain — +a per-provider timeout would make the worst case N × timeout. + +| Knob | Default | Bounds | +|---|---|---| +| `MCPCTL_LLM_CALL_BUDGET_MS` | 20s | one `ctx.llm.complete()`, failover included | +| `MCPCTL_LLM_PROVIDER_TIMEOUT_MS` | 10s | a single provider attempt inside that budget | +| `MCPCTL_STAGE_LLM_BUDGET_MS` | 30s | all LLM work in one stage invocation | +| `MCPCTL_GATE_LLM_TIMEOUT_MS` | 8s | the gate's `begin_session` prompt selection | +| `MCPCTL_PAGINATION_LLM_TIMEOUT_MS` | 10s | pagination's smart index (`llm/pagination.ts`) | + +The **stage** budget exists because a per-call timeout multiplies rather than +bounds when a stage loops: `summarize-tree` recurses to `maxDepth` (3) and loops +per section at every level, so hundreds of sequential calls are reachable. One +budget is shared across the whole recursion; when it is gone the stage switches +to first-line excerpts and says so. A **warm cache is never budget-gated** — a +cached summary costs nothing, so an exhausted budget must not degrade a result +we already hold. + +`read_prompts` is LLM-free by design. + +### Why this is written down twice + +This document stated the principle while `proxymodel/stages/paginate.ts` awaited +`ctx.llm.complete()` with no timeout at all. When a provider hung rather than +erroring, the promise never settled, the tool call never returned, and the +client waited out its own 1800s timeout — three such requests in production, +misdiagnosed twice as an upstream "transport fault". The doc named the gate and +`llm/pagination.ts` as the compliant sites, and the newer stages simply never +joined the list. Naming **one** helper here, rather than a list of call sites, +is what stops that drift recurring. Note: the gate's prompt-ranking uses the **heavy client provider's own model** — it deliberately does *not* force the project's vLLM model onto it (doing so made diff --git a/src/mcplocal/src/proxymodel/executor.ts b/src/mcplocal/src/proxymodel/executor.ts index 9b3f69d..6f8e208 100644 --- a/src/mcplocal/src/proxymodel/executor.ts +++ b/src/mcplocal/src/proxymodel/executor.ts @@ -6,6 +6,7 @@ import type { StageContext, StageResult, StageLogger, Section, ContentType, LLMProvider, CacheProvider, SystemPromptFetcher } from './types.js'; import type { ProxyModelDefinition } from './schema.js'; import { getStage } from './stage-registry.js'; +import { createStageBudget, STAGE_LLM_BUDGET_MS } from './stage-budget.js'; import type { AuditCollector } from '../audit/collector.js'; export interface ExecuteOptions { @@ -87,6 +88,11 @@ export async function executePipeline(opts: ExecuteOptions): Promise { controller.abort(); }, totalMs); + // Node keeps the process alive for pending timers; a stage budget must not. + if (typeof timer.unref === 'function') timer.unref(); + + return { + totalMs, + signal: controller.signal, + remainingMs: (): number => Math.max(0, totalMs - (Date.now() - startedAt)), + exhausted: (): boolean => Date.now() - startedAt >= totalMs, + dispose: (): void => { clearTimeout(timer); }, + }; +} + +/** Never-expiring budget, for tests and non-LLM call paths. */ +export function unlimitedStageBudget(): StageBudget { + return { + totalMs: Number.POSITIVE_INFINITY, + signal: new AbortController().signal, + remainingMs: (): number => Number.POSITIVE_INFINITY, + exhausted: (): boolean => false, + dispose: (): void => { /* nothing to clean up */ }, + }; +} diff --git a/src/mcplocal/src/proxymodel/stages/paginate.ts b/src/mcplocal/src/proxymodel/stages/paginate.ts index 7b522e0..664f474 100644 --- a/src/mcplocal/src/proxymodel/stages/paginate.ts +++ b/src/mcplocal/src/proxymodel/stages/paginate.ts @@ -9,6 +9,7 @@ * previewChars: number (chars per page sent to LLM for title generation, default 300) */ import type { StageHandler, StageContext, Section } from '../types.js'; +import { bounded, degradationNotice } from '../../util/degrade.js'; const handler: StageHandler = async (content, ctx) => { const pageSize = (ctx.config.pageSize as number | undefined) ?? 8000; @@ -24,7 +25,7 @@ const handler: StageHandler = async (content, ctx) => { return { content }; } - const titles = await generatePageTitles(pages, ctx); + const { titles, degradedReason } = await generatePageTitles(pages, ctx); const sections: Section[] = pages.map((page, i) => ({ id: `page-${i + 1}`, @@ -36,65 +37,104 @@ const handler: StageHandler = async (content, ctx) => { `[${s.id}] ${s.title} (${pages[i]!.length} chars)`, ).join('\n'); + // Loud, not silent: if the titles are numbered because the LLM was slow or + // down, say so rather than letting the model assume the pages are unnamed. + const notice = degradedReason === null + ? '' + : degradationNotice('Smart page titles', degradedReason, 'Pages are numbered instead.'); + return { // No navigation hint here — the caller (content-pipeline / router) appends // the authoritative _resultId/_section instruction once sections are stored. - content: `Content split into ${sections.length} pages (${content.length} total chars):\n${toc}`, + content: `${notice}Content split into ${sections.length} pages (${content.length} total chars):\n${toc}`, sections, + ...(degradedReason === null ? {} : { metadata: { degraded: true, degradedReason } }), }; }; /** - * Generate descriptive titles for each page using LLM. - * Falls back to generic "Page N" titles if LLM is unavailable or fails. + * Generate descriptive titles for each page using the LLM. + * + * Returns a degradation reason instead of throwing: the deterministic + * "Page N" fallback is unconditional. Before this was bounded, a provider that + * hung (rather than erroring) left the awaited promise unsettled forever — the + * tool call never returned and the client waited out its own 1800s timeout. */ -async function generatePageTitles(pages: string[], ctx: StageContext): Promise { +async function generatePageTitles( + pages: string[], + ctx: StageContext, +): Promise<{ titles: string[]; degradedReason: string | null }> { const fallback = pages.map((_, i) => `Page ${i + 1}`); + // No LLM configured is a deliberate choice, not a degradation — numbered + // pages are the expected output there, so don't warn about them. if (!ctx.llm.available()) { - return fallback; + return { titles: fallback, degradedReason: null }; } const previewChars = (ctx.config.previewChars as number | undefined) ?? 300; const cacheKey = `paginate-titles:${ctx.cache.hash(ctx.originalContent)}:${pages.length}`; - try { - const cached = await ctx.cache.getOrCompute(cacheKey, async () => { - const previews = pages.map((page, i) => { - const preview = page.slice(0, previewChars).trim(); - return `--- Page ${i + 1} (${page.length} chars) ---\n${preview}`; - }).join('\n\n'); + // Read the cache BEFORE consulting the budget — a warm cache costs nothing, + // so an exhausted budget must not degrade a result we already hold. + const cached = await ctx.cache.get(cacheKey); + if (cached !== null) { + try { + return { titles: JSON.parse(cached) as string[], degradedReason: null }; + } catch { + // Corrupt entry — fall through and recompute. + } + } - const DEFAULT_PROMPT = `Generate exactly {{pageCount}} short descriptive titles (max 60 chars each) for the following {{pageCount}} pages. Return ONLY a JSON array of {{pageCount}} strings. No markdown, no explanation.`; - const template = await ctx.getSystemPrompt('llm-paginate-titles', DEFAULT_PROMPT); - const prompt = template.replaceAll('{{pageCount}}', String(pages.length)); + if (ctx.budget.exhausted()) { + return { + titles: fallback, + degradedReason: `stage LLM budget of ${String(ctx.budget.totalMs)}ms exhausted`, + }; + } - const result = await ctx.llm.complete( - `${prompt}\n\n${previews}`, - { maxTokens: pages.length * 30 }, - ); + const previews = pages.map((page, i) => { + const preview = page.slice(0, previewChars).trim(); + return `--- Page ${i + 1} (${page.length} chars) ---\n${preview}`; + }).join('\n\n'); - // Parse JSON array from response, pad/truncate to match page count + const DEFAULT_PROMPT = `Generate exactly {{pageCount}} short descriptive titles (max 60 chars each) for the following {{pageCount}} pages. Return ONLY a JSON array of {{pageCount}} strings. No markdown, no explanation.`; + const template = await ctx.getSystemPrompt('llm-paginate-titles', DEFAULT_PROMPT); + const prompt = template.replaceAll('{{pageCount}}', String(pages.length)); + + const outcome = await bounded( + async (signal) => { + const result = await ctx.llm.complete(`${prompt}\n\n${previews}`, { + maxTokens: pages.length * 30, + budgetMs: ctx.budget.remainingMs(), + signal, + }); const match = result.match(/\[[\s\S]*\]/); if (!match) throw new Error('No JSON array in response'); const raw = JSON.parse(match[0]) as string[]; if (!Array.isArray(raw) || raw.length === 0) { throw new Error('Empty or invalid title array'); } - // Pad with generic titles if model returned fewer, truncate if more - const titles = pages.map((_, i) => + return pages.map((_, i) => (i < raw.length && typeof raw[i] === 'string' && raw[i]!.trim()) ? raw[i]!.trim().slice(0, 80) : `Page ${i + 1}`, ); - return JSON.stringify(titles); - }); + }, + { + label: 'paginate', + operation: 'Smart page titles', + timeoutMs: ctx.budget.remainingMs(), + signal: ctx.budget.signal, + fallback: 'to numbered pages', + }, + ); - return JSON.parse(cached) as string[]; - } catch (err) { - ctx.log.warn(`Smart page titles failed, using generic: ${(err as Error).message}`); - return fallback; + if (outcome.degraded) { + return { titles: fallback, degradedReason: outcome.reason }; } + await ctx.cache.set(cacheKey, JSON.stringify(outcome.value)); + return { titles: outcome.value, degradedReason: null }; } function splitPages(content: string, pageSize: number): string[] { diff --git a/src/mcplocal/src/proxymodel/stages/summarize-tree.ts b/src/mcplocal/src/proxymodel/stages/summarize-tree.ts index b75e419..312de5f 100644 --- a/src/mcplocal/src/proxymodel/stages/summarize-tree.ts +++ b/src/mcplocal/src/proxymodel/stages/summarize-tree.ts @@ -14,11 +14,13 @@ */ import type { StageHandler, Section } from '../types.js'; import { detectContentType } from '../content-type.js'; +import { bounded, degradationNotice } from '../../util/degrade.js'; const handler: StageHandler = async (content, ctx) => { const maxTokens = (ctx.config.maxSummaryTokens as number | undefined) ?? 200; const maxGroup = (ctx.config.maxGroupSize as number | undefined) ?? 5; const maxDepth = (ctx.config.maxDepth as number | undefined) ?? 3; + const deg = createDegradeSink(); // If content is small, just return it unchanged if (content.length < 2000) { @@ -37,7 +39,16 @@ const handler: StageHandler = async (content, ctx) => { return { content }; } - const summary = await cachedSummarize(ctx, ctx.originalContent, maxTokens); + const summary = await cachedSummarize(ctx, ctx.originalContent, maxTokens, deg); + if (summary === null) { + // Budget gone or the LLM failed: return the content untouched rather + // than a summary we could not produce. Previously this path had no + // try/catch at all, so a failure escaped the stage entirely. + return { + content: degradationNotice('Smart summaries', deg.reason ?? 'unavailable', 'Content is unsummarised.') + content, + metadata: { degraded: true, degradedReason: deg.reason }, + }; + } return { content: summary + '\n\nSection "full" holds the complete content.', sections: [{ id: 'full', title: 'Full Content', content: ctx.originalContent }], @@ -45,7 +56,7 @@ const handler: StageHandler = async (content, ctx) => { } // Build the summary tree - const tree = await buildTree(sections, ctx, { maxTokens, maxGroup, maxDepth, depth: 0 }); + const tree = await buildTree(sections, ctx, { maxTokens, maxGroup, maxDepth, depth: 0, deg }); // Format top-level ToC const toc = tree.map((s) => { @@ -55,9 +66,20 @@ const handler: StageHandler = async (content, ctx) => { return `[${s.id}] ${s.title}${childHint}`; }).join('\n'); + // One reason, plus how many summaries it cost — the recursion can hit the + // same exhausted budget dozens of times. + const notice = deg.reason === null + ? '' + : degradationNotice( + 'Smart summaries', + deg.reason, + `Showing first-line excerpts for ${String(deg.count)} of ${String(tree.length)} sections.`, + ); + return { - content: `${tree.length} sections:\n${toc}`, + content: `${notice}${tree.length} sections:\n${toc}`, sections: tree, + ...(deg.reason === null ? {} : { metadata: { degraded: true, degradedReason: deg.reason } }), }; }; @@ -66,6 +88,8 @@ interface TreeOpts { maxGroup: number; maxDepth: number; depth: number; + /** Shared across the whole recursion — one budget, one reported reason. */ + deg: DegradeSink; } async function buildTree( @@ -84,7 +108,8 @@ async function buildTree( summary = structuralSummary(section.content, contentType); } else if (ctx.llm.available()) { // LLM summary for prose/code - summary = await cachedSummarize(ctx, section.content, opts.maxTokens); + summary = await cachedSummarize(ctx, section.content, opts.maxTokens, opts.deg) + ?? (section.content.split('\n')[0] ?? '').slice(0, 200); } else { // No LLM — use first line as summary summary = (section.content.split('\n')[0] ?? '').slice(0, 200); @@ -125,18 +150,68 @@ async function cachedSummarize( ctx: import('../types.js').StageContext, content: string, maxTokens: number, -): Promise { + deg: DegradeSink, +): Promise { const key = `summary:${ctx.cache.hash(content)}:${maxTokens}`; - return ctx.cache.getOrCompute(key, async () => { - const DEFAULT_PROMPT = `Summarize the following in about {{maxTokens}} tokens. Preserve all items marked MUST, REQUIRED, or CRITICAL verbatim. Be specific — mention names, IDs, counts, key values.`; - const template = await ctx.getSystemPrompt('llm-summarize', DEFAULT_PROMPT); - const prompt = template.replaceAll('{{maxTokens}}', String(maxTokens)); - return ctx.llm.complete( - `${prompt}\n\n${content}`, - { maxTokens }, - ); - }); + // Read the cache BEFORE consulting the budget: a warm entry costs nothing, + // so an exhausted budget must never degrade a summary we already hold. + const cached = await ctx.cache.get(key); + if (cached !== null) return cached; + + if (ctx.budget.exhausted()) { + deg.note(`stage LLM budget of ${String(ctx.budget.totalMs)}ms exhausted`); + return null; + } + + const DEFAULT_PROMPT = `Summarize the following in about {{maxTokens}} tokens. Preserve all items marked MUST, REQUIRED, or CRITICAL verbatim. Be specific — mention names, IDs, counts, key values.`; + const template = await ctx.getSystemPrompt('llm-summarize', DEFAULT_PROMPT); + const prompt = template.replaceAll('{{maxTokens}}', String(maxTokens)); + + const outcome = await bounded( + (signal) => ctx.llm.complete(`${prompt}\n\n${content}`, { + maxTokens, + budgetMs: ctx.budget.remainingMs(), + signal, + }), + { + label: 'summarize-tree', + operation: 'Smart summaries', + timeoutMs: ctx.budget.remainingMs(), + signal: ctx.budget.signal, + fallback: 'to first-line excerpts', + }, + ); + + if (outcome.degraded) { + deg.note(outcome.reason); + return null; + } + await ctx.cache.set(key, outcome.value); + return outcome.value; +} + +/** + * Collects the first degradation reason and a count. + * + * buildTree recurses and loops, so one exhausted budget can produce dozens of + * identical reasons; the user needs the reason once plus how many summaries it + * cost them. + */ +export interface DegradeSink { + note(reason: string): void; + readonly reason: string | null; + readonly count: number; +} + +export function createDegradeSink(): DegradeSink { + let reason: string | null = null; + let count = 0; + return { + note(r: string): void { reason ??= r; count++; }, + get reason(): string | null { return reason; }, + get count(): number { return count; }, + }; } function structuralSummary(content: string, type: string): string { @@ -267,9 +342,10 @@ async function groupSections( const groupContent = chunk.map((s) => `[${s.id}] ${s.title}`).join('\n'); const groupId = `group-${Math.floor(i / opts.maxGroup) + 1}`; + const genericTitle = `Group ${Math.floor(i / opts.maxGroup) + 1} (${chunk.length} sections)`; const groupTitle = ctx.llm.available() - ? await cachedSummarize(ctx, groupContent, 50) - : `Group ${Math.floor(i / opts.maxGroup) + 1} (${chunk.length} sections)`; + ? (await cachedSummarize(ctx, groupContent, 50, opts.deg) ?? genericTitle) + : genericTitle; groups.push({ id: groupId, diff --git a/src/mcplocal/src/proxymodel/types.ts b/src/mcplocal/src/proxymodel/types.ts index da3b8a8..4ad62c1 100644 --- a/src/mcplocal/src/proxymodel/types.ts +++ b/src/mcplocal/src/proxymodel/types.ts @@ -10,6 +10,8 @@ * SessionController — method-level hooks with per-session state */ +import type { StageBudget } from './stage-budget.js'; + /** Fetches a system prompt by name, falling back to the provided default. */ export type SystemPromptFetcher = (name: string, fallback: string) => Promise; @@ -48,6 +50,13 @@ export interface StageContext { /** Stage-specific configuration from the proxymodel YAML */ config: Record; + + /** + * Shared wall-clock budget for this stage's LLM work. A stage that calls the + * LLM in a loop must check this between calls and fall back deterministically + * once it is gone — a per-call timeout alone multiplies rather than bounds. + */ + budget: StageBudget; } export interface StageResult { diff --git a/src/mcplocal/tests/proxymodel-stages.test.ts b/src/mcplocal/tests/proxymodel-stages.test.ts index d58a8b7..f11b076 100644 --- a/src/mcplocal/tests/proxymodel-stages.test.ts +++ b/src/mcplocal/tests/proxymodel-stages.test.ts @@ -5,6 +5,7 @@ import paginate from '../src/proxymodel/stages/paginate.js'; import sectionSplit from '../src/proxymodel/stages/section-split.js'; import summarizeTree from '../src/proxymodel/stages/summarize-tree.js'; import { BUILT_IN_STAGES } from '../src/proxymodel/stages/index.js'; +import { unlimitedStageBudget } from '../src/proxymodel/stage-budget.js'; function mockCtx(original: string, config: Record = {}, llmAvailable = false): StageContext { const llmResponses: string[] = []; @@ -48,6 +49,7 @@ function mockCtx(original: string, config: Record = {}, llmAvai cache: mockCache, log: mockLog, getSystemPrompt: async (_name: string, fallback: string) => fallback, + budget: unlimitedStageBudget(), config, }; } diff --git a/src/mcplocal/tests/proxymodel-types.test.ts b/src/mcplocal/tests/proxymodel-types.test.ts index 2983475..7febbba 100644 --- a/src/mcplocal/tests/proxymodel-types.test.ts +++ b/src/mcplocal/tests/proxymodel-types.test.ts @@ -12,6 +12,7 @@ import type { SessionContext, ContentType, } from '../src/proxymodel/index.js'; +import { unlimitedStageBudget } from '../src/proxymodel/stage-budget.js'; describe('ProxyModel type contract', () => { it('StageHandler can be implemented as a simple function', async () => { @@ -137,6 +138,7 @@ function createMockContext(original: string): StageContext { cache: mockCache, log: mockLog, getSystemPrompt: async (_name: string, fallback: string) => fallback, + budget: unlimitedStageBudget(), config: {}, }; } diff --git a/src/mcplocal/tests/stage-budget.test.ts b/src/mcplocal/tests/stage-budget.test.ts new file mode 100644 index 0000000..9e96177 --- /dev/null +++ b/src/mcplocal/tests/stage-budget.test.ts @@ -0,0 +1,122 @@ +import { describe, it, expect, vi, afterEach } from 'vitest'; +import { createStageBudget, unlimitedStageBudget, STAGE_LLM_BUDGET_MS } from '../src/proxymodel/stage-budget.js'; +import { getStage } from '../src/proxymodel/stage-registry.js'; +import '../src/proxymodel/stages/index.js'; +import type { StageContext, LLMProvider, CacheProvider, StageLogger } from '../src/proxymodel/types.js'; + +afterEach(() => { vi.restoreAllMocks(); }); + +function ctxWith(llm: LLMProvider, original: string, budgetMs: number): StageContext { + const store = new Map(); + const cache: CacheProvider = { + async getOrCompute(k, c) { + if (store.has(k)) return store.get(k)!; + const v = await c(); store.set(k, v); return v; + }, + hash: (c) => c.slice(0, 8), + async get(k) { return store.get(k) ?? null; }, + async set(k, v) { store.set(k, v); }, + }; + const log: StageLogger = { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() }; + return { + contentType: 'toolResult', sourceName: 'test/tool', projectName: 'test', sessionId: 's1', + originalContent: original, llm, cache, log, + getSystemPrompt: async (_n, f) => f, + config: {}, + budget: createStageBudget(budgetMs), + }; +} + +/** Never settles and ignores the signal — the shape that caused the outage. */ +const hangingLlm: LLMProvider = { + complete: () => new Promise(() => { /* never */ }), + available: () => true, +}; + +describe('StageBudget', () => { + it('reports remaining time and expires', async () => { + const b = createStageBudget(60); + expect(b.exhausted()).toBe(false); + expect(b.remainingMs()).toBeGreaterThan(0); + await new Promise((r) => setTimeout(r, 90)); + expect(b.exhausted()).toBe(true); + expect(b.remainingMs()).toBe(0); + expect(b.signal.aborted).toBe(true); + b.dispose(); + }); + + it('unlimited budget never expires', () => { + const b = unlimitedStageBudget(); + expect(b.exhausted()).toBe(false); + expect(b.remainingMs()).toBe(Number.POSITIVE_INFINITY); + b.dispose(); + }); + + it('defaults to the documented 30s', () => { + expect(STAGE_LLM_BUDGET_MS).toBe(30_000); + }); +}); + +describe('paginate under a hung LLM', () => { + it('returns numbered pages within the budget instead of hanging', async () => { + vi.spyOn(console, 'error').mockImplementation(() => { /* quiet */ }); + const handler = getStage('paginate')!; + const content = 'x'.repeat(30_000); + const ctx = ctxWith(hangingLlm, content, 200); + + const start = Date.now(); + const result = await handler(content, ctx); + const elapsed = Date.now() - start; + + // The whole point: this used to never return at all. + expect(elapsed).toBeLessThan(3000); + expect(result.content).toContain('⚠ Smart page titles unavailable'); + expect(result.content).toContain('page-1'); + expect(result.metadata?.['degraded']).toBe(true); + ctx.budget.dispose(); + }); + + it('serves a warm cache even with an exhausted budget', async () => { + const handler = getStage('paginate')!; + const content = 'y'.repeat(30_000); + + // Warm the cache with a working LLM and a generous budget. + const good: LLMProvider = { + complete: async () => '["Alpha","Beta","Gamma","Delta"]', + available: () => true, + }; + const warm = ctxWith(good, content, 5000); + await handler(content, warm); + + // Same cache, no budget left: the cached titles must still be used. + const cold: StageContext = { ...warm, llm: hangingLlm, budget: createStageBudget(1) }; + await new Promise((r) => setTimeout(r, 10)); + const result = await handler(content, cold); + + expect(result.content).toContain('Alpha'); + expect(result.content).not.toContain('⚠'); + warm.budget.dispose(); + cold.budget.dispose(); + }); +}); + +describe('summarize-tree under a hung LLM', () => { + it('bounds the whole recursion, not each call', async () => { + vi.spyOn(console, 'error').mockImplementation(() => { /* quiet */ }); + const handler = getStage('summarize-tree')!; + // Many sections: pre-fix, a per-call timeout would multiply across the + // sequential loop AND the recursion (maxDepth 3). + const sections = Array.from({ length: 20 }, (_, i) => + `## Section ${i + 1}\n${'prose '.repeat(400)}`).join('\n\n'); + const ctx = ctxWith(hangingLlm, sections, 300); + + const start = Date.now(); + const result = await handler(sections, ctx); + const elapsed = Date.now() - start; + + expect(elapsed).toBeLessThan(5000); + expect(result.content).toContain('⚠ Smart summaries unavailable'); + expect(result.metadata?.['degraded']).toBe(true); + ctx.budget.dispose(); + }); +}); diff --git a/src/mcplocal/tests/system-prompt-fetching.test.ts b/src/mcplocal/tests/system-prompt-fetching.test.ts index 545b34b..8730a38 100644 --- a/src/mcplocal/tests/system-prompt-fetching.test.ts +++ b/src/mcplocal/tests/system-prompt-fetching.test.ts @@ -2,6 +2,7 @@ import { describe, it, expect, vi } from 'vitest'; import type { StageContext, LLMProvider, CacheProvider, StageLogger, SystemPromptFetcher } from '../src/proxymodel/types.js'; import paginate from '../src/proxymodel/stages/paginate.js'; import summarizeTree from '../src/proxymodel/stages/summarize-tree.js'; +import { unlimitedStageBudget } from '../src/proxymodel/stage-budget.js'; function mockCtx( original: string, @@ -50,6 +51,7 @@ function mockCtx( cache: mockCache, log: mockLog, getSystemPrompt: opts.getSystemPrompt ?? (async (_name, fallback) => fallback), + budget: unlimitedStageBudget(), config, }; }