feat(ai): support thinking_token_budget on openai-completions (#7638)

vLLM caps the reasoning phase with a top-level thinking_token_budget. On
OpenAI-compatible endpoints reasoning and the answer share max_tokens, so a
reasoning-heavy turn can consume the entire response ceiling and return no
answer and no tool call — the agent loop then sees a message without tool
calls and ends the run as if it had finished.

thinkingBudgets from #529 already carries per-level budgets, but only the
Anthropic and Google paths read it. This wires it into openai-completions
behind a new opt-in compat flag, and clamps the budget so at least
MIN_ANSWER_TOKENS remain for the answer.

The flag is independent of thinkingFormat because one vLLM server can serve
zai, qwen and chat-template models.

Co-authored-by: Ben Shaharizad <ben.shaharizad@taboola.com>

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

@@ -28,6 +28,7 @@ import type {
 	StreamFunction,
 	StreamOptions,
 	TextContent,
+	ThinkingBudgets,
 	ThinkingContent,
 	Tool,
 	ToolCall,
@@ -51,7 +52,7 @@ import {
 } from "./constrained-sampling.ts";
 import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.ts";
 import { clampOpenAIPromptCacheKey } from "./openai-prompt-cache.ts";
-import { buildBaseOptions } from "./simple-options.ts";
+import { buildBaseOptions, clampReasoning, MIN_ANSWER_TOKENS } from "./simple-options.ts";
 import { transformMessages } from "./transform-messages.ts";
 
 /**
@@ -141,6 +142,8 @@ function isEncryptedReasoningDetail(detail: unknown): detail is OpenAIEncryptedR
 export interface OpenAICompletionsOptions extends StreamOptions {
 	toolChoice?: OpenAI.Chat.Completions.ChatCompletionToolChoiceOption;
 	reasoningEffort?: "minimal" | "low" | "medium" | "high" | "xhigh" | "max";
+	/** Token budgets per thinking level. Only used when `compat.supportsThinkingTokenBudget` is set. */
+	thinkingBudgets?: ThinkingBudgets;
 }
 
 export interface ConvertCompletionsMessagesOptions {
@@ -154,10 +157,11 @@ interface OpenAICompatCacheControl {
 
 type ResolvedOpenAICompletionsCompat = Omit<
 	Required<OpenAICompletionsCompat>,
-	"cacheControlFormat" | "deferredToolsMode"
+	"cacheControlFormat" | "deferredToolsMode" | "supportsThinkingTokenBudget"
 > & {
 	cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
 	deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
+	supportsThinkingTokenBudget?: OpenAICompletionsCompat["supportsThinkingTokenBudget"];
 };
 
 type ResolvedChatTemplateKwargValue = string | number | boolean | null;
@@ -625,6 +629,7 @@ export const streamSimple: StreamFunction<"openai-completions", SimpleStreamOpti
 		...base,
 		reasoningEffort,
 		toolChoice,
+		thinkingBudgets: options?.thinkingBudgets,
 	} satisfies OpenAICompletionsOptions);
 };
 
@@ -839,6 +844,27 @@ function buildParams(
 		}
 	}
 
+	// vLLM caps reasoning with a top-level thinking_token_budget. Independent of
+	// thinkingFormat: the same server can serve zai, qwen or chat-template models.
+	// Reasoning and the answer share max_tokens here, so an uncapped reasoning
+	// phase can consume the whole response and leave no answer and no tool call.
+	if (compat.supportsThinkingTokenBudget && options?.reasoningEffort && model.reasoning) {
+		const level = clampReasoning(options.reasoningEffort)!;
+		const budgets: ThinkingBudgets = {
+			minimal: 1024,
+			low: 2048,
+			medium: 8192,
+			high: 16384,
+			...options.thinkingBudgets,
+		};
+		const ceiling = (params as { max_tokens?: number }).max_tokens ?? params.max_completion_tokens ?? model.maxTokens;
+		// Always leave room for the answer, otherwise the budget recreates the bug it prevents.
+		const budget = Math.min(budgets[level]!, Math.max(0, ceiling - MIN_ANSWER_TOKENS));
+		if (budget > 0) {
+			(params as { thinking_token_budget?: number }).thinking_token_budget = budget;
+		}
+	}
+
 	// OpenRouter provider routing preferences
 	if (model.compat?.openRouterRouting) {
 		(params as any).provider = model.compat.openRouterRouting;
@@ -1492,6 +1518,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
 		chatTemplateKwargs: {},
 		chatTemplateArgs: {},
 		zaiToolStream: false,
+		supportsThinkingTokenBudget: false,
 		supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia,
 		supportsOpenAIGrammarTools: false,
 		cacheControlFormat,
@@ -1536,6 +1563,7 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion
 		chatTemplateKwargs: model.compat.chatTemplateKwargs ?? detected.chatTemplateKwargs,
 		chatTemplateArgs: model.compat.chatTemplateArgs ?? detected.chatTemplateArgs,
 		zaiToolStream: model.compat.zaiToolStream ?? detected.zaiToolStream,
+		supportsThinkingTokenBudget: model.compat.supportsThinkingTokenBudget ?? detected.supportsThinkingTokenBudget,
 		supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode,
 		supportsOpenAIGrammarTools: model.compat.supportsOpenAIGrammarTools ?? detected.supportsOpenAIGrammarTools,
 		cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat,

packages/ai/src/api/simple-options.ts

@@ -51,6 +51,9 @@ export function buildBaseOptions(
 	};
 }
 
