feat: enhance maxTokens handling in ZgsmAiHandler and SettingsView co…

src/api/providers/zgsm.ts

@@ -804,6 +804,9 @@ export class ZgsmAiHandler extends BaseProvider implements SingleCompletionHandl
 			| OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,
 		modelInfo: ModelInfo,
 	): void {
+		// Check if it's auto mode
+		const isAutoMode = modelInfo.id === "Auto" || modelInfo.id === "auto"
+
 		// Only add max_completion_tokens if includeMaxTokens is true
 		if (this.options.useZgsmCustomConfig) {
 			const maxTokens = this.options.modelMaxTokens || modelInfo.maxTokens
@@ -813,11 +816,26 @@ export class ZgsmAiHandler extends BaseProvider implements SingleCompletionHandl
 				Object.assign(requestOptions, {
 					[modelInfo.maxTokensKey || "max_completion_tokens"]: maxTokens,
 				})
+			} else {
+				Object.assign(requestOptions, {
+					max_tokens: modelInfo.maxTokens,
+					max_completion_tokens: modelInfo.maxTokens,
+				})
+			}
+		} else {
+			// If maxTokensKey exists, use it directly
+			if (modelInfo.maxTokensKey) {
+				Object.assign(requestOptions, {
+					[modelInfo.maxTokensKey]: modelInfo.maxTokens,
+				})
+			} else if (!isAutoMode) {
+				// Logic for non-auto mode
+				// If maxTokensKey doesn't exist, use both max_tokens and max_completion_tokens
+				Object.assign(requestOptions, {
+					max_tokens: modelInfo.maxTokens,
+					max_completion_tokens: modelInfo.maxTokens,
+				})
 			}
-		} else if (modelInfo.maxTokensKey) {
-			Object.assign(requestOptions, {
-				[modelInfo.maxTokensKey]: modelInfo.maxTokens,
-			})
 		}
 	}
 

webview-ui/src/components/settings/SettingsView.tsx

@@ -436,10 +436,19 @@ const SettingsView = forwardRef<SettingsViewRef, SettingsViewProps>(({ onDone, t
 			// by the `updateSettings` message.
 			vscode.postMessage({ type: "updateCondensingPrompt", text: customCondensingPrompt || "" })
 			// vscode.postMessage({ type: "upsertApiConfiguration", text: currentApiConfigName, apiConfiguration })
+			const finalApiConfig = { ...apiConfiguration, useZgsmCustomConfig, zgsmCodebaseIndexEnabled }
+
+			if (
+				finalApiConfig.useZgsmCustomConfig &&
+				finalApiConfig.includeMaxTokens === undefined &&
+				!Object.hasOwn(apiConfiguration, "includeMaxTokens")
+			) {
+				finalApiConfig.includeMaxTokens = true
+			}
 			vscode.postMessage({
 				type: "upsertApiConfiguration",
 				text: currentApiConfigName,
-				apiConfiguration: { ...apiConfiguration, useZgsmCustomConfig, zgsmCodebaseIndexEnabled },
+				apiConfiguration: finalApiConfig,
 			})
 			vscode.postMessage({ type: "telemetrySetting", text: telemetrySetting })
 

webview-ui/src/components/settings/providers/ZgsmAI.tsx

@@ -1,4 +1,4 @@
-import { useState, useCallback, useEffect } from "react"
+import { useState, useCallback, useEffect, useRef } from "react"
 import { useEvent } from "react-use"
 import { Checkbox } from "vscrui"
 import { VSCodeButton, VSCodeCheckbox, VSCodeLink, VSCodeTextField } from "@vscode/webview-ui-toolkit/react"
@@ -55,6 +55,9 @@ export const ZgsmAI = ({
 	const [azureApiVersionSelected, setAzureApiVersionSelected] = useState(!!apiConfiguration?.azureApiVersion)
 	const [openAiLegacyFormatSelected, setOpenAiLegacyFormatSelected] = useState(!!apiConfiguration?.openAiLegacyFormat)
 
+	// Use `ref` to track whether `includeMaxTokens` has been explicitly set by the user.
+	const includeMaxTokensInitializedRef = useRef(Object.hasOwn(apiConfiguration, "includeMaxTokens"))
+
 	const [openAiModels, setOpenAiModels] = useState<Record<string, ModelInfo> | null>(null)
 
 	const [customHeaders, setCustomHeaders] = useState<[string, string][]>(() => {
@@ -108,6 +111,26 @@ export const ZgsmAI = ({
 		return () => clearTimeout(timer)
 	}, [customHeaders, setApiConfigurationField])
 
+	useEffect(() => {
+		// Set the default value only when useZgsmCustomConfig is first enabled and includeMaxTokens has never been set before.
+		// Use ref to track whether includeMaxTokens has been explicitly set by the user, avoiding overriding the user's explicit selection.
+		if (
+			useZgsmCustomConfig &&
+			!includeMaxTokensInitializedRef.current &&
+			apiConfiguration?.includeMaxTokens === undefined
+		) {
+			setApiConfigurationField("includeMaxTokens", true)
+			includeMaxTokensInitializedRef.current = true
+		}
+	}, [useZgsmCustomConfig, apiConfiguration?.includeMaxTokens, setApiConfigurationField])
+
+	// Marked as initialized when the user manually modifies includeMaxTokens.
+	useEffect(() => {
+		if (Object.hasOwn(apiConfiguration, "includeMaxTokens")) {
+			includeMaxTokensInitializedRef.current = true
+		}
+	}, [apiConfiguration])
+
 	const handleInputChange = useCallback(
 		<K extends keyof ProviderSettings, E>(
 			field: K,
@@ -229,7 +252,9 @@ export const ZgsmAI = ({
 					<div>
 						<Checkbox
 							checked={apiConfiguration?.includeMaxTokens ?? true}
-							onChange={handleInputChange("includeMaxTokens", noTransform)}>
+							onChange={(checked: boolean) => {
+								setApiConfigurationField("includeMaxTokens", checked)
+							}}>
 							{t("settings:includeMaxOutputTokens")}
 						</Checkbox>
 						<div className="text-sm text-vscode-descriptionForeground ml-6">