diff --git a/lib/sandbox/__tests__/createSandboxFromSnapshot.test.ts b/lib/sandbox/__tests__/createSandboxFromSnapshot.test.ts index 887d136c..d7a97507 100644 --- a/lib/sandbox/__tests__/createSandboxFromSnapshot.test.ts +++ b/lib/sandbox/__tests__/createSandboxFromSnapshot.test.ts @@ -3,16 +3,16 @@ import type { Sandbox } from "@vercel/sandbox"; import { createSandboxFromSnapshot } from "../createSandboxFromSnapshot"; -const mockSelectAccountSnapshots = vi.fn(); +const mockGetValidSnapshotId = vi.fn(); const mockInsertAccountSandbox = vi.fn(); -const mockCreateSandbox = vi.fn(); +const mockCreateSandboxWithFallback = vi.fn(); -vi.mock("@/lib/sandbox/createSandbox", () => ({ - createSandbox: (...args: unknown[]) => mockCreateSandbox(...args), +vi.mock("@/lib/sandbox/createSandboxWithFallback", () => ({ + createSandboxWithFallback: (...args: unknown[]) => mockCreateSandboxWithFallback(...args), })); -vi.mock("@/lib/supabase/account_snapshots/selectAccountSnapshots", () => ({ - selectAccountSnapshots: (...args: unknown[]) => mockSelectAccountSnapshots(...args), +vi.mock("@/lib/sandbox/getValidSnapshotId", () => ({ + getValidSnapshotId: (...args: unknown[]) => mockGetValidSnapshotId(...args), })); vi.mock("@/lib/supabase/account_sandboxes/insertAccountSandbox", () => ({ @@ -28,7 +28,7 @@ describe("createSandboxFromSnapshot", () => { beforeEach(() => { vi.clearAllMocks(); - mockCreateSandbox.mockResolvedValue({ + mockCreateSandboxWithFallback.mockResolvedValue({ sandbox: mockSandbox, response: { sandboxId: "sbx_new", @@ -36,6 +36,7 @@ describe("createSandboxFromSnapshot", () => { timeout: 1800000, createdAt: "2024-01-01T00:00:00.000Z", }, + fromSnapshot: false, }); mockInsertAccountSandbox.mockResolvedValue({ data: { account_id: "acc_1", sandbox_id: "sbx_new" }, @@ -43,28 +44,24 @@ describe("createSandboxFromSnapshot", () => { }); }); - it("creates from snapshot when available", async () => { - mockSelectAccountSnapshots.mockResolvedValue([ - { snapshot_id: "snap_abc", account_id: "acc_1" }, - ]); + it("passes snapshotId to createSandboxWithFallback", async () => { + mockGetValidSnapshotId.mockResolvedValue("snap_abc"); await createSandboxFromSnapshot("acc_1"); - expect(mockCreateSandbox).toHaveBeenCalledWith({ - source: { type: "snapshot", snapshotId: "snap_abc" }, - }); + expect(mockCreateSandboxWithFallback).toHaveBeenCalledWith("snap_abc"); }); - it("creates fresh sandbox when no snapshot exists", async () => { - mockSelectAccountSnapshots.mockResolvedValue([]); + it("passes undefined when no snapshot exists", async () => { + mockGetValidSnapshotId.mockResolvedValue(undefined); await createSandboxFromSnapshot("acc_1"); - expect(mockCreateSandbox).toHaveBeenCalledWith({}); + expect(mockCreateSandboxWithFallback).toHaveBeenCalledWith(undefined); }); it("inserts account_sandbox record", async () => { - mockSelectAccountSnapshots.mockResolvedValue([]); + mockGetValidSnapshotId.mockResolvedValue(undefined); await createSandboxFromSnapshot("acc_1"); @@ -74,10 +71,18 @@ describe("createSandboxFromSnapshot", () => { }); }); - it("returns { sandbox, fromSnapshot: true } when snapshot exists", async () => { - mockSelectAccountSnapshots.mockResolvedValue([ - { snapshot_id: "snap_abc", account_id: "acc_1" }, - ]); + it("returns { sandbox, fromSnapshot: true } when snapshot used", async () => { + mockGetValidSnapshotId.mockResolvedValue("snap_abc"); + mockCreateSandboxWithFallback.mockResolvedValue({ + sandbox: mockSandbox, + response: { + sandboxId: "sbx_new", + sandboxStatus: "running", + timeout: 1800000, + createdAt: "2024-01-01T00:00:00.000Z", + }, + fromSnapshot: true, + }); const result = await createSandboxFromSnapshot("acc_1"); @@ -85,7 +90,34 @@ describe("createSandboxFromSnapshot", () => { }); it("returns { sandbox, fromSnapshot: false } when no snapshot", async () => { - mockSelectAccountSnapshots.mockResolvedValue([]); + mockGetValidSnapshotId.mockResolvedValue(undefined); + + const result = await createSandboxFromSnapshot("acc_1"); + + expect(result).toEqual({ sandbox: mockSandbox, fromSnapshot: false }); + }); + + it("returns { sandbox, fromSnapshot: false } for expired snapshot", async () => { + mockGetValidSnapshotId.mockResolvedValue(undefined); + + const result = await createSandboxFromSnapshot("acc_1"); + + expect(mockCreateSandboxWithFallback).toHaveBeenCalledWith(undefined); + expect(result).toEqual({ sandbox: mockSandbox, fromSnapshot: false }); + }); + + it("returns { sandbox, fromSnapshot: false } when snapshot creation fails", async () => { + mockGetValidSnapshotId.mockResolvedValue("snap_bad"); + mockCreateSandboxWithFallback.mockResolvedValue({ + sandbox: mockSandbox, + response: { + sandboxId: "sbx_new", + sandboxStatus: "running", + timeout: 1800000, + createdAt: "2024-01-01T00:00:00.000Z", + }, + fromSnapshot: false, + }); const result = await createSandboxFromSnapshot("acc_1"); diff --git a/lib/sandbox/__tests__/createSandboxPostHandler.test.ts b/lib/sandbox/__tests__/createSandboxPostHandler.test.ts index 0e1cfc00..03aafaaf 100644 --- a/lib/sandbox/__tests__/createSandboxPostHandler.test.ts +++ b/lib/sandbox/__tests__/createSandboxPostHandler.test.ts @@ -4,29 +4,14 @@ import { NextResponse } from "next/server"; import { createSandboxPostHandler } from "../createSandboxPostHandler"; import { validateSandboxBody } from "@/lib/sandbox/validateSandboxBody"; -import { createSandbox } from "@/lib/sandbox/createSandbox"; -import { insertAccountSandbox } from "@/lib/supabase/account_sandboxes/insertAccountSandbox"; -import { triggerPromptSandbox } from "@/lib/trigger/triggerPromptSandbox"; -import { selectAccountSnapshots } from "@/lib/supabase/account_snapshots/selectAccountSnapshots"; +import { processCreateSandbox } from "@/lib/sandbox/processCreateSandbox"; vi.mock("@/lib/sandbox/validateSandboxBody", () => ({ validateSandboxBody: vi.fn(), })); -vi.mock("@/lib/sandbox/createSandbox", () => ({ - createSandbox: vi.fn(), -})); - -vi.mock("@/lib/supabase/account_sandboxes/insertAccountSandbox", () => ({ - insertAccountSandbox: vi.fn(), -})); - -vi.mock("@/lib/trigger/triggerPromptSandbox", () => ({ - triggerPromptSandbox: vi.fn(), -})); - -vi.mock("@/lib/supabase/account_snapshots/selectAccountSnapshots", () => ({ - selectAccountSnapshots: vi.fn(), +vi.mock("@/lib/sandbox/processCreateSandbox", () => ({ + processCreateSandbox: vi.fn(), })); /** @@ -56,34 +41,17 @@ describe("createSandboxPostHandler", () => { expect(response.status).toBe(401); }); - it("returns runId when prompt is provided", async () => { + it("returns 200 with sandbox result when no prompt", async () => { vi.mocked(validateSandboxBody).mockResolvedValue({ accountId: "acc_123", orgId: null, authToken: "token", - prompt: "create a hello world page", }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_123", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - vi.mocked(triggerPromptSandbox).mockResolvedValue({ - id: "run_abc123", + vi.mocked(processCreateSandbox).mockResolvedValue({ + sandboxId: "sbx_123", + sandboxStatus: "running", + timeout: 600000, + createdAt: "2024-01-01T00:00:00.000Z", }); const request = createMockRequest(); @@ -99,181 +67,66 @@ describe("createSandboxPostHandler", () => { sandboxStatus: "running", timeout: 600000, createdAt: "2024-01-01T00:00:00.000Z", - runId: "run_abc123", }, ], }); }); - it("calls createSandbox with snapshotId when account has snapshot", async () => { + it("returns runId when prompt is provided", async () => { vi.mocked(validateSandboxBody).mockResolvedValue({ accountId: "acc_123", orgId: null, authToken: "token", - prompt: "say hello", - }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([ - { - id: "snap_record_123", - account_id: "acc_123", - snapshot_id: "snap_xyz", - created_at: "2024-01-01T00:00:00.000Z", - }, - ]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_456", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_456", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, + prompt: "create a hello world page", }); - vi.mocked(triggerPromptSandbox).mockResolvedValue({ - id: "run_def456", + vi.mocked(processCreateSandbox).mockResolvedValue({ + sandboxId: "sbx_123", + sandboxStatus: "running", + timeout: 600000, + createdAt: "2024-01-01T00:00:00.000Z", + runId: "run_abc123", }); const request = createMockRequest(); - await createSandboxPostHandler(request); + const response = await createSandboxPostHandler(request); - expect(createSandbox).toHaveBeenCalledWith({ - source: { type: "snapshot", snapshotId: "snap_xyz" }, - }); + expect(response.status).toBe(200); + const json = await response.json(); + expect(json.sandboxes[0].runId).toBe("run_abc123"); }); - it("calls createSandbox with empty params when account has no snapshot", async () => { + it("passes validated input to processCreateSandbox", async () => { vi.mocked(validateSandboxBody).mockResolvedValue({ accountId: "acc_123", orgId: null, authToken: "token", prompt: "say hello", }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_456", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_456", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - vi.mocked(triggerPromptSandbox).mockResolvedValue({ - id: "run_def456", + vi.mocked(processCreateSandbox).mockResolvedValue({ + sandboxId: "sbx_123", + sandboxStatus: "running", + timeout: 600000, + createdAt: "2024-01-01T00:00:00.000Z", }); const request = createMockRequest(); await createSandboxPostHandler(request); - expect(createSandbox).toHaveBeenCalledWith({}); - }); - - it("calls insertAccountSandbox with correct account_id and sandbox_id", async () => { - vi.mocked(validateSandboxBody).mockResolvedValue({ + expect(processCreateSandbox).toHaveBeenCalledWith({ accountId: "acc_123", orgId: null, authToken: "token", prompt: "say hello", }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_456", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_456", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - vi.mocked(triggerPromptSandbox).mockResolvedValue({ - id: "run_def456", - }); - - const request = createMockRequest(); - await createSandboxPostHandler(request); - - expect(insertAccountSandbox).toHaveBeenCalledWith({ - account_id: "acc_123", - sandbox_id: "sbx_456", - }); }); - it("calls triggerPromptSandbox with prompt, sandboxId, and accountId", async () => { + it("returns 400 when processCreateSandbox throws", async () => { vi.mocked(validateSandboxBody).mockResolvedValue({ accountId: "acc_123", orgId: null, authToken: "token", - prompt: "create a hello world page", - }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_789", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_789", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - vi.mocked(triggerPromptSandbox).mockResolvedValue({ - id: "run_ghi789", }); - - const request = createMockRequest(); - await createSandboxPostHandler(request); - - expect(triggerPromptSandbox).toHaveBeenCalledWith({ - prompt: "create a hello world page", - sandboxId: "sbx_789", - accountId: "acc_123", - }); - }); - - it("returns 400 with error status when createSandbox throws", async () => { - vi.mocked(validateSandboxBody).mockResolvedValue({ - accountId: "acc_123", - orgId: null, - authToken: "token", - prompt: "say hello", - }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockRejectedValue(new Error("Sandbox creation failed")); + vi.mocked(processCreateSandbox).mockRejectedValue(new Error("Sandbox creation failed")); const request = createMockRequest(); const response = await createSandboxPostHandler(request); @@ -285,126 +138,4 @@ describe("createSandboxPostHandler", () => { error: "Sandbox creation failed", }); }); - - it("returns 400 with error status when insertAccountSandbox throws", async () => { - vi.mocked(validateSandboxBody).mockResolvedValue({ - accountId: "acc_123", - orgId: null, - authToken: "token", - prompt: "say hello", - }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockRejectedValue(new Error("Database insert failed")); - - const request = createMockRequest(); - const response = await createSandboxPostHandler(request); - - expect(response.status).toBe(400); - const json = await response.json(); - expect(json).toEqual({ - status: "error", - error: "Database insert failed", - }); - }); - - it("returns 200 without runId when no prompt is provided", async () => { - vi.mocked(validateSandboxBody).mockResolvedValue({ - accountId: "acc_123", - orgId: null, - authToken: "token", - }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_123", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - - const request = createMockRequest(); - const response = await createSandboxPostHandler(request); - - expect(response.status).toBe(200); - const json = await response.json(); - expect(json).toEqual({ - status: "success", - sandboxes: [ - { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - ], - }); - expect(triggerPromptSandbox).not.toHaveBeenCalled(); - }); - - it("returns 200 without runId when triggerPromptSandbox throws", async () => { - vi.mocked(validateSandboxBody).mockResolvedValue({ - accountId: "acc_123", - orgId: null, - authToken: "token", - prompt: "say hello", - }); - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_123", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - vi.mocked(triggerPromptSandbox).mockRejectedValue(new Error("Task trigger failed")); - - const request = createMockRequest(); - const response = await createSandboxPostHandler(request); - - // Sandbox was created successfully, so return 200 even if prompt trigger fails - expect(response.status).toBe(200); - const json = await response.json(); - expect(json).toEqual({ - status: "success", - sandboxes: [ - { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - ], - }); - }); }); diff --git a/lib/sandbox/__tests__/createSandboxWithFallback.test.ts b/lib/sandbox/__tests__/createSandboxWithFallback.test.ts new file mode 100644 index 00000000..3af6d1fe --- /dev/null +++ b/lib/sandbox/__tests__/createSandboxWithFallback.test.ts @@ -0,0 +1,69 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +import { createSandboxWithFallback } from "../createSandboxWithFallback"; + +const mockCreateSandbox = vi.fn(); + +vi.mock("@/lib/sandbox/createSandbox", () => ({ + createSandbox: (...args: unknown[]) => mockCreateSandbox(...args), +})); + +const mockSandbox = { sandboxId: "sbx_123", status: "running" }; +const mockResponse = { + sandboxId: "sbx_123", + sandboxStatus: "running", + timeout: 1800000, + createdAt: "2024-01-01T00:00:00.000Z", +}; +const mockCreateResult = { sandbox: mockSandbox, response: mockResponse }; + +describe("createSandboxWithFallback", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockCreateSandbox.mockResolvedValue(mockCreateResult); + }); + + it("creates from snapshot when snapshotId is provided", async () => { + const result = await createSandboxWithFallback("snap_abc"); + + expect(mockCreateSandbox).toHaveBeenCalledWith({ + source: { type: "snapshot", snapshotId: "snap_abc" }, + }); + expect(result).toEqual({ ...mockCreateResult, fromSnapshot: true }); + }); + + it("creates fresh sandbox when snapshotId is undefined", async () => { + const result = await createSandboxWithFallback(undefined); + + expect(mockCreateSandbox).toHaveBeenCalledWith({}); + expect(result).toEqual({ ...mockCreateResult, fromSnapshot: false }); + }); + + it("falls back to fresh sandbox when snapshot creation fails", async () => { + const freshSandbox = { sandboxId: "sbx_fresh", status: "running" }; + const freshResponse = { ...mockResponse, sandboxId: "sbx_fresh" }; + const freshResult = { sandbox: freshSandbox, response: freshResponse }; + mockCreateSandbox + .mockRejectedValueOnce(new Error("Status code 400 is not ok")) + .mockResolvedValueOnce(freshResult); + + const result = await createSandboxWithFallback("snap_bad"); + + expect(mockCreateSandbox).toHaveBeenCalledTimes(2); + expect(result).toEqual({ ...freshResult, fromSnapshot: false }); + }); + + it("logs error when snapshot creation fails", async () => { + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + const snapshotError = new Error("Status code 400 is not ok"); + mockCreateSandbox.mockRejectedValueOnce(snapshotError).mockResolvedValueOnce(mockCreateResult); + + await createSandboxWithFallback("snap_bad"); + + expect(consoleSpy).toHaveBeenCalledWith( + "Snapshot sandbox creation failed, falling back to fresh sandbox:", + snapshotError, + ); + consoleSpy.mockRestore(); + }); +}); diff --git a/lib/sandbox/__tests__/getValidSnapshotId.test.ts b/lib/sandbox/__tests__/getValidSnapshotId.test.ts new file mode 100644 index 00000000..2ee3ad63 --- /dev/null +++ b/lib/sandbox/__tests__/getValidSnapshotId.test.ts @@ -0,0 +1,86 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +import { getValidSnapshotId } from "../getValidSnapshotId"; + +const mockSelectAccountSnapshots = vi.fn(); + +vi.mock("@/lib/supabase/account_snapshots/selectAccountSnapshots", () => ({ + selectAccountSnapshots: (...args: unknown[]) => mockSelectAccountSnapshots(...args), +})); + +describe("getValidSnapshotId", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("returns snapshot_id when snapshot exists and is not expired", async () => { + const futureDate = new Date(Date.now() + 86400000).toISOString(); + mockSelectAccountSnapshots.mockResolvedValue([ + { snapshot_id: "snap_abc", account_id: "acc_1", expires_at: futureDate }, + ]); + + const result = await getValidSnapshotId("acc_1"); + + expect(result).toBe("snap_abc"); + }); + + it("returns snapshot_id when snapshot has no expires_at", async () => { + mockSelectAccountSnapshots.mockResolvedValue([ + { snapshot_id: "snap_abc", account_id: "acc_1", expires_at: null }, + ]); + + const result = await getValidSnapshotId("acc_1"); + + expect(result).toBe("snap_abc"); + }); + + it("returns undefined when snapshot is expired", async () => { + const pastDate = new Date(Date.now() - 86400000).toISOString(); + mockSelectAccountSnapshots.mockResolvedValue([ + { snapshot_id: "snap_expired", account_id: "acc_1", expires_at: pastDate }, + ]); + + const result = await getValidSnapshotId("acc_1"); + + expect(result).toBeUndefined(); + }); + + it("returns undefined when no snapshots exist", async () => { + mockSelectAccountSnapshots.mockResolvedValue([]); + + const result = await getValidSnapshotId("acc_1"); + + expect(result).toBeUndefined(); + }); + + it("returns undefined when expires_at equals now", async () => { + const now = new Date().toISOString(); + mockSelectAccountSnapshots.mockResolvedValue([ + { snapshot_id: "snap_edge", account_id: "acc_1", expires_at: now }, + ]); + + const result = await getValidSnapshotId("acc_1"); + + expect(result).toBeUndefined(); + }); + + it("returns undefined when expires_at is unparsable", async () => { + mockSelectAccountSnapshots.mockResolvedValue([ + { snapshot_id: "snap_bad", account_id: "acc_1", expires_at: "not-a-date" }, + ]); + + const result = await getValidSnapshotId("acc_1"); + + expect(result).toBeUndefined(); + }); + + it("returns undefined when snapshot has no snapshot_id", async () => { + mockSelectAccountSnapshots.mockResolvedValue([ + { snapshot_id: null, account_id: "acc_1", expires_at: null }, + ]); + + const result = await getValidSnapshotId("acc_1"); + + expect(result).toBeUndefined(); + }); +}); diff --git a/lib/sandbox/__tests__/processCreateSandbox.test.ts b/lib/sandbox/__tests__/processCreateSandbox.test.ts index f70b40c2..5f8b33e2 100644 --- a/lib/sandbox/__tests__/processCreateSandbox.test.ts +++ b/lib/sandbox/__tests__/processCreateSandbox.test.ts @@ -1,53 +1,41 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; +import type { Sandbox } from "@vercel/sandbox"; import { processCreateSandbox } from "../processCreateSandbox"; -import { createSandbox } from "@/lib/sandbox/createSandbox"; -import { insertAccountSandbox } from "@/lib/supabase/account_sandboxes/insertAccountSandbox"; +import { createSandboxFromSnapshot } from "@/lib/sandbox/createSandboxFromSnapshot"; import { triggerPromptSandbox } from "@/lib/trigger/triggerPromptSandbox"; -import { selectAccountSnapshots } from "@/lib/supabase/account_snapshots/selectAccountSnapshots"; -vi.mock("@/lib/sandbox/createSandbox", () => ({ - createSandbox: vi.fn(), -})); - -vi.mock("@/lib/supabase/account_sandboxes/insertAccountSandbox", () => ({ - insertAccountSandbox: vi.fn(), +vi.mock("@/lib/sandbox/createSandboxFromSnapshot", () => ({ + createSandboxFromSnapshot: vi.fn(), })); vi.mock("@/lib/trigger/triggerPromptSandbox", () => ({ triggerPromptSandbox: vi.fn(), })); -vi.mock("@/lib/supabase/account_snapshots/selectAccountSnapshots", () => ({ - selectAccountSnapshots: vi.fn(), -})); +const mockSandbox = { + sandboxId: "sbx_123", + status: "running", + timeout: 600000, + createdAt: new Date("2024-01-01T00:00:00.000Z"), +} as unknown as Sandbox; describe("processCreateSandbox", () => { beforeEach(() => { vi.clearAllMocks(); + vi.mocked(createSandboxFromSnapshot).mockResolvedValue({ + sandbox: mockSandbox, + fromSnapshot: false, + }); }); - it("creates sandbox without prompt and returns result without runId", async () => { - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_123", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); + it("delegates to createSandboxFromSnapshot", async () => { + await processCreateSandbox({ accountId: "acc_123" }); + + expect(createSandboxFromSnapshot).toHaveBeenCalledWith("acc_123"); + }); + it("returns serializable response without runId when no prompt", async () => { const result = await processCreateSandbox({ accountId: "acc_123" }); expect(result).toEqual({ @@ -59,26 +47,7 @@ describe("processCreateSandbox", () => { expect(triggerPromptSandbox).not.toHaveBeenCalled(); }); - it("creates sandbox with prompt and returns result with runId", async () => { - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_123", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); + it("returns result with runId when prompt is provided", async () => { vi.mocked(triggerPromptSandbox).mockResolvedValue({ id: "run_prompt123", }); @@ -102,99 +71,8 @@ describe("processCreateSandbox", () => { }); }); - it("uses snapshot when account has one", async () => { - vi.mocked(selectAccountSnapshots).mockResolvedValue([ - { - id: "snap_record_123", - account_id: "acc_123", - snapshot_id: "snap_xyz", - created_at: "2024-01-01T00:00:00.000Z", - }, - ]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_456", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_456", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - - await processCreateSandbox({ accountId: "acc_123" }); - - expect(createSandbox).toHaveBeenCalledWith({ - source: { type: "snapshot", snapshotId: "snap_xyz" }, - }); - }); - - it("calls createSandbox with empty params when no snapshot", async () => { - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_456", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_456", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - - await processCreateSandbox({ accountId: "acc_123" }); - - expect(createSandbox).toHaveBeenCalledWith({}); - }); - - it("inserts account_sandbox record", async () => { - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_789", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_789", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); - - await processCreateSandbox({ accountId: "acc_123" }); - - expect(insertAccountSandbox).toHaveBeenCalledWith({ - account_id: "acc_123", - sandbox_id: "sbx_789", - }); - }); - - it("throws when createSandbox fails", async () => { - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockRejectedValue(new Error("Sandbox creation failed")); + it("throws when createSandboxFromSnapshot fails", async () => { + vi.mocked(createSandboxFromSnapshot).mockRejectedValue(new Error("Sandbox creation failed")); await expect(processCreateSandbox({ accountId: "acc_123" })).rejects.toThrow( "Sandbox creation failed", @@ -202,25 +80,6 @@ describe("processCreateSandbox", () => { }); it("returns result without runId when triggerPromptSandbox fails", async () => { - vi.mocked(selectAccountSnapshots).mockResolvedValue([]); - vi.mocked(createSandbox).mockResolvedValue({ - sandbox: {} as never, - response: { - sandboxId: "sbx_123", - sandboxStatus: "running", - timeout: 600000, - createdAt: "2024-01-01T00:00:00.000Z", - }, - }); - vi.mocked(insertAccountSandbox).mockResolvedValue({ - data: { - id: "record_123", - account_id: "acc_123", - sandbox_id: "sbx_123", - created_at: "2024-01-01T00:00:00.000Z", - }, - error: null, - }); vi.mocked(triggerPromptSandbox).mockRejectedValue(new Error("Task trigger failed")); const result = await processCreateSandbox({ diff --git a/lib/sandbox/createSandboxFromSnapshot.ts b/lib/sandbox/createSandboxFromSnapshot.ts index 98310a1f..1b6da57a 100644 --- a/lib/sandbox/createSandboxFromSnapshot.ts +++ b/lib/sandbox/createSandboxFromSnapshot.ts @@ -1,6 +1,6 @@ import type { Sandbox } from "@vercel/sandbox"; -import { createSandbox } from "@/lib/sandbox/createSandbox"; -import { selectAccountSnapshots } from "@/lib/supabase/account_snapshots/selectAccountSnapshots"; +import { createSandboxWithFallback } from "@/lib/sandbox/createSandboxWithFallback"; +import { getValidSnapshotId } from "@/lib/sandbox/getValidSnapshotId"; import { insertAccountSandbox } from "@/lib/supabase/account_sandboxes/insertAccountSandbox"; export interface CreateSandboxFromSnapshotResult { @@ -9,8 +9,8 @@ export interface CreateSandboxFromSnapshotResult { } /** - * Creates a new sandbox from the account's latest snapshot (or fresh if none) - * and records it in the database. + * Creates a new sandbox from the account's latest valid snapshot, + * falling back to a fresh sandbox if the snapshot is expired or fails. * * @param accountId - The account ID to create a sandbox for * @returns The created Sandbox instance and whether it was created from a snapshot @@ -18,17 +18,13 @@ export interface CreateSandboxFromSnapshotResult { export async function createSandboxFromSnapshot( accountId: string, ): Promise { - const snapshots = await selectAccountSnapshots(accountId); - const snapshotId = snapshots[0]?.snapshot_id; - - const { sandbox, response } = await createSandbox( - snapshotId ? { source: { type: "snapshot", snapshotId } } : {}, - ); + const snapshotId = await getValidSnapshotId(accountId); + const { sandbox, fromSnapshot } = await createSandboxWithFallback(snapshotId); await insertAccountSandbox({ account_id: accountId, - sandbox_id: response.sandboxId, + sandbox_id: sandbox.sandboxId, }); - return { sandbox, fromSnapshot: !!snapshotId }; + return { sandbox, fromSnapshot }; } diff --git a/lib/sandbox/createSandboxWithFallback.ts b/lib/sandbox/createSandboxWithFallback.ts new file mode 100644 index 00000000..93014ac4 --- /dev/null +++ b/lib/sandbox/createSandboxWithFallback.ts @@ -0,0 +1,25 @@ +import { createSandbox, type SandboxCreateResult } from "@/lib/sandbox/createSandbox"; + +export type SandboxWithFallbackResult = SandboxCreateResult & { fromSnapshot: boolean }; + +/** + * Attempts to create a sandbox from the given snapshot, falling back to a fresh sandbox on failure. + * If no snapshotId is provided, creates a fresh sandbox directly. + * + * @param snapshotId - Optional snapshot ID to restore from + * @returns The sandbox creation result with fromSnapshot flag + */ +export async function createSandboxWithFallback( + snapshotId: string | undefined, +): Promise { + if (snapshotId) { + try { + const result = await createSandbox({ source: { type: "snapshot", snapshotId } }); + return { ...result, fromSnapshot: true }; + } catch (error) { + console.error("Snapshot sandbox creation failed, falling back to fresh sandbox:", error); + } + } + const result = await createSandbox({}); + return { ...result, fromSnapshot: false }; +} diff --git a/lib/sandbox/getValidSnapshotId.ts b/lib/sandbox/getValidSnapshotId.ts new file mode 100644 index 00000000..17d7a6d5 --- /dev/null +++ b/lib/sandbox/getValidSnapshotId.ts @@ -0,0 +1,22 @@ +import { selectAccountSnapshots } from "@/lib/supabase/account_snapshots/selectAccountSnapshots"; + +/** + * Returns a valid (non-expired) snapshot ID for the account, or undefined. + * + * @param accountId - The account to look up + * @returns The snapshot ID if it exists and has not expired + */ +export async function getValidSnapshotId(accountId: string): Promise { + const accountSnapshots = await selectAccountSnapshots(accountId); + const snapshot = accountSnapshots[0]; + if (!snapshot?.snapshot_id) return undefined; + + if (snapshot.expires_at) { + const expiresAt = new Date(snapshot.expires_at).getTime(); + if (Number.isNaN(expiresAt) || expiresAt <= Date.now()) { + return undefined; + } + } + + return snapshot.snapshot_id; +} diff --git a/lib/sandbox/processCreateSandbox.ts b/lib/sandbox/processCreateSandbox.ts index be568c11..0ce50ed6 100644 --- a/lib/sandbox/processCreateSandbox.ts +++ b/lib/sandbox/processCreateSandbox.ts @@ -1,7 +1,6 @@ -import { createSandbox, type SandboxCreatedResponse } from "@/lib/sandbox/createSandbox"; -import { insertAccountSandbox } from "@/lib/supabase/account_sandboxes/insertAccountSandbox"; -import { selectAccountSnapshots } from "@/lib/supabase/account_snapshots/selectAccountSnapshots"; +import { createSandboxFromSnapshot } from "@/lib/sandbox/createSandboxFromSnapshot"; import { triggerPromptSandbox } from "@/lib/trigger/triggerPromptSandbox"; +import type { SandboxCreatedResponse } from "@/lib/sandbox/createSandbox"; type ProcessCreateSandboxInput = { accountId: string; @@ -21,19 +20,14 @@ export async function processCreateSandbox( ): Promise { const { accountId, prompt } = input; - // Get account's most recent snapshot if available - const accountSnapshots = await selectAccountSnapshots(accountId); - const snapshotId = accountSnapshots[0]?.snapshot_id; + const { sandbox } = await createSandboxFromSnapshot(accountId); - // Create sandbox (from snapshot if valid, otherwise fresh) - const { response: result } = await createSandbox( - snapshotId ? { source: { type: "snapshot", snapshotId } } : {}, - ); - - await insertAccountSandbox({ - account_id: accountId, - sandbox_id: result.sandboxId, - }); + const result: SandboxCreatedResponse = { + sandboxId: sandbox.sandboxId, + sandboxStatus: sandbox.status, + timeout: sandbox.timeout, + createdAt: sandbox.createdAt.toISOString(), + }; // Trigger the prompt execution task if a prompt was provided let runId: string | undefined; @@ -41,7 +35,7 @@ export async function processCreateSandbox( try { const handle = await triggerPromptSandbox({ prompt, - sandboxId: result.sandboxId, + sandboxId: sandbox.sandboxId, accountId, }); runId = handle.id;