fix(ai): repair evidenced prefill output-budget rejections

packages/ai/CHANGELOG.md

@@ -10,6 +10,8 @@
 
 ### Fixed
 
+- Recover once from an OpenAI-compatible prefill rejection that reports a consistent input/output context budget, reducing only the completion cap to the server-reported room while preserving conversation and reasoning settings ([#1616](https://github.com/code-yeongyu/senpi/pull/1616) by [@rlaope](https://github.com/rlaope)).
+
 ### Removed
 
 ## [2026.9.12-2] - 2026-09-12

packages/ai/src/api/openai-completions.ts

@@ -49,12 +49,13 @@ import { shortHash } from "../utils/hash.ts";
 import { headersToRecord } from "../utils/headers.ts";
 import { parseStreamingJson } from "../utils/json-parse.ts";
 import { getPiUserAgent } from "../utils/pi-user-agent.ts";
+import { repairedOutputBudget } from "../utils/prefill-budget-recovery.ts";
 import {
 	getOpenAICompletionsCompat as getCompat,
 	type ResolvedOpenAICompletionsCompat,
 } from "../utils/prompt-cache-ttl.ts";
 import { getProviderEnvValue } from "../utils/provider-env.ts";
-import { retryProviderStreamRequest } from "../utils/provider-retry.ts";
+import { retryProviderRequest, retryProviderStreamRequest } from "../utils/provider-retry.ts";
 import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts";
 import { isForcedToolChoiceUnsupportedError, omitToolChoiceParam } from "../utils/tool-choice-fallback.ts";
 import {
@@ -517,14 +518,37 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio
 					throw error;
 				}
 			};
-			const { stream: openaiStream } = await retryProviderStreamRequest(
-				async () => {
+			const startStream = () =>
+				retryProviderStreamRequest(async () => {
 					const { data, response } = await createRequest();
 					await options?.onResponse?.(
 						{ status: response.status, headers: headersToRecord(response.headers) },
 						model,
 					);
 					return { stream: data, metadata: response };
+				});
+			let repairedPrefill = false;
+			const { stream: openaiStream } = await retryProviderRequest(
+				async () => {
+					try {
+						return await startStream();
+					} catch (error) {
+						const maxTokens = repairedPrefill
+							? undefined
+							: repairedOutputBudget(error, {
+									requested: params.max_tokens ?? params.max_completion_tokens,
+									thinkingTokens: resolveClampedThinkingBudget(model, options, params) ?? 0,
+									signal: options?.signal,
+								});
+						if (maxTokens === undefined) throw error;
+						// Only the first-chunk prefetch can reach here: never replay visible output.
+						repairedPrefill = true;
+						params =
+							params.max_tokens != null
+								? { ...params, max_tokens: maxTokens }
+								: { ...params, max_completion_tokens: maxTokens };
+						return startStream();
+					}
 				},
 				{ maxRetries: options?.maxRetries, maxRetryDelayMs: options?.maxRetryDelayMs, signal: options?.signal },
 			);

packages/ai/src/changes.md

@@ -1,3 +1,24 @@
+## Evidence-based prefill output reservation repair (2026-09-12)
+
+### What changed
+
+- `packages/ai/src/api/openai-completions.ts`: before the first chunk is exposed, allow one completion-cap correction when a prefill 400 reports consistent context, input, total and completion counts matching the actual request. Preserve the built payload, hooks, context and reasoning; later stream failures never enter this repair.
+- `packages/ai/src/utils/prefill-budget-recovery.ts`: validate the observed OpenGateway prefill report and reserve the existing 4096 safety tokens plus the existing answer/reasoning minimum. Ambiguous counts, insufficient room and cancellation leave normal error handling in charge.
+
+### Why
+
+- A custom fallback target declared a 1048576-token window while its server enforced 294912. The existing target-window clamp therefore permitted 210744 input + 131072 completion tokens. The catalog default was already corrected in #1255; changing admission alone cannot fix stale custom metadata. The server's explicit counts can repair an output-only reservation overflow without dropping or compacting context. This does not prevent the initial rejection or mutate model metadata.
+
+### Why an extension could not handle it
+
+- `packages/ai/src/api/openai-completions.ts` owns the actual post-hook wire parameters and first-chunk prefetch. An extension cannot safely replay that same request before exposing output without duplicating the adapter.
+- `packages/ai/src/utils/prefill-budget-recovery.ts` needs the actual transmitted completion cap and the adapter's reasoning reservation, not merely the configured model window.
+
+### Expected merge conflict zones
+
+- `packages/ai/src/api/openai-completions.ts`: the request/first-chunk retry boundary and provider-retry import.
+- `packages/ai/src/utils/prefill-budget-recovery.ts`: fork-only helper; no expected upstream conflict.
+
 ## Cursor context ceilings come from the server, not the capability table (2026-09-12)
 
 ### What changed

packages/ai/src/utils/prefill-budget-recovery.ts

@@ -0,0 +1,26 @@
+import { CONTEXT_SAFETY_TOKENS, MIN_ANSWER_TOKENS } from "../api/context-room.ts";
+
+interface PrefillBudget {
+	readonly requested: number | null | undefined;
+	readonly thinkingTokens: number;
+	readonly signal?: AbortSignal;
+}
+
+/** OpenGateway's prefill rejection reports all four counts, including the actual wire completion cap. */
+const PREFILL_COUNTS =
+	/Prefill server error \(400 Bad Request\): .*Requested token count exceeds the model's maximum context length of (\d+) tokens\. You requested a total of (\d+) tokens: (\d+) tokens from the input messages and (\d+) tokens for the completion\./;
+
+/** Preserve the existing safety/answer reserve and reasoning budget; ambiguous reports remain errors. */
+export function repairedOutputBudget(error: unknown, budget: PrefillBudget): number | undefined {
+	if (!(error instanceof Error) || budget.signal?.aborted) return undefined;
+	const match = PREFILL_COUNTS.exec(error.message);
+	if (!match) return undefined;
+	const [window, total, input, completion] = match.slice(1).map(Number);
+	if (![window, total, input, completion].every((value) => Number.isSafeInteger(value) && value > 0)) {
+		return undefined;
+	}
+	if (completion !== budget.requested || input + completion !== total || total <= window) return undefined;
+	const available = window - input - CONTEXT_SAFETY_TOKENS;
+	if (available < MIN_ANSWER_TOKENS + budget.thinkingTokens || available >= completion) return undefined;
+	return available;
+}

packages/ai/test/openai-completions-context-budget.test.ts

@@ -0,0 +1,238 @@
+import { describe, expect, it, vi } from "vitest";
+import { streamSimple } from "../src/compat.ts";
+import type { AssistantMessageEvent, Context, Model } from "../src/types.ts";
+
+const INPUT_TOKENS = 210_744;
+const SERVER_WINDOW = 294_912;
+const OUTPUT_TOKENS = 131_072;
+const SAFE_OUTPUT = 80_072;
+const context: Context = {
+	messages: [{ role: "user", content: "abcd".repeat(INPUT_TOKENS), timestamp: 1 }],
+};
+
+function model(contextWindow = 1_048_576): Model<"openai-completions"> {
+	return {
+		id: "moonshotai/kimi-k3-ultrafast",
+		name: "Kimi fixture",
+		api: "openai-completions",
+		provider: "og",
+		baseUrl: "https://mock.invalid/v1",
+		reasoning: true,
+		input: ["text"],
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+		contextWindow,
+		maxTokens: OUTPUT_TOKENS,
+		compat: { maxTokensField: "max_tokens", supportsReasoningEffort: true },
+	};
+}
+
+function rejection(input: number, completion: number): string {
+	return `Prefill server error (400 Bad Request): ${JSON.stringify({
+		object: "error",
+		message: `Requested token count exceeds the model's maximum context length of ${SERVER_WINDOW} tokens. You requested a total of ${input + completion} tokens: ${input} tokens from the input messages and ${completion} tokens for the completion. Please reduce the number of tokens in the input messages or the completion to fit within the limit.`,
+		type: "BadRequestError",
+		param: null,
+		code: 400,
+	})}`;
+}
+
+function record(value: unknown): Record<string, unknown> {
+	if (typeof value !== "object" || value === null || Array.isArray(value)) throw new Error("Expected object");
+	return Object.fromEntries(Object.entries(value));
+}
+
+function provider(
+	options: { input?: number; alwaysReject?: boolean; partial?: boolean; error?: string; onRequest?: () => void } = {},
+) {
+	const requests: Record<string, unknown>[] = [];
+	const fetch: typeof globalThis.fetch = async (_url, init) => {
+		if (typeof init?.body !== "string") throw new Error("Expected serialized request body");
+		const body = record(JSON.parse(init.body));
+		requests.push(body);
+		const completion = body.max_tokens ?? body.max_completion_tokens;
+		if (typeof completion !== "number") throw new Error("Expected completion cap");
+		const input = options.input ?? INPUT_TOKENS;
+		const data =
+			options.alwaysReject || input + completion > SERVER_WINDOW
+				? { error: { message: options.error ?? rejection(input, completion), type: "BadRequestError", code: 400 } }
+				: {
+						id: "budget-fixture",
+						choices: [{ index: 0, delta: { content: "ok" }, finish_reason: "stop" }],
+						usage: { prompt_tokens: INPUT_TOKENS, completion_tokens: 1, total_tokens: INPUT_TOKENS + 1 },
+					};
+		const partial = options.partial
+			? `data: ${JSON.stringify({ id: "partial", choices: [{ index: 0, delta: { content: "partial" } }] })}\n\n`
+			: "";
+		options.onRequest?.();
+		return new Response(`${partial}data: ${JSON.stringify(data)}\n\ndata: [DONE]\n\n`, {
+			headers: { "content-type": "text/event-stream" },
+		});
+	};
+	return { requests, fetch };
+}
+
+describe("OpenAI-compatible prefill output budget", () => {
+	it("repairs the incident budget once when custom target metadata overstates the window", async () => {
+		// Given: the fallback target declares 1M, but prefill enforces 294912.
+		const remote = provider();
+		const events: AssistantMessageEvent[] = [];
+		// When: the real simple-stream adapter builds and sends the request.
+		const result = streamSimple(model(), context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+			reasoning: "high",
+		});
+		for await (const event of result) events.push(event);
+		const response = await result.result();
+		// Then: only the completion reservation changes, with a single exposed stream.
+		expect(response.stopReason, response.errorMessage).toBe("stop");
+		expect(remote.requests.map((request) => request.max_tokens)).toEqual([OUTPUT_TOKENS, SAFE_OUTPUT]);
+		expect(remote.requests[1]).toEqual({ ...remote.requests[0], max_tokens: SAFE_OUTPUT });
+		expect(events.filter((event) => event.type === "start")).toHaveLength(1);
+		expect(events.filter((event) => event.type === "error")).toHaveLength(0);
+	});
+
+	it("needs no repair when target metadata already gives the server window", async () => {
+		// Given: accurate target metadata.
+		const remote = provider();
+		// When: the existing admission clamp runs.
+		const response = await streamSimple(model(SERVER_WINDOW), context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+		}).result();
+		// Then: one safe request succeeds without a prefill rejection.
+		expect(response.stopReason).toBe("stop");
+		expect(remote.requests.map((request) => request.max_tokens)).toEqual([SAFE_OUTPUT]);
+	});
+
+	it.each([4096, 100_000])(
+		"preserves an explicit %i cap unless the reported window requires less",
+		async (maxTokens) => {
+			// Given: a caller-specified output cap, distinct from the model default.
+			const remote = provider();
+			// When: the provider enforces its actual window.
+			const response = await streamSimple(model(), context, {
+				apiKey: "fixture",
+				fetch: remote.fetch,
+				maxRetries: 0,
+				maxTokens,
+			}).result();
+			// Then: fitting caps are untouched; overflowing caps use the reported room.
+			expect(response.stopReason).toBe("stop");
+			expect(remote.requests.map((request) => request.max_tokens)).toEqual(
+				maxTokens > SAFE_OUTPUT ? [maxTokens, SAFE_OUTPUT] : [maxTokens],
+			);
+		},
+	);
+
+	it("uses the actual post-hook wire cap and does not run the payload hook twice", async () => {
+		// Given: a hook replaces the computed cap and adds an unrelated field.
+		const remote = provider();
+		const onPayload = vi.fn((payload: unknown) => ({ ...record(payload), max_tokens: 100_000, seed: 7 }));
+		// When: that wire request is rejected.
+		const response = await streamSimple(model(), context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+			onPayload,
+		}).result();
+		// Then: repair retains the exact payload other than its completion cap.
+		expect(response.stopReason).toBe("stop");
+		expect(onPayload).toHaveBeenCalledTimes(1);
+		expect(remote.requests[0]?.max_tokens).toBe(100_000);
+		expect(remote.requests[1]).toEqual({ ...remote.requests[0], max_tokens: SAFE_OUTPUT });
+	});
+
+	it("preserves the max_completion_tokens wire variant", async () => {
+		// Given: the other supported output field.
+		const target = model();
+		target.compat = { ...target.compat, maxTokensField: "max_completion_tokens" };
+		const remote = provider();
+		// When: prefill rejects its budget.
+		const response = await streamSimple(target, context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+		}).result();
+		// Then: the same field is corrected without adding max_tokens.
+		expect(response.stopReason).toBe("stop");
+		expect(remote.requests[1]).toEqual({ ...remote.requests[0], max_completion_tokens: SAFE_OUTPUT });
+		expect(remote.requests[1]?.max_tokens).toBeUndefined();
+	});
+
+	it.each([
+		{ name: "input itself exhausts the window", remote: { input: SERVER_WINDOW }, thinkingBudgets: undefined },
+		{ name: "explicit reasoning would lose room", remote: {}, thinkingBudgets: { high: 90_000 } },
+		{
+			name: "reported cap differs from the wire",
+			remote: { error: rejection(INPUT_TOKENS, 130_000) },
+			thinkingBudgets: undefined,
+		},
+		{
+			name: "error lacks complete budget evidence",
+			remote: { error: "context_length_exceeded" },
+			thinkingBudgets: undefined,
+		},
+	])("leaves normal error handling in charge when $name", async ({ remote: config, thinkingBudgets }) => {
+		// Given: reducing the output cannot safely satisfy the proven contract.
+		const remote = provider(config);
+		// When: the first prefill fails.
+		const response = await streamSimple(model(), context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+			reasoning: "high",
+			thinkingBudgets,
+		}).result();
+		// Then: no new send or reasoning downgrade is attempted.
+		expect(response.stopReason).toBe("error");
+		expect(remote.requests).toHaveLength(1);
+	});
+
+	it("propagates a second rejection without another budget repair", async () => {
+		// Given: the server continues rejecting after the evidenced correction.
+		const remote = provider({ alwaysReject: true });
+		// When: one correction is attempted.
+		const response = await streamSimple(model(), context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+		}).result();
+		// Then: bounded failure remains available to session overflow recovery.
+		expect(response.stopReason).toBe("error");
+		expect(remote.requests.map((request) => request.max_tokens)).toEqual([OUTPUT_TOKENS, SAFE_OUTPUT]);
+	});
+
+	it("never replays a response after the first chunk", async () => {
+		// Given: visible content precedes an otherwise matching error.
+		const remote = provider({ partial: true });
+		// When: the stream fails after content.
+		const response = await streamSimple(model(), context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+		}).result();
+		// Then: the partial response is retained without a duplicate request.
+		expect(response.stopReason).toBe("error");
+		expect(response.content).toEqual([{ type: "text", text: "partial" }]);
+		expect(remote.requests).toHaveLength(1);
+	});
+
+	it("does not repair after cancellation at the request boundary", async () => {
+		// Given: cancellation is signalled by the exact first-request event.
+		const controller = new AbortController();
+		const remote = provider({ onRequest: () => controller.abort() });
+		// When: the first request is aborted.
+		const response = await streamSimple(model(), context, {
+			apiKey: "fixture",
+			fetch: remote.fetch,
+			maxRetries: 0,
+			signal: controller.signal,
+		}).result();
+		// Then: no correction is sent.
+		expect(response.stopReason).toBe("aborted");
+		expect(remote.requests).toHaveLength(1);
+	});
+});