From 559f81fd059873a4034f5f3c2796ee4fded14604 Mon Sep 17 00:00:00 2001 From: syn Date: Thu, 23 Jul 2026 13:27:01 -0500 Subject: [PATCH 01/11] feat(billing): record cloud agent container usage --- dev/local/services.ts | 1 + pnpm-lock.yaml | 3 + services/cloud-agent-next/package.json | 1 + .../cloudflare-agent-sandbox.test.ts | 13 ++ .../cloudflare/cloudflare-agent-sandbox.ts | 24 +- .../src/container-usage-context.test.ts | 133 +++++++++++ .../src/container-usage-context.ts | 69 ++++++ .../src/container-usage.test.ts | 218 ++++++++++++++++++ .../cloud-agent-next/src/container-usage.ts | 133 +++++++++++ .../src/kilo-facade/session-proxy.ts | 8 + services/cloud-agent-next/src/router.test.ts | 8 + .../src/sandbox-outbound.test.ts | 1 + .../cloud-agent-next/src/sandbox-outbound.ts | 19 +- services/cloud-agent-next/src/types.ts | 3 + .../worker-configuration.d.ts | 60 +++-- services/cloud-agent-next/wrangler.jsonc | 10 + 16 files changed, 664 insertions(+), 40 deletions(-) create mode 100644 services/cloud-agent-next/src/container-usage-context.test.ts create mode 100644 services/cloud-agent-next/src/container-usage-context.ts create mode 100644 services/cloud-agent-next/src/container-usage.test.ts create mode 100644 services/cloud-agent-next/src/container-usage.ts diff --git a/dev/local/services.ts b/dev/local/services.ts index c31618c1fd..6aaedd6ac4 100644 --- a/dev/local/services.ts +++ b/dev/local/services.ts @@ -100,6 +100,7 @@ const serviceMeta: Record = { 'nextjs', 'cloudflare-session-ingest', 'cloudflare-git-token-service', + 'container-usage-meter', 'notifications', ], dir: 'services/cloud-agent-next', diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index c8d3cbd7bf..d756fbf068 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1784,6 +1784,9 @@ importers: '@kilocode/cloud-agent-profile': specifier: workspace:* version: link:../../packages/cloud-agent-profile + '@kilocode/container-usage': + specifier: workspace:* + version: link:../../packages/container-usage '@kilocode/db': specifier: workspace:* version: link:../../packages/db diff --git a/services/cloud-agent-next/package.json b/services/cloud-agent-next/package.json index e5df62130d..6c4d1a0e1f 100644 --- a/services/cloud-agent-next/package.json +++ b/services/cloud-agent-next/package.json @@ -33,6 +33,7 @@ "@cloudflare/sandbox": "0.12.1", "@hono/trpc-server": "0.4.2", "@kilocode/cloud-agent-profile": "workspace:*", + "@kilocode/container-usage": "workspace:*", "@kilocode/db": "workspace:*", "@kilocode/encryption": "workspace:*", "@kilocode/notifications": "workspace:*", diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts index 19a11f7f3a..b14ee1ecae 100644 --- a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts @@ -345,6 +345,7 @@ describe('CloudflareAgentSandbox', () => { const setOutboundHandler = vi.fn().mockResolvedValue(undefined); const exec = vi.fn().mockResolvedValue({ exitCode: 0, stdout: 'exists\n', stderr: '' }); const createSession = vi.fn().mockResolvedValue(bootstrapSession); + const configureBilling = vi.fn().mockResolvedValue(undefined); const ensureBootstrapWrapper = vi .spyOn(WrapperClient, 'ensureBootstrapWrapper') .mockResolvedValueOnce({ client: {} as WrapperClient }); @@ -357,6 +358,7 @@ describe('CloudflareAgentSandbox', () => { const sandbox = new CloudflareAgentSandbox(env, sessionMetadata, { resolveSandbox: () => ({ setOutboundHandler, exec, createSession }) as unknown as SandboxInstance, + configureBilling, }); await sandbox.ensureWrapper( @@ -364,6 +366,17 @@ describe('CloudflareAgentSandbox', () => { ); expect(setOutboundHandler).toHaveBeenCalledWith('managedScm'); + expect(configureBilling).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + subject: { type: 'org', id: 'org_cloudflare' }, + actor: { type: 'user', id: 'user_cloudflare' }, + metadata: { allocation: 'shared' }, + }) + ); + expect(configureBilling.mock.invocationCallOrder[0]).toBeLessThan( + setOutboundHandler.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); expect(setOutboundHandler.mock.invocationCallOrder[0]).toBeLessThan( exec.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY ); diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts index 9f8b5322ac..577afa088b 100644 --- a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts @@ -71,6 +71,14 @@ import { WorkspaceFilesystemPreparationError, } from '../../workspace-errors.js'; import { TOOL_CGROUP_ENV_KEYS, type ToolCgroupEnv } from '../../shared/tool-cgroup-env.js'; +import { + buildSandboxBillingInput, + type SandboxBillingInput, +} from '../../container-usage-context.js'; + +type BillingSandboxInstance = SandboxInstance & { + configureBilling(input: SandboxBillingInput): Promise; +}; const PREPARE_WORKSPACE_TIMEOUT_MS = 10 * 60 * 1000; const DEFAULT_STOP_OBSERVATION_DELAYS_MS = [100, 500, 1_000]; @@ -155,6 +163,7 @@ function withWorkspacePreparationTimeout(operation: Promise, step: string) export type CloudflareAgentSandboxDependencies = { resolveSandbox?: (sandboxId: SandboxId, options?: { sleepAfter?: number }) => SandboxInstance; + configureBilling?: (sandbox: SandboxInstance, input: SandboxBillingInput) => Promise; sessionService?: SessionService; stopObservedWrappers?: typeof stopObservedWrappers; sleep?: (ms: number) => Promise; @@ -170,6 +179,10 @@ export class CloudflareAgentSandbox implements AgentSandbox { private readonly stopObserved: typeof stopObservedWrappers; private readonly sleep: (ms: number) => Promise; private readonly stopObservationDelaysMs: number[]; + private readonly configureBilling: ( + sandbox: SandboxInstance, + input: SandboxBillingInput + ) => Promise; private sandboxIdPromise?: Promise; constructor( @@ -192,6 +205,12 @@ export class CloudflareAgentSandbox implements AgentSandbox { this.sleep = dependencies.sleep ?? (ms => new Promise(resolve => setTimeout(resolve, ms))); this.stopObservationDelaysMs = dependencies.stopObservationDelaysMs ?? DEFAULT_STOP_OBSERVATION_DELAYS_MS; + this.configureBilling = + dependencies.configureBilling ?? + (dependencies.resolveSandbox + ? async () => undefined + : async (sandbox, input) => + await (sandbox as BillingSandboxInstance).configureBilling(input)); } private resolveSandboxId(): Promise { @@ -213,7 +232,10 @@ export class CloudflareAgentSandbox implements AgentSandbox { } private async getSandbox(options?: { sleepAfter?: number }): Promise { - return this.resolveSandbox(await this.resolveSandboxId(), options); + const sandboxId = await this.resolveSandboxId(); + const sandbox = this.resolveSandbox(sandboxId, options); + await this.configureBilling(sandbox, buildSandboxBillingInput(this.metadata, sandboxId)); + return sandbox; } private async workspaceHasGit(sandbox: SandboxInstance, workspacePath: string): Promise { diff --git a/services/cloud-agent-next/src/container-usage-context.test.ts b/services/cloud-agent-next/src/container-usage-context.test.ts new file mode 100644 index 0000000000..c05239cfba --- /dev/null +++ b/services/cloud-agent-next/src/container-usage-context.test.ts @@ -0,0 +1,133 @@ +import { describe, expect, it } from 'vitest'; +import type { SessionMetadata } from './persistence/session-metadata.js'; +import { buildSandboxBillingInput, SANDBOX_USAGE_SKUS } from './container-usage-context.js'; + +function metadata(identity: SessionMetadata['identity']): SessionMetadata { + return { + metadataSchemaVersion: 2, + identity, + auth: {}, + repository: { type: 'github', repo: 'Kilo-Org/cloud' }, + lifecycle: { version: 1, timestamp: 1 }, + }; +} + +describe('container usage context', () => { + it('maps every concrete sandbox class to its immutable SKU', () => { + expect(SANDBOX_USAGE_SKUS).toEqual({ + Sandbox: 'cloud-agent-standard-2026-07', + SandboxContainment: 'cloud-agent-standard-2026-07', + SandboxSmall: 'cloud-agent-small-2026-07', + SandboxSmallContainment: 'cloud-agent-small-2026-07', + SandboxDIND: 'cloud-agent-dind-2026-07', + SandboxCodeReview: 'cloud-agent-code-review-2026-07', + SandboxCodeReviewContainment: 'cloud-agent-code-review-2026-07', + }); + }); + + it.each([ + { + name: 'personal human', + identity: { sessionId: 'agent_personal', userId: 'user_personal' }, + expected: { + subject: { type: 'user', id: 'user_personal' }, + actor: { type: 'user', id: 'user_personal' }, + }, + }, + { + name: 'organization human', + identity: { sessionId: 'agent_org', userId: 'user_org', orgId: 'org_1' }, + expected: { + subject: { type: 'org', id: 'org_1' }, + actor: { type: 'user', id: 'user_org' }, + }, + }, + { + name: 'personal bot', + identity: { sessionId: 'agent_bot', userId: 'user_bot', botId: 'bot_1' }, + expected: { + subject: { type: 'user', id: 'user_bot' }, + actor: { type: 'bot', id: 'bot_1' }, + onBehalfOf: { type: 'user', id: 'user_bot' }, + }, + }, + { + name: 'organization bot', + identity: { + sessionId: 'agent_org_bot', + userId: 'user_org_bot', + orgId: 'org_2', + botId: 'bot_2', + }, + expected: { + subject: { type: 'org', id: 'org_2' }, + actor: { type: 'bot', id: 'bot_2' }, + onBehalfOf: { type: 'org', id: 'org_2' }, + }, + }, + ])('derives trusted $name attribution', ({ identity, expected }) => { + expect(buildSandboxBillingInput(metadata(identity), 'ses-isolated')).toMatchObject(expected); + }); + + it('keeps isolated metadata bounded and normalizes automation origins', () => { + const input = buildSandboxBillingInput( + metadata({ + sessionId: 'agent_security', + userId: 'user_security', + orgId: 'org_security', + createdOnPlatform: 'security-remediation', + }), + 'crv-isolated' + ); + + expect(input).toMatchObject({ + sessionId: 'agent_security', + metadata: { + allocation: 'isolated', + origin: 'security-remediation', + repository_provider: 'github', + }, + }); + expect(JSON.stringify(input)).not.toContain('Kilo-Org/cloud'); + }); + + it('omits session, origin, and repository metadata for shared containers', () => { + const first = buildSandboxBillingInput( + metadata({ + sessionId: 'agent_first', + userId: 'user_shared', + orgId: 'org_shared', + createdOnPlatform: 'security-agent', + }), + 'org-shared' + ); + const second = buildSandboxBillingInput( + metadata({ + sessionId: 'agent_second', + userId: 'user_shared', + orgId: 'org_shared', + createdOnPlatform: 'cloud-agent-web', + }), + 'org-shared' + ); + + expect(first).toEqual(second); + expect(first).toEqual({ + subject: { type: 'org', id: 'org_shared' }, + actor: { type: 'user', id: 'user_shared' }, + metadata: { allocation: 'shared' }, + }); + }); + + it('maps unknown caller-provided origins to other', () => { + const input = buildSandboxBillingInput( + metadata({ + sessionId: 'agent_unknown', + userId: 'user_unknown', + createdOnPlatform: 'attacker-controlled-value', + }), + 'dind-isolated' + ); + expect(input.metadata?.origin).toBe('other'); + }); +}); diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts new file mode 100644 index 0000000000..db8855ee85 --- /dev/null +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -0,0 +1,69 @@ +import type { UsageContext } from '@kilocode/container-usage'; +import type { SessionMetadata } from './persistence/session-metadata.js'; +import type { SandboxId } from './types.js'; + +export const SANDBOX_USAGE_SKUS = { + Sandbox: 'cloud-agent-standard-2026-07', + SandboxContainment: 'cloud-agent-standard-2026-07', + SandboxSmall: 'cloud-agent-small-2026-07', + SandboxSmallContainment: 'cloud-agent-small-2026-07', + SandboxDIND: 'cloud-agent-dind-2026-07', + SandboxCodeReview: 'cloud-agent-code-review-2026-07', + SandboxCodeReviewContainment: 'cloud-agent-code-review-2026-07', +} as const; + +export type SandboxClassName = keyof typeof SANDBOX_USAGE_SKUS; +export type SandboxBillingInput = Omit; + +const KNOWN_ORIGINS = new Set([ + 'app-builder', + 'auto-triage', + 'autofix', + 'cloud-agent', + 'cloud-agent-web', + 'code-review', + 'discord', + 'github', + 'linear', + 'scheduled', + 'security-agent', + 'security-remediation', + 'slack', + 'webhook', +]); + +function normalizedOrigin(origin: string | undefined): string { + if (!origin) return 'cloud-agent'; + return KNOWN_ORIGINS.has(origin) ? origin : 'other'; +} + +function isIsolatedSandbox(sandboxId: SandboxId): boolean { + return /^(crv|dind|ses)-/.test(sandboxId); +} + +export function buildSandboxBillingInput( + metadata: SessionMetadata, + sandboxId: SandboxId +): SandboxBillingInput { + const subject = metadata.identity.orgId + ? { type: 'org' as const, id: metadata.identity.orgId } + : { type: 'user' as const, id: metadata.identity.userId }; + const actor = metadata.identity.botId + ? { type: 'bot' as const, id: metadata.identity.botId } + : { type: 'user' as const, id: metadata.identity.userId }; + const isolated = isIsolatedSandbox(sandboxId); + + return { + subject, + actor, + ...(actor.type === 'bot' ? { onBehalfOf: subject } : {}), + ...(isolated ? { sessionId: metadata.identity.sessionId } : {}), + metadata: isolated + ? { + allocation: 'isolated', + origin: normalizedOrigin(metadata.identity.createdOnPlatform), + ...(metadata.repository ? { repository_provider: metadata.repository.type } : {}), + } + : { allocation: 'shared' }, + }; +} diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts new file mode 100644 index 0000000000..ddfbc0344b --- /dev/null +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -0,0 +1,218 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { ContainerUsageRpcMethods } from '@kilocode/container-usage'; + +// oxlint-disable-next-line no-empty-object-type -- Matches the mocked Sandbox constructor. +type SandboxDurableObjectState = DurableObjectState<{}>; + +const sdk = vi.hoisted(() => { + class StockSandbox { + ctx: SandboxDurableObjectState; + env: unknown; + mockState: { status: string; exitCode?: number } = { status: 'stopped' }; + schedules: Array<{ when: number; callback: string; payload: unknown }> = []; + superStarted = false; + superStopped = false; + superActivityExpired = false; + superStopCalled = false; + + constructor(ctx: SandboxDurableObjectState, env: unknown) { + this.ctx = ctx; + this.env = env; + } + + deleteSchedules(callback: string): void { + this.schedules = this.schedules.filter(schedule => schedule.callback !== callback); + } + + async schedule(when: number, callback: string, payload?: unknown): Promise { + this.schedules.push({ when, callback, payload }); + return {}; + } + + async getState() { + return this.mockState; + } + + async onStart(): Promise { + this.superStarted = true; + } + + async onStop(): Promise { + this.superStopped = true; + } + + async onActivityExpired(): Promise { + this.superActivityExpired = true; + } + + async stop(): Promise { + this.superStopCalled = true; + } + } + return { StockSandbox }; +}); + +vi.mock('@cloudflare/sandbox', () => ({ Sandbox: sdk.StockSandbox })); + +import { getBillingContext } from '@kilocode/container-usage'; +import { MeteredSandbox } from './container-usage.js'; + +class MemoryStorage { + private readonly values = new Map(); + + async get(key: string): Promise { + return this.values.get(key) as T | undefined; + } + + async put(key: string, value: unknown): Promise { + this.values.set(key, value); + } + + async delete(key: string): Promise { + return this.values.delete(key); + } +} + +function ack(intervalId = 'interval-1') { + return { intervalId, durable: 'pg' as const, dedup: false }; +} + +function createRpc(): ContainerUsageRpcMethods { + return { + recordStart: vi.fn(async () => ({ + success: true, + ack: ack(), + })), + recordHeartbeat: vi.fn(async () => ({ + ...ack(), + budget: { verdict: 'continue' }, + })), + recordStop: vi.fn(async () => ack()), + }; +} + +type TestRuntime = MeteredSandbox & { + mockState: { status: string; exitCode?: number }; + schedules: Array<{ when: number; callback: string; payload: unknown }>; + superStarted: boolean; + superStopped: boolean; + superActivityExpired: boolean; + superStopCalled: boolean; +}; + +function createSandbox(rpc = createRpc()) { + const storage = new MemoryStorage(); + const ctx = { + id: { toString: () => 'do-id' }, + storage, + } as unknown as SandboxDurableObjectState; + class TestSandbox extends MeteredSandbox { + protected readonly sandboxClassName = 'SandboxSmallContainment' as const; + } + return { + rpc, + storage, + sandbox: new TestSandbox(ctx, { + CONTAINER_USAGE_METER: rpc, + } as never) as unknown as TestRuntime, + }; +} + +const billingInput = { + subject: { type: 'org' as const, id: 'org_1' }, + actor: { type: 'user' as const, id: 'user_1' }, + sessionId: 'agent_1', + metadata: { allocation: 'isolated' }, +}; + +describe('MeteredSandbox', () => { + beforeEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it('admits duplicate acquisition with one generation-stable start', async () => { + const { rpc, storage, sandbox } = createSandbox(); + vi.spyOn(Date, 'now').mockReturnValue(1_000); + + await sandbox.configureBilling(billingInput); + await sandbox.configureBilling(billingInput); + + expect(rpc.recordStart).toHaveBeenCalledTimes(2); + const starts = vi.mocked(rpc.recordStart).mock.calls.map(([input]) => input); + expect(starts[0]?.startEpochMs).toBe(1_000); + expect(starts[1]?.startEpochMs).toBe(1_000); + expect(starts[0]?.idempotencyKey).toBe(starts[1]?.idempotencyKey); + expect(starts[0]).toMatchObject({ + instanceId: 'SandboxSmallContainment:do-id', + sku: 'cloud-agent-small-2026-07', + metadata: { allocation: 'isolated', container_class: 'SandboxSmallContainment' }, + }); + expect((await getBillingContext(storage))?.measurementStarted).toBe(false); + }); + + it('starts five-minute heartbeat measurement after preserving the SDK start hook', async () => { + const { sandbox, storage } = createSandbox(); + vi.spyOn(Date, 'now').mockReturnValue(2_000); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + + await sandbox.onStart(); + + expect(sandbox.superStarted).toBe(true); + expect(sandbox.schedules).toEqual([ + expect.objectContaining({ when: 300, callback: 'billingHeartbeatTick' }), + ]); + expect((await getBillingContext(storage))?.measurementStarted).toBe(true); + }); + + it('closes a stopped generation before allocating a monotonic replacement', async () => { + const { rpc, sandbox } = createSandbox(); + vi.spyOn(Date, 'now').mockReturnValue(3_000); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + sandbox.mockState = { status: 'stopped_with_code', exitCode: 17 }; + + await sandbox.configureBilling(billingInput); + + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ reason: 'runtime_signal', exitCode: 17, startEpochMs: 3_000 }) + ); + expect(vi.mocked(rpc.recordStart).mock.calls.at(-1)?.[0].startEpochMs).toBe(3_001); + }); + + it('persists and retries the same durable stop intent while preserving SDK cleanup', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStop) + .mockRejectedValueOnce(new Error('postgres unavailable')) + .mockRejectedValueOnce(new Error('postgres unavailable')) + .mockRejectedValueOnce(new Error('postgres unavailable')) + .mockResolvedValue(ack()); + const { sandbox, storage } = createSandbox(rpc); + vi.spyOn(Date, 'now').mockReturnValue(4_000); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + + await expect(sandbox.onActivityExpired()).rejects.toThrow('postgres unavailable'); + expect(sandbox.superActivityExpired).toBe(true); + expect((await getBillingContext(storage))?.pendingStop?.reason).toBe('activity_expired'); + + await sandbox.onStop(); + expect(sandbox.superStopped).toBe(true); + expect(await getBillingContext(storage)).toBeUndefined(); + const stops = vi.mocked(rpc.recordStop).mock.calls.map(([input]) => input); + expect(stops.at(-1)?.reason).toBe('activity_expired'); + expect(new Set(stops.map(stop => stop.idempotencyKey))).toHaveLength(1); + }); + + it('stops an unadmitted runtime instead of silently running it', async () => { + const { sandbox } = createSandbox(); + await expect(sandbox.onStart()).rejects.toThrow( + 'Container started without an admitted billing context' + ); + expect(sandbox.superStarted).toBe(true); + expect(sandbox.superStopCalled).toBe(true); + }); +}); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts new file mode 100644 index 0000000000..14ae768362 --- /dev/null +++ b/services/cloud-agent-next/src/container-usage.ts @@ -0,0 +1,133 @@ +import { + createContainerUsageClient, + getBillingContext, + installBillingHeartbeat, + setBillingContext, + usageContextFromBillingContext, + type BillingContext, + type BillingHeartbeatController, + type ClientRecordStartInput, + type ContainerUsageClient, + type UsageContext, +} from '@kilocode/container-usage'; +import { Sandbox as StockSandbox } from '@cloudflare/sandbox'; +import type { Env } from './types.js'; +import { + SANDBOX_USAGE_SKUS, + type SandboxBillingInput, + type SandboxClassName, +} from './container-usage-context.js'; + +const SERVICE = 'cloud-agent-next'; +const LAST_START_EPOCH_STORAGE_KEY = 'container-usage:last-start-epoch:v1'; +// oxlint-disable-next-line no-empty-object-type -- Matches the Sandbox 0.12.1 constructor. +type SandboxDurableObjectState = DurableObjectState<{}>; + +function startInputFromContext(context: BillingContext): ClientRecordStartInput { + const { service: _service, ...usage } = usageContextFromBillingContext(context); + return { ...usage, startEpochMs: context.startEpochMs }; +} + +export abstract class MeteredSandbox extends StockSandbox { + protected abstract readonly sandboxClassName: SandboxClassName; + + private readonly usageClient: ContainerUsageClient; + private readonly billingHeartbeat: BillingHeartbeatController; + private billingAdmissionTail: Promise = Promise.resolve(); + + constructor(ctx: SandboxDurableObjectState, env: Env) { + super(ctx, env); + this.usageClient = createContainerUsageClient(env.CONTAINER_USAGE_METER, { + service: SERVICE, + }); + this.billingHeartbeat = installBillingHeartbeat(this, { + client: this.usageClient, + storage: this.ctx.storage, + // Shadow metering must never stop customer work. + enforceBudgetStop: async () => { + throw new Error('Container budget enforcement is disabled in shadow mode'); + }, + }); + } + + async configureBilling(input: SandboxBillingInput): Promise { + const operation = this.billingAdmissionTail.then( + () => this.configureBillingExclusive(input), + () => this.configureBillingExclusive(input) + ); + this.billingAdmissionTail = operation.then( + () => undefined, + () => undefined + ); + await operation; + } + + override async onStart(): Promise { + await super.onStart(); + const context = await getBillingContext(this.ctx.storage); + if (!context) { + await super.stop(); + throw new Error('Container started without an admitted billing context'); + } + await this.usageClient.recordStart(startInputFromContext(context)); + await this.billingHeartbeat.scheduleHeartbeat(); + } + + override async onStop(): Promise { + try { + await this.billingHeartbeat.recordStop({ reason: 'runtime_signal' }); + } finally { + await super.onStop(); + } + } + + override async onActivityExpired(): Promise { + try { + await this.billingHeartbeat.recordStop({ reason: 'activity_expired' }); + } finally { + await super.onActivityExpired(); + } + } + + private async configureBillingExclusive(input: SandboxBillingInput): Promise { + let context = await getBillingContext(this.ctx.storage); + if (context?.pendingStop) { + await this.billingHeartbeat.recordStop(context.pendingStop); + context = undefined; + } else if (context?.measurementStarted) { + const state = await this.getState(); + if (state.status === 'stopped' || state.status === 'stopped_with_code') { + await this.billingHeartbeat.recordStop({ + reason: 'runtime_signal', + ...(state.status === 'stopped_with_code' && state.exitCode !== undefined + ? { exitCode: state.exitCode } + : {}), + }); + context = undefined; + } + } + + const usageContext = { + service: SERVICE, + instanceId: `${this.sandboxClassName}:${this.ctx.id.toString()}`, + sku: SANDBOX_USAGE_SKUS[this.sandboxClassName], + ...input, + metadata: { ...input.metadata, container_class: this.sandboxClassName }, + } satisfies UsageContext; + + if (context) { + context = await setBillingContext(this.ctx.storage, { + ...usageContext, + startEpochMs: context.startEpochMs, + }); + } else { + const previousStartEpochMs = + (await this.ctx.storage.get(LAST_START_EPOCH_STORAGE_KEY)) ?? -1; + const startEpochMs = Math.max(Date.now(), previousStartEpochMs + 1); + await this.ctx.storage.put(LAST_START_EPOCH_STORAGE_KEY, startEpochMs); + context = await setBillingContext(this.ctx.storage, { ...usageContext, startEpochMs }); + } + + await this.usageClient.recordStart(startInputFromContext(context)); + } +} diff --git a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts index dfaa8a9677..677daab975 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts @@ -4,6 +4,11 @@ import { requiresContainmentSandbox } from '../persistence/session-metadata.js'; import { generateSandboxId, getSandboxNamespace } from '../sandbox-id.js'; import { fetchSessionMetadata } from '../session-service.js'; import type { Env, SandboxInstance, SandboxId, SessionId } from '../types.js'; +import { buildSandboxBillingInput, type SandboxBillingInput } from '../container-usage-context.js'; + +type BillingSandboxInstance = SandboxInstance & { + configureBilling(input: SandboxBillingInput): Promise; +}; export type SessionKiloFacadeDecision = | { kind: 'proxy-live-wrapper' } @@ -85,6 +90,9 @@ export async function resolveLiveWrapperTarget(params: { }), sandboxId ); + await (sandbox as BillingSandboxInstance).configureBilling( + buildSandboxBillingInput(metadata, sandboxId) + ); const wrapperInfo = await findWrapperForSession(sandbox, sessionId); if (!wrapperInfo) { return null; diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index b36a4af01b..d2b69e1cc9 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -368,6 +368,7 @@ describe('router sessionId validation', () => { SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], + CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], R2_BUCKET: {} as TRPCContext['env']['R2_BUCKET'], CLOUD_AGENT_REPORT_QUEUE: {} as TRPCContext['env']['CLOUD_AGENT_REPORT_QUEUE'], GIT_TOKEN_SERVICE: {} as Env['GIT_TOKEN_SERVICE'], @@ -385,6 +386,7 @@ describe('router sessionId validation', () => { // Mock sandbox with deleteSession method mockSandbox = { + configureBilling: vi.fn().mockResolvedValue(undefined), deleteSession: vi.fn().mockResolvedValue(undefined), } as unknown as ReturnType; @@ -799,6 +801,7 @@ describe('router sessionId validation', () => { SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], + CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], GIT_TOKEN_SERVICE: {} as Env['GIT_TOKEN_SERVICE'], R2_BUCKET: {} as TRPCContext['env']['R2_BUCKET'], CLOUD_AGENT_REPORT_QUEUE: {} as TRPCContext['env']['CLOUD_AGENT_REPORT_QUEUE'], @@ -917,6 +920,7 @@ describe('router sessionId validation', () => { SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], + CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], R2_BUCKET: {} as TRPCContext['env']['R2_BUCKET'], CLOUD_AGENT_REPORT_QUEUE: {} as TRPCContext['env']['CLOUD_AGENT_REPORT_QUEUE'], GIT_TOKEN_SERVICE: {} as Env['GIT_TOKEN_SERVICE'], @@ -1227,6 +1231,7 @@ describe('router sessionId validation', () => { SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], + CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], R2_BUCKET: {} as TRPCContext['env']['R2_BUCKET'], CLOUD_AGENT_REPORT_QUEUE: {} as TRPCContext['env']['CLOUD_AGENT_REPORT_QUEUE'], GIT_TOKEN_SERVICE: {} as Env['GIT_TOKEN_SERVICE'], @@ -1242,6 +1247,7 @@ describe('router sessionId validation', () => { }; cloudAgentSession = mockContext.env.CLOUD_AGENT_SESSION as unknown as MockCAS; vi.mocked(getSandbox).mockReturnValue({ + configureBilling: vi.fn().mockResolvedValue(undefined), listProcesses: mockListProcesses, } as unknown as ReturnType); caller = appRouter.createCaller(mockContext); @@ -1513,6 +1519,7 @@ describe('router sessionId validation', () => { SESSION_INGEST: { fetch: vi.fn(), } as unknown as TRPCContext['env']['SESSION_INGEST'], + CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], R2_BUCKET: {} as TRPCContext['env']['R2_BUCKET'], CLOUD_AGENT_REPORT_QUEUE: {} as TRPCContext['env']['CLOUD_AGENT_REPORT_QUEUE'], GIT_TOKEN_SERVICE: {} as Env['GIT_TOKEN_SERVICE'], @@ -1755,6 +1762,7 @@ describe('router question and permission controls', () => { }, ]); vi.mocked(getSandbox).mockReturnValue({ + configureBilling: vi.fn().mockResolvedValue(undefined), listProcesses, getSession, createSession, diff --git a/services/cloud-agent-next/src/sandbox-outbound.test.ts b/services/cloud-agent-next/src/sandbox-outbound.test.ts index e6332c2d2b..7ac7abe754 100644 --- a/services/cloud-agent-next/src/sandbox-outbound.test.ts +++ b/services/cloud-agent-next/src/sandbox-outbound.test.ts @@ -21,6 +21,7 @@ vi.mock('@cloudflare/sandbox', () => ({ Sandbox: sdk.StockSandbox, ContainerProxy: sdk.ContainerProxy, })); +vi.mock('./container-usage.js', () => ({ MeteredSandbox: sdk.StockSandbox })); vi.mock('./logger.js', () => ({ logger: logging.logger })); import { diff --git a/services/cloud-agent-next/src/sandbox-outbound.ts b/services/cloud-agent-next/src/sandbox-outbound.ts index d0cbd193ae..3db5b807c9 100644 --- a/services/cloud-agent-next/src/sandbox-outbound.ts +++ b/services/cloud-agent-next/src/sandbox-outbound.ts @@ -1,8 +1,10 @@ import { Buffer } from 'node:buffer'; -import { ContainerProxy, Sandbox as StockSandbox } from '@cloudflare/sandbox'; +import { ContainerProxy } from '@cloudflare/sandbox'; import { logger } from './logger.js'; import { MANAGED_SCM_OUTBOUND_HANDLER } from './sandbox-id.js'; import type { GitTokenService } from './types.js'; +import { MeteredSandbox } from './container-usage.js'; +import type { SandboxClassName } from './container-usage-context.js'; export { MANAGED_SCM_OUTBOUND_HANDLER } from './sandbox-id.js'; @@ -659,27 +661,32 @@ const managedScmOutboundHandlers = { [MANAGED_SCM_OUTBOUND_HANDLER]: handleManagedScmOutbound, }; -export class Sandbox extends StockSandbox { +export class Sandbox extends MeteredSandbox { + protected readonly sandboxClassName: SandboxClassName = 'Sandbox'; enableInternet = true; interceptHttps = false; } -export class SandboxSmall extends StockSandbox { +export class SandboxSmall extends MeteredSandbox { + protected readonly sandboxClassName: SandboxClassName = 'SandboxSmall'; enableInternet = true; interceptHttps = false; } -export class SandboxDIND extends StockSandbox { +export class SandboxDIND extends MeteredSandbox { + protected readonly sandboxClassName: SandboxClassName = 'SandboxDIND'; enableInternet = true; interceptHttps = false; } -export class SandboxCodeReview extends StockSandbox { +export class SandboxCodeReview extends MeteredSandbox { + protected readonly sandboxClassName: SandboxClassName = 'SandboxCodeReview'; enableInternet = true; interceptHttps = false; } export class SandboxContainment extends Sandbox { + protected override readonly sandboxClassName: SandboxClassName = 'SandboxContainment'; interceptHttps = true; } // Assignment (not a static class field) so it invokes the inherited Container.outboundHandlers @@ -688,11 +695,13 @@ export class SandboxContainment extends Sandbox { SandboxContainment.outboundHandlers = managedScmOutboundHandlers; export class SandboxSmallContainment extends SandboxSmall { + protected override readonly sandboxClassName: SandboxClassName = 'SandboxSmallContainment'; interceptHttps = true; } SandboxSmallContainment.outboundHandlers = managedScmOutboundHandlers; export class SandboxCodeReviewContainment extends SandboxCodeReview { + protected override readonly sandboxClassName: SandboxClassName = 'SandboxCodeReviewContainment'; interceptHttps = true; } SandboxCodeReviewContainment.outboundHandlers = managedScmOutboundHandlers; diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index e9824e2a23..95ed6a92d4 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -6,6 +6,7 @@ import type { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; import type { CallbackJob } from './callbacks/index.js'; import type { NotificationsBinding } from './notifications-binding.js'; import type { SessionIngestBinding } from './session-ingest-binding.js'; +import type { ContainerUsageRpcMethods } from '@kilocode/container-usage'; import type { SecretBinding } from './auth.js'; import * as z from 'zod'; import { Limits } from './schema.js'; @@ -473,6 +474,8 @@ export type Env = { SHARED_SANDBOX_OVERRIDES: KVNamespace; /** Service binding for the session ingest worker */ SESSION_INGEST: SessionIngestBinding; + /** Record-only container lifecycle usage meter. */ + CONTAINER_USAGE_METER: ContainerUsageRpcMethods; /** Shared secret for internal service-to-service authentication */ INTERNAL_API_SECRET_PROD: SecretsStoreSecret; /** R2 bucket for storing session logs */ diff --git a/services/cloud-agent-next/worker-configuration.d.ts b/services/cloud-agent-next/worker-configuration.d.ts index 2a910b806f..4b28db0973 100644 --- a/services/cloud-agent-next/worker-configuration.d.ts +++ b/services/cloud-agent-next/worker-configuration.d.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: dedf635481eae607f1cea623f29af2aa) +// Generated by Wrangler by running `wrangler types` (hash: 44ae1b44046f9f4e0c310861dcf7080d) // Runtime types generated with workerd@1.20260603.1 2026-06-03 nodejs_compat interface __BaseEnv_Env { SHARED_SANDBOX_OVERRIDES: KVNamespace; @@ -9,9 +9,13 @@ interface __BaseEnv_Env { CALLBACK_QUEUE: Queue; CLOUD_AGENT_REPORT_QUEUE: Queue; INTERNAL_API_SECRET_PROD: SecretsStoreSecret; + GITHUB_APP_SLUG: "kiloconnect-development" | "kiloconnect"; + GITHUB_APP_BOT_USER_ID: "242397087" | "240665456"; GITHUB_LITE_APP_SLUG: "" | "kiloconnect-lite"; GITHUB_LITE_APP_BOT_USER_ID: "" | "257753004"; SANDBOX_TRANSPORT?: "rpc"; + CLI_TIMEOUT_SECONDS: "900"; + REAPER_INTERVAL_MS: "300000"; R2_ATTACHMENTS_BUCKET: "cloud-agent-attachments-dev" | "cloud-agent-attachments"; BACKUP_BUCKET_NAME: "kilocode-sessions-dev" | "kilocode-sessions"; CLOUDFLARE_R2_ACCOUNT_ID: "e115e769bcdd4c3d66af59d3332cb394"; @@ -22,27 +26,18 @@ interface __BaseEnv_Env { REPO_SNAPSHOT_ORG_IDS?: ""; TOOL_CGROUP_ORG_IDS: "" | "*"; NEXTAUTH_SECRET: string; - KILO_SESSION_INGEST_URL: string; - WORKER_URL: string; - WS_ALLOWED_ORIGINS: string; + INTERNAL_API_SECRET: string; KILOCODE_BACKEND_BASE_URL: string; KILO_OPENROUTER_BASE: string; - INTERNAL_API_SECRET: string; + WORKER_URL: string; + AGENT_ENV_VARS_PRIVATE_KEY: string; R2_ENDPOINT: string; R2_ATTACHMENTS_READONLY_ACCESS_KEY_ID: string; R2_ATTACHMENTS_READONLY_SECRET_ACCESS_KEY: string; - AGENT_ENV_VARS_PRIVATE_KEY: string; - GITHUB_APP_ID: string; - GITHUB_APP_PRIVATE_KEY: string; - CLI_TIMEOUT_SECONDS: string; - REAPER_INTERVAL_MS: string; - STALE_THRESHOLD_MS: string; - PENDING_START_TIMEOUT_MS: string; - GITHUB_APP_SLUG: string; - GITHUB_APP_BOT_USER_ID: string; - GITHUB_LITE_APP_ID: string; - GITHUB_LITE_APP_PRIVATE_KEY: string; - KILOCODE_SANDBOX_BACKEND_BASE_URL: string; + R2_ACCESS_KEY_ID: string; + R2_SECRET_ACCESS_KEY: string; + KILO_SESSION_INGEST_URL: string; + WS_ALLOWED_ORIGINS: string; Sandbox: DurableObjectNamespace; SandboxSmall: DurableObjectNamespace; SandboxDIND: DurableObjectNamespace; @@ -55,6 +50,7 @@ interface __BaseEnv_Env { SESSION_INGEST: Service /* entrypoint SessionIngestRPC from session-ingest */; GIT_TOKEN_SERVICE: Service /* entrypoint GitTokenRPCEntrypoint from git-token-service-dev */ | Service /* entrypoint GitTokenRPCEntrypoint from git-token-service */; NOTIFICATIONS: Service /* entrypoint NotificationsService from notifications */; + CONTAINER_USAGE_METER: Service /* entrypoint ContainerUsageMeter from container-usage-meter */; TOOL_CGROUP_MODE?: "enforce"; TOOL_CGROUP_RESERVE_MB?: "1024"; TOOL_CGROUP_CPU_WEIGHT?: "50"; @@ -72,9 +68,13 @@ declare namespace Cloudflare { CALLBACK_QUEUE: Queue; CLOUD_AGENT_REPORT_QUEUE: Queue; INTERNAL_API_SECRET_PROD: SecretsStoreSecret; + GITHUB_APP_SLUG: "kiloconnect-development"; + GITHUB_APP_BOT_USER_ID: "242397087"; GITHUB_LITE_APP_SLUG: ""; GITHUB_LITE_APP_BOT_USER_ID: ""; SANDBOX_TRANSPORT: "rpc"; + CLI_TIMEOUT_SECONDS: "900"; + REAPER_INTERVAL_MS: "300000"; R2_ATTACHMENTS_BUCKET: "cloud-agent-attachments-dev"; BACKUP_BUCKET_NAME: "kilocode-sessions-dev"; CLOUDFLARE_R2_ACCOUNT_ID: "e115e769bcdd4c3d66af59d3332cb394"; @@ -85,27 +85,18 @@ declare namespace Cloudflare { REPO_SNAPSHOT_ORG_IDS: ""; TOOL_CGROUP_ORG_IDS: ""; NEXTAUTH_SECRET: string; - KILO_SESSION_INGEST_URL: string; - WORKER_URL: string; - WS_ALLOWED_ORIGINS: string; + INTERNAL_API_SECRET: string; KILOCODE_BACKEND_BASE_URL: string; KILO_OPENROUTER_BASE: string; - INTERNAL_API_SECRET: string; + WORKER_URL: string; + AGENT_ENV_VARS_PRIVATE_KEY: string; R2_ENDPOINT: string; R2_ATTACHMENTS_READONLY_ACCESS_KEY_ID: string; R2_ATTACHMENTS_READONLY_SECRET_ACCESS_KEY: string; - AGENT_ENV_VARS_PRIVATE_KEY: string; - GITHUB_APP_ID: string; - GITHUB_APP_PRIVATE_KEY: string; - CLI_TIMEOUT_SECONDS: string; - REAPER_INTERVAL_MS: string; - STALE_THRESHOLD_MS: string; - PENDING_START_TIMEOUT_MS: string; - GITHUB_APP_SLUG: string; - GITHUB_APP_BOT_USER_ID: string; - GITHUB_LITE_APP_ID: string; - GITHUB_LITE_APP_PRIVATE_KEY: string; - KILOCODE_SANDBOX_BACKEND_BASE_URL: string; + R2_ACCESS_KEY_ID: string; + R2_SECRET_ACCESS_KEY: string; + KILO_SESSION_INGEST_URL: string; + WS_ALLOWED_ORIGINS: string; Sandbox: DurableObjectNamespace; SandboxSmall: DurableObjectNamespace; SandboxDIND: DurableObjectNamespace; @@ -118,6 +109,7 @@ declare namespace Cloudflare { SESSION_INGEST: Service /* entrypoint SessionIngestRPC from session-ingest */; GIT_TOKEN_SERVICE: Service /* entrypoint GitTokenRPCEntrypoint from git-token-service-dev */; NOTIFICATIONS: Service /* entrypoint NotificationsService from notifications */; + CONTAINER_USAGE_METER: Service /* entrypoint ContainerUsageMeter from container-usage-meter */; } interface Env extends __BaseEnv_Env {} } @@ -126,7 +118,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } declare module "*.sql" { const value: string; diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index 5361132883..c9f2bc0f62 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -103,6 +103,11 @@ "service": "notifications", "entrypoint": "NotificationsService", }, + { + "binding": "CONTAINER_USAGE_METER", + "service": "container-usage-meter", + "entrypoint": "ContainerUsageMeter", + }, ], "secrets_store_secrets": [ { @@ -432,6 +437,11 @@ "service": "notifications", "entrypoint": "NotificationsService", }, + { + "binding": "CONTAINER_USAGE_METER", + "service": "container-usage-meter", + "entrypoint": "ContainerUsageMeter", + }, ], "secrets_store_secrets": [ { From 59b86ce4cc676e2b3b920b483194d20ee99fb769 Mon Sep 17 00:00:00 2001 From: syn Date: Thu, 23 Jul 2026 15:02:11 -0500 Subject: [PATCH 02/11] fix(billing): harden container usage lifecycle --- .../container-usage/src/heartbeat.test.ts | 88 +++++++ packages/container-usage/src/heartbeat.ts | 31 ++- .../cloudflare/cloudflare-agent-sandbox.ts | 10 +- .../src/container-usage-context.test.ts | 60 ++++- .../src/container-usage-context.ts | 83 +++++- .../src/container-usage.test.ts | 208 ++++++++++++--- .../cloud-agent-next/src/container-usage.ts | 237 +++++++++++++----- .../src/kilo-facade/session-proxy.ts | 10 +- .../src/persistence/session-metadata.ts | 1 + .../src/router/handlers/session-prepare.ts | 40 ++- .../src/router/handlers/session-start.ts | 18 +- .../src/session-prepare.test.ts | 47 ++++ .../src/session/session-registration.ts | 23 +- 13 files changed, 705 insertions(+), 151 deletions(-) diff --git a/packages/container-usage/src/heartbeat.test.ts b/packages/container-usage/src/heartbeat.test.ts index 7a470fdce7..346fa53c3a 100644 --- a/packages/container-usage/src/heartbeat.test.ts +++ b/packages/container-usage/src/heartbeat.test.ts @@ -129,6 +129,94 @@ describe('installBillingHeartbeat', () => { expect(schedule).toHaveBeenCalledWith(300, BILLING_HEARTBEAT_CALLBACK, expect.any(String)); }); + it('defers stopped-state closure when the producer owns an authoritative stop hook', async () => { + const storage = memoryStorage(); + await storedContext(storage); + const schedule = vi.fn(); + const recordStop = vi.fn(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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next' } + ); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(async () => ({ + status: 'stopped_with_code' as const, + lastChange: Date.now(), + exitCode: 17, + })), + schedule: schedule as Container['schedule'], + }, + { client, storage, stopOnStoppedState: false, enforceBudgetStop: vi.fn() } + ); + + await controller.billingHeartbeatTick(); + + expect(recordStop).not.toHaveBeenCalled(); + expect(await getBillingContext(storage)).toBeDefined(); + expect(schedule).toHaveBeenCalledWith(300, BILLING_HEARTBEAT_CALLBACK, expect.any(String)); + }); + + it('runs stop-delivery prerequisites before retrying a persisted stop', async () => { + const storage = memoryStorage(); + await storedContext(storage); + const beforeStopDelivery = vi.fn(async () => undefined); + const recordStop = vi.fn(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: 'continue' }, + }), + 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, beforeStopDelivery, enforceBudgetStop: vi.fn() } + ); + await controller.persistStop({ reason: 'exit', exitCode: 9 }); + + await controller.billingHeartbeatTick(); + + expect(beforeStopDelivery).toHaveBeenCalledOnce(); + expect(beforeStopDelivery.mock.invocationCallOrder[0]).toBeLessThan( + recordStop.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); + }); + it('immediately retries an unacknowledged heartbeat with the same segment payload', async () => { const storage = memoryStorage(); await storedContext(storage); diff --git a/packages/container-usage/src/heartbeat.ts b/packages/container-usage/src/heartbeat.ts index 0bdd83d644..019c71c956 100644 --- a/packages/container-usage/src/heartbeat.ts +++ b/packages/container-usage/src/heartbeat.ts @@ -20,6 +20,9 @@ export type BillingHeartbeatDependencies = { client: ContainerUsageClient; storage: BillingContextStorage; heartbeatSeconds?: number; + /** Defer stopped-state closure to the container's authoritative onStop hook. */ + stopOnStoppedState?: boolean; + beforeStopDelivery?: (context: BillingContext) => Promise; enforceBudgetStop: ( budget: BudgetVerdict, expected: { generation: string; startEpochMs: number } @@ -35,6 +38,10 @@ export type BillingHeartbeatController = { exitCode?: number; }) => Promise; cancelHeartbeat: () => void; + persistStop: (params: { + reason: 'exit' | 'runtime_signal' | 'activity_expired'; + exitCode?: number; + }) => Promise; }; function contextForHeartbeat(context: BillingContext) { @@ -120,6 +127,7 @@ export function installBillingHeartbeat( } const stopIntent = context.pendingStop; if (!stopIntent) throw new Error('Billing stop intent was not persisted'); + await dependencies.beforeStopDelivery?.(context); const ack = await dependencies.client.recordStop({ instanceId: context.instanceId, startEpochMs: context.startEpochMs, @@ -139,6 +147,23 @@ export function installBillingHeartbeat( const recordStop: BillingHeartbeatController['recordStop'] = params => runLifecycleExclusive(() => recordStopForGeneration(params)); + const persistStop: BillingHeartbeatController['persistStop'] = params => + runLifecycleExclusive(async () => { + const context = await getBillingContext(dependencies.storage); + if (!context) return undefined; + if (context.pendingStop) return context; + const pendingHeartbeat = context.pendingHeartbeat; + const elapsedMs = Math.max(0, Date.now() - context.usageMeasuredAtMs); + const stopSegment = pendingHeartbeat ?? { + seq: context.nextSeq, + usageSinceLast: Math.floor(elapsedMs / 1_000), + measuredAtMs: context.usageMeasuredAtMs + Math.floor(elapsedMs / 1_000) * 1_000, + }; + const updated = { ...context, pendingStop: { ...params, ...stopSegment } }; + await updateBillingContext(dependencies.storage, updated); + return updated; + }); + const billingHeartbeatTickForGeneration = async (generation?: string) => { let context = await getBillingContext(dependencies.storage); if (!context) { @@ -167,6 +192,10 @@ export function installBillingHeartbeat( if (!currentAfterState || !isSameBillingGeneration(currentAfterState, context)) return; context = currentAfterState; if (state.status === 'stopped' || state.status === 'stopped_with_code') { + if (dependencies.stopOnStoppedState === false) { + await rescheduleIfCurrent(context); + return; + } try { await recordStopForGeneration( { @@ -254,5 +283,5 @@ export function installBillingHeartbeat( value: billingHeartbeatTick, }); - return { scheduleHeartbeat, billingHeartbeatTick, recordStop, cancelHeartbeat }; + return { scheduleHeartbeat, billingHeartbeatTick, recordStop, persistStop, cancelHeartbeat }; } diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts index 577afa088b..6ca926f6ac 100644 --- a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts @@ -73,13 +73,10 @@ import { import { TOOL_CGROUP_ENV_KEYS, type ToolCgroupEnv } from '../../shared/tool-cgroup-env.js'; import { buildSandboxBillingInput, + configureSandboxBillingInput, type SandboxBillingInput, } from '../../container-usage-context.js'; -type BillingSandboxInstance = SandboxInstance & { - configureBilling(input: SandboxBillingInput): Promise; -}; - const PREPARE_WORKSPACE_TIMEOUT_MS = 10 * 60 * 1000; const DEFAULT_STOP_OBSERVATION_DELAYS_MS = [100, 500, 1_000]; @@ -207,10 +204,7 @@ export class CloudflareAgentSandbox implements AgentSandbox { dependencies.stopObservationDelaysMs ?? DEFAULT_STOP_OBSERVATION_DELAYS_MS; this.configureBilling = dependencies.configureBilling ?? - (dependencies.resolveSandbox - ? async () => undefined - : async (sandbox, input) => - await (sandbox as BillingSandboxInstance).configureBilling(input)); + (dependencies.resolveSandbox ? async () => undefined : configureSandboxBillingInput); } private resolveSandboxId(): Promise { diff --git a/services/cloud-agent-next/src/container-usage-context.test.ts b/services/cloud-agent-next/src/container-usage-context.test.ts index c05239cfba..66da8fb154 100644 --- a/services/cloud-agent-next/src/container-usage-context.test.ts +++ b/services/cloud-agent-next/src/container-usage-context.test.ts @@ -1,6 +1,10 @@ import { describe, expect, it } from 'vitest'; import type { SessionMetadata } from './persistence/session-metadata.js'; -import { buildSandboxBillingInput, SANDBOX_USAGE_SKUS } from './container-usage-context.js'; +import { + assertSandboxBillingAllocation, + buildSandboxBillingInput, + SANDBOX_USAGE_SKUS, +} from './container-usage-context.js'; function metadata(identity: SessionMetadata['identity']): SessionMetadata { return { @@ -75,7 +79,7 @@ describe('container usage context', () => { sessionId: 'agent_security', userId: 'user_security', orgId: 'org_security', - createdOnPlatform: 'security-remediation', + billingOrigin: 'security-remediation', }), 'crv-isolated' ); @@ -97,7 +101,7 @@ describe('container usage context', () => { sessionId: 'agent_first', userId: 'user_shared', orgId: 'org_shared', - createdOnPlatform: 'security-agent', + billingOrigin: 'security-agent', }), 'org-shared' ); @@ -106,7 +110,7 @@ describe('container usage context', () => { sessionId: 'agent_second', userId: 'user_shared', orgId: 'org_shared', - createdOnPlatform: 'cloud-agent-web', + billingOrigin: 'cloud-agent-web', }), 'org-shared' ); @@ -124,10 +128,56 @@ describe('container usage context', () => { metadata({ sessionId: 'agent_unknown', userId: 'user_unknown', - createdOnPlatform: 'attacker-controlled-value', + billingOrigin: 'attacker-controlled-value', }), 'dind-isolated' ); expect(input.metadata?.origin).toBe('other'); }); + + it('does not positively attribute legacy metadata without a trusted billing origin', () => { + const input = buildSandboxBillingInput( + metadata({ + sessionId: 'agent_legacy', + userId: 'user_legacy', + createdOnPlatform: 'code-review', + }), + 'crv-legacy' + ); + expect(input.metadata?.origin).toBe('other'); + }); + + it('does not trust the public createdOnPlatform label as billing origin', () => { + const input = buildSandboxBillingInput( + metadata({ + sessionId: 'agent_public', + userId: 'user_public', + createdOnPlatform: 'security-remediation', + billingOrigin: 'cloud-agent', + }), + 'ses-isolated' + ); + expect(input.metadata?.origin).toBe('cloud-agent'); + }); + + it('rejects session attribution and extra metadata for shared sandboxes', () => { + expect(() => + assertSandboxBillingAllocation('Sandbox', { + subject: { type: 'user', id: 'user_shared' }, + actor: { type: 'user', id: 'user_shared' }, + sessionId: 'agent_leak', + metadata: { allocation: 'shared', origin: 'cloud-agent' }, + }) + ).toThrow('Shared sandbox billing cannot contain session attribution'); + }); + + it('requires bounded isolated attribution for non-shared sandbox classes', () => { + expect(() => + assertSandboxBillingAllocation('SandboxSmall', { + subject: { type: 'user', id: 'user_isolated' }, + actor: { type: 'user', id: 'user_isolated' }, + metadata: { allocation: 'isolated' }, + }) + ).toThrow('Isolated sandbox billing requires session attribution'); + }); }); diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts index db8855ee85..7b888481a1 100644 --- a/services/cloud-agent-next/src/container-usage-context.ts +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -1,6 +1,12 @@ -import type { UsageContext } from '@kilocode/container-usage'; +import { + billingActorSchema, + billingSubjectSchema, + usageContextSchema, + type UsageContext, +} from '@kilocode/container-usage'; +import { z } from 'zod'; import type { SessionMetadata } from './persistence/session-metadata.js'; -import type { SandboxId } from './types.js'; +import type { SandboxId, SandboxInstance } from './types.js'; export const SANDBOX_USAGE_SKUS = { Sandbox: 'cloud-agent-standard-2026-07', @@ -14,6 +20,24 @@ export const SANDBOX_USAGE_SKUS = { export type SandboxClassName = keyof typeof SANDBOX_USAGE_SKUS; export type SandboxBillingInput = Omit; +export type MeteredSandboxInstance = SandboxInstance & { + configureBilling(input: unknown): Promise; +}; + +const sandboxBillingInputEnvelopeSchema = z + .object({ + subject: billingSubjectSchema, + actor: billingActorSchema, + onBehalfOf: billingSubjectSchema.optional(), + sessionId: z.string().min(1).max(256).optional(), + metadata: z + .record(z.string().min(1).max(64), z.string().max(512)) + .refine(metadata => Object.keys(metadata).length <= 16, { + message: 'Metadata may contain at most 16 entries', + }) + .optional(), + }) + .strict(); const KNOWN_ORIGINS = new Set([ 'app-builder', @@ -33,7 +57,7 @@ const KNOWN_ORIGINS = new Set([ ]); function normalizedOrigin(origin: string | undefined): string { - if (!origin) return 'cloud-agent'; + if (!origin) return 'other'; return KNOWN_ORIGINS.has(origin) ? origin : 'other'; } @@ -61,9 +85,60 @@ export function buildSandboxBillingInput( metadata: isolated ? { allocation: 'isolated', - origin: normalizedOrigin(metadata.identity.createdOnPlatform), + origin: normalizedOrigin(metadata.identity.billingOrigin), ...(metadata.repository ? { repository_provider: metadata.repository.type } : {}), } : { allocation: 'shared' }, }; } + +export function parseSandboxBillingInput(input: unknown): SandboxBillingInput { + const parsed = sandboxBillingInputEnvelopeSchema.parse(input); + const validated = usageContextSchema.parse({ + service: 'cloud-agent-next', + instanceId: 'validation', + sku: 'validation', + ...parsed, + }); + const { service: _service, instanceId: _instanceId, sku: _sku, ...billingInput } = validated; + return billingInput; +} + +export function assertSandboxBillingAllocation( + sandboxClassName: SandboxClassName, + input: SandboxBillingInput +): void { + const shared = sandboxClassName === 'Sandbox' || sandboxClassName === 'SandboxContainment'; + if (shared) { + if (input.sessionId !== undefined || input.metadata?.allocation !== 'shared') { + throw new Error('Shared sandbox billing cannot contain session attribution'); + } + if (Object.keys(input.metadata).some(key => key !== 'allocation')) { + throw new Error('Shared sandbox billing metadata must contain only allocation'); + } + return; + } + + if (!input.sessionId || input.metadata?.allocation !== 'isolated') { + throw new Error('Isolated sandbox billing requires session attribution'); + } + const allowedMetadata = new Set(['allocation', 'origin', 'repository_provider']); + if (Object.keys(input.metadata).some(key => !allowedMetadata.has(key))) { + throw new Error('Isolated sandbox billing metadata contains an unsupported field'); + } +} + +export async function configureSandboxBilling( + sandbox: SandboxInstance, + metadata: SessionMetadata, + sandboxId: SandboxId +): Promise { + await configureSandboxBillingInput(sandbox, buildSandboxBillingInput(metadata, sandboxId)); +} + +export async function configureSandboxBillingInput( + sandbox: SandboxInstance, + input: SandboxBillingInput +): Promise { + await (sandbox as MeteredSandboxInstance).configureBilling(input); +} diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index ddfbc0344b..f6a611bd65 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -1,5 +1,5 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; -import type { ContainerUsageRpcMethods } from '@kilocode/container-usage'; +import { getBillingContext, type ContainerUsageRpcMethods } from '@kilocode/container-usage'; // oxlint-disable-next-line no-empty-object-type -- Matches the mocked Sandbox constructor. type SandboxDurableObjectState = DurableObjectState<{}>; @@ -54,7 +54,6 @@ const sdk = vi.hoisted(() => { vi.mock('@cloudflare/sandbox', () => ({ Sandbox: sdk.StockSandbox })); -import { getBillingContext } from '@kilocode/container-usage'; import { MeteredSandbox } from './container-usage.js'; class MemoryStorage { @@ -98,6 +97,7 @@ type TestRuntime = MeteredSandbox & { superStopped: boolean; superActivityExpired: boolean; superStopCalled: boolean; + billingHeartbeatTick(generation?: string): Promise; }; function createSandbox(rpc = createRpc()) { @@ -131,58 +131,163 @@ describe('MeteredSandbox', () => { vi.restoreAllMocks(); }); - it('admits duplicate acquisition with one generation-stable start', async () => { + it('admits one start per physical generation and short-circuits active acquisition', async () => { const { rpc, storage, sandbox } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(1_000); await sandbox.configureBilling(billingInput); await sandbox.configureBilling(billingInput); + expect(rpc.recordStart).not.toHaveBeenCalled(); - expect(rpc.recordStart).toHaveBeenCalledTimes(2); - const starts = vi.mocked(rpc.recordStart).mock.calls.map(([input]) => input); - expect(starts[0]?.startEpochMs).toBe(1_000); - expect(starts[1]?.startEpochMs).toBe(1_000); - expect(starts[0]?.idempotencyKey).toBe(starts[1]?.idempotencyKey); - expect(starts[0]).toMatchObject({ - instanceId: 'SandboxSmallContainment:do-id', - sku: 'cloud-agent-small-2026-07', - metadata: { allocation: 'isolated', container_class: 'SandboxSmallContainment' }, + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + vi.mocked(rpc.recordStart).mockRejectedValue(new Error('meter unavailable')); + await sandbox.configureBilling(billingInput); + await sandbox.configureBilling(billingInput); + + expect(rpc.recordStart).toHaveBeenCalledOnce(); + expect(sandbox.schedules).toHaveLength(1); + expect(rpc.recordStart).toHaveBeenCalledWith( + expect.objectContaining({ + startEpochMs: 1_000, + instanceId: 'SandboxSmallContainment:do-id', + sku: 'cloud-agent-small-2026-07', + metadata: { allocation: 'isolated', container_class: 'SandboxSmallContainment' }, + }) + ); + expect((await getBillingContext(storage))?.measurementStarted).toBe(true); + }); + + it('adopts a physical container that predates shadow metering', async () => { + const { rpc, storage, sandbox } = createSandbox(); + vi.spyOn(Date, 'now').mockReturnValue(1_500); + sandbox.mockState = { status: 'healthy' }; + + await sandbox.configureBilling(billingInput); + + expect(rpc.recordStart).toHaveBeenCalledOnce(); + expect(await getBillingContext(storage)).toMatchObject({ + startEpochMs: 1_500, + measurementStarted: true, }); - expect((await getBillingContext(storage))?.measurementStarted).toBe(false); }); - it('starts five-minute heartbeat measurement after preserving the SDK start hook', async () => { - const { sandbox, storage } = createSandbox(); + it('retries an unacknowledged start before allowing active-generation work', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStart) + .mockRejectedValueOnce(new Error('ack lost')) + .mockRejectedValueOnce(new Error('ack lost')) + .mockRejectedValueOnce(new Error('ack lost')) + .mockResolvedValue({ success: true, ack: ack() }); + const { sandbox, storage } = createSandbox(rpc); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + + await expect(sandbox.onStart()).rejects.toThrow('ack lost'); + const context = await getBillingContext(storage); + expect(context?.measurementStarted).toBe(false); + + await sandbox.configureBilling(billingInput); + expect(rpc.recordStart).toHaveBeenCalledTimes(4); + expect((await getBillingContext(storage))?.measurementStarted).toBe(true); + }); + + it('keeps a delayed stop attached to the prior generation', async () => { + const { rpc, storage, sandbox } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(2_000); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + const first = await getBillingContext(storage); + + sandbox.mockState = { status: 'stopped_with_code', exitCode: 17 }; + await sandbox.configureBilling({ ...billingInput, sessionId: 'agent_2' }); + expect((await getBillingContext(storage))?.generation).toBe(first?.generation); + await sandbox.onStop({ reason: 'exit', exitCode: 17 }); + expect(await getBillingContext(storage)).toBeUndefined(); + + sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + const second = await getBillingContext(storage); - expect(sandbox.superStarted).toBe(true); - expect(sandbox.schedules).toEqual([ - expect.objectContaining({ when: 300, callback: 'billingHeartbeatTick' }), - ]); - expect((await getBillingContext(storage))?.measurementStarted).toBe(true); + expect(second?.generation).not.toBe(first?.generation); + expect(second?.startEpochMs).toBe(2_001); + expect(second?.sessionId).toBe('agent_2'); + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ reason: 'exit', exitCode: 17, startEpochMs: 2_000 }) + ); }); - it('closes a stopped generation before allocating a monotonic replacement', async () => { - const { rpc, sandbox } = createSandbox(); + it('treats duplicate start callbacks as one physical generation', async () => { + const { rpc, storage, sandbox } = createSandbox(); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + + await sandbox.onStart(); + const first = await getBillingContext(storage); + await sandbox.onStart(); + const second = await getBillingContext(storage); + + expect(second?.generation).toBe(first?.generation); + expect(rpc.recordStart).toHaveBeenCalledOnce(); + expect(rpc.recordStop).not.toHaveBeenCalled(); + }); + + it('does not admit a new physical generation until the prior stop is durable', async () => { + const rpc = createRpc(); + const { sandbox, storage } = createSandbox(rpc); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + const first = await getBillingContext(storage); + await sandbox.billingHeartbeatTick(first?.generation); + vi.mocked(rpc.recordStop).mockRejectedValue(new Error('meter unavailable')); + await sandbox.onStop({ reason: 'exit', exitCode: 1 }); + expect((await getBillingContext(storage))?.pendingStop).toBeDefined(); + + await expect(sandbox.onStart()).rejects.toThrow('meter unavailable'); + + expect((await getBillingContext(storage))?.generation).toBe(first?.generation); + expect(rpc.recordStart).toHaveBeenCalledOnce(); + }); + + it('defers activity-expiry closure until physical stop confirmation', async () => { + const { rpc, storage, sandbox } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(3_000); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); - sandbox.mockState = { status: 'stopped_with_code', exitCode: 17 }; + const active = await getBillingContext(storage); + + await sandbox.onActivityExpired(); + expect(sandbox.superActivityExpired).toBe(true); + expect(rpc.recordStop).not.toHaveBeenCalled(); + expect((await getBillingContext(storage))?.generation).toBe(active?.generation); + + await sandbox.onStop({ reason: 'exit', exitCode: 143 }); + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ reason: 'activity_expired', exitCode: 143 }) + ); + expect(await getBillingContext(storage)).toBeUndefined(); + }); + + it('preserves normal exit reason and exit code', async () => { + const { rpc, sandbox } = createSandbox(); await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + + await sandbox.onStop({ reason: 'exit', exitCode: 42 }); expect(rpc.recordStop).toHaveBeenCalledWith( - expect.objectContaining({ reason: 'runtime_signal', exitCode: 17, startEpochMs: 3_000 }) + expect.objectContaining({ reason: 'exit', exitCode: 42 }) ); - expect(vi.mocked(rpc.recordStart).mock.calls.at(-1)?.[0].startEpochMs).toBe(3_001); + expect(sandbox.superStopped).toBe(true); }); - it('persists and retries the same durable stop intent while preserving SDK cleanup', async () => { + it('persists a failed stop and retries it without blocking SDK cleanup', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStop) .mockRejectedValueOnce(new Error('postgres unavailable')) @@ -190,27 +295,56 @@ describe('MeteredSandbox', () => { .mockRejectedValueOnce(new Error('postgres unavailable')) .mockResolvedValue(ack()); const { sandbox, storage } = createSandbox(rpc); - vi.spyOn(Date, 'now').mockReturnValue(4_000); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); - await expect(sandbox.onActivityExpired()).rejects.toThrow('postgres unavailable'); - expect(sandbox.superActivityExpired).toBe(true); - expect((await getBillingContext(storage))?.pendingStop?.reason).toBe('activity_expired'); - - await sandbox.onStop(); + await expect(sandbox.onStop({ reason: 'exit', exitCode: 1 })).resolves.toBeUndefined(); expect(sandbox.superStopped).toBe(true); + const pending = await getBillingContext(storage); + expect(pending?.pendingStop).toMatchObject({ reason: 'exit', exitCode: 1 }); + + await sandbox.billingHeartbeatTick(pending?.generation); expect(await getBillingContext(storage)).toBeUndefined(); - const stops = vi.mocked(rpc.recordStop).mock.calls.map(([input]) => input); - expect(stops.at(-1)?.reason).toBe('activity_expired'); - expect(new Set(stops.map(stop => stop.idempotencyKey))).toHaveLength(1); + expect( + new Set(vi.mocked(rpc.recordStop).mock.calls.map(([input]) => input.idempotencyKey)) + ).toHaveLength(1); + }); + + it('persists authoritative stop details before recovering a missing start ack', async () => { + const rpc = createRpc(); + const { sandbox, storage } = createSandbox(rpc); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await storage.delete('container-usage:start-ack-generation:v1'); + vi.mocked(rpc.recordStart).mockRejectedValue(new Error('meter unavailable')); + + await sandbox.onStop({ reason: 'exit', exitCode: 9 }); + + expect(await getBillingContext(storage)).toMatchObject({ + pendingStop: { reason: 'exit', exitCode: 9 }, + }); + expect(sandbox.superStopped).toBe(true); + }); + + it('rejects meter-owned identity fields at the custom RPC boundary', async () => { + const { sandbox } = createSandbox(); + await expect( + sandbox.configureBilling({ ...billingInput, instanceId: 'forged-instance' }) + ).rejects.toThrow(); + await expect( + sandbox.configureBilling({ ...billingInput, sku: 'forged-sku' }) + ).rejects.toThrow(); + await expect( + sandbox.configureBilling({ ...billingInput, service: 'forged-service' }) + ).rejects.toThrow(); }); - it('stops an unadmitted runtime instead of silently running it', async () => { + it('stops a runtime that somehow starts without trusted attribution', async () => { const { sandbox } = createSandbox(); await expect(sandbox.onStart()).rejects.toThrow( - 'Container started without an admitted billing context' + 'Container started without pending billing attribution' ); expect(sandbox.superStarted).toBe(true); expect(sandbox.superStopCalled).toBe(true); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 14ae768362..0f1812f20e 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -11,17 +11,33 @@ import { type UsageContext, } from '@kilocode/container-usage'; import { Sandbox as StockSandbox } from '@cloudflare/sandbox'; +import { z } from 'zod'; +import { logger } from './logger.js'; import type { Env } from './types.js'; import { + assertSandboxBillingAllocation, + parseSandboxBillingInput, SANDBOX_USAGE_SKUS, type SandboxBillingInput, type SandboxClassName, } from './container-usage-context.js'; const SERVICE = 'cloud-agent-next'; +const PENDING_ATTRIBUTION_STORAGE_KEY = 'container-usage:pending-attribution:v1'; +const PENDING_STOP_REASON_STORAGE_KEY = 'container-usage:pending-stop-reason:v1'; +const START_ACK_GENERATION_STORAGE_KEY = 'container-usage:start-ack-generation:v1'; const LAST_START_EPOCH_STORAGE_KEY = 'container-usage:last-start-epoch:v1'; + // oxlint-disable-next-line no-empty-object-type -- Matches the Sandbox 0.12.1 constructor. type SandboxDurableObjectState = DurableObjectState<{}>; +type ContainerStopParams = { reason: 'exit' | 'runtime_signal'; exitCode?: number }; + +const pendingStopReasonSchema = z + .object({ + generation: z.uuid(), + reason: z.literal('activity_expired'), + }) + .strict(); function startInputFromContext(context: BillingContext): ClientRecordStartInput { const { service: _service, ...usage } = usageContextFromBillingContext(context); @@ -33,7 +49,7 @@ export abstract class MeteredSandbox extends StockSandbox { private readonly usageClient: ContainerUsageClient; private readonly billingHeartbeat: BillingHeartbeatController; - private billingAdmissionTail: Promise = Promise.resolve(); + private billingLifecycleTail: Promise = Promise.resolve(); constructor(ctx: SandboxDurableObjectState, env: Env) { super(ctx, env); @@ -43,91 +59,192 @@ export abstract class MeteredSandbox extends StockSandbox { this.billingHeartbeat = installBillingHeartbeat(this, { client: this.usageClient, storage: this.ctx.storage, - // Shadow metering must never stop customer work. + stopOnStoppedState: false, + beforeStopDelivery: context => this.ensureStartAcknowledged(context), + // The meter currently returns only `continue`; shadow mode must not enforce future verdicts. enforceBudgetStop: async () => { throw new Error('Container budget enforcement is disabled in shadow mode'); }, }); } - async configureBilling(input: SandboxBillingInput): Promise { - const operation = this.billingAdmissionTail.then( - () => this.configureBillingExclusive(input), - () => this.configureBillingExclusive(input) - ); - this.billingAdmissionTail = operation.then( - () => undefined, - () => undefined - ); - await operation; + async configureBilling(input: unknown): Promise { + const parsed = parseSandboxBillingInput(input); + assertSandboxBillingAllocation(this.sandboxClassName, parsed); + await this.runBillingExclusive(async () => { + await this.ctx.storage.put(PENDING_ATTRIBUTION_STORAGE_KEY, parsed); + let active = await getBillingContext(this.ctx.storage); + if (active?.pendingStop) { + await this.ensureStartAcknowledged(active); + await this.billingHeartbeat.recordStop(active.pendingStop); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + active = undefined; + } + if (active?.measurementStarted) { + await this.ensureStartAcknowledged(active); + return; + } + + // A start may have succeeded before the DO was evicted or a prior admission response failed. + // Retry the same idempotent start before allowing work to use that running generation. + if (active) { + const state = await this.getState(); + if (state.status !== 'stopped' && state.status !== 'stopped_with_code') { + await this.admitAndSchedule(active); + return; + } + await this.ensureStartAcknowledged(active); + await this.billingHeartbeat.recordStop({ + reason: 'runtime_signal', + ...(state.status === 'stopped_with_code' && state.exitCode !== undefined + ? { exitCode: state.exitCode } + : {}), + }); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + } + + // Adopt containers that were already running when shadow metering rolled out. + const state = await this.getState(); + if (state.status !== 'stopped' && state.status !== 'stopped_with_code') { + await this.startBillingGeneration(parsed); + } + }); } override async onStart(): Promise { await super.onStart(); - const context = await getBillingContext(this.ctx.storage); - if (!context) { - await super.stop(); - throw new Error('Container started without an admitted billing context'); - } - await this.usageClient.recordStart(startInputFromContext(context)); - await this.billingHeartbeat.scheduleHeartbeat(); + await this.runBillingExclusive(async () => { + const previous = await getBillingContext(this.ctx.storage); + if (previous) { + if (previous.pendingStop) { + await this.billingHeartbeat.recordStop(previous.pendingStop); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + } else if (!previous.measurementStarted) { + await this.admitAndSchedule(previous); + return; + } else { + // The SDK can dispatch onStart more than once for concurrent callers waiting on one + // physical start. Existing measured state is therefore already the current generation. + await this.ensureStartAcknowledged(previous); + return; + } + } + + const input = await this.getPendingAttribution(); + if (!input) { + await super.stop(); + throw new Error('Container started without pending billing attribution'); + } + + await this.startBillingGeneration(input); + }); } - override async onStop(): Promise { + override async onStop(params?: ContainerStopParams): Promise { try { - await this.billingHeartbeat.recordStop({ reason: 'runtime_signal' }); + await this.runBillingExclusive(async () => { + const context = await getBillingContext(this.ctx.storage); + if (!context) return; + const requestedReason = await this.getPendingStopReason(context.generation); + try { + const pending = await this.billingHeartbeat.persistStop({ + reason: requestedReason ?? params?.reason ?? 'runtime_signal', + exitCode: params?.exitCode, + }); + if (!pending) return; + await this.ensureStartAcknowledged(pending); + await this.billingHeartbeat.recordStop( + pending.pendingStop ?? { + reason: requestedReason ?? params?.reason ?? 'runtime_signal', + exitCode: params?.exitCode, + } + ); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); + } catch (error) { + // recordStop persists its intent before delivery. Keep the heartbeat schedule alive so + // the durable intent retries without blocking the SDK's physical stop transition. + try { + await this.billingHeartbeat.scheduleHeartbeat(); + } catch { + // The persisted stop intent remains recoverable on the next sandbox acquisition. + } + logger + .withFields({ + error: error instanceof Error ? error.message : String(error), + sandboxClass: this.sandboxClassName, + }) + .warn('Container usage stop delivery deferred'); + } + }); } finally { await super.onStop(); } } override async onActivityExpired(): Promise { - try { - await this.billingHeartbeat.recordStop({ reason: 'activity_expired' }); - } finally { - await super.onActivityExpired(); - } - } - - private async configureBillingExclusive(input: SandboxBillingInput): Promise { - let context = await getBillingContext(this.ctx.storage); - if (context?.pendingStop) { - await this.billingHeartbeat.recordStop(context.pendingStop); - context = undefined; - } else if (context?.measurementStarted) { - const state = await this.getState(); - if (state.status === 'stopped' || state.status === 'stopped_with_code') { - await this.billingHeartbeat.recordStop({ - reason: 'runtime_signal', - ...(state.status === 'stopped_with_code' && state.exitCode !== undefined - ? { exitCode: state.exitCode } - : {}), + await this.runBillingExclusive(async () => { + const context = await getBillingContext(this.ctx.storage); + if (context) { + await this.ctx.storage.put(PENDING_STOP_REASON_STORAGE_KEY, { + generation: context.generation, + reason: 'activity_expired', }); - context = undefined; } + }); + await super.onActivityExpired(); + } + + private runBillingExclusive(operation: () => Promise): Promise { + const result = this.billingLifecycleTail.then(operation, operation); + this.billingLifecycleTail = result.then( + () => undefined, + () => undefined + ); + return result; + } + + private async getPendingAttribution(): Promise { + const stored = await this.ctx.storage.get(PENDING_ATTRIBUTION_STORAGE_KEY); + return stored === undefined ? undefined : parseSandboxBillingInput(stored); + } + + private async getPendingStopReason(generation: string): Promise<'activity_expired' | undefined> { + const stored = await this.ctx.storage.get(PENDING_STOP_REASON_STORAGE_KEY); + if (stored === undefined) return undefined; + const parsed = pendingStopReasonSchema.parse(stored); + return parsed.generation === generation ? parsed.reason : undefined; + } + + private async admitAndSchedule(context: BillingContext): Promise { + await this.ensureStartAcknowledged(context); + await this.billingHeartbeat.scheduleHeartbeat(); + } + + private async ensureStartAcknowledged(context: BillingContext): Promise { + const acknowledgedGeneration = await this.ctx.storage.get( + START_ACK_GENERATION_STORAGE_KEY + ); + if (acknowledgedGeneration !== context.generation) { + await this.usageClient.recordStart(startInputFromContext(context)); + await this.ctx.storage.put(START_ACK_GENERATION_STORAGE_KEY, context.generation); } + } - const usageContext = { + private async startBillingGeneration(input: SandboxBillingInput): Promise { + const previousStartEpochMs = + (await this.ctx.storage.get(LAST_START_EPOCH_STORAGE_KEY)) ?? -1; + const startEpochMs = Math.max(Date.now(), previousStartEpochMs + 1); + await this.ctx.storage.put(LAST_START_EPOCH_STORAGE_KEY, startEpochMs); + const context = await setBillingContext(this.ctx.storage, { + ...input, service: SERVICE, instanceId: `${this.sandboxClassName}:${this.ctx.id.toString()}`, sku: SANDBOX_USAGE_SKUS[this.sandboxClassName], - ...input, metadata: { ...input.metadata, container_class: this.sandboxClassName }, - } satisfies UsageContext; - - if (context) { - context = await setBillingContext(this.ctx.storage, { - ...usageContext, - startEpochMs: context.startEpochMs, - }); - } else { - const previousStartEpochMs = - (await this.ctx.storage.get(LAST_START_EPOCH_STORAGE_KEY)) ?? -1; - const startEpochMs = Math.max(Date.now(), previousStartEpochMs + 1); - await this.ctx.storage.put(LAST_START_EPOCH_STORAGE_KEY, startEpochMs); - context = await setBillingContext(this.ctx.storage, { ...usageContext, startEpochMs }); - } - - await this.usageClient.recordStart(startInputFromContext(context)); + startEpochMs, + } satisfies UsageContext & { startEpochMs: number }); + await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); + await this.admitAndSchedule(context); } } diff --git a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts index 677daab975..eed0a1ed3b 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts @@ -4,11 +4,7 @@ import { requiresContainmentSandbox } from '../persistence/session-metadata.js'; import { generateSandboxId, getSandboxNamespace } from '../sandbox-id.js'; import { fetchSessionMetadata } from '../session-service.js'; import type { Env, SandboxInstance, SandboxId, SessionId } from '../types.js'; -import { buildSandboxBillingInput, type SandboxBillingInput } from '../container-usage-context.js'; - -type BillingSandboxInstance = SandboxInstance & { - configureBilling(input: SandboxBillingInput): Promise; -}; +import { configureSandboxBilling } from '../container-usage-context.js'; export type SessionKiloFacadeDecision = | { kind: 'proxy-live-wrapper' } @@ -90,9 +86,7 @@ export async function resolveLiveWrapperTarget(params: { }), sandboxId ); - await (sandbox as BillingSandboxInstance).configureBilling( - buildSandboxBillingInput(metadata, sandboxId) - ); + await configureSandboxBilling(sandbox, metadata, sandboxId); const wrapperInfo = await findWrapperForSession(sandbox, sessionId); if (!wrapperInfo) { return null; diff --git a/services/cloud-agent-next/src/persistence/session-metadata.ts b/services/cloud-agent-next/src/persistence/session-metadata.ts index 1e2c073204..33ce5c0e2e 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.ts @@ -36,6 +36,7 @@ const MetadataIdentitySchema = z orgId: z.string().optional(), botId: z.string().optional(), createdOnPlatform: z.string().max(100).optional(), + billingOrigin: z.string().max(100).optional(), }) .strip(); diff --git a/services/cloud-agent-next/src/router/handlers/session-prepare.ts b/services/cloud-agent-next/src/router/handlers/session-prepare.ts index 1be3f2fa9a..c5c674c23e 100644 --- a/services/cloud-agent-next/src/router/handlers/session-prepare.ts +++ b/services/cloud-agent-next/src/router/handlers/session-prepare.ts @@ -43,6 +43,7 @@ import type { SessionProfileBundle } from '../../session-profile.js'; import type { SessionCreateRequest } from '../../session/session-requests.js'; import { assertKiloModelAvailable } from '../../model-validation.js'; import { assertBitbucketRepositoryAccessBeforeSessionCreation } from '../../session/validate-repository-access.js'; +import { assertOrganizationMembership } from './session-start.js'; type SessionPrepareHandlers = { prepareSession: typeof prepareSessionHandler; @@ -313,6 +314,13 @@ const prepareSessionHandler = internalApiProtectedProcedure .mutation(async ({ input, ctx }) => { return withLogTags({ source: 'prepareSession' }, async () => { const request = prepareInputToSessionCreateRequest(input); + if (input.kilocodeOrganizationId) { + await assertOrganizationMembership( + getPgDb(ctx.env), + ctx.userId, + input.kilocodeOrganizationId + ); + } await assertBitbucketRepositoryAccessBeforeSessionCreation({ env: ctx.env, userId: ctx.userId, @@ -356,18 +364,26 @@ const prepareSessionHandler = internalApiProtectedProcedure const result = input.autoInitiate === true - ? await startNewSession(requestWithProfile, { - env: ctx.env, - userId: ctx.userId, - authToken: ctx.authToken, - botId: ctx.botId, - }) - : await registerNewSession(requestWithProfile, { - env: ctx.env, - userId: ctx.userId, - authToken: ctx.authToken, - botId: ctx.botId, - }); + ? await startNewSession( + requestWithProfile, + { + env: ctx.env, + userId: ctx.userId, + authToken: ctx.authToken, + botId: ctx.botId, + }, + { billingOrigin: input.createdOnPlatform } + ) + : await registerNewSession( + requestWithProfile, + { + env: ctx.env, + userId: ctx.userId, + authToken: ctx.authToken, + botId: ctx.botId, + }, + { billingOrigin: input.createdOnPlatform } + ); return { cloudAgentSessionId: result.cloudAgentSessionId, diff --git a/services/cloud-agent-next/src/router/handlers/session-start.ts b/services/cloud-agent-next/src/router/handlers/session-start.ts index 56cfae0f34..ea5124ba81 100644 --- a/services/cloud-agent-next/src/router/handlers/session-start.ts +++ b/services/cloud-agent-next/src/router/handlers/session-start.ts @@ -91,7 +91,7 @@ function startInputToSessionCreateRequest( }; } -async function assertOrganizationMembership( +export async function assertOrganizationMembership( db: WorkerDb, userId: string, organizationId: string @@ -156,12 +156,16 @@ const startSessionHandler = protectedProcedure procedure: 'start', }); - const registration = await startNewSession(requestWithProfile, { - env: ctx.env, - userId: ctx.userId, - authToken: ctx.authToken, - botId: ctx.botId, - }); + const registration = await startNewSession( + requestWithProfile, + { + env: ctx.env, + userId: ctx.userId, + authToken: ctx.authToken, + botId: ctx.botId, + }, + { billingOrigin: 'cloud-agent' } + ); const ack = registration.admission; logger diff --git a/services/cloud-agent-next/src/session-prepare.test.ts b/services/cloud-agent-next/src/session-prepare.test.ts index b91521575a..92168d1345 100644 --- a/services/cloud-agent-next/src/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session-prepare.test.ts @@ -285,6 +285,7 @@ describe('prepareSession endpoint', () => { recordInitialAdmissionMock.mockResolvedValue(undefined); recordInternalCompensationMock.mockResolvedValue(undefined); mergeProfileConfigurationMock.mockResolvedValue({}); + organizationMembershipLimitMock.mockResolvedValue([{ id: 'membership-123' }]); assertKiloModelAvailableMock.mockResolvedValue(undefined); }); @@ -499,6 +500,7 @@ describe('prepareSession endpoint', () => { orgId: 'f47ac10b-58cc-4372-a567-0e02b2c3d479', botId: undefined, createdOnPlatform: 'code-review', + billingOrigin: 'code-review', }, auth: { kiloSessionId: 'cli-session-abc123', @@ -551,6 +553,23 @@ describe('prepareSession endpoint', () => { ); }); + it('rejects organization attribution when the internal caller user is not a member', async () => { + organizationMembershipLimitMock.mockResolvedValueOnce([]); + const doStub = createMockDOStub(); + const caller = appRouter.createCaller(createInternalApiContext({ doStub })); + + await expect( + caller.prepareSession({ + prompt: 'Attempt unrelated organization attribution', + mode: 'code', + model: 'claude-3', + githubRepo: 'acme/repo', + kilocodeOrganizationId: 'f47ac10b-58cc-4372-a567-0e02b2c3d479', + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(doStub.registerSession).not.toHaveBeenCalled(); + }); + it('retains split legacy preparation as registration-only', async () => { const doStub = createMockDOStub(); const caller = appRouter.createCaller(createInternalApiContext({ doStub })); @@ -774,6 +793,7 @@ describe('prepareSession endpoint', () => { expect(overrideStore.get).toHaveBeenCalledWith(`shared-sandbox-route:${routeKey}`); expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( expect.objectContaining({ + identity: expect.objectContaining({ billingOrigin: 'cloud-agent' }), workspace: { sandboxId: failoverSandboxId, sandboxProvider: 'cloudflare', @@ -793,6 +813,33 @@ describe('prepareSession endpoint', () => { ); }); + it('does not let public createdOnPlatform select the Code Review sandbox class', async () => { + const doStub = createMockDOStub(); + const caller = appRouter.createCaller(createInternalApiContext({ doStub })); + + await caller.start({ + message: { prompt: 'Attempt to select a reserved class' }, + agent: { mode: 'code', model: 'anthropic/claude-sonnet-4-20250514' }, + repository: { type: 'github', repo: 'acme/repo' }, + profile: { id: 'f47ac10b-58cc-4372-a567-0e02b2c3d479' }, + options: { createdOnPlatform: 'code-review' }, + }); + + expect(generateSandboxRoutingTargetMock).toHaveBeenCalledWith( + undefined, + undefined, + 'test-user-123', + expect.any(String), + undefined, + expect.objectContaining({ createdOnPlatform: undefined }) + ); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + identity: expect.objectContaining({ billingOrigin: 'cloud-agent' }), + }) + ); + }); + it('creates auto-initiated devcontainer sessions with grouped DIND sandbox intent', async () => { generateSandboxRoutingTargetMock.mockResolvedValueOnce({ kind: 'isolated', diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index ab3062b327..90f47d6dfb 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -145,7 +145,8 @@ async function recordPostSetupFailure(record: () => Promise): Promise { const sessionService = new SessionService(); const initialTurn = acceptInitialTurn(input.initialTurn); @@ -184,7 +185,7 @@ async function allocateNewSession( ctx.botId, { devcontainer: input.runtime?.devcontainer, - createdOnPlatform: input.options?.createdOnPlatform, + createdOnPlatform: options?.billingOrigin === 'code-review' ? 'code-review' : undefined, } ); if (target.kind === 'shared') { @@ -282,7 +283,8 @@ async function allocateNewSession( function buildSessionRegistrationCommand( input: SessionRegistrationInput, ctx: SessionRegistrationContext, - allocation: NewSessionAllocation + allocation: NewSessionAllocation, + options?: { billingOrigin?: string } ) { return { identity: { @@ -291,6 +293,7 @@ function buildSessionRegistrationCommand( orgId: input.options?.kilocodeOrganizationId, botId: ctx.botId, createdOnPlatform: input.options?.createdOnPlatform, + billingOrigin: options?.billingOrigin, }, auth: { kiloSessionId: allocation.kiloSessionId, @@ -328,9 +331,10 @@ function buildSessionRegistrationCommand( */ export async function registerNewSession( input: SessionRegistrationInput, - ctx: SessionRegistrationContext + ctx: SessionRegistrationContext, + options?: { billingOrigin?: string } ): Promise { - const allocation = await allocateNewSession(input, ctx); + const allocation = await allocateNewSession(input, ctx, options); const doId = ctx.env.CLOUD_AGENT_SESSION.idFromName( `${ctx.userId}:${allocation.cloudAgentSessionId}` ); @@ -338,7 +342,7 @@ export async function registerNewSession( let registerResult: Awaited>; try { registerResult = await stub.registerSession( - buildSessionRegistrationCommand(input, ctx, allocation) + buildSessionRegistrationCommand(input, ctx, allocation, options) ); } catch (error) { await recordPostSetupFailure(() => @@ -384,9 +388,10 @@ export async function registerNewSession( */ export async function startNewSession( input: SessionRegistrationInput, - ctx: SessionRegistrationContext + ctx: SessionRegistrationContext, + options?: { billingOrigin?: string } ): Promise { - const allocation = await allocateNewSession(input, ctx); + const allocation = await allocateNewSession(input, ctx, options); const doId = ctx.env.CLOUD_AGENT_SESSION.idFromName( `${ctx.userId}:${allocation.cloudAgentSessionId}` ); @@ -399,7 +404,7 @@ export async function startNewSession( () => ctx.env.CLOUD_AGENT_SESSION.get(doId), stub => stub.createSessionWithInitialAdmission({ - ...buildSessionRegistrationCommand(input, ctx, allocation), + ...buildSessionRegistrationCommand(input, ctx, allocation, options), message: { initialTurn: allocation.initialTurn }, }), 'createSessionWithInitialAdmission' From f14e2b09e46a8fe1731ebed29f512322cec25705 Mon Sep 17 00:00:00 2001 From: syn Date: Thu, 23 Jul 2026 16:43:46 -0500 Subject: [PATCH 03/11] fix(billing): follow physical container lifecycle Use ctx.container.running as the authoritative signal for adopting or reusing a usage generation. Persisted Sandbox health can remain stale after the physical process exits, which could otherwise create phantom intervals or carry one interval across two container instances. Mark heartbeat measurement active only after Container scheduling succeeds. This keeps scheduling failures recoverable instead of leaving a running generation permanently marked as metered without heartbeat segments. --- .../container-usage/src/heartbeat.test.ts | 18 +++++++ packages/container-usage/src/heartbeat.ts | 6 +-- .../src/container-usage.test.ts | 54 ++++++++++++++++++- .../cloud-agent-next/src/container-usage.ts | 17 ++++-- 4 files changed, 87 insertions(+), 8 deletions(-) diff --git a/packages/container-usage/src/heartbeat.test.ts b/packages/container-usage/src/heartbeat.test.ts index 346fa53c3a..21f78dc15e 100644 --- a/packages/container-usage/src/heartbeat.test.ts +++ b/packages/container-usage/src/heartbeat.test.ts @@ -107,6 +107,24 @@ describe('installBillingHeartbeat', () => { expect(Object.hasOwn(container, BILLING_HEARTBEAT_CALLBACK)).toBe(true); }); + it('does not mark measurement started when initial scheduling fails', async () => { + const storage = memoryStorage(); + await storedContext(storage); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(), + schedule: vi.fn(async () => { + throw new Error('schedule unavailable'); + }) as Container['schedule'], + }, + { client: usageClient('continue'), storage, enforceBudgetStop: vi.fn() } + ); + + await expect(controller.scheduleHeartbeat()).rejects.toThrow('schedule unavailable'); + expect((await getBillingContext(storage))?.measurementStarted).toBe(false); + }); + it('keeps a stopped-state probe scheduled until stop is durably acknowledged', async () => { const storage = memoryStorage(); await storedContext(storage); diff --git a/packages/container-usage/src/heartbeat.ts b/packages/container-usage/src/heartbeat.ts index 019c71c956..4aa4df5f38 100644 --- a/packages/container-usage/src/heartbeat.ts +++ b/packages/container-usage/src/heartbeat.ts @@ -81,15 +81,15 @@ export function installBillingHeartbeat( const startedContext = context.measurementStarted ? context : { ...context, measurementStarted: true, usageMeasuredAtMs: Date.now() }; - if (!context.measurementStarted) { - await updateBillingContext(dependencies.storage, startedContext); - } cancelHeartbeat(); await container.schedule( heartbeatSeconds, BILLING_HEARTBEAT_CALLBACK, startedContext.generation ); + if (!context.measurementStarted) { + await updateBillingContext(dependencies.storage, startedContext); + } }; const rescheduleIfCurrent = async (expected: BillingContext): Promise => { diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index f6a611bd65..c3744f8337 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -34,10 +34,16 @@ const sdk = vi.hoisted(() => { } async onStart(): Promise { + if (this.ctx.container) { + Object.defineProperty(this.ctx.container, 'running', { value: true, configurable: true }); + } this.superStarted = true; } async onStop(): Promise { + if (this.ctx.container) { + Object.defineProperty(this.ctx.container, 'running', { value: false, configurable: true }); + } this.superStopped = true; } @@ -97,17 +103,28 @@ type TestRuntime = MeteredSandbox & { superStopped: boolean; superActivityExpired: boolean; superStopCalled: boolean; + setPhysicalRunning(running: boolean): void; billingHeartbeatTick(generation?: string): Promise; }; -function createSandbox(rpc = createRpc()) { +function createSandbox(rpc = createRpc(), containerRunning = false) { const storage = new MemoryStorage(); const ctx = { id: { toString: () => 'do-id' }, storage, + container: { running: containerRunning }, } as unknown as SandboxDurableObjectState; class TestSandbox extends MeteredSandbox { protected readonly sandboxClassName = 'SandboxSmallContainment' as const; + + setPhysicalRunning(running: boolean): void { + if (this.ctx.container) { + Object.defineProperty(this.ctx.container, 'running', { + value: running, + configurable: true, + }); + } + } } return { rpc, @@ -159,7 +176,7 @@ describe('MeteredSandbox', () => { }); it('adopts a physical container that predates shadow metering', async () => { - const { rpc, storage, sandbox } = createSandbox(); + const { rpc, storage, sandbox } = createSandbox(createRpc(), true); vi.spyOn(Date, 'now').mockReturnValue(1_500); sandbox.mockState = { status: 'healthy' }; @@ -172,6 +189,39 @@ describe('MeteredSandbox', () => { }); }); + it('does not adopt stale healthy state when no physical container is running', async () => { + const { rpc, storage, sandbox } = createSandbox(); + sandbox.mockState = { status: 'healthy' }; + + await sandbox.configureBilling(billingInput); + + expect(rpc.recordStart).not.toHaveBeenCalled(); + expect(await getBillingContext(storage)).toBeUndefined(); + }); + + it('closes a missed-stop generation before the next physical start', async () => { + const { rpc, storage, sandbox } = createSandbox(); + vi.spyOn(Date, 'now').mockReturnValue(1_750); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + const first = await getBillingContext(storage); + + sandbox.setPhysicalRunning(false); + sandbox.mockState = { status: 'stopped' }; + await sandbox.configureBilling({ ...billingInput, sessionId: 'agent_2' }); + + expect(await getBillingContext(storage)).toBeUndefined(); + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ startEpochMs: first?.startEpochMs, reason: 'runtime_signal' }) + ); + + await sandbox.onStart(); + const second = await getBillingContext(storage); + expect(second?.generation).not.toBe(first?.generation); + expect(second?.startEpochMs).toBe(1_751); + }); + it('retries an unacknowledged start before allowing active-generation work', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStart) diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 0f1812f20e..31af8b98c9 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -81,8 +81,20 @@ export abstract class MeteredSandbox extends StockSandbox { active = undefined; } if (active?.measurementStarted) { + if (this.ctx.container?.running === true) { + await this.ensureStartAcknowledged(active); + return; + } + const state = await this.getState(); await this.ensureStartAcknowledged(active); - return; + await this.billingHeartbeat.recordStop({ + reason: 'runtime_signal', + ...(state.status === 'stopped_with_code' && state.exitCode !== undefined + ? { exitCode: state.exitCode } + : {}), + }); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + active = undefined; } // A start may have succeeded before the DO was evicted or a prior admission response failed. @@ -104,8 +116,7 @@ export abstract class MeteredSandbox extends StockSandbox { } // Adopt containers that were already running when shadow metering rolled out. - const state = await this.getState(); - if (state.status !== 'stopped' && state.status !== 'stopped_with_code') { + if (this.ctx.container?.running === true) { await this.startBillingGeneration(parsed); } }); From a42184b56be0b492a1871fa2c34a3a43c6a6baf4 Mon Sep 17 00:00:00 2001 From: syn Date: Thu, 23 Jul 2026 16:54:50 -0500 Subject: [PATCH 04/11] fix(billing): keep shadow metering non-blocking Treat exhausted meter retries as telemetry delivery failures: retain durable generation state, arm heartbeat recovery, and warn without rejecting sandbox acquisition or physical container start. This preserves the record-only shadow contract while keeping local attribution validation strict. Bound deferred stopped-state observation to the existing 15-minute stale grace. If authoritative onStop never arrives, close the interval at the first observed stopped timestamp so alarms do not wake forever and the grace period is not billed. --- packages/container-usage/src/context.ts | 2 + .../container-usage/src/heartbeat.test.ts | 58 +++++++++ packages/container-usage/src/heartbeat.ts | 33 ++++- .../src/container-usage-context.test.ts | 28 ++++- .../src/container-usage-context.ts | 14 ++- .../src/container-usage.test.ts | 13 +- .../cloud-agent-next/src/container-usage.ts | 113 +++++++++++++----- 7 files changed, 219 insertions(+), 42 deletions(-) diff --git a/packages/container-usage/src/context.ts b/packages/container-usage/src/context.ts index 0099aba683..c4698ca634 100644 --- a/packages/container-usage/src/context.ts +++ b/packages/container-usage/src/context.ts @@ -10,6 +10,7 @@ export const billingContextSchema = usageContextSchema measurementStarted: z.boolean(), nextSeq: z.number().int().positive().max(2_147_483_647).default(1), usageMeasuredAtMs: z.number().int().nonnegative().finite(), + stoppedObservedAtMs: z.number().int().nonnegative().finite().optional(), pendingHeartbeat: z .object({ seq: z.number().int().positive().finite(), @@ -69,6 +70,7 @@ export async function setBillingContext( measurementStarted: false, nextSeq: 1, usageMeasuredAtMs: Date.now(), + stoppedObservedAtMs: undefined, pendingHeartbeat: undefined, pendingStop: undefined, }); diff --git a/packages/container-usage/src/heartbeat.test.ts b/packages/container-usage/src/heartbeat.test.ts index 21f78dc15e..8baabef08f 100644 --- a/packages/container-usage/src/heartbeat.test.ts +++ b/packages/container-usage/src/heartbeat.test.ts @@ -192,6 +192,64 @@ describe('installBillingHeartbeat', () => { expect(schedule).toHaveBeenCalledWith(300, BILLING_HEARTBEAT_CALLBACK, expect.any(String)); }); + it('closes a repeatedly stopped generation after the deferred stale grace', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + const recordStop = vi.fn(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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next' } + ); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(async () => ({ status: 'stopped' as const, lastChange: Date.now() })), + schedule: vi.fn() as Container['schedule'], + }, + { + client, + storage, + stopOnStoppedState: false, + stoppedStateGraceSeconds: 900, + enforceBudgetStop: vi.fn(), + } + ); + + now.mockReturnValue(10_000); + await controller.billingHeartbeatTick(); + expect(recordStop).not.toHaveBeenCalled(); + expect((await getBillingContext(storage))?.stoppedObservedAtMs).toBe(10_000); + + now.mockReturnValue(910_000); + await controller.billingHeartbeatTick(); + + expect(recordStop).toHaveBeenCalledWith( + expect.objectContaining({ reason: 'runtime_signal', usageSinceLast: 9 }) + ); + expect(await getBillingContext(storage)).toBeUndefined(); + } finally { + now.mockRestore(); + } + }); + it('runs stop-delivery prerequisites before retrying a persisted stop', async () => { const storage = memoryStorage(); await storedContext(storage); diff --git a/packages/container-usage/src/heartbeat.ts b/packages/container-usage/src/heartbeat.ts index 4aa4df5f38..2493a3e516 100644 --- a/packages/container-usage/src/heartbeat.ts +++ b/packages/container-usage/src/heartbeat.ts @@ -13,6 +13,7 @@ import { export const BILLING_HEARTBEAT_CALLBACK = 'billingHeartbeatTick'; export const DEFAULT_BILLING_HEARTBEAT_SECONDS = 5 * 60; +export const DEFAULT_STOPPED_STATE_GRACE_SECONDS = 15 * 60; type BillingContainer = Pick; @@ -22,6 +23,8 @@ export type BillingHeartbeatDependencies = { heartbeatSeconds?: number; /** Defer stopped-state closure to the container's authoritative onStop hook. */ stopOnStoppedState?: boolean; + stoppedStateGraceSeconds?: number; + beforeHeartbeatDelivery?: (context: BillingContext) => Promise; beforeStopDelivery?: (context: BillingContext) => Promise; enforceBudgetStop: ( budget: BudgetVerdict, @@ -54,9 +57,14 @@ export function installBillingHeartbeat( dependencies: BillingHeartbeatDependencies ): BillingHeartbeatController { const heartbeatSeconds = dependencies.heartbeatSeconds ?? DEFAULT_BILLING_HEARTBEAT_SECONDS; + const stoppedStateGraceSeconds = + dependencies.stoppedStateGraceSeconds ?? DEFAULT_STOPPED_STATE_GRACE_SECONDS; if (heartbeatSeconds <= 0) { throw new Error('Billing heartbeat interval must be positive'); } + if (stoppedStateGraceSeconds <= 0) { + throw new Error('Stopped-state grace interval must be positive'); + } let lifecycleTail: Promise = Promise.resolve(); const runLifecycleExclusive = (operation: () => Promise): Promise => { @@ -102,7 +110,8 @@ export function installBillingHeartbeat( const recordStopForGeneration = async ( params: Parameters[0], - expectedGeneration?: string + expectedGeneration?: string, + usageEndedAtMs = Date.now() ): Promise => { let context = await getBillingContext(dependencies.storage); if (!context) return undefined; @@ -111,7 +120,7 @@ export function installBillingHeartbeat( } if (!context.pendingStop) { const pendingHeartbeat = context.pendingHeartbeat; - const elapsedMs = Math.max(0, Date.now() - context.usageMeasuredAtMs); + const elapsedMs = Math.max(0, usageEndedAtMs - context.usageMeasuredAtMs); const stopSegment = pendingHeartbeat ?? { seq: context.nextSeq, usageSinceLast: Math.floor(elapsedMs / 1_000), @@ -193,8 +202,15 @@ export function installBillingHeartbeat( context = currentAfterState; if (state.status === 'stopped' || state.status === 'stopped_with_code') { if (dependencies.stopOnStoppedState === false) { - await rescheduleIfCurrent(context); - return; + const stoppedObservedAtMs = context.stoppedObservedAtMs ?? Date.now(); + if (context.stoppedObservedAtMs === undefined) { + context = { ...context, stoppedObservedAtMs }; + await updateBillingContext(dependencies.storage, context); + } + if (Date.now() - stoppedObservedAtMs < stoppedStateGraceSeconds * 1_000) { + await rescheduleIfCurrent(context); + return; + } } try { await recordStopForGeneration( @@ -202,7 +218,8 @@ export function installBillingHeartbeat( reason: 'runtime_signal', exitCode: state.status === 'stopped_with_code' ? state.exitCode : undefined, }, - context.generation + context.generation, + context.stoppedObservedAtMs ); } catch (error) { await rescheduleIfCurrent(context); @@ -216,6 +233,11 @@ export function installBillingHeartbeat( return; } + if (context.stoppedObservedAtMs !== undefined) { + context = { ...context, stoppedObservedAtMs: undefined }; + await updateBillingContext(dependencies.storage, context); + } + const pendingHeartbeat = context.pendingHeartbeat ?? (() => { @@ -233,6 +255,7 @@ export function installBillingHeartbeat( await updateBillingContext(dependencies.storage, { ...context, pendingHeartbeat }); } try { + await dependencies.beforeHeartbeatDelivery?.(context); const ack = await dependencies.client.recordHeartbeat({ instanceId: context.instanceId, startEpochMs: context.startEpochMs, diff --git a/services/cloud-agent-next/src/container-usage-context.test.ts b/services/cloud-agent-next/src/container-usage-context.test.ts index 66da8fb154..0b4186383e 100644 --- a/services/cloud-agent-next/src/container-usage-context.test.ts +++ b/services/cloud-agent-next/src/container-usage-context.test.ts @@ -1,8 +1,10 @@ -import { describe, expect, it } from 'vitest'; +import { describe, expect, it, vi } from 'vitest'; +import type { SandboxInstance } from './types.js'; import type { SessionMetadata } from './persistence/session-metadata.js'; import { assertSandboxBillingAllocation, buildSandboxBillingInput, + configureSandboxBillingInput, SANDBOX_USAGE_SKUS, } from './container-usage-context.js'; @@ -180,4 +182,28 @@ describe('container usage context', () => { }) ).toThrow('Isolated sandbox billing requires session attribution'); }); + + it('skips shadow configuration when a sandbox does not expose the metering RPC', async () => { + await expect( + configureSandboxBillingInput({} as SandboxInstance, { + subject: { type: 'user', id: 'user_1' }, + actor: { type: 'user', id: 'user_1' }, + sessionId: 'agent_1', + metadata: { allocation: 'isolated' }, + }) + ).resolves.toBeUndefined(); + }); + + it('does not propagate shadow configuration delivery failures', async () => { + const configureBilling = vi.fn().mockRejectedValue(new Error('meter unavailable')); + await expect( + configureSandboxBillingInput({ configureBilling } as unknown as SandboxInstance, { + subject: { type: 'user', id: 'user_1' }, + actor: { type: 'user', id: 'user_1' }, + sessionId: 'agent_1', + metadata: { allocation: 'isolated' }, + }) + ).resolves.toBeUndefined(); + expect(configureBilling).toHaveBeenCalledOnce(); + }); }); diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts index 7b888481a1..74a90ed598 100644 --- a/services/cloud-agent-next/src/container-usage-context.ts +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -5,6 +5,7 @@ import { type UsageContext, } from '@kilocode/container-usage'; import { z } from 'zod'; +import { logger } from './logger.js'; import type { SessionMetadata } from './persistence/session-metadata.js'; import type { SandboxId, SandboxInstance } from './types.js'; @@ -140,5 +141,16 @@ export async function configureSandboxBillingInput( sandbox: SandboxInstance, input: SandboxBillingInput ): Promise { - await (sandbox as MeteredSandboxInstance).configureBilling(input); + const configureBilling = (sandbox as Partial).configureBilling; + if (typeof configureBilling !== 'function') { + logger.warn('Container usage shadow metering is unavailable for sandbox'); + return; + } + try { + await configureBilling.call(sandbox, input); + } catch (error) { + logger + .withFields({ error: error instanceof Error ? error.message : String(error) }) + .warn('Container usage shadow configuration deferred'); + } } diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index c3744f8337..40cc7aa10f 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -222,7 +222,7 @@ describe('MeteredSandbox', () => { expect(second?.startEpochMs).toBe(1_751); }); - it('retries an unacknowledged start before allowing active-generation work', async () => { + it('keeps physical start non-fatal while retrying an unacknowledged shadow start', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStart) .mockRejectedValueOnce(new Error('ack lost')) @@ -233,13 +233,14 @@ describe('MeteredSandbox', () => { await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; - await expect(sandbox.onStart()).rejects.toThrow('ack lost'); + await expect(sandbox.onStart()).resolves.toBeUndefined(); const context = await getBillingContext(storage); - expect(context?.measurementStarted).toBe(false); + expect(context?.measurementStarted).toBe(true); + expect(sandbox.superStarted).toBe(true); - await sandbox.configureBilling(billingInput); + await sandbox.billingHeartbeatTick(context?.generation); expect(rpc.recordStart).toHaveBeenCalledTimes(4); - expect((await getBillingContext(storage))?.measurementStarted).toBe(true); + expect(rpc.recordHeartbeat).toHaveBeenCalledOnce(); }); it('keeps a delayed stop attached to the prior generation', async () => { @@ -296,7 +297,7 @@ describe('MeteredSandbox', () => { await sandbox.onStop({ reason: 'exit', exitCode: 1 }); expect((await getBillingContext(storage))?.pendingStop).toBeDefined(); - await expect(sandbox.onStart()).rejects.toThrow('meter unavailable'); + await expect(sandbox.onStart()).resolves.toBeUndefined(); expect((await getBillingContext(storage))?.generation).toBe(first?.generation); expect(rpc.recordStart).toHaveBeenCalledOnce(); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 31af8b98c9..252e19d4fd 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -60,6 +60,7 @@ export abstract class MeteredSandbox extends StockSandbox { client: this.usageClient, storage: this.ctx.storage, stopOnStoppedState: false, + beforeHeartbeatDelivery: context => this.ensureStartAcknowledged(context), beforeStopDelivery: context => this.ensureStartAcknowledged(context), // The meter currently returns only `continue`; shadow mode must not enforce future verdicts. enforceBudgetStop: async () => { @@ -75,25 +76,37 @@ export abstract class MeteredSandbox extends StockSandbox { await this.ctx.storage.put(PENDING_ATTRIBUTION_STORAGE_KEY, parsed); let active = await getBillingContext(this.ctx.storage); if (active?.pendingStop) { - await this.ensureStartAcknowledged(active); - await this.billingHeartbeat.recordStop(active.pendingStop); - await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + try { + await this.billingHeartbeat.recordStop(active.pendingStop); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + } catch (error) { + await this.deferBillingDelivery(error, 'pending stop recovery'); + return; + } active = undefined; } if (active?.measurementStarted) { if (this.ctx.container?.running === true) { - await this.ensureStartAcknowledged(active); + try { + await this.ensureStartAcknowledged(active); + } catch (error) { + await this.deferBillingDelivery(error, 'active start acknowledgement'); + } return; } const state = await this.getState(); - await this.ensureStartAcknowledged(active); - await this.billingHeartbeat.recordStop({ - reason: 'runtime_signal', - ...(state.status === 'stopped_with_code' && state.exitCode !== undefined - ? { exitCode: state.exitCode } - : {}), - }); - await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + try { + await this.billingHeartbeat.recordStop({ + reason: 'runtime_signal', + ...(state.status === 'stopped_with_code' && state.exitCode !== undefined + ? { exitCode: state.exitCode } + : {}), + }); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + } catch (error) { + await this.deferBillingDelivery(error, 'missed stop recovery'); + return; + } active = undefined; } @@ -102,17 +115,20 @@ export abstract class MeteredSandbox extends StockSandbox { if (active) { const state = await this.getState(); if (state.status !== 'stopped' && state.status !== 'stopped_with_code') { - await this.admitAndSchedule(active); + await this.admitAndScheduleBestEffort(active); return; } - await this.ensureStartAcknowledged(active); - await this.billingHeartbeat.recordStop({ - reason: 'runtime_signal', - ...(state.status === 'stopped_with_code' && state.exitCode !== undefined - ? { exitCode: state.exitCode } - : {}), - }); - await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + try { + await this.billingHeartbeat.recordStop({ + reason: 'runtime_signal', + ...(state.status === 'stopped_with_code' && state.exitCode !== undefined + ? { exitCode: state.exitCode } + : {}), + }); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + } catch (error) { + await this.deferBillingDelivery(error, 'unmeasured stop recovery'); + } } // Adopt containers that were already running when shadow metering rolled out. @@ -128,15 +144,24 @@ export abstract class MeteredSandbox extends StockSandbox { const previous = await getBillingContext(this.ctx.storage); if (previous) { if (previous.pendingStop) { - await this.billingHeartbeat.recordStop(previous.pendingStop); - await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + try { + await this.billingHeartbeat.recordStop(previous.pendingStop); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + } catch (error) { + await this.deferBillingDelivery(error, 'start blocked by pending stop'); + return; + } } else if (!previous.measurementStarted) { - await this.admitAndSchedule(previous); + await this.admitAndScheduleBestEffort(previous); return; } else { // The SDK can dispatch onStart more than once for concurrent callers waiting on one // physical start. Existing measured state is therefore already the current generation. - await this.ensureStartAcknowledged(previous); + try { + await this.ensureStartAcknowledged(previous); + } catch (error) { + await this.deferBillingDelivery(error, 'duplicate start acknowledgement'); + } return; } } @@ -227,9 +252,39 @@ export abstract class MeteredSandbox extends StockSandbox { return parsed.generation === generation ? parsed.reason : undefined; } - private async admitAndSchedule(context: BillingContext): Promise { - await this.ensureStartAcknowledged(context); - await this.billingHeartbeat.scheduleHeartbeat(); + private async admitAndScheduleBestEffort(context: BillingContext): Promise { + try { + await this.ensureStartAcknowledged(context); + } catch (error) { + await this.deferBillingDelivery(error, 'start acknowledgement'); + return; + } + try { + await this.billingHeartbeat.scheduleHeartbeat(); + } catch (error) { + await this.deferBillingDelivery(error, 'heartbeat scheduling', false); + } + } + + private async deferBillingDelivery( + error: unknown, + operation: string, + scheduleRetry = true + ): Promise { + if (scheduleRetry) { + try { + await this.billingHeartbeat.scheduleHeartbeat(); + } catch { + // A later sandbox acquisition retries persisted shadow state. + } + } + logger + .withFields({ + error: error instanceof Error ? error.message : String(error), + operation, + sandboxClass: this.sandboxClassName, + }) + .warn('Container usage shadow delivery deferred'); } private async ensureStartAcknowledged(context: BillingContext): Promise { @@ -256,6 +311,6 @@ export abstract class MeteredSandbox extends StockSandbox { startEpochMs, } satisfies UsageContext & { startEpochMs: number }); await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); - await this.admitAndSchedule(context); + await this.admitAndScheduleBestEffort(context); } } From 61f7b4001ebe6c75117680345ddff4f9c4ce2f06 Mon Sep 17 00:00:00 2001 From: syn Date: Thu, 23 Jul 2026 20:07:25 -0500 Subject: [PATCH 05/11] fix(billing): isolate shadow metering from containers Run acquisition and Sandbox lifecycle telemetry outside the real container control path. Start, stop, idle expiry, and wrapper lookup now complete without waiting for meter RPCs or shadow storage, while Durable Object waitUntil tasks retain best-effort recording and failure logs. Bound missed-stop retries, preserve the first observed stop cutoff, use trusted fallback routing metadata, validate origins at the Sandbox boundary, and remove the organization-handler import cycle. Shadow failures can no longer stop, reset, delay, or prevent reaping a customer container. --- .../container-usage/src/heartbeat.test.ts | 54 ++++++ packages/container-usage/src/heartbeat.ts | 65 ++++--- .../cloudflare-agent-sandbox.test.ts | 28 ++++ .../cloudflare/cloudflare-agent-sandbox.ts | 14 +- .../src/container-usage-context.test.ts | 11 ++ .../src/container-usage-context.ts | 6 +- .../src/container-usage.test.ts | 158 +++++++++++++++--- .../cloud-agent-next/src/container-usage.ts | 113 +++++++------ .../src/kilo-facade/session-proxy.ts | 4 +- .../handlers/organization-membership.ts | 28 ++++ .../src/router/handlers/session-prepare.ts | 2 +- .../src/router/handlers/session-start.ts | 28 +--- 12 files changed, 380 insertions(+), 131 deletions(-) create mode 100644 services/cloud-agent-next/src/router/handlers/organization-membership.ts diff --git a/packages/container-usage/src/heartbeat.test.ts b/packages/container-usage/src/heartbeat.test.ts index 8baabef08f..333f42c340 100644 --- a/packages/container-usage/src/heartbeat.test.ts +++ b/packages/container-usage/src/heartbeat.test.ts @@ -250,6 +250,60 @@ describe('installBillingHeartbeat', () => { } }); + it('abandons local stopped-state retries after the hard ceiling', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + const recordStop = vi.fn(async () => { + throw new Error('meter unavailable'); + }); + 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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next', retry: { attempts: 1 } } + ); + const deleteSchedules = vi.fn(); + const controller = installBillingHeartbeat( + { + deleteSchedules, + getState: vi.fn(async () => ({ status: 'stopped' as const, lastChange: Date.now() })), + schedule: vi.fn() as Container['schedule'], + }, + { + client, + storage, + stopOnStoppedState: false, + stoppedStateGraceSeconds: 900, + stoppedStateAbandonSeconds: 3_600, + enforceBudgetStop: vi.fn(), + } + ); + + now.mockReturnValue(10_000); + await controller.billingHeartbeatTick(); + now.mockReturnValue(3_610_000); + await controller.billingHeartbeatTick(); + + expect(recordStop).toHaveBeenCalledOnce(); + expect(await getBillingContext(storage)).toBeUndefined(); + expect(deleteSchedules).toHaveBeenCalledWith(BILLING_HEARTBEAT_CALLBACK); + } finally { + now.mockRestore(); + } + }); + it('runs stop-delivery prerequisites before retrying a persisted stop', async () => { const storage = memoryStorage(); await storedContext(storage); diff --git a/packages/container-usage/src/heartbeat.ts b/packages/container-usage/src/heartbeat.ts index 2493a3e516..9ad5722561 100644 --- a/packages/container-usage/src/heartbeat.ts +++ b/packages/container-usage/src/heartbeat.ts @@ -14,6 +14,7 @@ import { export const BILLING_HEARTBEAT_CALLBACK = 'billingHeartbeatTick'; export const DEFAULT_BILLING_HEARTBEAT_SECONDS = 5 * 60; export const DEFAULT_STOPPED_STATE_GRACE_SECONDS = 15 * 60; +export const DEFAULT_STOPPED_STATE_ABANDON_SECONDS = 60 * 60; type BillingContainer = Pick; @@ -24,6 +25,7 @@ export type BillingHeartbeatDependencies = { /** Defer stopped-state closure to the container's authoritative onStop hook. */ stopOnStoppedState?: boolean; stoppedStateGraceSeconds?: number; + stoppedStateAbandonSeconds?: number; beforeHeartbeatDelivery?: (context: BillingContext) => Promise; beforeStopDelivery?: (context: BillingContext) => Promise; enforceBudgetStop: ( @@ -36,10 +38,13 @@ export type BillingHeartbeatDependencies = { export type BillingHeartbeatController = { scheduleHeartbeat: () => Promise; billingHeartbeatTick: (generation?: string) => Promise; - recordStop: (params: { - reason: 'exit' | 'runtime_signal' | 'activity_expired'; - exitCode?: number; - }) => Promise; + recordStop: ( + params: { + reason: 'exit' | 'runtime_signal' | 'activity_expired'; + exitCode?: number; + }, + usageEndedAtMs?: number + ) => Promise; cancelHeartbeat: () => void; persistStop: (params: { reason: 'exit' | 'runtime_signal' | 'activity_expired'; @@ -59,12 +64,17 @@ export function installBillingHeartbeat( const heartbeatSeconds = dependencies.heartbeatSeconds ?? DEFAULT_BILLING_HEARTBEAT_SECONDS; const stoppedStateGraceSeconds = dependencies.stoppedStateGraceSeconds ?? DEFAULT_STOPPED_STATE_GRACE_SECONDS; + const stoppedStateAbandonSeconds = + dependencies.stoppedStateAbandonSeconds ?? DEFAULT_STOPPED_STATE_ABANDON_SECONDS; if (heartbeatSeconds <= 0) { throw new Error('Billing heartbeat interval must be positive'); } if (stoppedStateGraceSeconds <= 0) { throw new Error('Stopped-state grace interval must be positive'); } + if (stoppedStateAbandonSeconds < stoppedStateGraceSeconds) { + throw new Error('Stopped-state abandon interval must not be shorter than its grace interval'); + } let lifecycleTail: Promise = Promise.resolve(); const runLifecycleExclusive = (operation: () => Promise): Promise => { @@ -108,6 +118,17 @@ export function installBillingHeartbeat( return true; }; + const computeStopSegment = (context: BillingContext, usageEndedAtMs: number) => { + if (context.pendingHeartbeat) return context.pendingHeartbeat; + const elapsedMs = Math.max(0, usageEndedAtMs - context.usageMeasuredAtMs); + const usageSinceLast = Math.floor(elapsedMs / 1_000); + return { + seq: context.nextSeq, + usageSinceLast, + measuredAtMs: context.usageMeasuredAtMs + usageSinceLast * 1_000, + }; + }; + const recordStopForGeneration = async ( params: Parameters[0], expectedGeneration?: string, @@ -119,13 +140,7 @@ export function installBillingHeartbeat( return undefined; } if (!context.pendingStop) { - const pendingHeartbeat = context.pendingHeartbeat; - const elapsedMs = Math.max(0, usageEndedAtMs - context.usageMeasuredAtMs); - const stopSegment = pendingHeartbeat ?? { - seq: context.nextSeq, - usageSinceLast: Math.floor(elapsedMs / 1_000), - measuredAtMs: context.usageMeasuredAtMs + Math.floor(elapsedMs / 1_000) * 1_000, - }; + const stopSegment = computeStopSegment(context, usageEndedAtMs); await updateBillingContext(dependencies.storage, { ...context, pendingStop: { ...params, ...stopSegment }, @@ -153,21 +168,15 @@ export function installBillingHeartbeat( return ack; }; - const recordStop: BillingHeartbeatController['recordStop'] = params => - runLifecycleExclusive(() => recordStopForGeneration(params)); + const recordStop: BillingHeartbeatController['recordStop'] = (params, usageEndedAtMs) => + runLifecycleExclusive(() => recordStopForGeneration(params, undefined, usageEndedAtMs)); const persistStop: BillingHeartbeatController['persistStop'] = params => runLifecycleExclusive(async () => { const context = await getBillingContext(dependencies.storage); if (!context) return undefined; if (context.pendingStop) return context; - const pendingHeartbeat = context.pendingHeartbeat; - const elapsedMs = Math.max(0, Date.now() - context.usageMeasuredAtMs); - const stopSegment = pendingHeartbeat ?? { - seq: context.nextSeq, - usageSinceLast: Math.floor(elapsedMs / 1_000), - measuredAtMs: context.usageMeasuredAtMs + Math.floor(elapsedMs / 1_000) * 1_000, - }; + const stopSegment = computeStopSegment(context, Date.now()); const updated = { ...context, pendingStop: { ...params, ...stopSegment } }; await updateBillingContext(dependencies.storage, updated); return updated; @@ -184,6 +193,14 @@ export function installBillingHeartbeat( try { await recordStopForGeneration(context.pendingStop, context.generation); } catch (error) { + if ( + context.stoppedObservedAtMs !== undefined && + Date.now() - context.stoppedObservedAtMs >= stoppedStateAbandonSeconds * 1_000 + ) { + cancelHeartbeat(); + await clearBillingContext(dependencies.storage); + return; + } await rescheduleIfCurrent(context); throw error; } @@ -222,6 +239,14 @@ export function installBillingHeartbeat( context.stoppedObservedAtMs ); } catch (error) { + if ( + context.stoppedObservedAtMs !== undefined && + Date.now() - context.stoppedObservedAtMs >= stoppedStateAbandonSeconds * 1_000 + ) { + cancelHeartbeat(); + await clearBillingContext(dependencies.storage); + return; + } await rescheduleIfCurrent(context); throw error; } diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts index b14ee1ecae..89a055530c 100644 --- a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts @@ -254,6 +254,34 @@ describe('deriveSetupEnvironment', () => { }); describe('CloudflareAgentSandbox', () => { + it('keeps default billing configuration when the sandbox resolver is injected', async () => { + const configureBilling = vi.fn().mockResolvedValue(undefined); + const renewActivityTimeout = vi.fn(); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => + ({ configureBilling, renewActivityTimeout }) as unknown as SandboxInstance, + }); + + await sandbox.keepAlive(); + + expect(configureBilling).toHaveBeenCalledOnce(); + expect(renewActivityTimeout).toHaveBeenCalledOnce(); + }); + + it('does not await shadow configuration before using an injected sandbox', async () => { + const configureBilling = vi.fn(() => new Promise(() => undefined)); + const renewActivityTimeout = vi.fn(); + const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { + resolveSandbox: () => ({ renewActivityTimeout }) as unknown as SandboxInstance, + configureBilling, + }); + + await expect(sandbox.keepAlive()).resolves.toBeUndefined(); + + expect(configureBilling).toHaveBeenCalledOnce(); + expect(renewActivityTimeout).toHaveBeenCalledOnce(); + }); + it('starts an ordinary bootstrap wrapper through the adapter', async () => { const bootstrapSession = {}; const createSession = vi.fn().mockResolvedValue(bootstrapSession); diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts index 6ca926f6ac..cc8d5685ae 100644 --- a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts @@ -202,9 +202,7 @@ export class CloudflareAgentSandbox implements AgentSandbox { this.sleep = dependencies.sleep ?? (ms => new Promise(resolve => setTimeout(resolve, ms))); this.stopObservationDelaysMs = dependencies.stopObservationDelaysMs ?? DEFAULT_STOP_OBSERVATION_DELAYS_MS; - this.configureBilling = - dependencies.configureBilling ?? - (dependencies.resolveSandbox ? async () => undefined : configureSandboxBillingInput); + this.configureBilling = dependencies.configureBilling ?? configureSandboxBillingInput; } private resolveSandboxId(): Promise { @@ -218,7 +216,7 @@ export class CloudflareAgentSandbox implements AgentSandbox { this.metadata.identity.sessionId, this.metadata.identity.botId, { - createdOnPlatform: this.metadata.identity.createdOnPlatform, + createdOnPlatform: this.metadata.identity.billingOrigin, } ); } @@ -228,7 +226,13 @@ export class CloudflareAgentSandbox implements AgentSandbox { private async getSandbox(options?: { sleepAfter?: number }): Promise { const sandboxId = await this.resolveSandboxId(); const sandbox = this.resolveSandbox(sandboxId, options); - await this.configureBilling(sandbox, buildSandboxBillingInput(this.metadata, sandboxId)); + void this.configureBilling(sandbox, buildSandboxBillingInput(this.metadata, sandboxId)).catch( + error => { + logger + .withFields({ error: error instanceof Error ? error.message : String(error) }) + .warn('Container usage shadow configuration deferred'); + } + ); return sandbox; } diff --git a/services/cloud-agent-next/src/container-usage-context.test.ts b/services/cloud-agent-next/src/container-usage-context.test.ts index 0b4186383e..a943f40455 100644 --- a/services/cloud-agent-next/src/container-usage-context.test.ts +++ b/services/cloud-agent-next/src/container-usage-context.test.ts @@ -183,6 +183,17 @@ describe('container usage context', () => { ).toThrow('Isolated sandbox billing requires session attribution'); }); + it('rejects unsupported isolated origins at the sandbox RPC boundary', () => { + expect(() => + assertSandboxBillingAllocation('SandboxSmall', { + subject: { type: 'user', id: 'user_isolated' }, + actor: { type: 'user', id: 'user_isolated' }, + sessionId: 'agent_1', + metadata: { allocation: 'isolated', origin: 'forged-origin' }, + }) + ).toThrow('Isolated sandbox billing origin is unsupported'); + }); + it('skips shadow configuration when a sandbox does not expose the metering RPC', async () => { await expect( configureSandboxBillingInput({} as SandboxInstance, { diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts index 74a90ed598..28d7036fcd 100644 --- a/services/cloud-agent-next/src/container-usage-context.ts +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -123,6 +123,10 @@ export function assertSandboxBillingAllocation( if (!input.sessionId || input.metadata?.allocation !== 'isolated') { throw new Error('Isolated sandbox billing requires session attribution'); } + const origin = input.metadata.origin; + if (origin === undefined || normalizedOrigin(origin) !== origin) { + throw new Error('Isolated sandbox billing origin is unsupported'); + } const allowedMetadata = new Set(['allocation', 'origin', 'repository_provider']); if (Object.keys(input.metadata).some(key => !allowedMetadata.has(key))) { throw new Error('Isolated sandbox billing metadata contains an unsupported field'); @@ -147,7 +151,7 @@ export async function configureSandboxBillingInput( return; } try { - await configureBilling.call(sandbox, input); + await (sandbox as MeteredSandboxInstance).configureBilling(input); } catch (error) { logger .withFields({ error: error instanceof Error ? error.message : String(error) }) diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index 40cc7aa10f..fd44233f90 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -1,5 +1,9 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { getBillingContext, type ContainerUsageRpcMethods } from '@kilocode/container-usage'; +import { + getBillingContext, + updateBillingContext, + type ContainerUsageRpcMethods, +} from '@kilocode/container-usage'; // oxlint-disable-next-line no-empty-object-type -- Matches the mocked Sandbox constructor. type SandboxDurableObjectState = DurableObjectState<{}>; @@ -64,12 +68,16 @@ import { MeteredSandbox } from './container-usage.js'; class MemoryStorage { private readonly values = new Map(); + failWrites = false; + hangReads = false; async get(key: string): Promise { + if (this.hangReads) return await new Promise(() => undefined); return this.values.get(key) as T | undefined; } async put(key: string, value: unknown): Promise { + if (this.failWrites) throw new Error('storage unavailable'); this.values.set(key, value); } @@ -109,10 +117,12 @@ type TestRuntime = MeteredSandbox & { function createSandbox(rpc = createRpc(), containerRunning = false) { const storage = new MemoryStorage(); + const shadowTasks: Promise[] = []; const ctx = { id: { toString: () => 'do-id' }, storage, container: { running: containerRunning }, + waitUntil: (promise: Promise) => shadowTasks.push(promise), } as unknown as SandboxDurableObjectState; class TestSandbox extends MeteredSandbox { protected readonly sandboxClassName = 'SandboxSmallContainment' as const; @@ -129,6 +139,7 @@ function createSandbox(rpc = createRpc(), containerRunning = false) { return { rpc, storage, + flushShadowTasks: () => Promise.all(shadowTasks), sandbox: new TestSandbox(ctx, { CONTAINER_USAGE_METER: rpc, } as never) as unknown as TestRuntime, @@ -139,7 +150,7 @@ const billingInput = { subject: { type: 'org' as const, id: 'org_1' }, actor: { type: 'user' as const, id: 'user_1' }, sessionId: 'agent_1', - metadata: { allocation: 'isolated' }, + metadata: { allocation: 'isolated', origin: 'cloud-agent' }, }; describe('MeteredSandbox', () => { @@ -149,7 +160,7 @@ describe('MeteredSandbox', () => { }); it('admits one start per physical generation and short-circuits active acquisition', async () => { - const { rpc, storage, sandbox } = createSandbox(); + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(1_000); await sandbox.configureBilling(billingInput); @@ -158,6 +169,7 @@ describe('MeteredSandbox', () => { sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); vi.mocked(rpc.recordStart).mockRejectedValue(new Error('meter unavailable')); await sandbox.configureBilling(billingInput); await sandbox.configureBilling(billingInput); @@ -169,18 +181,23 @@ describe('MeteredSandbox', () => { startEpochMs: 1_000, instanceId: 'SandboxSmallContainment:do-id', sku: 'cloud-agent-small-2026-07', - metadata: { allocation: 'isolated', container_class: 'SandboxSmallContainment' }, + metadata: { + allocation: 'isolated', + origin: 'cloud-agent', + container_class: 'SandboxSmallContainment', + }, }) ); expect((await getBillingContext(storage))?.measurementStarted).toBe(true); }); it('adopts a physical container that predates shadow metering', async () => { - const { rpc, storage, sandbox } = createSandbox(createRpc(), true); + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(createRpc(), true); vi.spyOn(Date, 'now').mockReturnValue(1_500); sandbox.mockState = { status: 'healthy' }; await sandbox.configureBilling(billingInput); + await flushShadowTasks(); expect(rpc.recordStart).toHaveBeenCalledOnce(); expect(await getBillingContext(storage)).toMatchObject({ @@ -200,11 +217,12 @@ describe('MeteredSandbox', () => { }); it('closes a missed-stop generation before the next physical start', async () => { - const { rpc, storage, sandbox } = createSandbox(); + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(1_750); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); const first = await getBillingContext(storage); sandbox.setPhysicalRunning(false); @@ -217,11 +235,33 @@ describe('MeteredSandbox', () => { ); await sandbox.onStart(); + await flushShadowTasks(); const second = await getBillingContext(storage); expect(second?.generation).not.toBe(first?.generation); expect(second?.startEpochMs).toBe(1_751); }); + it('uses the first stopped observation as the re-acquisition usage cutoff', async () => { + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + const active = await getBillingContext(storage); + if (!active) throw new Error('Expected active billing context'); + await updateBillingContext(storage, { ...active, stoppedObservedAtMs: 10_000 }); + + now.mockReturnValue(500_000); + sandbox.setPhysicalRunning(false); + sandbox.mockState = { status: 'stopped' }; + await sandbox.configureBilling(billingInput); + + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ usageSinceLast: 9, reason: 'runtime_signal' }) + ); + }); + it('keeps physical start non-fatal while retrying an unacknowledged shadow start', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStart) @@ -229,11 +269,12 @@ describe('MeteredSandbox', () => { .mockRejectedValueOnce(new Error('ack lost')) .mockRejectedValueOnce(new Error('ack lost')) .mockResolvedValue({ success: true, ack: ack() }); - const { sandbox, storage } = createSandbox(rpc); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await expect(sandbox.onStart()).resolves.toBeUndefined(); + await flushShadowTasks(); const context = await getBillingContext(storage); expect(context?.measurementStarted).toBe(true); expect(sandbox.superStarted).toBe(true); @@ -243,12 +284,25 @@ describe('MeteredSandbox', () => { expect(rpc.recordHeartbeat).toHaveBeenCalledOnce(); }); + it('does not await an unresolved meter during physical start', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStart).mockImplementation(() => new Promise(() => undefined)); + const { sandbox } = createSandbox(rpc); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + + await expect(sandbox.onStart()).resolves.toBeUndefined(); + + expect(sandbox.superStarted).toBe(true); + }); + it('keeps a delayed stop attached to the prior generation', async () => { - const { rpc, storage, sandbox } = createSandbox(); + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(2_000); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); const first = await getBillingContext(storage); sandbox.mockState = { status: 'stopped_with_code', exitCode: 17 }; @@ -256,10 +310,12 @@ describe('MeteredSandbox', () => { expect((await getBillingContext(storage))?.generation).toBe(first?.generation); await sandbox.onStop({ reason: 'exit', exitCode: 17 }); + await flushShadowTasks(); expect(await getBillingContext(storage)).toBeUndefined(); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); const second = await getBillingContext(storage); expect(second?.generation).not.toBe(first?.generation); @@ -271,13 +327,15 @@ describe('MeteredSandbox', () => { }); it('treats duplicate start callbacks as one physical generation', async () => { - const { rpc, storage, sandbox } = createSandbox(); + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); const first = await getBillingContext(storage); await sandbox.onStart(); + await flushShadowTasks(); const second = await getBillingContext(storage); expect(second?.generation).toBe(first?.generation); @@ -287,14 +345,16 @@ describe('MeteredSandbox', () => { it('does not admit a new physical generation until the prior stop is durable', async () => { const rpc = createRpc(); - const { sandbox, storage } = createSandbox(rpc); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); const first = await getBillingContext(storage); await sandbox.billingHeartbeatTick(first?.generation); vi.mocked(rpc.recordStop).mockRejectedValue(new Error('meter unavailable')); await sandbox.onStop({ reason: 'exit', exitCode: 1 }); + await flushShadowTasks(); expect((await getBillingContext(storage))?.pendingStop).toBeDefined(); await expect(sandbox.onStart()).resolves.toBeUndefined(); @@ -304,20 +364,23 @@ describe('MeteredSandbox', () => { }); it('defers activity-expiry closure until physical stop confirmation', async () => { - const { rpc, storage, sandbox } = createSandbox(); + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(3_000); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); const active = await getBillingContext(storage); await sandbox.onActivityExpired(); + await flushShadowTasks(); expect(sandbox.superActivityExpired).toBe(true); expect(rpc.recordStop).not.toHaveBeenCalled(); expect((await getBillingContext(storage))?.generation).toBe(active?.generation); await sandbox.onStop({ reason: 'exit', exitCode: 143 }); + await flushShadowTasks(); expect(rpc.recordStop).toHaveBeenCalledWith( expect.objectContaining({ reason: 'activity_expired', exitCode: 143 }) ); @@ -325,12 +388,14 @@ describe('MeteredSandbox', () => { }); it('preserves normal exit reason and exit code', async () => { - const { rpc, sandbox } = createSandbox(); + const { rpc, sandbox, flushShadowTasks } = createSandbox(); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); await sandbox.onStop({ reason: 'exit', exitCode: 42 }); + await flushShadowTasks(); expect(rpc.recordStop).toHaveBeenCalledWith( expect.objectContaining({ reason: 'exit', exitCode: 42 }) @@ -338,6 +403,20 @@ describe('MeteredSandbox', () => { expect(sandbox.superStopped).toBe(true); }); + it('does not await an unresolved meter during physical stop', async () => { + const rpc = createRpc(); + const { sandbox, flushShadowTasks } = createSandbox(rpc); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + vi.mocked(rpc.recordStop).mockImplementation(() => new Promise(() => undefined)); + + await expect(sandbox.onStop({ reason: 'exit', exitCode: 0 })).resolves.toBeUndefined(); + + expect(sandbox.superStopped).toBe(true); + }); + it('persists a failed stop and retries it without blocking SDK cleanup', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStop) @@ -345,12 +424,14 @@ describe('MeteredSandbox', () => { .mockRejectedValueOnce(new Error('postgres unavailable')) .mockRejectedValueOnce(new Error('postgres unavailable')) .mockResolvedValue(ack()); - const { sandbox, storage } = createSandbox(rpc); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); await expect(sandbox.onStop({ reason: 'exit', exitCode: 1 })).resolves.toBeUndefined(); + await flushShadowTasks(); expect(sandbox.superStopped).toBe(true); const pending = await getBillingContext(storage); expect(pending?.pendingStop).toMatchObject({ reason: 'exit', exitCode: 1 }); @@ -364,14 +445,16 @@ describe('MeteredSandbox', () => { it('persists authoritative stop details before recovering a missing start ack', async () => { const rpc = createRpc(); - const { sandbox, storage } = createSandbox(rpc); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); await sandbox.configureBilling(billingInput); sandbox.mockState = { status: 'healthy' }; await sandbox.onStart(); + await flushShadowTasks(); await storage.delete('container-usage:start-ack-generation:v1'); vi.mocked(rpc.recordStart).mockRejectedValue(new Error('meter unavailable')); await sandbox.onStop({ reason: 'exit', exitCode: 9 }); + await flushShadowTasks(); expect(await getBillingContext(storage)).toMatchObject({ pendingStop: { reason: 'exit', exitCode: 9 }, @@ -392,12 +475,47 @@ describe('MeteredSandbox', () => { ).rejects.toThrow(); }); - it('stops a runtime that somehow starts without trusted attribution', async () => { - const { sandbox } = createSandbox(); - await expect(sandbox.onStart()).rejects.toThrow( - 'Container started without pending billing attribution' - ); + it('does not stop a runtime that starts without shadow attribution', async () => { + const { sandbox, flushShadowTasks } = createSandbox(); + await expect(sandbox.onStart()).resolves.toBeUndefined(); + await flushShadowTasks(); expect(sandbox.superStarted).toBe(true); - expect(sandbox.superStopCalled).toBe(true); + expect(sandbox.superStopCalled).toBe(false); + }); + + it('does not fail physical start when persisted shadow state is corrupt', async () => { + const { sandbox, storage, flushShadowTasks } = createSandbox(); + await storage.put('container-usage:billing-context:v1', { invalid: true }); + + await expect(sandbox.onStart()).resolves.toBeUndefined(); + await flushShadowTasks(); + expect(sandbox.superStarted).toBe(true); + expect(sandbox.superStopCalled).toBe(false); + }); + + it('always performs real activity expiry when shadow storage fails', async () => { + const { sandbox, storage, flushShadowTasks } = createSandbox(); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + storage.failWrites = true; + + await expect(sandbox.onActivityExpired()).resolves.toBeUndefined(); + await flushShadowTasks(); + expect(sandbox.superActivityExpired).toBe(true); + }); + + it('does not await shadow persistence during real activity expiry', async () => { + const { sandbox, storage, flushShadowTasks } = createSandbox(); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + storage.hangReads = true; + + await expect(sandbox.onActivityExpired()).resolves.toBeUndefined(); + + expect(sandbox.superActivityExpired).toBe(true); }); }); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 252e19d4fd..539570fd43 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -50,6 +50,7 @@ export abstract class MeteredSandbox extends StockSandbox { private readonly usageClient: ContainerUsageClient; private readonly billingHeartbeat: BillingHeartbeatController; private billingLifecycleTail: Promise = Promise.resolve(); + private activityExpiryRequested = false; constructor(ctx: SandboxDurableObjectState, env: Env) { super(ctx, env); @@ -96,12 +97,15 @@ export abstract class MeteredSandbox extends StockSandbox { } const state = await this.getState(); try { - await this.billingHeartbeat.recordStop({ - reason: 'runtime_signal', - ...(state.status === 'stopped_with_code' && state.exitCode !== undefined - ? { exitCode: state.exitCode } - : {}), - }); + await this.billingHeartbeat.recordStop( + { + reason: 'runtime_signal', + ...(state.status === 'stopped_with_code' && state.exitCode !== undefined + ? { exitCode: state.exitCode } + : {}), + }, + active.stoppedObservedAtMs + ); await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); } catch (error) { await this.deferBillingDelivery(error, 'missed stop recovery'); @@ -119,12 +123,15 @@ export abstract class MeteredSandbox extends StockSandbox { return; } try { - await this.billingHeartbeat.recordStop({ - reason: 'runtime_signal', - ...(state.status === 'stopped_with_code' && state.exitCode !== undefined - ? { exitCode: state.exitCode } - : {}), - }); + await this.billingHeartbeat.recordStop( + { + reason: 'runtime_signal', + ...(state.status === 'stopped_with_code' && state.exitCode !== undefined + ? { exitCode: state.exitCode } + : {}), + }, + active.stoppedObservedAtMs + ); await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); } catch (error) { await this.deferBillingDelivery(error, 'unmeasured stop recovery'); @@ -140,7 +147,7 @@ export abstract class MeteredSandbox extends StockSandbox { override async onStart(): Promise { await super.onStart(); - await this.runBillingExclusive(async () => { + this.runShadowTask('start lifecycle', async () => { const previous = await getBillingContext(this.ctx.storage); if (previous) { if (previous.pendingStop) { @@ -168,8 +175,10 @@ export abstract class MeteredSandbox extends StockSandbox { const input = await this.getPendingAttribution(); if (!input) { - await super.stop(); - throw new Error('Container started without pending billing attribution'); + logger + .withFields({ sandboxClass: this.sandboxClassName }) + .warn('Container usage shadow start has no attribution'); + return; } await this.startBillingGeneration(input); @@ -177,49 +186,33 @@ export abstract class MeteredSandbox extends StockSandbox { } override async onStop(params?: ContainerStopParams): Promise { - try { - await this.runBillingExclusive(async () => { - const context = await getBillingContext(this.ctx.storage); - if (!context) return; - const requestedReason = await this.getPendingStopReason(context.generation); - try { - const pending = await this.billingHeartbeat.persistStop({ - reason: requestedReason ?? params?.reason ?? 'runtime_signal', - exitCode: params?.exitCode, - }); - if (!pending) return; - await this.ensureStartAcknowledged(pending); - await this.billingHeartbeat.recordStop( - pending.pendingStop ?? { - reason: requestedReason ?? params?.reason ?? 'runtime_signal', - exitCode: params?.exitCode, - } - ); - await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); - await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); - } catch (error) { - // recordStop persists its intent before delivery. Keep the heartbeat schedule alive so - // the durable intent retries without blocking the SDK's physical stop transition. - try { - await this.billingHeartbeat.scheduleHeartbeat(); - } catch { - // The persisted stop intent remains recoverable on the next sandbox acquisition. - } - logger - .withFields({ - error: error instanceof Error ? error.message : String(error), - sandboxClass: this.sandboxClassName, - }) - .warn('Container usage stop delivery deferred'); - } + await super.onStop(); + this.runShadowTask('stop lifecycle', async () => { + const context = await getBillingContext(this.ctx.storage); + if (!context) return; + const requestedReason = this.activityExpiryRequested + ? 'activity_expired' + : await this.getPendingStopReason(context.generation); + const pending = await this.billingHeartbeat.persistStop({ + reason: requestedReason ?? params?.reason ?? 'runtime_signal', + exitCode: params?.exitCode, }); - } finally { - await super.onStop(); - } + if (!pending) return; + await this.ensureStartAcknowledged(pending); + await this.billingHeartbeat.recordStop({ + reason: requestedReason ?? params?.reason ?? 'runtime_signal', + exitCode: params?.exitCode, + }); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); + this.activityExpiryRequested = false; + }); } override async onActivityExpired(): Promise { - await this.runBillingExclusive(async () => { + this.activityExpiryRequested = true; + await super.onActivityExpired(); + this.runShadowTask('activity expiry', async () => { const context = await getBillingContext(this.ctx.storage); if (context) { await this.ctx.storage.put(PENDING_STOP_REASON_STORAGE_KEY, { @@ -228,7 +221,6 @@ export abstract class MeteredSandbox extends StockSandbox { }); } }); - await super.onActivityExpired(); } private runBillingExclusive(operation: () => Promise): Promise { @@ -240,6 +232,13 @@ export abstract class MeteredSandbox extends StockSandbox { return result; } + private runShadowTask(operation: string, task: () => Promise): void { + const promise = this.runBillingExclusive(task).catch(error => { + this.logShadowFailure(error, operation); + }); + this.ctx.waitUntil(promise); + } + private async getPendingAttribution(): Promise { const stored = await this.ctx.storage.get(PENDING_ATTRIBUTION_STORAGE_KEY); return stored === undefined ? undefined : parseSandboxBillingInput(stored); @@ -278,6 +277,10 @@ export abstract class MeteredSandbox extends StockSandbox { // A later sandbox acquisition retries persisted shadow state. } } + this.logShadowFailure(error, operation); + } + + private logShadowFailure(error: unknown, operation: string): void { logger .withFields({ error: error instanceof Error ? error.message : String(error), diff --git a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts index eed0a1ed3b..dd1865d68d 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts @@ -76,7 +76,7 @@ export async function resolveLiveWrapperTarget(params: { metadata.identity.sessionId, metadata.identity.botId, { - createdOnPlatform: metadata.identity.createdOnPlatform, + createdOnPlatform: metadata.identity.billingOrigin, } )); @@ -86,7 +86,7 @@ export async function resolveLiveWrapperTarget(params: { }), sandboxId ); - await configureSandboxBilling(sandbox, metadata, sandboxId); + void configureSandboxBilling(sandbox, metadata, sandboxId); const wrapperInfo = await findWrapperForSession(sandbox, sessionId); if (!wrapperInfo) { return null; diff --git a/services/cloud-agent-next/src/router/handlers/organization-membership.ts b/services/cloud-agent-next/src/router/handlers/organization-membership.ts new file mode 100644 index 0000000000..5f39f38988 --- /dev/null +++ b/services/cloud-agent-next/src/router/handlers/organization-membership.ts @@ -0,0 +1,28 @@ +import { TRPCError } from '@trpc/server'; +import type { WorkerDb } from '@kilocode/db/client'; +import { organization_memberships } from '@kilocode/db/schema'; +import { and, eq } from 'drizzle-orm'; + +export async function assertOrganizationMembership( + db: WorkerDb, + userId: string, + organizationId: string +): Promise { + const [membership] = await db + .select({ id: organization_memberships.id }) + .from(organization_memberships) + .where( + and( + eq(organization_memberships.organization_id, organizationId), + eq(organization_memberships.kilo_user_id, userId) + ) + ) + .limit(1); + + if (!membership) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'You do not have access to this organization', + }); + } +} diff --git a/services/cloud-agent-next/src/router/handlers/session-prepare.ts b/services/cloud-agent-next/src/router/handlers/session-prepare.ts index c5c674c23e..4bfc46c15f 100644 --- a/services/cloud-agent-next/src/router/handlers/session-prepare.ts +++ b/services/cloud-agent-next/src/router/handlers/session-prepare.ts @@ -43,7 +43,7 @@ import type { SessionProfileBundle } from '../../session-profile.js'; import type { SessionCreateRequest } from '../../session/session-requests.js'; import { assertKiloModelAvailable } from '../../model-validation.js'; import { assertBitbucketRepositoryAccessBeforeSessionCreation } from '../../session/validate-repository-access.js'; -import { assertOrganizationMembership } from './session-start.js'; +import { assertOrganizationMembership } from './organization-membership.js'; type SessionPrepareHandlers = { prepareSession: typeof prepareSessionHandler; diff --git a/services/cloud-agent-next/src/router/handlers/session-start.ts b/services/cloud-agent-next/src/router/handlers/session-start.ts index ea5124ba81..16fe0aa496 100644 --- a/services/cloud-agent-next/src/router/handlers/session-start.ts +++ b/services/cloud-agent-next/src/router/handlers/session-start.ts @@ -11,10 +11,7 @@ * session ownership state is created. */ import { protectedProcedure } from '../auth.js'; -import { TRPCError } from '@trpc/server'; -import { organization_memberships } from '@kilocode/db/schema'; import type { WorkerDb } from '@kilocode/db/client'; -import { and, eq } from 'drizzle-orm'; import { logger, withLogTags } from '../../logger.js'; import { getPgDb } from '../../db/pg.js'; import type * as z from 'zod'; @@ -28,6 +25,7 @@ import { import type { SessionCreateRequest } from '../../session/session-requests.js'; import { assertKiloModelAvailable } from '../../model-validation.js'; import { assertBitbucketRepositoryAccessBeforeSessionCreation } from '../../session/validate-repository-access.js'; +import { assertOrganizationMembership } from './organization-membership.js'; type SessionStartHandlers = { start: typeof startSessionHandler; @@ -91,30 +89,6 @@ function startInputToSessionCreateRequest( }; } -export async function assertOrganizationMembership( - db: WorkerDb, - userId: string, - organizationId: string -): Promise { - const [membership] = await db - .select({ id: organization_memberships.id }) - .from(organization_memberships) - .where( - and( - eq(organization_memberships.organization_id, organizationId), - eq(organization_memberships.kilo_user_id, userId) - ) - ) - .limit(1); - - if (!membership) { - throw new TRPCError({ - code: 'FORBIDDEN', - message: 'You do not have access to this organization', - }); - } -} - const startSessionHandler = protectedProcedure .input(StartSessionInput) .output(StartSessionOutput) From abeb18897ff45e686c953399a4a08602fee88706 Mon Sep 17 00:00:00 2001 From: syn Date: Thu, 23 Jul 2026 21:24:55 -0500 Subject: [PATCH 06/11] fix(billing): align instance IDs with Cloudflare Record the exact Cloudflare Sandbox ID as the usage instanceId so reconciliation can pass recorder values directly to containersUsageAdaptiveGroups and match the Containers dashboard. Keep the opaque Durable Object ID as diagnostic metadata, retain trusted origin for isolated workloads, and remove allocation and repository-provider metadata that are not needed for billing reconciliation. --- .../src/container-usage-context.test.ts | 48 +++++++++------ .../src/container-usage-context.ts | 59 +++++++++++++------ .../src/container-usage.test.ts | 46 +++++++++++++-- .../cloud-agent-next/src/container-usage.ts | 13 +++- 4 files changed, 124 insertions(+), 42 deletions(-) diff --git a/services/cloud-agent-next/src/container-usage-context.test.ts b/services/cloud-agent-next/src/container-usage-context.test.ts index a943f40455..9070d2d411 100644 --- a/services/cloud-agent-next/src/container-usage-context.test.ts +++ b/services/cloud-agent-next/src/container-usage-context.test.ts @@ -72,7 +72,7 @@ describe('container usage context', () => { }, }, ])('derives trusted $name attribution', ({ identity, expected }) => { - expect(buildSandboxBillingInput(metadata(identity), 'ses-isolated')).toMatchObject(expected); + expect(buildSandboxBillingInput(metadata(identity), 'ses-abcdef')).toMatchObject(expected); }); it('keeps isolated metadata bounded and normalizes automation origins', () => { @@ -83,16 +83,13 @@ describe('container usage context', () => { orgId: 'org_security', billingOrigin: 'security-remediation', }), - 'crv-isolated' + 'crv-abcdef' ); expect(input).toMatchObject({ + sandboxId: 'crv-abcdef', sessionId: 'agent_security', - metadata: { - allocation: 'isolated', - origin: 'security-remediation', - repository_provider: 'github', - }, + metadata: { origin: 'security-remediation' }, }); expect(JSON.stringify(input)).not.toContain('Kilo-Org/cloud'); }); @@ -105,7 +102,7 @@ describe('container usage context', () => { orgId: 'org_shared', billingOrigin: 'security-agent', }), - 'org-shared' + 'org-abcdef' ); const second = buildSandboxBillingInput( metadata({ @@ -114,14 +111,14 @@ describe('container usage context', () => { orgId: 'org_shared', billingOrigin: 'cloud-agent-web', }), - 'org-shared' + 'org-abcdef' ); expect(first).toEqual(second); expect(first).toEqual({ + sandboxId: 'org-abcdef', subject: { type: 'org', id: 'org_shared' }, actor: { type: 'user', id: 'user_shared' }, - metadata: { allocation: 'shared' }, }); }); @@ -132,7 +129,7 @@ describe('container usage context', () => { userId: 'user_unknown', billingOrigin: 'attacker-controlled-value', }), - 'dind-isolated' + 'dind-abcdef' ); expect(input.metadata?.origin).toBe('other'); }); @@ -157,7 +154,7 @@ describe('container usage context', () => { createdOnPlatform: 'security-remediation', billingOrigin: 'cloud-agent', }), - 'ses-isolated' + 'ses-abcdef' ); expect(input.metadata?.origin).toBe('cloud-agent'); }); @@ -165,10 +162,11 @@ describe('container usage context', () => { it('rejects session attribution and extra metadata for shared sandboxes', () => { expect(() => assertSandboxBillingAllocation('Sandbox', { + sandboxId: 'org-abcdef', subject: { type: 'user', id: 'user_shared' }, actor: { type: 'user', id: 'user_shared' }, sessionId: 'agent_leak', - metadata: { allocation: 'shared', origin: 'cloud-agent' }, + metadata: { origin: 'cloud-agent' }, }) ).toThrow('Shared sandbox billing cannot contain session attribution'); }); @@ -176,9 +174,10 @@ describe('container usage context', () => { it('requires bounded isolated attribution for non-shared sandbox classes', () => { expect(() => assertSandboxBillingAllocation('SandboxSmall', { + sandboxId: 'ses-abcdef', subject: { type: 'user', id: 'user_isolated' }, actor: { type: 'user', id: 'user_isolated' }, - metadata: { allocation: 'isolated' }, + metadata: { origin: 'cloud-agent' }, }) ).toThrow('Isolated sandbox billing requires session attribution'); }); @@ -186,21 +185,35 @@ describe('container usage context', () => { it('rejects unsupported isolated origins at the sandbox RPC boundary', () => { expect(() => assertSandboxBillingAllocation('SandboxSmall', { + sandboxId: 'ses-abcdef', subject: { type: 'user', id: 'user_isolated' }, actor: { type: 'user', id: 'user_isolated' }, sessionId: 'agent_1', - metadata: { allocation: 'isolated', origin: 'forged-origin' }, + metadata: { origin: 'forged-origin' }, }) ).toThrow('Isolated sandbox billing origin is unsupported'); }); + it('rejects a sandbox ID that does not match the concrete container class', () => { + expect(() => + assertSandboxBillingAllocation('SandboxDIND', { + sandboxId: 'ses-abcdef', + subject: { type: 'user', id: 'user_1' }, + actor: { type: 'user', id: 'user_1' }, + sessionId: 'agent_1', + metadata: { origin: 'cloud-agent' }, + }) + ).toThrow('SandboxDIND billing requires a dind- sandbox ID'); + }); + it('skips shadow configuration when a sandbox does not expose the metering RPC', async () => { await expect( configureSandboxBillingInput({} as SandboxInstance, { + sandboxId: 'ses-abcdef', subject: { type: 'user', id: 'user_1' }, actor: { type: 'user', id: 'user_1' }, sessionId: 'agent_1', - metadata: { allocation: 'isolated' }, + metadata: { origin: 'cloud-agent' }, }) ).resolves.toBeUndefined(); }); @@ -209,10 +222,11 @@ describe('container usage context', () => { const configureBilling = vi.fn().mockRejectedValue(new Error('meter unavailable')); await expect( configureSandboxBillingInput({ configureBilling } as unknown as SandboxInstance, { + sandboxId: 'ses-abcdef', subject: { type: 'user', id: 'user_1' }, actor: { type: 'user', id: 'user_1' }, sessionId: 'agent_1', - metadata: { allocation: 'isolated' }, + metadata: { origin: 'cloud-agent' }, }) ).resolves.toBeUndefined(); expect(configureBilling).toHaveBeenCalledOnce(); diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts index 28d7036fcd..4111880539 100644 --- a/services/cloud-agent-next/src/container-usage-context.ts +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -20,13 +20,24 @@ export const SANDBOX_USAGE_SKUS = { } as const; export type SandboxClassName = keyof typeof SANDBOX_USAGE_SKUS; -export type SandboxBillingInput = Omit; +export type SandboxBillingInput = Omit & { + sandboxId: SandboxId; +}; export type MeteredSandboxInstance = SandboxInstance & { configureBilling(input: unknown): Promise; }; const sandboxBillingInputEnvelopeSchema = z .object({ + sandboxId: z + .string() + .min(1) + .max(63) + .refine( + value => /^(ses|crv|dind|org|usr|bot|ubt)-[0-9a-f]+$/.test(value) || value.includes('__'), + 'Invalid sandboxId format' + ) + .transform(value => value as SandboxId), subject: billingSubjectSchema, actor: billingActorSchema, onBehalfOf: billingSubjectSchema.optional(), @@ -79,30 +90,28 @@ export function buildSandboxBillingInput( const isolated = isIsolatedSandbox(sandboxId); return { + sandboxId, subject, actor, ...(actor.type === 'bot' ? { onBehalfOf: subject } : {}), ...(isolated ? { sessionId: metadata.identity.sessionId } : {}), - metadata: isolated - ? { - allocation: 'isolated', - origin: normalizedOrigin(metadata.identity.billingOrigin), - ...(metadata.repository ? { repository_provider: metadata.repository.type } : {}), - } - : { allocation: 'shared' }, + ...(isolated + ? { metadata: { origin: normalizedOrigin(metadata.identity.billingOrigin) } } + : {}), }; } export function parseSandboxBillingInput(input: unknown): SandboxBillingInput { const parsed = sandboxBillingInputEnvelopeSchema.parse(input); + const { sandboxId, ...usageInput } = parsed; const validated = usageContextSchema.parse({ service: 'cloud-agent-next', instanceId: 'validation', sku: 'validation', - ...parsed, + ...usageInput, }); const { service: _service, instanceId: _instanceId, sku: _sku, ...billingInput } = validated; - return billingInput; + return { sandboxId, ...billingInput }; } export function assertSandboxBillingAllocation( @@ -111,24 +120,40 @@ export function assertSandboxBillingAllocation( ): void { const shared = sandboxClassName === 'Sandbox' || sandboxClassName === 'SandboxContainment'; if (shared) { - if (input.sessionId !== undefined || input.metadata?.allocation !== 'shared') { + if (input.sessionId !== undefined) { throw new Error('Shared sandbox billing cannot contain session attribution'); } - if (Object.keys(input.metadata).some(key => key !== 'allocation')) { - throw new Error('Shared sandbox billing metadata must contain only allocation'); + if (!/^(org|usr|bot|ubt)-/.test(input.sandboxId) && !input.sandboxId.includes('__')) { + throw new Error('Shared sandbox billing requires a shared sandbox ID'); + } + if (input.metadata !== undefined && Object.keys(input.metadata).length > 0) { + throw new Error('Shared sandbox billing cannot contain metadata'); } return; } - if (!input.sessionId || input.metadata?.allocation !== 'isolated') { + const expectedPrefix = + sandboxClassName === 'SandboxDIND' + ? 'dind-' + : sandboxClassName === 'SandboxSmall' || sandboxClassName === 'SandboxSmallContainment' + ? 'ses-' + : 'crv-'; + if (!input.sandboxId.startsWith(expectedPrefix)) { + throw new Error(`${sandboxClassName} billing requires a ${expectedPrefix} sandbox ID`); + } + if (!input.sessionId) { throw new Error('Isolated sandbox billing requires session attribution'); } - const origin = input.metadata.origin; + const metadata = input.metadata; + if (!metadata) { + throw new Error('Isolated sandbox billing origin is unsupported'); + } + const origin = metadata.origin; if (origin === undefined || normalizedOrigin(origin) !== origin) { throw new Error('Isolated sandbox billing origin is unsupported'); } - const allowedMetadata = new Set(['allocation', 'origin', 'repository_provider']); - if (Object.keys(input.metadata).some(key => !allowedMetadata.has(key))) { + const allowedMetadata = new Set(['origin']); + if (Object.keys(metadata).some(key => !allowedMetadata.has(key))) { throw new Error('Isolated sandbox billing metadata contains an unsupported field'); } } diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index fd44233f90..823645671a 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -115,7 +115,11 @@ type TestRuntime = MeteredSandbox & { billingHeartbeatTick(generation?: string): Promise; }; -function createSandbox(rpc = createRpc(), containerRunning = false) { +function createSandbox( + rpc = createRpc(), + containerRunning = false, + sandboxClassName: 'SandboxSmallContainment' | 'SandboxDIND' = 'SandboxSmallContainment' +) { const storage = new MemoryStorage(); const shadowTasks: Promise[] = []; const ctx = { @@ -125,7 +129,7 @@ function createSandbox(rpc = createRpc(), containerRunning = false) { waitUntil: (promise: Promise) => shadowTasks.push(promise), } as unknown as SandboxDurableObjectState; class TestSandbox extends MeteredSandbox { - protected readonly sandboxClassName = 'SandboxSmallContainment' as const; + protected readonly sandboxClassName = sandboxClassName; setPhysicalRunning(running: boolean): void { if (this.ctx.container) { @@ -147,10 +151,11 @@ function createSandbox(rpc = createRpc(), containerRunning = false) { } const billingInput = { + sandboxId: 'ses-abcdef' as const, subject: { type: 'org' as const, id: 'org_1' }, actor: { type: 'user' as const, id: 'user_1' }, sessionId: 'agent_1', - metadata: { allocation: 'isolated', origin: 'cloud-agent' }, + metadata: { origin: 'cloud-agent' }, }; describe('MeteredSandbox', () => { @@ -179,12 +184,12 @@ describe('MeteredSandbox', () => { expect(rpc.recordStart).toHaveBeenCalledWith( expect.objectContaining({ startEpochMs: 1_000, - instanceId: 'SandboxSmallContainment:do-id', + instanceId: 'ses-abcdef', sku: 'cloud-agent-small-2026-07', metadata: { - allocation: 'isolated', origin: 'cloud-agent', container_class: 'SandboxSmallContainment', + durable_object_id: 'do-id', }, }) ); @@ -206,6 +211,36 @@ describe('MeteredSandbox', () => { }); }); + it('records a DIND instance using its Cloudflare instance ID', async () => { + const rpc = createRpc(); + const { sandbox, flushShadowTasks } = createSandbox(rpc, false, 'SandboxDIND'); + await sandbox.configureBilling({ + ...billingInput, + sandboxId: 'dind-abcdef', + metadata: { origin: 'cloud-agent' }, + }); + sandbox.mockState = { status: 'healthy' }; + + await sandbox.onStart(); + await flushShadowTasks(); + + expect(rpc.recordStart).toHaveBeenCalledWith( + expect.objectContaining({ + service: 'cloud-agent-next', + instanceId: 'dind-abcdef', + sku: 'cloud-agent-dind-2026-07', + subject: { type: 'org', id: 'org_1' }, + actor: { type: 'user', id: 'user_1' }, + sessionId: 'agent_1', + metadata: { + container_class: 'SandboxDIND', + durable_object_id: 'do-id', + origin: 'cloud-agent', + }, + }) + ); + }); + it('does not adopt stale healthy state when no physical container is running', async () => { const { rpc, storage, sandbox } = createSandbox(); sandbox.mockState = { status: 'healthy' }; @@ -238,6 +273,7 @@ describe('MeteredSandbox', () => { await flushShadowTasks(); const second = await getBillingContext(storage); expect(second?.generation).not.toBe(first?.generation); + expect(second?.instanceId).toBe('ses-abcdef'); expect(second?.startEpochMs).toBe(1_751); }); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 539570fd43..5f916ad957 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -306,11 +306,18 @@ export abstract class MeteredSandbox extends StockSandbox { const startEpochMs = Math.max(Date.now(), previousStartEpochMs + 1); await this.ctx.storage.put(LAST_START_EPOCH_STORAGE_KEY, startEpochMs); const context = await setBillingContext(this.ctx.storage, { - ...input, + subject: input.subject, + actor: input.actor, + ...(input.onBehalfOf ? { onBehalfOf: input.onBehalfOf } : {}), + ...(input.sessionId ? { sessionId: input.sessionId } : {}), service: SERVICE, - instanceId: `${this.sandboxClassName}:${this.ctx.id.toString()}`, + instanceId: input.sandboxId, sku: SANDBOX_USAGE_SKUS[this.sandboxClassName], - metadata: { ...input.metadata, container_class: this.sandboxClassName }, + metadata: { + container_class: this.sandboxClassName, + durable_object_id: this.ctx.id.toString(), + ...(input.metadata?.origin ? { origin: input.metadata.origin } : {}), + }, startEpochMs, } satisfies UsageContext & { startEpochMs: number }); await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); From 77e262b818834419b2f206b6778ab24227ce27ac Mon Sep 17 00:00:00 2001 From: syn Date: Thu, 23 Jul 2026 21:38:52 -0500 Subject: [PATCH 07/11] test(billing): update shared sandbox payload expectation Shared container billing now carries the Cloudflare sandbox ID directly and intentionally omits allocation metadata. Keep the containment test aligned with that recorder contract while preserving its ordering assertions for billing dispatch, containment activation, and workspace probing. --- .../cloudflare/cloudflare-agent-sandbox.test.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts index 89a055530c..720966e845 100644 --- a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts @@ -396,11 +396,11 @@ describe('CloudflareAgentSandbox', () => { expect(setOutboundHandler).toHaveBeenCalledWith('managedScm'); expect(configureBilling).toHaveBeenCalledWith( expect.anything(), - expect.objectContaining({ + { + sandboxId: 'usr-shared', subject: { type: 'org', id: 'org_cloudflare' }, actor: { type: 'user', id: 'user_cloudflare' }, - metadata: { allocation: 'shared' }, - }) + } ); expect(configureBilling.mock.invocationCallOrder[0]).toBeLessThan( setOutboundHandler.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY From fdf40753e64d4277d4d4315e0f5c5f0402c42f53 Mon Sep 17 00:00:00 2001 From: syn Date: Fri, 24 Jul 2026 10:11:50 -0500 Subject: [PATCH 08/11] style(billing): format shared sandbox assertion --- .../cloudflare/cloudflare-agent-sandbox.test.ts | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts index 720966e845..2f4ef845fb 100644 --- a/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts +++ b/services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts @@ -394,14 +394,11 @@ describe('CloudflareAgentSandbox', () => { ); expect(setOutboundHandler).toHaveBeenCalledWith('managedScm'); - expect(configureBilling).toHaveBeenCalledWith( - expect.anything(), - { - sandboxId: 'usr-shared', - subject: { type: 'org', id: 'org_cloudflare' }, - actor: { type: 'user', id: 'user_cloudflare' }, - } - ); + expect(configureBilling).toHaveBeenCalledWith(expect.anything(), { + sandboxId: 'usr-shared', + subject: { type: 'org', id: 'org_cloudflare' }, + actor: { type: 'user', id: 'user_cloudflare' }, + }); expect(configureBilling.mock.invocationCallOrder[0]).toBeLessThan( setOutboundHandler.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY ); From 0cd4ce2fe6b06b72f07d47b01b88fe924ae9861f Mon Sep 17 00:00:00 2001 From: syn Date: Fri, 24 Jul 2026 14:57:59 -0500 Subject: [PATCH 09/11] fix(billing): fence container usage recovery Stop failed prior-generation recovery from falling through into replacement generation creation, consume activity-expiry intent at the physical stop boundary, and reject isolated-prefixed legacy IDs for shared sandbox classes. --- .../src/container-usage-context.test.ts | 10 +++++ .../src/container-usage-context.ts | 4 +- .../src/container-usage.test.ts | 39 +++++++++++++++++++ .../cloud-agent-next/src/container-usage.ts | 6 ++- 4 files changed, 56 insertions(+), 3 deletions(-) diff --git a/services/cloud-agent-next/src/container-usage-context.test.ts b/services/cloud-agent-next/src/container-usage-context.test.ts index 9070d2d411..bd9a0deb95 100644 --- a/services/cloud-agent-next/src/container-usage-context.test.ts +++ b/services/cloud-agent-next/src/container-usage-context.test.ts @@ -171,6 +171,16 @@ describe('container usage context', () => { ).toThrow('Shared sandbox billing cannot contain session attribution'); }); + it('rejects isolated-prefixed legacy IDs for shared sandbox classes', () => { + expect(() => + assertSandboxBillingAllocation('Sandbox', { + sandboxId: 'ses-abcdef__legacy', + subject: { type: 'user', id: 'user_shared' }, + actor: { type: 'user', id: 'user_shared' }, + }) + ).toThrow('Shared sandbox billing requires a shared sandbox ID'); + }); + it('requires bounded isolated attribution for non-shared sandbox classes', () => { expect(() => assertSandboxBillingAllocation('SandboxSmall', { diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts index 4111880539..02fcebbecf 100644 --- a/services/cloud-agent-next/src/container-usage-context.ts +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -123,7 +123,9 @@ export function assertSandboxBillingAllocation( if (input.sessionId !== undefined) { throw new Error('Shared sandbox billing cannot contain session attribution'); } - if (!/^(org|usr|bot|ubt)-/.test(input.sandboxId) && !input.sandboxId.includes('__')) { + const legacySharedId = + !/^(ses|crv|dind)-/.test(input.sandboxId) && input.sandboxId.includes('__'); + if (!/^(org|usr|bot|ubt)-/.test(input.sandboxId) && !legacySharedId) { throw new Error('Shared sandbox billing requires a shared sandbox ID'); } if (input.metadata !== undefined && Object.keys(input.metadata).length > 0) { diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index 823645671a..3f4bc72ec0 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -399,6 +399,26 @@ describe('MeteredSandbox', () => { expect(rpc.recordStart).toHaveBeenCalledOnce(); }); + it('does not replace an unmeasured generation when stop recovery fails', async () => { + const rpc = createRpc(); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + const first = await getBillingContext(storage); + if (!first) throw new Error('Expected active billing context'); + await updateBillingContext(storage, { ...first, measurementStarted: false }); + sandbox.setPhysicalRunning(true); + sandbox.mockState = { status: 'stopped' }; + vi.mocked(rpc.recordStop).mockRejectedValue(new Error('meter unavailable')); + + await sandbox.configureBilling({ ...billingInput, sessionId: 'agent_2' }); + + expect((await getBillingContext(storage))?.generation).toBe(first.generation); + expect(rpc.recordStart).toHaveBeenCalledOnce(); + }); + it('defers activity-expiry closure until physical stop confirmation', async () => { const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(3_000); @@ -423,6 +443,25 @@ describe('MeteredSandbox', () => { expect(await getBillingContext(storage)).toBeUndefined(); }); + it('does not carry an activity-expiry reason across generations without context', async () => { + const { rpc, sandbox, flushShadowTasks } = createSandbox(); + await sandbox.onActivityExpired(); + await flushShadowTasks(); + await sandbox.onStop({ reason: 'exit', exitCode: 0 }); + await flushShadowTasks(); + + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + await sandbox.onStop({ reason: 'exit', exitCode: 42 }); + await flushShadowTasks(); + + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ reason: 'exit', exitCode: 42 }) + ); + }); + it('preserves normal exit reason and exit code', async () => { const { rpc, sandbox, flushShadowTasks } = createSandbox(); await sandbox.configureBilling(billingInput); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 5f916ad957..93b8a82130 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -135,6 +135,7 @@ export abstract class MeteredSandbox extends StockSandbox { await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); } catch (error) { await this.deferBillingDelivery(error, 'unmeasured stop recovery'); + return; } } @@ -187,10 +188,12 @@ export abstract class MeteredSandbox extends StockSandbox { override async onStop(params?: ContainerStopParams): Promise { await super.onStop(); + const activityExpiryRequested = this.activityExpiryRequested; + this.activityExpiryRequested = false; this.runShadowTask('stop lifecycle', async () => { const context = await getBillingContext(this.ctx.storage); if (!context) return; - const requestedReason = this.activityExpiryRequested + const requestedReason = activityExpiryRequested ? 'activity_expired' : await this.getPendingStopReason(context.generation); const pending = await this.billingHeartbeat.persistStop({ @@ -205,7 +208,6 @@ export abstract class MeteredSandbox extends StockSandbox { }); await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); - this.activityExpiryRequested = false; }); } From 87f30b7d1eaf62aef5a2c4f930292ec199bbf1b2 Mon Sep 17 00:00:00 2001 From: syn Date: Mon, 27 Jul 2026 12:17:36 -0500 Subject: [PATCH 10/11] fix(billing): preserve container generation boundaries Namespace recorder services by concrete Sandbox class while retaining raw Cloudflare instance IDs, and preserve authoritative isolated routing targets for code-review and DIND sessions. Capture physical stop cutoffs, deliver pending heartbeats before final stop remainders, use container transition timestamps for missed stops, bound authoritative stop retries, and start already-running replacements after prior generations close. --- .../container-usage/src/heartbeat.test.ts | 170 +++++++++++++++++- packages/container-usage/src/heartbeat.ts | 77 ++++++-- .../src/container-usage.test.ts | 78 +++++++- .../cloud-agent-next/src/container-usage.ts | 55 ++++-- .../cloud-agent-next/src/sandbox-outbound.ts | 28 ++- .../src/session-prepare.test.ts | 22 +-- .../src/session/session-registration.ts | 19 +- 7 files changed, 385 insertions(+), 64 deletions(-) diff --git a/packages/container-usage/src/heartbeat.test.ts b/packages/container-usage/src/heartbeat.test.ts index 333f42c340..4e764ed7c4 100644 --- a/packages/container-usage/src/heartbeat.test.ts +++ b/packages/container-usage/src/heartbeat.test.ts @@ -1,7 +1,12 @@ import { describe, expect, it, vi } from 'vitest'; import type { Container } from '@cloudflare/containers'; import { ContainerUsageClient } from './client'; -import { getBillingContext, setBillingContext, type BillingContextStorage } from './context'; +import { + getBillingContext, + setBillingContext, + updateBillingContext, + type BillingContextStorage, +} from './context'; import type { ContainerUsageRpcMethods, HeartbeatAck, RecordAck } from './contracts'; import { BILLING_HEARTBEAT_CALLBACK, installBillingHeartbeat } from './heartbeat'; @@ -250,6 +255,169 @@ describe('installBillingHeartbeat', () => { } }); + it('uses the container transition time when stopped state is detected late', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + const recordStop = vi.fn(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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next' } + ); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(async () => ({ status: 'stopped' as const, lastChange: 2_000 })), + schedule: vi.fn() as Container['schedule'], + }, + { + client, + storage, + stopOnStoppedState: false, + stoppedStateGraceSeconds: 900, + enforceBudgetStop: vi.fn(), + } + ); + + now.mockReturnValue(10_000); + await controller.billingHeartbeatTick(); + expect((await getBillingContext(storage))?.stoppedObservedAtMs).toBe(2_000); + now.mockReturnValue(902_000); + await controller.billingHeartbeatTick(); + + expect(recordStop).toHaveBeenCalledWith( + expect.objectContaining({ reason: 'runtime_signal', usageSinceLast: 1 }) + ); + } finally { + now.mockRestore(); + } + }); + + it('delivers a pending heartbeat before the final stop remainder', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + const context = await getBillingContext(storage); + if (!context) throw new Error('Expected billing context'); + await updateBillingContext(storage, { + ...context, + pendingHeartbeat: { seq: 1, usageSinceLast: 4, measuredAtMs: 5_000 }, + }); + const recordHeartbeat = vi.fn(async input => ({ + intervalId: `${input.instanceId}:${input.startEpochMs}`, + durable: 'pg', + dedup: false, + budget: { verdict: 'continue' }, + })); + const recordStop = vi.fn(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, + recordStop, + }, + { service: 'cloud-agent-next' } + ); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(), + schedule: vi.fn() as Container['schedule'], + }, + { client, storage, enforceBudgetStop: vi.fn() } + ); + + await controller.persistStop({ reason: 'exit', exitCode: 0 }, 7_000); + await controller.recordStop({ reason: 'exit', exitCode: 0 }, 7_000); + + expect(recordHeartbeat).toHaveBeenCalledWith( + expect.objectContaining({ seq: 1, usageSinceLast: 4 }) + ); + expect(recordStop).toHaveBeenCalledWith( + expect.objectContaining({ seq: 2, usageSinceLast: 2, reason: 'exit' }) + ); + expect(recordHeartbeat.mock.invocationCallOrder[0]).toBeLessThan( + recordStop.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); + } finally { + now.mockRestore(); + } + }); + + it('abandons a failed authoritative stop after its captured cutoff', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + const recordStop = vi.fn(async () => { + throw new Error('meter unavailable'); + }); + 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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next', retry: { attempts: 1 } } + ); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(), + schedule: vi.fn() as Container['schedule'], + }, + { + client, + storage, + stoppedStateAbandonSeconds: 3_600, + enforceBudgetStop: vi.fn(), + } + ); + await controller.persistStop({ reason: 'exit', exitCode: 1 }, 10_000); + + now.mockReturnValue(3_610_000); + await controller.billingHeartbeatTick(); + + expect(recordStop).toHaveBeenCalledOnce(); + expect(await getBillingContext(storage)).toBeUndefined(); + } finally { + now.mockRestore(); + } + }); + it('abandons local stopped-state retries after the hard ceiling', async () => { const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); try { diff --git a/packages/container-usage/src/heartbeat.ts b/packages/container-usage/src/heartbeat.ts index 9ad5722561..48215b217d 100644 --- a/packages/container-usage/src/heartbeat.ts +++ b/packages/container-usage/src/heartbeat.ts @@ -28,6 +28,7 @@ export type BillingHeartbeatDependencies = { stoppedStateAbandonSeconds?: number; beforeHeartbeatDelivery?: (context: BillingContext) => Promise; beforeStopDelivery?: (context: BillingContext) => Promise; + onGenerationClosed?: (context: BillingContext) => void; enforceBudgetStop: ( budget: BudgetVerdict, expected: { generation: string; startEpochMs: number } @@ -46,10 +47,13 @@ export type BillingHeartbeatController = { usageEndedAtMs?: number ) => Promise; cancelHeartbeat: () => void; - persistStop: (params: { - reason: 'exit' | 'runtime_signal' | 'activity_expired'; - exitCode?: number; - }) => Promise; + persistStop: ( + params: { + reason: 'exit' | 'runtime_signal' | 'activity_expired'; + exitCode?: number; + }, + usageEndedAtMs?: number + ) => Promise; }; function contextForHeartbeat(context: BillingContext) { @@ -118,8 +122,17 @@ export function installBillingHeartbeat( return true; }; + const contextAfterPendingHeartbeat = (context: BillingContext): BillingContext => { + if (!context.pendingHeartbeat) return context; + return { + ...context, + nextSeq: context.pendingHeartbeat.seq + 1, + usageMeasuredAtMs: context.pendingHeartbeat.measuredAtMs, + pendingHeartbeat: undefined, + }; + }; + const computeStopSegment = (context: BillingContext, usageEndedAtMs: number) => { - if (context.pendingHeartbeat) return context.pendingHeartbeat; const elapsedMs = Math.max(0, usageEndedAtMs - context.usageMeasuredAtMs); const usageSinceLast = Math.floor(elapsedMs / 1_000); return { @@ -129,6 +142,32 @@ export function installBillingHeartbeat( }; }; + const acknowledgePendingHeartbeat = async (context: BillingContext): Promise => { + const pendingHeartbeat = context.pendingHeartbeat; + if (!pendingHeartbeat) return context; + await dependencies.beforeHeartbeatDelivery?.(context); + await dependencies.client.recordHeartbeat({ + instanceId: context.instanceId, + startEpochMs: context.startEpochMs, + seq: pendingHeartbeat.seq, + usageSinceLast: pendingHeartbeat.usageSinceLast, + context: contextForHeartbeat(context), + }); + const current = await getBillingContext(dependencies.storage); + if (!current || !isSameBillingGeneration(current, context)) return context; + if ( + !current.pendingHeartbeat || + current.pendingHeartbeat.seq !== pendingHeartbeat.seq || + current.pendingHeartbeat.measuredAtMs !== pendingHeartbeat.measuredAtMs || + current.pendingHeartbeat.usageSinceLast !== pendingHeartbeat.usageSinceLast + ) { + return current; + } + const updated = contextAfterPendingHeartbeat(current); + await updateBillingContext(dependencies.storage, updated); + return updated; + }; + const recordStopForGeneration = async ( params: Parameters[0], expectedGeneration?: string, @@ -140,15 +179,17 @@ export function installBillingHeartbeat( return undefined; } if (!context.pendingStop) { - const stopSegment = computeStopSegment(context, usageEndedAtMs); + const stopSegment = computeStopSegment(contextAfterPendingHeartbeat(context), usageEndedAtMs); await updateBillingContext(dependencies.storage, { ...context, + stoppedObservedAtMs: context.stoppedObservedAtMs ?? usageEndedAtMs, pendingStop: { ...params, ...stopSegment }, }); const current = await getBillingContext(dependencies.storage); if (!current || !isSameBillingGeneration(current, context)) return undefined; context = current; } + context = await acknowledgePendingHeartbeat(context); const stopIntent = context.pendingStop; if (!stopIntent) throw new Error('Billing stop intent was not persisted'); await dependencies.beforeStopDelivery?.(context); @@ -165,19 +206,27 @@ export function installBillingHeartbeat( if (!current || !isSameBillingGeneration(current, context)) return ack; cancelHeartbeat(); await clearBillingContext(dependencies.storage); + dependencies.onGenerationClosed?.(context); return ack; }; const recordStop: BillingHeartbeatController['recordStop'] = (params, usageEndedAtMs) => runLifecycleExclusive(() => recordStopForGeneration(params, undefined, usageEndedAtMs)); - const persistStop: BillingHeartbeatController['persistStop'] = params => + const persistStop: BillingHeartbeatController['persistStop'] = ( + params, + usageEndedAtMs = Date.now() + ) => runLifecycleExclusive(async () => { const context = await getBillingContext(dependencies.storage); if (!context) return undefined; if (context.pendingStop) return context; - const stopSegment = computeStopSegment(context, Date.now()); - const updated = { ...context, pendingStop: { ...params, ...stopSegment } }; + const stopSegment = computeStopSegment(contextAfterPendingHeartbeat(context), usageEndedAtMs); + const updated = { + ...context, + stoppedObservedAtMs: context.stoppedObservedAtMs ?? usageEndedAtMs, + pendingStop: { ...params, ...stopSegment }, + }; await updateBillingContext(dependencies.storage, updated); return updated; }); @@ -199,6 +248,7 @@ export function installBillingHeartbeat( ) { cancelHeartbeat(); await clearBillingContext(dependencies.storage); + dependencies.onGenerationClosed?.(context); return; } await rescheduleIfCurrent(context); @@ -219,7 +269,13 @@ export function installBillingHeartbeat( context = currentAfterState; if (state.status === 'stopped' || state.status === 'stopped_with_code') { if (dependencies.stopOnStoppedState === false) { - const stoppedObservedAtMs = context.stoppedObservedAtMs ?? Date.now(); + const observedAtMs = Date.now(); + const stateLastChange = Number.isFinite(state.lastChange) ? state.lastChange : observedAtMs; + const stoppedObservedAtMs = + context.stoppedObservedAtMs ?? + (stateLastChange >= 0 && stateLastChange <= observedAtMs + ? stateLastChange + : observedAtMs); if (context.stoppedObservedAtMs === undefined) { context = { ...context, stoppedObservedAtMs }; await updateBillingContext(dependencies.storage, context); @@ -245,6 +301,7 @@ export function installBillingHeartbeat( ) { cancelHeartbeat(); await clearBillingContext(dependencies.storage); + dependencies.onGenerationClosed?.(context); return; } await rescheduleIfCurrent(context); diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index 3f4bc72ec0..f681156879 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -118,7 +118,11 @@ type TestRuntime = MeteredSandbox & { function createSandbox( rpc = createRpc(), containerRunning = false, - sandboxClassName: 'SandboxSmallContainment' | 'SandboxDIND' = 'SandboxSmallContainment' + sandboxClassName: + | 'Sandbox' + | 'SandboxContainment' + | 'SandboxSmallContainment' + | 'SandboxDIND' = 'SandboxSmallContainment' ) { const storage = new MemoryStorage(); const shadowTasks: Promise[] = []; @@ -129,7 +133,9 @@ function createSandbox( waitUntil: (promise: Promise) => shadowTasks.push(promise), } as unknown as SandboxDurableObjectState; class TestSandbox extends MeteredSandbox { - protected readonly sandboxClassName = sandboxClassName; + protected get sandboxClassName() { + return sandboxClassName; + } setPhysicalRunning(running: boolean): void { if (this.ctx.container) { @@ -226,7 +232,7 @@ describe('MeteredSandbox', () => { expect(rpc.recordStart).toHaveBeenCalledWith( expect.objectContaining({ - service: 'cloud-agent-next', + service: 'cloud-agent-next-sandbox-dind', instanceId: 'dind-abcdef', sku: 'cloud-agent-dind-2026-07', subject: { type: 'org', id: 'org_1' }, @@ -241,6 +247,40 @@ describe('MeteredSandbox', () => { ); }); + it('uses distinct recorder services for standard and containment namespaces', async () => { + const standardRpc = createRpc(); + const containmentRpc = createRpc(); + const standard = createSandbox(standardRpc, false, 'Sandbox'); + const containment = createSandbox(containmentRpc, false, 'SandboxContainment'); + const sharedInput = { + sandboxId: 'usr-abcdef' as const, + subject: { type: 'user' as const, id: 'user_1' }, + actor: { type: 'user' as const, id: 'user_1' }, + }; + await standard.sandbox.configureBilling(sharedInput); + await containment.sandbox.configureBilling(sharedInput); + standard.sandbox.mockState = { status: 'healthy' }; + containment.sandbox.mockState = { status: 'healthy' }; + + await standard.sandbox.onStart(); + await containment.sandbox.onStart(); + await standard.flushShadowTasks(); + await containment.flushShadowTasks(); + + expect(standardRpc.recordStart).toHaveBeenCalledWith( + expect.objectContaining({ + service: 'cloud-agent-next-sandbox', + instanceId: 'usr-abcdef', + }) + ); + expect(containmentRpc.recordStart).toHaveBeenCalledWith( + expect.objectContaining({ + service: 'cloud-agent-next-sandbox-containment', + instanceId: 'usr-abcdef', + }) + ); + }); + it('does not adopt stale healthy state when no physical container is running', async () => { const { rpc, storage, sandbox } = createSandbox(); sandbox.mockState = { status: 'healthy' }; @@ -286,11 +326,11 @@ describe('MeteredSandbox', () => { await flushShadowTasks(); const active = await getBillingContext(storage); if (!active) throw new Error('Expected active billing context'); - await updateBillingContext(storage, { ...active, stoppedObservedAtMs: 10_000 }); now.mockReturnValue(500_000); sandbox.setPhysicalRunning(false); sandbox.mockState = { status: 'stopped' }; + Object.assign(sandbox.mockState, { lastChange: 10_000 }); await sandbox.configureBilling(billingInput); expect(rpc.recordStop).toHaveBeenCalledWith( @@ -399,6 +439,36 @@ describe('MeteredSandbox', () => { expect(rpc.recordStart).toHaveBeenCalledOnce(); }); + it('starts a running replacement after the prior pending stop is recovered', async () => { + const rpc = createRpc(); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + const first = await getBillingContext(storage); + if (!first) throw new Error('Expected first billing generation'); + + vi.mocked(rpc.recordStop).mockRejectedValue(new Error('meter unavailable')); + await sandbox.onStop({ reason: 'exit', exitCode: 1 }); + await flushShadowTasks(); + await sandbox.configureBilling({ ...billingInput, sessionId: 'agent_2' }); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + expect((await getBillingContext(storage))?.generation).toBe(first.generation); + + vi.mocked(rpc.recordStop).mockResolvedValue(ack()); + await sandbox.billingHeartbeatTick(first.generation); + await flushShadowTasks(); + await flushShadowTasks(); + + const second = await getBillingContext(storage); + expect(second?.generation).not.toBe(first.generation); + expect(second?.sessionId).toBe('agent_2'); + expect(rpc.recordStart).toHaveBeenCalledTimes(2); + }); + it('does not replace an unmeasured generation when stop recovery fails', async () => { const rpc = createRpc(); const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 93b8a82130..8a265786d0 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -44,8 +44,21 @@ function startInputFromContext(context: BillingContext): ClientRecordStartInput return { ...usage, startEpochMs: context.startEpochMs }; } +function usageServiceForSandboxClass(sandboxClassName: SandboxClassName): string { + const suffix = sandboxClassName.replace(/([a-z])([A-Z])/g, '$1-$2').toLowerCase(); + return `${SERVICE}-${suffix}`; +} + +function stoppedAtFromState(state: { lastChange: number }, observedAtMs = Date.now()): number { + return Number.isFinite(state.lastChange) && + state.lastChange >= 0 && + state.lastChange <= observedAtMs + ? state.lastChange + : observedAtMs; +} + export abstract class MeteredSandbox extends StockSandbox { - protected abstract readonly sandboxClassName: SandboxClassName; + protected abstract get sandboxClassName(): SandboxClassName; private readonly usageClient: ContainerUsageClient; private readonly billingHeartbeat: BillingHeartbeatController; @@ -54,15 +67,14 @@ export abstract class MeteredSandbox extends StockSandbox { constructor(ctx: SandboxDurableObjectState, env: Env) { super(ctx, env); - this.usageClient = createContainerUsageClient(env.CONTAINER_USAGE_METER, { - service: SERVICE, - }); + this.usageClient = this.createUsageClient(env); this.billingHeartbeat = installBillingHeartbeat(this, { client: this.usageClient, storage: this.ctx.storage, stopOnStoppedState: false, beforeHeartbeatDelivery: context => this.ensureStartAcknowledged(context), beforeStopDelivery: context => this.ensureStartAcknowledged(context), + onGenerationClosed: () => this.schedulePendingGenerationIfRunning(), // The meter currently returns only `continue`; shadow mode must not enforce future verdicts. enforceBudgetStop: async () => { throw new Error('Container budget enforcement is disabled in shadow mode'); @@ -70,6 +82,12 @@ export abstract class MeteredSandbox extends StockSandbox { }); } + private createUsageClient(env: Env): ContainerUsageClient { + return createContainerUsageClient(env.CONTAINER_USAGE_METER, { + service: usageServiceForSandboxClass(this.sandboxClassName), + }); + } + async configureBilling(input: unknown): Promise { const parsed = parseSandboxBillingInput(input); assertSandboxBillingAllocation(this.sandboxClassName, parsed); @@ -96,6 +114,7 @@ export abstract class MeteredSandbox extends StockSandbox { return; } const state = await this.getState(); + const stoppedAtMs = active.stoppedObservedAtMs ?? stoppedAtFromState(state); try { await this.billingHeartbeat.recordStop( { @@ -104,7 +123,7 @@ export abstract class MeteredSandbox extends StockSandbox { ? { exitCode: state.exitCode } : {}), }, - active.stoppedObservedAtMs + stoppedAtMs ); await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); } catch (error) { @@ -122,6 +141,7 @@ export abstract class MeteredSandbox extends StockSandbox { await this.admitAndScheduleBestEffort(active); return; } + const stoppedAtMs = active.stoppedObservedAtMs ?? stoppedAtFromState(state); try { await this.billingHeartbeat.recordStop( { @@ -130,7 +150,7 @@ export abstract class MeteredSandbox extends StockSandbox { ? { exitCode: state.exitCode } : {}), }, - active.stoppedObservedAtMs + stoppedAtMs ); await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); } catch (error) { @@ -187,6 +207,7 @@ export abstract class MeteredSandbox extends StockSandbox { } override async onStop(params?: ContainerStopParams): Promise { + const stoppedAtMs = Date.now(); await super.onStop(); const activityExpiryRequested = this.activityExpiryRequested; this.activityExpiryRequested = false; @@ -196,10 +217,13 @@ export abstract class MeteredSandbox extends StockSandbox { const requestedReason = activityExpiryRequested ? 'activity_expired' : await this.getPendingStopReason(context.generation); - const pending = await this.billingHeartbeat.persistStop({ - reason: requestedReason ?? params?.reason ?? 'runtime_signal', - exitCode: params?.exitCode, - }); + const pending = await this.billingHeartbeat.persistStop( + { + reason: requestedReason ?? params?.reason ?? 'runtime_signal', + exitCode: params?.exitCode, + }, + stoppedAtMs + ); if (!pending) return; await this.ensureStartAcknowledged(pending); await this.billingHeartbeat.recordStop({ @@ -241,6 +265,15 @@ export abstract class MeteredSandbox extends StockSandbox { this.ctx.waitUntil(promise); } + private schedulePendingGenerationIfRunning(): void { + this.runShadowTask('replacement generation', async () => { + if (this.ctx.container?.running !== true) return; + if (await getBillingContext(this.ctx.storage)) return; + const input = await this.getPendingAttribution(); + if (input) await this.startBillingGeneration(input); + }); + } + private async getPendingAttribution(): Promise { const stored = await this.ctx.storage.get(PENDING_ATTRIBUTION_STORAGE_KEY); return stored === undefined ? undefined : parseSandboxBillingInput(stored); @@ -312,7 +345,7 @@ export abstract class MeteredSandbox extends StockSandbox { actor: input.actor, ...(input.onBehalfOf ? { onBehalfOf: input.onBehalfOf } : {}), ...(input.sessionId ? { sessionId: input.sessionId } : {}), - service: SERVICE, + service: usageServiceForSandboxClass(this.sandboxClassName), instanceId: input.sandboxId, sku: SANDBOX_USAGE_SKUS[this.sandboxClassName], metadata: { diff --git a/services/cloud-agent-next/src/sandbox-outbound.ts b/services/cloud-agent-next/src/sandbox-outbound.ts index 3db5b807c9..a871d667a7 100644 --- a/services/cloud-agent-next/src/sandbox-outbound.ts +++ b/services/cloud-agent-next/src/sandbox-outbound.ts @@ -662,31 +662,41 @@ const managedScmOutboundHandlers = { }; export class Sandbox extends MeteredSandbox { - protected readonly sandboxClassName: SandboxClassName = 'Sandbox'; + protected get sandboxClassName(): SandboxClassName { + return 'Sandbox'; + } enableInternet = true; interceptHttps = false; } export class SandboxSmall extends MeteredSandbox { - protected readonly sandboxClassName: SandboxClassName = 'SandboxSmall'; + protected get sandboxClassName(): SandboxClassName { + return 'SandboxSmall'; + } enableInternet = true; interceptHttps = false; } export class SandboxDIND extends MeteredSandbox { - protected readonly sandboxClassName: SandboxClassName = 'SandboxDIND'; + protected get sandboxClassName(): SandboxClassName { + return 'SandboxDIND'; + } enableInternet = true; interceptHttps = false; } export class SandboxCodeReview extends MeteredSandbox { - protected readonly sandboxClassName: SandboxClassName = 'SandboxCodeReview'; + protected get sandboxClassName(): SandboxClassName { + return 'SandboxCodeReview'; + } enableInternet = true; interceptHttps = false; } export class SandboxContainment extends Sandbox { - protected override readonly sandboxClassName: SandboxClassName = 'SandboxContainment'; + protected override get sandboxClassName(): SandboxClassName { + return 'SandboxContainment'; + } interceptHttps = true; } // Assignment (not a static class field) so it invokes the inherited Container.outboundHandlers @@ -695,13 +705,17 @@ export class SandboxContainment extends Sandbox { SandboxContainment.outboundHandlers = managedScmOutboundHandlers; export class SandboxSmallContainment extends SandboxSmall { - protected override readonly sandboxClassName: SandboxClassName = 'SandboxSmallContainment'; + protected override get sandboxClassName(): SandboxClassName { + return 'SandboxSmallContainment'; + } interceptHttps = true; } SandboxSmallContainment.outboundHandlers = managedScmOutboundHandlers; export class SandboxCodeReviewContainment extends SandboxCodeReview { - protected override readonly sandboxClassName: SandboxClassName = 'SandboxCodeReviewContainment'; + protected override get sandboxClassName(): SandboxClassName { + return 'SandboxCodeReviewContainment'; + } interceptHttps = true; } SandboxCodeReviewContainment.outboundHandlers = managedScmOutboundHandlers; diff --git a/services/cloud-agent-next/src/session-prepare.test.ts b/services/cloud-agent-next/src/session-prepare.test.ts index 92168d1345..a9d084069e 100644 --- a/services/cloud-agent-next/src/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session-prepare.test.ts @@ -453,6 +453,10 @@ describe('prepareSession endpoint', () => { }); it('registers full lazy-prep metadata in one DO call', async () => { + generateSandboxRoutingTargetMock.mockResolvedValueOnce({ + kind: 'isolated', + sandboxId: 'crv-abcdef', + }); const doStub = createMockDOStub(); const orgId = 'f47ac10b-58cc-4372-a567-0e02b2c3d479'; const caller = appRouter.createCaller( @@ -544,13 +548,14 @@ describe('prepareSession endpoint', () => { target: { url: 'https://example.com/callback' }, }, workspace: { - sandboxId: 'sb-test-123', + sandboxId: 'crv-abcdef', sandboxProvider: 'cloudflare', shallow: true, credentialContainment: { github: true, gitlab: false, kilocode: false }, }, }) ); + expect(selectSandboxForNewSessionMock).not.toHaveBeenCalled(); }); it('rejects organization attribution when the internal caller user is not a member', async () => { @@ -845,10 +850,6 @@ describe('prepareSession endpoint', () => { kind: 'isolated', sandboxId: 'dind-abcdef', }); - selectSandboxForNewSessionMock.mockResolvedValueOnce({ - sandboxId: 'dind-abcdef', - provider: 'cloudflare', - }); const doStub = createMockDOStub(); const caller = appRouter.createCaller(createInternalApiContext({ doStub })); @@ -872,16 +873,7 @@ describe('prepareSession endpoint', () => { createdOnPlatform: undefined, } ); - expect(selectSandboxForNewSessionMock).toHaveBeenCalledWith( - expect.objectContaining({ - env: expect.any(Object), - orgId: undefined, - userId: 'test-user-123', - sessionId: 'agent_12345678-1234-1234-1234-123456789abc', - botId: undefined, - devcontainer: true, - }) - ); + expect(selectSandboxForNewSessionMock).not.toHaveBeenCalled(); expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( expect.objectContaining({ workspace: { diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index 90f47d6dfb..9b4f659ad5 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -27,12 +27,7 @@ import { recordCloudAgentSandboxIdentity, recordCloudAgentSessionFailure, } from '../telemetry/session-reports.js'; -import { - generateSandboxRoutingTarget, - isOrgInList, - selectSandboxForNewSession, - type SandboxSelection, -} from '../sandbox-id.js'; +import { generateSandboxRoutingTarget, isOrgInList, type SandboxSelection } from '../sandbox-id.js'; import { resolveSharedSandboxAssignment } from '../shared-sandbox-route.js'; import { generateKiloSessionId } from '../utils/kilo-session-id.js'; import { createMessageId } from './message-id.js'; @@ -200,16 +195,8 @@ async function allocateNewSession( ...(assignment.suffix ? { suffix: assignment.suffix } : {}), }; } else { - const selection = await selectSandboxForNewSession({ - env: ctx.env, - orgId, - userId: ctx.userId, - sessionId: cloudAgentSessionId, - botId: ctx.botId, - devcontainer: input.runtime?.devcontainer, - }); - sandboxId = selection.sandboxId; - sandboxProvider = selection.provider; + sandboxId = target.sandboxId; + sandboxProvider = 'cloudflare'; } } catch (error) { await recordCloudAgentSessionFailure( From 386dee0efcf86d520b14c9ce4c5329b9b3f8cf83 Mon Sep 17 00:00:00 2001 From: syn Date: Mon, 27 Jul 2026 12:48:01 -0500 Subject: [PATCH 11/11] fix(billing): fence stale generation races --- .../container-usage/src/heartbeat.test.ts | 191 ++++++++++++++++++ packages/container-usage/src/heartbeat.ts | 29 ++- .../src/container-usage.test.ts | 20 ++ .../cloud-agent-next/src/container-usage.ts | 7 +- 4 files changed, 237 insertions(+), 10 deletions(-) diff --git a/packages/container-usage/src/heartbeat.test.ts b/packages/container-usage/src/heartbeat.test.ts index 4e764ed7c4..b3f1427d66 100644 --- a/packages/container-usage/src/heartbeat.test.ts +++ b/packages/container-usage/src/heartbeat.test.ts @@ -418,6 +418,67 @@ describe('installBillingHeartbeat', () => { } }); + it('does not abandon a replacement generation after a pending-stop failure', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + let replacement: Awaited> | undefined; + const recordStop = vi.fn(async () => { + replacement = await setBillingContext(storage, { + service: 'cloud-agent-next', + instanceId: 'instance-1', + startEpochMs: 456, + sku: 'cloud-agent-next:Sandbox', + subject: { type: 'user', id: 'user-1' }, + actor: { type: 'user', id: 'user-1' }, + }); + throw new Error('meter unavailable'); + }); + const deleteSchedules = vi.fn(); + const onGenerationClosed = vi.fn(); + const controller = installBillingHeartbeat( + { + deleteSchedules, + getState: vi.fn(), + schedule: vi.fn() as Container['schedule'], + }, + { + 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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next', retry: { attempts: 1 } } + ), + storage, + stoppedStateAbandonSeconds: 3_600, + onGenerationClosed, + enforceBudgetStop: vi.fn(), + } + ); + await controller.persistStop({ reason: 'exit', exitCode: 1 }, 10_000); + + now.mockReturnValue(3_610_000); + await controller.billingHeartbeatTick(); + + expect(await getBillingContext(storage)).toEqual(replacement); + expect(deleteSchedules).not.toHaveBeenCalled(); + expect(onGenerationClosed).not.toHaveBeenCalled(); + } finally { + now.mockRestore(); + } + }); + it('abandons local stopped-state retries after the hard ceiling', async () => { const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); try { @@ -472,6 +533,70 @@ describe('installBillingHeartbeat', () => { } }); + it('does not abandon a replacement generation after a stopped-state failure', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + let replacement: Awaited> | undefined; + const recordStop = vi.fn(async () => { + replacement = await setBillingContext(storage, { + service: 'cloud-agent-next', + instanceId: 'instance-1', + startEpochMs: 456, + sku: 'cloud-agent-next:Sandbox', + subject: { type: 'user', id: 'user-1' }, + actor: { type: 'user', id: 'user-1' }, + }); + throw new Error('meter unavailable'); + }); + const deleteSchedules = vi.fn(); + const onGenerationClosed = vi.fn(); + const controller = installBillingHeartbeat( + { + deleteSchedules, + getState: vi.fn(async () => ({ status: 'stopped' as const, lastChange: Date.now() })), + schedule: vi.fn() as Container['schedule'], + }, + { + 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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next', retry: { attempts: 1 } } + ), + storage, + stopOnStoppedState: false, + stoppedStateGraceSeconds: 900, + stoppedStateAbandonSeconds: 3_600, + onGenerationClosed, + enforceBudgetStop: vi.fn(), + } + ); + + now.mockReturnValue(10_000); + await controller.billingHeartbeatTick(); + now.mockReturnValue(3_610_000); + await controller.billingHeartbeatTick(); + + expect(await getBillingContext(storage)).toEqual(replacement); + expect(deleteSchedules).toHaveBeenCalledTimes(1); + expect(onGenerationClosed).not.toHaveBeenCalled(); + } finally { + now.mockRestore(); + } + }); + it('runs stop-delivery prerequisites before retrying a persisted stop', async () => { const storage = memoryStorage(); await storedContext(storage); @@ -598,6 +723,72 @@ describe('installBillingHeartbeat', () => { } }); + it('does not let late scheduling restore an abandoned generation', async () => { + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + try { + const storage = memoryStorage(); + await storedContext(storage); + let resolveSchedule = (): void => undefined; + const schedule = vi.fn( + () => + new Promise(resolve => { + resolveSchedule = () => resolve({}); + }) + ); + const recordStop = vi.fn(async () => { + throw new Error('meter unavailable'); + }); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(), + schedule: schedule as Container['schedule'], + }, + { + 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: 'continue' }, + }), + recordStop, + }, + { service: 'cloud-agent-next', retry: { attempts: 1 } } + ), + storage, + stoppedStateAbandonSeconds: 3_600, + enforceBudgetStop: vi.fn(), + } + ); + await controller.persistStop({ reason: 'exit', exitCode: 1 }, 1_000); + const scheduling = controller.scheduleHeartbeat(); + await vi.waitFor(() => expect(schedule).toHaveBeenCalledOnce()); + + now.mockReturnValue(3_601_000); + await controller.billingHeartbeatTick(); + const replacement = await setBillingContext(storage, { + service: 'cloud-agent-next', + instanceId: 'instance-1', + startEpochMs: 456, + sku: 'cloud-agent-next:Sandbox', + subject: { type: 'user', id: 'user-1' }, + actor: { type: 'user', id: 'user-1' }, + }); + resolveSchedule(); + await scheduling; + + expect(await getBillingContext(storage)).toEqual(replacement); + } finally { + now.mockRestore(); + } + }); + it('carries subsecond remainder into the next acknowledged heartbeat', async () => { const now = vi.spyOn(Date, 'now'); try { diff --git a/packages/container-usage/src/heartbeat.ts b/packages/container-usage/src/heartbeat.ts index 48215b217d..7dd84d9ef9 100644 --- a/packages/container-usage/src/heartbeat.ts +++ b/packages/container-usage/src/heartbeat.ts @@ -110,7 +110,13 @@ export function installBillingHeartbeat( startedContext.generation ); if (!context.measurementStarted) { - await updateBillingContext(dependencies.storage, startedContext); + const current = await getBillingContext(dependencies.storage); + if (!current || !isSameBillingGeneration(current, context)) return; + await updateBillingContext(dependencies.storage, { + ...current, + measurementStarted: true, + usageMeasuredAtMs: startedContext.usageMeasuredAtMs, + }); } }; @@ -132,6 +138,15 @@ export function installBillingHeartbeat( }; }; + const abandonIfCurrent = async (expected: BillingContext): Promise => { + const current = await getBillingContext(dependencies.storage); + if (!current || !isSameBillingGeneration(current, expected)) return false; + await clearBillingContext(dependencies.storage); + cancelHeartbeat(); + dependencies.onGenerationClosed?.(expected); + return true; + }; + const computeStopSegment = (context: BillingContext, usageEndedAtMs: number) => { const elapsedMs = Math.max(0, usageEndedAtMs - context.usageMeasuredAtMs); const usageSinceLast = Math.floor(elapsedMs / 1_000); @@ -246,12 +261,10 @@ export function installBillingHeartbeat( context.stoppedObservedAtMs !== undefined && Date.now() - context.stoppedObservedAtMs >= stoppedStateAbandonSeconds * 1_000 ) { - cancelHeartbeat(); - await clearBillingContext(dependencies.storage); - dependencies.onGenerationClosed?.(context); + await abandonIfCurrent(context); return; } - await rescheduleIfCurrent(context); + if (!(await rescheduleIfCurrent(context))) return; throw error; } return; @@ -299,12 +312,10 @@ export function installBillingHeartbeat( context.stoppedObservedAtMs !== undefined && Date.now() - context.stoppedObservedAtMs >= stoppedStateAbandonSeconds * 1_000 ) { - cancelHeartbeat(); - await clearBillingContext(dependencies.storage); - dependencies.onGenerationClosed?.(context); + await abandonIfCurrent(context); return; } - await rescheduleIfCurrent(context); + if (!(await rescheduleIfCurrent(context))) return; throw error; } return; diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index f681156879..3e105f11d9 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -338,6 +338,26 @@ describe('MeteredSandbox', () => { ); }); + it('uses observation time when re-acquiring against stale healthy SDK state', async () => { + const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + await sandbox.configureBilling(billingInput); + sandbox.mockState = { status: 'healthy' }; + await sandbox.onStart(); + await flushShadowTasks(); + expect((await getBillingContext(storage))?.measurementStarted).toBe(true); + + now.mockReturnValue(500_000); + sandbox.setPhysicalRunning(false); + sandbox.mockState = { status: 'healthy' }; + Object.assign(sandbox.mockState, { lastChange: 10_000 }); + await sandbox.configureBilling(billingInput); + + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ usageSinceLast: 499, reason: 'runtime_signal' }) + ); + }); + it('keeps physical start non-fatal while retrying an unacknowledged shadow start', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStart) diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 8a265786d0..41e0c0d5f7 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -49,8 +49,13 @@ function usageServiceForSandboxClass(sandboxClassName: SandboxClassName): string return `${SERVICE}-${suffix}`; } -function stoppedAtFromState(state: { lastChange: number }, observedAtMs = Date.now()): number { +function stoppedAtFromState( + state: { status: string; lastChange?: number }, + observedAtMs = Date.now() +): number { + if (state.status !== 'stopped' && state.status !== 'stopped_with_code') return observedAtMs; return Number.isFinite(state.lastChange) && + state.lastChange !== undefined && state.lastChange >= 0 && state.lastChange <= observedAtMs ? state.lastChange