+/** Tokens always left for the answer when a thinking budget shares the response ceiling. */
+export const MIN_ANSWER_TOKENS = 1024;
+
 export function clampReasoning(effort: ThinkingLevel | undefined): Exclude<ThinkingLevel, "xhigh" | "max"> | undefined {
 	return effort === "xhigh" || effort === "max" ? "high" : effort;
 }
@@ -70,14 +73,13 @@ export function adjustMaxTokensForThinking(
 	};
 	const budgets = { ...defaultBudgets, ...customBudgets };
 
-	const minOutputTokens = 1024;
 	const level = clampReasoning(reasoningLevel)!;
 	let thinkingBudget = budgets[level]!;
 	const maxTokens =
 		baseMaxTokens === undefined ? modelMaxTokens : Math.min(baseMaxTokens + thinkingBudget, modelMaxTokens);
 
 	if (maxTokens <= thinkingBudget) {
-		thinkingBudget = Math.max(0, maxTokens - minOutputTokens);
+		thinkingBudget = Math.max(0, maxTokens - MIN_ANSWER_TOKENS);
 	}
 
 	return { maxTokens, thinkingBudget };

packages/ai/src/types.ts

@@ -578,6 +578,8 @@ export interface OpenAICompletionsCompat {
 	vercelGatewayRouting?: VercelGatewayRouting;
 	/** Whether z.ai supports top-level `tool_stream: true` for streaming tool call deltas. Default: false. */
 	zaiToolStream?: boolean;
+	/** Whether the provider supports top-level `thinking_token_budget` to cap reasoning tokens (vLLM). Reasoning and the answer share `max_tokens` on these endpoints, so without a budget a reasoning-heavy turn can consume the whole response and emit no answer. Default: false. */
+	supportsThinkingTokenBudget?: boolean;
 	/** Whether the provider supports OpenAI custom tools with Lark/regex grammar formats. When false, grammar-constrained tools fall back to normal function tools. Default: false; the generated model catalog enables it for capable models. */
 	supportsOpenAIGrammarTools?: boolean;
 	/** Whether the provider supports the `strict` field in tool definitions. Default: true. */

packages/ai/test/openai-completions-thinking-as-text.test.ts

@@ -38,6 +38,7 @@ const compat = {
 	chatTemplateKwargs: {},
 	chatTemplateArgs: {},
 	zaiToolStream: false,
+	supportsThinkingTokenBudget: false,
 	supportsStrictMode: true,
 	supportsOpenAIGrammarTools: false,
 	cacheControlFormat: undefined,

packages/ai/test/openai-completions-thinking-token-budget.test.ts

@@ -0,0 +1,124 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { streamSimple } from "../src/compat.ts";
+import type { Model, SimpleStreamOptions, ThinkingBudgets } from "../src/types.ts";
+
+const mockState = vi.hoisted(() => ({
+	lastParams: undefined as unknown,
+}));
+
+vi.mock("openai", () => {
+	class FakeOpenAI {
+		chat = {
+			completions: {
+				create: (params: unknown) => {
+					mockState.lastParams = params;
+					const stream = {
+						async *[Symbol.asyncIterator]() {
+							yield {
+								choices: [{ delta: {}, finish_reason: "stop" }],
+								usage: {
+									prompt_tokens: 1,
+									completion_tokens: 1,
+									prompt_tokens_details: { cached_tokens: 0 },
+									completion_tokens_details: { reasoning_tokens: 0 },
+								},
+							};
+						},
+					};
+					const promise = Promise.resolve(stream) as Promise<typeof stream> & {
+						withResponse: () => Promise<{ data: typeof stream; response: { status: number; headers: Headers } }>;
+					};
+					promise.withResponse = async () => ({
+						data: stream,
+						response: { status: 200, headers: new Headers() },
+					});
+					return promise;
+				},
+			},
+		};
+	}
+
+	return { default: FakeOpenAI };
+});
+
+// vLLM-served reasoning model: reasoning and the answer share max_tokens.
+const vllmModel: Model<"openai-completions"> = {
+	id: "zai-org/glm-5.2",
+	name: "GLM 5.2 (local vLLM)",
+	api: "openai-completions",
+	provider: "local-vllm",
+	baseUrl: "http://localhost:8000/v1",
+	reasoning: true,
+	input: ["text"],
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+	contextWindow: 262144,
+	maxTokens: 16384,
+	compat: { thinkingFormat: "zai", supportsThinkingTokenBudget: true },
+};
+
+async function capture(
+	model: Model<"openai-completions">,
+	options?: {
+		reasoning?: SimpleStreamOptions["reasoning"];
+		thinkingBudgets?: ThinkingBudgets;
+		maxTokens?: number;
+	},
+): Promise<{ thinking_token_budget?: number; thinking?: unknown }> {
+	let payload: unknown;
+
+	await streamSimple(
+		model,
+		{ messages: [{ role: "user", content: "Hi", timestamp: Date.now() }] },
+		{
+			apiKey: "test",
+			reasoning: options?.reasoning,
+			thinkingBudgets: options?.thinkingBudgets,
+			maxTokens: options?.maxTokens,
+			onPayload: (params: unknown) => {
+				payload = params;
+			},
+		},
+	).result();
+
+	return (payload ?? mockState.lastParams) as { thinking_token_budget?: number; thinking?: unknown };
+}
+
+describe("openai-completions thinking_token_budget", () => {
+	beforeEach(() => {
+		mockState.lastParams = undefined;
+	});
+
+	it("sends the configured budget for the requested level", async () => {
+		const params = await capture(vllmModel, { reasoning: "medium", thinkingBudgets: { medium: 4096 } });
+		expect(params.thinking_token_budget).toBe(4096);
+	});
+
+	it("omits the budget when the compat flag is not set", async () => {
+		const model = { ...vllmModel, compat: { thinkingFormat: "zai" } } as Model<"openai-completions">;
+		const params = await capture(model, { reasoning: "medium", thinkingBudgets: { medium: 4096 } });
+		expect(params.thinking_token_budget).toBeUndefined();
+	});
+
+	it("omits the budget when thinking is off", async () => {
+		const params = await capture(vllmModel, { reasoning: undefined, thinkingBudgets: { high: 8192 } });
+		expect(params.thinking_token_budget).toBeUndefined();
+	});
+
+	it("clamps xhigh and max to the high budget", async () => {
+		const xhigh = await capture(vllmModel, { reasoning: "xhigh", thinkingBudgets: { high: 8192 } });
+		const max = await capture(vllmModel, { reasoning: "max", thinkingBudgets: { high: 8192 } });
+		expect(xhigh.thinking_token_budget).toBe(8192);
+		expect(max.thinking_token_budget).toBe(8192);
+	});
+
+	it("leaves room for the answer when the budget meets the response ceiling", async () => {
+		// Default high budget (16384) equals the model ceiling, which would leave no answer.
+		const params = await capture(vllmModel, { reasoning: "high" });
+		expect(params.thinking_token_budget).toBe(16384 - 1024);
+	});
+
+	it("uses the caller max_tokens as the ceiling when it is lower than the model cap", async () => {
+		const params = await capture(vllmModel, { reasoning: "high", thinkingBudgets: { high: 8192 }, maxTokens: 4096 });
+		expect(params.thinking_token_budget).toBe(4096 - 1024);
+	});
+});

packages/ai/test/openai-completions-tool-result-images.test.ts

@@ -38,6 +38,7 @@ const compat: Omit<Required<OpenAICompletionsCompat>, "deferredToolsMode"> & {
 	chatTemplateKwargs: {},
 	chatTemplateArgs: {},
 	zaiToolStream: false,
+	supportsThinkingTokenBudget: false,
 	supportsStrictMode: true,
 	supportsOpenAIGrammarTools: false,
 	cacheControlFormat: "anthropic",