|
| 1 | +/** |
| 2 | + * The persistent workbench's input history is what lets a scratch file reach the model. This runs |
| 3 | + * the real code boundary against the real history script in a disposable Redis, so a mount the |
| 4 | + * caller could not classify has to leave the machine uncertified. Only the sandbox provider is a |
| 5 | + * stand-in: it hands back an existing machine that executes nothing. |
| 6 | + * |
| 7 | + * Set TEST_REDIS_URL to an isolated local Redis service. |
| 8 | + */ |
| 9 | +import { createHash } from 'node:crypto' |
| 10 | +import { |
| 11 | + remoteSandboxProviderMock, |
| 12 | + remoteSandboxProviderMockFns, |
| 13 | +} from '@sim/testing/mocks/remote-sandbox-provider.mock' |
| 14 | +import { generateShortId } from '@sim/utils/id' |
| 15 | +import { afterAll, describe, expect, it, vi } from 'vitest' |
| 16 | + |
| 17 | +const { redisUrl, inheritedRedisUrl, mockFindSessionSandbox } = await vi.hoisted(async () => { |
| 18 | + const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') |
| 19 | + const url = readTestRedisUrl() |
| 20 | + const inheritedRedisUrl = process.env.REDIS_URL |
| 21 | + /** The real Redis module reads this at import. */ |
| 22 | + if (url) process.env.REDIS_URL = url |
| 23 | + return { redisUrl: url, inheritedRedisUrl, mockFindSessionSandbox: vi.fn() } |
| 24 | +}) |
| 25 | + |
| 26 | +vi.mock('@/lib/execution/remote-sandbox/provider', () => remoteSandboxProviderMock) |
| 27 | +vi.mock('@/lib/execution/remote-sandbox/resolve', () => ({ |
| 28 | + resolveWorkspaceSandbox: async () => null, |
| 29 | + provisionRuntimeDependencies: async () => {}, |
| 30 | + repairMissingSandboxImage: async () => null, |
| 31 | + RUNTIME_INSTALL_TIMEOUT_MS: 60_000, |
| 32 | +})) |
| 33 | + |
| 34 | +import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis' |
| 35 | +import { CodeLanguage } from '@/lib/execution/languages' |
| 36 | +import { |
| 37 | + executeInSandbox, |
| 38 | + executeShellInSandbox, |
| 39 | + SIM_RESULT_PREFIX, |
| 40 | +} from '@/lib/execution/remote-sandbox' |
| 41 | +import { observeSandboxSessionInputs } from '@/lib/execution/remote-sandbox/execution-observer' |
| 42 | +import { |
| 43 | + initializeSessionFileProvenance, |
| 44 | + isSessionFileProvenanceClean, |
| 45 | +} from '@/lib/execution/remote-sandbox/session-file-provenance' |
| 46 | +import type { SandboxHandle, SandboxProvider } from '@/lib/execution/remote-sandbox/types' |
| 47 | + |
| 48 | +remoteSandboxProviderMockFns.mockResolveProvider.mockImplementation( |
| 49 | + () => |
| 50 | + ({ |
| 51 | + id: 'e2b', |
| 52 | + dependencyStrategy: 'prebuilt', |
| 53 | + resolveLifetimeMs: (ms: number) => ms, |
| 54 | + create: async () => { |
| 55 | + throw new Error('The certification fixture only reuses an existing machine') |
| 56 | + }, |
| 57 | + findSessionSandbox: mockFindSessionSandbox, |
| 58 | + }) satisfies SandboxProvider |
| 59 | +) |
| 60 | + |
| 61 | +function machine(sandboxId: string): SandboxHandle { |
| 62 | + return { |
| 63 | + sandboxId, |
| 64 | + runCode: async () => ({ text: `${SIM_RESULT_PREFIX}{"ok":true}`, stdout: '', stderr: '' }), |
| 65 | + runCommand: async () => ({ stdout: '', stderr: '', exitCode: 0 }), |
| 66 | + extendLifetime: async () => {}, |
| 67 | + getFileSize: async () => 0, |
| 68 | + readFile: async () => '', |
| 69 | + readFileWithLimit: async () => ({ content: '', byteLength: 0 }), |
| 70 | + writeFile: async () => {}, |
| 71 | + removeFile: async () => {}, |
| 72 | + listFiles: async () => [], |
| 73 | + kill: async () => {}, |
| 74 | + } |
| 75 | +} |
| 76 | + |
| 77 | +/** Machine-history keys this suite created, so cleanup never touches another suite's state. */ |
| 78 | +const createdKeys: string[] = [] |
| 79 | + |
| 80 | +afterAll(async () => { |
| 81 | + if (redisUrl && createdKeys.length) { |
| 82 | + const redis = getRedisClient() |
| 83 | + if (redis) await redis.del(...createdKeys) |
| 84 | + } |
| 85 | + if (redisUrl) await closeRedisConnection() |
| 86 | + // Only restore what the hoisted setup changed; assigning undefined would store the string "undefined". |
| 87 | + if (!redisUrl) return |
| 88 | + if (inheritedRedisUrl === undefined) Reflect.deleteProperty(process.env, 'REDIS_URL') |
| 89 | + else process.env.REDIS_URL = inheritedRedisUrl |
| 90 | +}) |
| 91 | + |
| 92 | +describe.skipIf(!redisUrl)('workbench certification at the code boundary', () => { |
| 93 | + it.each([ |
| 94 | + ['code', false], |
| 95 | + ['code', true], |
| 96 | + ['shell', false], |
| 97 | + ['shell', true], |
| 98 | + ] as const)('%s with unprovenanced mounts %s', async (kind, unprovenanced) => { |
| 99 | + const sandboxId = `machine-${generateShortId(12)}` |
| 100 | + const key = `certification-${generateShortId(12)}` |
| 101 | + const identity = { providerId: 'e2b', sandboxId } as const |
| 102 | + createdKeys.push( |
| 103 | + `mothership:workbench-provenance:v1:${createHash('sha256') |
| 104 | + .update(JSON.stringify([key, identity.providerId, sandboxId])) |
| 105 | + .digest('hex')}` |
| 106 | + ) |
| 107 | + mockFindSessionSandbox.mockResolvedValue(machine(sandboxId)) |
| 108 | + await initializeSessionFileProvenance(key, identity) |
| 109 | + expect(await isSessionFileProvenanceClean(key, identity)).toBe(true) |
| 110 | + |
| 111 | + const request = { |
| 112 | + code: 'print(1)', |
| 113 | + language: CodeLanguage.Python, |
| 114 | + timeoutMs: 30_000, |
| 115 | + session: { key, ...(unprovenanced ? { unprovenancedInputs: true } : {}) }, |
| 116 | + } |
| 117 | + await observeSandboxSessionInputs( |
| 118 | + () => true, |
| 119 | + () => |
| 120 | + kind === 'code' |
| 121 | + ? executeInSandbox(request) |
| 122 | + : executeShellInSandbox({ ...request, envs: {} }) |
| 123 | + ) |
| 124 | + expect(await isSessionFileProvenanceClean(key, identity)).toBe(!unprovenanced) |
| 125 | + }) |
| 126 | +}) |
0 commit comments