From 55be662840c070ce361deed2b571ad04310c4ec2 Mon Sep 17 00:00:00 2001 From: syn Date: Tue, 18 Aug 2026 17:12:45 -0500 Subject: [PATCH 1/6] feat(cloud-agent-next): enforce container billing --- packages/container-usage/src/client.test.ts | 37 ++ packages/container-usage/src/client.ts | 29 +- .../container-usage/src/contracts.test.ts | 28 + packages/container-usage/src/contracts.ts | 44 ++ .../cloudflare-agent-sandbox.test.ts | 32 + .../cloudflare/cloudflare-agent-sandbox.ts | 101 ++- .../src/agent-sandbox/protocol.ts | 4 + .../src/container-billing-rollout.test.ts | 41 ++ .../src/container-billing-rollout.ts | 26 + .../src/container-usage-context.ts | 53 +- .../src/container-usage.test.ts | 127 ++++ .../cloud-agent-next/src/container-usage.ts | 281 +++++++- .../cloud-agent-next/src/execution/types.ts | 8 +- .../src/kilo-facade/session-proxy.test.ts | 81 +++ .../src/kilo-facade/session-proxy.ts | 21 +- .../src/kilo-facade/user-kilo-facade.ts | 2 + .../src/persistence/CloudAgentSession.ts | 46 +- .../src/session/agent-runtime.test.ts | 24 + .../src/session/agent-runtime.ts | 4 + .../src/session/pending-messages.ts | 3 + .../src/session/queue-message.ts | 1 + .../src/session/session-message-queue.test.ts | 24 + .../src/session/session-message-queue.ts | 8 + .../src/terminal/access.test.ts | 2 + services/cloud-agent-next/src/types.ts | 3 + .../src/websocket/ingest.test.ts | 26 + .../cloud-agent-next/src/websocket/ingest.ts | 8 + .../worker-configuration.d.ts | 608 ++++++++++++------ services/cloud-agent-next/wrangler.jsonc | 3 + .../cloud-agent-next/wrapper/src/lifecycle.ts | 60 +- services/cloud-agent-next/wrapper/src/main.ts | 7 +- .../src/billing-config.test.ts | 78 +++ .../src/billing-config.ts | 37 +- .../container-usage-meter/src/meter.test.ts | 42 ++ services/container-usage-meter/src/meter.ts | 67 +- .../container-usage-meter/src/postgres.ts | 48 +- .../test/postgres.test.ts | 69 +- .../worker-configuration.d.ts | 6 +- services/container-usage-meter/wrangler.jsonc | 2 + .../src/session-ingest-rpc.test.ts | 70 ++ .../session-ingest/src/session-ingest-rpc.ts | 10 +- 41 files changed, 1865 insertions(+), 306 deletions(-) create mode 100644 services/cloud-agent-next/src/container-billing-rollout.test.ts create mode 100644 services/cloud-agent-next/src/container-billing-rollout.ts create mode 100644 services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts diff --git a/packages/container-usage/src/client.test.ts b/packages/container-usage/src/client.test.ts index 1a7bf75cc2..fcb0198aff 100644 --- a/packages/container-usage/src/client.test.ts +++ b/packages/container-usage/src/client.test.ts @@ -12,6 +12,16 @@ function binding(overrides: Partial = {}): ContainerUs dedup: false, }, })), + recordStartV2: vi.fn(async input => ({ + success: true as const, + ack: { + intervalId: `${input.instanceId}:${input.startEpochMs}`, + durable: 'pg' as const, + dedup: false, + }, + billingMode: 'paid' as const, + remainingMicrodollars: 10_000_001, + })), recordHeartbeat: vi.fn(async input => ({ intervalId: `${input.instanceId}:${input.startEpochMs}`, durable: 'pg' as const, @@ -176,4 +186,31 @@ describe('ContainerUsageClient', () => { }); expect(recordStart).toHaveBeenCalledOnce(); }); + + it('returns versioned paid admission details without changing v1 start behavior', async () => { + const rpc = binding(); + const client = new ContainerUsageClient(rpc, { service: 'cloud-agent-next-sandbox' }); + + await expect(client.recordStartV2({ ...context, startEpochMs: 123 })).resolves.toMatchObject({ + success: true, + billingMode: 'paid', + remainingMicrodollars: 10_000_001, + }); + expect(rpc.recordStartV2).toHaveBeenCalledWith( + expect.objectContaining({ + service: 'cloud-agent-next-sandbox', + idempotencyKey: 'v1:cloud-agent-next-sandbox:instance-1:123:start', + }) + ); + }); + + it('reports unavailable versioned admission support explicitly', async () => { + const client = new ContainerUsageClient(binding({ recordStartV2: undefined }), { + service: 'cloud-agent-next-sandbox', + }); + + await expect(client.recordStartV2({ ...context, startEpochMs: 123 })).rejects.toThrow( + 'does not support recordStartV2' + ); + }); }); diff --git a/packages/container-usage/src/client.ts b/packages/container-usage/src/client.ts index b75bedfb71..858fec9bd1 100644 --- a/packages/container-usage/src/client.ts +++ b/packages/container-usage/src/client.ts @@ -3,6 +3,7 @@ import { heartbeatAckSchema, recordAckSchema, recordStartResultSchema, + recordStartV2ResultSchema, startIdempotencyKey, stopIdempotencyKey, type ContainerUsageRpcMethods, @@ -10,6 +11,7 @@ import { type RecordAck, type RecordHeartbeatInput, type RecordStartInput, + type RecordStartV2Result, type RecordStartFailureCode, type RecordStopInput, type UsageContext, @@ -90,11 +92,7 @@ export class ContainerUsageClient { } async recordStart(input: ClientRecordStartInput): Promise { - const request = { - ...input, - service: this.service, - idempotencyKey: startIdempotencyKey(this.service, input.instanceId, input.startEpochMs), - } satisfies RecordStartInput; + const request = this.startRequest(input); const result = await this.withRetry(async () => recordStartResultSchema.parse(await this.binding.recordStart(request)) ); @@ -104,6 +102,19 @@ export class ContainerUsageClient { return result.ack; } + async recordStartV2(input: ClientRecordStartInput): Promise { + // A deployed Worker RPC proxy may expose an absent method as callable and reject only + // when invoked. Keep this check for direct bindings; either failure mode remains fail-closed. + const recordStartV2 = this.binding.recordStartV2; + if (!recordStartV2) { + throw new Error('Container usage meter does not support recordStartV2'); + } + const request = this.startRequest(input); + return await this.withRetry(async () => + recordStartV2ResultSchema.parse(await recordStartV2.call(this.binding, request)) + ); + } + async recordHeartbeat(input: ClientRecordHeartbeatInput): Promise { const request = { ...input, @@ -148,6 +159,14 @@ export class ContainerUsageClient { } throw lastError; } + + private startRequest(input: ClientRecordStartInput): RecordStartInput { + return { + ...input, + service: this.service, + idempotencyKey: startIdempotencyKey(this.service, input.instanceId, input.startEpochMs), + }; + } } export function createContainerUsageClient( diff --git a/packages/container-usage/src/contracts.test.ts b/packages/container-usage/src/contracts.test.ts index 8f9f3679df..8a41315c2f 100644 --- a/packages/container-usage/src/contracts.test.ts +++ b/packages/container-usage/src/contracts.test.ts @@ -2,6 +2,8 @@ import { describe, expect, it } from 'vitest'; import { intervalId, recordHeartbeatInputSchema, + recordStartResultSchema, + recordStartV2ResultSchema, recordStopInputSchema, usageContextSchema, } from './contracts'; @@ -121,4 +123,30 @@ describe('container usage contracts', () => { intervalId('gastown', 'shared-instance', 123) ); }); + + it('keeps v1 starts stable while v2 reports paid admission details', () => { + const v1 = { + success: true, + ack: { intervalId: 'interval-1', durable: 'pg', dedup: false }, + }; + expect(recordStartResultSchema.parse(v1)).toEqual(v1); + expect( + recordStartV2ResultSchema.parse({ + ...v1, + billingMode: 'paid', + remainingMicrodollars: 10_000_001, + }) + ).toMatchObject({ success: true, billingMode: 'paid', remainingMicrodollars: 10_000_001 }); + expect( + recordStartV2ResultSchema.parse({ + success: false, + error: { + code: 'insufficient_credits', + message: 'Insufficient credits', + remainingMicrodollars: 5_000_000, + minimumRequiredMicrodollars: 5_000_000, + }, + }) + ).toMatchObject({ success: false }); + }); }); diff --git a/packages/container-usage/src/contracts.ts b/packages/container-usage/src/contracts.ts index a27fa904c5..5161c51dc2 100644 --- a/packages/container-usage/src/contracts.ts +++ b/packages/container-usage/src/contracts.ts @@ -148,6 +148,49 @@ export const recordStartResultSchema = z.discriminatedUnion('success', [ ]); export type RecordStartResult = z.infer; +export const recordStartV2ResultSchema = z.union([ + z + .object({ + success: z.literal(true), + ack: recordAckSchema, + billingMode: z.literal('shadow'), + }) + .strict(), + z + .object({ + success: z.literal(true), + ack: recordAckSchema, + billingMode: z.literal('paid'), + remainingMicrodollars: z.number().int(), + }) + .strict(), + z + .object({ + success: z.literal(false), + error: z + .object({ + code: z.literal('insufficient_credits'), + message: z.string().min(1), + remainingMicrodollars: z.number().int(), + minimumRequiredMicrodollars: z.number().int().positive(), + }) + .strict(), + }) + .strict(), + z + .object({ + success: z.literal(false), + error: z + .object({ + code: z.enum(['sku_not_found', 'sku_unit_mismatch', 'sku_not_accepting_new_usage']), + message: z.string().min(1), + }) + .strict(), + }) + .strict(), +]); +export type RecordStartV2Result = z.infer; + export const budgetVerdictSchema = z.discriminatedUnion('verdict', [ z.object({ verdict: z.literal('continue'), remaining: z.number().int().optional() }).strict(), z @@ -180,6 +223,7 @@ export type HeartbeatAck = z.infer; export type ContainerUsageRpcMethods = { recordStart: (input: RecordStartInput) => Promise; + recordStartV2?: (input: RecordStartInput) => Promise; recordHeartbeat: (input: RecordHeartbeatInput) => Promise; recordStop: (input: RecordStopInput) => Promise; }; 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 e770a4445f..6835298a96 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 @@ -437,6 +437,38 @@ describe('CloudflareAgentSandbox', () => { ensureBootstrapWrapper.mockRestore(); }); + it('requests paid admission for an organization canary before sandbox work', async () => { + const ensureBillingAdmission = vi.fn().mockResolvedValue({ + success: false, + code: 'insufficient_credits', + message: 'Low balance', + }); + const sandbox = new CloudflareAgentSandbox( + { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: '', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: 'org_cloudflare', + } as Env, + metadata({ sandboxId: 'usr-shared' }), + { + resolveSandbox: () => ({}) as SandboxInstance, + ensureBillingAdmission, + isBillingBlocked: vi.fn().mockResolvedValue(false), + } + ); + + await expect(sandbox.ensureBillingAdmission()).resolves.toMatchObject({ + success: false, + code: 'insufficient_credits', + }); + expect(ensureBillingAdmission).toHaveBeenCalledWith(expect.anything(), { + sandboxId: 'usr-shared', + subject: { type: 'org', id: 'org_cloudflare' }, + actor: { type: 'user', id: 'user_cloudflare' }, + enforcementRequested: true, + }); + }); + it('reports malformed worker URLs before degrading to a cold bootstrap', async () => { const request = ensureRequest({ cacheEligible: true }); const bucket = { get: vi.fn(), put: vi.fn() }; 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 0b5d330340..40ec543c79 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 @@ -1,11 +1,12 @@ -import type { - AgentSandbox, - EnsureWrapperRequest, - StopWrappersResult, - TerminalClientResult, - WrapperLogs, - WrapperObservation, - WrapperStopTarget, +import { + AgentSandboxUnavailableError, + type AgentSandbox, + type EnsureWrapperRequest, + type StopWrappersResult, + type TerminalClientResult, + type WrapperLogs, + type WrapperObservation, + type WrapperStopTarget, } from '../protocol.js'; import type { Env, @@ -75,9 +76,12 @@ import { TOOL_CGROUP_ENV_KEYS, type ToolCgroupEnv } from '../../shared/tool-cgro import { buildSandboxBillingInput, configureSandboxBillingInput, + ensureSandboxBillingAdmissionInput, + isSandboxBillingBlocked, isSandboxContainerRunning, type SandboxBillingInput, } from '../../container-usage-context.js'; +import { isCloudAgentContainerBillingEnabled } from '../../container-billing-rollout.js'; const PREPARE_WORKSPACE_TIMEOUT_MS = 10 * 60 * 1000; const DEFAULT_STOP_OBSERVATION_DELAYS_MS = [100, 500, 1_000]; @@ -188,6 +192,8 @@ function withWorkspacePreparationTimeout(operation: Promise, step: string) export type CloudflareAgentSandboxDependencies = { resolveSandbox?: (sandboxId: SandboxId, options?: { sleepAfter?: number }) => SandboxInstance; configureBilling?: (sandbox: SandboxInstance, input: SandboxBillingInput) => Promise; + ensureBillingAdmission?: typeof ensureSandboxBillingAdmissionInput; + isBillingBlocked?: typeof isSandboxBillingBlocked; sessionService?: SessionService; stopObservedWrappers?: typeof stopObservedWrappers; sleep?: (ms: number) => Promise; @@ -207,6 +213,8 @@ export class CloudflareAgentSandbox implements AgentSandbox { sandbox: SandboxInstance, input: SandboxBillingInput ) => Promise; + private readonly ensureSandboxBillingAdmission: typeof ensureSandboxBillingAdmissionInput; + private readonly sandboxBillingBlocked: typeof isSandboxBillingBlocked; private sandboxIdPromise?: Promise; constructor( @@ -230,6 +238,9 @@ export class CloudflareAgentSandbox implements AgentSandbox { this.stopObservationDelaysMs = dependencies.stopObservationDelaysMs ?? DEFAULT_STOP_OBSERVATION_DELAYS_MS; this.configureBilling = dependencies.configureBilling ?? configureSandboxBillingInput; + this.ensureSandboxBillingAdmission = + dependencies.ensureBillingAdmission ?? ensureSandboxBillingAdmissionInput; + this.sandboxBillingBlocked = dependencies.isBillingBlocked ?? isSandboxBillingBlocked; } private resolveSandboxId(): Promise { @@ -250,16 +261,60 @@ export class CloudflareAgentSandbox implements AgentSandbox { return this.sandboxIdPromise; } - private async getSandbox(options?: { sleepAfter?: number }): Promise { + private billingInput(sandboxId: SandboxId): SandboxBillingInput { + return buildSandboxBillingInput( + this.metadata, + sandboxId, + isCloudAgentContainerBillingEnabled(this.env, this.metadata.identity) + ); + } + + async ensureBillingAdmission() { const sandboxId = await this.resolveSandboxId(); - const sandbox = this.resolveSandbox(sandboxId, options); - 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'); - } + const sandbox = this.resolveSandbox(sandboxId); + const input = this.billingInput(sandboxId); + const blocked = await this.sandboxBillingBlocked(sandbox); + if (!input.enforcementRequested && !blocked) { + return { success: true as const, billingMode: 'shadow' as const }; + } + return this.ensureSandboxBillingAdmission(sandbox, input); + } + + async isBillingBlocked(): Promise { + const sandboxId = await this.resolveSandboxId(); + return this.sandboxBillingBlocked(this.resolveSandbox(sandboxId)); + } + + private async getSandbox(options?: { + sleepAfter?: number; + bypassBilling?: boolean; + }): Promise { + const sandboxId = await this.resolveSandboxId(); + const sandbox = this.resolveSandbox( + sandboxId, + options?.sleepAfter === undefined ? undefined : { sleepAfter: options.sleepAfter } ); + const input = this.billingInput(sandboxId); + if (!options?.bypassBilling) { + const blocked = await this.sandboxBillingBlocked(sandbox); + if (input.enforcementRequested || blocked) { + const admission = await this.ensureSandboxBillingAdmission(sandbox, input); + if (!admission.success) { + throw new AgentSandboxUnavailableError( + admission.code === 'insufficient_credits' || admission.code === 'stopping' + ? 'Container billing requires additional credits' + : 'Container billing admission is temporarily unavailable', + 'billing_blocked' + ); + } + } else { + void this.configureBilling(sandbox, input).catch(error => { + logger + .withFields({ error: error instanceof Error ? error.message : String(error) }) + .warn('Container usage shadow configuration deferred'); + }); + } + } return sandbox; } @@ -765,9 +820,13 @@ export class CloudflareAgentSandbox implements AgentSandbox { } async discoverSessionWrappers(): Promise { - return discoverSessionWrappers(await this.getSandbox(), this.metadata.identity.sessionId, { - inspectContainers: this.usesDevcontainerRuntime(), - }); + return discoverSessionWrappers( + await this.getSandbox({ bypassBilling: true }), + this.metadata.identity.sessionId, + { + inspectContainers: this.usesDevcontainerRuntime(), + } + ); } async observeWrappersWithoutWaking(): Promise { @@ -787,7 +846,7 @@ export class CloudflareAgentSandbox implements AgentSandbox { attemptId: string; reason: WrapperStopReason; }): Promise { - const sandbox = await this.getSandbox(); + const sandbox = await this.getSandbox({ bypassBilling: true }); // Inspecting is a container fetch, so it boots a sleeping container. A wrapper is a // process, and a process cannot outlive its container (activity expiry SIGTERMs the @@ -934,7 +993,7 @@ export class CloudflareAgentSandbox implements AgentSandbox { } async delete(reason: SandboxDeleteReason): Promise { - const sandbox = await this.getSandbox(); + const sandbox = await this.getSandbox({ bypassBilling: true }); if (reason === 'recovery') { await sandbox.destroy(); return; diff --git a/services/cloud-agent-next/src/agent-sandbox/protocol.ts b/services/cloud-agent-next/src/agent-sandbox/protocol.ts index c13a968f05..153e2f15a7 100644 --- a/services/cloud-agent-next/src/agent-sandbox/protocol.ts +++ b/services/cloud-agent-next/src/agent-sandbox/protocol.ts @@ -7,6 +7,7 @@ import type { WorkspaceReady, } from '../execution/types.js'; import type { SessionMetadata } from '../persistence/session-metadata.js'; +import type { SandboxBillingAdmissionResult } from '../container-usage-context.js'; export type SandboxDeleteReason = 'explicit' | 'retention-expired' | 'recovery'; @@ -24,6 +25,7 @@ export type AgentSandboxFailure = | 'runtime_deleted_during_active_work' | 'runtime_max_duration_reached' | 'runtime_infrastructure_failed' + | 'billing_blocked' | 'capability_unavailable'; export class AgentSandboxUnavailableError extends Error { @@ -129,6 +131,8 @@ export type EnsuredWrapper = * Provider process, filesystem, and raw sandbox APIs remain private to adapters. */ export type AgentSandbox = { + ensureBillingAdmission(): Promise; + isBillingBlocked(): Promise; ensureWrapper(request: EnsureWrapperRequest): Promise; discoverSessionWrappers(): Promise; /** diff --git a/services/cloud-agent-next/src/container-billing-rollout.test.ts b/services/cloud-agent-next/src/container-billing-rollout.test.ts new file mode 100644 index 0000000000..fb209664e4 --- /dev/null +++ b/services/cloud-agent-next/src/container-billing-rollout.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it } from 'vitest'; +import { isCloudAgentContainerBillingEnabled } from './container-billing-rollout.js'; + +const enabled = { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: 'user-1', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: 'org-1', +}; + +describe('Cloud Agent container billing rollout', () => { + it('selects personal and organization payers independently', () => { + expect(isCloudAgentContainerBillingEnabled(enabled, { userId: 'user-1' })).toBe(true); + expect( + isCloudAgentContainerBillingEnabled(enabled, { userId: 'user-1', orgId: 'other-org' }) + ).toBe(false); + expect( + isCloudAgentContainerBillingEnabled(enabled, { userId: 'other-user', orgId: 'org-1' }) + ).toBe(true); + }); + + it('requires the exact global switch and fails malformed lists closed', () => { + expect( + isCloudAgentContainerBillingEnabled( + { ...enabled, CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'TRUE' }, + { userId: 'user-1' } + ) + ).toBe(false); + expect( + isCloudAgentContainerBillingEnabled( + { ...enabled, CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: 'user-1,' }, + { userId: 'user-1' } + ) + ).toBe(false); + expect( + isCloudAgentContainerBillingEnabled( + { ...enabled, CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: '*' }, + { userId: 'user-2', orgId: 'org-1' } + ) + ).toBe(false); + }); +}); diff --git a/services/cloud-agent-next/src/container-billing-rollout.ts b/services/cloud-agent-next/src/container-billing-rollout.ts new file mode 100644 index 0000000000..1161178104 --- /dev/null +++ b/services/cloud-agent-next/src/container-billing-rollout.ts @@ -0,0 +1,26 @@ +import type { Env } from './types.js'; + +type BillingIdentity = { userId: string; orgId?: string }; + +function parseAllowlist(value: string | undefined): ReadonlySet | null { + if (value === undefined || value === '') return new Set(); + const values = value.split(',').map(item => item.trim()); + if (values.some(item => item.length === 0 || item === '*')) return null; + return new Set(values); +} + +export function isCloudAgentContainerBillingEnabled( + env: Pick< + Env, + | 'CLOUD_AGENT_CONTAINER_BILLING_ENABLED' + | 'CLOUD_AGENT_CONTAINER_BILLING_USER_IDS' + | 'CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS' + >, + identity: BillingIdentity +): boolean { + if (env.CLOUD_AGENT_CONTAINER_BILLING_ENABLED !== 'true') return false; + const userIds = parseAllowlist(env.CLOUD_AGENT_CONTAINER_BILLING_USER_IDS); + const orgIds = parseAllowlist(env.CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS); + if (!userIds || !orgIds) return false; + return identity.orgId ? orgIds.has(identity.orgId) : userIds.has(identity.userId); +} diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts index 72410211a7..47d0abe29a 100644 --- a/services/cloud-agent-next/src/container-usage-context.ts +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -38,9 +38,21 @@ export const SANDBOX_CAPACITIES: Record< }; export type SandboxBillingInput = Omit & { sandboxId: SandboxId; + enforcementRequested?: boolean; }; +export type SandboxBillingAdmissionResult = + | { success: true; billingMode: 'shadow' | 'paid'; remainingMicrodollars?: number } + | { + success: false; + code: 'insufficient_credits' | 'meter_unavailable' | 'configuration_mismatch' | 'stopping'; + message: string; + remainingMicrodollars?: number; + minimumRequiredMicrodollars?: number; + }; export type MeteredSandboxInstance = SandboxInstance & { configureBilling(input: unknown): Promise; + ensureBillingAdmission(input: unknown): Promise; + isBillingBlocked(): Promise; isContainerRunning(): Promise; }; @@ -65,6 +77,7 @@ const sandboxBillingInputEnvelopeSchema = z message: 'Metadata may contain at most 16 entries', }) .optional(), + enforcementRequested: z.boolean().default(false), }) .strict(); @@ -96,7 +109,8 @@ function isIsolatedSandbox(sandboxId: SandboxId): boolean { export function buildSandboxBillingInput( metadata: SessionMetadata, - sandboxId: SandboxId + sandboxId: SandboxId, + enforcementRequested = false ): SandboxBillingInput { const subject = metadata.identity.orgId ? { type: 'org' as const, id: metadata.identity.orgId } @@ -108,6 +122,7 @@ export function buildSandboxBillingInput( return { sandboxId, + ...(enforcementRequested ? { enforcementRequested: true } : {}), subject, actor, ...(actor.type === 'bot' ? { onBehalfOf: subject } : {}), @@ -120,7 +135,7 @@ export function buildSandboxBillingInput( export function parseSandboxBillingInput(input: unknown): SandboxBillingInput { const parsed = sandboxBillingInputEnvelopeSchema.parse(input); - const { sandboxId, ...usageInput } = parsed; + const { sandboxId, enforcementRequested, ...usageInput } = parsed; const validated = usageContextSchema.parse({ service: 'cloud-agent-next', instanceId: 'validation', @@ -128,7 +143,7 @@ export function parseSandboxBillingInput(input: unknown): SandboxBillingInput { ...usageInput, }); const { service: _service, instanceId: _instanceId, sku: _sku, ...billingInput } = validated; - return { sandboxId, ...billingInput }; + return { sandboxId, enforcementRequested, ...billingInput }; } export function assertSandboxBillingAllocation( @@ -185,6 +200,38 @@ export async function configureSandboxBilling( await configureSandboxBillingInput(sandbox, buildSandboxBillingInput(metadata, sandboxId)); } +export async function ensureSandboxBillingAdmissionInput( + sandbox: SandboxInstance, + input: SandboxBillingInput +): Promise { + const ensureBillingAdmission = (sandbox as Partial) + .ensureBillingAdmission; + if (typeof ensureBillingAdmission !== 'function') { + return input.enforcementRequested + ? { + success: false, + code: 'meter_unavailable', + message: 'Container billing admission is unavailable', + } + : { success: true, billingMode: 'shadow' }; + } + try { + return await (sandbox as MeteredSandboxInstance).ensureBillingAdmission(input); + } catch (error) { + return { + success: false, + code: 'meter_unavailable', + message: error instanceof Error ? error.message : 'Container billing admission failed', + }; + } +} + +export async function isSandboxBillingBlocked(sandbox: SandboxInstance): Promise { + const isBillingBlocked = (sandbox as Partial).isBillingBlocked; + if (typeof isBillingBlocked !== 'function') return false; + return await (sandbox as MeteredSandboxInstance).isBillingBlocked(); +} + /** * Whether the sandbox's container is currently running, read over Durable Object RPC. * diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index 3f28a36098..106e189b3b 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -18,6 +18,7 @@ const sdk = vi.hoisted(() => { superStopped = false; superActivityExpired = false; superStopCalled = false; + superDestroyCalled = false; constructor(ctx: SandboxDurableObjectState, env: unknown) { this.ctx = ctx; @@ -58,6 +59,10 @@ const sdk = vi.hoisted(() => { async stop(): Promise { this.superStopCalled = true; } + + async destroy(): Promise { + this.superDestroyCalled = true; + } } return { StockSandbox }; }); @@ -96,6 +101,11 @@ function createRpc(): ContainerUsageRpcMethods { success: true, ack: ack(), })), + recordStartV2: vi.fn>(async () => ({ + success: true, + ack: ack(), + billingMode: 'shadow', + })), recordHeartbeat: vi.fn(async () => ({ ...ack(), budget: { verdict: 'continue' }, @@ -111,6 +121,7 @@ type TestRuntime = MeteredSandbox & { superStopped: boolean; superActivityExpired: boolean; superStopCalled: boolean; + superDestroyCalled: boolean; setPhysicalRunning(running: boolean): void; billingHeartbeatTick(generation?: string): Promise; }; @@ -170,6 +181,122 @@ describe('MeteredSandbox', () => { vi.restoreAllMocks(); }); + it('requires a paid meter admission before a selected cold start', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStartV2!).mockResolvedValue({ + success: true, + ack: ack(), + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }); + const { sandbox, flushShadowTasks } = createSandbox(rpc); + + await expect( + sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }) + ).resolves.toEqual({ + success: true, + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }); + expect(rpc.recordStartV2).toHaveBeenCalledOnce(); + expect(sandbox.schedules).toHaveLength(0); + await sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }); + expect(rpc.recordStartV2).toHaveBeenCalledOnce(); + + await sandbox.onStart(); + await flushShadowTasks(); + expect(sandbox.schedules).toEqual([ + expect.objectContaining({ callback: 'billingHeartbeatTick' }), + ]); + }); + + it('fails selected admission closed for low balance or a shadow meter mismatch', async () => { + const lowBalanceRpc = createRpc(); + vi.mocked(lowBalanceRpc.recordStartV2!).mockResolvedValue({ + success: false, + error: { + code: 'insufficient_credits', + message: 'Low balance', + remainingMicrodollars: 5_000_000, + minimumRequiredMicrodollars: 5_000_000, + }, + }); + const { sandbox: lowBalance } = createSandbox(lowBalanceRpc); + await expect( + lowBalance.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }) + ).resolves.toMatchObject({ + success: false, + code: 'insufficient_credits', + remainingMicrodollars: 5_000_000, + }); + + const shadowRpc = createRpc(); + const { sandbox: mismatch } = createSandbox(shadowRpc); + await expect( + mismatch.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }) + ).resolves.toMatchObject({ success: false, code: 'configuration_mismatch' }); + }); + + it('fails selected admission closed for an already-running shadow generation', async () => { + const { sandbox, flushShadowTasks } = createSandbox(createRpc()); + await sandbox.configureBilling(billingInput); + await sandbox.onStart(); + await flushShadowTasks(); + + await expect( + sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }) + ).resolves.toMatchObject({ success: false, code: 'configuration_mismatch' }); + }); + + it('persists stop enforcement, reports final usage, and resumes only after new admission', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStartV2!).mockResolvedValue({ + success: true, + ack: ack(), + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }); + vi.mocked(rpc.recordHeartbeat).mockResolvedValue({ + ...ack(), + budget: { + verdict: 'stop', + remainingMicrodollars: 5_000_000, + minimumRequiredMicrodollars: 5_000_000, + }, + }); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + vi.spyOn(Date, 'now').mockReturnValue(1_000); + await sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }); + await sandbox.onStart(); + await flushShadowTasks(); + const active = await getBillingContext(storage); + if (!active) throw new Error('Expected active paid billing context'); + expect(active.measurementStarted).toBe(true); + sandbox.mockState = { status: 'running' }; + + vi.spyOn(Date, 'now').mockReturnValue(301_000); + await sandbox.billingHeartbeatTick(active.generation); + expect(sandbox.superStopCalled).toBe(true); + expect(await sandbox.isBillingBlocked()).toBe(true); + expect(sandbox.schedules).toContainEqual({ + when: 120, + callback: 'billingForceStop', + payload: active.generation, + }); + expect(rpc.recordStop).toHaveBeenCalledOnce(); + await sandbox.billingForceStop(active.generation); + expect(sandbox.superDestroyCalled).toBe(true); + expect(await sandbox.isBillingBlocked()).toBe(true); + + sandbox.setPhysicalRunning(false); + vi.spyOn(Date, 'now').mockReturnValue(302_000); + await expect( + sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }) + ).resolves.toMatchObject({ success: true, billingMode: 'paid' }); + expect(await sandbox.isBillingBlocked()).toBe(false); + expect(rpc.recordStartV2).toHaveBeenCalledTimes(2); + }); + it('admits one start per physical generation and short-circuits active acquisition', async () => { const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(1_000); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index c54fd561d8..b017104027 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -1,4 +1,5 @@ import { + clearBillingContext, createContainerUsageClient, getBillingContext, installBillingHeartbeat, @@ -20,6 +21,7 @@ import { SANDBOX_CAPACITIES, SANDBOX_USAGE_SKUS, type SandboxBillingInput, + type SandboxBillingAdmissionResult, type SandboxClassName, } from './container-usage-context.js'; @@ -28,6 +30,9 @@ 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'; +const START_ADMISSION_STORAGE_KEY = 'container-usage:start-admission:v1'; +const BILLING_BLOCK_STORAGE_KEY = 'container-usage:budget-block:v1'; +const BILLING_FORCE_STOP_SECONDS = 120; // oxlint-disable-next-line no-empty-object-type -- Matches the Sandbox 0.12.1 constructor. type SandboxDurableObjectState = DurableObjectState<{}>; @@ -47,6 +52,24 @@ const pendingStopReasonSchema = z }) .strict(); +const startAdmissionSchema = z + .object({ + generation: z.uuid(), + billingMode: z.enum(['shadow', 'paid']), + remainingMicrodollars: z.number().int().optional(), + }) + .strict(); + +const billingBlockSchema = z + .object({ + generation: z.uuid(), + startEpochMs: z.number().int().nonnegative(), + blockedAt: z.number().int().nonnegative(), + forceStopAt: z.number().int().nonnegative(), + remainingMicrodollars: z.number().int().optional(), + }) + .strict(); + function startInputFromContext(context: BillingContext): ClientRecordStartInput { const { service: _service, ...usage } = usageContextFromBillingContext(context); return { ...usage, startEpochMs: context.startEpochMs }; @@ -88,10 +111,8 @@ export abstract class MeteredSandbox extends StockSandbox { 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'); - }, + onBudgetWarning: budget => this.logBudgetWarning(budget), + enforceBudgetStop: (budget, expected) => this.enforceBudgetStop(budget, expected), }); } @@ -112,6 +133,152 @@ export abstract class MeteredSandbox extends StockSandbox { return this.ctx.container?.running === true; } + async isBillingBlocked(): Promise { + return (await this.getBillingBlock()) !== undefined; + } + + async ensureBillingAdmission(input: unknown): Promise { + const parsed = parseSandboxBillingInput(input); + assertSandboxBillingAllocation(this.sandboxClassName, parsed); + return this.runBillingExclusive(async () => { + await this.ctx.storage.put(PENDING_ATTRIBUTION_STORAGE_KEY, parsed); + const block = await this.getBillingBlock(); + let active = await getBillingContext(this.ctx.storage); + + if (!block && !parsed.enforcementRequested) { + return { success: true, billingMode: 'shadow' }; + } + + if (active && !active.measurementStarted && !block) { + const admission = await this.getStartAdmission(active.generation); + if (admission) { + if (parsed.enforcementRequested && admission.billingMode !== 'paid') { + return { + success: false, + code: 'configuration_mismatch', + message: 'Container billing rollout is not enabled in the usage meter', + }; + } + return { + success: true, + billingMode: admission.billingMode, + ...(admission.remainingMicrodollars === undefined + ? {} + : { remainingMicrodollars: admission.remainingMicrodollars }), + }; + } + } + + if (active?.measurementStarted && this.ctx.container?.running === true) { + const admission = await this.getStartAdmission(active.generation); + if (block) { + return { + success: false, + code: 'stopping', + message: 'Container is stopping because its billing balance is too low', + remainingMicrodollars: block.remainingMicrodollars, + }; + } + if (parsed.enforcementRequested && admission?.billingMode !== 'paid') { + return { + success: false, + code: 'configuration_mismatch', + message: 'Container billing rollout is not enabled for the active usage generation', + }; + } + return admission + ? { + success: true, + billingMode: admission.billingMode, + remainingMicrodollars: admission.remainingMicrodollars, + } + : { success: true, billingMode: 'shadow' }; + } + + if (this.ctx.container?.running === true) { + return { + success: false, + code: 'stopping', + message: 'Container billing admission is waiting for the previous run to stop', + }; + } + + if (active) { + try { + await this.billingHeartbeat.recordStop( + { reason: 'runtime_signal' }, + active.stoppedObservedAtMs ?? Date.now() + ); + await this.ctx.storage.delete(START_ACK_GENERATION_STORAGE_KEY); + await this.ctx.storage.delete(START_ADMISSION_STORAGE_KEY); + } catch (error) { + return { + success: false, + code: 'meter_unavailable', + message: + error instanceof Error ? error.message : 'Final usage settlement is unavailable', + }; + } + active = undefined; + } + + const context = await this.createBillingGeneration(parsed, 'attribution-adoption'); + let result; + try { + result = await this.usageClient.recordStartV2(startInputFromContext(context)); + } catch (error) { + await clearBillingContext(this.ctx.storage); + return { + success: false, + code: 'meter_unavailable', + message: + error instanceof Error ? error.message : 'Container billing meter is unavailable', + }; + } + if (!result.success) { + await clearBillingContext(this.ctx.storage); + return { + success: false, + code: + result.error.code === 'insufficient_credits' + ? 'insufficient_credits' + : 'configuration_mismatch', + message: result.error.message, + ...(result.error.code === 'insufficient_credits' + ? { + remainingMicrodollars: result.error.remainingMicrodollars, + minimumRequiredMicrodollars: result.error.minimumRequiredMicrodollars, + } + : {}), + }; + } + if (parsed.enforcementRequested && result.billingMode !== 'paid') { + await clearBillingContext(this.ctx.storage); + return { + success: false, + code: 'configuration_mismatch', + message: 'Container billing rollout is not enabled in the usage meter', + }; + } + await this.ctx.storage.put(START_ACK_GENERATION_STORAGE_KEY, context.generation); + await this.ctx.storage.put(START_ADMISSION_STORAGE_KEY, { + generation: context.generation, + billingMode: result.billingMode, + ...(result.billingMode === 'paid' + ? { remainingMicrodollars: result.remainingMicrodollars } + : {}), + }); + await this.ctx.storage.delete(BILLING_BLOCK_STORAGE_KEY); + return result.billingMode === 'paid' + ? { + success: true, + billingMode: 'paid', + remainingMicrodollars: result.remainingMicrodollars, + } + : { success: true, billingMode: 'shadow' }; + }); + } + async configureBilling(input: unknown): Promise { const parsed = parseSandboxBillingInput(input); assertSandboxBillingAllocation(this.sandboxClassName, parsed); @@ -193,6 +360,12 @@ export abstract class MeteredSandbox extends StockSandbox { override async onStart(): Promise { await super.onStart(); this.runShadowTask('start lifecycle', async () => { + const block = await this.getBillingBlock(); + if (block) { + await this.scheduleForceStop(block); + await this.stop(); + return; + } const previous = await getBillingContext(this.ctx.storage); if (previous) { if (previous.pendingStop) { @@ -287,6 +460,12 @@ export abstract class MeteredSandbox extends StockSandbox { }); } + override async destroy(): Promise { + const block = await this.getBillingBlock(); + await super.destroy(); + if (block) await this.ctx.storage.put(BILLING_BLOCK_STORAGE_KEY, block); + } + private runBillingExclusive(operation: () => Promise): Promise { const result = this.billingLifecycleTail.then(operation, operation); this.billingLifecycleTail = result.then( @@ -305,6 +484,7 @@ export abstract class MeteredSandbox extends StockSandbox { private schedulePendingGenerationIfRunning(): void { this.runShadowTask('replacement generation', async () => { + if (await this.getBillingBlock()) return; if (this.ctx.container?.running !== true) return; if (await getBillingContext(this.ctx.storage)) return; const input = await this.getPendingAttribution(); @@ -373,10 +553,87 @@ export abstract class MeteredSandbox extends StockSandbox { } } - private async startBillingGeneration( + private async getStartAdmission(generation: string) { + const stored = await this.ctx.storage.get(START_ADMISSION_STORAGE_KEY); + if (stored === undefined) return undefined; + const parsed = startAdmissionSchema.parse(stored); + return parsed.generation === generation ? parsed : undefined; + } + + private async getBillingBlock() { + const stored = await this.ctx.storage.get(BILLING_BLOCK_STORAGE_KEY); + return stored === undefined ? undefined : billingBlockSchema.parse(stored); + } + + private async logBudgetWarning(budget: { + verdict: string; + remainingMicrodollars?: number; + }): Promise { + const context = await getBillingContext(this.ctx.storage); + logger + .withTags({ logTag: 'container_billing_warning' }) + .withFields({ + sandboxClass: this.sandboxClassName, + billingMode: 'paid', + verdict: budget.verdict, + remainingMicrodollars: budget.remainingMicrodollars, + generation: context?.generation, + subjectType: context?.subject.type, + sessionId: context?.sessionId, + }) + .warn('Container billing balance is approaching the stop threshold'); + } + + private async scheduleForceStop(block: z.infer): Promise { + const delaySeconds = Math.max(0, Math.ceil((block.forceStopAt - Date.now()) / 1_000)); + this.deleteSchedules('billingForceStop'); + await this.schedule(delaySeconds, 'billingForceStop', block.generation); + } + + private async enforceBudgetStop( + budget: { verdict: string; remainingMicrodollars?: number }, + expected: { generation: string; startEpochMs: number } + ): Promise { + const now = Date.now(); + const block = { + generation: expected.generation, + startEpochMs: expected.startEpochMs, + blockedAt: now, + forceStopAt: now + BILLING_FORCE_STOP_SECONDS * 1_000, + remainingMicrodollars: budget.remainingMicrodollars, + }; + await this.ctx.storage.put(BILLING_BLOCK_STORAGE_KEY, block); + await this.scheduleForceStop(block); + logger + .withTags({ logTag: 'container_billing_stop' }) + .withFields({ + sandboxClass: this.sandboxClassName, + generation: expected.generation, + remainingMicrodollars: budget.remainingMicrodollars, + forceStopAt: block.forceStopAt, + }) + .warn('Container billing stop initiated'); + await this.stop(); + } + + async billingForceStop(generation: string): Promise { + const block = await this.getBillingBlock(); + if (!block || block.generation !== generation || this.ctx.container?.running !== true) return; + logger + .withTags({ logTag: 'container_billing_force_stop' }) + .withFields({ + sandboxClass: this.sandboxClassName, + generation, + stopLatencyMs: Date.now() - block.blockedAt, + }) + .error('Force-destroying container after billing stop deadline'); + await this.destroy(); + } + + private async createBillingGeneration( input: SandboxBillingInput, trigger: ContainerStartTrigger - ): Promise { + ): Promise { const capacity = SANDBOX_CAPACITIES[this.sandboxClassName]; const previousStartEpochMs = (await this.ctx.storage.get(LAST_START_EPOCH_STORAGE_KEY)) ?? -1; @@ -401,8 +658,6 @@ export abstract class MeteredSandbox extends StockSandbox { startEpochMs, } satisfies UsageContext & { startEpochMs: number }); await this.ctx.storage.delete(PENDING_STOP_REASON_STORAGE_KEY); - // The only worker-side record that a container run began. `service:sandboxId:startEpochMs` - // is the usage `intervalId`, so these fields join a log line to its usage row. logger .withTags({ logTag: 'container_started', sandboxId: input.sandboxId }) .withFields({ @@ -413,7 +668,15 @@ export abstract class MeteredSandbox extends StockSandbox { sessionId: input.sessionId, durableObjectId: this.ctx.id.toString(), }) - .info('Container started'); + .info('Container billing generation created'); + return context; + } + + private async startBillingGeneration( + input: SandboxBillingInput, + trigger: ContainerStartTrigger + ): Promise { + const context = await this.createBillingGeneration(input, trigger); await this.admitAndScheduleBestEffort(context); } } diff --git a/services/cloud-agent-next/src/execution/types.ts b/services/cloud-agent-next/src/execution/types.ts index 75d49fc4d7..721b2ad737 100644 --- a/services/cloud-agent-next/src/execution/types.ts +++ b/services/cloud-agent-next/src/execution/types.ts @@ -232,7 +232,13 @@ export type PermanentDeliveryResultCode = 'SANDBOX_CAPABILITY_UNAVAILABLE'; export type AdmissionFailure = { success: false; - code: 'NOT_FOUND' | 'BAD_REQUEST' | 'INTERNAL' | 'PENDING_QUEUE_FULL' | RetryableResultCode; + code: + | 'NOT_FOUND' + | 'BAD_REQUEST' + | 'INTERNAL' + | 'PAYMENT_REQUIRED' + | 'PENDING_QUEUE_FULL' + | RetryableResultCode; error: string; failureBoundary?: 'registration' | 'admission'; }; diff --git a/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts b/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts new file mode 100644 index 0000000000..3912c983a0 --- /dev/null +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts @@ -0,0 +1,81 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import type { Env } from '../types.js'; +import type * as SandboxIdModule from '../sandbox-id.js'; + +const mocks = vi.hoisted(() => ({ + getSandbox: vi.fn(), + findWrapperForSession: vi.fn(), + fetchSessionMetadata: vi.fn(), +})); + +vi.mock('@cloudflare/sandbox', () => ({ getSandbox: mocks.getSandbox })); +vi.mock('../kilo/wrapper-manager.js', () => ({ + findWrapperForSession: mocks.findWrapperForSession, +})); +vi.mock('../session-service.js', () => ({ fetchSessionMetadata: mocks.fetchSessionMetadata })); +vi.mock('../sandbox-id.js', async importOriginal => ({ + ...(await importOriginal()), + generateSandboxId: vi.fn(), + getSandboxNamespace: vi.fn().mockReturnValue({}), +})); + +import { resolveLiveWrapperTarget } from './session-proxy.js'; + +const metadata = { + metadataSchemaVersion: 2, + identity: { + sessionId: 'agent_facade', + userId: 'user_facade', + orgId: 'org_facade', + }, + auth: { kiloSessionId: 'kilo_facade' }, + lifecycle: { version: 1, timestamp: 1 }, + workspace: { + sandboxId: 'ses-facade', + sandboxProvider: 'cloudflare', + workspacePath: '/workspace/facade', + sessionHome: '/home/agent_facade', + branchName: 'main', + }, +} satisfies SessionMetadata; + +describe('resolveLiveWrapperTarget billing admission', () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.fetchSessionMetadata.mockResolvedValue(metadata); + }); + + it('does not inspect a selected live wrapper when paid admission is rejected', async () => { + const ensureBillingAdmission = vi.fn().mockResolvedValue({ + success: false, + code: 'insufficient_credits', + message: 'Low balance', + }); + mocks.getSandbox.mockReturnValue({ + isBillingBlocked: vi.fn().mockResolvedValue(false), + ensureBillingAdmission, + }); + + await expect( + resolveLiveWrapperTarget({ + env: { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: '', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: 'org_facade', + } as Env, + userId: 'user_facade', + cloudAgentSessionId: 'agent_facade', + }) + ).resolves.toBeNull(); + expect(ensureBillingAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + sandboxId: 'ses-facade', + subject: { type: 'org', id: 'org_facade' }, + actor: { type: 'user', id: 'user_facade' }, + enforcementRequested: true, + }) + ); + expect(mocks.findWrapperForSession).not.toHaveBeenCalled(); + }); +}); 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 dd1865d68d..9d2152eece 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts @@ -4,7 +4,13 @@ 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 { configureSandboxBilling } from '../container-usage-context.js'; +import { + buildSandboxBillingInput, + configureSandboxBillingInput, + ensureSandboxBillingAdmissionInput, + isSandboxBillingBlocked, +} from '../container-usage-context.js'; +import { isCloudAgentContainerBillingEnabled } from '../container-billing-rollout.js'; export type SessionKiloFacadeDecision = | { kind: 'proxy-live-wrapper' } @@ -86,7 +92,18 @@ export async function resolveLiveWrapperTarget(params: { }), sandboxId ); - void configureSandboxBilling(sandbox, metadata, sandboxId); + const billingInput = buildSandboxBillingInput( + metadata, + sandboxId, + isCloudAgentContainerBillingEnabled(env, metadata.identity) + ); + const billingBlocked = await isSandboxBillingBlocked(sandbox); + if (billingInput.enforcementRequested || billingBlocked) { + const admission = await ensureSandboxBillingAdmissionInput(sandbox, billingInput); + if (!admission.success) return null; + } else { + void configureSandboxBillingInput(sandbox, billingInput); + } const wrapperInfo = await findWrapperForSession(sandbox, sessionId); if (!wrapperInfo) { return null; diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 3ff7d2ba41..0fef97db60 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -633,6 +633,8 @@ function promptAdmissionError( return missingRootKiloSessionResponse(); case 'PENDING_QUEUE_FULL': return facadeError(429, 'KILO_PROMPT_QUEUE_FULL', result.error); + case 'PAYMENT_REQUIRED': + return facadeError(402, 'KILO_PROMPT_PAYMENT_REQUIRED', result.error); case 'SANDBOX_CONNECT_FAILED': case 'WORKSPACE_SETUP_FAILED': case 'KILO_SERVER_FAILED': diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index aea565ba2d..3a3a97364e 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -155,6 +155,7 @@ import { } from '../session/wrapper-supervisor.js'; import { emitRunStateReport } from '../telemetry/queue-reports.js'; import { createAgentSandbox, createAgentSandboxLifecycle } from '../agent-sandbox/factory.js'; +import { isCloudAgentContainerBillingEnabled } from '../container-billing-rollout.js'; import type { AgentSandboxLifecycle, SandboxDeleteReason, @@ -731,6 +732,7 @@ export class CloudAgentSession extends DurableObject { ? Promise.resolve({ status: 'absent' }) : createAgentSandbox(this.env, metadata).discoverSessionWrappers(), requestAlarmAtOrBefore: deadline => this.scheduleAlarmAtOrBefore(deadline), + checkBillingAdmission: () => this.containerBillingAdmissionFailure(), }); } @@ -745,6 +747,34 @@ export class CloudAgentSession extends DurableObject { return !(await this.hasDeletionIntent()); } + private async containerBillingAdmissionFailure(): Promise | null> { + const metadata = await this.getMetadata(); + if (!metadata) return null; + if (!isCloudAgentContainerBillingEnabled(this.env, metadata.identity)) return null; + const admission = await createAgentSandbox(this.env, metadata).ensureBillingAdmission(); + if (admission.success) return null; + const paymentRequired = + admission.code === 'insufficient_credits' || admission.code === 'stopping'; + return { + success: false, + code: paymentRequired ? 'PAYMENT_REQUIRED' : 'INTERNAL', + error: paymentRequired + ? 'Container billing requires additional credits' + : 'Container billing admission is temporarily unavailable', + failureBoundary: 'admission', + }; + } + + private async isContainerBillingBlocked(): Promise { + const metadata = await this.getMetadata(); + if (!metadata || !isCloudAgentContainerBillingEnabled(this.env, metadata.identity)) + return false; + return createAgentSandbox(this.env, metadata).isBillingBlocked(); + } + private getSandboxLifecycle(): AgentSandboxLifecycle { if (!this.sandboxLifecycle) { this.sandboxLifecycle = createAgentSandboxLifecycle(this.env, { @@ -885,6 +915,7 @@ export class CloudAgentSession extends DurableObject { deliver: plan => this.executeDirectly(plan), isDeliveryHeld: async () => isWrapperRunFinalizing(await getWrapperRuntimeState(this.ctx.storage)), + checkBillingAdmission: () => this.containerBillingAdmissionFailure(), ensureQueuedMessageEvent: event => { this.ensureQueuedMessageEvent({ executionId: '' as EventSourceId, @@ -941,8 +972,9 @@ export class CloudAgentSession extends DurableObject { wrapperSupervisor: this.getWrapperSupervisor(), handleWrapperTerminalEvent: params => this.handleWrapperTerminalEvent(params), keepContainerAlive: () => { - void this.keepContainerAlive(); + void this.keepContainerAliveIfBillingAllowed(); }, + isBillingBlocked: () => this.isContainerBillingBlocked(), observeCorrelatedAgentActivity: messageId => this.recordCorrelatedAgentActivity(messageId), terminalizeSessionMessageOnce: async (messageId, params, wrapperRunId) => { await this.ensureAcceptedMessageBeforeTerminal(messageId, wrapperRunId); @@ -979,6 +1011,18 @@ export class CloudAgentSession extends DurableObject { return this.ingestHandler; } + private async keepContainerAliveIfBillingAllowed(): Promise { + try { + if (await this.isContainerBillingBlocked()) return; + } catch { + logger + .withFields({ sessionId: this.sessionId, skipped: 'billing-state-unavailable' }) + .warn('Cloud agent skipped sandbox keepalive'); + return; + } + await this.keepContainerAlive(); + } + // --------------------------------------------------------------------------- // HTTP/WebSocket Routing // --------------------------------------------------------------------------- diff --git a/services/cloud-agent-next/src/session/agent-runtime.test.ts b/services/cloud-agent-next/src/session/agent-runtime.test.ts index d8678b93d9..822a19b60e 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.test.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.test.ts @@ -105,6 +105,30 @@ function createWorkspaceReady(): WorkspaceReady { } describe('AgentRuntime', () => { + it('rechecks billing immediately before physical delivery', async () => { + const createSandbox = vi.fn(); + const runtime = createAgentRuntime({ + storage: createMemoryStorage(), + env: { WORKER_URL: 'http://worker.test' } as Env, + getMetadata: async () => createMetadata(), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: vi.fn(), + createAgentSandbox: createSandbox, + checkBillingAdmission: async () => ({ + success: false, + code: 'PAYMENT_REQUIRED', + error: 'Container billing balance is too low', + failureBoundary: 'admission', + }), + }); + + await expect(runtime.send(createPlan())).resolves.toMatchObject({ + success: false, + code: 'PAYMENT_REQUIRED', + }); + expect(createSandbox).not.toHaveBeenCalled(); + }); + it('returns a permanent delivery failure for an intentionally unavailable sandbox capability', async () => { const createSandbox = vi.fn( () => diff --git a/services/cloud-agent-next/src/session/agent-runtime.ts b/services/cloud-agent-next/src/session/agent-runtime.ts index b37bc18d4c..4a02d12cbd 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.ts @@ -12,6 +12,7 @@ import type { FencedWrapperDispatchRequest, MessageDeliveryRequest, MessageDeliveryResult, + AdmissionFailure, WorkspaceReady, } from '../execution/types.js'; import { logger } from '../logger.js'; @@ -91,6 +92,7 @@ export type AgentRuntimeDependencies = { createAgentSandbox?: (metadata: SessionMetadata) => AgentSandbox; discoverSessionWrappers?: (metadata: SessionMetadata) => Promise; requestAlarmAtOrBefore?: (deadline: number) => Promise; + checkBillingAdmission?: () => Promise; }; function cleanupBlockedError(lease: WrapperLease): WrapperCleanupBlockedError { @@ -319,6 +321,8 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen if (canUseSandboxRuntime && !(await canUseSandboxRuntime())) { return { success: false, code: 'INTERNAL', error: 'Session deletion is in progress' }; } + const billingFailure = await dependencies.checkBillingAdmission?.(); + if (billingFailure) return billingFailure; const { sessionId } = plan.scope; const { turn, agent } = plan; const currentRuntimeState = await getWrapperRuntimeState(storage); diff --git a/services/cloud-agent-next/src/session/pending-messages.ts b/services/cloud-agent-next/src/session/pending-messages.ts index 04a2bbcefc..cda48d25d9 100644 --- a/services/cloud-agent-next/src/session/pending-messages.ts +++ b/services/cloud-agent-next/src/session/pending-messages.ts @@ -96,6 +96,7 @@ const PendingFlushFailureCodeSchema = z.enum([ 'SANDBOX_CAPABILITY_UNAVAILABLE', 'NOT_FOUND', 'BAD_REQUEST', + 'PAYMENT_REQUIRED', 'INTERNAL', 'PENDING_QUEUE_FULL', 'MODEL_MISSING', @@ -535,6 +536,7 @@ export async function recordPendingFlushFailure( | 'WRAPPER_CLEANUP_EXHAUSTED' | 'NOT_FOUND' | 'BAD_REQUEST' + | 'PAYMENT_REQUIRED' | 'INTERNAL' | 'PENDING_QUEUE_FULL' | 'MODEL_MISSING' @@ -630,6 +632,7 @@ function isRetryableFlushCode( | 'WRAPPER_CLEANUP_EXHAUSTED' | 'NOT_FOUND' | 'BAD_REQUEST' + | 'PAYMENT_REQUIRED' | 'INTERNAL' | 'PENDING_QUEUE_FULL' | 'MODEL_MISSING' diff --git a/services/cloud-agent-next/src/session/queue-message.ts b/services/cloud-agent-next/src/session/queue-message.ts index e42a07e063..0deef65b42 100644 --- a/services/cloud-agent-next/src/session/queue-message.ts +++ b/services/cloud-agent-next/src/session/queue-message.ts @@ -43,6 +43,7 @@ type TRPCCodeName = ConstructorParameters[0]['code']; const ADMISSION_CODE_TO_TRPC: Record = { NOT_FOUND: 'NOT_FOUND', BAD_REQUEST: 'BAD_REQUEST', + PAYMENT_REQUIRED: 'PAYMENT_REQUIRED', PENDING_QUEUE_FULL: 'TOO_MANY_REQUESTS', INTERNAL: 'INTERNAL_SERVER_ERROR', }; diff --git a/services/cloud-agent-next/src/session/session-message-queue.test.ts b/services/cloud-agent-next/src/session/session-message-queue.test.ts index 9c5254eb75..282884801e 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.test.ts @@ -12,6 +12,7 @@ import { buildCloudMessageFailedPayload } from './message-settlement-outbox.js'; import { createSessionMessageQueue, flushNextPendingSessionMessage, + type SessionMessageQueueDependencies, type SessionMessageQueueStorage, } from './session-message-queue.js'; import { @@ -130,6 +131,7 @@ function createQueueHarness(options?: { ensureAcceptedMessageEffects?: (messageId: string) => Promise; getDeliveryBlock?: () => Promise; recoverExhaustedDeliveryBlock?: () => Promise; + checkBillingAdmission?: SessionMessageQueueDependencies['checkBillingAdmission']; }) { const storage = options?.storage ?? createMemoryStorage(); const events: QueueEvent[] = []; @@ -167,6 +169,7 @@ function createQueueHarness(options?: { getDeliveryContext: async () => (metadata ? createContext(metadata) : null), getDeliveryBlock: options?.getDeliveryBlock ?? (async () => null), recoverExhaustedDeliveryBlock: options?.recoverExhaustedDeliveryBlock, + checkBillingAdmission: options?.checkBillingAdmission, deliver, ensureQueuedMessageEvent: event => { if (failQueuedEvent) { @@ -620,6 +623,27 @@ describe('flushNextPendingSessionMessage', () => { }); describe('SessionMessageQueue', () => { + it('rejects new work before persisting it when container billing is blocked', async () => { + const harness = createQueueHarness({ + checkBillingAdmission: async () => ({ + success: false, + code: 'PAYMENT_REQUIRED', + error: 'Container billing balance is too low', + failureBoundary: 'admission', + }), + }); + + await expect( + harness.queue.admitSubmittedMessage({ + userId: 'user_test' as UserId, + turn: { type: 'prompt', id: FIRST_MESSAGE_ID, prompt: 'do not enqueue this prompt' }, + }) + ).resolves.toMatchObject({ success: false, code: 'PAYMENT_REQUIRED' }); + expect(await listPendingSessionMessages(harness.storage)).toHaveLength(0); + expect(harness.events).toHaveLength(0); + expect(harness.alarmDeadlines).toHaveLength(0); + }); + it('reports whether a message identity already has durable admission state', async () => { const harness = createQueueHarness(); diff --git a/services/cloud-agent-next/src/session/session-message-queue.ts b/services/cloud-agent-next/src/session/session-message-queue.ts index b5b5ca8c0b..060674b390 100644 --- a/services/cloud-agent-next/src/session/session-message-queue.ts +++ b/services/cloud-agent-next/src/session/session-message-queue.ts @@ -148,6 +148,10 @@ export type SessionMessageQueueDependencies = { recoverExhaustedDeliveryBlock?: () => Promise; deliver: (plan: MessageDeliveryRequest) => Promise; isDeliveryHeld?: () => Promise; + checkBillingAdmission?: () => Promise | null>; ensureQueuedMessageEvent: (event: PersistedQueuedMessageEvent & { entityId: string }) => void; reportQueuedState?: (state: SessionMessageState) => void; ensureAcceptedMessageEffects: (messageId: string) => Promise; @@ -253,6 +257,8 @@ function classifyDeliveryFailure(code: PendingFlushFailureCode | undefined): { case 'BAD_REQUEST': case 'PENDING_QUEUE_FULL': return { failureStage: 'pre_dispatch', failureCode: 'invalid_delivery_request' }; + case 'PAYMENT_REQUIRED': + return { failureStage: 'pre_dispatch', failureCode: 'payment_required' }; case 'MODEL_MISSING': return { failureStage: 'pre_dispatch', failureCode: 'model_missing' }; case 'WRAPPER_CLEANUP_EXHAUSTED': @@ -834,6 +840,8 @@ export function createSessionMessageQueue( const capacityError = await checkPendingQueueCapacity(); if (capacityError) return capacityError; + const billingError = await dependencies.checkBillingAdmission?.(); + if (billingError) return billingError; const metadata = await getMetadata(); const callbackTarget = metadata?.callback?.target; diff --git a/services/cloud-agent-next/src/terminal/access.test.ts b/services/cloud-agent-next/src/terminal/access.test.ts index 63e5f3ddb1..06cc67f5e7 100644 --- a/services/cloud-agent-next/src/terminal/access.test.ts +++ b/services/cloud-agent-next/src/terminal/access.test.ts @@ -99,6 +99,8 @@ function sandboxWithTerminalResult( getRunningTerminalClient: AgentSandbox['getRunningTerminalClient'] ): AgentSandbox { return { + ensureBillingAdmission: vi.fn().mockResolvedValue({ success: true, billingMode: 'shadow' }), + isBillingBlocked: vi.fn().mockResolvedValue(false), ensureWrapper: vi.fn(), discoverSessionWrappers: vi.fn(), observeWrappersWithoutWaking: vi.fn(), diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index 7f0f65329f..f6833a3cd2 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -497,6 +497,9 @@ export type GitTokenService = { }; export type Env = { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED?: string; + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS?: string; + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS?: string; Sandbox: DurableObjectNamespace; /** Durable Object namespace for shared sandbox containers with SCM credential containment */ SandboxContainment: DurableObjectNamespace; diff --git a/services/cloud-agent-next/src/websocket/ingest.test.ts b/services/cloud-agent-next/src/websocket/ingest.test.ts index c350ca8176..945d1f2535 100644 --- a/services/cloud-agent-next/src/websocket/ingest.test.ts +++ b/services/cloud-agent-next/src/websocket/ingest.test.ts @@ -638,6 +638,32 @@ describe('createIngestHandler', () => { expect(state.acceptWebSocket).not.toHaveBeenCalled(); }); + it('rejects wrapper reconnect while container billing is blocked', async () => { + const state = createFakeState(); + const doContext = createFakeDOContext(); + doContext.isBillingBlocked = vi.fn().mockResolvedValue(true); + const handler = createIngestHandler( + state, + createFakeEventQueries(), + SESSION_ID, + vi.fn(), + doContext + ); + + const response = await handler.handleIngestRequest( + makeIngestRequest({ + wrapperRunId: WRAPPER_RUN_ID, + wrapperGeneration: '2', + wrapperConnectionId: 'conn_current', + sessionId: SESSION_ID, + }) + ); + + expect(response.status).toBe(409); + expect(doContext.wrapperSupervisor.checkReconnect).not.toHaveBeenCalled(); + expect(state.acceptWebSocket).not.toHaveBeenCalled(); + }); + itWithWebSocketPair( 'accepts current fenced connection and cancels matching grace', async () => { diff --git a/services/cloud-agent-next/src/websocket/ingest.ts b/services/cloud-agent-next/src/websocket/ingest.ts index 67df9abcc3..bcbaf08d92 100644 --- a/services/cloud-agent-next/src/websocket/ingest.ts +++ b/services/cloud-agent-next/src/websocket/ingest.ts @@ -257,6 +257,7 @@ export type IngestDOContext = { >; handleWrapperTerminalEvent: (params: WrapperTerminalEvent) => Promise; keepContainerAlive?: () => void; + isBillingBlocked?: () => Promise; observeCorrelatedAgentActivity?: (messageId: string) => Promise; terminalizeSessionMessageOnce: ( messageId: string, @@ -406,6 +407,13 @@ export function createIngestHandler( ); } + if (await doContext.isBillingBlocked?.()) { + logger + .withFields({ sessionId, wrapperRunId, wrapperGeneration, wrapperConnectionId }) + .warn('Wrapper ingest rejected: container billing is blocked'); + return new Response('Container billing is blocked', { status: 409 }); + } + const reconnectDecision = await doContext.wrapperSupervisor.checkReconnect({ wrapperRunId, wrapperGeneration, diff --git a/services/cloud-agent-next/worker-configuration.d.ts b/services/cloud-agent-next/worker-configuration.d.ts index cb4dd29e5d..2a8e273cb3 100644 --- a/services/cloud-agent-next/worker-configuration.d.ts +++ b/services/cloud-agent-next/worker-configuration.d.ts @@ -1,6 +1,6 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: 44ae1b44046f9f4e0c310861dcf7080d) -// Runtime types generated with workerd@1.20260603.1 2026-06-03 nodejs_compat +// Generated by Wrangler by running `wrangler types` (hash: e09c335f8af290cc1fa17a4d7f7ab90a) +// Runtime types generated with workerd@1.20260714.1 2026-06-03 nodejs_compat interface __BaseEnv_Env { SHARED_SANDBOX_OVERRIDES: KVNamespace; R2_BUCKET: R2Bucket; @@ -22,6 +22,7 @@ interface __BaseEnv_Env { PER_SESSION_SANDBOX_ORG_IDS?: "*"; GITHUB_TOKEN_CONTAINMENT_ORG_IDS?: ""; GITLAB_TOKEN_CONTAINMENT_ORG_IDS?: ""; + BITBUCKET_TOKEN_CONTAINMENT_ORG_IDS?: ""; KILOCODE_TOKEN_CONTAINMENT_ORG_IDS?: ""; REPO_SNAPSHOT_ORG_IDS?: ""; TOOL_CGROUP_ORG_IDS: "" | "*"; @@ -37,6 +38,7 @@ interface __BaseEnv_Env { R2_ACCESS_KEY_ID: string; R2_SECRET_ACCESS_KEY: string; KILO_SESSION_INGEST_URL: string; + LOG_REJECTED_KILO_URLS: string; WS_ALLOWED_ORIGINS: string; Sandbox: DurableObjectNamespace; SandboxSmall: DurableObjectNamespace; @@ -53,13 +55,16 @@ interface __BaseEnv_Env { 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_RESERVE_MB?: "2048"; TOOL_CGROUP_CPU_WEIGHT?: "50"; + CLOUD_AGENT_CONTAINER_BILLING_ENABLED?: "false"; + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS?: ""; + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS?: ""; } declare namespace Cloudflare { interface GlobalProps { mainModule: typeof import("./src/index"); - durableNamespaces: "Sandbox" | "CloudAgentSession" | "SandboxSmall" | "SandboxDIND" | "UserKiloFacade" | "SandboxCodeReview" | "SandboxContainment" | "SandboxSmallContainment" | "SandboxCodeReviewContainment"; + durableNamespaces: "Sandbox" | "CloudAgentSession" | "SandboxSmall" | "SandboxDIND" | "UserKiloFacade" | "SandboxCodeReview" | "SandboxContainment" | "SandboxSmallContainment" | "SandboxCodeReviewContainment" | "StreamTicketNonceDO"; } interface DevEnv { SHARED_SANDBOX_OVERRIDES: KVNamespace; @@ -82,6 +87,7 @@ declare namespace Cloudflare { PER_SESSION_SANDBOX_ORG_IDS: "*"; GITHUB_TOKEN_CONTAINMENT_ORG_IDS: ""; GITLAB_TOKEN_CONTAINMENT_ORG_IDS: ""; + BITBUCKET_TOKEN_CONTAINMENT_ORG_IDS: ""; KILOCODE_TOKEN_CONTAINMENT_ORG_IDS: ""; REPO_SNAPSHOT_ORG_IDS: ""; TOOL_CGROUP_ORG_IDS: ""; @@ -97,6 +103,7 @@ declare namespace Cloudflare { R2_ACCESS_KEY_ID: string; R2_SECRET_ACCESS_KEY: string; KILO_SESSION_INGEST_URL: string; + LOG_REJECTED_KILO_URLS: string; WS_ALLOWED_ORIGINS: string; Sandbox: DurableObjectNamespace; SandboxSmall: DurableObjectNamespace; @@ -120,7 +127,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; @@ -549,7 +556,8 @@ interface ExecutionContext { readonly exports: Cloudflare.Exports; readonly props: Props; cache?: CacheContext; - tracing?: Tracing; + readonly access?: CloudflareAccessContext; + tracing: Tracing; } type ExportedHandlerFetchHandler = (request: Request>, env: Env, ctx: ExecutionContext) => Response | Promise; type ExportedHandlerConnectHandler = (socket: Socket, env: Env, ctx: ExecutionContext) => void | Promise; @@ -605,6 +613,10 @@ interface CachePurgeOptions { interface CacheContext { purge(options: CachePurgeOptions): Promise; } +interface CloudflareAccessContext { + readonly aud: string; + getIdentity(): Promise; +} declare abstract class ColoLocalActorNamespace { get(actorId: string): Fetcher; } @@ -634,11 +646,11 @@ declare abstract class DurableObjectNamespace; jurisdiction(jurisdiction: DurableObjectJurisdiction): DurableObjectNamespace; } -type DurableObjectJurisdiction = "eu" | "fedramp" | "fedramp-high"; +type DurableObjectJurisdiction = "eu" | "fedramp" | "fedramp-high" | "us"; interface DurableObjectNamespaceNewUniqueIdOptions { jurisdiction?: DurableObjectJurisdiction; } -type DurableObjectLocationHint = "wnam" | "enam" | "sam" | "weur" | "eeur" | "apac" | "oc" | "afr" | "me"; +type DurableObjectLocationHint = "wnam" | "enam" | "sam" | "weur" | "eeur" | "apac" | "apac-ne" | "apac-se" | "oc" | "afr" | "me"; type DurableObjectRoutingMode = "primary-only"; interface DurableObjectNamespaceGetDurableObjectOptions { locationHint?: DurableObjectLocationHint; @@ -734,6 +746,7 @@ interface DurableObjectFacets { get(name: string, getStartupOptions: () => FacetStartupOptions | Promise>): Fetcher; abort(name: string, reason: any): void; delete(name: string): void; + clone(src: string, dst: string): void; } interface FacetStartupOptions { id?: DurableObjectId | string; @@ -3409,6 +3422,28 @@ interface EventSourceEventSourceInit { withCredentials?: boolean; fetcher?: Fetcher; } +interface ExecOutput { + readonly stdout: ArrayBuffer; + readonly stderr: ArrayBuffer; + readonly exitCode: number; +} +interface ContainerExecOptions { + cwd?: string; + env?: Record; + user?: string; + stdin?: ReadableStream | "pipe"; + stdout?: "pipe" | "ignore"; + stderr?: "pipe" | "ignore" | "combined"; +} +interface ExecProcess { + readonly stdin: WritableStream | null; + readonly stdout: ReadableStream | null; + readonly stderr: ReadableStream | null; + readonly pid: number; + readonly exitCode: Promise; + output(): Promise; + kill(signal?: number): void; +} interface Container { get running(): boolean; start(options?: ContainerStartupOptions): void; @@ -3422,6 +3457,7 @@ interface Container { snapshotDirectory(options: ContainerDirectorySnapshotOptions): Promise; snapshotContainer(options: ContainerSnapshotOptions): Promise; interceptOutboundHttps(addr: string, binding: Fetcher): Promise; + exec(cmd: string[], options?: ContainerExecOptions): Promise; } interface ContainerDirectorySnapshot { id: string; @@ -3592,11 +3628,58 @@ declare abstract class Performance { } interface Tracing { enterSpan(name: string, callback: (span: Span, ...args: A) => T, ...args: A): T; + startActiveSpan(name: string, callback: (span: Span, ...args: A) => T, ...args: A): T; Span: typeof Span; } declare abstract class Span { get isTraced(): boolean; setAttribute(key: string, value?: (boolean | number | string)): void; + end(): void; +} +/** + * Represents the identity of a user authenticated via Cloudflare Access. + * This matches the result of calling /cdn-cgi/access/get-identity. + * + * The exact structure of the returned object depends on the identity provider + * configuration for the Access application. The fields below represent commonly + * available properties, but additional provider-specific fields may be present. + */ +interface CloudflareAccessIdentity extends Record { + /** The user's email address, if available from the identity provider. */ + email?: string; + /** The user's display name. */ + name?: string; + /** The user's unique identifier. */ + user_uuid?: string; + /** The Cloudflare account ID. */ + account_id?: string; + /** Login timestamp (Unix epoch seconds). */ + iat?: number; + /** The user's IP address at authentication time. */ + ip?: string; + /** Authentication methods used (e.g., "pwd"). */ + amr?: string[]; + /** Identity provider information. */ + idp?: { + id: string; + type: string; + }; + /** Geographic information about where the user authenticated. */ + geo?: { + country: string; + }; + /** Group memberships from the identity provider. */ + groups?: Array<{ + id: string; + name: string; + email?: string; + }>; + /** Device posture check results, keyed by check ID. */ + devicePosture?: Record; + /** True if the user connected via Cloudflare WARP. */ + is_warp?: boolean; + /** True if the user is authenticated via Cloudflare Gateway. */ + is_gateway?: boolean; } // ============================================================================ // Agent Memory @@ -11189,78 +11272,6 @@ declare abstract class BrowserRun { */ quickAction(action: 'markdown', options: BrowserRunMarkdownOptions): Promise; } -interface BasicImageTransformations { - /** - * Maximum width in image pixels. The value must be an integer. - */ - width?: number; - /** - * Maximum height in image pixels. The value must be an integer. - */ - height?: number; - /** - * Resizing mode as a string. It affects interpretation of width and height - * options: - * - scale-down: Similar to contain, but the image is never enlarged. If - * the image is larger than given width or height, it will be resized. - * Otherwise its original size will be kept. - * - contain: Resizes to maximum size that fits within the given width and - * height. If only a single dimension is given (e.g. only width), the - * image will be shrunk or enlarged to exactly match that dimension. - * Aspect ratio is always preserved. - * - cover: Resizes (shrinks or enlarges) to fill the entire area of width - * and height. If the image has an aspect ratio different from the ratio - * of width and height, it will be cropped to fit. - * - crop: The image will be shrunk and cropped to fit within the area - * specified by width and height. The image will not be enlarged. For images - * smaller than the given dimensions it's the same as scale-down. For - * images larger than the given dimensions, it's the same as cover. - * See also trim. - * - pad: Resizes to the maximum size that fits within the given width and - * height, and then fills the remaining area with a background color - * (white by default). Use of this mode is not recommended, as the same - * effect can be more efficiently achieved with the contain mode and the - * CSS object-fit: contain property. - * - squeeze: Stretches and deforms to the width and height given, even if it - * breaks aspect ratio - */ - fit?: "scale-down" | "contain" | "cover" | "crop" | "pad" | "squeeze"; - /** - * Image segmentation using artificial intelligence models. Sets pixels not - * within selected segment area to transparent e.g "foreground" sets every - * background pixel as transparent. - */ - segment?: "foreground"; - /** - * When cropping with fit: "cover", this defines the side or point that should - * be left uncropped. The value is either a string - * "left", "right", "top", "bottom", "auto", or "center" (the default), - * or an object {x, y} containing focal point coordinates in the original - * image expressed as fractions ranging from 0.0 (top or left) to 1.0 - * (bottom or right), 0.5 being the center. {fit: "cover", gravity: "top"} will - * crop bottom or left and right sides as necessary, but won’t crop anything - * from the top. {fit: "cover", gravity: {x:0.5, y:0.2}} will crop each side to - * preserve as much as possible around a point at 20% of the height of the - * source image. - */ - gravity?: 'face' | 'left' | 'right' | 'top' | 'bottom' | 'center' | 'auto' | 'entropy' | BasicImageTransformationsGravityCoordinates; - /** - * Background color to add underneath the image. Applies only to images with - * transparency (such as PNG). Accepts any CSS color (#RRGGBB, rgba(…), - * hsl(…), etc.) - */ - background?: string; - /** - * Number of degrees (90, 180, 270) to rotate the image by. width and height - * options refer to axes after rotation. - */ - rotate?: 0 | 90 | 180 | 270 | 360; -} -interface BasicImageTransformationsGravityCoordinates { - x?: number; - y?: number; - mode?: 'remainder' | 'box-center'; -} /** * In addition to the properties you can set in the RequestInit dict * that you pass as an argument to the Request constructor, you can @@ -11299,6 +11310,8 @@ interface RequestInitCfProperties extends Record { * (e.g. { '200-299': 86400, '404': 1, '500-599': 0 }) */ cacheTtlByStatus?: Record; + /** Controls how responses with a `Vary` header are cached for this request. */ + vary?: RequestInitCfPropertiesVary; /** * Explicit Cache-Control header value to set on the response stored in cache. * This gives full control over cache directives (e.g. 'public, max-age=3600, s-maxage=86400'). @@ -11336,6 +11349,17 @@ interface RequestInitCfProperties extends Record { cacheReserveMinimumFileSize?: number; scrapeShield?: boolean; apps?: boolean; + /** + * Controls whether an outbound gRPC-web subrequest from this Worker is + * converted to gRPC at the Cloudflare edge. + * + * - `"passthrough"`: forward the subrequest unchanged as gRPC-web (default). + * - `"convert"`: convert the gRPC-web subrequest to gRPC at the edge. + * + * Provides per-request control over the same edge conversion behavior + * gated by the `auto_grpc_convert` compatibility flag. + */ + grpcWeb?: "passthrough" | "convert"; image?: RequestInitCfPropertiesImage; minify?: RequestInitCfPropertiesImageMinify; mirage?: boolean; @@ -11356,6 +11380,244 @@ interface RequestInitCfProperties extends Record { */ resolveOverride?: string; } +/** + * Controls how Workers Standard Vary handles a request header listed by an + * origin `Vary` response header: + * + * - `"normalize"`: normalize the request header value before it is used in the + * cache variance key. + * - `"passthrough"`: use the raw request header value in the cache variance + * key. + * - `"bypass"`: bypass cache when the header appears in the origin `Vary` + * response header. + */ +type RequestInitCfPropertiesVaryAction = "normalize" | "passthrough" | "bypass"; +/** Configuration for Workers Standard Vary support. */ +interface RequestInitCfPropertiesVary { + /** The fallback action for varied request headers not listed in `headers`. */ + default: RequestInitCfPropertiesVaryHeader; + /** + * Lowercase request header names and their Vary configuration. + * + * The `accept` header can include `media_types`, the `accept-language` + * header can include `languages`, and other headers support only `action`. + */ + headers?: RequestInitCfPropertiesVaryHeaders; +} +/** Common Vary behavior for a single request header. */ +interface RequestInitCfPropertiesVaryHeader { + /** How this request header contributes to cache variance. */ + action: RequestInitCfPropertiesVaryAction; +} +/** Vary behavior for the `accept` request header. */ +interface RequestInitCfPropertiesVaryAcceptHeader extends RequestInitCfPropertiesVaryHeader { + /** + * Media types to keep when normalizing the `Accept` request header. + * + * Named `media_types` to match the serialized `cf.vary` configuration. + */ + media_types?: string[]; +} +/** Vary behavior for the `accept-language` request header. */ +interface RequestInitCfPropertiesVaryAcceptLanguageHeader extends RequestInitCfPropertiesVaryHeader { + /** + * Language tags to keep when normalizing the `Accept-Language` request + * header. + */ + languages?: string[]; +} +/** + * Lowercase request header names and their Vary behavior. + * + * The index signature allows arbitrary custom request headers beyond the + * well-known `accept` and `accept-language` specializations. + */ +interface RequestInitCfPropertiesVaryHeaders { + accept?: RequestInitCfPropertiesVaryAcceptHeader; + "accept-language"?: RequestInitCfPropertiesVaryAcceptLanguageHeader; + [header: string]: RequestInitCfPropertiesVaryHeader | RequestInitCfPropertiesVaryAcceptHeader | RequestInitCfPropertiesVaryAcceptLanguageHeader | undefined; +} +interface BasicImageTransformations { + /** + * Maximum width in image pixels. The value must be an integer. + */ + width?: number; + /** + * Maximum height in image pixels. The value must be an integer. + */ + height?: number; + /** + * When cropping with fit: "cover", this defines the side or point that should + * be left uncropped. The value is either a string + * "left", "right", "top", "bottom", "auto", or "center" (the default), + * or an object {x, y} containing focal point coordinates in the original + * image expressed as fractions ranging from 0.0 (top or left) to 1.0 + * (bottom or right), 0.5 being the center. {fit: "cover", gravity: "top"} will + * crop bottom or left and right sides as necessary, but won’t crop anything + * from the top. {fit: "cover", gravity: {x:0.5, y:0.2}} will crop each side to + * preserve as much as possible around a point at 20% of the height of the + * source image. + */ + gravity?: 'face' | 'left' | 'right' | 'top' | 'bottom' | 'center' | 'auto' | 'entropy' | BasicImageTransformationsGravityCoordinates; + /** + * Specifies how closely the image is cropped toward detected faces when combined + * with the gravity=face option. Accepts a valid range between 0.0 (includes as much + * of the background as possible) and 1.0 (crops the image as closely to the face as + * possible). The default is 0. + */ + zoom?: number; + /** + * Resizing mode as a string. It affects interpretation of width and height + * options: + * - scale-down: Similar to contain, but the image is never enlarged. If + * the image is larger than given width or height, it will be resized. + * Otherwise its original size will be kept. + * - scale-up: Similar to contain, but the image is never shrunk. If the + * image is smaller than the given width or height, it will be resized. + * Otherwise its original size will be kept. + * - contain: Resizes to maximum size that fits within the given width and + * height. If only a single dimension is given (e.g. only width), the + * image will be shrunk or enlarged to exactly match that dimension. + * Aspect ratio is always preserved. + * - cover: Resizes (shrinks or enlarges) to fill the entire area of width + * and height. If the image has an aspect ratio different from the ratio + * of width and height, it will be cropped to fit. + * - crop: The image will be shrunk and cropped to fit within the area + * specified by width and height. The image will not be enlarged. For images + * smaller than the given dimensions it's the same as scale-down. For + * images larger than the given dimensions, it's the same as cover. + * See also trim. + * - pad: Resizes to the maximum size that fits within the given width and + * height, and then fills the remaining area with a background color + * (white by default). Use of this mode is not recommended, as the same + * effect can be more efficiently achieved with the contain mode and the + * CSS object-fit: contain property. + * - squeeze: Stretches and deforms to the width and height given, even if it + * breaks aspect ratio + */ + fit?: "scale-down" | "scale-up" | "contain" | "cover" | "crop" | "pad" | "squeeze"; + /** + * Allows you to trim your image. Takes dpr into account and is performed before + * resizing or rotation. + * + * It can be used as: + * - left, top, right, bottom - it will specify the number of pixels to cut + * off each side + * - width, height - the width/height you'd like to end up with - can be used + * in combination with the properties above + * - border - this will automatically trim the surroundings of an image based on + * it's color. It consists of three properties: + * - color: rgb or hex representation of the color you wish to trim (todo: verify the rgba bit) + * - tolerance: difference from color to treat as color + * - keep: the number of pixels of border to keep + */ + trim?: "border" | { + top?: number; + bottom?: number; + left?: number; + right?: number; + width?: number; + height?: number; + border?: boolean | { + color?: string; + tolerance?: number; + keep?: number; + }; + }; + /** + * Background color to add underneath the image. Applies only to images with + * transparency (such as PNG). Accepts any CSS color (#RRGGBB, rgba(…), + * hsl(…), etc.) + */ + background?: string; + /** + * Flips the images horizontally, vertically, or both. Flipping is applied before + * rotation, so if you apply flip=h,rotate=90 then the image will be flipped + * horizontally, then rotated by 90 degrees. + */ + flip?: 'h' | 'v' | 'hv'; + /** + * Number of degrees (90, 180, 270) to rotate the image by. width and height + * options refer to axes after rotation. + */ + rotate?: 0 | 90 | 180 | 270 | 360; + /** + * Strength of sharpening filter to apply to the image. Floating-point + * number between 0 (no sharpening, default) and 10 (maximum). 1.0 is a + * recommended value for downscaled images. + */ + sharpen?: number; + /** + * Radius of a blur filter (approximate gaussian). Maximum supported radius + * is 250. + */ + blur?: number; + /** + * Increase contrast by a factor. A value of 1.0 equals no change, a value of + * 0.5 equals low contrast, and a value of 2.0 equals high contrast. 0 is + * ignored. + */ + contrast?: number; + /** + * Increase brightness by a factor. A value of 1.0 equals no change, a value + * of 0.5 equals half brightness, and a value of 2.0 equals twice as bright. + * 0 is ignored. + */ + brightness?: number; + /** + * Increase exposure by a factor. A value of 1.0 equals no change, a value of + * 0.5 darkens the image, and a value of 2.0 lightens the image. 0 is ignored. + */ + gamma?: number; + /** + * Increase contrast by a factor. A value of 1.0 equals no change, a value of + * 0.5 equals low contrast, and a value of 2.0 equals high contrast. 0 is + * ignored. + */ + saturation?: number; + /** + * Device Pixel Ratio. Default 1. Multiplier for width/height that makes it + * easier to specify higher-DPI sizes in . + */ + dpr?: number; + /** + * Adds a border around the image. The border is added after resizing. Border + * width takes dpr into account, and can be specified either using a single + * width property, or individually for each side. + */ + border?: { + color: string; + width: number; + } | { + color: string; + top: number; + right: number; + bottom: number; + left: number; + }; + /** + * Image segmentation using artificial intelligence models. Sets pixels not + * within selected segment area to transparent e.g "foreground" sets every + * background pixel as transparent. + */ + segment?: "foreground"; + /** + * Controls the algorithm used when an image needs to be enlarged. This + * parameter works with any fit mode that upscales, such as `contain`, + * `cover`, and `scale-up`. It has no effect when `fit=scale-down` or when + * the target dimensions are smaller than the source. + * - interpolate: Uses bicubic interpolation, which may reduce image quality. + * This is the default behavior when `upscale` is not specified. + * - generate: Uses AI upscaling to produce sharper, more detailed results + * when enlarging images. + */ + upscale?: "interpolate" | "generate"; +} +interface BasicImageTransformationsGravityCoordinates { + x?: number; + y?: number; + mode?: 'remainder' | 'box-center'; +} interface RequestInitCfPropertiesImageDraw extends BasicImageTransformations { /** * Absolute URL of the image file to use for the drawing. It can be any of @@ -11410,39 +11672,6 @@ interface RequestInitCfPropertiesImageDraw extends BasicImageTransformations { right?: number; } interface RequestInitCfPropertiesImage extends BasicImageTransformations { - /** - * Device Pixel Ratio. Default 1. Multiplier for width/height that makes it - * easier to specify higher-DPI sizes in . - */ - dpr?: number; - /** - * Allows you to trim your image. Takes dpr into account and is performed before - * resizing or rotation. - * - * It can be used as: - * - left, top, right, bottom - it will specify the number of pixels to cut - * off each side - * - width, height - the width/height you'd like to end up with - can be used - * in combination with the properties above - * - border - this will automatically trim the surroundings of an image based on - * it's color. It consists of three properties: - * - color: rgb or hex representation of the color you wish to trim (todo: verify the rgba bit) - * - tolerance: difference from color to treat as color - * - keep: the number of pixels of border to keep - */ - trim?: "border" | { - top?: number; - bottom?: number; - left?: number; - right?: number; - width?: number; - height?: number; - border?: boolean | { - color?: string; - tolerance?: number; - keep?: number; - }; - }; /** * Quality setting from 1-100 (useful values are in 60-90 range). Lower values * make images look worse, but load faster. The default is 85. It applies only @@ -11484,17 +11713,6 @@ interface RequestInitCfPropertiesImage extends BasicImageTransformations { * output formats always discard metadata. */ metadata?: "keep" | "copyright" | "none"; - /** - * Strength of sharpening filter to apply to the image. Floating-point - * number between 0 (no sharpening, default) and 10 (maximum). 1.0 is a - * recommended value for downscaled images. - */ - sharpen?: number; - /** - * Radius of a blur filter (approximate gaussian). Maximum supported radius - * is 250. - */ - blur?: number; /** * Overlays are drawn in the order they appear in the array (last array * entry is the topmost layer). @@ -11506,50 +11724,6 @@ interface RequestInitCfPropertiesImage extends BasicImageTransformations { * the origin. */ "origin-auth"?: "share-publicly"; - /** - * Adds a border around the image. The border is added after resizing. Border - * width takes dpr into account, and can be specified either using a single - * width property, or individually for each side. - */ - border?: { - color: string; - width: number; - } | { - color: string; - top: number; - right: number; - bottom: number; - left: number; - }; - /** - * Increase brightness by a factor. A value of 1.0 equals no change, a value - * of 0.5 equals half brightness, and a value of 2.0 equals twice as bright. - * 0 is ignored. - */ - brightness?: number; - /** - * Increase contrast by a factor. A value of 1.0 equals no change, a value of - * 0.5 equals low contrast, and a value of 2.0 equals high contrast. 0 is - * ignored. - */ - contrast?: number; - /** - * Increase exposure by a factor. A value of 1.0 equals no change, a value of - * 0.5 darkens the image, and a value of 2.0 lightens the image. 0 is ignored. - */ - gamma?: number; - /** - * Increase contrast by a factor. A value of 1.0 equals no change, a value of - * 0.5 equals low contrast, and a value of 2.0 equals high contrast. 0 is - * ignored. - */ - saturation?: number; - /** - * Flips the images horizontally, vertically, or both. Flipping is applied before - * rotation, so if you apply flip=h,rotate=90 then the image will be flipped - * horizontally, then rotated by 90 degrees. - */ - flip?: 'h' | 'v' | 'hv'; /** * Slightly reduces latency on a cache miss by selecting a * quickest-to-compress file format, at a cost of increased file size and @@ -12167,6 +12341,13 @@ interface ForwardableEmailMessage extends EmailMessage { * @returns A promise that resolves when the email message is replied. */ reply(message: EmailMessage): Promise; + /** + * Reply to the sender of this email message with a message built from the given + * fields. Threading headers (In-Reply-To/References) are set automatically. + * @param builder The reply message contents. + * @returns A promise that resolves when the email message is replied. + */ + reply(builder: EmailReplyMessageBuilder): Promise; } /** A file attachment for an email message */ type EmailAttachment = { @@ -12187,23 +12368,46 @@ interface EmailAddress { name: string; email: string; } +/** + * Recipient fields for `SendEmail.send()`. At least one of `to`, `cc`, or + * `bcc` must be provided. + */ +type EmailDestinations = { + to?: string | EmailAddress | (string | EmailAddress)[]; + cc?: string | EmailAddress | (string | EmailAddress)[]; + bcc?: string | EmailAddress | (string | EmailAddress)[]; +} & ({ + to: string | EmailAddress | (string | EmailAddress)[]; +} | { + cc: string | EmailAddress | (string | EmailAddress)[]; +} | { + bcc: string | EmailAddress | (string | EmailAddress)[]; +}); +/** + * Fields shared by all composed emails (no recipients). Used directly by + * `ForwardableEmailMessage.reply()`, which always replies to the original + * sender, and extended by `EmailMessageBuilder` for `SendEmail.send()`. + */ +interface EmailReplyMessageBuilder { + from: string | EmailAddress; + subject: string; + replyTo?: string | EmailAddress; + headers?: Record; + text?: string; + html?: string; + attachments?: EmailAttachment[]; +} +/** + * Fields for composing an email without constructing raw MIME, for + * `SendEmail.send()`. Requires at least one of `to`, `cc`, or `bcc`. + */ +type EmailMessageBuilder = EmailReplyMessageBuilder & EmailDestinations; /** * A binding that allows a Worker to send email messages. */ interface SendEmail { send(message: EmailMessage): Promise; - send(builder: { - from: string | EmailAddress; - to: string | EmailAddress | (string | EmailAddress)[]; - subject: string; - replyTo?: string | EmailAddress; - cc?: string | EmailAddress | (string | EmailAddress)[]; - bcc?: string | EmailAddress | (string | EmailAddress)[]; - headers?: Record; - text?: string; - html?: string; - attachments?: EmailAttachment[]; - }): Promise; + send(builder: EmailMessageBuilder): Promise; } declare abstract class EmailEvent extends ExtendableEvent { readonly message: ForwardableEmailMessage; @@ -13025,6 +13229,11 @@ declare namespace CloudflareWorkersModule { export type WorkflowDurationLabel = 'second' | 'minute' | 'hour' | 'day' | 'week' | 'month' | 'year'; export type WorkflowSleepDuration = `${number} ${WorkflowDurationLabel}${'s' | ''}` | number; export type WorkflowDelayDuration = WorkflowSleepDuration; + export type WorkflowDynamicDelayContext = { + ctx: WorkflowStepContext; + error: Error; + }; + export type WorkflowDelayFunction = (input: WorkflowDynamicDelayContext) => WorkflowDelayDuration | Promise; export type WorkflowTimeoutDuration = WorkflowSleepDuration; export type WorkflowRetentionDuration = WorkflowSleepDuration; export type WorkflowBackoff = 'constant' | 'linear' | 'exponential'; @@ -13032,12 +13241,13 @@ declare namespace CloudflareWorkersModule { export type WorkflowStepConfig = { retries?: { limit: number; - delay: WorkflowDelayDuration | number; + delay: WorkflowDelayDuration | number | WorkflowDelayFunction; backoff?: WorkflowBackoff; }; timeout?: WorkflowTimeoutDuration | number; sensitive?: WorkflowStepSensitivity; }; + export type WorkflowStepRollbackConfig = Pick; export type WorkflowCronSchedule = { /** Cron expression that triggered this event. */ cron: string; @@ -13057,27 +13267,40 @@ declare namespace CloudflareWorkersModule { type: string; sensitive?: WorkflowStepSensitivity; }; - export type WorkflowStepContext = { + export type WorkflowStepContext = { step: { name: string; count: number; }; attempt: number; - config: WorkflowStepConfig; + config: { + retries?: { + limit: number; + backoff?: WorkflowBackoff; + } & (Delay extends WorkflowDelayFunction ? {} : { + delay: WorkflowDelayDuration | number; + }); + timeout?: WorkflowTimeoutDuration | number; + sensitive?: WorkflowStepSensitivity; + }; }; export type WorkflowRollbackContext = { + ctx: WorkflowStepContext; error: Error; output: T | undefined; + /** @deprecated Use `ctx.step.name` and `ctx.step.count` instead. */ stepName: string; }; export type WorkflowRollbackHandler = (ctx: WorkflowRollbackContext) => Promise; export type WorkflowStepRollbackOptions = { - rollback?: WorkflowRollbackHandler; - rollbackConfig?: WorkflowStepConfig; + rollback: WorkflowRollbackHandler; + rollbackConfig?: WorkflowStepRollbackConfig; }; export abstract class WorkflowStep { do>(name: string, callback: (ctx: WorkflowStepContext) => Promise, rollbackOptions?: WorkflowStepRollbackOptions): Promise; - do>(name: string, config: WorkflowStepConfig, callback: (ctx: WorkflowStepContext) => Promise, rollbackOptions?: WorkflowStepRollbackOptions): Promise; + do, const C extends WorkflowStepConfig>(name: string, config: C, callback: (ctx: WorkflowStepContext) => Promise, rollbackOptions?: WorkflowStepRollbackOptions): Promise; sleep: (name: string, duration: WorkflowSleepDuration) => Promise; sleepUntil: (name: string, timestamp: Date | number) => Promise; waitForEvent>(name: string, options: { @@ -13968,7 +14191,7 @@ declare namespace TailStream { interface ConnectEventInfo { readonly type: "connect"; } - type EventOutcome = "ok" | "canceled" | "exception" | "unknown" | "killSwitch" | "daemonDown" | "exceededCpu" | "exceededMemory" | "loadShed" | "responseStreamDisconnected" | "scriptNotFound" | "internalError"; + type EventOutcome = "ok" | "canceled" | "exception" | "unknown" | "killSwitch" | "daemonDown" | "exceededCpu" | "exceededMemory" | "loadShed" | "responseStreamDisconnected" | "scriptNotFound" | "internalError" | "exceededWallTime"; interface ScriptVersion { readonly id: string; readonly tag?: string; @@ -13987,6 +14210,7 @@ declare namespace TailStream { readonly dispatchNamespace?: string; readonly entrypoint?: string; readonly executionModel: string; + readonly durableObjectId?: string; readonly scriptName?: string; readonly scriptTags?: string[]; readonly scriptVersion?: ScriptVersion; @@ -14541,6 +14765,13 @@ interface WorkflowError { code?: number; message: string; } +interface WorkflowInstanceTerminateOptions { + /** + * If true, run registered rollback handlers before terminating the instance. + * Only steps that registered rollback handlers are rolled back. + */ + rollback?: boolean; +} interface WorkflowInstanceRestartOptions { /** * Restart from a specific step. If omitted, the instance restarts from the beginning. @@ -14574,8 +14805,9 @@ declare abstract class WorkflowInstance { public resume(): Promise; /** * Terminate the instance. If it is errored, terminated or complete, an error will be thrown. + * @param options Options for termination, including whether registered rollback handlers should run. */ - public terminate(): Promise; + public terminate(options?: WorkflowInstanceTerminateOptions): Promise; /** * Restart the instance. Optionally restart from a specific step, preserving * cached results for all steps before it. diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index d8e49a6725..71919d0243 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -61,6 +61,9 @@ "TOOL_CGROUP_MODE": "enforce", "TOOL_CGROUP_RESERVE_MB": "2048", "TOOL_CGROUP_CPU_WEIGHT": "50", + "CLOUD_AGENT_CONTAINER_BILLING_ENABLED": "false", + "CLOUD_AGENT_CONTAINER_BILLING_USER_IDS": "", + "CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS": "", }, "placement": { "mode": "smart" }, /** diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.ts b/services/cloud-agent-next/wrapper/src/lifecycle.ts index 791011eed5..422f40dd5c 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -29,6 +29,7 @@ export type LifecycleManager = { onDeliveryAcknowledged: (kind: 'async-prompt' | 'sync-command' | 'failed') => void; onConnectionRestored: () => void; triggerDrainAndClose: () => void; + drainAndClose: () => Promise; signalCompletion: () => void; setAborted: () => void; reset: () => void; @@ -42,12 +43,12 @@ export function createLifecycleManager( const { state, kiloClient } = deps; let sseTransportTimer: ReturnType | null = null; let stableIdleTimer: ReturnType | null = null; - let drainTimeout: ReturnType | null = null; - let isDraining = false; let isAborted = false; let rootIdleCandidatePresent = false; let idleObservedDuringDelivery = false; let postProcessingResolve: (() => void) | null = null; + let drainPromise: Promise | null = null; + let lifecycleGeneration = 0; let postProcessingCompleted = false; function clearSseTransportTimer(): void { @@ -68,7 +69,7 @@ export function createLifecycleManager( clearSseTransportTimer(); // Idle is the last expected SSE event. Reconnecting during drain races // auto-commit and used to abort the complete event. - if (!state.hasSession || isDraining) return; + if (!state.hasSession || drainPromise) return; sseTransportTimer = setTimeout(() => { logToFile('SSE transport timeout — reconnecting event subscription'); deps.reconnectEventSubscription(); @@ -153,10 +154,10 @@ export function createLifecycleManager( } } - function triggerDrainAndClose(): void { + function drainAndClose(): Promise { state.blockAdmissions(); - if (isDraining) return; - isDraining = true; + if (drainPromise) return drainPromise; + const drainGeneration = lifecycleGeneration; clearStableIdleCandidate(); clearSseTransportTimer(); const sealedMessageIds = state.pendingMessageIds; @@ -173,7 +174,7 @@ export function createLifecycleManager( }); } - void (async () => { + drainPromise = (async () => { try { await runPostCompletionTasks(); const uploader = state.logUploader; @@ -188,6 +189,7 @@ export function createLifecycleManager( uploader.stop(); } } finally { + if (drainGeneration !== lifecycleGeneration) return; const currentSession = state.currentSession; if (completeSession && currentSession) { const currentBranch = await getCurrentBranch(config.workspacePath, 10_000).catch( @@ -207,22 +209,30 @@ export function createLifecycleManager( }); } - if (isDraining) { - drainTimeout = setTimeout(() => { - void deps - .closeConnections() - .catch(error => - logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) - ) - .finally(() => { - isDraining = false; - drainTimeout = null; - state.clearSession(); - }); - }, DRAIN_DELAY_MS); - } + await new Promise(resolve => setTimeout(resolve, DRAIN_DELAY_MS)); + if (drainGeneration !== lifecycleGeneration) return; + await deps + .closeConnections() + .catch(error => + logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) + ); + state.clearSession(); } })(); + const currentDrain = drainPromise; + void currentDrain.then( + () => { + if (drainPromise === currentDrain) drainPromise = null; + }, + () => { + if (drainPromise === currentDrain) drainPromise = null; + } + ); + return currentDrain; + } + + function triggerDrainAndClose(): void { + void drainAndClose(); } function trySealIdleBatch(): void { @@ -256,8 +266,6 @@ export function createLifecycleManager( isAborted = true; clearSseTransportTimer(); clearStableIdleCandidate(); - if (drainTimeout) clearTimeout(drainTimeout); - drainTimeout = null; }, onSessionIdle: () => { rootIdleCandidatePresent = true; @@ -284,6 +292,7 @@ export function createLifecycleManager( }, onConnectionRestored: armStableIdleCandidate, triggerDrainAndClose, + drainAndClose, signalCompletion, setAborted: () => { isAborted = true; @@ -291,14 +300,13 @@ export function createLifecycleManager( clearStableIdleCandidate(); }, reset: () => { + lifecycleGeneration += 1; isAborted = false; - isDraining = false; clearStableIdleCandidate(); postProcessingCompleted = false; postProcessingResolve = null; clearSseTransportTimer(); - if (drainTimeout) clearTimeout(drainTimeout); - drainTimeout = null; + drainPromise = null; }, onSseEvent: resetSseTransportTimer, }; diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index c882013b7e..5b7b005cfb 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -54,8 +54,8 @@ import { // Constants // --------------------------------------------------------------------------- -/** Grace period before force exit during shutdown (20 seconds) */ -const SHUTDOWN_TIMEOUT_MS = 20_000; +/** Grace period before force exit during shutdown (110 seconds) */ +const SHUTDOWN_TIMEOUT_MS = 110_000; /** Timeout for createKilo() server startup */ const KILO_STARTUP_TIMEOUT_MS = 30_000; @@ -973,7 +973,8 @@ async function main() { logToFile('shutdown startup cleanup finished'); } - // Stop lifecycle timers + await lifecycleManager?.drainAndClose(); + // Stop lifecycle timers after finalization has flushed. lifecycleManager?.stop(); globalFeedManager.close(); toolCgroup?.stop(); diff --git a/services/container-usage-meter/src/billing-config.test.ts b/services/container-usage-meter/src/billing-config.test.ts index 5c850361d6..b7b4c29123 100644 --- a/services/container-usage-meter/src/billing-config.test.ts +++ b/services/container-usage-meter/src/billing-config.test.ts @@ -59,4 +59,82 @@ describe('container billing configuration', () => { 'shadow' ); }); + + it('uses Cloud Agent payer lists only for Cloud Agent service names', () => { + const config = billingConfigFromEnv( + env({ + CONTAINER_BILLING_SERVICES: 'gastown,cloud-agent-next-sandbox', + CONTAINER_BILLING_USER_IDS: 'gastown-user', + CONTAINER_BILLING_ORG_IDS: '', + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS: 'cloud-agent-user', + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS: 'cloud-agent-org', + CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS: '10000000', + }) + ); + + expect( + billingModeFor(config, 'cloud-agent-next-sandbox', { type: 'user', id: 'cloud-agent-user' }) + ).toBe('paid'); + expect( + billingModeFor(config, 'cloud-agent-next-sandbox', { type: 'user', id: 'gastown-user' }) + ).toBe('shadow'); + expect(billingModeFor(config, 'gastown', { type: 'user', id: 'cloud-agent-user' })).toBe( + 'shadow' + ); + }); + + it('requires Cloud Agent service names as well as Cloud Agent payer lists', () => { + const config = billingConfigFromEnv( + env({ + CONTAINER_BILLING_SERVICES: 'gastown', + CONTAINER_BILLING_USER_IDS: '', + CONTAINER_BILLING_ORG_IDS: '', + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS: 'cloud-agent-user', + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS: '', + CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS: '10000000', + }) + ); + expect( + billingModeFor(config, 'cloud-agent-next-sandbox', { type: 'user', id: 'cloud-agent-user' }) + ).toBe('shadow'); + }); + + it('fails Cloud Agent billing closed for empty or malformed Cloud Agent lists', () => { + const empty = billingConfigFromEnv( + env({ + CONTAINER_BILLING_SERVICES: 'gastown,cloud-agent-next-sandbox', + CONTAINER_BILLING_USER_IDS: 'gastown-user', + CONTAINER_BILLING_ORG_IDS: '', + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS: '', + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS: 'cloud-agent-org', + CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS: '10000000', + }) + ); + expect( + billingModeFor(empty, 'cloud-agent-next-sandbox', { type: 'user', id: 'gastown-user' }) + ).toBe('shadow'); + + const malformed = billingConfigFromEnv( + env({ + CONTAINER_BILLING_SERVICES: 'gastown,cloud-agent-next-sandbox', + CONTAINER_BILLING_USER_IDS: 'gastown-user', + CONTAINER_BILLING_ORG_IDS: '', + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS: 'cloud-agent-user,', + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS: '', + CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS: '10000000', + }) + ); + expect( + billingModeFor(malformed, 'cloud-agent-next-sandbox', { + type: 'user', + id: 'cloud-agent-user', + }) + ).toBe('shadow'); + expect( + billingModeFor(malformed, 'cloud-agent-next-sandbox', { + type: 'org', + id: 'cloud-agent-org', + }) + ).toBe('shadow'); + }); }); diff --git a/services/container-usage-meter/src/billing-config.ts b/services/container-usage-meter/src/billing-config.ts index eed55a10c0..e1739bd747 100644 --- a/services/container-usage-meter/src/billing-config.ts +++ b/services/container-usage-meter/src/billing-config.ts @@ -1,10 +1,13 @@ export const MINIMUM_REMAINING_MICRODOLLARS = 5_000_000; const DEFAULT_WARN_REMAINING_MICRODOLLARS = 10_000_000; +const CLOUD_AGENT_SERVICE_PREFIX = 'cloud-agent-next-'; export type BillingConfig = { services: ReadonlySet; userIds: ReadonlySet; orgIds: ReadonlySet; + cloudAgentUserIds: ReadonlySet; + cloudAgentOrgIds: ReadonlySet; warnRemainingMicrodollars: number; enabled: boolean; }; @@ -13,6 +16,8 @@ export const SHADOW_ONLY_BILLING_CONFIG: BillingConfig = { services: new Set(), userIds: new Set(), orgIds: new Set(), + cloudAgentUserIds: new Set(), + cloudAgentOrgIds: new Set(), warnRemainingMicrodollars: DEFAULT_WARN_REMAINING_MICRODOLLARS, enabled: false, }; @@ -21,6 +26,8 @@ type BillingEnvironment = { CONTAINER_BILLING_SERVICES?: string; CONTAINER_BILLING_USER_IDS?: string; CONTAINER_BILLING_ORG_IDS?: string; + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS?: string; + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS?: string; CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS?: string; }; @@ -45,24 +52,39 @@ export function billingConfigFromEnv(env: BillingEnvironment): BillingConfig { const services = parseRequiredAllowlist(env.CONTAINER_BILLING_SERVICES); const userIds = parsePayerAllowlist(env.CONTAINER_BILLING_USER_IDS); const orgIds = parsePayerAllowlist(env.CONTAINER_BILLING_ORG_IDS); + const cloudAgentUserIds = parsePayerAllowlist(env.CONTAINER_BILLING_CLOUD_AGENT_USER_IDS); + const cloudAgentOrgIds = parsePayerAllowlist(env.CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS); const warnRemainingMicrodollars = Number(env.CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS); const validThreshold = Number.isSafeInteger(warnRemainingMicrodollars) && warnRemainingMicrodollars >= MINIMUM_REMAINING_MICRODOLLARS; - const enabled = + const gastownEnabled = services !== null && userIds !== null && orgIds !== null && (userIds.size > 0 || orgIds.size > 0) && validThreshold; + const cloudAgentEnabled = + services !== null && + cloudAgentUserIds !== null && + cloudAgentOrgIds !== null && + (cloudAgentUserIds.size > 0 || cloudAgentOrgIds.size > 0) && + validThreshold; + const cloudAgentListsAreValid = cloudAgentUserIds !== null && cloudAgentOrgIds !== null; return { services: services ?? SHADOW_ONLY_BILLING_CONFIG.services, userIds: userIds ?? SHADOW_ONLY_BILLING_CONFIG.userIds, orgIds: orgIds ?? SHADOW_ONLY_BILLING_CONFIG.orgIds, + cloudAgentUserIds: cloudAgentListsAreValid + ? cloudAgentUserIds + : SHADOW_ONLY_BILLING_CONFIG.cloudAgentUserIds, + cloudAgentOrgIds: cloudAgentListsAreValid + ? cloudAgentOrgIds + : SHADOW_ONLY_BILLING_CONFIG.cloudAgentOrgIds, warnRemainingMicrodollars: validThreshold ? warnRemainingMicrodollars : DEFAULT_WARN_REMAINING_MICRODOLLARS, - enabled, + enabled: gastownEnabled || cloudAgentEnabled, }; } @@ -72,7 +94,12 @@ export function billingModeFor( subject: { type: 'user' | 'org'; id: string } ): 'shadow' | 'paid' { if (!config.enabled || !config.services.has(service)) return 'shadow'; - return (subject.type === 'user' ? config.userIds : config.orgIds).has(subject.id) - ? 'paid' - : 'shadow'; + const payerIds = service.startsWith(CLOUD_AGENT_SERVICE_PREFIX) + ? subject.type === 'user' + ? config.cloudAgentUserIds + : config.cloudAgentOrgIds + : subject.type === 'user' + ? config.userIds + : config.orgIds; + return payerIds?.has(subject.id) ? 'paid' : 'shadow'; } diff --git a/services/container-usage-meter/src/meter.test.ts b/services/container-usage-meter/src/meter.test.ts index 1ce94f922c..f1a7300009 100644 --- a/services/container-usage-meter/src/meter.test.ts +++ b/services/container-usage-meter/src/meter.test.ts @@ -138,6 +138,48 @@ describe('ContainerUsageMeter', () => { ); }); + it('returns paid admission details from the versioned start RPC', async () => { + vi.mocked(applyStart).mockResolvedValue({ + kind: 'applied', + dedup: false, + billingMode: 'paid', + remainingMicrodollars: 10_000_001, + }); + + await expect(createMeter().recordStartV2(validStart())).resolves.toEqual({ + success: true, + ack: { intervalId: 'cloud-agent-next:instance-1:123', durable: 'pg', dedup: false }, + billingMode: 'paid', + remainingMicrodollars: 10_000_001, + }); + }); + + it('returns detailed insufficient-credit rejections from the versioned start RPC', async () => { + vi.mocked(applyStart).mockResolvedValue({ + kind: 'rejected', + code: 'insufficient_credits', + message: 'Container billing requires at least $5.00 in remaining credits', + remainingMicrodollars: 5_000_000, + minimumRequiredMicrodollars: 5_000_000, + }); + + await expect(createMeter().recordStartV2(validStart())).resolves.toMatchObject({ + success: false, + error: { + code: 'insufficient_credits', + remainingMicrodollars: 5_000_000, + minimumRequiredMicrodollars: 5_000_000, + }, + }); + await expect(createMeter().recordStart(validStart())).resolves.toEqual({ + success: false, + error: { + code: 'insufficient_credits', + message: 'Container billing requires at least $5.00 in remaining credits', + }, + }); + }); + it('writes heartbeats directly and returns the shadow budget verdict', async () => { await expect(createMeter().recordHeartbeat(validHeartbeat())).resolves.toEqual({ intervalId: 'cloud-agent-next:instance-1:123', diff --git a/services/container-usage-meter/src/meter.ts b/services/container-usage-meter/src/meter.ts index f61d37df3e..b6904de079 100644 --- a/services/container-usage-meter/src/meter.ts +++ b/services/container-usage-meter/src/meter.ts @@ -14,6 +14,7 @@ import { type RecordHeartbeatInput, type RecordStartInput, type RecordStartResult, + type RecordStartV2Result, type RecordStopInput, type UsageContext, } from '@kilocode/container-usage'; @@ -71,6 +72,54 @@ export class ContainerUsageMeter implements ContainerUsageRpcMethods { async recordStart(input: RecordStartInput): Promise { + const result = await this.applyStart(input); + if (result.kind === 'rejected') { + return { success: false, error: { code: result.code, message: result.message } }; + } + return { + success: true, + ack: { + intervalId: intervalId(input.service, input.instanceId, input.startEpochMs), + durable: 'pg', + dedup: result.dedup, + }, + }; + } + + async recordStartV2(input: RecordStartInput): Promise { + const result = await this.applyStart(input); + if (result.kind === 'rejected') { + if (result.code === 'insufficient_credits') { + return { + success: false, + error: { + code: result.code, + message: result.message, + remainingMicrodollars: result.remainingMicrodollars, + minimumRequiredMicrodollars: result.minimumRequiredMicrodollars, + }, + }; + } + return { success: false, error: { code: result.code, message: result.message } }; + } + const ack = { + intervalId: intervalId(input.service, input.instanceId, input.startEpochMs), + durable: 'pg' as const, + dedup: result.dedup, + }; + return result.billingMode === 'paid' + ? { + success: true, + ack, + billingMode: 'paid', + remainingMicrodollars: result.remainingMicrodollars, + } + : { success: true, ack, billingMode: 'shadow' }; + } + + private async applyStart( + input: RecordStartInput + ): Promise>> { const parsed = recordStartInputSchema.parse(input); assertIdempotencyKey( parsed.idempotencyKey, @@ -94,17 +143,13 @@ export class ContainerUsageMeter }); throw error; } - switch (result.kind) { - case 'rejected': - logRpcOutcome('start', parsed.service, 'rejected', { rejectionCode: result.code }); - return { success: false, error: { code: result.code, message: result.message } }; - case 'applied': - logRpcOutcome('start', parsed.service, 'accepted', { dedup: result.dedup }); - return { - success: true, - ack: { intervalId: id, durable: 'pg', dedup: result.dedup }, - }; - } + logRpcOutcome( + 'start', + parsed.service, + result.kind === 'rejected' ? 'rejected' : 'accepted', + result.kind === 'rejected' ? { rejectionCode: result.code } : { dedup: result.dedup } + ); + return result; } async recordHeartbeat(input: RecordHeartbeatInput): Promise { diff --git a/services/container-usage-meter/src/postgres.ts b/services/container-usage-meter/src/postgres.ts index b8a73fd53f..095062140c 100644 --- a/services/container-usage-meter/src/postgres.ts +++ b/services/container-usage-meter/src/postgres.ts @@ -50,8 +50,20 @@ export function getContainerUsageDb(env: Cloudflare.Env): WorkerDb { } export type StartSkuAdmission = - | { kind: 'applied'; dedup: boolean; billingMode: 'shadow' | 'paid' } - | { kind: 'rejected'; code: RecordStartFailureCode; message: string }; + | { kind: 'applied'; dedup: boolean; billingMode: 'shadow' } + | { kind: 'applied'; dedup: boolean; billingMode: 'paid'; remainingMicrodollars: number } + | { + kind: 'rejected'; + code: Exclude; + message: string; + } + | { + kind: 'rejected'; + code: 'insufficient_credits'; + message: string; + remainingMicrodollars: number; + minimumRequiredMicrodollars: number; + }; export type ApplyResult = { kind: 'applied'; @@ -365,7 +377,15 @@ export async function applyStartWithDb( .limit(1); if (existing) { assertMatchingContext(existing, input, contextFingerprint); - return { kind: 'applied', dedup: true, billingMode: existing.billing_mode }; + if (existing.billing_mode === 'paid') { + return { + kind: 'applied', + dedup: true, + billingMode: 'paid', + remainingMicrodollars: await balanceForSubject(tx, input.subject, true), + }; + } + return { kind: 'applied', dedup: true, billingMode: 'shadow' }; } const [sku] = await tx @@ -420,6 +440,8 @@ export async function applyStartWithDb( kind: 'rejected', code: 'insufficient_credits', message: 'Container billing requires at least $5.00 in remaining credits', + remainingMicrodollars: remaining, + minimumRequiredMicrodollars: MINIMUM_REMAINING_MICRODOLLARS, }; } } @@ -459,9 +481,25 @@ export async function applyStartWithDb( .limit(1); if (!winner) throw new Error('Container usage interval insert lost without a winner'); assertMatchingContext(winner, input, contextFingerprint); - return { kind: 'applied', dedup: true, billingMode: winner.billing_mode }; + if (winner.billing_mode === 'paid') { + return { + kind: 'applied', + dedup: true, + billingMode: 'paid', + remainingMicrodollars: await balanceForSubject(tx, input.subject, true), + }; + } + return { kind: 'applied', dedup: true, billingMode: 'shadow' }; + } + if (billingMode === 'paid') { + return { + kind: 'applied', + dedup: false, + billingMode: 'paid', + remainingMicrodollars: await balanceForSubject(tx, input.subject), + }; } - return { kind: 'applied', dedup: false, billingMode }; + return { kind: 'applied', dedup: false, billingMode: 'shadow' }; }); return operation.catch(mapSingleOpenIntervalConflict); } diff --git a/services/container-usage-meter/test/postgres.test.ts b/services/container-usage-meter/test/postgres.test.ts index eb1d64b077..c679f4418d 100644 --- a/services/container-usage-meter/test/postgres.test.ts +++ b/services/container-usage-meter/test/postgres.test.ts @@ -44,6 +44,8 @@ const paidBillingConfig: BillingConfig = { services: new Set(['cloud-agent-next']), userIds: new Set([userId]), orgIds: new Set([organizationId]), + cloudAgentUserIds: new Set([userId]), + cloudAgentOrgIds: new Set([organizationId]), warnRemainingMicrodollars: 10_000_000, enabled: true, }; @@ -247,7 +249,27 @@ describe('container usage PostgreSQL application', () => { startEpochMs, paidBillingConfig ) - ).resolves.toEqual({ kind: 'applied', dedup: false, billingMode: 'paid' }); + ).resolves.toEqual({ + kind: 'applied', + dedup: false, + billingMode: 'paid', + remainingMicrodollars: 10_050_000, + }); + await expect( + applyStartWithDb( + client.db, + start, + paidIntervalId, + paidFingerprint, + startEpochMs + 1, + paidBillingConfig + ) + ).resolves.toEqual({ + kind: 'applied', + dedup: true, + billingMode: 'paid', + remainingMicrodollars: 10_050_000, + }); const heartbeat = { service: paidContext.service, @@ -344,7 +366,12 @@ describe('container usage PostgreSQL application', () => { startEpochMs + 1, paidBillingConfig ) - ).resolves.toMatchObject({ kind: 'rejected', code: 'insufficient_credits' }); + ).resolves.toMatchObject({ + kind: 'rejected', + code: 'insufficient_credits', + remainingMicrodollars: MINIMUM_REMAINING_MICRODOLLARS, + minimumRequiredMicrodollars: MINIMUM_REMAINING_MICRODOLLARS, + }); }); it('settles paid organization usage against its aggregate wallet', async () => { @@ -365,22 +392,29 @@ describe('container usage PostgreSQL application', () => { startEpochMs, 1 ); - await applyStartWithDb( - client.db, - { - ...organizationContext, + await expect( + applyStartWithDb( + client.db, + { + ...organizationContext, + startEpochMs, + idempotencyKey: startIdempotencyKey( + organizationContext.service, + organizationContext.instanceId, + startEpochMs + ), + }, + organizationIntervalId, + organizationFingerprint, startEpochMs, - idempotencyKey: startIdempotencyKey( - organizationContext.service, - organizationContext.instanceId, - startEpochMs - ), - }, - organizationIntervalId, - organizationFingerprint, - startEpochMs, - paidBillingConfig - ); + paidBillingConfig + ) + ).resolves.toEqual({ + kind: 'applied', + dedup: false, + billingMode: 'paid', + remainingMicrodollars: 5_100_000, + }); await expect( applyHeartbeatWithDb( client.db, @@ -517,6 +551,7 @@ describe('container usage PostgreSQL application', () => { kind: 'applied', dedup: false, billingMode: 'paid', + remainingMicrodollars: 20_000_000, }); await client.db .delete(container_usage_interval) diff --git a/services/container-usage-meter/worker-configuration.d.ts b/services/container-usage-meter/worker-configuration.d.ts index 78a2582487..b1775fff30 100644 --- a/services/container-usage-meter/worker-configuration.d.ts +++ b/services/container-usage-meter/worker-configuration.d.ts @@ -1,10 +1,12 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types --include-runtime=false` (hash: c944f9baaa02912ff6b6072cd23f5e9e) +// Generated by Wrangler by running `wrangler types --include-runtime=false` (hash: b75aee7fe5ef6dc5a5520a8bb8de8470) interface __BaseEnv_Env { HYPERDRIVE: Hyperdrive; CONTAINER_BILLING_SERVICES: "gastown"; CONTAINER_BILLING_USER_IDS: "daef8451-f3f3-490e-93f1-21fafc2b005e"; CONTAINER_BILLING_ORG_IDS: ""; + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS: ""; + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS: ""; CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS: "10000000"; } declare namespace Cloudflare { @@ -18,5 +20,5 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } diff --git a/services/container-usage-meter/wrangler.jsonc b/services/container-usage-meter/wrangler.jsonc index f48b8e6352..661c291ab5 100644 --- a/services/container-usage-meter/wrangler.jsonc +++ b/services/container-usage-meter/wrangler.jsonc @@ -23,6 +23,8 @@ "CONTAINER_BILLING_SERVICES": "gastown", "CONTAINER_BILLING_USER_IDS": "daef8451-f3f3-490e-93f1-21fafc2b005e", "CONTAINER_BILLING_ORG_IDS": "", + "CONTAINER_BILLING_CLOUD_AGENT_USER_IDS": "", + "CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS": "", "CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS": "10000000", }, "triggers": { diff --git a/services/session-ingest/src/session-ingest-rpc.test.ts b/services/session-ingest/src/session-ingest-rpc.test.ts index db218ad920..750b77f1f2 100644 --- a/services/session-ingest/src/session-ingest-rpc.test.ts +++ b/services/session-ingest/src/session-ingest-rpc.test.ts @@ -225,6 +225,76 @@ describe('createSessionForCloudAgent', () => { ); expect(fake.values).not.toHaveBeenCalled(); }); + + it('refuses to change the payer organization of an existing root', async () => { + const fake = makeRootWriteDb({ + existing: { + session_id: params.sessionId, + kilo_user_id: params.kiloUserId, + cloud_agent_session_id: params.cloudAgentSessionId, + cloud_agent_session_scope_id: params.cloudAgentSessionId, + organization_id: null, + parent_session_id: null, + }, + }); + const rpc = makeRpc(fake.db as never); + + await expect(rpc.createSessionForCloudAgent(params)).rejects.toThrow( + 'Cloud Agent root session identity conflict' + ); + }); + + it('refuses to remove the payer organization of an existing root', async () => { + const fake = makeRootWriteDb({ + existing: { + session_id: params.sessionId, + kilo_user_id: params.kiloUserId, + cloud_agent_session_id: params.cloudAgentSessionId, + cloud_agent_session_scope_id: params.cloudAgentSessionId, + organization_id: params.organizationId, + parent_session_id: null, + }, + }); + const rpc = makeRpc(fake.db as never); + const { organizationId: _organizationId, ...personalParams } = params; + + await expect(rpc.createSessionForCloudAgent(personalParams)).rejects.toThrow( + 'Cloud Agent root session identity conflict' + ); + }); + + it('accepts an idempotent retry with the same payer organization', async () => { + const existing = { + session_id: params.sessionId, + kilo_user_id: params.kiloUserId, + cloud_agent_session_id: params.cloudAgentSessionId, + cloud_agent_session_scope_id: params.cloudAgentSessionId, + organization_id: params.organizationId, + parent_session_id: null, + }; + const fake = makeRootWriteDb({ existing }); + const rpc = makeRpc(fake.db as never); + + await expect(rpc.createSessionForCloudAgent(params)).resolves.toBeUndefined(); + expect(fake.values).toHaveBeenCalledTimes(1); + }); + + it('accepts an idempotent retry for a personal root', async () => { + const { organizationId: _organizationId, ...personalParams } = params; + const existing = { + session_id: personalParams.sessionId, + kilo_user_id: personalParams.kiloUserId, + cloud_agent_session_id: personalParams.cloudAgentSessionId, + cloud_agent_session_scope_id: personalParams.cloudAgentSessionId, + organization_id: null, + parent_session_id: null, + }; + const fake = makeRootWriteDb({ existing }); + const rpc = makeRpc(fake.db as never); + + await expect(rpc.createSessionForCloudAgent(personalParams)).resolves.toBeUndefined(); + expect(fake.values).toHaveBeenCalledTimes(1); + }); }); describe('Kilo SDK persisted identity schemas', () => { diff --git a/services/session-ingest/src/session-ingest-rpc.ts b/services/session-ingest/src/session-ingest-rpc.ts index 7ce7ed9b07..ef0496e724 100644 --- a/services/session-ingest/src/session-ingest-rpc.ts +++ b/services/session-ingest/src/session-ingest-rpc.ts @@ -117,7 +117,8 @@ export class SessionIngestRPC extends WorkerEntrypoint implements SessionIn existing.parent_session_id !== null || existing.cloud_agent_session_id !== parsed.cloudAgentSessionId || (existing.cloud_agent_session_scope_id !== null && - existing.cloud_agent_session_scope_id !== parsed.cloudAgentSessionId) + existing.cloud_agent_session_scope_id !== parsed.cloudAgentSessionId) || + existing.organization_id !== (parsed.organizationId ?? null) ) { throw new Error('Cloud Agent root session identity conflict'); } @@ -126,9 +127,6 @@ export class SessionIngestRPC extends WorkerEntrypoint implements SessionIn .update(cli_sessions_v2) .set({ cloud_agent_session_scope_id: parsed.cloudAgentSessionId, - ...(parsed.organizationId !== undefined - ? { organization_id: parsed.organizationId } - : {}), }) .where( and( @@ -141,9 +139,7 @@ export class SessionIngestRPC extends WorkerEntrypoint implements SessionIn }); const hasMeaningfulChange = existingRow - ? existingRow.cloud_agent_session_scope_id !== parsed.cloudAgentSessionId || - (parsed.organizationId !== undefined && - existingRow.organization_id !== parsed.organizationId) + ? existingRow.cloud_agent_session_scope_id !== parsed.cloudAgentSessionId : true; if (existingRow && hasMeaningfulChange && persistedRow) { From 628dc95756431035080fd898d99cd240cc3839bc Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Wed, 19 Aug 2026 11:42:39 -0500 Subject: [PATCH 2/6] fix(cloud-agent-next): settle billing on physical stop --- .../container-usage/src/heartbeat.test.ts | 83 +++++++++ packages/container-usage/src/heartbeat.ts | 9 +- .../src/container-usage.test.ts | 171 +++++++++++++++++- .../cloud-agent-next/src/container-usage.ts | 101 +++++++++-- .../wrapper/src/lifecycle.test.ts | 38 +++- .../cloud-agent-next/wrapper/src/lifecycle.ts | 63 ++++--- 6 files changed, 413 insertions(+), 52 deletions(-) diff --git a/packages/container-usage/src/heartbeat.test.ts b/packages/container-usage/src/heartbeat.test.ts index b3f1427d66..d5cca5fa27 100644 --- a/packages/container-usage/src/heartbeat.test.ts +++ b/packages/container-usage/src/heartbeat.test.ts @@ -152,6 +152,89 @@ describe('installBillingHeartbeat', () => { expect(schedule).toHaveBeenCalledWith(300, BILLING_HEARTBEAT_CALLBACK, expect.any(String)); }); + it('keeps the default controller behavior of settling immediately after a budget stop', async () => { + const storage = memoryStorage(); + await storedContext(storage); + const recordStop = vi.fn(async input => ({ + intervalId: `${input.instanceId}:${input.startEpochMs}`, + durable: 'pg', + dedup: false, + })); + const client = new ContainerUsageClient( + { + recordStart: async () => ({ + success: true, + ack: { intervalId: 'instance-1:123', durable: 'pg', dedup: false }, + }), + recordHeartbeat: async () => ({ + intervalId: 'instance-1:123', + durable: 'pg', + dedup: false, + budget: { verdict: 'stop' }, + }), + recordStop, + }, + { service: 'gastown' } + ); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(async () => ({ status: 'running' as const, lastChange: Date.now() })), + schedule: vi.fn() as Container['schedule'], + }, + { client, storage, enforceBudgetStop: vi.fn() } + ); + + await controller.billingHeartbeatTick(); + + expect(recordStop).toHaveBeenCalledOnce(); + expect(await getBillingContext(storage)).toBeUndefined(); + }); + + it('defers budget-stop settlement when the producer owns physical shutdown', async () => { + const storage = memoryStorage(); + await storedContext(storage); + const recordStop = vi.fn(async input => ({ + intervalId: `${input.instanceId}:${input.startEpochMs}`, + durable: 'pg', + dedup: false, + })); + const client = new ContainerUsageClient( + { + recordStart: async () => ({ + success: true, + ack: { intervalId: 'instance-1:123', durable: 'pg', dedup: false }, + }), + recordHeartbeat: async () => ({ + intervalId: 'instance-1:123', + durable: 'pg', + dedup: false, + budget: { verdict: 'stop' }, + }), + recordStop, + }, + { service: 'cloud-agent-next' } + ); + const controller = installBillingHeartbeat( + { + deleteSchedules: vi.fn(), + getState: vi.fn(async () => ({ status: 'running' as const, lastChange: Date.now() })), + schedule: vi.fn() as Container['schedule'], + }, + { + client, + storage, + deferBudgetStopFinalSettlement: true, + enforceBudgetStop: vi.fn(), + } + ); + + await controller.billingHeartbeatTick(); + + expect(recordStop).not.toHaveBeenCalled(); + expect(await getBillingContext(storage)).toBeDefined(); + }); + it('defers stopped-state closure when the producer owns an authoritative stop hook', async () => { const storage = memoryStorage(); await storedContext(storage); diff --git a/packages/container-usage/src/heartbeat.ts b/packages/container-usage/src/heartbeat.ts index 7dd84d9ef9..02755f1683 100644 --- a/packages/container-usage/src/heartbeat.ts +++ b/packages/container-usage/src/heartbeat.ts @@ -24,6 +24,11 @@ export type BillingHeartbeatDependencies = { heartbeatSeconds?: number; /** Defer stopped-state closure to the container's authoritative onStop hook. */ stopOnStoppedState?: boolean; + /** + * Let a producer that owns physical shutdown settle from its authoritative + * onStop hook instead of closing usage at a budget verdict. + */ + deferBudgetStopFinalSettlement?: boolean; stoppedStateGraceSeconds?: number; stoppedStateAbandonSeconds?: number; beforeHeartbeatDelivery?: (context: BillingContext) => Promise; @@ -379,7 +384,9 @@ export function installBillingHeartbeat( generation: context.generation, startEpochMs: context.startEpochMs, }); - await recordStopForGeneration({ reason: 'runtime_signal' }, context.generation); + if (!dependencies.deferBudgetStopFinalSettlement) { + await recordStopForGeneration({ reason: 'runtime_signal' }, context.generation); + } return; } if (ack.budget.verdict === 'warn') { diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index 106e189b3b..7f834a8b26 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -62,6 +62,8 @@ const sdk = vi.hoisted(() => { async destroy(): Promise { this.superDestroyCalled = true; + const storage = this.ctx.storage as unknown as MemoryStorage; + if (storage.clearOnDestroy) storage.clear(); } } return { StockSandbox }; @@ -75,6 +77,7 @@ class MemoryStorage { private readonly values = new Map(); failWrites = false; hangReads = false; + clearOnDestroy = false; async get(key: string): Promise { if (this.hangReads) return await new Promise(() => undefined); @@ -89,6 +92,10 @@ class MemoryStorage { async delete(key: string): Promise { return this.values.delete(key); } + + clear(): void { + this.values.clear(); + } } function ack(intervalId = 'interval-1') { @@ -248,7 +255,7 @@ describe('MeteredSandbox', () => { ).resolves.toMatchObject({ success: false, code: 'configuration_mismatch' }); }); - it('persists stop enforcement, reports final usage, and resumes only after new admission', async () => { + it('settles a graceful budget stop from physical onStop and resumes only after paid admission', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStartV2!).mockResolvedValue({ success: true, @@ -283,9 +290,16 @@ describe('MeteredSandbox', () => { callback: 'billingForceStop', payload: active.generation, }); + expect(rpc.recordStop).not.toHaveBeenCalled(); + vi.spyOn(Date, 'now').mockReturnValue(361_000); + await sandbox.onStop({ reason: 'runtime_signal' }); + await flushShadowTasks(); + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ usageSinceLast: 60, reason: 'runtime_signal' }) + ); + await sandbox.onStop({ reason: 'runtime_signal' }); + await flushShadowTasks(); expect(rpc.recordStop).toHaveBeenCalledOnce(); - await sandbox.billingForceStop(active.generation); - expect(sandbox.superDestroyCalled).toBe(true); expect(await sandbox.isBillingBlocked()).toBe(true); sandbox.setPhysicalRunning(false); @@ -297,6 +311,157 @@ describe('MeteredSandbox', () => { expect(rpc.recordStartV2).toHaveBeenCalledTimes(2); }); + it('settles force-stop usage at the physical stop after the deadline', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStartV2!).mockResolvedValue({ + success: true, + ack: ack(), + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }); + vi.mocked(rpc.recordHeartbeat).mockResolvedValue({ + ...ack(), + budget: { verdict: 'stop', remainingMicrodollars: 5_000_000 }, + }); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + await sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }); + await sandbox.onStart(); + await flushShadowTasks(); + const active = await getBillingContext(storage); + if (!active) throw new Error('Expected active paid billing context'); + sandbox.mockState = { status: 'running' }; + + now.mockReturnValue(301_000); + await sandbox.billingHeartbeatTick(active.generation); + now.mockReturnValue(421_000); + await sandbox.billingForceStop(active.generation); + expect(sandbox.superDestroyCalled).toBe(true); + expect(rpc.recordStop).not.toHaveBeenCalled(); + + await sandbox.onStop({ reason: 'runtime_signal' }); + await flushShadowTasks(); + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ usageSinceLast: 120, reason: 'runtime_signal' }) + ); + expect(await sandbox.isBillingBlocked()).toBe(true); + }); + + it('uses the container stop transition rather than a late onStop callback time', async () => { + const rpc = createRpc(); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + 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'); + + now.mockReturnValue(400_000); + sandbox.mockState = { status: 'stopped' }; + Object.assign(sandbox.mockState, { lastChange: 301_000 }); + await sandbox.onStop({ reason: 'runtime_signal' }); + await flushShadowTasks(); + + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ usageSinceLast: 300, startEpochMs: active.startEpochMs }) + ); + }); + + it('restores the durable billing block when destroy clears storage', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStartV2!).mockResolvedValue({ + success: true, + ack: ack(), + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }); + vi.mocked(rpc.recordHeartbeat).mockResolvedValue({ + ...ack(), + budget: { verdict: 'stop', remainingMicrodollars: 5_000_000 }, + }); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + await sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }); + await sandbox.onStart(); + await flushShadowTasks(); + const active = await getBillingContext(storage); + if (!active) throw new Error('Expected active paid billing context'); + sandbox.mockState = { status: 'running' }; + await sandbox.billingHeartbeatTick(active.generation); + + storage.clearOnDestroy = true; + await sandbox.billingForceStop(active.generation); + + expect(await sandbox.isBillingBlocked()).toBe(true); + }); + + it('reissues a failed durable force-destroy without clearing the billing block', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStartV2!).mockResolvedValue({ + success: true, + ack: ack(), + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }); + vi.mocked(rpc.recordHeartbeat).mockResolvedValue({ + ...ack(), + budget: { verdict: 'stop', remainingMicrodollars: 5_000_000 }, + }); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + await sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }); + await sandbox.onStart(); + await flushShadowTasks(); + const active = await getBillingContext(storage); + if (!active) throw new Error('Expected active paid billing context'); + sandbox.mockState = { status: 'running' }; + await sandbox.billingHeartbeatTick(active.generation); + const destroy = vi.spyOn(sandbox, 'destroy').mockRejectedValueOnce(new Error('unavailable')); + + await expect(sandbox.billingForceStop(active.generation)).rejects.toThrow('unavailable'); + + expect(await sandbox.isBillingBlocked()).toBe(true); + expect(sandbox.schedules).toContainEqual({ + when: 5, + callback: 'billingForceStop', + payload: active.generation, + }); + await sandbox.billingForceStop(active.generation); + expect(destroy).toHaveBeenCalledTimes(2); + }); + + it('does not clear a budget block for a shadow admission', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStartV2!) + .mockResolvedValueOnce({ + success: true, + ack: ack(), + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }) + .mockResolvedValueOnce({ success: true, ack: ack(), billingMode: 'shadow' }); + vi.mocked(rpc.recordHeartbeat).mockResolvedValue({ + ...ack(), + budget: { verdict: 'stop', remainingMicrodollars: 5_000_000 }, + }); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + await sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }); + await sandbox.onStart(); + await flushShadowTasks(); + const active = await getBillingContext(storage); + if (!active) throw new Error('Expected active paid billing context'); + sandbox.mockState = { status: 'running' }; + await sandbox.billingHeartbeatTick(active.generation); + await sandbox.onStop({ reason: 'runtime_signal' }); + await flushShadowTasks(); + + await expect(sandbox.ensureBillingAdmission(billingInput)).resolves.toMatchObject({ + success: false, + code: 'configuration_mismatch', + }); + expect(await sandbox.isBillingBlocked()).toBe(true); + }); + it('admits one start per physical generation and short-circuits active acquisition', async () => { const { rpc, storage, sandbox, flushShadowTasks } = createSandbox(); vi.spyOn(Date, 'now').mockReturnValue(1_000); diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index b017104027..58350f145c 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -33,6 +33,7 @@ const LAST_START_EPOCH_STORAGE_KEY = 'container-usage:last-start-epoch:v1'; const START_ADMISSION_STORAGE_KEY = 'container-usage:start-admission:v1'; const BILLING_BLOCK_STORAGE_KEY = 'container-usage:budget-block:v1'; const BILLING_FORCE_STOP_SECONDS = 120; +const BILLING_FORCE_STOP_RETRY_SECONDS = 5; // oxlint-disable-next-line no-empty-object-type -- Matches the Sandbox 0.12.1 constructor. type SandboxDurableObjectState = DurableObjectState<{}>; @@ -108,6 +109,7 @@ export abstract class MeteredSandbox extends StockSandbox { client: this.usageClient, storage: this.ctx.storage, stopOnStoppedState: false, + deferBudgetStopFinalSettlement: true, beforeHeartbeatDelivery: context => this.ensureStartAcknowledged(context), beforeStopDelivery: context => this.ensureStartAcknowledged(context), onGenerationClosed: () => this.schedulePendingGenerationIfRunning(), @@ -260,6 +262,14 @@ export abstract class MeteredSandbox extends StockSandbox { message: 'Container billing rollout is not enabled in the usage meter', }; } + if (block && result.billingMode !== 'paid') { + await clearBillingContext(this.ctx.storage); + return { + success: false, + code: 'configuration_mismatch', + message: 'Container billing remains stopped until funded paid admission succeeds', + }; + } await this.ctx.storage.put(START_ACK_GENERATION_STORAGE_KEY, context.generation); await this.ctx.storage.put(START_ADMISSION_STORAGE_KEY, { generation: context.generation, @@ -268,7 +278,11 @@ export abstract class MeteredSandbox extends StockSandbox { ? { remainingMicrodollars: result.remainingMicrodollars } : {}), }); - await this.ctx.storage.delete(BILLING_BLOCK_STORAGE_KEY); + // A budget block is cleared only by a later paid admission. Shadow + // configuration must not resume work that was stopped for insufficient funds. + if (result.billingMode === 'paid') { + await this.ctx.storage.delete(BILLING_BLOCK_STORAGE_KEY); + } return result.billingMode === 'paid' ? { success: true, @@ -404,8 +418,11 @@ export abstract class MeteredSandbox extends StockSandbox { } override async onStop(params?: ContainerStopParams): Promise { - const stoppedAtMs = Date.now(); await super.onStop(); + // `onStop` is the first durable lifecycle signal after the container has + // actually stopped. Do not use the earlier budget verdict or force-destroy + // request as the usage boundary. + const stoppedAtMs = await this.getObservedStopTime(); const activityExpiryRequested = this.activityExpiryRequested; this.activityExpiryRequested = false; this.runShadowTask('stop lifecycle', async () => { @@ -460,12 +477,6 @@ export abstract class MeteredSandbox extends StockSandbox { }); } - override async destroy(): Promise { - const block = await this.getBillingBlock(); - await super.destroy(); - if (block) await this.ctx.storage.put(BILLING_BLOCK_STORAGE_KEY, block); - } - private runBillingExclusive(operation: () => Promise): Promise { const result = this.billingLifecycleTail.then(operation, operation); this.billingLifecycleTail = result.then( @@ -482,6 +493,12 @@ export abstract class MeteredSandbox extends StockSandbox { this.ctx.waitUntil(promise); } + override async destroy(): Promise { + const block = await this.getBillingBlock(); + await super.destroy(); + if (block) await this.ctx.storage.put(BILLING_BLOCK_STORAGE_KEY, block); + } + private schedulePendingGenerationIfRunning(): void { this.runShadowTask('replacement generation', async () => { if (await this.getBillingBlock()) return; @@ -565,6 +582,16 @@ export abstract class MeteredSandbox extends StockSandbox { return stored === undefined ? undefined : billingBlockSchema.parse(stored); } + private async getObservedStopTime(): Promise { + try { + return stoppedAtFromState(await this.getState()); + } catch { + // The lifecycle callback itself is still authoritative when the control + // plane cannot provide a state transition timestamp. + return Date.now(); + } + } + private async logBudgetWarning(budget: { verdict: string; remainingMicrodollars?: number; @@ -594,15 +621,26 @@ export abstract class MeteredSandbox extends StockSandbox { budget: { verdict: string; remainingMicrodollars?: number }, expected: { generation: string; startEpochMs: number } ): Promise { - const now = Date.now(); - const block = { - generation: expected.generation, - startEpochMs: expected.startEpochMs, - blockedAt: now, - forceStopAt: now + BILLING_FORCE_STOP_SECONDS * 1_000, - remainingMicrodollars: budget.remainingMicrodollars, - }; - await this.ctx.storage.put(BILLING_BLOCK_STORAGE_KEY, block); + const active = await getBillingContext(this.ctx.storage); + if ( + !active || + active.generation !== expected.generation || + active.startEpochMs !== expected.startEpochMs + ) { + return; + } + const existing = await this.getBillingBlock(); + if (existing && existing.generation !== expected.generation) return; + const block = + existing ?? + ({ + generation: expected.generation, + startEpochMs: expected.startEpochMs, + blockedAt: Date.now(), + forceStopAt: Date.now() + BILLING_FORCE_STOP_SECONDS * 1_000, + remainingMicrodollars: budget.remainingMicrodollars, + } satisfies z.infer); + if (!existing) await this.ctx.storage.put(BILLING_BLOCK_STORAGE_KEY, block); await this.scheduleForceStop(block); logger .withTags({ logTag: 'container_billing_stop' }) @@ -617,8 +655,19 @@ export abstract class MeteredSandbox extends StockSandbox { } async billingForceStop(generation: string): Promise { + return this.runBillingExclusive(() => this.forceStopBillingGeneration(generation)); + } + + private async forceStopBillingGeneration(generation: string): Promise { const block = await this.getBillingBlock(); - if (!block || block.generation !== generation || this.ctx.container?.running !== true) return; + if (!block || block.generation !== generation) return; + const active = await getBillingContext(this.ctx.storage); + if ( + !active || + active.generation !== block.generation || + active.startEpochMs !== block.startEpochMs + ) + return; logger .withTags({ logTag: 'container_billing_force_stop' }) .withFields({ @@ -627,7 +676,21 @@ export abstract class MeteredSandbox extends StockSandbox { stopLatencyMs: Date.now() - block.blockedAt, }) .error('Force-destroying container after billing stop deadline'); - await this.destroy(); + try { + // The control plane may fail after the deadline. Keep the durable block and + // reissue destroy until the physical stop hook settles the generation. + await this.destroy(); + } catch (error) { + logger + .withFields({ + error: error instanceof Error ? error.message : String(error), + sandboxClass: this.sandboxClassName, + generation, + }) + .warn('Billing force-destroy issuance failed; retrying'); + await this.schedule(BILLING_FORCE_STOP_RETRY_SECONDS, 'billingForceStop', generation); + throw error; + } } private async createBillingGeneration( diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts index 9ad2363204..3a3e192f16 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.test.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.test.ts @@ -25,12 +25,15 @@ describe('wrapper lifecycle drain races', () => { state.bindSession(sessionContext); state.setSendToIngestFn(event => events.push(event)); + let closeCalls = 0; const lifecycle = createLifecycleManager( { workspacePath: '/tmp' }, { state, kiloClient: {} as WrapperKiloClient, - closeConnections: async () => {}, + closeConnections: async () => { + closeCalls += 1; + }, isConnected: () => true, reconnectEventSubscription: () => {}, } @@ -50,6 +53,7 @@ describe('wrapper lifecycle drain races', () => { condenseOnComplete: false, }); await wait(300); + expect(closeCalls).toBe(0); lifecycle.onSessionIdle(); await wait(3_050); @@ -57,6 +61,38 @@ describe('wrapper lifecycle drain races', () => { expect(events.map(event => event.streamEventType)).toContain('complete'); }); + it('does not complete, close, or clear a session when reset interrupts an active drain', async () => { + const state = new WrapperState(); + const events: IngestEvent[] = []; + let closeCalls = 0; + state.bindSession(sessionContext); + state.setSendToIngestFn(event => events.push(event)); + const lifecycle = createLifecycleManager( + { workspacePath: '/tmp' }, + { + state, + kiloClient: {} as WrapperKiloClient, + closeConnections: async () => { + closeCalls += 1; + }, + isConnected: () => true, + reconnectEventSubscription: () => {}, + } + ); + + state.acceptMessage('message-1', { autoCommit: false, condenseOnComplete: false }); + state.clearAllMessages(); + const drain = lifecycle.drainAndClose(); + expect(events.map(event => event.streamEventType)).toContain('wrapper_finalizing'); + + lifecycle.reset(); + await drain; + + expect(events.map(event => event.streamEventType)).not.toContain('complete'); + expect(closeCalls).toBe(0); + expect(state.currentSession).toEqual(sessionContext); + }); + it('waits for three seconds of stable root idle before completing', async () => { const state = new WrapperState(); const events: IngestEvent[] = []; diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.ts b/services/cloud-agent-next/wrapper/src/lifecycle.ts index 422f40dd5c..ca1ebb2f26 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -154,6 +154,40 @@ export function createLifecycleManager( } } + async function finalizeDrain( + drainGeneration: number, + completeSession: typeof state.currentSession | undefined, + sealedMessageIds: string[] + ): Promise { + if (drainGeneration !== lifecycleGeneration) return; + const currentSession = state.currentSession; + if (completeSession && currentSession) { + const currentBranch = await getCurrentBranch(config.workspacePath, 10_000).catch(() => ''); + if (drainGeneration !== lifecycleGeneration) return; + const gateResult = state.consumeObservedGateResult(); + state.sendToIngest({ + streamEventType: 'complete', + data: { + exitCode: 0, + kiloSessionId: currentSession.kiloSessionId, + messageIds: sealedMessageIds, + ...(currentBranch ? { currentBranch } : {}), + ...(gateResult ? { gateResult } : {}), + }, + timestamp: new Date().toISOString(), + }); + } + + await new Promise(resolve => setTimeout(resolve, DRAIN_DELAY_MS)); + if (drainGeneration !== lifecycleGeneration) return; + await deps + .closeConnections() + .catch(error => + logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) + ); + if (drainGeneration === lifecycleGeneration) state.clearSession(); + } + function drainAndClose(): Promise { state.blockAdmissions(); if (drainPromise) return drainPromise; @@ -189,34 +223,7 @@ export function createLifecycleManager( uploader.stop(); } } finally { - if (drainGeneration !== lifecycleGeneration) return; - const currentSession = state.currentSession; - if (completeSession && currentSession) { - const currentBranch = await getCurrentBranch(config.workspacePath, 10_000).catch( - () => '' - ); - const gateResult = state.consumeObservedGateResult(); - state.sendToIngest({ - streamEventType: 'complete', - data: { - exitCode: 0, - kiloSessionId: currentSession.kiloSessionId, - messageIds: sealedMessageIds, - ...(currentBranch ? { currentBranch } : {}), - ...(gateResult ? { gateResult } : {}), - }, - timestamp: new Date().toISOString(), - }); - } - - await new Promise(resolve => setTimeout(resolve, DRAIN_DELAY_MS)); - if (drainGeneration !== lifecycleGeneration) return; - await deps - .closeConnections() - .catch(error => - logToFile(`close failed: ${error instanceof Error ? error.message : String(error)}`) - ); - state.clearSession(); + await finalizeDrain(drainGeneration, completeSession, sealedMessageIds); } })(); const currentDrain = drainPromise; From 5dff18e3587c9061e85011d96cdf19f3b9819806 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Wed, 19 Aug 2026 12:39:55 -0500 Subject: [PATCH 3/6] fix(cloud-agent-next): preserve paid shutdown state --- .../cloudflare-agent-sandbox.test.ts | 6 +- .../cloudflare/cloudflare-agent-sandbox.ts | 9 +-- .../src/container-usage-context.ts | 14 ++++- .../src/container-usage.test.ts | 41 ++++++++++++ .../cloud-agent-next/src/container-usage.ts | 18 ++++++ .../src/kilo-facade/session-proxy.test.ts | 62 +++++++++++++++++++ .../src/kilo-facade/session-proxy.ts | 2 +- services/cloud-agent-next/wrapper/src/main.ts | 10 ++- .../wrapper/src/shutdown.test.ts | 53 ++++++++++++++++ .../cloud-agent-next/wrapper/src/shutdown.ts | 13 ++++ 10 files changed, 210 insertions(+), 18 deletions(-) create mode 100644 services/cloud-agent-next/wrapper/src/shutdown.test.ts create mode 100644 services/cloud-agent-next/wrapper/src/shutdown.ts 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 6835298a96..a1af8023e7 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 @@ -1752,7 +1752,7 @@ describe('CloudflareAgentSandbox', () => { expect(listProcesses).toHaveBeenCalled(); }); - it('inspects a stopped container for stop reasons other than idle-timeout', async () => { + it('does not wake a confirmed stopped container for session deletion cleanup', async () => { const listProcesses = vi.fn().mockResolvedValue([]); const isContainerRunning = vi.fn().mockResolvedValue(false); const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), { @@ -1766,8 +1766,8 @@ describe('CloudflareAgentSandbox', () => { reason: 'session-delete', }) ).resolves.toEqual({ status: 'absent' }); - expect(listProcesses).toHaveBeenCalled(); - expect(isContainerRunning).not.toHaveBeenCalled(); + expect(isContainerRunning).toHaveBeenCalledOnce(); + expect(listProcesses).not.toHaveBeenCalled(); }); it('falls back to inspection when the sandbox cannot report container state', async () => { 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 40ec543c79..242846dea7 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 @@ -852,12 +852,9 @@ export class CloudflareAgentSandbox implements AgentSandbox { // process, and a process cannot outlive its container (activity expiry SIGTERMs the // whole container), so a stopped container cannot be hiding a leaked wrapper. // - // Scoped to idle-timeout: that sweep already established via DO state that no wrapper - // runtime or pending work remains, and it is the path that was waking cold containers - // for nothing. Every other stop reason keeps inspecting, which preserves the leaked - // wrapper recovery those paths were built for. - const skipsInspection = - request.reason === 'idle-timeout' && (await isSandboxContainerRunning(sandbox)) === false; + // A stopped container cannot hide a leaked wrapper, regardless of why cleanup was + // requested. Unknown remains inspectable so recovery paths keep their existing fence. + const skipsInspection = (await isSandboxContainerRunning(sandbox)) === false; const initial: StopInspection = skipsInspection ? { status: 'absent-no-container' } : await this.observeTarget(request.target); diff --git a/services/cloud-agent-next/src/container-usage-context.ts b/services/cloud-agent-next/src/container-usage-context.ts index 47d0abe29a..a4d45bf4e4 100644 --- a/services/cloud-agent-next/src/container-usage-context.ts +++ b/services/cloud-agent-next/src/container-usage-context.ts @@ -226,10 +226,20 @@ export async function ensureSandboxBillingAdmissionInput( } } -export async function isSandboxBillingBlocked(sandbox: SandboxInstance): Promise { +export async function isSandboxBillingBlocked( + sandbox: SandboxInstance, + enforcementRequested = false +): Promise { const isBillingBlocked = (sandbox as Partial).isBillingBlocked; if (typeof isBillingBlocked !== 'function') return false; - return await (sandbox as MeteredSandboxInstance).isBillingBlocked(); + try { + return await (sandbox as MeteredSandboxInstance).isBillingBlocked(); + } catch (error) { + logger + .withFields({ error: error instanceof Error ? error.message : String(error) }) + .warn('Container billing block check failed'); + return enforcementRequested; + } } /** diff --git a/services/cloud-agent-next/src/container-usage.test.ts b/services/cloud-agent-next/src/container-usage.test.ts index 7f834a8b26..3408d57635 100644 --- a/services/cloud-agent-next/src/container-usage.test.ts +++ b/services/cloud-agent-next/src/container-usage.test.ts @@ -396,6 +396,47 @@ describe('MeteredSandbox', () => { expect(await sandbox.isBillingBlocked()).toBe(true); }); + it('restores the force-stopped generation so onStop settles it exactly once', async () => { + const rpc = createRpc(); + vi.mocked(rpc.recordStartV2!).mockResolvedValue({ + success: true, + ack: ack(), + billingMode: 'paid', + remainingMicrodollars: 20_000_000, + }); + vi.mocked(rpc.recordHeartbeat).mockResolvedValue({ + ...ack(), + budget: { verdict: 'stop', remainingMicrodollars: 5_000_000 }, + }); + const { sandbox, storage, flushShadowTasks } = createSandbox(rpc); + const now = vi.spyOn(Date, 'now').mockReturnValue(1_000); + await sandbox.ensureBillingAdmission({ ...billingInput, enforcementRequested: true }); + await sandbox.onStart(); + await flushShadowTasks(); + const active = await getBillingContext(storage); + if (!active) throw new Error('Expected active paid billing context'); + sandbox.mockState = { status: 'running' }; + now.mockReturnValue(301_000); + await sandbox.billingHeartbeatTick(active.generation); + + storage.clearOnDestroy = true; + now.mockReturnValue(421_000); + await sandbox.billingForceStop(active.generation); + expect(await getBillingContext(storage)).toMatchObject({ generation: active.generation }); + + sandbox.mockState = { status: 'stopped' }; + await sandbox.onStop({ reason: 'runtime_signal' }); + await flushShadowTasks(); + await sandbox.onStop({ reason: 'runtime_signal' }); + await flushShadowTasks(); + + expect(rpc.recordStop).toHaveBeenCalledOnce(); + expect(rpc.recordStop).toHaveBeenCalledWith( + expect.objectContaining({ startEpochMs: active.startEpochMs, usageSinceLast: 120 }) + ); + expect(await sandbox.isBillingBlocked()).toBe(true); + }); + it('reissues a failed durable force-destroy without clearing the billing block', async () => { const rpc = createRpc(); vi.mocked(rpc.recordStartV2!).mockResolvedValue({ diff --git a/services/cloud-agent-next/src/container-usage.ts b/services/cloud-agent-next/src/container-usage.ts index 58350f145c..a8e4a0455c 100644 --- a/services/cloud-agent-next/src/container-usage.ts +++ b/services/cloud-agent-next/src/container-usage.ts @@ -4,6 +4,7 @@ import { getBillingContext, installBillingHeartbeat, setBillingContext, + updateBillingContext, usageContextFromBillingContext, type BillingContext, type BillingHeartbeatController, @@ -494,9 +495,26 @@ export abstract class MeteredSandbox extends StockSandbox { } override async destroy(): Promise { + const context = await getBillingContext(this.ctx.storage); const block = await this.getBillingBlock(); + const startAcknowledgement = await this.ctx.storage.get( + START_ACK_GENERATION_STORAGE_KEY + ); + const pendingStopReason = context + ? await this.getPendingStopReason(context.generation) + : undefined; await super.destroy(); + if (context) await updateBillingContext(this.ctx.storage, context); if (block) await this.ctx.storage.put(BILLING_BLOCK_STORAGE_KEY, block); + if (context && startAcknowledgement === context.generation) { + await this.ctx.storage.put(START_ACK_GENERATION_STORAGE_KEY, startAcknowledgement); + } + if (pendingStopReason) { + await this.ctx.storage.put(PENDING_STOP_REASON_STORAGE_KEY, { + generation: context?.generation, + reason: pendingStopReason, + }); + } } private schedulePendingGenerationIfRunning(): void { diff --git a/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts b/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts index 3912c983a0..07cd10bfb7 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts @@ -78,4 +78,66 @@ describe('resolveLiveWrapperTarget billing admission', () => { ); expect(mocks.findWrapperForSession).not.toHaveBeenCalled(); }); + + it('allows shadow acquisition when the billing block method is unavailable', async () => { + mocks.getSandbox.mockReturnValue({}); + mocks.findWrapperForSession.mockResolvedValue({ port: 5000 }); + + await expect( + resolveLiveWrapperTarget({ + env: { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'false', + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: '', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: '', + } as Env, + userId: 'user_facade', + cloudAgentSessionId: 'agent_facade', + }) + ).resolves.toMatchObject({ port: 5000 }); + }); + + it('allows shadow acquisition when the callable billing block proxy rejects', async () => { + mocks.getSandbox.mockReturnValue({ + isBillingBlocked: vi.fn().mockRejectedValue(new Error('RPC unavailable')), + }); + mocks.findWrapperForSession.mockResolvedValue({ port: 5000 }); + + await expect( + resolveLiveWrapperTarget({ + env: { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'false', + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: '', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: '', + } as Env, + userId: 'user_facade', + cloudAgentSessionId: 'agent_facade', + }) + ).resolves.toMatchObject({ port: 5000 }); + }); + + it('fails enforced acquisition closed when the callable billing block proxy rejects', async () => { + const ensureBillingAdmission = vi.fn().mockResolvedValue({ + success: false, + code: 'meter_unavailable', + message: 'Container billing admission is unavailable', + }); + mocks.getSandbox.mockReturnValue({ + isBillingBlocked: vi.fn().mockRejectedValue(new Error('RPC unavailable')), + ensureBillingAdmission, + }); + + await expect( + resolveLiveWrapperTarget({ + env: { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: '', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: 'org_facade', + } as Env, + userId: 'user_facade', + cloudAgentSessionId: 'agent_facade', + }) + ).resolves.toBeNull(); + expect(ensureBillingAdmission).toHaveBeenCalledOnce(); + expect(mocks.findWrapperForSession).not.toHaveBeenCalled(); + }); }); 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 9d2152eece..f04f8b8cd9 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts @@ -97,7 +97,7 @@ export async function resolveLiveWrapperTarget(params: { sandboxId, isCloudAgentContainerBillingEnabled(env, metadata.identity) ); - const billingBlocked = await isSandboxBillingBlocked(sandbox); + const billingBlocked = await isSandboxBillingBlocked(sandbox, billingInput.enforcementRequested); if (billingInput.enforcementRequested || billingBlocked) { const admission = await ensureSandboxBillingAdmissionInput(sandbox, billingInput); if (!admission.success) return null; diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index 5b7b005cfb..286c750775 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -31,6 +31,7 @@ import { openKiloGlobalFeed } from './global-feed.js'; import { createGlobalFeedManager, type SessionBoundFeedPolicy } from './global-feed-manager.js'; import { logToFile } from './utils.js'; import { startToolCgroup } from './tool-cgroup.js'; +import { abortKiloSessionForShutdown } from './shutdown.js'; import { kiloServerBootstrapError, kiloServerStartupError, @@ -940,6 +941,8 @@ async function main() { async function handleShutdown(signal: string): Promise { if (isShuttingDown) return; isShuttingDown = true; + const activeKiloSessionId = state.currentSession?.kiloSessionId; + lifecycleManager?.setAborted(); logToFile(`shutdown signal: ${signal}`); console.error(`Received ${signal}, shutting down...`); @@ -960,6 +963,7 @@ async function main() { }, timestamp: new Date().toISOString(), }); + await abortKiloSessionForShutdown({ activeKiloSessionId, kiloClient }); workspaceBootstrapController.abort(); const workspaceBootstraps = [...activeWorkspaceBootstraps]; @@ -987,12 +991,6 @@ async function main() { uploader.stop(); } - // Abort kilo session if running - const session = state.currentSession; - if (session && kiloClient) { - kiloClient.abortSession({ sessionId: session.kiloSessionId }).catch(() => {}); - } - // Close connections void connectionManager?.close(); diff --git a/services/cloud-agent-next/wrapper/src/shutdown.test.ts b/services/cloud-agent-next/wrapper/src/shutdown.test.ts new file mode 100644 index 0000000000..d647ae9336 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/shutdown.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it } from 'bun:test'; +import type { WrapperKiloClient } from './kilo-api'; +import { createLifecycleManager } from './lifecycle'; +import { abortKiloSessionForShutdown } from './shutdown'; +import { WrapperState } from './state'; +import type { IngestEvent } from '../../src/shared/protocol'; + +describe('abortKiloSessionForShutdown', () => { + it('aborts the active Kilo session before draining and never emits complete', async () => { + const events: IngestEvent[] = []; + const state = new WrapperState(); + state.bindSession({ + kiloSessionId: 'kilo_sess_test', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + }); + state.setSendToIngestFn(event => events.push(event)); + state.acceptMessage('message-1', { autoCommit: false, condenseOnComplete: false }); + const calls: string[] = []; + const lifecycleManager = createLifecycleManager( + { workspacePath: '/tmp' }, + { + state, + kiloClient: {} as WrapperKiloClient, + closeConnections: async () => { + calls.push('close'); + }, + isConnected: () => true, + reconnectEventSubscription: () => {}, + } + ); + const kiloClient: Pick = { + abortSession: async () => { + calls.push('abort'); + return true; + }, + }; + + state.sendToIngest({ + streamEventType: 'interrupted', + data: { reason: 'Container shutdown' }, + timestamp: new Date().toISOString(), + }); + lifecycleManager.setAborted(); + const activeKiloSessionId = state.currentSession?.kiloSessionId; + await abortKiloSessionForShutdown({ activeKiloSessionId, kiloClient }); + await lifecycleManager.drainAndClose(); + + expect(calls).toEqual(['abort', 'close']); + expect(events.map(event => event.streamEventType)).toEqual(['interrupted']); + expect(state.currentSession).toBeNull(); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/shutdown.ts b/services/cloud-agent-next/wrapper/src/shutdown.ts new file mode 100644 index 0000000000..ff090cccee --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/shutdown.ts @@ -0,0 +1,13 @@ +import type { WrapperKiloClient } from './kilo-api'; + +export async function abortKiloSessionForShutdown({ + activeKiloSessionId, + kiloClient, +}: { + activeKiloSessionId: string | undefined; + kiloClient: Pick | undefined; +}): Promise { + if (activeKiloSessionId && kiloClient) { + await kiloClient.abortSession({ sessionId: activeKiloSessionId }).catch(() => {}); + } +} From 5a8aa1c86de2f9e8e8ebda4b5ddbae2d3083b121 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Wed, 19 Aug 2026 12:51:24 -0500 Subject: [PATCH 4/6] docs(specs): remove private billing plan reference --- .specs/gastown-usage-based-billing.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/.specs/gastown-usage-based-billing.md b/.specs/gastown-usage-based-billing.md index 45a9104b81..897c17f1fb 100644 --- a/.specs/gastown-usage-based-billing.md +++ b/.specs/gastown-usage-based-billing.md @@ -2,9 +2,8 @@ ## Role of This Document -This is retained design history for the Gastown producer. The authoritative rollout and -settlement design is `fd-plans/research/container-billing-charge-and-enforce.md`; where this -document differs from that plan, the plan wins. +This is retained design history for the Gastown producer. It is superseded for rollout and +settlement decisions. ## Status From 77f42f7889a5add5a6173a380ee04dc39af976e7 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Wed, 19 Aug 2026 13:04:54 -0500 Subject: [PATCH 5/6] fix(cloud-agent-next): suppress completion after shutdown --- .../cloud-agent-next/wrapper/src/lifecycle.ts | 26 ++++---- .../wrapper/src/shutdown.test.ts | 66 +++++++++++++++++++ 2 files changed, 80 insertions(+), 12 deletions(-) diff --git a/services/cloud-agent-next/wrapper/src/lifecycle.ts b/services/cloud-agent-next/wrapper/src/lifecycle.ts index ca1ebb2f26..cefca3cae5 100644 --- a/services/cloud-agent-next/wrapper/src/lifecycle.ts +++ b/services/cloud-agent-next/wrapper/src/lifecycle.ts @@ -164,18 +164,20 @@ export function createLifecycleManager( if (completeSession && currentSession) { const currentBranch = await getCurrentBranch(config.workspacePath, 10_000).catch(() => ''); if (drainGeneration !== lifecycleGeneration) return; - const gateResult = state.consumeObservedGateResult(); - state.sendToIngest({ - streamEventType: 'complete', - data: { - exitCode: 0, - kiloSessionId: currentSession.kiloSessionId, - messageIds: sealedMessageIds, - ...(currentBranch ? { currentBranch } : {}), - ...(gateResult ? { gateResult } : {}), - }, - timestamp: new Date().toISOString(), - }); + if (!isAborted) { + const gateResult = state.consumeObservedGateResult(); + state.sendToIngest({ + streamEventType: 'complete', + data: { + exitCode: 0, + kiloSessionId: currentSession.kiloSessionId, + messageIds: sealedMessageIds, + ...(currentBranch ? { currentBranch } : {}), + ...(gateResult ? { gateResult } : {}), + }, + timestamp: new Date().toISOString(), + }); + } } await new Promise(resolve => setTimeout(resolve, DRAIN_DELAY_MS)); diff --git a/services/cloud-agent-next/wrapper/src/shutdown.test.ts b/services/cloud-agent-next/wrapper/src/shutdown.test.ts index d647ae9336..da4fc06c61 100644 --- a/services/cloud-agent-next/wrapper/src/shutdown.test.ts +++ b/services/cloud-agent-next/wrapper/src/shutdown.test.ts @@ -50,4 +50,70 @@ describe('abortKiloSessionForShutdown', () => { expect(events.map(event => event.streamEventType)).toEqual(['interrupted']); expect(state.currentSession).toBeNull(); }); + + it('does not complete a drain interrupted while final log upload is pending', async () => { + const events: IngestEvent[] = []; + const state = new WrapperState(); + state.bindSession({ + kiloSessionId: 'kilo_sess_test', + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + }); + state.setSendToIngestFn(event => events.push(event)); + state.acceptMessage('message-1', { autoCommit: false, condenseOnComplete: false }); + + let resolveUpload: (() => void) | undefined; + let signalUploadStarted: (() => void) | undefined; + const uploadStarted = new Promise(resolve => { + signalUploadStarted = resolve; + }); + let uploaderStopped = false; + state.setLogUploader({ + start: () => {}, + uploadNow: async () => { + signalUploadStarted?.(); + await new Promise(resolve => { + resolveUpload = resolve; + }); + }, + stop: () => { + uploaderStopped = true; + }, + }); + + let closeCalls = 0; + const lifecycleManager = createLifecycleManager( + { workspacePath: '/tmp' }, + { + state, + kiloClient: {} as WrapperKiloClient, + closeConnections: async () => { + closeCalls += 1; + }, + isConnected: () => true, + reconnectEventSubscription: () => {}, + } + ); + + const drain = lifecycleManager.drainAndClose(); + await uploadStarted; + + state.sendToIngest({ + streamEventType: 'interrupted', + data: { reason: 'Container shutdown' }, + timestamp: new Date().toISOString(), + }); + lifecycleManager.setAborted(); + if (!resolveUpload) throw new Error('Expected final log upload to be pending'); + resolveUpload(); + await drain; + + expect(events.map(event => event.streamEventType)).toEqual([ + 'wrapper_finalizing', + 'interrupted', + ]); + expect(uploaderStopped).toBe(true); + expect(closeCalls).toBe(1); + expect(state.currentSession).toBeNull(); + }); }); From 3fa5f6eeda0945b343905336d06305a86796f1c2 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Wed, 19 Aug 2026 14:59:30 -0500 Subject: [PATCH 6/6] fix(container-billing): expose facade admission failures --- .../src/kilo-facade/session-proxy.test.ts | 18 +++-- .../src/kilo-facade/session-proxy.ts | 19 ++++-- .../src/kilo-facade/user-kilo-facade.test.ts | 59 +++++++++++++++-- .../src/kilo-facade/user-kilo-facade.ts | 66 ++++++++++++++++--- .../src/billing-config.test.ts | 55 ++++++++++++++++ .../src/billing-config.ts | 8 ++- 6 files changed, 201 insertions(+), 24 deletions(-) diff --git a/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts b/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts index 07cd10bfb7..d58836f16f 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.test.ts @@ -67,7 +67,14 @@ describe('resolveLiveWrapperTarget billing admission', () => { userId: 'user_facade', cloudAgentSessionId: 'agent_facade', }) - ).resolves.toBeNull(); + ).resolves.toEqual({ + kind: 'billing-rejected', + admission: { + success: false, + code: 'insufficient_credits', + message: 'Low balance', + }, + }); expect(ensureBillingAdmission).toHaveBeenCalledWith( expect.objectContaining({ sandboxId: 'ses-facade', @@ -93,7 +100,7 @@ describe('resolveLiveWrapperTarget billing admission', () => { userId: 'user_facade', cloudAgentSessionId: 'agent_facade', }) - ).resolves.toMatchObject({ port: 5000 }); + ).resolves.toMatchObject({ kind: 'available', target: { port: 5000 } }); }); it('allows shadow acquisition when the callable billing block proxy rejects', async () => { @@ -112,7 +119,7 @@ describe('resolveLiveWrapperTarget billing admission', () => { userId: 'user_facade', cloudAgentSessionId: 'agent_facade', }) - ).resolves.toMatchObject({ port: 5000 }); + ).resolves.toMatchObject({ kind: 'available', target: { port: 5000 } }); }); it('fails enforced acquisition closed when the callable billing block proxy rejects', async () => { @@ -136,7 +143,10 @@ describe('resolveLiveWrapperTarget billing admission', () => { userId: 'user_facade', cloudAgentSessionId: 'agent_facade', }) - ).resolves.toBeNull(); + ).resolves.toMatchObject({ + kind: 'billing-rejected', + admission: { code: 'meter_unavailable' }, + }); expect(ensureBillingAdmission).toHaveBeenCalledOnce(); expect(mocks.findWrapperForSession).not.toHaveBeenCalled(); }); 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 f04f8b8cd9..aa67648367 100644 --- a/services/cloud-agent-next/src/kilo-facade/session-proxy.ts +++ b/services/cloud-agent-next/src/kilo-facade/session-proxy.ts @@ -10,6 +10,7 @@ import { ensureSandboxBillingAdmissionInput, isSandboxBillingBlocked, } from '../container-usage-context.js'; +import type { SandboxBillingAdmissionResult } from '../container-usage-context.js'; import { isCloudAgentContainerBillingEnabled } from '../container-billing-rollout.js'; export type SessionKiloFacadeDecision = @@ -30,6 +31,14 @@ export type LiveWrapperTarget = { port: number; }; +export type LiveWrapperResolution = + | { kind: 'available'; target: LiveWrapperTarget } + | { kind: 'unavailable' } + | { + kind: 'billing-rejected'; + admission: Extract; + }; + export function decideSessionKiloFacadeRoute( input: SessionKiloFacadePolicyInput ): SessionKiloFacadeDecision { @@ -65,11 +74,11 @@ export async function resolveLiveWrapperTarget(params: { env: Env; userId: string; cloudAgentSessionId: string; -}): Promise { +}): Promise { const { env, userId, cloudAgentSessionId } = params; const metadata = await fetchSessionMetadata(env, userId, cloudAgentSessionId); if (!metadata) { - return null; + return { kind: 'unavailable' }; } const sessionId = cloudAgentSessionId as SessionId; @@ -100,14 +109,14 @@ export async function resolveLiveWrapperTarget(params: { const billingBlocked = await isSandboxBillingBlocked(sandbox, billingInput.enforcementRequested); if (billingInput.enforcementRequested || billingBlocked) { const admission = await ensureSandboxBillingAdmissionInput(sandbox, billingInput); - if (!admission.success) return null; + if (!admission.success) return { kind: 'billing-rejected', admission }; } else { void configureSandboxBillingInput(sandbox, billingInput); } const wrapperInfo = await findWrapperForSession(sandbox, sessionId); if (!wrapperInfo) { - return null; + return { kind: 'unavailable' }; } - return { sandbox, port: wrapperInfo.port }; + return { kind: 'available', target: { sandbox, port: wrapperInfo.port } }; } diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts index 39ae833945..29b0135e09 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts @@ -39,7 +39,7 @@ import { rewriteGlobalEventDirectory, UserKiloFacade, } from './user-kilo-facade'; -import type { LiveWrapperTarget, SessionKiloFacadeDecision } from './session-proxy'; +import type { LiveWrapperResolution, SessionKiloFacadeDecision } from './session-proxy'; const kiloSessionId = 'ses_12345678901234567890123456'; type ContainerFetch = (request: Request, port: number) => Promise; @@ -364,10 +364,16 @@ function projectedSdkMessageHistory(): KiloSdkStoredMessage[] { ]; } -function liveWrapperTarget(containerFetch: ContainerFetch): LiveWrapperTarget { +function liveWrapperTarget(containerFetch: ContainerFetch): LiveWrapperResolution { return { - port: 5123, - sandbox: { containerFetch } as unknown as LiveWrapperTarget['sandbox'], + kind: 'available', + target: { + port: 5123, + sandbox: { containerFetch } as unknown as Extract< + LiveWrapperResolution, + { kind: 'available' } + >['target']['sandbox'], + }, }; } @@ -728,7 +734,11 @@ describe('handleKiloFacadeRequest', () => { resolveRootSessionForKiloSession: vi.fn(async () => ({ cloudAgentSessionId: 'agent_cold', })), - resolveLiveWrapper: vi.fn(async () => null), + resolveLiveWrapper: vi.fn( + async (): Promise => ({ + kind: 'unavailable', + }) + ), }, }); @@ -740,6 +750,45 @@ describe('handleKiloFacadeRequest', () => { }); }); + it('returns a billing rejection instead of reading a persisted detail snapshot', async () => { + const env = envStub(); + const getCloudAgentRootSessionSnapshot = vi.mocked( + env.SESSION_INGEST.getCloudAgentRootSessionSnapshot + ); + + const response = await handleKiloFacadeRequest({ + request: new Request(`http://worker.test/kilo/session/${kiloSessionId}`), + env, + userId: 'usr_1', + deps: { + resolveRootSessionForKiloSession: vi.fn(async () => ({ + cloudAgentSessionId: 'agent_billing_blocked', + })), + resolveLiveWrapper: vi.fn( + async (): Promise => ({ + kind: 'billing-rejected', + admission: { + success: false, + code: 'insufficient_credits', + message: 'internal admission detail', + remainingMicrodollars: 1, + minimumRequiredMicrodollars: 5, + }, + }) + ), + }, + }); + + expect(response.status).toBe(402); + await expect(response.json()).resolves.toEqual({ + error: 'KILO_BILLING_PAYMENT_REQUIRED', + message: 'Container billing requires additional credits', + remainingMicrodollars: 1, + minimumRequiredMicrodollars: 5, + }); + expect(getCloudAgentRootSessionSnapshot).not.toHaveBeenCalled(); + }); + it('falls back to the persisted detail snapshot when the wrapper has no private Kilo runtime', async () => { const env = envStub(); vi.mocked(env.SESSION_INGEST.getCloudAgentRootSessionSnapshot).mockResolvedValue({ diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 0fef97db60..ebae8405c9 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -42,7 +42,7 @@ import { buildWrapperKiloProxyUrl, decideSessionKiloFacadeRoute, resolveLiveWrapperTarget, - type LiveWrapperTarget, + type LiveWrapperResolution, type SessionKiloFacadeDecision, type SessionKiloFacadePolicyInput, } from './session-proxy.js'; @@ -115,7 +115,7 @@ export type KiloFacadeRequestDeps = { env: Env; userId: string; cloudAgentSessionId: string; - }) => Promise; + }) => Promise; admitPrompt?: (params: { env: Env; userId: string; @@ -142,8 +142,50 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } -function facadeError(status: number, code: string, message: string): Response { - return Response.json({ error: code, message }, { status }); +function facadeError( + status: number, + code: string, + message: string, + details?: Record +): Response { + return Response.json({ error: code, message, ...details }, { status }); +} + +function billingAdmissionRejectedResponse( + admission: Extract['admission'] +): Response { + const details = { + ...(admission.remainingMicrodollars === undefined + ? {} + : { remainingMicrodollars: admission.remainingMicrodollars }), + ...(admission.minimumRequiredMicrodollars === undefined + ? {} + : { minimumRequiredMicrodollars: admission.minimumRequiredMicrodollars }), + }; + switch (admission.code) { + case 'insufficient_credits': + return facadeError( + 402, + 'KILO_BILLING_PAYMENT_REQUIRED', + 'Container billing requires additional credits', + details + ); + case 'stopping': + return facadeError( + 409, + 'KILO_BILLING_BLOCKED', + 'Container billing is unavailable while the session is stopping', + details + ); + case 'meter_unavailable': + case 'configuration_mismatch': + return facadeError( + 503, + 'KILO_BILLING_UNAVAILABLE', + 'Container billing is temporarily unavailable', + details + ); + } } function kiloRelativePath(pathname: string): string { @@ -1099,7 +1141,7 @@ async function proxyOwnedKiloSessionRequest(params: { }) : null; - let liveWrapper: LiveWrapperTarget | null; + let liveWrapper: LiveWrapperResolution | null; try { liveWrapper = await (deps?.resolveLiveWrapper ?? resolveLiveWrapperTarget)({ env, @@ -1111,7 +1153,13 @@ async function proxyOwnedKiloSessionRequest(params: { if (fallback) return fallback; throw error; } - if (!liveWrapper) { + if (liveWrapper === null) { + liveWrapper = { kind: 'unavailable' }; + } + if (liveWrapper.kind === 'billing-rejected') { + return billingAdmissionRejectedResponse(liveWrapper.admission); + } + if (liveWrapper.kind === 'unavailable') { const fallback = persistedFallback(); if (fallback) return fallback; return facadeError( @@ -1121,18 +1169,20 @@ async function proxyOwnedKiloSessionRequest(params: { ); } + const target = liveWrapper.target; + const upstreamSearchParams = new URLSearchParams(url.searchParams); upstreamSearchParams.delete('directory'); const upstreamSearch = upstreamSearchParams.size > 0 ? `?${upstreamSearchParams.toString()}` : ''; const targetUrl = buildWrapperKiloProxyUrl({ - wrapperPort: liveWrapper.port, + wrapperPort: target.port, kiloRelativePath: kiloPath, search: upstreamSearch, }); const proxyRequest = createProxyRequest(request, targetUrl); let response: Response; try { - response = await liveWrapper.sandbox.containerFetch(proxyRequest, liveWrapper.port); + response = await target.sandbox.containerFetch(proxyRequest, target.port); } catch (error) { const fallback = persistedFallback(); if (fallback) return fallback; diff --git a/services/container-usage-meter/src/billing-config.test.ts b/services/container-usage-meter/src/billing-config.test.ts index b7b4c29123..77b40d85bc 100644 --- a/services/container-usage-meter/src/billing-config.test.ts +++ b/services/container-usage-meter/src/billing-config.test.ts @@ -83,6 +83,61 @@ describe('container billing configuration', () => { ); }); + it('allows the Cloud Agent family token for current and future service classes only', () => { + const config = billingConfigFromEnv( + env({ + CONTAINER_BILLING_SERVICES: 'gastown,cloud-agent-next', + CONTAINER_BILLING_USER_IDS: 'gastown-user', + CONTAINER_BILLING_ORG_IDS: '', + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS: 'cloud-agent-user', + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS: '', + CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS: '10000000', + }) + ); + + expect( + billingModeFor(config, 'cloud-agent-next-sandbox', { type: 'user', id: 'cloud-agent-user' }) + ).toBe('paid'); + expect( + billingModeFor(config, 'cloud-agent-next-future-class', { + type: 'user', + id: 'cloud-agent-user', + }) + ).toBe('paid'); + expect( + billingModeFor(config, 'cloud-agent-next', { type: 'user', id: 'cloud-agent-user' }) + ).toBe('shadow'); + expect( + billingModeFor(config, 'cloud-agent-other-sandbox', { + type: 'user', + id: 'cloud-agent-user', + }) + ).toBe('shadow'); + }); + + it('keeps exact Cloud Agent class tokens available for per-class rollout', () => { + const config = billingConfigFromEnv( + env({ + CONTAINER_BILLING_SERVICES: 'cloud-agent-next-sandbox', + CONTAINER_BILLING_USER_IDS: '', + CONTAINER_BILLING_ORG_IDS: '', + CONTAINER_BILLING_CLOUD_AGENT_USER_IDS: 'cloud-agent-user', + CONTAINER_BILLING_CLOUD_AGENT_ORG_IDS: '', + CONTAINER_BILLING_WARN_REMAINING_MICRODOLLARS: '10000000', + }) + ); + + expect( + billingModeFor(config, 'cloud-agent-next-sandbox', { type: 'user', id: 'cloud-agent-user' }) + ).toBe('paid'); + expect( + billingModeFor(config, 'cloud-agent-next-future-class', { + type: 'user', + id: 'cloud-agent-user', + }) + ).toBe('shadow'); + }); + it('requires Cloud Agent service names as well as Cloud Agent payer lists', () => { const config = billingConfigFromEnv( env({ diff --git a/services/container-usage-meter/src/billing-config.ts b/services/container-usage-meter/src/billing-config.ts index e1739bd747..3292a5b534 100644 --- a/services/container-usage-meter/src/billing-config.ts +++ b/services/container-usage-meter/src/billing-config.ts @@ -93,8 +93,12 @@ export function billingModeFor( service: string, subject: { type: 'user' | 'org'; id: string } ): 'shadow' | 'paid' { - if (!config.enabled || !config.services.has(service)) return 'shadow'; - const payerIds = service.startsWith(CLOUD_AGENT_SERVICE_PREFIX) + const isCloudAgentService = service.startsWith(CLOUD_AGENT_SERVICE_PREFIX); + const serviceEnabled = + config.services.has(service) || + (isCloudAgentService && config.services.has('cloud-agent-next')); + if (!config.enabled || !serviceEnabled) return 'shadow'; + const payerIds = isCloudAgentService ? subject.type === 'user' ? config.cloudAgentUserIds : config.cloudAgentOrgIds