From dead07bd466a17dab4bc5b2b3520312fa52b0e03 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Fri, 11 Sep 2026 16:36:29 -0400 Subject: [PATCH] perf: reuse authorized input versions and make egress accounting atomic (#180) * perf: reuse authorized input versions and make egress accounting atomic * perf: resolve authorized input manifests once per execution * test: preserve fetch signature in revocation fixture * fix: isolate shared-grant failures and prevent ledger replay --- .github/workflows/ci.yml | 3 + api/src/config.ts | 3 + api/src/download.test.ts | 44 +++ api/src/http-input-cache.test.ts | 163 +++++++++++ api/src/http-input-cache.ts | 177 +++++++++++ api/src/input-manifest.test.ts | 124 ++++++++ api/src/job.ts | 55 ++++ api/src/metrics.ts | 6 + api/src/session-inputs.ts | 36 ++- docs/INPUT_REUSE.md | 91 ++++++ .../templates/egress-gateway-deployment.yaml | 8 + .../templates/file-server-deployment.yaml | 4 + .../templates/worker-sandbox-deployment.yaml | 8 + helm/codeapi/values.yaml | 11 + launcher/src/main.rs | 4 + packages/code/src/relay.test.ts | 65 +++++ packages/code/src/relay.ts | 16 +- service/src/config.ts | 6 + service/src/egress-gateway.test.ts | 207 ++++++++++--- service/src/egress-gateway.ts | 97 +++++- service/src/egress-ledger-reconnect.test.ts | 63 ++++ service/src/egress-ledger-script.ts | 123 ++++++++ service/src/egress-ledger.test.ts | 81 ++++- service/src/egress-ledger.ts | 276 ++++-------------- service/src/file-download.test.ts | 49 ++++ service/src/file-download.ts | 27 ++ service/src/file-object-resolver.test.ts | 50 ++++ service/src/file-object-resolver.ts | 75 +++++ service/src/file-server.ts | 160 ++++------ service/src/test/redis.ts | 45 +++ 30 files changed, 1690 insertions(+), 387 deletions(-) create mode 100644 api/src/http-input-cache.test.ts create mode 100644 api/src/http-input-cache.ts create mode 100644 api/src/input-manifest.test.ts create mode 100644 docs/INPUT_REUSE.md create mode 100644 service/src/egress-ledger-reconnect.test.ts create mode 100644 service/src/egress-ledger-script.ts create mode 100644 service/src/file-download.test.ts create mode 100644 service/src/file-download.ts create mode 100644 service/src/file-object-resolver.test.ts create mode 100644 service/src/file-object-resolver.ts create mode 100644 service/src/test/redis.ts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 71536caa..f041e2d6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -139,6 +139,9 @@ jobs: - name: Install dependencies run: bun ci + - name: Install Redis for ledger integration tests + run: sudo apt-get update && sudo apt-get install -y redis-server + - name: Build service run: bun run build diff --git a/api/src/config.ts b/api/src/config.ts index bbc2c184..f3b935ff 100644 --- a/api/src/config.ts +++ b/api/src/config.ts @@ -122,6 +122,9 @@ export const config = { /* Ceiling for the pushed input cache (session-inputs.ts). Eviction is * always safe — a miss simply re-pushes on the next probe — so this is a * disk guard, not a correctness knob. */ + http_input_cache_enabled: process.env.SANDBOX_HTTP_INPUT_CACHE_ENABLED === 'true', + http_input_cache_max_objects: safeInt(process.env.SANDBOX_HTTP_INPUT_CACHE_MAX_OBJECTS, 4096), + http_input_cache_max_inflight: safeInt(process.env.SANDBOX_HTTP_INPUT_CACHE_MAX_INFLIGHT, 16), input_cache_max_bytes: safeInt( process.env.SANDBOX_INPUT_CACHE_MAX_BYTES, 512 * 1024 * 1024, diff --git a/api/src/download.test.ts b/api/src/download.test.ts index bb8ec7e9..da877dad 100644 --- a/api/src/download.test.ts +++ b/api/src/download.test.ts @@ -2,6 +2,8 @@ import { describe, it, expect, beforeEach, afterEach, beforeAll, afterAll, spyOn import * as fsp from 'fs/promises'; import * as path from 'path'; import * as os from 'os'; +import { createHash, randomUUID } from 'node:crypto'; +import { SESSION_INPUT_CACHE_DIR } from './session-inputs'; import * as semver from 'semver'; import { Job, SessionWorkspaceDirtyError, type TFile } from './job'; import type { Runtime } from './runtime'; @@ -511,6 +513,48 @@ describe('downloadAndWriteFile / RFC 5987 round-trip', () => { } }); + it('reuses versioned bytes in fresh workspaces without bypassing a later denial', async () => { + const previousCache = config.http_input_cache_enabled; + const version = randomUUID(); + const cacheKey = createHash('sha256').update(version).digest('hex'); + const otherDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'codeapi-cache-second-')); + config.http_input_cache_enabled = true; + config.egress_gateway_url = `http://127.0.0.1:${serverPort}`; + let reads = 0; + let checks = 0; + const meta: Route = { status: 200, body: JSON.stringify({ cacheable: true, cacheKey, version, size: 8, readOnly: false }), + onRequest: () => { checks++; }, + }; + routes.set('/sessions/previous/objects/cached/metadata', meta); + routes.set('/sessions/previous/objects/cached', { status: 200, body: 'original', + headers: { 'X-CodeAPI-Input-Version': version }, + onRequest: request => { reads++; expect(request.headers.get('x-codeapi-input-version')).toBe(version); }, + }); + const file: TFile = { id: 'cached', storage_session_id: 'previous', name: 'data.txt', input_cache_key: cacheKey }; + try { + const first = makeJob([file]); + asInternals(first).submissionDir = tmpDir; + await first.downloadAndWriteFile(file); + await fsp.writeFile(path.join(tmpDir, 'data.txt'), 'sandbox changed this'); + const second = makeJob([file]); + asInternals(second).submissionDir = otherDir; + await second.downloadAndWriteFile(file); + expect(await fsp.readFile(path.join(otherDir, 'data.txt'), 'utf8')).toBe('original'); + expect(reads).toBe(1); + expect(checks).toBe(2); + meta.status = 403; + meta.headers = { 'X-CodeAPI-Error-Code': 'scope_mismatch' }; + await expect(second.downloadAndWriteFile(file)).rejects.toThrow('HTTP error: 403'); + expect(checks).toBe(3); + expect(reads).toBe(1); + } finally { + config.http_input_cache_enabled = previousCache; + await fsp.rm(otherDir, { recursive: true, force: true }); + await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, cacheKey), { force: true }); + await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, `${cacheKey}.json`), { force: true }); + } + }); + it('does not retry an unclassified direct file-server denial', async () => { config.egress_gateway_url = ''; let requests = 0; diff --git a/api/src/http-input-cache.test.ts b/api/src/http-input-cache.test.ts new file mode 100644 index 00000000..674a03ab --- /dev/null +++ b/api/src/http-input-cache.test.ts @@ -0,0 +1,163 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { createHash, randomUUID } from 'node:crypto'; +import { rm } from 'node:fs/promises'; +import path from 'node:path'; +import { fetchCachedHttpInput } from './http-input-cache'; +import { hasCachedInput, SESSION_INPUT_CACHE_DIR } from './session-inputs'; + +const keys = new Set(); +afterEach(async () => { + for (const key of keys) { + await rm(path.join(SESSION_INPUT_CACHE_DIR, key), { force: true }); + await rm(path.join(SESSION_INPUT_CACHE_DIR, `${key}.json`), { force: true }); + } + keys.clear(); +}); +function fixture(body = 'input', principal = 'tenant/user') { + const version = randomUUID(); + const cacheKey = createHash('sha256').update(principal + version).digest('hex'); + keys.add(cacheKey); + const meta = { cacheable: true, version, cacheKey, size: Buffer.byteLength(body), readOnly: false, name: 'input.txt' }; + let reads = 0; + let authorizations = 0; + return { + meta, + counts: () => ({ reads, authorizations }), + args: { + maxBytes: 8192, maxObjects: 2, maxFileBytes: 8192, maxInflight: 4, + metadata: async () => { authorizations++; return Response.json(meta); }, + download: async (expected: string, _signal: AbortSignal) => { + reads++; + expect(expected).toBe(version); + return new Response(body, { headers: { 'X-CodeAPI-Input-Version': version } }); + }, + }, + }; +} + +function gate() { + let release!: () => void; + const promise = new Promise(resolve => { release = resolve; }); + return { promise, release }; +} + +describe('authorized HTTP input cache', () => { + test('fresh executions reuse bytes but authorize every hit', async () => { + const f = fixture(); + for (let i = 0; i < 3; i++) { + const response = await fetchCachedHttpInput(f.args); + expect(await response?.text()).toBe('input'); + expect(response?.headers.get('content-disposition')).toContain('input.txt'); + } + expect(f.counts()).toEqual({ reads: 1, authorizations: 3 }); + expect(await hasCachedInput('', '', f.meta.cacheKey)).toBe(false); // Cannot bypass preflight using a pushed key. + expect(await hasCachedInput('', '', f.meta.cacheKey, 'http')).toBe(true); + const denied = await fetchCachedHttpInput({ ...f.args, + metadata: async () => new Response(null, { status: 403, headers: { 'X-CodeAPI-Error-Code': 'scope_mismatch' } }), + }); + expect(denied?.status).toBe(403); + expect(f.counts().reads).toBe(1); + }); + + test('new versions and principals never reuse an existing version key', async () => { + for (const [body, principal] of [['old', 'tenant/user'], ['new', 'tenant/user'], ['private', 'another-tenant/user']]) { + const f = fixture(body, principal); + expect(await (await fetchCachedHttpInput(f.args))?.text()).toBe(body); + expect(f.counts().reads).toBe(1); + } + }); + + test('coalesces misses while one cancelled caller leaves the remaining reader intact', async () => { + const f = fixture(); + const started = gate(); + const finish = gate(); + let downloads = 0; + let sharedSignal: AbortSignal | undefined; + const args = { ...f.args, download: async (version: string, signal: AbortSignal) => { + downloads++; sharedSignal = signal; started.release(); + await finish.promise; + return f.args.download(version, signal); + } }; + const controller = new AbortController(); + const first = fetchCachedHttpInput({ ...args, signal: controller.signal }); + const second = fetchCachedHttpInput(args); + await started.promise; + await new Promise(resolve => setTimeout(resolve, 10)); + controller.abort(new Error('first cancelled')); + await expect(first).rejects.toThrow('first cancelled'); + expect(sharedSignal?.aborted).toBe(false); + finish.release(); + expect(await (await second)?.text()).toBe('input'); + expect(downloads).toBe(1); + expect(f.counts().authorizations).toBe(2); + }); + + test('a shared fill does not propagate its creator grant denial to a valid waiter', async () => { + const f = fixture(); + const started = gate(); + const finish = gate(); + const denied = fetchCachedHttpInput({ ...f.args, download: async () => { + started.release(); + await finish.promise; + return new Response(null, { status: 403, headers: { 'X-CodeAPI-Error-Code': 'request_budget_exceeded' } }); + } }); + await started.promise; + const valid = fetchCachedHttpInput(f.args); + await Bun.sleep(10); + finish.release(); + expect((await denied)?.status).toBe(403); + // Job.fetchInputObject uses the waiter's own normal download on undefined. + expect(await valid).toBeUndefined(); + expect(await (await f.args.download(f.meta.version, new AbortController().signal)).text()).toBe('input'); + expect(f.counts().authorizations).toBe(2); + }); + + test('last-reader cancellation aborts the upstream fill without publishing', async () => { + const f = fixture(); + const started = gate(); + const aborted = gate(); + const controller = new AbortController(); + const pending = fetchCachedHttpInput({ ...f.args, signal: controller.signal, + download: async (_version, signal) => { + started.release(); + return new Promise((_resolve, reject) => signal.addEventListener('abort', () => { + aborted.release(); reject(signal.reason); + }, { once: true })); + }, + }); + await started.promise; + controller.abort(new Error('cancel fill')); + await expect(pending).rejects.toThrow('cancel fill'); + await aborted.promise; + expect(await hasCachedInput('', '', f.meta.cacheKey, 'http')).toBe(false); + }); + + test('changed-version and oversized responses are never published', async () => { + const f = fixture(); + const changed = await fetchCachedHttpInput({ ...f.args, download: async () => new Response('changed', { + headers: { 'X-CodeAPI-Input-Version': randomUUID() }, + }) }); + expect(changed).toBeUndefined(); + expect(await hasCachedInput('', '', f.meta.cacheKey, 'http')).toBe(false); + await expect(fetchCachedHttpInput({ ...f.args, download: async () => new Response('too many bytes', { + headers: { 'X-CodeAPI-Input-Version': f.meta.version }, + }) })).rejects.toThrow(); + expect(await hasCachedInput('', '', f.meta.cacheKey, 'http')).toBe(false); + }); + + test('cache quotas evict old entries and preserve an already-open reader', async () => { + const first = fixture('a'.repeat(4000)); + const second = fixture('b'.repeat(4000)); + const response = await fetchCachedHttpInput({ ...first.args, maxObjects: 1 }); + expect(await (await fetchCachedHttpInput({ ...second.args, maxObjects: 1 }))?.text()).toBe('b'.repeat(4000)); + expect(await hasCachedInput('', '', first.meta.cacheKey, 'http')).toBe(false); + expect(await response?.text()).toBe('a'.repeat(4000)); + }); + + test('legacy metadata protocols fall back without a cache read or fill', async () => { + const f = fixture(); + expect(await fetchCachedHttpInput({ ...f.args, metadata: async () => new Response(null, { status: 404 }) })).toBeUndefined(); + expect(await fetchCachedHttpInput({ ...f.args, metadata: async () => Response.json({ cacheable: false }) })).toBeUndefined(); + expect(f.counts().reads).toBe(0); + }); +}); diff --git a/api/src/http-input-cache.ts b/api/src/http-input-cache.ts new file mode 100644 index 00000000..c785bdf0 --- /dev/null +++ b/api/src/http-input-cache.ts @@ -0,0 +1,177 @@ +import { Readable } from 'node:stream'; +import { createGzip } from 'node:zlib'; +import { httpInputCacheEvents } from './metrics'; +import { cachedInputResponse, openCachedInput, storeCachedInputs } from './session-inputs'; + +type Metadata = { cacheable: true; cacheKey: string; version: string; size: number; name?: string; readOnly: boolean }; +type FillResult = { stored: boolean; status?: number; headers?: Headers }; +type Fill = { controller: AbortController; users: number; result: Promise }; +const fills = new Map(); + +function validMetadata(value: unknown, maxBytes: number): value is Metadata { + if (!value || typeof value !== 'object') return false; + const m = value as Metadata; + return m.cacheable === true && typeof m.cacheKey === 'string' && /^[0-9a-f]{64}$/.test(m.cacheKey) && + typeof m.version === 'string' && /^[0-9a-f-]{36}$/.test(m.version) && + Number.isSafeInteger(m.size) && m.size >= 0 && m.size + 1024 <= maxBytes && + typeof m.readOnly === 'boolean' && (m.name === undefined || (typeof m.name === 'string' && m.name.length <= 4096)); +} + +function tarHeader(name: string, bytes: number): Buffer { + const header = Buffer.alloc(512); + header.write(name, 0, 100, 'utf8'); + header.write('0000600\0', 100, 8, 'ascii'); + header.write(bytes.toString(8).padStart(11, '0') + '\0', 124, 12, 'ascii'); + header.fill(32, 148, 156); + header[156] = 48; + header.write('ustar\0', 257, 6, 'ascii'); + const checksum = header.reduce((sum, byte) => sum + byte, 0); + header.write(checksum.toString(8).padStart(6, '0') + '\0 ', 148, 8, 'ascii'); + return header; +} + +/** Reuse the pushed-cache writer's staging, quota, no-follow and atomic commit + * rules. No workspace pathname or sandbox-visible file is used as cache input. */ +async function fillCache(response: Response, meta: Metadata, maxBytes: number, maxObjects: number): Promise { + if (!response.body) throw new Error('Input response has no body'); + const body = response.body; + const sidecar = Buffer.from(JSON.stringify({ readOnly: meta.readOnly, source: 'http' })); + async function* archive(): AsyncGenerator { + yield tarHeader(meta.cacheKey, meta.size); + const reader = body.getReader(); + let bytes = 0; + try { + for (;;) { + const part = await reader.read(); + if (part.done) break; + bytes += part.value.byteLength; + if (bytes > meta.size) throw new Error('Input exceeded its authorized metadata size'); + yield Buffer.from(part.value); + } + if (bytes !== meta.size) throw new Error('Input size changed during preparation'); + } finally { + await reader.cancel().catch(() => {}); + reader.releaseLock(); + } + yield Buffer.alloc((512 - meta.size % 512) % 512); + yield tarHeader(`${meta.cacheKey}.json`, sidecar.length); + yield sidecar; + yield Buffer.alloc((512 - sidecar.length % 512) % 512); + yield Buffer.alloc(1024); + } + const source = Readable.from(archive()); + const compressed = createGzip(); + compressed.on('error', () => {}); // Queue admission checks an already-failed stream. + source.on('error', error => compressed.destroy(error)); + source.pipe(compressed); + try { + await storeCachedInputs(compressed, maxBytes, meta.size + sidecar.length, maxObjects); + } finally { + source.destroy(); + compressed.destroy(); + await body.cancel().catch(() => {}); + } +} + +async function waitForFill(fill: Fill, signal?: AbortSignal): Promise { + if (signal?.aborted) { + if (fill.users === 0) fill.controller.abort(signal.reason); + signal.throwIfAborted(); + } + fill.users++; + let abort: (() => void) | undefined; + try { + const cancelled = new Promise((_resolve, reject) => { + abort = () => reject(signal?.reason ?? new Error('Input preparation cancelled')); + signal?.addEventListener('abort', abort, { once: true }); + if (signal?.aborted) abort(); + }); + return await Promise.race([fill.result, cancelled]); + } finally { + if (abort) signal?.removeEventListener('abort', abort); + if (--fill.users === 0) fill.controller.abort(new Error('No input-cache consumers remain')); + } +} + +/** Every caller performs its own scoped preflight, even for hits or shared fills. + * Only opaque version keys returned by that authorized gateway enter this cache. */ +export async function fetchCachedHttpInput(args: { + metadata(): Promise; + download(version: string, signal: AbortSignal): Promise; + signal?: AbortSignal; + maxBytes: number; + maxFileBytes: number; + maxInflight: number; + maxObjects: number; +}): Promise { + args.signal?.throwIfAborted(); + const preflight = await args.metadata(); + if (preflight.status === 404 || preflight.status === 405) { + await preflight.body?.cancel(); + httpInputCacheEvents.inc({ event: 'legacy_bypass' }); + return undefined; // Older gateway/relay: retain the uncached protocol. + } + if (!preflight.ok) { httpInputCacheEvents.inc({ event: 'preflight_failure' }); return preflight; } + const value: unknown = await preflight.json(); + if (!validMetadata(value, args.maxBytes) || value.size > args.maxFileBytes) { + httpInputCacheEvents.inc({ event: 'uncacheable' }); + return undefined; + } + const meta = value; + args.signal?.throwIfAborted(); + let cached = await openCachedInput('', '', meta.cacheKey, 'http'); + if (cached) httpInputCacheEvents.inc({ event: 'hit' }); + if (!cached) { + let fill = fills.get(meta.cacheKey); + const joinedExistingFill = fill !== undefined; + if (fill) httpInputCacheEvents.inc({ event: 'coalesced' }); + if (!fill) { + if (fills.size >= args.maxInflight) { + httpInputCacheEvents.inc({ event: 'capacity_bypass' }); + return undefined; + } + httpInputCacheEvents.inc({ event: 'fill' }); + const controller = new AbortController(); + fill = { controller, users: 0, result: Promise.resolve({ stored: false }) }; + const ownFill = fill; + fills.set(meta.cacheKey, ownFill); + ownFill.result = (async (): Promise => { + const response = await args.download(meta.version, controller.signal); + if (!response.ok) { + await response.body?.cancel(); + return { stored: false, status: response.status, headers: response.headers }; + } + if (response.headers.get('x-codeapi-input-version') !== meta.version || + (response.headers.get('x-read-only')?.toLowerCase() === 'true') !== meta.readOnly) { + await response.body?.cancel(); + return { stored: false }; // Old file server or inconsistent metadata: never publish. + } + await fillCache(response, meta, args.maxBytes, args.maxObjects); + return { stored: true }; + })().finally(() => { + if (fills.get(meta.cacheKey) === ownFill) fills.delete(meta.cacheKey); + }); + // A caller can cancel between admission and waiting; avoid an unhandled rejection. + void ownFill.result.catch(() => {}); + } + const result = await waitForFill(fill, args.signal); + if (result.status) { + // The transfer used the creator's grant. Its denial/budget must not reject + // another caller whose own preflight succeeded; use that caller's fetch. + if (joinedExistingFill) return undefined; + const headers = new Headers(result.headers); + headers.delete('content-length'); + return new Response(null, { status: result.status, headers }); + } + if (!result.stored) return undefined; + cached = await openCachedInput('', '', meta.cacheKey, 'http'); + } + if (!cached) return undefined; // Evicted between commit and open: normal download remains correct. + if (args.signal?.aborted) { + await cached.handle.close(); + args.signal.throwIfAborted(); + } + const response = cachedInputResponse(cached); + if (meta.name) response.headers.set('content-disposition', `attachment; filename*=UTF-8''${encodeURIComponent(meta.name)}`); + return response; +} diff --git a/api/src/input-manifest.test.ts b/api/src/input-manifest.test.ts new file mode 100644 index 00000000..db55ca4a --- /dev/null +++ b/api/src/input-manifest.test.ts @@ -0,0 +1,124 @@ +import { afterEach, expect, test } from 'bun:test'; +import { createHash, randomUUID } from 'node:crypto'; +import * as fsp from 'node:fs/promises'; +import * as os from 'node:os'; +import * as path from 'node:path'; +import { config } from './config'; +import { Job } from './job'; +import { SESSION_INPUT_CACHE_DIR } from './session-inputs'; +import { fallbackSandboxIdentity } from './workspace-isolation'; + +const originalFetch = globalThis.fetch; +const originalConfig = { http_input_cache_enabled: config.http_input_cache_enabled, egress_gateway_url: config.egress_gateway_url }; +const dirs: string[] = []; +const keys: string[] = []; +afterEach(async () => { + globalThis.fetch = originalFetch; + Object.assign(config, originalConfig); + await Promise.all(dirs.splice(0).map(dir => fsp.rm(dir, { recursive: true, force: true }))); + await Promise.all(keys.splice(0).flatMap(key => [key, `${key}.json`]).map(key => fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, key), { force: true }))); +}); + +async function fixture(count: number, mode: 'batch' | 'legacy' | 'race' = 'batch') { + config.http_input_cache_enabled = true; + config.egress_gateway_url = 'http://manifest.test'; + let manifests = 0, singlePreflights = 0, downloads = 0; + const version = randomUUID(); + const freshVersion = randomUUID(); + const metadata = (id: string, v = version) => { + const key = createHash('sha256').update(id + v).digest('hex'); + keys.push(key); + return { cacheable: true, cacheKey: key, version: v, size: 5, readOnly: false }; + }; + globalThis.fetch = (async (url: RequestInfo | URL, init?: RequestInit) => { + const pathname = new URL(String(url)).pathname; + if (pathname.endsWith('/objects')) return Response.json([]); + if (pathname === '/input-manifest') { + manifests++; + if (mode === 'legacy') return new Response(null, { status: 404 }); + const request = JSON.parse(init!.body as string) as { files: { objectHandle: string }[] }; + expect(request.files).toHaveLength(count); + return Response.json({ files: request.files.map(file => metadata(file.objectHandle)) }); + } + if (pathname.endsWith('/metadata')) { + singlePreflights++; + return Response.json(metadata(pathname.split('/').slice(-2)[0], mode === 'race' ? freshVersion : version)); + } + downloads++; + if (mode === 'race' && new Headers(init?.headers).get('X-CodeAPI-Input-Version') === version) { + return new Response(null, { status: 409 }); + } + return new Response('bytes', { headers: { 'X-CodeAPI-Input-Version': mode === 'race' ? freshVersion : version } }); + }) as typeof fetch; + async function prime(egressGrant = 'test-grant') { + const dir = await fsp.mkdtemp(path.join(os.tmpdir(), 'manifest-prime-')); + dirs.push(dir); + const session = { + runtimeSessionId: 'test', acquire: async () => ({ dir, workspaceId: 'test', identity: fallbackSandboxIdentity() }), + primedInputId: () => undefined, markPrimed: () => {}, markDirty: () => {}, + }; + const job = new Job({ + session_id: 'test', egress_grant: egressGrant, runtime: { language: 'bash', version: '5.0.0', aliases: [] }, + files: Array.from({ length: count }, (_, i) => ({ id: `f${i}`, storage_session_id: 's', name: `file${i}.txt` })), + args: [], stdin: '', timeouts: { run: 5000, compile: 5000 }, cpu_times: { run: 5000, compile: 5000 }, + memory_limits: { run: 128e6, compile: 128e6 }, session, + } as never); + (job as unknown as { log: { level: string } }).log.level = 'silent'; + await job.prime(); + expect(await fsp.readFile(path.join(dir, 'file0.txt'), 'utf8')).toBe('bytes'); + } + return { prime, counts: () => ({ manifests, singlePreflights, downloads }) }; +} + +test('240 inputs use one authorized manifest per fresh workspace and reuse only content', async () => { + const f = await fixture(240); + await f.prime(); + await f.prime(); + expect(f.counts()).toEqual({ manifests: 2, singlePreflights: 0, downloads: 240 }); +}, 30000); + +test('an older gateway falls back to independently authorized preflights', async () => { + const f = await fixture(2, 'legacy'); + await f.prime(); + expect(f.counts()).toEqual({ manifests: 1, singlePreflights: 2, downloads: 2 }); +}); + +test('a raced version consumes the batch entry and retries against fresh metadata', async () => { + const f = await fixture(1, 'race'); + await f.prime(); + expect(f.counts()).toEqual({ manifests: 1, singlePreflights: 1, downloads: 2 }); +}); + + +test('one execution grant denial cannot fail a coalesced execution with its own valid grant', async () => { + const f = await fixture(1); + const underlying = globalThis.fetch; + let release!: () => void; + let started!: () => void; + const blocked = new Promise(resolve => { release = resolve; }); + const creatorStarted = new Promise(resolve => { started = resolve; }); + let deniedDownloads = 0, validDownloads = 0; + globalThis.fetch = (async (url: RequestInfo | URL, init?: RequestInit) => { + if (String(url).endsWith('/objects/f0')) { + if (new Headers(init?.headers).get('X-CodeAPI-Egress-Grant') === 'denied') { + deniedDownloads++; + started(); await blocked; + return new Response(null, { status: 403, headers: { 'X-CodeAPI-Error-Code': 'scope_mismatch' } }); + } + validDownloads++; + } + return underlying(url, init); + }) as typeof fetch; + const creator = f.prime('denied'); + void creator.catch(() => {}); + await creatorStarted; + const waiter = f.prime('valid'); + try { + while (f.counts().manifests < 2) await Bun.sleep(1); + await Bun.sleep(20); + } finally { release(); } + const result = await Promise.allSettled([creator, waiter]); + expect(result.map(item => item.status)).toEqual(['rejected', 'fulfilled']); + expect(deniedDownloads).toBe(1); + expect(validDownloads).toBe(1); +}); diff --git a/api/src/job.ts b/api/src/job.ts index bdcd5061..d2504de9 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -42,6 +42,7 @@ import { validateFilePath, isValidFilePath, } from './validation'; +import { fetchCachedHttpInput } from './http-input-cache'; import { cachedInputResponse, inputCacheKey, openCachedInput } from './session-inputs'; export { @@ -733,6 +734,7 @@ export class Job { private sessionFiles: FileRef[] = []; private inheritedRefs: FileRef[] = []; private inputFileHashes = new Map(); + private inputManifest = new Map(); private inputDestinations = new Map(); private entryPointName: string | undefined; private chmoddedDirs = new Set(); @@ -938,6 +940,7 @@ export class Job { async prime(): Promise { this.inputDestinations.clear(); + this.inputManifest.clear(); const requestedDestinations = new Map(); for (const file of this.files) { validateFilePath(file.name, '/tmp/codeapi-request-validation'); @@ -990,6 +993,8 @@ export class Job { await this.autoLoadDirkeep(); } + await this.prepareInputManifest(); + /* Promise.all rejects as soon as one operation fails, while its siblings * keep running. The route's finally then calls cleanup(), which clears the * session path/identity. A delayed sibling used to resume afterward and @@ -1431,6 +1436,38 @@ export class Job { throw lastError ?? new Error(`Failed to download input ${file.id}`); } + private async prepareInputManifest(): Promise { + if (!config.http_input_cache_enabled || !config.egress_gateway_url) return; + const files = this.files.filter(file => file.id && file.storage_session_id); + if (!files.length) return; + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), AUTO_LOAD_DIRKEEP_TIMEOUT_MS); + try { + const response = await fetch(`${this.fileEgressBaseUrl()}/input-manifest`, { + method: 'POST', headers: this.fileEgressHeaders({ 'Content-Type': 'application/json' }), + body: JSON.stringify({ files: files.map(file => ({ + sessionHandle: file.storage_session_id, objectHandle: file.id, + })) }), signal: controller.signal, + }); + if (!response.ok) { + await response.body?.cancel(); + return; // Older gateways/relays and transient failures use per-file preflight. + } + const manifest = await response.json() as { files?: unknown[] }; + if (!Array.isArray(manifest.files) || manifest.files.length !== files.length) return; + manifest.files.forEach((metadata, index) => { + if (metadata && typeof metadata === 'object' && !('retry' in metadata)) { + this.inputManifest.set(files[index], metadata); + } + }); + } catch { + // This is an optimization. Individual reads still authorize and report failures. + } finally { + clearTimeout(timeout); + controller.abort(); + } + } + /** * Resolves an input object's bytes, preferring the runner-local cache the * control plane pushes into on backends whose sandbox cannot reach the file @@ -1462,6 +1499,24 @@ export class Job { `Input ${file.id} was not delivered to the sandbox and no file server is reachable`, ); } + if (config.http_input_cache_enabled && config.egress_gateway_url) { + const response = await fetchCachedHttpInput({ + metadata: () => { + const metadata = this.inputManifest.get(file); + // A version-race retry must obtain a new authorized storage version. + this.inputManifest.delete(file); + return metadata === undefined + ? fetch(`${this.buildDownloadUrl(file)}/metadata`, { headers: this.fileEgressHeaders(), signal }) + : Promise.resolve(Response.json(metadata)); + }, + download: (version, sharedSignal) => fetch(this.buildDownloadUrl(file), { + headers: this.fileEgressHeaders({ 'X-CodeAPI-Input-Version': version }), signal: sharedSignal, + }), + signal, maxBytes: config.input_cache_max_bytes, maxFileBytes: config.max_file_size, + maxInflight: config.http_input_cache_max_inflight, maxObjects: config.http_input_cache_max_objects, + }); + if (response) return response; + } return fetch(this.buildDownloadUrl(file), { headers: this.fileEgressHeaders(), signal, diff --git a/api/src/metrics.ts b/api/src/metrics.ts index 2da33b27..328b4871 100644 --- a/api/src/metrics.ts +++ b/api/src/metrics.ts @@ -16,6 +16,12 @@ const httpRequestDuration = new Histogram({ buckets: [0.01, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60], }); +export const httpInputCacheEvents = new Counter({ + name: 'codeapi_sandbox_http_input_cache_events_total', + help: 'Authorized HTTP input cache events; fills and failures may both occur for one input', + labelNames: ['event'] as const, +}); + export const sandboxExecutions = new Counter({ name: 'codeapi_sandbox_executions_total', help: 'Total number of sandbox execution attempts by outcome', diff --git a/api/src/session-inputs.ts b/api/src/session-inputs.ts index 0f1b7f3a..cc9f4b5a 100644 --- a/api/src/session-inputs.ts +++ b/api/src/session-inputs.ts @@ -147,6 +147,8 @@ export interface CachedInputMeta { * ref resolve to the first ref's path — which then overwrote a file the * sandbox had edited. */ readOnly: boolean; + /** HTTP entries require a fresh gateway preflight and cannot satisfy pushed-input probes. */ + source?: 'http'; } export interface CachedInput { @@ -176,7 +178,9 @@ function parseCachedInputMeta(raw: string): CachedInputMeta | null { ) { return null; } - return { readOnly: (parsed as { readOnly: boolean }).readOnly }; + const source = (parsed as { source?: unknown }).source; + if (source !== undefined && source !== 'http') return null; + return { readOnly: (parsed as { readOnly: boolean }).readOnly, ...(source === 'http' ? { source } : {}) }; } catch { return null; } @@ -186,8 +190,9 @@ export async function hasCachedInput( storageSessionId: string, id: string, cacheKey?: string, + source: 'push' | 'http' = 'push', ): Promise { - const opened = await openCachedInput(storageSessionId, id, cacheKey); + const opened = await openCachedInput(storageSessionId, id, cacheKey, source); if (!opened) return false; await opened.handle.close(); return true; @@ -197,6 +202,7 @@ export async function openCachedInput( storageSessionId: string, id: string, cacheKey?: string, + source: 'push' | 'http' = 'push', ): Promise { const key = cacheKey ?? inputCacheKey(storageSessionId, id); if (!/^[0-9a-f]{64}$/.test(key)) return null; @@ -229,7 +235,7 @@ export async function openCachedInput( await metaHandle?.close().catch(() => {}); } const meta = raw === null ? null : parseCachedInputMeta(raw); - if (!meta) { + if (!meta || (meta.source ?? 'push') !== source) { logger.warn({ key }, 'Ignoring session input with missing or invalid metadata'); await handle.close(); return null; @@ -325,7 +331,13 @@ async function extractInputArchive( }; compressedGuard.on('error', forwardCompressedError); body.once('error', forwardBodyError); - body.pipe(compressedGuard).pipe(gunzip); + const sourceState = body as NodeJS.ReadableStream & { errored?: Error; destroyed?: boolean }; + compressedGuard.pipe(gunzip); + if (sourceState.errored || sourceState.destroyed) { + forwardBodyError(sourceState.errored ?? new Error('Input stream was cancelled before extraction')); + } else { + body.pipe(compressedGuard); + } let buffered = Buffer.alloc(0); let current: @@ -517,7 +529,10 @@ async function storeCachedInputsOnce( body: NodeJS.ReadableStream, maxBytes = Number.MAX_SAFE_INTEGER, expectedBytes?: number, + maxObjects = Number.MAX_SAFE_INTEGER, ): Promise { + const readable = body as NodeJS.ReadableStream & { errored?: Error; destroyed?: boolean }; + if (readable.errored || readable.destroyed) throw readable.errored ?? new Error('Input stream was cancelled before cache admission'); await fsp.mkdir(SESSION_INPUT_CACHE_DIR, { recursive: true, mode: 0o700 }); const cacheStat = await fsp.lstat(SESSION_INPUT_CACHE_DIR); if (!cacheStat.isDirectory() || cacheStat.isSymbolicLink()) { @@ -571,6 +586,8 @@ async function storeCachedInputsOnce( } } + if (keys.length > maxObjects) throw new Error('Input batch exceeds cache object limit'); + await pruneInputCache(Math.max(0, maxBytes - stagedBytes), maxObjects - keys.length); let stored = 0; /* Commit sidecars before data. A new key remains a probe miss until both * exist; replacing an immutable key can only expose its new validated @@ -596,6 +613,7 @@ export async function storeCachedInputs( body: NodeJS.ReadableStream, maxBytes = Number.MAX_SAFE_INTEGER, expectedBytes?: number, + maxObjects = Number.MAX_SAFE_INTEGER, ): Promise { /* Concurrent pushes otherwise each budget only its own staging tree and can * collectively recreate the same transient disk spike. Queue extraction; @@ -607,7 +625,7 @@ export async function storeCachedInputs( }); await previous; try { - return await storeCachedInputsOnce(body, maxBytes, expectedBytes); + return await storeCachedInputsOnce(body, maxBytes, expectedBytes, maxObjects); } finally { release(); } @@ -615,7 +633,7 @@ export async function storeCachedInputs( /** Drops least-recently-used entries until the cache fits `maxBytes`. Eviction * is always safe: a miss simply re-pushes on the next probe. */ -export async function pruneInputCache(maxBytes: number): Promise { +export async function pruneInputCache(maxBytes: number, maxObjects = Number.MAX_SAFE_INTEGER): Promise { const names = await fsp.readdir(SESSION_INPUT_CACHE_DIR).catch(() => [] as string[]); const nameSet = new Set(names.filter(name => ENTRY_PATTERN.test(name))); const pairs: Array<{ key: string; size: number; atime: number }> = []; @@ -644,14 +662,16 @@ export async function pruneInputCache(maxBytes: number): Promise { await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, orphan), { force: true }).catch(() => {}); } } - if (total <= maxBytes) return; + let objects = pairs.length; + if (total <= maxBytes && objects <= maxObjects) return; pairs.sort((a, b) => a.atime - b.atime); for (const pair of pairs) { - if (total <= maxBytes) break; + if (total <= maxBytes && objects <= maxObjects) break; await fsp.rm(path.join(SESSION_INPUT_CACHE_DIR, pair.key), { force: true }).catch(() => {}); await fsp .rm(path.join(SESSION_INPUT_CACHE_DIR, `${pair.key}${META_SUFFIX}`), { force: true }) .catch(() => {}); total -= pair.size; + objects -= 1; } } diff --git a/docs/INPUT_REUSE.md b/docs/INPUT_REUSE.md new file mode 100644 index 00000000..40e277c4 --- /dev/null +++ b/docs/INPUT_REUSE.md @@ -0,0 +1,91 @@ +# Bounded input reuse for stateless executions + +Fresh execution workspaces can reuse input contents without retaining a mutable conversation sandbox. Each reader first authorizes a metadata request through the egress gateway. New file-server uploads carry a random `codeapi-version` metadata value that changes on every PUT, including overwrites with identical contents. The gateway binds a cache key to that version, storage identity, tenant, user, size, filename, and read-only flag. + +```mermaid +sequenceDiagram + participant R as Runner + participant G as Egress gateway + participant F as File server + participant C as Protected input cache + R->>G: POST bounded input manifest (one per execution) + G->>G: Verify grant, scope, expiry, revocation, budget + G->>F: Resolve current metadata with bounded concurrency + F-->>G: Current upload version and metadata + G-->>R: Ordered principal-scoped version keys + R->>C: Open authorized version + alt Cache miss + R->>G: Download with expected version + G->>G: Authorize and account download + G->>F: Forward expected version + F-->>R: Exact GET metadata and bytes, or 409 if changed + R->>C: Stage, validate size, atomically publish + end + R->>R: Copy into fresh workspace using existing priming rules +``` + +## Invariants + +- A cache hit never authorizes an input. Every execution performs its own preflight, including readers joining a shared fill. Denied or revoked grants cannot use cached data. +- Every manifest handle is scope-checked before storage access. Revocation is checked again before returning resolved metadata. A deadline and disconnect cancel storage work; failures and older gateways fall back to independently authorized per-file preflights. A version-race retry discards its manifest entry. +- HTTP entries are marked separately from pushed inputs. Supplying an HTTP key in `input_cache_key` cannot bypass preflight through the older pushed-cache path. +- Cache files stay outside execution workspaces and sandbox mounts. Priming copies bytes; it never hard-links a writable workspace to trusted cache contents. Existing no-follow, read-only, hashing, atomic rename, and descriptor-pinning behavior remains in use. +- Concurrent authorized misses for the same version can share one download. Cancelling one reader does not cancel remaining readers; cancelling the last reader aborts the shared request. The number of fills, cached bytes, and object count are bounded. +- The downloader checks the version from the **actual GET**, rather than labeling bytes with metadata from an earlier HEAD. A raced overwrite returns 409 and preparation retries from current metadata. Legacy objects without a version use the uncached path. +- All writers of input objects must assign a fresh version on every overwrite. The file server does so for both upload routes. Checkpoint storage uses a separate path. Direct bucket writes that preserve an old version marker are outside this protocol. +- The optional Redis object-key index stores only a locator hint. It is not an authorization or metadata cache. Preflights still read current storage metadata; indexed keys must match the exact session and object identity. +- Shared download errors belong to the initiating grant. A coalesced caller falls back to its own authorized download rather than inheriting that grant's denial or exhausted budget. +- Redis reconnects never replay unfulfilled ledger mutations. A lost reply fails closed and may leave a conservatively charged counter/reservation until grant expiry; automatically refunding an ambiguous mutation could over-credit its budget. +- Full grant policy is no longer returned to the gateway for each authorization check. Atomic Redis scripts serialize accounting with revocation. Duplicate releases cannot repeatedly refund unrelated counters. Newly created compact ledgers keep immutable policy separate from mutable counters. + +## Configuration + +| Helm value | Environment variable | Default | +|---|---|---| +| `egressGrant.ledgerCompact` | `CODEAPI_EGRESS_LEDGER_COMPACT` | `false` | +| `egressGrant.inputManifestMaxFiles` | `CODEAPI_INPUT_MANIFEST_MAX_FILES` | `512` | +| `egressGrant.inputManifestConcurrency` | `CODEAPI_INPUT_MANIFEST_CONCURRENCY` | `8` | +| `egressGrant.inputManifestTimeoutMs` | `CODEAPI_INPUT_MANIFEST_TIMEOUT_MS` | `10000` | +| `fileServer.objectIndexEnabled` | `CODEAPI_FILE_OBJECT_INDEX_ENABLED` | `false` | +| `fileServer.metadataConcurrency` | `CODEAPI_FILE_METADATA_CONCURRENCY` | `1` | +| `workerSandbox.sandbox.httpInputCacheEnabled` | `SANDBOX_HTTP_INPUT_CACHE_ENABLED` | `false` | +| `workerSandbox.sandbox.httpInputCacheMaxInflight` | `SANDBOX_HTTP_INPUT_CACHE_MAX_INFLIGHT` | `16` | +| `workerSandbox.sandbox.httpInputCacheMaxObjects` | `SANDBOX_HTTP_INPUT_CACHE_MAX_OBJECTS` | `4096` | +| `workerSandbox.sandbox.inputCacheMaxBytes` | `SANDBOX_INPUT_CACHE_MAX_BYTES` | `536870912` | + +HTTP reuse requires a configured egress gateway. Cacheable objects are also bounded by the existing runner maximum file size. Cache capacity is local to each runner; eviction, restart, or routing to another runner causes a safe cache miss. The cache does not require persistent-session affinity. + +The manifest accepts at most 512 entries and a 4 MiB JSON body (protocol safety ceilings), with configured concurrency capped at 64. The runner bounds its opportunistic manifest request to 10 seconds, matching directory preparation, then uses per-file authorization if it cannot obtain a complete response. Oversized batches also fall back. Manifest requests remove repeated grant-header transfer and decoding, but still read current storage metadata for each file. + +Metadata listing concurrency preserves order and is capped at 64. A canary can use 8 after measuring storage load. This applies to directory-marker preparation as well; marker listings still happen and are not a retained conversation manifest. + +## Rollout and rollback + +1. Deploy the new binaries with feature flags off. Atomic accounting supports existing JSON ledgers, and legacy downloads retain metadata compatibility. The new file server stamps future uploads with versions. +2. Update **all** egress-gateway replicas before enabling compact ledgers. New binaries read both formats regardless of the creation flag. Older binaries cannot read compact hashes. To roll back to an older binary, disable compact creation, drain active grants, and wait their maximum TTL plus grace; never delete active ledgers to force a rollback. +3. Update all file-server writers before enabling the object-key index. Otherwise an older writer can change a locator without updating the index. Keep file-server replicas consistent during an indexed rollout. +4. Update the gateway, relay, runner, and launcher before enabling HTTP reuse on a small runner canary. Older gateway/relay metadata routes return 404/405 and fall back safely. Older files return `cacheable: false`. Keep the feature disabled for storage adapters that cannot return user metadata on GET. +5. Observe `codeapi_sandbox_http_input_cache_events_total` (bounded event labels, no identities), cold and warm preparation latency, storage/Redis operations, admission fairness, request budgets, and memory/disk pressure before widening the rollout. A successful manifest consumes one read request for the batch, matching the existing list-request accounting unit. Per-file compatibility preflights each consume a read request; each cold miss consumes an additional download request. Do not disable budget enforcement to accommodate a workload. +6. Disable HTTP reuse to return to normal downloads immediately. Cached files can age out normally; no workspace deletion or migration is needed. + +The creation flags default off. No deployment or object retention policy is changed by this code. Command grouping and persistent sessions remain independent options, not prerequisites for content reuse. Nothing deletes user inputs or infers shell dependencies. + +## Validation + +Ledger tests use an isolated real `redis-server` on a Unix socket with persistence disabled. Install Redis before running service tests. They cover legacy/compact formats, 240 concurrent reads against a strict budget, revocation, expiry, rejected uploads, duplicate releases, and format changes without resetting state. + +Focused commands: + +```sh +cd api +bun test src/input-manifest.test.ts src/http-input-cache.test.ts src/session-inputs.test.ts src/session-inputs.prime.test.ts src/download.test.ts src/inline-prime-atomicity.test.ts src/job-cleanup.test.ts +npx tsc --noEmit +``` + +```sh +cd service +bun test src/egress-ledger.test.ts src/egress-ledger-reconnect.test.ts src/egress-gateway.test.ts src/file-object-resolver.test.ts src/file-download.test.ts src/file-metadata.test.ts +npx tsc --noEmit +``` + +Run the code-package tests with Node (its supported test runner), plus launcher and deployment checks in CI. Cache regressions cover fresh-workspace reuse, cross-principal/version separation, denied preflights, pushed-key bypass prevention, coalesced cancellation, changed or oversized responses, and descriptor-safe eviction. Production latency targets must be validated with real regional storage latency and workload sizes; local synthetic results are not a production SLO. diff --git a/helm/codeapi/templates/egress-gateway-deployment.yaml b/helm/codeapi/templates/egress-gateway-deployment.yaml index 6e394909..49c865f0 100644 --- a/helm/codeapi/templates/egress-gateway-deployment.yaml +++ b/helm/codeapi/templates/egress-gateway-deployment.yaml @@ -40,6 +40,14 @@ spec: value: {{ .Values.hardenedSandboxMode | quote }} - name: CODEAPI_EGRESS_LEDGER_REQUIRED value: {{ .Values.egressGrant.ledgerRequired | quote }} + - name: CODEAPI_INPUT_MANIFEST_MAX_FILES + value: {{ .Values.egressGrant.inputManifestMaxFiles | quote }} + - name: CODEAPI_INPUT_MANIFEST_CONCURRENCY + value: {{ .Values.egressGrant.inputManifestConcurrency | quote }} + - name: CODEAPI_INPUT_MANIFEST_TIMEOUT_MS + value: {{ .Values.egressGrant.inputManifestTimeoutMs | quote }} + - name: CODEAPI_EGRESS_LEDGER_COMPACT + value: {{ .Values.egressGrant.ledgerCompact | quote }} - name: CODEAPI_EGRESS_LEDGER_TTL_GRACE_SECONDS value: {{ .Values.egressGrant.ledgerTtlGraceSeconds | quote }} - name: EGRESS_GATEWAY_PORT diff --git a/helm/codeapi/templates/file-server-deployment.yaml b/helm/codeapi/templates/file-server-deployment.yaml index 1dd95166..1a9e8531 100644 --- a/helm/codeapi/templates/file-server-deployment.yaml +++ b/helm/codeapi/templates/file-server-deployment.yaml @@ -51,6 +51,10 @@ spec: containerPort: {{ .Values.fileServer.service.port }} protocol: TCP env: + - name: CODEAPI_FILE_METADATA_CONCURRENCY + value: {{ .Values.fileServer.metadataConcurrency | quote }} + - name: CODEAPI_FILE_OBJECT_INDEX_ENABLED + value: {{ .Values.fileServer.objectIndexEnabled | quote }} {{ include "codeapi.otel.env" (dict "root" . "serviceName" "aiml-codeapi-file-server") | nindent 12 }} {{- if $useS3 }} # AWS S3 configuration (IRSA or static credentials) diff --git a/helm/codeapi/templates/worker-sandbox-deployment.yaml b/helm/codeapi/templates/worker-sandbox-deployment.yaml index 39c2e197..2819d998 100644 --- a/helm/codeapi/templates/worker-sandbox-deployment.yaml +++ b/helm/codeapi/templates/worker-sandbox-deployment.yaml @@ -314,6 +314,14 @@ spec: value: {{ (.Values.workerSandbox.sandbox.jobUidCount | default (.Values.workerSandbox.sandbox.maxConcurrentJobs | default (mul (.Values.workerSandbox.launcher.vcpus | default 2) 4))) | quote }} - name: SANDBOX_WORKSPACE_REAPER_MAX_AGE_SECONDS value: {{ (.Values.workerSandbox.sandbox.workspaceReaperMaxAgeSeconds | default 3600) | quote }} + - name: SANDBOX_HTTP_INPUT_CACHE_ENABLED + value: {{ .Values.workerSandbox.sandbox.httpInputCacheEnabled | quote }} + - name: SANDBOX_HTTP_INPUT_CACHE_MAX_INFLIGHT + value: {{ .Values.workerSandbox.sandbox.httpInputCacheMaxInflight | quote }} + - name: SANDBOX_HTTP_INPUT_CACHE_MAX_OBJECTS + value: {{ .Values.workerSandbox.sandbox.httpInputCacheMaxObjects | quote }} + - name: SANDBOX_INPUT_CACHE_MAX_BYTES + value: {{ .Values.workerSandbox.sandbox.inputCacheMaxBytes | quote }} - name: SANDBOX_EXECUTE_BODY_LIMIT value: {{ .Values.workerSandbox.sandbox.executeBodyLimit | quote }} - name: SANDBOX_DISABLE_NETWORKING diff --git a/helm/codeapi/values.yaml b/helm/codeapi/values.yaml index ee997ce8..3ecfbef8 100644 --- a/helm/codeapi/values.yaml +++ b/helm/codeapi/values.yaml @@ -46,6 +46,11 @@ egressGrant: ttlSeconds: 900 ledgerRequired: true ledgerTtlGraceSeconds: 300 + # Enable only after all gateway replicas support compact ledgers. + ledgerCompact: false + inputManifestMaxFiles: 512 + inputManifestConcurrency: 8 + inputManifestTimeoutMs: 10000 # Worker signs sandbox execute requests with this private key; sandbox-runner # receives only the public verifier so a runner compromise cannot mint new @@ -268,6 +273,10 @@ workerSandbox: # Defaults to maxConcurrentJobs when unset. jobUidCount: null workspaceReaperMaxAgeSeconds: 3600 + httpInputCacheEnabled: false + httpInputCacheMaxInflight: 16 + httpInputCacheMaxObjects: 4096 + inputCacheMaxBytes: 536870912 # Language runtime package delivery packages: @@ -327,6 +336,8 @@ workerSandbox: # FILE SERVER (S3/MinIO integration, stateless) # ============================================================================= fileServer: + objectIndexEnabled: false + metadataConcurrency: 1 enabled: true replicaCount: 1 diff --git a/launcher/src/main.rs b/launcher/src/main.rs index 02d26418..7771b9b9 100644 --- a/launcher/src/main.rs +++ b/launcher/src/main.rs @@ -421,6 +421,10 @@ fn is_allowed_guest_env_key(key: &str, egress_gateway_enabled: bool) -> bool { "SANDBOX_COMPILE_TIMEOUT", "SANDBOX_DATA_DIRECTORY", "SANDBOX_DISABLE_NETWORKING", + "SANDBOX_HTTP_INPUT_CACHE_ENABLED", + "SANDBOX_HTTP_INPUT_CACHE_MAX_INFLIGHT", + "SANDBOX_HTTP_INPUT_CACHE_MAX_OBJECTS", + "SANDBOX_INPUT_CACHE_MAX_BYTES", "SANDBOX_EXECUTE_BODY_LIMIT", "SANDBOX_EXECUTION_MANIFEST_PUBLIC_KEY", "SANDBOX_FORWARD_TARGET", diff --git a/packages/code/src/relay.test.ts b/packages/code/src/relay.test.ts index a9b81fd0..c4d92288 100644 --- a/packages/code/src/relay.test.ts +++ b/packages/code/src/relay.test.ts @@ -416,3 +416,68 @@ for (const [status, reason] of [[403, 'scope_mismatch'], [503, 'ledger_conflict' } }); } + +test('file relay carries version preflights and download preconditions without opening metadata writes', async () => { + let requests = 0; + const upstream = createServer((req, res) => { + requests++; + if (req.url?.endsWith('/metadata')) { + res.setHeader('Content-Type', 'application/json'); + res.end(JSON.stringify({ cacheable: true, version: 'opaque-version' })); + } else { + assert.equal(req.headers['x-codeapi-input-version'], 'opaque-version'); + res.writeHead(200, { 'X-CodeAPI-Input-Version': 'opaque-version' }).end('bytes'); + } + }); + const upstreamUrl = await listen(upstream); + const relay = await startFileRelay({ host: '127.0.0.1', port: 0, upstreamUrl, token: 'relay-secret', maxBytes: 1024, timeoutMs: 1000 }); + const headers = { 'X-LibreChat-Code-Relay-Token': 'relay-secret', 'X-CodeAPI-Egress-Grant': 'grant' }; + try { + const metadata = await fetch(`${relay.url}/sessions/s/objects/o/metadata`, { headers }); + assert.equal(metadata.status, 200); + assert.equal((await metadata.json() as { version: string }).version, 'opaque-version'); + const input = await fetch(`${relay.url}/sessions/s/objects/o`, { headers: { ...headers, 'X-CodeAPI-Input-Version': 'opaque-version' } }); + assert.equal(input.headers.get('x-codeapi-input-version'), 'opaque-version'); + assert.equal(await input.text(), 'bytes'); + const denied = await fetch(`${relay.url}/sessions/s/objects/o/metadata`, { method: 'PUT', headers, body: '' }); + assert.equal(denied.status, 404); + assert.equal(requests, 2); + } finally { + await relay.close(); + await new Promise((resolve, reject) => upstream.close(error => error ? reject(error) : resolve())); + } +}); + + +test('file relay forwards bounded manifest POSTs and rejects alternate manifest methods', async () => { + let requests = 0; + const upstream = createServer(async (req, res) => { + requests++; + assert.equal(req.method, 'POST'); + assert.equal(req.url, '/input-manifest'); + assert.equal(req.headers['x-codeapi-egress-grant'], 'grant'); + const chunks = []; + for await (const chunk of req) chunks.push(chunk); + assert.deepEqual(JSON.parse(Buffer.concat(chunks).toString()), { files: [] }); + res.setHeader('Content-Type', 'application/json'); + res.end(JSON.stringify({ files: [] })); + }); + const upstreamUrl = await listen(upstream); + const relay = await startFileRelay({ host: '127.0.0.1', port: 0, upstreamUrl, token: 'relay-secret', maxBytes: 128, timeoutMs: 1000 }); + const headers = { 'X-LibreChat-Code-Relay-Token': 'relay-secret', 'X-CodeAPI-Egress-Grant': 'grant', 'Content-Type': 'application/json' }; + try { + const response = await fetch(`${relay.url}/input-manifest`, { method: 'POST', headers, body: JSON.stringify({ files: [] }) }); + assert.equal(response.status, 200); + assert.deepEqual(await response.json(), { files: [] }); + for (const method of ['GET', 'PUT']) { + const denied = await fetch(`${relay.url}/input-manifest`, { method, headers }); + assert.equal(denied.status, 404); + } + const oversized = await fetch(`${relay.url}/input-manifest`, { method: 'POST', headers, body: 'x'.repeat(129) }); + assert.equal(oversized.status, 413); + assert.equal(requests, 1); + } finally { + await relay.close(); + await new Promise((resolve, reject) => upstream.close(error => error ? reject(error) : resolve())); + } +}); diff --git a/packages/code/src/relay.ts b/packages/code/src/relay.ts index 14e91280..07aa8714 100644 --- a/packages/code/src/relay.ts +++ b/packages/code/src/relay.ts @@ -20,6 +20,7 @@ export interface FileRelayHandle { } const OBJECT_PATH = /^\/sessions\/[^/]+\/objects\/[^/]+$/; +const OBJECT_METADATA_PATH = /^\/sessions\/[^/]+\/objects\/[^/]+\/metadata$/; const OBJECT_LIST_PATH = /^\/sessions\/[^/]+\/objects$/; const MAX_RELAY_HEADER_BYTES = 512 * 1024; const LOCAL_HTTP_HOSTS = new Set([ @@ -160,16 +161,19 @@ export async function startFileRelay( response.end('{"status":"ok"}'); return; } + const manifestRequest = request.method === 'POST' && requestUrl.pathname === '/input-manifest' && requestUrl.search.length === 0; const objectRequest = OBJECT_PATH.test(requestUrl.pathname) && requestUrl.search.length === 0; + const metadataRequest = request.method === 'GET' && + OBJECT_METADATA_PATH.test(requestUrl.pathname) && requestUrl.search.length === 0; const normalizedListRequest = request.method === 'GET' && OBJECT_LIST_PATH.test(requestUrl.pathname) && requestUrl.searchParams.size === 1 && requestUrl.searchParams.get('detail') === 'normalized'; if ( - (request.method !== 'GET' && request.method !== 'PUT') || - (!objectRequest && !normalizedListRequest) + (request.method !== 'GET' && request.method !== 'PUT' && !manifestRequest) || + (!objectRequest && !normalizedListRequest && !metadataRequest && !manifestRequest) ) { response.writeHead(404).end(); return; @@ -191,7 +195,7 @@ export async function startFileRelay( }`; target.search = requestUrl.search; const requestBody = - request.method === 'PUT' + (request.method === 'PUT' || manifestRequest) ? await readRequestBody(request, options.maxBytes) : undefined; const upstreamResponse = await fetch(target, { @@ -200,7 +204,9 @@ export async function startFileRelay( ...(typeof grant === 'string' ? { 'X-CodeAPI-Egress-Grant': grant } : {}), - ...(request.method === 'PUT' + ...(typeof request.headers['x-codeapi-input-version'] === 'string' + ? { 'X-CodeAPI-Input-Version': request.headers['x-codeapi-input-version'] } : {}), + ...((request.method === 'PUT' || manifestRequest) ? { 'Content-Length': String(requestBody?.length ?? 0), ...(typeof request.headers['content-type'] === 'string' @@ -246,6 +252,8 @@ export async function startFileRelay( ...(upstreamResponse.headers.has('retry-after') ? { 'Retry-After': upstreamResponse.headers.get('retry-after')! } : {}), + ...(upstreamResponse.headers.has('x-codeapi-input-version') + ? { 'X-CodeAPI-Input-Version': upstreamResponse.headers.get('x-codeapi-input-version')! } : {}), 'Content-Length': String(body.length), }); response.end(body); diff --git a/service/src/config.ts b/service/src/config.ts index 94daf5a3..90df6b58 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -318,6 +318,12 @@ export const env = { EGRESS_GATEWAY_REQUEST_TIMEOUT_MS: Number(process.env.EGRESS_GATEWAY_REQUEST_TIMEOUT_MS) || 30_000, EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: Number(process.env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS) || 5_000, EGRESS_LEDGER_REQUIRED: process.env.CODEAPI_EGRESS_LEDGER_REQUIRED === 'true' || process.env.CODEAPI_HARDENED_SANDBOX_MODE === 'true', + FILE_METADATA_CONCURRENCY: Math.min(64, Math.max(1, Math.floor(Number(process.env.CODEAPI_FILE_METADATA_CONCURRENCY) || 1))), + FILE_OBJECT_INDEX_ENABLED: process.env.CODEAPI_FILE_OBJECT_INDEX_ENABLED === 'true', + INPUT_MANIFEST_MAX_FILES: Math.min(512, Math.max(1, Math.floor(Number(process.env.CODEAPI_INPUT_MANIFEST_MAX_FILES) || 512))), + INPUT_MANIFEST_CONCURRENCY: Math.min(64, Math.max(1, Math.floor(Number(process.env.CODEAPI_INPUT_MANIFEST_CONCURRENCY) || 8))), + INPUT_MANIFEST_TIMEOUT_MS: Math.max(1, Math.floor(Number(process.env.CODEAPI_INPUT_MANIFEST_TIMEOUT_MS) || 10000)), + EGRESS_LEDGER_COMPACT: process.env.CODEAPI_EGRESS_LEDGER_COMPACT === 'true', EGRESS_LEDGER_TTL_GRACE_SECONDS: Number(process.env.CODEAPI_EGRESS_LEDGER_TTL_GRACE_SECONDS) || 300, EGRESS_GRANT_SECRET: process.env.CODEAPI_EGRESS_GRANT_SECRET ?? '', EGRESS_GRANT_TTL_SECONDS: resolveEgressGrantTtlSeconds(process.env.EGRESS_GRANT_TTL_SECONDS, defaultJobTimeoutMs), diff --git a/service/src/egress-gateway.test.ts b/service/src/egress-gateway.test.ts index 717309d7..c3cd4da4 100644 --- a/service/src/egress-gateway.test.ts +++ b/service/src/egress-gateway.test.ts @@ -1,14 +1,15 @@ process.env.CODEAPI_EGRESS_GATEWAY_AUTOSTART = 'false'; -import { afterAll, beforeAll, beforeEach, describe, expect, test, spyOn } from 'bun:test'; +import { afterAll, beforeAll, beforeEach, describe, expect, test } from 'bun:test'; import crypto from 'crypto'; -import RedisMock from 'ioredis-mock'; +import { startTestRedis } from './test/redis'; import type { Server } from 'http'; import type { AddressInfo } from 'net'; import { env } from './config'; import { assertEgressGrantActive, createEgressLedger, + revokeEgressLedger, setEgressLedgerRedisForTest, } from './egress-ledger'; import { @@ -435,39 +436,171 @@ describe('egress gateway routes', () => { } }); - test('reports exhausted ledger conflicts as retryable without forwarding the read', async () => { - const redis = new RedisMock(); + test('batch preflight scopes every entry before reading and preserves order under a concurrency bound', async () => { + const files = Array.from({ length: 17 }, (_, i) => ({ id: `file_${i}`, session_id: 'sess_input', name: `file_${i}.csv` })); + const grant = claims({ input_files: files }); + const sid = sessionHandle({ dir: 'read', sessionId: 'sess_input' }); + const body = { files: files.map(file => ({ sessionHandle: sid, objectHandle: objectHandle({ fileId: file.id, name: file.name }) })) }; + let active = 0, peak = 0, calls = 0; + const width = env.INPUT_MANIFEST_CONCURRENCY; + env.INPUT_MANIFEST_CONCURRENCY = 3; + globalThis.fetch = (async (input: RequestInfo | URL) => { + calls++; active++; peak = Math.max(peak, active); + await Bun.sleep(5); + active--; + const id = String(input).split('/').at(-2)!; + return Response.json({ version: crypto.randomUUID(), size: 5, originalFilename: `${id}.csv` }); + }) as typeof fetch; + const redis = await startTestRedis(); + setEgressLedgerRedisForTest(redis); env.EGRESS_LEDGER_REQUIRED = true; - setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); - const duplicate = redis.duplicate.bind(redis); - const duplication = spyOn(redis, 'duplicate').mockImplementation(() => { - const connection = duplicate(); - const transaction = { - set: () => transaction, - exec: async () => null, - }; - spyOn(connection, 'multi').mockImplementation(() => transaction as never); - return connection; + try { + await createEgressLedger(grant); + const response = await gatewayFetch('/input-manifest', { + method: 'POST', headers: { ...grantHeader(grant), 'Content-Type': 'application/json' }, body: JSON.stringify(body), + }); + expect(response.status).toBe(200); + const result = await response.json() as { files: { cacheKey: string; name: string }[] }; + expect(result.files.map(file => file.name)).toEqual(files.map(file => file.name)); + expect(result.files.every(file => /^[0-9a-f]{64}$/.test(file.cacheKey))).toBe(true); + expect(calls).toBe(17); + expect(peak).toBe(3); + expect((await assertEgressGrantActive(grant)).request_count).toBe(1); + body.files.push({ sessionHandle: sid, objectHandle: objectHandle({ fileId: 'outside_scope' }) }); + const denied = await gatewayFetch('/input-manifest', { + method: 'POST', headers: { ...grantHeader(grant), 'Content-Type': 'application/json' }, body: JSON.stringify(body), + }); + expect(denied.status).toBe(403); + expect(calls).toBe(17); + } finally { + env.INPUT_MANIFEST_CONCURRENCY = width; + env.EGRESS_LEDGER_REQUIRED = false; + setEgressLedgerRedisForTest(null); + await redis.closeTestServer(); + } + }); + + test('batch preflight withholds resolved metadata if the grant is revoked during storage access', async () => { + const redis = await startTestRedis(); + setEgressLedgerRedisForTest(redis); + env.EGRESS_LEDGER_REQUIRED = true; + try { + const grant = claims(); + await createEgressLedger(grant); + globalThis.fetch = (async (_input: RequestInfo | URL) => { + await revokeEgressLedger(grant.grant_id!, 'test revocation'); + return Response.json({ version: crypto.randomUUID(), size: 5 }); + }) as typeof fetch; + const response = await gatewayFetch('/input-manifest', { + method: 'POST', headers: { ...grantHeader(grant), 'Content-Type': 'application/json' }, + body: JSON.stringify({ files: [{ sessionHandle: sessionHandle({ dir: 'read', sessionId: 'sess_input' }), objectHandle: objectHandle({}) }] }), + }); + expect(response.status).toBe(403); + expect(await response.text()).not.toContain('cacheKey'); + } finally { + env.EGRESS_LEDGER_REQUIRED = false; + setEgressLedgerRedisForTest(null); + await redis.closeTestServer(); + } + }); + + test('batch deadline aborts in-flight storage requests without starting queued inputs', async () => { + const timeout = env.INPUT_MANIFEST_TIMEOUT_MS; + const width = env.INPUT_MANIFEST_CONCURRENCY; + env.INPUT_MANIFEST_TIMEOUT_MS = 20; + env.INPUT_MANIFEST_CONCURRENCY = 1; + let started = 0, aborted = 0; + globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + started++; + return new Promise((_resolve, reject) => { + init!.signal!.addEventListener('abort', () => { aborted++; reject(init!.signal!.reason); }, { once: true }); + }); + }) as typeof fetch; + try { + const file = { sessionHandle: sessionHandle({ dir: 'read', sessionId: 'sess_input' }), objectHandle: objectHandle({}) }; + const response = await gatewayFetch('/input-manifest', { + method: 'POST', headers: { ...grantHeader(), 'Content-Type': 'application/json' }, + body: JSON.stringify({ files: [file, file] }), + }); + expect(response.ok).toBe(false); + expect(started).toBe(1); + expect(aborted).toBe(1); + } finally { + env.INPUT_MANIFEST_TIMEOUT_MS = timeout; + env.INPUT_MANIFEST_CONCURRENCY = width; + } + }); + + test('batch preflight rejects oversized and malformed manifests before storage access', async () => { + for (const files of [Array.from({ length: env.INPUT_MANIFEST_MAX_FILES + 1 }, () => ({})), [null], [{}]]) { + const response = await gatewayFetch('/input-manifest', { + method: 'POST', headers: { ...grantHeader(), 'Content-Type': 'application/json' }, body: JSON.stringify({ files }), + }); + expect(response.status).toBe(400); + } + expect(upstreamCalls).toHaveLength(0); + }); + + test('preflight authorizes scope and returns version keys scoped to the principal', async () => { + const version = crypto.randomUUID(); + upstreamResponse = Response.json({ version, size: 5, originalFilename: 'inputs/data.csv', readOnly: true }); + const sid = sessionHandle({ dir: 'read', sessionId: 'sess_input' }); + const object = objectHandle({}); + const response = await gatewayFetch(`/sessions/${sid}/objects/${object}/metadata`, { headers: grantHeader() }); + expect(response.status).toBe(200); + const metadata = await response.json() as { cacheKey: string; version: string; readOnly: boolean }; + expect(metadata.cacheKey).toMatch(/^[0-9a-f]{64}$/); + expect(metadata.version).toBe(version); + expect(metadata.readOnly).toBe(true); + expect(upstreamCalls[0].url).toEndWith('/sessions/sess_input/objects/file_123/metadata'); + const second = await gatewayFetch(`/sessions/${sid}/objects/${object}/metadata`, { + headers: grantHeader(claims({ tenant_id: 'another_tenant' })), }); + expect((await second.json() as { cacheKey: string }).cacheKey).not.toBe(metadata.cacheKey); + const before = upstreamCalls.length; + const denied = await gatewayFetch(`/sessions/${sid}/objects/${objectHandle({ fileId: 'outside_scope' })}/metadata`, { + headers: grantHeader(), + }); + expect(denied.status).toBe(403); + expect(upstreamCalls).toHaveLength(before); + }); + + test('preflight denies revoked grants and old metadata stays uncached', async () => { + const sid = sessionHandle({ dir: 'read', sessionId: 'sess_input' }); + const object = objectHandle({}); + upstreamResponse = Response.json({ size: 5 }); + const legacy = await gatewayFetch(`/sessions/${sid}/objects/${object}/metadata`, { headers: grantHeader() }); + expect(await legacy.json()).toEqual({ cacheable: false }); + const redis = await startTestRedis(); + setEgressLedgerRedisForTest(redis); + env.EGRESS_LEDGER_REQUIRED = true; try { await createEgressLedger(claims()); - const readSession = sessionHandle({ dir: 'read', sessionId: 'sess_input' }); - const response = await gatewayFetch(`/sessions/${readSession}/objects?detail=normalized`, { - headers: grantHeader(), - }); - expect(response.status).toBe(503); - expect(response.headers.get('X-CodeAPI-Error-Code')).toBe('ledger_conflict'); - expect(response.headers.get('Retry-After')).toBe('1'); - expect(upstreamCalls).toHaveLength(0); - expect((await assertEgressGrantActive(claims())).request_count).toBe(0); + await redis.del(`codeapi:egress:grant:${claims().grant_id}`); + const before = upstreamCalls.length; + const denied = await gatewayFetch(`/sessions/${sid}/objects/${object}/metadata`, { headers: grantHeader() }); + expect(denied.status).toBe(403); + expect(upstreamCalls).toHaveLength(before); } finally { - duplication.mockRestore(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); - redis.disconnect(); env.EGRESS_LEDGER_REQUIRED = false; } }); + test('forwards the authorized input-version precondition on downloads', async () => { + const version = crypto.randomUUID(); + upstreamResponse = new Response('bytes', { headers: { 'X-CodeAPI-Input-Version': version } }); + const sid = sessionHandle({ dir: 'read', sessionId: 'sess_input' }); + const response = await gatewayFetch(`/sessions/${sid}/objects/${objectHandle({})}`, { + headers: { ...grantHeader(), 'X-CodeAPI-Input-Version': version }, + }); + expect(response.status).toBe(200); + expect(response.headers.get('x-codeapi-input-version')).toBe(version); + expect(new Headers(upstreamCalls[0].init.headers).get('x-codeapi-input-version')).toBe(version); + await response.text(); + }); + test('lists only scoped objects and injects internal credentials', async () => { upstreamResponse = Response.json([ { id: 'file_123', name: 'inputs/data.csv', storage_session_id: 'sess_input' }, @@ -495,7 +628,7 @@ describe('egress gateway routes', () => { }); test('accepts legacy rollout grants and handles while ledger-required mode is enabled', async () => { - const redis = new RedisMock(); + const redis = await startTestRedis(); env.EGRESS_LEDGER_REQUIRED = true; setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); try { @@ -526,14 +659,14 @@ describe('egress gateway routes', () => { expect(record.max_output_files).toBe(50); expect(record.max_requests).toBe(1000); } finally { - await redis.disconnect(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); env.EGRESS_LEDGER_REQUIRED = false; } }); test('restores token-only legacy grants and creates ledger state before returning handles', async () => { - const redis = new RedisMock(); + const redis = await startTestRedis(); env.EGRESS_LEDGER_REQUIRED = true; setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); try { @@ -569,14 +702,14 @@ describe('egress gateway routes', () => { expect(record.grant_id).toBe(legacyGrant.grant_id); expect(record.exec_id).toBe('exec_123'); } finally { - await redis.disconnect(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); env.EGRESS_LEDGER_REQUIRED = false; } }); test('rejects grantless handles for non-legacy grants in ledger-required mode', async () => { - const redis = new RedisMock(); + const redis = await startTestRedis(); env.EGRESS_LEDGER_REQUIRED = true; setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); try { @@ -593,7 +726,7 @@ describe('egress gateway routes', () => { expect(response.headers.get('X-CodeAPI-Error-Code')).toBe('scope_mismatch'); expect(upstreamCalls).toHaveLength(0); } finally { - await redis.disconnect(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); env.EGRESS_LEDGER_REQUIRED = false; } @@ -849,7 +982,7 @@ describe('egress gateway routes', () => { }); test('rolls back upload reservations when upstream PUT throws', async () => { - const redis = new RedisMock(); + const redis = await startTestRedis(); env.EGRESS_LEDGER_REQUIRED = true; setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); const grant = claims({ max_output_files: 1, max_requests: 3 }); @@ -894,14 +1027,14 @@ describe('egress gateway routes', () => { expect(retried.status).toBe(201); } finally { - await redis.disconnect(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); env.EGRESS_LEDGER_REQUIRED = false; } }); test('does not roll back ledger state when upload reservation is rejected', async () => { - const redis = new RedisMock(); + const redis = await startTestRedis(); env.EGRESS_LEDGER_REQUIRED = true; setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { @@ -933,14 +1066,14 @@ describe('egress gateway routes', () => { expect(await upload('bbbbbbbbbbbbbbbbbbbbb')).toBe(403); expect(upstreamCalls).toHaveLength(1); } finally { - await redis.disconnect(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); env.EGRESS_LEDGER_REQUIRED = false; } }); test('enforces output budgets per turn when grants reuse an output session', async () => { - const redis = new RedisMock(); + const redis = await startTestRedis(); env.EGRESS_LEDGER_REQUIRED = true; setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { @@ -989,7 +1122,7 @@ describe('egress gateway routes', () => { expect(await upload(secondTurn, 'ddddddddddddddddddddd')).toBe(201); expect(await upload(secondTurn, 'eeeeeeeeeeeeeeeeeeeee')).toBe(403); } finally { - await redis.disconnect(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); env.EGRESS_LEDGER_REQUIRED = false; } diff --git a/service/src/egress-gateway.ts b/service/src/egress-gateway.ts index 13a6de95..f86ce656 100644 --- a/service/src/egress-gateway.ts +++ b/service/src/egress-gateway.ts @@ -26,7 +26,7 @@ import { isSyntheticInternalRequestHeader, } from './internal-synthetic'; import { - assertEgressGrantActive, + checkEgressGrantActive, createEgressLedger, ensureEgressLedger, pingEgressLedger, @@ -44,6 +44,7 @@ import logger from './logger'; import { parseBoundedContentLength } from './http-limits'; import { validateEgressGatewayHardenedConfig } from './secure-startup'; import { isOpaqueObjectContentDisposition } from './file-metadata'; +import { mapObjectDetails } from './file-object-resolver'; export const app: Express = express(); app.disable('x-powered-by'); @@ -89,6 +90,8 @@ function routeFamily(req: Request): string { if (req.path === '/tool-call') return 'ptc-tool-call'; if (req.path.startsWith('/sessions/')) { if (req.method === 'PUT') return 'file-upload'; + if (req.method === 'POST' && req.path === '/input-manifest') return 'input-manifest'; + if (req.method === 'GET' && req.path.endsWith('/metadata')) return 'input-metadata'; if (req.method === 'GET' && req.path.includes('/objects/')) return 'file-download'; if (req.method === 'GET' && req.path.endsWith('/objects')) return 'file-list'; return 'file-unknown'; @@ -199,7 +202,7 @@ async function getGrant(req: Request, res: Response): Promise if (grant.legacy_grant) { await ensureEgressLedger(grant); } - await assertEgressGrantActive(grant); + await checkEgressGrantActive(grant); return grant; } @@ -438,7 +441,7 @@ async function restoreInternalSandboxResult(args: { if (grant.legacy_grant) { await ensureEgressLedger(grant); } - await assertEgressGrantActive(grant); + await checkEgressGrantActive(grant); const restored = restoreSandboxExecuteResult( args.result as Parameters[0], args.egressGrantToken, @@ -624,6 +627,88 @@ app.get('/sessions/:sessionHandle/objects', async (req, res) => { } }); +function inputMetadata( + metadata: { version?: unknown; size?: unknown; originalFilename?: unknown; readOnly?: unknown }, + grant: EgressGrantClaims, sessionId: string, objectId: string, +) { + if (typeof metadata.version !== 'string' || !/^[0-9a-f-]{36}$/.test(metadata.version) || + !Number.isSafeInteger(metadata.size) || (metadata.size as number) < 0) { + return { cacheable: false }; + } + const name = typeof metadata.originalFilename === 'string' ? metadata.originalFilename : undefined; + const readOnly = metadata.readOnly === true; + const cacheKey = crypto.createHash('sha256').update(JSON.stringify([ + 'authorized-http-input-v1', grant.tenant_id, grant.user_id, sessionId, objectId, + metadata.version, metadata.size, name, readOnly, + ])).digest('hex'); + return { cacheable: true, cacheKey, version: metadata.version, size: metadata.size, name, readOnly }; +} + +/** One budgeted HTTP read, with every handle checked before any storage access. + * Results belong only to this execution; the manifest is not an auth token. */ +app.post('/input-manifest', express.json({ limit: '4mb' }), async (req, res) => { + const controller = new AbortController(); + const cancel = () => controller.abort(); + const timeout = setTimeout(cancel, env.INPUT_MANIFEST_TIMEOUT_MS); + res.once('close', cancel); + try { + if (Object.keys(req.query).length || !Array.isArray(req.body?.files) || + req.body.files.length > env.INPUT_MANIFEST_MAX_FILES || + req.body.files.some((file: unknown) => !file || typeof file !== 'object' || + typeof (file as { sessionHandle?: unknown }).sessionHandle !== 'string' || + typeof (file as { objectHandle?: unknown }).objectHandle !== 'string')) { + return res.status(400).json({ error: 'Invalid input manifest' }); + } + const grant = await getGrant(req, res); + const files: Array<{ sessionId: string; objectId: string }> = req.body.files.map((file: { sessionHandle: string; objectHandle: string }) => { + const sessionId = openSessionParam(file.sessionHandle, grant, 'read'); + const object = openObjectParam(file.objectHandle, grant, sessionId); + return { sessionId, objectId: object.id }; + }); + await recordEgressRead(grant); + async function* inputs() { yield* files; } + const metadata = await mapObjectDetails(inputs(), async ({ sessionId, objectId }) => { + controller.signal.throwIfAborted(); + const upstream = await fetch(forwardUrl(env.EGRESS_GATEWAY_FILE_SERVER_URL, + `/sessions/${encodeURIComponent(sessionId)}/objects/${encodeURIComponent(objectId)}/metadata`), + { headers: injectTraceHeaders(internalServiceHeaders()), signal: controller.signal }); + if (!upstream.ok) { + await upstream.body?.cancel(); + // Recheck failures individually through the existing retry/classification path. + return { cacheable: false, retry: true }; + } + return inputMetadata(await upstream.json(), grant, sessionId, objectId); + }, env.INPUT_MANIFEST_CONCURRENCY); + // Do not publish a manifest after revocation/expiry during storage resolution. + await checkEgressGrantActive(grant); + return res.json({ files: metadata }); + } catch (error) { + return sendEgressError(req, res, error); + } finally { + clearTimeout(timeout); + res.removeListener('close', cancel); + controller.abort(); + } +}); + +/** Cache preflight is a scoped, budgeted read, never a reusable authorization grant. */ +app.get('/sessions/:sessionHandle/objects/:objectHandle/metadata', async (req, res) => { + try { + if (Object.keys(req.query).length) return res.status(400).json({ error: 'Metadata query parameters are not supported' }); + const grant = await getGrant(req, res); + const sessionId = openSessionParam(req.params.sessionHandle, grant, 'read'); + const object = openObjectParam(req.params.objectHandle, grant, sessionId); + await recordEgressRead(grant); + const upstream = await fetch(forwardUrl(env.EGRESS_GATEWAY_FILE_SERVER_URL, + `/sessions/${encodeURIComponent(sessionId)}/objects/${encodeURIComponent(object.id)}/metadata`), + { headers: injectTraceHeaders(internalServiceHeaders()) }); + if (!upstream.ok) return pipeFetchResponse(upstream, res); + return res.json(inputMetadata(await upstream.json(), grant, sessionId, object.id)); + } catch (error) { + return sendEgressError(req, res, error); + } +}); + app.get('/sessions/:sessionHandle/objects/:objectHandle', async (req, res) => { try { if (Object.keys(req.query).length > 0) { @@ -633,12 +718,16 @@ app.get('/sessions/:sessionHandle/objects/:objectHandle', async (req, res) => { const sessionId = openSessionParam(req.params.sessionHandle, grant, 'read'); const object = openObjectParam(req.params.objectHandle, grant, sessionId); await recordEgressRead(grant); + const expectedVersion = req.header('x-codeapi-input-version'); + if (expectedVersion && !/^[0-9a-f-]{36}$/.test(expectedVersion)) { + return res.status(400).json({ error: 'Invalid input version' }); + } const upstream = await fetch( forwardUrl( env.EGRESS_GATEWAY_FILE_SERVER_URL, `/sessions/${encodeURIComponent(sessionId)}/objects/${encodeURIComponent(object.id)}`, ), - { headers: injectTraceHeaders(internalServiceHeaders()) }, + { headers: injectTraceHeaders(internalServiceHeaders(expectedVersion ? { 'X-CodeAPI-Input-Version': expectedVersion } : {})) }, ); const headerOverrides = isOpaqueObjectContentDisposition( upstream.headers.get('content-disposition'), diff --git a/service/src/egress-ledger-reconnect.test.ts b/service/src/egress-ledger-reconnect.test.ts new file mode 100644 index 00000000..0cf4bb78 --- /dev/null +++ b/service/src/egress-ledger-reconnect.test.ts @@ -0,0 +1,63 @@ +import { expect, test } from 'bun:test'; +import { createConnection, createServer, type Socket, type AddressInfo } from 'node:net'; +import IORedis from 'ioredis'; +import { env } from './config'; +import { startTestRedis } from './test/redis'; +import { + EGRESS_LEDGER_REDIS_RETRY_OPTIONS, createEgressLedger, recordEgressRead, + assertEgressGrantActive, setEgressLedgerRedisForTest, +} from './egress-ledger'; +import type { EgressGrantClaims } from './egress-grant'; + +test('a lost Redis mutation reply rejects without replaying its applied counter after reconnect', async () => { + const redis = await startTestRedis(); + const sockets = new Set(); + let dropReply = false; + const proxy = createServer(downstream => { + const upstream = createConnection(redis.options.path!); + sockets.add(downstream); sockets.add(upstream); + downstream.pipe(upstream); + upstream.on('data', data => { + if (dropReply) { + dropReply = false; + downstream.destroy(); upstream.destroy(); + } else downstream.write(data); + }); + downstream.on('error', () => {}); + upstream.on('error', () => downstream.destroy()); + downstream.on('close', () => { sockets.delete(downstream); upstream.destroy(); }); + upstream.on('close', () => { sockets.delete(upstream); downstream.destroy(); }); + }); + await new Promise(resolve => proxy.listen(0, '127.0.0.1', resolve)); + const client = new IORedis({ host: '127.0.0.1', port: (proxy.address() as AddressInfo).port, + ...EGRESS_LEDGER_REDIS_RETRY_OPTIONS, retryStrategy: () => 10, lazyConnect: true }); + client.on('error', () => {}); + const required = env.EGRESS_LEDGER_REQUIRED; + const compact = env.EGRESS_LEDGER_COMPACT; + env.EGRESS_LEDGER_REQUIRED = true; + env.EGRESS_LEDGER_COMPACT = true; + setEgressLedgerRedisForTest(client); + try { + await client.connect(); + const now = Math.floor(Date.now() / 1000); + const grant: EgressGrantClaims = { v: 1, typ: 'grant', grant_id: 'reconnect', exec_id: 'exec', + tenant_id: 'tenant', user_id: 'user', session_key: 'session', input_files: [], read_sessions: [], + output_session_id: 'output', max_upload_bytes: 100, max_output_files: 10, max_requests: 10, iat: now, exp: now + 300 }; + await createEgressLedger(grant); + const reconnected = new Promise(resolve => client.once('ready', resolve)); + dropReply = true; + await expect(recordEgressRead(grant)).rejects.toThrow(); + await reconnected; + expect((await assertEgressGrantActive(grant)).request_count).toBe(1); + await recordEgressRead(grant); + expect((await assertEgressGrantActive(grant)).request_count).toBe(2); + } finally { + env.EGRESS_LEDGER_REQUIRED = required; + env.EGRESS_LEDGER_COMPACT = compact; + setEgressLedgerRedisForTest(null); + client.disconnect(); + for (const socket of sockets) socket.destroy(); + await new Promise(resolve => proxy.close(() => resolve())); + await redis.closeTestServer(); + } +}, 10000); diff --git a/service/src/egress-ledger-script.ts b/service/src/egress-ledger-script.ts new file mode 100644 index 00000000..c63346ad --- /dev/null +++ b/service/src/egress-ledger-script.ts @@ -0,0 +1,123 @@ +/** One-key transactions work on Redis Cluster and serialize admission with revocation. + * Legacy JSON is retained for mixed-version rollout; compact hashes avoid decoding + * the immutable input policy on the hot path. Never retry an ambiguous EVAL result: + * the operation may already have consumed its budget. */ +export const EGRESS_LEDGER_SCRIPT = ` +local key = KEYS[1] +local op = ARGV[1] +local kind = redis.call('TYPE', key) +if type(kind) == 'table' then kind = kind.ok end +local now = tonumber(ARGV[3]) +local function denied(message) return {'error', 'scope_mismatch', message} end +if op == 'create' then + if kind ~= 'none' then return {'ok'} end + local policy = cjson.decode(ARGV[4]) + if ARGV[5] == 'compact' then + redis.call('HSET', key, 'policy', ARGV[4], 'status', 'active', + 'exec_id', policy.exec_id, 'exp', policy.exp, + 'max_requests', policy.max_requests, 'max_upload_bytes', policy.max_upload_bytes, + 'max_output_files', policy.max_output_files, + 'request_count', 0, 'read_count', 0, 'upload_count', 0, 'tool_call_count', 0, 'uploaded_bytes', 0) + else + redis.call('SET', key, ARGV[4]) + end + redis.call('EXPIRE', key, tonumber(ARGV[6])) + return {'ok'} +end +if kind == 'none' then + if op == 'revoke' then return {'ok'} end + return denied('Egress grant ledger record is missing') +end +if kind ~= 'hash' and kind ~= 'string' then return denied('Invalid egress ledger representation') end +local compact = kind == 'hash' +local record = nil +if not compact then record = cjson.decode(redis.call('GET', key)) end +local function get(field) + if compact then return redis.call('HGET', key, field) end + return record[field] +end +local function number(field) return tonumber(get(field)) end +local function put(field, value) + if compact then redis.call('HSET', key, field, value) else record[field] = value end +end +local function add(field, value) + if compact then redis.call('HINCRBY', key, field, value) else record[field] = record[field] + value end +end +local function encodeRecord(value) + local encoded = cjson.encode(value) + -- Preserve the array contract for old gateways when cjson sees empty lists. + for _, field in ipairs({'input_files', 'read_sessions', 'output_file_ids'}) do + encoded = string.gsub(encoded, '"' .. field .. '":{}', '"' .. field .. '":[]') + end + return encoded +end +local function save() + if not compact then + -- Keep the original expiration, including a revocation tombstone's lifetime. + local ttl = redis.call('PTTL', key) + redis.call('SET', key, encodeRecord(record)) + if ttl >= 0 then redis.call('PEXPIRE', key, math.max(1, ttl)) end + end +end +if op == 'revoke' then + put('status', 'revoked') + put('revoked_at', now) + put('revoke_reason', ARGV[4]) + save() + return {'ok'} +end +if get('exec_id') ~= ARGV[2] then return denied('Egress grant ledger record does not match token') end +if get('status') ~= 'active' then return denied('Egress grant has been revoked') end +if number('exp') <= now then return {'error', 'expired', 'Egress grant is expired'} end +if op == 'check' then return {'ok'} end +if op == 'snapshot' then + if compact then + record = cjson.decode(redis.call('HGET', key, 'policy')) + for _, field in ipairs({'request_count', 'read_count', 'upload_count', 'tool_call_count', 'uploaded_bytes'}) do + record[field] = number(field) + end + record.output_file_ids = {} + local fields = redis.call('HKEYS', key) + for _, field in ipairs(fields) do + if string.sub(field, 1, 7) == 'output:' then table.insert(record.output_file_ids, string.sub(field, 8)) end + end + end + return {'ok', encodeRecord(record)} +end +local file = ARGV[4] +local bytes = tonumber(ARGV[5]) +local outputField = 'output:' .. file +local outputIndex = nil +if not compact and (op == 'reserve' or op == 'release') then + for i, id in ipairs(record.output_file_ids) do if id == file then outputIndex = i end end +end +local existing = compact and redis.call('HGET', key, outputField) or outputIndex +if op == 'release' then + -- A retried release must not refund another operation's request or byte budget. + if not existing then return {'ok'} end + local reservedBytes = compact and tonumber(existing) or bytes + if compact and reservedBytes ~= bytes then return denied('Upload release does not match reservation') end + add('uploaded_bytes', -math.min(number('uploaded_bytes'), reservedBytes)) + add('upload_count', -1) + add('request_count', -1) + if compact then redis.call('HDEL', key, outputField) else table.remove(record.output_file_ids, outputIndex) end +elseif op == 'reserve' or op == 'read' or op == 'tool' then + if number('request_count') >= number('max_requests') then return denied('Egress grant request budget exceeded') end + if op == 'reserve' then + local maxBytes = math.min(number('max_upload_bytes'), tonumber(ARGV[6])) + if not bytes or bytes < 0 or bytes ~= math.floor(bytes) or bytes > maxBytes then + return denied('Upload exceeds per-file egress byte limit') + end + if existing then return denied('Output file id has already been used for this grant') end + if number('upload_count') >= number('max_output_files') then return denied('Output file count budget exceeded') end + if number('uploaded_bytes') + bytes > maxBytes * number('max_output_files') then return denied('Aggregate upload byte budget exceeded') end + add('uploaded_bytes', bytes) + add('upload_count', 1) + if compact then redis.call('HSET', key, outputField, bytes) else table.insert(record.output_file_ids, file) end + elseif op == 'read' then add('read_count', 1) + else add('tool_call_count', 1) end + add('request_count', 1) +else return denied('Unknown egress ledger operation') end +save() +return {'ok'} +`; diff --git a/service/src/egress-ledger.test.ts b/service/src/egress-ledger.test.ts index 48b69d71..87fbec8d 100644 --- a/service/src/egress-ledger.test.ts +++ b/service/src/egress-ledger.test.ts @@ -1,5 +1,5 @@ import { afterEach, beforeEach, describe, expect, test } from 'bun:test'; -import RedisMock from 'ioredis-mock'; +import { startTestRedis } from './test/redis'; import { env } from './config'; import type { EgressGrantClaims } from './egress-grant'; import { EgressGrantError } from './egress-grant'; @@ -9,6 +9,9 @@ import { ensureEgressLedger, releaseEgressUpload, reserveEgressUpload, + recordEgressRead, + recordEgressToolCall, + checkEgressGrantActive, revokeEgressLedger, setEgressLedgerRedisForTest, } from './egress-ledger'; @@ -49,26 +52,30 @@ function expectEgressError(fn: () => Promise, reason: EgressGrantError[ ); } -describe('egress Redis ledger', () => { - let redis: InstanceType; +describe.each([false, true])('egress Redis ledger compact=%s', compact => { + let redis: Awaited>; let previousRequired: boolean; + let previousCompact: boolean; let previousMaxFileBytes: number; let previousTtlGraceSeconds: number; - beforeEach(() => { + beforeEach(async () => { previousRequired = env.EGRESS_LEDGER_REQUIRED; + previousCompact = env.EGRESS_LEDGER_COMPACT; + env.EGRESS_LEDGER_COMPACT = compact; previousMaxFileBytes = env.EGRESS_GATEWAY_MAX_FILE_BYTES; previousTtlGraceSeconds = env.EGRESS_LEDGER_TTL_GRACE_SECONDS; env.EGRESS_LEDGER_REQUIRED = true; env.EGRESS_GATEWAY_MAX_FILE_BYTES = 10; - redis = new RedisMock(); + redis = await startTestRedis(); setEgressLedgerRedisForTest(redis as unknown as Parameters[0]); }); afterEach(async () => { - await redis.disconnect(); + await redis.closeTestServer(); setEgressLedgerRedisForTest(null); env.EGRESS_LEDGER_REQUIRED = previousRequired; + env.EGRESS_LEDGER_COMPACT = previousCompact; env.EGRESS_GATEWAY_MAX_FILE_BYTES = previousMaxFileBytes; env.EGRESS_LEDGER_TTL_GRACE_SECONDS = previousTtlGraceSeconds; }); @@ -103,7 +110,7 @@ describe('egress Redis ledger', () => { ); }); - test('clears Redis WATCH after rejected mutations so later valid updates can proceed', async () => { + test('leaves counters unchanged after a rejected mutation', async () => { const claims = grant({ max_output_files: 2, max_requests: 5 }); await createEgressLedger(claims); @@ -143,7 +150,7 @@ describe('egress Redis ledger', () => { await expectEgressError(() => assertEgressGrantActive(claims), 'scope_mismatch'); }); - test('keeps concurrent WATCH mutations isolated on dedicated Redis connections', async () => { + test('accounts concurrent operations without WATCH connections', async () => { const claims = grant({ max_output_files: 16, max_requests: 16, @@ -170,7 +177,7 @@ describe('egress Redis ledger', () => { ); const record = await assertEgressGrantActive(claims); - expect(duplicateCount).toBe(8); + expect(duplicateCount).toBe(0); expect(record.request_count).toBe(12); expect(record.upload_count).toBe(12); expect(record.uploaded_bytes).toBe(12); @@ -184,4 +191,60 @@ describe('egress Redis ledger', () => { redis.duplicate = duplicate as typeof redis.duplicate; } }); + test('admits exactly the request budget under a 240-file burst', async () => { + const claims = grant({ max_requests: 100, input_files: Array.from({ length: 240 }, (_, i) => ({ + id: `file_${i}`, session_id: 'inputs', name: `${i}.txt`, + })) }); + await createEgressLedger(claims); + const results = await Promise.allSettled(Array.from({ length: 240 }, () => recordEgressRead(claims))); + expect(results.filter(result => result.status === 'fulfilled')).toHaveLength(100); + expect((await assertEgressGrantActive(claims)).request_count).toBe(100); + expect((await assertEgressGrantActive(claims)).read_count).toBe(100); + }); + + test('rejects wrong execution, expired grants and mutations after revocation', async () => { + const claims = grant({ max_requests: 1000 }); + await createEgressLedger(claims); + await expectEgressError(() => recordEgressRead({ ...claims, exec_id: 'wrong' }), 'scope_mismatch'); + await recordEgressToolCall(claims.grant_id, claims.exec_id); + await Promise.all([ + ...Array.from({ length: 40 }, () => recordEgressRead(claims).catch(() => {})), + revokeEgressLedger(claims.grant_id, 'done'), + ]); + await createEgressLedger(claims); + await expectEgressError(() => checkEgressGrantActive(claims), 'scope_mismatch'); + await expectEgressError(() => recordEgressRead(claims), 'scope_mismatch'); + const expired = grant({ grant_id: 'expired', exp: nowSeconds() - 1 }); + await createEgressLedger(expired); + await expectEgressError(() => checkEgressGrantActive(expired), 'expired'); + }); + + test('duplicate releases cannot refund another upload or read', async () => { + const claims = grant({ max_requests: 10, max_output_files: 2 }); + await createEgressLedger(claims); + await recordEgressRead(claims); + await reserveEgressUpload({ grant: claims, fileId: 'a', bytes: 3 }); + await reserveEgressUpload({ grant: claims, fileId: 'b', bytes: 4 }); + await Promise.all(Array.from({ length: 10 }, () => releaseEgressUpload({ grant: claims, fileId: 'a', bytes: 3 }))); + expect(await assertEgressGrantActive(claims)).toMatchObject({ + request_count: 2, upload_count: 1, uploaded_bytes: 4, output_file_ids: ['b'], + }); + }); + + test('format selection affects new grants only and never resets existing state', async () => { + const claims = grant(); + await createEgressLedger(claims); + await recordEgressRead(claims); + env.EGRESS_LEDGER_COMPACT = !compact; + await ensureEgressLedger(claims); + await recordEgressRead(claims); + expect(await redis.type(`codeapi:egress:grant:${claims.grant_id}`)).toBe(compact ? 'hash' : 'string'); + expect((await assertEgressGrantActive(claims)).request_count).toBe(2); + if (!compact) { + const legacy = JSON.parse((await redis.get(`codeapi:egress:grant:${claims.grant_id}`))!); + expect(legacy.output_file_ids).toEqual([]); + expect(Array.isArray(legacy.input_files)).toBe(true); + } + }); + }); diff --git a/service/src/egress-ledger.ts b/service/src/egress-ledger.ts index b6bbf4fb..c635dcc4 100644 --- a/service/src/egress-ledger.ts +++ b/service/src/egress-ledger.ts @@ -6,6 +6,7 @@ import type { EgressGrantClaims } from './egress-grant'; import { EgressGrantError } from './egress-grant'; import logger from './logger'; import { redisKeepAliveOptions } from './redis-options'; +import { EGRESS_LEDGER_SCRIPT } from './egress-ledger-script'; type LedgerStatus = 'active' | 'revoked'; @@ -30,21 +31,17 @@ export interface EgressLedgerRecord { output_file_ids: string[]; } -let redis: IORedis | null = null; -const LEDGER_MUTATION_ATTEMPTS = 32; -const LEDGER_MUTATION_POOL_SIZE = Math.max(1, Number(process.env.CODEAPI_EGRESS_LEDGER_MUTATION_CONNECTIONS) || 32); - -type MutationConnectionWaiter = { - resolve: (client: IORedis) => void; - reject: (error: Error) => void; -}; - -const mutationConnections = new Set(); -let idleMutationConnections: IORedis[] = []; -let mutationConnectionWaiters: MutationConnectionWaiter[] = []; +/** A lost reply is ambiguous: never replay a possibly applied mutation. + * maxRetries=0 also rejects the pending promise on disconnect instead of leaving + * it unresolved when ioredis discards its unfulfilled-command queue. */ +export const EGRESS_LEDGER_REDIS_RETRY_OPTIONS = { + autoResendUnfulfilledCommands: false, + maxRetriesPerRequest: 0, +} as const; +let redis: IORedis | null = null; +const scriptClients = new WeakSet(); export function setEgressLedgerRedisForTest(client: IORedis | null): void { - resetMutationConnections(); redis = client; } @@ -66,7 +63,7 @@ function redisConnection(): IORedis { host: process.env.REDIS_HOST ?? 'redis', port: Number(process.env.REDIS_PORT) || 6379, password: process.env.REDIS_PASSWORD, - maxRetriesPerRequest: 1, + ...EGRESS_LEDGER_REDIS_RETRY_OPTIONS, retryStrategy, enableReadyCheck: true, connectTimeout: 10000, @@ -82,71 +79,6 @@ function redisConnection(): IORedis { return redis; } -function resetMutationConnections(): void { - const resetError = new Error('Egress ledger Redis connection reset'); - for (const waiter of mutationConnectionWaiters) { - waiter.reject(resetError); - } - mutationConnectionWaiters = []; - idleMutationConnections = []; - for (const client of mutationConnections) { - client.disconnect(); - } - mutationConnections.clear(); -} - -async function dedicatedMutationConnection(): Promise { - while (idleMutationConnections.length > 0) { - const client = idleMutationConnections.pop()!; - if (client.status !== 'end') { - return client; - } - mutationConnections.delete(client); - } - - if (mutationConnections.size < LEDGER_MUTATION_POOL_SIZE) { - return createMutationConnection(); - } - - return new Promise((resolve, reject) => { - mutationConnectionWaiters.push({ resolve, reject }); - }); -} - -function createMutationConnection(): IORedis { - const client = redisConnection().duplicate(); - mutationConnections.add(client); - client.on('error', error => logger.error('Egress ledger mutation Redis error', { error })); - return client; -} - -function releaseMutationConnection(client: IORedis): void { - if (!mutationConnections.has(client) || client.status === 'end') { - mutationConnections.delete(client); - const waiter = mutationConnectionWaiters.shift(); - if (waiter) { - try { - waiter.resolve(createMutationConnection()); - } catch (error) { - waiter.reject(error instanceof Error ? error : new Error(String(error))); - } - } - return; - } - - const waiter = mutationConnectionWaiters.shift(); - if (waiter) { - waiter.resolve(client); - return; - } - - idleMutationConnections.push(client); -} - -function sleep(ms: number): Promise { - return new Promise(resolve => setTimeout(resolve, ms)); -} - export async function pingEgressLedger(): Promise { if (!env.EGRESS_LEDGER_REQUIRED) return; await redisConnection().ping(); @@ -173,171 +105,79 @@ function recordFromGrant(grant: EgressGrantClaims): EgressLedgerRecord { }; } -export async function createEgressLedger(grant: EgressGrantClaims): Promise { - if (!grant.grant_id) { - throw new EgressGrantError('malformed', 'Egress grant id is required'); +async function executeLedger( + operation: string, + grantId: string, + executionId = '', + extra: Array = [], +): Promise { + const client = redisConnection() as IORedis & { + executeEgressLedger: (...args: Array) => Promise; + }; + if (!scriptClients.has(client)) { + client.defineCommand('executeEgressLedger', { numberOfKeys: 1, lua: EGRESS_LEDGER_SCRIPT }); + scriptClients.add(client); } - if (!env.EGRESS_LEDGER_REQUIRED) return; - await redisConnection().set( - ledgerKey(grant.grant_id), - JSON.stringify(recordFromGrant(grant)), - 'EX', - ttlSeconds(grant.exp), - ); + const result = await client.executeEgressLedger( + ledgerKey(grantId), operation, executionId, + Math.floor(Date.now() / 1000), ...extra, + ) as string[]; + if (result[0] === 'error') { + throw new EgressGrantError(result[1] as EgressGrantError['reason'], result[2]); + } + return result[1]; } -export async function ensureEgressLedger(grant: EgressGrantClaims): Promise { - if (!grant.grant_id) { - throw new EgressGrantError('malformed', 'Egress grant id is required'); - } +export async function createEgressLedger(grant: EgressGrantClaims): Promise { + if (!grant.grant_id) throw new EgressGrantError('malformed', 'Egress grant id is required'); if (!env.EGRESS_LEDGER_REQUIRED) return; - await redisConnection().set( - ledgerKey(grant.grant_id), - JSON.stringify(recordFromGrant(grant)), - 'EX', - ttlSeconds(grant.exp), - 'NX', - ); + await executeLedger('create', grant.grant_id, grant.exec_id, [ + JSON.stringify(recordFromGrant(grant)), env.EGRESS_LEDGER_COMPACT ? 'compact' : 'legacy', ttlSeconds(grant.exp), + ]); } -async function loadRecord(grantId: string): Promise { - const raw = await redisConnection().get(ledgerKey(grantId)); - if (!raw) { - throw new EgressGrantError('scope_mismatch', 'Egress grant ledger record is missing'); - } - return JSON.parse(raw) as EgressLedgerRecord; -} +/** Admission is idempotent: neither replay nor rolling deployment resets budgets or revocation. */ +export const ensureEgressLedger = createEgressLedger; -function assertActive(record: EgressLedgerRecord, grant: Pick): void { - if (record.grant_id !== grant.grant_id || record.exec_id !== grant.exec_id) { - throw new EgressGrantError('scope_mismatch', 'Egress grant ledger record does not match token'); - } - if (record.status !== 'active') { - throw new EgressGrantError('scope_mismatch', 'Egress grant has been revoked'); - } - if (record.exp <= Math.floor(Date.now() / 1000)) { - throw new EgressGrantError('expired', 'Egress grant is expired'); - } -} - -async function mutateRecord( - grant: EgressGrantClaims, - mutate: (record: EgressLedgerRecord) => void, -): Promise { - if (!env.EGRESS_LEDGER_REQUIRED) { - return recordFromGrant(grant); - } - const client = await dedicatedMutationConnection(); - const key = ledgerKey(grant.grant_id); - try { - for (let i = 0; i < LEDGER_MUTATION_ATTEMPTS; i++) { - await client.watch(key); - let record: EgressLedgerRecord; - try { - const raw = await client.get(key); - if (!raw) { - throw new EgressGrantError('scope_mismatch', 'Egress grant ledger record is missing'); - } - record = JSON.parse(raw) as EgressLedgerRecord; - assertActive(record, grant); - mutate(record); - if (record.request_count > record.max_requests) { - throw new EgressGrantError('scope_mismatch', 'Egress grant request budget exceeded'); - } - } catch (error) { - await client.unwatch().catch(unwatchError => { - logger.warn('Failed to clear egress ledger WATCH after rejected mutation', { error: unwatchError }); - }); - throw error; - } - const result = await client.multi() - .set(key, JSON.stringify(record), 'EX', ttlSeconds(record.exp)) - .exec(); - if (result) return record; - if (i < LEDGER_MUTATION_ATTEMPTS - 1) { - await sleep(Math.min(25, i + 1)); - } - } - } finally { - await client.unwatch().catch(error => { - logger.warn('Failed to clear egress ledger WATCH before returning mutation connection', { error }); - }); - releaseMutationConnection(client); - } - throw new EgressGrantError('ledger_conflict', 'Egress grant ledger update conflicted'); +/** Authorization hot path deliberately does not return the potentially large input policy. */ +export async function checkEgressGrantActive(grant: Pick): Promise { + if (!env.EGRESS_LEDGER_REQUIRED) return; + await executeLedger('check', grant.grant_id, grant.exec_id); } export async function assertEgressGrantActive(grant: EgressGrantClaims): Promise { if (!env.EGRESS_LEDGER_REQUIRED) return recordFromGrant(grant); - const record = await loadRecord(grant.grant_id); - assertActive(record, grant); + const record = JSON.parse((await executeLedger('snapshot', grant.grant_id, grant.exec_id))!) as EgressLedgerRecord; + // Redis cjson represents empty Lua arrays as objects. + if (!Array.isArray(record.output_file_ids)) record.output_file_ids = []; + if (!Array.isArray(record.input_files)) record.input_files = []; + if (!Array.isArray(record.read_sessions)) record.read_sessions = []; return record; } export async function recordEgressRead(grant: EgressGrantClaims): Promise { - await mutateRecord(grant, record => { - record.request_count += 1; - record.read_count += 1; - }); + if (!env.EGRESS_LEDGER_REQUIRED) return; + await executeLedger('read', grant.grant_id, grant.exec_id, ['', 0]); } -export async function reserveEgressUpload(args: { - grant: EgressGrantClaims; - fileId: string; - bytes: number; -}): Promise { - await mutateRecord(args.grant, record => { - if (args.bytes > Math.min(record.max_upload_bytes, env.EGRESS_GATEWAY_MAX_FILE_BYTES)) { - throw new EgressGrantError('scope_mismatch', 'Upload exceeds per-file egress byte limit'); - } - if (record.output_file_ids.includes(args.fileId)) { - throw new EgressGrantError('scope_mismatch', 'Output file id has already been used for this grant'); - } - if (record.output_file_ids.length >= record.max_output_files) { - throw new EgressGrantError('scope_mismatch', 'Output file count budget exceeded'); - } - const aggregateLimit = Math.min(record.max_upload_bytes, env.EGRESS_GATEWAY_MAX_FILE_BYTES) * record.max_output_files; - if (record.uploaded_bytes + args.bytes > aggregateLimit) { - throw new EgressGrantError('scope_mismatch', 'Aggregate upload byte budget exceeded'); - } - record.request_count += 1; - record.upload_count += 1; - record.uploaded_bytes += args.bytes; - record.output_file_ids.push(args.fileId); - }); +export async function reserveEgressUpload(args: { grant: EgressGrantClaims; fileId: string; bytes: number }): Promise { + if (!env.EGRESS_LEDGER_REQUIRED) return; + if (!Number.isSafeInteger(args.bytes) || args.bytes < 0) throw new EgressGrantError('scope_mismatch', 'Invalid upload byte count'); + await executeLedger('reserve', args.grant.grant_id, args.grant.exec_id, [args.fileId, args.bytes, env.EGRESS_GATEWAY_MAX_FILE_BYTES]); } -export async function releaseEgressUpload(args: { - grant: EgressGrantClaims; - fileId: string; - bytes: number; -}): Promise { +export async function releaseEgressUpload(args: { grant: EgressGrantClaims; fileId: string; bytes: number }): Promise { if (!env.EGRESS_LEDGER_REQUIRED) return; - await mutateRecord(args.grant, record => { - record.uploaded_bytes = Math.max(0, record.uploaded_bytes - args.bytes); - record.upload_count = Math.max(0, record.upload_count - 1); - record.request_count = Math.max(0, record.request_count - 1); - record.output_file_ids = record.output_file_ids.filter(id => id !== args.fileId); - }); + if (!Number.isSafeInteger(args.bytes) || args.bytes < 0) throw new EgressGrantError('scope_mismatch', 'Invalid upload byte count'); + await executeLedger('release', args.grant.grant_id, args.grant.exec_id, [args.fileId, args.bytes]); } export async function recordEgressToolCall(grantId: string | undefined, executionId: string): Promise { if (!env.EGRESS_LEDGER_REQUIRED || !grantId) return; - const grant = { grant_id: grantId, exec_id: executionId } as EgressGrantClaims; - await mutateRecord(grant, record => { - record.request_count += 1; - record.tool_call_count += 1; - }); + await executeLedger('tool', grantId, executionId, ['', 0]); } export async function revokeEgressLedger(grantId: string, reason: string): Promise { if (!env.EGRESS_LEDGER_REQUIRED) return; - const key = ledgerKey(grantId); - const raw = await redisConnection().get(key); - if (!raw) return; - const record = JSON.parse(raw) as EgressLedgerRecord; - record.status = 'revoked'; - record.revoked_at = Math.floor(Date.now() / 1000); - record.revoke_reason = reason; - await redisConnection().set(key, JSON.stringify(record), 'EX', ttlSeconds(record.exp)); + await executeLedger('revoke', grantId, '', [reason]); } diff --git a/service/src/file-download.test.ts b/service/src/file-download.test.ts new file mode 100644 index 00000000..cb32d771 --- /dev/null +++ b/service/src/file-download.test.ts @@ -0,0 +1,49 @@ +import { expect, test } from 'bun:test'; +import { Readable, Writable } from 'node:stream'; +import { createServer } from 'node:http'; +import express from 'express'; +import { sendFileDownload } from './file-download'; + +async function serverFor(stream: Readable) { + const app = express(); + app.get('/', (req, res) => { void sendFileDownload(stream, res, req.header('x-codeapi-input-version')).catch(() => res.destroy()); }); + const server = createServer(app); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const address = server.address() as { port: number }; + return { url: `http://127.0.0.1:${address.port}`, close: () => new Promise(resolve => server.close(() => resolve())) }; +} + +test('serves metadata from the downloaded version and rejects a stale preflight', async () => { + for (const expected of ['current', 'stale']) { + const stream = Object.assign(Readable.from(['bytes']), { headers: { + 'x-amz-meta-codeapi-version': 'current', 'x-amz-meta-read-only': 'true', + 'x-amz-meta-original-filename': 'verified.txt', + } }); + const server = await serverFor(stream); + try { + const response = await fetch(server.url, { headers: { 'X-CodeAPI-Input-Version': expected } }); + expect(response.status).toBe(expected === 'current' ? 200 : 409); + if (expected === 'current') { + expect(response.headers.get('x-read-only')).toBe('true'); + expect(response.headers.get('content-disposition')).toContain('verified.txt'); + expect(await response.text()).toBe('bytes'); + } else await response.text(); + } finally { await server.close(); } + } +}); + +test('downstream cancellation stops the storage stream under backpressure', async () => { + let produced = 0; + const stream = new Readable({ read() { if (++produced <= 1000) this.push(Buffer.alloc(64 * 1024)); else this.push(null); } }); + const response = Object.assign(new Writable({ + highWaterMark: 1, + write(_chunk, _encoding, callback) { setTimeout(callback, 10); }, + }), { setHeader() {} }); + const transfer = sendFileDownload(stream, response as unknown as express.Response); + const timer = setTimeout(() => response.destroy(), 25); + try { + await expect(transfer).rejects.toThrow(); + expect(stream.destroyed).toBe(true); + expect(produced).toBeLessThan(1000); + } finally { clearTimeout(timer); } +}); diff --git a/service/src/file-download.ts b/service/src/file-download.ts new file mode 100644 index 00000000..c46f80a1 --- /dev/null +++ b/service/src/file-download.ts @@ -0,0 +1,27 @@ +import type { Readable } from 'node:stream'; +import { pipeline } from 'node:stream/promises'; +import type { Response } from 'express'; +import { contentDispositionForOriginalFilename, originalFilenameFromMetadata } from './file-metadata'; + +export async function sendFileDownload(dataStream: Readable, res: Response, expectedVersion?: string): Promise { + // MinIO returns the HTTP response stream. Read metadata from this exact GET, + // avoiding both a redundant HEAD and metadata/content races on overwrite. + const headers = (dataStream as Readable & { headers?: Record }).headers ?? {}; + if (expectedVersion && headers['x-amz-meta-codeapi-version'] !== expectedVersion) { + dataStream.destroy(); + res.status(409).json({ error: 'Input changed during preparation; retry with current metadata' }); + return; + } + const metadata: Record = {}; + for (const [key, value] of Object.entries(headers)) { + if (key.startsWith('x-amz-meta-')) metadata[key.slice(11)] = value; + } + res.setHeader('Content-Disposition', contentDispositionForOriginalFilename(originalFilenameFromMetadata(metadata))); + if (headers['content-type']) res.setHeader('Content-Type', headers['content-type']); + if (headers['content-length']) res.setHeader('Content-Length', headers['content-length']); + if (metadata['read-only'] === 'true') res.setHeader('X-Read-Only', 'true'); + if (metadata['codeapi-version']) res.setHeader('X-CodeAPI-Input-Version', metadata['codeapi-version']); + const cancel = (): void => { if (!res.writableFinished) dataStream.destroy(new Error('Download client disconnected')); }; + res.once('close', cancel); + try { await pipeline(dataStream, res); } finally { res.off('close', cancel); } +} diff --git a/service/src/file-object-resolver.test.ts b/service/src/file-object-resolver.test.ts new file mode 100644 index 00000000..deeb3265 --- /dev/null +++ b/service/src/file-object-resolver.test.ts @@ -0,0 +1,50 @@ +import { describe, expect, test } from 'bun:test'; +import { FileObjectResolver, mapObjectDetails } from './file-object-resolver'; +import type { BucketItemStat } from 'minio'; + +describe('storage object resolution', () => { + test('indexes exact identities while reading fresh version metadata on every request', async () => { + const index = new Map(); + let lists = 0; + let heads = 0; + let version = 'first'; + const resolver = new FileObjectResolver({ + bucket: 'files', + list: async function* () { lists++; yield { name: 's/identifier.txt' }; yield { name: 's/id.txt' }; }, + stat: async key => { heads++; expect(key).toBe('s/id.txt'); return { size: 5, etag: 'etag', lastModified: new Date(), metaData: { 'codeapi-version': version } } as BucketItemStat; }, + index: { get: async k => index.get(k) ?? null, set: async (k, v) => index.set(k, v), forget: async k => index.delete(k) }, + }); + expect((await resolver.metadata('s', 'id'))?.stat.metaData['codeapi-version']).toBe('first'); + version = 'second'; + expect((await resolver.metadata('s', 'id'))?.stat.metaData['codeapi-version']).toBe('second'); + expect(lists).toBe(1); + expect(heads).toBe(2); + }); + + test('ignores foreign-session index entries and does not cache absence', async () => { + let present = false; + const resolver = new FileObjectResolver({ + bucket: 'files', + list: async function* () { yield { name: 's2/id.txt' }; if (present) yield { name: 's/id.txt' }; }, + stat: async () => ({ metaData: {} } as BucketItemStat), + index: { get: async () => 's2/id.txt', set: async () => {}, forget: async () => {} }, + }); + expect(await resolver.resolve('s', 'id')).toBeUndefined(); + present = true; + expect(await resolver.resolve('s', 'id')).toBe('s/id.txt'); + }); +}); + +test('metadata listing stays bounded and ordered across 240 objects', async () => { + let active = 0; + let maximum = 0; + async function* objects() { for (let i = 0; i < 240; i++) yield i; } + const result = await mapObjectDetails(objects(), async value => { + maximum = Math.max(maximum, ++active); + await new Promise(resolve => setTimeout(resolve, value % 3)); + active--; + return value; + }, 8); + expect(maximum).toBe(8); + expect(result).toEqual(Array.from({ length: 240 }, (_, i) => i)); +}); diff --git a/service/src/file-object-resolver.ts b/service/src/file-object-resolver.ts new file mode 100644 index 00000000..a41e0dfc --- /dev/null +++ b/service/src/file-object-resolver.ts @@ -0,0 +1,75 @@ +import { createHash } from 'node:crypto'; +import path from 'node:path'; +import type { BucketItemStat } from 'minio'; + +export interface ObjectResolverDependencies { + bucket: string; + list(prefix: string): AsyncIterable<{ name?: string }>; + stat(key: string): Promise; + index?: { + get(key: string): Promise; + set(key: string, value: string, replace: boolean): Promise; + forget(key: string, value: string): Promise; + }; +} + +/** Storage-key index is a hint, never metadata or authorization. A fresh HEAD + * proves existence and supplies the current version even on index/cache hits. */ +export class FileObjectResolver { + constructor(private readonly deps: ObjectResolverDependencies) {} + + private indexKey(session: string, id: string): string { + return `codeapi:file-key:${createHash('sha256').update(JSON.stringify([this.deps.bucket, session, id])).digest('hex')}`; + } + + private matches(key: string, session: string, id: string): boolean { + return path.posix.dirname(key) === session && + (path.posix.basename(key) === id || path.posix.basename(key, path.posix.extname(key)) === id); + } + + async remember(session: string, id: string, key: string, replace = true): Promise { + if (!this.matches(key, session, id)) throw new Error('Object key does not match storage identity'); + await this.deps.index?.set(this.indexKey(session, id), key, replace); + } + + async resolve(session: string, id: string): Promise { + const cached = await this.deps.index?.get(this.indexKey(session, id)); + if (cached && this.matches(cached, session, id)) return cached; + for await (const object of this.deps.list(`${session}/${id}`)) { + if (object.name && this.matches(object.name, session, id)) { + await this.remember(session, id, object.name, false); + return object.name; + } + } + return undefined; + } + + async metadata(session: string, id: string): Promise<{ key: string; stat: BucketItemStat } | undefined> { + const key = await this.resolve(session, id); + if (!key) return undefined; + try { + return { key, stat: await this.deps.stat(key) }; + } catch (error) { + if (!['NoSuchKey', 'NotFound', 'NoSuchObject'].includes((error as { code?: string }).code ?? '')) throw error; + await this.deps.index?.forget(this.indexKey(session, id), key); + // Do not cache absence: a later upload can publish this identity again. + return undefined; + } + } +} + +/** Bound storage metadata requests while preserving listing order. */ +export async function mapObjectDetails(objects: AsyncIterable, describe: (object: T) => Promise, concurrency: number): Promise { + const results: R[] = []; + const batch: T[] = []; + const width = Math.max(1, Math.min(64, Math.floor(concurrency) || 1)); + for await (const object of objects) { + batch.push(object); + if (batch.length === width) { + results.push(...await Promise.all(batch.map(describe))); + batch.length = 0; + } + } + results.push(...await Promise.all(batch.map(describe))); + return results; +} diff --git a/service/src/file-server.ts b/service/src/file-server.ts index f9293e48..9b326ca2 100644 --- a/service/src/file-server.ts +++ b/service/src/file-server.ts @@ -1,4 +1,8 @@ import b from 'busboy'; +import { randomUUID } from 'node:crypto'; +import { mapObjectDetails } from './file-object-resolver'; +import { sendFileDownload } from './file-download'; +import { FileObjectResolver } from './file-object-resolver'; import path from 'path'; import IORedis from 'ioredis'; import express from 'express'; @@ -18,7 +22,6 @@ import logger from './fileServerLogger'; import { env } from './config'; import { redisKeepAliveOptions } from './redis-options'; import { - contentDispositionForOriginalFilename, decodeOriginalFilename, originalFilenameFromMetadata, } from './file-metadata'; @@ -144,6 +147,21 @@ redisClient.on('ready', () => { logger.info('Redis Client Ready'); }); +const objectResolver = new FileObjectResolver({ + bucket: bucketName, + list: prefix => minioClient.listObjects(bucketName, prefix, true), + stat: key => minioClient.statObject(bucketName, key), + ...(env.FILE_OBJECT_INDEX_ENABLED ? { index: { + get: (key: string) => redisClient.get(key), + set: (key: string, value: string, replace: boolean) => replace + ? redisClient.set(key, value, 'EX', env.SESSION_CACHE_TTL) + : redisClient.set(key, value, 'EX', env.SESSION_CACHE_TTL, 'NX'), + forget: (key: string, value: string) => redisClient.eval( + "if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('DEL', KEYS[1]) end return 0", 1, key, value, + ), + } } : {}), +}); + const minioRegion = process.env.MINIO_REGION ?? process.env.AWS_REGION ?? 'us-east-1'; async function ensureBucketExists(retries = 10, delay = 1000): Promise { @@ -250,6 +268,8 @@ async function uploadFile( * `getObject` / `statObject` without a separate Redis lookup. */ const metaData: Record = { 'Content-Type': mimetype, + // New marker on every PUT, including same-ID overwrites and metadata changes. + 'X-Amz-Meta-Codeapi-Version': randomUUID(), 'X-Amz-Meta-Original-Filename': encodedFilename, 'X-Amz-Meta-Original-Filename-Encoded': 'base64', }; @@ -267,6 +287,7 @@ async function uploadFile( } else { await minioClient.putObject(bucketName, objectName, peeked.body, undefined, metaData); } + await objectResolver.remember(session_id, fileId, objectName); logger.info(`[${INSTANCE_ID}] File ID: ${fileId} | Filename: ${filename} | Session key: ${sessionKey}`); await redisClient.set(`upload:${sessionKey}${session_id}${fileId}`, 'true', 'EX', env.SESSION_CACHE_TTL); fileUploads.inc(); @@ -443,30 +464,14 @@ app.get('/sessions/:session_id/objects/:objectId/metadata', async (req, res) => const { session_id, objectId } = req.params; try { - const stream = minioClient.listObjects(bucketName, `${session_id}/${objectId}`, true); - let objectName = ''; - - for await (const obj of stream) { - if (obj.name.startsWith(`${session_id}/${objectId}`) === true) { - objectName = obj.name; - break; - } - } - - if (!objectName) { - return res.status(404).json({ - error: 'File not found', - details: 'No matching file found', - session_id, - objectId, - }); - } - - const stat: Partial = await minioClient.statObject(bucketName, objectName); + const resolved = await objectResolver.metadata(session_id, objectId); + if (!resolved) return res.status(404).json({ error: 'File not found' }); + const { key: objectName, stat } = resolved; const originalFilename = originalFilenameFromMetadata(stat.metaData); return res.status(200).json({ name: objectName, + version: stat.metaData?.['codeapi-version'], ...(originalFilename ? { originalFilename } : {}), size: stat.size, lastModified: stat.lastModified, @@ -487,87 +492,33 @@ app.get('/sessions/:session_id/objects/:objectId', async (req, res) => { const { session_id, objectId } = req.params; try { - // List objects to find the correct file with extension - const stream = minioClient.listObjects(bucketName, `${session_id}/${objectId}`, true); - let objectName = ''; - - for await (const obj of stream) { - if (obj.name.startsWith(`${session_id}/${objectId}`) === true) { - objectName = obj.name; - break; - } - } - - if (!objectName) { - logger.warn('File not found', { session_id, objectId, bucketName }); - return res.status(404).json({ - error: 'File not found', - details: 'No matching file found', - session_id, - objectId, - bucketName - }); - } - - logger.info(`[${INSTANCE_ID}] Attempting to download: ${objectName}`); - - const stat: Partial = await minioClient.statObject(bucketName, objectName); - - const originalFilename = originalFilenameFromMetadata(stat.metaData); - - logger.info(`[${INSTANCE_ID}] File found: ${objectName}`); - - // Explicitly remove problematic headers that might be duplicated - res.removeHeader('Transfer-Encoding'); - res.removeHeader('Date'); - - /* An object-key basename is only a storage identifier, not an original - * filename. If an S3-compatible backend drops user metadata, retain - * attachment semantics but omit the filename so the runner uses its - * caller-supplied destination. */ - res.setHeader('Content-Disposition', contentDispositionForOriginalFilename(originalFilename)); - if (stat.metaData?.['content-type'] != null) { - res.setHeader('Content-Type', stat.metaData['content-type']); - } - /* Surface the read-only flag on download so the sandbox can plumb it - * onto its in-memory file metadata without a separate metadata fetch. - * MinIO normalizes `X-Amz-Meta-Read-Only` to `read-only` in stat.metaData. */ - if (stat.metaData?.['read-only'] === 'true') { - res.setHeader('X-Read-Only', 'true'); - } - + const objectName = await objectResolver.resolve(session_id, objectId); + if (!objectName) return res.status(404).json({ error: 'File not found' }); const dataStream = await minioClient.getObject(bucketName, objectName); - fileDownloads.inc(); - - dataStream.on('data', (chunk) => { - res.write(chunk); - }); - - dataStream.on('end', () => { - res.end(); - }); - - dataStream.on('error', (err) => { - logger.error('Error streaming file:', { error: err, session_id, objectId, bucketName }); - // Only send error if headers haven't been sent yet - if (!res.headersSent) { - res.status(500).json({ - error: 'Error streaming file', - details: err.message - }); - } else { - res.end(); + try { + const headers = (dataStream as Readable & { headers?: Record }).headers ?? {}; + if (!headers['x-amz-meta-codeapi-version'] || !headers['x-amz-meta-original-filename']) { + // Preserve legacy/S3-compatible metadata behavior without promoting a + // later HEAD's version marker onto bytes from an earlier GET. + const stat = await minioClient.statObject(bucketName, objectName); + if (headers.etag?.replace(/^"|"$/g, '') !== stat.etag) { + return res.status(409).json({ error: 'Input changed during metadata lookup' }); + } + for (const [key, value] of Object.entries(stat.metaData ?? {})) { + if (key !== 'codeapi-version') headers[`x-amz-meta-${key}`] ??= value; + } } - }); + fileDownloads.inc(); + await sendFileDownload(dataStream, res, req.header('x-codeapi-input-version')); + } finally { + dataStream.destroy(); + } } catch (err) { - logger.error('Error downloading file:', { error: err, session_id, objectId, bucketName }); - return res.status(500).json({ - error: 'Error downloading file', - details: (err as Error | undefined)?.message, - session_id, - objectId, - bucketName - }); + logger.error('Error downloading file', { error: err, session_id, objectId }); + if (!res.headersSent && !res.destroyed) { + const missing = ['NoSuchKey', 'NotFound', 'NoSuchObject'].includes((err as { code?: string }).code ?? ''); + return res.status(missing ? 404 : 500).json({ error: 'Error downloading file' }); + } } }); @@ -585,7 +536,7 @@ function parseObjectName(objectName: string | undefined): { session_id: string; return { session_id, file_id }; } -const detailLevels: Record Promise> | undefined> = { +const detailLevels: Record Promise>> = { simple: async (obj: BucketItem): Promise> => obj.name ?? '', summary: async (obj: BucketItem): Promise> => ({ name: obj.name, @@ -639,14 +590,9 @@ app.get('/sessions/:session_id/objects', async (req, res) => { const { detail = 'simple' } = req.query; try { - const stream = minioClient.listObjects(bucketName, session_id, true); - const objects: (t.ObjectTypes | Partial | undefined)[] = []; - + const stream = minioClient.listObjects(bucketName, `${session_id}/`, true); const getDetail = detailLevels[detail as string] ?? detailLevels.simple; - - for await (const obj of stream) { - objects.push(await getDetail(obj)); - } + const objects = await mapObjectDetails(stream, getDetail, env.FILE_METADATA_CONCURRENCY); res.json(objects); } catch (err) { diff --git a/service/src/test/redis.ts b/service/src/test/redis.ts new file mode 100644 index 00000000..f243f657 --- /dev/null +++ b/service/src/test/redis.ts @@ -0,0 +1,45 @@ +import { spawn } from 'node:child_process'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import IORedis from 'ioredis'; + +/** Real Lua semantics, isolated Unix socket, no TCP listener or durable data. */ +export async function startTestRedis(): Promise }> { + const dir = await mkdtemp(path.join(tmpdir(), 'codeapi-redis-')); + const socket = path.join(dir, 'redis.sock'); + const process = spawn('redis-server', [ + '--port', '0', '--unixsocket', socket, '--unixsocketperm', '700', + '--save', '', '--appendonly', 'no', '--dir', dir, + ], { stdio: 'ignore' }); + let failure: Error | undefined; + process.on('error', error => { failure = error; }); + const exited = new Promise(resolve => { + process.once('exit', () => resolve()); + process.once('error', () => resolve()); + }); + const client = new IORedis(socket, { lazyConnect: true, retryStrategy: () => null, maxRetriesPerRequest: 0 }); + client.on('error', () => {}); + const closeTestServer = async (): Promise => { + client.disconnect(); + process.kill('SIGTERM'); + await exited; + await rm(dir, { recursive: true, force: true }); + }; + try { + for (let attempt = 0; attempt < 100; attempt++) { + if (failure) throw failure; + try { + await client.connect(); + await client.ping(); + return Object.assign(client, { closeTestServer }); + } catch { + await new Promise(resolve => setTimeout(resolve, 20)); + } + } + throw new Error('Test Redis did not start; install redis-server'); + } catch (error) { + await closeTestServer(); + throw error; + } +}