From d59376691cf20892d83f15d32b91d6dbe016ac4e Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Thu, 3 Sep 2026 14:54:31 +0200 Subject: [PATCH] feat(cloud-agent-next): add organization and personal sandbox selection Add organization and personal sandbox destination selection for cloud-agent-next sessions. The capabilities contract now carries only the allocations an owner can use: presence in options means available, unavailable Vercel rows are omitted instead of listed greyed-out, and the picker no longer renders disabled rows or an unavailable-reason sub-label. Worker submission still rejects unavailable allocations. SANDBOX_SELECTION_IDS replaces the org-only allowlist so personal owners can also pick a destination, and the last-used allocation is persisted per owner. Explicit Dedicated Standard allocations require organization membership and SANDBOX_SELECTION_ORG_IDS, matching the other explicit destinations. The picker and sandbox status show container capacity. --- ENVIRONMENT.md | 1 + .../cloud-agent-next/NewSessionPanel.tsx | 439 ++++++--- .../SandboxStatusIndicator.tsx | 6 + .../model-preferences.test.ts | 12 + .../cloud-agent-next/model-preferences.ts | 24 + .../sandbox-selection.test.ts | 492 ++++++++++ .../cloud-agent-next/sandbox-selection.ts | 192 ++++ .../cloud-agent-next/sandbox-status.test.ts | 76 ++ .../cloud-agent-next/sandbox-status.ts | 35 + .../webhook-triggers/TriggerForm.tsx | 4 +- .../cloud-agent-client.test.ts | 142 +++ .../cloud-agent-next/cloud-agent-client.ts | 22 + .../src/lib/cloudflare/container-capacity.ts | 4 + .../routers/cloud-agent-next-router.test.ts | 138 ++- .../src/routers/cloud-agent-next-router.ts | 11 + .../routers/cloud-agent-next-schemas.test.ts | 174 ++++ .../src/routers/cloud-agent-next-schemas.ts | 30 +- ...ganization-cloud-agent-next-router.test.ts | 275 +++++- .../organization-cloud-agent-next-router.ts | 26 +- .../webhook-triggers-router.schema.test.ts | 152 +-- .../src/routers/webhook-triggers-router.ts | 29 +- packages/worker-utils/package.json | 1 + .../src/sandbox-allocation.test.ts | 164 ++++ .../worker-utils/src/sandbox-allocation.ts | 176 ++++ services/cloud-agent-next/.dev.vars.example | 2 + .../vercel/vercel-runtime-config.test.ts | 44 + .../vercel/vercel-runtime-config.ts | 27 +- .../vercel/vercel-sandbox-rest-client.test.ts | 73 ++ .../vercel/vercel-sandbox-rest-client.ts | 19 +- .../src/callbacks/queue-config.test.ts | 4 +- .../src/persistence/SandboxControl.ts | 104 ++- .../src/persistence/session-metadata.test.ts | 74 ++ .../src/persistence/session-metadata.ts | 66 +- services/cloud-agent-next/src/router.ts | 2 + .../src/router/handlers/sandbox-selection.ts | 31 + .../src/router/handlers/session-start.ts | 1 + .../router/handlers/session-worktree.test.ts | 192 +++- .../src/router/handlers/session-worktree.ts | 24 +- .../src/router/schemas.test.ts | 89 ++ .../cloud-agent-next/src/router/schemas.ts | 35 +- .../src/sandbox-control/lifecycle.test.ts | 139 ++- .../src/sandbox-control/physical-lifecycle.ts | 13 +- .../sandbox-control/vercel-provider.test.ts | 21 + .../src/sandbox-control/vercel-provider.ts | 2 + .../cloud-agent-next/src/sandbox-id.test.ts | 74 ++ services/cloud-agent-next/src/sandbox-id.ts | 121 ++- .../src/sandbox-selection.test.ts | 270 ++++++ .../cloud-agent-next/src/sandbox-selection.ts | 84 ++ .../src/sandbox-session/SandboxSession.ts | 2 + .../src/sandbox-session/control-rpc.ts | 2 + .../session-message-queue.test.ts | 35 + .../src/session-prepare.test.ts | 274 ++++++ .../src/session/session-prepare.test.ts | 873 +++++++++++++++++- .../src/session/session-registration.ts | 165 +++- .../src/session/session-requests.ts | 3 +- services/cloud-agent-next/src/types.ts | 2 + .../test/integration/sandbox-control.test.ts | 211 ++++- .../worker-configuration.d.ts | 4 +- services/cloud-agent-next/wrangler.jsonc | 1 + 59 files changed, 5393 insertions(+), 315 deletions(-) create mode 100644 apps/web/src/components/cloud-agent-next/sandbox-selection.test.ts create mode 100644 apps/web/src/components/cloud-agent-next/sandbox-selection.ts create mode 100644 packages/worker-utils/src/sandbox-allocation.test.ts create mode 100644 packages/worker-utils/src/sandbox-allocation.ts create mode 100644 services/cloud-agent-next/src/router/handlers/sandbox-selection.ts create mode 100644 services/cloud-agent-next/src/sandbox-selection.test.ts create mode 100644 services/cloud-agent-next/src/sandbox-selection.ts diff --git a/ENVIRONMENT.md b/ENVIRONMENT.md index f4e6547e2e..f2c335d212 100644 --- a/ENVIRONMENT.md +++ b/ENVIRONMENT.md @@ -394,6 +394,7 @@ The key is team-scoped for all topics and valid in both the sandbox and producti - `SESSION_ID` - Reserved session identifier for the `cloud-agent-next` runtime; reserved in `RESERVED_ENV_VARS`. [SERVER] - `CONTROL_PLANE_IDS` - Comma-separated user or org IDs admitted to the call-home control plane at interactive web (`cloud-agent-web`) session creation. Empty admits nobody. `*` includes personal accounts. Omitted from production `wrangler.jsonc` so the Cloudflare dashboard value survives deploy; unset admits nobody. Wrangler `dev` and `.dev.vars.example` default to `*`. Non-interactive origins (Slack, scheduled, code review, and similar) keep legacy `agent_` sessions even when enrolled. Does not enable new worktree creation by itself; that also requires `WORKTREE_CREATION_ENABLED_IDS` enrollment. [SERVER] - `WORKTREE_CREATION_ENABLED_IDS` - Comma-separated user or org IDs allowed to create new worktrees, or `*` for all, including personal accounts. Omitted from production `wrangler.jsonc` so the Cloudflare dashboard value survives deploy; unset is off. Wrangler `dev` and `.dev.vars.example` default to `*`. Also requires enrollment in `CONTROL_PLANE_IDS`. Disabling it does not block existing worktrees or sibling chats in them. [SERVER] +- `SANDBOX_SELECTION_IDS` - Comma-separated user or org IDs allowed to pick a Cloud Agent sandbox destination on the new-session page. Empty admits nobody. `*` includes personal accounts. Omitted from production `wrangler.jsonc` so the Cloudflare dashboard value survives deploy; unset admits nobody. Wrangler `dev` and `.dev.vars.example` default to `*`. [SERVER] - `VERCEL_SANDBOX_ORG_IDS` - Comma-separated org IDs routed to Vercel sandboxes. Empty is off. `*` includes personal accounts. [SERVER] - `HOME` - Reserved in `RESERVED_ENV_VARS` for cloud-agent-next session home management. [SYSTEM] diff --git a/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx b/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx index 14d9779cf3..89c8560445 100644 --- a/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx +++ b/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx @@ -24,6 +24,28 @@ import { startOfDay, subDays } from 'date-fns'; import { useTRPC, useRawTRPCClient } from '@/lib/trpc/utils'; import { SetPageTitle } from '@/components/SetPageTitle'; import { Badge } from '@/components/ui/badge'; +import { + Select, + SelectContent, + SelectGroup, + SelectItem, + SelectLabel, + SelectSeparator, + SelectTrigger, + SelectValue, +} from '@/components/ui/select'; +import { getSandboxAllocationKey } from '@kilocode/worker-utils/sandbox-allocation'; +import { + formatSandboxDestination, + formatSandboxDestinationWithoutTier, + formatSandboxInstance, + getSandboxSelectionGroups, + getSandboxSelectionOptions, + getPreferredInitialSandboxAllocation, + resolveSandboxSelection, + resolveSandboxSelectionSubmissionError, + type SandboxSelectionDraft, +} from './sandbox-selection'; import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip'; import { MobileSidebarToggle } from './MobileSidebarToggle'; import { MobileToolbarPopover } from './MobileToolbarPopover'; @@ -81,7 +103,10 @@ import { CloudAgentBillingError } from './CloudAgentBillingError'; import { billingPayerPresentation } from './billing-payer-presentation'; import type { OrganizationRole } from '@/lib/organizations/organization-types'; import { generateMessageId } from '@kilocode/cloud-agent-sdk/message-id'; -import { useCloudAgentAttachmentUpload } from '@/hooks/useCloudAgentAttachmentUpload'; +import { + buildCloudAgentAttachments, + useCloudAgentAttachmentUpload, +} from '@/hooks/useCloudAgentAttachmentUpload'; import { AttachmentPreviewStrip } from './AttachmentPreviewStrip'; import { CLOUD_AGENT_ATTACHMENT_MAX_COUNT, @@ -91,6 +116,7 @@ import { getDevcontainerEnabled, getLastUsedModel, getLastUsedRepo, + getLastUsedSandboxAllocationKey, getLastUsedVariant, getPreferredInitialModel, getPreferredInitialRepo, @@ -98,6 +124,7 @@ import { setDevcontainerEnabled, setLastUsedModel, setLastUsedRepo, + setLastUsedSandboxAllocationKey, setLastUsedVariant, } from '@/components/cloud-agent-next/model-preferences'; import { @@ -153,7 +180,7 @@ export function NewSessionPanel({ const textareaRef = useRef(null); const fileInputRef = useRef(null); const commandListRef = useRef(null); - const firstChatCreationOperationRef = useRef< + const [firstChatCreationOperation, setFirstChatCreationOperation] = useState< (CloudSessionCreationOperation & { initialMessageId: string }) | null >(null); const [devcontainer, setDevcontainer] = useState(false); @@ -291,6 +318,69 @@ export function NewSessionPanel({ const effectiveDevcontainer = isDevcontainerAvailable && devcontainer; const availableVariants = modelOptions.find(m => m.id === model)?.variants ?? []; + const [sandboxSelection, setSandboxSelection] = useState({ + organizationId, + }); + if (sandboxSelection.organizationId !== organizationId) { + setSandboxSelection({ organizationId }); + } + const sandboxPreferenceOwnerRef = useRef(undefined); + const sandboxSelectionQuery = useQuery({ + ...(organizationId + ? trpc.organizations.cloudAgentNext.getSandboxSelectionOptions.queryOptions({ + organizationId, + ...(effectiveDevcontainer ? { devcontainer: true } : {}), + }) + : trpc.cloudAgentNext.getSandboxSelectionOptions.queryOptions({ + ...(effectiveDevcontainer ? { devcontainer: true } : {}), + })), + retry: false, + }); + const sandboxCapabilities = sandboxSelectionQuery.isSuccess + ? sandboxSelectionQuery.data + : undefined; + const showSandboxSelector = + sandboxCapabilities?.enabled === true || + (sandboxSelection.organizationId === organizationId && !!sandboxSelection.allocation); + const { sandboxAllocation, error: sandboxAvailabilityError } = resolveSandboxSelection({ + organizationId, + draft: sandboxSelection, + capabilities: sandboxCapabilities, + devcontainer: effectiveDevcontainer, + }); + const sandboxOptions = getSandboxSelectionOptions(sandboxCapabilities); + const sandboxGroups = getSandboxSelectionGroups(sandboxOptions); + const sandboxDestination = sandboxAllocation ?? sandboxCapabilities?.defaultDestination; + const sandboxDestinationLabel = formatSandboxDestination(sandboxDestination); + const sandboxDestinationWithoutTier = formatSandboxDestinationWithoutTier(sandboxDestination); + const defaultSandboxLabel = sandboxCapabilities?.defaultDestination + ? `Default · ${formatSandboxDestination(sandboxCapabilities.defaultDestination)}` + : 'Default'; + + // --------------------------------------------------------------------------- + // Sandbox destination auto-selection + // --------------------------------------------------------------------------- + // Each owner restores its own saved allocation. Clear the previous owner's + // guard here so auto-selection cannot skip restoration when returning to an + // owner whose capabilities were previously unavailable. + useEffect(() => { + sandboxPreferenceOwnerRef.current = undefined; + }, [organizationId]); + + useEffect(() => { + if (!sandboxCapabilities?.enabled) return; + + const owner = organizationId ?? null; + if (sandboxPreferenceOwnerRef.current === owner) return; + sandboxPreferenceOwnerRef.current = owner; + + const allocation = getPreferredInitialSandboxAllocation({ + options: sandboxCapabilities.options, + lastUsedKey: getLastUsedSandboxAllocationKey(organizationId), + }); + if (allocation) setSandboxSelection({ organizationId, allocation }); + }, [sandboxCapabilities, organizationId]); + // --------------------------------------------------------------------------- // Model auto-selection // --------------------------------------------------------------------------- @@ -361,7 +451,10 @@ export function NewSessionPanel({ setVariant(newVariant); if (model) { setLastUsedVariant(model, newVariant, organizationId); - persistServerLastSelected({ model, ...(newVariant ? { variant: newVariant } : {}) }); + persistServerLastSelected({ + model, + ...(newVariant ? { variant: newVariant } : {}), + }); } }, [model, organizationId, persistServerLastSelected] @@ -729,7 +822,9 @@ export function NewSessionPanel({ } void queryClient.invalidateQueries({ - queryKey: trpc.organizations.bitbucket.getStatus.queryKey({ organizationId }), + queryKey: trpc.organizations.bitbucket.getStatus.queryKey({ + organizationId, + }), }); }, }) @@ -776,7 +871,9 @@ export function NewSessionPanel({ ); } void queryClient.invalidateQueries({ - queryKey: trpc.organizations.bitbucket.getStatus.queryKey({ organizationId }), + queryKey: trpc.organizations.bitbucket.getStatus.queryKey({ + organizationId, + }), }); if (result.status !== 'available') { throw new Error(getBitbucketRepositoryRefreshFailureMessage(result.status)); @@ -989,65 +1086,110 @@ export function NewSessionPanel({ !!selectedModelOption && (selectedModelOption.isFree || selectedModelOption.hasUserByokAvailable); + const bitbucketRepo = useMemo(() => { + if (!organizationId || selectedPlatform !== 'bitbucket') return undefined; + const repository = unifiedRepositories.find( + repository => repository.fullName === selectedRepo && repository.platform === 'bitbucket' + ); + if (!repository || typeof repository.id !== 'string' || !repository.workspaceUuid) { + return undefined; + } + return { + fullName: selectedRepo, + workspaceUuid: repository.workspaceUuid, + repositoryUuid: repository.id, + }; + }, [organizationId, selectedPlatform, selectedRepo, unifiedRepositories]); + const creationInput = useMemo(() => { + const trimmed = prompt.trim(); + const slashMatch = /^\s*\/([\w.-]+)(?:\s+([\s\S]*))?\s*$/.exec(trimmed); + const slashCommand = + slashMatch && slashCommands.some(command => command.trigger === slashMatch[1]) + ? { command: slashMatch[1], args: slashMatch[2]?.trim() ?? '' } + : null; + return { + prompt: trimmed, + mode, + model: displayModel, + variant: displayVariant, + profileId: selectedProfileId ?? undefined, + autoCommit: true, + autoInitiate: true, + attachments: buildCloudAgentAttachments(attachmentMessageUuid, attachmentUpload.attachments), + ...(slashCommand + ? { + initialPayload: { + type: 'command' as const, + command: slashCommand.command, + arguments: slashCommand.args, + }, + } + : {}), + ...(effectiveDevcontainer ? { devcontainer: true } : {}), + }; + }, [ + prompt, + slashCommands, + mode, + displayModel, + displayVariant, + selectedProfileId, + attachmentMessageUuid, + attachmentUpload.attachments, + effectiveDevcontainer, + ]); + const creationIntent = JSON.stringify({ + ...creationInput, + organizationId: organizationId ?? null, + repository: selectedRepo, + platform: selectedPlatform, + ...(organizationId && selectedPlatform === 'github' + ? { githubIntegrationId: selectedGitHubIntegrationId } + : {}), + bitbucketRepo, + sandboxAllocation: sandboxAllocation ? getSandboxAllocationKey(sandboxAllocation) : undefined, + }); + const sandboxSelectionError = resolveSandboxSelectionSubmissionError({ + error: sandboxAvailabilityError, + intent: creationIntent, + pendingOperation: firstChatCreationOperation, + }); + const sandboxDescriptionId = sandboxSelectionError + ? 'new-session-sandbox-error' + : effectiveDevcontainer + ? 'new-session-sandbox-devcontainer' + : undefined; + const isFormValid = prompt.trim().length > 0 && !isPromptTooLong && model.length > 0 && !isPreparing && + !sandboxSelectionError && (!hasInsufficientBalance || limitedAccessModelIsAllowed) && !attachmentUpload.hasUploadingAttachments; const handleStartSession = useCallback(async () => { - if (!prompt.trim() || attachmentUpload.hasUploadingAttachments) return; + if ( + !prompt.trim() || + attachmentUpload.hasUploadingAttachments || + isPreparing || + sandboxSelectionError + ) + return; if (!selectedRepo) { setShowRepositoryRequiredMessage(true); return; } - const selectedRepository = unifiedRepositories.find( - repository => - repository.fullName === selectedRepo && - repository.platform === selectedPlatform && - (selectedPlatform !== 'github' || - repository.platformIntegrationId === selectedGitHubIntegrationId) - ); - if ( - selectedPlatform === 'bitbucket' && - (!organizationId || - !selectedRepository || - typeof selectedRepository.id !== 'string' || - !selectedRepository.workspaceUuid) - ) { + if (selectedPlatform === 'bitbucket' && !bitbucketRepo) { toast.error('Select the Bitbucket repository again.'); return; } - const bitbucketRepo = - organizationId && - selectedPlatform === 'bitbucket' && - selectedRepository && - typeof selectedRepository.id === 'string' && - selectedRepository.workspaceUuid - ? { - fullName: selectedRepository.fullName, - workspaceUuid: selectedRepository.workspaceUuid, - repositoryUuid: selectedRepository.id, - } - : undefined; setIsPreparing(true); try { - const trimmed = prompt.trim(); - - // Parse slash command: if the input matches a known command, send a - // structured initialPayload so the backend dispatches a command rather - // than treating the text as a free-text prompt. - const slashMatch = /^\s*\/([\w.-]+)(?:\s+([\s\S]*))?\s*$/.exec(trimmed); - const slashCommand = - slashMatch && slashCommands.some(c => c.trigger === slashMatch[1]) - ? { command: slashMatch[1], args: slashMatch[2]?.trim() ?? '' } - : null; - - if (slashCommand && attachmentUpload.attachments.length > 0) { + if (creationInput.initialPayload && attachmentUpload.attachments.length > 0) { toast.error('Files cannot be attached to slash commands', { description: 'Remove the files or type a plain prompt instead.', }); @@ -1055,40 +1197,17 @@ export function NewSessionPanel({ return; } - const creationInput = { - prompt: trimmed, - mode, - model: displayModel, - variant: displayVariant, - profileId: selectedProfileId ?? undefined, - autoCommit: true, - autoInitiate: true, - attachments: await attachmentUpload.finalizeAttachments(), - ...(slashCommand - ? { - initialPayload: { - type: 'command' as const, - command: slashCommand.command, - arguments: slashCommand.args, - }, - } - : {}), - ...(effectiveDevcontainer ? { devcontainer: true } : {}), - }; - const intent = JSON.stringify({ - ...creationInput, - organizationId: organizationId ?? null, - repository: selectedRepo, - platform: selectedPlatform, - bitbucketRepo, - }); - const previousOperation = firstChatCreationOperationRef.current; - const pendingOperation = getCloudSessionCreationOperation(previousOperation, intent, uuidv4); + const previousOperation = firstChatCreationOperation; + const pendingOperation = getCloudSessionCreationOperation( + previousOperation, + creationIntent, + uuidv4 + ); const operation = previousOperation?.operationKey === pendingOperation.operationKey ? previousOperation : { ...pendingOperation, initialMessageId: generateMessageId() }; - firstChatCreationOperationRef.current = operation; + setFirstChatCreationOperation(operation); const baseInput = { ...creationInput, initialMessageId: operation.initialMessageId, @@ -1097,41 +1216,45 @@ export function NewSessionPanel({ let result: { kiloSessionId: string; cloudAgentSessionId: string }; if (organizationId) { + const organizationInput = { + ...baseInput, + organizationId, + ...(sandboxAllocation ? { sandboxAllocation } : {}), + }; if (selectedPlatform === 'gitlab') { result = await trpcClient.organizations.cloudAgentNext.prepareSession.mutate({ - ...baseInput, + ...organizationInput, gitlabProject: selectedRepo, - organizationId, }); } else if (selectedPlatform === 'bitbucket' && bitbucketRepo) { result = await trpcClient.organizations.cloudAgentNext.prepareSession.mutate({ - ...baseInput, + ...organizationInput, bitbucketRepo, - organizationId, }); } else { result = await trpcClient.organizations.cloudAgentNext.prepareSession.mutate({ - ...baseInput, + ...organizationInput, githubRepo: selectedRepo, githubIntegrationId: selectedGitHubIntegrationId, - organizationId, }); } } else if (selectedPlatform === 'gitlab') { result = await trpcClient.cloudAgentNext.prepareSession.mutate({ ...baseInput, + ...(sandboxAllocation ? { sandboxAllocation } : {}), gitlabProject: selectedRepo, }); } else { result = await trpcClient.cloudAgentNext.prepareSession.mutate({ ...baseInput, + ...(sandboxAllocation ? { sandboxAllocation } : {}), githubRepo: selectedRepo, }); } - if (firstChatCreationOperationRef.current?.operationKey === operation.operationKey) { - firstChatCreationOperationRef.current = null; - } + setFirstChatCreationOperation(current => + current?.operationKey === operation.operationKey ? null : current + ); if (!hasAgentModelOverride) { setLastUsedModel(model, organizationId); @@ -1157,43 +1280,40 @@ export function NewSessionPanel({ setBillingFailure(null); } catch (error) { if (!isAmbiguousCloudSessionCreationError(error)) { - firstChatCreationOperationRef.current = null; + setFirstChatCreationOperation(null); } const failure = parseCustomerBillingFailure(error); setBillingFailure(failure); console.error('Failed to prepare session:', error); if (!failure) { - toast.error('Failed to create session', { description: formatSessionError(error) }); + toast.error('Failed to create session', { + description: formatSessionError(error), + }); } } finally { setIsPreparing(false); } }, [ - effectiveDevcontainer, attachmentUpload, - displayModel, - // `displayVariant` is what we actually submit; raw `variant` is only read - // inside the `!hasAgentModelOverride` branch for last-used persistence, so - // keeping `displayVariant` (which equals `variant` in that branch) here is - // sufficient and avoids the stale-variant race when the agent-provided - // override changes while `variant`/`model`/`mode`/`hasAgentModelOverride` - // stay the same. - displayVariant, + bitbucketRepo, + creationInput, + creationIntent, + firstChatCreationOperation, hasAgentModelOverride, + isPreparing, + sandboxAllocation, + sandboxSelectionError, model, - mode, + variant, organizationId, prompt, queryClient, router, selectedPlatform, - selectedProfileId, selectedRepo, selectedGitHubIntegrationId, - slashCommands, trpc.cliSessionsV2.list, trpcClient, - unifiedRepositories, ]); // --------------------------------------------------------------------------- @@ -1299,7 +1419,11 @@ export function NewSessionPanel({ currentUserId, organization: organizationId && organizationName && organizationRole - ? { id: organizationId, name: organizationName, role: organizationRole } + ? { + id: organizationId, + name: organizationName, + role: organizationRole, + } : undefined, })} /> @@ -1491,7 +1615,8 @@ export function NewSessionPanel({ - Model is locked by agent “{selectedCustomAgent?.name}” + Model is locked by agent “{selectedCustomAgent?.name} + ” ) : ( @@ -1515,7 +1640,8 @@ export function NewSessionPanel({ - Locked by agent “{selectedCustomAgent?.name}” + Locked by agent “{selectedCustomAgent?.name} + ” ) @@ -1579,7 +1705,7 @@ export function NewSessionPanel({ {/* Repo + Settings row (outside prompt box) */} -
+
{/* Repo — bottom left */} @@ -1755,12 +1881,83 @@ export function NewSessionPanel({ -
+
{effectiveDevcontainer && ( Dev container on )} + {showSandboxSelector && ( + + )}
+ {showSandboxSelector && effectiveDevcontainer && ( +

+ Dev containers use the Default sandbox. Turn off dev containers in Profile to choose a + sandbox. +

+ )} + {sandboxSelectionError && ( +
+ + { + setSandboxSelection({ organizationId }); + setLastUsedSandboxAllocationKey(undefined, organizationId); + }} + > + Use Default + + {sandboxSelectionQuery.isError && ( + void sandboxSelectionQuery.refetch()} + > + Retry + + )} +
+ )} {githubIdentityHint && ( )} diff --git a/apps/web/src/components/cloud-agent-next/SandboxStatusIndicator.tsx b/apps/web/src/components/cloud-agent-next/SandboxStatusIndicator.tsx index 6518c1035e..be736aafa7 100644 --- a/apps/web/src/components/cloud-agent-next/SandboxStatusIndicator.tsx +++ b/apps/web/src/components/cloud-agent-next/SandboxStatusIndicator.tsx @@ -57,6 +57,12 @@ function SandboxStatusDetails({ view }: { view: SandboxStatusPresentation }) {
{view.provider}
Sandbox type
{view.sandboxType}
+ {view.capacity !== null && ( + <> +
Capacity
+
{view.capacity}
+ + )} {hasTiming && ( diff --git a/apps/web/src/components/cloud-agent-next/model-preferences.test.ts b/apps/web/src/components/cloud-agent-next/model-preferences.test.ts index cbc29b973b..7aa466e1f0 100644 --- a/apps/web/src/components/cloud-agent-next/model-preferences.test.ts +++ b/apps/web/src/components/cloud-agent-next/model-preferences.test.ts @@ -3,6 +3,7 @@ import { getDevcontainerEnabledStorageKey, getLastUsedModelStorageKey, getLastUsedRepoStorageKey, + getLastUsedSandboxAllocationStorageKey, getLastUsedVariantsStorageKey, getPreferredInitialModel, getPreferredInitialRepo, @@ -78,6 +79,17 @@ describe('getLastUsedVariantsStorageKey', () => { }); }); +describe('getLastUsedSandboxAllocationStorageKey', () => { + it('uses separate keys for personal and organization contexts', () => { + expect(getLastUsedSandboxAllocationStorageKey()).toBe( + 'cloud-agent:last-used-sandbox-allocation:personal' + ); + expect(getLastUsedSandboxAllocationStorageKey('org_123')).toBe( + 'cloud-agent:last-used-sandbox-allocation:organization:org_123' + ); + }); +}); + describe('repository preference', () => { it('uses separate keys for personal and organization contexts', () => { expect(getLastUsedRepoStorageKey()).toBe('cloud-agent:last-used-repo:personal'); diff --git a/apps/web/src/components/cloud-agent-next/model-preferences.ts b/apps/web/src/components/cloud-agent-next/model-preferences.ts index 37514f4ca7..6aa304a9c0 100644 --- a/apps/web/src/components/cloud-agent-next/model-preferences.ts +++ b/apps/web/src/components/cloud-agent-next/model-preferences.ts @@ -6,6 +6,7 @@ const MODEL_STORAGE_KEY_PREFIX = 'cloud-agent:last-used-model'; const VARIANTS_STORAGE_KEY_PREFIX = 'cloud-agent:last-used-variants'; const DEVCONTAINER_ENABLED_STORAGE_KEY = 'cloud-agent:devcontainer-enabled'; const REPO_STORAGE_KEY_PREFIX = 'cloud-agent:last-used-repo'; +const SANDBOX_ALLOCATION_STORAGE_KEY_PREFIX = 'cloud-agent:last-used-sandbox-allocation'; type LastUsedRepo = { fullName: string; @@ -194,6 +195,29 @@ export function setLastUsedVariant( safeLocalStorage.setItem(getLastUsedVariantsStorageKey(organizationId), JSON.stringify(map)); } +export function getLastUsedSandboxAllocationStorageKey(organizationId?: string) { + return organizationId + ? `${SANDBOX_ALLOCATION_STORAGE_KEY_PREFIX}:organization:${organizationId}` + : `${SANDBOX_ALLOCATION_STORAGE_KEY_PREFIX}:personal`; +} + +export function getLastUsedSandboxAllocationKey(organizationId?: string): string | null { + const stored = safeLocalStorage.getItem(getLastUsedSandboxAllocationStorageKey(organizationId)); + return stored && stored.trim().length > 0 ? stored : null; +} + +export function setLastUsedSandboxAllocationKey( + allocationKey: string | undefined, + organizationId?: string +): void { + const storageKey = getLastUsedSandboxAllocationStorageKey(organizationId); + if (allocationKey) { + safeLocalStorage.setItem(storageKey, allocationKey); + } else { + safeLocalStorage.removeItem(storageKey); + } +} + export function getPreferredInitialVariant({ availableVariants, lastUsedVariant, diff --git a/apps/web/src/components/cloud-agent-next/sandbox-selection.test.ts b/apps/web/src/components/cloud-agent-next/sandbox-selection.test.ts new file mode 100644 index 0000000000..c1006a0ad1 --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/sandbox-selection.test.ts @@ -0,0 +1,492 @@ +import { describe, expect, it, jest } from '@jest/globals'; +import { + getSandboxAllocationKey, + getSandboxAllocationRequest, + SELECTABLE_SANDBOX_ALLOCATIONS, + type SandboxDestination, + type SandboxSelectionCapabilities, + type SelectableSandboxAllocationRequest, +} from '@kilocode/worker-utils/sandbox-allocation'; +import { + formatSandboxCapacity, + formatSandboxDestination, + formatSandboxDestinationWithoutTier, + formatSandboxInstance, + getPreferredInitialSandboxAllocation, + getSandboxSelectionGroups, + getSandboxSelectionOptions, + resolveSandboxSelection, + resolveSandboxSelectionSubmissionError, + type SandboxSelectionDraft, +} from './sandbox-selection'; +import { getCloudSessionCreationOperation } from './types'; + +const vercelLarge = getSandboxAllocationRequest('vercel-large'); +const byocLarge = { + provider: { id: 'vercel', account: 'byoc' }, + instanceType: 'large', +} satisfies SelectableSandboxAllocationRequest; +const draft: SandboxSelectionDraft = { + organizationId: 'organization-a', + allocation: vercelLarge, +}; +const capabilities: SandboxSelectionCapabilities = { + enabled: true, + options: [{ allocation: structuredClone(vercelLarge) }], +}; +const input = { + organizationId: 'organization-a', + draft, + capabilities, + devcontainer: false, +}; + +describe('formatSandboxDestination', () => { + it.each([ + [undefined, 'Default'], + [ + getSandboxAllocationRequest('cloudflare-single'), + 'Kilo · Cloudflare · 2 vCPU / 6 GiB · Single', + ], + [ + getSandboxAllocationRequest('cloudflare-shared'), + 'Kilo · Cloudflare · 4 vCPU / 12 GiB · Shared', + ], + [ + getSandboxAllocationRequest('isolated-standard'), + 'Kilo · Cloudflare · 4 vCPU / 12 GiB · Dedicated Standard', + ], + [getSandboxAllocationRequest('vercel-small'), 'Kilo · Vercel · 2 vCPU / 4 GiB'], + [vercelLarge, 'Kilo · Vercel · 4 vCPU / 8 GiB'], + [byocLarge, 'BYOC · Vercel · 4 vCPU / 8 GiB'], + [{ ...byocLarge, instanceType: 'small' }, 'BYOC · Vercel · 2 vCPU / 4 GiB'], + [ + { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'devcontainer' }, + 'Kilo · Cloudflare · 2 vCPU / 6 GiB · Dev container', + ], + [ + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + 'Kilo · Vercel · Provider default', + ], + ] satisfies Array<[SandboxDestination | undefined, string]>)( + 'formats %j as %s', + (destination, label) => { + expect(formatSandboxDestination(destination)).toBe(label); + } + ); +}); + +describe('formatSandboxDestinationWithoutTier', () => { + it.each([ + [undefined, 'Default'], + [getSandboxAllocationRequest('cloudflare-single'), 'Kilo · Cloudflare · 2 vCPU / 6 GiB'], + [getSandboxAllocationRequest('cloudflare-shared'), 'Kilo · Cloudflare · 4 vCPU / 12 GiB'], + [ + { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'devcontainer' }, + 'Kilo · Cloudflare · 2 vCPU / 6 GiB', + ], + [getSandboxAllocationRequest('vercel-small'), 'Kilo · Vercel · 2 vCPU / 4 GiB'], + [byocLarge, 'BYOC · Vercel · 4 vCPU / 8 GiB'], + [ + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + 'Kilo · Vercel · Provider default', + ], + ] satisfies Array<[SandboxDestination | undefined, string]>)( + 'states the account, provider and capacity without the tenancy tier: %j', + (destination, label) => { + expect(formatSandboxDestinationWithoutTier(destination)).toBe(label); + } + ); + + it('keeps every selectable destination distinguishable while closed', () => { + const labels = SELECTABLE_SANDBOX_ALLOCATIONS.map(allocation => + formatSandboxDestinationWithoutTier(getSandboxAllocationRequest(allocation)) + ); + expect(new Set(labels).size).toBe(labels.length); + }); + + it('stays shorter than the expanded label wherever a tenancy tier exists', () => { + const cloudflare = getSandboxAllocationRequest('cloudflare-single'); + expect(formatSandboxDestinationWithoutTier(cloudflare).length).toBeLessThan( + formatSandboxDestination(cloudflare).length + ); + }); +}); + +describe('formatSandboxCapacity', () => { + it.each([ + [getSandboxAllocationRequest('cloudflare-single'), '2 vCPU / 6 GiB'], + [getSandboxAllocationRequest('cloudflare-shared'), '4 vCPU / 12 GiB'], + [getSandboxAllocationRequest('isolated-standard'), '4 vCPU / 12 GiB'], + [getSandboxAllocationRequest('vercel-small'), '2 vCPU / 4 GiB'], + [byocLarge, '4 vCPU / 8 GiB'], + [{ provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, 'Provider default'], + ] satisfies Array<[SandboxDestination, string]>)( + 'reports the vCPU and memory a destination runs with: %j', + (destination, label) => { + expect(formatSandboxCapacity(destination)).toBe(label); + } + ); +}); + +describe('formatSandboxInstance', () => { + it.each([ + [getSandboxAllocationRequest('cloudflare-single'), '2 vCPU / 6 GiB · Single'], + [getSandboxAllocationRequest('cloudflare-shared'), '4 vCPU / 12 GiB · Shared'], + [getSandboxAllocationRequest('vercel-small'), '2 vCPU / 4 GiB'], + [byocLarge, '4 vCPU / 8 GiB'], + [{ provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, 'Provider default'], + ] satisfies Array<[SandboxDestination, string]>)( + 'formats an instance without repeating its account/provider heading: %j', + (destination, label) => { + expect(formatSandboxInstance(destination)).toBe(label); + } + ); + + it('adds the tenancy tier that capacity alone cannot separate', () => { + const shared = getSandboxAllocationRequest('cloudflare-shared'); + const dedicated = getSandboxAllocationRequest('isolated-standard'); + expect(formatSandboxCapacity(shared)).toBe(formatSandboxCapacity(dedicated)); + expect(formatSandboxInstance(shared)).not.toBe(formatSandboxInstance(dedicated)); + }); +}); + +describe('getSandboxSelectionGroups', () => { + it('builds one heading per account and provider without changing instance order', () => { + const options: SandboxSelectionCapabilities['options'] = [ + { allocation: vercelLarge }, + { allocation: byocLarge }, + { allocation: getSandboxAllocationRequest('cloudflare-single') }, + { allocation: getSandboxAllocationRequest('cloudflare-shared') }, + { allocation: { ...byocLarge, instanceType: 'small' } }, + ]; + const original = structuredClone(options); + const groups = getSandboxSelectionGroups( + getSandboxSelectionOptions({ enabled: true, options }) + ); + expect(groups).toEqual([ + { + key: 'byoc:vercel', + label: 'BYOC · Vercel', + options: [options[1], options[4]], + }, + { + key: 'kilo:cloudflare', + label: 'Kilo · Cloudflare', + options: [options[2], options[3]], + }, + { key: 'kilo:vercel', label: 'Kilo · Vercel', options: [options[0]] }, + ]); + expect(groups[0].options[0]).toBe(options[1]); + expect(options).toEqual(original); + }); + + it('does not render empty provider groups', () => { + expect(getSandboxSelectionGroups([])).toEqual([]); + expect(getSandboxSelectionGroups(getSandboxSelectionOptions(capabilities))).toEqual([ + { key: 'kilo:vercel', label: 'Kilo · Vercel', options: capabilities.options }, + ]); + }); +}); + +describe('getSandboxSelectionOptions', () => { + it('places BYOC before Kilo and gathers each provider without reordering instances', () => { + const options: SandboxSelectionCapabilities['options'] = [ + { allocation: vercelLarge }, + { allocation: byocLarge }, + { allocation: getSandboxAllocationRequest('cloudflare-shared') }, + { allocation: { ...byocLarge, instanceType: 'small' } }, + { allocation: getSandboxAllocationRequest('vercel-small') }, + { allocation: getSandboxAllocationRequest('cloudflare-single') }, + ]; + const originalOrder = [...options]; + expect(getSandboxSelectionOptions({ enabled: true, options })).toEqual([ + options[1], + options[3], + options[2], + options[5], + options[0], + options[4], + ]); + expect(options).toEqual(originalOrder); + }); + + it('does not invent BYOC options when the Worker only offers Kilo', () => { + expect(getSandboxSelectionOptions(capabilities)).toEqual(capabilities.options); + expect(getSandboxSelectionOptions(undefined)).toEqual([]); + expect(getSandboxSelectionOptions({ enabled: false, options: [] })).toEqual([]); + }); +}); + +describe('resolveSandboxSelection', () => { + it('forwards the original draft when an equal descriptor is available in its organization', () => { + expect(resolveSandboxSelection(input)).toEqual({ + sandboxAllocation: vercelLarge, + }); + expect(resolveSandboxSelection(input).sandboxAllocation).toBe(draft.allocation); + }); + + it('forwards a personal draft when selection is enabled', () => { + expect( + resolveSandboxSelection({ + organizationId: undefined, + draft: { organizationId: undefined, allocation: vercelLarge }, + capabilities, + devcontainer: false, + }) + ).toEqual({ sandboxAllocation: vercelLarge }); + }); + + it.each(['organization-b', undefined])( + 'omits a previous organization preset in context %s', + organizationId => { + expect(resolveSandboxSelection({ ...input, organizationId })).toEqual({}); + } + ); + + it('uses Default with dev containers without changing the draft', () => { + expect(resolveSandboxSelection({ ...input, devcontainer: true })).toEqual({}); + expect(draft.allocation).toEqual(vercelLarge); + expect(resolveSandboxSelection(input)).toEqual({ + sandboxAllocation: vercelLarge, + }); + }); + + it.each([ + getSandboxAllocationRequest('cloudflare-shared'), + getSandboxAllocationRequest('vercel-large'), + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + { + provider: { id: 'cloudflare', account: 'kilo' }, + instanceType: 'devcontainer', + }, + ] satisfies SandboxDestination[])( + 'does not submit the resolved Default destination: %j', + defaultDestination => { + expect( + resolveSandboxSelection({ + ...input, + draft: { organizationId: input.organizationId }, + capabilities: { ...capabilities, defaultDestination }, + }) + ).toEqual({}); + } + ); + + it('does not substitute a Kilo destination for an unavailable BYOC draft', () => { + expect( + resolveSandboxSelection({ + ...input, + draft: { ...draft, allocation: byocLarge }, + }) + ).toEqual({ + sandboxAllocation: byocLarge, + error: 'This sandbox is unavailable. Choose another sandbox or Default.', + }); + }); + + it('preserves an available BYOC destination without converting its account', () => { + expect( + resolveSandboxSelection({ + ...input, + draft: { ...draft, allocation: byocLarge }, + capabilities: { + enabled: true, + options: [{ allocation: structuredClone(byocLarge) }], + }, + }) + ).toEqual({ sandboxAllocation: byocLarge }); + }); + + it.each([undefined, { enabled: false, options: [] }])( + 'allows Default when capabilities are unavailable: %j', + capabilities => { + expect( + resolveSandboxSelection({ + ...input, + draft: { organizationId: input.organizationId }, + capabilities, + }) + ).toEqual({}); + } + ); + + it.each([undefined, { enabled: false, options: [] }])( + 'blocks an explicit choice without silently switching to Default: %j', + capabilities => { + expect(resolveSandboxSelection({ ...input, capabilities })).toEqual({ + sandboxAllocation: vercelLarge, + error: expect.stringContaining('choose Default'), + }); + } + ); + + it('blocks a preset missing from the allowed options', () => { + expect( + resolveSandboxSelection({ + ...input, + capabilities: { enabled: true, options: [] }, + }) + ).toEqual({ + sandboxAllocation: vercelLarge, + error: 'This sandbox is unavailable. Choose another sandbox or Default.', + }); + }); +}); + +describe('getPreferredInitialSandboxAllocation', () => { + const options: SandboxSelectionCapabilities['options'] = [ + { + allocation: getSandboxAllocationRequest('cloudflare-single'), + }, + { + allocation: getSandboxAllocationRequest('vercel-large'), + }, + { + allocation: getSandboxAllocationRequest('vercel-small'), + }, + ]; + + it('restores the matching destination', () => { + expect( + getPreferredInitialSandboxAllocation({ + options, + lastUsedKey: getSandboxAllocationKey(getSandboxAllocationRequest('vercel-large')), + }) + ).toEqual(vercelLarge); + }); + + it('restores nothing without a stored destination', () => { + expect(getPreferredInitialSandboxAllocation({ options, lastUsedKey: null })).toBeUndefined(); + }); + + it('ignores a stored destination that is no longer offered', () => { + expect( + getPreferredInitialSandboxAllocation({ + options, + lastUsedKey: getSandboxAllocationKey(getSandboxAllocationRequest('cloudflare-shared')), + }) + ).toBeUndefined(); + }); +}); + +describe('resolveSandboxSelectionSubmissionError', () => { + const creation = { + organizationId: draft.organizationId, + sandboxAllocation: getSandboxAllocationKey(vercelLarge), + prompt: 'Build the feature', + model: 'test-model', + repository: 'acme/repo', + githubIntegrationId: 'integration-a', + attachments: { path: 'upload', files: ['notes.md'] }, + }; + const intent = JSON.stringify(creation); + const pendingOperation = { intent, operationKey: 'original-operation' }; + const unavailableCapabilities: Array = [ + undefined, + { enabled: false, options: [] }, + { enabled: true, options: [] }, + ]; + + it.each(unavailableCapabilities)( + 'allows an unchanged pending operation after availability changes: %j', + capabilities => { + const { error } = resolveSandboxSelection({ ...input, capabilities }); + expect(error).toBeDefined(); + expect( + resolveSandboxSelectionSubmissionError({ + error, + intent, + pendingOperation, + }) + ).toBeUndefined(); + const createOperationKey = jest.fn(() => 'new-operation'); + expect(getCloudSessionCreationOperation(pendingOperation, intent, createOperationKey)).toBe( + pendingOperation + ); + expect(createOperationKey).not.toHaveBeenCalled(); + expect( + resolveSandboxSelectionSubmissionError({ + error, + intent, + pendingOperation: null, + }) + ).toBe(error); + } + ); + + it('reuses the pending operation when descriptor keys are reordered', () => { + const allocation = { + instanceType: vercelLarge.instanceType, + provider: { + account: vercelLarge.provider.account, + id: vercelLarge.provider.id, + }, + }; + expect(JSON.stringify(allocation)).not.toBe(JSON.stringify(vercelLarge)); + const { sandboxAllocation, error } = resolveSandboxSelection({ + ...input, + draft: { ...draft, allocation }, + capabilities: undefined, + }); + const retryIntent = JSON.stringify({ + ...creation, + sandboxAllocation: sandboxAllocation ? getSandboxAllocationKey(sandboxAllocation) : undefined, + }); + expect(retryIntent).toBe(intent); + expect(error).toBeDefined(); + expect( + resolveSandboxSelectionSubmissionError({ + error, + intent: retryIntent, + pendingOperation, + }) + ).toBeUndefined(); + const createOperationKey = jest.fn(() => 'new-operation'); + expect( + getCloudSessionCreationOperation(pendingOperation, retryIntent, createOperationKey) + ).toBe(pendingOperation); + expect(createOperationKey).not.toHaveBeenCalled(); + }); + + it.each([ + { prompt: 'A different prompt' }, + { model: 'another-model' }, + { repository: 'acme/other-repo' }, + { githubIntegrationId: 'integration-b' }, + { organizationId: 'organization-b' }, + { + sandboxAllocation: getSandboxAllocationKey(getSandboxAllocationRequest('vercel-small')), + }, + { + sandboxAllocation: getSandboxAllocationKey(getSandboxAllocationRequest('cloudflare-single')), + }, + { sandboxAllocation: getSandboxAllocationKey(byocLarge) }, + { sandboxAllocation: undefined }, + { devcontainer: true }, + { attachments: { path: 'upload', files: ['other-notes.md'] } }, + ])('does not bypass validation for a changed creation intent: %j', change => { + const { error } = resolveSandboxSelection({ + ...input, + capabilities: undefined, + }); + expect( + resolveSandboxSelectionSubmissionError({ + error, + intent: JSON.stringify({ ...creation, ...change }), + pendingOperation, + }) + ).toBe(error); + }); + + it('does not block an available new selection', () => { + const { error } = resolveSandboxSelection(input); + expect( + resolveSandboxSelectionSubmissionError({ + error, + intent, + pendingOperation: null, + }) + ).toBeUndefined(); + }); +}); diff --git a/apps/web/src/components/cloud-agent-next/sandbox-selection.ts b/apps/web/src/components/cloud-agent-next/sandbox-selection.ts new file mode 100644 index 0000000000..f31e24bcbc --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/sandbox-selection.ts @@ -0,0 +1,192 @@ +import { + getSandboxAllocationKey, + type SelectableSandboxAllocationRequest, + type SandboxDestination, + type SandboxSelectionCapabilities, +} from '@kilocode/worker-utils/sandbox-allocation'; +import { + containerCapacityForService, + formatContainerCapacity, +} from '@/lib/cloudflare/container-capacity'; +import type { CloudSessionCreationOperation } from './types'; + +const accountLabels: Record = { + kilo: 'Kilo', + byoc: 'BYOC', +}; + +const providerLabels: Record = { + cloudflare: 'Cloudflare', + vercel: 'Vercel', +}; + +const cloudflareServices: Partial> = { + single: 'cloud-agent-next-sandbox-small', + shared: 'cloud-agent-next-sandbox', + 'isolated-standard': 'cloud-agent-next-sandbox', + devcontainer: 'cloud-agent-next-sandbox-dind', +}; + +const instanceLabels: Record = { + single: 'Single', + shared: 'Shared', + 'isolated-standard': 'Dedicated Standard', + devcontainer: 'Dev container', + small: '2 vCPU / 4 GiB', + large: '4 vCPU / 8 GiB', + default: 'Provider default', +}; + +/** + * Names the account whose cloud runs the sandbox alongside the provider. Both + * halves stay visible because a BYOC account can offer the same provider as Kilo. + */ +export function formatSandboxProvider(provider: SandboxDestination['provider']): string { + return `${accountLabels[provider.account]} · ${providerLabels[provider.id]}`; +} + +/** + * Splits an instance into the capacity every provider can be compared by and the + * tenancy tier only Cloudflare names on top of it. Vercel publishes no container + * class, so its instance name already states the capacity. + */ +function getSandboxInstanceLabels(destination: SandboxDestination): { + capacity: string; + tier?: string; +} { + const name = instanceLabels[destination.instanceType]; + const service = cloudflareServices[destination.instanceType]; + const capacity = service ? containerCapacityForService(service) : null; + return capacity + ? { capacity: formatContainerCapacity(capacity), tier: name } + : { capacity: name }; +} + +export function formatSandboxCapacity(destination: SandboxDestination): string { + return getSandboxInstanceLabels(destination).capacity; +} + +export function formatSandboxInstance(destination: SandboxDestination): string { + const { capacity, tier } = getSandboxInstanceLabels(destination); + return tier ? `${capacity} · ${tier}` : capacity; +} + +export function formatSandboxDestination(destination: SandboxDestination | undefined): string { + if (!destination) return 'Default'; + return `${formatSandboxProvider(destination.provider)} · ${formatSandboxInstance(destination)}`; +} + +/** + * Trigger label. Capacity is the part users weigh when picking a sandbox, so it + * stays; the tenancy tier waits for the open menu. Providers that name no tier — + * every Vercel instance — get the same string as `formatSandboxDestination`. + */ +export function formatSandboxDestinationWithoutTier( + destination: SandboxDestination | undefined +): string { + if (!destination) return 'Default'; + return `${formatSandboxProvider(destination.provider)} · ${formatSandboxCapacity(destination)}`; +} + +type SandboxSelectionGroup = { + key: string; + label: string; + options: SandboxSelectionCapabilities['options']; +}; + +export function getSandboxSelectionGroups( + options: SandboxSelectionCapabilities['options'] +): SandboxSelectionGroup[] { + const groups: SandboxSelectionGroup[] = []; + for (const option of options) { + const { provider } = option.allocation; + const key = `${provider.account}:${provider.id}`; + let group = groups.find(group => group.key === key); + if (!group) { + group = { key, label: formatSandboxProvider(provider), options: [] }; + groups.push(group); + } + group.options.push(option); + } + return groups; +} + +export function getSandboxSelectionOptions( + capabilities: SandboxSelectionCapabilities | undefined +): SandboxSelectionCapabilities['options'] { + const options = capabilities?.enabled ? [...capabilities.options] : []; + + // Display order for the picker: BYOC accounts first, then one block per + // provider, keeping each provider's instances in the order the Worker offers. + return options.sort( + (a, b) => + Number(a.allocation.provider.account !== 'byoc') - + Number(b.allocation.provider.account !== 'byoc') || + a.allocation.provider.id.localeCompare(b.allocation.provider.id) + ); +} + +export function getPreferredInitialSandboxAllocation({ + options, + lastUsedKey, +}: { + options: SandboxSelectionCapabilities['options']; + lastUsedKey: string | null; +}): SelectableSandboxAllocationRequest | undefined { + if (!lastUsedKey) return undefined; + return options.find(option => getSandboxAllocationKey(option.allocation) === lastUsedKey) + ?.allocation; +} + +export function resolveSandboxSelectionSubmissionError({ + error, + intent, + pendingOperation, +}: { + error: string | undefined; + intent: string; + pendingOperation: CloudSessionCreationOperation | null; +}): string | undefined { + return pendingOperation?.intent === intent ? undefined : error; +} + +export type SandboxSelectionDraft = { + organizationId: string | undefined; + allocation?: SelectableSandboxAllocationRequest; +}; + +export function resolveSandboxSelection({ + organizationId, + draft, + capabilities, + devcontainer, +}: { + organizationId: string | undefined; + draft: SandboxSelectionDraft; + capabilities: SandboxSelectionCapabilities | undefined; + devcontainer: boolean; +}): { sandboxAllocation?: SelectableSandboxAllocationRequest; error?: string } { + if (draft.organizationId !== organizationId || devcontainer || !draft.allocation) { + return {}; + } + + if (!capabilities?.enabled) { + return { + sandboxAllocation: draft.allocation, + error: 'Sandbox selection is unavailable. Retry or choose Default to continue.', + }; + } + + const selectedKey = getSandboxAllocationKey(draft.allocation); + const option = capabilities.options.find( + option => getSandboxAllocationKey(option.allocation) === selectedKey + ); + if (!option) { + return { + sandboxAllocation: draft.allocation, + error: 'This sandbox is unavailable. Choose another sandbox or Default.', + }; + } + + return { sandboxAllocation: draft.allocation }; +} diff --git a/apps/web/src/components/cloud-agent-next/sandbox-status.test.ts b/apps/web/src/components/cloud-agent-next/sandbox-status.test.ts index 97e989b1e4..0bd08dd601 100644 --- a/apps/web/src/components/cloud-agent-next/sandbox-status.test.ts +++ b/apps/web/src/components/cloud-agent-next/sandbox-status.test.ts @@ -1,11 +1,14 @@ import { describe, expect, it, jest } from '@jest/globals'; import type { CloudAgentSessionId, KiloSessionId } from '@kilocode/cloud-agent-sdk'; import type { SandboxStatusSnapshot } from '@/routers/cloud-agent-next-schemas'; +import { getSandboxAllocationRequest } from '@kilocode/worker-utils/sandbox-allocation'; import { isSandboxStatusEligible, observeSandboxStatus, sandboxStatusPresentation, + sandboxTypeCapacity, } from './sandbox-status'; +import { formatSandboxCapacity } from './sandbox-selection'; const now = 1_800_000_000_000; const snapshot: SandboxStatusSnapshot = { @@ -650,3 +653,76 @@ describe('sandbox status presentation', () => { expect(JSON.stringify(view)).not.toContain('PRIVATE_SENTINEL'); }); }); + +describe('sandboxTypeCapacity', () => { + it.each([ + ['shared', '4 vCPU / 12 GiB'], + ['isolated-standard', '4 vCPU / 12 GiB'], + ['isolated-small', '2 vCPU / 6 GiB'], + ['code-review', '1 vCPU / 4 GiB'], + ['devcontainer', '2 vCPU / 6 GiB'], + ] as const)('reports the container capacity behind %s', (sandboxType, capacity) => { + expect(sandboxTypeCapacity(sandboxType)).toBe(capacity); + }); + + it.each([undefined, null, 'unknown'] as const)( + 'reports no capacity for an unidentified sandbox: %j', + sandboxType => { + expect(sandboxTypeCapacity(sandboxType)).toBeNull(); + } + ); + + it('agrees with the destination picker wherever both name the same container', () => { + expect(sandboxTypeCapacity('isolated-standard')).toBe( + formatSandboxCapacity(getSandboxAllocationRequest('cloudflare-shared')) + ); + expect(sandboxTypeCapacity('isolated-small')).toBe( + formatSandboxCapacity(getSandboxAllocationRequest('cloudflare-single')) + ); + expect(sandboxTypeCapacity('devcontainer')).toBe( + formatSandboxCapacity({ + provider: { id: 'cloudflare', account: 'kilo' }, + instanceType: 'devcontainer', + }) + ); + }); +}); + +describe('sandbox status capacity', () => { + it('reports the capacity of the observed sandbox, not the requested one', () => { + expect( + sandboxStatusPresentation({ + ...observation, + data: { ...snapshot, runtime: { ...runtime, sandboxType: 'isolated-small' } }, + }).capacity + ).toBe('2 vCPU / 6 GiB'); + }); + + it('keeps the capacity of a sleeping sandbox, which still names its container', () => { + expect( + sandboxStatusPresentation({ + ...observation, + data: { + ...snapshot, + status: 'sleeping', + detailCode: 'sandbox_stopped', + estimatedSleepAt: null, + runtime: { ...runtime, startedAt: null, stoppedAt: now - 60_000 }, + }, + }).capacity + ).toBe('4 vCPU / 12 GiB'); + }); + + it.each(['paused', 'unavailable', 'checking'] as const)( + 'reports no capacity without a usable observation: %s', + observationState => { + expect( + sandboxStatusPresentation({ ...observation, observation: observationState }).capacity + ).toBeNull(); + } + ); + + it('reports no capacity when the snapshot carries no runtime', () => { + expect(sandboxStatusPresentation(observation).capacity).toBeNull(); + }); +}); diff --git a/apps/web/src/components/cloud-agent-next/sandbox-status.ts b/apps/web/src/components/cloud-agent-next/sandbox-status.ts index 18a5d4b163..06a301555e 100644 --- a/apps/web/src/components/cloud-agent-next/sandbox-status.ts +++ b/apps/web/src/components/cloud-agent-next/sandbox-status.ts @@ -4,7 +4,12 @@ import { SANDBOX_STATUS_DETAIL_MESSAGES, type SandboxLifecycleStatus, type SandboxProviderLabel, + type SandboxStatusSnapshot, } from '@/routers/cloud-agent-next-schemas'; +import { + containerCapacityForService, + formatContainerCapacity, +} from '@/lib/cloudflare/container-capacity'; import type { FetchedSessionData, ResolvedSession } from '@kilocode/cloud-agent-sdk'; export const SANDBOX_STATUS_POLL_INTERVAL_MS = 5_000; @@ -60,6 +65,8 @@ const statusLabels = { unknown: 'Unknown', } satisfies Record; +type SandboxType = NonNullable['sandboxType']>; + const sandboxTypes = { shared: 'Shared', 'isolated-small': 'Small', @@ -67,14 +74,40 @@ const sandboxTypes = { 'code-review': 'Code review', devcontainer: 'Custom environment', unknown: 'Unknown', +} satisfies Record; + +/** + * The container class each observed sandbox type runs on, mirroring + * `expectedSandboxClassName` in the Worker. Containment variants share their base + * class's capacity, so the classification alone fixes the vCPU and memory. + * + * `sandbox-selection.ts` keys a similar map off a requested destination instead; + * the two stay separate because a request and an observed sandbox can disagree + * after a fallback. `sandbox-status.test.ts` locks the values they share. + */ +const sandboxTypeServices: Record = { + shared: 'cloud-agent-next-sandbox', + 'isolated-small': 'cloud-agent-next-sandbox-small', + 'isolated-standard': 'cloud-agent-next-sandbox', + 'code-review': 'cloud-agent-next-sandbox-code-review', + devcontainer: 'cloud-agent-next-sandbox-dind', + unknown: null, }; +/** The vCPU and memory an observed sandbox type runs with, or null when unknown. */ +export function sandboxTypeCapacity(sandboxType: SandboxType | null | undefined): string | null { + const service = sandboxTypeServices[sandboxType ?? 'unknown']; + const capacity = service ? containerCapacityForService(service) : null; + return capacity ? formatContainerCapacity(capacity) : null; +} + export type SandboxStatusPresentation = { status: SandboxLifecycleStatus | 'sleeping-soon'; label: string; detail: string; provider: SandboxProviderLabel; sandboxType: string; + capacity: string | null; kiloCliVersion: string | null; wrapperVersion: string | null; startedAt: number | null; @@ -109,6 +142,7 @@ export function sandboxStatusPresentation({ detail: SANDBOX_STATUS_DETAIL_MESSAGES.status_unavailable, provider: 'Unknown', sandboxType: 'Unknown', + capacity: null, kiloCliVersion: null, wrapperVersion: null, startedAt: null, @@ -193,6 +227,7 @@ export function sandboxStatusPresentation({ : SANDBOX_STATUS_DETAIL_MESSAGES[snapshot.detailCode], provider: snapshot.provider, sandboxType: sandboxTypes[runtime?.sandboxType ?? 'unknown'], + capacity: sandboxTypeCapacity(runtime?.sandboxType), kiloCliVersion: runtime?.kiloCliVersion ?? null, wrapperVersion: runtime?.wrapperVersion ?? null, startedAt: runtime?.startedAt ?? null, diff --git a/apps/web/src/components/webhook-triggers/TriggerForm.tsx b/apps/web/src/components/webhook-triggers/TriggerForm.tsx index c2fcdf6efd..725619cf1a 100644 --- a/apps/web/src/components/webhook-triggers/TriggerForm.tsx +++ b/apps/web/src/components/webhook-triggers/TriggerForm.tsx @@ -689,8 +689,8 @@ export function TriggerForm({

{!canSetSandboxAllocation && hasSavedSandboxAllocation && (

- Only Kilo admins can newly enable Dedicated Standard. You can keep this - allocation or clear it. + Dedicated Standard requires an enabled organization. Choose Automatic to use + default routing.

)}
diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts index d4ea8e00ef..3a10002805 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts @@ -1,6 +1,11 @@ import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals'; import type * as TrpcClientModule from '@trpc/client'; import type * as CloudAgentClientModule from './cloud-agent-client'; +import { + getSandboxAllocationRequest, + type SandboxAllocationInput, + type SandboxSelectionCapabilities, +} from '@kilocode/worker-utils/sandbox-allocation'; import type { ComputeBillingStatus, CreateWorktreeChatInput, @@ -412,6 +417,143 @@ describe('createAppBuilderCloudAgentNextClient', () => { }); }); +describe('CloudAgentNextClient sandbox selection', () => { + const kilocodeOrganizationId = '9a283301-b75d-4375-a1ba-e319a02e18b7'; + + it('forwards personal capability discovery without an organization', async () => { + const query = jest + .fn() + .mockResolvedValue({ enabled: true, options: [] }); + mockCreateTRPCClient.mockReturnValueOnce({ getSandboxSelectionOptions: { query } }); + + await expect( + new CloudAgentNextClient('auth-token').getSandboxSelectionOptions({}) + ).resolves.toEqual({ enabled: true, options: [] }); + expect(query).toHaveBeenCalledWith({}); + }); + + it.each([undefined, false, true])( + 'forwards optional devcontainer context: %j', + async devcontainer => { + const input = { + kilocodeOrganizationId, + ...(devcontainer !== undefined ? { devcontainer } : {}), + }; + const capabilities: SandboxSelectionCapabilities = { + enabled: true, + defaultDestination: { + provider: { id: 'vercel', account: 'kilo' }, + instanceType: 'default', + }, + options: [ + { allocation: getSandboxAllocationRequest('vercel-large') }, + { + allocation: { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + }, + ], + }; + const query = jest + .fn() + .mockResolvedValue(capabilities); + mockCreateTRPCClient.mockReturnValueOnce({ getSandboxSelectionOptions: { query } }); + + await expect( + new CloudAgentNextClient('auth-token').getSandboxSelectionOptions(input) + ).resolves.toEqual(capabilities); + expect(query).toHaveBeenCalledWith(input); + } + ); + + it('normalizes older capabilities without inventing a default destination', async () => { + mockCreateTRPCClient.mockReturnValueOnce({ + getSandboxSelectionOptions: { + query: jest.fn(async () => ({ + enabled: true, + options: [{ allocation: 'cloudflare-single', available: true }], + })), + }, + }); + await expect( + new CloudAgentNextClient('auth-token').getSandboxSelectionOptions({ kilocodeOrganizationId }) + ).resolves.toEqual({ + enabled: true, + options: [{ allocation: getSandboxAllocationRequest('cloudflare-single') }], + }); + }); + + it('does not turn capability failure into an available default', async () => { + const error = new Error('Worker unavailable'); + mockCreateTRPCClient.mockReturnValueOnce({ + getSandboxSelectionOptions: { + query: jest.fn(async () => { + throw error; + }), + }, + }); + await expect( + new CloudAgentNextClient('auth-token').getSandboxSelectionOptions({ kilocodeOrganizationId }) + ).rejects.toBe(error); + }); + + it('rejects invalid capability descriptors', async () => { + mockCreateTRPCClient.mockReturnValueOnce({ + getSandboxSelectionOptions: { + query: jest.fn(async () => ({ + enabled: true, + options: [ + { + allocation: { + provider: { id: 'cloudflare', account: 'byoc' }, + instanceType: 'single', + }, + }, + ], + })), + }, + }); + await expect( + new CloudAgentNextClient('auth-token').getSandboxSelectionOptions({ kilocodeOrganizationId }) + ).rejects.toThrow(); + }); + + it.each([ + undefined, + 'isolated-standard', + 'vercel-small', + getSandboxAllocationRequest('cloudflare-shared'), + getSandboxAllocationRequest('vercel-large'), + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + ] satisfies Array)( + 'forwards the prepare wire allocation unchanged: %j', + async sandboxAllocation => { + const input: PrepareSessionInput = { + kilocodeOrganizationId, + githubRepo: 'acme/repo', + prompt: 'Build the feature', + mode: 'code', + model: 'kilo/test-model', + operationKey: '12345678-1234-4234-9234-123456789abc', + autoInitiate: true, + ...(sandboxAllocation ? { sandboxAllocation } : {}), + }; + const output = { + kiloSessionId: 'ses_12345678901234567890123456', + cloudAgentSessionId: 'agent_123', + replayed: true, + }; + const mutate = jest + .fn() + .mockResolvedValue(output); + mockCreateTRPCClient.mockReturnValueOnce({ prepareSession: { mutate } }); + + await expect(new CloudAgentNextClient('auth-token').prepareSession(input)).resolves.toEqual( + output + ); + expect(mutate).toHaveBeenCalledWith(input); + } + ); +}); + describe('CloudAgentNextClient sensitive error reporting', () => { const error = new Error('Worker unavailable'); diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts index 7f24fc06f6..1d2167d228 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts @@ -3,6 +3,11 @@ import { createTRPCClient, httpLink, TRPCClientError } from '@trpc/client'; import { TRPCError } from '@trpc/server'; import * as z from 'zod'; import type { AgentConfig } from '@kilocode/db/schema-types'; +import { + sandboxSelectionCapabilitiesSchema, + type SandboxAllocationInput, + type SandboxSelectionCapabilities, +} from '@kilocode/worker-utils/sandbox-allocation'; import type { EncryptedEnvelope } from '@/lib/encryption'; import type { CloudAgentAttachments } from '@/lib/cloud-agent/constants'; import type { Images } from '@/lib/images-schema'; @@ -169,6 +174,7 @@ type PrepareSessionSharedFields = { gateThreshold?: 'off' | 'all' | 'warning' | 'critical'; /** When true, route the session to a Docker-in-Docker sandbox that supports devcontainer runtimes */ devcontainer?: boolean; + sandboxAllocation?: SandboxAllocationInput; }; /** Non-clone prepare input: required `prompt` and the current optional initial fields. */ @@ -193,6 +199,11 @@ export type PrepareSessionCloneInput = PrepareSessionSharedFields & { /** Input for prepareSession procedure */ export type PrepareSessionInput = PrepareSessionNonCloneInput | PrepareSessionCloneInput; +export type GetSandboxSelectionOptionsInput = { + kilocodeOrganizationId?: string; + devcontainer?: boolean; +}; + /** Output from prepareSession procedure */ export type PrepareSessionOutput = { /** The Kilo CLI session ID */ @@ -608,6 +619,9 @@ type CloudAgentNextTRPCClient = { getComputeBillingStatus: { query: (input: GetSessionInput) => Promise; }; + getSandboxSelectionOptions: { + query: (input: GetSandboxSelectionOptionsInput) => Promise; + }; prepareSession: { mutate: (input: PrepareSessionInput) => Promise; }; @@ -920,6 +934,14 @@ export class CloudAgentNextClient { } } + async getSandboxSelectionOptions( + input: GetSandboxSelectionOptionsInput + ): Promise { + return sandboxSelectionCapabilitiesSchema.parse( + await this.client.getSandboxSelectionOptions.query(input) + ); + } + /** * Prepare a new cloud agent session. */ diff --git a/apps/web/src/lib/cloudflare/container-capacity.ts b/apps/web/src/lib/cloudflare/container-capacity.ts index 3e5027ee57..7b655c355b 100644 --- a/apps/web/src/lib/cloudflare/container-capacity.ts +++ b/apps/web/src/lib/cloudflare/container-capacity.ts @@ -34,6 +34,10 @@ export function containerCapacityForService(service: string): ContainerCapacity } } +export function formatContainerCapacity(capacity: ContainerCapacity): string { + return `${capacity.vcpu} vCPU / ${capacity.memoryBytes / 1024 ** 3} GiB`; +} + export function sharedContainerCapacity(services: Set): ContainerCapacity | null { let shared: ContainerCapacity | null = null; for (const service of services) { diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index 83fdcf1c6d..2734b895be 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -17,6 +17,15 @@ import type { personalPrepareSessionNextSchema, SandboxStatusSnapshot, } from '@/routers/cloud-agent-next-schemas'; +import { + getSandboxAllocationRequest, + SELECTABLE_SANDBOX_ALLOCATIONS, + type SandboxAllocationInput, + type SandboxSelectionCapabilities, + type SelectableSandboxAllocation, + type SelectableSandboxAllocationRequest, +} from '@kilocode/worker-utils/sandbox-allocation'; +import type { GetSandboxSelectionOptionsInput } from '@/lib/cloud-agent-next/cloud-agent-client'; import { TRPCError } from '@trpc/server'; import type { verifyUserOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; @@ -26,6 +35,7 @@ const mockPrepareSession = jest.fn< (input: { githubRepo?: string; devcontainer?: boolean; + sandboxAllocation?: SandboxAllocationInput; attachments?: AttachmentReference; }) => Promise<{ cloudAgentSessionId: string; @@ -84,7 +94,11 @@ const mockGetWorktreeFile = (input: WorktreeFileQuery & { cloudAgentSessionId: string }) => Promise >(); +const mockGetSandboxSelectionOptions = + jest.fn<(input: GetSandboxSelectionOptionsInput) => Promise>(); + const mockCreateCloudAgentNextClient = jest.fn((_authToken: string) => ({ + getSandboxSelectionOptions: mockGetSandboxSelectionOptions, prepareSession: mockPrepareSession, sendMessage: mockSendMessage, getSession: mockGetSession, @@ -203,7 +217,10 @@ jest.mock('@/lib/cloud-agent/session-ownership', () => ({ })); let createCaller: (ctx: { user: User; headersList?: Headers }) => { - prepareSession: (input: z.infer) => Promise<{ + getSandboxSelectionOptions: (input: { + devcontainer?: boolean; + }) => Promise; + prepareSession: (input: z.input) => Promise<{ cloudAgentSessionId: string; kiloSessionId: string; }>; @@ -854,6 +871,50 @@ describe('cloudAgentNextRouter helper procedures', () => { }); }); +describe('cloudAgentNextRouter.getSandboxSelectionOptions', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('forwards personal capability discovery with the caller token', async () => { + const capabilities: SandboxSelectionCapabilities = { + enabled: true, + defaultDestination: getSandboxAllocationRequest('cloudflare-single'), + options: [{ allocation: getSandboxAllocationRequest('cloudflare-single') }], + }; + mockGetSandboxSelectionOptions.mockResolvedValueOnce(capabilities); + const caller = createCaller({ user: { id: 'oauth/user', is_admin: false } as User }); + + await expect(caller.getSandboxSelectionOptions({})).resolves.toEqual(capabilities); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('cloud-agent-token'); + expect(mockGetSandboxSelectionOptions).toHaveBeenCalledWith({}); + }); + + it.each([false, true])( + 'forwards explicit personal devcontainer context %s', + async devcontainer => { + mockGetSandboxSelectionOptions.mockResolvedValueOnce({ enabled: false, options: [] }); + const caller = createCaller({ user: { id: 'oauth/user', is_admin: false } as User }); + + await expect(caller.getSandboxSelectionOptions({ devcontainer })).resolves.toEqual({ + enabled: false, + options: [], + }); + expect(mockGetSandboxSelectionOptions).toHaveBeenCalledWith({ devcontainer }); + } + ); + + it('propagates Worker-disabled selection without granting a capability', async () => { + mockGetSandboxSelectionOptions.mockResolvedValueOnce({ enabled: false, options: [] }); + const caller = createCaller({ user: { id: 'admin-user', is_admin: true } as User }); + + await expect(caller.getSandboxSelectionOptions({})).resolves.toEqual({ + enabled: false, + options: [], + }); + }); +}); + describe('cloudAgentNextRouter.prepareSession', () => { beforeEach(() => { jest.clearAllMocks(); @@ -981,6 +1042,81 @@ describe('cloudAgentNextRouter.prepareSession', () => { expect(mockPrepareSession).not.toHaveBeenCalled(); }); + const sandboxInput = { + prompt: 'Test prompt', + mode: 'code', + model: 'kilo/test-model', + githubRepo: 'acme/repo', + }; + + it.each(SELECTABLE_SANDBOX_ALLOCATIONS)( + 'forwards the normalized legacy personal sandbox preset %s', + async sandboxAllocation => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + + await caller.prepareSession({ ...sandboxInput, sandboxAllocation }); + + expect(mockPrepareSession).toHaveBeenCalledWith( + expect.objectContaining({ + sandboxAllocation: getSandboxAllocationRequest(sandboxAllocation), + }) + ); + } + ); + + it.each([ + ...SELECTABLE_SANDBOX_ALLOCATIONS.map(allocation => getSandboxAllocationRequest(allocation)), + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'large' }, + ] satisfies SelectableSandboxAllocationRequest[])( + 'forwards a structured personal sandbox destination without changing its account: %j', + async sandboxAllocation => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + + await caller.prepareSession({ ...sandboxInput, sandboxAllocation }); + + expect(mockPrepareSession).toHaveBeenCalledWith( + expect.objectContaining({ sandboxAllocation }) + ); + expect(mockGetSandboxSelectionOptions).not.toHaveBeenCalled(); + } + ); + + it('keeps Default omitted on personal prepares', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + + await caller.prepareSession(sandboxInput); + + expect(mockPrepareSession).toHaveBeenCalledTimes(1); + expect(mockPrepareSession.mock.calls[0][0]).not.toHaveProperty('sandboxAllocation'); + }); + + it('rejects an invalid personal preset before calling the Worker', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + + await expect( + caller.prepareSession({ + ...sandboxInput, + sandboxAllocation: 'vercel-medium' as SelectableSandboxAllocation, + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockPrepareSession).not.toHaveBeenCalled(); + }); + + it('rejects explicit personal presets with dev containers before calling the Worker', async () => { + const caller = createCaller({ user: { id: 'user-1', is_admin: false } as User }); + + await expect( + caller.prepareSession({ + ...sandboxInput, + sandboxAllocation: 'cloudflare-single', + devcontainer: true, + }) + ).rejects.toThrow('Sandbox selection is not available with dev containers'); + expect(mockPrepareSession).not.toHaveBeenCalled(); + expect(mockIsFeatureFlagEnabledOrDevelopment).not.toHaveBeenCalled(); + }); + it('forwards devcontainer sessions when the feature flag is enabled', async () => { mockIsFeatureFlagEnabledOrDevelopment.mockResolvedValue(true); const caller = createCaller({ diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index bcc31460a3..a343511c5c 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -1,5 +1,6 @@ import 'server-only'; import { baseProcedure, createTRPCRouter } from '@/lib/trpc/init'; +import { sandboxSelectionCapabilitiesSchema } from '@kilocode/worker-utils/sandbox-allocation'; import { createCloudAgentNextClient, createCloudAgentNextClientForModel, @@ -145,6 +146,16 @@ async function createCloudAgentControlToken(user: User, headersList?: Headers): * separately via WebSocket connection. */ export const cloudAgentNextRouter = createTRPCRouter({ + getSandboxSelectionOptions: baseProcedure + .input(z.object({ devcontainer: z.boolean().optional() })) + .output(sandboxSelectionCapabilitiesSchema) + .query(async ({ ctx, input }) => { + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); + return await createCloudAgentNextClient(authToken).getSandboxSelectionOptions({ + ...(input.devcontainer !== undefined ? { devcontainer: input.devcontainer } : {}), + }); + }), + /** * Prepare a new cloud agent session. * diff --git a/apps/web/src/routers/cloud-agent-next-schemas.test.ts b/apps/web/src/routers/cloud-agent-next-schemas.test.ts index 5a0e53f6ca..87bc2d1f78 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.test.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.test.ts @@ -1,10 +1,17 @@ import { describe, expect, it } from '@jest/globals'; +import { + getSandboxAllocationRequest, + SELECTABLE_SANDBOX_ALLOCATIONS, + type SelectableSandboxAllocationRequest, +} from '@kilocode/worker-utils/sandbox-allocation'; import { baseCreateWorktreeChatNextOutputSchema, baseCreateWorktreeChatNextSchema, baseGetSandboxStatusNextOutputSchema, baseGetSandboxStatusNextSchema, basePrepareSessionNextSchema, + organizationPrepareSessionNextSchema, + personalPrepareSessionNextSchema, baseCancelQueuedMessageNextSchema, SANDBOX_STATUS_DETAIL_MESSAGES, type SandboxStatusSnapshot, @@ -571,6 +578,21 @@ describe('createWorktreeChat schemas', () => { } }); + it.each([ + ...SELECTABLE_SANDBOX_ALLOCATIONS, + ...SELECTABLE_SANDBOX_ALLOCATIONS.map(allocation => getSandboxAllocationRequest(allocation)), + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'large' }, + ])('rejects attempts to change an existing worktree to %j', sandboxAllocation => { + expect( + baseCreateWorktreeChatNextSchema.safeParse({ + sourceKiloSessionId: KILO_SESSION_ID, + operationKey, + sandboxAllocation, + }).success + ).toBe(false); + }); + it('requires canonical workspace/worktree output and rejects private runtime paths', () => { const output = { kiloSessionId: KILO_SESSION_ID, @@ -591,6 +613,158 @@ describe('createWorktreeChat schemas', () => { }); }); +describe('prepare session sandbox selection', () => { + const baseInput = { + githubRepo: 'acme/repo', + prompt: 'Test prompt', + mode: 'code', + model: 'kilo/test-model', + }; + const organizationInput = { ...baseInput, organizationId: MESSAGE_UUID }; + const allocations = SELECTABLE_SANDBOX_ALLOCATIONS; + const requests: SelectableSandboxAllocationRequest[] = [ + ...allocations.map(allocation => getSandboxAllocationRequest(allocation)), + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'large' }, + ]; + + it.each(['isolated-standard', getSandboxAllocationRequest('isolated-standard')])( + 'rejects isolated-standard on the organization schema: %j', + sandboxAllocation => { + expect( + organizationPrepareSessionNextSchema.safeParse({ ...organizationInput, sandboxAllocation }) + .success + ).toBe(false); + } + ); + + it.each(allocations)('normalizes the legacy organization allocation %s', sandboxAllocation => { + expect( + organizationPrepareSessionNextSchema.parse({ ...organizationInput, sandboxAllocation }) + .sandboxAllocation + ).toEqual(getSandboxAllocationRequest(sandboxAllocation)); + }); + + it.each(allocations)('normalizes the legacy personal allocation %s', sandboxAllocation => { + expect( + personalPrepareSessionNextSchema.parse({ ...baseInput, sandboxAllocation }).sandboxAllocation + ).toEqual(getSandboxAllocationRequest(sandboxAllocation)); + }); + + it.each(requests)('preserves the structured organization allocation %j', sandboxAllocation => { + expect( + organizationPrepareSessionNextSchema.parse({ ...organizationInput, sandboxAllocation }) + .sandboxAllocation + ).toEqual(sandboxAllocation); + }); + + it.each(requests)('preserves the structured personal allocation %j', sandboxAllocation => { + expect( + personalPrepareSessionNextSchema.parse({ ...baseInput, sandboxAllocation }).sandboxAllocation + ).toEqual(sandboxAllocation); + }); + + it.each([ + 'default', + 'vercel-medium', + '', + null, + 2, + { vcpus: 4 }, + { provider: { id: 'cloudflare', account: 'byoc' }, instanceType: 'single' }, + { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'single' }, + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'devcontainer' }, + ])('rejects an invalid sandbox allocation: %j', sandboxAllocation => { + expect( + organizationPrepareSessionNextSchema.safeParse({ ...organizationInput, sandboxAllocation }) + .success + ).toBe(false); + expect( + personalPrepareSessionNextSchema.safeParse({ ...baseInput, sandboxAllocation }).success + ).toBe(false); + }); + + it('keeps Default omitted in personal and organization requests', () => { + expect(personalPrepareSessionNextSchema.parse(baseInput)).not.toHaveProperty( + 'sandboxAllocation' + ); + expect(organizationPrepareSessionNextSchema.parse(organizationInput)).not.toHaveProperty( + 'sandboxAllocation' + ); + }); + + it.each([...allocations, ...requests])( + 'rejects dev containers combined with %j', + sandboxAllocation => { + for (const result of [ + organizationPrepareSessionNextSchema.safeParse({ + ...organizationInput, + devcontainer: true, + sandboxAllocation, + }), + personalPrepareSessionNextSchema.safeParse({ + ...baseInput, + devcontainer: true, + sandboxAllocation, + }), + ]) { + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error.issues).toEqual( + expect.arrayContaining([expect.objectContaining({ path: ['sandboxAllocation'] })]) + ); + } + } + } + ); + + it('accepts dev containers with Default', () => { + expect( + organizationPrepareSessionNextSchema.safeParse({ ...organizationInput, devcontainer: true }) + .success + ).toBe(true); + expect( + personalPrepareSessionNextSchema.safeParse({ ...baseInput, devcontainer: true }).success + ).toBe(true); + }); + + it('preserves clone-only validation while accepting a preset', () => { + const personalClone = { + ...baseInput, + prompt: undefined, + cloneFromKiloSessionId: KILO_SESSION_ID, + autoInitiate: true, + operationKey: MESSAGE_UUID, + sandboxAllocation: 'cloudflare-single', + }; + expect(personalPrepareSessionNextSchema.parse(personalClone).sandboxAllocation).toEqual( + getSandboxAllocationRequest('cloudflare-single') + ); + expect( + personalPrepareSessionNextSchema.safeParse({ ...personalClone, prompt: 'Not allowed' }) + .success + ).toBe(false); + + const cloneInput = { + ...organizationInput, + prompt: undefined, + cloneFromKiloSessionId: KILO_SESSION_ID, + autoInitiate: true, + operationKey: MESSAGE_UUID, + sandboxAllocation: 'cloudflare-single', + }; + expect(organizationPrepareSessionNextSchema.parse(cloneInput).sandboxAllocation).toEqual( + getSandboxAllocationRequest('cloudflare-single') + ); + expect( + organizationPrepareSessionNextSchema.safeParse({ ...cloneInput, prompt: 'Not allowed' }) + .success + ).toBe(false); + }); +}); + describe('baseCancelQueuedMessageNextSchema', () => { const VALID_MESSAGE_ID = 'msg_123456789abc123456789ABCDE'; diff --git a/apps/web/src/routers/cloud-agent-next-schemas.ts b/apps/web/src/routers/cloud-agent-next-schemas.ts index 6d8067aada..25d1e66501 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.ts @@ -1,5 +1,6 @@ import * as z from 'zod'; import { SandboxStatusSessionIdSchema } from '../../../../services/cloud-agent-next/src/shared/sandbox-status'; +import { selectableSandboxAllocationInputSchema } from '@kilocode/worker-utils/sandbox-allocation'; import { cloudAgentWorktreeIdSchema, sessionIdSchema as kiloSessionIdSchema, @@ -478,13 +479,32 @@ export const basePrepareSessionNextSchema = z path: ['attachments'], }); -export const personalPrepareSessionNextSchema = basePrepareSessionNextSchema.refine( - data => data.bitbucketRepo === undefined, - { +export const personalPrepareSessionNextSchema = basePrepareSessionNextSchema + .and( + z.object({ + sandboxAllocation: selectableSandboxAllocationInputSchema.optional(), + }) + ) + .refine(data => data.bitbucketRepo === undefined, { message: 'Bitbucket repositories require an organization', path: ['bitbucketRepo'], - } -); + }) + .refine(data => !data.devcontainer || data.sandboxAllocation === undefined, { + message: 'Sandbox selection is not available with dev containers. Choose Default.', + path: ['sandboxAllocation'], + }); + +export const organizationPrepareSessionNextSchema = basePrepareSessionNextSchema + .and( + z.object({ + organizationId: z.uuid(), + sandboxAllocation: selectableSandboxAllocationInputSchema.optional(), + }) + ) + .refine(data => !data.devcontainer || data.sandboxAllocation === undefined, { + message: 'Sandbox selection is not available with dev containers. Choose Default.', + path: ['sandboxAllocation'], + }); // Output schema for prepareSession export const basePrepareSessionNextOutputSchema = z.object({ diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index 4a68159a01..2585343d52 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -30,9 +30,18 @@ import type { BitbucketOrganizationRepositoryListResult } from '@/lib/cloud-agen import { TRPCError } from '@trpc/server'; import type { verifyOrgOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; import type { - basePrepareSessionNextSchema, + organizationPrepareSessionNextSchema, SandboxStatusSnapshot, } from '@/routers/cloud-agent-next-schemas'; +import { + getSandboxAllocationRequest, + SELECTABLE_SANDBOX_ALLOCATIONS, + type SandboxAllocationInput, + type SandboxSelectionCapabilities, + type SelectableSandboxAllocation, + type SelectableSandboxAllocationRequest, +} from '@kilocode/worker-utils/sandbox-allocation'; +import type { GetSandboxSelectionOptionsInput } from '@/lib/cloud-agent-next/cloud-agent-client'; const ORGANIZATION_ID = '9a283301-b75d-4375-a1ba-e319a02e18b7'; @@ -46,6 +55,7 @@ const mockPrepareSession = jest.fn< bitbucketWorkspaceUuid?: string; bitbucketRepositoryUuid?: string; devcontainer?: boolean; + sandboxAllocation?: SandboxAllocationInput; kilocodeOrganizationId?: string; attachments?: AttachmentReference; }) => Promise<{ @@ -97,7 +107,11 @@ const mockGetWorktreeFile = (input: WorktreeFileQuery & { cloudAgentSessionId: string }) => Promise >(); +const mockGetSandboxSelectionOptions = + jest.fn<(input: GetSandboxSelectionOptionsInput) => Promise>(); + const mockCreateCloudAgentNextClient = jest.fn((_authToken: string) => ({ + getSandboxSelectionOptions: mockGetSandboxSelectionOptions, prepareSession: mockPrepareSession, sendMessage: mockSendMessage, getSession: mockGetSession, @@ -235,6 +249,11 @@ jest.mock('@/lib/r2/cloud-agent-attachments', () => ({ generateCloudAgentAttachmentUploadUrl: mockGenerateCloudAgentAttachmentUploadUrl, })); +jest.mock('@/lib/r2/cloud-agent-pending-uploads', () => ({ + linkPendingUploads: jest.fn(), + releasePendingUploads: jest.fn(), +})); + jest.mock('@/routers/organizations/utils', () => { const trpcInit = jest.requireActual('@/lib/trpc/init'); const zod = jest.requireActual('zod'); @@ -256,9 +275,11 @@ jest.mock('@/routers/organizations/utils', () => { }); let createCaller: (ctx: { user: User; headersList?: Headers }) => { - prepareSession: ( - input: z.infer & { organizationId: string } - ) => Promise<{ + getSandboxSelectionOptions: (input: { + organizationId: string; + devcontainer?: boolean; + }) => Promise; + prepareSession: (input: z.input) => Promise<{ cloudAgentSessionId: string; kiloSessionId: string; }>; @@ -1192,6 +1213,112 @@ describe('organizationCloudAgentNextRouter terminal ownership', () => { }); }); +describe('organizationCloudAgentNextRouter.getSandboxSelectionOptions', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('authorizes membership and forwards the organization with the caller token', async () => { + const capabilities: SandboxSelectionCapabilities = { + enabled: true, + defaultDestination: { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + options: [ + { allocation: getSandboxAllocationRequest('cloudflare-single') }, + { + allocation: getSandboxAllocationRequest('vercel-large'), + }, + ], + }; + mockGetSandboxSelectionOptions.mockResolvedValueOnce(capabilities); + const caller = createCaller({ user: { id: 'oauth/member', is_admin: false } as User }); + + await expect( + caller.getSandboxSelectionOptions({ organizationId: ORGANIZATION_ID }) + ).resolves.toEqual(capabilities); + expect(mockEnsureOrganizationAccess).toHaveBeenCalledWith( + expect.objectContaining({ user: { id: 'oauth/member', is_admin: false } }), + ORGANIZATION_ID + ); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('cloud-agent-token'); + expect(mockGetSandboxSelectionOptions).toHaveBeenCalledWith({ + kilocodeOrganizationId: ORGANIZATION_ID, + }); + expect(mockIsFeatureFlagEnabledOrDevelopment).not.toHaveBeenCalled(); + expect(mockCreateCloudAgentNextClientForModel).not.toHaveBeenCalled(); + }); + + it.each([false, true])('forwards explicit devcontainer context %s', async devcontainer => { + const capabilities: SandboxSelectionCapabilities = { + enabled: true, + defaultDestination: devcontainer + ? { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'devcontainer' } + : getSandboxAllocationRequest('cloudflare-shared'), + options: [], + }; + mockGetSandboxSelectionOptions.mockResolvedValueOnce(capabilities); + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await expect( + caller.getSandboxSelectionOptions({ organizationId: ORGANIZATION_ID, devcontainer }) + ).resolves.toEqual(capabilities); + expect(mockGetSandboxSelectionOptions).toHaveBeenCalledWith({ + kilocodeOrganizationId: ORGANIZATION_ID, + devcontainer, + }); + }); + + it('rejects non-members before calling the Worker', async () => { + mockEnsureOrganizationAccess.mockImplementationOnce(() => { + throw new TRPCError({ code: 'FORBIDDEN', message: 'Membership required' }); + }); + const caller = createCaller({ user: { id: 'non-member', is_admin: false } as User }); + + await expect( + caller.getSandboxSelectionOptions({ organizationId: ORGANIZATION_ID }) + ).rejects.toThrow('Membership required'); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + expect(mockGetSandboxSelectionOptions).not.toHaveBeenCalled(); + }); + + it('keeps Worker-disabled selection disabled for administrators', async () => { + mockGetSandboxSelectionOptions.mockResolvedValueOnce({ enabled: false, options: [] }); + const caller = createCaller({ user: { id: 'admin-member', is_admin: true } as User }); + + await expect( + caller.getSandboxSelectionOptions({ organizationId: ORGANIZATION_ID }) + ).resolves.toEqual({ enabled: false, options: [] }); + expect(mockIsFeatureFlagEnabledOrDevelopment).not.toHaveBeenCalled(); + }); + + it('propagates Worker failure without granting a capability', async () => { + mockGetSandboxSelectionOptions.mockRejectedValueOnce(new Error('Worker unavailable')); + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await expect( + caller.getSandboxSelectionOptions({ organizationId: ORGANIZATION_ID }) + ).rejects.toThrow('Worker unavailable'); + }); + + it('rejects a malformed Worker capability response', async () => { + mockGetSandboxSelectionOptions.mockResolvedValueOnce({ + enabled: true, + options: [ + { + allocation: { + provider: { id: 'vercel', account: 'kilo' }, + instanceType: 'medium' as 'small', + }, + }, + ], + }); + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await expect( + caller.getSandboxSelectionOptions({ organizationId: ORGANIZATION_ID }) + ).rejects.toThrow('Output validation failed'); + }); +}); + describe('organizationCloudAgentNextRouter.prepareSession', () => { beforeEach(() => { jest.clearAllMocks(); @@ -1247,7 +1374,7 @@ describe('organizationCloudAgentNextRouter.prepareSession', () => { githubRepo: 'acme/repo', autoInitiate: true, clientProvenance: 'browser', - } as z.infer & { organizationId: string }); + } as z.input); expect(mockPrepareSession).toHaveBeenCalledWith( expect.objectContaining({ @@ -1258,6 +1385,121 @@ describe('organizationCloudAgentNextRouter.prepareSession', () => { ); }); + const sandboxInput = { + organizationId: ORGANIZATION_ID, + prompt: 'Test prompt', + mode: 'code', + model: 'kilo/test-model', + githubRepo: 'acme/repo', + }; + + it.each(SELECTABLE_SANDBOX_ALLOCATIONS)( + 'forwards the normalized legacy organization sandbox preset %s', + async sandboxAllocation => { + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await caller.prepareSession({ ...sandboxInput, sandboxAllocation }); + + expect(mockEnsureOrganizationAccess).toHaveBeenCalledWith( + expect.objectContaining({ user: { id: 'member', is_admin: false } }), + ORGANIZATION_ID + ); + expect(mockPrepareSession).toHaveBeenCalledWith( + expect.objectContaining({ + sandboxAllocation: getSandboxAllocationRequest(sandboxAllocation), + kilocodeOrganizationId: ORGANIZATION_ID, + }) + ); + } + ); + + it.each([ + ...SELECTABLE_SANDBOX_ALLOCATIONS.map(allocation => getSandboxAllocationRequest(allocation)), + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'large' }, + ] satisfies SelectableSandboxAllocationRequest[])( + 'forwards a structured sandbox destination without changing its account: %j', + async sandboxAllocation => { + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await caller.prepareSession({ ...sandboxInput, sandboxAllocation }); + + expect(mockPrepareSession).toHaveBeenCalledWith( + expect.objectContaining({ sandboxAllocation, kilocodeOrganizationId: ORGANIZATION_ID }) + ); + expect(mockGetSandboxSelectionOptions).not.toHaveBeenCalled(); + } + ); + + it('forwards the first-chat operation and attachments with the preset and server provenance', async () => { + const caller = createCaller({ + user: { id: 'organization-browser', is_admin: false } as User, + headersList: new Headers({ 'x-kilo-client': 'web' }), + }); + const input = { + ...sandboxInput, + sandboxAllocation: getSandboxAllocationRequest('vercel-small'), + operationKey: '12345678-1234-4234-9234-123456789abc', + initialMessageId: 'msg_123456789abc123456789ABCDE', + autoInitiate: true, + attachments: { + path: '12345678-1234-4234-9234-123456789abc', + files: ['87654321-4321-4321-8321-cba987654321.md'], + }, + }; + + await caller.prepareSession(input); + + expect(mockPrepareSession).toHaveBeenCalledWith( + expect.objectContaining({ + operationKey: input.operationKey, + initialMessageId: input.initialMessageId, + autoInitiate: true, + attachments: input.attachments, + sandboxAllocation: input.sandboxAllocation, + kilocodeOrganizationId: ORGANIZATION_ID, + clientProvenance: 'browser', + }) + ); + }); + + it('keeps Default omitted and independent of the capability query', async () => { + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await caller.prepareSession(sandboxInput); + + expect(mockPrepareSession).toHaveBeenCalledTimes(1); + expect(mockPrepareSession.mock.calls[0][0]).not.toHaveProperty('sandboxAllocation'); + expect(mockGetSandboxSelectionOptions).not.toHaveBeenCalled(); + }); + + it('rejects an invalid preset before calling the Worker', async () => { + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await expect( + caller.prepareSession({ + ...sandboxInput, + sandboxAllocation: 'vercel-medium' as SelectableSandboxAllocation, + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockPrepareSession).not.toHaveBeenCalled(); + expect(mockComputeCloudAgentNextBalanceCheckEligibility).not.toHaveBeenCalled(); + }); + + it('rejects explicit presets with dev containers before calling the Worker', async () => { + const caller = createCaller({ user: { id: 'member', is_admin: false } as User }); + + await expect( + caller.prepareSession({ + ...sandboxInput, + sandboxAllocation: 'cloudflare-single', + devcontainer: true, + }) + ).rejects.toThrow('Sandbox selection is not available with dev containers'); + expect(mockPrepareSession).not.toHaveBeenCalled(); + expect(mockIsFeatureFlagEnabledOrDevelopment).not.toHaveBeenCalled(); + }); + it('rejects devcontainer sessions when the feature flag is disabled', async () => { mockIsFeatureFlagEnabledOrDevelopment.mockResolvedValue(false); const caller = createCaller({ @@ -1506,8 +1748,31 @@ describe('organizationCloudAgentNextRouter.createWorktreeChat', () => { operationKey: uuid, organizationId: ORGANIZATION_ID, }); + expect(mockGetSandboxSelectionOptions).not.toHaveBeenCalled(); + expect(mockIsFeatureFlagEnabledOrDevelopment).not.toHaveBeenCalled(); }); + it.each([ + ...SELECTABLE_SANDBOX_ALLOCATIONS, + ...SELECTABLE_SANDBOX_ALLOCATIONS.map(allocation => getSandboxAllocationRequest(allocation)), + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'large' }, + ])( + 'rejects a sibling chat sandbox override %j before invoking the operation', + async sandboxAllocation => { + const caller = createCaller({ user: { id: 'organization-owner', is_admin: false } as User }); + const input = { + organizationId: ORGANIZATION_ID, + sourceKiloSessionId, + operationKey: uuid, + sandboxAllocation, + }; + + await expect(caller.createWorktreeChat(input)).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockCreateWorktreeChat).not.toHaveBeenCalled(); + } + ); + it('rejects a revoked organization member before resolving the source session', async () => { mockEnsureOrganizationAccess.mockImplementation(() => { throw new TRPCError({ diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index 9425cd6e39..c97e0729ce 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -1,5 +1,6 @@ import 'server-only'; import { baseProcedure, createTRPCRouter } from '@/lib/trpc/init'; +import { sandboxSelectionCapabilitiesSchema } from '@kilocode/worker-utils/sandbox-allocation'; import { createCloudAgentNextClient, createCloudAgentNextClientForModel, @@ -28,7 +29,7 @@ import { } from '@/lib/cloud-agent/gitlab-integration-helpers'; import { orderRepositoriesByUsage } from '@/lib/cloud-agent/order-repositories'; import { - basePrepareSessionNextSchema, + organizationPrepareSessionNextSchema, basePrepareSessionNextOutputSchema, baseCreateWorktreeChatNextSchema, baseCreateWorktreeChatNextOutputSchema, @@ -153,12 +154,6 @@ async function assertOrganizationOwnsSession(params: { } // Extend base schemas with organizationId for organization context -const PrepareSessionInput = basePrepareSessionNextSchema.and( - z.object({ - organizationId: z.uuid(), - }) -); - const CreateWorktreeChatInput = baseCreateWorktreeChatNextSchema.extend({ organizationId: z.uuid(), }); @@ -266,6 +261,21 @@ const ListBitbucketRepositoriesInput = z.object({ * separately via WebSocket connection. */ export const organizationCloudAgentNextRouter = createTRPCRouter({ + getSandboxSelectionOptions: organizationMemberProcedure + .input(z.object({ devcontainer: z.boolean().optional() })) + .output(sandboxSelectionCapabilitiesSchema) + .query(async ({ ctx, input }) => { + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); + return await createCloudAgentNextClient(authToken).getSandboxSelectionOptions({ + kilocodeOrganizationId: input.organizationId, + ...(input.devcontainer !== undefined ? { devcontainer: input.devcontainer } : {}), + }); + }), + /** * Prepare a new cloud agent session (organization context). * @@ -274,7 +284,7 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ * initiateFromPreparedSession. */ prepareSession: organizationMemberMutationProcedure - .input(PrepareSessionInput) + .input(organizationPrepareSessionNextSchema) .output(basePrepareSessionNextOutputSchema) .mutation(async ({ ctx, input }) => { if ( diff --git a/apps/web/src/routers/webhook-triggers-router.schema.test.ts b/apps/web/src/routers/webhook-triggers-router.schema.test.ts index 180a9f896f..a043f077be 100644 --- a/apps/web/src/routers/webhook-triggers-router.schema.test.ts +++ b/apps/web/src/routers/webhook-triggers-router.schema.test.ts @@ -2,6 +2,7 @@ import { beforeAll, beforeEach, describe, expect, it, jest } from '@jest/globals import { TRPCError } from '@trpc/server'; import { createCallerFactory } from '@/lib/trpc/init'; import type { User } from '@kilocode/db/schema'; +import type { CloudAgentNextClient } from '@/lib/cloud-agent-next/cloud-agent-client'; import type { ensureOrganizationAccess } from '@/routers/organizations/utils'; import type { createWorkerTrigger as createWorkerTriggerType, @@ -16,6 +17,18 @@ import type { } from './webhook-triggers-router'; const mockEnsureOrganizationAccess = jest.fn(); +const mockGetSandboxSelectionOptions = + jest.fn(); + +jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ + createCloudAgentNextClient: () => ({ + getSandboxSelectionOptions: mockGetSandboxSelectionOptions, + }), +})); + +jest.mock('@/lib/tokens', () => ({ + generateCloudAgentToken: () => 'test-cloud-agent-token', +})); const mockCreateWorkerTrigger = jest.fn(); const mockUpdateWorkerTrigger = jest.fn(); @@ -122,6 +135,7 @@ describe('webhook trigger variant inputs', () => { requestId: '00000000-0000-4000-8000-000000000005', }); mockEnsureOrganizationAccess.mockResolvedValue('owner'); + mockGetSandboxSelectionOptions.mockResolvedValue({ enabled: false, options: [] }); }); it.each(['high', 'High', 'a'.repeat(50)])('accepts valid create variant %s', variant => { @@ -243,59 +257,55 @@ describe('webhook trigger variant inputs', () => { ); it.each(['owner', 'admin', 'member'] as const)( - 'reports actual platform-admin capability for an organization %s', + 'uses Worker organization capabilities for an organization %s', async role => { const organizationId = '00000000-0000-4000-8000-000000000003'; - mockEnsureOrganizationAccess.mockResolvedValueOnce(role); + mockEnsureOrganizationAccess.mockResolvedValue(role); await expect(createCaller({ user }).capabilities({ organizationId })).resolves.toEqual({ canSetSandboxAllocation: false, }); + mockGetSandboxSelectionOptions.mockResolvedValueOnce({ enabled: true, options: [] }); + await expect(createCaller({ user }).capabilities({ organizationId })).resolves.toEqual({ + canSetSandboxAllocation: true, + }); expect(mockEnsureOrganizationAccess).toHaveBeenCalledWith(expect.anything(), organizationId); + expect(mockGetSandboxSelectionOptions).toHaveBeenCalledWith({ + kilocodeOrganizationId: organizationId, + }); } ); - it('reports platform-admin capability and preserves organization access checks', async () => { - await expect(createCaller({ user }).capabilities({})).resolves.toEqual({ + it.each([false, true])('never enables personal selection for is_admin=%s', async is_admin => { + await expect(createCaller({ user: { ...user, is_admin } }).capabilities({})).resolves.toEqual({ canSetSandboxAllocation: false, }); - await expect( - createCaller({ user: { ...user, is_admin: true } }).capabilities({}) - ).resolves.toEqual({ canSetSandboxAllocation: true }); + expect(mockGetSandboxSelectionOptions).not.toHaveBeenCalled(); + }); + it('checks organization access before requesting Worker capabilities', async () => { const organizationId = '00000000-0000-4000-8000-000000000003'; mockEnsureOrganizationAccess.mockRejectedValueOnce(new TRPCError({ code: 'UNAUTHORIZED' })); await expect(createCaller({ user }).capabilities({ organizationId })).rejects.toMatchObject({ code: 'UNAUTHORIZED', }); + expect(mockGetSandboxSelectionOptions).not.toHaveBeenCalled(); }); - it.each([ - { scope: 'personal', activationMode: 'webhook', organizationId: undefined }, - { - scope: 'organization', - activationMode: 'webhook', - organizationId: '00000000-0000-4000-8000-000000000003', - }, - { scope: 'personal', activationMode: 'scheduled', organizationId: undefined }, - { - scope: 'organization', - activationMode: 'scheduled', - organizationId: '00000000-0000-4000-8000-000000000003', - }, - ] as const)( - 'allows a Kilo admin to create and update Dedicated Standard for $scope $activationMode triggers', - async ({ activationMode, organizationId }) => { - const admin = { ...user, is_admin: true }; - const triggerInput = { + it.each(['webhook', 'scheduled'] as const)( + 'allows an enrolled organization member to set Dedicated Standard on %s triggers', + async activationMode => { + const organizationId = '00000000-0000-4000-8000-000000000003'; + mockEnsureOrganizationAccess.mockResolvedValue('member'); + mockGetSandboxSelectionOptions.mockResolvedValue({ enabled: true, options: [] }); + await createCaller({ user }).create({ ...createInput, - ...(organizationId ? { organizationId } : {}), + organizationId, activationMode, ...(activationMode === 'scheduled' ? { cronExpression: '* * * * *' } : {}), - sandboxAllocation: 'isolated-standard' as const, - }; - await createCaller({ user: admin }).create(triggerInput); + sandboxAllocation: 'isolated-standard', + }); expect(mockCreateWorkerTrigger).toHaveBeenCalledWith( - organizationId ? undefined : 'user-1', + undefined, organizationId, 'trigger-id', expect.objectContaining({ sandboxAllocation: 'isolated-standard' }) @@ -308,13 +318,13 @@ describe('webhook trigger variant inputs', () => { target_type: 'cloud_agent', }, ]); - await createCaller({ user: admin }).update({ + await createCaller({ user }).update({ triggerId: 'trigger-id', - ...(organizationId ? { organizationId } : {}), + organizationId, sandboxAllocation: 'isolated-standard', }); expect(mockUpdateWorkerTrigger).toHaveBeenCalledWith( - organizationId ? undefined : 'user-1', + undefined, organizationId, 'trigger-id', expect.objectContaining({ sandboxAllocation: 'isolated-standard' }) @@ -361,28 +371,33 @@ describe('webhook trigger variant inputs', () => { } ); - it.each(['webhook', 'scheduled'] as const)( - 'rejects a non-admin Dedicated Standard create and update for %s triggers before writes', - async activationMode => { - const input = { - ...createInput, - activationMode, - ...(activationMode === 'scheduled' ? { cronExpression: '* * * * *' } : {}), - sandboxAllocation: 'isolated-standard' as const, - }; - await expect(createCaller({ user }).create(input)).rejects.toMatchObject({ - code: 'FORBIDDEN', - message: 'Kilo admin access is required to select Dedicated Standard', - }); - await expect( - createCaller({ user }).update({ - triggerId: 'trigger-id', - sandboxAllocation: 'isolated-standard', - }) - ).rejects.toMatchObject({ - code: 'FORBIDDEN', - message: 'Kilo admin access is required to select Dedicated Standard', - }); + it.each([ + { activationMode: 'webhook', is_admin: false }, + { activationMode: 'webhook', is_admin: true }, + { activationMode: 'scheduled', is_admin: false }, + { activationMode: 'scheduled', is_admin: true }, + ] as const)( + 'rejects unenrolled Dedicated Standard $activationMode triggers for is_admin=$is_admin before writes', + async ({ activationMode, is_admin }) => { + const caller = createCaller({ user: { ...user, is_admin } }); + for (const organizationId of [undefined, '00000000-0000-4000-8000-000000000003']) { + await expect( + caller.create({ + ...createInput, + organizationId, + activationMode, + ...(activationMode === 'scheduled' ? { cronExpression: '* * * * *' } : {}), + sandboxAllocation: 'isolated-standard', + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + await expect( + caller.update({ + triggerId: 'trigger-id', + organizationId, + sandboxAllocation: 'isolated-standard', + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + } expect(mockDbInsert).not.toHaveBeenCalled(); expect(mockDbUpdate).not.toHaveBeenCalled(); expect(mockCreateWorkerTrigger).not.toHaveBeenCalled(); @@ -390,19 +405,24 @@ describe('webhook trigger variant inputs', () => { } ); - it('does not grant allocation privileges to an organization owner or member', async () => { + it('fails closed before writes when Worker capabilities are unavailable', async () => { const organizationId = '00000000-0000-4000-8000-000000000003'; - for (const role of ['owner', 'member'] as const) { - mockEnsureOrganizationAccess.mockResolvedValueOnce(role); - await expect( - createCaller({ user }).create({ - ...createInput, - organizationId, - sandboxAllocation: 'isolated-standard', - }) - ).rejects.toMatchObject({ code: 'FORBIDDEN' }); - } + mockGetSandboxSelectionOptions.mockRejectedValue(new Error('Capabilities unavailable')); + const caller = createCaller({ user }); + await expect( + caller.create({ ...createInput, organizationId, sandboxAllocation: 'isolated-standard' }) + ).rejects.toThrow('Capabilities unavailable'); + await expect( + caller.update({ + triggerId: 'trigger-id', + organizationId, + sandboxAllocation: 'isolated-standard', + }) + ).rejects.toThrow('Capabilities unavailable'); + expect(mockDbInsert).not.toHaveBeenCalled(); + expect(mockDbUpdate).not.toHaveBeenCalled(); expect(mockCreateWorkerTrigger).not.toHaveBeenCalled(); + expect(mockUpdateWorkerTrigger).not.toHaveBeenCalled(); }); it.each([null, 'isolated-standard'] as const)( diff --git a/apps/web/src/routers/webhook-triggers-router.ts b/apps/web/src/routers/webhook-triggers-router.ts index 9d5786d83e..c4c29c5f76 100644 --- a/apps/web/src/routers/webhook-triggers-router.ts +++ b/apps/web/src/routers/webhook-triggers-router.ts @@ -1,4 +1,6 @@ import { createTRPCRouter, baseProcedure } from '@/lib/trpc/init'; +import { createCloudAgentNextClient } from '@/lib/cloud-agent-next/cloud-agent-client'; +import { generateCloudAgentToken } from '@/lib/tokens'; import { ensureOrganizationAccess } from '@/routers/organizations/utils'; import { TRPCError } from '@trpc/server'; import { and, eq, isNull } from 'drizzle-orm'; @@ -8,6 +10,7 @@ import { cloud_agent_webhook_triggers, agent_environment_profiles, kiloclaw_instances, + type User, } from '@kilocode/db/schema'; import { resolveCloudAgentSessionIds } from '@/lib/webhook-session-resolution'; import { triggerIdSchema, triggerIdCreateSchema } from '@/lib/webhook-trigger-validation'; @@ -276,6 +279,14 @@ async function assertProfileOwnership( } } +async function canSetSandboxAllocation(user: User, organizationId?: string): Promise { + if (!organizationId) return false; + const capabilities = await createCloudAgentNextClient( + generateCloudAgentToken(user) + ).getSandboxSelectionOptions({ kilocodeOrganizationId: organizationId }); + return capabilities.enabled; +} + export const webhookTriggersRouter = createTRPCRouter({ capabilities: baseProcedure .input(z.object({ organizationId: z.string().uuid().optional() })) @@ -283,7 +294,9 @@ export const webhookTriggersRouter = createTRPCRouter({ if (input.organizationId) { await ensureOrganizationAccess(ctx, input.organizationId); } - return { canSetSandboxAllocation: ctx.user.is_admin }; + return { + canSetSandboxAllocation: await canSetSandboxAllocation(ctx.user, input.organizationId), + }; }), /** @@ -424,10 +437,13 @@ export const webhookTriggersRouter = createTRPCRouter({ await ensureOrganizationAccess(ctx, input.organizationId, ['owner', 'member']); } - if (input.sandboxAllocation === 'isolated-standard' && !ctx.user.is_admin) { + if ( + input.sandboxAllocation && + !(await canSetSandboxAllocation(ctx.user, input.organizationId)) + ) { throw new TRPCError({ code: 'FORBIDDEN', - message: 'Kilo admin access is required to select Dedicated Standard', + message: 'Sandbox selection requires an enabled organization', }); } @@ -609,10 +625,13 @@ export const webhookTriggersRouter = createTRPCRouter({ }); } - if (input.sandboxAllocation === 'isolated-standard' && !ctx.user.is_admin) { + if ( + input.sandboxAllocation && + !(await canSetSandboxAllocation(ctx.user, input.organizationId)) + ) { throw new TRPCError({ code: 'FORBIDDEN', - message: 'Kilo admin access is required to select Dedicated Standard', + message: 'Sandbox selection requires an enabled organization', }); } diff --git a/packages/worker-utils/package.json b/packages/worker-utils/package.json index ef2afd39f7..6b001e02f6 100644 --- a/packages/worker-utils/package.json +++ b/packages/worker-utils/package.json @@ -18,6 +18,7 @@ "./runtime-proxy-attestation": "./src/runtime-proxy-attestation.ts", "./kilo-auth-middleware": "./src/kilo-auth-middleware.ts", "./sandbox-id": "./src/sandbox-id.ts", + "./sandbox-allocation": "./src/sandbox-allocation.ts", "./hostname-label": "./src/hostname-label.ts", "./deployment-slug": "./src/deployment-slug.ts", "./redact-headers": "./src/redact-headers.ts", diff --git a/packages/worker-utils/src/sandbox-allocation.test.ts b/packages/worker-utils/src/sandbox-allocation.test.ts new file mode 100644 index 0000000000..a1dc02efa0 --- /dev/null +++ b/packages/worker-utils/src/sandbox-allocation.test.ts @@ -0,0 +1,164 @@ +import { describe, expect, it } from 'vitest'; +import { + SELECTABLE_SANDBOX_ALLOCATIONS, + getSandboxAllocationProvider, + getSandboxAllocationResources, + getSandboxAllocationRequest, + getSandboxAllocationKey, + getKiloSandboxAllocation, + sandboxAllocationInputSchema, + sandboxAllocationRequestSchema, + selectableSandboxAllocationInputSchema, + isSelectableSandboxAllocation, + sandboxAllocationRequiresControlPlane, + sandboxAllocationSchema, + sandboxSelectionCapabilitiesSchema, + vercelSandboxResourcesSchema, +} from './sandbox-allocation.js'; + +describe('sandbox allocation contract', () => { + it.each(sandboxAllocationSchema.options)( + 'normalizes legacy %s to its structured request', + allocation => { + const request = getSandboxAllocationRequest(allocation); + expect(sandboxAllocationInputSchema.parse(allocation)).toEqual(request); + expect(sandboxAllocationInputSchema.parse(request)).toEqual(request); + expect(getKiloSandboxAllocation(request)).toBe(allocation); + } + ); + + it('keeps BYOC distinct from Kilo compute with the same instance type', () => { + const byoc = sandboxAllocationRequestSchema.parse({ + provider: { id: 'vercel', account: 'byoc' }, + instanceType: 'small', + }); + const kilo = getSandboxAllocationRequest('vercel-small'); + expect(getSandboxAllocationKey(byoc)).not.toBe(getSandboxAllocationKey(kilo)); + expect(getKiloSandboxAllocation(byoc)).toBeUndefined(); + expect(selectableSandboxAllocationInputSchema.parse(byoc)).toEqual(byoc); + }); + + it.each([ + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'shared' }, + { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'large' }, + { provider: { id: 'cloudflare', account: 'byoc' }, instanceType: 'single' }, + { provider: { id: 'vercel', account: 'platform' }, instanceType: 'small' }, + { provider: { id: 'vercel' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'kilo', token: 'not-allowed' }, instanceType: 'small' }, + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'small', vcpus: 2 }, + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'devcontainer' }, + ])('rejects unsupported provider metadata and instance combinations %j', request => { + expect(sandboxAllocationInputSchema.safeParse(request).success).toBe(false); + }); + + it('keeps structured Dedicated Standard out of the manual choices', () => { + const request = getSandboxAllocationRequest('isolated-standard'); + expect(sandboxAllocationInputSchema.parse(request)).toEqual(request); + expect(selectableSandboxAllocationInputSchema.safeParse(request).success).toBe(false); + }); + + it('reads legacy capabilities as structured choices without inventing a default', () => { + const capabilities = sandboxSelectionCapabilitiesSchema.parse({ + enabled: true, + options: [{ allocation: 'vercel-small', available: true }], + }); + expect(capabilities).toEqual({ + enabled: true, + options: [{ allocation: getSandboxAllocationRequest('vercel-small') }], + }); + }); + + it('accepts provider-default and devcontainer descriptions without making them selectable', () => { + for (const defaultDestination of [ + { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'devcontainer' }, + ]) { + const capabilities = sandboxSelectionCapabilitiesSchema.parse({ + enabled: true, + options: [], + defaultDestination, + }); + expect(capabilities.defaultDestination).toEqual(defaultDestination); + } + }); + + it.each([ + ['vercel-small', { vcpus: 2, memory: 4096 }], + ['vercel-large', { vcpus: 4, memory: 8192 }], + ['cloudflare-single', undefined], + ['cloudflare-shared', undefined], + ['isolated-standard', undefined], + [undefined, undefined], + ] as const)('maps %s to fixed provider resources', (preset, resources) => { + expect(getSandboxAllocationResources(preset)).toEqual(resources); + if (resources) expect(vercelSandboxResourcesSchema.parse(resources)).toEqual(resources); + }); + + it.each([ + { vcpus: 2, memory: 8192 }, + { vcpus: 4, memory: 4096 }, + { vcpus: 1, memory: 2048 }, + { vcpus: 2, memory: 4096, disk: 20 }, + { vcpus: '2', memory: 4096 }, + ])('rejects unsupported resource pairs %j', resources => { + expect(vercelSandboxResourcesSchema.safeParse(resources).success).toBe(false); + }); + + it.each([ + ['isolated-standard', 'cloudflare', false], + ['cloudflare-single', 'cloudflare', false], + ['cloudflare-shared', 'cloudflare', false], + ['vercel-small', 'vercel', true], + ['vercel-large', 'vercel', true], + ] as const)( + 'maps %s to provider %s and control-plane requirement %s', + (preset, provider, forcesPlane) => { + expect(getSandboxAllocationProvider(preset)).toBe(provider); + expect(sandboxAllocationRequiresControlPlane(preset)).toBe(forcesPlane); + } + ); + + it('leaves an omitted selection to the existing plane decision', () => { + expect(sandboxAllocationRequiresControlPlane(undefined)).toBe(false); + }); + + it('rejects arbitrary allocations and capability options', () => { + expect(sandboxAllocationSchema.safeParse('vercel-custom').success).toBe(false); + expect( + sandboxSelectionCapabilitiesSchema.safeParse({ + enabled: true, + options: [{ allocation: 'custom' }], + }).success + ).toBe(false); + }); + + it('keeps isolated-standard in the field but out of the selectable set', () => { + expect(sandboxAllocationSchema.safeParse('isolated-standard').success).toBe(true); + expect(isSelectableSandboxAllocation('isolated-standard')).toBe(false); + expect(SELECTABLE_SANDBOX_ALLOCATIONS).not.toContain('isolated-standard'); + expect( + sandboxSelectionCapabilitiesSchema.safeParse({ + enabled: true, + options: [{ allocation: 'isolated-standard' }], + }).success + ).toBe(false); + }); + + it.each(SELECTABLE_SANDBOX_ALLOCATIONS)('treats %s as selectable', allocation => { + expect(isSelectableSandboxAllocation(allocation)).toBe(true); + }); + + it('requires disabled capability results to omit all options', () => { + expect(sandboxSelectionCapabilitiesSchema.parse({ enabled: false, options: [] })).toEqual({ + enabled: false, + options: [], + }); + expect( + sandboxSelectionCapabilitiesSchema.safeParse({ + enabled: false, + options: [{ allocation: 'vercel-small' }], + }).success + ).toBe(false); + }); +}); diff --git a/packages/worker-utils/src/sandbox-allocation.ts b/packages/worker-utils/src/sandbox-allocation.ts new file mode 100644 index 0000000000..9ab27c4e74 --- /dev/null +++ b/packages/worker-utils/src/sandbox-allocation.ts @@ -0,0 +1,176 @@ +import { z } from 'zod'; + +export const sandboxAllocationSchema = z.enum([ + 'isolated-standard', + 'cloudflare-single', + 'cloudflare-shared', + 'vercel-small', + 'vercel-large', +]); + +export type SandboxAllocation = z.infer; + +export const SELECTABLE_SANDBOX_ALLOCATIONS = [ + 'cloudflare-single', + 'cloudflare-shared', + 'vercel-small', + 'vercel-large', +] as const satisfies readonly SandboxAllocation[]; + +export type SelectableSandboxAllocation = (typeof SELECTABLE_SANDBOX_ALLOCATIONS)[number]; + +const cloudflareAllocationRequestSchema = z + .object({ + provider: z.object({ id: z.literal('cloudflare'), account: z.literal('kilo') }).strict(), + instanceType: z.enum(['single', 'shared', 'isolated-standard']), + }) + .strict(); + +const vercelAllocationRequestSchema = z + .object({ + provider: z.object({ id: z.literal('vercel'), account: z.enum(['kilo', 'byoc']) }).strict(), + instanceType: z.enum(['small', 'large']), + }) + .strict(); + +export const sandboxAllocationRequestSchema = z.union([ + cloudflareAllocationRequestSchema, + vercelAllocationRequestSchema, +]); + +export type SandboxAllocationRequest = z.infer; + +const selectableSandboxAllocationRequestSchema = z.union([ + cloudflareAllocationRequestSchema.extend({ instanceType: z.enum(['single', 'shared']) }), + vercelAllocationRequestSchema, +]); + +export type SelectableSandboxAllocationRequest = z.infer< + typeof selectableSandboxAllocationRequestSchema +>; + +const allocationRequests = { + 'isolated-standard': { + provider: { id: 'cloudflare', account: 'kilo' }, + instanceType: 'isolated-standard', + }, + 'cloudflare-single': { + provider: { id: 'cloudflare', account: 'kilo' }, + instanceType: 'single', + }, + 'cloudflare-shared': { + provider: { id: 'cloudflare', account: 'kilo' }, + instanceType: 'shared', + }, + 'vercel-small': { + provider: { id: 'vercel', account: 'kilo' }, + instanceType: 'small', + }, + 'vercel-large': { + provider: { id: 'vercel', account: 'kilo' }, + instanceType: 'large', + }, +} as const satisfies Record; + +export function getSandboxAllocationRequest(allocation: T) { + return allocationRequests[allocation]; +} + +export function getSandboxAllocationKey(allocation: SandboxAllocationRequest): string { + return `${allocation.provider.id}:${allocation.provider.account}:${allocation.instanceType}`; +} + +export function getKiloSandboxAllocation( + request: SandboxAllocationRequest +): SandboxAllocation | undefined { + const key = getSandboxAllocationKey(request); + return sandboxAllocationSchema.options.find( + allocation => getSandboxAllocationKey(allocationRequests[allocation]) === key + ); +} + +export const sandboxAllocationInputSchema = z.union([ + sandboxAllocationRequestSchema, + sandboxAllocationSchema.transform(allocation => getSandboxAllocationRequest(allocation)), +]); + +export type SandboxAllocationInput = z.input; + +export const selectableSandboxAllocationInputSchema = z.union([ + selectableSandboxAllocationRequestSchema, + z + .enum(SELECTABLE_SANDBOX_ALLOCATIONS) + .transform(allocation => getSandboxAllocationRequest(allocation)), +]); + +export const sandboxDestinationSchema = z.union([ + cloudflareAllocationRequestSchema.extend({ + instanceType: z.enum(['single', 'shared', 'isolated-standard', 'devcontainer']), + }), + vercelAllocationRequestSchema.extend({ instanceType: z.enum(['small', 'large', 'default']) }), +]); + +export type SandboxDestination = z.infer; + +export function isSelectableSandboxAllocation( + allocation: SandboxAllocation | undefined +): allocation is SelectableSandboxAllocation { + return ( + allocation !== undefined && + (SELECTABLE_SANDBOX_ALLOCATIONS as readonly string[]).includes(allocation) + ); +} + +export const vercelSandboxResourcesSchema = z.union([ + z.object({ vcpus: z.literal(2), memory: z.literal(4096) }).strict(), + z.object({ vcpus: z.literal(4), memory: z.literal(8192) }).strict(), +]); + +export type VercelSandboxResources = z.infer; + +export function getSandboxAllocationProvider( + allocation: SandboxAllocation +): 'cloudflare' | 'vercel' { + return allocation.startsWith('vercel-') ? 'vercel' : 'cloudflare'; +} + +/** + * Vercel sandboxes exist only on the control plane, so a Vercel allocation forces a + * control-plane session regardless of `CONTROL_PLANE_IDS`. Cloudflare allocations pick + * the sandbox shape only and leave the plane decision to that allowlist. + */ +export function sandboxAllocationRequiresControlPlane( + allocation: SandboxAllocation | undefined +): boolean { + return allocation !== undefined && getSandboxAllocationProvider(allocation) === 'vercel'; +} + +export function getSandboxAllocationResources( + allocation: SandboxAllocation | undefined +): VercelSandboxResources | undefined { + switch (allocation) { + case 'vercel-small': + return { vcpus: 2, memory: 4096 }; + case 'vercel-large': + return { vcpus: 4, memory: 8192 }; + default: + return undefined; + } +} + +export const sandboxSelectionCapabilitiesSchema = z + .object({ + enabled: z.boolean(), + defaultDestination: sandboxDestinationSchema.optional(), + options: z.array( + z.object({ + allocation: selectableSandboxAllocationInputSchema, + }) + ), + }) + .refine(value => value.enabled || value.options.length === 0, { + message: 'Disabled sandbox selection must not expose options', + path: ['options'], + }); + +export type SandboxSelectionCapabilities = z.infer; diff --git a/services/cloud-agent-next/.dev.vars.example b/services/cloud-agent-next/.dev.vars.example index dcb43c52e8..a2be954522 100644 --- a/services/cloud-agent-next/.dev.vars.example +++ b/services/cloud-agent-next/.dev.vars.example @@ -46,6 +46,8 @@ PER_SESSION_SANDBOX_ORG_IDS= # Local wrangler `dev` also defaults these to `*` in wrangler.jsonc. CONTROL_PLANE_IDS=* WORKTREE_CREATION_ENABLED_IDS=* +# Comma-separated user or org IDs allowed to pick a sandbox destination, or `*` for all +SANDBOX_SELECTION_IDS=* # Non-secret rollout control for per-session Kilo runtime isolation. RUNTIME_ISOLATION_ENABLED=true diff --git a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.test.ts b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.test.ts index 761c006197..2af9118c2a 100644 --- a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.test.ts +++ b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.test.ts @@ -3,6 +3,7 @@ import { parseVercelSandboxCredentials, parseVercelSandboxEnrollment, parseVercelSandboxRuntimeConfig, + resolveVercelSandboxRuntimeConfig, } from './vercel-runtime-config.js'; const completeRuntimeEnv = { @@ -67,6 +68,49 @@ describe('parseVercelSandboxRuntimeConfig', () => { }); }); +describe('resolveVercelSandboxRuntimeConfig', () => { + const persisted = { + projectId: 'old-project', + snapshotId: 'old-snapshot', + runtimeBuildId: 'old-build', + runtime: 'node24', + }; + + it.each([ + { vcpus: 2, memory: 4096 }, + { vcpus: 4, memory: 8192 }, + ] as const)( + 'preserves persisted $vcpus vCPU resources across operational configuration changes', + resources => { + expect( + resolveVercelSandboxRuntimeConfig(completeRuntimeEnv, { ...persisted, resources }) + ).toMatchObject({ + ...persisted, + resources, + accessToken: completeRuntimeEnv.VERCEL_TOKEN, + }); + } + ); + + it('does not add sizing to older persisted runtime identities', () => { + expect(resolveVercelSandboxRuntimeConfig(completeRuntimeEnv, persisted)).not.toHaveProperty( + 'resources' + ); + expect(resolveVercelSandboxRuntimeConfig(completeRuntimeEnv)).not.toHaveProperty('resources'); + }); + + it.each([{ vcpus: 2, memory: 8192 }, { vcpus: 8, memory: 16384 }, null])( + 'rejects invalid persisted resources even when operational configuration is unavailable: %j', + resources => { + const input = { ...persisted, resources } as Parameters< + typeof resolveVercelSandboxRuntimeConfig + >[1]; + expect(() => resolveVercelSandboxRuntimeConfig(completeRuntimeEnv, input)).toThrow(); + expect(() => resolveVercelSandboxRuntimeConfig({}, input)).toThrow(); + } + ); +}); + describe('parseVercelSandboxEnrollment', () => { it('is disabled when the list is empty', () => { expect(parseVercelSandboxEnrollment({})).toEqual({ diff --git a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.ts b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.ts index 48c0a1b52d..edf659e342 100644 --- a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.ts +++ b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-runtime-config.ts @@ -1,3 +1,7 @@ +import { + vercelSandboxResourcesSchema, + type VercelSandboxResources, +} from '@kilocode/worker-utils/sandbox-allocation'; import { z } from 'zod'; const positiveIntegerString = z @@ -24,6 +28,7 @@ export type VercelSandboxRuntimeConfig = { snapshotId: string; runtimeBuildId: string; runtime: 'node24'; + resources?: VercelSandboxResources; initialTimeoutMs: number; extendDurationMs: number; }; @@ -82,22 +87,26 @@ export function resolveVercelSandboxRuntimeConfig( snapshotId?: string; runtimeBuildId?: string; runtime?: string; + resources?: VercelSandboxResources; } ): VercelSandboxRuntimeConfig | undefined { + const resources = vercelSandboxResourcesSchema.optional().parse(persisted?.resources); const configured = parseVercelSandboxRuntimeConfig(env); if (!configured) return undefined; - return persisted?.projectId && + const resolved: VercelSandboxRuntimeConfig = + persisted?.projectId && persisted.snapshotId && persisted.runtimeBuildId && persisted.runtime === 'node24' - ? { - ...configured, - projectId: persisted.projectId, - snapshotId: persisted.snapshotId, - runtimeBuildId: persisted.runtimeBuildId, - runtime: persisted.runtime, - } - : configured; + ? { + ...configured, + projectId: persisted.projectId, + snapshotId: persisted.snapshotId, + runtimeBuildId: persisted.runtimeBuildId, + runtime: persisted.runtime, + } + : configured; + return resources === undefined ? resolved : { ...resolved, resources }; } export type VercelSandboxEnrollmentEnv = { diff --git a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.test.ts b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.test.ts index 1422bd5cbb..39dd62fbef 100644 --- a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.test.ts +++ b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.test.ts @@ -149,6 +149,79 @@ describe('VercelSandboxRestClient', () => { }); }); + it.each([ + { vcpus: 2, memory: 4096 }, + { vcpus: 4, memory: 8192 }, + ] as const)('creates and inspects explicitly sized $vcpus vCPU sandboxes', async resources => { + const providerFetch = vi.fn().mockImplementation(async () => + jsonResponse({ + resumed: false, + sandbox: sandbox(), + session: session(resources), + routes: [], + }) + ); + const client = clientFor(providerFetch); + const input = { ...createInput(), resources }; + await expect(client.createSandbox(input)).resolves.toMatchObject({ session: resources }); + expect(JSON.parse(providerFetch.mock.calls[0][1].body as string).resources).toEqual(resources); + await expect(client.inspectByName(input)).resolves.toMatchObject({ session: resources }); + }); + + describe.each(['createSandbox', 'inspectByName'] as const)('%s resources', operation => { + it.each([{ vcpus: 4 }, { memory: 8192 }])( + 'rejects a mismatched response %j', + async mismatch => { + const client = clientFor( + vi.fn().mockResolvedValue( + jsonResponse({ + resumed: false, + sandbox: sandbox(), + session: session(mismatch), + routes: [], + }) + ) + ); + await expect( + client[operation]({ + ...createInput(), + resources: { vcpus: 2, memory: 4096 }, + }) + ).rejects.toMatchObject({ kind: 'correlation_mismatch' }); + } + ); + + it('retains provider-default behavior when resources are omitted', async () => { + const client = clientFor( + vi.fn().mockResolvedValue( + jsonResponse({ + resumed: false, + sandbox: sandbox(), + session: session({ vcpus: 8, memory: 16384 }), + routes: [], + }) + ) + ); + await expect(client[operation](createInput())).resolves.toMatchObject({ + session: { vcpus: 8, memory: 16384 }, + }); + }); + + it.each([{ vcpus: 2, memory: 8192 }, { vcpus: 8, memory: 16384 }, null])( + 'rejects invalid resources before provider I/O: %j', + async resources => { + const providerFetch = vi.fn(); + const input = { ...createInput(), resources } as Parameters< + VercelSandboxRestClient['createSandbox'] + >[0]; + await expect(clientFor(providerFetch)[operation](input)).rejects.toMatchObject({ + kind: 'invalid_request', + }); + expect(providerFetch).not.toHaveBeenCalled(); + } + ); + }); + it('creates a contained sandbox with a nested REST-native policy and redirects disabled', async () => { const providerFetch = vi .fn() diff --git a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.ts b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.ts index fec9df0660..024ef378a5 100644 --- a/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.ts +++ b/services/cloud-agent-next/src/agent-sandbox/vercel/vercel-sandbox-rest-client.ts @@ -1,3 +1,7 @@ +import { + vercelSandboxResourcesSchema, + type VercelSandboxResources, +} from '@kilocode/worker-utils/sandbox-allocation'; import { z } from 'zod'; const VERCEL_SANDBOX_API_BASE_URL = 'https://api.vercel.com'; @@ -123,6 +127,7 @@ export type CreateSandboxInput = { snapshotId: string; runtime: VercelSandboxRuntime; timeoutMs: number; + resources?: VercelSandboxResources; networkPolicy?: VercelSandboxNetworkPolicy; }; @@ -292,7 +297,10 @@ export class VercelSandboxRestClient { requireIdentifier(input.runtimeBuildId, operation); requireIdentifier(input.snapshotId, operation); requirePositiveDuration(input.timeoutMs, operation); - if (!runtimeSchema.safeParse(input.runtime).success) { + if ( + !runtimeSchema.safeParse(input.runtime).success || + !vercelSandboxResourcesSchema.optional().safeParse(input.resources).success + ) { throw new VercelSandboxRestError('invalid_request', operation); } @@ -311,6 +319,7 @@ export class VercelSandboxRestClient { [VERCEL_CLOUD_AGENT_CREATE_OPERATION_TAG]: input.operationId, [VERCEL_CLOUD_AGENT_RUNTIME_BUILD_TAG]: input.runtimeBuildId, }, + ...(input.resources === undefined ? {} : { resources: input.resources }), ...(input.networkPolicy === undefined ? {} : { networkPolicy: input.networkPolicy }), }), }); @@ -329,7 +338,10 @@ export class VercelSandboxRestClient { requireIdentifier(input.operationId, operation); requireIdentifier(input.runtimeBuildId, operation); requireIdentifier(input.snapshotId, operation); - if (!runtimeSchema.safeParse(input.runtime).success) { + if ( + !runtimeSchema.safeParse(input.runtime).success || + !vercelSandboxResourcesSchema.optional().safeParse(input.resources).success + ) { throw new VercelSandboxRestError('invalid_request', operation); } const url = this.namedSandboxUrl(input.name, projectId); @@ -557,6 +569,9 @@ export class VercelSandboxRestClient { envelope.session.projectId !== this.config.projectId || envelope.session.sourceSnapshotId !== input.snapshotId || envelope.session.runtime !== input.runtime || + (input.resources !== undefined && + (envelope.session.vcpus !== input.resources.vcpus || + envelope.session.memory !== input.resources.memory)) || tags[VERCEL_CLOUD_AGENT_RESOURCE_TAG] !== VERCEL_CLOUD_AGENT_RESOURCE_TAG_VALUE || tags[VERCEL_CLOUD_AGENT_CREATE_OPERATION_TAG] !== input.operationId || tags[VERCEL_CLOUD_AGENT_RUNTIME_BUILD_TAG] !== input.runtimeBuildId diff --git a/services/cloud-agent-next/src/callbacks/queue-config.test.ts b/services/cloud-agent-next/src/callbacks/queue-config.test.ts index 8277db3ac3..7da96d8738 100644 --- a/services/cloud-agent-next/src/callbacks/queue-config.test.ts +++ b/services/cloud-agent-next/src/callbacks/queue-config.test.ts @@ -37,14 +37,16 @@ describe('callback queue retry configuration', () => { expect(dev?.max_retries).toBe(CONFIGURED_REDELIVERIES); }); - it('omits production control-plane enrollment from wrangler and enables it on wrangler dev', () => { + it('omits production enrollment from wrangler and enables control-plane, worktree, and sandbox selection on wrangler dev', () => { const config = readWranglerConfig(); expect(config.vars?.VERCEL_SANDBOX_ORG_IDS).toBe(''); expect(config.vars?.CONTROL_PLANE_IDS).toBeUndefined(); expect(config.vars?.WORKTREE_CREATION_ENABLED_IDS).toBeUndefined(); + expect(config.vars?.SANDBOX_SELECTION_IDS).toBeUndefined(); expect(config.env?.dev?.vars?.VERCEL_SANDBOX_ORG_IDS).toBe(''); expect(config.env?.dev?.vars?.CONTROL_PLANE_IDS).toBe('*'); expect(config.env?.dev?.vars?.WORKTREE_CREATION_ENABLED_IDS).toBe('*'); + expect(config.env?.dev?.vars?.SANDBOX_SELECTION_IDS).toBe('*'); }); }); diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index 12efb305aa..f62126ef8d 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -18,6 +18,10 @@ import { } from '../sandbox-control/worktree-deletion.js'; import { getSandbox } from '@cloudflare/sandbox'; import { DEFAULT_DO_RETRY_CONFIG, withTimeout } from '@kilocode/worker-utils'; +import { + getSandboxAllocationResources, + type VercelSandboxResources, +} from '@kilocode/worker-utils/sandbox-allocation'; import { z } from 'zod'; import type { Env } from '../types.js'; import { resolveSecret } from '../auth.js'; @@ -99,6 +103,8 @@ import { recordStopAttempt, sameAllocation, getWorktreeCredentialContainment, + sandboxProviderConfigurationSchema, + type SandboxProviderConfiguration, WORKTREE_CREDENTIAL_CONTAINMENT, type CredentialContainmentRequirements, type ObserveResult, @@ -244,6 +250,7 @@ const ACTIVE_WRAPPER_RUNTIME_KEY = 'active_wrapper_runtime'; const DIAGNOSTIC_BUNDLE_KEY = 'diagnostic_bundle'; const PROVIDER_KIND_KEY = 'provider_kind'; const PROVIDER_LOCATOR_KEY = 'provider_locator'; +const PROVIDER_CONFIGURATION_KEY = 'provider_configuration'; const BILLING_INPUT_KEY = 'billing_input'; const ACQUISITION_RECEIPTS_KEY = 'acquisition_receipts'; const CREDENTIAL_POLICY_DIRTY_KEY = 'credential_policy_dirty'; @@ -330,6 +337,7 @@ export class SandboxControl extends DurableObject { private activeConnection: SandboxControlConnectionIdentity | null = null; private readyConnectionId: string | null = null; private providerKind: AgentSandboxProvider = 'cloudflare'; + private vercelResources: VercelSandboxResources | undefined; private readonly sessionForwarding = createSessionForwarding(); private readonly forwarding = { enqueued: 0, @@ -449,17 +457,19 @@ export class SandboxControl extends DurableObject { private ensureOperationalInitialized(): Promise { return (this.operationalInitialization ??= this.ctx.blockConcurrencyWhile(async () => { const ctx = this.ctx; - const [readyAt, runtime, kind, physical, recovery] = await Promise.all([ + const [readyAt, runtime, configuration, physical, recovery] = await Promise.all([ ctx.storage.get(WRAPPER_READY_AT_KEY), ctx.storage.get(ACTIVE_WRAPPER_RUNTIME_KEY), - ctx.storage.get(PROVIDER_KIND_KEY), + this.readProviderConfiguration(), loadPhysicalRecord(ctx.storage), loadRecoveryDecisions(ctx.storage), ]); this.vercelLocator = vercelProviderLocatorSchema .optional() .parse(await ctx.storage.get(PROVIDER_LOCATOR_KEY)); - this.providerKind = kind ?? 'cloudflare'; + this.providerKind = configuration?.provider ?? 'cloudflare'; + this.vercelResources = + configuration?.provider === 'vercel' ? configuration.resources : undefined; this.provider = this.createProviderAdapter(this.providerKind, physical); this.runtimeDeleted = (await ctx.storage.get(RUNTIME_DELETED_KEY)) === true; this.exclusiveDeletionWorktreeId = cloudAgentWorktreeIdSchema @@ -1344,7 +1354,9 @@ export class SandboxControl extends DurableObject { } const metadata = await this.readCredentialMetadata(input); const provider = getSandboxProvider(metadata); - await this.pinProvider(provider); + await this.pinProvider(provider, { + resources: getSandboxAllocationResources(metadata.workspace?.sandboxAllocation), + }); const physical = await loadPhysicalRecord(this.ctx.storage); const requiredContainment = getWorktreeCredentialContainment( requiresContainmentSandbox(metadata) @@ -1512,6 +1524,7 @@ export class SandboxControl extends DurableObject { ownerId: string; sessionId: string; provider?: AgentSandboxProvider; + resources?: VercelSandboxResources; allowCreate?: boolean; acquisition?: SandboxAcquisition; billing?: SandboxBillingInput; @@ -1553,7 +1566,7 @@ export class SandboxControl extends DurableObject { await this.ctx.storage.delete(RUNTIME_DELETED_KEY); this.runtimeDeleted = false; } - await this.pinProvider(input.provider); + await this.pinProvider(input.provider, { resources: input.resources }); if (acquisition && this.providerKind !== 'cloudflare') { throw new Error('Sandbox acquisition is only supported for Cloudflare'); } @@ -2606,7 +2619,13 @@ export class SandboxControl extends DurableObject { const { projectId, snapshotId, runtimeBuildId, runtime } = vercel; next.createIntent = { ...next.createIntent, - vercel: { projectId, snapshotId, runtimeBuildId, runtime }, + vercel: { + projectId, + snapshotId, + runtimeBuildId, + runtime, + ...(this.vercelResources === undefined ? {} : { resources: this.vercelResources }), + }, }; this.vercelLocator = vercelProviderLocatorSchema.parse({ teamId: vercel.teamId, @@ -2863,12 +2882,14 @@ export class SandboxControl extends DurableObject { ACTIVE_WRAPPER_RUNTIME_KEY, DIAGNOSTIC_BUNDLE_KEY, PROVIDER_KIND_KEY, + PROVIDER_CONFIGURATION_KEY, BILLING_INPUT_KEY, CREDENTIAL_POLICY_DIRTY_KEY, PROVIDER_LOCATOR_KEY, ...(options?.preserveAcquisitionReceipts ? [] : [ACQUISITION_RECEIPTS_KEY]), ]); this.vercelLocator = undefined; + this.vercelResources = undefined; this.activeConnection = null; this.readyConnectionId = null; this.kiloReady = false; @@ -3019,20 +3040,67 @@ export class SandboxControl extends DurableObject { }); } - private async pinProvider(requested?: AgentSandboxProvider): Promise { - const stored = await this.ctx.storage.get(PROVIDER_KIND_KEY); - const kind = stored ?? requested ?? 'cloudflare'; - if (stored !== undefined && requested !== undefined && stored !== requested) { + private async readProviderConfiguration(): Promise { + const [raw, legacyKind] = await Promise.all([ + this.ctx.storage.get(PROVIDER_CONFIGURATION_KEY), + this.ctx.storage.get(PROVIDER_KIND_KEY), + ]); + const legacy = + legacyKind === undefined + ? undefined + : sandboxProviderConfigurationSchema.parse({ provider: legacyKind }); + if (raw === undefined) return legacy; + const configuration = sandboxProviderConfigurationSchema.parse(raw); + if (legacy && legacy.provider !== configuration.provider) { throw new Error('Sandbox provider mismatch'); } - if (kind === 'vercel' && parseVercelSandboxRuntimeConfig(this.env) === undefined) { - throw new Error('Vercel sandbox runtime configuration is unavailable'); - } - if (stored === undefined) { - await this.ctx.storage.put(PROVIDER_KIND_KEY, kind); - } - this.providerKind = kind; - this.provider = this.createProviderAdapter(kind, await loadPhysicalRecord(this.ctx.storage)); + return configuration; + } + + private async pinProvider( + requested?: AgentSandboxProvider, + allocation?: { resources?: VercelSandboxResources } + ): Promise { + const configuration = await this.ctx.storage.transaction(async () => { + const stored = await this.readProviderConfiguration(); + const resources = + allocation !== undefined + ? allocation.resources + : stored?.provider === 'vercel' + ? stored.resources + : undefined; + const provider = requested ?? stored?.provider ?? 'cloudflare'; + if (stored && stored.provider !== provider) { + throw new Error('Sandbox provider mismatch'); + } + const next = sandboxProviderConfigurationSchema.parse({ + provider, + ...(resources === undefined ? {} : { resources }), + }); + if ( + stored?.provider === 'vercel' && + next.provider === 'vercel' && + (stored.resources?.vcpus !== next.resources?.vcpus || + stored.resources?.memory !== next.resources?.memory) + ) { + throw new Error('Sandbox resources mismatch'); + } + if (next.provider === 'vercel' && parseVercelSandboxRuntimeConfig(this.env) === undefined) { + throw new Error('Vercel sandbox runtime configuration is unavailable'); + } + await this.ctx.storage.put({ + [PROVIDER_KIND_KEY]: next.provider, + [PROVIDER_CONFIGURATION_KEY]: next, + }); + return next; + }); + this.providerKind = configuration.provider; + this.vercelResources = + configuration.provider === 'vercel' ? configuration.resources : undefined; + this.provider = this.createProviderAdapter( + configuration.provider, + await loadPhysicalRecord(this.ctx.storage) + ); } private async billingInput( diff --git a/services/cloud-agent-next/src/persistence/session-metadata.test.ts b/services/cloud-agent-next/src/persistence/session-metadata.test.ts index e9c5036a3b..816b6b8789 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.test.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.test.ts @@ -432,6 +432,80 @@ describe('session metadata boundary', () => { ).toThrow(); }); + it.each(['cloudflare-single', 'cloudflare-shared', 'vercel-small', 'vercel-large'] as const)( + 'round-trips the immutable %s preset and rejects inconsistent allocations', + sandboxAllocation => { + const sandboxId = `${sandboxAllocation === 'cloudflare-shared' ? 'org' : 'ses'}-${'a'.repeat(48)}`; + const current = { + metadataSchemaVersion: 2, + identity: { sessionId: 'workspace_preset', userId: 'oauth/user', orgId: 'org-id' }, + auth: {}, + workspace: { + sandboxId, + sandboxAllocation, + sandboxProvider: sandboxAllocation.startsWith('vercel-') ? 'vercel' : 'cloudflare', + ...(sandboxAllocation === 'cloudflare-shared' + ? { sandboxRoute: { kind: 'shared', routeKey: sandboxId } } + : {}), + }, + lifecycle: { version: 1, timestamp: 1 }, + }; + expect(serializeSessionMetadata(parseSessionMetadata(current))).toEqual(current); + const vercel = sandboxAllocation.startsWith('vercel-'); + for (const workspace of [ + { ...current.workspace, sandboxId: `dind-${'a'.repeat(48)}` }, + { + ...current.workspace, + sandboxProvider: current.workspace.sandboxProvider === 'vercel' ? 'cloudflare' : 'vercel', + }, + { ...current.workspace, devcontainerRequested: true }, + // `istd-` is the only identity Isolated Standard accepts. + { ...current.workspace, sandboxAllocation: 'isolated-standard' }, + { ...current.workspace, sandboxAllocation: 'custom' }, + ...(vercel ? [{ ...current.workspace, sandboxProvider: undefined }] : []), + ]) { + expect(() => parseSessionMetadata({ ...current, workspace })).toThrow(); + } + if (!vercel) { + // Metadata written before the explicit provider field defaults to Cloudflare. + const implicit = { + ...current, + workspace: { ...current.workspace, sandboxProvider: undefined }, + }; + expect(parseSessionMetadata(implicit).workspace?.sandboxAllocation).toBe(sandboxAllocation); + } + for (const identity of [{ ...current.identity, billingOrigin: 'code-review' }]) { + expect(() => parseSessionMetadata({ ...current, identity })).toThrow(); + } + // Selectable allocations are authorized per owner (user or organization) at + // creation, so a personal session without an org is valid metadata. + expect( + parseSessionMetadata({ + ...current, + identity: { ...current.identity, orgId: undefined }, + }).workspace?.sandboxAllocation + ).toBe(sandboxAllocation); + // Cloudflare presets keep the owner's plane; Vercel presets exist only on control. + const legacy = { ...current, identity: { ...current.identity, sessionId: 'agent_legacy' } }; + if (sandboxAllocation.startsWith('vercel-')) { + expect(() => parseSessionMetadata(legacy)).toThrow('control-plane session'); + } else { + expect(serializeSessionMetadata(parseSessionMetadata(legacy))).toEqual(legacy); + } + expect(() => + parseSessionMetadata({ + ...current, + devcontainer: { + workspacePath: '/repo', + innerWorkspaceFolder: '/repo', + wrapperPort: 3000, + configPath: '/repo/devcontainer.json', + }, + }) + ).toThrow(); + } + ); + it('accepts legacy current metadata without an explicit sandbox provider as Cloudflare', () => { const current = { metadataSchemaVersion: 2 as const, diff --git a/services/cloud-agent-next/src/persistence/session-metadata.ts b/services/cloud-agent-next/src/persistence/session-metadata.ts index 63f1f55353..7f3eedf8c7 100644 --- a/services/cloud-agent-next/src/persistence/session-metadata.ts +++ b/services/cloud-agent-next/src/persistence/session-metadata.ts @@ -3,9 +3,20 @@ import { cloudAgentWorktreeIdSchema, sessionIdSchema as kiloSessionIdSchema, } from '@kilocode/session-ingest-contracts'; +import { + getSandboxAllocationProvider, + sandboxAllocationRequiresControlPlane, + sandboxAllocationSchema, + type SandboxAllocation, +} from '@kilocode/worker-utils/sandbox-allocation'; import { PROVIDER_CAPABILITIES } from '../agent-sandbox/capabilities.js'; -import { isGeneratedSharedSandboxId, isValidSandboxId } from '../sandbox-id.js'; +import { + classifySandboxId, + isGeneratedSharedSandboxId, + isValidSandboxId, + type SandboxIdClass, +} from '../sandbox-id.js'; import { sessionPlaneFromId } from '../session-plane.js'; import { SHARED_SANDBOX_FAILOVER_SUFFIX } from '../shared-sandbox-route.js'; import { MESSAGE_ID_FORMAT_DESCRIPTION, MESSAGE_ID_PATTERN } from '../session/message-id.js'; @@ -233,11 +244,23 @@ const CredentialContainmentSchema = z }) .strip(); +/** Sandbox-ID class each isolated allocation must have in persisted metadata. */ +const SANDBOX_ALLOCATION_ID_CLASS: Record< + Exclude, + SandboxIdClass +> = { + 'isolated-standard': 'isolated-standard', + 'cloudflare-single': 'isolated-small', + 'vercel-small': 'isolated-small', + 'vercel-large': 'isolated-small', +}; + const MetadataWorkspaceSchema = z .object({ sandboxId: SandboxIdSchema.optional(), sandboxRoute: MetadataSharedSandboxRouteSchema.optional(), sandboxProvider: SandboxProviderSchema.optional(), + sandboxAllocation: sandboxAllocationSchema.optional(), providerRuntime: ProviderRuntimeSchema.optional(), worktreeId: cloudAgentWorktreeIdSchema.optional(), workspacePath: z.string().optional(), @@ -247,10 +270,28 @@ const MetadataWorkspaceSchema = z credentialContainment: CredentialContainmentSchema.optional(), managedScmContainment: z.boolean().optional(), devcontainerRequested: z.boolean().optional(), - sandboxAllocation: z.literal('isolated-standard').optional(), }) .strip() .superRefine((workspace, context) => { + const allocation = workspace.sandboxAllocation; + if (allocation !== undefined) { + const shared = allocation === 'cloudflare-shared'; + if ( + // Metadata written before an explicit provider defaults to Cloudflare. + (workspace.sandboxProvider ?? 'cloudflare') !== getSandboxAllocationProvider(allocation) || + !workspace.sandboxId || + (shared + ? !isGeneratedSharedSandboxId(workspace.sandboxId) || !workspace.sandboxRoute + : classifySandboxId(workspace.sandboxId) !== SANDBOX_ALLOCATION_ID_CLASS[allocation]) || + workspace.devcontainerRequested === true + ) { + context.addIssue({ + code: 'custom', + path: ['sandboxAllocation'], + message: 'Sandbox allocation conflicts with workspace identity', + }); + } + } const route = workspace.sandboxRoute; if (!route) return; const sandboxId = workspace.sandboxId; @@ -346,6 +387,27 @@ export const CurrentSessionMetadataSchema = z PROVIDER_CAPABILITIES[metadata.workspace?.sandboxProvider ?? 'cloudflare'].devcontainer || !metadata.devcontainer, 'Sandbox provider metadata cannot contain a devcontainer runtime' + ) + .refine( + metadata => + metadata.workspace?.sandboxAllocation === undefined || + (!metadata.devcontainer && + metadata.identity.billingOrigin !== 'code-review' && + metadata.identity.createdOnPlatform !== 'code-review'), + 'Sandbox allocations cannot be combined with specialized routing' + ) + .refine( + metadata => + // `isolated-standard` remains legacy-plane only; Vercel is control-plane only. + metadata.workspace?.sandboxAllocation !== 'isolated-standard' || + sessionPlaneFromId(metadata.identity.sessionId) === 'legacy', + 'Isolated Standard allocation is not supported for control-plane sessions' + ) + .refine( + metadata => + !sandboxAllocationRequiresControlPlane(metadata.workspace?.sandboxAllocation) || + sessionPlaneFromId(metadata.identity.sessionId) === 'control', + 'Vercel sandbox allocations require a control-plane session' ); export type SessionMetadata = z.infer; diff --git a/services/cloud-agent-next/src/router.ts b/services/cloud-agent-next/src/router.ts index f0a512863d..301d69a16c 100644 --- a/services/cloud-agent-next/src/router.ts +++ b/services/cloud-agent-next/src/router.ts @@ -15,9 +15,11 @@ import { createSessionSendHandlers } from './router/handlers/session-send.js'; import { createSessionWorktreeHandlers } from './router/handlers/session-worktree.js'; import { deleteWorktree } from './router/handlers/worktree-deletion.js'; import { createSessionWorktreeChangesHandlers } from './router/handlers/session-worktree-changes.js'; +import { getSandboxSelectionOptions } from './router/handlers/sandbox-selection.js'; export const appRouter = router({ deleteWorktree, + getSandboxSelectionOptions, ...createSessionManagementHandlers(), ...createSessionPrepareHandlers(), ...createSessionExecutionV2Handlers(), diff --git a/services/cloud-agent-next/src/router/handlers/sandbox-selection.ts b/services/cloud-agent-next/src/router/handlers/sandbox-selection.ts new file mode 100644 index 0000000000..8206892e77 --- /dev/null +++ b/services/cloud-agent-next/src/router/handlers/sandbox-selection.ts @@ -0,0 +1,31 @@ +import { z } from 'zod'; +import { sandboxSelectionCapabilitiesSchema } from '@kilocode/worker-utils/sandbox-allocation'; +import { getPgDb } from '../../db/pg.js'; +import { getSandboxSelectionCapabilities } from '../../sandbox-selection.js'; +import { protectedProcedure } from '../auth.js'; +import { assertOrganizationMembership } from './organization-membership.js'; + +export const getSandboxSelectionOptions = protectedProcedure + .input( + z + .object({ + kilocodeOrganizationId: z.string().uuid().optional(), + devcontainer: z.boolean().optional(), + }) + .strict() + ) + .output(sandboxSelectionCapabilitiesSchema) + .query(async ({ input, ctx }) => { + if (input.kilocodeOrganizationId) { + await assertOrganizationMembership( + getPgDb(ctx.env), + ctx.userId, + input.kilocodeOrganizationId + ); + } + return getSandboxSelectionCapabilities( + ctx.env, + { userId: ctx.userId, orgId: input.kilocodeOrganizationId }, + input.devcontainer + ); + }); diff --git a/services/cloud-agent-next/src/router/handlers/session-start.ts b/services/cloud-agent-next/src/router/handlers/session-start.ts index 7c2ffd1200..054f3f44aa 100644 --- a/services/cloud-agent-next/src/router/handlers/session-start.ts +++ b/services/cloud-agent-next/src/router/handlers/session-start.ts @@ -79,6 +79,7 @@ function startInputToSessionCreateRequest( }, agent: input.agent, repository, + ...(input.runtime ? { runtime: input.runtime } : {}), profile: profile ? { id: profile.id, diff --git a/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts b/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts index 235c939912..a062fd7513 100644 --- a/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts +++ b/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts @@ -2,9 +2,13 @@ import { TRPCError } from '@trpc/server'; import type { WorkerDb } from '@kilocode/db/client'; import type { OperationLedgerRow } from '@kilocode/db/schema'; import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { + SELECTABLE_SANDBOX_ALLOCATIONS, + type SandboxAllocation, +} from '@kilocode/worker-utils/sandbox-allocation'; import { t } from '../auth.js'; -import type { SessionMetadata } from '../../persistence/session-metadata.js'; +import { parseSessionMetadata, type SessionMetadata } from '../../persistence/session-metadata.js'; import type * as SessionPlane from '../../session-plane.js'; import type { TRPCContext } from '../../types.js'; import { sha256Hex } from '../../utils/sha256.js'; @@ -161,6 +165,33 @@ function sourceMetadata(options?: { }; } +function sourceMetadataWithPreset(sandboxAllocation: SandboxAllocation): SessionMetadata { + const metadata = sourceMetadata({ organizationId: ORGANIZATION_ID }); + const sandboxId = + sandboxAllocation === 'cloudflare-shared' + ? 'org-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa' + : 'ses-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'; + const sandboxProvider = sandboxAllocation.startsWith('vercel-') ? 'vercel' : 'cloudflare'; + return { + ...metadata, + workspace: { + ...metadata.workspace, + sandboxId, + sandboxProvider, + sandboxAllocation, + sandboxRoute: + sandboxAllocation === 'cloudflare-shared' + ? { kind: 'shared', routeKey: sandboxId } + : undefined, + providerRuntime: + sandboxProvider === 'vercel' + ? { provider: 'vercel', sessionId: 'vercel-runtime' } + : undefined, + credentialContainment: { github: false, gitlab: false, bitbucket: false, kilocode: false }, + }, + }; +} + function ownershipRow(overrides: Partial = {}): OwnershipFixture { return { kiloSessionId: SOURCE_KILO_SESSION_ID, @@ -538,11 +569,166 @@ describe('createWorktreeChat request validation and authorization', () => { expect(admitOperationMock).not.toHaveBeenCalled(); }); - it('rejects a grouped owner that is no longer enrolled in the control plane', async () => { + it('creates a sibling for a grouped owner no longer enrolled in the control plane', async () => { const { caller, input } = fixture({ controlPlaneIds: 'another-owner' }); - await expect(caller.createWorktreeChat(input)).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + await expect(caller.createWorktreeChat(input)).resolves.toEqual({ + cloudAgentSessionId: DESTINATION_WORKSPACE_ID, + kiloSessionId: DESTINATION_KILO_SESSION_ID, + worktreeId: WORKTREE_ID, + }); + expect(generateSessionIdMock).toHaveBeenCalledWith('control'); + }); +}); + +describe('createWorktreeChat sandbox preset inheritance', () => { + it.each(SELECTABLE_SANDBOX_ALLOCATIONS)( + 'inherits %s after control-plane and selection rollouts are disabled', + async sandboxAllocation => { + const metadata = sourceMetadataWithPreset(sandboxAllocation); + const { caller, context, input, destinationStub, sandboxControlNamespace } = fixture({ + metadata, + organizationId: ORGANIZATION_ID, + }); + context.env.SANDBOX_SELECTION_IDS = ''; + context.env.CONTROL_PLANE_IDS = ''; + context.env.PER_SESSION_SANDBOX_ORG_IDS = + sandboxAllocation === 'cloudflare-shared' ? '*' : ''; + context.env.VERCEL_SANDBOX_ORG_IDS = sandboxAllocation.startsWith('vercel-') ? '' : '*'; + context.env.CREDENTIAL_CONTAINMENT_ENABLED = 'true'; + await caller.createWorktreeChat(input); + const registration = destinationStub.registerSession.mock.calls[0]?.[0]; + const registered = parseSessionMetadata({ + ...registration, + metadataSchemaVersion: 2, + lifecycle: { version: 1, timestamp: 1 }, + }); + expect(registered.identity.sessionId).toBe(DESTINATION_WORKSPACE_ID); + expect(registered.workspace).toEqual({ ...metadata.workspace, providerRuntime: undefined }); + expect(registered.workspace?.worktreeId).toBe(WORKTREE_ID); + expect(registered.workspace?.sandboxAllocation).toBe(sandboxAllocation); + expect(registered.workspace).not.toHaveProperty('resources'); + expect(recordOperationProgressMock).toHaveBeenCalledWith( + expect.anything(), + LEDGER_ROW_ID, + expect.objectContaining({ sandboxAllocation }) + ); + expect(sandboxControlNamespace.get).not.toHaveBeenCalled(); + } + ); + + it.each([ + { sandboxAllocation: 'vercel-large' }, + { runtime: { sandboxAllocation: 'vercel-large' } }, + { workspace: { sandboxAllocation: 'vercel-large' } }, + { sandboxProvider: 'vercel' }, + { resources: { vcpus: 4, memory: 8192 } }, + ])('rejects direct sibling allocation overrides before side effects: %j', async override => { + const { caller, input, destinationStub } = fixture({ + metadata: sourceMetadataWithPreset('cloudflare-single'), + organizationId: ORGANIZATION_ID, + }); + await expect(caller.createWorktreeChat({ ...input, ...override })).rejects.toMatchObject({ + code: 'BAD_REQUEST', + }); + expect(assertOrganizationMembershipMock).not.toHaveBeenCalled(); expect(admitOperationMock).not.toHaveBeenCalled(); + expect(createSessionForCloudAgentMock).not.toHaveBeenCalled(); + expect(destinationStub.registerSession).not.toHaveBeenCalled(); + }); + + it.each(['vercel-large', undefined] as const)( + 'rejects same-key replay if the source preset changes to %s', + async sandboxAllocation => { + const metadata = sourceMetadataWithPreset('vercel-small'); + const source = ownershipRow({ organizationId: ORGANIZATION_ID }); + const { caller, input, destinationStub } = fixture({ + metadata, + organizationId: ORGANIZATION_ID, + ownershipResults: [[source], [source]], + }); + await caller.createWorktreeChat(input); + const progress = recordOperationProgressMock.mock.calls[0]?.[2] as Record; + metadata.workspace = { ...metadata.workspace, sandboxAllocation }; + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_settled', + row: ledgerRow({ + organization_id: ORGANIZATION_ID, + status: 'completed', + canonical_result: progress, + }), + }); + await expect(caller.createWorktreeChat(input)).rejects.toMatchObject({ + code: 'CONFLICT', + message: 'operation_key_reuse_mismatch', + }); + expect(createSessionForCloudAgentMock).toHaveBeenCalledTimes(1); + expect(destinationStub.registerSession).toHaveBeenCalledTimes(1); + expect(destinationStub.getMetadata).not.toHaveBeenCalled(); + } + ); + + it('rejects recovered sibling metadata with a different size on the same provider', async () => { + const metadata = sourceMetadataWithPreset('vercel-small'); + const source = ownershipRow({ organizationId: ORGANIZATION_ID }); + const { caller, input, destinationStub } = fixture({ + metadata, + organizationId: ORGANIZATION_ID, + ownershipResults: [[source], [source]], + }); + destinationStub.registerSession.mockRejectedValueOnce(new Error('registration response lost')); + await expect(caller.createWorktreeChat(input)).rejects.toThrow('registration response lost'); + const progress = recordOperationProgressMock.mock.calls[0]?.[2] as Record; + const registered = destinationMetadata(metadata); + registered.workspace = { ...registered.workspace, sandboxAllocation: 'vercel-large' }; + destinationStub.getMetadata.mockResolvedValueOnce(registered); + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_reconcile_pending', + row: ledgerRow({ + organization_id: ORGANIZATION_ID, + status: 'reconcile_pending', + canonical_result: progress, + }), + }); + await expect(caller.createWorktreeChat(input)).rejects.toMatchObject({ + code: 'CONFLICT', + message: 'operation_key_reuse_mismatch', + }); + expect(createSessionForCloudAgentMock).toHaveBeenCalledTimes(1); + expect(destinationStub.registerSession).toHaveBeenCalledTimes(1); + expect(settleOperationMock).not.toHaveBeenCalled(); + }); + + it('replays the inherited preset after rollouts are disabled without re-registering', async () => { + const metadata = sourceMetadataWithPreset('vercel-large'); + const source = ownershipRow({ organizationId: ORGANIZATION_ID }); + const { caller, context, input, destinationStub } = fixture({ + metadata, + organizationId: ORGANIZATION_ID, + ownershipResults: [[source], [source]], + }); + await caller.createWorktreeChat(input); + const progress = recordOperationProgressMock.mock.calls[0]?.[2] as Record; + context.env.SANDBOX_SELECTION_IDS = ''; + context.env.CONTROL_PLANE_IDS = ''; + context.env.VERCEL_TOKEN = ''; + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_settled', + row: ledgerRow({ + organization_id: ORGANIZATION_ID, + status: 'completed', + canonical_result: progress, + }), + }); + await expect(caller.createWorktreeChat(input)).resolves.toMatchObject({ + cloudAgentSessionId: DESTINATION_WORKSPACE_ID, + kiloSessionId: DESTINATION_KILO_SESSION_ID, + worktreeId: WORKTREE_ID, + replayed: true, + }); + expect(createSessionForCloudAgentMock).toHaveBeenCalledTimes(1); + expect(destinationStub.registerSession).toHaveBeenCalledTimes(1); + expect(assertOrganizationMembershipMock).toHaveBeenCalledTimes(2); }); }); diff --git a/services/cloud-agent-next/src/router/handlers/session-worktree.ts b/services/cloud-agent-next/src/router/handlers/session-worktree.ts index 3ecfeb9b52..19265e1f7e 100644 --- a/services/cloud-agent-next/src/router/handlers/session-worktree.ts +++ b/services/cloud-agent-next/src/router/handlers/session-worktree.ts @@ -18,6 +18,7 @@ import { sealRuntimeAuthorization, } from '@kilocode/worker-utils/runtime-authorization'; import { verifyKiloTokenForPolicy } from '@kilocode/worker-utils/kilo-token-policy'; +import { sandboxAllocationSchema } from '@kilocode/worker-utils/sandbox-allocation'; import { and, eq } from 'drizzle-orm'; import { z } from 'zod'; @@ -28,7 +29,7 @@ import { } from '../../persistence/session-metadata.js'; import { logControlDiagnostic } from '../../sandbox-control/diagnostics.js'; import { getSandboxSessionStub } from '../../sandbox-session/session-stub.js'; -import { generateSessionId, isControlPlaneOwner } from '../../session-plane.js'; +import { generateSessionId } from '../../session-plane.js'; import { assertSessionOperationIdentity, assertRuntimeIsolationAdmission, @@ -81,6 +82,7 @@ const ownershipRowSchema = z const operationProgressSchema = CreateWorktreeChatOutput.omit({ replayed: true }) .extend({ [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: z.string().regex(/^[a-f0-9]{64}$/), + sandboxAllocation: sandboxAllocationSchema.optional(), }) .strict(); @@ -206,11 +208,7 @@ async function loadWorktreeSource( !worktreeId || ownership.parentSessionId !== null || ownership.cloudAgentSessionScopeId !== input.sourceCloudAgentSessionId || - ownership.createdOnPlatform !== 'cloud-agent-web' || - !isControlPlaneOwner(ctx.env, { - userId: ctx.userId, - orgId: input.kilocodeOrganizationId, - }) + ownership.createdOnPlatform !== 'cloud-agent-web' ) { throw sourceRejected(); } @@ -265,15 +263,18 @@ async function loadWorktreeSource( async function worktreeIntentFingerprint( input: WorktreeInput, - worktreeId: CloudAgentWorktreeId + source: WorktreeSource ): Promise { return sha256Hex( JSON.stringify({ sourceKiloSessionId: input.sourceKiloSessionId, sourceCloudAgentSessionId: input.sourceCloudAgentSessionId, organizationId: input.kilocodeOrganizationId ?? null, - worktreeId, + worktreeId: source.worktreeId, clientProvenance: input.clientProvenance, + ...(source.workspace.sandboxAllocation + ? { sandboxAllocation: source.workspace.sandboxAllocation } + : {}), }) ); } @@ -288,6 +289,7 @@ function readOperationProgress( if ( !progress.success || progress.data.worktreeId !== source.worktreeId || + progress.data.sandboxAllocation !== source.workspace.sandboxAllocation || progress.data[SESSION_CREATE_INTENT_FINGERPRINT_KEY] !== fingerprint ) { throw operationConflict(); @@ -444,6 +446,7 @@ function assertRegisteredMetadata( workspace.workspacePath !== source.workspace.workspacePath || workspace.sandboxId !== source.workspace.sandboxId || workspace.sandboxProvider !== source.workspace.sandboxProvider || + workspace.sandboxAllocation !== source.workspace.sandboxAllocation || workspace.branchName !== sourceWorktreeBranchName(source) || JSON.stringify(workspace.sandboxRoute) !== JSON.stringify(source.workspace.sandboxRoute) || !metadata.repository || @@ -812,6 +815,9 @@ async function executeWorktreeCreate( kiloSessionId: generateKiloSessionId(), worktreeId: source.worktreeId, [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: fingerprint, + ...(source.workspace.sandboxAllocation + ? { sandboxAllocation: source.workspace.sandboxAllocation } + : {}), }); const diagnostic = { operationRowId: row.id, @@ -982,7 +988,7 @@ const createWorktreeChatHandler = internalApiProtectedProcedure const startedAt = Date.now(); const db = getPgDb(ctx.env); const source = await loadWorktreeSource(db, ctx, input); - const fingerprint = await worktreeIntentFingerprint(input, source.worktreeId); + const fingerprint = await worktreeIntentFingerprint(input, source); const admission = await admitOperation(db, { userId: ctx.userId, orgId: input.kilocodeOrganizationId, diff --git a/services/cloud-agent-next/src/router/schemas.test.ts b/services/cloud-agent-next/src/router/schemas.test.ts index 462cae2fe0..e4e98d8b39 100644 --- a/services/cloud-agent-next/src/router/schemas.test.ts +++ b/services/cloud-agent-next/src/router/schemas.test.ts @@ -1,4 +1,8 @@ import { describe, expect, it } from 'vitest'; +import { + getSandboxAllocationRequest, + sandboxAllocationSchema, +} from '@kilocode/worker-utils/sandbox-allocation'; import { ExecutionResponse, GetMessageResultInput, @@ -225,6 +229,91 @@ describe('grouped unified session input contracts', () => { }); }); +describe('sandbox preset input boundaries', () => { + const prepare = { ...basePromptInput, githubRepo: 'acme/repo' }; + + it.each(sandboxAllocationSchema.options)( + 'normalizes structured %s requests to the existing canonical allocation', + allocation => { + const sandboxAllocation = getSandboxAllocationRequest(allocation); + expect(PrepareSessionInput.parse({ ...prepare, sandboxAllocation })).toEqual( + PrepareSessionInput.parse({ ...prepare, sandboxAllocation: allocation }) + ); + expect( + StartSessionInput.parse({ ...baseStartInput, runtime: { sandboxAllocation } }) + ).toEqual( + StartSessionInput.parse({ ...baseStartInput, runtime: { sandboxAllocation: allocation } }) + ); + } + ); + + it.each(['small', 'large'] as const)( + 'rejects unsupported BYOC %s rather than allocating Kilo compute', + instanceType => { + const sandboxAllocation = { provider: { id: 'vercel', account: 'byoc' }, instanceType }; + for (const result of [ + PrepareSessionInput.safeParse({ ...prepare, sandboxAllocation }), + StartSessionInput.safeParse({ ...baseStartInput, runtime: { sandboxAllocation } }), + ]) { + expect(result.success).toBe(false); + if (!result.success) expect(result.error.message).toContain('BYOC'); + } + } + ); + + it('keeps omitted allocations omitted', () => { + expect(PrepareSessionInput.parse(prepare).sandboxAllocation).toBeUndefined(); + expect(StartSessionInput.parse(baseStartInput).runtime).toBeUndefined(); + }); + + it.each(['devcontainer', 'code-review'] as const)( + 'rejects structured allocation combined with %s routing', + routing => { + const sandboxAllocation = getSandboxAllocationRequest('vercel-small'); + expect( + PrepareSessionInput.safeParse({ + ...prepare, + sandboxAllocation, + ...(routing === 'devcontainer' ? { devcontainer: true } : { createdOnPlatform: routing }), + }).success + ).toBe(false); + } + ); + + it.each(['cloudflare-single', 'cloudflare-shared', 'vercel-small', 'vercel-large'] as const)( + 'preserves the finite %s preset in both prepare and grouped start', + sandboxAllocation => { + expect(PrepareSessionInput.parse({ ...prepare, sandboxAllocation }).sandboxAllocation).toBe( + sandboxAllocation + ); + expect( + StartSessionInput.parse({ ...baseStartInput, runtime: { sandboxAllocation } }).runtime + ).toEqual({ sandboxAllocation }); + } + ); + + it('rejects invalid allocations and specialized prepare allocation combinations', () => { + for (const override of [ + { sandboxAllocation: 'custom' }, + { sandboxAllocation: 'vercel-large', devcontainer: true }, + { sandboxAllocation: 'vercel-medium' }, + { sandboxAllocation: 'isolated-standard', devcontainer: true }, + { sandboxAllocation: 'cloudflare-single', createdOnPlatform: 'code-review' }, + { sandboxAllocation: 'isolated-standard', createdOnPlatform: 'code-review' }, + ]) { + expect(PrepareSessionInput.safeParse({ ...prepare, ...override }).success).toBe(false); + } + for (const runtime of [ + { sandboxAllocation: 'custom' }, + { sandboxAllocation: 'vercel-small', resources: { vcpus: 8, memory: 16384 } }, + { sandboxAllocation: 'cloudflare-single', devcontainer: true }, + { sandboxAllocation: 'vercel-medium' }, + ]) { + expect(StartSessionInput.safeParse({ ...baseStartInput, runtime }).success).toBe(false); + } + }); +}); + describe('legacy live attachment input compatibility', () => { it('accepts only the supported isolated Standard allocation', () => { const input = { diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index ae201c9ae4..63c5b60641 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/services/cloud-agent-next/src/router/schemas.ts @@ -1,4 +1,8 @@ import * as z from 'zod'; +import { + getKiloSandboxAllocation, + sandboxAllocationInputSchema, +} from '@kilocode/worker-utils/sandbox-allocation'; import { sessionIdSchema as kiloSessionIdSchema } from '@kilocode/session-ingest-contracts'; import { worktreeFileQuerySchema } from '@kilocode/worker-utils/cloud-agent-worktree-changes'; import { @@ -53,6 +57,15 @@ export { export const MessageIdSchema = z.string().regex(MESSAGE_ID_PATTERN, MESSAGE_ID_FORMAT_DESCRIPTION); +const ManagedSandboxAllocationInput = sandboxAllocationInputSchema.transform((request, ctx) => { + const allocation = getKiloSandboxAllocation(request); + if (allocation === undefined) { + ctx.addIssue({ code: 'custom', message: 'BYOC sandbox allocation is not available' }); + return z.NEVER; + } + return allocation; +}); + // Re-export types export type { EncryptedSecretEnvelope, @@ -578,10 +591,9 @@ const PrepareSessionSharedFields = { .describe( 'When true, route the session to a Docker-in-Docker sandbox that supports devcontainer runtimes' ), - sandboxAllocation: z - .literal('isolated-standard') - .optional() - .describe('Allocate a dedicated Standard Cloudflare container for this session'), + sandboxAllocation: ManagedSandboxAllocationInput.optional().describe( + 'Select a provider account and instance type instead of default routing' + ), }; const PrepareSessionNonCloneVariant = z.object({ @@ -652,21 +664,18 @@ export const PrepareSessionInput = z path: ['githubRepo'], }) .superRefine((data, ctx) => { - if (data.sandboxAllocation === 'isolated-standard' && data.devcontainer) { + if (data.sandboxAllocation !== undefined && data.devcontainer) { ctx.addIssue({ code: 'custom', path: ['sandboxAllocation'], - message: 'Isolated Standard allocation cannot be combined with devcontainer', + message: 'Sandbox allocation cannot be combined with devcontainer', }); } - if ( - data.sandboxAllocation === 'isolated-standard' && - data.createdOnPlatform === 'code-review' - ) { + if (data.sandboxAllocation !== undefined && data.createdOnPlatform === 'code-review') { ctx.addIssue({ code: 'custom', path: ['sandboxAllocation'], - message: 'Isolated Standard allocation cannot be combined with code review', + message: 'Sandbox allocation cannot be combined with code review', }); } @@ -908,6 +917,10 @@ export const StartSessionInput = z .optional(), repository: RepositoryInputSchema, profile: ProfileInputSchema, + runtime: z + .object({ sandboxAllocation: ManagedSandboxAllocationInput.optional() }) + .strict() + .optional(), options: z .object({ kilocodeOrganizationId: z.string().uuid().optional(), diff --git a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts index 579d4d05a2..9459be82ec 100644 --- a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts @@ -2,6 +2,11 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { SandboxSession } from '../sandbox-session/SandboxSession.js'; import { createMemoryEventQueries } from '../session/preparation-test-helpers.js'; import type { BillingContext } from '@kilocode/container-usage'; +import { + getSandboxAllocationResources, + type SandboxAllocation, + type VercelSandboxResources, +} from '@kilocode/worker-utils/sandbox-allocation'; import { SandboxControl, type SandboxAcquisition } from '../persistence/SandboxControl.js'; import { SESSION_DELIVERY_TIMEOUT_MS } from '../sandbox-session/control-dispatch.js'; import { @@ -79,7 +84,7 @@ vi.mock('drizzle-orm/durable-sqlite/migrator', () => ({ migrate: vi.fn(async () vi.mock('../../drizzle/migrations', () => ({ default: {} })); vi.mock('../session/queries/index.js', () => ({ createEventQueries: mocks.eventQueries })); -const SANDBOX_ID = 'ses-abcdef'; +const SANDBOX_ID = `ses-${'a'.repeat(48)}`; const OWNER = 'owner_1'; const ROUTE = { ownerId: OWNER, @@ -165,6 +170,7 @@ async function harness( options: { containmentEnabled?: boolean; env?: Partial; + sandboxAllocation?: SandboxAllocation; configureAllocation?: (value: ReturnType, id: string) => void; } = {} ) { @@ -284,7 +290,11 @@ async function harness( getCredentialMetadata: vi.fn(async () => parseSessionMetadata({ metadataSchemaVersion: 2, - identity: { sessionId: ROUTE.sessionId, userId: OWNER }, + identity: { + sessionId: ROUTE.sessionId, + userId: OWNER, + ...(options.sandboxAllocation ? { orgId: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa' } : {}), + }, auth: { kiloSessionId: ROUTE.kiloSessionId, kilocodeToken: 'test-token' }, workspace: { sandboxId: SANDBOX_ID, @@ -300,6 +310,7 @@ async function harness( kilocode: containmentEnabled, }, }), + ...(options.sandboxAllocation ? { sandboxAllocation: options.sandboxAllocation } : {}), }, lifecycle: { version: 1, timestamp: Date.now() }, }) @@ -374,6 +385,7 @@ async function harness( ownerId: OWNER, sessionId: ROUTE.sessionId, allowCreate: true, + resources: getSandboxAllocationResources(options.sandboxAllocation), billing: BILLING, }); }, @@ -1230,11 +1242,14 @@ describe('SandboxControl lifecycle boundaries', () => { const h = await harness(); const acquisition = { id: 'attempt_a', deadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS }; const transaction = h.storage.transaction.bind(h.storage); - vi.spyOn(h.storage, 'transaction').mockImplementationOnce(operation => + const transactionSpy = vi.spyOn(h.storage, 'transaction').mockImplementation(operation => transaction(async storage => { - await operation(storage); - expect(h.records.has('acquisition_receipts')).toBe(true); - throw new Error('acquisition transaction failed'); + const result = await operation(storage); + if (h.records.has('acquisition_receipts')) { + transactionSpy.mockRestore(); + throw new Error('acquisition transaction failed'); + } + return result; }) ); await expect(h.acquire(acquisition)).rejects.toThrow('acquisition transaction failed'); @@ -1250,10 +1265,16 @@ describe('SandboxControl lifecycle boundaries', () => { const h = await harness(); const acquisition = { id: 'attempt_a', deadlineAt: Date.now() + SESSION_DELIVERY_TIMEOUT_MS }; const transaction = h.storage.transaction.bind(h.storage); - vi.spyOn(h.storage, 'transaction').mockImplementationOnce(async operation => { - await transaction(operation); - throw new Error('reset after acquisition commit'); - }); + const transactionSpy = vi + .spyOn(h.storage, 'transaction') + .mockImplementation(async operation => { + const result = await transaction(operation); + if (h.records.has('acquisition_receipts')) { + transactionSpy.mockRestore(); + throw new Error('reset after acquisition commit'); + } + return result; + }); await expect(h.acquire(acquisition)).rejects.toThrow('reset after acquisition commit'); const claimed = await h.control.getPhysicalRecord(); expect(claimed.state).toBe('creating'); @@ -2257,6 +2278,73 @@ describe('SandboxControl lifecycle boundaries', () => { expect(h.sendRequest).not.toHaveBeenCalled(); }); + it.each([undefined, 'vercel-small', 'vercel-large'] as const)( + 'pins %s effective resources and rejects incompatible reuse after eviction', + async preset => { + const h = await harness({ + env: { + VERCEL_TOKEN: 'test-token', + VERCEL_TEAM_ID: 'team_1', + VERCEL_PROJECT_ID: 'project_1', + VERCEL_SANDBOX_RUNTIME_BUILD_ID: 'build_1', + VERCEL_SANDBOX_SNAPSHOT_ID: 'snapshot_1', + VERCEL_SANDBOX_RUNTIME: 'node24', + VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: '300000', + VERCEL_SANDBOX_EXTEND_DURATION_MS: '120000', + }, + }); + if (preset === undefined) h.records.set('provider_kind', 'vercel'); + const input = { + ownerId: OWNER, + sessionId: ROUTE.sessionId, + provider: 'vercel' as const, + resources: getSandboxAllocationResources(preset), + }; + await h.control.ensureReady(input); + await h.evict(); + await expect(h.control.ensureReady(input)).resolves.toMatchObject({ physical: 'stopped' }); + for (const other of [undefined, 'vercel-small', 'vercel-large'] as const) { + if (other === preset) continue; + await expect( + h.control.ensureReady({ ...input, resources: getSandboxAllocationResources(other) }) + ).rejects.toThrow('Sandbox resources mismatch'); + } + await expect( + h.control.ensureReady({ ...input, provider: 'cloudflare', resources: undefined }) + ).rejects.toThrow('Sandbox provider mismatch'); + const claimed = await h.control.claimCreate('after-reset'); + expect(claimed.createIntent?.vercel?.resources).toEqual(input.resources); + expect(mocks.getSandbox).not.toHaveBeenCalled(); + } + ); + + it.each([ + { provider: 'vercel', resources: { vcpus: 2, memory: 8192 } }, + { provider: 'cloudflare', resources: { vcpus: 2, memory: 4096 } }, + { provider: 'unknown' }, + ])('rejects invalid persisted provider configuration %j on restart', async configuration => { + const h = await harness(); + h.records.set('provider_configuration', configuration); + await expect(h.evict()).rejects.toThrow(); + expect(mocks.getSandbox).not.toHaveBeenCalled(); + }); + + it('rejects resources on a Cloudflare readiness request before pinning or creating', async () => { + const h = await harness(); + await expect( + h.control.ensureReady({ + ownerId: OWNER, + sessionId: ROUTE.sessionId, + provider: 'cloudflare', + resources: { vcpus: 2, memory: 4096 }, + allowCreate: true, + }) + ).rejects.toThrow(); + expect(h.records.has('provider_configuration')).toBe(false); + expect((await h.control.getPhysicalRecord()).state).toBe('stopped'); + expect(mocks.getSandbox).not.toHaveBeenCalled(); + }); + it('rejects missing provider configuration and mismatched billing without an illegal transition', async () => { const h = await harness(); await expect( @@ -3669,9 +3757,17 @@ describe('SandboxControl lifecycle boundaries', () => { await h.flush(); }); - it.each([true, false])( - 'reconciles a lost Vercel create response across runtime config rotation with containment %s', - async containmentEnabled => { + it.each( + [true, false].flatMap(containmentEnabled => + ([undefined, 'vercel-small', 'vercel-large'] as const).map(sandboxAllocation => ({ + containmentEnabled, + sandboxAllocation, + })) + ) + )( + 'reconciles a lost Vercel create response and retains $sandboxAllocation sizing with containment $containmentEnabled across restart and replacement', + async ({ containmentEnabled, sandboxAllocation }) => { + const resources = getSandboxAllocationResources(sandboxAllocation); const remote = new Map(); const inspected: URL[] = []; const policyUpdates: URL[] = []; @@ -3688,10 +3784,12 @@ describe('SandboxControl lifecycle boundaries', () => { projectId: string; runtime: string; timeout: number; + resources?: VercelSandboxResources; source: { snapshotId: string }; tags: Record; networkPolicy?: unknown; }; + expect(body.resources).toEqual(resources); if (!containmentEnabled) expect(body.networkPolicy).toBeUndefined(); if (remote.has(body.name)) throw new Error('Name is retained'); const sessionId = `vsess_${remote.size + 1}`; @@ -3712,8 +3810,8 @@ describe('SandboxControl lifecycle boundaries', () => { sourceSnapshotId: body.source.snapshotId, runtime: body.runtime, status: 'running', - memory: 2048, - vcpus: 2, + memory: body.resources?.memory ?? 2048, + vcpus: body.resources?.vcpus ?? 2, region: 'iad1', timeout: body.timeout, requestedAt: Date.now(), @@ -3773,11 +3871,13 @@ describe('SandboxControl lifecycle boundaries', () => { VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: '300000', VERCEL_SANDBOX_EXTEND_DURATION_MS: '120000', }, + sandboxAllocation, }); const creating = h.control.ensureReady({ ownerId: OWNER, sessionId: ROUTE.sessionId, provider: 'vercel', + resources, allowCreate: true, billing: BILLING, }); @@ -3790,6 +3890,11 @@ describe('SandboxControl lifecycle boundaries', () => { allocationName: [...remote.keys()][0], vercel: { runtimeBuildId: 'build_1', snapshotId: 'snapshot_1' }, }); + expect(uncertain.createIntent?.vercel?.resources).toEqual(resources); + expect(h.records.get('provider_configuration')).toEqual({ + provider: 'vercel', + ...(resources ? { resources } : {}), + }); h.env.VERCEL_SANDBOX_RUNTIME_BUILD_ID = 'build_2'; h.env.VERCEL_SANDBOX_SNAPSHOT_ID = 'snapshot_2'; h.env.CREDENTIAL_CONTAINMENT_ENABLED = containmentEnabled ? 'false' : 'true'; @@ -3799,7 +3904,11 @@ describe('SandboxControl lifecycle boundaries', () => { expect(inspected[0]?.searchParams.get('resume')).toBe('false'); expect((await h.control.getPhysicalRecord()).state).toBe('stopped'); expect([...remote.values()][0]?.session.status).toBe('stopped'); + await h.evict(); await h.create(); + expect((await h.control.getPhysicalRecord()).createIntent?.vercel?.resources).toEqual( + resources + ); await h.ready(); expect(remote.size).toBe(2); expect([...remote.values()][1]?.session.sourceSnapshotId).toBe('snapshot_2'); diff --git a/services/cloud-agent-next/src/sandbox-control/physical-lifecycle.ts b/services/cloud-agent-next/src/sandbox-control/physical-lifecycle.ts index 391b6ab3b6..b7d13b9dc1 100644 --- a/services/cloud-agent-next/src/sandbox-control/physical-lifecycle.ts +++ b/services/cloud-agent-next/src/sandbox-control/physical-lifecycle.ts @@ -1,5 +1,16 @@ +import { vercelSandboxResourcesSchema } from '@kilocode/worker-utils/sandbox-allocation'; +import { z } from 'zod'; import type { VercelSandboxRuntimeConfig } from '../agent-sandbox/vercel/vercel-runtime-config.js'; +export const sandboxProviderConfigurationSchema = z.discriminatedUnion('provider', [ + z.object({ provider: z.literal('cloudflare') }).strict(), + z + .object({ provider: z.literal('vercel'), resources: vercelSandboxResourcesSchema.optional() }) + .strict(), +]); + +export type SandboxProviderConfiguration = z.infer; + export type PhysicalState = 'stopped' | 'creating' | 'running' | 'stopping' | 'failed' | 'unknown'; export type CredentialContainmentRequirements = { @@ -28,7 +39,7 @@ export type CreateIntent = { allocationName?: string; vercel?: Pick< VercelSandboxRuntimeConfig, - 'projectId' | 'snapshotId' | 'runtimeBuildId' | 'runtime' + 'projectId' | 'snapshotId' | 'runtimeBuildId' | 'runtime' | 'resources' >; containment?: CredentialContainmentRequirements; }; diff --git a/services/cloud-agent-next/src/sandbox-control/vercel-provider.test.ts b/services/cloud-agent-next/src/sandbox-control/vercel-provider.test.ts index af9019bdb8..b537e3c7cd 100644 --- a/services/cloud-agent-next/src/sandbox-control/vercel-provider.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/vercel-provider.test.ts @@ -123,6 +123,27 @@ describe('vercel provider adapter', () => { }); }); + it.each([ + { vcpus: 2, memory: 4096 }, + { vcpus: 4, memory: 8192 }, + ] as const)( + 'propagates $vcpus vCPU resources for creation and uncertain-create inspection', + async resources => { + const createSandbox = vi.fn().mockRejectedValue(new Error('response lost')); + const inspectByName = vi.fn(fakeClient().inspectByName); + const provider = createVercelProviderAdapter({ + sandboxName: intent.allocationName, + config: { ...config, resources }, + restClient: fakeClient({ createSandbox, inspectByName }), + }); + await expect(provider.create(intent)).rejects.toThrow('response lost'); + await expect(provider.observe(null, intent)).resolves.toMatchObject({ status: 'active' }); + expect(createSandbox).toHaveBeenCalledWith(expect.objectContaining({ resources })); + expect(inspectByName).toHaveBeenCalledWith(expect.objectContaining({ resources })); + expect(createSandbox).toHaveBeenCalledTimes(1); + } + ); + it('installs the creation policy before launching the control wrapper', async () => { const createSandbox = vi.fn(fakeClient().createSandbox); const executeCommand = vi.fn(fakeClient().executeCommand); diff --git a/services/cloud-agent-next/src/sandbox-control/vercel-provider.ts b/services/cloud-agent-next/src/sandbox-control/vercel-provider.ts index 6cf2f1909e..f943a4be2c 100644 --- a/services/cloud-agent-next/src/sandbox-control/vercel-provider.ts +++ b/services/cloud-agent-next/src/sandbox-control/vercel-provider.ts @@ -144,6 +144,7 @@ export function createVercelProviderAdapter(deps: { snapshotId: config.snapshotId, runtime: config.runtime, timeoutMs: config.initialTimeoutMs, + ...(config.resources === undefined ? {} : { resources: config.resources }), ...(intent.networkPolicy === undefined ? {} : { networkPolicy: intent.networkPolicy }), }); return { providerRef: encodeVercelProviderRef(created.runtime) }; @@ -178,6 +179,7 @@ export function createVercelProviderAdapter(deps: { runtimeBuildId: config.runtimeBuildId, snapshotId: config.snapshotId, runtime: config.runtime, + ...(config.resources === undefined ? {} : { resources: config.resources }), }); if (!inspected) { return { diff --git a/services/cloud-agent-next/src/sandbox-id.test.ts b/services/cloud-agent-next/src/sandbox-id.test.ts index 9656e9090c..626b89cecf 100644 --- a/services/cloud-agent-next/src/sandbox-id.test.ts +++ b/services/cloud-agent-next/src/sandbox-id.test.ts @@ -532,6 +532,80 @@ describe('selectSandboxForNewSession', () => { const controlSessionId = 'workspace_12345678-1234-1234-1234-123456789abc'; const legacySessionId = 'agent_12345678-1234-1234-1234-123456789abc'; + it.each(['cloudflare-single', 'cloudflare-shared', 'vercel-small', 'vercel-large'] as const)( + 'routes explicit %s independently of conflicting allocation and provider rollouts', + async sandboxAllocation => { + for (const rollout of ['', '*']) { + const selection = await selectSandboxForNewSession({ + env: { + ...completeVercelConfiguration, + PER_SESSION_SANDBOX_ORG_IDS: rollout, + VERCEL_SANDBOX_ORG_IDS: rollout, + }, + orgId: 'org-id', + userId: 'user-id', + sessionId: controlSessionId, + sandboxAllocation, + }); + expect(selection.provider).toBe( + sandboxAllocation.startsWith('vercel-') ? 'vercel' : 'cloudflare' + ); + expect(selection.sandboxId).toMatch( + sandboxAllocation === 'cloudflare-shared' ? /^org-/ : /^ses-/ + ); + } + } + ); + + it('retains the default shared identity for an explicit shared preset', async () => { + const shared = await generateSandboxRoutingTarget('', 'org-id', 'user-id', controlSessionId); + expect( + await generateSandboxRoutingTarget('*', 'org-id', 'user-id', controlSessionId, undefined, { + sandboxAllocation: 'cloudflare-shared', + }) + ).toEqual(shared); + expect( + await generateSandboxRoutingTarget('*', 'org-id', 'other-user', controlSessionId, undefined, { + sandboxAllocation: 'cloudflare-shared', + }) + ).not.toEqual(shared); + expect( + await generateSandboxRoutingTarget('', 'org-id', 'user-id', legacySessionId) + ).not.toEqual(shared); + }); + + it.each(['cloudflare-single', 'cloudflare-shared'] as const)( + 'routes an explicit %s on a legacy session without changing its plane', + async sandboxAllocation => { + const selection = await selectSandboxForNewSession({ + env: { ...completeVercelConfiguration, VERCEL_SANDBOX_ORG_IDS: '*' }, + orgId: 'org-id', + userId: 'user-id', + sessionId: legacySessionId, + sandboxAllocation, + }); + expect(selection.provider).toBe('cloudflare'); + expect(selection.sandboxId).toMatch( + sandboxAllocation === 'cloudflare-shared' ? /^org-/ : /^ses-/ + ); + } + ); + + it.each(['vercel-small', 'vercel-large'] as const)( + 'rejects %s on a legacy session because Vercel exists only on the control plane', + async sandboxAllocation => { + await expect( + selectSandboxForNewSession({ + env: completeVercelConfiguration, + orgId: 'org-id', + userId: 'user-id', + sessionId: legacySessionId, + sandboxAllocation, + }) + ).rejects.toThrow('require a control-plane session'); + } + ); + it('selects Cloudflare by default while preserving shared sandbox allocation', async () => { const selection = await selectSandboxForNewSession({ env: {}, diff --git a/services/cloud-agent-next/src/sandbox-id.ts b/services/cloud-agent-next/src/sandbox-id.ts index 03837deeda..7ffaee49a3 100644 --- a/services/cloud-agent-next/src/sandbox-id.ts +++ b/services/cloud-agent-next/src/sandbox-id.ts @@ -1,5 +1,17 @@ +import { + getSandboxAllocationProvider, + getSandboxAllocationRequest, + sandboxAllocationRequiresControlPlane, + type SandboxAllocation, + type SandboxDestination, +} from '@kilocode/worker-utils/sandbox-allocation'; import type { AgentSandboxProvider, SandboxId, Env } from './types.js'; -import { sessionPlaneFromId } from './session-plane.js'; +import { + sessionPlaneForNewOwner, + sessionPlaneFromId, + type ControlPlaneOwnerEnv, + type SessionPlane, +} from './session-plane.js'; import type { Sandbox } from '@cloudflare/sandbox'; import { parseVercelSandboxEnrollment, @@ -42,11 +54,30 @@ export type SandboxRoutingTarget = }; export type SandboxRoutingOptions = { + sandboxAllocation?: SandboxAllocation; devcontainer?: boolean; createdOnPlatform?: string; - sandboxAllocation?: 'isolated-standard'; }; +/** + * Isolated sandbox-ID prefix per allocation. `cloudflare-shared` maps to no prefix + * because it routes to the shared sandbox rather than an isolated identity. + */ +const SANDBOX_ALLOCATION_ID_PREFIX: Record = { + 'isolated-standard': 'istd', + 'cloudflare-single': 'ses', + 'cloudflare-shared': undefined, + 'vercel-small': 'ses', + 'vercel-large': 'ses', +}; + +function sandboxIdMatchesAllocation(sandboxId: string, allocation: SandboxAllocation): boolean { + const prefix = SANDBOX_ALLOCATION_ID_PREFIX[allocation]; + return prefix === undefined + ? isGeneratedSharedSandboxId(sandboxId) + : new RegExp(`^${prefix}-[0-9a-f]{48}$`).test(sandboxId); +} + export type SandboxIdClass = | 'shared' | 'legacy-shared' @@ -199,6 +230,7 @@ type SelectSandboxForNewSessionInput = { sessionId: string; botId?: string; devcontainer?: boolean; + sandboxAllocation?: SandboxAllocation; }; /** @@ -213,6 +245,39 @@ export function selectSandboxProvider(input: { sandboxId: SandboxId; sessionId: string; devcontainer?: boolean; + sandboxAllocation?: SandboxAllocation; +}): AgentSandboxProvider { + const allocation = input.sandboxAllocation; + if (allocation !== undefined) { + if (input.devcontainer) { + throw new Error('Sandbox allocations cannot be combined with specialized sandbox routing'); + } + if ( + sandboxAllocationRequiresControlPlane(allocation) && + sessionPlaneFromId(input.sessionId) !== 'control' + ) { + throw new Error('Vercel sandbox allocations require a control-plane session'); + } + if (!sandboxIdMatchesAllocation(input.sandboxId, allocation)) { + throw new Error('Sandbox allocation does not match the sandbox identity'); + } + return getSandboxAllocationProvider(allocation); + } + return selectDefaultSandboxProvider({ + env: input.env, + orgId: input.orgId, + plane: sessionPlaneFromId(input.sessionId), + isolated: input.sandboxId.startsWith('ses-'), + devcontainer: input.devcontainer, + }); +} + +function selectDefaultSandboxProvider(input: { + env: SandboxSelectionEnv; + orgId?: string; + plane: SessionPlane; + isolated: boolean; + devcontainer?: boolean; }): AgentSandboxProvider { const enrollment = parseVercelSandboxEnrollment(input.env); const runtimeConfig = parseVercelSandboxRuntimeConfig(input.env); @@ -221,9 +286,9 @@ export function selectSandboxProvider(input: { ? enrollment.orgIds.has('*') || enrollment.orgIds.has(input.orgId) : enrollment.allowPersonal; const useVercel = - sessionPlaneFromId(input.sessionId) === 'control' && + input.plane === 'control' && !input.devcontainer && - input.sandboxId.startsWith('ses-') && + input.isolated && enrollment.enabled && enrolled && runtimeConfig !== undefined; @@ -231,6 +296,30 @@ export function selectSandboxProvider(input: { return useVercel ? 'vercel' : 'cloudflare'; } +export function getDefaultSandboxDestination( + env: SandboxSelectionEnv & ControlPlaneOwnerEnv, + owner: { userId: string; orgId?: string }, + devcontainer = false +): SandboxDestination { + if (devcontainer) { + return { + provider: { id: 'cloudflare', account: 'kilo' }, + instanceType: 'devcontainer', + }; + } + const isolated = isOrgInList(env.PER_SESSION_SANDBOX_ORG_IDS, owner.orgId); + const provider = selectDefaultSandboxProvider({ + env, + orgId: owner.orgId, + plane: sessionPlaneForNewOwner(env, owner, { createdOnPlatform: 'cloud-agent-web' }), + isolated, + }); + if (provider === 'vercel') { + return { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }; + } + return getSandboxAllocationRequest(isolated ? 'cloudflare-single' : 'cloudflare-shared'); +} + export async function selectSandboxForNewSession( input: SelectSandboxForNewSessionInput ): Promise { @@ -240,7 +329,7 @@ export async function selectSandboxForNewSession( input.userId, input.sessionId, input.botId, - input.devcontainer + { devcontainer: input.devcontainer, sandboxAllocation: input.sandboxAllocation } ); const provider = selectSandboxProvider({ env: input.env, @@ -248,6 +337,7 @@ export async function selectSandboxForNewSession( sandboxId, sessionId: input.sessionId, devcontainer: input.devcontainer, + sandboxAllocation: input.sandboxAllocation, }); return { sandboxId, provider }; @@ -279,17 +369,28 @@ export async function generateSandboxRoutingTarget( options?: boolean | SandboxRoutingOptions ): Promise { const routingOptions = typeof options === 'boolean' ? { devcontainer: options } : (options ?? {}); - const perSessionOrgs = parseOrgIdList(perSessionOrgIds); + const allocation = routingOptions.sandboxAllocation; + if (allocation !== undefined) { + if (routingOptions.devcontainer || routingOptions.createdOnPlatform === 'code-review') { + throw new Error('Sandbox allocations cannot be combined with specialized sandbox routing'); + } + if ( + sandboxAllocationRequiresControlPlane(allocation) && + sessionPlaneFromId(sessionId) !== 'control' + ) { + throw new Error('Vercel sandbox allocations require a control-plane session'); + } + const prefix = SANDBOX_ALLOCATION_ID_PREFIX[allocation]; + // `cloudflare-shared` has no isolated prefix: it falls through to the shared route. + if (prefix) return { kind: 'isolated', sandboxId: await hashToSandboxId(sessionId, prefix) }; + } if (routingOptions.devcontainer) { return { kind: 'isolated', sandboxId: await hashToSandboxId(sessionId, 'dind') }; } if (routingOptions.createdOnPlatform === 'code-review') { return { kind: 'isolated', sandboxId: await hashToSandboxId(sessionId, 'crv') }; } - if (routingOptions.sandboxAllocation === 'isolated-standard') { - return { kind: 'isolated', sandboxId: await hashToSandboxId(sessionId, 'istd') }; - } - if (perSessionOrgs.has('*') || (orgId !== undefined && perSessionOrgs.has(orgId))) { + if (allocation !== 'cloudflare-shared' && isOrgInList(perSessionOrgIds, orgId)) { return { kind: 'isolated', sandboxId: await hashToSandboxId(sessionId, 'ses') }; } diff --git a/services/cloud-agent-next/src/sandbox-selection.test.ts b/services/cloud-agent-next/src/sandbox-selection.test.ts new file mode 100644 index 0000000000..818470a325 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-selection.test.ts @@ -0,0 +1,270 @@ +import { describe, expect, it } from 'vitest'; +import { + getSandboxAllocationRequest, + sandboxAllocationSchema, +} from '@kilocode/worker-utils/sandbox-allocation'; +import { + assertSandboxAllocationAvailable, + getSandboxSelectionCapabilities, + isSandboxAllocationAvailable, +} from './sandbox-selection.js'; +import { classifySandboxId, selectSandboxForNewSession } from './sandbox-id.js'; +import { sessionPlaneForNewOwner } from './session-plane.js'; +import type { Env } from './types.js'; + +const configured = { + SANDBOX_SELECTION_IDS: 'org-id', + CONTROL_PLANE_IDS: 'org-id', + VERCEL_TOKEN: 'test-token', + VERCEL_TEAM_ID: 'team-id', + VERCEL_PROJECT_ID: 'project-id', + VERCEL_SANDBOX_SNAPSHOT_ID: 'snapshot-id', + VERCEL_SANDBOX_RUNTIME_BUILD_ID: 'build-id', + VERCEL_SANDBOX_RUNTIME: 'node24', + VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: '300000', + VERCEL_SANDBOX_EXTEND_DURATION_MS: '600000', +} satisfies Partial; +const owner = { userId: 'oauth/user', orgId: 'org-id' }; + +describe('sandbox selection policy', () => { + it.each([ + { + name: 'shared Cloudflare', + overrides: {}, + devcontainer: false, + expected: getSandboxAllocationRequest('cloudflare-shared'), + }, + { + name: 'isolated Cloudflare', + overrides: { PER_SESSION_SANDBOX_ORG_IDS: owner.orgId }, + devcontainer: false, + expected: getSandboxAllocationRequest('cloudflare-single'), + }, + { + name: 'Vercel with provider-default resources', + overrides: { PER_SESSION_SANDBOX_ORG_IDS: owner.orgId, VERCEL_SANDBOX_ORG_IDS: owner.orgId }, + devcontainer: false, + expected: { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + }, + { + name: 'Vercel with user-level control-plane enrollment', + overrides: { + CONTROL_PLANE_IDS: owner.userId, + PER_SESSION_SANDBOX_ORG_IDS: '*', + VERCEL_SANDBOX_ORG_IDS: '*', + }, + devcontainer: false, + expected: { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + }, + { + name: 'legacy isolation despite Vercel enrollment', + overrides: { + CONTROL_PLANE_IDS: '', + PER_SESSION_SANDBOX_ORG_IDS: owner.orgId, + VERCEL_SANDBOX_ORG_IDS: owner.orgId, + }, + devcontainer: false, + expected: getSandboxAllocationRequest('cloudflare-single'), + }, + { + name: 'shared routing despite Vercel enrollment', + overrides: { VERCEL_SANDBOX_ORG_IDS: owner.orgId }, + devcontainer: false, + expected: getSandboxAllocationRequest('cloudflare-shared'), + }, + { + name: 'missing Vercel configuration', + overrides: { + PER_SESSION_SANDBOX_ORG_IDS: owner.orgId, + VERCEL_SANDBOX_ORG_IDS: owner.orgId, + VERCEL_TOKEN: undefined, + }, + devcontainer: false, + expected: getSandboxAllocationRequest('cloudflare-single'), + }, + { + name: 'devcontainer instead of the normal Vercel default', + overrides: { PER_SESSION_SANDBOX_ORG_IDS: owner.orgId, VERCEL_SANDBOX_ORG_IDS: owner.orgId }, + devcontainer: true, + expected: { provider: { id: 'cloudflare', account: 'kilo' }, instanceType: 'devcontainer' }, + }, + { + name: 'unchanged default when explicit Vercel is unavailable for compute billing', + overrides: { + PER_SESSION_SANDBOX_ORG_IDS: owner.orgId, + VERCEL_SANDBOX_ORG_IDS: owner.orgId, + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: owner.orgId, + }, + devcontainer: false, + expected: { provider: { id: 'vercel', account: 'kilo' }, instanceType: 'default' }, + }, + ])( + 'previews $name consistently with actual routing', + async ({ overrides, devcontainer, expected }) => { + const env = { ...configured, ...overrides } as Env; + const capabilities = getSandboxSelectionCapabilities(env, owner, devcontainer); + expect(capabilities.defaultDestination).toEqual(expected); + const plane = sessionPlaneForNewOwner(env, owner, { createdOnPlatform: 'cloud-agent-web' }); + const actual = await selectSandboxForNewSession({ + env, + ...owner, + sessionId: `${plane === 'control' ? 'workspace' : 'agent'}_12345678-1234-1234-1234-123456789012`, + devcontainer, + }); + expect(actual.provider).toBe(expected.provider.id); + if (actual.provider === 'cloudflare') { + expect(classifySandboxId(actual.sandboxId)).toBe( + expected.instanceType === 'shared' + ? 'shared' + : expected.instanceType === 'devcontainer' + ? 'devcontainer' + : 'isolated-small' + ); + } + } + ); + + it('does not authorize a Kilo allocation from a BYOC capability with the same size', () => { + expect( + isSandboxAllocationAvailable( + { + enabled: true, + options: [ + { + allocation: { provider: { id: 'vercel', account: 'byoc' }, instanceType: 'small' }, + }, + ], + }, + 'vercel-small' + ) + ).toBe(false); + }); + + it.each([ + { SANDBOX_SELECTION_IDS: undefined }, + { SANDBOX_SELECTION_IDS: '' }, + { SANDBOX_SELECTION_IDS: 'other-org' }, + { SANDBOX_SELECTION_IDS: '', NODE_ENV: 'development' }, + ])('disables selection without the owner allowlist: %j', overrides => { + const env = { ...configured, ...overrides } as Env; + expect(getSandboxSelectionCapabilities(env, owner)).toEqual({ enabled: false, options: [] }); + for (const preset of sandboxAllocationSchema.options) { + expect(() => assertSandboxAllocationAvailable(env, owner, preset)).toThrow('not enabled'); + } + }); + + it('keeps the trigger-only allocation out of the manual picker', () => { + const capabilities = getSandboxSelectionCapabilities(configured as Env, owner); + expect(capabilities.options.map(option => option.allocation)).toEqual([ + getSandboxAllocationRequest('cloudflare-single'), + getSandboxAllocationRequest('cloudflare-shared'), + getSandboxAllocationRequest('vercel-small'), + getSandboxAllocationRequest('vercel-large'), + ]); + }); + + it('enables personal selection when the user is listed', () => { + const env = { ...configured, SANDBOX_SELECTION_IDS: owner.userId } as Env; + const capabilities = getSandboxSelectionCapabilities(env, { userId: owner.userId }); + expect(capabilities.enabled).toBe(true); + expect(capabilities.options).toHaveLength(4); + }); + + it('enables personal selection for wildcard rollouts', () => { + const env = { ...configured, SANDBOX_SELECTION_IDS: '*', CONTROL_PLANE_IDS: '*' } as Env; + expect(getSandboxSelectionCapabilities(env, { userId: owner.userId }).enabled).toBe(true); + }); + + it('does not enable personal selection from an organization-only allowlist', () => { + expect(getSandboxSelectionCapabilities(configured as Env, { userId: owner.userId })).toEqual({ + enabled: false, + options: [], + }); + }); + + it('enables an organization session when the user is listed', () => { + const env = { ...configured, SANDBOX_SELECTION_IDS: owner.userId } as Env; + expect(getSandboxSelectionCapabilities(env, owner).enabled).toBe(true); + }); + + it.each([ + { CONTROL_PLANE_IDS: 'org-id' }, + { CONTROL_PLANE_IDS: owner.userId }, + // Plane enrollment is not an availability condition: a legacy owner may still + // choose, and a Vercel choice plane-forces that one session. + { CONTROL_PLANE_IDS: '' }, + { CONTROL_PLANE_IDS: undefined }, + { CONTROL_PLANE_IDS: 'other-owner' }, + ])('enables selection regardless of plane enrollment: %j', overrides => { + const env = { + ...configured, + ...overrides, + VERCEL_SANDBOX_ORG_IDS: '', + } as Env; + const capabilities = getSandboxSelectionCapabilities(env, owner); + expect(capabilities.enabled).toBe(true); + expect(capabilities.options).toHaveLength(4); + expect(capabilities.options.every(option => !('available' in option))).toBe(true); + expect(JSON.stringify(capabilities)).not.toContain('test-token'); + for (const preset of sandboxAllocationSchema.options) { + expect(() => assertSandboxAllocationAvailable(env, owner, preset)).not.toThrow(); + } + }); + + it.each(Object.keys(configured).filter(key => key.startsWith('VERCEL_')))( + 'disables only Vercel when operational configuration %s is absent', + key => { + const env = { ...configured, [key]: undefined } as Env; + const capabilities = getSandboxSelectionCapabilities(env, owner); + expect(capabilities.enabled).toBe(true); + expect(capabilities.options.map(option => option.allocation)).toEqual([ + getSandboxAllocationRequest('cloudflare-single'), + getSandboxAllocationRequest('cloudflare-shared'), + ]); + for (const preset of ['vercel-small', 'vercel-large'] as const) { + expect(() => assertSandboxAllocationAvailable(env, owner, preset)).toThrow( + 'not configured' + ); + } + expect(() => assertSandboxAllocationAvailable(env, owner, 'isolated-standard')).not.toThrow(); + } + ); + + it('fails Vercel closed for enforced organization billing', () => { + const env = { + ...configured, + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: owner.orgId, + } as Env; + const capabilities = getSandboxSelectionCapabilities(env, owner); + expect(capabilities.options.map(option => option.allocation)).toEqual([ + getSandboxAllocationRequest('cloudflare-single'), + getSandboxAllocationRequest('cloudflare-shared'), + ]); + for (const preset of ['vercel-small', 'vercel-large'] as const) { + expect(() => assertSandboxAllocationAvailable(env, owner, preset)).toThrow( + 'enforced compute billing' + ); + } + expect(() => assertSandboxAllocationAvailable(env, owner, 'cloudflare-single')).not.toThrow(); + }); + + it('admits the trigger-only allocation without listing it', () => { + expect(isSandboxAllocationAvailable({ enabled: true, options: [] }, 'isolated-standard')).toBe( + true + ); + expect( + isSandboxAllocationAvailable( + { + enabled: true, + options: [{ allocation: getSandboxAllocationRequest('cloudflare-single') }], + }, + 'isolated-standard' + ) + ).toBe(true); + expect(isSandboxAllocationAvailable({ enabled: false, options: [] }, 'isolated-standard')).toBe( + false + ); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-selection.ts b/services/cloud-agent-next/src/sandbox-selection.ts new file mode 100644 index 0000000000..5fd47589c0 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-selection.ts @@ -0,0 +1,84 @@ +import { TRPCError } from '@trpc/server'; +import { + SELECTABLE_SANDBOX_ALLOCATIONS, + getKiloSandboxAllocation, + getSandboxAllocationProvider, + getSandboxAllocationRequest, + type SandboxAllocation, + type SandboxSelectionCapabilities, +} from '@kilocode/worker-utils/sandbox-allocation'; +import { parseVercelSandboxRuntimeConfig } from './agent-sandbox/vercel/vercel-runtime-config.js'; +import { isCloudAgentContainerBillingEnabled } from './container-billing-rollout.js'; +import { getDefaultSandboxDestination, isOrgInList } from './sandbox-id.js'; +import type { Env } from './types.js'; + +type SelectionOwner = { userId: string; orgId?: string }; + +function sandboxAllocationUnavailableReason( + env: Env, + owner: SelectionOwner, + allocation: SandboxAllocation +): string | undefined { + if (getSandboxAllocationProvider(allocation) !== 'vercel') return undefined; + if (!parseVercelSandboxRuntimeConfig(env)) return 'Vercel sandboxes are not configured'; + if (isCloudAgentContainerBillingEnabled(env, owner)) { + return 'Vercel sandboxes do not support enforced compute billing'; + } + return undefined; +} + +export function getSandboxSelectionCapabilities( + env: Env, + owner: SelectionOwner, + devcontainer = false +): SandboxSelectionCapabilities { + if ( + !isOrgInList(env.SANDBOX_SELECTION_IDS, owner.userId) && + !isOrgInList(env.SANDBOX_SELECTION_IDS, owner.orgId) + ) { + return { enabled: false, options: [] }; + } + + return { + enabled: true, + defaultDestination: getDefaultSandboxDestination(env, owner, devcontainer), + options: SELECTABLE_SANDBOX_ALLOCATIONS.filter( + allocation => sandboxAllocationUnavailableReason(env, owner, allocation) === undefined + ).map(allocation => ({ allocation: getSandboxAllocationRequest(allocation) })), + }; +} + +export function isSandboxAllocationAvailable( + capabilities: SandboxSelectionCapabilities, + allocation: SandboxAllocation +): boolean { + return ( + capabilities.enabled && + (allocation === 'isolated-standard' || + capabilities.options.some( + option => getKiloSandboxAllocation(option.allocation) === allocation + )) + ); +} + +export function assertSandboxAllocationAvailable( + env: Env, + owner: SelectionOwner, + allocation: SandboxAllocation +): void { + const capabilities = getSandboxSelectionCapabilities(env, owner); + if (!capabilities.enabled) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Sandbox selection is not enabled for this owner', + }); + } + if (!isSandboxAllocationAvailable(capabilities, allocation)) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: + sandboxAllocationUnavailableReason(env, owner, allocation) ?? + 'Sandbox allocation is unavailable', + }); + } +} diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index c93d9dcc91..cd52fd5891 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -17,6 +17,7 @@ import { } from '@kilocode/worker-utils/runtime-authorization'; import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; import { RuntimeAuthorizationSchema } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { getSandboxAllocationResources } from '@kilocode/worker-utils/sandbox-allocation'; import { resolveSecret } from '../auth.js'; import { issuePersistedRuntimeProxyGrant, @@ -3360,6 +3361,7 @@ export class SandboxSession extends DurableObject { ownerId: metadata.identity.userId, sessionId, provider, + resources: getSandboxAllocationResources(metadata.workspace?.sandboxAllocation), ...(acquisition ? { acquisition } : { allowCreate }), ...(metadata.workspace?.worktreeId ? { worktreeId: metadata.workspace.worktreeId } diff --git a/services/cloud-agent-next/src/sandbox-session/control-rpc.ts b/services/cloud-agent-next/src/sandbox-session/control-rpc.ts index 6541e525f5..c4063ce40a 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-rpc.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-rpc.ts @@ -1,3 +1,4 @@ +import type { VercelSandboxResources } from '@kilocode/worker-utils/sandbox-allocation'; import type { VercelSandboxNetworkPolicy } from '../agent-sandbox/vercel/vercel-sandbox-rest-client.js'; import type { CredentialContainmentRequirements } from '../sandbox-control/physical-lifecycle.js'; import { @@ -34,6 +35,7 @@ type SandboxControlRpc = { ownerId: string; sessionId: string; provider?: 'cloudflare' | 'vercel'; + resources?: VercelSandboxResources; allowCreate?: boolean; acquisition?: SandboxAcquisition; billing?: SandboxBillingInput; diff --git a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts index 292fd950cd..332d393668 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts @@ -1858,6 +1858,41 @@ describe('SandboxSession orchestration', () => { } ); + it.each([ + ['vercel-small', { vcpus: 2, memory: 4096 }], + ['vercel-large', { vcpus: 4, memory: 8192 }], + ] as const)( + 'forwards persisted %s resources through readiness after a session reset', + async (sandboxAllocation, resources) => { + const fixture = sessionFixture({ + identity: { + sessionId: SESSION_ID, + userId: 'user_1', + orgId: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + }, + workspace: { + sandboxId: SANDBOX_ID, + workspacePath: DIRECTORY, + sandboxProvider: 'vercel', + sandboxAllocation, + }, + }); + const signing = deferred<[]>(); + orchestrationMocks.signedAttachments.mockImplementationOnce(() => signing.promise); + await fixture.admit('sized'); + await fixture.flush(); + expect(fixture.control.ensureReady).not.toHaveBeenCalled(); + fixture.reload(); + await fixture.fireAlarm(); + expect(fixture.control.ensureReady).toHaveBeenCalledWith( + expect.objectContaining({ provider: 'vercel', resources }) + ); + expect(fixture.record('sized')?.state).toBe('accepted'); + signing.resolve([]); + await fixture.flush(); + } + ); + it('persists the alarm and head budget before the first RPC and wakes the head on a fresh ID after reset', async () => { const fixture = sessionFixture(); const firstReady = deferred(); diff --git a/services/cloud-agent-next/src/session-prepare.test.ts b/services/cloud-agent-next/src/session-prepare.test.ts index 3ab9c579a9..c867e3cfb0 100644 --- a/services/cloud-agent-next/src/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session-prepare.test.ts @@ -1,6 +1,10 @@ import type * as CloudAgentProfile from '@kilocode/cloud-agent-profile'; import type * as SandboxIdModule from './sandbox-id.js'; import { TRPCError } from '@trpc/server'; +import { + getSandboxAllocationRequest, + sandboxAllocationSchema, +} from '@kilocode/worker-utils/sandbox-allocation'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import * as schemas from './router/schemas.js'; @@ -247,6 +251,276 @@ function createInternalApiContext(options: { } as TRPCContext; } +describe('sandbox selection Worker API', () => { + const orgId = 'f47ac10b-58cc-4372-a567-0e02b2c3d479'; + const prepareInput = { + prompt: 'Test prompt', + mode: 'code', + model: 'claude-3', + githubRepo: 'acme/repo', + kilocodeOrganizationId: orgId, + sandboxAllocation: 'cloudflare-single' as const, + }; + const startInput = { + message: { prompt: 'Test prompt' }, + agent: { mode: 'code', model: 'claude-3' }, + repository: { type: 'github' as const, repo: 'acme/repo' }, + options: { kilocodeOrganizationId: orgId }, + runtime: { sandboxAllocation: 'cloudflare-single' as const }, + }; + const allocationInputs = sandboxAllocationSchema.options.flatMap(allocation => [ + { name: allocation, sandboxAllocation: allocation }, + { + name: `structured ${allocation}`, + sandboxAllocation: getSandboxAllocationRequest(allocation), + }, + ]); + + beforeEach(() => { + vi.clearAllMocks(); + organizationMembershipLimitMock.mockResolvedValue([{ id: 'membership' }]); + mergeProfileConfigurationMock.mockResolvedValue({}); + assertKiloModelAvailableMock.mockResolvedValue(undefined); + }); + + it('requires authentication for capability discovery', async () => { + const caller = appRouter.createCaller( + createInternalApiContext({ userId: null, authToken: null }) + ); + await expect( + caller.getSandboxSelectionOptions({ kilocodeOrganizationId: orgId }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + expect(organizationMembershipLimitMock).not.toHaveBeenCalled(); + }); + + it.each(allocationInputs)( + 'requires membership for $name even for disabled options and skip-balance callers', + async ({ sandboxAllocation }) => { + organizationMembershipLimitMock.mockResolvedValue([]); + const caller = appRouter.createCaller(createInternalApiContext({ skipBalanceCheck: true })); + await expect( + caller.getSandboxSelectionOptions({ kilocodeOrganizationId: orgId }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + await expect( + caller.prepareSession({ ...prepareInput, sandboxAllocation }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + await expect( + caller.start({ ...startInput, runtime: { sandboxAllocation } }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mergeProfileConfigurationMock).not.toHaveBeenCalled(); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(createCliSessionMock).not.toHaveBeenCalled(); + } + ); + + it('exposes only authoritative capability options and keeps the disabled result empty', async () => { + const ctx = createInternalApiContext({}); + const caller = appRouter.createCaller(ctx); + await expect( + caller.getSandboxSelectionOptions({ kilocodeOrganizationId: orgId }) + ).resolves.toEqual({ enabled: false, options: [] }); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + ctx.env.CONTROL_PLANE_IDS = orgId; + const result = await caller.getSandboxSelectionOptions({ kilocodeOrganizationId: orgId }); + expect(result.enabled).toBe(true); + expect(result.options.map(option => option.allocation)).toEqual([ + getSandboxAllocationRequest('cloudflare-single'), + getSandboxAllocationRequest('cloudflare-shared'), + ]); + }); + + it.each(allocationInputs)( + 'rejects $name in prepare and public start when selection is disabled', + async ({ sandboxAllocation }) => { + const caller = appRouter.createCaller(createInternalApiContext({ skipBalanceCheck: true })); + await expect( + caller.prepareSession({ ...prepareInput, sandboxAllocation }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + await expect( + caller.start({ ...startInput, runtime: { sandboxAllocation } }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(createCliSessionMock).not.toHaveBeenCalled(); + expect(generateSandboxRoutingTargetMock).not.toHaveBeenCalled(); + } + ); + + it('exposes personal capabilities without membership when the user is enrolled', async () => { + const ctx = createInternalApiContext({}); + ctx.env.SANDBOX_SELECTION_IDS = 'test-user-123'; + const caller = appRouter.createCaller(ctx); + const result = await caller.getSandboxSelectionOptions({}); + expect(result.enabled).toBe(true); + expect(organizationMembershipLimitMock).not.toHaveBeenCalled(); + }); + + it.each(allocationInputs)( + 'rejects personal $name when only an organization is enrolled', + async ({ sandboxAllocation }) => { + const ctx = createInternalApiContext({ skipBalanceCheck: true }); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + const caller = appRouter.createCaller(ctx); + await expect( + caller.prepareSession({ + ...prepareInput, + kilocodeOrganizationId: undefined, + sandboxAllocation, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + await expect( + caller.start({ ...startInput, options: {}, runtime: { sandboxAllocation } }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(createCliSessionMock).not.toHaveBeenCalled(); + expect(generateSandboxRoutingTargetMock).not.toHaveBeenCalled(); + } + ); + + it.each(['*', 'test-user-123'] as const)( + 'authorizes personal Cloudflare allocation when SANDBOX_SELECTION_IDS is %s', + async allowlist => { + const doStub = createMockDOStub(); + const ctx = createInternalApiContext({ doStub, skipBalanceCheck: true }); + ctx.env.SANDBOX_SELECTION_IDS = allowlist; + generateSessionIdMock.mockReturnValue('agent_12345678-1234-1234-1234-123456789abc'); + const caller = appRouter.createCaller(ctx); + await caller.prepareSession({ + ...prepareInput, + kilocodeOrganizationId: undefined, + sandboxAllocation: 'cloudflare-single', + }); + await caller.start({ + ...startInput, + options: {}, + runtime: { sandboxAllocation: 'cloudflare-single' }, + }); + expect(doStub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + workspace: expect.objectContaining({ sandboxAllocation: 'cloudflare-single' }), + }) + ); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + workspace: expect.objectContaining({ sandboxAllocation: 'cloudflare-single' }), + }) + ); + } + ); + + it('returns default destination metadata without allocating, including devcontainer context', async () => { + const ctx = createInternalApiContext({}); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + ctx.env.PER_SESSION_SANDBOX_ORG_IDS = orgId; + const caller = appRouter.createCaller(ctx); + const normal = await caller.getSandboxSelectionOptions({ kilocodeOrganizationId: orgId }); + expect(normal.defaultDestination).toEqual(getSandboxAllocationRequest('cloudflare-single')); + const devcontainer = await caller.getSandboxSelectionOptions({ + kilocodeOrganizationId: orgId, + devcontainer: true, + }); + expect(devcontainer.defaultDestination).toEqual({ + provider: { id: 'cloudflare', account: 'kilo' }, + instanceType: 'devcontainer', + }); + expect(generateSessionIdMock).not.toHaveBeenCalled(); + expect(generateSandboxRoutingTargetMock).not.toHaveBeenCalled(); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(createCliSessionMock).not.toHaveBeenCalled(); + }); + + it.each(['small', 'large'] as const)( + 'rejects BYOC %s before allocation side effects', + async instanceType => { + const ctx = createInternalApiContext({ skipBalanceCheck: true }); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + const caller = appRouter.createCaller(ctx); + const sandboxAllocation = { + provider: { id: 'vercel', account: 'byoc' }, + instanceType, + } as const; + await expect( + caller.prepareSession({ ...prepareInput, sandboxAllocation }) + ).rejects.toMatchObject({ + code: 'BAD_REQUEST', + message: expect.stringContaining('BYOC'), + }); + await expect( + caller.start({ ...startInput, runtime: { sandboxAllocation } }) + ).rejects.toMatchObject({ + code: 'BAD_REQUEST', + message: expect.stringContaining('BYOC'), + }); + expect(mergeProfileConfigurationMock).not.toHaveBeenCalled(); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(createCliSessionMock).not.toHaveBeenCalled(); + expect(generateSandboxRoutingTargetMock).not.toHaveBeenCalled(); + } + ); + + it.each(['webhook', 'scheduled'])( + 'authorizes Dedicated Standard for an enrolled organization %s trigger', + async createdOnPlatform => { + const doStub = createMockDOStub(); + const ctx = createInternalApiContext({ doStub }); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + generateSessionIdMock.mockReturnValue('agent_12345678-1234-1234-1234-123456789abc'); + await appRouter.createCaller(ctx).prepareSession({ + ...prepareInput, + sandboxAllocation: 'isolated-standard', + createdOnPlatform, + }); + expect(doStub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + identity: expect.objectContaining({ orgId, createdOnPlatform }), + workspace: expect.objectContaining({ sandboxAllocation: 'isolated-standard' }), + }) + ); + } + ); + + it('authorizes Dedicated Standard in public start without an internal API key', async () => { + const doStub = createMockDOStub(); + const ctx = createInternalApiContext({ doStub, requestInternalApiKey: null }); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + generateSessionIdMock.mockReturnValue('agent_12345678-1234-1234-1234-123456789abc'); + await appRouter.createCaller(ctx).start({ + ...startInput, + runtime: { sandboxAllocation: 'isolated-standard' }, + }); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + identity: expect.objectContaining({ orgId }), + workspace: expect.objectContaining({ sandboxAllocation: 'isolated-standard' }), + }) + ); + }); + + it.each(['legacy', 'structured'] as const)( + 'keeps canonical workspace metadata for authorized %s starts and prepares', + async format => { + const doStub = createMockDOStub(); + const ctx = createInternalApiContext({ doStub }); + ctx.env.SANDBOX_SESSION = ctx.env + .CLOUD_AGENT_SESSION as unknown as typeof ctx.env.SANDBOX_SESSION; + ctx.env.SANDBOX_SELECTION_IDS = orgId; + ctx.env.CONTROL_PLANE_IDS = orgId; + generateSessionIdMock.mockReturnValue('workspace_12345678-1234-1234-1234-123456789abc'); + const sandboxAllocation = + format === 'legacy' + ? 'cloudflare-single' + : getSandboxAllocationRequest('cloudflare-single'); + const caller = appRouter.createCaller(ctx); + await caller.start({ ...startInput, runtime: { sandboxAllocation } }); + await caller.prepareSession({ ...prepareInput, sandboxAllocation }); + const metadata = expect.objectContaining({ + workspace: expect.objectContaining({ sandboxAllocation: 'cloudflare-single' }), + }); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith(metadata); + expect(doStub.registerSession).toHaveBeenCalledWith(metadata); + } + ); +}); + describe('effective session profile policy', () => { it.each([ { origin: 'cloud-agent-web', expected: 'include-web-defaults' }, diff --git a/services/cloud-agent-next/src/session/session-prepare.test.ts b/services/cloud-agent-next/src/session/session-prepare.test.ts index 3a17b56098..cb218c4cb9 100644 --- a/services/cloud-agent-next/src/session/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session/session-prepare.test.ts @@ -15,6 +15,16 @@ import { } from '@kilocode/worker-utils/runtime-authorization'; import { assertKiloModelAvailable } from '../model-validation.js'; import type { WorkerDb } from '@kilocode/db/client'; +import { + SELECTABLE_SANDBOX_ALLOCATIONS, + getSandboxAllocationRequest, + type SandboxAllocation, +} from '@kilocode/worker-utils/sandbox-allocation'; +import { parseSessionMetadata } from '../persistence/session-metadata.js'; +import { + resolveSharedSandboxAssignment, + SHARED_SANDBOX_FAILOVER_SUFFIX, +} from '../shared-sandbox-route.js'; import type { OperationLedgerRow } from '@kilocode/db/schema'; import type { Env } from '../types.js'; @@ -39,6 +49,7 @@ import { type SessionRegistrationContext, } from './session-registration.js'; import { prepareInputToSessionCreateRequest } from '../router/handlers/session-prepare.js'; +import { StartSessionInput } from '../router/schemas.js'; const { admitOperationMock, @@ -196,8 +207,12 @@ function makeLedgerRow(overrides: Partial = {}): OperationLe } /** Fake Drizzle db: `.limit(1)` returns the next queued result per query. */ -function makeDb(limitResults: unknown[][]): WorkerDb { - const limit = vi.fn(async () => limitResults.shift() ?? []); +function makeDb(limitResults: (unknown[] | Error)[]): WorkerDb { + const limit = vi.fn(async () => { + const result = limitResults.shift() ?? []; + if (result instanceof Error) throw result; + return result; + }); const select = vi.fn(() => ({ from: vi.fn(() => ({ where: vi.fn(() => ({ limit })), @@ -343,6 +358,768 @@ describe('assertSessionOperationIdentity', () => { }); }); +describe('explicit sandbox session creation', () => { + const orgId = 'f47ac10b-58cc-4372-a567-0e02b2c3d479'; + const vercel = { + VERCEL_TOKEN: 'test-token', + VERCEL_TEAM_ID: 'team-id', + VERCEL_PROJECT_ID: 'project-id', + VERCEL_SANDBOX_SNAPSHOT_ID: 'snapshot-id', + VERCEL_SANDBOX_RUNTIME_BUILD_ID: 'build-id', + VERCEL_SANDBOX_RUNTIME: 'node24', + VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: '300000', + VERCEL_SANDBOX_EXTEND_DURATION_MS: '600000', + }; + const requestForPreset = (sandboxAllocation: SandboxAllocation) => + makeRequest({ + runtime: { sandboxAllocation }, + options: { kilocodeOrganizationId: orgId, operationKey: OPERATION_KEY }, + }); + function requestFromStructuredAllocation(allocation: SandboxAllocation): SessionCreateRequest { + const runtime = StartSessionInput.parse({ + message: { prompt: 'Build the feature' }, + agent: { mode: 'code', model: 'claude-3' }, + repository: { type: 'github', repo: 'acme/repo' }, + runtime: { sandboxAllocation: getSandboxAllocationRequest(allocation) }, + }).runtime; + return { ...requestForPreset(allocation), runtime }; + } + function selectedContext(doStub = makeDoStub()) { + const ctx = makeContext(doStub); + Object.assign(ctx.env, vercel, { + CONTROL_PLANE_IDS: orgId, + SANDBOX_SELECTION_IDS: orgId, + PER_SESSION_SANDBOX_ORG_IDS: '*', + VERCEL_SANDBOX_ORG_IDS: '*', + }); + return ctx; + } + function expectNoAllocation(doStub: ReturnType) { + expect(generateSessionIdMock).not.toHaveBeenCalled(); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(recordSandboxIdentityMock).not.toHaveBeenCalled(); + expect(createCliSessionMock).not.toHaveBeenCalled(); + expect(doStub.createSessionWithInitialAdmission).not.toHaveBeenCalled(); + expect(doStub.registerSession).not.toHaveBeenCalled(); + } + + beforeEach(async () => { + vi.clearAllMocks(); + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [{ email: 'test@example.com' }], [{ id: 'member' }]]) + ); + generateSessionIdMock.mockReturnValue(WORKSPACE_SESSION_ID); + generateKiloSessionIdMock.mockReturnValue(KILO_SESSION_ID); + const routing = await vi.importActual('../sandbox-id.js'); + generateSandboxRoutingTargetMock.mockImplementation(routing.generateSandboxRoutingTarget); + vi.mocked(resolveSharedSandboxAssignment).mockImplementation(async (_store, routeKey) => ({ + sandboxId: await routing.deriveSharedSandboxId(routeKey, SHARED_SANDBOX_FAILOVER_SUFFIX), + suffix: SHARED_SANDBOX_FAILOVER_SUFFIX, + })); + admitOperationMock.mockResolvedValue({ + admission: 'admitted', + row: makeLedgerRow({ organization_id: orgId }), + }); + settleOperationMock.mockResolvedValue({ settled: true }); + recordOperationProgressMock.mockResolvedValue(undefined); + }); + + it.each( + SELECTABLE_SANDBOX_ALLOCATIONS.flatMap(preset => [ + { preset, format: 'legacy' }, + { preset, format: 'structured' }, + ]) + )( + 'persists $format $preset with its canonical allocation and valid metadata', + async ({ preset, format }) => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + if (preset.startsWith('vercel-')) { + ctx.env.VERCEL_SANDBOX_ORG_IDS = ''; + ctx.env.PER_SESSION_SANDBOX_ORG_IDS = ''; + } + await runCreate( + ctx, + format === 'structured' ? requestFromStructuredAllocation(preset) : requestForPreset(preset) + ); + expect(recordOperationProgressMock).toHaveBeenCalledWith( + expect.anything(), + ROW_ID, + expect.objectContaining({ + sandboxAllocation: preset, + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: expect.any(String), + }) + ); + const command = doStub.createSessionWithInitialAdmission.mock.calls[0]?.[0]; + const metadata = parseSessionMetadata({ + ...command, + metadataSchemaVersion: 2, + lifecycle: { version: 1, timestamp: 1 }, + }); + expect(metadata.workspace?.sandboxAllocation).toBe(preset); + expect(metadata.workspace?.sandboxProvider).toBe( + preset.startsWith('vercel-') ? 'vercel' : 'cloudflare' + ); + expect(metadata.workspace).not.toHaveProperty('resources'); + if (preset === 'cloudflare-shared') { + expect(resolveSharedSandboxAssignment).toHaveBeenCalledOnce(); + expect(metadata.workspace?.sandboxRoute?.suffix).toBe(SHARED_SANDBOX_FAILOVER_SUFFIX); + expect(metadata.workspace?.sandboxId).not.toBe(metadata.workspace?.sandboxRoute?.routeKey); + } + } + ); + + it.each([ + { preset: 'cloudflare-single', sandboxId: /^ses-/ }, + { preset: 'cloudflare-shared', sandboxId: /^org-/ }, + ] as const)( + 'keeps a legacy owner on the legacy plane for $preset', + async ({ preset, sandboxId }) => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.CONTROL_PLANE_IDS = ''; + generateSessionIdMock.mockReturnValue(CLOUD_AGENT_SESSION_ID); + await runCreate(ctx, requestForPreset(preset)); + expect(generateSessionIdMock).toHaveBeenCalledWith('legacy'); + const command = doStub.createSessionWithInitialAdmission.mock.calls[0]?.[0]; + expect(command?.workspace).toMatchObject({ sandboxAllocation: preset }); + expect(command?.workspace?.sandboxProvider).toBe('cloudflare'); + expect(command?.workspace?.sandboxId).toMatch(sandboxId); + } + ); + + it.each(['vercel-small', 'vercel-large'] as const)( + 'forces the control plane for %s even when the owner is not enrolled', + async preset => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.CONTROL_PLANE_IDS = ''; + ctx.env.VERCEL_SANDBOX_ORG_IDS = ''; + await runCreate(ctx, requestForPreset(preset)); + expect(generateSessionIdMock).toHaveBeenCalledWith('control'); + const command = doStub.createSessionWithInitialAdmission.mock.calls[0]?.[0]; + expect(command?.workspace).toMatchObject({ + sandboxAllocation: preset, + sandboxProvider: 'vercel', + }); + expect(command?.workspace?.sandboxId).toMatch(/^ses-/); + } + ); + + it.each(SELECTABLE_SANDBOX_ALLOCATIONS)( + 'recovers %s on same-key takeover after membership lookup fails before progress', + async sandboxAllocation => { + const request = requestForPreset(sandboxAllocation); + const row = makeLedgerRow({ organization_id: orgId }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + const membershipError = new Error('Membership lookup temporarily unavailable'); + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [{ email: 'test@example.com' }], membershipError]) + ); + admitOperationMock + .mockResolvedValueOnce({ admission: 'admitted', row }) + .mockResolvedValueOnce({ admission: 'takeover', row }); + + await expect(runCreate(ctx, request)).rejects.toBe(membershipError); + expect(admitOperationMock).toHaveBeenCalledOnce(); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expect(settleOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [{ email: 'test@example.com' }], [{ id: 'member' }]]) + ); + await expect(runCreate(ctx, request)).resolves.toEqual({ + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + }); + expect(admitOperationMock).toHaveBeenCalledTimes(2); + expect(admitOperationMock.mock.calls[1]?.[1]).toMatchObject({ + operationKey: OPERATION_KEY, + }); + expect(generateSessionIdMock).toHaveBeenCalledOnce(); + expect(recordOperationProgressMock).toHaveBeenCalledWith( + expect.anything(), + ROW_ID, + expect.objectContaining({ + sandboxAllocation, + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(request), + }) + ); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledOnce(); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( + expect.objectContaining({ workspace: expect.objectContaining({ sandboxAllocation }) }) + ); + expect(settleOperationMock).toHaveBeenCalledOnce(); + expect(settleOptions(0)?.outboxEvent).toMatchObject({ + properties: { outcome: 'completed', admission: 'takeover' }, + }); + } + ); + + it.each([ + { SANDBOX_SELECTION_IDS: '' }, + { SANDBOX_SELECTION_IDS: undefined }, + { SANDBOX_SELECTION_IDS: 'other-org' }, + ])('rejects new overrides before allocating or admitting the ledger: %j', async overrides => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + Object.assign(ctx.env, overrides); + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], []])); + await expect(runCreate(ctx, requestForPreset('cloudflare-single'))).rejects.toMatchObject({ + code: 'FORBIDDEN', + }); + expect(admitOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + }); + + it('requires membership on direct registration and replay', async () => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + getPgDbMock.mockReturnValue(makeDb([])); + await expect( + registerNewSession(requestForPreset('cloudflare-single'), ctx) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + await expect(runCreate(ctx, requestForPreset('cloudflare-single'))).rejects.toMatchObject({ + code: 'FORBIDDEN', + }); + expect(admitOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + }); + + it.each([ + { runtime: { sandboxAllocation: 'custom' } }, + { runtime: { sandboxAllocation: 'cloudflare-single', devcontainer: true } }, + { + runtime: { sandboxAllocation: 'cloudflare-single' }, + options: { createdOnPlatform: 'code-review' }, + }, + ])('rejects invalid or conflicting explicit intent before side effects: %j', async overrides => { + const doStub = makeDoStub(); + await expect( + runCreate(selectedContext(doStub), makeRequest(overrides as Partial)) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(admitOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + }); + + it.each(['vercel-small', 'vercel-large'] as const)( + 'rejects unavailable %s before creation side effects', + async preset => { + for (const unavailable of [ + { VERCEL_TOKEN: '' }, + { + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: orgId, + }, + ]) { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + Object.assign(ctx.env, unavailable); + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], []])); + await expect(runCreate(ctx, requestForPreset(preset))).rejects.toMatchObject({ + code: 'BAD_REQUEST', + }); + expectNoAllocation(doStub); + expect(admitOperationMock).not.toHaveBeenCalled(); + } + } + ); + + it.each(SELECTABLE_SANDBOX_ALLOCATIONS)( + 'creates the first worktree chat with %s and upstream finalization/containment semantics', + async sandboxAllocation => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.WORKTREE_CREATION_ENABLED_IDS = orgId; + ctx.env.CREDENTIAL_CONTAINMENT_ENABLED = 'false'; + const request = requestForPreset(sandboxAllocation); + request.options = { + ...request.options, + createdOnPlatform: 'cloud-agent-web', + clientProvenance: 'browser', + }; + request.finalization = { autoCommit: true }; + await runCreate(ctx, request); + const command = doStub.createSessionWithInitialAdmission.mock.calls[0]?.[0]; + const metadata = parseSessionMetadata({ + ...command, + metadataSchemaVersion: 2, + lifecycle: { version: 1, timestamp: 1 }, + }); + expect(metadata.workspace).toMatchObject({ + worktreeId: WORKTREE_ID, + sandboxAllocation, + credentialContainment: { github: false, gitlab: false, bitbucket: false, kilocode: false }, + }); + expect(metadata.finalization?.autoCommit).toBe(true); + expect(recordOperationProgressMock).toHaveBeenCalledWith( + expect.anything(), + ROW_ID, + expect.objectContaining({ + [SESSION_CREATE_WORKTREE_ENABLED_KEY]: true, + [SESSION_CREATE_FINALIZATION_VERSION_KEY]: 2, + sandboxAllocation, + }) + ); + expect(createCliSessionMock.mock.calls[0]?.slice(-2)).toEqual([ + WORKTREE_ID, + { sandboxId: metadata.workspace?.sandboxId, provider: metadata.workspace?.sandboxProvider }, + ]); + } + ); + + it.each([1, 2] as const)( + 'preserves worktree finalization version %s while replaying after selection is disabled', + async finalizationVersion => { + const request = requestForPreset('vercel-large'); + request.options = { + ...request.options, + createdOnPlatform: 'cloud-agent-web', + clientProvenance: 'browser', + }; + request.finalization = { autoCommit: true }; + const canonicalRequest = + finalizationVersion === 1 ? { ...request, finalization: { autoCommit: false } } : request; + const row = makeLedgerRow({ + organization_id: orgId, + status: 'completed', + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + sandboxAllocation: 'vercel-large', + [SESSION_CREATE_WORKTREE_ENABLED_KEY]: true, + [SESSION_CREATE_FINALIZATION_VERSION_KEY]: finalizationVersion, + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: + await sessionCreateIntentFingerprint(canonicalRequest), + }, + }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = ''; + ctx.env.CONTROL_PLANE_IDS = ''; + ctx.env.WORKTREE_CREATION_ENABLED_IDS = ''; + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], [row]])); + admitOperationMock.mockResolvedValue({ admission: 'duplicate_settled', row }); + await expect(runCreate(ctx, request)).resolves.toMatchObject({ + cloudAgentSessionId: WORKSPACE_SESSION_ID, + replayed: true, + }); + expectNoAllocation(doStub); + } + ); + + it.each([1, 2] as const)( + 'rebuilds a recorded worktree allocation with its preset and finalization version %s', + async finalizationVersion => { + const request = requestForPreset('vercel-large'); + request.options = { + ...request.options, + createdOnPlatform: 'cloud-agent-web', + clientProvenance: 'browser', + }; + request.finalization = { autoCommit: true }; + const canonicalRequest = + finalizationVersion === 1 ? { ...request, finalization: { autoCommit: false } } : request; + const row = makeLedgerRow({ + organization_id: orgId, + status: 'reconcile_pending', + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + initialMessageId: INITIAL_MESSAGE_ID, + sandboxId: `ses-${'a'.repeat(48)}`, + sandboxProvider: 'vercel', + sandboxAllocation: 'vercel-large', + [SESSION_CREATE_WORKTREE_ENABLED_KEY]: true, + [SESSION_CREATE_FINALIZATION_VERSION_KEY]: finalizationVersion, + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: + await sessionCreateIntentFingerprint(canonicalRequest), + }, + }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = ''; + ctx.env.CONTROL_PLANE_IDS = ''; + ctx.env.WORKTREE_CREATION_ENABLED_IDS = ''; + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [row], [], [{ email: 'test@example.com' }]]) + ); + admitOperationMock.mockResolvedValue({ admission: 'duplicate_reconcile_pending', row }); + createCliSessionMock.mockResolvedValueOnce({ status: 'ready' }); + await expect(runCreate(ctx, request)).resolves.toMatchObject({ + cloudAgentSessionId: WORKSPACE_SESSION_ID, + replayed: true, + }); + expect(generateSessionIdMock).not.toHaveBeenCalled(); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + workspace: expect.objectContaining({ + worktreeId: WORKTREE_ID, + sandboxAllocation: 'vercel-large', + sandboxProvider: 'vercel', + }), + finalization: { autoCommit: finalizationVersion === 2 }, + }) + ); + } + ); + + it.each(SELECTABLE_SANDBOX_ALLOCATIONS)( + 'replays a structured %s request against a legacy stored fingerprint', + async sandboxAllocation => { + const legacyRequest = requestForPreset(sandboxAllocation); + const request = requestFromStructuredAllocation(sandboxAllocation); + const fingerprint = await sessionCreateIntentFingerprint(legacyRequest); + expect(await sessionCreateIntentFingerprint(request)).toBe(fingerprint); + const row = makeLedgerRow({ + organization_id: orgId, + status: 'completed', + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + sandboxAllocation, + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: fingerprint, + }, + }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = ''; + ctx.env.CONTROL_PLANE_IDS = ''; + ctx.env.VERCEL_TOKEN = ''; + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], [row]])); + admitOperationMock.mockResolvedValue({ admission: 'duplicate_settled', row }); + await expect(runCreate(ctx, request)).resolves.toEqual({ + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + replayed: true, + }); + expectNoAllocation(doStub); + } + ); + + it('returns the canonical result after both rollout flags are disabled', async () => { + const request = requestForPreset('vercel-large'); + const row = makeLedgerRow({ + organization_id: orgId, + status: 'completed', + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + sandboxAllocation: 'vercel-large', + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(request), + }, + }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = ''; + ctx.env.CONTROL_PLANE_IDS = ''; + ctx.env.VERCEL_TOKEN = ''; + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], [row]])); + admitOperationMock.mockResolvedValue({ admission: 'duplicate_settled', row }); + await expect(runCreate(ctx, request)).resolves.toEqual({ + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + replayed: true, + }); + expectNoAllocation(doStub); + }); + + it.each(['vercel-small', 'cloudflare-single', undefined] as const)( + 'rejects a same-key replay changing or removing the preset to %s', + async sandboxAllocation => { + const row = makeLedgerRow({ + organization_id: orgId, + status: 'completed', + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + sandboxAllocation: 'vercel-large', + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint( + requestForPreset('vercel-large') + ), + }, + }); + const doStub = makeDoStub(); + admitOperationMock.mockResolvedValue({ admission: 'duplicate_settled', row }); + const request = { + ...requestForPreset('vercel-large'), + runtime: sandboxAllocation ? { sandboxAllocation } : undefined, + }; + await expect(runCreate(selectedContext(doStub), request)).rejects.toThrow( + 'session_creation_failed' + ); + expectNoAllocation(doStub); + } + ); + + it('rechecks the flag after admission and before allocation side effects', async () => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + admitOperationMock.mockImplementationOnce(async () => { + ctx.env.SANDBOX_SELECTION_IDS = ''; + return { admission: 'admitted', row: makeLedgerRow({ organization_id: orgId }) }; + }); + await expect(runCreate(ctx, requestForPreset('cloudflare-single'))).rejects.toMatchObject({ + code: 'FORBIDDEN', + }); + expectNoAllocation(doStub); + }); + + it('rejects personal explicit creates when only an organization is enrolled', async () => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + ctx.env.CONTROL_PLANE_IDS = '*'; + getPgDbMock.mockReturnValue(makeDb([[]])); + await expect( + runCreate(ctx, makeRequest({ runtime: { sandboxAllocation: 'cloudflare-single' } })) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(admitOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + }); + + it.each(['*', USER_ID] as const)( + 'authorizes personal explicit creates when SANDBOX_SELECTION_IDS is %s', + async allowlist => { + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = allowlist; + ctx.env.CONTROL_PLANE_IDS = '*'; + generateSessionIdMock.mockReturnValue(CLOUD_AGENT_SESSION_ID); + getPgDbMock.mockReturnValue(makeDb([[{ email: 'test@example.com' }]])); + admitOperationMock.mockResolvedValue({ + admission: 'admitted', + row: makeLedgerRow(), + }); + await expect( + runCreate(ctx, makeRequest({ runtime: { sandboxAllocation: 'cloudflare-single' } })) + ).resolves.toMatchObject({ cloudAgentSessionId: CLOUD_AGENT_SESSION_ID }); + expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + workspace: expect.objectContaining({ sandboxAllocation: 'cloudflare-single' }), + }) + ); + } + ); + + it('does not let an explicit preset inherit a pre-preset canonical result', async () => { + const doStub = makeDoStub(); + admitOperationMock.mockResolvedValue({ + admission: 'duplicate_settled', + row: makeLedgerRow({ + organization_id: orgId, + status: 'completed', + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + }, + }), + }); + await expect( + runCreate(selectedContext(doStub), requestForPreset('cloudflare-single')) + ).rejects.toThrow('session_creation_failed'); + expectNoAllocation(doStub); + }); + + it('rebuilds a clone using its persisted preset after both rollouts change', async () => { + const request = { + ...requestForPreset('vercel-large'), + initialTurn: undefined, + clone: { cloneFromKiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa' }, + }; + const row = makeLedgerRow({ + organization_id: orgId, + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + sandboxId: `ses-${'a'.repeat(48)}`, + sandboxProvider: 'vercel', + sandboxAllocation: 'vercel-large', + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(request), + }, + }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = ''; + ctx.env.CONTROL_PLANE_IDS = ''; + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [row], [], [{ email: 'test@example.com' }]]) + ); + admitOperationMock.mockResolvedValue({ admission: 'takeover', row }); + createCliSessionMock.mockResolvedValueOnce({ status: 'ready' }); + await expect(runCreate(ctx, request)).resolves.toMatchObject({ + cloudAgentSessionId: WORKSPACE_SESSION_ID, + replayed: true, + }); + expect(generateSessionIdMock).not.toHaveBeenCalled(); + expect(doStub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + workspace: expect.objectContaining({ + sandboxAllocation: 'vercel-large', + sandboxProvider: 'vercel', + }), + }) + ); + }); + + it('rejects corrupted clone preset allocation before resuming ownership side effects', async () => { + const request = { + ...requestForPreset('vercel-large'), + initialTurn: undefined, + clone: { cloneFromKiloSessionId: 'ses_aaaaaaaaaaaaaaaaaaaaaaaaaa' }, + }; + const row = makeLedgerRow({ + organization_id: orgId, + canonical_result: { + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + sandboxId: `ses-${'a'.repeat(48)}`, + sandboxProvider: 'cloudflare', + sandboxAllocation: 'vercel-large', + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(request), + }, + }); + const doStub = makeDoStub(); + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], []])); + admitOperationMock.mockResolvedValue({ admission: 'takeover', row }); + await expect(runCreate(selectedContext(doStub), request)).rejects.toThrow( + 'creation_in_progress' + ); + expectNoAllocation(doStub); + }); + + it.each([ + ['empty result', {}], + ['partial session IDs', { cloudAgentSessionId: WORKSPACE_SESSION_ID }], + [ + 'recorded session IDs without intent', + { cloudAgentSessionId: WORKSPACE_SESSION_ID, kiloSessionId: KILO_SESSION_ID }, + ], + ['preset without fingerprint', { sandboxAllocation: 'cloudflare-single' }], + [ + 'preset with empty fingerprint', + { sandboxAllocation: 'cloudflare-single', [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: '' }, + ], + ] as const)( + 'rejects incomplete recorded intent on takeover: %s', + async (_name, canonicalResult) => { + const doStub = makeDoStub(); + const row = makeLedgerRow({ organization_id: orgId, canonical_result: canonicalResult }); + admitOperationMock.mockResolvedValueOnce({ admission: 'takeover', row }); + + await expect( + runCreate(selectedContext(doStub), requestForPreset('cloudflare-single')) + ).rejects.toMatchObject({ code: 'BAD_REQUEST', message: 'session_creation_failed' }); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expect(settleOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + } + ); + + it.each([undefined, 'vercel-large'] as const)( + 'rejects recorded preset %s before session IDs are recorded, even with a matching fingerprint', + async sandboxAllocation => { + const request = requestForPreset('cloudflare-single'); + const row = makeLedgerRow({ + organization_id: orgId, + canonical_result: { + ...(sandboxAllocation ? { sandboxAllocation } : {}), + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(request), + }, + }); + const doStub = makeDoStub(); + admitOperationMock.mockResolvedValueOnce({ admission: 'takeover', row }); + + await expect(runCreate(selectedContext(doStub), request)).rejects.toMatchObject({ + code: 'BAD_REQUEST', + message: 'session_creation_failed', + }); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expect(settleOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + } + ); + + it('rejects a conflicting fingerprint before session IDs are recorded', async () => { + const request = requestForPreset('cloudflare-single'); + const row = makeLedgerRow({ + organization_id: orgId, + canonical_result: { + sandboxAllocation: 'cloudflare-single', + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(request), + }, + }); + const doStub = makeDoStub(); + admitOperationMock.mockResolvedValueOnce({ admission: 'takeover', row }); + + await expect( + runCreate(selectedContext(doStub), { + ...request, + initialTurn: { type: 'prompt', prompt: 'A different prompt' }, + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST', message: 'session_creation_failed' }); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expect(settleOperationMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + }); + + it.each([{ SANDBOX_SELECTION_IDS: '' }, { SANDBOX_SELECTION_IDS: 'other-org' }])( + 'rejects a no-progress takeover after allocation is disabled: %j', + async overrides => { + const row = makeLedgerRow({ organization_id: orgId }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + Object.assign(ctx.env, overrides); + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [row], [{ email: 'test@example.com' }], [{ id: 'member' }]]) + ); + admitOperationMock.mockResolvedValueOnce({ admission: 'takeover', row }); + + await expect(runCreate(ctx, requestForPreset('cloudflare-single'))).rejects.toMatchObject({ + code: 'FORBIDDEN', + }); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + } + ); + + it('rechecks membership before a no-progress takeover allocates a new session', async () => { + const row = makeLedgerRow({ organization_id: orgId }); + const doStub = makeDoStub(); + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], [{ email: 'test@example.com' }], []])); + admitOperationMock.mockResolvedValueOnce({ admission: 'takeover', row }); + + await expect( + runCreate(selectedContext(doStub), requestForPreset('cloudflare-single')) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expectNoAllocation(doStub); + }); + + it('rechecks authorization before a takeover can allocate a new session', async () => { + const request = requestForPreset('cloudflare-single'); + const row = makeLedgerRow({ + organization_id: orgId, + canonical_result: { + sandboxAllocation: 'cloudflare-single', + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(request), + }, + }); + const doStub = makeDoStub(); + const ctx = selectedContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = ''; + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [row], [{ email: 'test@example.com' }], [{ id: 'member' }]]) + ); + admitOperationMock.mockResolvedValue({ admission: 'takeover', row }); + await expect(runCreate(ctx, request)).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expectNoAllocation(doStub); + }); +}); + describe('createSessionWithLedger admission ladder', () => { beforeEach(() => { vi.clearAllMocks(); @@ -575,23 +1352,36 @@ describe('createSessionWithLedger admission ladder', () => { } ); - it('routes and persists isolated Standard allocation for an agent session', async () => { + it('routes and persists isolated Standard allocation for an enrolled organization', async () => { + const orgId = 'f47ac10b-58cc-4372-a567-0e02b2c3d479'; const sandboxId = `istd-${'a'.repeat(48)}` as const; generateSandboxRoutingTargetMock.mockResolvedValueOnce({ kind: 'isolated', sandboxId }); const doStub = makeDoStub(); const ctx = makeContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = orgId; + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [{ id: 'member' }], [{ email: 'test@example.com' }]]) + ); + admitOperationMock.mockResolvedValueOnce({ + admission: 'admitted', + row: makeLedgerRow({ organization_id: orgId }), + }); await runCreate( ctx, makeRequest({ runtime: { sandboxAllocation: 'isolated-standard' }, - options: { operationKey: OPERATION_KEY, createdOnPlatform: 'webhook' }, + options: { + operationKey: OPERATION_KEY, + createdOnPlatform: 'webhook', + kilocodeOrganizationId: orgId, + }, }) ); expect(generateSandboxRoutingTargetMock).toHaveBeenCalledWith( undefined, - undefined, + orgId, USER_ID, CLOUD_AGENT_SESSION_ID, undefined, @@ -599,7 +1389,7 @@ describe('createSessionWithLedger admission ladder', () => { ); expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( expect.objectContaining({ - identity: expect.objectContaining({ createdOnPlatform: 'webhook' }), + identity: expect.objectContaining({ orgId, createdOnPlatform: 'webhook' }), workspace: expect.objectContaining({ sandboxId, sandboxProvider: 'cloudflare', @@ -609,6 +1399,40 @@ describe('createSessionWithLedger admission ladder', () => { ); }); + it('replays a settled organization Standard allocation after enrollment is removed', async () => { + const orgId = 'f47ac10b-58cc-4372-a567-0e02b2c3d479'; + const input = makeRequest({ + runtime: { sandboxAllocation: 'isolated-standard' }, + options: { + operationKey: OPERATION_KEY, + createdOnPlatform: 'webhook', + kilocodeOrganizationId: orgId, + }, + }); + const row = makeLedgerRow({ + organization_id: orgId, + status: 'completed', + outcome_code: 'ok', + canonical_result: { + cloudAgentSessionId: CLOUD_AGENT_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + initialMessageId: INITIAL_MESSAGE_ID, + [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(input), + }, + }); + const doStub = makeDoStub(); + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }], [row]])); + admitOperationMock.mockResolvedValueOnce({ admission: 'duplicate_settled', row }); + await expect(runCreate(makeContext(doStub), input)).resolves.toMatchObject({ + cloudAgentSessionId: CLOUD_AGENT_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + replayed: true, + }); + expect(createCliSessionMock).not.toHaveBeenCalled(); + expect(generateSandboxRoutingTargetMock).not.toHaveBeenCalled(); + expect(doStub.createSessionWithInitialAdmission).not.toHaveBeenCalled(); + }); + it.each([ 'admitted', 'duplicate_settled', @@ -2881,6 +3705,8 @@ describe('createSessionWithLedger changed-intent rejection', () => { ) { const doStub = makeDoStub(); const ctx = makeContext(doStub); + ctx.env.SANDBOX_SELECTION_IDS = ORIGINAL_OPTIONS.kilocodeOrganizationId; + getPgDbMock.mockReturnValue(makeDb([[{ id: 'member' }]])); await expect(runCreate(ctx, request)).rejects.toMatchObject( identityConflict @@ -3419,24 +4245,37 @@ describe('createSessionWithLedger clone allocation outcomes', () => { }); it('allows isolated Standard allocation for enrolled non-interactive sessions', async () => { + const orgId = 'f47ac10b-58cc-4372-a567-0e02b2c3d479'; const sandboxId = `istd-${'a'.repeat(48)}` as const; generateSandboxRoutingTargetMock.mockResolvedValueOnce({ kind: 'isolated', sandboxId }); const doStub = makeDoStub(); const ctx = makeContext(doStub); ctx.env.CONTROL_PLANE_IDS = USER_ID; + ctx.env.SANDBOX_SELECTION_IDS = orgId; + getPgDbMock.mockReturnValue( + makeDb([[{ id: 'member' }], [{ id: 'member' }], [{ email: 'test@example.com' }]]) + ); + admitOperationMock.mockResolvedValue({ + admission: 'admitted', + row: makeLedgerRow({ organization_id: orgId }), + }); await runCreate( ctx, makeRequest({ runtime: { sandboxAllocation: 'isolated-standard' }, - options: { operationKey: OPERATION_KEY, createdOnPlatform: 'slack' }, + options: { + operationKey: OPERATION_KEY, + createdOnPlatform: 'slack', + kilocodeOrganizationId: orgId, + }, }) ); expect(generateSessionIdMock).toHaveBeenCalledWith('legacy'); expect(doStub.createSessionWithInitialAdmission).toHaveBeenCalledWith( expect.objectContaining({ - identity: expect.objectContaining({ sessionId: CLOUD_AGENT_SESSION_ID }), + identity: expect.objectContaining({ sessionId: CLOUD_AGENT_SESSION_ID, orgId }), workspace: expect.objectContaining({ sandboxAllocation: 'isolated-standard' }), }) ); @@ -3470,7 +4309,7 @@ describe('createSessionWithLedger clone allocation outcomes', () => { ) ).rejects.toMatchObject({ code: 'BAD_REQUEST', - message: 'Isolated Standard allocation is incompatible with specialized sandbox routing', + message: 'Sandbox allocations cannot be combined with specialized sandbox routing', }); expect(admitOperationMock).not.toHaveBeenCalled(); @@ -4364,6 +5203,22 @@ describe('prepareInputToSessionCreateRequest clone mapping', () => { expect(request.clone).toBeUndefined(); }); + it.each([...SELECTABLE_SANDBOX_ALLOCATIONS, 'isolated-standard', undefined] as const)( + 'maps %s into grouped runtime without introducing provider resources', + sandboxAllocation => { + const request = prepareInputToSessionCreateRequest({ + prompt: 'Continue', + mode: 'code', + model: 'claude-3', + githubRepo: 'acme/repo', + shallow: false, + devcontainer: false, + sandboxAllocation, + }); + expect(request.runtime).toEqual(sandboxAllocation ? { sandboxAllocation } : undefined); + } + ); + it('maps isolated Standard allocation into grouped runtime intent', () => { const request = prepareInputToSessionCreateRequest({ prompt: 'Continue', diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index 6444bd4c78..ac00e0abb2 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -18,7 +18,19 @@ import { TRPCError } from '@trpc/server'; import { and, eq } from 'drizzle-orm'; import { z } from 'zod'; import type { WorkerDb } from '@kilocode/db/client'; -import { cli_sessions_v2, kilocode_users } from '@kilocode/db/schema'; +import { cli_sessions_v2, kilocode_users, operation_ledgers } from '@kilocode/db/schema'; +import { + isSelectableSandboxAllocation, + sandboxAllocationRequiresControlPlane, + sandboxAllocationSchema, + type SandboxAllocation, +} from '@kilocode/worker-utils/sandbox-allocation'; +import { + assertSandboxAllocationAvailable, + getSandboxSelectionCapabilities, + isSandboxAllocationAvailable, +} from '../sandbox-selection.js'; +import { assertOrganizationMembership } from '../router/handlers/organization-membership.js'; import { admitOperation, markReconcilePending, @@ -43,6 +55,7 @@ import type { Env, SandboxId } from '../types.js'; import type { CloudAgentSession } from '../persistence/CloudAgentSession.js'; import { getControlPlaneCredentialContainment, + CurrentSessionMetadataSchema, type CredentialContainment, type SessionMetadata, } from '../persistence/session-metadata.js'; @@ -84,24 +97,50 @@ type SharedSandboxRouteMetadata = NonNullable< NonNullable['sandboxRoute'] >; +/** + * The plane a new session will be created on. Vercel sandboxes exist only on the + * control plane, so those allocations force it; every other request — including a + * Cloudflare allocation — defers to `sessionPlaneForNewOwner`. Single source of + * truth: the allocation checks and the session-ID generation must agree. + */ +function sessionPlaneForCreate( + input: SessionRegistrationInput, + ctx: SessionRegistrationContext +): SessionPlane { + return sandboxAllocationRequiresControlPlane(input.runtime?.sandboxAllocation) + ? 'control' + : sessionPlaneForNewOwner( + ctx.env, + { + userId: ctx.userId, + orgId: input.options?.kilocodeOrganizationId, + }, + { createdOnPlatform: input.options?.createdOnPlatform } + ); +} + function assertSupportedSandboxAllocation( input: SessionRegistrationInput, ctx: SessionRegistrationContext, options?: { billingOrigin?: string } ): void { + const allocation = input.runtime?.sandboxAllocation; + if (allocation === undefined) return; + if (!sandboxAllocationSchema.safeParse(allocation).success) { + throw new TRPCError({ code: 'BAD_REQUEST', message: 'Invalid sandbox allocation' }); + } if ( - input.runtime?.sandboxAllocation === 'isolated-standard' && - (input.runtime.devcontainer === true || options?.billingOrigin === 'code-review') + input.runtime?.devcontainer === true || + options?.billingOrigin === 'code-review' || + input.options?.createdOnPlatform === 'code-review' ) { throw new TRPCError({ code: 'BAD_REQUEST', - message: 'Isolated Standard allocation is incompatible with specialized sandbox routing', + message: 'Sandbox allocations cannot be combined with specialized sandbox routing', }); } - if ( - input.runtime?.sandboxAllocation === 'isolated-standard' && - sessionPlaneForCreate(input, ctx) === 'control' - ) { + // Isolated Standard predates the selectable allocations and remains legacy-plane only. + if (allocation === 'isolated-standard' && sessionPlaneForCreate(input, ctx) === 'control') { throw new TRPCError({ code: 'BAD_REQUEST', message: 'Isolated Standard allocation is not supported for control-plane sessions', @@ -123,6 +162,7 @@ export type SessionRegistrationResult = { sandboxRoute?: SharedSandboxRouteMetadata; sandboxProvider: SandboxSelection['provider']; worktreeId?: CloudAgentWorktreeId; + sandboxAllocation?: SandboxAllocation; /** * Canonical initial turn reserved for a later legacy initiation request. * Omitted for a clone-only create, which has no synthetic initial turn. @@ -486,17 +526,6 @@ function worktreeEnabledForCreate( return sessionPlaneForCreate(input, ctx) === 'control' && isWorktreeOwner(ctx.env, owner); } -function sessionPlaneForCreate( - input: SessionRegistrationInput, - ctx: SessionRegistrationContext -): SessionPlane { - return sessionPlaneForNewOwner( - ctx.env, - { userId: ctx.userId, orgId: input.options?.kilocodeOrganizationId }, - { createdOnPlatform: input.options?.createdOnPlatform } - ); -} - export function assertRuntimeIsolationAdmission(env: Pick): void { if (env.RUNTIME_ISOLATION_ENABLED === 'true') return; throw new TRPCError({ code: 'PRECONDITION_FAILED', message: 'runtime_isolation_unavailable' }); @@ -572,15 +601,30 @@ async function issueSessionRuntimeAuthorization( return runtimeAuthorization; } +async function assertSandboxAllocationMembership( + input: SessionRegistrationInput, + ctx: SessionRegistrationContext +): Promise { + if (input.runtime?.sandboxAllocation === undefined) return; + const orgId = input.options?.kilocodeOrganizationId; + if (!orgId) return; + await assertOrganizationMembership(getPgDb(ctx.env), ctx.userId, orgId); +} + async function allocateNewSession( input: SessionRegistrationInput, ctx: SessionRegistrationContext, options?: { billingOrigin?: string }, ledger?: SessionCreationLedgerHooks ): Promise { + await assertSandboxAllocationMembership(input, ctx); + const sandboxAllocation = input.runtime?.sandboxAllocation; + const orgId = input.options?.kilocodeOrganizationId; + if (sandboxAllocation !== undefined) { + assertSandboxAllocationAvailable(ctx.env, { userId: ctx.userId, orgId }, sandboxAllocation); + } const sessionService = new SessionService(); const initialTurn = input.initialTurn ? acceptInitialTurn(input.initialTurn) : undefined; - const orgId = input.options?.kilocodeOrganizationId; const cloudAgentSessionId = generateSessionId(sessionPlaneForCreate(input, ctx)); const kiloSessionId = generateKiloSessionId(); const reportingCreatedAt = @@ -619,6 +663,7 @@ async function allocateNewSession( [SESSION_CREATE_WORKTREE_ENABLED_KEY]: worktreeId !== undefined, [SESSION_CREATE_FINALIZATION_VERSION_KEY]: ledger.finalizationVersion, [SESSION_CREATE_INTENT_FINGERPRINT_KEY]: await sessionCreateIntentFingerprint(input), + ...(isSelectableSandboxAllocation(sandboxAllocation) ? { sandboxAllocation } : {}), }); } @@ -646,7 +691,7 @@ async function allocateNewSession( { devcontainer: input.runtime?.devcontainer, createdOnPlatform: options?.billingOrigin === 'code-review' ? 'code-review' : undefined, - sandboxAllocation: input.runtime?.sandboxAllocation, + sandboxAllocation, } ); if (target.kind === 'shared') { @@ -668,6 +713,7 @@ async function allocateNewSession( sandboxId, sessionId: cloudAgentSessionId, devcontainer: input.runtime?.devcontainer, + sandboxAllocation, }); } } catch (error) { @@ -804,6 +850,7 @@ async function allocateNewSession( sandboxRoute, sandboxProvider, ...(worktreeId ? { worktreeId } : {}), + ...(sandboxAllocation ? { sandboxAllocation } : {}), initialTurn, reportingCreatedAt, credentialContainment, @@ -851,6 +898,27 @@ function rebuildRecordedSessionAllocation( const sandboxId = canonical.sandboxId; const sandboxProvider = canonical.sandboxProvider; const sandboxRoute = canonical.sandboxRoute; + const requested = isSelectableSandboxAllocation(input.runtime?.sandboxAllocation) + ? input.runtime?.sandboxAllocation + : undefined; + const recorded = sandboxAllocationSchema.optional().safeParse(canonical.sandboxAllocation); + if (!recorded.success || recorded.data !== requested) { + throw creationInProgressError(); + } + if (recorded.data !== undefined) { + const metadata = CurrentSessionMetadataSchema.safeParse({ + metadataSchemaVersion: 2, + identity: { + sessionId: cloudAgentSessionId, + userId: ctx.userId, + orgId: input.options?.kilocodeOrganizationId, + }, + auth: {}, + workspace: { sandboxId, sandboxProvider, sandboxRoute, sandboxAllocation: recorded.data }, + lifecycle: { version: 1, timestamp: Date.now() }, + }); + if (!metadata.success) throw creationInProgressError(); + } const reportingCreatedAt = z .string() .datetime({ offset: true }) @@ -917,6 +985,7 @@ function rebuildRecordedSessionAllocation( sandboxRoute: route, sandboxProvider, ...(worktreeId ? { worktreeId } : {}), + ...(recorded.data ? { sandboxAllocation: recorded.data } : {}), initialTurn, reportingCreatedAt: !initialTurn && cloudAgentSessionId.startsWith('agent_') @@ -994,6 +1063,7 @@ function buildSessionRegistrationCommand( workspace: { sandboxId: allocation.sandboxId, sandboxProvider: allocation.sandboxProvider, + ...(allocation.sandboxAllocation ? { sandboxAllocation: allocation.sandboxAllocation } : {}), shallow: input.options?.shallow, ...(allocation.worktreeId ? { @@ -1153,6 +1223,7 @@ async function registerAndAdmitInitialTurn( kiloSessionId: allocation.kiloSessionId, sandboxId: allocation.sandboxId, sandboxProvider: allocation.sandboxProvider, + ...(allocation.sandboxAllocation ? { sandboxAllocation: allocation.sandboxAllocation } : {}), admission, }; if (ledger) { @@ -1474,7 +1545,17 @@ async function assertCreateIntentUnchanged( input: SessionRegistrationInput, row: OperationLedgerRow ): Promise { + if (row.canonical_result === null) return; const stored = row.canonical_result?.[SESSION_CREATE_INTENT_FINGERPRINT_KEY]; + const selectable = isSelectableSandboxAllocation(input.runtime?.sandboxAllocation) + ? input.runtime?.sandboxAllocation + : undefined; + if ( + row.canonical_result?.sandboxAllocation !== selectable || + (selectable !== undefined && (typeof stored !== 'string' || stored.length === 0)) + ) { + throw new TRPCError({ code: 'BAD_REQUEST', message: 'session_creation_failed' }); + } if (typeof stored !== 'string' || stored.length === 0) { return; } @@ -1497,7 +1578,47 @@ export async function createSessionWithLedger( options: SessionLedgerCreateOptions ): Promise { assertSupportedSandboxAllocation(input, ctx, { billingOrigin: options.billingOrigin }); + await assertSandboxAllocationMembership(input, ctx); const db = getPgDb(ctx.env); + const allocation = input.runtime?.sandboxAllocation; + if (allocation !== undefined) { + const owner = { userId: ctx.userId, orgId: input.options?.kilocodeOrganizationId }; + const available = isSandboxAllocationAvailable( + getSandboxSelectionCapabilities(ctx.env, owner), + allocation + ); + if (!available) { + const [existing] = await db + .select() + .from(operation_ledgers) + .where( + and( + eq(operation_ledgers.kilo_user_id, ctx.userId), + eq(operation_ledgers.domain, 'session'), + eq(operation_ledgers.operation_key, options.operationKey) + ) + ) + .limit(1); + if (!existing || new Date(existing.expires_at).getTime() <= Date.now()) { + assertSandboxAllocationAvailable(ctx.env, owner, allocation); + } else { + assertSessionOperationIdentity(existing, { + userId: ctx.userId, + intent: 'create_cloud', + organizationId: input.options?.kilocodeOrganizationId, + resourceKey: null, + }); + await assertCreateIntentUnchanged( + effectiveSessionRegistrationInput( + input, + worktreeEnabledForCreate(input, ctx, existing), + finalizationVersionForCreate(existing) + ), + existing + ); + } + } + } const admission = await admitOperation(db, { userId: ctx.userId, orgId: input.options?.kilocodeOrganizationId, @@ -1737,6 +1858,7 @@ async function resumeCloneCreate( throw creationInProgressError(); } + const allocation = rebuildRecordedSessionAllocation(input, ctx, row); const hooks = await buildLedgerHooks(input, ctx, options, db, row, 'takeover'); const sessionService = new SessionService(); const createdOnPlatform = input.options?.createdOnPlatform ?? 'cloud-agent'; @@ -1804,7 +1926,6 @@ async function resumeCloneCreate( } // `ready` continues. - const allocation = rebuildRecordedSessionAllocation(input, ctx, row); allocation.runtimeAuthorization = await issueSessionRuntimeAuthorization( input, ctx, diff --git a/services/cloud-agent-next/src/session/session-requests.ts b/services/cloud-agent-next/src/session/session-requests.ts index c7798df89f..409ea24535 100644 --- a/services/cloud-agent-next/src/session/session-requests.ts +++ b/services/cloud-agent-next/src/session/session-requests.ts @@ -1,3 +1,4 @@ +import type { SandboxAllocation } from '@kilocode/worker-utils/sandbox-allocation'; import type { CallbackTarget } from '../callbacks/index.js'; import type { AgentSelection, @@ -44,8 +45,8 @@ export type SessionRepositoryRequest = }; export type SessionRuntimeIntent = { + sandboxAllocation?: SandboxAllocation; devcontainer?: boolean; - sandboxAllocation?: 'isolated-standard'; }; export type SessionCreateRequest = { diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index 719c44a8da..bfb59709f7 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -605,6 +605,8 @@ export type Env = { CONTROL_PLANE_IDS?: string; WORKTREE_CREATION_ENABLED_IDS?: string; RUNTIME_ISOLATION_ENABLED?: string; + /** Comma-separated user or org IDs allowed to pick a sandbox destination. `*` includes personal. */ + SANDBOX_SELECTION_IDS?: string; CREDENTIAL_CONTAINMENT_ENABLED?: string; /** Comma-separated org IDs that receive workspace repo snapshots, or '*' for all */ REPO_SNAPSHOT_ORG_IDS?: string; diff --git a/services/cloud-agent-next/test/integration/sandbox-control.test.ts b/services/cloud-agent-next/test/integration/sandbox-control.test.ts index b6a09dd4a1..ff669044db 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -15,6 +15,10 @@ import { type WorktreeFileRecord, type WorktreeSnapshotCapture, } from '@kilocode/worker-utils/cloud-agent-worktree-changes'; +import { + getSandboxAllocationResources, + type SandboxAllocation, +} from '@kilocode/worker-utils/sandbox-allocation'; import { drizzle } from 'drizzle-orm/durable-sqlite'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import type { BillingContext } from '@kilocode/container-usage'; @@ -27,7 +31,10 @@ import type { VercelSandboxNetworkPolicy, VercelSandboxSession, } from '../../src/agent-sandbox/vercel/vercel-sandbox-rest-client.js'; -import { parseVercelSandboxRuntimeConfig } from '../../src/agent-sandbox/vercel/vercel-runtime-config.js'; +import { + parseVercelSandboxRuntimeConfig, + resolveVercelSandboxRuntimeConfig, +} from '../../src/agent-sandbox/vercel/vercel-runtime-config.js'; import { TRPCError } from '@trpc/server'; import { router } from '../../src/router/auth.js'; import { createSessionManagementHandlers } from '../../src/router/handlers/session-management.js'; @@ -966,6 +973,10 @@ function fakeCloudflareContainers(readPhysical: () => Promise) { function fakeVercelRuntime(sandboxName: string, readPhysical: () => Promise) { const runtime = { creates: 0, + createInputs: [] as Parameters[0][], + inspectInputs: [] as Parameters[0][], + loseCreateResponse: false, + readPhysical, launches: [] as WrapperLaunch[], policy: undefined as VercelSandboxNetworkPolicy | undefined, stoppedSessions: [] as string[], @@ -994,6 +1005,7 @@ function fakeVercelRuntime(sandboxName: string, readPhysical: () => Promise Promise Promise Promise - createVercelProviderAdapter({ sandboxName: allocationName, config, restClient: client }), + createAdapter: ( + allocationName: string, + persisted?: NonNullable['vercel'] + ) => + createVercelProviderAdapter({ + sandboxName: allocationName, + config: resolveVercelSandboxRuntimeConfig(VERCEL_ENV, persisted), + restClient: client, + }), }; } @@ -1097,7 +1119,8 @@ async function registerCredentialSession(registration: CredentialRegistration) { async function credentialFixture( provider: AgentSandboxProvider = 'cloudflare', - id: SandboxId = `${provider === 'vercel' ? 'ses' : 'usr'}-${crypto.randomUUID().replaceAll('-', '')}` + id: SandboxId = `${provider === 'vercel' ? 'ses' : 'usr'}-${crypto.randomUUID().replaceAll('-', '').padEnd(48, '0')}`, + sandboxAllocation?: SandboxAllocation ) { const control = env.SANDBOX_CONTROL.getByName(id); const broker = fakeCredentialBroker(); @@ -1122,7 +1145,10 @@ async function credentialFixture( const runtime = vercel; Object.assign(instance, { createProviderAdapter: (_kind: AgentSandboxProvider, physical?: PhysicalRecord) => - runtime.createAdapter(physical?.createIntent?.allocationName ?? id), + runtime.createAdapter( + physical?.createIntent?.allocationName ?? id, + physical?.createIntent?.vercel + ), }); } }); @@ -1144,6 +1170,10 @@ async function credentialFixture( workspace: { sandboxId: id, sandboxProvider: provider, + ...(sandboxAllocation ? { sandboxAllocation } : {}), + ...(sandboxAllocation === 'cloudflare-shared' + ? { sandboxRoute: { kind: 'shared' as const, routeKey: id } } + : {}), worktreeId: WORKTREE_ID, workspacePath: '/workspace/joined', }, @@ -1968,6 +1998,175 @@ describe('SandboxControl Vercel network policy updates', () => { }); describe('SandboxControl contained Vercel lifecycle', () => { + it.each(['vercel-small', 'vercel-large'] as const)( + 'cleans up an exclusive %s worktree using its pinned resources', + async sandboxAllocation => { + const { control, registration, environment, sandboxId, vercel } = await credentialFixture( + 'vercel', + undefined, + sandboxAllocation + ); + Object.assign(environment, { + SESSION_INGEST: { + canDestroyCloudAgentWorktreeSandbox: async () => ({ kind: 'exclusive' }), + }, + }); + await control.ensureReady({ + ...credentialInput(registration), + provider: 'vercel', + resources: getSandboxAllocationResources(sandboxAllocation), + allowCreate: true, + }); + const physical = await control.getPhysicalRecord(); + const clock = vi + .spyOn(Date, 'now') + .mockReturnValue((physical.createIntent?.createdAt ?? 0) + DEADLINE_MS.createSettle + 1); + try { + await expect( + control.deleteWorktreeResources({ + worktreeId: WORKTREE_ID, + kiloUserId: registration.identity.userId, + organizationId: registration.identity.orgId, + location: { sandboxId, provider: 'vercel' }, + sessionIds: [ROOT_ID], + }) + ).resolves.toEqual({ deleted: true, sessionIds: [ROOT_ID] }); + } finally { + clock.mockRestore(); + } + expect(vercel.runtime.creates).toBe(1); + expect(vercel.runtime.stoppedSessions).toEqual(['vsess_joined_1']); + expect((await control.getPhysicalRecord()).state).toBe('stopped'); + await runInDurableObject(control, async (_instance, state) => { + expect(await state.storage.get('provider_configuration')).toBeUndefined(); + expect(await state.storage.get('provider_locator')).toBeUndefined(); + }); + } + ); + + it.each([undefined, 'vercel-small', 'vercel-large'] as const)( + 'retains %s resources through an uncertain create, object resets, inspection, and replacement', + async sandboxAllocation => { + const fixture = await credentialFixture('vercel', undefined, sandboxAllocation); + const { vercel, registration, environment, sandboxId } = fixture; + let control = fixture.control; + const resources = getSandboxAllocationResources(sandboxAllocation); + const input = { + ...credentialInput(registration), + provider: 'vercel' as const, + resources, + allowCreate: true, + }; + vercel.runtime.loseCreateResponse = true; + await expect(control.ensureReady(input)).resolves.toMatchObject({ physical: 'failed' }); + const uncertain = await control.getPhysicalRecord(); + expect(uncertain.providerRef).toBeNull(); + expect(uncertain.createIntent?.vercel?.resources).toEqual(resources); + expect(vercel.runtime.createInputs[0]?.resources).toEqual(resources); + expect(vercel.runtime.launches).toHaveLength(0); + + const restart = async () => { + await abortAllDurableObjects(); + control = env.SANDBOX_CONTROL.getByName(sandboxId); + await runInDurableObject(control, async (instance, state) => { + const physical = await instance.getPhysicalRecord(); + expect(await state.storage.get('provider_configuration')).toEqual({ + provider: 'vercel', + ...(resources ? { resources } : {}), + }); + vercel.runtime.readPhysical = () => instance.getPhysicalRecord(); + const createProviderAdapter = (_kind: AgentSandboxProvider, value?: PhysicalRecord) => + vercel.createAdapter( + value?.createIntent?.allocationName ?? sandboxId, + value?.createIntent?.vercel + ); + Object.assign(instance, { + env: environment, + createProviderAdapter, + provider: createProviderAdapter('vercel', physical), + }); + }); + }; + await restart(); + expect((await control.getPhysicalRecord()).createIntent).toEqual(uncertain.createIntent); + const clock = vi + .spyOn(Date, 'now') + .mockReturnValue((uncertain.createIntent?.createdAt ?? 0) + DEADLINE_MS.createSettle + 1); + try { + await fireControlDeadline(control, 'stopAttempt'); + } finally { + clock.mockRestore(); + } + expect(vercel.runtime.inspectInputs).toHaveLength(1); + expect(vercel.runtime.inspectInputs[0]).toMatchObject({ + name: uncertain.createIntent?.allocationName, + operationId: uncertain.createIntent?.intentId, + }); + expect(vercel.runtime.inspectInputs[0]?.resources).toEqual(resources); + expect(vercel.runtime.creates).toBe(1); + expect(vercel.runtime.stoppedSessions).toEqual(['vsess_joined_1']); + await expect(control.getPhysicalRecord()).resolves.toMatchObject({ + state: 'stopped', + createIntent: null, + }); + await restart(); + await expect(async () => + control.ensureReady({ + ...input, + resources: + sandboxAllocation === 'vercel-large' + ? { vcpus: 2, memory: 4096 } + : { vcpus: 4, memory: 8192 }, + }) + ).rejects.toThrow('Sandbox resources mismatch'); + vercel.runtime.loseCreateResponse = false; + await expect(control.ensureReady(input)).resolves.toMatchObject({ physical: 'running' }); + const replacement = await control.getPhysicalRecord(); + expect(replacement.createIntent?.vercel?.resources).toEqual(resources); + expect(replacement.createIntent?.allocationName).not.toBe( + uncertain.createIntent?.allocationName + ); + expect(vercel.runtime.createInputs).toHaveLength(2); + expect(vercel.runtime.createInputs[1]?.resources).toEqual(resources); + expect(vercel.runtime.launches).toHaveLength(1); + } + ); + + it.each([false, true])( + 'shares compatible Cloudflare allocation when explicit selection comes first: %s', + async explicitFirst => { + const fixture = await credentialFixture( + 'cloudflare', + undefined, + explicitFirst ? 'cloudflare-shared' : undefined + ); + const { control, registration, containers } = fixture; + await control.ensureReady({ + ...credentialInput(registration), + provider: 'cloudflare', + allowCreate: true, + }); + const original = await control.getPhysicalRecord(); + const sibling = await registerSiblingWorktree({ + ...registration, + workspace: { + ...registration.workspace, + sandboxAllocation: explicitFirst ? undefined : 'cloudflare-shared', + sandboxRoute: { kind: 'shared', routeKey: fixture.sandboxId }, + }, + }); + await expect( + control.ensureReady({ + ...credentialInput(sibling), + provider: 'cloudflare', + allowCreate: true, + }) + ).resolves.toMatchObject({ physical: 'running' }); + expect((await control.getPhysicalRecord()).createIntent).toEqual(original.createIntent); + expect(containers.launches).toHaveLength(1); + } + ); + it.each(['malformed', 'cross-sandbox'] as const)( 'rejects a %s Vercel handshake before binding a creating instance', async identityKind => { diff --git a/services/cloud-agent-next/worker-configuration.d.ts b/services/cloud-agent-next/worker-configuration.d.ts index b488f3b0d7..53db547f91 100644 --- a/services/cloud-agent-next/worker-configuration.d.ts +++ b/services/cloud-agent-next/worker-configuration.d.ts @@ -26,6 +26,7 @@ interface __BaseEnv_Env { CONTROL_PLANE_IDS?: "*"; WORKTREE_CREATION_ENABLED_IDS?: "*"; RUNTIME_ISOLATION_ENABLED: "true"; + SANDBOX_SELECTION_IDS?: "*"; VERCEL_SANDBOX_ORG_IDS: ""; VERCEL_PROJECT_ID: ""; VERCEL_TEAM_ID: ""; @@ -100,6 +101,7 @@ declare namespace Cloudflare { CONTROL_PLANE_IDS: "*"; WORKTREE_CREATION_ENABLED_IDS: "*"; RUNTIME_ISOLATION_ENABLED: "true"; + SANDBOX_SELECTION_IDS: "*"; VERCEL_SANDBOX_ORG_IDS: ""; VERCEL_PROJECT_ID: ""; VERCEL_TEAM_ID: ""; @@ -149,7 +151,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } declare module "*.sql" { const value: string; diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index 97d29c855b..b72fbd946c 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -457,6 +457,7 @@ "CONTROL_PLANE_IDS": "*", "WORKTREE_CREATION_ENABLED_IDS": "*", "RUNTIME_ISOLATION_ENABLED": "true", + "SANDBOX_SELECTION_IDS": "*", "VERCEL_SANDBOX_ORG_IDS": "", "VERCEL_PROJECT_ID": "", "VERCEL_TEAM_ID": "",