diff --git a/server/entry-node.ts b/server/entry-node.ts index 554d14f2..fd13b208 100644 --- a/server/entry-node.ts +++ b/server/entry-node.ts @@ -7,6 +7,7 @@ import { createLibsqlPlatform } from './platform/libsql' import { createNodePlatform } from './platform/node' import { syncPendingCloudTrafficReports } from './services/cloud-traffic-metering' import { runLicensingRefresh } from './services/licensing-refresh-runner' +import { syncPendingRemoteDownloadUsageReports } from './services/remote-download-usage' const REFRESH_INTERVAL_MS = 6 * 60 * 60 * 1000 // 6 hours const TRAFFIC_SYNC_INTERVAL_MS = 10 * 60 * 1000 // 10 minutes @@ -40,4 +41,5 @@ setInterval(() => { console.log('traffic.sync.scheduler.started interval=10m') setInterval(() => { void syncPendingCloudTrafficReports({ db: platform.db, cloudBaseUrl }) + void syncPendingRemoteDownloadUsageReports({ db: platform.db, cloudBaseUrl }) }, TRAFFIC_SYNC_INTERVAL_MS) diff --git a/server/routes/cloud-traffic-metering.integration.test.ts b/server/routes/cloud-traffic-metering.integration.test.ts index 20c99256..e1d2247d 100644 --- a/server/routes/cloud-traffic-metering.integration.test.ts +++ b/server/routes/cloud-traffic-metering.integration.test.ts @@ -160,7 +160,7 @@ describe('object download cloud traffic reporting', () => { ]) }) - it('records a failed report without refunding local traffic when Cloud returns a mismatched event id', async () => { + it('records a reported report without refunding local traffic when Cloud returns its own usage event id', async () => { const { app, db } = await createTestApp() await seedTrafficBinding(db) vi.stubGlobal( @@ -181,7 +181,7 @@ describe('object download cloud traffic reporting', () => { sql`SELECT traffic_used AS trafficUsed FROM org_quotas WHERE org_id = ${orgId}`, ) expect(rows[0].trafficUsed).toBe(125) - await expect(db.select().from(cloudTrafficReports)).resolves.toMatchObject([{ status: 'failed' }]) + await expect(db.select().from(cloudTrafficReports)).resolves.toMatchObject([{ status: 'reported', error: null }]) }) it('keeps the pre-presign Cloud report and refunds local traffic when presign fails', async () => { diff --git a/server/routes/download-tasks.integration.test.ts b/server/routes/download-tasks.integration.test.ts index d315d6b3..284c2b89 100644 --- a/server/routes/download-tasks.integration.test.ts +++ b/server/routes/download-tasks.integration.test.ts @@ -1,12 +1,16 @@ import type { Downloader, DownloadTask } from '@shared/types' import { sql } from 'drizzle-orm' import { beforeEach, describe, expect, it, vi } from 'vitest' +import { remoteDownloadUsageReports } from '../db/schema' +import { createLicenseBinding } from '../licensing/license-state' +import type { Database } from '../platform/interface' import { S3Service } from '../services/s3.js' import { adminHeaders, authedHeaders, createTestApp } from '../test/setup.js' type DownloadTaskList = { items: DownloadTask[] } beforeEach(() => { + vi.unstubAllGlobals() vi.restoreAllMocks() vi.spyOn(S3Service.prototype, 'presignUpload').mockResolvedValue('https://presigned-upload.example.com') vi.spyOn(S3Service.prototype, 'createMultipartUpload').mockResolvedValue('upload-1') @@ -59,6 +63,15 @@ function transferProgress(input: { } } +function makeCloudResponse(body: unknown, status = 200): Response { + return { + ok: status >= 200 && status < 300, + status, + json: async () => body, + text: async () => JSON.stringify(body), + } as unknown as Response +} + async function insertStorage(db: Awaited>['db']) { const now = Date.now() await db.run(sql` @@ -75,6 +88,19 @@ async function insertStorage(db: Awaited>['db'] `) } +async function seedCloudBinding(db: Database) { + await createLicenseBinding(db, { + cloudBindingId: 'download-billing-binding', + cloudStoreId: 'store-download-billing', + instanceId: 'test-instance', + cloudAccountId: 'test-account', + refreshToken: 'test-refresh-token', + cachedCert: 'test-certificate', + cachedExpiresAt: Math.floor(Date.now() / 1000) + 3600, + lastRefreshAt: Math.floor(Date.now() / 1000), + }) +} + async function registerDownloaderThroughDeviceLogin( app: Awaited>['app'], name: string, @@ -596,6 +622,82 @@ describe('Download tasks API integration', () => { expect(task.status.progress.upload.bytes).toBe(10 * 1024 * 1024) }) + it('accepts Cloud usage event ids that differ from local remote download idempotency keys', async () => { + const { app, db } = await createTestApp({ + DOWNLOAD_TOKEN_SECRET: 'test-download-token-secret', + ZPAN_CLOUD_URL: 'https://cloud.example', + }) + await insertStorage(db) + await seedCloudBinding(db) + vi.stubGlobal( + 'fetch', + vi + .fn() + .mockResolvedValue( + makeCloudResponse({ data: { accepted: true, duplicate: false, eventId: 'different-event-id' } }), + ), + ) + + const createdDownloader = await registerDownloaderThroughDeviceLogin(app, 'billing-transient-downloader') + await db.run(sql` + UPDATE downloaders + SET remote_download_credit_billing_enabled = 1, + remote_download_credit_unit_bytes = ${5 * 1024 * 1024}, + remote_download_credit_per_unit = 1 + WHERE id = ${createdDownloader.downloader.id} + `) + const downloaderHeaders = { + Authorization: `Bearer ${createdDownloader.token}`, + 'Content-Type': 'application/json', + } + const heartbeatRes = await app.request('/api/downloader/heartbeat', { + method: 'POST', + headers: downloaderHeaders, + body: JSON.stringify({ ...heartbeat, currentTasks: 0 }), + }) + expect(heartbeatRes.status).toBe(200) + + const user = await authedHeaders(app, 'download-billing-transient-user@example.com') + const createTaskRes = await app.request('/api/download-tasks', { + method: 'POST', + headers: { ...user, 'Content-Type': 'application/json' }, + body: JSON.stringify({ + source: { type: 'http', uri: 'https://example.com/billing-transient.bin' }, + targetFolder: 'Remote Downloads', + }), + }) + expect(createTaskRes.status).toBe(201) + const task = (await createTaskRes.json()) as DownloadTask + + const patchRes = await app.request(`/api/download-tasks/${task.id}`, { + method: 'PATCH', + headers: downloaderHeaders, + body: JSON.stringify({ + status: 'downloading', + ...transferProgress({ + downloadBytes: 5 * 1024 * 1024, + totalBytes: 10 * 1024 * 1024, + downloadBps: 512_000, + }), + }), + }) + + expect(patchRes.status).toBe(200) + await expect(patchRes.json()).resolves.toMatchObject({ + status: { + state: 'downloading', + billing: { state: 'ok', chargedBytes: 5 * 1024 * 1024, chargedCredits: 1 }, + }, + }) + await expect(db.select().from(remoteDownloadUsageReports)).resolves.toMatchObject([ + { + eventId: `remote_download:${task.id}:1`, + status: 'reported', + error: null, + }, + ]) + }) + it('stores downloader runtime reports as snapshots while progress remains patchable', async () => { const { app, db } = await createTestApp({ DOWNLOAD_TOKEN_SECRET: 'test-download-token-secret' }) await insertStorage(db) diff --git a/server/routes/licensing.ts b/server/routes/licensing.ts index 4955d1ae..d0fbaae7 100644 --- a/server/routes/licensing.ts +++ b/server/routes/licensing.ts @@ -8,6 +8,7 @@ import { normalizeHost } from '../licensing/verify' import type { Env } from '../middleware/platform' import { syncPendingCloudTrafficReports } from '../services/cloud-traffic-metering' import { runLicensingRefresh } from '../services/licensing-refresh-runner' +import { syncPendingRemoteDownloadUsageReports } from '../services/remote-download-usage' function configuredPublicHost(c: Context): string | null { const value = c.get('platform').getEnv('ZPAN_PUBLIC_ORIGIN') ?? c.get('platform').getEnv('BETTER_AUTH_URL') @@ -63,9 +64,12 @@ const app = new Hono() const db = c.get('platform').db const cloudBaseUrl = c.get('platform').getEnv('ZPAN_CLOUD_URL') ?? ZPAN_CLOUD_URL_DEFAULT - const result = await syncPendingCloudTrafficReports({ db, cloudBaseUrl }) + const [traffic, remoteDownload] = await Promise.all([ + syncPendingCloudTrafficReports({ db, cloudBaseUrl }), + syncPendingRemoteDownloadUsageReports({ db, cloudBaseUrl }), + ]) - return c.json({ ok: true, ...result }) + return c.json({ ok: true, ...traffic, remoteDownload }) }) export default app diff --git a/server/scheduled-worker.test.ts b/server/scheduled-worker.test.ts index 121a8d79..e0f830e9 100644 --- a/server/scheduled-worker.test.ts +++ b/server/scheduled-worker.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { syncPendingCloudTrafficReports } from '../server/services/cloud-traffic-metering' import { runLicensingRefresh } from '../server/services/licensing-refresh-runner' +import { syncPendingRemoteDownloadUsageReports } from '../server/services/remote-download-usage' import { handleScheduled } from '../workers/scheduled' vi.mock('../server/platform/cloudflare', () => ({ @@ -17,16 +18,25 @@ vi.mock('../server/services/licensing-refresh-runner', () => ({ runLicensingRefresh: vi.fn(), })) +vi.mock('../server/services/remote-download-usage', () => ({ + syncPendingRemoteDownloadUsageReports: vi.fn(), +})) + describe('handleScheduled', () => { beforeEach(() => { vi.mocked(syncPendingCloudTrafficReports).mockReset() + vi.mocked(syncPendingRemoteDownloadUsageReports).mockReset() vi.mocked(runLicensingRefresh).mockReset() }) - it('syncs traffic reports on the traffic cron only', async () => { + it('syncs usage reports on the traffic cron only', async () => { await handleScheduled({ cron: '*/10 * * * *' }, { DB: {} as D1Database, ZPAN_CLOUD_URL: 'https://cloud.example' }) expect(syncPendingCloudTrafficReports).toHaveBeenCalledWith({ db: 'db', cloudBaseUrl: 'https://cloud.example' }) + expect(syncPendingRemoteDownloadUsageReports).toHaveBeenCalledWith({ + db: 'db', + cloudBaseUrl: 'https://cloud.example', + }) expect(runLicensingRefresh).not.toHaveBeenCalled() }) @@ -35,5 +45,6 @@ describe('handleScheduled', () => { expect(runLicensingRefresh).toHaveBeenCalledWith('db', 'https://cloud.example') expect(syncPendingCloudTrafficReports).not.toHaveBeenCalled() + expect(syncPendingRemoteDownloadUsageReports).not.toHaveBeenCalled() }) }) diff --git a/server/services/cloud-traffic-metering.test.ts b/server/services/cloud-traffic-metering.test.ts index ccdc30d0..129fe2a3 100644 --- a/server/services/cloud-traffic-metering.test.ts +++ b/server/services/cloud-traffic-metering.test.ts @@ -248,15 +248,15 @@ describe('cloud traffic metering', () => { ]) }) - it('marks reports failed when Cloud does not accept the same event id', async () => { - const { db, platform } = await createTestApp() + it('marks reports reported when Cloud returns its own usage event id', async () => { + const { db, platform } = await createTestApp({ ZPAN_CLOUD_URL: 'https://cloud.example' }) await seedTrafficBinding(db) vi.stubGlobal( 'fetch', vi.fn().mockResolvedValue(makeResponse({ data: { accepted: true, duplicate: false, eventId: 'other_evt' } })), ) - await reportTrafficEgress({ + const result = await reportTrafficEgress({ platform, orgId: 'org_1', bytes: 1024, @@ -265,15 +265,10 @@ describe('cloud traffic metering', () => { eventId: 'evt_mismatch', ...meteredStorage, }) - await expect(syncPendingCloudTrafficReports({ db, cloudBaseUrl: 'https://cloud.example' })).resolves.toEqual({ - attempted: 1, - reported: 0, - blocked: 0, - failed: 1, - }) + expect(result).toMatchObject({ status: 'reported', eventId: 'evt_mismatch', duplicate: false }) await expect(db.select().from(cloudTrafficReports)).resolves.toMatchObject([ - { eventId: 'evt_mismatch', status: 'failed', error: 'cloud_usage_report_rejected' }, + { eventId: 'evt_mismatch', status: 'reported', error: null }, ]) }) diff --git a/server/services/cloud-traffic-metering.ts b/server/services/cloud-traffic-metering.ts index bed9778c..2b3565b5 100644 --- a/server/services/cloud-traffic-metering.ts +++ b/server/services/cloud-traffic-metering.ts @@ -173,7 +173,7 @@ async function syncTrafficReport(params: { }, ) const response = usageResponseSchema.parse(data) - if (!response.accepted || response.eventId !== report.eventId) throw new Error('cloud_usage_report_rejected') + if (!response.accepted) throw new Error('cloud_usage_report_rejected') await updateTrafficReport(db, report.eventId, 'reported', null, now) return 'reported' } catch (error) { diff --git a/server/services/remote-download-usage.ts b/server/services/remote-download-usage.ts index 8d7eab87..cb007b10 100644 --- a/server/services/remote-download-usage.ts +++ b/server/services/remote-download-usage.ts @@ -1,9 +1,10 @@ -import { eq } from 'drizzle-orm' +import { asc, eq, inArray } from 'drizzle-orm' import { nanoid } from 'nanoid' +import { z } from 'zod' import { ZPAN_CLOUD_URL_DEFAULT } from '../../shared/constants' import { remoteDownloadUsageReports } from '../db/schema' import { loadActiveLicenseBinding } from '../licensing/license-state' -import type { Platform } from '../platform/interface' +import type { Database, Platform } from '../platform/interface' import { postBoundCloudJson } from './licensing-cloud' export class RemoteDownloadBillingBlockedError extends Error { @@ -13,6 +14,15 @@ export class RemoteDownloadBillingBlockedError extends Error { } } +const usageResponseSchema = z.object({ + accepted: z.boolean(), + duplicate: z.boolean().optional(), + eventId: z.string().min(1), +}) + +type RemoteDownloadUsageStatus = 'pending' | 'reported' | 'skipped_unbound' | 'blocked' | 'failed' +type RemoteDownloadUsageReport = typeof remoteDownloadUsageReports.$inferSelect + export async function reportRemoteDownloadUnit(params: { platform: Platform orgId: string @@ -22,15 +32,15 @@ export async function reportRemoteDownloadUnit(params: { unitBytes: number creditsPerUnit: number enabled: boolean -}): Promise { - if (!params.enabled) return +}): Promise<{ status: RemoteDownloadUsageStatus; eventId: string }> { + if (!params.enabled) return { status: 'reported', eventId: '' } const eventId = `remote_download:${params.taskId}:${params.unitIndex}` const existing = await params.platform.db .select() .from(remoteDownloadUsageReports) .where(eq(remoteDownloadUsageReports.eventId, eventId)) .limit(1) - if (existing[0]?.status === 'reported') return + if (existing[0]?.status === 'reported') return { status: 'reported', eventId } if (existing[0]?.status === 'blocked') throw new RemoteDownloadBillingBlockedError() const now = new Date() @@ -51,46 +61,100 @@ export async function reportRemoteDownloadUnit(params: { }) } - const binding = await loadActiveLicenseBinding(params.platform.db) + const status = await syncRemoteDownloadUsageReport({ + db: params.platform.db, + cloudBaseUrl: params.platform.getEnv('ZPAN_CLOUD_URL') ?? ZPAN_CLOUD_URL_DEFAULT, + report: (await loadRemoteDownloadUsageReport(params.platform.db, eventId))!, + now, + }) + if (status === 'blocked') throw new RemoteDownloadBillingBlockedError() + return { status, eventId } +} + +export async function syncPendingRemoteDownloadUsageReports(params: { + db: Database + cloudBaseUrl: string + limit?: number + now?: Date +}): Promise<{ attempted: number; reported: number; blocked: number; failed: number }> { + const { db, cloudBaseUrl, limit = 100, now = new Date() } = params + const binding = await loadActiveLicenseBinding(db) + if (!binding?.refreshToken || !binding.cloudStoreId) return { attempted: 0, reported: 0, blocked: 0, failed: 0 } + + const reports = await db + .select() + .from(remoteDownloadUsageReports) + .where(inArray(remoteDownloadUsageReports.status, ['pending', 'failed'])) + .orderBy(asc(remoteDownloadUsageReports.createdAt)) + .limit(limit) + + const result = { attempted: reports.length, reported: 0, blocked: 0, failed: 0 } + for (const report of reports) { + const status = await syncRemoteDownloadUsageReport({ db, cloudBaseUrl, report, now }) + result[status] += 1 + } + return result +} + +async function syncRemoteDownloadUsageReport(params: { + db: Database + cloudBaseUrl: string + report: RemoteDownloadUsageReport + now: Date +}): Promise<'reported' | 'blocked' | 'failed'> { + const { db, cloudBaseUrl, report, now } = params + const binding = await loadActiveLicenseBinding(db) if (!binding?.refreshToken || !binding.cloudStoreId) { - await mark(params.platform, eventId, 'reported', null) - return + await mark(db, report.eventId, 'skipped_unbound', null, now) + return 'reported' } try { - const response = (await postBoundCloudJson( - params.platform.getEnv('ZPAN_CLOUD_URL') ?? ZPAN_CLOUD_URL_DEFAULT, - `/api/stores/${encodeURIComponent(binding.cloudStoreId)}/billing/usage-events`, - binding.refreshToken, - { - resource: 'remote_download', - unit: 'byte', - bytes: params.unitBytes, - eventId, - idempotencyKey: eventId, - customerId: params.orgId, - source: 'remote_download', - sourceId: params.taskId, - usageContext: { downloaderId: params.downloaderId }, - pricing: { unitQuantity: params.unitBytes, creditsPerUnit: params.creditsPerUnit }, - }, - )) as { accepted?: boolean; eventId?: string; error?: { code?: string } } - if (!response.accepted || response.eventId !== eventId) throw new Error('cloud_usage_report_rejected') - await mark(params.platform, eventId, 'reported', null) + const response = usageResponseSchema.parse( + await postBoundCloudJson( + cloudBaseUrl, + `/api/stores/${encodeURIComponent(binding.cloudStoreId)}/billing/usage-events`, + binding.refreshToken, + { + resource: 'remote_download', + unit: 'byte', + bytes: report.unitBytes, + eventId: report.eventId, + idempotencyKey: report.eventId, + customerId: report.orgId, + source: 'remote_download', + sourceId: report.taskId, + usageContext: { downloaderId: report.downloaderId }, + pricing: { unitQuantity: report.unitBytes, creditsPerUnit: report.creditsPerUnit }, + }, + ), + ) + if (!response.accepted) throw new Error('cloud_usage_report_rejected') + await mark(db, report.eventId, 'reported', null, now) + return 'reported' } catch (error) { const message = error instanceof Error ? error.message : 'cloud_usage_report_failed' if (message === 'insufficient_credits' || message === 'overage_cap_exceeded') { - await mark(params.platform, eventId, 'blocked', message) - throw new RemoteDownloadBillingBlockedError() + await mark(db, report.eventId, 'blocked', message, now) + return 'blocked' } - await mark(params.platform, eventId, 'failed', message) - throw error + await mark(db, report.eventId, 'failed', message, now) + return 'failed' } } -async function mark(platform: Platform, eventId: string, status: string, error: string | null) { - await platform.db +async function loadRemoteDownloadUsageReport(db: Database, eventId: string) { + const rows = await db + .select() + .from(remoteDownloadUsageReports) + .where(eq(remoteDownloadUsageReports.eventId, eventId)) + .limit(1) + return rows[0] +} + +async function mark(db: Database, eventId: string, status: string, error: string | null, now: Date) { + await db .update(remoteDownloadUsageReports) - .set({ status, error, updatedAt: new Date() }) + .set({ status, error, updatedAt: now }) .where(eq(remoteDownloadUsageReports.eventId, eventId)) } diff --git a/workers/scheduled.ts b/workers/scheduled.ts index f867f9a3..f831dd08 100644 --- a/workers/scheduled.ts +++ b/workers/scheduled.ts @@ -3,6 +3,7 @@ import { createCloudflarePlatform } from '../server/platform/cloudflare' import { syncPendingCloudTrafficReports } from '../server/services/cloud-traffic-metering' import { runLicensingRefresh } from '../server/services/licensing-refresh-runner' +import { syncPendingRemoteDownloadUsageReports } from '../server/services/remote-download-usage' import { ZPAN_CLOUD_URL_DEFAULT } from '../shared/constants' // Subset of the worker Env used by the scheduled handler. @@ -21,6 +22,7 @@ export async function handleScheduled(event: ScheduledTrigger, env: ScheduledEnv const cloudBaseUrl = env.ZPAN_CLOUD_URL ?? ZPAN_CLOUD_URL_DEFAULT if (event.cron === TRAFFIC_SYNC_CRON) { await syncPendingCloudTrafficReports({ db: platform.db, cloudBaseUrl }) + await syncPendingRemoteDownloadUsageReports({ db: platform.db, cloudBaseUrl }) return }