From b7bb5a7053511cb0e8bb0432945b0df40aa150cb Mon Sep 17 00:00:00 2001
From: Roo Code <roomote@roocode.com>
Date: Tue, 11 Nov 2025 19:23:12 +0000
Subject: [PATCH] fix: improve Z AI connection handling with timeout and retry
 configuration

- Add timeout configuration using getApiRequestTimeout() utility
- Add maxRetries: 3 for transient connection issues
- Enhance error messages for connection failures (ECONNRESET, ECONNREFUSED, ETIMEDOUT)
- Add specific error handling for SSL/TLS certificate issues
- Update tests to properly mock vscode workspace configuration

Fixes #9174
---
 src/api/providers/__tests__/zai.spec.ts |  8 +++-
 src/api/providers/zai.ts                | 53 +++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/api/providers/__tests__/zai.spec.ts b/src/api/providers/__tests__/zai.spec.ts
index 9db5350080..c1c9c85ac1 100644
--- a/src/api/providers/__tests__/zai.spec.ts
+++ b/src/api/providers/__tests__/zai.spec.ts
@@ -1,7 +1,13 @@
 // npx vitest run src/api/providers/__tests__/zai.spec.ts
 
 // Mock vscode first to avoid import errors
-vitest.mock("vscode", () => ({}))
+vitest.mock("vscode", () => ({
+	workspace: {
+		getConfiguration: () => ({
+			get: (key: string, defaultValue: any) => defaultValue,
+		}),
+	},
+}))
 
 import OpenAI from "openai"
 import { Anthropic } from "@anthropic-ai/sdk"
diff --git a/src/api/providers/zai.ts b/src/api/providers/zai.ts
index cc83945e48..4e05435ec8 100644
--- a/src/api/providers/zai.ts
+++ b/src/api/providers/zai.ts
@@ -18,10 +18,14 @@ import { getModelMaxOutputTokens } from "../../shared/api"
 import { convertToOpenAiMessages } from "../transform/openai-format"
 import type { ApiHandlerCreateMessageMetadata } from "../index"
 import { handleOpenAIError } from "./utils/openai-error-handler"
+import { getApiRequestTimeout } from "./utils/timeout-config"
+import { DEFAULT_HEADERS } from "./constants"
 
 import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider"
 
 export class ZAiHandler extends BaseOpenAiCompatibleProvider<string> {
+	protected override client: OpenAI
+
 	constructor(options: ApiHandlerOptions) {
 		const isChina = zaiApiLineConfigs[options.zaiApiLine ?? "international_coding"].isChina
 		const models = (isChina ? mainlandZAiModels : internationalZAiModels) as unknown as Record<string, ModelInfo>
@@ -36,6 +40,19 @@ export class ZAiHandler extends BaseOpenAiCompatibleProvider<string> {
 			providerModels: models,
 			defaultTemperature: ZAI_DEFAULT_TEMPERATURE,
 		})
+
+		// Override the client with proper timeout and retry configuration
+		const timeout = getApiRequestTimeout()
+		const baseURL = zaiApiLineConfigs[options.zaiApiLine ?? "international_coding"].baseUrl
+		const apiKey = options.zaiApiKey ?? "not-provided"
+
+		this.client = new OpenAI({
+			baseURL,
+			apiKey,
+			defaultHeaders: DEFAULT_HEADERS,
+			timeout,
+			maxRetries: 3, // Add retry logic for transient connection issues
+		})
 	}
 
 	protected override createStream(
@@ -75,6 +92,24 @@ export class ZAiHandler extends BaseOpenAiCompatibleProvider<string> {
 		try {
 			return this.client.chat.completions.create(params, requestOptions)
 		} catch (error) {
+			// Enhanced error handling for Z AI connection issues
+			if (error instanceof Error) {
+				const errorMessage = error.message.toLowerCase()
+				if (
+					errorMessage.includes("econnreset") ||
+					errorMessage.includes("econnrefused") ||
+					errorMessage.includes("etimedout")
+				) {
+					throw new Error(
+						`Z AI connection error: Unable to connect to Z AI API. Please check your network connection and API endpoint configuration. Original error: ${error.message}`,
+					)
+				}
+				if (errorMessage.includes("certificate") || errorMessage.includes("ssl")) {
+					throw new Error(
+						`Z AI SSL/TLS error: Certificate validation failed. This may be due to network proxy settings or firewall restrictions. Original error: ${error.message}`,
+					)
+				}
+			}
 			throw handleOpenAIError(error, this.providerName)
 		}
 	}
@@ -97,6 +132,24 @@ export class ZAiHandler extends BaseOpenAiCompatibleProvider<string> {
 			const response = await this.client.chat.completions.create(params)
 			return response.choices[0]?.message.content || ""
 		} catch (error) {
+			// Enhanced error handling for Z AI connection issues
+			if (error instanceof Error) {
+				const errorMessage = error.message.toLowerCase()
+				if (
+					errorMessage.includes("econnreset") ||
+					errorMessage.includes("econnrefused") ||
+					errorMessage.includes("etimedout")
+				) {
+					throw new Error(
+						`Z AI connection error: Unable to connect to Z AI API. Please check your network connection and API endpoint configuration. Original error: ${error.message}`,
+					)
+				}
+				if (errorMessage.includes("certificate") || errorMessage.includes("ssl")) {
+					throw new Error(
+						`Z AI SSL/TLS error: Certificate validation failed. This may be due to network proxy settings or firewall restrictions. Original error: ${error.message}`,
+					)
+				}
+			}
 			throw handleOpenAIError(error, this.providerName)
 		}
 	}