From a803cd72910ea7d4e0443ee3b0a54228f6185319 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 1 Sep 2026 15:51:10 -0500 Subject: [PATCH] fix(cloud-agent-next): retry compute status connection resets Fixes KILOCODE-WEB-27TA --- .../cloud-agent-client.test.ts | 119 ++++++++++++++++++ .../cloud-agent-next/cloud-agent-client.ts | 22 +++- 2 files changed, 140 insertions(+), 1 deletion(-) diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts index 581d5e4288..2661665c88 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts @@ -2,10 +2,12 @@ import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals import type * as TrpcClientModule from '@trpc/client'; import type { CloudAgentNextClient as CloudAgentNextClientType, + ComputeBillingStatus, CreateWorktreeChatInput, CreateWorktreeChatOutput, DeleteWorktreeInput, DeleteWorktreeOutput, + GetSessionInput, PrepareSessionInput, SendMessageInput, } from './cloud-agent-client'; @@ -235,6 +237,123 @@ describe('CloudAgentNextClient sensitive error reporting', () => { }); }); +describe('CloudAgentNextClient.getComputeBillingStatus', () => { + type StatusQuery = (input: GetSessionInput) => Promise; + const sessionId = 'agent_12345678-1234-4234-9234-123456789abc'; + const status: ComputeBillingStatus = { + payer: { type: 'user', id: 'user-123' }, + attribution: 'session', + phase: 'idle', + estimatedHourlyRateMicrodollars: null, + estimatedIntervalAmountMicrodollars: null, + billingMode: null, + interval: null, + }; + const { TRPCClientError } = jest.requireActual('@trpc/client'); + const connectionReset = () => + TRPCClientError.from( + new TypeError('fetch failed', { + cause: Object.assign(new Error('read ECONNRESET'), { code: 'ECONNRESET' }), + }) + ); + + beforeEach(() => { + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('returns successful status without retrying', async () => { + const query = jest.fn().mockResolvedValue(status); + mockCreateTRPCClient.mockReturnValueOnce({ getComputeBillingStatus: { query } }); + + await expect( + new CloudAgentNextClient('token').getComputeBillingStatus(sessionId) + ).resolves.toBe(status); + expect(query).toHaveBeenCalledTimes(1); + expect(query).toHaveBeenCalledWith({ cloudAgentSessionId: sessionId }); + expect(jest.getTimerCount()).toBe(0); + }); + + it('retries a nested connection reset once after a short delay', async () => { + const query = jest + .fn() + .mockRejectedValueOnce(connectionReset()) + .mockResolvedValueOnce(status); + mockCreateTRPCClient.mockReturnValueOnce({ getComputeBillingStatus: { query } }); + + const result = new CloudAgentNextClient('token').getComputeBillingStatus(sessionId); + await Promise.all([ + expect(result).resolves.toBe(status), + (async () => { + await jest.advanceTimersByTimeAsync(99); + expect(query).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(101); + })(), + ]); + + expect(query).toHaveBeenCalledTimes(2); + expect(query).toHaveBeenNthCalledWith(1, { cloudAgentSessionId: sessionId }); + expect(query).toHaveBeenNthCalledWith(2, { cloudAgentSessionId: sessionId }); + expect(mockCaptureException).not.toHaveBeenCalled(); + }); + + it.each([connectionReset(), new Error('Worker unavailable')])( + 'propagates the second failure unchanged without a third attempt: %s', + async error => { + const query = jest + .fn() + .mockRejectedValueOnce(connectionReset()) + .mockRejectedValueOnce(error); + mockCreateTRPCClient.mockReturnValueOnce({ getComputeBillingStatus: { query } }); + + const result = new CloudAgentNextClient('token').getComputeBillingStatus(sessionId); + await Promise.all([expect(result).rejects.toBe(error), jest.runAllTimersAsync()]); + + expect(query).toHaveBeenCalledTimes(2); + expect(jest.getTimerCount()).toBe(0); + } + ); + + it.each([ + new TRPCClientError('Forbidden', { + result: { + error: { code: -32003, message: 'Forbidden', data: { code: 'FORBIDDEN', httpStatus: 403 } }, + }, + }), + new TypeError('fetch failed'), + new Error('read ECONNRESET'), + new TypeError('fetch failed', { cause: { code: 'ETIMEDOUT' } }), + new Error('Malformed cause', { cause: 'ECONNRESET' }), + null, + undefined, + ])('does not retry unrelated or unstructured errors: %s', async error => { + const query = jest.fn().mockRejectedValue(error); + mockCreateTRPCClient.mockReturnValueOnce({ getComputeBillingStatus: { query } }); + + await expect(new CloudAgentNextClient('token').getComputeBillingStatus(sessionId)).rejects.toBe( + error + ); + expect(query).toHaveBeenCalledTimes(1); + expect(jest.getTimerCount()).toBe(0); + }); + + it('terminates cause inspection for cyclic errors', async () => { + const error = new Error('Cyclic cause'); + error.cause = error; + const query = jest.fn().mockRejectedValue(error); + mockCreateTRPCClient.mockReturnValueOnce({ getComputeBillingStatus: { query } }); + + await expect(new CloudAgentNextClient('token').getComputeBillingStatus(sessionId)).rejects.toBe( + error + ); + expect(query).toHaveBeenCalledTimes(1); + expect(jest.getTimerCount()).toBe(0); + }); +}); + describe('CloudAgentNextClient.deleteSession', () => { const sessionId = 'workspace_12345678-1234-4234-9234-123456789abc'; diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts index 7ccd741260..2f752d9ce5 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts @@ -26,6 +26,18 @@ export type { SendMessagePayload } from './types.js'; // TODO: Update this URL when the new cloud-agent-next worker is deployed const CLOUD_AGENT_NEXT_API_URL = getEnvVariable('CLOUD_AGENT_NEXT_API_URL') || ''; +function isConnectionResetError(error: unknown): boolean { + const seen = new Set(); + while (typeof error === 'object' && error !== null && !seen.has(error)) { + seen.add(error); + if ('code' in error && error.code === 'ECONNRESET') { + return true; + } + error = 'cause' in error ? error.cause : undefined; + } + return false; +} + // MCP server config types — CLI-native local/remote format. // Each env/header value is either a plain string (passed through verbatim) // or an RSA+AES envelope (decrypted per-value by the worker when @@ -786,7 +798,15 @@ export class CloudAgentNextClient { } async getComputeBillingStatus(cloudAgentSessionId: string): Promise { - return await this.client.getComputeBillingStatus.query({ cloudAgentSessionId }); + try { + return await this.client.getComputeBillingStatus.query({ cloudAgentSessionId }); + } catch (error) { + if (!isConnectionResetError(error)) { + throw error; + } + await new Promise(resolve => setTimeout(resolve, 100 + Math.random() * 100)); + return await this.client.getComputeBillingStatus.query({ cloudAgentSessionId }); + } } async createWorktreeChat(input: CreateWorktreeChatInput): Promise {