feat(cloud-agent-next): enforce container billing

.specs/gastown-usage-based-billing.md

@@ -2,9 +2,8 @@
 
 ## Role of This Document
 
-This is retained design history for the Gastown producer. The authoritative rollout and
-settlement design is `fd-plans/research/container-billing-charge-and-enforce.md`; where this
-document differs from that plan, the plan wins.
+This is retained design history for the Gastown producer. It is superseded for rollout and
+settlement decisions.
 
 ## Status
 

packages/container-usage/src/client.test.ts

@@ -12,6 +12,16 @@ function binding(overrides: Partial<ContainerUsageRpcMethods> = {}): ContainerUs
         dedup: false,
       },
     })),
+    recordStartV2: vi.fn(async input => ({
+      success: true as const,
+      ack: {
+        intervalId: `${input.instanceId}:${input.startEpochMs}`,
+        durable: 'pg' as const,
+        dedup: false,
+      },
+      billingMode: 'paid' as const,
+      remainingMicrodollars: 10_000_001,
+    })),
     recordHeartbeat: vi.fn(async input => ({
       intervalId: `${input.instanceId}:${input.startEpochMs}`,
       durable: 'pg' as const,
@@ -176,4 +186,31 @@ describe('ContainerUsageClient', () => {
     });
     expect(recordStart).toHaveBeenCalledOnce();
   });
+
+  it('returns versioned paid admission details without changing v1 start behavior', async () => {
+    const rpc = binding();
+    const client = new ContainerUsageClient(rpc, { service: 'cloud-agent-next-sandbox' });
+
+    await expect(client.recordStartV2({ ...context, startEpochMs: 123 })).resolves.toMatchObject({
+      success: true,
+      billingMode: 'paid',
+      remainingMicrodollars: 10_000_001,
+    });
+    expect(rpc.recordStartV2).toHaveBeenCalledWith(
+      expect.objectContaining({
+        service: 'cloud-agent-next-sandbox',
+        idempotencyKey: 'v1:cloud-agent-next-sandbox:instance-1:123:start',
+      })
+    );
+  });
+
+  it('reports unavailable versioned admission support explicitly', async () => {
+    const client = new ContainerUsageClient(binding({ recordStartV2: undefined }), {
+      service: 'cloud-agent-next-sandbox',
+    });
+
+    await expect(client.recordStartV2({ ...context, startEpochMs: 123 })).rejects.toThrow(
+      'does not support recordStartV2'
+    );
+  });
 });

packages/container-usage/src/client.ts

@@ -3,13 +3,15 @@ import {
   heartbeatAckSchema,
   recordAckSchema,
   recordStartResultSchema,
+  recordStartV2ResultSchema,
   startIdempotencyKey,
   stopIdempotencyKey,
   type ContainerUsageRpcMethods,
   type HeartbeatAck,
   type RecordAck,
   type RecordHeartbeatInput,
   type RecordStartInput,
+  type RecordStartV2Result,
   type RecordStartFailureCode,
   type RecordStopInput,
   type UsageContext,
@@ -90,11 +92,7 @@ export class ContainerUsageClient {
   }
 
   async recordStart(input: ClientRecordStartInput): Promise<RecordAck> {
-    const request = {
-      ...input,
-      service: this.service,
-      idempotencyKey: startIdempotencyKey(this.service, input.instanceId, input.startEpochMs),
-    } satisfies RecordStartInput;
+    const request = this.startRequest(input);
     const result = await this.withRetry(async () =>
       recordStartResultSchema.parse(await this.binding.recordStart(request))
     );
@@ -104,6 +102,19 @@ export class ContainerUsageClient {
     return result.ack;
   }
 
+  async recordStartV2(input: ClientRecordStartInput): Promise<RecordStartV2Result> {
+    // A deployed Worker RPC proxy may expose an absent method as callable and reject only
+    // when invoked. Keep this check for direct bindings; either failure mode remains fail-closed.
+    const recordStartV2 = this.binding.recordStartV2;
+    if (!recordStartV2) {
+      throw new Error('Container usage meter does not support recordStartV2');
+    }
+    const request = this.startRequest(input);
+    return await this.withRetry(async () =>
+      recordStartV2ResultSchema.parse(await recordStartV2.call(this.binding, request))
+    );
+  }
+
   async recordHeartbeat(input: ClientRecordHeartbeatInput): Promise<HeartbeatAck> {
     const request = {
       ...input,
@@ -148,6 +159,14 @@ export class ContainerUsageClient {
     }
     throw lastError;
   }
+
+  private startRequest(input: ClientRecordStartInput): RecordStartInput {
+    return {
+      ...input,
+      service: this.service,
+      idempotencyKey: startIdempotencyKey(this.service, input.instanceId, input.startEpochMs),
+    };
+  }
 }
 
 export function createContainerUsageClient(

packages/container-usage/src/contracts.test.ts

@@ -2,6 +2,8 @@ import { describe, expect, it } from 'vitest';
 import {
   intervalId,
   recordHeartbeatInputSchema,
+  recordStartResultSchema,
+  recordStartV2ResultSchema,
   recordStopInputSchema,
   usageContextSchema,
 } from './contracts';
@@ -121,4 +123,30 @@ describe('container usage contracts', () => {
       intervalId('gastown', 'shared-instance', 123)
     );
   });
+
+  it('keeps v1 starts stable while v2 reports paid admission details', () => {
+    const v1 = {
+      success: true,
+      ack: { intervalId: 'interval-1', durable: 'pg', dedup: false },
+    };
+    expect(recordStartResultSchema.parse(v1)).toEqual(v1);
+    expect(
+      recordStartV2ResultSchema.parse({
+        ...v1,
+        billingMode: 'paid',
+        remainingMicrodollars: 10_000_001,
+      })
+    ).toMatchObject({ success: true, billingMode: 'paid', remainingMicrodollars: 10_000_001 });
+    expect(
+      recordStartV2ResultSchema.parse({
+        success: false,
+        error: {
+          code: 'insufficient_credits',
+          message: 'Insufficient credits',
+          remainingMicrodollars: 5_000_000,
+          minimumRequiredMicrodollars: 5_000_000,
+        },
+      })
+    ).toMatchObject({ success: false });
+  });
 });

