diff --git a/src/glmt/glmt-proxy.ts b/src/glmt/glmt-proxy.ts index d1170f6e..1a07d163 100644 --- a/src/glmt/glmt-proxy.ts +++ b/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,21 @@ 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 + this.retryConfig = { + maxRetries: parseInt(process.env.GLMT_MAX_RETRIES || '3', 10), + baseDelay: parseInt(process.env.GLMT_RETRY_BASE_DELAY || '1000', 10), + 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 +275,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 +335,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 +368,146 @@ export class GlmtProxy { }); } + /** + * Sleep for specified milliseconds + */ + private sleep(ms: number): Promise { + 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) + const exponentialDelay = Math.pow(2, attempt) * this.retryConfig.baseDelay; + 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')) { + // Try to extract Retry-After from error message + 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 + ): Promise { + 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, + clientRes: http.ServerResponse, + thinkingConfig: ThinkingConfig, + startTime: number + ): Promise { + 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 +528,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), @@ -438,6 +605,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), @@ -537,6 +705,10 @@ export class GlmtProxy { this.log('Stopping proxy server'); this.server.close(); } + // Destroy connection pool + if (this.httpsAgent) { + this.httpsAgent.destroy(); + } } /** diff --git a/tests/unit/glmt/retry-logic.test.ts b/tests/unit/glmt/retry-logic.test.ts new file mode 100644 index 00000000..ab289fd1 --- /dev/null +++ b/tests/unit/glmt/retry-logic.test.ts @@ -0,0 +1,276 @@ +/** + * GLMT Retry Logic Unit Tests + * + * Tests for exponential backoff retry behavior on 429 rate limit errors + */ + +import { describe, it, expect, beforeEach, afterEach } from 'bun:test'; + +// Store original env vars +const originalEnv = { ...process.env }; + +// Reset env before each test +beforeEach(() => { + delete process.env.GLMT_MAX_RETRIES; + delete process.env.GLMT_RETRY_BASE_DELAY; + delete process.env.GLMT_DISABLE_RETRY; +}); + +afterEach(() => { + // Restore original env + Object.keys(process.env).forEach((key) => { + if (key.startsWith('GLMT_')) { + delete process.env[key]; + } + }); + Object.assign(process.env, originalEnv); +}); + +// Helper to create proxy instance with specific config +async function createTestableProxy(config: { + maxRetries?: number; + baseDelay?: number; + enabled?: boolean; +} = {}) { + // Set env vars before import + if (config.maxRetries !== undefined) { + process.env.GLMT_MAX_RETRIES = String(config.maxRetries); + } + if (config.baseDelay !== undefined) { + process.env.GLMT_RETRY_BASE_DELAY = String(config.baseDelay); + } + if (config.enabled === false) { + process.env.GLMT_DISABLE_RETRY = '1'; + } + + // Dynamic import to pick up env changes + const module = await import('../../../src/glmt/glmt-proxy'); + return new module.GlmtProxy({ verbose: false }); +} + +describe('GLMT Retry Logic', () => { + describe('RetryConfig initialization', () => { + it('should use default values when env vars not set', async () => { + const proxy = await createTestableProxy(); + // Access private via type assertion for testing + const config = (proxy as unknown as { retryConfig: { maxRetries: number; baseDelay: number; enabled: boolean } }).retryConfig; + expect(config.maxRetries).toBe(3); + expect(config.baseDelay).toBe(1000); + expect(config.enabled).toBe(true); + }); + + it('should respect GLMT_MAX_RETRIES env var', async () => { + const proxy = await createTestableProxy({ maxRetries: 5 }); + const config = (proxy as unknown as { retryConfig: { maxRetries: number } }).retryConfig; + expect(config.maxRetries).toBe(5); + }); + + it('should respect GLMT_RETRY_BASE_DELAY env var', async () => { + const proxy = await createTestableProxy({ baseDelay: 2000 }); + const config = (proxy as unknown as { retryConfig: { baseDelay: number } }).retryConfig; + expect(config.baseDelay).toBe(2000); + }); + + it('should disable retry when GLMT_DISABLE_RETRY=1', async () => { + const proxy = await createTestableProxy({ enabled: false }); + const config = (proxy as unknown as { retryConfig: { enabled: boolean } }).retryConfig; + expect(config.enabled).toBe(false); + }); + }); + + describe('calculateRetryDelay', () => { + it('should calculate exponential delay with jitter', async () => { + const proxy = await createTestableProxy({ baseDelay: 1000 }); + const calcDelay = (proxy as unknown as { calculateRetryDelay: (attempt: number, retryAfter?: string) => number }).calculateRetryDelay.bind(proxy); + + // Attempt 0: 2^0 * 1000 = 1000 + jitter (0-500) + const delay0 = calcDelay(0); + expect(delay0).toBeGreaterThanOrEqual(1000); + expect(delay0).toBeLessThan(1500); + + // Attempt 1: 2^1 * 1000 = 2000 + jitter (0-500) + const delay1 = calcDelay(1); + expect(delay1).toBeGreaterThanOrEqual(2000); + expect(delay1).toBeLessThan(2500); + + // Attempt 2: 2^2 * 1000 = 4000 + jitter (0-500) + const delay2 = calcDelay(2); + expect(delay2).toBeGreaterThanOrEqual(4000); + expect(delay2).toBeLessThan(4500); + }); + + it('should honor Retry-After header in seconds', async () => { + const proxy = await createTestableProxy(); + const calcDelay = (proxy as unknown as { calculateRetryDelay: (attempt: number, retryAfter?: string) => number }).calculateRetryDelay.bind(proxy); + + // Retry-After: 5 seconds → 5000ms + const delay = calcDelay(0, '5'); + expect(delay).toBe(5000); + }); + + it('should ignore invalid Retry-After header and fallback to exponential', async () => { + const proxy = await createTestableProxy({ baseDelay: 1000 }); + const calcDelay = (proxy as unknown as { calculateRetryDelay: (attempt: number, retryAfter?: string) => number }).calculateRetryDelay.bind(proxy); + + // Invalid header falls back to exponential + const delay = calcDelay(0, 'invalid'); + expect(delay).toBeGreaterThanOrEqual(1000); + expect(delay).toBeLessThan(1500); + }); + + it('should ignore zero or negative Retry-After', async () => { + const proxy = await createTestableProxy({ baseDelay: 1000 }); + const calcDelay = (proxy as unknown as { calculateRetryDelay: (attempt: number, retryAfter?: string) => number }).calculateRetryDelay.bind(proxy); + + const delay = calcDelay(0, '0'); + expect(delay).toBeGreaterThanOrEqual(1000); + expect(delay).toBeLessThan(1500); + }); + }); + + describe('isRetryableError', () => { + it('should return true for 429 status code', async () => { + const proxy = await createTestableProxy(); + const isRetryable = (proxy as unknown as { isRetryableError: (error: Error) => { retryable: boolean; retryAfter?: string } }).isRetryableError.bind(proxy); + + const result = isRetryable(new Error('Upstream error: 429 Too Many Requests')); + expect(result.retryable).toBe(true); + }); + + it('should return true for rate limit message', async () => { + const proxy = await createTestableProxy(); + const isRetryable = (proxy as unknown as { isRetryableError: (error: Error) => { retryable: boolean; retryAfter?: string } }).isRetryableError.bind(proxy); + + const result = isRetryable(new Error('Rate limit exceeded')); + expect(result.retryable).toBe(true); + }); + + it('should return false for non-retryable errors', async () => { + const proxy = await createTestableProxy(); + const isRetryable = (proxy as unknown as { isRetryableError: (error: Error) => { retryable: boolean; retryAfter?: string } }).isRetryableError.bind(proxy); + + expect(isRetryable(new Error('Connection refused')).retryable).toBe(false); + expect(isRetryable(new Error('Timeout')).retryable).toBe(false); + expect(isRetryable(new Error('500 Internal Server Error')).retryable).toBe(false); + expect(isRetryable(new Error('401 Unauthorized')).retryable).toBe(false); + }); + + it('should extract Retry-After from error message', async () => { + const proxy = await createTestableProxy(); + const isRetryable = (proxy as unknown as { isRetryableError: (error: Error) => { retryable: boolean; retryAfter?: string } }).isRetryableError.bind(proxy); + + const result = isRetryable(new Error('429 Too Many Requests, Retry-After: 10')); + expect(result.retryable).toBe(true); + expect(result.retryAfter).toBe('10'); + }); + }); + + describe('forwardWithRetry behavior', () => { + it('should succeed on first attempt without retry', async () => { + const proxy = await createTestableProxy(); + let attempts = 0; + + // Mock forwardToUpstream + (proxy as unknown as { forwardToUpstream: () => Promise }).forwardToUpstream = async () => { + attempts++; + return { choices: [{ message: { content: 'success' } }] }; + }; + + const forwardWithRetry = (proxy as unknown as { forwardWithRetry: (req: unknown, headers: unknown) => Promise }).forwardWithRetry.bind(proxy); + const result = await forwardWithRetry({}, {}); + + expect(attempts).toBe(1); + expect((result as { choices: Array<{ message: { content: string } }> }).choices[0].message.content).toBe('success'); + }); + + it('should retry on 429 and succeed eventually', async () => { + const proxy = await createTestableProxy({ baseDelay: 10 }); // Fast for tests + let attempts = 0; + + (proxy as unknown as { forwardToUpstream: () => Promise }).forwardToUpstream = async () => { + attempts++; + if (attempts < 3) { + throw new Error('Upstream error: 429 Too Many Requests'); + } + return { choices: [{ message: { content: 'success after retry' } }] }; + }; + + const forwardWithRetry = (proxy as unknown as { forwardWithRetry: (req: unknown, headers: unknown) => Promise }).forwardWithRetry.bind(proxy); + const result = await forwardWithRetry({}, {}); + + expect(attempts).toBe(3); + expect((result as { choices: Array<{ message: { content: string } }> }).choices[0].message.content).toBe('success after retry'); + }); + + it('should fail after max retries exhausted', async () => { + const proxy = await createTestableProxy({ maxRetries: 2, baseDelay: 10 }); + let attempts = 0; + + (proxy as unknown as { forwardToUpstream: () => Promise }).forwardToUpstream = async () => { + attempts++; + throw new Error('Upstream error: 429 Too Many Requests'); + }; + + const forwardWithRetry = (proxy as unknown as { forwardWithRetry: (req: unknown, headers: unknown) => Promise }).forwardWithRetry.bind(proxy); + + await expect(forwardWithRetry({}, {})).rejects.toThrow('429'); + expect(attempts).toBe(3); // Initial + 2 retries + }); + + it('should not retry when disabled', async () => { + const proxy = await createTestableProxy({ enabled: false }); + let attempts = 0; + + (proxy as unknown as { forwardToUpstream: () => Promise }).forwardToUpstream = async () => { + attempts++; + throw new Error('Upstream error: 429 Too Many Requests'); + }; + + const forwardWithRetry = (proxy as unknown as { forwardWithRetry: (req: unknown, headers: unknown) => Promise }).forwardWithRetry.bind(proxy); + + await expect(forwardWithRetry({}, {})).rejects.toThrow('429'); + expect(attempts).toBe(1); + }); + + it('should not retry non-429 errors', async () => { + const proxy = await createTestableProxy({ baseDelay: 10 }); + let attempts = 0; + + (proxy as unknown as { forwardToUpstream: () => Promise }).forwardToUpstream = async () => { + attempts++; + throw new Error('Upstream error: 500 Internal Server Error'); + }; + + const forwardWithRetry = (proxy as unknown as { forwardWithRetry: (req: unknown, headers: unknown) => Promise }).forwardWithRetry.bind(proxy); + + await expect(forwardWithRetry({}, {})).rejects.toThrow('500'); + expect(attempts).toBe(1); + }); + }); + + describe('connection pooling', () => { + it('should create https.Agent with keepAlive enabled', async () => { + const proxy = await createTestableProxy(); + const agent = (proxy as unknown as { httpsAgent: { options?: { keepAlive?: boolean } } }).httpsAgent; + + expect(agent).toBeDefined(); + // Agent should have keepAlive behavior (internal property) + expect(agent.options?.keepAlive).toBe(true); + }); + + it('should destroy agent on stop', async () => { + const proxy = await createTestableProxy(); + const agent = (proxy as unknown as { httpsAgent: { destroy: () => void; destroyed?: boolean } }).httpsAgent; + + let destroyed = false; + const originalDestroy = agent.destroy.bind(agent); + agent.destroy = () => { + destroyed = true; + originalDestroy(); + }; + + proxy.stop(); + expect(destroyed).toBe(true); + }); + }); +});