From 60da0ed5aab4a6c8018a51816cc3ca5e388d5b3e Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Sat, 24 Oct 2026 22:55:53 +0200 Subject: [PATCH 1/2] fix(cloud-agent): complete native container billing and recovery --- .../compute-balance-alert/route.test.ts | 62 + .../compute-balance-alert/route.ts | 33 + .../CloudAgentBillingError.test.ts | 14 +- .../CloudAgentBillingError.tsx | 13 +- .../cloud-agent-next/ComputePrice.tsx | 31 + .../cloud-agent-next/NewSessionPanel.tsx | 9 +- .../NewSessionSandboxSelect.tsx | 10 +- .../cloud-agent-next/SessionComputeRate.tsx | 31 + .../cloud-agent-next/SessionInfoDialog.tsx | 2 + .../usage-analytics/ComputeUsageSection.tsx | 149 + .../UsageAnalyticsDashboard.tsx | 15 + .../src/lib/cloud-agent/compute-billing.ts | 127 + apps/web/src/lib/cloud-agent/compute-price.ts | 11 + .../src/lib/kilo-pass-org/consumption.test.ts | 38 +- .../src/lib/spend-alerts/compute-backfill.ts | 28 + .../lib/spend-alerts/compute-sweep.test.ts | 228 + apps/web/src/lib/spend-alerts/sweep.ts | 46 +- apps/web/src/lib/spend-writer-audit.test.ts | 4 +- .../routers/compute-billing-router.test.ts | 191 + .../web/src/routers/compute-billing-router.ts | 42 + apps/web/src/routers/root-router.ts | 2 + .../scripts/backfill-compute-spend-alerts.ts | 15 + dev/local/env-sync/plan.test.ts | 22 + .../cloud-agent-sdk/src/normalizer.test.ts | 37 + packages/cloud-agent-sdk/src/normalizer.ts | 6 + packages/cloud-agent-sdk/src/schemas.ts | 20 +- .../src/session-manager.test.ts | 25 + .../cloud-agent-sdk/src/session-manager.ts | 9 + .../container-usage/src/heartbeat.test.ts | 3 +- packages/container-usage/src/heartbeat.ts | 6 +- .../src/migrations/0297_windy_changeling.sql | 1 + .../db/src/migrations/meta/0297_snapshot.json | 38741 ++++++++++++++++ packages/db/src/migrations/meta/_journal.json | 7 + packages/db/src/schema.ts | 1 + .../src/cloud-agent-billing.ts | 13 + .../session-ingest-contracts/src/index.ts | 1 + .../src/lib/kilo-pass-org/consumption.ts | 264 +- packages/web-shared/src/lib/microdollars.ts | 4 +- .../organizations/organization-usage.test.ts | 32 +- .../lib/organizations/organization-usage.ts | 210 +- packages/worker-utils/package.json | 7 +- .../worker-utils/src/cloud-agent-compute.ts | 32 + .../worker-utils/src/compute-member.test.ts | 25 + packages/worker-utils/src/compute-member.ts | 34 + packages/worker-utils/src/microdollars.ts | 3 + .../src/organization-consumption.ts | 258 + .../worker-utils/src/organization-usage.ts | 115 + pnpm-lock.yaml | 3 + .../cloud-agent-next/docs/control-plane.md | 49 +- .../cloud-agent-next/src/billing-failure.ts | 23 + .../src/container-usage-context.test.ts | 31 + .../src/container-usage-context.ts | 45 +- .../generated-migrations.test.ts | 17 + .../control-plane/sandbox/allocation.test.ts | 17 + .../src/control-plane/sandbox/allocation.ts | 14 +- .../drizzle/0003_damp_sharon_carter.sql | 1 + .../sandbox/drizzle/meta/0003_snapshot.json | 244 + .../sandbox/drizzle/meta/_journal.json | 7 + .../sandbox/drizzle/migrations.js | 2 + .../sandbox/native-billing-admission.test.ts | 71 + .../sandbox/native-billing-admission.ts | 52 + .../src/control-plane/sandbox/routes.ts | 35 +- .../src/control-plane/sandbox/sandbox-do.ts | 221 +- .../control-plane/sandbox/sqlite-schema.ts | 1 + .../session/drizzle/0002_nifty_shen.sql | 1 + .../session/drizzle/meta/0002_snapshot.json | 187 + .../session/drizzle/meta/_journal.json | 9 +- .../session/drizzle/migrations.js | 2 + .../src/control-plane/session/messages.ts | 10 +- .../src/control-plane/session/session-do.ts | 32 +- .../control-plane/session/sqlite-schema.ts | 1 + .../src/metered-billing-lifecycle.ts | 64 +- .../customer-billing-failure.test.ts | 21 + .../persistence/customer-billing-failure.ts | 16 + .../sandbox-containers/SandboxContainers.ts | 54 +- .../containers-billing.test.ts | 195 +- .../sandbox-containers/containers-billing.ts | 9 +- .../cloudflare-containers-provider.ts | 8 +- .../sandbox-control/provider-creation.test.ts | 61 +- .../src/sandbox-control/provider.ts | 6 +- .../src/session/message-settlement-outbox.ts | 2 + .../shared/cloud-agent-billing-mirror.test.ts | 15 + .../src/shared/cloud-agent-billing.ts | 13 + .../src/shared/control-plane-protocol.ts | 11 + .../src/shared/control-plane-timers.ts | 5 + services/cloud-agent-next/src/trpc-error.ts | 13 +- .../control-plane-session-v2.test.ts | 131 +- .../sandbox-control-v2-migrations.test.ts | 13 +- .../integration/sandbox-control-v2.test.ts | 297 +- .../test/sdk-provenance/diagnostic.test.ts | 32 +- .../sdk-provenance/reconciliation.test.ts | 18 +- .../container-usage-meter/.dev.vars.example | 3 + services/container-usage-meter/RUNBOOK.md | 172 + services/container-usage-meter/package.json | 3 +- .../src/balance-alert.ts | 24 + .../src/billing-config.test.ts | 133 +- .../src/billing-config.ts | 43 +- services/container-usage-meter/src/meter.ts | 5 + .../container-usage-meter/src/postgres.ts | 133 +- .../test/organization-consumption.test.ts | 535 + .../test/postgres.test.ts | 212 +- .../vitest.postgres.config.ts | 6 +- .../worker-configuration.d.ts | 6 +- services/container-usage-meter/wrangler.jsonc | 8 + 104 files changed, 43600 insertions(+), 719 deletions(-) create mode 100644 apps/web/src/app/api/internal/organizations/compute-balance-alert/route.test.ts create mode 100644 apps/web/src/app/api/internal/organizations/compute-balance-alert/route.ts create mode 100644 apps/web/src/components/cloud-agent-next/ComputePrice.tsx create mode 100644 apps/web/src/components/cloud-agent-next/SessionComputeRate.tsx create mode 100644 apps/web/src/components/usage-analytics/ComputeUsageSection.tsx create mode 100644 apps/web/src/lib/cloud-agent/compute-billing.ts create mode 100644 apps/web/src/lib/cloud-agent/compute-price.ts create mode 100644 apps/web/src/lib/spend-alerts/compute-backfill.ts create mode 100644 apps/web/src/lib/spend-alerts/compute-sweep.test.ts create mode 100644 apps/web/src/routers/compute-billing-router.test.ts create mode 100644 apps/web/src/routers/compute-billing-router.ts create mode 100644 apps/web/src/scripts/backfill-compute-spend-alerts.ts create mode 100644 packages/db/src/migrations/0297_windy_changeling.sql create mode 100644 packages/db/src/migrations/meta/0297_snapshot.json create mode 100644 packages/session-ingest-contracts/src/cloud-agent-billing.ts create mode 100644 packages/worker-utils/src/cloud-agent-compute.ts create mode 100644 packages/worker-utils/src/compute-member.test.ts create mode 100644 packages/worker-utils/src/compute-member.ts create mode 100644 packages/worker-utils/src/microdollars.ts create mode 100644 packages/worker-utils/src/organization-consumption.ts create mode 100644 packages/worker-utils/src/organization-usage.ts create mode 100644 services/cloud-agent-next/src/billing-failure.ts create mode 100644 services/cloud-agent-next/src/control-plane/generated-migrations.test.ts create mode 100644 services/cloud-agent-next/src/control-plane/sandbox/drizzle/0003_damp_sharon_carter.sql create mode 100644 services/cloud-agent-next/src/control-plane/sandbox/drizzle/meta/0003_snapshot.json create mode 100644 services/cloud-agent-next/src/control-plane/sandbox/native-billing-admission.test.ts create mode 100644 services/cloud-agent-next/src/control-plane/sandbox/native-billing-admission.ts create mode 100644 services/cloud-agent-next/src/control-plane/session/drizzle/0002_nifty_shen.sql create mode 100644 services/cloud-agent-next/src/control-plane/session/drizzle/meta/0002_snapshot.json create mode 100644 services/cloud-agent-next/src/persistence/customer-billing-failure.test.ts create mode 100644 services/cloud-agent-next/src/persistence/customer-billing-failure.ts create mode 100644 services/cloud-agent-next/src/shared/cloud-agent-billing-mirror.test.ts create mode 100644 services/cloud-agent-next/src/shared/cloud-agent-billing.ts create mode 100644 services/container-usage-meter/.dev.vars.example create mode 100644 services/container-usage-meter/RUNBOOK.md create mode 100644 services/container-usage-meter/src/balance-alert.ts create mode 100644 services/container-usage-meter/test/organization-consumption.test.ts diff --git a/apps/web/src/app/api/internal/organizations/compute-balance-alert/route.test.ts b/apps/web/src/app/api/internal/organizations/compute-balance-alert/route.test.ts new file mode 100644 index 00000000000..aa0ab630ac2 --- /dev/null +++ b/apps/web/src/app/api/internal/organizations/compute-balance-alert/route.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it } from '@jest/globals'; +import { after, NextRequest } from 'next/server'; +import { INTERNAL_API_SECRET } from '@kilocode/web-shared/lib/config.server'; +import { insertTestUser } from '@kilocode/web-shared/tests/helpers/user.helper'; +import { createTestOrganization } from '@kilocode/web-shared/tests/helpers/organization.helper'; +import { sendBalanceAlertEmail } from '@kilocode/web-shared/lib/email'; +import { dispatchLowBalancePush } from '@kilocode/web-shared/lib/notifications-worker-client'; +import { POST } from './route'; + +jest.mock('next/server', () => ({ ...jest.requireActual('next/server'), after: jest.fn() })); + +jest.mock('@kilocode/web-shared/lib/email', () => ({ + ...jest.requireActual('@kilocode/web-shared/lib/email'), + sendBalanceAlertEmail: jest.fn().mockResolvedValue(undefined), +})); +jest.mock('@kilocode/web-shared/lib/notifications-worker-client', () => ({ + dispatchLowBalancePush: jest.fn().mockResolvedValue(undefined), +})); + +function request(body: unknown, secret = INTERNAL_API_SECRET) { + return new NextRequest('http://localhost/api/internal/organizations/compute-balance-alert', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Internal-Secret': secret }, + body: JSON.stringify(body), + }); +} + +describe('internal organization compute balance alert', () => { + it('requires service authentication and a valid organization ID', async () => { + expect((await POST(request({ organizationId: crypto.randomUUID() }, 'wrong'))).status).toBe( + 401 + ); + expect((await POST(request({ organizationId: 'invalid' }))).status).toBe(400); + expect(sendBalanceAlertEmail).not.toHaveBeenCalled(); + }); + + it('uses the existing email and push delivery with configured recipients', async () => { + const user = await insertTestUser(); + const org = await createTestOrganization('Compute alert recipients', user.id, 1_000_000, { + minimum_balance: 18, + minimum_balance_alert_email: [user.google_user_email], + }); + expect((await POST(request({ organizationId: org.id }))).status).toBe(200); + expect(sendBalanceAlertEmail).not.toHaveBeenCalled(); + const task = jest.mocked(after).mock.calls[0]?.[0]; + expect(typeof task).toBe('function'); + if (typeof task !== 'function') throw new Error('Missing deferred alert'); + await task(); + expect(sendBalanceAlertEmail).toHaveBeenCalledWith({ + organizationId: org.id, + minimum_balance: 18, + to: [user.google_user_email], + }); + expect(dispatchLowBalancePush).toHaveBeenCalledWith( + expect.objectContaining({ + organizationId: org.id, + recipientUserIds: [user.id], + minimumBalanceUsd: 18, + }) + ); + }); +}); diff --git a/apps/web/src/app/api/internal/organizations/compute-balance-alert/route.ts b/apps/web/src/app/api/internal/organizations/compute-balance-alert/route.ts new file mode 100644 index 00000000000..11e6f6ca644 --- /dev/null +++ b/apps/web/src/app/api/internal/organizations/compute-balance-alert/route.ts @@ -0,0 +1,33 @@ +import { after, NextResponse, type NextRequest } from 'next/server'; +import { z } from 'zod'; +import { eq } from 'drizzle-orm'; +import { organizations } from '@kilocode/db/schema'; +import { db } from '@kilocode/web-shared/lib/drizzle'; +import { sendOrganizationLowBalanceAlert } from '@kilocode/web-shared/lib/organizations/organization-usage'; +import { toMicrodollars } from '@kilocode/web-shared/lib/microdollars'; +import { authorizeInternalApiRequest } from '@/lib/internal-api-auth'; + +const bodySchema = z.object({ organizationId: z.uuid() }).strict(); + +export async function POST(request: NextRequest) { + const unauthorized = authorizeInternalApiRequest(request); + if (unauthorized) return unauthorized; + const body = bodySchema.safeParse(await request.json().catch(() => null)); + if (!body.success) return NextResponse.json({ error: 'Invalid body' }, { status: 400 }); + const [organization] = await db + .select({ settings: organizations.settings }) + .from(organizations) + .where(eq(organizations.id, body.data.organizationId)) + .limit(1); + const minimum = organization?.settings?.minimum_balance; + if (minimum) { + after(() => + sendOrganizationLowBalanceAlert( + body.data.organizationId, + toMicrodollars(minimum), + organization.settings?.minimum_balance_alert_email ?? [] + ) + ); + } + return NextResponse.json({ ok: true }); +} diff --git a/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.test.ts b/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.test.ts index f2116788aba..0f0565176ef 100644 --- a/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.test.ts +++ b/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.test.ts @@ -39,15 +39,17 @@ describe('CloudAgentBillingError', () => { presentation: personal, }); expect(html).toContain('role="alert"'); - expect(html).toContain('Your account needs more credits to start Cloud Agent compute.'); - expect(html).toContain('Available: $1.25 · Need more than $2.00'); - expect(html).toContain('Your prompt did not start.'); + expect(html).toContain( + 'Your account needs more credits for Cloud Agent compute. Add credits, then send a new message.' + ); + expect(html).toContain('Available: $1.25 · Need at least $2.00'); + expect(html).not.toContain('Your prompt did not start.'); expect(html).toContain('Add credits'); }); it.each([ ['remainingMicrodollars', 'Available: $1.00'], - ['minimumRequiredMicrodollars', 'Need more than $2.00'], + ['minimumRequiredMicrodollars', 'Need at least $2.00'], ] as const)('renders %s independently', (key, text) => { const html = renderBillingError({ failure: { @@ -71,7 +73,7 @@ describe('CloudAgentBillingError', () => { ], [ 'BILLING_UNAVAILABLE', - 'Cloud Agent cannot verify compute billing right now. Your prompt has not started and you have not been charged.', + 'Cloud Agent cannot verify compute billing right now. Try again shortly.', ], ] as const)('renders exact %s core copy without an action', (code, copy) => { const html = renderBillingError({ @@ -100,6 +102,6 @@ describe('CloudAgentBillingError', () => { remainingMicrodollars: 5_000_000, minimumRequiredMicrodollars: 5_000_000, }) - ).toBe('Available: $5.00 · Need more than $5.00'); + ).toBe('Available: $5.00 · Need at least $5.00'); }); }); diff --git a/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.tsx b/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.tsx index d441294c1ff..b85afdc98fa 100644 --- a/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.tsx +++ b/apps/web/src/components/cloud-agent-next/CloudAgentBillingError.tsx @@ -30,7 +30,7 @@ export function billingBalanceCopy(failure: CustomerBillingFailure): string | nu const required = failure.minimumRequiredMicrodollars === undefined ? undefined - : `Need more than ${formatBillingMoney(failure.minimumRequiredMicrodollars)}`; + : `Need at least ${formatBillingMoney(failure.minimumRequiredMicrodollars)}`; return [available, required].filter(Boolean).join(' · ') || null; } @@ -41,10 +41,12 @@ export function CloudAgentBillingError({ failure, presentation }: Props) { const balanceCopy = billingBalanceCopy(failure); const content = failure.code === 'INSUFFICIENT_CREDITS' - ? `${payerName} needs more credits to start Cloud Agent compute.` + ? failure.payer.type === 'org' + ? `${payerName} cannot cover Cloud Agent compute. Add credits or ask an owner or admin to check your daily spending limit, then send a new message.` + : `${payerName} needs more credits for Cloud Agent compute. Add credits, then send a new message.` : failure.code === 'COMPUTE_STOPPING' ? 'Cloud Agent is saving and stopping compute. Your prompt has not started. Try again after shutdown completes.' - : 'Cloud Agent cannot verify compute billing right now. Your prompt has not started and you have not been charged.'; + : 'Cloud Agent cannot verify compute billing right now. Try again shortly.'; return (