feat(zai): expose configurable max output tokens for GLM models (#161) (#274)

* feat(zai): expose configurable max output tokens for GLM models (#161)

Add a standalone "max output tokens" slider for models that advertise
`supportsMaxTokens` (e.g. Z.ai GLM) but do not surface a reasoning budget, and
send the chosen value to the provider as `max_tokens`.

- ThinkingBudget: extract the max-tokens slider into a small render helper shared
  by the standalone control and the reasoning-budget branch, and surface it for
  binary-reasoning models that also support a configurable max.
- getModelMaxOutputTokens: honor the user's `modelMaxTokens` override for
  `supportsMaxTokens` models (capped at the model ceiling) instead of the 20%
  context-window clamp, so the runtime budget matches what is sent to the provider.
- ProviderSettingsManager.export(): preserve `modelMaxTokens` for configurable-max
  models while dropping it for models that support neither reasoning budgets nor a
  configurable max.

* Update src/shared/api.ts

* Update webview-ui/src/components/settings/ThinkingBudget.tsx

---------

Co-authored-by: Armando Vaquera <263793884+proyectoauraorg@users.noreply.github.com>
Co-authored-by: edelauna <54631123+edelauna@users.noreply.github.com>

Author: proyectoauraorg

packages/types/src/model.ts

@@ -81,6 +81,10 @@ export const modelInfoSchema = z.object({
 	promptCacheRetention: z.enum(["in_memory", "24h"]).optional(),
 	// Capability flag to indicate whether the model supports an output verbosity parameter
 	supportsVerbosity: z.boolean().optional(),
+	// Capability flag to indicate whether the model exposes a user-configurable max output
+	// tokens control in settings. When set, the settings UI surfaces a slider that persists
+	// `modelMaxTokens`; when the user leaves it unset, the default output clamp is used.
+	supportsMaxTokens: z.boolean().optional(),
 	supportsReasoningBudget: z.boolean().optional(),
 	// Capability flag to indicate whether the model supports simple on/off binary reasoning
 	supportsReasoningBinary: z.boolean().optional(),

packages/types/src/providers/zai.ts

@@ -142,6 +142,7 @@ export const internationalZAiModels = {
 		contextWindow: 200_000,
 		supportsImages: false,
 		supportsPromptCache: true,
+		supportsMaxTokens: true,
 		supportsReasoningEffort: ["disable", "medium"],
 		reasoningEffort: "medium",
 		preserveReasoning: true,
@@ -157,6 +158,7 @@ export const internationalZAiModels = {
 		contextWindow: 202_752,
 		supportsImages: false,
 		supportsPromptCache: true,
+		supportsMaxTokens: true,
 		supportsReasoningEffort: ["disable", "medium"],
 		reasoningEffort: "medium",
 		preserveReasoning: true,
@@ -348,6 +350,7 @@ export const mainlandZAiModels = {
 		contextWindow: 204_800,
 		supportsImages: false,
 		supportsPromptCache: true,
+		supportsMaxTokens: true,
 		supportsReasoningEffort: ["disable", "medium"],
 		reasoningEffort: "medium",
 		preserveReasoning: true,
@@ -363,6 +366,7 @@ export const mainlandZAiModels = {
 		contextWindow: 202_752,
 		supportsImages: false,
 		supportsPromptCache: true,
+		supportsMaxTokens: true,
 		supportsReasoningEffort: ["disable", "medium"],
 		reasoningEffort: "medium",
 		preserveReasoning: true,

src/api/providers/__tests__/zai.spec.ts

@@ -503,6 +503,46 @@ describe("ZAiHandler", () => {
 			)
 		})
 
+		it("should advertise supportsMaxTokens for configurable GLM models", () => {
+			expect(internationalZAiModels["glm-5.1"].supportsMaxTokens).toBe(true)
+			expect(internationalZAiModels["glm-5-turbo"].supportsMaxTokens).toBe(true)
+			expect(mainlandZAiModels["glm-5.1"].supportsMaxTokens).toBe(true)
+			expect(mainlandZAiModels["glm-5-turbo"].supportsMaxTokens).toBe(true)
+			// Models without a configurable output budget should not advertise the flag.
+			expect((internationalZAiModels["glm-4.7"] as { supportsMaxTokens?: boolean }).supportsMaxTokens).toBe(
+				undefined,
+			)
+		})
+
+		it("should honor an explicit modelMaxTokens override instead of the 20% clamp", async () => {
+			const handlerWithModel = new ZAiHandler({
+				apiModelId: "glm-5.1",
+				zaiApiKey: "test-zai-api-key",
+				zaiApiLine: "international_coding",
+				modelMaxTokens: 100_000,
+			})
+
+			mockCreate.mockImplementationOnce(() => {
+				return {
+					[Symbol.asyncIterator]: () => ({
+						async next() {
+							return { done: true }
+						},
+					}),
+				}
+			})
+
+			const messageGenerator = handlerWithModel.createMessage("system prompt", [])
+			await messageGenerator.next()
+
+			expect(mockCreate).toHaveBeenCalledWith(
+				expect.objectContaining({
+					model: "glm-5.1",
+					max_tokens: 100_000,
+				}),
+			)
+		})
+
 		it("should enable thinking by default for GLM-4.7 (default reasoningEffort is medium)", async () => {
 			const handlerWithModel = new ZAiHandler({
 				apiModelId: "glm-4.7",

src/core/config/ProviderSettingsManager.ts

@@ -567,11 +567,18 @@ export class ProviderSettingsManager {
 						const supportsReasoningBudget =
 							modelInfo.supportsReasoningBudget || modelInfo.requiredReasoningBudget
 
-						// If the model doesn't support reasoning budgets, remove the token fields
+						// modelMaxThinkingTokens only applies to reasoning budgets, but modelMaxTokens
+						// also caps output on models that expose a configurable max (e.g. GLM), so keep
+						// it whenever the model supports either feature.
+						const supportsMaxTokens = supportsReasoningBudget || modelInfo.supportsMaxTokens
+
 						if (!supportsReasoningBudget) {
-							delete configs[name].modelMaxTokens
 							delete configs[name].modelMaxThinkingTokens
 						}
+
+						if (!supportsMaxTokens) {
+							delete configs[name].modelMaxTokens
+						}
 					} catch (error) {
 						// If we can't build the API handler or get model info, skip filtering
 						// to avoid accidental data loss from incomplete configurations

src/core/config/__tests__/ProviderSettingsManager.spec.ts

@@ -6,6 +6,33 @@ import type { ProviderSettings } from "@roo-code/types"
 
 import { ProviderSettingsManager, ProviderProfiles, SyncCloudProfilesResult } from "../ProviderSettingsManager"
 
+// `export()` builds an API handler per profile to read model capabilities. Mock
+// buildApiHandler with the real @roo-code/types model definitions so the token-field
+// filtering is driven by real capability flags, and so this suite stays isolated from
+// sibling specs that also mock "../../../api" (avoids a cross-file mock leak under
+// Vitest's singleFork pool).
+vi.mock("../../../api", async () => {
+	const types = await vi.importActual<typeof import("@roo-code/types")>("@roo-code/types")
+	const zaiModels = { ...types.internationalZAiModels, ...types.mainlandZAiModels } as Record<string, unknown>
+	const anthropicModels = types.anthropicModels as Record<string, unknown>
+	const modelInfoFor = (config: { apiProvider?: string; apiModelId?: string }) => {
+		const id = config?.apiModelId ?? ""
+		switch (config?.apiProvider) {
+			case "zai":
+				return zaiModels[id] ?? {}
+			case "anthropic":
+				return anthropicModels[id] ?? {}
+			default:
+				return {}
+		}
+	}
+	return {
+		buildApiHandler: (config: any) => ({
+			getModel: () => ({ id: config?.apiModelId ?? "", info: modelInfoFor(config) }),
+		}),
+	}
+})
+
 // Mock VSCode ExtensionContext
 const mockSecrets = {
 	get: vi.fn(),
@@ -874,6 +901,52 @@ describe("ProviderSettingsManager", () => {
 			expect(exported.apiConfigs.retired.modelMaxTokens).toBe(4096)
 			expect(exported.apiConfigs.retired.modelMaxThinkingTokens).toBe(2048)
 		})
+
+		it("should preserve modelMaxTokens for models that support a configurable max output (e.g. GLM)", async () => {
+			const existingConfig: ProviderProfiles = {
+				currentApiConfigName: "glm",
+				apiConfigs: {
+					glm: {
+						id: "glm-id",
+						apiProvider: "zai",
+						apiModelId: "glm-5.1",
+						modelMaxTokens: 8192,
+						modelMaxThinkingTokens: 2048,
+					},
+				},
+			}
+
+			mockSecrets.get.mockResolvedValue(JSON.stringify(existingConfig))
+
+			const exported = await providerSettingsManager.export()
+
+			// GLM exposes a configurable max output (supportsMaxTokens) but no reasoning budget,
+			// so modelMaxTokens must survive the export while modelMaxThinkingTokens is dropped.
+			expect(exported.apiConfigs.glm.modelMaxTokens).toBe(8192)
+			expect(exported.apiConfigs.glm.modelMaxThinkingTokens).toBeUndefined()
+		})
+
+		it("should strip both token fields for models that support neither reasoning budgets nor a configurable max", async () => {
+			const existingConfig: ProviderProfiles = {
+				currentApiConfigName: "anthropic",
+				apiConfigs: {
+					anthropic: {
+						id: "anthropic-id",
+						apiProvider: "anthropic",
+						apiModelId: "claude-3-5-haiku-20241022",
+						modelMaxTokens: 8192,
+						modelMaxThinkingTokens: 2048,
+					},
+				},
+			}
+
+			mockSecrets.get.mockResolvedValue(JSON.stringify(existingConfig))
+
+			const exported = await providerSettingsManager.export()
+
+			expect(exported.apiConfigs.anthropic.modelMaxTokens).toBeUndefined()
+			expect(exported.apiConfigs.anthropic.modelMaxThinkingTokens).toBeUndefined()
+		})
 	})
 
 	describe("ResetAllConfigs", () => {

src/shared/__tests__/api.spec.ts

@@ -240,6 +240,41 @@ describe("getModelMaxOutputTokens", () => {
 		})
 	})
 
+	test("should honor the user's modelMaxTokens override for supportsMaxTokens models (e.g. Z.ai GLM)", () => {
+		const model: ModelInfo = {
+			contextWindow: 200_000,
+			supportsPromptCache: false,
+			supportsMaxTokens: true,
+			maxTokens: 98_304, // model ceiling, well above the 20% clamp (40k)
+		}
+
+		const settings: ProviderSettings = {
+			apiProvider: "zai",
+			modelMaxTokens: 64_000, // user override, above 20% of the context window (40k)
+		}
+
+		const result = getModelMaxOutputTokens({ modelId: "glm-4.6", model, settings, format: "openai" })
+		// Honored instead of clamped to 20% of the context window.
+		expect(result).toBe(64_000)
+	})
+
+	test("should cap the supportsMaxTokens override at the model's own maxTokens ceiling", () => {
+		const model: ModelInfo = {
+			contextWindow: 200_000,
+			supportsPromptCache: false,
+			supportsMaxTokens: true,
+			maxTokens: 32_000,
+		}
+
+		const settings: ProviderSettings = {
+			apiProvider: "zai",
+			modelMaxTokens: 999_999, // beyond the model ceiling
+		}
+
+		const result = getModelMaxOutputTokens({ modelId: "glm-4.6", model, settings, format: "openai" })
+		expect(result).toBe(32_000)
+	})
+
 	test("should still apply 20% cap to non-GPT-5 models", () => {
 		const model: ModelInfo = {
 			contextWindow: 200_000,

src/shared/api.ts

@@ -132,6 +132,13 @@ export const getModelMaxOutputTokens = ({
 		return ANTHROPIC_DEFAULT_MAX_TOKENS
 	}
 
+	// Models that expose a configurable max-output slider (e.g. Z.ai GLM) honor the user's
+	// explicit override instead of the default 20% context-window clamp, capped at the model's
+	// own ceiling. This keeps the runtime budget consistent with the value sent to the provider.
+	if (model.supportsMaxTokens && settings?.modelMaxTokens != null && settings.modelMaxTokens > 0) {
+		return model.maxTokens ? Math.min(settings.modelMaxTokens, model.maxTokens) : settings.modelMaxTokens
+	}
+
 	// If model has explicit maxTokens, clamp it to 20% of the context window
 	// Exception: GPT-5 models should use their exact configured max output tokens
 	if (model.maxTokens) {