diff --git a/.dependency-cruiser.cjs b/.dependency-cruiser.cjs index 31d28dee..9930d7ef 100644 --- a/.dependency-cruiser.cjs +++ b/.dependency-cruiser.cjs @@ -40,6 +40,17 @@ module.exports = { from: { path: "^apps/" }, to: { path: "^packages/db/(src|dist)/schema(/|$)" }, }, + { + name: "gateway-quota-runtime-only-through-do-shell", + severity: "error", + from: { + path: "^apps/gateway-worker/src/", + pathNot: "^apps/gateway-worker/src/durable-objects/quota-tracker\\.ts$", + }, + to: { + path: "^(@cheatcode/billing/quota-runtime|packages/billing/(src|dist)/quota-runtime\\.(js|ts|d\\.ts))$", + }, + }, ], options: { doNotFollow: { path: "node_modules" }, diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-account-deletion.ts b/apps/agent-worker/src/durable-objects/project-sandbox-account-deletion.ts new file mode 100644 index 00000000..96e2915f --- /dev/null +++ b/apps/agent-worker/src/durable-objects/project-sandbox-account-deletion.ts @@ -0,0 +1,132 @@ +import { APIError } from "@cheatcode/observability"; +import type { DaytonaClient, DaytonaSandbox, SandboxDestroyResult } from "@cheatcode/tools-code"; +import { z } from "zod"; +import { clearWorkspaceCommand } from "./project-sandbox-cleanup-script"; +import type { ProjectSandboxIdentityState } from "./project-sandbox-identity-state"; +import { + ACCOUNT_DELETION_TOMBSTONE_KEY, + DAYTONA_ID_KEY, + type ProjectSandboxEnv, + parseSandboxJson, + toUpstreamError, + uniqueSandboxes, +} from "./project-sandbox-lifecycle-support"; +import { + clearSandboxMeterState, + finalizeSandboxUsageBestEffort, + type SandboxMeteringContext, +} from "./project-sandbox-metering"; +import type { ProjectSandboxProvisioning } from "./project-sandbox-provisioning"; +import type { ProjectSandboxWorkspaceState } from "./project-sandbox-workspace-state"; + +const ClearWorkspaceEvidenceSchema = z.object({ cleared: z.literal(true) }).strict(); + +interface AccountDeletionState { + accountDeletionCompleted: boolean; + activeOperationCount: number; + activeOperationDrainWaiters: Set<() => void>; + ctx: DurableObjectState; + env: ProjectSandboxEnv; + identity: ProjectSandboxIdentityState; + provisioning: ProjectSandboxProvisioning; + workspaceState: ProjectSandboxWorkspaceState | undefined; +} + +interface AccountDeletionRuntime { + clearCachedSandbox: () => void; + ensureClient: () => Promise; + meteringContext: () => SandboxMeteringContext; + withSandboxMutation: (operation: () => Promise) => Promise; +} + +export async function performAccountDeletion( + state: AccountDeletionState, + runtime: AccountDeletionRuntime, +): Promise { + await state.ctx.storage.put(ACCOUNT_DELETION_TOMBSTONE_KEY, true); + await waitForActiveSandboxOperations(state); + await runtime.withSandboxMutation(async () => { + await finalizeSandboxUsageBestEffort(runtime.meteringContext()); + await destroySandboxExclusive(state, runtime); + await clearDurableStateForDeletedAccount(state, runtime); + }); + state.accountDeletionCompleted = true; +} + +function waitForActiveSandboxOperations(state: AccountDeletionState): Promise { + if (state.activeOperationCount === 0) { + return Promise.resolve(); + } + return new Promise((resolve) => { + state.activeOperationDrainWaiters.add(resolve); + }); +} + +async function destroySandboxExclusive( + state: AccountDeletionState, + runtime: AccountDeletionRuntime, +): Promise { + const daytona = await runtime.ensureClient(); + try { + const canonical = await state.provisioning.findExisting(daytona); + const owned = await state.provisioning.findOwned(daytona); + const sandboxes = uniqueSandboxes(canonical ? [canonical, ...owned] : owned); + const volumeSandbox = sandboxes.find( + (sandbox) => sandbox.labels["workspaceVolumeName"] === state.env.DAYTONA_WORKSPACE_VOLUME, + ); + if (volumeSandbox) { + await clearWorkspaceVolume(state, daytona, volumeSandbox); + } + for (const sandbox of sandboxes) { + await daytona.deleteSandbox(sandbox.id); + } + } catch (error) { + throw toUpstreamError(error, "Daytona sandbox deletion failed.", state.identity.sandboxName()); + } + return finishSandboxDestruction(state, runtime); +} + +async function clearWorkspaceVolume( + state: AccountDeletionState, + daytona: DaytonaClient, + sandbox: DaytonaSandbox, +): Promise { + if (!(await state.provisioning.ensureStarted(daytona, sandbox))) { + throw new APIError(502, "upstream_sandbox_failed", "Daytona workspace disappeared", { + retriable: true, + }); + } + const cleared = await daytona.execute(sandbox.id, { + command: clearWorkspaceCommand(), + timeout: 480, + }); + if ( + cleared.exitCode !== 0 || + !ClearWorkspaceEvidenceSchema.safeParse(parseSandboxJson(cleared.result)).success + ) { + throw new APIError(502, "upstream_sandbox_failed", "Daytona workspace deletion failed", { + details: { sandboxId: state.identity.sandboxName() }, + retriable: true, + }); + } +} + +async function finishSandboxDestruction( + state: AccountDeletionState, + runtime: AccountDeletionRuntime, +): Promise { + runtime.clearCachedSandbox(); + await state.ctx.storage.delete(DAYTONA_ID_KEY); + await clearSandboxMeterState(state.ctx.storage); + return { deleted: true, sandboxId: state.identity.sandboxName() }; +} + +async function clearDurableStateForDeletedAccount( + state: AccountDeletionState, + runtime: AccountDeletionRuntime, +): Promise { + await state.ctx.storage.deleteAll(); + state.workspaceState = undefined; + state.identity.clearRegisteredOwner(); + runtime.clearCachedSandbox(); +} diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-content.ts b/apps/agent-worker/src/durable-objects/project-sandbox-content.ts index 5596721d..6354bd71 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-content.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-content.ts @@ -41,7 +41,8 @@ import { shellQuote, timeoutSeconds, } from "./project-sandbox-process-support"; -import { ProjectSandboxProjectFiles } from "./project-sandbox-project-files"; +import type { CoordinatedProcessOps, ProcessOps } from "./project-sandbox-processes"; +import type { FileOps } from "./project-sandbox-project-files"; import { type ProjectArchiveInput, ProjectArchiveInputSchema, @@ -62,6 +63,7 @@ import { ProjectPreviewStatusInputSchema, type ProjectReadFileInput, ProjectReadFileInputSchema, + type ProjectSandboxRuntimeState, type ProjectSearchFilesInput, ProjectSearchFilesInputSchema, type ProjectSignedPreviewUrlInput, @@ -72,6 +74,7 @@ import { type ProjectWriteFileInput, ProjectWriteFileInputSchema, } from "./project-sandbox-runtime"; +import type { SandboxRuntime } from "./project-sandbox-runtime-handle"; const PREVIEW_STATUS_PROBE_TIMEOUT_MS = 3_000; const PREVIEW_WAKE_TIMEOUT_MS = 90_000; @@ -81,439 +84,613 @@ const BROWSER_TAKEOVER_PORT_MIN = 60_000; const BROWSER_TAKEOVER_PORT_MAX = 60_999; const BROWSER_TAKEOVER_SCRIPT = "/opt/cheatcode/start-browser-takeover.sh"; -export abstract class ProjectSandboxContent extends ProjectSandboxProjectFiles { - public downloadProjectArchive(input: ProjectArchiveInput): Promise { - return this.downloadProjectArchiveForRpc(input, () => undefined); - } +export interface ContentOps { + cleanupProjectWorkspace: (input: ProjectCleanupWorkspaceInput) => Promise; + deleteFile: (input: ProjectDeleteFileInput) => Promise; + downloadProjectArchive: (input: ProjectArchiveInput, onFinished: () => void) => Promise; + exposeBrowserTakeover: ( + input: ProjectBrowserTakeoverInput, + ) => Promise; + exposeCodeServer: (input: ProjectCodeServerInput) => Promise<{ + expiresAt: string; + port: number; + url: string; + workspacePath: string; + }>; + getSignedPreviewUrl: ( + input: ProjectSignedPreviewUrlInput, + ) => Promise<{ token: string; url: string }>; + listFiles: (input: ProjectListFilesInput) => Promise; + projectPreviewStatus: ( + input: ProjectPreviewStatusInput, + ) => Promise<{ running: boolean; state: string }>; + readFile: (input: ProjectReadFileInput) => Promise; + searchFiles: (input: ProjectSearchFilesInput) => Promise; + stopBrowserTakeover: (input: ProjectBrowserTakeoverStopInput) => Promise; + wakePreview: (input: ProjectWakePreviewInput) => Promise; + writeFile: (input: ProjectWriteFileInput) => Promise; +} - protected async downloadProjectArchiveForRpc( - input: ProjectArchiveInput, - onFinished: () => void, - ): Promise { - const parsed = ProjectArchiveInputSchema.parse(input); - const archivePath = `/tmp/cheatcode-project-${crypto.randomUUID()}.zip`; - const workspacePath = `${WORKSPACE_DIR}/${parsed.workspaceSlug}`; - const result = await this.exec({ - command: [ - "python3", - "-c", - PROJECT_ARCHIVE_SCRIPT, - workspacePath, - archivePath, - String(PROJECT_ARCHIVE_MAX_BYTES), - String(PROJECT_ARCHIVE_MAX_FILES), - String(PROJECT_ARCHIVE_MAX_OUTPUT_BYTES), - ], - cwd: workspacePath, - timeoutMs: 300_000, +type ContentRuntime = Pick< + SandboxRuntime, + | "client" + | "deleteProjectWorkspace" + | "ensureExistingSandboxStarted" + | "ensureSandbox" + | "previewHostname" + | "previewSecret" + | "toUpstreamError" +>; + +type ContentProcessOps = Pick< + ProcessOps, + | "deleteProcessRecord" + | "deleteProcessesOnPort" + | "freeProjectPort" + | "httpPortReady" + | "isPortAlive" + | "killAllProcesses" + | "processRecord" + | "relaunchDevServer" + | "terminateUntrackedSandboxProcesses" + | "waitForPort" +>; + +type ContentCoordinatedProcessOps = Pick< + CoordinatedProcessOps, + "allocateProcessPort" | "exec" | "killProcess" | "startProcess" +>; + +interface ContentDependencies { + coordinatedProcess: ContentCoordinatedProcessOps; + deleteUploadedFileMetadata: FileOps["deleteUploadedFileMetadata"]; + process: ContentProcessOps; + sandboxRuntimeState: () => Promise; +} + +interface ContentContext { + dependencies: ContentDependencies; + runtime: ContentRuntime; +} + +export function createContentOps( + runtime: ContentRuntime, + dependencies: ContentDependencies, +): ContentOps { + const context = { dependencies, runtime }; + return { + cleanupProjectWorkspace: (input) => cleanupProjectWorkspace(context, input), + deleteFile: (input) => deleteFile(runtime, input), + downloadProjectArchive: (input, onFinished) => + downloadProjectArchive(context, input, onFinished), + exposeBrowserTakeover: (input) => exposeBrowserTakeover(context, input), + exposeCodeServer: (input) => exposeCodeServer(context, input), + getSignedPreviewUrl: (input) => getSignedPreviewUrl(runtime, input), + listFiles: (input) => listFiles(runtime, input), + projectPreviewStatus: (input) => projectPreviewStatus(context, input), + readFile: (input) => readFile(runtime, input), + searchFiles: (input) => searchFiles(runtime, input), + stopBrowserTakeover: (input) => stopBrowserTakeover(context, input), + wakePreview: (input) => wakePreview(context, input), + writeFile: (input) => writeFile(runtime, input), + }; +} + +async function downloadProjectArchive( + context: ContentContext, + input: ProjectArchiveInput, + onFinished: () => void, +): Promise { + const parsed = ProjectArchiveInputSchema.parse(input); + const archivePath = `/tmp/cheatcode-project-${crypto.randomUUID()}.zip`; + const workspacePath = `${WORKSPACE_DIR}/${parsed.workspaceSlug}`; + const result = await context.dependencies.coordinatedProcess.exec({ + command: [ + "python3", + "-c", + PROJECT_ARCHIVE_SCRIPT, + workspacePath, + archivePath, + String(PROJECT_ARCHIVE_MAX_BYTES), + String(PROJECT_ARCHIVE_MAX_FILES), + String(PROJECT_ARCHIVE_MAX_OUTPUT_BYTES), + ], + cwd: workspacePath, + timeoutMs: 300_000, + }); + const id = await context.runtime.ensureSandbox(); + if (!result.success) { + await context.runtime + .client() + .deleteFilePath(id, archivePath, false) + .catch(() => undefined); + throw new APIError(422, "sandbox_command_failed", "Unable to prepare this project download", { + hint: result.stdout.trim().slice(-300) || "Check the project files and try again.", + retriable: true, }); - const id = await this.ensureSandbox(); - if (!result.success) { - await this.client() - .deleteFilePath(id, archivePath, false) - .catch(() => undefined); - throw new APIError(422, "sandbox_command_failed", "Unable to prepare this project download", { - hint: result.stdout.trim().slice(-300) || "Check the project files and try again.", - retriable: true, - }); - } - const client = this.client(); - const cleanup = async (): Promise => { - try { - await client.deleteFilePath(id, archivePath, false).catch(() => undefined); - } finally { - onFinished(); - } - }; - try { - const upstream = await client.downloadFileResponse(id, archivePath); - return await projectArchiveResponse(upstream, cleanup); - } catch (error) { - await cleanup(); - throw error; - } } + return streamProjectArchive(context.runtime.client(), id, archivePath, onFinished); +} - public async readFile(input: ProjectReadFileInput): Promise { - const parsed = ProjectReadFileInputSchema.parse(input); - const id = await this.ensureSandbox(); - let bytes: Uint8Array; +async function streamProjectArchive( + client: ReturnType, + sandboxId: string, + archivePath: string, + onFinished: () => void, +): Promise { + const cleanup = async (): Promise => { try { - bytes = await this.client().downloadFile(id, parsed.path, SANDBOX_READ_FILE_MAX_BYTES); - } catch (error) { - if (isDaytonaResponseTooLarge(error)) { - throw new APIError(422, "sandbox_command_failed", "File is too large to read inline", { - hint: "Use file search or a shell command to inspect a smaller range.", - retriable: false, - }); - } - throw error; + await client.deleteFilePath(sandboxId, archivePath, false).catch(() => undefined); + } finally { + onFinished(); } - if (parsed.encoding === "base64") { - return { - content: encodeBase64(bytes), - encoding: "base64", - path: parsed.path, - size: bytes.byteLength, - }; + }; + try { + const upstream = await client.downloadFileResponse(sandboxId, archivePath); + return await projectArchiveResponse(upstream, cleanup); + } catch (error) { + await cleanup(); + throw error; + } +} + +async function readFile( + runtime: ContentRuntime, + input: ProjectReadFileInput, +): Promise { + const parsed = ProjectReadFileInputSchema.parse(input); + const id = await runtime.ensureSandbox(); + let bytes: Uint8Array; + try { + bytes = await runtime.client().downloadFile(id, parsed.path, SANDBOX_READ_FILE_MAX_BYTES); + } catch (error) { + if (isDaytonaResponseTooLarge(error)) { + throw new APIError(422, "sandbox_command_failed", "File is too large to read inline", { + hint: "Use file search or a shell command to inspect a smaller range.", + retriable: false, + }); } + throw error; + } + if (parsed.encoding === "base64") { return { - content: new TextDecoder().decode(bytes), - encoding: "utf8", + content: encodeBase64(bytes), + encoding: "base64", path: parsed.path, size: bytes.byteLength, }; } + return { + content: new TextDecoder().decode(bytes), + encoding: "utf8", + path: parsed.path, + size: bytes.byteLength, + }; +} - public async writeFile(input: ProjectWriteFileInput): Promise { - const parsed = ProjectWriteFileInputSchema.parse(input); - assertMutableWorkspacePath(parsed.path); - const id = await this.ensureSandbox(); - await this.client().createFolder(id, dirname(parsed.path)); - const bytes = - parsed.encoding === "base64" - ? decodeBase64(parsed.content) - : new TextEncoder().encode(parsed.content); - await this.client().uploadFile(id, parsed.path, bytes); - return { path: parsed.path, success: true }; - } +async function writeFile( + runtime: ContentRuntime, + input: ProjectWriteFileInput, +): Promise { + const parsed = ProjectWriteFileInputSchema.parse(input); + assertMutableWorkspacePath(parsed.path); + const id = await runtime.ensureSandbox(); + await runtime.client().createFolder(id, dirname(parsed.path)); + const bytes = + parsed.encoding === "base64" + ? decodeBase64(parsed.content) + : new TextEncoder().encode(parsed.content); + await runtime.client().uploadFile(id, parsed.path, bytes); + return { path: parsed.path, success: true }; +} - public async listFiles(input: ProjectListFilesInput): Promise { - const parsed = ProjectListFilesInputSchema.parse(input); - const id = await this.ensureSandbox(); - const files = await listSandboxFiles({ - client: this.client(), - includeHidden: parsed.includeHidden, - path: parsed.path, - recursive: parsed.recursive, - sandboxId: id, - }); - return { files, path: parsed.path }; - } +async function listFiles( + runtime: ContentRuntime, + input: ProjectListFilesInput, +): Promise { + const parsed = ProjectListFilesInputSchema.parse(input); + const id = await runtime.ensureSandbox(); + const files = await listSandboxFiles({ + client: runtime.client(), + includeHidden: parsed.includeHidden, + path: parsed.path, + recursive: parsed.recursive, + sandboxId: id, + }); + return { files, path: parsed.path }; +} - public async searchFiles(input: ProjectSearchFilesInput): Promise { - const parsed = ProjectSearchFilesInputSchema.parse(input); - const id = await this.ensureSandbox(); - const completed = await this.client().execute(id, { - command: buildGrepCommand(parsed), - cwd: WORKSPACE_DIR, - timeout: timeoutSeconds(60_000), - }); - const matches = parseGrepOutput(completed.result ?? "", parsed.maxResults); - return { - matches, - query: parsed.query, - total: matches.length, - truncated: matches.length >= parsed.maxResults, - }; - } +async function searchFiles( + runtime: ContentRuntime, + input: ProjectSearchFilesInput, +): Promise { + const parsed = ProjectSearchFilesInputSchema.parse(input); + const id = await runtime.ensureSandbox(); + const completed = await runtime.client().execute(id, { + command: buildGrepCommand(parsed), + cwd: WORKSPACE_DIR, + timeout: timeoutSeconds(60_000), + }); + const matches = parseGrepOutput(completed.result ?? "", parsed.maxResults); + return { + matches, + query: parsed.query, + total: matches.length, + truncated: matches.length >= parsed.maxResults, + }; +} - public async deleteFile(input: ProjectDeleteFileInput): Promise { - const parsed = ProjectDeleteFileInputSchema.parse(input); - assertDeletableWorkspacePath(parsed.path); - const id = await this.ensureSandbox(); - await this.client().deleteFilePath(id, parsed.path, parsed.recursive); - return { path: parsed.path, success: true }; - } +async function deleteFile( + runtime: ContentRuntime, + input: ProjectDeleteFileInput, +): Promise { + const parsed = ProjectDeleteFileInputSchema.parse(input); + assertDeletableWorkspacePath(parsed.path); + const id = await runtime.ensureSandbox(); + await runtime.client().deleteFilePath(id, parsed.path, parsed.recursive); + return { path: parsed.path, success: true }; +} - public async getSignedPreviewUrl( - input: ProjectSignedPreviewUrlInput, - ): Promise<{ token: string; url: string }> { - const parsed = ProjectSignedPreviewUrlInputSchema.parse(input); - const id = await this.ensureSandbox(); - const link = await this.client().getSignedPreviewUrl(id, parsed.port, parsed.expiresInSeconds); - return { token: link.token, url: link.url }; - } +async function getSignedPreviewUrl( + runtime: ContentRuntime, + input: ProjectSignedPreviewUrlInput, +): Promise<{ token: string; url: string }> { + const parsed = ProjectSignedPreviewUrlInputSchema.parse(input); + const id = await runtime.ensureSandbox(); + const link = await runtime.client().getSignedPreviewUrl(id, parsed.port, parsed.expiresInSeconds); + return { token: link.token, url: link.url }; +} - public async exposeBrowserTakeover( - input: ProjectBrowserTakeoverInput, - ): Promise { - const parsed = ProjectBrowserTakeoverInputSchema.parse(input); - const browserDriver = await this.processRecord(browserDriverProcessId(parsed.runId)); - if (!browserDriver) { - throw new APIError(409, "conflict_state_invalid", "No live browser session is available", { - hint: "Let Cheatcode open a website before taking over the browser.", - retriable: true, - }); - } - const processId = browserTakeoverProcessId(parsed.runId); - const port = await this.allocateProcessPort({ - maxPort: BROWSER_TAKEOVER_PORT_MAX, - minPort: BROWSER_TAKEOVER_PORT_MIN, - processId, - }); - const password = crypto.randomUUID().replaceAll("-", ""); - await this.startProcess({ - command: ["sh", BROWSER_TAKEOVER_SCRIPT], - env: { TAKEOVER_PASSWORD: password, TAKEOVER_PORT: String(port) }, - keepAliveTimeoutMs: parsed.expiresInSeconds * 1_000, - maxRestarts: 0, - processId, - restartOnFailure: false, - waitForPort: { path: "/vnc.html", port, timeoutMs: 30_000 }, +async function exposeBrowserTakeover( + context: ContentContext, + input: ProjectBrowserTakeoverInput, +): Promise { + const parsed = ProjectBrowserTakeoverInputSchema.parse(input); + const browserDriver = await context.dependencies.process.processRecord( + browserDriverProcessId(parsed.runId), + ); + if (!browserDriver) { + throw new APIError(409, "conflict_state_invalid", "No live browser session is available", { + hint: "Let Cheatcode open a website before taking over the browser.", + retriable: true, }); - const id = await this.ensureSandbox(); - const signed = await this.client().getSignedPreviewUrl(id, port, parsed.expiresInSeconds); - const url = noVncSessionUrl(signed.url, password); - return { - expiresAt: new Date(Date.now() + parsed.expiresInSeconds * 1_000).toISOString(), - takeoverId: parsed.takeoverId, - url, - }; } + const processId = browserTakeoverProcessId(parsed.runId); + const port = await context.dependencies.coordinatedProcess.allocateProcessPort({ + maxPort: BROWSER_TAKEOVER_PORT_MAX, + minPort: BROWSER_TAKEOVER_PORT_MIN, + processId, + }); + const password = crypto.randomUUID().replaceAll("-", ""); + await context.dependencies.coordinatedProcess.startProcess({ + command: ["sh", BROWSER_TAKEOVER_SCRIPT], + env: { TAKEOVER_PASSWORD: password, TAKEOVER_PORT: String(port) }, + keepAliveTimeoutMs: parsed.expiresInSeconds * 1_000, + maxRestarts: 0, + processId, + restartOnFailure: false, + waitForPort: { path: "/vnc.html", port, timeoutMs: 30_000 }, + }); + return browserTakeoverResult(context.runtime, parsed, port, password); +} - public async stopBrowserTakeover(input: ProjectBrowserTakeoverStopInput): Promise { - const parsed = ProjectBrowserTakeoverStopInputSchema.parse(input); - await this.killProcess({ processId: browserTakeoverProcessId(parsed.runId) }); - } +async function browserTakeoverResult( + runtime: ContentRuntime, + input: ReturnType, + port: number, + password: string, +): Promise { + const id = await runtime.ensureSandbox(); + const signed = await runtime.client().getSignedPreviewUrl(id, port, input.expiresInSeconds); + return { + expiresAt: new Date(Date.now() + input.expiresInSeconds * 1_000).toISOString(), + takeoverId: input.takeoverId, + url: noVncSessionUrl(signed.url, password), + }; +} - public async exposeCodeServer(input: ProjectCodeServerInput): Promise<{ - expiresAt: string; - port: number; - url: string; - workspacePath: string; - }> { - const parsed = ProjectCodeServerInputSchema.parse(input); - const id = await this.ensureSandbox(); - await this.ensureCodeServer(id); - const displayFolder = - parsed.workspacePath === WORKSPACE_DIR - ? await this.ensureCodeServerDisplayFolder(id, parsed.workspacePath) - : parsed.workspacePath; - await this.client() - .getPreviewLink(id, CODE_SERVER_PORT) +async function stopBrowserTakeover( + context: ContentContext, + input: ProjectBrowserTakeoverStopInput, +): Promise { + const parsed = ProjectBrowserTakeoverStopInputSchema.parse(input); + await context.dependencies.coordinatedProcess.killProcess({ + processId: browserTakeoverProcessId(parsed.runId), + }); +} + +async function exposeCodeServer( + context: ContentContext, + input: ProjectCodeServerInput, +): Promise<{ + expiresAt: string; + port: number; + url: string; + workspacePath: string; +}> { + const parsed = ProjectCodeServerInputSchema.parse(input); + const id = await context.runtime.ensureSandbox(); + await ensureCodeServer(context, id); + const displayFolder = + parsed.workspacePath === WORKSPACE_DIR + ? await ensureCodeServerDisplayFolder(context.runtime, id, parsed.workspacePath) + : parsed.workspacePath; + await context.runtime + .client() + .getPreviewLink(id, CODE_SERVER_PORT) + .catch(() => undefined); + const built = await buildPreviewUrl({ + hostname: context.runtime.previewHostname(), + port: CODE_SERVER_PORT, + sandboxId: id, + secret: await context.runtime.previewSecret(), + useSubdomain: true, + }); + return { + expiresAt: built.expiresAt, + port: CODE_SERVER_PORT, + url: codeServerFolderUrl(built.url, displayFolder, parsed.initialFilePath), + workspacePath: parsed.workspacePath, + }; +} + +async function wakePreview( + context: ContentContext, + input: ProjectWakePreviewInput, +): Promise { + const parsed = ProjectWakePreviewInputSchema.parse(input); + const id = await context.runtime.ensureSandbox(); + const slot = parsed.workspaceSlug + ? `${APP_PREVIEW_SLOT_PREFIX}${parsed.workspaceSlug}` + : "app-preview"; + const record = await context.dependencies.process.processRecord(slot); + if (!record?.port) return { running: false, state: "started" }; + const mobile = await mobileExpoProxy(context.runtime, id, record); + const repaired = record.isMobile + ? await ensureMobileMetroForwardedHostConfig(context.runtime, id, record.cwd) + : false; + let running = await context.dependencies.process.isPortAlive(id, record.port); + if (!running || repaired) { + await context.dependencies.process.relaunchDevServer( + id, + slot, + record, + mobile?.restartEnv ?? restartEnvironment(slot, record), + ); + await context.dependencies.process + .waitForPort(id, record.port, "/", PREVIEW_WAKE_TIMEOUT_MS) .catch(() => undefined); - const built = await buildPreviewUrl({ - hostname: this.previewHostname(), - port: CODE_SERVER_PORT, - sandboxId: id, - secret: await this.previewSecret(), - useSubdomain: true, - }); - return { - expiresAt: built.expiresAt, - port: CODE_SERVER_PORT, - url: codeServerFolderUrl(built.url, displayFolder, parsed.initialFilePath), - workspacePath: parsed.workspacePath, - }; + running = await context.dependencies.process.isPortAlive(id, record.port); } + return wakePreviewResult(context.runtime, id, record, record.port, running, mobile?.expoUrl); +} - public async wakePreview(input: ProjectWakePreviewInput): Promise { - const parsed = ProjectWakePreviewInputSchema.parse(input); - const id = await this.ensureSandbox(); - const slot = parsed.workspaceSlug - ? `${APP_PREVIEW_SLOT_PREFIX}${parsed.workspaceSlug}` - : "app-preview"; - const record = await this.processRecord(slot); - if (!record?.port) { - return { running: false, state: "started" }; - } - const mobile = await this.mobileExpoProxy(id, record); - const repairedMobileConfig = record.isMobile - ? await this.ensureMobileMetroForwardedHostConfig(id, record.cwd) - : false; - let running = await this.isPortAlive(id, record.port); - if (!running || repairedMobileConfig) { - await this.relaunchDevServer( - id, - slot, - record, - mobile?.restartEnv ?? restartEnvironment(slot, record), - ); - await this.waitForPort(id, record.port, "/", PREVIEW_WAKE_TIMEOUT_MS).catch(() => undefined); - running = await this.isPortAlive(id, record.port); - } - await this.client() - .getPreviewLink(id, record.port) - .catch(() => undefined); - const built = await buildPreviewUrl({ - hostname: this.previewHostname(), - isMobile: record.isMobile === true, - port: record.port, - sandboxId: id, - secret: await this.previewSecret(), - }); - return { - expiresAt: built.expiresAt, - port: record.port, - running, - state: "started", - url: built.url, - ...(mobile?.expoUrl ? { expoUrl: mobile.expoUrl } : {}), - }; +async function wakePreviewResult( + runtime: ContentRuntime, + sandboxId: string, + record: ProcessRecord, + port: number, + running: boolean, + expoUrl: string | undefined, +): Promise { + await runtime + .client() + .getPreviewLink(sandboxId, port) + .catch(() => undefined); + const built = await buildPreviewUrl({ + hostname: runtime.previewHostname(), + isMobile: record.isMobile === true, + port, + sandboxId, + secret: await runtime.previewSecret(), + }); + return { + expiresAt: built.expiresAt, + port, + running, + state: "started", + url: built.url, + ...(expoUrl ? { expoUrl } : {}), + }; +} + +async function projectPreviewStatus( + context: ContentContext, + input: ProjectPreviewStatusInput, +): Promise<{ running: boolean; state: string }> { + const parsed = ProjectPreviewStatusInputSchema.parse(input); + const runtimeState = await context.dependencies.sandboxRuntimeState(); + if (runtimeState.state !== "started" || !runtimeState.sandboxId) { + return { running: false, state: runtimeState.state }; } + const record = await context.dependencies.process.processRecord( + `${APP_PREVIEW_SLOT_PREFIX}${parsed.workspaceSlug}`, + ); + if (!record?.port) return { running: false, state: runtimeState.state }; + const running = await context.dependencies.process.httpPortReady( + runtimeState.sandboxId, + record.port, + "/", + PREVIEW_STATUS_PROBE_TIMEOUT_MS, + ); + return { running, state: runtimeState.state }; +} - public async projectPreviewStatus( - input: ProjectPreviewStatusInput, - ): Promise<{ running: boolean; state: string }> { - const parsed = ProjectPreviewStatusInputSchema.parse(input); - const runtime = await this.sandboxRuntimeState(); - if (runtime.state !== "started" || !runtime.sandboxId) { - return { running: false, state: runtime.state }; - } - const record = await this.processRecord(`${APP_PREVIEW_SLOT_PREFIX}${parsed.workspaceSlug}`); - if (!record?.port) { - return { running: false, state: runtime.state }; - } - const running = await this.httpPortReady( - runtime.sandboxId, - record.port, - "/", - PREVIEW_STATUS_PROBE_TIMEOUT_MS, - ); - return { running, state: runtime.state }; +function cleanupProjectWorkspace( + context: ContentContext, + input: ProjectCleanupWorkspaceInput, +): Promise { + const parsed = ProjectCleanupWorkspaceInputSchema.parse(input); + return context.runtime.deleteProjectWorkspace(parsed, () => + performProjectWorkspaceCleanup(context, parsed.projectId, parsed.workspaceSlug), + ); +} + +async function performProjectWorkspaceCleanup( + context: ContentContext, + projectId: string, + workspaceSlug: string, +): Promise { + const id = await context.runtime.ensureExistingSandboxStarted(); + // Explicit raw bypass: cleanup has already fenced and drained workspace operations. + await context.dependencies.process.killAllProcesses(); + if (id) { + await context.dependencies.process.terminateUntrackedSandboxProcesses(id); + await removeWorkspaceFolder(context.runtime, id, workspaceSlug); } + await context.dependencies.process.freeProjectPort(workspaceSlug); + await context.dependencies.deleteUploadedFileMetadata(projectId); +} - public cleanupProjectWorkspace(input: ProjectCleanupWorkspaceInput): Promise { - const parsed = ProjectCleanupWorkspaceInputSchema.parse(input); - return this.deleteProjectWorkspace(parsed, () => - this.performProjectWorkspaceCleanup(parsed.projectId, parsed.workspaceSlug), +async function mobileExpoProxy( + runtime: ContentRuntime, + id: string, + record: ProcessRecord, +): Promise<{ expoUrl: string; restartEnv: Record } | null> { + if (!record.isMobile || record.port === undefined) return null; + const signed = await runtime + .client() + .getSignedPreviewUrl(id, record.port, SIGNED_PREVIEW_TTL_SECONDS) + .catch(() => null); + if (!signed) return null; + return { + expoUrl: signedUrlToExpo(signed.url), + restartEnv: { + CI: "1", + EXPO_NO_TELEMETRY: "1", + EXPO_PACKAGER_PROXY_URL: signed.url, + PORT: String(record.port), + }, + }; +} + +async function ensureMobileMetroForwardedHostConfig( + runtime: ContentRuntime, + id: string, + cwd: string, +): Promise { + const current = await runtime + .client() + .execute(id, { + command: 'test -f metro.config.js && grep -q "x-forwarded-host" metro.config.js', + cwd, + timeout: 5, + }) + .catch(() => null); + if (current?.exitCode === 0) return false; + const repair = await runtime.client().execute(id, { + command: `bash -lc ${shellQuote(metroForwardedHostFixScript())}`, + cwd, + timeout: 15, + }); + if (repair.exitCode !== 0) { + throw new APIError( + 502, + "upstream_sandbox_failed", + "Could not repair the mobile preview proxy configuration.", + { retriable: true }, ); } + return true; +} - private async performProjectWorkspaceCleanup( - projectId: string, - workspaceSlug: string, - ): Promise { - const id = await this.ensureExistingSandboxStarted(); - await super.killAllProcesses(); - if (id) { - await this.terminateUntrackedSandboxProcesses(id); - await this.removeWorkspaceFolder(id, workspaceSlug); - } - await this.freeProjectPort(workspaceSlug); - await this.deleteUploadedFileMetadata(projectId); +async function ensureCodeServer(context: ContentContext, id: string): Promise { + if ( + (await context.dependencies.process.httpPortReady(id, CODE_SERVER_PORT, "/", 5_000)) && + (await hasCodeServerSettingsMarker(context.runtime, id)) + ) { + return; } - - private async mobileExpoProxy( - id: string, - record: ProcessRecord, - ): Promise<{ expoUrl: string; restartEnv: Record } | null> { - if (!record.isMobile || record.port === undefined) { - return null; - } - const signed = await this.client() - .getSignedPreviewUrl(id, record.port, SIGNED_PREVIEW_TTL_SECONDS) - .catch(() => null); - if (!signed) { - return null; - } - return { - expoUrl: signedUrlToExpo(signed.url), - restartEnv: { - CI: "1", - EXPO_NO_TELEMETRY: "1", - EXPO_PACKAGER_PROXY_URL: signed.url, - PORT: String(record.port), - }, - }; + if (!(await hasCodeServerRuntime(context.runtime, id))) { + throw new APIError(502, "sandbox_failed_to_start", "code-server is not installed", { + hint: "Start a new project sandbox from the current Daytona snapshot to use the Files viewer.", + retriable: false, + }); } - - private async ensureMobileMetroForwardedHostConfig(id: string, cwd: string): Promise { - const current = await this.client() - .execute(id, { - command: 'test -f metro.config.js && grep -q "x-forwarded-host" metro.config.js', - cwd, - timeout: 5, - }) - .catch(() => null); - if (current?.exitCode === 0) { - return false; - } - const repair = await this.client().execute(id, { - command: `bash -lc ${shellQuote(metroForwardedHostFixScript())}`, - cwd, - timeout: 15, + await context.dependencies.process.deleteProcessRecord(id, CODE_SERVER_PROCESS_ID); + await context.dependencies.process.deleteProcessesOnPort( + id, + CODE_SERVER_PORT, + CODE_SERVER_PROCESS_ID, + ); + await context.runtime + .client() + .execute(id, { command: "pkill -f code-server || true", timeout: 5 }) + .catch(() => null); + await startCodeServer(context); + if (!(await context.dependencies.process.httpPortReady(id, CODE_SERVER_PORT, "/", 5_000))) { + throw new APIError(502, "sandbox_failed_to_start", "Unable to start code-server", { + hint: "Rebuild the Daytona sandbox snapshot with code-server, then retry the Files tab.", + retriable: true, }); - if (repair.exitCode !== 0) { - throw new APIError( - 502, - "upstream_sandbox_failed", - "Could not repair the mobile preview proxy configuration.", - { retriable: true }, - ); - } - return true; } +} - private async ensureCodeServer(id: string): Promise { - if ( - (await this.httpPortReady(id, CODE_SERVER_PORT, "/", 5_000)) && - (await this.hasCodeServerSettingsMarker(id)) - ) { - return; - } - if (!(await this.hasCodeServerRuntime(id))) { - throw new APIError(502, "sandbox_failed_to_start", "code-server is not installed", { - hint: "Start a new project sandbox from the current Daytona snapshot to use the Files viewer.", - retriable: false, - }); - } - await this.deleteProcessRecord(id, CODE_SERVER_PROCESS_ID); - await this.deleteProcessesOnPort(id, CODE_SERVER_PORT, CODE_SERVER_PROCESS_ID); - await this.client() - .execute(id, { command: "pkill -f code-server || true", timeout: 5 }) - .catch(() => null); - await this.startProcess({ - command: ["bash", "-lc", codeServerStartCommand()], - cwd: WORKSPACE_DIR, - env: { - CODE_SERVER_PORT: String(CODE_SERVER_PORT), - CODE_SERVER_TRUSTED_ORIGINS: codeServerTrustedOrigins(this.previewHostname()), - CODE_SERVER_WORKSPACE: WORKSPACE_DIR, - }, - keepAliveTimeoutMs: 0, - maxRestarts: 3, - processId: CODE_SERVER_PROCESS_ID, - restartOnFailure: true, +async function startCodeServer(context: ContentContext): Promise { + await context.dependencies.coordinatedProcess.startProcess({ + command: ["bash", "-lc", codeServerStartCommand()], + cwd: WORKSPACE_DIR, + env: { + CODE_SERVER_PORT: String(CODE_SERVER_PORT), + CODE_SERVER_TRUSTED_ORIGINS: codeServerTrustedOrigins(context.runtime.previewHostname()), + CODE_SERVER_WORKSPACE: WORKSPACE_DIR, + }, + keepAliveTimeoutMs: 0, + maxRestarts: 3, + processId: CODE_SERVER_PROCESS_ID, + restartOnFailure: true, + timeoutMs: CODE_SERVER_START_TIMEOUT_MS, + waitForPort: { + path: "/", + port: CODE_SERVER_PORT, timeoutMs: CODE_SERVER_START_TIMEOUT_MS, - waitForPort: { - path: "/", - port: CODE_SERVER_PORT, - timeoutMs: CODE_SERVER_START_TIMEOUT_MS, - }, - }); - if (!(await this.httpPortReady(id, CODE_SERVER_PORT, "/", 5_000))) { - throw new APIError(502, "sandbox_failed_to_start", "Unable to start code-server", { - hint: "Rebuild the Daytona sandbox snapshot with code-server, then retry the Files tab.", - retriable: true, - }); - } - } + }, + }); +} - private async hasCodeServerRuntime(id: string): Promise { - const probe = await this.client() - .execute(id, { command: "command -v code-server >/dev/null", timeout: 5 }) - .catch(() => null); - return probe?.exitCode === 0; - } +async function hasCodeServerRuntime(runtime: ContentRuntime, id: string): Promise { + const probe = await runtime + .client() + .execute(id, { command: "command -v code-server >/dev/null", timeout: 5 }) + .catch(() => null); + return probe?.exitCode === 0; +} - private async hasCodeServerSettingsMarker(id: string): Promise { - const probe = await this.client() - .execute(id, { - command: `test -f ${shellQuote(CODE_SERVER_SETTINGS_MARKER)}`, - timeout: 5, - }) - .catch(() => null); - return probe?.exitCode === 0; - } +async function hasCodeServerSettingsMarker(runtime: ContentRuntime, id: string): Promise { + const probe = await runtime + .client() + .execute(id, { + command: `test -f ${shellQuote(CODE_SERVER_SETTINGS_MARKER)}`, + timeout: 5, + }) + .catch(() => null); + return probe?.exitCode === 0; +} - private async ensureCodeServerDisplayFolder(id: string, workspacePath: string): Promise { - const probe = await this.client() - .execute(id, { - command: `ln -sfn ${shellQuote(workspacePath)} ${shellQuote(CODE_SERVER_DISPLAY_DIR)} && test -d ${shellQuote(CODE_SERVER_DISPLAY_DIR)}`, - timeout: 10, - }) - .catch(() => null); - return probe?.exitCode === 0 ? CODE_SERVER_DISPLAY_DIR : workspacePath; - } +async function ensureCodeServerDisplayFolder( + runtime: ContentRuntime, + id: string, + workspacePath: string, +): Promise { + const probe = await runtime + .client() + .execute(id, { + command: `ln -sfn ${shellQuote(workspacePath)} ${shellQuote(CODE_SERVER_DISPLAY_DIR)} && test -d ${shellQuote(CODE_SERVER_DISPLAY_DIR)}`, + timeout: 10, + }) + .catch(() => null); + return probe?.exitCode === 0 ? CODE_SERVER_DISPLAY_DIR : workspacePath; +} - private async removeWorkspaceFolder(id: string, workspaceSlug: string): Promise { - try { - await this.client().deleteFilePath(id, `${WORKSPACE_DIR}/${workspaceSlug}`, true); - } catch (error) { - throw this.toUpstreamError(error, "Project workspace removal failed."); - } +async function removeWorkspaceFolder( + runtime: ContentRuntime, + id: string, + workspaceSlug: string, +): Promise { + try { + await runtime.client().deleteFilePath(id, `${WORKSPACE_DIR}/${workspaceSlug}`, true); + } catch (error) { + throw runtime.toUpstreamError(error, "Project workspace removal failed."); } } @@ -556,9 +733,7 @@ async function projectArchiveResponse( throw archiveDownloadError("The sandbox returned an oversized project archive."); } const headers = new Headers({ "Content-Type": "application/zip" }); - if (declaredLength !== null) { - headers.set("Content-Length", String(declaredLength)); - } + if (declaredLength !== null) headers.set("Content-Length", String(declaredLength)); return new Response( archiveStreamWithCleanup(upstream.body, PROJECT_ARCHIVE_MAX_OUTPUT_BYTES, cleanup), { headers }, @@ -606,9 +781,7 @@ function archiveStreamWithCleanup( } function boundedArchiveContentLength(value: string | null): number | null | "too-large" { - if (!value || !/^\d+$/u.test(value)) { - return null; - } + if (!value || !/^\d+$/u.test(value)) return null; const length = Number(value); if (!Number.isSafeInteger(length) || length > PROJECT_ARCHIVE_MAX_OUTPUT_BYTES) { return "too-large"; diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-lease-policy.ts b/apps/agent-worker/src/durable-objects/project-sandbox-lease-policy.ts index 8220f70b..32ac13b1 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-lease-policy.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-lease-policy.ts @@ -3,6 +3,7 @@ import { ProjectWorkspaceSlugSchema, workspaceSlugFromPath } from "./project-san type LeasePolicy = readonly [ kind: + | "account-deletion-control" | "cleanup-signal" | "owner-registration" | "project-cleanup" @@ -15,6 +16,7 @@ type LeasePolicy = readonly [ // biome-ignore format: Group the complete RPC policy surface by lease behavior for compact auditing. const LEASE_POLICIES = { + deleteAccountState: ["account-deletion-control"], registerOwner: ["owner-registration"], setQuotaPeriod: ["sandbox"], beginRun: ["sandbox"], renewRun: ["cleanup-signal"], endRun: ["cleanup-signal"], alarm: ["cleanup-signal"], diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-lease-runtime.ts b/apps/agent-worker/src/durable-objects/project-sandbox-lease-runtime.ts new file mode 100644 index 00000000..c1055d5d --- /dev/null +++ b/apps/agent-worker/src/durable-objects/project-sandbox-lease-runtime.ts @@ -0,0 +1,236 @@ +import type { ProjectSandboxIdentityState } from "./project-sandbox-identity-state"; +import { + accountSandboxDeletedError, + type ProjectSandboxEnv, + sandboxRuntimeUpdatePending, +} from "./project-sandbox-lifecycle-support"; +import { assertProjectSandboxOwnerActive } from "./project-sandbox-owner-admission"; +import { + initializeProjectSandboxStorage, + ProjectSandboxWorkspaceState, +} from "./project-sandbox-workspace-state"; + +interface SandboxLeaseState { + accountDeletionInProgress: boolean; + activeOperationCount: number; + activeOperationDrainWaiters: Set<() => void>; + ctx: DurableObjectState; + env: ProjectSandboxEnv; + identity: ProjectSandboxIdentityState; + sandboxRuntimeUpdateInProgress: boolean; + workspaceState: ProjectSandboxWorkspaceState | undefined; +} + +export function workspaceState(state: SandboxLeaseState): ProjectSandboxWorkspaceState { + if (!state.workspaceState) { + initializeProjectSandboxStorage(state.ctx); + state.workspaceState = new ProjectSandboxWorkspaceState(state.ctx); + } + return state.workspaceState; +} + +export function withOwnerRegistration( + state: SandboxLeaseState, + userId: string, + operation: () => Promise, +): Promise { + if (state.identity.hasRegisteredOwner()) { + return withActiveOperation(state, null, operation); + } + let release: (() => void) | undefined; + try { + release = acquireActiveSandboxOperation(state, true); + return runOwnerRegistrationPreflight(state, userId, operation).finally(release); + } catch (error) { + release?.(); + return Promise.reject(error); + } +} + +async function runOwnerRegistrationPreflight( + state: SandboxLeaseState, + userId: string, + operation: () => Promise, +): Promise { + await assertProjectSandboxOwnerActive(state.env, userId); + if (state.accountDeletionInProgress) { + throw accountSandboxDeletedError(); + } + const result = await operation(); + if (state.accountDeletionInProgress) { + throw accountSandboxDeletedError(); + } + workspaceState(state); + return result; +} + +export function withSharedWorkspaceMutation( + state: SandboxLeaseState, + operation: () => Promise, +): Promise { + return withActiveOperation( + state, + null, + async () => { + await workspaceState(state).waitForWorkspaceDrain(); + return operation(); + }, + true, + ); +} + +export function withActiveOperation( + state: SandboxLeaseState, + scope: string | readonly string[] | null, + operation: () => Promise, + isSharedMutation = false, + shouldLeaseUnknownWorkspace = false, + allowRuntimeUpdate = false, +): Promise { + let release: (() => void) | undefined; + try { + release = acquireActiveOperation( + state, + scope, + isSharedMutation, + shouldLeaseUnknownWorkspace, + allowRuntimeUpdate, + ); + return operation().finally(release); + } catch (error) { + release?.(); + return Promise.reject(error); + } +} + +export function withStreamingOperation( + state: SandboxLeaseState, + scope: string | readonly string[] | null, + operation: (release: () => void) => Promise, +): Promise { + let release: (() => void) | undefined; + try { + release = acquireActiveOperation(state, scope, false, true, false); + return operation(release).catch((error: unknown) => { + release?.(); + throw error; + }); + } catch (error) { + release?.(); + return Promise.reject(error); + } +} + +export function withCleanupSignal( + state: SandboxLeaseState, + operation: () => Promise, +): Promise { + return state.accountDeletionInProgress || !state.identity.hasRegisteredOwner() + ? Promise.resolve(undefined) + : withActiveOperation(state, null, operation, false, false, true); +} + +export function withProjectCleanup( + state: SandboxLeaseState, + operation: () => Promise, +): Promise { + let release: (() => void) | undefined; + try { + // Cleanup holds only the sandbox lease after its tombstone drains workspace work. + release = acquireActiveSandboxOperation(state, false, true); + return operation().finally(release); + } catch (error) { + release?.(); + return Promise.reject(error); + } +} + +function acquireActiveSandboxOperation( + state: SandboxLeaseState, + allowUnregisteredOwner = false, + allowWorkspaceCleanup = false, + allowRuntimeUpdate = false, +): () => void { + assertSandboxOperationAllowed( + state, + allowUnregisteredOwner, + allowWorkspaceCleanup, + allowRuntimeUpdate, + ); + state.activeOperationCount += 1; + let isReleased = false; + return () => { + if (isReleased) return; + isReleased = true; + finishActiveSandboxOperation(state); + }; +} + +function assertSandboxOperationAllowed( + state: SandboxLeaseState, + allowUnregisteredOwner: boolean, + allowWorkspaceCleanup: boolean, + allowRuntimeUpdate: boolean, +): void { + if (state.accountDeletionInProgress) { + throw accountSandboxDeletedError(); + } + if (state.sandboxRuntimeUpdateInProgress && !allowRuntimeUpdate) { + throw sandboxRuntimeUpdatePending(state.env.DAYTONA_SANDBOX_SNAPSHOT); + } + if (!allowUnregisteredOwner && !state.identity.hasRegisteredOwner()) { + throw accountSandboxDeletedError(); + } + const currentWorkspaceState = + allowUnregisteredOwner && !state.identity.hasRegisteredOwner() + ? state.workspaceState + : workspaceState(state); + currentWorkspaceState?.assertOperationAllowed(allowWorkspaceCleanup); +} + +function acquireActiveOperation( + state: SandboxLeaseState, + scope: string | readonly string[] | null, + isSharedMutation: boolean, + shouldLeaseUnknownWorkspace: boolean, + allowRuntimeUpdate: boolean, +): () => void { + const releaseSandbox = acquireActiveSandboxOperation(state, false, false, allowRuntimeUpdate); + let releaseWorkspace: (() => void) | undefined; + try { + const slugs = typeof scope === "string" ? [scope] : (scope ?? []); + releaseWorkspace = workspaceRelease( + workspaceState(state), + slugs, + isSharedMutation, + shouldLeaseUnknownWorkspace, + ); + } catch (error) { + releaseSandbox(); + throw error; + } + return () => { + releaseWorkspace?.(); + releaseSandbox(); + }; +} + +function workspaceRelease( + state: ProjectSandboxWorkspaceState, + slugs: readonly string[], + isSharedMutation: boolean, + shouldLeaseUnknownWorkspace: boolean, +): (() => void) | undefined { + if (isSharedMutation) return state.acquireSharedMutation(); + if (slugs.length > 0) return state.acquire(slugs); + return shouldLeaseUnknownWorkspace ? state.acquireUnscoped() : undefined; +} + +function finishActiveSandboxOperation(state: SandboxLeaseState): void { + state.activeOperationCount -= 1; + if (state.activeOperationCount !== 0) return; + for (const resolve of state.activeOperationDrainWaiters) { + resolve(); + } + state.activeOperationDrainWaiters.clear(); +} diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle-support.ts b/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle-support.ts index c3f5031e..cf0e0c64 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle-support.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle-support.ts @@ -33,7 +33,7 @@ export const STALE_RUN_LEASE_MS = 20 * 60 * 1_000; export const STARTED_REVERIFY_MS = 30_000; export const ENSURE_STARTED_ATTEMPTS = 30; export const ENSURE_STARTED_DELAY_MS = 2_000; -export const RunLeasesSchema = z +const RunLeasesSchema = z .array(z.object({ runId: z.string(), startedMs: z.number() }).strict()) .default([]); @@ -69,3 +69,41 @@ export function parseSandboxJson(value: string | null | undefined): unknown { return null; } } + +export async function runLeases( + storage: DurableObjectStorage, +): Promise> { + return RunLeasesSchema.parse((await storage.get(RUN_LEASES_KEY)) ?? []); +} + +export function toUpstreamError(error: unknown, fallback: string, sandboxId: string): APIError { + if (error instanceof APIError) { + return error; + } + const status = error instanceof DaytonaApiError ? error.status : 502; + const retriable = error instanceof DaytonaApiError ? error.retriable : true; + return new APIError(status >= 500 ? 502 : status, "upstream_sandbox_failed", fallback, { + cause: error, + details: { sandboxId }, + hint: "Retry. If it persists, check Daytona sandbox lifecycle status.", + retriable, + }); +} + +export async function storedDaytonaId(storage: DurableObjectStorage): Promise { + const value = await storage.get(DAYTONA_ID_KEY); + return typeof value === "string" ? value : null; +} + +export function sandboxRuntimeUpdatePending(expectedSnapshot: string): APIError { + return new APIError( + 503, + "unavailable_maintenance", + "This computer is updating to the current runtime.", + { + details: { expectedSnapshot }, + hint: "Retry after the active operation finishes.", + retriable: true, + }, + ); +} diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle.ts b/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle.ts index d4cf97cd..1528275c 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-lifecycle.ts @@ -1,804 +1,182 @@ -import { DurableObject } from "cloudflare:workers"; -import { PreviewHostnameSchema, resolveWorkerSecret } from "@cheatcode/env"; -import { APIError, createLogger } from "@cheatcode/observability"; +import { createLogger } from "@cheatcode/observability"; import { - DaytonaApiError, - DaytonaClient, - type DaytonaSandbox, - type SandboxDestroyResult, -} from "@cheatcode/tools-code"; -import { z } from "zod"; -import { type SandboxExecAuditEntry, writeExecAudit } from "./project-sandbox-audit"; -import { clearWorkspaceCommand } from "./project-sandbox-cleanup-script"; -import { ProjectSandboxIdentityState } from "./project-sandbox-identity-state"; -import { - ACCOUNT_DELETION_TOMBSTONE_KEY, - accountSandboxDeletedError, - DAYTONA_ID_KEY, DEFAULT_IDLE_STOP_MIN, KEEPALIVE_ALARM_MS, - type ProjectSandboxEnv, - parseSandboxJson, RUN_LEASES_KEY, - RunLeasesSchema, + runLeases, STALE_RUN_LEASE_MS, - STARTED_REVERIFY_MS, - uniqueSandboxes, } from "./project-sandbox-lifecycle-support"; import { beginSandboxUsageBestEffort, - clearSandboxMeterState, finalizeSandboxUsageBestEffort, recordSandboxUsageBestEffort, - type SandboxMeteringContext, setSandboxQuotaPeriod, } from "./project-sandbox-metering"; -import { assertProjectSandboxOwnerActive } from "./project-sandbox-owner-admission"; -import { PROC_PREFIX, PROCESS_PORT_ALLOC_KEY } from "./project-sandbox-process-support"; -import { ProjectSandboxProvisioning } from "./project-sandbox-provisioning"; -import type { - ParsedProjectCleanupWorkspaceInput, - ProjectSandboxRuntimeState, -} from "./project-sandbox-runtime"; -import { - initializeProjectSandboxStorage, - openProjectSandboxWorkspaceState, - ProjectSandboxWorkspaceState, -} from "./project-sandbox-workspace-state"; - -const ClearWorkspaceEvidenceSchema = z.object({ cleared: z.literal(true) }).strict(); -const RUNTIME_RESET_PENDING_KEY = "sandbox_runtime_reset_pending"; -const SKILL_RUNTIME_DIRECTORY = "/workspace/.cheatcode/runtime"; - -export abstract class ProjectSandboxLifecycle extends DurableObject { - private accountDeletionCompleted = false; - private accountDeletionInProgress = false; - private accountDeletionPromise: Promise | undefined; - private activeOperationCount = 0; - private readonly activeOperationDrainWaiters = new Set<() => void>(); - private daytonaClient: DaytonaClient | undefined; - private daytonaId: string | undefined; - private sandboxMutationTail: Promise = Promise.resolve(); - private sandboxRuntimeUpdateInProgress = false; - private startedVerifiedAtMs = 0; - private readonly identityState: ProjectSandboxIdentityState; - private readonly provisioning: ProjectSandboxProvisioning; - private workspaceStateValue: ProjectSandboxWorkspaceState | undefined; - - constructor(ctx: DurableObjectState, env: ProjectSandboxEnv) { - super(ctx, env); - this.identityState = new ProjectSandboxIdentityState(ctx); - this.provisioning = new ProjectSandboxProvisioning({ - cachedSandboxId: async () => this.daytonaId ?? this.storedDaytonaId(), - env, - sandboxName: () => this.sandboxName(), - toUpstreamError: (error, fallback) => this.toUpstreamError(error, fallback), - }); - this.workspaceStateValue = openProjectSandboxWorkspaceState(ctx); - void ctx.blockConcurrencyWhile(() => this.initializeIdentityState()); - } - - public deleteAccountState(): Promise { - if (this.accountDeletionCompleted) { - return Promise.resolve(); - } - if (this.accountDeletionPromise !== undefined) { - return this.accountDeletionPromise; - } - - this.accountDeletionInProgress = true; - const deletion = this.performAccountDeletion(); - const tracked = deletion.finally(() => { - if (this.accountDeletionPromise === tracked) { - this.accountDeletionPromise = undefined; - } - }); - this.accountDeletionPromise = tracked; - return tracked; - } - - protected withActiveSandboxOperation(operation: () => Promise): Promise { - return this.withActiveOperation(null, operation, false, true); - } - protected withActiveOwnerRegistration( - userId: string, - operation: () => Promise, - ): Promise { - if (this.identityState.hasRegisteredOwner()) { - return this.withActiveOperation(null, operation); - } - let release: (() => void) | undefined; - try { - release = this.acquireActiveSandboxOperation(true); - return assertProjectSandboxOwnerActive(this.env, userId) - .then(() => { - if (this.accountDeletionInProgress) { - throw accountSandboxDeletedError(); - } - return operation().then((result) => { - if (this.accountDeletionInProgress) { - throw accountSandboxDeletedError(); - } - this.ensureWorkspaceState(); - return result; - }); - }) - .finally(release); - } catch (error) { - release?.(); - return Promise.reject(error); - } - } - protected withActiveSharedWorkspaceMutation(operation: () => Promise): Promise { - return this.withActiveOperation( - null, - async () => { - await this.workspaceState.waitForWorkspaceDrain(); - return operation(); - }, - true, - ); - } - protected withActiveProjectWorkspaceOperation( - workspaceScope: string | readonly string[] | null, - operation: () => Promise, - ): Promise { - return this.withActiveOperation(workspaceScope, operation, false, true); - } - private withActiveOperation( - workspaceScope: string | readonly string[] | null, - operation: () => Promise, - isSharedMutation = false, - shouldLeaseUnknownWorkspace = false, - allowRuntimeUpdate = false, - ): Promise { - let release: (() => void) | undefined; - try { - release = this.acquireActiveOperation( - workspaceScope, - isSharedMutation, - shouldLeaseUnknownWorkspace, - allowRuntimeUpdate, - ); - return operation().finally(release); - } catch (error) { - release?.(); - return Promise.reject(error); - } - } - - protected withActiveProjectWorkspaceStreamingOperation( - workspaceScope: string | readonly string[] | null, - operation: (release: () => void) => Promise, - ): Promise { - return this.withActiveStreamingOperation(workspaceScope, operation); - } - private withActiveStreamingOperation( - workspaceScope: string | readonly string[] | null, - operation: (release: () => void) => Promise, - ): Promise { - let release: (() => void) | undefined; - try { - release = this.acquireActiveOperation(workspaceScope, false, true); - return operation(release).catch((error: unknown) => { - release?.(); - throw error; - }); - } catch (error) { - release?.(); - return Promise.reject(error); - } - } - - protected withActiveSandboxCleanupSignal(operation: () => Promise): Promise { - return this.accountDeletionInProgress || !this.identityState.hasRegisteredOwner() - ? Promise.resolve() - : this.withActiveOperation(null, operation, false, false, true); - } - - private async initializeIdentityState(): Promise { - const isAccountDeleted = (await this.ctx.storage.get(ACCOUNT_DELETION_TOMBSTONE_KEY)) === true; - if (isAccountDeleted) { - this.accountDeletionInProgress = true; - } - await this.identityState.initialize(); - } - - private get workspaceState(): ProjectSandboxWorkspaceState { - return this.ensureWorkspaceState(); - } - - private ensureWorkspaceState(): ProjectSandboxWorkspaceState { - if (!this.workspaceStateValue) { - initializeProjectSandboxStorage(this.ctx); - this.workspaceStateValue = new ProjectSandboxWorkspaceState(this.ctx); - } - return this.workspaceStateValue; - } - - protected deleteProjectWorkspace( - input: ParsedProjectCleanupWorkspaceInput, - cleanup: () => Promise, - ): Promise { - return this.workspaceState.deleteWorkspace(input, cleanup); - } - - protected withActiveProjectWorkspaceCleanup(operation: () => Promise): Promise { - let release: (() => void) | undefined; - try { - // Cleanup itself must not take a workspace lease: its durable tombstone - // blocks new work, then it drains every existing lease before killing - // sandbox-wide processes. - release = this.acquireActiveSandboxOperation(false, true); - return operation().finally(release); - } catch (error) { - release?.(); - return Promise.reject(error); - } - } +import type { ProjectSandboxRuntimeState } from "./project-sandbox-runtime"; +import type { SandboxRuntime } from "./project-sandbox-runtime-handle"; + +export interface LifecycleOps { + alarm: () => Promise; + beginRun: (runId: string) => Promise; + deleteAccountState: () => Promise; + endRun: (runId: string) => Promise; + existingDaytonaId: () => Promise; + registerOwner: (userId: string, sandboxName?: string) => Promise; + renewRun: (runId: string) => Promise; + runtimeSandboxId: () => Promise; + sandboxRuntimeState: () => Promise; + setQuotaPeriod: (periodEndIso: string) => Promise; +} - public async registerOwner(userId: string, sandboxName?: string): Promise { - await this.identityState.registerOwner(userId, sandboxName); - } +type LifecycleRuntime = Pick< + SandboxRuntime, + | "client" + | "deleteAccountState" + | "ensureSandbox" + | "existingSandboxId" + | "meteringContext" + | "registerOwner" + | "sandboxName" + | "storage" +>; + +export function createLifecycleOps(runtime: LifecycleRuntime): LifecycleOps { + return { + alarm: () => handleAlarm(runtime), + beginRun: (runId) => beginRun(runtime, runId), + deleteAccountState: runtime.deleteAccountState, + endRun: (runId) => endRun(runtime, runId), + existingDaytonaId: runtime.existingSandboxId, + registerOwner: runtime.registerOwner, + renewRun: (runId) => renewRun(runtime, runId), + runtimeSandboxId: () => runtime.ensureSandbox(), + sandboxRuntimeState: () => sandboxRuntimeState(runtime), + setQuotaPeriod: (periodEndIso) => setSandboxQuotaPeriod(runtime.storage, periodEndIso), + }; +} - public async setQuotaPeriod(periodEndIso: string): Promise { - await setSandboxQuotaPeriod(this.ctx.storage, periodEndIso); +async function beginRun(runtime: LifecycleRuntime, runId: string): Promise { + const leases = await runLeases(runtime.storage); + const remaining = leases.filter((lease) => lease.runId !== runId); + remaining.push({ runId, startedMs: Date.now() }); + await runtime.storage.put(RUN_LEASES_KEY, remaining); + try { + const id = await runtime.ensureSandbox(runId); + await runtime + .client() + .setAutoStopInterval(id, 0) + .catch(() => undefined); + await beginSandboxUsageBestEffort(await runtime.meteringContext()); + await runtime.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); + } catch (error) { + await compensateFailedRunStart(runtime, runId); + throw error; } +} - public async beginRun(runId: string): Promise { - const leases = await this.runLeases(); - const remaining = leases.filter((lease) => lease.runId !== runId); - remaining.push({ runId, startedMs: Date.now() }); - await this.ctx.storage.put(RUN_LEASES_KEY, remaining); - try { - const id = await this.ensureSandbox(runId); - await this.client() - .setAutoStopInterval(id, 0) - .catch(() => undefined); - await beginSandboxUsageBestEffort(await this.meteringContext()); - await this.ctx.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); - } catch (error) { - await this.compensateFailedRunStart(runId); - throw error; - } +async function renewRun(runtime: LifecycleRuntime, runId: string): Promise { + const leases = await runLeases(runtime.storage); + const lease = leases.find((candidate) => candidate.runId === runId); + if (!lease) { + return; } + const renewed = leases.filter((candidate) => candidate.runId !== runId); + renewed.push({ runId, startedMs: Date.now() }); + await runtime.storage.put(RUN_LEASES_KEY, renewed); + await runtime.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); +} - public async renewRun(runId: string): Promise { - const leases = await this.runLeases(); - const lease = leases.find((candidate) => candidate.runId === runId); - if (!lease) { - return; - } - const renewed = leases.filter((candidate) => candidate.runId !== runId); - renewed.push({ runId, startedMs: Date.now() }); - await this.ctx.storage.put(RUN_LEASES_KEY, renewed); - await this.ctx.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); +async function endRun(runtime: LifecycleRuntime, runId: string): Promise { + const remaining = (await runLeases(runtime.storage)).filter((lease) => lease.runId !== runId); + await runtime.storage.put(RUN_LEASES_KEY, remaining); + if (remaining.length > 0) { + await recordSandboxUsageBestEffort(await runtime.meteringContext()); + return; } + await finalizeLastRunLease(runtime); +} - public async endRun(runId: string): Promise { - const remaining = (await this.runLeases()).filter((lease) => lease.runId !== runId); - await this.ctx.storage.put(RUN_LEASES_KEY, remaining); +async function compensateFailedRunStart(runtime: LifecycleRuntime, runId: string): Promise { + try { + const remaining = (await runLeases(runtime.storage)).filter((lease) => lease.runId !== runId); + await runtime.storage.put(RUN_LEASES_KEY, remaining); if (remaining.length > 0) { - await recordSandboxUsageBestEffort(await this.meteringContext()); - return; - } - await this.finalizeLastRunLease(); - } - - private async compensateFailedRunStart(runId: string): Promise { - try { - const remaining = (await this.runLeases()).filter((lease) => lease.runId !== runId); - await this.ctx.storage.put(RUN_LEASES_KEY, remaining); - if (remaining.length > 0) { - await this.ctx.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); - return; - } - await this.finalizeLastRunLease(); - } catch (error) { - createLogger().error("sandbox_run_lease_rollback_failed", { - error, - runId, - sandboxId: this.sandboxName(), - }); - await this.ctx.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS).catch(() => undefined); - } - } - - private async finalizeLastRunLease(): Promise { - await finalizeSandboxUsageBestEffort(await this.meteringContext()); - if (await this.restoreIdleAutoStop()) { - await this.ctx.storage.deleteAlarm(); - return; - } - await this.ctx.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); - } - - override async alarm(): Promise { - const leases = (await this.runLeases()).filter( - (lease) => Date.now() - lease.startedMs < STALE_RUN_LEASE_MS, - ); - await this.ctx.storage.put(RUN_LEASES_KEY, leases); - if (leases.length === 0) { - await this.finalizeLastRunLease(); - return; - } - await this.refreshActiveSandboxBestEffort(); - try { - await recordSandboxUsageBestEffort(await this.meteringContext()); - } finally { - await this.ctx.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); - } - } - - private async refreshActiveSandboxBestEffort(): Promise { - try { - const id = await this.existingSandboxId(); - if (id) { - await this.client().refreshActivity(id); - } - } catch (error) { - createLogger().warn("sandbox_keepalive_refresh_failed", { - error, - sandboxId: this.sandboxName(), - }); - } - } - - public async runtimeSandboxId(): Promise { - return this.ensureSandbox(); - } - - public async existingDaytonaId(): Promise { - return this.existingSandboxId(); - } - - public async sandboxRuntimeState(): Promise { - const existing = await this.existingSandboxId(); - if (!existing) { - return { state: "none" }; - } - const sandbox = await this.client().getSandbox(existing); - return { sandboxId: existing, state: sandbox?.state ?? "unknown" }; - } - - private async performAccountDeletion(): Promise { - await this.ctx.storage.put(ACCOUNT_DELETION_TOMBSTONE_KEY, true); - await this.waitForActiveSandboxOperations(); - await this.withSandboxMutation(async () => { - await finalizeSandboxUsageBestEffort(await this.meteringContext()); - await this.destroySandboxExclusive(); - await this.clearDurableStateForDeletedAccount(); - }); - this.accountDeletionCompleted = true; - } - - private async destroySandboxExclusive(): Promise { - const client = await this.ensureClient(); - try { - const canonical = await this.provisioning.findExisting(client); - const owned = await this.provisioning.findOwned(client); - const sandboxes = uniqueSandboxes(canonical ? [canonical, ...owned] : owned); - const volumeSandbox = sandboxes.find( - (sandbox) => sandbox.labels["workspaceVolumeName"] === this.env.DAYTONA_WORKSPACE_VOLUME, - ); - if (volumeSandbox) { - if (!(await this.provisioning.ensureStarted(client, volumeSandbox))) { - throw new APIError(502, "upstream_sandbox_failed", "Daytona workspace disappeared", { - retriable: true, - }); - } - const cleared = await client.execute(volumeSandbox.id, { - command: clearWorkspaceCommand(), - timeout: 480, - }); - if ( - cleared.exitCode !== 0 || - !ClearWorkspaceEvidenceSchema.safeParse(parseSandboxJson(cleared.result)).success - ) { - throw new APIError(502, "upstream_sandbox_failed", "Daytona workspace deletion failed", { - details: { sandboxId: this.sandboxName() }, - retriable: true, - }); - } - } - for (const sandbox of sandboxes) { - await client.deleteSandbox(sandbox.id); - } - } catch (error) { - throw this.toUpstreamError(error, "Daytona sandbox deletion failed."); - } - return this.finishSandboxDestruction(); - } - - private async finishSandboxDestruction(): Promise { - this.clearCachedSandbox(); - await this.ctx.storage.delete(DAYTONA_ID_KEY); - await clearSandboxMeterState(this.ctx.storage); - return { deleted: true, sandboxId: this.sandboxName() }; - } - - private async clearDurableStateForDeletedAccount(): Promise { - await this.ctx.storage.deleteAll(); - this.workspaceStateValue = undefined; - this.identityState.clearRegisteredOwner(); - this.clearCachedSandbox(); - } - - private waitForActiveSandboxOperations(): Promise { - if (this.activeOperationCount === 0) { - return Promise.resolve(); - } - return new Promise((resolve) => { - this.activeOperationDrainWaiters.add(resolve); - }); - } - - private finishActiveSandboxOperation(): void { - this.activeOperationCount -= 1; - if (this.activeOperationCount !== 0) { + await runtime.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); return; } - for (const resolve of this.activeOperationDrainWaiters) { - resolve(); - } - this.activeOperationDrainWaiters.clear(); - } - - private acquireActiveSandboxOperation( - allowUnregisteredOwner = false, - allowWorkspaceCleanup = false, - allowRuntimeUpdate = false, - ): () => void { - if (this.accountDeletionInProgress) { - throw accountSandboxDeletedError(); - } - if (this.sandboxRuntimeUpdateInProgress && !allowRuntimeUpdate) { - throw sandboxRuntimeUpdatePending(this.env.DAYTONA_SANDBOX_SNAPSHOT); - } - if (!allowUnregisteredOwner && !this.identityState.hasRegisteredOwner()) { - throw accountSandboxDeletedError(); - } - const workspaceState = - allowUnregisteredOwner && !this.identityState.hasRegisteredOwner() - ? this.workspaceStateValue - : this.ensureWorkspaceState(); - workspaceState?.assertOperationAllowed(allowWorkspaceCleanup); - this.activeOperationCount += 1; - let isReleased = false; - return () => { - if (isReleased) { - return; - } - isReleased = true; - this.finishActiveSandboxOperation(); - }; - } - - private acquireActiveOperation( - workspaceScope: string | readonly string[] | null, - isSharedMutation = false, - shouldLeaseUnknownWorkspace = false, - allowRuntimeUpdate = false, - ): () => void { - const releaseSandbox = this.acquireActiveSandboxOperation(false, false, allowRuntimeUpdate); - let releaseWorkspace: (() => void) | undefined; - try { - const workspaceSlugs = - typeof workspaceScope === "string" ? [workspaceScope] : (workspaceScope ?? []); - releaseWorkspace = isSharedMutation - ? this.workspaceState.acquireSharedMutation() - : workspaceSlugs.length > 0 - ? this.workspaceState.acquire(workspaceSlugs) - : shouldLeaseUnknownWorkspace - ? this.workspaceState.acquireUnscoped() - : undefined; - } catch (error) { - releaseSandbox(); - throw error; - } - return () => { - releaseWorkspace?.(); - releaseSandbox(); - }; - } - - protected client(): DaytonaClient { - if (!this.daytonaClient) { - throw new Error("Daytona client accessed before initialization."); - } - return this.daytonaClient; - } - - protected async ensureClient(): Promise { - if (this.daytonaClient) { - return this.daytonaClient; - } - const apiKey = await resolveWorkerSecret(this.env.DAYTONA_API_KEY); - if (!apiKey) { - throw new APIError(503, "unavailable_maintenance", "DAYTONA_API_KEY is not configured", { - retriable: false, - }); - } - this.daytonaClient = new DaytonaClient({ - apiKey, - apiUrl: this.env.DAYTONA_API_URL, - target: this.env.DAYTONA_TARGET, - ...(this.env.DAYTONA_ORG_ID ? { organizationId: this.env.DAYTONA_ORG_ID } : {}), - ...(this.env.DAYTONA_PREVIEW_HOST_SUFFIXES - ? { previewHostSuffixes: this.env.DAYTONA_PREVIEW_HOST_SUFFIXES } - : {}), - }); - return this.daytonaClient; - } - - protected sandboxProvisioning(): ProjectSandboxProvisioning { - return this.provisioning; - } - - protected adoptDaytonaId(sandboxId: string): void { - this.daytonaId = sandboxId; - this.startedVerifiedAtMs = Date.now(); - } - - protected async ensureSandbox(startingRunId?: string): Promise { - return this.withSandboxMutation(async () => { - if (this.daytonaId && Date.now() - this.startedVerifiedAtMs < STARTED_REVERIFY_MS) { - return this.daytonaId; - } - return this.resolveStartedSandbox(startingRunId); - }); - } - - protected async restartSandboxForWorkspaceRecovery(sandboxId: string): Promise { - await this.withSandboxMutation(async () => { - const client = await this.ensureClient(); - try { - await this.provisioning.restart(client, sandboxId); - } catch (error) { - throw this.toUpstreamError(error, "Daytona workspace recovery failed."); - } - this.daytonaId = sandboxId; - this.startedVerifiedAtMs = Date.now(); + await finalizeLastRunLease(runtime); + } catch (error) { + createLogger().error("sandbox_run_lease_rollback_failed", { + error, + runId, + sandboxId: runtime.sandboxName(), }); + await runtime.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS).catch(() => undefined); } +} - protected async ensureExistingSandboxStarted(): Promise { - return this.withSandboxMutation(async () => { - const client = await this.ensureClient(); - let existing: DaytonaSandbox | null; - try { - existing = await this.provisioning.findExisting(client); - if (!existing) { - return null; - } - if (!(await this.provisioning.ensureStarted(client, existing))) { - return null; - } - } catch (error) { - throw this.toUpstreamError(error, "Daytona sandbox cleanup startup failed."); - } - this.daytonaId = existing.id; - await this.ctx.storage.put(DAYTONA_ID_KEY, existing.id); - this.startedVerifiedAtMs = Date.now(); - return existing.id; - }); - } - - private async withSandboxMutation(operation: () => Promise): Promise { - const previous = this.sandboxMutationTail; - let release = (): void => undefined; - const gate = new Promise((resolve) => { - release = resolve; - }); - this.sandboxMutationTail = previous.catch(() => undefined).then(() => gate); - await previous.catch(() => undefined); - try { - return await operation(); - } finally { - release(); - } - } - - private async resolveStartedSandbox(startingRunId?: string): Promise { - const client = await this.ensureClient(); - let resolved: DaytonaSandbox; - try { - resolved = await this.provisioning.resolve(client); - if (!this.provisioning.isDesired(resolved)) { - resolved = await this.replaceSandboxRuntime(client, resolved, startingRunId); - } - } catch (error) { - throw this.toUpstreamError(error, "Daytona sandbox lookup failed."); - } - this.daytonaId = resolved.id; - await this.ctx.storage.put(DAYTONA_ID_KEY, resolved.id); - if (!(await this.provisioning.ensureStarted(client, resolved))) { - throw new APIError(502, "upstream_sandbox_failed", "Daytona sandbox disappeared", { - retriable: true, - }); - } - await this.clearPersistedRuntimeProjection(client, resolved.id); - this.startedVerifiedAtMs = Date.now(); - return resolved.id; - } - - private async replaceSandboxRuntime( - client: DaytonaClient, - current: DaytonaSandbox, - startingRunId?: string, - ): Promise { - this.sandboxRuntimeUpdateInProgress = true; - try { - await this.assertSandboxReplacementAllowed(startingRunId); - this.provisioning.assertRuntimeReplacementSafe(current); - await this.prepareForSandboxReplacement(); - await this.provisioning.deleteForReplacement(client, current); - const replacement = await this.provisioning.create(client); - if (!this.provisioning.isDesired(replacement)) { - throw sandboxRuntimeUpdatePending(this.env.DAYTONA_SANDBOX_SNAPSHOT); - } - createLogger().info("sandbox_runtime_replaced", { - sandboxId: this.sandboxName(), - snapshot: this.env.DAYTONA_SANDBOX_SNAPSHOT, - }); - return replacement; - } finally { - this.sandboxRuntimeUpdateInProgress = false; - } - } - - private async assertSandboxReplacementAllowed(startingRunId?: string): Promise { - const leases = await this.runLeases(); - const activeRunLeases = leases.filter( - (lease) => Date.now() - lease.startedMs < STALE_RUN_LEASE_MS, - ); - if (activeRunLeases.length !== leases.length) { - await this.ctx.storage.put(RUN_LEASES_KEY, activeRunLeases); - } - const otherRunLeases = activeRunLeases.filter((lease) => lease.runId !== startingRunId); - if (this.activeOperationCount > 1 || otherRunLeases.length > 0) { - throw sandboxRuntimeUpdatePending(this.env.DAYTONA_SANDBOX_SNAPSHOT); - } - } - - private async prepareForSandboxReplacement(): Promise { - this.daytonaId = undefined; - this.startedVerifiedAtMs = 0; - await this.ctx.storage.delete(DAYTONA_ID_KEY); - await this.ctx.storage.put(RUNTIME_RESET_PENDING_KEY, true); - const processRecords = await this.ctx.storage.list({ prefix: PROC_PREFIX }); - if (processRecords.size > 0) { - await this.ctx.storage.delete([...processRecords.keys()]); - } - await this.ctx.storage.delete(PROCESS_PORT_ALLOC_KEY); - } - - private async clearPersistedRuntimeProjection( - client: DaytonaClient, - sandboxId: string, - ): Promise { - if ((await this.ctx.storage.get(RUNTIME_RESET_PENDING_KEY)) !== true) { - return; - } - try { - await client.deleteFilePath(sandboxId, SKILL_RUNTIME_DIRECTORY, true); - await this.ctx.storage.delete(RUNTIME_RESET_PENDING_KEY); - } catch (error) { - throw this.toUpstreamError(error, "Daytona runtime reset failed."); - } - } - - protected async existingSandboxId(): Promise { - const client = await this.ensureClient(); - try { - const existing = await this.provisioning.findExisting(client); - if (existing) { - this.daytonaId = existing.id; - return existing.id; - } - return null; - } catch (error) { - throw this.toUpstreamError(error, "Daytona sandbox lookup failed."); - } - } - - protected async meteringContext(): Promise { - return { - env: this.env, - ownerUserId: this.identityState.ownerUserId(), - sandboxId: this.sandboxName(), - storage: this.ctx.storage, - }; - } - - protected async previewSecret(): Promise { - const secret = await resolveWorkerSecret(this.env.PREVIEW_TOKEN_SECRET); - if (!secret) { - throw new APIError(503, "unavailable_maintenance", "PREVIEW_TOKEN_SECRET is not configured", { - retriable: false, - }); - } - return secret; - } - - protected previewHostname(): string { - return PreviewHostnameSchema.parse(this.env.PREVIEW_HOSTNAME); - } - - protected sandboxName(): string { - return this.identityState.sandboxName(); +async function finalizeLastRunLease(runtime: LifecycleRuntime): Promise { + await finalizeSandboxUsageBestEffort(await runtime.meteringContext()); + if (await restoreIdleAutoStop(runtime)) { + await runtime.storage.deleteAlarm(); + return; } + await runtime.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); +} - protected ownerUserId(): string { - const userId = this.identityState.ownerUserId(); - if (!userId) { - throw new APIError(500, "internal_error", "ProjectSandbox owner is not registered", { - retriable: false, - }); - } - return userId; +async function handleAlarm(runtime: LifecycleRuntime): Promise { + const leases = (await runLeases(runtime.storage)).filter( + (lease) => Date.now() - lease.startedMs < STALE_RUN_LEASE_MS, + ); + await runtime.storage.put(RUN_LEASES_KEY, leases); + if (leases.length === 0) { + await finalizeLastRunLease(runtime); + return; } - - protected writeExecAudit(entry: SandboxExecAuditEntry): Promise { - return writeExecAudit(this.env.R2_AUDIT, entry); + await refreshActiveSandboxBestEffort(runtime); + try { + await recordSandboxUsageBestEffort(await runtime.meteringContext()); + } finally { + await runtime.storage.setAlarm(Date.now() + KEEPALIVE_ALARM_MS); } +} - protected toUpstreamError(error: unknown, fallback: string): APIError { - if (error instanceof APIError) { - return error; +async function refreshActiveSandboxBestEffort(runtime: LifecycleRuntime): Promise { + try { + const id = await runtime.existingSandboxId(); + if (id) { + await runtime.client().refreshActivity(id); } - const status = error instanceof DaytonaApiError ? error.status : 502; - const retriable = error instanceof DaytonaApiError ? error.retriable : true; - return new APIError(status >= 500 ? 502 : status, "upstream_sandbox_failed", fallback, { - cause: error, - details: { sandboxId: this.sandboxName() }, - hint: "Retry. If it persists, check Daytona sandbox lifecycle status.", - retriable, + } catch (error) { + createLogger().warn("sandbox_keepalive_refresh_failed", { + error, + sandboxId: runtime.sandboxName(), }); } +} - private async storedDaytonaId(): Promise { - const value = await this.ctx.storage.get(DAYTONA_ID_KEY); - return typeof value === "string" ? value : null; - } - - private async runLeases(): Promise> { - return RunLeasesSchema.parse((await this.ctx.storage.get(RUN_LEASES_KEY)) ?? []); +async function sandboxRuntimeState(runtime: LifecycleRuntime): Promise { + const existing = await runtime.existingSandboxId(); + if (!existing) { + return { state: "none" }; } + const sandbox = await runtime.client().getSandbox(existing); + return { sandboxId: existing, state: sandbox?.state ?? "unknown" }; +} - private async restoreIdleAutoStop(): Promise { - try { - const id = await this.existingSandboxId(); - if (!id) { - return true; - } - await this.client().setAutoStopInterval(id, DEFAULT_IDLE_STOP_MIN); +async function restoreIdleAutoStop(runtime: LifecycleRuntime): Promise { + try { + const id = await runtime.existingSandboxId(); + if (!id) { return true; - } catch (error) { - createLogger().warn("sandbox_autostop_restore_failed", { - error, - sandboxId: this.sandboxName(), - }); - return false; } + await runtime.client().setAutoStopInterval(id, DEFAULT_IDLE_STOP_MIN); + return true; + } catch (error) { + createLogger().warn("sandbox_autostop_restore_failed", { + error, + sandboxId: runtime.sandboxName(), + }); + return false; } - - private clearCachedSandbox(): void { - this.daytonaClient = undefined; - this.daytonaId = undefined; - this.startedVerifiedAtMs = 0; - } -} - -function sandboxRuntimeUpdatePending(expectedSnapshot: string): APIError { - return new APIError( - 503, - "unavailable_maintenance", - "This computer is updating to the current runtime.", - { - details: { expectedSnapshot }, - hint: "Retry after the active operation finishes.", - retriable: true, - }, - ); } diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-process-cleanup.ts b/apps/agent-worker/src/durable-objects/project-sandbox-process-cleanup.ts index 60ddb39a..a61dd366 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-process-cleanup.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-process-cleanup.ts @@ -1,72 +1,3 @@ -/** Kills every sandbox process whose real cwd is the exact project directory or a child. */ -export const WORKSPACE_PROCESS_TERMINATION_SCRIPT = ` -import os -import signal -import stat -import sys -import time - -root = sys.argv[1] -try: - metadata = os.lstat(root) -except FileNotFoundError: - metadata = None -if metadata is not None and stat.S_ISLNK(metadata.st_mode): - print("Project workspace is a symbolic link", file=sys.stderr) - raise SystemExit(2) - -root = os.path.realpath(root) if metadata is not None else os.path.abspath(root) -prefix = root + os.sep -self_pid = os.getpid() -self_uid = os.getuid() - -def matching_pids(): - matches = [] - for name in os.listdir("/proc"): - if not name.isdigit(): - continue - pid = int(name) - if pid == self_pid: - continue - try: - if os.stat(f"/proc/{pid}").st_uid != self_uid: - continue - cwd = os.readlink(f"/proc/{pid}/cwd") - except (FileNotFoundError, ProcessLookupError): - continue - except OSError as error: - raise RuntimeError(f"Could not inspect cwd for process {pid}: {error}") from error - if cwd.endswith(" (deleted)"): - cwd = cwd.removesuffix(" (deleted)") - if cwd == root or cwd.startswith(prefix): - matches.append(pid) - return matches - -def send(pids, sig): - denied = [] - for pid in pids: - try: - os.kill(pid, sig) - except ProcessLookupError: - continue - except PermissionError: - denied.append(pid) - return denied - -denied = send(matching_pids(), signal.SIGTERM) -deadline = time.monotonic() + 3 -remaining = matching_pids() -while remaining and time.monotonic() < deadline: - time.sleep(0.1) - remaining = matching_pids() -denied.extend(send(remaining, signal.SIGKILL)) -time.sleep(0.1) -survivors = matching_pids() -if denied or survivors: - print(f"Could not terminate workspace processes: denied={denied} survivors={survivors}", file=sys.stderr) - raise SystemExit(1) -`; - /** Kills every same-user sandbox process except the cleanup command's control-plane ancestry. */ export const SANDBOX_PROCESS_TERMINATION_SCRIPT = ` import os diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-process-control.ts b/apps/agent-worker/src/durable-objects/project-sandbox-process-control.ts new file mode 100644 index 00000000..ecac63d6 --- /dev/null +++ b/apps/agent-worker/src/durable-objects/project-sandbox-process-control.ts @@ -0,0 +1,507 @@ +import { APIError } from "@cheatcode/observability"; +import type { DaytonaSessionExecResponse } from "@cheatcode/tools-code"; +import { WORKSPACE_DIR } from "./project-sandbox-content-support"; +import { SANDBOX_PROCESS_TERMINATION_SCRIPT } from "./project-sandbox-process-cleanup"; +import { + ENV_FILE_DIR, + MAX_TRACKED_PROCESSES, + type ParsedProcessStartInput, + PORT_ALLOC_KEY, + PortAllocationSchema, + PROC_PREFIX, + PROCESS_PORT_ALLOC_KEY, + ProcessPortReservationsSchema, + type ProcessRecord, + ProcessRecordSchema, + restartEnvironment, + shellQuote, + sleep, + supervisedProcessCommand, + timeoutSeconds, + withoutProcessReservation, +} from "./project-sandbox-process-support"; +import type { SandboxRuntime } from "./project-sandbox-runtime-handle"; + +type ProcessControlRuntime = Pick; + +export interface ProcessControl { + cleanupLaunchedProcess: (id: string, sessionId: string, name: string) => Promise; + deleteProcessRecord: (id: string, name: string, keepPortReservation?: boolean) => Promise; + deleteProcessesOnPort: (id: string, port: number, exceptName: string) => Promise; + freeProjectPort: (workspaceSlug: string) => Promise; + httpPortReady: (id: string, port: number, path: string, timeoutMs: number) => Promise; + isPortAlive: (id: string, port: number) => Promise; + launchSessionProcess: ( + id: string, + sessionId: string, + name: string, + cwd: string, + command: string, + env: Record | undefined, + stdin?: string, + ) => Promise; + persistProcessOwnershipIntent: (name: string, record: ProcessRecord) => Promise; + persistStartedProcess: ( + id: string, + name: string, + record: ProcessRecord, + waitForPort: ParsedProcessStartInput["waitForPort"], + ) => Promise; + prepareProcessSlot: (id: string, name: string, input: ParsedProcessStartInput) => Promise; + processRecord: (name: string) => Promise; + relaunchDevServer: ( + id: string, + name: string, + record: ProcessRecord, + restartEnv?: Record, + ) => Promise; + releaseProcessPort: (processId: string) => Promise; + terminateUntrackedSandboxProcesses: (id: string) => Promise; + waitForPort: ( + id: string, + port: number, + path: string | undefined, + timeoutMs: number | undefined, + process?: { cmdId: string; sessionId: string }, + ) => Promise; +} + +export function createProcessControl(runtime: ProcessControlRuntime): ProcessControl { + const control: ProcessControl = { + cleanupLaunchedProcess: (id, sessionId, name) => + cleanupLaunchedProcess(runtime, id, sessionId, name), + deleteProcessRecord: (id, name, keepPortReservation) => + deleteProcessRecord(runtime, id, name, keepPortReservation), + deleteProcessesOnPort: (id, port, exceptName) => + deleteProcessesOnPort(runtime, control, id, port, exceptName), + freeProjectPort: (workspaceSlug) => freeProjectPort(runtime, workspaceSlug), + httpPortReady: (id, port, path, timeoutMs) => httpPortReady(runtime, id, port, path, timeoutMs), + isPortAlive: (id, port) => isPortAlive(runtime, id, port), + launchSessionProcess: (id, sessionId, name, cwd, command, env, stdin) => + launchSessionProcess(runtime, id, sessionId, name, cwd, command, env, stdin), + persistProcessOwnershipIntent: (name, record) => + persistProcessOwnershipIntent(runtime, name, record), + persistStartedProcess: (id, name, record, waitForPort) => + persistStartedProcess(runtime, control, id, name, record, waitForPort), + prepareProcessSlot: (id, name, input) => prepareProcessSlot(runtime, control, id, name, input), + processRecord: (name) => processRecord(runtime, name), + relaunchDevServer: (id, name, record, restartEnv) => + relaunchDevServer(runtime, control, id, name, record, restartEnv), + releaseProcessPort: (processId) => releaseProcessPort(runtime, processId), + terminateUntrackedSandboxProcesses: (id) => terminateUntrackedSandboxProcesses(runtime, id), + waitForPort: (id, port, path, timeoutMs, process) => + waitForPort(runtime, id, port, path, timeoutMs, process), + }; + return control; +} + +async function isPortAlive( + runtime: ProcessControlRuntime, + id: string, + port: number, +): Promise { + const probe = await runtime + .client() + .execute(id, { + command: `curl -sf -o /dev/null --max-time 3 http://localhost:${port}/`, + timeout: 5, + }) + .catch(() => null); + return probe?.exitCode === 0; +} + +async function relaunchDevServer( + runtime: ProcessControlRuntime, + control: ProcessControl, + id: string, + name: string, + record: ProcessRecord, + restartEnv?: Record, +): Promise { + const sessionId = record.sessionId || `cc-${name}`; + await runtime.client().deleteSession(id, sessionId); + const exec = await control.launchSessionProcess( + id, + sessionId, + name, + record.cwd, + supervisedProcessCommand(record.command, record), + restartEnv ?? restartEnvironment(name, record), + ); + try { + await runtime.storage.put(`${PROC_PREFIX}${name}`, { + ...record, + cmdId: exec.cmdId ?? sessionId, + startedAtMs: Date.now(), + } satisfies ProcessRecord); + } catch (error) { + await control.cleanupLaunchedProcess(id, sessionId, name); + throw error; + } +} + +async function waitForPort( + runtime: ProcessControlRuntime, + id: string, + port: number, + path: string | undefined, + timeoutMs: number | undefined, + process?: { cmdId: string; sessionId: string }, +): Promise { + const deadline = Date.now() + (timeoutMs ?? 120_000); + const url = `http://localhost:${port}${path ?? "/"}`; + while (Date.now() < deadline) { + const probe = await runtime + .client() + .execute(id, { + command: `curl -sf -o /dev/null --max-time 3 ${shellQuote(url)}`, + timeout: 5, + }) + .catch(() => null); + if (probe?.exitCode === 0) return; + if (process) { + await throwIfProcessExited(runtime, id, port, process); + } + await sleep(1_500); + } + throw new APIError(504, "upstream_timeout_sandbox", "Sandbox process did not become ready.", { + details: { port, timeoutMs: timeoutMs ?? 120_000, url }, + hint: "Inspect the process command and logs, then retry.", + retriable: true, + }); +} + +async function httpPortReady( + runtime: ProcessControlRuntime, + id: string, + port: number, + path: string, + timeoutMs: number, +): Promise { + const deadline = Date.now() + timeoutMs; + const url = `http://localhost:${port}${path}`; + while (Date.now() < deadline) { + const probe = await runtime + .client() + .execute(id, { + command: `curl -sf -o /dev/null --max-time 3 ${shellQuote(url)}`, + timeout: 5, + }) + .catch(() => null); + if (probe?.exitCode === 0) { + return true; + } + await sleep(1_000); + } + return false; +} + +async function processRecord( + runtime: ProcessControlRuntime, + name: string, +): Promise { + const value = await runtime.storage.get(`${PROC_PREFIX}${name}`); + const parsed = ProcessRecordSchema.safeParse(value); + return parsed.success ? parsed.data : null; +} + +async function deleteProcessRecord( + runtime: ProcessControlRuntime, + id: string, + name: string, + keepPortReservation = false, +): Promise { + const record = await processRecord(runtime, name); + if (record) { + await runtime.client().deleteSession(id, record.sessionId); + await deleteSessionEnvironment(runtime, id, record.sessionId); + } + if (!keepPortReservation) { + await releaseProcessPort(runtime, name); + } + await runtime.storage.delete(`${PROC_PREFIX}${name}`); +} + +async function terminateUntrackedSandboxProcesses( + runtime: ProcessControlRuntime, + id: string, +): Promise { + const result = await runtime + .client() + .execute(id, { + command: `python3 -c ${shellQuote(SANDBOX_PROCESS_TERMINATION_SCRIPT)}`, + cwd: WORKSPACE_DIR, + timeout: timeoutSeconds(15_000), + }) + .catch((error: unknown) => { + throw runtime.toUpstreamError(error, "Sandbox process termination failed."); + }); + if (result.exitCode !== 0) { + throw new APIError( + 502, + "upstream_sandbox_failed", + "Sandbox processes could not be terminated.", + { + details: { output: (result.result ?? "").slice(-1_000) }, + retriable: true, + }, + ); + } +} + +async function deleteProcessesOnPort( + runtime: ProcessControlRuntime, + control: ProcessControl, + id: string, + port: number, + exceptName: string, +): Promise { + const records = await runtime.storage.list({ prefix: PROC_PREFIX }); + for (const [key, value] of records) { + const parsed = ProcessRecordSchema.safeParse(value); + const name = key.slice(PROC_PREFIX.length); + if (parsed.success && name !== exceptName && parsed.data.port === port) { + await control.deleteProcessRecord(id, name); + } + } +} + +async function freeProjectPort( + runtime: ProcessControlRuntime, + workspaceSlug: string, +): Promise { + await runtime.storage.transaction(async (transaction) => { + const allocation = PortAllocationSchema.parse((await transaction.get(PORT_ALLOC_KEY)) ?? {}); + if (allocation.ports[workspaceSlug] === undefined) return; + const ports = Object.fromEntries( + Object.entries(allocation.ports).filter(([slug]) => slug !== workspaceSlug), + ); + await transaction.put(PORT_ALLOC_KEY, { ...allocation, ports }); + }); +} + +async function prepareProcessSlot( + runtime: ProcessControlRuntime, + control: ProcessControl, + id: string, + name: string, + input: ParsedProcessStartInput, +): Promise { + await control.deleteProcessRecord(id, name, true); + await ensureProcessCapacity(runtime, control, id); + if (input.waitForPort) { + await control.deleteProcessesOnPort(id, input.waitForPort.port, name); + } +} + +async function ensureProcessCapacity( + runtime: ProcessControlRuntime, + control: ProcessControl, + id: string, +): Promise { + let records = await runtime.storage.list({ prefix: PROC_PREFIX }); + if (records.size < MAX_TRACKED_PROCESSES) return; + await pruneCompletedProcessRecords(runtime, control, id, records); + records = await runtime.storage.list({ prefix: PROC_PREFIX }); + if (records.size >= MAX_TRACKED_PROCESSES) { + throw new APIError( + 429, + "sandbox_process_limit_reached", + "The sandbox has no available managed process slot.", + { + hint: "Stop an existing managed process or reuse its stable process ID.", + retriable: false, + }, + ); + } +} + +async function pruneCompletedProcessRecords( + runtime: ProcessControlRuntime, + control: ProcessControl, + id: string, + records: Map, +): Promise { + for (const [key, value] of records) { + const parsed = ProcessRecordSchema.safeParse(value); + if (!parsed.success) { + await control.deleteProcessRecord(id, key.slice(PROC_PREFIX.length)); + continue; + } + const session = await runtime.client().getSession(id, parsed.data.sessionId); + const command = session?.commands.find((candidate) => candidate.id === parsed.data.cmdId); + if (session === null || typeof command?.exitCode === "number") { + await control.deleteProcessRecord(id, key.slice(PROC_PREFIX.length)); + } + } +} + +async function persistStartedProcess( + runtime: ProcessControlRuntime, + control: ProcessControl, + id: string, + name: string, + record: ProcessRecord, + waitForPort: ParsedProcessStartInput["waitForPort"], +): Promise { + try { + await runtime.storage.put(`${PROC_PREFIX}${name}`, record); + if (waitForPort) { + await control.waitForPort(id, waitForPort.port, waitForPort.path, waitForPort.timeoutMs, { + cmdId: record.cmdId, + sessionId: record.sessionId, + }); + } + } catch (error) { + await control.cleanupLaunchedProcess(id, record.sessionId, name); + throw error; + } +} + +async function persistProcessOwnershipIntent( + runtime: ProcessControlRuntime, + name: string, + record: ProcessRecord, +): Promise { + try { + await runtime.storage.put(`${PROC_PREFIX}${name}`, record); + } catch (error) { + await releaseProcessPort(runtime, name); + throw error; + } +} + +async function buildSessionCommand( + runtime: ProcessControlRuntime, + id: string, + sessionId: string, + cwd: string, + rawCommand: string, + env: Record | undefined, +): Promise { + if (!env || Object.keys(env).length === 0) { + return `cd ${shellQuote(cwd)} && ${rawCommand}`; + } + const envPath = `${ENV_FILE_DIR}/${sessionId}.env`; + const body = Object.entries(env) + .map(([key, value]) => `export ${key}=${shellQuote(value)}`) + .join("\n"); + await runtime.client().createFolder(id, ENV_FILE_DIR, "700"); + await runtime.client().uploadFile(id, envPath, new TextEncoder().encode(`${body}\n`)); + const permissions = await runtime.client().execute(id, { + command: `chmod 600 ${shellQuote(envPath)}`, + timeout: 10, + }); + if (permissions.exitCode !== 0) { + throw new Error("Could not secure the transient process environment."); + } + return processShellCommand(envPath, cwd, rawCommand); +} + +function processShellCommand(envPath: string, cwd: string, rawCommand: string): string { + return `set -a +. ${shellQuote(envPath)} +env_status=$? +rm -f ${shellQuote(envPath)} +[ "$env_status" -eq 0 ] || exit "$env_status" +set +a +cd ${shellQuote(cwd)} && ${rawCommand}`; +} + +async function launchSessionProcess( + runtime: ProcessControlRuntime, + id: string, + sessionId: string, + name: string, + cwd: string, + rawCommand: string, + env: Record | undefined, + stdin?: string, +): Promise { + try { + const command = await buildSessionCommand(runtime, id, sessionId, cwd, rawCommand, env); + await runtime.client().createSession(id, sessionId); + const execution = await runtime.client().execSessionCommand(id, sessionId, command, true); + if (stdin !== undefined) { + await sendBootstrapInput(runtime, id, sessionId, execution, stdin); + } + return execution; + } catch (error) { + await runtime.client().deleteSession(id, sessionId); + await deleteSessionEnvironment(runtime, id, sessionId); + await releaseProcessPort(runtime, name); + throw error; + } +} + +async function sendBootstrapInput( + runtime: ProcessControlRuntime, + id: string, + sessionId: string, + execution: DaytonaSessionExecResponse, + stdin: string, +): Promise { + if (!execution.cmdId) { + throw new Error("Sandbox process did not return a command ID for bootstrap input."); + } + await runtime.client().sendSessionCommandInput(id, sessionId, execution.cmdId, stdin); +} + +async function cleanupLaunchedProcess( + runtime: ProcessControlRuntime, + id: string, + sessionId: string, + name: string, +): Promise { + await runtime.client().deleteSession(id, sessionId); + await deleteSessionEnvironment(runtime, id, sessionId); + await releaseProcessPort(runtime, name); + await runtime.storage.delete(`${PROC_PREFIX}${name}`); +} + +async function deleteSessionEnvironment( + runtime: ProcessControlRuntime, + id: string, + sessionId: string, +): Promise { + await runtime.client().deleteFilePath(id, `${ENV_FILE_DIR}/${sessionId}.env`, false); +} + +async function throwIfProcessExited( + runtime: ProcessControlRuntime, + id: string, + port: number, + process: { cmdId: string; sessionId: string }, +): Promise { + const session = await runtime + .client() + .getSession(id, process.sessionId) + .catch(() => null); + const command = session?.commands.find((candidate) => candidate.id === process.cmdId); + if (typeof command?.exitCode !== "number") return; + const logs = await runtime + .client() + .getSessionCommandLogs(id, process.sessionId, process.cmdId) + .catch(() => ""); + throw new APIError(502, "sandbox_command_failed", "Sandbox process exited before readiness.", { + details: { exitCode: command.exitCode, logs: logs.slice(-2_000), port }, + hint: "Inspect the process logs, fix the start command, and retry.", + retriable: false, + }); +} + +async function releaseProcessPort( + runtime: ProcessControlRuntime, + processId: string, +): Promise { + await runtime.storage.transaction(async (transaction) => { + const reservations = ProcessPortReservationsSchema.parse( + (await transaction.get(PROCESS_PORT_ALLOC_KEY)) ?? {}, + ); + if (reservations[processId] === undefined) return; + await transaction.put( + PROCESS_PORT_ALLOC_KEY, + withoutProcessReservation(reservations, processId), + ); + }); +} diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-process-support.ts b/apps/agent-worker/src/durable-objects/project-sandbox-process-support.ts index 177b9e03..4435d277 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-process-support.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-process-support.ts @@ -32,7 +32,6 @@ export const ProcessRecordSchema = z }) .strict(); export type ProcessRecord = z.infer; -export type NamedProcessRecord = { name: string; record: ProcessRecord }; export type ParsedProcessStartInput = z.infer; export type ProcessPolicy = Pick< ProcessRecord, diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-processes.ts b/apps/agent-worker/src/durable-objects/project-sandbox-processes.ts index 323d7d3f..9931616c 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-processes.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-processes.ts @@ -5,23 +5,16 @@ import type { SandboxProcessResult, SandboxRunCodeResult, } from "@cheatcode/sandbox-contracts"; -import type { DaytonaSessionExecResponse } from "@cheatcode/tools-code"; import type { SandboxConsoleSnapshot } from "@cheatcode/types"; import { sandboxExecProcessName } from "./project-sandbox-audit"; import { WORKSPACE_DIR } from "./project-sandbox-content-support"; -import { ProjectSandboxLifecycle } from "./project-sandbox-lifecycle"; import { recordSandboxUsageBestEffort } from "./project-sandbox-metering"; -import { - SANDBOX_PROCESS_TERMINATION_SCRIPT, - WORKSPACE_PROCESS_TERMINATION_SCRIPT, -} from "./project-sandbox-process-cleanup"; +import { createProcessControl, type ProcessControl } from "./project-sandbox-process-control"; import { emptyConsoleSnapshot, sliceProcessLogs } from "./project-sandbox-process-logs"; import { assertValidProcessStart, ENV_FILE_DIR, firstAvailablePort, - MAX_TRACKED_PROCESSES, - type NamedProcessRecord, type ParsedProcessStartInput, PORT_ALLOC_KEY, PortAllocationSchema, @@ -34,13 +27,9 @@ import { ProcessRecordSchema, processRecordFromLaunch, pruneExpiredProcessPortReservations, - restartEnvironment, - shellQuote, - sleep, supervisedProcessCommand, timeoutSeconds, usedProcessPorts, - withoutProcessReservation, } from "./project-sandbox-process-support"; import { commandToShellString, @@ -59,689 +48,409 @@ import { type ProjectStartProcessInput, ProjectStartProcessInputSchema, } from "./project-sandbox-runtime"; +import type { SandboxRuntime } from "./project-sandbox-runtime-handle"; const DEFAULT_EXEC_TIMEOUT_MS = 60_000; -interface ProjectSandboxStatus { +export interface ProjectSandboxStatus { healthy: boolean; ping: string; sandboxId: string; } -export abstract class ProjectSandboxProcesses extends ProjectSandboxLifecycle { - private readonly processMutations = new ProcessMutationQueue(); - - public async ensureReady(): Promise { - const result = await this.runCode({ code: "print('ready')", language: "python" }); - return { - healthy: result.success, - ping: result.stdout.trim(), - sandboxId: this.sandboxName(), - }; - } +export type ProcessOps = ProcessControl & { + allocateProcessPort: (input: ProjectAllocateProcessPortInput) => Promise; + allocateProjectPort: (input: ProjectAllocatePortInput) => Promise; + ensureReady: () => Promise; + exec: (input: ProjectExecInput) => Promise; + getStatus: () => Promise; + killAllProcesses: () => Promise; + killProcess: (input: ProjectKillProcessInput) => Promise; + readDevServerLogs: (input: ProjectReadDevServerLogsInput) => Promise; + runCode: (input: ProjectRunCodeInput) => Promise; + startProcess: (input: ProjectStartProcessInput) => Promise; +}; + +export interface CoordinatedProcessOps { + allocateProcessPort: (input: ProjectAllocateProcessPortInput) => Promise; + ensureReady: () => Promise; + exec: (input: ProjectExecInput) => Promise; + killProcess: (input: ProjectKillProcessInput) => Promise; + runCode: (input: ProjectRunCodeInput) => Promise; + startProcess: (input: ProjectStartProcessInput) => Promise; +} - public async getStatus(): Promise { - return this.ensureReady(); - } +type ProcessRuntime = Pick< + SandboxRuntime, + | "client" + | "ensureSandbox" + | "existingSandboxId" + | "meteringContext" + | "sandboxName" + | "storage" + | "toUpstreamError" + | "writeExecAudit" +>; + +interface ProcessContext { + control: ProcessControl; + coordinated: CoordinatedProcessOps; + mutations: ProcessMutationQueue; + runtime: ProcessRuntime; +} - public async runCode(input: ProjectRunCodeInput): Promise { - const parsed = ProjectRunCodeInputSchema.parse(input); - const command = - parsed.language === "python" - ? ["python3", "-c", parsed.code] - : ["node", "--input-type=module", "-e", parsed.code]; - const result = await this.exec({ - command, - cwd: parsed.cwd ?? WORKSPACE_DIR, - env: parsed.env, - timeoutMs: parsed.timeoutMs ?? DEFAULT_EXEC_TIMEOUT_MS, - }); - return { - exitCode: result.exitCode, - stderr: result.stderr, - stdout: result.stdout, - success: result.success, - }; - } +export function createProcessOps( + runtime: ProcessRuntime, + coordinated: CoordinatedProcessOps, +): ProcessOps { + const control = createProcessControl(runtime); + const context: ProcessContext = { + control, + coordinated, + mutations: new ProcessMutationQueue(), + runtime, + }; + return { + ...control, + allocateProcessPort: (input) => allocateProcessPort(runtime, input), + allocateProjectPort: (input) => allocateProjectPort(runtime, input), + ensureReady: () => ensureReady(context), + exec: (input) => exec(runtime, input), + getStatus: () => coordinated.ensureReady(), + killAllProcesses: () => context.mutations.run(() => killAllProcesses(context)), + killProcess: (input) => killProcess(context, input), + readDevServerLogs: (input) => readDevServerLogs(context, input), + runCode: (input) => runCode(context, input), + startProcess: (input) => context.mutations.run(() => startProcess(context, input)), + }; +} - public async exec(input: ProjectExecInput): Promise { - const parsed = ProjectExecInputSchema.parse(input); - const id = await this.ensureSandbox(); - const startedAt = Date.now(); - const command = commandToShellString(parsed.command); - const cwd = parsed.cwd ?? WORKSPACE_DIR; - const timeoutMs = parsed.timeoutMs ?? DEFAULT_EXEC_TIMEOUT_MS; - try { - const completed = await this.client().execute(id, { - command, - cwd, - timeout: timeoutSeconds(timeoutMs), - ...(parsed.env === undefined ? {} : { env: parsed.env }), - }); - const result: SandboxExecResult = { - command, - durationMs: Date.now() - startedAt, - exitCode: completed.exitCode, - stderr: "", - stdout: completed.result ?? "", - success: completed.exitCode === 0, - }; - const processName = sandboxExecProcessName(parsed.command[0] ?? "process"); - await this.writeExecAudit({ - argc: parsed.command.length, - argv0: processName, - cwd, - durationMs: result.durationMs ?? 0, - exitCode: completed.exitCode, - processName, - sandboxId: this.sandboxName(), - status: result.success ? "completed" : "failed", - success: result.success, - timestamp: new Date(startedAt).toISOString(), - type: "sandbox_exec", - }).catch((error: unknown) => { - createLogger().warn("sandbox_exec_audit_failed", { - error, - processName, - sandboxId: this.sandboxName(), - }); - }); - await recordSandboxUsageBestEffort(await this.meteringContext()); - return result; - } catch (error) { - throw this.toUpstreamError(error, "Sandbox command failed."); - } - } +async function ensureReady(context: ProcessContext): Promise { + const result = await context.coordinated.runCode({ + code: "print('ready')", + language: "python", + }); + return { + healthy: result.success, + ping: result.stdout.trim(), + sandboxId: context.runtime.sandboxName(), + }; +} - public async startProcess(input: ProjectStartProcessInput): Promise { - return this.processMutations.run(() => this.startProcessExclusive(input)); - } +async function runCode( + context: ProcessContext, + input: ProjectRunCodeInput, +): Promise { + const parsed = ProjectRunCodeInputSchema.parse(input); + const command = + parsed.language === "python" + ? ["python3", "-c", parsed.code] + : ["node", "--input-type=module", "-e", parsed.code]; + const result = await context.coordinated.exec({ + command, + cwd: parsed.cwd ?? WORKSPACE_DIR, + env: parsed.env, + timeoutMs: parsed.timeoutMs ?? DEFAULT_EXEC_TIMEOUT_MS, + }); + return { + exitCode: result.exitCode, + stderr: result.stderr, + stdout: result.stdout, + success: result.success, + }; +} - private async startProcessExclusive( - input: ProjectStartProcessInput, - ): Promise { - const parsed = ProjectStartProcessInputSchema.parse(input); - assertValidProcessStart(parsed); - const id = await this.ensureSandbox(); - const name = parsed.processId; - const sessionId = `cc-${name}`; - await this.prepareProcessSlot(id, name, parsed); - const cwd = parsed.cwd ?? WORKSPACE_DIR; - const rawCommand = commandToShellString(parsed.command); - const policy: ProcessPolicy = { - keepAliveTimeoutMs: parsed.keepAliveTimeoutMs ?? 0, - maxRestarts: parsed.maxRestarts ?? 0, - restartOnFailure: parsed.restartOnFailure ?? false, - }; - const provisionalRecord = processRecordFromLaunch(parsed, policy, { - cmdId: sessionId, - command: rawCommand, +async function exec(runtime: ProcessRuntime, input: ProjectExecInput): Promise { + const parsed = ProjectExecInputSchema.parse(input); + const id = await runtime.ensureSandbox(); + const startedAt = Date.now(); + const command = commandToShellString(parsed.command); + const cwd = parsed.cwd ?? WORKSPACE_DIR; + const timeoutMs = parsed.timeoutMs ?? DEFAULT_EXEC_TIMEOUT_MS; + try { + const completed = await runtime.client().execute(id, { + command, cwd, - sessionId, + timeout: timeoutSeconds(timeoutMs), + ...(parsed.env === undefined ? {} : { env: parsed.env }), }); - await this.persistProcessOwnershipIntent(name, provisionalRecord); - let exec: DaytonaSessionExecResponse; - try { - exec = await this.launchSessionProcess( - id, - sessionId, - name, - cwd, - supervisedProcessCommand(rawCommand, policy), - parsed.env, - parsed.stdin, - ); - } catch (error) { - await this.cleanupLaunchedProcess(id, sessionId, name); - throw error; - } - const record = { ...provisionalRecord, cmdId: exec.cmdId ?? sessionId }; - await this.persistStartedProcess(id, name, record, parsed.waitForPort); - await recordSandboxUsageBestEffort(await this.meteringContext()); - return { command: record.command, id: name, status: "running" }; + const result = execResult(command, completed, startedAt); + await recordExecAudit(runtime, parsed.command, cwd, result, completed.exitCode, startedAt); + await recordSandboxUsageBestEffort(await runtime.meteringContext()); + return result; + } catch (error) { + throw runtime.toUpstreamError(error, "Sandbox command failed."); } +} - public async allocateProjectPort(input: ProjectAllocatePortInput): Promise { - const parsed = ProjectAllocatePortInputSchema.parse(input); - return this.ctx.storage.transaction(async (transaction) => { - const allocation = PortAllocationSchema.parse((await transaction.get(PORT_ALLOC_KEY)) ?? {}); - const existing = allocation.ports[parsed.projectId]; - if (existing !== undefined) { - return existing; - } - const used = new Set(Object.values(allocation.ports)); - let candidate = parsed.stack === "mobile" ? allocation.mobileNext : allocation.webNext; - while (used.has(candidate)) { - candidate += 1; - } - allocation.ports[parsed.projectId] = candidate; - if (parsed.stack === "mobile") { - allocation.mobileNext = candidate + 1; - } else { - allocation.webNext = candidate + 1; - } - await transaction.put(PORT_ALLOC_KEY, allocation); - return candidate; - }); - } +function execResult( + command: string, + completed: { exitCode: number; result?: string | null | undefined }, + startedAt: number, +): SandboxExecResult { + return { + command, + durationMs: Date.now() - startedAt, + exitCode: completed.exitCode, + stderr: "", + stdout: completed.result ?? "", + success: completed.exitCode === 0, + }; +} - public async allocateProcessPort(input: ProjectAllocateProcessPortInput): Promise { - const parsed = ProjectAllocateProcessPortInputSchema.parse(input); - return this.ctx.storage.transaction(async (transaction) => { - const now = Date.now(); - let reservations = ProcessPortReservationsSchema.parse( - (await transaction.get(PROCESS_PORT_ALLOC_KEY)) ?? {}, - ); - const records = await transaction.list({ prefix: PROC_PREFIX }); - reservations = pruneExpiredProcessPortReservations(reservations, records, now); - const used = usedProcessPorts(reservations, records, parsed.processId); - const existing = reservations[parsed.processId]; - if ( - existing && - existing.port >= parsed.minPort && - existing.port <= parsed.maxPort && - !used.has(existing.port) - ) { - reservations[parsed.processId] = { ...existing, reservedAtMs: now }; - await transaction.put(PROCESS_PORT_ALLOC_KEY, reservations); - return existing.port; - } - const port = firstAvailablePort(used, parsed.minPort, parsed.maxPort); - if (port === null) { - throw new APIError( - 503, - "sandbox_failed_to_start", - "No sandbox process port is available.", - { - retriable: true, - }, - ); - } - reservations[parsed.processId] = { port, reservedAtMs: now }; - await transaction.put(PROCESS_PORT_ALLOC_KEY, reservations); - return port; +async function recordExecAudit( + runtime: ProcessRuntime, + argv: string[], + cwd: string, + result: SandboxExecResult, + exitCode: number, + startedAt: number, +): Promise { + const processName = sandboxExecProcessName(argv[0] ?? "process"); + await runtime + .writeExecAudit({ + argc: argv.length, + argv0: processName, + cwd, + durationMs: result.durationMs ?? 0, + exitCode, + processName, + sandboxId: runtime.sandboxName(), + status: result.success ? "completed" : "failed", + success: result.success, + timestamp: new Date(startedAt).toISOString(), + type: "sandbox_exec", + }) + .catch((error: unknown) => { + createLogger().warn("sandbox_exec_audit_failed", { + error, + processName, + sandboxId: runtime.sandboxName(), + }); }); - } - - public async killAllProcesses(): Promise { - return this.processMutations.run(() => this.killAllProcessesExclusive()); - } - - private async killAllProcessesExclusive(): Promise { - const id = await this.existingSandboxId(); - const records = await this.ctx.storage.list({ prefix: PROC_PREFIX }); - let killed = 0; - for (const [key, value] of records) { - const name = key.slice(PROC_PREFIX.length); - if (id && ProcessRecordSchema.safeParse(value).success) { - await this.deleteProcessRecord(id, name); - killed += 1; - } else { - await this.ctx.storage.delete(key); - } - } - if (id) { - await this.client().deleteFilePath(id, ENV_FILE_DIR, true); - } - await this.ctx.storage.delete(PROCESS_PORT_ALLOC_KEY); - return killed; - } - - public async killProcess(input: ProjectKillProcessInput): Promise { - const parsed = ProjectKillProcessInputSchema.parse(input); - return this.processMutations.run(() => this.killProcessExclusive(parsed.processId)); - } - - private async killProcessExclusive(processId: string): Promise { - const record = await this.processRecord(processId); - const id = record ? await this.existingSandboxId() : null; - if (record && id) { - await this.deleteProcessRecord(id, processId); - } else { - await this.ctx.storage.delete(`${PROC_PREFIX}${processId}`); - await this.releaseProcessPort(processId); - } - return { processId, status: "killed", success: true }; - } +} - public async readDevServerLogs( - input: ProjectReadDevServerLogsInput, - ): Promise { - const parsed = ProjectReadDevServerLogsInputSchema.parse(input); - const id = await this.existingSandboxId(); - const processes = await this.processRecordsForRead(parsed.processId); - if (id === null || processes.length === 0) { - return emptyConsoleSnapshot({ stderr: parsed.stderrCursor, stdout: parsed.stdoutCursor }); - } - for (const process of processes) { - const snapshot = await this.readProcessLogs(id, process, parsed).catch(async (error) => { - if (isMissingDaytonaProcessError(error)) { - await this.ctx.storage.delete(`${PROC_PREFIX}${process.name}`); - await this.releaseProcessPort(process.name); - return null; - } - throw this.toUpstreamError(error, "Sandbox console read failed."); - }); - if (snapshot) { - return snapshot; - } - } - return emptyConsoleSnapshot({ stderr: parsed.stderrCursor, stdout: parsed.stdoutCursor }); - } +async function startProcess( + context: ProcessContext, + input: ProjectStartProcessInput, +): Promise { + const parsed = ProjectStartProcessInputSchema.parse(input); + assertValidProcessStart(parsed); + const id = await context.runtime.ensureSandbox(); + const name = parsed.processId; + const sessionId = `cc-${name}`; + await context.control.prepareProcessSlot(id, name, parsed); + const cwd = parsed.cwd ?? WORKSPACE_DIR; + const rawCommand = commandToShellString(parsed.command); + const policy = processPolicy(parsed); + const provisional = processRecordFromLaunch(parsed, policy, { + cmdId: sessionId, + command: rawCommand, + cwd, + sessionId, + }); + await context.control.persistProcessOwnershipIntent(name, provisional); + const record = await launchProcess(context, id, name, sessionId, parsed, provisional); + await context.control.persistStartedProcess(id, name, record, parsed.waitForPort); + await recordSandboxUsageBestEffort(await context.runtime.meteringContext()); + return { command: record.command, id: name, status: "running" }; +} - protected async isPortAlive(id: string, port: number): Promise { - const probe = await this.client() - .execute(id, { - command: `curl -sf -o /dev/null --max-time 3 http://localhost:${port}/`, - timeout: 5, - }) - .catch(() => null); - return probe?.exitCode === 0; - } +function processPolicy(input: ParsedProcessStartInput): ProcessPolicy { + return { + keepAliveTimeoutMs: input.keepAliveTimeoutMs ?? 0, + maxRestarts: input.maxRestarts ?? 0, + restartOnFailure: input.restartOnFailure ?? false, + }; +} - protected async relaunchDevServer( - id: string, - name: string, - record: ProcessRecord, - restartEnv?: Record, - ): Promise { - const sessionId = record.sessionId || `cc-${name}`; - await this.client().deleteSession(id, sessionId); - const exec = await this.launchSessionProcess( +async function launchProcess( + context: ProcessContext, + id: string, + name: string, + sessionId: string, + input: ParsedProcessStartInput, + provisional: ProcessRecord, +): Promise { + try { + const execution = await context.control.launchSessionProcess( id, sessionId, name, - record.cwd, - supervisedProcessCommand(record.command, record), - restartEnv ?? restartEnvironment(name, record), + provisional.cwd, + supervisedProcessCommand(provisional.command, provisional), + input.env, + input.stdin, ); - try { - await this.ctx.storage.put(`${PROC_PREFIX}${name}`, { - ...record, - cmdId: exec.cmdId ?? sessionId, - startedAtMs: Date.now(), - } satisfies ProcessRecord); - } catch (error) { - await this.cleanupLaunchedProcess(id, sessionId, name); - throw error; - } - } - - protected async waitForPort( - id: string, - port: number, - path: string | undefined, - timeoutMs: number | undefined, - process?: { cmdId: string; sessionId: string }, - ): Promise { - const deadline = Date.now() + (timeoutMs ?? 120_000); - const url = `http://localhost:${port}${path ?? "/"}`; - while (Date.now() < deadline) { - const probe = await this.client() - .execute(id, { - command: `curl -sf -o /dev/null --max-time 3 ${shellQuote(url)}`, - timeout: 5, - }) - .catch(() => null); - if (probe?.exitCode === 0) { - return; - } - if (process) { - await this.throwIfProcessExited(id, port, process); - } - await sleep(1_500); - } - throw new APIError(504, "upstream_timeout_sandbox", "Sandbox process did not become ready.", { - details: { port, timeoutMs: timeoutMs ?? 120_000, url }, - hint: "Inspect the process command and logs, then retry.", - retriable: true, - }); - } - - protected async httpPortReady( - id: string, - port: number, - path: string, - timeoutMs: number, - ): Promise { - const deadline = Date.now() + timeoutMs; - const url = `http://localhost:${port}${path}`; - while (Date.now() < deadline) { - const probe = await this.client() - .execute(id, { - command: `curl -sf -o /dev/null --max-time 3 ${shellQuote(url)}`, - timeout: 5, - }) - .catch(() => null); - if (probe?.exitCode === 0) { - return true; - } - await sleep(1_000); - } - return false; - } - - protected async processRecord(name: string): Promise { - const value = await this.ctx.storage.get(`${PROC_PREFIX}${name}`); - const parsed = ProcessRecordSchema.safeParse(value); - return parsed.success ? parsed.data : null; - } - - protected async deleteProcessRecord( - id: string, - name: string, - keepPortReservation = false, - ): Promise { - const record = await this.processRecord(name); - if (record) { - await this.client().deleteSession(id, record.sessionId); - await this.deleteSessionEnvironment(id, record.sessionId); - } - if (!keepPortReservation) { - await this.releaseProcessPort(name); - } - await this.ctx.storage.delete(`${PROC_PREFIX}${name}`); - } - - protected async terminateUntrackedWorkspaceProcesses( - id: string, - workspaceSlug: string, - ): Promise { - const workspacePath = `${WORKSPACE_DIR}/${workspaceSlug}`; - const result = await this.client() - .execute(id, { - command: `python3 -c ${shellQuote(WORKSPACE_PROCESS_TERMINATION_SCRIPT)} ${shellQuote(workspacePath)}`, - cwd: WORKSPACE_DIR, - timeout: timeoutSeconds(15_000), - }) - .catch((error: unknown) => { - throw this.toUpstreamError(error, "Project workspace process termination failed."); - }); - if (result.exitCode !== 0) { - throw new APIError( - 502, - "upstream_sandbox_failed", - "Project workspace processes could not be terminated.", - { - details: { output: (result.result ?? "").slice(-1_000), workspaceSlug }, - retriable: true, - }, - ); - } - } - - protected async terminateUntrackedSandboxProcesses(id: string): Promise { - const result = await this.client() - .execute(id, { - command: `python3 -c ${shellQuote(SANDBOX_PROCESS_TERMINATION_SCRIPT)}`, - cwd: WORKSPACE_DIR, - timeout: timeoutSeconds(15_000), - }) - .catch((error: unknown) => { - throw this.toUpstreamError(error, "Sandbox process termination failed."); - }); - if (result.exitCode !== 0) { - throw new APIError( - 502, - "upstream_sandbox_failed", - "Sandbox processes could not be terminated.", - { - details: { output: (result.result ?? "").slice(-1_000) }, - retriable: true, - }, - ); - } - } - - protected async deleteProcessesOnPort( - id: string, - port: number, - exceptName: string, - ): Promise { - const records = await this.ctx.storage.list({ prefix: PROC_PREFIX }); - for (const [key, value] of records) { - const parsed = ProcessRecordSchema.safeParse(value); - const name = key.slice(PROC_PREFIX.length); - if (parsed.success && name !== exceptName && parsed.data.port === port) { - await this.deleteProcessRecord(id, name); - } - } - } - - protected async freeProjectPort(workspaceSlug: string): Promise { - await this.ctx.storage.transaction(async (transaction) => { - const allocation = PortAllocationSchema.parse((await transaction.get(PORT_ALLOC_KEY)) ?? {}); - if (allocation.ports[workspaceSlug] === undefined) { - return; - } - const ports = Object.fromEntries( - Object.entries(allocation.ports).filter(([slug]) => slug !== workspaceSlug), - ); - await transaction.put(PORT_ALLOC_KEY, { ...allocation, ports }); - }); + return { ...provisional, cmdId: execution.cmdId ?? sessionId }; + } catch (error) { + await context.control.cleanupLaunchedProcess(id, sessionId, name); + throw error; } +} - private async prepareProcessSlot( - id: string, - name: string, - input: ParsedProcessStartInput, - ): Promise { - await this.deleteProcessRecord(id, name, true); - await this.ensureProcessCapacity(id); - if (input.waitForPort) { - await this.deleteProcessesOnPort(id, input.waitForPort.port, name); +async function allocateProjectPort( + runtime: ProcessRuntime, + input: ProjectAllocatePortInput, +): Promise { + const parsed = ProjectAllocatePortInputSchema.parse(input); + return runtime.storage.transaction(async (transaction) => { + const allocation = PortAllocationSchema.parse((await transaction.get(PORT_ALLOC_KEY)) ?? {}); + const existing = allocation.ports[parsed.projectId]; + if (existing !== undefined) return existing; + const used = new Set(Object.values(allocation.ports)); + let candidate = parsed.stack === "mobile" ? allocation.mobileNext : allocation.webNext; + while (used.has(candidate)) { + candidate += 1; } - } + allocation.ports[parsed.projectId] = candidate; + if (parsed.stack === "mobile") allocation.mobileNext = candidate + 1; + else allocation.webNext = candidate + 1; + await transaction.put(PORT_ALLOC_KEY, allocation); + return candidate; + }); +} - private async ensureProcessCapacity(id: string): Promise { - let records = await this.ctx.storage.list({ prefix: PROC_PREFIX }); - if (records.size < MAX_TRACKED_PROCESSES) { - return; - } - await this.pruneCompletedProcessRecords(id, records); - records = await this.ctx.storage.list({ prefix: PROC_PREFIX }); - if (records.size >= MAX_TRACKED_PROCESSES) { - throw new APIError( - 429, - "sandbox_process_limit_reached", - "The sandbox has no available managed process slot.", - { - hint: "Stop an existing managed process or reuse its stable process ID.", - retriable: false, - }, - ); +async function allocateProcessPort( + runtime: ProcessRuntime, + input: ProjectAllocateProcessPortInput, +): Promise { + const parsed = ProjectAllocateProcessPortInputSchema.parse(input); + return runtime.storage.transaction(async (transaction) => { + const now = Date.now(); + let reservations = ProcessPortReservationsSchema.parse( + (await transaction.get(PROCESS_PORT_ALLOC_KEY)) ?? {}, + ); + const records = await transaction.list({ prefix: PROC_PREFIX }); + reservations = pruneExpiredProcessPortReservations(reservations, records, now); + const used = usedProcessPorts(reservations, records, parsed.processId); + const existing = reservations[parsed.processId]; + if (isReusableReservation(existing, used, parsed.minPort, parsed.maxPort)) { + reservations[parsed.processId] = { ...existing, reservedAtMs: now }; + await transaction.put(PROCESS_PORT_ALLOC_KEY, reservations); + return existing.port; } - } + const port = firstAvailablePort(used, parsed.minPort, parsed.maxPort); + if (port === null) throw noProcessPortAvailable(); + reservations[parsed.processId] = { port, reservedAtMs: now }; + await transaction.put(PROCESS_PORT_ALLOC_KEY, reservations); + return port; + }); +} - private async pruneCompletedProcessRecords( - id: string, - records: Map, - ): Promise { - for (const [key, value] of records) { - const parsed = ProcessRecordSchema.safeParse(value); - if (!parsed.success) { - await this.deleteProcessRecord(id, key.slice(PROC_PREFIX.length)); - continue; - } - const session = await this.client().getSession(id, parsed.data.sessionId); - const command = session?.commands.find((candidate) => candidate.id === parsed.data.cmdId); - if (session === null || typeof command?.exitCode === "number") { - await this.deleteProcessRecord(id, key.slice(PROC_PREFIX.length)); - } - } - } +function isReusableReservation( + reservation: { port: number; reservedAtMs: number } | undefined, + used: Set, + minPort: number, + maxPort: number, +): reservation is { port: number; reservedAtMs: number } { + return ( + reservation !== undefined && + reservation.port >= minPort && + reservation.port <= maxPort && + !used.has(reservation.port) + ); +} - private async persistStartedProcess( - id: string, - name: string, - record: ProcessRecord, - waitForPort: ParsedProcessStartInput["waitForPort"], - ): Promise { - try { - await this.ctx.storage.put(`${PROC_PREFIX}${name}`, record); - if (waitForPort) { - await this.waitForPort(id, waitForPort.port, waitForPort.path, waitForPort.timeoutMs, { - cmdId: record.cmdId, - sessionId: record.sessionId, - }); - } - } catch (error) { - await this.cleanupLaunchedProcess(id, record.sessionId, name); - throw error; - } - } +function noProcessPortAvailable(): APIError { + return new APIError(503, "sandbox_failed_to_start", "No sandbox process port is available.", { + retriable: true, + }); +} - private async persistProcessOwnershipIntent(name: string, record: ProcessRecord): Promise { - try { - await this.ctx.storage.put(`${PROC_PREFIX}${name}`, record); - } catch (error) { - await this.releaseProcessPort(name); - throw error; +async function killAllProcesses(context: ProcessContext): Promise { + const id = await context.runtime.existingSandboxId(); + const records = await context.runtime.storage.list({ prefix: PROC_PREFIX }); + let killed = 0; + for (const [key, value] of records) { + const name = key.slice(PROC_PREFIX.length); + if (id && ProcessRecordSchema.safeParse(value).success) { + await context.control.deleteProcessRecord(id, name); + killed += 1; + } else { + await context.runtime.storage.delete(key); } } - - private async buildSessionCommand( - id: string, - sessionId: string, - cwd: string, - rawCommand: string, - env: Record | undefined, - ): Promise { - if (!env || Object.keys(env).length === 0) { - return `cd ${shellQuote(cwd)} && ${rawCommand}`; - } - const envPath = `${ENV_FILE_DIR}/${sessionId}.env`; - const body = Object.entries(env) - .map(([key, value]) => `export ${key}=${shellQuote(value)}`) - .join("\n"); - await this.client().createFolder(id, ENV_FILE_DIR, "700"); - await this.client().uploadFile(id, envPath, new TextEncoder().encode(`${body}\n`)); - const permissions = await this.client().execute(id, { - command: `chmod 600 ${shellQuote(envPath)}`, - timeout: 10, - }); - if (permissions.exitCode !== 0) { - throw new Error("Could not secure the transient process environment."); - } - return `set -a -. ${shellQuote(envPath)} -env_status=$? -rm -f ${shellQuote(envPath)} -[ "$env_status" -eq 0 ] || exit "$env_status" -set +a -cd ${shellQuote(cwd)} && ${rawCommand}`; + if (id) { + await context.runtime.client().deleteFilePath(id, ENV_FILE_DIR, true); } + await context.runtime.storage.delete(PROCESS_PORT_ALLOC_KEY); + return killed; +} - private async launchSessionProcess( - id: string, - sessionId: string, - name: string, - cwd: string, - rawCommand: string, - env: Record | undefined, - stdin?: string, - ): Promise { - try { - const command = await this.buildSessionCommand(id, sessionId, cwd, rawCommand, env); - await this.client().createSession(id, sessionId); - const execution = await this.client().execSessionCommand(id, sessionId, command, true); - if (stdin !== undefined) { - if (!execution.cmdId) { - throw new Error("Sandbox process did not return a command ID for bootstrap input."); - } - await this.client().sendSessionCommandInput(id, sessionId, execution.cmdId, stdin); - } - return execution; - } catch (error) { - await this.client().deleteSession(id, sessionId); - await this.deleteSessionEnvironment(id, sessionId); - await this.releaseProcessPort(name); - throw error; - } - } +async function killProcess( + context: ProcessContext, + input: ProjectKillProcessInput, +): Promise { + const parsed = ProjectKillProcessInputSchema.parse(input); + return context.mutations.run(() => killProcessExclusive(context, parsed.processId)); +} - private async cleanupLaunchedProcess(id: string, sessionId: string, name: string): Promise { - await this.client().deleteSession(id, sessionId); - await this.deleteSessionEnvironment(id, sessionId); - await this.releaseProcessPort(name); - await this.ctx.storage.delete(`${PROC_PREFIX}${name}`); - } +async function killProcessExclusive( + context: ProcessContext, + processId: string, +): Promise { + const record = await context.control.processRecord(processId); + const id = record ? await context.runtime.existingSandboxId() : null; + if (record && id) { + await context.control.deleteProcessRecord(id, processId); + } else { + await context.runtime.storage.delete(`${PROC_PREFIX}${processId}`); + await context.control.releaseProcessPort(processId); + } + return { processId, status: "killed", success: true }; +} - private async deleteSessionEnvironment(id: string, sessionId: string): Promise { - await this.client().deleteFilePath(id, `${ENV_FILE_DIR}/${sessionId}.env`, false); +async function readDevServerLogs( + context: ProcessContext, + input: ProjectReadDevServerLogsInput, +): Promise { + const parsed = ProjectReadDevServerLogsInputSchema.parse(input); + const id = await context.runtime.existingSandboxId(); + const record = await context.control.processRecord(parsed.processId); + if (id === null || !record) { + return emptyConsoleSnapshot({ stderr: parsed.stderrCursor, stdout: parsed.stdoutCursor }); } - - private async throwIfProcessExited( - id: string, - port: number, - process: { cmdId: string; sessionId: string }, - ): Promise { - const session = await this.client() - .getSession(id, process.sessionId) - .catch(() => null); - const command = session?.commands.find((candidate) => candidate.id === process.cmdId); - if (typeof command?.exitCode !== "number") { - return; + try { + return await readProcessLogs(context, id, parsed.processId, record, parsed); + } catch (error) { + if (isMissingDaytonaProcessError(error)) { + await context.runtime.storage.delete(`${PROC_PREFIX}${parsed.processId}`); + await context.control.releaseProcessPort(parsed.processId); + return emptyConsoleSnapshot({ stderr: parsed.stderrCursor, stdout: parsed.stdoutCursor }); } - const logs = await this.client() - .getSessionCommandLogs(id, process.sessionId, process.cmdId) - .catch(() => ""); - throw new APIError(502, "sandbox_command_failed", "Sandbox process exited before readiness.", { - details: { exitCode: command.exitCode, logs: logs.slice(-2_000), port }, - hint: "Inspect the process logs, fix the start command, and retry.", - retriable: false, - }); - } - - private async readProcessLogs( - id: string, - process: NamedProcessRecord, - input: ReturnType, - ): Promise { - const buffer = await this.client().getSessionCommandLogs( - id, - process.record.sessionId, - process.record.cmdId, - ); - const sliced = sliceProcessLogs({ - lastPid: input.lastPid, - pid: process.record.cmdId, - stderrCursor: input.stderrCursor, - stderrText: "", - stdoutCursor: input.stdoutCursor, - stdoutText: buffer, - tail: input.tail, - }); - return { - ...sliced, - process: { - command: process.record.command, - id: process.name, - pid: process.record.cmdId, - status: "running", - }, - }; - } - - private async processRecordsForRead(name: string): Promise { - const exact = await this.processRecord(name); - return exact ? [{ name, record: exact }] : []; + throw context.runtime.toUpstreamError(error, "Sandbox console read failed."); } +} - private async releaseProcessPort(processId: string): Promise { - await this.ctx.storage.transaction(async (transaction) => { - const reservations = ProcessPortReservationsSchema.parse( - (await transaction.get(PROCESS_PORT_ALLOC_KEY)) ?? {}, - ); - if (reservations[processId] === undefined) { - return; - } - await transaction.put( - PROCESS_PORT_ALLOC_KEY, - withoutProcessReservation(reservations, processId), - ); - }); - } +async function readProcessLogs( + context: ProcessContext, + id: string, + name: string, + record: ProcessRecord, + input: ReturnType, +): Promise { + const buffer = await context.runtime + .client() + .getSessionCommandLogs(id, record.sessionId, record.cmdId); + const sliced = sliceProcessLogs({ + lastPid: input.lastPid, + pid: record.cmdId, + stderrCursor: input.stderrCursor, + stderrText: "", + stdoutCursor: input.stdoutCursor, + stdoutText: buffer, + tail: input.tail, + }); + return { + ...sliced, + process: { + command: record.command, + id: name, + pid: record.cmdId, + status: "running", + }, + }; } function isMissingDaytonaProcessError(error: unknown): boolean { diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-project-files.ts b/apps/agent-worker/src/durable-objects/project-sandbox-project-files.ts index 117039a5..843ad76b 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox-project-files.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox-project-files.ts @@ -12,7 +12,6 @@ import { } from "@cheatcode/types"; import { z } from "zod"; import { sleep } from "./project-sandbox-process-support"; -import { ProjectSandboxProcesses } from "./project-sandbox-processes"; import { type ProjectListUploadedFilesInput, ProjectListUploadedFilesInputSchema, @@ -21,6 +20,7 @@ import { type ProjectUploadFileInput, ProjectUploadFileInputSchema, } from "./project-sandbox-runtime"; +import type { SandboxRuntime } from "./project-sandbox-runtime-handle"; const FILE_DIGEST_DOMAIN = "cheatcode:project-file:v2"; const VERSION_DIGEST_DOMAIN = "cheatcode:project-file-version:v2"; @@ -73,427 +73,480 @@ interface CurrentProjectFile { version: ProjectFileVersion; } -export abstract class ProjectSandboxProjectFiles extends ProjectSandboxProcesses { - private projectFileMutationTail: Promise = Promise.resolve(); +export interface FileOps { + deleteUploadedFileMetadata: (projectId: string) => Promise; + listUploadedFiles: (input: ProjectListUploadedFilesInput) => Promise<{ files: ProjectFile[] }>; + restoreUploadedFiles: ( + input: ProjectRestoreUploadedFilesInput, + ) => Promise<{ restoredFileCount: number }>; + uploadProjectFile: (input: ProjectUploadFileInput) => Promise; +} - public listUploadedFiles( - input: ProjectListUploadedFilesInput, - ): Promise<{ files: ProjectFile[] }> { - const parsed = ProjectListUploadedFilesInputSchema.parse(input); - return this.listProjectFileRecords(parsed.projectId); - } +type FileRuntime = Pick< + SandboxRuntime, + | "client" + | "ensureSandbox" + | "outputBucket" + | "ownerUserId" + | "restartSandboxForWorkspaceRecovery" + | "storage" +>; + +interface FileContext { + mutationTail: Promise; + runtime: FileRuntime; +} - public uploadProjectFile(input: ProjectUploadFileInput): Promise { - const parsed = ProjectUploadFileInputSchema.parse(input); - return this.enqueueProjectFileMutation(() => this.persistProjectFile(parsed)); - } +export function createFileOps(runtime: FileRuntime): FileOps { + const context: FileContext = { mutationTail: Promise.resolve(), runtime }; + return { + deleteUploadedFileMetadata: (projectId) => deleteUploadedFileMetadata(runtime, projectId), + listUploadedFiles: (input) => + listProjectFileRecords(runtime, ProjectListUploadedFilesInputSchema.parse(input).projectId), + restoreUploadedFiles: (input) => { + const parsed = ProjectRestoreUploadedFilesInputSchema.parse(input); + return enqueueProjectFileMutation(context, () => restoreProjectFiles(context, parsed)); + }, + uploadProjectFile: (input) => { + const parsed = ProjectUploadFileInputSchema.parse(input); + return enqueueProjectFileMutation(context, () => persistProjectFile(context, parsed)); + }, + }; +} - public restoreUploadedFiles( - input: ProjectRestoreUploadedFilesInput, - ): Promise<{ restoredFileCount: number }> { - const parsed = ProjectRestoreUploadedFilesInputSchema.parse(input); - return this.enqueueProjectFileMutation(() => this.restoreProjectFiles(parsed)); - } +function deleteUploadedFileMetadata(runtime: FileRuntime, projectId: string): Promise { + return Promise.all([ + deleteStoragePrefix(runtime, fileRecordProjectPrefix(projectId)), + deleteStoragePrefix(runtime, materializationRecordProjectPrefix(projectId)), + deleteStoragePrefix(runtime, versionRecordProjectPrefix(projectId)), + ]).then(() => undefined); +} - protected deleteUploadedFileMetadata(projectId: string): Promise { - return Promise.all([ - this.deleteStoragePrefix(fileRecordProjectPrefix(projectId)), - this.deleteStoragePrefix(materializationRecordProjectPrefix(projectId)), - this.deleteStoragePrefix(versionRecordProjectPrefix(projectId)), - ]).then(() => undefined); +async function persistProjectFile( + context: FileContext, + input: z.output, +): Promise { + const existing = await currentFile(context.runtime, input.projectId, input.path); + await enforceFileCount(context.runtime, input.projectId, existing !== null); + const prepared = await prepareProjectFile(input, context.runtime.ownerUserId()); + const persistedAt = new Date().toISOString(); + const version = projectFileVersion(input, prepared, persistedAt); + const previousVersion = await storedVersion(context.runtime, version); + const status = fileUploadStatus(existing, prepared.versionId); + await writeAndVerifyObject(context.runtime, input, version); + try { + const materialization = await materializeProjectFile(context, input, prepared); + const file = await commitProjectFile( + context.runtime, + input, + prepared, + existing, + previousVersion, + version, + persistedAt, + materialization, + ); + return ProjectFileUploadResponseSchema.parse({ file, status }); + } catch (error) { + if (!previousVersion) { + await context.runtime.outputBucket.delete(prepared.r2Key).catch(() => undefined); + } + throw error; } +} - private async persistProjectFile( - input: z.output, - ): Promise { - const existing = await this.currentFile(input.projectId, input.path); - await this.enforceFileCount(input.projectId, existing !== null); - const prepared = await prepareProjectFile(input, this.ownerUserId()); - const persistedAt = new Date().toISOString(); - const version = projectFileVersion(input, prepared, persistedAt); - const previousVersion = await this.storedVersion(version); - const status = existing - ? existing.versionId === prepared.versionId - ? "unchanged" - : "updated" - : "created"; - await this.writeAndVerifyObject(input, version); - try { - const materialization = await this.materializeProjectFile(input, prepared); - const file = await this.commitProjectFile( - input, - prepared, - existing, - previousVersion, +function fileUploadStatus( + existing: ProjectFile | null, + versionId: string, +): "created" | "unchanged" | "updated" { + if (!existing) return "created"; + return existing.versionId === versionId ? "unchanged" : "updated"; +} + +async function commitProjectFile( + runtime: FileRuntime, + input: z.output, + prepared: PreparedProjectFile, + existing: ProjectFile | null, + previousVersion: ProjectFileVersion | null, + version: ProjectFileVersion, + persistedAt: string, + materialization: ProjectFileMaterialization, +): Promise { + const file = currentProjectFile(input, prepared, existing, previousVersion, persistedAt); + await runtime.storage.transaction(async (transaction) => { + if (!previousVersion) { + await transaction.put( + versionRecordKey(input.projectId, prepared.fileId, prepared.versionId), version, - persistedAt, - materialization, ); - return ProjectFileUploadResponseSchema.parse({ file, status }); - } catch (error) { - if (!previousVersion) { - await this.env.R2_OUTPUTS.delete(prepared.r2Key).catch(() => undefined); - } - throw error; } - } + await transaction.put(fileRecordKey(input.projectId, prepared.fileId), file); + await transaction.put( + materializationRecordKey(input.projectId, prepared.fileId), + materialization, + ); + }); + return file; +} - private async commitProjectFile( - input: z.output, - prepared: PreparedProjectFile, - existing: ProjectFile | null, - previousVersion: ProjectFileVersion | null, - version: ProjectFileVersion, - persistedAt: string, - materialization: ProjectFileMaterialization, - ): Promise { - const file = ProjectFileSchema.parse({ - contentType: input.contentType, - createdAt: existing?.createdAt ?? persistedAt, - fileId: prepared.fileId, - name: input.name, - path: input.path, - projectId: input.projectId, - sha256: prepared.contentSha256, - sizeBytes: input.bytes.byteLength, - updatedAt: existing?.versionId === prepared.versionId ? existing.updatedAt : persistedAt, - versionCount: (existing?.versionCount ?? 0) + (previousVersion ? 0 : 1), - versionId: prepared.versionId, - }); - await this.ctx.storage.transaction(async (transaction) => { - if (!previousVersion) { - await transaction.put( - versionRecordKey(input.projectId, prepared.fileId, prepared.versionId), - version, - ); - } - await transaction.put(fileRecordKey(input.projectId, prepared.fileId), file); - await transaction.put( - materializationRecordKey(input.projectId, prepared.fileId), - materialization, - ); - }); - return file; - } +function currentProjectFile( + input: z.output, + prepared: PreparedProjectFile, + existing: ProjectFile | null, + previousVersion: ProjectFileVersion | null, + persistedAt: string, +): ProjectFile { + return ProjectFileSchema.parse({ + contentType: input.contentType, + createdAt: existing?.createdAt ?? persistedAt, + fileId: prepared.fileId, + name: input.name, + path: input.path, + projectId: input.projectId, + sha256: prepared.contentSha256, + sizeBytes: input.bytes.byteLength, + updatedAt: existing?.versionId === prepared.versionId ? existing.updatedAt : persistedAt, + versionCount: (existing?.versionCount ?? 0) + (previousVersion ? 0 : 1), + versionId: prepared.versionId, + }); +} - private async materializeProjectFile( - input: z.output, - prepared: PreparedProjectFile, - ): Promise { - const sandboxId = await this.ensureSandbox(); - const projectRoot = workspacePathForSlug(input.workspaceSlug); - const uploadsPath = `${projectRoot}/uploads`; - const workspacePath = `${projectRoot}/${input.path}`; - const written = await this.writeProjectFileWithRecovery( - sandboxId, - projectRoot, - workspacePath, - input.bytes, +async function materializeProjectFile( + context: FileContext, + input: z.output, + prepared: PreparedProjectFile, +): Promise { + const sandboxId = await context.runtime.ensureSandbox(); + const projectRoot = workspacePathForSlug(input.workspaceSlug); + const uploadsPath = `${projectRoot}/uploads`; + const workspacePath = `${projectRoot}/${input.path}`; + const written = await writeProjectFileWithRecovery( + context.runtime, + sandboxId, + projectRoot, + workspacePath, + input.bytes, + ); + if ( + written.byteLength !== input.bytes.byteLength || + (await sha256Hex(written)) !== prepared.contentSha256 + ) { + throw new APIError( + 502, + "upstream_sandbox_failed", + "Project file could not be verified in the workspace", + { retriable: true }, ); - if ( - written.byteLength !== input.bytes.byteLength || - (await sha256Hex(written)) !== prepared.contentSha256 - ) { - throw new APIError( - 502, - "upstream_sandbox_failed", - "Project file could not be verified in the workspace", - { retriable: true }, - ); - } - return this.readProjectFileMaterialization(sandboxId, uploadsPath, input, prepared); } + return readProjectFileMaterialization(context.runtime, sandboxId, uploadsPath, input, prepared); +} - private async writeProjectFileWithRecovery( - sandboxId: string, - projectRoot: string, - workspacePath: string, - bytes: Uint8Array, - ): Promise { - try { - return await this.writeProjectFileToWorkspace(sandboxId, projectRoot, workspacePath, bytes); - } catch (error) { - if (!isRecoverableWorkspaceMountError(error)) { - throw error; - } - await this.restartSandboxForWorkspaceRecovery(sandboxId); - return this.writeProjectFileToWorkspace(sandboxId, projectRoot, workspacePath, bytes); - } +async function writeProjectFileWithRecovery( + runtime: FileRuntime, + sandboxId: string, + projectRoot: string, + workspacePath: string, + bytes: Uint8Array, +): Promise { + try { + return await writeProjectFileToWorkspace(runtime, sandboxId, projectRoot, workspacePath, bytes); + } catch (error) { + if (!isRecoverableWorkspaceMountError(error)) throw error; + await runtime.restartSandboxForWorkspaceRecovery(sandboxId); + return writeProjectFileToWorkspace(runtime, sandboxId, projectRoot, workspacePath, bytes); } +} - private async writeProjectFileToWorkspace( - sandboxId: string, - projectRoot: string, - workspacePath: string, - bytes: Uint8Array, - ): Promise { - const uploadsPath = `${projectRoot}/uploads`; - await this.client().createFolder(sandboxId, uploadsPath); - await this.client().uploadFile(sandboxId, workspacePath, bytes); - return this.downloadUploadedFile(sandboxId, workspacePath, bytes.byteLength); - } +async function writeProjectFileToWorkspace( + runtime: FileRuntime, + sandboxId: string, + projectRoot: string, + workspacePath: string, + bytes: Uint8Array, +): Promise { + await runtime.client().createFolder(sandboxId, `${projectRoot}/uploads`); + await runtime.client().uploadFile(sandboxId, workspacePath, bytes); + return downloadUploadedFile(runtime, sandboxId, workspacePath, bytes.byteLength); +} - private async downloadUploadedFile( - sandboxId: string, - workspacePath: string, - maxBytes: number, - ): Promise { - let lastError: unknown; - for (let attempt = 0; attempt < WORKSPACE_FILE_VISIBILITY_ATTEMPTS; attempt += 1) { - try { - return await this.client().downloadFile(sandboxId, workspacePath, maxBytes); - } catch (error) { - if (!(error instanceof DaytonaApiError) || (error.status !== 404 && !error.retriable)) { - throw error; - } - lastError = error; - await sleep(WORKSPACE_FILE_VISIBILITY_DELAY_MS); +async function downloadUploadedFile( + runtime: FileRuntime, + sandboxId: string, + workspacePath: string, + maxBytes: number, +): Promise { + let lastError: unknown; + for (let attempt = 0; attempt < WORKSPACE_FILE_VISIBILITY_ATTEMPTS; attempt += 1) { + try { + return await runtime.client().downloadFile(sandboxId, workspacePath, maxBytes); + } catch (error) { + if (!(error instanceof DaytonaApiError) || (error.status !== 404 && !error.retriable)) { + throw error; } + lastError = error; + await sleep(WORKSPACE_FILE_VISIBILITY_DELAY_MS); } - throw lastError ?? new Error("Uploaded project file did not become readable"); } + throw lastError ?? new Error("Uploaded project file did not become readable"); +} - private async writeAndVerifyObject( - input: z.output, - version: ProjectFileVersion, - ): Promise { - const stored = await this.env.R2_OUTPUTS.put(version.r2Key, input.bytes, { - customMetadata: { - contentSha256: version.sha256, - fileId: version.fileId, - projectId: version.projectId, - versionId: version.versionId, - }, - httpMetadata: { contentType: version.contentType }, - onlyIf: { etagDoesNotMatch: "*" }, - sha256: version.sha256, - }); - assertStoredProjectFile(stored ?? (await this.env.R2_OUTPUTS.head(version.r2Key)), version); - } +async function writeAndVerifyObject( + runtime: FileRuntime, + input: z.output, + version: ProjectFileVersion, +): Promise { + const stored = await runtime.outputBucket.put(version.r2Key, input.bytes, { + customMetadata: { + contentSha256: version.sha256, + fileId: version.fileId, + projectId: version.projectId, + versionId: version.versionId, + }, + httpMetadata: { contentType: version.contentType }, + onlyIf: { etagDoesNotMatch: "*" }, + sha256: version.sha256, + }); + assertStoredProjectFile(stored ?? (await runtime.outputBucket.head(version.r2Key)), version); +} - private async currentFile(projectId: string, path: string): Promise { - const fileId = await deterministicUuid([ - FILE_DIGEST_DOMAIN, - this.ownerUserId(), - projectId, - path, - ]); - const value = await this.ctx.storage.get(fileRecordKey(projectId, fileId)); - const parsed = ProjectFileSchema.safeParse(value); - return parsed.success ? parsed.data : null; - } +async function currentFile( + runtime: FileRuntime, + projectId: string, + path: string, +): Promise { + const fileId = await deterministicUuid([ + FILE_DIGEST_DOMAIN, + runtime.ownerUserId(), + projectId, + path, + ]); + const value = await runtime.storage.get(fileRecordKey(projectId, fileId)); + const parsed = ProjectFileSchema.safeParse(value); + return parsed.success ? parsed.data : null; +} - private async storedVersion(version: ProjectFileVersion): Promise { - const value = await this.ctx.storage.get( - versionRecordKey(version.projectId, version.fileId, version.versionId), - ); - const parsed = ProjectFileVersionSchema.safeParse(value); - return parsed.success ? parsed.data : null; - } +async function storedVersion( + runtime: FileRuntime, + version: ProjectFileVersion, +): Promise { + const value = await runtime.storage.get( + versionRecordKey(version.projectId, version.fileId, version.versionId), + ); + const parsed = ProjectFileVersionSchema.safeParse(value); + return parsed.success ? parsed.data : null; +} - private async enforceFileCount(projectId: string, fileExists: boolean): Promise { - if (fileExists) return; - const files = await this.ctx.storage.list({ prefix: fileRecordProjectPrefix(projectId) }); - if (files.size >= PROJECT_FILE_MAX_CURRENT_FILES) { - throw new APIError( - 409, - "conflict_state_invalid", - "This project has too many uploaded files", - { - hint: "Remove an older project file before uploading another one.", - retriable: false, - }, - ); - } +async function enforceFileCount( + runtime: FileRuntime, + projectId: string, + fileExists: boolean, +): Promise { + if (fileExists) return; + const files = await runtime.storage.list({ prefix: fileRecordProjectPrefix(projectId) }); + if (files.size >= PROJECT_FILE_MAX_CURRENT_FILES) { + throw new APIError(409, "conflict_state_invalid", "This project has too many uploaded files", { + hint: "Remove an older project file before uploading another one.", + retriable: false, + }); } +} - private enqueueProjectFileMutation(operation: () => Promise): Promise { - const pending = this.projectFileMutationTail.catch(() => undefined).then(operation); - this.projectFileMutationTail = pending.then( - () => undefined, - () => undefined, - ); - return pending; - } +function enqueueProjectFileMutation( + context: FileContext, + operation: () => Promise, +): Promise { + const pending = context.mutationTail.catch(() => undefined).then(operation); + context.mutationTail = pending.then( + () => undefined, + () => undefined, + ); + return pending; +} - private async restoreProjectFiles( - input: z.output, - ): Promise<{ restoredFileCount: number }> { - const currentFiles = await this.currentProjectFiles(input.projectId); - if (currentFiles.length === 0) { - return { restoredFileCount: 0 }; - } - const sandboxId = await this.ensureSandbox(); - const projectRoot = workspacePathForSlug(input.workspaceSlug); - const uploadsPath = `${projectRoot}/uploads`; - const workspaceFiles = await this.workspaceUploadedFiles(sandboxId, uploadsPath); - const workspaceFilesByName = new Map(workspaceFiles.map((file) => [file.name, file])); - const restoredFiles: CurrentProjectFile[] = []; - let restoredFileCount = 0; - for (const current of currentFiles) { - const workspaceFile = workspaceFilesByName.get(current.file.name); - if (!workspaceProjectFileNeedsRestore(workspaceFile, current.file, current.materialization)) { - continue; - } - const bytes = await this.readStoredProjectFile(current.version); - const written = await this.writeProjectFileWithRecovery( - sandboxId, - projectRoot, - `${projectRoot}/${current.file.path}`, - bytes, - ); - await assertWorkspaceProjectFile(written, current.file); +async function restoreProjectFiles( + context: FileContext, + input: z.output, +): Promise<{ restoredFileCount: number }> { + const currentFiles = await currentProjectFiles(context.runtime, input.projectId); + if (currentFiles.length === 0) return { restoredFileCount: 0 }; + const sandboxId = await context.runtime.ensureSandbox(); + const projectRoot = workspacePathForSlug(input.workspaceSlug); + const uploadsPath = `${projectRoot}/uploads`; + const workspaceFiles = await workspaceUploadedFiles(context.runtime, sandboxId, uploadsPath); + const workspaceFilesByName = new Map(workspaceFiles.map((file) => [file.name, file])); + const restoredFiles: CurrentProjectFile[] = []; + for (const current of currentFiles) { + const workspaceFile = workspaceFilesByName.get(current.file.name); + if (workspaceProjectFileNeedsRestore(workspaceFile, current.file, current.materialization)) { + await restoreProjectFile(context.runtime, sandboxId, projectRoot, current); restoredFiles.push(current); - restoredFileCount += 1; } - await this.recordRestoredProjectFiles(sandboxId, uploadsPath, restoredFiles); - return { restoredFileCount }; } + await recordRestoredProjectFiles(context.runtime, sandboxId, uploadsPath, restoredFiles); + return { restoredFileCount: restoredFiles.length }; +} - private async recordRestoredProjectFiles( - sandboxId: string, - uploadsPath: string, - restoredFiles: CurrentProjectFile[], - ): Promise { - if (restoredFiles.length === 0) return; - const workspaceFiles = await this.workspaceUploadedFiles(sandboxId, uploadsPath); - const workspaceFilesByName = new Map(workspaceFiles.map((file) => [file.name, file])); - for (const current of restoredFiles) { - const workspaceFile = workspaceFilesByName.get(current.file.name); - if (!workspaceFile || workspaceFile.isDir || workspaceFile.size !== current.file.sizeBytes) { - throw new APIError( - 502, - "upstream_sandbox_failed", - "Restored project file metadata is invalid", - { retriable: true }, - ); - } - const materialization = ProjectFileMaterializationSchema.parse({ - fileId: current.file.fileId, - modifiedAt: workspaceFile.modifiedAt, - projectId: current.file.projectId, - sizeBytes: current.file.sizeBytes, - versionId: current.file.versionId, - }); - await this.ctx.storage.put( - materializationRecordKey(materialization.projectId, materialization.fileId), - materialization, +async function restoreProjectFile( + runtime: FileRuntime, + sandboxId: string, + projectRoot: string, + current: CurrentProjectFile, +): Promise { + const bytes = await readStoredProjectFile(runtime, current.version); + const written = await writeProjectFileWithRecovery( + runtime, + sandboxId, + projectRoot, + `${projectRoot}/${current.file.path}`, + bytes, + ); + await assertWorkspaceProjectFile(written, current.file); +} + +async function recordRestoredProjectFiles( + runtime: FileRuntime, + sandboxId: string, + uploadsPath: string, + restoredFiles: CurrentProjectFile[], +): Promise { + if (restoredFiles.length === 0) return; + const workspaceFiles = await workspaceUploadedFiles(runtime, sandboxId, uploadsPath); + const workspaceFilesByName = new Map(workspaceFiles.map((file) => [file.name, file])); + for (const current of restoredFiles) { + const workspaceFile = workspaceFilesByName.get(current.file.name); + if (!workspaceFile || workspaceFile.isDir || workspaceFile.size !== current.file.sizeBytes) { + throw new APIError( + 502, + "upstream_sandbox_failed", + "Restored project file metadata is invalid", + { retriable: true }, ); } + const materialization = ProjectFileMaterializationSchema.parse({ + fileId: current.file.fileId, + modifiedAt: workspaceFile.modifiedAt, + projectId: current.file.projectId, + sizeBytes: current.file.sizeBytes, + versionId: current.file.versionId, + }); + await runtime.storage.put( + materializationRecordKey(materialization.projectId, materialization.fileId), + materialization, + ); } +} - private async currentProjectFiles(projectId: string): Promise { - const records = await this.ctx.storage.list({ - prefix: fileRecordProjectPrefix(projectId), +async function currentProjectFiles( + runtime: FileRuntime, + projectId: string, +): Promise { + const records = await runtime.storage.list({ prefix: fileRecordProjectPrefix(projectId) }); + const currentFiles: CurrentProjectFile[] = []; + for (const value of records.values()) { + const file = ProjectFileSchema.parse(value); + const version = ProjectFileVersionSchema.parse( + await runtime.storage.get(versionRecordKey(file.projectId, file.fileId, file.versionId)), + ); + assertCurrentProjectFile(file, version); + const materialization = ProjectFileMaterializationSchema.safeParse( + await runtime.storage.get(materializationRecordKey(file.projectId, file.fileId)), + ); + currentFiles.push({ + file, + materialization: materialization.success ? materialization.data : null, + version, }); - const currentFiles: CurrentProjectFile[] = []; - for (const value of records.values()) { - const file = ProjectFileSchema.parse(value); - const versionValue = await this.ctx.storage.get( - versionRecordKey(file.projectId, file.fileId, file.versionId), - ); - const version = ProjectFileVersionSchema.parse(versionValue); - assertCurrentProjectFile(file, version); - const materializationValue = await this.ctx.storage.get( - materializationRecordKey(file.projectId, file.fileId), - ); - const materialization = ProjectFileMaterializationSchema.safeParse(materializationValue); - currentFiles.push({ - file, - materialization: materialization.success ? materialization.data : null, - version, - }); - } - currentFiles.sort((left, right) => left.file.path.localeCompare(right.file.path)); - return currentFiles; } + currentFiles.sort((left, right) => left.file.path.localeCompare(right.file.path)); + return currentFiles; +} - private async workspaceUploadedFiles(sandboxId: string, uploadsPath: string) { +async function workspaceUploadedFiles( + runtime: FileRuntime, + sandboxId: string, + uploadsPath: string, +) { + try { + return await runtime.client().listFiles(sandboxId, uploadsPath); + } catch (error) { + if (error instanceof DaytonaApiError && error.status === 404) return []; + if (!isRecoverableWorkspaceMountError(error)) throw error; + await runtime.restartSandboxForWorkspaceRecovery(sandboxId); try { - return await this.client().listFiles(sandboxId, uploadsPath); - } catch (error) { - if (error instanceof DaytonaApiError && error.status === 404) { - return []; - } - if (!isRecoverableWorkspaceMountError(error)) { - throw error; - } - await this.restartSandboxForWorkspaceRecovery(sandboxId); - try { - return await this.client().listFiles(sandboxId, uploadsPath); - } catch (retryError) { - if (retryError instanceof DaytonaApiError && retryError.status === 404) { - return []; - } - throw retryError; - } + return await runtime.client().listFiles(sandboxId, uploadsPath); + } catch (retryError) { + if (retryError instanceof DaytonaApiError && retryError.status === 404) return []; + throw retryError; } } +} - private async readStoredProjectFile(version: ProjectFileVersion): Promise { - const object = await this.env.R2_OUTPUTS.get(version.r2Key); - if (!object) { - throw new APIError(409, "conflict_state_invalid", "Stored project file is missing", { - retriable: false, - }); - } - assertStoredProjectFile(object, version); - const bytes = new Uint8Array(await object.arrayBuffer()); - if (bytes.byteLength !== version.sizeBytes || (await sha256Hex(bytes)) !== version.sha256) { - throw new APIError( - 409, - "conflict_state_invalid", - "Stored project file contents are invalid", - { retriable: false }, - ); - } - return bytes; +async function readStoredProjectFile( + runtime: FileRuntime, + version: ProjectFileVersion, +): Promise { + const object = await runtime.outputBucket.get(version.r2Key); + if (!object) { + throw new APIError(409, "conflict_state_invalid", "Stored project file is missing", { + retriable: false, + }); } - - private async readProjectFileMaterialization( - sandboxId: string, - uploadsPath: string, - input: z.output, - prepared: PreparedProjectFile, - ): Promise { - const workspaceFiles = await this.workspaceUploadedFiles(sandboxId, uploadsPath); - const workspaceFile = workspaceFiles.find((file) => file.name === input.name); - if (!workspaceFile || workspaceFile.isDir || workspaceFile.size !== input.bytes.byteLength) { - throw new APIError( - 502, - "upstream_sandbox_failed", - "Project file workspace metadata is invalid", - { retriable: true }, - ); - } - return ProjectFileMaterializationSchema.parse({ - fileId: prepared.fileId, - modifiedAt: workspaceFile.modifiedAt, - projectId: input.projectId, - sizeBytes: input.bytes.byteLength, - versionId: prepared.versionId, + assertStoredProjectFile(object, version); + const bytes = new Uint8Array(await object.arrayBuffer()); + if (bytes.byteLength !== version.sizeBytes || (await sha256Hex(bytes)) !== version.sha256) { + throw new APIError(409, "conflict_state_invalid", "Stored project file contents are invalid", { + retriable: false, }); } + return bytes; +} - private async listProjectFileRecords(projectId: string): Promise<{ files: ProjectFile[] }> { - const records = await this.ctx.storage.list({ - prefix: fileRecordProjectPrefix(projectId), - }); - const files = Array.from(records.values()).flatMap((value) => { - const parsed = ProjectFileSchema.safeParse(value); - return parsed.success ? [parsed.data] : []; - }); - files.sort((left, right) => left.path.localeCompare(right.path)); - return ProjectFileListSchema.parse({ files }); +async function readProjectFileMaterialization( + runtime: FileRuntime, + sandboxId: string, + uploadsPath: string, + input: z.output, + prepared: PreparedProjectFile, +): Promise { + const workspaceFiles = await workspaceUploadedFiles(runtime, sandboxId, uploadsPath); + const workspaceFile = workspaceFiles.find((file) => file.name === input.name); + if (!workspaceFile || workspaceFile.isDir || workspaceFile.size !== input.bytes.byteLength) { + throw new APIError( + 502, + "upstream_sandbox_failed", + "Project file workspace metadata is invalid", + { retriable: true }, + ); } + return ProjectFileMaterializationSchema.parse({ + fileId: prepared.fileId, + modifiedAt: workspaceFile.modifiedAt, + projectId: input.projectId, + sizeBytes: input.bytes.byteLength, + versionId: prepared.versionId, + }); +} - private async deleteStoragePrefix(prefix: string): Promise { - while (true) { - const records = await this.ctx.storage.list({ limit: DELETE_BATCH_SIZE, prefix }); - const keys = [...records.keys()]; - if (keys.length === 0) return; - await this.ctx.storage.delete(keys); - } +async function listProjectFileRecords( + runtime: FileRuntime, + projectId: string, +): Promise<{ files: ProjectFile[] }> { + const records = await runtime.storage.list({ prefix: fileRecordProjectPrefix(projectId) }); + const files = Array.from(records.values()).flatMap((value) => { + const parsed = ProjectFileSchema.safeParse(value); + return parsed.success ? [parsed.data] : []; + }); + files.sort((left, right) => left.path.localeCompare(right.path)); + return ProjectFileListSchema.parse({ files }); +} + +async function deleteStoragePrefix(runtime: FileRuntime, prefix: string): Promise { + while (true) { + const records = await runtime.storage.list({ limit: DELETE_BATCH_SIZE, prefix }); + const keys = [...records.keys()]; + if (keys.length === 0) return; + await runtime.storage.delete(keys); } } @@ -545,16 +598,12 @@ function assertCurrentProjectFile(file: ProjectFile, version: ProjectFileVersion } async function assertWorkspaceProjectFile(bytes: Uint8Array, file: ProjectFile): Promise { - if (bytes.byteLength === file.sizeBytes && (await sha256Hex(bytes)) === file.sha256) { - return; - } + if (bytes.byteLength === file.sizeBytes && (await sha256Hex(bytes)) === file.sha256) return; throw new APIError( 502, "upstream_sandbox_failed", "Project file restoration could not be verified", - { - retriable: true, - }, + { retriable: true }, ); } diff --git a/apps/agent-worker/src/durable-objects/project-sandbox-runtime-handle.ts b/apps/agent-worker/src/durable-objects/project-sandbox-runtime-handle.ts new file mode 100644 index 00000000..29c853cc --- /dev/null +++ b/apps/agent-worker/src/durable-objects/project-sandbox-runtime-handle.ts @@ -0,0 +1,463 @@ +import { PreviewHostnameSchema, resolveWorkerSecret } from "@cheatcode/env"; +import { APIError, createLogger } from "@cheatcode/observability"; +import { DaytonaClient, type DaytonaSandbox } from "@cheatcode/tools-code"; +import { performAccountDeletion } from "./project-sandbox-account-deletion"; +import { type SandboxExecAuditEntry, writeExecAudit } from "./project-sandbox-audit"; +import { ProjectSandboxIdentityState } from "./project-sandbox-identity-state"; +import { + withActiveOperation, + withCleanupSignal, + withOwnerRegistration, + withProjectCleanup, + withSharedWorkspaceMutation, + withStreamingOperation, + workspaceState, +} from "./project-sandbox-lease-runtime"; +import { + ACCOUNT_DELETION_TOMBSTONE_KEY, + DAYTONA_ID_KEY, + type ProjectSandboxEnv, + RUN_LEASES_KEY, + runLeases, + STALE_RUN_LEASE_MS, + STARTED_REVERIFY_MS, + sandboxRuntimeUpdatePending, + storedDaytonaId, + toUpstreamError, +} from "./project-sandbox-lifecycle-support"; +import type { SandboxMeteringContext } from "./project-sandbox-metering"; +import { PROC_PREFIX, PROCESS_PORT_ALLOC_KEY } from "./project-sandbox-process-support"; +import { ProjectSandboxProvisioning } from "./project-sandbox-provisioning"; +import type { ParsedProjectCleanupWorkspaceInput } from "./project-sandbox-runtime"; +import { + openProjectSandboxWorkspaceState, + type ProjectSandboxWorkspaceState, +} from "./project-sandbox-workspace-state"; + +const RUNTIME_RESET_PENDING_KEY = "sandbox_runtime_reset_pending"; +const SKILL_RUNTIME_DIRECTORY = "/workspace/.cheatcode/runtime"; +interface RuntimeCache { + client: DaytonaClient | undefined; + sandboxId: string | undefined; + startedVerifiedAtMs: number; +} + +interface RuntimeState { + accountDeletionCompleted: boolean; + accountDeletionInProgress: boolean; + accountDeletionPromise: Promise | undefined; + activeOperationCount: number; + activeOperationDrainWaiters: Set<() => void>; + cache: RuntimeCache; + ctx: DurableObjectState; + env: ProjectSandboxEnv; + identity: ProjectSandboxIdentityState; + provisioning: ProjectSandboxProvisioning; + sandboxMutationTail: Promise; + sandboxRuntimeUpdateInProgress: boolean; + workspaceState: ProjectSandboxWorkspaceState | undefined; +} + +interface SandboxLeaseRuntime { + withCleanupSignal: (operation: () => Promise) => Promise; + withOwnerRegistration: ( + userId: string, + operation: () => Promise, + ) => Promise; + withProjectCleanup: (operation: () => Promise) => Promise; + withSandboxOperation: (operation: () => Promise) => Promise; + withSharedWorkspaceMutation: (operation: () => Promise) => Promise; + withStreamingOperation: ( + workspaceScope: string | readonly string[] | null, + operation: (release: () => void) => Promise, + ) => Promise; + withWorkspaceOperation: ( + workspaceScope: string | readonly string[] | null, + operation: () => Promise, + ) => Promise; +} + +export interface SandboxRuntime { + readonly client: () => DaytonaClient; + readonly deleteAccountState: () => Promise; + readonly deleteProjectWorkspace: ( + input: ParsedProjectCleanupWorkspaceInput, + cleanup: () => Promise, + ) => Promise; + readonly ensureExistingSandboxStarted: () => Promise; + readonly ensureSandbox: (startingRunId?: string) => Promise; + readonly existingSandboxId: () => Promise; + readonly lease: SandboxLeaseRuntime; + readonly meteringContext: () => Promise; + readonly outputBucket: R2Bucket; + readonly ownerUserId: () => string; + readonly previewHostname: () => string; + readonly previewSecret: () => Promise; + readonly registerOwner: (userId: string, sandboxName?: string) => Promise; + readonly restartSandboxForWorkspaceRecovery: (sandboxId: string) => Promise; + readonly sandboxName: () => string; + readonly storage: DurableObjectStorage; + readonly toUpstreamError: (error: unknown, fallback: string) => APIError; + readonly writeExecAudit: (entry: SandboxExecAuditEntry) => Promise; +} + +export function createSandboxRuntime( + ctx: DurableObjectState, + env: ProjectSandboxEnv, +): SandboxRuntime { + const identity = new ProjectSandboxIdentityState(ctx); + const cache: RuntimeCache = { + client: undefined, + sandboxId: undefined, + startedVerifiedAtMs: 0, + }; + const provisioning = createProvisioning(env, identity, cache, ctx); + const state: RuntimeState = { + accountDeletionCompleted: false, + accountDeletionInProgress: false, + accountDeletionPromise: undefined, + activeOperationCount: 0, + activeOperationDrainWaiters: new Set(), + cache, + ctx, + env, + identity, + provisioning, + sandboxMutationTail: Promise.resolve(), + sandboxRuntimeUpdateInProgress: false, + workspaceState: openProjectSandboxWorkspaceState(ctx), + }; + void ctx.blockConcurrencyWhile(() => initializeIdentityState(state)); + return runtimeHandle(state); +} + +function createProvisioning( + env: ProjectSandboxEnv, + identity: ProjectSandboxIdentityState, + cache: RuntimeCache, + ctx: DurableObjectState, +): ProjectSandboxProvisioning { + return new ProjectSandboxProvisioning({ + cachedSandboxId: async () => cache.sandboxId ?? storedDaytonaId(ctx.storage), + env, + sandboxName: () => identity.sandboxName(), + toUpstreamError: (error, fallback) => toUpstreamError(error, fallback, identity.sandboxName()), + }); +} + +function runtimeHandle(state: RuntimeState): SandboxRuntime { + return { + client: () => client(state), + deleteAccountState: () => deleteAccountState(state), + deleteProjectWorkspace: (input, cleanup) => + workspaceState(state).deleteWorkspace(input, cleanup), + ensureExistingSandboxStarted: () => ensureExistingSandboxStarted(state), + ensureSandbox: (startingRunId) => ensureSandbox(state, startingRunId), + existingSandboxId: () => existingSandboxId(state), + lease: leaseRuntime(state), + meteringContext: () => Promise.resolve(meteringContext(state)), + outputBucket: state.env.R2_OUTPUTS, + ownerUserId: () => ownerUserId(state), + previewHostname: () => PreviewHostnameSchema.parse(state.env.PREVIEW_HOSTNAME), + previewSecret: () => previewSecret(state), + registerOwner: (userId, sandboxName) => state.identity.registerOwner(userId, sandboxName), + restartSandboxForWorkspaceRecovery: (sandboxId) => + restartSandboxForWorkspaceRecovery(state, sandboxId), + sandboxName: () => state.identity.sandboxName(), + storage: state.ctx.storage, + toUpstreamError: (error, fallback) => + toUpstreamError(error, fallback, state.identity.sandboxName()), + writeExecAudit: (entry) => writeExecAudit(state.env.R2_AUDIT, entry), + }; +} + +function leaseRuntime(state: RuntimeState): SandboxLeaseRuntime { + return { + withCleanupSignal: (operation) => withCleanupSignal(state, operation), + withOwnerRegistration: (userId, operation) => withOwnerRegistration(state, userId, operation), + withProjectCleanup: (operation) => withProjectCleanup(state, operation), + withSandboxOperation: (operation) => withActiveOperation(state, null, operation, false, true), + withSharedWorkspaceMutation: (operation) => withSharedWorkspaceMutation(state, operation), + withStreamingOperation: (scope, operation) => withStreamingOperation(state, scope, operation), + withWorkspaceOperation: (scope, operation) => + withActiveOperation(state, scope, operation, false, true), + }; +} + +async function initializeIdentityState(state: RuntimeState): Promise { + const isDeleted = (await state.ctx.storage.get(ACCOUNT_DELETION_TOMBSTONE_KEY)) === true; + if (isDeleted) { + state.accountDeletionInProgress = true; + } + await state.identity.initialize(); +} + +function deleteAccountState(state: RuntimeState): Promise { + if (state.accountDeletionCompleted) { + return Promise.resolve(); + } + if (state.accountDeletionPromise !== undefined) { + return state.accountDeletionPromise; + } + state.accountDeletionInProgress = true; + const deletion = performAccountDeletion(state, { + clearCachedSandbox: () => clearCachedSandbox(state), + ensureClient: () => ensureClient(state), + meteringContext: () => meteringContext(state), + withSandboxMutation: (operation) => withSandboxMutation(state, operation), + }); + const tracked = deletion.finally(() => { + if (state.accountDeletionPromise === tracked) { + state.accountDeletionPromise = undefined; + } + }); + state.accountDeletionPromise = tracked; + return tracked; +} + +function client(state: RuntimeState): DaytonaClient { + if (!state.cache.client) { + throw new Error("Daytona client accessed before initialization."); + } + return state.cache.client; +} + +async function ensureClient(state: RuntimeState): Promise { + if (state.cache.client) { + return state.cache.client; + } + const apiKey = await resolveWorkerSecret(state.env.DAYTONA_API_KEY); + if (!apiKey) { + throw new APIError(503, "unavailable_maintenance", "DAYTONA_API_KEY is not configured", { + retriable: false, + }); + } + state.cache.client = new DaytonaClient({ + apiKey, + apiUrl: state.env.DAYTONA_API_URL, + target: state.env.DAYTONA_TARGET, + ...(state.env.DAYTONA_ORG_ID ? { organizationId: state.env.DAYTONA_ORG_ID } : {}), + ...(state.env.DAYTONA_PREVIEW_HOST_SUFFIXES + ? { previewHostSuffixes: state.env.DAYTONA_PREVIEW_HOST_SUFFIXES } + : {}), + }); + return state.cache.client; +} + +function setCachedSandboxId(state: RuntimeState, sandboxId: string): void { + state.cache.sandboxId = sandboxId; + state.cache.startedVerifiedAtMs = Date.now(); +} + +async function ensureSandbox(state: RuntimeState, startingRunId?: string): Promise { + return withSandboxMutation(state, async () => { + if ( + state.cache.sandboxId && + Date.now() - state.cache.startedVerifiedAtMs < STARTED_REVERIFY_MS + ) { + return state.cache.sandboxId; + } + return resolveStartedSandbox(state, startingRunId); + }); +} + +async function restartSandboxForWorkspaceRecovery( + state: RuntimeState, + sandboxId: string, +): Promise { + await withSandboxMutation(state, async () => { + const daytona = await ensureClient(state); + try { + await state.provisioning.restart(daytona, sandboxId); + } catch (error) { + throw toUpstreamError( + error, + "Daytona workspace recovery failed.", + state.identity.sandboxName(), + ); + } + setCachedSandboxId(state, sandboxId); + }); +} + +async function ensureExistingSandboxStarted(state: RuntimeState): Promise { + return withSandboxMutation(state, async () => { + const daytona = await ensureClient(state); + let existing: DaytonaSandbox | null; + try { + existing = await state.provisioning.findExisting(daytona); + if (!existing || !(await state.provisioning.ensureStarted(daytona, existing))) { + return null; + } + } catch (error) { + throw toUpstreamError( + error, + "Daytona sandbox cleanup startup failed.", + state.identity.sandboxName(), + ); + } + state.cache.sandboxId = existing.id; + await state.ctx.storage.put(DAYTONA_ID_KEY, existing.id); + state.cache.startedVerifiedAtMs = Date.now(); + return existing.id; + }); +} + +async function withSandboxMutation( + state: RuntimeState, + operation: () => Promise, +): Promise { + const previous = state.sandboxMutationTail; + let release = (): void => undefined; + const gate = new Promise((resolve) => { + release = resolve; + }); + state.sandboxMutationTail = previous.catch(() => undefined).then(() => gate); + await previous.catch(() => undefined); + try { + return await operation(); + } finally { + release(); + } +} + +async function resolveStartedSandbox(state: RuntimeState, startingRunId?: string): Promise { + const daytona = await ensureClient(state); + let resolved: DaytonaSandbox; + try { + resolved = await state.provisioning.resolve(daytona); + if (!state.provisioning.isDesired(resolved)) { + resolved = await replaceSandboxRuntime(state, daytona, resolved, startingRunId); + } + } catch (error) { + throw toUpstreamError(error, "Daytona sandbox lookup failed.", state.identity.sandboxName()); + } + state.cache.sandboxId = resolved.id; + await state.ctx.storage.put(DAYTONA_ID_KEY, resolved.id); + if (!(await state.provisioning.ensureStarted(daytona, resolved))) { + throw new APIError(502, "upstream_sandbox_failed", "Daytona sandbox disappeared", { + retriable: true, + }); + } + await clearPersistedRuntimeProjection(state, daytona, resolved.id); + state.cache.startedVerifiedAtMs = Date.now(); + return resolved.id; +} + +async function replaceSandboxRuntime( + state: RuntimeState, + daytona: DaytonaClient, + current: DaytonaSandbox, + startingRunId?: string, +): Promise { + state.sandboxRuntimeUpdateInProgress = true; + try { + await assertSandboxReplacementAllowed(state, startingRunId); + state.provisioning.assertRuntimeReplacementSafe(current); + await prepareForSandboxReplacement(state); + await state.provisioning.deleteForReplacement(daytona, current); + const replacement = await state.provisioning.create(daytona); + if (!state.provisioning.isDesired(replacement)) { + throw sandboxRuntimeUpdatePending(state.env.DAYTONA_SANDBOX_SNAPSHOT); + } + createLogger().info("sandbox_runtime_replaced", { + sandboxId: state.identity.sandboxName(), + snapshot: state.env.DAYTONA_SANDBOX_SNAPSHOT, + }); + return replacement; + } finally { + state.sandboxRuntimeUpdateInProgress = false; + } +} + +async function assertSandboxReplacementAllowed( + state: RuntimeState, + startingRunId?: string, +): Promise { + const leases = await runLeases(state.ctx.storage); + const active = leases.filter((lease) => Date.now() - lease.startedMs < STALE_RUN_LEASE_MS); + if (active.length !== leases.length) { + await state.ctx.storage.put(RUN_LEASES_KEY, active); + } + const otherRuns = active.filter((lease) => lease.runId !== startingRunId); + if (state.activeOperationCount > 1 || otherRuns.length > 0) { + throw sandboxRuntimeUpdatePending(state.env.DAYTONA_SANDBOX_SNAPSHOT); + } +} + +async function prepareForSandboxReplacement(state: RuntimeState): Promise { + state.cache.sandboxId = undefined; + state.cache.startedVerifiedAtMs = 0; + await state.ctx.storage.delete(DAYTONA_ID_KEY); + await state.ctx.storage.put(RUNTIME_RESET_PENDING_KEY, true); + const processRecords = await state.ctx.storage.list({ prefix: PROC_PREFIX }); + if (processRecords.size > 0) { + await state.ctx.storage.delete([...processRecords.keys()]); + } + await state.ctx.storage.delete(PROCESS_PORT_ALLOC_KEY); +} + +async function clearPersistedRuntimeProjection( + state: RuntimeState, + daytona: DaytonaClient, + sandboxId: string, +): Promise { + if ((await state.ctx.storage.get(RUNTIME_RESET_PENDING_KEY)) !== true) { + return; + } + try { + await daytona.deleteFilePath(sandboxId, SKILL_RUNTIME_DIRECTORY, true); + await state.ctx.storage.delete(RUNTIME_RESET_PENDING_KEY); + } catch (error) { + throw toUpstreamError(error, "Daytona runtime reset failed.", state.identity.sandboxName()); + } +} + +async function existingSandboxId(state: RuntimeState): Promise { + const daytona = await ensureClient(state); + try { + const existing = await state.provisioning.findExisting(daytona); + if (existing) { + state.cache.sandboxId = existing.id; + return existing.id; + } + return null; + } catch (error) { + throw toUpstreamError(error, "Daytona sandbox lookup failed.", state.identity.sandboxName()); + } +} + +function meteringContext(state: RuntimeState): SandboxMeteringContext { + return { + env: state.env, + ownerUserId: state.identity.ownerUserId(), + sandboxId: state.identity.sandboxName(), + storage: state.ctx.storage, + }; +} + +async function previewSecret(state: RuntimeState): Promise { + const secret = await resolveWorkerSecret(state.env.PREVIEW_TOKEN_SECRET); + if (!secret) { + throw new APIError(503, "unavailable_maintenance", "PREVIEW_TOKEN_SECRET is not configured", { + retriable: false, + }); + } + return secret; +} + +function ownerUserId(state: RuntimeState): string { + const userId = state.identity.ownerUserId(); + if (!userId) { + throw new APIError(500, "internal_error", "ProjectSandbox owner is not registered", { + retriable: false, + }); + } + return userId; +} + +function clearCachedSandbox(state: RuntimeState): void { + state.cache.client = undefined; + state.cache.sandboxId = undefined; + state.cache.startedVerifiedAtMs = 0; +} diff --git a/apps/agent-worker/src/durable-objects/project-sandbox.ts b/apps/agent-worker/src/durable-objects/project-sandbox.ts index d99f193a..e709db27 100644 --- a/apps/agent-worker/src/durable-objects/project-sandbox.ts +++ b/apps/agent-worker/src/durable-objects/project-sandbox.ts @@ -1,86 +1,330 @@ -import { ProjectSandboxContent } from "./project-sandbox-content"; +import { DurableObject } from "cloudflare:workers"; +import type { + SandboxDeleteFileResult, + SandboxExecResult, + SandboxKillProcessResult, + SandboxListFilesResult, + SandboxProcessResult, + SandboxReadFileResult, + SandboxRunCodeResult, + SandboxSearchFilesResult, + SandboxWriteFileResult, +} from "@cheatcode/sandbox-contracts"; +import type { + ProjectFile, + ProjectFileUploadResponse, + SandboxConsoleSnapshot, +} from "@cheatcode/types"; +import { type ContentOps, createContentOps } from "./project-sandbox-content"; import { type LeaseMethod, leaseKind, workspaceScope } from "./project-sandbox-lease-policy"; +import { createLifecycleOps, type LifecycleOps } from "./project-sandbox-lifecycle"; +import type { ProjectSandboxEnv } from "./project-sandbox-lifecycle-support"; +import { + type CoordinatedProcessOps, + createProcessOps, + type ProcessOps, + type ProjectSandboxStatus, +} from "./project-sandbox-processes"; +import { createFileOps, type FileOps } from "./project-sandbox-project-files"; +import type { + ProjectAllocatePortInput, + ProjectAllocateProcessPortInput, + ProjectArchiveInput, + ProjectBrowserTakeoverInput, + ProjectBrowserTakeoverResult, + ProjectBrowserTakeoverStopInput, + ProjectCleanupWorkspaceInput, + ProjectCodeServerInput, + ProjectDeleteFileInput, + ProjectExecInput, + ProjectKillProcessInput, + ProjectListFilesInput, + ProjectListUploadedFilesInput, + ProjectPreviewStatusInput, + ProjectReadDevServerLogsInput, + ProjectReadFileInput, + ProjectRestoreUploadedFilesInput, + ProjectRunCodeInput, + ProjectSandboxRuntimeState, + ProjectSearchFilesInput, + ProjectSignedPreviewUrlInput, + ProjectStartProcessInput, + ProjectUploadFileInput, + ProjectWakePreviewInput, + ProjectWakePreviewResult, + ProjectWriteFileInput, +} from "./project-sandbox-runtime"; +import { createSandboxRuntime, type SandboxRuntime } from "./project-sandbox-runtime-handle"; -type MethodArgs = ProjectSandboxContent[Method] extends ( - ...args: infer Args -) => unknown - ? Args - : never; +type ProjectSandboxOperations = LifecycleOps & ProcessOps & FileOps & ContentOps; /** - * Public Durable Object facade. Every operational RPC takes its table-selected - * in-memory lease before the first await so account deletion can drain safely. + * Public Durable Object facade. The policy table is interpreted only here; + * collaborators expose raw operations and request coordinated nested calls + * through callbacks built in this constructor. */ -// biome-ignore format: Keep the explicit Durable Object RPC facade as an auditable one-line policy map. -export class ProjectSandbox extends ProjectSandboxContent { - public override registerOwner(...args: MethodArgs<"registerOwner">) { return this.withLease("registerOwner", args[0], () => super.registerOwner(...args)); } - public override setQuotaPeriod(...args: MethodArgs<"setQuotaPeriod">) { return this.withLease("setQuotaPeriod", args[0], () => super.setQuotaPeriod(...args)); } - public override beginRun(...args: MethodArgs<"beginRun">) { return this.withLease("beginRun", args[0], () => super.beginRun(...args)); } - public override renewRun(...args: MethodArgs<"renewRun">) { return this.withLease("renewRun", args[0], () => super.renewRun(...args)); } - public override endRun(...args: MethodArgs<"endRun">) { return this.withLease("endRun", args[0], () => super.endRun(...args)); } - public override alarm(...args: MethodArgs<"alarm">) { return this.withLease("alarm", undefined, () => super.alarm(...args)); } - public override runtimeSandboxId(...args: MethodArgs<"runtimeSandboxId">) { return this.withLease("runtimeSandboxId", undefined, () => super.runtimeSandboxId(...args)); } - public override existingDaytonaId(...args: MethodArgs<"existingDaytonaId">) { return this.withLease("existingDaytonaId", undefined, () => super.existingDaytonaId(...args)); } - public override sandboxRuntimeState(...args: MethodArgs<"sandboxRuntimeState">) { return this.withLease("sandboxRuntimeState", undefined, () => super.sandboxRuntimeState(...args)); } - public override ensureReady(...args: MethodArgs<"ensureReady">) { return this.withLease("ensureReady", undefined, () => super.ensureReady(...args)); } - public override getStatus(...args: MethodArgs<"getStatus">) { return this.withLease("getStatus", undefined, () => super.getStatus(...args)); } - public override runCode(...args: MethodArgs<"runCode">) { return this.withLease("runCode", args[0], () => super.runCode(...args)); } - public override exec(...args: MethodArgs<"exec">) { return this.withLease("exec", args[0], () => super.exec(...args)); } - public override startProcess(...args: MethodArgs<"startProcess">) { return this.withLease("startProcess", args[0], () => super.startProcess(...args)); } - public override allocateProjectPort(...args: MethodArgs<"allocateProjectPort">) { return this.withLease("allocateProjectPort", args[0], () => super.allocateProjectPort(...args)); } - public override allocateProcessPort(...args: MethodArgs<"allocateProcessPort">) { return this.withLease("allocateProcessPort", args[0], () => super.allocateProcessPort(...args)); } - public override killAllProcesses(...args: MethodArgs<"killAllProcesses">) { return this.withLease("killAllProcesses", undefined, () => super.killAllProcesses(...args)); } - public override killProcess(...args: MethodArgs<"killProcess">) { return this.withLease("killProcess", args[0], () => super.killProcess(...args)); } - public override readDevServerLogs(...args: MethodArgs<"readDevServerLogs">) { return this.withLease("readDevServerLogs", args[0], () => super.readDevServerLogs(...args)); } - public override downloadProjectArchive(...args: MethodArgs<"downloadProjectArchive">) { return this.withActiveProjectWorkspaceStreamingOperation(workspaceScope("downloadProjectArchive", args[0]), (release) => super.downloadProjectArchiveForRpc(args[0], release)); } - public override readFile(...args: MethodArgs<"readFile">) { return this.withLease("readFile", args[0], () => super.readFile(...args)); } - public override listUploadedFiles(...args: MethodArgs<"listUploadedFiles">) { return this.withLease("listUploadedFiles", undefined, () => super.listUploadedFiles(...args)); } - public override uploadProjectFile(...args: MethodArgs<"uploadProjectFile">) { return this.withLease("uploadProjectFile", args[0], () => super.uploadProjectFile(...args)); } - public override restoreUploadedFiles(...args: MethodArgs<"restoreUploadedFiles">) { return this.withLease("restoreUploadedFiles", args[0], () => super.restoreUploadedFiles(...args)); } - public override writeFile(...args: MethodArgs<"writeFile">) { return this.withLease("writeFile", args[0], () => super.writeFile(...args)); } - public override listFiles(...args: MethodArgs<"listFiles">) { return this.withLease("listFiles", args[0], () => super.listFiles(...args)); } - public override searchFiles(...args: MethodArgs<"searchFiles">) { return this.withLease("searchFiles", args[0], () => super.searchFiles(...args)); } - public override deleteFile(...args: MethodArgs<"deleteFile">) { return this.withLease("deleteFile", args[0], () => super.deleteFile(...args)); } - public override getSignedPreviewUrl(...args: MethodArgs<"getSignedPreviewUrl">) { return this.withLease("getSignedPreviewUrl", args[0], () => super.getSignedPreviewUrl(...args)); } - public override exposeBrowserTakeover(...args: MethodArgs<"exposeBrowserTakeover">) { return this.withLease("exposeBrowserTakeover", args[0], () => super.exposeBrowserTakeover(...args)); } - public override stopBrowserTakeover(...args: MethodArgs<"stopBrowserTakeover">) { return this.withLease("stopBrowserTakeover", args[0], () => super.stopBrowserTakeover(...args)); } - public override exposeCodeServer(...args: MethodArgs<"exposeCodeServer">) { return this.withLease("exposeCodeServer", args[0], () => super.exposeCodeServer(...args)); } - public override wakePreview(...args: MethodArgs<"wakePreview">) { return this.withLease("wakePreview", args[0], () => super.wakePreview(...args)); } - public override projectPreviewStatus(...args: MethodArgs<"projectPreviewStatus">) { return this.withLease("projectPreviewStatus", args[0], () => super.projectPreviewStatus(...args)); } - public override cleanupProjectWorkspace(...args: MethodArgs<"cleanupProjectWorkspace">) { return this.withLease("cleanupProjectWorkspace", args[0], () => super.cleanupProjectWorkspace(...args)); } +export class ProjectSandbox extends DurableObject { + private readonly operations: ProjectSandboxOperations; + private readonly runtime: SandboxRuntime; + + public constructor(ctx: DurableObjectState, env: ProjectSandboxEnv) { + super(ctx, env); + this.runtime = createSandboxRuntime(ctx, env); + const lifecycle = createLifecycleOps(this.runtime); + // TDZ invariant: collaborator factories must not invoke coordinated ops during construction. + let process: ProcessOps; + const coordinated = this.coordinatedProcessOps(() => process); + process = createProcessOps(this.runtime, coordinated); + const files = createFileOps(this.runtime); + const content = createContentOps(this.runtime, { + coordinatedProcess: coordinated, + deleteUploadedFileMetadata: files.deleteUploadedFileMetadata, + process, + sandboxRuntimeState: () => + this.withLease("sandboxRuntimeState", undefined, lifecycle.sandboxRuntimeState), + }); + this.operations = { ...lifecycle, ...process, ...files, ...content }; + } + + // Account deletion fences and drains active operations; taking a lease would deadlock. + public deleteAccountState(): Promise { + return this.operations.deleteAccountState(); + } + + public registerOwner(userId: string, sandboxName?: string): Promise { + return this.withLease("registerOwner", userId, () => + this.operations.registerOwner(userId, sandboxName), + ); + } + + public setQuotaPeriod(periodEndIso: string): Promise { + return this.withLease("setQuotaPeriod", periodEndIso, () => + this.operations.setQuotaPeriod(periodEndIso), + ); + } + + public beginRun(runId: string): Promise { + return this.withLease("beginRun", runId, () => this.operations.beginRun(runId)); + } + + public renewRun(runId: string): Promise { + return this.withLease("renewRun", runId, () => this.operations.renewRun(runId)); + } + + public endRun(runId: string): Promise { + return this.withLease("endRun", runId, () => this.operations.endRun(runId)); + } + + public override alarm(): Promise { + return this.withLease("alarm", undefined, this.operations.alarm); + } + + public runtimeSandboxId(): Promise { + return this.withLease("runtimeSandboxId", undefined, this.operations.runtimeSandboxId); + } + + public existingDaytonaId(): Promise { + return this.withLease("existingDaytonaId", undefined, this.operations.existingDaytonaId); + } + + public sandboxRuntimeState(): Promise { + return this.withLease("sandboxRuntimeState", undefined, this.operations.sandboxRuntimeState); + } + + public ensureReady(): Promise { + return this.withLease("ensureReady", undefined, this.operations.ensureReady); + } + + public getStatus(): Promise { + return this.withLease("getStatus", undefined, this.operations.getStatus); + } + + public runCode(input: ProjectRunCodeInput): Promise { + return this.withLease("runCode", input, () => this.operations.runCode(input)); + } + + public exec(input: ProjectExecInput): Promise { + return this.withLease("exec", input, () => this.operations.exec(input)); + } + + public startProcess(input: ProjectStartProcessInput): Promise { + return this.withLease("startProcess", input, () => this.operations.startProcess(input)); + } + + public allocateProjectPort(input: ProjectAllocatePortInput): Promise { + return this.withLease("allocateProjectPort", input, () => + this.operations.allocateProjectPort(input), + ); + } + + public allocateProcessPort(input: ProjectAllocateProcessPortInput): Promise { + return this.withLease("allocateProcessPort", input, () => + this.operations.allocateProcessPort(input), + ); + } + + public killAllProcesses(): Promise { + return this.withLease("killAllProcesses", undefined, this.operations.killAllProcesses); + } + + public killProcess(input: ProjectKillProcessInput): Promise { + return this.withLease("killProcess", input, () => this.operations.killProcess(input)); + } + + public readDevServerLogs(input: ProjectReadDevServerLogsInput): Promise { + return this.withLease("readDevServerLogs", input, () => + this.operations.readDevServerLogs(input), + ); + } + + public downloadProjectArchive(input: ProjectArchiveInput): Promise { + return this.runtime.lease.withStreamingOperation( + workspaceScope("downloadProjectArchive", input), + (release) => this.operations.downloadProjectArchive(input, release), + ); + } + + public readFile(input: ProjectReadFileInput): Promise { + return this.withLease("readFile", input, () => this.operations.readFile(input)); + } + + public listUploadedFiles( + input: ProjectListUploadedFilesInput, + ): Promise<{ files: ProjectFile[] }> { + return this.withLease("listUploadedFiles", input, () => + this.operations.listUploadedFiles(input), + ); + } + + public uploadProjectFile(input: ProjectUploadFileInput): Promise { + return this.withLease("uploadProjectFile", input, () => + this.operations.uploadProjectFile(input), + ); + } + + public restoreUploadedFiles( + input: ProjectRestoreUploadedFilesInput, + ): Promise<{ restoredFileCount: number }> { + return this.withLease("restoreUploadedFiles", input, () => + this.operations.restoreUploadedFiles(input), + ); + } + + public writeFile(input: ProjectWriteFileInput): Promise { + return this.withLease("writeFile", input, () => this.operations.writeFile(input)); + } + + public listFiles(input: ProjectListFilesInput): Promise { + return this.withLease("listFiles", input, () => this.operations.listFiles(input)); + } + + public searchFiles(input: ProjectSearchFilesInput): Promise { + return this.withLease("searchFiles", input, () => this.operations.searchFiles(input)); + } + + public deleteFile(input: ProjectDeleteFileInput): Promise { + return this.withLease("deleteFile", input, () => this.operations.deleteFile(input)); + } + + public getSignedPreviewUrl( + input: ProjectSignedPreviewUrlInput, + ): Promise<{ token: string; url: string }> { + return this.withLease("getSignedPreviewUrl", input, () => + this.operations.getSignedPreviewUrl(input), + ); + } + + public exposeBrowserTakeover( + input: ProjectBrowserTakeoverInput, + ): Promise { + return this.withLease("exposeBrowserTakeover", input, () => + this.operations.exposeBrowserTakeover(input), + ); + } + + public stopBrowserTakeover(input: ProjectBrowserTakeoverStopInput): Promise { + return this.withLease("stopBrowserTakeover", input, () => + this.operations.stopBrowserTakeover(input), + ); + } + + public exposeCodeServer(input: ProjectCodeServerInput): Promise<{ + expiresAt: string; + port: number; + url: string; + workspacePath: string; + }> { + return this.withLease("exposeCodeServer", input, () => this.operations.exposeCodeServer(input)); + } + + public wakePreview(input: ProjectWakePreviewInput): Promise { + return this.withLease("wakePreview", input, () => this.operations.wakePreview(input)); + } + + public projectPreviewStatus( + input: ProjectPreviewStatusInput, + ): Promise<{ running: boolean; state: string }> { + return this.withLease("projectPreviewStatus", input, () => + this.operations.projectPreviewStatus(input), + ); + } + + public cleanupProjectWorkspace(input: ProjectCleanupWorkspaceInput): Promise { + return this.withLease("cleanupProjectWorkspace", input, () => + this.operations.cleanupProjectWorkspace(input), + ); + } + + private coordinatedProcessOps(process: () => ProcessOps): CoordinatedProcessOps { + return { + allocateProcessPort: (input) => + this.withLease("allocateProcessPort", input, () => process().allocateProcessPort(input)), + ensureReady: () => this.withLease("ensureReady", undefined, () => process().ensureReady()), + exec: (input) => this.withLease("exec", input, () => process().exec(input)), + killProcess: (input) => + this.withLease("killProcess", input, () => process().killProcess(input)), + runCode: (input) => this.withLease("runCode", input, () => process().runCode(input)), + startProcess: (input) => + this.withLease("startProcess", input, () => process().startProcess(input)), + }; + } private withLease( - method: Exclude, + method: Exclude, input: unknown, operation: () => Promise, ): Promise; private withLease( - method: Exclude, + method: Exclude, input: unknown, operation: () => Promise, ): Promise { const kind = leaseKind(method); + if (kind === "account-deletion-control" || kind === "streaming") { + throw new TypeError(`Lease kind "${kind}" must not route through withLease`); + } if (kind === "owner-registration") { if (typeof input !== "string") throw new TypeError("Expected string RPC argument"); - return this.withActiveOwnerRegistration(input, operation); + return this.runtime.lease.withOwnerRegistration(input, operation); } if (kind === "cleanup-signal") { - // The signal wrapper is void-typed; capture the value so the overload's - // Promise contract holds if a signal method ever returns one. let result: unknown; - return this.withActiveSandboxCleanupSignal(async () => { - result = await operation(); - }).then(() => result); + return this.runtime.lease + .withCleanupSignal(async () => { + result = await operation(); + return result; + }) + .then(() => result); } if (kind === "shared-workspace") { - return this.withActiveSharedWorkspaceMutation(operation); + return this.runtime.lease.withSharedWorkspaceMutation(operation); } if (kind === "project-cleanup") { - return this.withActiveProjectWorkspaceCleanup(operation); + return this.runtime.lease.withProjectCleanup(operation); } if (kind === "workspace") { - return this.withActiveProjectWorkspaceOperation(workspaceScope(method, input), operation); + return this.runtime.lease.withWorkspaceOperation(workspaceScope(method, input), operation); } - return this.withActiveSandboxOperation(operation); + return this.runtime.lease.withSandboxOperation(operation); } } diff --git a/apps/gateway-worker/src/durable-objects/quota-tracker.ts b/apps/gateway-worker/src/durable-objects/quota-tracker.ts index bd2d3454..c67eea80 100644 --- a/apps/gateway-worker/src/durable-objects/quota-tracker.ts +++ b/apps/gateway-worker/src/durable-objects/quota-tracker.ts @@ -1,404 +1,62 @@ import { DurableObject } from "cloudflare:workers"; -import { - type QuotaFeature, - QuotaFeatureSchema, - type QuotaHistoryResult, - QuotaHistoryResultSchema, - type QuotaSnapshotResult, - QuotaSnapshotResultSchema, - type QuotaTryConsumeResponse, - QuotaTryConsumeResponseSchema, - type QuotaUsageResponse, - QuotaUsageResponseSchema, +import { QuotaTrackerRuntime } from "@cheatcode/billing/quota-runtime"; +import type { + QuotaFeature, + QuotaHistoryResult, + QuotaSnapshotResult, + QuotaTryConsumeResponse, + QuotaUsageResponse, } from "@cheatcode/types/quota"; -import { ensureQuotaTrackerStorage, hasQuotaTrackerStorage } from "./quota-tracker-storage"; -import { nextGatewayDurableObjectAlarm, QUOTA_TRACKER_RETENTION_MS } from "./retention"; -interface CounterRow { - used: number; -} - -interface LimitRow { - feature: QuotaFeature; - limit_val: number; -} - -interface HistoryRow { - amount: number; - recorded_at: number; -} - -interface OperationRow { - allowed: number; - amount: number; - event_id: string; - feature: string; - limit_val: number; - operation: string; - period_key: string; - remaining: number; - used: number; -} - -interface QuotaOperationInput { - amount: number; - eventId: string; - feature: QuotaFeature; - operation: "record" | "try-consume"; - periodKey: string; -} - -function isCounterRow(value: unknown): value is CounterRow { - return isRecord(value) && typeof value["used"] === "number"; -} - -function isLimitRow(value: unknown): value is LimitRow { - return ( - isRecord(value) && - QuotaFeatureSchema.safeParse(value["feature"]).success && - typeof value["limit_val"] === "number" - ); -} +/** Gateway-owned Durable Object facade over the worker-only billing runtime. */ +export class QuotaTracker extends DurableObject { + private readonly runtime: QuotaTrackerRuntime; -function isHistoryRow(value: unknown): value is HistoryRow { - return ( - isRecord(value) && - typeof value["amount"] === "number" && - typeof value["recorded_at"] === "number" - ); -} - -function isOperationRow(value: unknown): value is OperationRow { - if (!isRecord(value)) { - return false; + public constructor(ctx: DurableObjectState, env: unknown) { + super(ctx, env); + this.runtime = new QuotaTrackerRuntime(ctx); } - return ( - typeof value["allowed"] === "number" && - typeof value["amount"] === "number" && - typeof value["event_id"] === "string" && - typeof value["feature"] === "string" && - typeof value["limit_val"] === "number" && - typeof value["operation"] === "string" && - typeof value["period_key"] === "string" && - typeof value["remaining"] === "number" && - typeof value["used"] === "number" - ); -} -function isRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null; -} - -export class QuotaTracker extends DurableObject { - private isStorageInitialized = false; - - public async tryConsume( + public tryConsume( feature: QuotaFeature, amount: number, periodEnd: Date, eventId: string, ): Promise { - this.ensureStorage(); - const periodKey = periodKeyFromDate(periodEnd); - const input: QuotaOperationInput = { - amount, - eventId, - feature, - operation: "try-consume", - periodKey, - }; - const result = this.ctx.storage.transactionSync(() => this.consumeOnce(input)); - await this.ensureCleanupAlarm(); - return result; + return this.runtime.tryConsume(feature, amount, periodEnd, eventId); } - public async peek(feature: QuotaFeature, periodEnd: Date): Promise { - this.ensureStorage(); - const limit = this.readLimit(feature); - const used = this.readUsed(feature, periodKeyFromDate(periodEnd)); - return QuotaUsageResponseSchema.parse({ - limit, - remaining: Math.max(0, limit - used), - used, - }); + public peek(feature: QuotaFeature, periodEnd: Date): Promise { + return this.runtime.peek(feature, periodEnd); } - public async record( + public record( feature: QuotaFeature, amount: number, periodEnd: Date, eventId: string, recordedAt: Date, ): Promise { - this.ensureStorage(); - const periodKey = periodKeyFromDate(periodEnd); - const input: QuotaOperationInput = { - amount, - eventId, - feature, - operation: "record", - periodKey, - }; - const result = this.ctx.storage.transactionSync(() => - this.recordOnce(input, recordedAt.getTime()), - ); - await this.ensureCleanupAlarm(); - return result; - } - - public async history(feature: QuotaFeature, from: Date): Promise { - this.ensureStorage(); - const events = this.ctx.storage.sql - .exec( - `SELECT SUM(amount) AS amount, - (recorded_at / 86400000) * 86400000 AS recorded_at - FROM usage_event - WHERE feature = ? AND recorded_at >= ? - GROUP BY recorded_at / 86400000 - ORDER BY recorded_at`, - feature, - from.getTime(), - ) - .toArray(); - return historyResult(events); - } - - public async setLimit( - feature: QuotaFeature, - limit: number, - entitlementVersion: number, - ): Promise { - this.ensureStorage(); - this.ctx.storage.sql.exec( - `INSERT INTO limit_override (feature, limit_val, entitlement_version) - VALUES (?, ?, ?) - ON CONFLICT(feature) DO UPDATE SET - limit_val = excluded.limit_val, - entitlement_version = excluded.entitlement_version - WHERE excluded.entitlement_version >= limit_override.entitlement_version`, - feature, - limit, - entitlementVersion, - ); - } - - public async deleteAllState(): Promise { - await this.ctx.storage.deleteAll(); - this.isStorageInitialized = false; + return this.runtime.record(feature, amount, periodEnd, eventId, recordedAt); } - public async snapshot(periodEnd: Date): Promise { - this.ensureStorage(); - const periodKey = periodKeyFromDate(periodEnd); - const rawRows = this.ctx.storage.sql - .exec("SELECT feature, limit_val FROM limit_override ORDER BY feature") - .toArray(); - const rows: LimitRow[] = []; - for (const row of rawRows) { - if (isLimitRow(row)) { - rows.push(row); - } - } - const snapshot: QuotaSnapshotResult = {}; - for (const row of rows) { - snapshot[row.feature] = { - limit: row.limit_val, - used: this.readUsed(row.feature, periodKey), - }; - } - return QuotaSnapshotResultSchema.parse(snapshot); + public history(feature: QuotaFeature, from: Date): Promise { + return this.runtime.history(feature, from); } - public override async alarm(): Promise { - if (!hasQuotaTrackerStorage(this.ctx)) { - await this.ctx.storage.deleteAlarm(); - return; - } - this.ensureStorage(); - this.ctx.storage.sql.exec( - "DELETE FROM counter WHERE updated_at < ?", - Date.now() - QUOTA_TRACKER_RETENTION_MS, - ); - this.ctx.storage.sql.exec( - "DELETE FROM usage_event WHERE recorded_at < ?", - Date.now() - QUOTA_TRACKER_RETENTION_MS, - ); - this.ctx.storage.sql.exec( - "DELETE FROM quota_operation WHERE recorded_at < ?", - Date.now() - QUOTA_TRACKER_RETENTION_MS, - ); - await this.refreshCleanupAlarm(); + public setLimit(feature: QuotaFeature, limit: number, entitlementVersion: number): Promise { + return this.runtime.setLimit(feature, limit, entitlementVersion); } - private ensureStorage(): void { - if (this.isStorageInitialized) { - return; - } - ensureQuotaTrackerStorage(this.ctx); - this.isStorageInitialized = true; + public deleteAllState(): Promise { + return this.runtime.deleteAllState(); } - private consumeOnce(input: QuotaOperationInput): QuotaTryConsumeResponse { - const existing = this.readOperation(input); - if (existing) { - return operationConsumeResult(existing); - } - const limit = this.readLimit(input.feature); - const used = this.readUsed(input.feature, input.periodKey); - const allowed = used + input.amount <= limit; - const nextUsed = allowed ? used + input.amount : used; - const remaining = Math.max(0, limit - nextUsed); - if (allowed) { - this.writeUsage(input.feature, input.periodKey, input.amount, nextUsed, Date.now()); - } - this.insertOperation(input, { allowed, limit, remaining, used: nextUsed }); - return QuotaTryConsumeResponseSchema.parse({ allowed, limit, remaining }); + public snapshot(periodEnd: Date): Promise { + return this.runtime.snapshot(periodEnd); } - private recordOnce(input: QuotaOperationInput, recordedAt: number): QuotaUsageResponse { - const existing = this.readOperation(input); - if (existing) { - return operationPeekResult(existing); - } - const nextUsed = this.readUsed(input.feature, input.periodKey) + input.amount; - const limit = this.readLimit(input.feature); - const remaining = Math.max(0, limit - nextUsed); - this.writeUsage(input.feature, input.periodKey, input.amount, nextUsed, recordedAt); - this.insertOperation(input, { allowed: true, limit, remaining, used: nextUsed }); - return QuotaUsageResponseSchema.parse({ limit, remaining, used: nextUsed }); + public override alarm(): Promise { + return this.runtime.alarm(); } - - private readOperation(input: QuotaOperationInput): OperationRow | null { - const [row] = this.ctx.storage.sql - .exec("SELECT * FROM quota_operation WHERE event_id = ?", input.eventId) - .toArray(); - if (!isOperationRow(row)) { - return null; - } - if ( - row.operation !== input.operation || - row.feature !== input.feature || - row.period_key !== input.periodKey || - row.amount !== input.amount - ) { - throw new Error("Quota event id was reused with different operation data."); - } - return row; - } - - private insertOperation( - input: QuotaOperationInput, - result: { allowed: boolean; limit: number; remaining: number; used: number }, - ): void { - this.ctx.storage.sql.exec( - `INSERT INTO quota_operation - (event_id, operation, feature, period_key, amount, allowed, limit_val, remaining, used, recorded_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - input.eventId, - input.operation, - input.feature, - input.periodKey, - input.amount, - result.allowed ? 1 : 0, - result.limit, - result.remaining, - result.used, - Date.now(), - ); - } - - private async ensureCleanupAlarm(): Promise { - const currentAlarm = await this.ctx.storage.getAlarm(); - if (currentAlarm === null) { - await this.ctx.storage.setAlarm(nextGatewayDurableObjectAlarm(Date.now())); - } - } - - private async refreshCleanupAlarm(): Promise { - if (this.hasRetainedUsage()) { - await this.ensureCleanupAlarm(); - return; - } - await this.ctx.storage.deleteAlarm(); - } - - private hasRetainedUsage(): boolean { - for (const table of ["counter", "usage_event", "quota_operation"] as const) { - const [row] = this.ctx.storage.sql - .exec(`SELECT 1 AS present FROM ${table} LIMIT 1`) - .toArray(); - if (isRecord(row) && row["present"] === 1) { - return true; - } - } - return false; - } - - private readLimit(feature: QuotaFeature): number { - const [rawRow] = this.ctx.storage.sql - .exec("SELECT limit_val FROM limit_override WHERE feature = ?", feature) - .toArray(); - return isRecord(rawRow) && typeof rawRow["limit_val"] === "number" ? rawRow["limit_val"] : 0; - } - - private readUsed(feature: QuotaFeature, periodKey: string): number { - const [rawRow] = this.ctx.storage.sql - .exec("SELECT used FROM counter WHERE feature = ? AND period_key = ?", feature, periodKey) - .toArray(); - return isCounterRow(rawRow) ? rawRow.used : 0; - } - - private writeUsage( - feature: QuotaFeature, - periodKey: string, - amount: number, - used: number, - recordedAt: number, - ): void { - this.ctx.storage.sql.exec( - `INSERT INTO counter (feature, period_key, used, updated_at) - VALUES (?, ?, ?, ?) - ON CONFLICT(feature, period_key) DO UPDATE SET - used = excluded.used, - updated_at = excluded.updated_at`, - feature, - periodKey, - used, - Date.now(), - ); - this.ctx.storage.sql.exec( - "INSERT INTO usage_event (feature, amount, recorded_at) VALUES (?, ?, ?)", - feature, - amount, - recordedAt, - ); - } -} - -function operationConsumeResult(row: OperationRow): QuotaTryConsumeResponse { - return QuotaTryConsumeResponseSchema.parse({ - allowed: row.allowed === 1, - limit: row.limit_val, - remaining: row.remaining, - }); -} - -function operationPeekResult(row: OperationRow): QuotaUsageResponse { - return QuotaUsageResponseSchema.parse({ - limit: row.limit_val, - remaining: row.remaining, - used: row.used, - }); -} - -function historyResult(rows: unknown[]): QuotaHistoryResult { - return QuotaHistoryResultSchema.parse( - rows.filter(isHistoryRow).map((row) => ({ amount: row.amount, recordedAt: row.recorded_at })), - ); -} - -function periodKeyFromDate(date: Date): string { - const year = date.getUTCFullYear(); - const month = String(date.getUTCMonth() + 1).padStart(2, "0"); - return `${year}-${month}`; } diff --git a/apps/gateway-worker/src/durable-objects/retention.ts b/apps/gateway-worker/src/durable-objects/retention.ts index 664e7b92..15ceb3ec 100644 --- a/apps/gateway-worker/src/durable-objects/retention.ts +++ b/apps/gateway-worker/src/durable-objects/retention.ts @@ -1,7 +1,6 @@ const DAY_MS = 86_400_000; export const RATE_LIMITER_RETENTION_MS = DAY_MS; -export const QUOTA_TRACKER_RETENTION_MS = 366 * DAY_MS; export function nextGatewayDurableObjectAlarm(nowMs: number): number { return nowMs + DAY_MS; diff --git a/apps/web/src/components/chat/use-chat-panel-controller.ts b/apps/web/src/components/chat/use-chat-panel-controller.ts index c4ac029e..80ce07b9 100644 --- a/apps/web/src/components/chat/use-chat-panel-controller.ts +++ b/apps/web/src/components/chat/use-chat-panel-controller.ts @@ -26,10 +26,11 @@ import { import { useChatSubmission } from "@/components/chat/use-chat-submission"; import type { OlderMessagesLoadResult } from "@/components/chat/use-message-list-scroll"; import { - applySandboxStatus, - isBrowserToolName, type SandboxStatusActions, useSandboxSurfaceSync, + useWorkspaceSurfaceApplier, + type WorkspaceSurfaceApplier, + workspaceSurfaceEffect, } from "@/components/chat/use-sandbox-surface-sync"; import { agentModelRequestValue } from "@/lib/agent-models"; import { cancelRun, getThread } from "@/lib/api/project-thread"; @@ -66,7 +67,14 @@ export function useChatPanelController(input: ChatPanelProps) { function useChatPanelRuntime(input: ChatPanelProps) { const store = useChatPanelStore(input.threadId); - const runtime = useChatRuntimeBase(input, store); + const surfaceApplier = useWorkspaceSurfaceApplier({ + projectId: input.project?.id ?? null, + setActivePreviewTab: store.setActivePreviewTab, + setPreviewPanelOpen: store.setPreviewPanelOpen, + setSandboxStatus: store.setSandboxStatus, + threadId: input.threadId, + }); + const runtime = useChatRuntimeBase(input, store, surfaceApplier); const resumeStream = useSerializedResume(runtime.chat.resumeStream); const messages = useDeferredValue(runtime.chat.messages); const loadOlderMessages = useOlderMessageLoader( @@ -80,11 +88,16 @@ function useChatPanelRuntime(input: ChatPanelProps) { pendingSubmissionRef: runtime.pendingSubmissionRef, queryClient: runtime.queryClient, resumeStream, + surfaceApplier, }); return { ...runtime, loadOlderMessages, messages, store }; } -function useChatRuntimeBase(input: ChatPanelProps, store: ReturnType) { +function useChatRuntimeBase( + input: ChatPanelProps, + store: ReturnType, + surfaceApplier: WorkspaceSurfaceApplier, +) { const { getToken } = useAuth(); const router = useRouter(); const queryClient = useQueryClient(); @@ -105,6 +118,7 @@ function useChatRuntimeBase(input: ChatPanelProps, store: ReturnType; sandboxActions: SandboxStatusActions; + surfaceApplier: WorkspaceSurfaceApplier; threadId: string; transport: ReturnType; }) { @@ -319,11 +334,9 @@ function handleStreamData( if (part.type === "data-seq") { handleSequenceData(part.data, input.threadId); } - if (part.type === "data-sandbox-status") { - handleSandboxStatusData(part.data, input.sandboxActions); - } - if (part.type === "data-tool") { - handleToolData(part.data, input.sandboxActions); + const surfaceCommand = workspaceSurfaceEffect(part); + if (surfaceCommand) { + input.surfaceApplier.apply(surfaceCommand); } if (part.type === "data-project-created") { handleProjectCreatedData(part.data, input); @@ -340,21 +353,6 @@ function handleSequenceData(data: unknown, threadId: string): void { } } -function handleSandboxStatusData(data: unknown, actions: SandboxStatusActions): void { - const parsed = CHEATCODE_DATA_SCHEMAS["sandbox-status"].safeParse(data); - if (parsed.success) { - applySandboxStatus(parsed.data, actions); - } -} - -function handleToolData(data: unknown, actions: SandboxStatusActions): void { - const parsed = CHEATCODE_DATA_SCHEMAS.tool.safeParse(data); - if (parsed.success && isBrowserToolName(parsed.data.toolName)) { - actions.setActivePreviewTab("app"); - actions.setPreviewPanelOpen(true); - } -} - function handleProjectCreatedData( data: unknown, input: Parameters[0], @@ -460,6 +458,7 @@ function useChatPanelEffects( pendingSubmissionRef: { current: PendingSubmission | null }; queryClient: ReturnType; resumeStream: () => Promise; + surfaceApplier: WorkspaceSurfaceApplier; }, ): void { useSandboxSurfaceSync({ @@ -473,6 +472,7 @@ function useChatPanelEffects( setPreviewPanelOpen: store.setPreviewPanelOpen, setPreviewUrl: store.setPreviewUrl, setSandboxStatus: store.setSandboxStatus, + surfaceApplier: shared.surfaceApplier, }); useConnectionStateSync(); useVisibleStreamResume({ diff --git a/apps/web/src/components/chat/use-sandbox-surface-sync.ts b/apps/web/src/components/chat/use-sandbox-surface-sync.ts index c725b060..8eb1efe3 100644 --- a/apps/web/src/components/chat/use-sandbox-surface-sync.ts +++ b/apps/web/src/components/chat/use-sandbox-surface-sync.ts @@ -1,11 +1,12 @@ import { + CHEATCODE_DATA_SCHEMAS, type CheatcodeUIMessage, type ProjectSummary, reconstructedTranscriptUIMessage, type SandboxState, } from "@cheatcode/types"; import type { ChatStatus } from "ai"; -import { useEffect, useMemo, useRef } from "react"; +import { useCallback, useEffect, useMemo, useRef } from "react"; import { type PreviewTab, useAppStore } from "@/lib/store/app-store"; type SandboxStatusData = Extract< @@ -28,14 +29,63 @@ interface SandboxSurfaceSyncInput extends SandboxStatusActions { setExpoUrl: (url: null | string) => void; setPreviewPanelOpen: (open: boolean) => void; setPreviewUrl: (url: null | string) => void; + surfaceApplier: WorkspaceSurfaceApplier; +} + +export type SurfaceCommand = + | { status: SandboxStatusData["status"]; type: "status" } + | { toolCallId: string; type: "open-browser-preview" }; + +export interface WorkspaceSurfaceApplier { + apply: (command: SurfaceCommand) => void; + reset: () => void; +} + +interface WorkspaceSurfaceApplierInput extends SandboxStatusActions { + projectId: string | null; + threadId: string; +} + +interface SurfaceCommandState { + openedBrowserToolKeys: Set; + scopeKey: string; +} + +interface HydratedSurfaceCommands { + browser: SurfaceCommand | null; + status: SurfaceCommand | null; +} + +export function useWorkspaceSurfaceApplier( + input: WorkspaceSurfaceApplierInput, +): WorkspaceSurfaceApplier { + const scopeKey = `${input.threadId}:${input.projectId ?? ""}`; + const stateRef = useRef(createSurfaceCommandState(scopeKey)); + if (stateRef.current.scopeKey !== scopeKey) { + stateRef.current = createSurfaceCommandState(scopeKey); + } + const actions = useMemo( + () => ({ + setActivePreviewTab: input.setActivePreviewTab, + setPreviewPanelOpen: input.setPreviewPanelOpen, + setSandboxStatus: input.setSandboxStatus, + }), + [input.setActivePreviewTab, input.setPreviewPanelOpen, input.setSandboxStatus], + ); + const apply = useCallback( + (command: SurfaceCommand) => + applyWorkspaceSurfaceCommand(command, input.threadId, stateRef.current, actions), + [actions, input.threadId], + ); + const reset = useCallback(() => { + stateRef.current = createSurfaceCommandState(scopeKey); + }, [scopeKey]); + return useMemo(() => ({ apply, reset }), [apply, reset]); } export function useSandboxSurfaceSync(input: SandboxSurfaceSyncInput): void { - const latestStatus = latestSandboxStatusFromMessages(input.messages); - const browserActivityKey = latestBrowserActivityKeyFromMessages(input.messages); - const status = latestStatus?.status ?? null; - const appliedSnapshotRef = useRef(undefined); - const openedBrowserActivityRef = useRef(null); + const commands = useMemo(() => hydratedSurfaceCommands(input.messages), [input.messages]); + const status = commands.status?.type === "status" ? commands.status.status : null; const defaultedProjectFilesRef = useRef(null); const previousStatusRef = useRef(input.chatStatus); const actions = useMemo( @@ -48,58 +98,58 @@ export function useSandboxSurfaceSync(input: SandboxSurfaceSyncInput): void { ); useResetSandboxSurface(input); - useMessageStatusSync(status, actions, appliedSnapshotRef); + useHydratedSurfaceCommand(commands.status, input.surfaceApplier); useProjectFilesDefault(input.project, status, actions, defaultedProjectFilesRef); - useBrowserActivityDefault(browserActivityKey, actions, openedBrowserActivityRef); + useHydratedSurfaceCommand(commands.browser, input.surfaceApplier); useCompletionPreview(input.chatStatus, previousStatusRef); } -export function applySandboxStatus(data: SandboxStatusData, actions: SandboxStatusActions): void { - actions.setSandboxStatus(data.status); -} - -export function isBrowserToolName(toolName: string): boolean { - return ( - toolName === "browser_act" || - toolName === "browser_extract" || - toolName === "browser_observe" || - toolName === "browser_open" || - toolName === "browser_screenshot" - ); +export function workspaceSurfaceEffect(part: unknown): SurfaceCommand | null { + if (!isRecord(part)) { + return null; + } + if (part["type"] === "data-sandbox-status") { + const parsed = CHEATCODE_DATA_SCHEMAS["sandbox-status"].safeParse(part["data"]); + return parsed.success ? { status: parsed.data.status, type: "status" } : null; + } + if (part["type"] !== "data-tool") { + return null; + } + const parsed = CHEATCODE_DATA_SCHEMAS.tool.safeParse(part["data"]); + return parsed.success && isBrowserToolName(parsed.data.toolName) + ? { toolCallId: parsed.data.toolCallId, type: "open-browser-preview" } + : null; } function useResetSandboxSurface(input: SandboxSurfaceSyncInput): void { useEffect(() => { + input.surfaceApplier.reset(); input.resetConsole(); input.resetPreviewNavigation(); input.setPreviewUrl(null); input.setExpoUrl(null); input.setPreviewPanelOpen(false); + input.setSandboxStatus("cold"); }, [ input.resetConsole, input.resetPreviewNavigation, input.setExpoUrl, input.setPreviewPanelOpen, input.setPreviewUrl, + input.setSandboxStatus, + input.surfaceApplier, ]); } -function useMessageStatusSync( - status: null | SandboxStatusData["status"], - actions: SandboxStatusActions, - appliedSnapshotRef: { current: string | null | undefined }, +function useHydratedSurfaceCommand( + command: SurfaceCommand | null, + applier: WorkspaceSurfaceApplier, ): void { useEffect(() => { - if (appliedSnapshotRef.current === status) { - return; - } - appliedSnapshotRef.current = status; - if (!status) { - actions.setSandboxStatus("cold"); - return; + if (command) { + applier.apply(command); } - applySandboxStatus({ v: 1, status }, actions); - }, [actions, appliedSnapshotRef, status]); + }, [applier, command]); } function useProjectFilesDefault( @@ -117,21 +167,6 @@ function useProjectFilesDefault( }, [actions, defaultedProjectFilesRef, project, status]); } -function useBrowserActivityDefault( - activityKey: string | null, - actions: SandboxStatusActions, - openedBrowserActivityRef: { current: string | null }, -): void { - useEffect(() => { - if (!activityKey || openedBrowserActivityRef.current === activityKey) { - return; - } - openedBrowserActivityRef.current = activityKey; - actions.setActivePreviewTab("app"); - actions.setPreviewPanelOpen(true); - }, [actions, activityKey, openedBrowserActivityRef]); -} - function useCompletionPreview( status: ChatStatus, previousStatusRef: { current: ChatStatus }, @@ -150,40 +185,66 @@ function useCompletionPreview( }, [previousStatusRef, status]); } -function latestSandboxStatusFromMessages( - messages: readonly CheatcodeUIMessage[], -): SandboxStatusData | null { - for (let messageIndex = messages.length - 1; messageIndex >= 0; messageIndex -= 1) { +function hydratedSurfaceCommands(messages: readonly CheatcodeUIMessage[]): HydratedSurfaceCommands { + let browser: SurfaceCommand | null = null; + let status: SurfaceCommand | null = null; + for (let messageIndex = 0; messageIndex < messages.length; messageIndex += 1) { const message = messages[messageIndex]; if (!message) { continue; } const parts = reconstructedTranscriptUIMessage(message).parts; - for (let partIndex = parts.length - 1; partIndex >= 0; partIndex -= 1) { + for (let partIndex = 0; partIndex < parts.length; partIndex += 1) { const part = parts[partIndex]; - if (part?.type === "data-sandbox-status") { - return part.data; + if (!part) { + continue; + } + const command = workspaceSurfaceEffect(part); + if (command?.type === "status") { + status = command; + } else if (command?.type === "open-browser-preview") { + browser = command; } } } - return null; + return { browser, status }; } -function latestBrowserActivityKeyFromMessages( - messages: readonly CheatcodeUIMessage[], -): string | null { - for (let messageIndex = messages.length - 1; messageIndex >= 0; messageIndex -= 1) { - const message = messages[messageIndex]; - if (!message) { - continue; - } - const parts = reconstructedTranscriptUIMessage(message).parts; - for (let partIndex = parts.length - 1; partIndex >= 0; partIndex -= 1) { - const part = parts[partIndex]; - if (part?.type === "data-tool" && isBrowserToolName(part.data.toolName)) { - return `${message.id}:${part.data.toolCallId}`; - } +function applyWorkspaceSurfaceCommand( + command: SurfaceCommand, + threadId: string, + state: SurfaceCommandState, + actions: SandboxStatusActions, +): void { + if (command.type === "status") { + if (useAppStore.getState().sandboxStatus !== command.status) { + actions.setSandboxStatus(command.status); } + return; } - return null; + const onceKey = `${threadId}:${command.toolCallId}`; + if (state.openedBrowserToolKeys.has(onceKey)) { + return; + } + state.openedBrowserToolKeys.add(onceKey); + actions.setActivePreviewTab("app"); + actions.setPreviewPanelOpen(true); +} + +function createSurfaceCommandState(scopeKey: string): SurfaceCommandState { + return { openedBrowserToolKeys: new Set(), scopeKey }; +} + +function isBrowserToolName(toolName: string): boolean { + return ( + toolName === "browser_act" || + toolName === "browser_extract" || + toolName === "browser_observe" || + toolName === "browser_open" || + toolName === "browser_screenshot" + ); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; } diff --git a/packages/billing/README.md b/packages/billing/README.md index 086c6f5d..a5da8844 100644 --- a/packages/billing/README.md +++ b/packages/billing/README.md @@ -16,6 +16,9 @@ plan catalog, and resource-entitlement helpers. - `entitlementValuesForTier` - `PLAN_CATALOG` - sandbox-hour quota helpers +- `@cheatcode/billing/quota-runtime`: worker-only `QuotaTrackerRuntime`; this + subpath owns quota storage, retention, and RPC input validation and is not + re-exported from the Node-safe package root Current tiers are `free`, `pro`, `premium`, `ultra`, and `max`. Entitlements cover sandbox hours, active projects, BYOK provider slots, and Composio calls. diff --git a/packages/billing/package.json b/packages/billing/package.json index 9cbcc96f..5ba7bea3 100644 --- a/packages/billing/package.json +++ b/packages/billing/package.json @@ -9,6 +9,10 @@ ".": { "types": "./dist/index.d.ts", "import": "./dist/index.js" + }, + "./quota-runtime": { + "types": "./dist/quota-runtime.d.ts", + "import": "./dist/quota-runtime.js" } }, "scripts": { @@ -17,6 +21,7 @@ "typecheck": "tsc -p tsconfig.json --noEmit" }, "dependencies": { + "@cheatcode/durable-storage": "workspace:*", "@cheatcode/observability": "workspace:*", "@cheatcode/types": "workspace:*", "@polar-sh/sdk": "catalog:", @@ -24,6 +29,7 @@ }, "devDependencies": { "@cheatcode/tsconfig": "workspace:*", + "@cloudflare/workers-types": "catalog:", "typescript": "catalog:" } } diff --git a/packages/billing/src/quota-runtime-retention.ts b/packages/billing/src/quota-runtime-retention.ts new file mode 100644 index 00000000..0a7a2622 --- /dev/null +++ b/packages/billing/src/quota-runtime-retention.ts @@ -0,0 +1,7 @@ +const DAY_MS = 86_400_000; + +export const QUOTA_TRACKER_RETENTION_MS = 366 * DAY_MS; + +export function nextQuotaTrackerAlarm(nowMs: number): number { + return nowMs + DAY_MS; +} diff --git a/apps/gateway-worker/src/durable-objects/quota-tracker-storage.ts b/packages/billing/src/quota-runtime-storage.ts similarity index 100% rename from apps/gateway-worker/src/durable-objects/quota-tracker-storage.ts rename to packages/billing/src/quota-runtime-storage.ts diff --git a/packages/billing/src/quota-runtime.ts b/packages/billing/src/quota-runtime.ts new file mode 100644 index 00000000..38a6c62c --- /dev/null +++ b/packages/billing/src/quota-runtime.ts @@ -0,0 +1,483 @@ +import { APIError } from "@cheatcode/observability"; +import { + type QuotaFeature, + QuotaFeatureSchema, + type QuotaHistoryResult, + QuotaHistoryResultSchema, + type QuotaSnapshotResult, + QuotaSnapshotResultSchema, + type QuotaTryConsumeResponse, + QuotaTryConsumeResponseSchema, + type QuotaUsageResponse, + QuotaUsageResponseSchema, +} from "@cheatcode/types/quota"; +import { z } from "zod"; +import { nextQuotaTrackerAlarm, QUOTA_TRACKER_RETENTION_MS } from "./quota-runtime-retention"; +import { ensureQuotaTrackerStorage, hasQuotaTrackerStorage } from "./quota-runtime-storage"; + +const QuotaDateSchema = z.date(); +const QuotaAmountSchema = z.number().finite().positive(); +const QuotaEventIdSchema = z.string().min(1).max(200); +const QuotaLimitSchema = z.number().finite().nonnegative(); +const EntitlementVersionSchema = z.number().int().nonnegative(); + +const FeatureAndPeriodSchema = z + .object({ + feature: QuotaFeatureSchema, + periodEnd: QuotaDateSchema, + }) + .strict(); + +const QuotaOperationSchema = FeatureAndPeriodSchema.extend({ + amount: QuotaAmountSchema, + eventId: QuotaEventIdSchema, +}).strict(); + +const QuotaRecordSchema = QuotaOperationSchema.extend({ + recordedAt: QuotaDateSchema, +}).strict(); + +const QuotaHistorySchema = z + .object({ + feature: QuotaFeatureSchema, + from: QuotaDateSchema, + }) + .strict(); + +const QuotaLimitInputSchema = z + .object({ + entitlementVersion: EntitlementVersionSchema, + feature: QuotaFeatureSchema, + limit: QuotaLimitSchema, + }) + .strict(); + +interface CounterRow { + used: number; +} + +interface LimitRow { + feature: QuotaFeature; + limit_val: number; +} + +interface HistoryRow { + amount: number; + recorded_at: number; +} + +interface OperationRow { + allowed: number; + amount: number; + event_id: string; + feature: string; + limit_val: number; + operation: string; + period_key: string; + remaining: number; + used: number; +} + +interface QuotaOperationInput { + amount: number; + eventId: string; + feature: QuotaFeature; + operation: "record" | "try-consume"; + periodKey: string; +} + +function isCounterRow(value: unknown): value is CounterRow { + return isRecord(value) && typeof value["used"] === "number"; +} + +function isLimitRow(value: unknown): value is LimitRow { + return ( + isRecord(value) && + QuotaFeatureSchema.safeParse(value["feature"]).success && + typeof value["limit_val"] === "number" + ); +} + +function isHistoryRow(value: unknown): value is HistoryRow { + return ( + isRecord(value) && + typeof value["amount"] === "number" && + typeof value["recorded_at"] === "number" + ); +} + +function isOperationRow(value: unknown): value is OperationRow { + if (!isRecord(value)) { + return false; + } + return ( + typeof value["allowed"] === "number" && + typeof value["amount"] === "number" && + typeof value["event_id"] === "string" && + typeof value["feature"] === "string" && + typeof value["limit_val"] === "number" && + typeof value["operation"] === "string" && + typeof value["period_key"] === "string" && + typeof value["remaining"] === "number" && + typeof value["used"] === "number" + ); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +/** + * Worker-only QuotaTracker implementation. The Durable Object facade delegates + * every public operation here so validation and storage ownership cannot drift. + */ +export class QuotaTrackerRuntime { + private isStorageInitialized = false; + + public constructor(private readonly ctx: DurableObjectState) {} + + public async tryConsume( + feature: QuotaFeature, + amount: number, + periodEnd: Date, + eventId: string, + ): Promise { + const parsed = parseQuotaInput( + QuotaOperationSchema, + { amount, eventId, feature, periodEnd }, + "tryConsume", + ); + this.ensureStorage(); + const input: QuotaOperationInput = { + amount: parsed.amount, + eventId: parsed.eventId, + feature: parsed.feature, + operation: "try-consume", + periodKey: periodKeyFromDate(parsed.periodEnd), + }; + const result = this.ctx.storage.transactionSync(() => this.consumeOnce(input)); + await this.ensureCleanupAlarm(); + return result; + } + + public async peek(feature: QuotaFeature, periodEnd: Date): Promise { + const parsed = parseQuotaInput(FeatureAndPeriodSchema, { feature, periodEnd }, "peek"); + this.ensureStorage(); + const limit = this.readLimit(parsed.feature); + const used = this.readUsed(parsed.feature, periodKeyFromDate(parsed.periodEnd)); + return QuotaUsageResponseSchema.parse({ + limit, + remaining: Math.max(0, limit - used), + used, + }); + } + + public async record( + feature: QuotaFeature, + amount: number, + periodEnd: Date, + eventId: string, + recordedAt: Date, + ): Promise { + const parsed = parseQuotaInput( + QuotaRecordSchema, + { amount, eventId, feature, periodEnd, recordedAt }, + "record", + ); + this.ensureStorage(); + const input: QuotaOperationInput = { + amount: parsed.amount, + eventId: parsed.eventId, + feature: parsed.feature, + operation: "record", + periodKey: periodKeyFromDate(parsed.periodEnd), + }; + const result = this.ctx.storage.transactionSync(() => + this.recordOnce(input, parsed.recordedAt.getTime()), + ); + await this.ensureCleanupAlarm(); + return result; + } + + public async history(feature: QuotaFeature, from: Date): Promise { + const parsed = parseQuotaInput(QuotaHistorySchema, { feature, from }, "history"); + this.ensureStorage(); + const events = this.ctx.storage.sql + .exec( + `SELECT SUM(amount) AS amount, + (recorded_at / 86400000) * 86400000 AS recorded_at + FROM usage_event + WHERE feature = ? AND recorded_at >= ? + GROUP BY recorded_at / 86400000 + ORDER BY recorded_at`, + parsed.feature, + parsed.from.getTime(), + ) + .toArray(); + return historyResult(events); + } + + public async setLimit( + feature: QuotaFeature, + limit: number, + entitlementVersion: number, + ): Promise { + const parsed = parseQuotaInput( + QuotaLimitInputSchema, + { entitlementVersion, feature, limit }, + "setLimit", + ); + this.ensureStorage(); + this.ctx.storage.sql.exec( + `INSERT INTO limit_override (feature, limit_val, entitlement_version) + VALUES (?, ?, ?) + ON CONFLICT(feature) DO UPDATE SET + limit_val = excluded.limit_val, + entitlement_version = excluded.entitlement_version + WHERE excluded.entitlement_version >= limit_override.entitlement_version`, + parsed.feature, + parsed.limit, + parsed.entitlementVersion, + ); + } + + public async deleteAllState(): Promise { + await this.ctx.storage.deleteAll(); + this.isStorageInitialized = false; + } + + public async snapshot(periodEnd: Date): Promise { + const parsed = parseQuotaInput( + z.object({ periodEnd: QuotaDateSchema }).strict(), + { + periodEnd, + }, + "snapshot", + ); + this.ensureStorage(); + const periodKey = periodKeyFromDate(parsed.periodEnd); + const rawRows = this.ctx.storage.sql + .exec("SELECT feature, limit_val FROM limit_override ORDER BY feature") + .toArray(); + const rows: LimitRow[] = []; + for (const row of rawRows) { + if (isLimitRow(row)) { + rows.push(row); + } + } + const snapshot: QuotaSnapshotResult = {}; + for (const row of rows) { + snapshot[row.feature] = { + limit: row.limit_val, + used: this.readUsed(row.feature, periodKey), + }; + } + return QuotaSnapshotResultSchema.parse(snapshot); + } + + public async alarm(): Promise { + if (!hasQuotaTrackerStorage(this.ctx)) { + await this.ctx.storage.deleteAlarm(); + return; + } + this.ensureStorage(); + this.ctx.storage.sql.exec( + "DELETE FROM counter WHERE updated_at < ?", + Date.now() - QUOTA_TRACKER_RETENTION_MS, + ); + this.ctx.storage.sql.exec( + "DELETE FROM usage_event WHERE recorded_at < ?", + Date.now() - QUOTA_TRACKER_RETENTION_MS, + ); + this.ctx.storage.sql.exec( + "DELETE FROM quota_operation WHERE recorded_at < ?", + Date.now() - QUOTA_TRACKER_RETENTION_MS, + ); + await this.refreshCleanupAlarm(); + } + + private ensureStorage(): void { + if (this.isStorageInitialized) { + return; + } + ensureQuotaTrackerStorage(this.ctx); + this.isStorageInitialized = true; + } + + private consumeOnce(input: QuotaOperationInput): QuotaTryConsumeResponse { + const existing = this.readOperation(input); + if (existing) { + return operationConsumeResult(existing); + } + const limit = this.readLimit(input.feature); + const used = this.readUsed(input.feature, input.periodKey); + const allowed = used + input.amount <= limit; + const nextUsed = allowed ? used + input.amount : used; + const remaining = Math.max(0, limit - nextUsed); + if (allowed) { + this.writeUsage(input.feature, input.periodKey, input.amount, nextUsed, Date.now()); + } + this.insertOperation(input, { allowed, limit, remaining, used: nextUsed }); + return QuotaTryConsumeResponseSchema.parse({ allowed, limit, remaining }); + } + + private recordOnce(input: QuotaOperationInput, recordedAt: number): QuotaUsageResponse { + const existing = this.readOperation(input); + if (existing) { + return operationPeekResult(existing); + } + const nextUsed = this.readUsed(input.feature, input.periodKey) + input.amount; + const limit = this.readLimit(input.feature); + const remaining = Math.max(0, limit - nextUsed); + this.writeUsage(input.feature, input.periodKey, input.amount, nextUsed, recordedAt); + this.insertOperation(input, { allowed: true, limit, remaining, used: nextUsed }); + return QuotaUsageResponseSchema.parse({ limit, remaining, used: nextUsed }); + } + + private readOperation(input: QuotaOperationInput): OperationRow | null { + const [row] = this.ctx.storage.sql + .exec("SELECT * FROM quota_operation WHERE event_id = ?", input.eventId) + .toArray(); + if (!isOperationRow(row)) { + return null; + } + if ( + row.operation !== input.operation || + row.feature !== input.feature || + row.period_key !== input.periodKey || + row.amount !== input.amount + ) { + throw new Error("Quota event id was reused with different operation data."); + } + return row; + } + + private insertOperation( + input: QuotaOperationInput, + result: { allowed: boolean; limit: number; remaining: number; used: number }, + ): void { + this.ctx.storage.sql.exec( + `INSERT INTO quota_operation + (event_id, operation, feature, period_key, amount, allowed, limit_val, remaining, used, recorded_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + input.eventId, + input.operation, + input.feature, + input.periodKey, + input.amount, + result.allowed ? 1 : 0, + result.limit, + result.remaining, + result.used, + Date.now(), + ); + } + + private async ensureCleanupAlarm(): Promise { + const currentAlarm = await this.ctx.storage.getAlarm(); + if (currentAlarm === null) { + await this.ctx.storage.setAlarm(nextQuotaTrackerAlarm(Date.now())); + } + } + + private async refreshCleanupAlarm(): Promise { + if (this.hasRetainedUsage()) { + await this.ensureCleanupAlarm(); + return; + } + await this.ctx.storage.deleteAlarm(); + } + + private hasRetainedUsage(): boolean { + for (const table of ["counter", "usage_event", "quota_operation"] as const) { + const [row] = this.ctx.storage.sql + .exec(`SELECT 1 AS present FROM ${table} LIMIT 1`) + .toArray(); + if (isRecord(row) && row["present"] === 1) { + return true; + } + } + return false; + } + + private readLimit(feature: QuotaFeature): number { + const [rawRow] = this.ctx.storage.sql + .exec("SELECT limit_val FROM limit_override WHERE feature = ?", feature) + .toArray(); + return isRecord(rawRow) && typeof rawRow["limit_val"] === "number" ? rawRow["limit_val"] : 0; + } + + private readUsed(feature: QuotaFeature, periodKey: string): number { + const [rawRow] = this.ctx.storage.sql + .exec("SELECT used FROM counter WHERE feature = ? AND period_key = ?", feature, periodKey) + .toArray(); + return isCounterRow(rawRow) ? rawRow.used : 0; + } + + private writeUsage( + feature: QuotaFeature, + periodKey: string, + amount: number, + used: number, + recordedAt: number, + ): void { + this.ctx.storage.sql.exec( + `INSERT INTO counter (feature, period_key, used, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(feature, period_key) DO UPDATE SET + used = excluded.used, + updated_at = excluded.updated_at`, + feature, + periodKey, + used, + Date.now(), + ); + this.ctx.storage.sql.exec( + "INSERT INTO usage_event (feature, amount, recorded_at) VALUES (?, ?, ?)", + feature, + amount, + recordedAt, + ); + } +} + +function parseQuotaInput(schema: z.ZodType, value: unknown, method: string): T { + const result = schema.safeParse(value); + if (result.success) { + return result.data; + } + throw new APIError(400, "invalid_request_body", `Invalid QuotaTracker ${method} input`, { + details: { + fields: result.error.issues.map((issue) => issue.path.map(String).join(".")), + }, + retriable: false, + }); +} + +function operationConsumeResult(row: OperationRow): QuotaTryConsumeResponse { + return QuotaTryConsumeResponseSchema.parse({ + allowed: row.allowed === 1, + limit: row.limit_val, + remaining: row.remaining, + }); +} + +function operationPeekResult(row: OperationRow): QuotaUsageResponse { + return QuotaUsageResponseSchema.parse({ + limit: row.limit_val, + remaining: row.remaining, + used: row.used, + }); +} + +function historyResult(rows: unknown[]): QuotaHistoryResult { + return QuotaHistoryResultSchema.parse( + rows.filter(isHistoryRow).map((row) => ({ amount: row.amount, recordedAt: row.recorded_at })), + ); +} + +function periodKeyFromDate(date: Date): string { + const year = date.getUTCFullYear(); + const month = String(date.getUTCMonth() + 1).padStart(2, "0"); + return `${year}-${month}`; +} diff --git a/packages/billing/tsconfig.json b/packages/billing/tsconfig.json index d74d0577..27b1334e 100644 --- a/packages/billing/tsconfig.json +++ b/packages/billing/tsconfig.json @@ -3,7 +3,8 @@ "compilerOptions": { "lib": ["ES2023", "DOM"], "outDir": "dist", - "rootDir": "src" + "rootDir": "src", + "types": ["@cloudflare/workers-types"] }, "include": ["src/**/*"] } diff --git a/packages/sandbox-contracts/src/runtime.ts b/packages/sandbox-contracts/src/runtime.ts index f7f9fd1b..6b642680 100644 --- a/packages/sandbox-contracts/src/runtime.ts +++ b/packages/sandbox-contracts/src/runtime.ts @@ -1,3 +1,4 @@ +import type { SandboxExecResultBase, SandboxFileEntry } from "@cheatcode/types"; import type { ArtifactKind } from "@cheatcode/types/artifacts"; import { z } from "zod"; @@ -36,14 +37,7 @@ export interface SandboxExecInput { timeoutMs?: number; } -export interface SandboxExecResult { - command: string; - durationMs?: number; - exitCode: number; - stderr: string; - stdout: string; - success: boolean; -} +export type SandboxExecResult = SandboxExecResultBase; export interface SandboxReadFileInput { encoding?: "utf8" | "base64"; @@ -68,15 +62,6 @@ export interface SandboxWriteFileResult { success: boolean; } -interface SandboxFileEntry { - modifiedAt: string; - name: string; - path: string; - relativePath: string; - size: number; - type: "file" | "directory" | "symlink" | "other"; -} - export interface SandboxListFilesInput { includeHidden?: boolean; path: string; diff --git a/packages/tools-code/src/files.ts b/packages/tools-code/src/files.ts index 9ca092ef..482f3565 100644 --- a/packages/tools-code/src/files.ts +++ b/packages/tools-code/src/files.ts @@ -1,5 +1,6 @@ import { APIError } from "@cheatcode/observability"; import { callSandboxMethod, type getCodeRuntimeContext } from "@cheatcode/sandbox-contracts"; +import { sandboxFileEntryShape } from "@cheatcode/types"; import { z } from "zod"; import { resolveProjectWorkspacePath, @@ -53,16 +54,7 @@ export const ListFilesInputSchema = z }) .strict(); -const FileEntrySchema = z - .object({ - name: z.string(), - path: z.string(), - relativePath: z.string(), - type: z.enum(["file", "directory", "symlink", "other"]), - size: z.number().int().nonnegative(), - modifiedAt: z.string(), - }) - .strict(); +const FileEntrySchema = z.object(sandboxFileEntryShape(z.string(), "list-files")).strict(); export const ListFilesOutputSchema = z .object({ diff --git a/packages/types/README.md b/packages/types/README.md index 9c8a80d4..ec0a1182 100644 --- a/packages/types/README.md +++ b/packages/types/README.md @@ -21,6 +21,8 @@ capability discovery contracts, error codes, and UI message types. - `models.ts`: catalog IDs plus the open provider-prefixed logical-model schema - `@cheatcode/types/quota`: strict cross-Worker QuotaTracker request/response contracts and canonical quota feature identifiers +- `sandbox-wire.ts`: canonical sandbox file-entry fields and exec-result base + used by API, runtime-port, and code-tool schemas - `skill-runtime.ts`: canonical skill-runtime capability scopes and schema - `ui-message.ts`: the exact AI SDK UI message data-part contract persisted in Postgres and replayed to the web client diff --git a/packages/types/src/api.ts b/packages/types/src/api.ts index 70e88651..090c842b 100644 --- a/packages/types/src/api.ts +++ b/packages/types/src/api.ts @@ -1,6 +1,7 @@ import { z } from "zod"; import { IntegrationNameSchema } from "./integrations"; import { LogicalModelIdSchema } from "./models"; +import { extendSandboxExecResultShape, sandboxFileEntryShape } from "./sandbox-wire"; import { MessagePartsSchema } from "./ui-message"; /** Canonical total character budget for one submitted user message, including inline attachments. */ @@ -319,16 +320,7 @@ export const SandboxFilePathSchema = z "Path must be canonical and stay under /workspace.", ); -const SandboxFileEntrySchema = z - .object({ - modifiedAt: z.string(), - name: z.string(), - path: SandboxFilePathSchema, - relativePath: z.string(), - size: z.number().int().nonnegative(), - type: z.enum(["file", "directory", "symlink", "other"]), - }) - .strict(); +const SandboxFileEntrySchema = z.object(sandboxFileEntryShape(SandboxFilePathSchema)).strict(); export const SandboxTerminalCommandSchema = z .object({ @@ -339,15 +331,11 @@ export const SandboxTerminalCommandSchema = z .strict(); export const SandboxTerminalResultSchema = z - .object({ - command: z.string(), - cwd: SandboxFilePathSchema.optional(), - durationMs: z.number().int().nonnegative().optional(), - exitCode: z.number().int(), - stderr: z.string(), - stdout: z.string(), - success: z.boolean(), - }) + .object( + extendSandboxExecResultShape({ + cwd: SandboxFilePathSchema.optional(), + }), + ) .strict(); export const SandboxTerminalContextSchema = z diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts index 16f5a4cd..32b55c5d 100644 --- a/packages/types/src/index.ts +++ b/packages/types/src/index.ts @@ -178,6 +178,8 @@ export { UserProfileSchema, } from "./profile"; export { RunStatusSnapshotSchema } from "./run-control"; +export type { SandboxExecResultBase } from "./sandbox-wire"; +export { sandboxFileEntryShape } from "./sandbox-wire"; export type { SkillRuntimeScope } from "./skill-runtime"; export { SkillRuntimeScopeSchema } from "./skill-runtime"; export { diff --git a/packages/types/src/sandbox-wire.ts b/packages/types/src/sandbox-wire.ts new file mode 100644 index 00000000..04ca3c6f --- /dev/null +++ b/packages/types/src/sandbox-wire.ts @@ -0,0 +1,57 @@ +import { z } from "zod"; + +const SandboxFileEntryShape = { + modifiedAt: z.string(), + name: z.string(), + path: z.string(), + relativePath: z.string(), + size: z.number().int().nonnegative(), + type: z.enum(["file", "directory", "symlink", "other"]), +} satisfies z.ZodRawShape; + +type SandboxFileEntryShapeWithPath> = Omit< + typeof SandboxFileEntryShape, + "path" +> & { + path: PathSchema; +}; + +/** + * Builds every sandbox file-entry schema from one field definition. The list + * order option preserves the existing tools-code JSON Schema serialization. + */ +export function sandboxFileEntryShape>( + path: PathSchema, + order: "canonical" | "list-files" = "canonical", +): SandboxFileEntryShapeWithPath { + const shape = { ...SandboxFileEntryShape, path }; + if (order === "canonical") { + return shape; + } + const { modifiedAt, size, type, ...leadingFields } = shape; + return { ...leadingFields, type, size, modifiedAt }; +} + +const SandboxExecResultBaseShape = { + command: z.string(), + durationMs: z.number().int().nonnegative().optional(), + exitCode: z.number().int(), + stderr: z.string(), + stdout: z.string(), + success: z.boolean(), +} satisfies z.ZodRawShape; + +const SandboxExecResultBaseSchema = z.object(SandboxExecResultBaseShape).strict(); + +export type SandboxExecResultBase = z.infer; + +/** + * Extends the canonical exec-result base while preserving the terminal wire + * schema's existing placement of `cwd` immediately after `command`. + */ +export function extendSandboxExecResultShape( + extension: Extension, +): typeof SandboxExecResultBaseShape & Extension { + const { command, ...remainingBase } = SandboxExecResultBaseShape; + return { command, ...extension, ...remainingBase }; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 980ccfb2..e9d36279 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -647,6 +647,9 @@ importers: packages/billing: dependencies: + '@cheatcode/durable-storage': + specifier: workspace:* + version: link:../durable-storage '@cheatcode/observability': specifier: workspace:* version: link:../observability @@ -663,6 +666,9 @@ importers: '@cheatcode/tsconfig': specifier: workspace:* version: link:../tsconfig + '@cloudflare/workers-types': + specifier: 'catalog:' + version: 5.20260716.1 typescript: specifier: 'catalog:' version: 6.0.3