packages/container-usage/src/contracts.ts

@@ -148,6 +148,49 @@ export const recordStartResultSchema = z.discriminatedUnion('success', [
 ]);
 export type RecordStartResult = z.infer<typeof recordStartResultSchema>;
 
+export const recordStartV2ResultSchema = z.union([
+  z
+    .object({
+      success: z.literal(true),
+      ack: recordAckSchema,
+      billingMode: z.literal('shadow'),
+    })
+    .strict(),
+  z
+    .object({
+      success: z.literal(true),
+      ack: recordAckSchema,
+      billingMode: z.literal('paid'),
+      remainingMicrodollars: z.number().int(),
+    })
+    .strict(),
+  z
+    .object({
+      success: z.literal(false),
+      error: z
+        .object({
+          code: z.literal('insufficient_credits'),
+          message: z.string().min(1),
+          remainingMicrodollars: z.number().int(),
+          minimumRequiredMicrodollars: z.number().int().positive(),
+        })
+        .strict(),
+    })
+    .strict(),
+  z
+    .object({
+      success: z.literal(false),
+      error: z
+        .object({
+          code: z.enum(['sku_not_found', 'sku_unit_mismatch', 'sku_not_accepting_new_usage']),
+          message: z.string().min(1),
+        })
+        .strict(),
+    })
+    .strict(),
+]);
+export type RecordStartV2Result = z.infer<typeof recordStartV2ResultSchema>;
+
 export const budgetVerdictSchema = z.discriminatedUnion('verdict', [
   z.object({ verdict: z.literal('continue'), remaining: z.number().int().optional() }).strict(),
   z
@@ -180,6 +223,7 @@ export type HeartbeatAck = z.infer<typeof heartbeatAckSchema>;
 
 export type ContainerUsageRpcMethods = {
   recordStart: (input: RecordStartInput) => Promise<RecordStartResult>;
+  recordStartV2?: (input: RecordStartInput) => Promise<RecordStartV2Result>;
   recordHeartbeat: (input: RecordHeartbeatInput) => Promise<HeartbeatAck>;
   recordStop: (input: RecordStopInput) => Promise<RecordAck>;
 };

packages/container-usage/src/heartbeat.test.ts

@@ -152,6 +152,89 @@ describe('installBillingHeartbeat', () => {
     expect(schedule).toHaveBeenCalledWith(300, BILLING_HEARTBEAT_CALLBACK, expect.any(String));
   });
 
+  it('keeps the default controller behavior of settling immediately after a budget stop', async () => {
+    const storage = memoryStorage();
+    await storedContext(storage);
+    const recordStop = vi.fn<ContainerUsageRpcMethods['recordStop']>(async input => ({
+      intervalId: `${input.instanceId}:${input.startEpochMs}`,
+      durable: 'pg',
+      dedup: false,
+    }));
+    const client = new ContainerUsageClient(
+      {
+        recordStart: async () => ({
+          success: true,
+          ack: { intervalId: 'instance-1:123', durable: 'pg', dedup: false },
+        }),
+        recordHeartbeat: async () => ({
+          intervalId: 'instance-1:123',
+          durable: 'pg',
+          dedup: false,
+          budget: { verdict: 'stop' },
+        }),
+        recordStop,
+      },
+      { service: 'gastown' }
+    );
+    const controller = installBillingHeartbeat(
+      {
+        deleteSchedules: vi.fn(),
+        getState: vi.fn(async () => ({ status: 'running' as const, lastChange: Date.now() })),
+        schedule: vi.fn() as Container['schedule'],
+      },
+      { client, storage, enforceBudgetStop: vi.fn() }
+    );
+
+    await controller.billingHeartbeatTick();
+
+    expect(recordStop).toHaveBeenCalledOnce();
+    expect(await getBillingContext(storage)).toBeUndefined();
+  });
+
+  it('defers budget-stop settlement when the producer owns physical shutdown', async () => {
+    const storage = memoryStorage();
+    await storedContext(storage);
+    const recordStop = vi.fn<ContainerUsageRpcMethods['recordStop']>(async input => ({
+      intervalId: `${input.instanceId}:${input.startEpochMs}`,
+      durable: 'pg',
+      dedup: false,
+    }));
+    const client = new ContainerUsageClient(
+      {
+        recordStart: async () => ({
+          success: true,
+          ack: { intervalId: 'instance-1:123', durable: 'pg', dedup: false },
+        }),
+        recordHeartbeat: async () => ({
+          intervalId: 'instance-1:123',
+          durable: 'pg',
+          dedup: false,
+          budget: { verdict: 'stop' },
+        }),
+        recordStop,
+      },
+      { service: 'cloud-agent-next' }
+    );
+    const controller = installBillingHeartbeat(
+      {
+        deleteSchedules: vi.fn(),
+        getState: vi.fn(async () => ({ status: 'running' as const, lastChange: Date.now() })),
+        schedule: vi.fn() as Container['schedule'],
+      },
+      {
+        client,
+        storage,
+        deferBudgetStopFinalSettlement: true,
+        enforceBudgetStop: vi.fn(),
+      }
+    );
+
+    await controller.billingHeartbeatTick();
+
+    expect(recordStop).not.toHaveBeenCalled();
+    expect(await getBillingContext(storage)).toBeDefined();
+  });
+
   it('defers stopped-state closure when the producer owns an authoritative stop hook', async () => {
     const storage = memoryStorage();
     await storedContext(storage);

packages/container-usage/src/heartbeat.ts

@@ -24,6 +24,11 @@ export type BillingHeartbeatDependencies = {
   heartbeatSeconds?: number;
   /** Defer stopped-state closure to the container's authoritative onStop hook. */
   stopOnStoppedState?: boolean;
+  /**
+   * Let a producer that owns physical shutdown settle from its authoritative
+   * onStop hook instead of closing usage at a budget verdict.
+   */
+  deferBudgetStopFinalSettlement?: boolean;
   stoppedStateGraceSeconds?: number;
   stoppedStateAbandonSeconds?: number;
   beforeHeartbeatDelivery?: (context: BillingContext) => Promise<void>;
@@ -379,7 +384,9 @@ export function installBillingHeartbeat(
           generation: context.generation,
           startEpochMs: context.startEpochMs,
         });
-        await recordStopForGeneration({ reason: 'runtime_signal' }, context.generation);
+        if (!dependencies.deferBudgetStopFinalSettlement) {
+          await recordStopForGeneration({ reason: 'runtime_signal' }, context.generation);
+        }
         return;
       }
       if (ack.budget.verdict === 'warn') {

services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts

@@ -437,6 +437,38 @@ describe('CloudflareAgentSandbox', () => {
     ensureBootstrapWrapper.mockRestore();
   });
 
+  it('requests paid admission for an organization canary before sandbox work', async () => {
+    const ensureBillingAdmission = vi.fn().mockResolvedValue({
+      success: false,
+      code: 'insufficient_credits',
+      message: 'Low balance',
+    });
+    const sandbox = new CloudflareAgentSandbox(
+      {
+        CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true',
+        CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: '',
+        CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: 'org_cloudflare',
+      } as Env,
+      metadata({ sandboxId: 'usr-shared' }),
+      {
+        resolveSandbox: () => ({}) as SandboxInstance,
+        ensureBillingAdmission,
+        isBillingBlocked: vi.fn().mockResolvedValue(false),
+      }
+    );
+
+    await expect(sandbox.ensureBillingAdmission()).resolves.toMatchObject({
+      success: false,
+      code: 'insufficient_credits',
+    });
+    expect(ensureBillingAdmission).toHaveBeenCalledWith(expect.anything(), {
+      sandboxId: 'usr-shared',
+      subject: { type: 'org', id: 'org_cloudflare' },
+      actor: { type: 'user', id: 'user_cloudflare' },
+      enforcementRequested: true,
+    });
+  });
+
   it('reports malformed worker URLs before degrading to a cold bootstrap', async () => {
     const request = ensureRequest({ cacheEligible: true });
     const bucket = { get: vi.fn(), put: vi.fn() };
@@ -1720,7 +1752,7 @@ describe('CloudflareAgentSandbox', () => {
     expect(listProcesses).toHaveBeenCalled();
   });
 
-  it('inspects a stopped container for stop reasons other than idle-timeout', async () => {
+  it('does not wake a confirmed stopped container for session deletion cleanup', async () => {
     const listProcesses = vi.fn().mockResolvedValue([]);
     const isContainerRunning = vi.fn().mockResolvedValue(false);
     const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), {
@@ -1734,8 +1766,8 @@ describe('CloudflareAgentSandbox', () => {
         reason: 'session-delete',
       })
     ).resolves.toEqual({ status: 'absent' });
-    expect(listProcesses).toHaveBeenCalled();
-    expect(isContainerRunning).not.toHaveBeenCalled();
+    expect(isContainerRunning).toHaveBeenCalledOnce();
+    expect(listProcesses).not.toHaveBeenCalled();
   });
 
   it('falls back to inspection when the sandbox cannot report container state', async () => {