Merge pull request #407 from kaitranntt/feat/glmt-rate-limit-resilience

feat(glmt): add rate limit resilience with exponential backoff retry

src/glmt/glmt-proxy.ts

@@ -29,6 +29,15 @@ interface GlmtProxyConfig {
   timeout?: number;
 }
 
+/**
+ * Configuration for retry behavior on rate limit errors
+ */
+interface RetryConfig {
+  maxRetries: number;
+  baseDelay: number;
+  enabled: boolean;
+}
+
 interface ThinkingConfig {
   thinking: boolean;
   effort: string;
@@ -79,6 +88,8 @@ export class GlmtProxy {
   private port: number | null;
   private verbose: boolean;
   private timeout: number;
+  private retryConfig: RetryConfig;
+  private httpsAgent: https.Agent;
 
   constructor(config: GlmtProxyConfig = {}) {
     this.transformer = new GlmtTransformer({
@@ -93,6 +104,23 @@ export class GlmtProxy {
     this.port = null;
     this.verbose = config.verbose || false;
     this.timeout = config.timeout || 120000; // 120s default
+
+    // Retry configuration for handling 429 rate limit errors
+    const maxRetries = parseInt(process.env.GLMT_MAX_RETRIES || '3', 10);
+    const baseDelay = parseInt(process.env.GLMT_RETRY_BASE_DELAY || '1000', 10);
+    this.retryConfig = {
+      maxRetries: isNaN(maxRetries) || maxRetries < 0 ? 3 : Math.min(maxRetries, 10),
+      baseDelay: isNaN(baseDelay) || baseDelay < 0 ? 1000 : baseDelay,
+      enabled: process.env.GLMT_DISABLE_RETRY !== '1',
+    };
+
+    // Connection pooling for better performance
+    this.httpsAgent = new https.Agent({
+      keepAlive: true,
+      maxSockets: 5,
+      maxFreeSockets: 2,
+      timeout: this.timeout,
+    });
   }
 
   /**
@@ -249,8 +277,8 @@ export class GlmtProxy {
 
     this.log(`Transformed request, thinking: ${thinkingConfig.thinking}`);
 
-    // Forward to Z.AI
-    const openaiResponse = (await this.forwardToUpstream(
+    // Forward to Z.AI with retry on rate limit
+    const openaiResponse = (await this.forwardWithRetry(
       openaiRequest as unknown as OpenAIRequest,
       {}
     )) as OpenAIResponse;
@@ -309,8 +337,8 @@ export class GlmtProxy {
 
     this.log('Starting SSE stream to Claude CLI (socket buffering disabled)');
 
-    // Forward and stream
-    await this.forwardAndStreamUpstream(
+    // Forward and stream with retry on rate limit
+    await this.forwardAndStreamWithRetry(
       openaiRequest as unknown as OpenAIRequest,
       {},
       res,
@@ -342,6 +370,154 @@ export class GlmtProxy {
     });
   }
 
+  /**
+   * Sleep for specified milliseconds
+   */
+  private sleep(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+  }
+
+  /**
+   * Calculate retry delay with exponential backoff and jitter
+   * Honors Retry-After header if provided
+   */
+  private calculateRetryDelay(attempt: number, retryAfterHeader?: string): number {
+    // Honor Retry-After header if present and valid
+    if (retryAfterHeader) {
+      const retryAfter = parseInt(retryAfterHeader, 10);
+      if (!isNaN(retryAfter) && retryAfter > 0) {
+        return retryAfter * 1000; // Convert seconds to ms
+      }
+    }
+    // Exponential backoff: 2^attempt * baseDelay + random jitter (0-500ms), capped at 30s
+    const MAX_DELAY_MS = 30000;
+    const exponentialDelay = Math.min(
+      Math.pow(2, attempt) * this.retryConfig.baseDelay,
+      MAX_DELAY_MS
+    );
+    const jitter = Math.random() * 500;
+    return exponentialDelay + jitter;
+  }
+
+  /**
+   * Check if error is retryable (429 rate limit)
+   * Returns retryable status and optional Retry-After value
+   */
+  private isRetryableError(error: Error): { retryable: boolean; retryAfter?: string } {
+    const message = error.message;
+    if (message.includes('429') || message.toLowerCase().includes('rate limit')) {
+      // Prefer Retry-After from error property (captured from HTTP header), fallback to parsing message
+      const errWithRetryAfter = error as Error & { retryAfter?: string };
+      if (errWithRetryAfter.retryAfter) {
+        return { retryable: true, retryAfter: errWithRetryAfter.retryAfter };
+      }
+      const retryAfterMatch = message.match(/retry-after:\s*(\d+)/i);
+      return {
+        retryable: true,
+        retryAfter: retryAfterMatch?.[1],
+      };
+    }
+    return { retryable: false };
+  }
+
+  /**
+   * Forward request with retry logic for rate limit errors
+   */
+  private async forwardWithRetry(
+    openaiRequest: OpenAIRequest,
+    headers: Record<string, string | undefined>
+  ): Promise<unknown> {
+    if (!this.retryConfig.enabled) {
+      return this.forwardToUpstream(openaiRequest, headers);
+    }
+
+    let lastError: Error | null = null;
+
+    for (let attempt = 0; attempt <= this.retryConfig.maxRetries; attempt++) {
+      try {
+        return await this.forwardToUpstream(openaiRequest, headers);
+      } catch (error) {
+        const err = error as Error;
+        const { retryable, retryAfter } = this.isRetryableError(err);
+
+        if (!retryable || attempt >= this.retryConfig.maxRetries) {
+          throw err;
+        }
+
+        lastError = err;
+        const delay = this.calculateRetryDelay(attempt, retryAfter);
+
+        console.error(
+          `[glmt-proxy] Rate limited, retry ${attempt + 1}/${this.retryConfig.maxRetries} after ${Math.round(delay)}ms`
+        );
+
+        await this.sleep(delay);
+      }
+    }
+
+    throw lastError ?? new Error('Retry failed');
+  }
+
+  /**
+   * Forward streaming request with retry logic for rate limit errors
+   * Only retries if headers haven't been sent yet
+   */
+  private async forwardAndStreamWithRetry(
+    openaiRequest: OpenAIRequest,
+    headers: Record<string, string | undefined>,
+    clientRes: http.ServerResponse,
+    thinkingConfig: ThinkingConfig,
+    startTime: number
+  ): Promise<void> {
+    if (!this.retryConfig.enabled) {
+      return this.forwardAndStreamUpstream(
+        openaiRequest,
+        headers,
+        clientRes,
+        thinkingConfig,
+        startTime
+      );
+    }
+
+    let lastError: Error | null = null;
+
+    for (let attempt = 0; attempt <= this.retryConfig.maxRetries; attempt++) {
+      try {
+        return await this.forwardAndStreamUpstream(
+          openaiRequest,
+          headers,
+          clientRes,
+          thinkingConfig,
+          startTime
+        );
+      } catch (error) {
+        const err = error as Error;
+
+        // Don't retry if headers already sent (streaming started)
+        if (clientRes.headersSent) {
+          throw err;
+        }
+
+        const { retryable, retryAfter } = this.isRetryableError(err);
+
+        if (!retryable || attempt >= this.retryConfig.maxRetries) {
+          throw err;
+        }
+
+        lastError = err;
+        const delay = this.calculateRetryDelay(attempt, retryAfter);
+
+        console.error(
+          `[glmt-proxy] Rate limited, retry ${attempt + 1}/${this.retryConfig.maxRetries} after ${Math.round(delay)}ms`
+        );
+
+        await this.sleep(delay);
+      }
+    }
+
+    throw lastError ?? new Error('Retry failed');
+  }
+
   /**
    * Forward request to Z.AI upstream
    */
@@ -362,6 +538,7 @@ export class GlmtProxy {
         port: url.port || 443,
         path: url.pathname || '/api/coding/paas/v4/chat/completions',
         method: 'POST',
+        agent: this.httpsAgent,
         headers: {
           'Content-Type': 'application/json',
           'Content-Length': Buffer.byteLength(requestBody),
@@ -392,7 +569,17 @@ export class GlmtProxy {
 
             // Check for non-200 status
             if (res.statusCode !== 200) {
-              reject(new Error(`Upstream error: ${res.statusCode} ${res.statusMessage}\n${body}`));
+              const err = new Error(
+                `Upstream error: ${res.statusCode} ${res.statusMessage}\n${body}`
+              ) as Error & { retryAfter?: string };
+              // Capture Retry-After header for rate limit handling
+              const retryAfterHeader = res.headers['retry-after'];
+              if (retryAfterHeader) {
+                err.retryAfter = Array.isArray(retryAfterHeader)
+                  ? retryAfterHeader[0]
+                  : retryAfterHeader;
+              }
+              reject(err);
               return;
             }
 
@@ -438,6 +625,7 @@ export class GlmtProxy {
         port: url.port || 443,
         path: url.pathname || '/api/coding/paas/v4/chat/completions',
         method: 'POST',
+        agent: this.httpsAgent,
         headers: {
           'Content-Type': 'application/json',
           'Content-Length': Buffer.byteLength(requestBody),
@@ -461,7 +649,17 @@ export class GlmtProxy {
           let body = '';
           upstreamRes.on('data', (chunk: Buffer) => (body += chunk.toString()));
           upstreamRes.on('end', () => {
-            reject(new Error(`Upstream error: ${upstreamRes.statusCode}\n${body}`));
+            const err = new Error(`Upstream error: ${upstreamRes.statusCode}\n${body}`) as Error & {
+              retryAfter?: string;
+            };
+            // Capture Retry-After header for rate limit handling
+            const retryAfterHeader = upstreamRes.headers['retry-after'];
+            if (retryAfterHeader) {
+              err.retryAfter = Array.isArray(retryAfterHeader)
+                ? retryAfterHeader[0]
+                : retryAfterHeader;
+            }
+            reject(err);
           });
           return;
         }
@@ -537,6 +735,10 @@ export class GlmtProxy {
       this.log('Stopping proxy server');
       this.server.close();
     }
+    // Destroy connection pool
+    if (this.httpsAgent) {
+      this.httpsAgent.destroy();
+    }
   }
 
   /**