diff --git a/src/__tests__/usageMetering.test.ts b/src/__tests__/usageMetering.test.ts index c9a32834..488a6661 100644 --- a/src/__tests__/usageMetering.test.ts +++ b/src/__tests__/usageMetering.test.ts @@ -117,7 +117,7 @@ beforeEach(async () => { billing.clear(); billing.setBalance(TEST_DEVELOPER_ID, 1000); rateLimiter.reset(); - currentProxyConfig = { timeoutMs: 2000 }; + currentProxyConfig = { timeoutMs: 2000, allowedHosts: ['localhost', '127.0.0.1'] }; await startProxy(); setUpstreamHandler((_req, res) => { @@ -215,7 +215,11 @@ describe('Usage Metering & Billing (Post-Proxy)', () => { }); it('records usage on 500 if configured in recordableStatuses', async () => { - currentProxyConfig = { timeoutMs: 2000, recordableStatuses: () => true }; + currentProxyConfig = { + timeoutMs: 2000, + allowedHosts: ['localhost', '127.0.0.1'], + recordableStatuses: () => true, + }; await startProxy(); setUpstreamHandler((_req, res) => { @@ -239,6 +243,61 @@ describe('Usage Metering & Billing (Post-Proxy)', () => { expect(billing.getBalance(TEST_DEVELOPER_ID)).toBeCloseTo(999.95); }); + it('records usage reliably for multi-chunk streamed responses', async () => { + setUpstreamHandler((_req, res) => { + res.status(200); + res.setHeader('Content-Type', 'text/plain'); + res.write('chunk-1:'); + setTimeout(() => { + res.write('chunk-2:'); + res.end('chunk-3'); + }, 15); + }); + + const res = await fetch(`${proxyUrl}/v1/call/${TEST_API_SLUG}/data`, { + method: 'GET', + headers: { 'x-api-key': TEST_API_KEY }, + }); + expect(res.status).toBe(200); + const text = await res.text(); + expect(text).toBe('chunk-1:chunk-2:chunk-3'); + + await yieldTick(); + + const events = usageStore.getEvents(TEST_API_KEY); + expect(events).toHaveLength(1); + expect(events[0].endpointId).toBe('ep_data'); + expect(billing.getBalance(TEST_DEVELOPER_ID)).toBeCloseTo(999.95); + }); + + it('does NOT record usage or deduct billing when upstream stream aborts mid-response', async () => { + setUpstreamHandler((_req, res) => { + res.writeHead(200, { + 'Content-Type': 'text/plain', + 'Content-Length': '1000', + }); + res.write('partial-body'); + setTimeout(() => { + res.destroy(new Error('upstream connection reset')); + }, 10); + }); + + try { + const res = await fetch(`${proxyUrl}/v1/call/${TEST_API_SLUG}/data`, { + method: 'GET', + headers: { 'x-api-key': TEST_API_KEY }, + }); + await res.text(); + } catch { + // Expected mid-stream termination + } + + await yieldTick(); + + expect(usageStore.getEvents()).toHaveLength(0); + expect(billing.getBalance(TEST_DEVELOPER_ID)).toBe(1000); + }); + it('records neither usage nor deduction when anchored charging fails', async () => { billing.failNextUsageCharge('pending billing write failed', true); const errorSpy = jest.spyOn(console, 'error').mockImplementation(() => undefined); @@ -249,6 +308,8 @@ describe('Usage Metering & Billing (Post-Proxy)', () => { }); expect(res.status).toBe(200); + await yieldTick(); + expect(usageStore.getEvents()).toHaveLength(0); expect(billing.getBalance(TEST_DEVELOPER_ID)).toBe(1000); expect(errorSpy).toHaveBeenCalledWith( @@ -272,26 +333,30 @@ describe('Usage Metering & Billing (Post-Proxy)', () => { throw new Error('usage store offline'); }); - const res = await fetch(`${proxyUrl}/v1/call/${TEST_API_SLUG}/data`, { - method: 'GET', - headers: { 'x-api-key': TEST_API_KEY }, - }); - expect(res.status).toBe(200); - - expect(billing.getBalance(TEST_DEVELOPER_ID)).toBeCloseTo(999.95); - expect(errorSpy).toHaveBeenCalledWith( - '[proxy billing reconciliation] Usage view write threw after successful billing charge', - expect.objectContaining({ - requestId: expect.any(String), - apiId: TEST_API_ID, - endpointId: 'ep_data', - developerId: TEST_DEVELOPER_ID, - error: 'usage store offline', - }), - ); - - recordSpy.mockRestore(); - errorSpy.mockRestore(); + try { + const res = await fetch(`${proxyUrl}/v1/call/${TEST_API_SLUG}/data`, { + method: 'GET', + headers: { 'x-api-key': TEST_API_KEY }, + }); + expect(res.status).toBe(200); + + await yieldTick(); + + expect(billing.getBalance(TEST_DEVELOPER_ID)).toBeCloseTo(999.95); + expect(errorSpy).toHaveBeenCalledWith( + '[proxy billing reconciliation] Usage view write threw after successful billing charge', + expect.objectContaining({ + requestId: expect.any(String), + apiId: TEST_API_ID, + endpointId: 'ep_data', + developerId: TEST_DEVELOPER_ID, + error: 'usage store offline', + }), + ); + } finally { + recordSpy.mockRestore(); + errorSpy.mockRestore(); + } }); it('increments gateway_usage_record_failures_total when usageStore.record throws', async () => { @@ -302,18 +367,20 @@ describe('Usage Metering & Billing (Post-Proxy)', () => { throw new Error('usage store offline'); }); - const res = await fetch(`${proxyUrl}/v1/call/${TEST_API_SLUG}/data`, { - method: 'GET', - headers: { 'x-api-key': TEST_API_KEY }, - }); - expect(res.status).toBe(200); - - await yieldTick(); + try { + const res = await fetch(`${proxyUrl}/v1/call/${TEST_API_SLUG}/data`, { + method: 'GET', + headers: { 'x-api-key': TEST_API_KEY }, + }); + expect(res.status).toBe(200); - const after = await getUsageRecordFailuresMetric(); - expect(after).toBeGreaterThan(before); + await yieldTick(); - recordSpy.mockRestore(); + const after = await getUsageRecordFailuresMetric(); + expect(after).toBeGreaterThan(before); + } finally { + recordSpy.mockRestore(); + } }); it('exposes gateway_usage_record_failures_total at /api/metrics', async () => { @@ -365,7 +432,8 @@ describe('Usage Metering & Billing (Post-Proxy)', () => { expect(res.status).toBe(402); const body = await res.json(); - expect(body.message ?? body.error).toMatch(/insufficient balance/i); + const errorMsg = typeof body.error === 'string' ? body.error : body.error?.message ?? body.message; + expect(errorMsg).toMatch(/insufficient balance/i); await yieldTick(); diff --git a/src/routes/proxyRoutes.ts b/src/routes/proxyRoutes.ts index cafcf43d..b69333bb 100644 --- a/src/routes/proxyRoutes.ts +++ b/src/routes/proxyRoutes.ts @@ -264,6 +264,134 @@ export function createProxyRouter(deps: ProxyDeps): Router { }); res.set('x-request-id', requestId); + // 8. Keep metering and billing consistent — attach 'finish' and 'close' + // listeners BEFORE streaming or sending the response body so we never + // miss a synchronous completion or record usage on a mid-stream abort. + let responseFinished = false; + let streamErrored = false; + + if (config.recordableStatuses(upstreamStatus)) { + res.once('finish', () => { + if (streamErrored) { + return; + } + responseFinished = true; + + void (async () => { + const timestamp = new Date().toISOString(); + let chargeSucceeded = false; + + try { + if (endpoint.priceUsdc > 0) { + if (typeof billing.chargeUsage === 'function') { + const chargeResult = await billing.chargeUsage({ + requestId, + developerId: keyRecord.userId, + apiKey: apiKeyHeader, + apiId: String(apiEntry.id), + endpointId: endpoint.endpointId, + amountUsdc: endpoint.priceUsdc, + statusCode: upstreamStatus, + timestamp, + }); + + if (!chargeResult.success) { + if (chargeResult.reconciliationRequired) { + console.error( + '[proxy billing reconciliation] Billing anchor failed after usage write phase started', + { + requestId, + apiId: String(apiEntry.id), + endpointId: endpoint.endpointId, + developerId: keyRecord.userId, + error: chargeResult.error, + }, + ); + } + return; + } + chargeSucceeded = true; + } + } + + const recorded = await usageStore.record({ + id: randomUUID(), // ID of the usage event itself + requestId, // Idempotency key — prevents double-counts + apiKey: apiKeyHeader, + apiKeyId: keyRecord.id, + apiId: String(apiEntry.id), + endpointId: endpoint.endpointId, + userId: keyRecord.userId, + amountUsdc: endpoint.priceUsdc, + statusCode: upstreamStatus, + timestamp, + }); + + if (recorded) { + defaultUsageSseBroadcaster.emitForUser(keyRecord.userId, { + id: randomUUID(), + requestId, + apiKeyId: keyRecord.id, + apiKeyPrefix: maskApiKey(apiKeyHeader), + apiId: String(apiEntry.id), + endpointId: endpoint.endpointId, + userId: keyRecord.userId, + amountUsdc: endpoint.priceUsdc, + statusCode: upstreamStatus, + timestamp, + }); + } + + recordEndpointThroughputSaturation({ + apiId: String(apiEntry.id), + endpointId: endpoint.endpointId, + endpointPath: endpoint.path, + advertisedLimitPerMinute: Number(keyRecord?.rateLimitPerMinute ?? 0), + observedAt: Date.now(), + }); + + if ( + recorded && + endpoint.priceUsdc > 0 && + typeof billing.chargeUsage !== 'function' + ) { + billing.deductCredit(keyRecord.userId, endpoint.priceUsdc).catch((err) => { + console.error('Background billing deduction failed:', err); + }); + } + } catch (err) { + recordUsageRecordFailure(); + if (chargeSucceeded) { + console.error( + '[proxy billing reconciliation] Usage view write threw after successful billing charge', + { + requestId, + apiId: String(apiEntry.id), + endpointId: endpoint.endpointId, + developerId: keyRecord.userId, + error: err instanceof Error ? err.message : String(err), + }, + ); + } + logger.error('Background usage recording failed', { + requestId, + apiId: String(apiEntry.id), + error: err, + }); + } + })(); + }); + + res.once('close', () => { + // 'close' fires after 'finish' on a normal response, or on its own + // when the socket is destroyed prematurely. Only treat it as an + // abort when 'finish' has NOT already fired. + if (!responseFinished) { + recordProxyPrematureAbort(); + } + }); + } + // Stream body back res.status(upstreamStatus); if (upstreamRes.body) { @@ -287,6 +415,7 @@ export function createProxyRouter(deps: ProxyDeps): Router { try { await pump(); } catch (streamErr) { + streamErrored = true; // Mid-stream failure: headers are already flushed, so the error // handler cannot send an envelope. Cancel the upstream body so its // socket is freed; errorHandler then destroys the client response @@ -359,99 +488,6 @@ export function createProxyRouter(deps: ProxyDeps): Router { throw new BadGatewayError('Bad Gateway: upstream unreachable'); } } - - // 8. Keep metering and billing consistent — but ONLY after the response - // has been fully delivered to the caller. - // - // We distinguish two response lifecycle events: - // • 'finish' — Node/Express has flushed all data and ended the - // response normally. This is the success path; we - // record usage here. - // • 'close' — The underlying socket was torn down. When this - // fires WITHOUT a prior 'finish' it means the client - // disconnected mid-stream (premature abort). In that - // case we must NOT record usage because the caller - // never received the response. - // - // Using a one-shot 'finish' listener (registered before we start - // streaming) ensures we capture the event even if the stream - // completes synchronously. The 'close' listener is a guard that - // cancels the deferred work when the socket drops first. - if (config.recordableStatuses(upstreamStatus)) { - // Track whether the response finished cleanly before the socket closed. - let responseFinished = false; - - res.once('finish', () => { - responseFinished = true; - - // Run usage recording in a non-blocking microtask so it does not - // delay the event loop that is already handling the next request. - setImmediate(() => { - void (async () => { - try { - const recorded = await usageStore.record({ - id: randomUUID(), // ID of the usage event itself - requestId, // Idempotency key — prevents double-counts - apiKeyId: keyRecord.id, - apiId: String(apiEntry.id), - endpointId: endpoint.endpointId, - userId: keyRecord.userId, - amountUsdc: endpoint.priceUsdc, - statusCode: upstreamStatus, - timestamp: new Date().toISOString(), - }); - - if (recorded) { - defaultUsageSseBroadcaster.emitForUser(keyRecord.userId, { - id: randomUUID(), - requestId, - apiKeyId: keyRecord.id, - apiKeyPrefix: maskApiKey(apiKeyHeader), - apiId: String(apiEntry.id), - endpointId: endpoint.endpointId, - userId: keyRecord.userId, - amountUsdc: endpoint.priceUsdc, - statusCode: upstreamStatus, - timestamp: new Date().toISOString(), - }); - } - - recordEndpointThroughputSaturation({ - apiId: String(apiEntry.id), - endpointId: endpoint.endpointId, - endpointPath: endpoint.path, - advertisedLimitPerMinute: Number(keyRecord?.rateLimitPerMinute ?? 0), - observedAt: Date.now(), - }); - - // Only deduct billing if this requestId hasn't been processed - // before (idempotency guard inside usageStore.record). - if (recorded && endpoint.priceUsdc > 0) { - billing.deductCredit(keyRecord.userId, endpoint.priceUsdc).catch((err) => { - console.error('Background billing deduction failed:', err); - }); - } - } catch (err) { - recordUsageRecordFailure(); - logger.error('Background usage recording failed', { - requestId, - apiId: String(apiEntry.id), - error: err, - }); - } - })(); - }); - }); - - res.once('close', () => { - // 'close' fires after 'finish' on a normal response, or on its own - // when the socket is destroyed prematurely. Only treat it as an - // abort when 'finish' has NOT already fired. - if (!responseFinished) { - recordProxyPrematureAbort(); - } - }); - } } catch (error) { next(error); }