diff --git a/.env.example b/.env.example index a4b018bf50..ddd9b4c6b1 100644 --- a/.env.example +++ b/.env.example @@ -726,8 +726,14 @@ TTS_API_KEY= # CODE_SANDBOX_PREWARM=true # Time in milliseconds before LibreChat treats a tracked sandbox as cold (default: 2100000 / 35 minutes). # CODE_SANDBOX_COLD_AFTER_MS=2100000 -# Bytes read per sandbox image chunk; lower this if the runner has a small stdout limit (default: 32768). -# LIBRECHAT_CODE_IMAGE_CHUNK_BYTES=32768 +# Sandbox stdout a single /exec response may carry (default: 65536). Sandbox image reads are +# windowed to fill this budget, so matching it to the runner's SANDBOX_OUTPUT_MAX_SIZE keeps the +# round-trip count — and the load on the Code API's execution rate limit — as low as possible. +# A runner with a smaller cap is also detected at runtime, at the cost of one discarded read. +# LIBRECHAT_CODE_SANDBOX_OUTPUT_MAX_SIZE=65536 +# Bytes read per sandbox image window, overriding the size derived from the stdout budget above. +# Prefer setting the budget; use this only to pin a window size exactly. +# LIBRECHAT_CODE_IMAGE_CHUNK_BYTES= #==================================================# # RAG # diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index 311d1c76cd..a3c2e5283f 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -11,7 +11,11 @@ const { flattenArtifactPath, createAxiosInstance, getCodeApiAuthHeaders, + withCodeApiRateLimit, classifyCodeArtifact, + parseSandboxImageChunk, + readWindowedSandboxImage, + createCodeApiRateLimitBudget, codeServerHttpAgent, codeServerHttpsAgent, extractCodeArtifactText, @@ -1493,12 +1497,16 @@ async function readSandboxFile({ * Reads a small image file out of the code-execution sandbox as base64 so * `read_file` can surface it to vision-capable models. `readSandboxFile`'s * `cat` round-trips stdout through codeapi's JSON transport, which lossily - * replaces non-UTF-8 bytes and corrupts image data. Here a tiny Python - * reader stats the file, refuses (without transferring) anything over - * `maxBytes`, and otherwise base64-encodes the bytes IN the sandbox so the - * payload stays ASCII-safe across the JSON `/exec` transport. Session - * forwarding mirrors `readSandboxFile` so the read lands in the same - * sandbox session that holds the agent's prior-turn artifacts. + * replaces non-UTF-8 bytes and corrupts image data. The in-sandbox reader + * base64-encodes the bytes instead, so the payload stays ASCII-safe across + * the JSON `/exec` transport. Session forwarding mirrors `readSandboxFile` + * so the read lands in the same sandbox session that holds the agent's + * prior-turn artifacts. + * + * Windowing, window sizing, and assembly live in `@librechat/api` + * (`readWindowedSandboxImage`); this function is the `/exec` transport it + * calls, plus the rate-limit wait that keeps a multi-window read from + * discarding the bytes it already pulled. * * @param {Object} params * @param {string} params.file_path - Path inside the sandbox (e.g. `/mnt/data/chart.png`). @@ -1507,7 +1515,8 @@ async function readSandboxFile({ * @param {string} [params.runtime_session_hint] - Per-conversation stateful runtime-session hint. * @param {number} [params.maxBytes] - In-sandbox size cap; larger files return `{ tooLarge, bytes }`. * @param {ServerRequest} [params.req] - Current authenticated request, used to mint Code API auth. - * @returns {Promise<{base64: string, bytes: number} | {tooLarge: true, bytes: number} | null>} + * @returns {Promise<{base64: string, bytes: number} + * | {tooLarge: true, reason: 'size' | 'round_trips', bytes: number} | null>} * `null` when codeapi is unavailable; throws on transport / read errors. */ async function readSandboxImage({ @@ -1525,124 +1534,35 @@ async function readSandboxImage({ return null; } - const limit = typeof maxBytes === 'number' && maxBytes > 0 ? maxBytes : 5 * megabyte; - const chunkBytes = getImageChunkBytes(); - const maxChunks = Math.ceil(limit / chunkBytes) + 1; - - /** @type {Buffer[]} */ - const parts = []; - let offset = 0; - let total = null; - - for (let i = 0; i < maxChunks; i++) { - const payload = Buffer.from( - JSON.stringify({ file_path, limit, offset, chunk: chunkBytes }), - 'utf8', - ).toString('base64'); - const code = [ - "python3 - <<'PY'", - 'import base64, json, os, stat', - `payload = ${JSON.stringify(payload)}`, - "data = json.loads(base64.b64decode(payload).decode('utf-8'))", - "p = data['file_path']", - "limit = data['limit']", - "offset = data['offset']", - "chunk = data['chunk']", - 'try:', - ' st = os.stat(p)', - 'except OSError as e:', - ' print(json.dumps({"error": str(e)}))', - ' raise SystemExit(0)', - // Reject FIFOs, sockets, and device files (e.g. a symlink to /dev/zero): - // os.stat can report a small/zero size while an unbounded read blocks or - // streams forever until the request times out. - 'if not stat.S_ISREG(st.st_mode):', - ' print(json.dumps({"error": "not a regular file"}))', - ' raise SystemExit(0)', - 'if st.st_size > limit:', - ' print(json.dumps({"too_large": True, "bytes": st.st_size}))', - ' raise SystemExit(0)', - // Read only this window. The whole base64 payload cannot be emitted in - // one shot: the runner caps stdout at SANDBOX_OUTPUT_MAX_SIZE (1024 - // bytes by default) and SIGKILLs the job on overflow, which truncates - // the JSON mid-string. Windowing keeps every response under that cap. - "with open(p, 'rb') as f:", - ' f.seek(offset)', - ' raw = f.read(chunk)', - 'print(json.dumps({"total": st.st_size, "n": len(raw), "b64": base64.b64encode(raw).decode("ascii")}))', - 'PY', - ].join('\n'); - - const parsed = await execSandboxImageChunk({ - baseURL, - code, - file_path, - session_id, - runtime_session_hint, - executionProfile, - files, - req, - chunkBytes, - }); - - if (parsed.error) { - throw new Error(String(parsed.error)); - } - if (parsed.too_large === true) { - return { tooLarge: true, bytes: Number(parsed.bytes) || 0 }; - } - if (typeof parsed.b64 !== 'string' || typeof parsed.n !== 'number') { - return null; - } - - if (total == null) { - total = Number(parsed.total) || 0; - if (total > limit) { - return { tooLarge: true, bytes: total }; - } - } else if (Number(parsed.total) !== total) { - /* The file changed underneath us; a spliced-together buffer would be - * a mix of two versions rather than any real image. */ - throw new Error(`"${file_path}" changed while being read from the sandbox`); - } - - parts.push(Buffer.from(parsed.b64, 'base64')); - offset += parsed.n; - - if (parsed.n === 0 || offset >= total) { - break; - } - } - - const buffer = Buffer.concat(parts); - if (total == null) { - return null; - } - if (buffer.length !== total) { - /* Ran out of chunk budget (or short reads); returning a partial image - * would render as a corrupt file, so surface it as unreadable-inline. */ - return { tooLarge: true, bytes: total }; - } - return { base64: buffer.toString('base64'), bytes: buffer.length }; + /** Every window is one `/exec` call against the Code API's per-user + * execution limiter, so the read shares one wait budget: a window that + * resets mid-read is worth pausing for, an exhausted budget is not. */ + const rateLimit = createCodeApiRateLimitBudget(); + return readWindowedSandboxImage({ + filePath: file_path, + baseUrl: baseURL, + limit: typeof maxBytes === 'number' && maxBytes > 0 ? maxBytes : 5 * megabyte, + readChunk: ({ code }) => + execSandboxImageChunk({ + baseURL, + code, + file_path, + session_id, + runtime_session_hint, + executionProfile, + files, + req, + rateLimit, + }), + }); } /** - * Raw bytes pulled per `/exec` round-trip when inlining a sandbox image. - * Each chunk is base64-encoded (~1.33x) into the response's stdout, which - * the runner truncates + SIGKILLs past `SANDBOX_OUTPUT_MAX_SIZE`. The - * default leaves headroom for a runner configured at 64KB; deployments - * with a smaller cap must lower this, and a larger cap can raise it to cut - * round-trips. - * @returns {number} - */ -function getImageChunkBytes() { - const parsed = Number(process.env.LIBRECHAT_CODE_IMAGE_CHUNK_BYTES); - return Number.isFinite(parsed) && parsed > 0 ? Math.floor(parsed) : 32 * 1024; -} - -/** - * Runs one image-chunk read over `/exec` and parses its JSON line. - * @returns {Promise>} + * Runs one image-window read over `/exec` and hands the response to the + * shared parser. Rate limits are waited out inside the shared budget; a + * truncated response comes back as a chunk the reader narrows for, not an + * error, so it is neither logged nor thrown here. + * @returns {Promise} */ async function execSandboxImageChunk({ baseURL, @@ -1653,7 +1573,7 @@ async function execSandboxImageChunk({ executionProfile, files, req, - chunkBytes, + rateLimit, }) { /** @type {Record} */ const postData = { lang: 'bash', code }; @@ -1668,52 +1588,32 @@ async function execSandboxImageChunk({ } try { - const authHeaders = await getCodeApiAuthHeaders(req); - const response = await axios({ - method: 'post', - url: `${baseURL}/exec`, - data: postData, - headers: { - 'Content-Type': 'application/json', - 'User-Agent': 'LibreChat/1.0', - ...authHeaders, - ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), + const response = await withCodeApiRateLimit({ + label: `reading "${file_path}" from the sandbox`, + budget: rateLimit, + onWait: (waitMs) => + logger.warn( + `[readSandboxImage] Rate-limited reading "${file_path}"; retrying in ${waitMs}ms`, + ), + attempt: async () => { + const authHeaders = await getCodeApiAuthHeaders(req); + return axios({ + method: 'post', + url: `${baseURL}/exec`, + data: postData, + headers: { + 'Content-Type': 'application/json', + 'User-Agent': 'LibreChat/1.0', + ...authHeaders, + ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), + }, + httpAgent: codeServerHttpAgent, + httpsAgent: codeServerHttpsAgent, + timeout: 15000, + }); }, - httpAgent: codeServerHttpAgent, - httpsAgent: codeServerHttpsAgent, - timeout: 15000, }); - const result = response?.data ?? {}; - /* The runner truncates stdout at SANDBOX_OUTPUT_MAX_SIZE and SIGKILLs the - * job (status `OL`). Detect that explicitly: the surviving stdout is a - * base64 string cut mid-flight, so parsing it yields a misleading - * "unexpected output" instead of naming the real, fixable cause. */ - if (result.status === 'OL') { - throw new Error( - `Reading "${file_path}" exceeded the sandbox stdout limit (chunk ${chunkBytes} bytes). ` + - 'Lower LIBRECHAT_CODE_IMAGE_CHUNK_BYTES or raise SANDBOX_OUTPUT_MAX_SIZE on the runner.', - ); - } - if (result.stderr && (result.stdout == null || result.stdout === '')) { - throw new Error(String(result.stderr).trim()); - } - if (result.stdout == null || String(result.stdout).trim() === '') { - return {}; - } - /* Parse the LAST non-empty line: the reader's JSON is the final thing it - * prints, so anything a shell profile or library emitted ahead of it - * (banners, warnings) must not break the read. */ - const lines = String(result.stdout) - .split('\n') - .map((line) => line.trim()) - .filter(Boolean); - try { - return JSON.parse(lines[lines.length - 1]); - } catch { - throw new Error( - `Unexpected output while reading image bytes from the sandbox: ${String(result.stdout).slice(0, 120)}`, - ); - } + return parseSandboxImageChunk(response?.data ?? {}); } catch (error) { logAxiosError({ message: `Error reading sandbox image "${file_path}"`, diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index 89e56d584e..e951c485e0 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -55,6 +55,9 @@ const mockGetExtractedTextFormat = jest.fn((_name, _mime, text) => (text == null * legacy single-phase tests below (txt/png/etc) exercise the inline path * unchanged. The dedicated office/finalize describe block toggles it on. */ const mockHasOfficeHtmlPath = jest.fn(() => false); +/* Stands in for the real chunk parser so a transport test can decide what + * one `/exec` response means without reaching into `packages/api`. */ +const mockParseSandboxImageChunk = jest.fn((response) => response); /* Pass-through `withTimeout`: tests don't drive timeouts here (those live * in promise.spec.ts and the finalizePreview unit tests below). */ const passthroughWithTimeout = async (promise) => promise; @@ -68,6 +71,19 @@ jest.mock('@librechat/api', () => { flattenArtifactPath: jest.fn((name) => name.replace(/\//g, '__')), createAxiosInstance: jest.fn(() => mockAxios), getCodeApiAuthHeaders: jest.fn(async () => ({})), + /* Windowing, sizing and rate-limit policy are real code in + * `packages/api` with their own tests (`files/code/image.spec.ts`, + * `utils/code.spec.ts`). These stand-ins are deliberately inert + * passthroughs — they assert nothing about that behavior, so the + * transport tests below cover only what this file still owns. */ + createCodeApiRateLimitBudget: jest.fn(() => ({ remainingMs: 20_000 })), + withCodeApiRateLimit: jest.fn(({ attempt }) => attempt()), + buildSandboxImageReaderCode: jest.fn( + ({ filePath, limit, offset, chunkBytes }) => + `read ${filePath} ${limit} ${offset} ${chunkBytes}`, + ), + parseSandboxImageChunk: jest.fn((response) => mockParseSandboxImageChunk(response)), + readWindowedSandboxImage: jest.fn(({ readChunk }) => readChunk({ code: 'read-one-window' })), getCodeExecutionBaseUrl: jest.fn((profile) => profile === 'stateful' ? 'https://code-stateful.example.com' : 'https://code-api.example.com', ), @@ -2776,111 +2792,87 @@ describe('Code Process', () => { }); /** - * These drive the REAL reader against a mocked `/exec` transport (rather - * than mocking `readSandboxImage` itself), because the bug this covers - * lived entirely in the transport: base64 leaves the sandbox on stdout, - * which the runner truncates + SIGKILLs past `SANDBOX_OUTPUT_MAX_SIZE`. + * Windowing, window sizing, rate-limit waiting and byte assembly moved to + * `packages/api` (`files/code/image.ts`, `utils/code.ts`) and are tested + * there against the real implementations. What remains here is the piece + * this file still owns: the `/exec` transport — session forwarding, the + * profile header, and handing the response to the shared parser. */ - describe('readSandboxImage', () => { - const crypto = require('crypto'); - const PNG_HEADER = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); - /** Reply as the sandbox would: serve `buffer` through the windowed reader. */ - const serveFile = (buffer) => - mockAxios.mockImplementation(async ({ data }) => { - const payload = JSON.parse( - Buffer.from(JSON.parse(/payload = ("[^"]+")/.exec(data.code)[1]), 'base64').toString(), - ); - const slice = buffer.subarray(payload.offset, payload.offset + payload.chunk); - return { - data: { - stdout: JSON.stringify({ - total: buffer.length, - n: slice.length, - b64: slice.toString('base64'), - }), - }, - }; - }); - + describe('readSandboxImage transport', () => { beforeEach(() => { process.env.LIBRECHAT_CODE_BASEURL = 'http://code.test/v1'; - delete process.env.LIBRECHAT_CODE_IMAGE_CHUNK_BYTES; mockAxios.mockReset(); + mockParseSandboxImageChunk.mockReset(); + mockParseSandboxImageChunk.mockImplementation((response) => response); }); - it('reassembles an image larger than one chunk, byte-for-byte', async () => { - /* 200KB of PNG-headed noise: > 6 chunks at the 32KB default, and the - * exact shape that used to blow the stdout cap and SIGKILL the job. */ - const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(200 * 1024)]); - serveFile(source); + it('forwards the sandbox session, runtime hint and profile header', async () => { + mockAxios.mockResolvedValue({ data: { stdout: '{}' } }); - const result = await readSandboxImage({ file_path: '/mnt/data/big.png' }); - - expect(mockAxios.mock.calls.length).toBeGreaterThan(1); - expect(result.bytes).toBe(source.length); - expect(Buffer.from(result.base64, 'base64').equals(source)).toBe(true); - }); - - it('reads a single-chunk image in one round-trip', async () => { - const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(1024)]); - serveFile(source); - - const result = await readSandboxImage({ file_path: '/mnt/data/small.png' }); - - expect(mockAxios).toHaveBeenCalledTimes(1); - expect(Buffer.from(result.base64, 'base64').equals(source)).toBe(true); - }); - - it('names the real cause when a chunk overflows the runner stdout cap', async () => { - /* The runner truncates stdout and SIGKILLs with status `OL`; the old - * reader parsed the truncated base64 and reported a misleading - * "unexpected output" instead of the fixable limit. */ - mockAxios.mockResolvedValue({ - data: { stdout: '{"total":999999,"n":32768,"b64":"iVBORw0KGg', status: 'OL', code: 137 }, + await readSandboxImage({ + file_path: '/mnt/data/chart.png', + session_id: 'session-abc', + runtime_session_hint: 'hint-xyz', + files: [{ id: 'file-1', name: 'seed.csv' }], + codeApiBaseUrl: 'https://code-stateful.example.com', + executionProfile: 'stateful', }); - await expect(readSandboxImage({ file_path: '/mnt/data/big.png' })).rejects.toThrow( - /exceeded the sandbox stdout limit/, + expect(mockAxios).toHaveBeenCalledTimes(1); + const call = mockAxios.mock.calls[0][0]; + expect(call.url).toBe('https://code-stateful.example.com/exec'); + expect(call.data).toMatchObject({ + lang: 'bash', + session_id: 'session-abc', + runtime_session_hint: 'hint-xyz', + files: [{ id: 'file-1', name: 'seed.csv' }], + }); + expect(call.headers['X-CodeAPI-Expected-Profile']).toBe('stateful'); + }); + + it('omits the profile header and optional session fields when unset', async () => { + mockAxios.mockResolvedValue({ data: { stdout: '{}' } }); + + await readSandboxImage({ file_path: '/mnt/data/chart.png' }); + + const call = mockAxios.mock.calls[0][0]; + expect(call.data.session_id).toBeUndefined(); + expect(call.data.runtime_session_hint).toBeUndefined(); + expect(call.data.files).toBeUndefined(); + expect(call.headers['X-CodeAPI-Expected-Profile']).toBeUndefined(); + }); + + it('hands the raw `/exec` response to the shared chunk parser', async () => { + const data = { stdout: '{"total":3,"n":3,"b64":"AAAA"}', status: null }; + mockAxios.mockResolvedValue({ data }); + mockParseSandboxImageChunk.mockReturnValue({ total: 3, n: 3, b64: 'AAAA' }); + + const result = await readSandboxImage({ file_path: '/mnt/data/x.png' }); + + expect(mockParseSandboxImageChunk).toHaveBeenCalledWith(data); + expect(result).toEqual({ total: 3, n: 3, b64: 'AAAA' }); + }); + + it('returns null when no code base URL is configured', async () => { + const { getCodeBaseURL } = require('@librechat/agents'); + getCodeBaseURL.mockReturnValueOnce(''); + + await expect(readSandboxImage({ file_path: '/mnt/data/x.png' })).resolves.toBeNull(); + expect(mockAxios).not.toHaveBeenCalled(); + }); + + it('logs and rethrows a transport failure', async () => { + const { logAxiosError } = require('@librechat/api'); + mockAxios.mockRejectedValue(new Error('codeapi unreachable')); + + await expect(readSandboxImage({ file_path: '/mnt/data/x.png' })).rejects.toThrow( + 'codeapi unreachable', + ); + expect(logAxiosError).toHaveBeenCalledWith( + expect.objectContaining({ + message: 'Error reading sandbox image "/mnt/data/x.png"', + }), ); - }); - - it('honors LIBRECHAT_CODE_IMAGE_CHUNK_BYTES', async () => { - process.env.LIBRECHAT_CODE_IMAGE_CHUNK_BYTES = '1024'; - const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(4 * 1024)]); - serveFile(source); - - const result = await readSandboxImage({ file_path: '/mnt/data/x.png' }); - - expect(mockAxios.mock.calls.length).toBe(5); - expect(Buffer.from(result.base64, 'base64').equals(source)).toBe(true); - }); - - it('parses the reader JSON even when the shell emits a banner first', async () => { - const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(64)]); - mockAxios.mockResolvedValue({ - data: { - stdout: `motd banner\n${JSON.stringify({ - total: source.length, - n: source.length, - b64: source.toString('base64'), - })}`, - }, - }); - - const result = await readSandboxImage({ file_path: '/mnt/data/x.png' }); - - expect(Buffer.from(result.base64, 'base64').equals(source)).toBe(true); - }); - - it('refuses an oversize file in-sandbox without transferring bytes', async () => { - mockAxios.mockResolvedValue({ - data: { stdout: JSON.stringify({ too_large: true, bytes: 9 * 1024 * 1024 }) }, - }); - - const result = await readSandboxImage({ file_path: '/mnt/data/huge.png' }); - - expect(result).toEqual({ tooLarge: true, bytes: 9 * 1024 * 1024 }); - expect(mockAxios).toHaveBeenCalledTimes(1); }); }); diff --git a/packages/api/src/agents/handlers.spec.ts b/packages/api/src/agents/handlers.spec.ts index 025f75414a..2a8790dd7d 100644 --- a/packages/api/src/agents/handlers.spec.ts +++ b/packages/api/src/agents/handlers.spec.ts @@ -5601,10 +5601,43 @@ describe('createToolExecuteHandler', () => { expect(result.content).toContain('bash_tool'); }); - it('degrades to the image hint when decoded bytes are truncated (integrity guard)', async () => { + it('reports a round-trip-bound image as unreadable inline, not oversize', async () => { + /* Within the byte cap but needing more windowed `/exec` reads than + * one call may spend on the Code API's execution limiter. Saying + * "over the inline limit" here would misstate a fixable cause. */ + const readSandboxImage = jest.fn(async () => ({ + tooLarge: true as const, + reason: 'round_trips' as const, + bytes: 900_000, + inlineCeiling: 489_600, + })); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + readSandboxImage, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_trips', + name: Constants.READ_FILE, + args: { path: '/mnt/data/wide.png' }, + }, + ]); + + expect(result.status).toBe('success'); + /* Names the size that would actually work, so the model has a + * downscale target instead of a guess. */ + expect(result.content).toContain('489600'); + expect(result.content).not.toContain('inline limit'); + expect(result.content).toContain('Downscale'); + }); + + it('reports a truncated transfer as a failed read, not an unreadable format', async () => { /* Simulate codeapi clipping a large `/exec` stdout: the reported - * size does not match the decoded base64 length, so the bytes are - * unsafe to forward and we fall back to the bash hint. */ + * size does not match the decoded base64 length. The bytes are + * unsafe to forward, but the read is retryable — saying the file + * "cannot be read as text" would report a permanent limit. */ const readSandboxImage = jest.fn(async () => ({ base64: PNG_B64, bytes: pngBytes + 100 })); const handler = makeReadFileHandler({ codeEnvAvailable: true, @@ -5622,11 +5655,39 @@ describe('createToolExecuteHandler', () => { expect(result.status).toBe('error'); expect(result.artifact).toBeUndefined(); - expect(result.errorMessage).toContain('image file'); - expect(result.errorMessage).toContain('bash_tool'); + expect(result.errorMessage).toContain('truncated transfer'); + expect(result.errorMessage).toContain('Retry the read'); + expect(result.errorMessage).not.toContain('cannot be read as text'); }); - it('degrades to the image hint when the image reader throws', async () => { + it('reports a missing interpreter as itself, not as a missing image path', async () => { + /* The sandbox reader surfaces `python3: not found` on stderr. A + * generic "not found" match would send the model to `ls /mnt/data` + * and hide the runner dependency the operator has to fix. */ + const readSandboxImage = jest.fn(async () => { + throw new Error('python3: not found'); + }); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + readSandboxImage, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_nopython', + name: Constants.READ_FILE, + args: { path: '/mnt/data/chart.png' }, + }, + ]); + + expect(result.status).toBe('error'); + expect(result.errorMessage).toContain('python3: not found'); + expect(result.errorMessage).not.toContain('was not found in the code-execution sandbox'); + expect(result.errorMessage).not.toContain('ls /mnt/data'); + }); + + it('surfaces the transport failure when the image reader throws', async () => { const readSandboxImage = jest.fn(async () => { throw new Error('codeapi unreachable'); }); @@ -5645,8 +5706,62 @@ describe('createToolExecuteHandler', () => { ]); expect(result.status).toBe('error'); - expect(result.errorMessage).toContain('image file'); - expect(result.errorMessage).toContain('bash_tool'); + expect(result.errorMessage).toContain('codeapi unreachable'); + expect(result.errorMessage).toContain('Retry the read'); + expect(result.errorMessage).not.toContain('cannot be read as text'); + }); + + it('tells the model to wait when the sandbox rate-limited the read', async () => { + /* Each window is one `/exec` call against a per-user limiter, so a + * chart-heavy turn can exhaust it. The old catch-all told the model + * images are unreadable, which stopped it from ever retrying. */ + const readSandboxImage = jest.fn(async () => { + throw new Error( + 'Code API rate limit reached while reading "/mnt/data/7_interest_gap.png" from the sandbox (retry in 17s).', + ); + }); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + readSandboxImage, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_429', + name: Constants.READ_FILE, + args: { path: '/mnt/data/7_interest_gap.png' }, + }, + ]); + + expect(result.status).toBe('error'); + expect(result.errorMessage).toContain('rate limit reached'); + expect(result.errorMessage).toContain('read it once more'); + expect(result.errorMessage).not.toContain('cannot be read as text'); + }); + + it('points a missing image at the directory listing instead of the bytes', async () => { + const readSandboxImage = jest.fn(async () => { + throw new Error("[Errno 2] No such file or directory: '/mnt/data/gone.png'"); + }); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + readSandboxImage, + }); + + const [result] = await invokeHandler(handler, [ + { + id: 'call_missing', + name: Constants.READ_FILE, + args: { path: '/mnt/data/gone.png' }, + }, + ]); + + expect(result.status).toBe('error'); + expect(result.errorMessage).toContain('was not found'); + expect(result.errorMessage).toContain('ls /mnt/data'); + expect(result.errorMessage).not.toContain('cannot be read as text'); }); it('rejects non-image binary types with a bash-pointing message (not the image path)', async () => { diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 3e8ede7dc4..71099d8aa0 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -473,7 +473,14 @@ export interface ToolExecuteOptions { /** In-sandbox size cap; files larger than this return `tooLarge` without transferring bytes. */ maxBytes?: number; req?: ServerRequest; - }) => Promise<{ base64: string; bytes: number } | { tooLarge: true; bytes: number } | null>; + }) => Promise< + | { base64: string; bytes: number } + /** `size`: over `maxBytes`. `round_trips`: within the byte cap, but more + * windowed `/exec` reads than one call may spend on the Code API's + * per-user execution limiter. */ + | { tooLarge: true; reason?: 'size' | 'round_trips'; bytes: number; inlineCeiling?: number } + | null + >; /** * Writes a UTF-8 text file into the code-execution sandbox via the * sandbox `/exec` endpoint. Mirrors `readSandboxFile` session forwarding @@ -1848,15 +1855,56 @@ function looksBinary(content: string): boolean { return false; } +/** + * True for the errors a sandbox raises about the requested PATH, and only + * those. Deliberately narrower than {@link isSandboxMissingFileError}: that + * predicate also accepts a bare "not found", which the sandbox reader emits + * for a missing interpreter (`python3: not found`). Reporting that as a + * missing image would send the model to `ls /mnt/data` while hiding a + * runner dependency the operator needs to see. + */ +function isMissingSandboxPathError(reason: string): boolean { + const message = reason.toLowerCase(); + return ( + message.includes('no such file or directory') || + message.includes('cannot access') || + message.includes('cannot find the path') || + message.includes('enoent') + ); +} + +/** + * Model-visible error for an image the sandbox could not hand back. The + * read is a supported operation that FAILED, so the message must not reuse + * the "images cannot be read as text" phrasing — that reads as a permanent + * capability limit and stops the model from ever retrying. State the real + * cause and the affordance that matches it: a rate-limited or truncated + * read is worth retrying, a missing path is worth listing, and only a + * genuine transport dead end falls back to `bash_tool`. Classification is + * by message, matching how `isSandboxMissingFileError` already reads + * sandbox failures; the rate-limit wording is the one `readSandboxImage` + * throws when the Code API limiter turns a chunk away. + */ +function buildImageReadError(filePath: string, reason: string): string { + const detail = reason.replace(/\.$/, ''); + if (isMissingSandboxPathError(reason)) { + return `"${filePath}" was not found in the code-execution sandbox (${detail}). List the directory with \`bash_tool\` (e.g. \`ls /mnt/data\`) to find the correct path.`; + } + if (/rate limit/i.test(reason)) { + return `Could not read image "${filePath}": ${detail}. Wait for the sandbox to accept requests again, then read it once more.`; + } + return `Could not read image "${filePath}" from the code-execution sandbox: ${detail}. Retry the read; if it keeps failing, inspect the file with \`bash_tool\` (e.g. \`file ${filePath}\`).`; +} + /** * Reads a sandbox image as a viewable artifact so `read_file` can hand the * bytes to vision-capable models instead of refusing them. Fetches the file * base64-encoded from the sandbox (`readSandboxImage`), verifies the decoded * length matches the size the sandbox reported (guards against codeapi * truncating a large `/exec` stdout into a corrupt image), sniffs the real - * MIME, and returns the shared image-artifact result. Degrades to the - * text-oriented binary hint when the reader is unavailable, the image is - * over the inline cap, or the read fails — never throws. + * MIME, and returns the shared image-artifact result. Never throws: a + * mislabeled or corrupt image degrades to the binary hint, while a failed + * read reports what actually went wrong (see {@link buildImageReadError}). */ async function handleSandboxImageRead( tc: ToolCallRequest, @@ -1879,12 +1927,21 @@ async function handleSandboxImageRead( content: '', errorMessage: buildBinaryFileError(filePath, ext), }); + const readFailure = (reason: string): ToolExecuteResult => ({ + toolCallId: tc.id, + status: 'error', + content: '', + errorMessage: buildImageReadError(filePath, reason), + }); if (!readSandboxImage) { return binaryHint(); } const ctx = tc.codeSessionContext as SandboxSessionContext | undefined; - let read: { base64: string; bytes: number } | { tooLarge: true; bytes: number } | null; + let read: + | { base64: string; bytes: number } + | { tooLarge: true; reason?: 'size' | 'round_trips'; bytes: number; inlineCeiling?: number } + | null; try { read = await readSandboxImage({ file_path: filePath, @@ -1895,9 +1952,9 @@ async function handleSandboxImageRead( ...(req ? { req } : {}), }); } catch (error) { - const message = error instanceof Error ? error.message : String(error); + const message = getThrownValueMessage(error); logger.warn(`[handleReadFileCall] Sandbox image read failed for "${filePath}": ${message}`); - return binaryHint(); + return readFailure(message); } if (!read) { @@ -1905,10 +1962,22 @@ async function handleSandboxImageRead( } if ('tooLarge' in read) { onSuccess?.(); + /* Name the size that would actually work: each window costs one sandbox + * execution, so what can be inlined depends on the runner's stdout + * budget, not only on the byte cap. Without a target the model can only + * guess how far to downscale. */ + const ceiling = + read.reason === 'round_trips' && read.inlineCeiling != null + ? read.inlineCeiling + : MAX_SANDBOX_INLINE_IMAGE_BYTES; + const overBudget = + read.reason === 'round_trips' + ? `more than this sandbox can return inline (about ${ceiling} bytes)` + : `over the ${MAX_SANDBOX_INLINE_IMAGE_BYTES}-byte inline limit`; return { toolCallId: tc.id, status: 'success', - content: `Image "${filePath}" is ${read.bytes} bytes, over the ${MAX_SANDBOX_INLINE_IMAGE_BYTES}-byte inline limit. Use \`bash_tool\` to process it (e.g. \`file ${filePath}\` for metadata).`, + content: `Image "${filePath}" is ${read.bytes} bytes, ${overBudget}. Downscale it under ${ceiling} bytes in the sandbox with \`bash_tool\` and read the smaller copy to view it, or inspect it with \`bash_tool\` (e.g. \`file ${filePath}\` for metadata).`, }; } @@ -1917,7 +1986,9 @@ async function handleSandboxImageRead( logger.warn( `[handleReadFileCall] Sandbox image byte mismatch for "${filePath}" (decoded ${buffer.length} != reported ${read.bytes})`, ); - return binaryHint(); + return readFailure( + `the sandbox returned ${buffer.length} of ${read.bytes} bytes (truncated transfer)`, + ); } // Resolve the MIME from the actual bytes, never the extension: a file // routed here by its `.png`/`.jpg`/... name whose header matches none of diff --git a/packages/api/src/agents/skillFiles.ts b/packages/api/src/agents/skillFiles.ts index 27bef398b1..8d8bda9c19 100644 --- a/packages/api/src/agents/skillFiles.ts +++ b/packages/api/src/agents/skillFiles.ts @@ -1,5 +1,4 @@ import { Readable } from 'stream'; -import { isAxiosError } from 'axios'; import { Constants } from '@librechat/agents'; import { logger } from '@librechat/data-schemas'; import { @@ -19,10 +18,10 @@ import { inspectContent, UninspectableFileError, } from '~/protection'; +import { createConcurrencyLimiter, getCodeApiRetryAfterMs, getSafeErrorMetadata } from '~/utils'; import { seedCodeFilesIntoSessions, type CodeExecutionProfileRoute } from './codeFilesSession'; import { ContentFilterError, isContentFilterError } from '~/middleware/contentFilter'; import { getCodeExecutionRouteKey, type CodeExecutionContext } from './execution'; -import { createConcurrencyLimiter, getSafeErrorMetadata } from '~/utils'; import { assertSkillFileContentAllowed } from '~/skills/protection'; import { createSkillContentDigest } from './compatibility'; import { extractInvokedSkillsFromPayload } from './run'; @@ -146,25 +145,13 @@ function isCurrentSkillRef( return ref?.kind === 'skill' && ref.version === skillVersion; } -function getRetryAfterMs(error: unknown): number | null { - if (!isAxiosError(error) || error.response?.status !== 429) { - return null; - } - const header = error.response.headers?.['retry-after']; - const seconds = Number(Array.isArray(header) ? header[0] : header); - if (!Number.isFinite(seconds) || seconds < 0) { - return null; - } - return seconds * 1000; -} - /** Single retry on 429, honoring Retry-After up to MAX_RETRY_AFTER_MS. * Runs inside an upload slot so the wait also brakes queued uploads. */ async function retryOn429(attempt: () => Promise, label: string): Promise { try { return await attempt(); } catch (error) { - const retryAfterMs = getRetryAfterMs(error); + const retryAfterMs = getCodeApiRetryAfterMs(error); if (retryAfterMs == null || retryAfterMs > MAX_RETRY_AFTER_MS) { throw error; } diff --git a/packages/api/src/files/code/image.spec.ts b/packages/api/src/files/code/image.spec.ts new file mode 100644 index 0000000000..fba9e88869 --- /dev/null +++ b/packages/api/src/files/code/image.spec.ts @@ -0,0 +1,419 @@ +import crypto from 'crypto'; +import type { SandboxImageChunk } from './image'; +import { + buildSandboxImageReaderCode, + getSandboxImageChunkBytes, + narrowSandboxImageChunkBytes, + parseSandboxImageChunk, + readWindowedSandboxImage, + MAX_SANDBOX_IMAGE_EXEC_CALLS, +} from './image'; + +const PNG_HEADER = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); + +/** Reads the window parameters back out of the generated reader script, + * the same way the sandbox would. */ +function windowParams(code: string): { offset: number; chunk: number; limit: number } { + const encoded = /payload = ("[^"]+")/.exec(code); + if (!encoded) { + throw new Error('reader code carried no payload'); + } + return JSON.parse(Buffer.from(JSON.parse(encoded[1]), 'base64').toString()); +} + +/** Serves `buffer` through the windowed reader, recording each window size. */ +function serveFile( + buffer: Buffer, + options: { acceptedWindow?: number; reportsCap?: boolean } = {}, +) { + const windows: number[] = []; + const readChunk = async ({ code }: { code: string }): Promise => { + const { offset, chunk, limit } = windowParams(code); + windows.push(chunk); + if (options.acceptedWindow != null && chunk > options.acceptedWindow) { + /* Report the cap the way a runner does: by truncating at it. */ + return { + outputOverflow: true, + observedStdoutBytes: options.reportsCap === false ? undefined : options.acceptedWindow, + }; + } + if (buffer.length > limit) { + return { too_large: true, bytes: buffer.length }; + } + const slice = buffer.subarray(offset, offset + chunk); + return { total: buffer.length, n: slice.length, b64: slice.toString('base64') }; + }; + return { readChunk, windows }; +} + +describe('sandbox image window sizing', () => { + const envKeys = ['LIBRECHAT_CODE_IMAGE_CHUNK_BYTES', 'LIBRECHAT_CODE_SANDBOX_OUTPUT_MAX_SIZE']; + const saved: Record = {}; + + beforeEach(() => { + for (const key of envKeys) { + saved[key] = process.env[key]; + delete process.env[key]; + } + }); + + afterEach(() => { + for (const key of envKeys) { + if (saved[key] == null) { + delete process.env[key]; + } else { + process.env[key] = saved[key]; + } + } + }); + + it('fills the default 64KB stdout budget without exceeding it', () => { + const window = getSandboxImageChunkBytes('https://default.example.com'); + /* base64 is 4 bytes per 3, and the reader wraps it in a JSON envelope. */ + const encoded = Math.ceil(window / 3) * 4; + expect(window % 3).toBe(0); + expect(encoded).toBeLessThan(64 * 1024); + }); + + it('derives a window a small runner can actually emit', () => { + /* An 8KB cap used to floor at an 8KB window, whose base64 alone is + * ~10.9KB — every read overflowed, and narrowing was pinned to the + * same floor, so no retry could ever succeed. */ + process.env.LIBRECHAT_CODE_SANDBOX_OUTPUT_MAX_SIZE = String(8 * 1024); + const window = getSandboxImageChunkBytes('https://small.example.com'); + expect(Math.ceil(window / 3) * 4).toBeLessThan(8 * 1024); + }); + + it('uses an explicit chunk override verbatim', () => { + process.env.LIBRECHAT_CODE_IMAGE_CHUNK_BYTES = '1024'; + expect(getSandboxImageChunkBytes('https://override.example.com')).toBe(1024); + }); + + it('narrows per base URL, and never past an explicit override', () => { + const url = 'https://narrowing.example.com'; + const before = getSandboxImageChunkBytes(url); + const narrowed = narrowSandboxImageChunkBytes(before, url); + + expect(narrowed).not.toBeNull(); + expect(narrowed).toBeLessThan(before); + expect(getSandboxImageChunkBytes(url)).toBe(narrowed); + expect(getSandboxImageChunkBytes('https://untouched.example.com')).toBe(before); + + process.env.LIBRECHAT_CODE_IMAGE_CHUNK_BYTES = '512'; + expect(getSandboxImageChunkBytes(url)).toBe(512); + }); + + it("halves the caller's own failing window, not the shared learned value", () => { + /* Reads run concurrently (one per tool call). Two reads failing at the + * same size used to halve the shared value twice, skipping a size the + * runner would have accepted and persisting it for later reads. */ + const url = 'https://concurrent.example.com'; + const window = getSandboxImageChunkBytes(url); + + const first = narrowSandboxImageChunkBytes(window, url); + const second = narrowSandboxImageChunkBytes(window, url); + + expect(second).toBe(first); + expect(getSandboxImageChunkBytes(url)).toBe(first); + }); + + it('sizes the retry from the cap a truncated response revealed', () => { + /* Blind halving from ~49KB needs eight round-trips to reach a 1KB + * runner — most of the call budget spent on discovery. The truncated + * response already measured the cap, so one retry lands. */ + const url = 'https://measured.example.com'; + const window = getSandboxImageChunkBytes(url); + + const narrowed = narrowSandboxImageChunkBytes(window, url, 1024); + + expect(narrowed).not.toBeNull(); + expect(Math.ceil((narrowed as number) / 3) * 4).toBeLessThan(1024); + }); + + it('keeps narrowing until an unconfigured small runner is reachable', () => { + /* A ~1KB stdout cap is a real runner default. A floor above it declared + * every image unreadable even though small ones fit. */ + const url = 'https://tiny-cap.example.com'; + let window: number | null = getSandboxImageChunkBytes(url); + let attempts = 0; + while (window != null && Math.ceil(window / 3) * 4 + 256 > 1024) { + attempts++; + window = narrowSandboxImageChunkBytes(window, url); + } + expect(window).not.toBeNull(); + expect(attempts).toBeLessThan(MAX_SANDBOX_IMAGE_EXEC_CALLS); + }); + + it('stops once no smaller window exists', () => { + expect(narrowSandboxImageChunkBytes(3, 'https://floor.example.com')).toBeNull(); + }); +}); + +describe('parseSandboxImageChunk', () => { + it('reports a truncated response as an overflow, not garbled output', () => { + /* The runner truncates stdout and SIGKILLs with status `OL`; parsing + * the clipped base64 would report a misleading "unexpected output" + * instead of the narrowable cause. */ + const stdout = '{"total":999999,"n":32768,"b64":"iVBORw0KGg'; + const chunk = parseSandboxImageChunk({ stdout, status: 'OL' }); + /* What survived is exactly what the runner allows, so it measures the + * cap this deployment never declared. */ + expect(chunk).toEqual({ outputOverflow: true, observedStdoutBytes: stdout.length }); + }); + + it('parses the reader JSON even when the shell emits a banner first', () => { + const chunk = parseSandboxImageChunk({ + stdout: `motd banner\n${JSON.stringify({ total: 3, n: 3, b64: 'AAAA' })}`, + }); + expect(chunk).toEqual({ total: 3, n: 3, b64: 'AAAA' }); + }); + + it('surfaces stderr when the runner produced no stdout', () => { + expect(() => parseSandboxImageChunk({ stderr: 'python3: not found', stdout: '' })).toThrow( + /python3: not found/, + ); + }); + + it('names unparseable output rather than returning junk', () => { + expect(() => parseSandboxImageChunk({ stdout: 'not json at all' })).toThrow( + /Unexpected output/, + ); + }); + + it('returns an empty chunk when the runner printed nothing', () => { + expect(parseSandboxImageChunk({ stdout: ' ' })).toEqual({}); + }); +}); + +describe('readWindowedSandboxImage', () => { + const limit = 1024 * 1024; + + it('reassembles an image larger than one window, byte-for-byte', async () => { + const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(200 * 1024)]); + const { readChunk, windows } = serveFile(source); + + const result = await readWindowedSandboxImage({ + filePath: '/mnt/data/big.png', + baseUrl: 'https://assembly.example.com', + limit, + readChunk, + }); + + expect(windows.length).toBeGreaterThan(1); + expect(result).not.toBeNull(); + expect(result).toMatchObject({ bytes: source.length }); + const base64 = (result as { base64: string }).base64; + expect(Buffer.from(base64, 'base64').equals(source)).toBe(true); + }); + + it('reads a single-window image in one round-trip', async () => { + const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(1024)]); + const { readChunk, windows } = serveFile(source); + + const result = await readWindowedSandboxImage({ + filePath: '/mnt/data/small.png', + baseUrl: 'https://single.example.com', + limit, + readChunk, + }); + + expect(windows).toHaveLength(1); + expect(Buffer.from((result as { base64: string }).base64, 'base64').equals(source)).toBe(true); + }); + + it('narrows the window and re-reads the same offset when the runner truncates', async () => { + const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(8 * 1024)]); + const acceptedWindow = 32 * 1024; + const { readChunk, windows } = serveFile(source, { acceptedWindow }); + + const result = await readWindowedSandboxImage({ + filePath: '/mnt/data/x.png', + baseUrl: 'https://truncating.example.com', + limit, + readChunk, + }); + + expect(windows[0]).toBeGreaterThan(acceptedWindow); + expect(windows[1]).toBeLessThanOrEqual(acceptedWindow); + expect(Buffer.from((result as { base64: string }).base64, 'base64').equals(source)).toBe(true); + }); + + it('reads a small image from a runner with a ~1KB stdout cap', async () => { + /* A 1KB cap is a real runner default. Discovery has to reach a window + * that small, and the truncated response says how small in one step. */ + const cap = 1024; + const source = Buffer.concat([PNG_HEADER, crypto.randomBytes(3 * 1024)]); + let calls = 0; + const readChunk = async ({ code }: { code: string }): Promise => { + calls++; + const { offset, chunk } = windowParams(code); + const slice = source.subarray(offset, offset + chunk); + const encoded = Math.ceil(slice.length / 3) * 4 + 64; + if (encoded > cap) { + return { outputOverflow: true, observedStdoutBytes: cap }; + } + return { total: source.length, n: slice.length, b64: slice.toString('base64') }; + }; + + const result = await readWindowedSandboxImage({ + filePath: '/mnt/data/icon.png', + baseUrl: 'https://kilobyte.example.com', + limit, + readChunk, + }); + + expect(Buffer.from((result as { base64: string }).base64, 'base64').equals(source)).toBe(true); + expect(calls).toBeLessThanOrEqual(MAX_SANDBOX_IMAGE_EXEC_CALLS); + }); + + it('names the stdout limit when every narrowed window still overflows', async () => { + const readChunk = async (): Promise => ({ outputOverflow: true }); + + await expect( + readWindowedSandboxImage({ + filePath: '/mnt/data/big.png', + baseUrl: 'https://always-overflows.example.com', + limit, + readChunk, + }), + ).rejects.toThrow(/exceeded the sandbox stdout limit/); + }); + + it('refuses an oversize file in-sandbox without transferring bytes', async () => { + let calls = 0; + const readChunk = async (): Promise => { + calls++; + return { too_large: true, bytes: 9 * 1024 * 1024 }; + }; + + const result = await readWindowedSandboxImage({ + filePath: '/mnt/data/huge.png', + baseUrl: 'https://oversize.example.com', + limit, + readChunk, + }); + + expect(result).toEqual({ tooLarge: true, reason: 'size', bytes: 9 * 1024 * 1024 }); + expect(calls).toBe(1); + }); + + it('gives up from the first window rather than draining the limiter', async () => { + /* The first response reveals the file size, so a read that cannot + * finish within the round-trip ceiling is known immediately — spending + * the rest of the limiter window to rediscover it would leave the turn + * with no executions left. */ + const source = crypto.randomBytes(64 * 1024); + let calls = 0; + const readChunk = async ({ code }: { code: string }): Promise => { + calls++; + const { offset } = windowParams(code); + const slice = source.subarray(offset, offset + 300); + return { total: source.length, n: slice.length, b64: slice.toString('base64') }; + }; + + const result = await readWindowedSandboxImage({ + filePath: '/mnt/data/slow.png', + baseUrl: 'https://ceiling.example.com', + limit, + readChunk, + }); + + expect(calls).toBe(1); + expect(result).toEqual({ + tooLarge: true, + reason: 'round_trips', + bytes: source.length, + /* What this deployment could actually deliver: 300 bytes a call for + * the whole round-trip budget. The caller names it as a downscale + * target instead of leaving the model to guess. */ + inlineCeiling: 300 * MAX_SANDBOX_IMAGE_EXEC_CALLS, + }); + }); + + it('reads a file that exactly fills the round-trip ceiling', async () => { + const window = getSandboxImageChunkBytes('https://exact.example.com'); + const source = crypto.randomBytes(window * MAX_SANDBOX_IMAGE_EXEC_CALLS); + const { readChunk, windows } = serveFile(source); + + const result = await readWindowedSandboxImage({ + filePath: '/mnt/data/exact.png', + baseUrl: 'https://exact.example.com', + limit: source.length, + readChunk, + }); + + expect(windows).toHaveLength(MAX_SANDBOX_IMAGE_EXEC_CALLS); + expect(Buffer.from((result as { base64: string }).base64, 'base64').equals(source)).toBe(true); + }); + + it('refuses to splice a file that changed mid-read', async () => { + let call = 0; + const readChunk = async ({ code }: { code: string }): Promise => { + const { chunk } = windowParams(code); + call++; + /* Second window reports a different size: the assembled buffer would + * be a mix of two versions rather than any real image. */ + return { + total: call === 1 ? chunk * 3 : chunk * 4, + n: chunk, + b64: Buffer.alloc(chunk).toString('base64'), + }; + }; + + await expect( + readWindowedSandboxImage({ + filePath: '/mnt/data/moving.png', + baseUrl: 'https://changing.example.com', + limit, + readChunk, + }), + ).rejects.toThrow(/changed while being read/); + }); + + it('surfaces the in-sandbox reader error for a path it refused', async () => { + const readChunk = async (): Promise => ({ + error: "[Errno 2] No such file or directory: '/mnt/data/gone.png'", + }); + + await expect( + readWindowedSandboxImage({ + filePath: '/mnt/data/gone.png', + baseUrl: 'https://missing.example.com', + limit, + readChunk, + }), + ).rejects.toThrow(/No such file or directory/); + }); + + it('returns null when the runner produced no output', async () => { + const readChunk = async (): Promise => ({}); + + await expect( + readWindowedSandboxImage({ + filePath: '/mnt/data/quiet.png', + baseUrl: 'https://silent.example.com', + limit, + readChunk, + }), + ).resolves.toBeNull(); + }); +}); + +describe('buildSandboxImageReaderCode', () => { + it('carries the window parameters base64-encoded, never as shell syntax', () => { + const code = buildSandboxImageReaderCode({ + filePath: `/mnt/data/'; rm -rf /; '.png`, + limit: 1024, + offset: 96, + chunkBytes: 48, + }); + + expect(code).not.toContain('rm -rf'); + expect(windowParams(code)).toEqual({ + file_path: `/mnt/data/'; rm -rf /; '.png`, + limit: 1024, + offset: 96, + chunk: 48, + }); + }); +}); diff --git a/packages/api/src/files/code/image.ts b/packages/api/src/files/code/image.ts new file mode 100644 index 0000000000..7fc53ede01 --- /dev/null +++ b/packages/api/src/files/code/image.ts @@ -0,0 +1,366 @@ +import { logger } from '@librechat/data-schemas'; + +/** + * Windowed base64 reader for images living in a code-execution sandbox. + * + * The bytes can only come back through `/exec` stdout, which the runner + * truncates and SIGKILLs past its own `SANDBOX_OUTPUT_MAX_SIZE`. One read + * is therefore a series of `/exec` calls, each spending a request against + * the Code API's per-user execution limiter (20 per 30s in the reference + * deployment). Everything here exists to keep that count down and bounded: + * the window fills the runner's stdout budget, narrows itself if a runner + * rejects it, and a file that still needs more calls than one read may + * spend is reported as unreadable-inline rather than stalling a chat turn + * across limiter windows. + * + * The transport stays with the caller ({@link SandboxImageChunkReader}) so + * this module owns only the sizing, windowing, and assembly decisions. + */ + +/** + * Sandbox stdout a single `/exec` response may carry, in bytes — 64KB in + * the reference deployment. Deployments that changed the runner's cap set + * `LIBRECHAT_CODE_SANDBOX_OUTPUT_MAX_SIZE` to match; a smaller cap is also + * discovered at runtime (see {@link narrowSandboxImageChunkBytes}). + */ +const DEFAULT_SANDBOX_OUTPUT_MAX_SIZE = 64 * 1024; + +/** JSON envelope around the base64 window (`{"total":…,"n":…,"b64":"…"}`), + * plus room for a trailing newline and a stray shell banner line. */ +const IMAGE_CHUNK_ENVELOPE_BYTES = 256; + +/** Base64 encodes in 3-byte groups; windowing on a multiple of 3 keeps + * every response padding-free. */ +const BASE64_GROUP_BYTES = 3; + +/** + * Hard round-trip ceiling for one image read, matched to a single window of + * the Code API's per-user execution limiter (20 requests per 30s in the + * reference deployment). Beyond it a read would have to wait out limiter + * windows to finish, stalling a chat turn; a file that needs more windows + * than this degrades to the same "too large to inline" result as one over + * the byte cap — and {@link readWindowedSandboxImage} detects that from the + * first window rather than spending the whole budget discovering it. + */ +export const MAX_SANDBOX_IMAGE_EXEC_CALLS = 20; + +/** Per-base-URL window size, narrowed in place the first time a runner + * reports truncation so later reads start at a size it accepts. */ +const learnedChunkBytes = new Map(); + +/** + * One `/exec` window's outcome, as the transport observed it. Fields are + * optional because the body is whatever the in-sandbox reader printed: + * `error` when it refused the path, `too_large` when the file is over the + * caller's cap, `total`/`n`/`b64` for a window of bytes, `outputOverflow` + * when the runner truncated the response, and nothing at all when it + * produced no output. + */ +export interface SandboxImageChunk { + outputOverflow?: boolean; + /** Bytes of stdout the runner did emit before truncating — its actual + * cap, and a far better next window than another blind halving. */ + observedStdoutBytes?: number; + error?: string; + too_large?: boolean; + bytes?: number; + total?: number; + n?: number; + b64?: string; +} + +export type SandboxImageChunkReader = (params: { code: string }) => Promise; + +export type SandboxImageReadResult = + | { base64: string; bytes: number } + /** `size`: over the caller's cap. `round_trips`: within the byte cap, but + * more windows than {@link MAX_SANDBOX_IMAGE_EXEC_CALLS} allows — + * `inlineCeiling` is the largest file this deployment's window size can + * actually deliver, which the caller can name as a downscale target. */ + | { tooLarge: true; reason: 'size' | 'round_trips'; bytes: number; inlineCeiling?: number } + | null; + +/** + * Raw bytes to pull per `/exec` round-trip. Each window is base64-encoded + * (~1.33x) into the response's stdout, so the auto-derived size is the + * largest multiple of 3 whose encoding plus envelope fits the runner's + * stdout budget. `LIBRECHAT_CODE_IMAGE_CHUNK_BYTES` overrides it outright + * and is used verbatim, multiple of 3 or not. + */ +export function getSandboxImageChunkBytes(baseUrl?: string): number { + const baseline = baselineChunkBytes(); + const learned = baseUrl == null ? undefined : learnedChunkBytes.get(baseUrl); + /* A learned narrowing only ever caps the configured size: an operator who + * sets a smaller window explicitly still wins. */ + return learned == null ? baseline : Math.min(learned, baseline); +} + +function baselineChunkBytes(): number { + const override = Number(process.env.LIBRECHAT_CODE_IMAGE_CHUNK_BYTES); + if (Number.isFinite(override) && override > 0) { + return Math.floor(override); + } + const configured = Number(process.env.LIBRECHAT_CODE_SANDBOX_OUTPUT_MAX_SIZE); + return chunkBytesForBudget( + Number.isFinite(configured) && configured > 0 + ? Math.floor(configured) + : DEFAULT_SANDBOX_OUTPUT_MAX_SIZE, + ); +} + +/** Largest window whose base64 encoding plus envelope fits `budget`. */ +function chunkBytesForBudget(budget: number): number { + const usable = Math.max(budget - IMAGE_CHUNK_ENVELOPE_BYTES, 0); + return alignToBase64Group(Math.floor((usable * 3) / 4)); +} + +function alignToBase64Group(bytes: number): number { + return Math.max(Math.floor(bytes / BASE64_GROUP_BYTES) * BASE64_GROUP_BYTES, BASE64_GROUP_BYTES); +} + +/** + * Returns the window to retry a truncated read with, or `null` when no + * smaller one exists. + * + * A truncated response reveals the runner's real cap — it emitted exactly + * as much as it allows — so `observedStdoutBytes` sizes the retry directly + * and one round-trip is usually enough to land on a size that fits. Blind + * halving is the fallback, and bounds the result either way so every retry + * is strictly smaller than the window that just failed. + * + * The reduction is computed from the caller's own failing window, never + * from the shared learned value: reads run concurrently (one per tool + * call), so two reads failing at the same size would otherwise reduce the + * shared value twice and skip a size the runner would have accepted. For + * the same reason the learned value only ever moves down to the narrowest + * size any read has needed. + */ +export function narrowSandboxImageChunkBytes( + failedChunkBytes: number, + baseUrl?: string, + observedStdoutBytes?: number, +): number | null { + const halved = alignToBase64Group(Math.floor(failedChunkBytes / 2)); + const fromObserved = + observedStdoutBytes != null && observedStdoutBytes > 0 + ? chunkBytesForBudget(observedStdoutBytes) + : Number.POSITIVE_INFINITY; + const narrowed = Math.min(halved, fromObserved); + /* Already as small as a base64 group: this runner cannot emit an image + * window at all, and another attempt would only be truncated again. */ + if (narrowed >= failedChunkBytes) { + return null; + } + if (baseUrl != null) { + const learned = learnedChunkBytes.get(baseUrl); + learnedChunkBytes.set(baseUrl, learned == null ? narrowed : Math.min(learned, narrowed)); + } + logger.warn( + `[readSandboxImage] Sandbox stdout limit exceeded at ${failedChunkBytes} bytes; retrying with ${narrowed} for ${baseUrl}. ` + + "Set LIBRECHAT_CODE_SANDBOX_OUTPUT_MAX_SIZE to this runner's stdout cap to skip the discovery reads.", + ); + return narrowed; +} + +/** + * The in-sandbox reader for one window. Stats the file, refuses (without + * transferring) anything over `limit` or any non-regular file, and prints + * a single JSON line the transport hands back as a {@link SandboxImageChunk}. + * The parameters travel base64-encoded so neither the path nor the numbers + * are interpolated into shell syntax. + */ +export function buildSandboxImageReaderCode(params: { + filePath: string; + limit: number; + offset: number; + chunkBytes: number; +}): string { + const payload = Buffer.from( + JSON.stringify({ + file_path: params.filePath, + limit: params.limit, + offset: params.offset, + chunk: params.chunkBytes, + }), + 'utf8', + ).toString('base64'); + return [ + "python3 - <<'PY'", + 'import base64, json, os, stat', + `payload = ${JSON.stringify(payload)}`, + "data = json.loads(base64.b64decode(payload).decode('utf-8'))", + "p = data['file_path']", + "limit = data['limit']", + "offset = data['offset']", + "chunk = data['chunk']", + 'try:', + ' st = os.stat(p)', + 'except OSError as e:', + ' print(json.dumps({"error": str(e)}))', + ' raise SystemExit(0)', + // Reject FIFOs, sockets, and device files (e.g. a symlink to /dev/zero): + // os.stat can report a small/zero size while an unbounded read blocks or + // streams forever until the request times out. + 'if not stat.S_ISREG(st.st_mode):', + ' print(json.dumps({"error": "not a regular file"}))', + ' raise SystemExit(0)', + 'if st.st_size > limit:', + ' print(json.dumps({"too_large": True, "bytes": st.st_size}))', + ' raise SystemExit(0)', + // Read only this window. The whole base64 payload cannot be emitted in + // one shot: the runner caps stdout at SANDBOX_OUTPUT_MAX_SIZE and + // SIGKILLs the job on overflow, which truncates the JSON mid-string. + "with open(p, 'rb') as f:", + ' f.seek(offset)', + ' raw = f.read(chunk)', + 'print(json.dumps({"total": st.st_size, "n": len(raw), "b64": base64.b64encode(raw).decode("ascii")}))', + 'PY', + ].join('\n'); +} + +/** + * Turns one `/exec` response into a {@link SandboxImageChunk}. Throws only + * for output the reader could not have produced, so the caller can tell a + * broken transport from a file the sandbox declined to serve. + */ +export function parseSandboxImageChunk(response: { + stdout?: unknown; + stderr?: unknown; + status?: unknown; +}): SandboxImageChunk { + /* The runner truncates stdout at SANDBOX_OUTPUT_MAX_SIZE and SIGKILLs the + * job (status `OL`). Detect that explicitly: the surviving stdout is a + * base64 string cut mid-flight, so parsing it yields a misleading + * "unexpected output" instead of the narrowable, fixable cause. */ + const stdout = response.stdout == null ? '' : String(response.stdout); + if (response.status === 'OL') { + /* What survived is exactly what the runner allows, so it doubles as a + * measurement of the cap this deployment never declared. */ + return { outputOverflow: true, observedStdoutBytes: stdout.length }; + } + if (response.stderr && stdout === '') { + throw new Error(String(response.stderr).trim()); + } + if (stdout.trim() === '') { + return {}; + } + /* Parse the LAST non-empty line: the reader's JSON is the final thing it + * prints, so anything a shell profile or library emitted ahead of it + * (banners, warnings) must not break the read. */ + const lines = stdout + .split('\n') + .map((line) => line.trim()) + .filter(Boolean); + try { + return JSON.parse(lines[lines.length - 1]); + } catch { + throw new Error( + `Unexpected output while reading image bytes from the sandbox: ${stdout.slice(0, 120)}`, + ); + } +} + +/** + * Pulls a sandbox image through as many windows as it takes, narrowing the + * window if the runner truncates and stopping at the round-trip ceiling. + * Returns `null` when the runner produced no output at all; throws what + * the transport throws. + */ +export async function readWindowedSandboxImage(params: { + filePath: string; + limit: number; + baseUrl?: string; + readChunk: SandboxImageChunkReader; +}): Promise { + const { filePath, limit, baseUrl, readChunk } = params; + let chunkBytes = getSandboxImageChunkBytes(baseUrl); + let sawOverflow = false; + const parts: Buffer[] = []; + let offset = 0; + let total: number | null = null; + + for (let call = 0; call < MAX_SANDBOX_IMAGE_EXEC_CALLS; call++) { + const chunk = await readChunk({ + code: buildSandboxImageReaderCode({ filePath, limit, offset, chunkBytes }), + }); + + if (chunk.outputOverflow === true) { + /* The runner's stdout cap is smaller than this deployment assumed. + * Halve the window and re-read the same offset rather than failing + * the whole image; keep halving until the runner accepts a size or + * the floor says it never will. */ + sawOverflow = true; + const narrowed = narrowSandboxImageChunkBytes(chunkBytes, baseUrl, chunk.observedStdoutBytes); + if (narrowed == null) { + throw new Error( + `Reading "${filePath}" exceeded the sandbox stdout limit (window ${chunkBytes} bytes).`, + ); + } + chunkBytes = narrowed; + continue; + } + if (chunk.error) { + throw new Error(String(chunk.error)); + } + if (chunk.too_large === true) { + return { tooLarge: true, reason: 'size', bytes: Number(chunk.bytes) || 0 }; + } + if (typeof chunk.b64 !== 'string' || typeof chunk.n !== 'number') { + return null; + } + + if (total == null) { + total = Number(chunk.total) || 0; + if (total > limit) { + return { tooLarge: true, reason: 'size', bytes: total }; + } + } else if (Number(chunk.total) !== total) { + /* The file changed underneath us; a spliced-together buffer would be + * a mix of two versions rather than any real image. */ + throw new Error(`"${filePath}" changed while being read from the sandbox`); + } + + parts.push(Buffer.from(chunk.b64, 'base64')); + offset += chunk.n; + + if (chunk.n === 0 || offset >= total) { + break; + } + /* The first window reveals both the file's size and what a call really + * delivers, so a read that cannot finish is known now rather than after + * draining the limiter. Project from the bytes actually returned, not + * the window requested: a runner that serves short reads would other- + * wise look like it was keeping up. */ + if (Math.ceil((total - offset) / chunk.n) > MAX_SANDBOX_IMAGE_EXEC_CALLS - call - 1) { + return { + tooLarge: true, + reason: 'round_trips', + bytes: total, + inlineCeiling: chunk.n * MAX_SANDBOX_IMAGE_EXEC_CALLS, + }; + } + } + + if (total == null) { + /* Every call was spent narrowing: the runner never accepted a window. */ + if (sawOverflow) { + throw new Error( + `Reading "${filePath}" exceeded the sandbox stdout limit (window ${chunkBytes} bytes).`, + ); + } + return null; + } + const buffer = Buffer.concat(parts); + if (buffer.length !== total) { + /* Short reads: returning a partial image would render as a corrupt + * file, so surface it as unreadable-inline with the reason the caller + * should report. */ + return { + tooLarge: true, + reason: 'round_trips', + bytes: total, + inlineCeiling: buffer.length, + }; + } + return { base64: buffer.toString('base64'), bytes: buffer.length }; +} diff --git a/packages/api/src/files/code/index.ts b/packages/api/src/files/code/index.ts index d24bd420a4..6414ff2564 100644 --- a/packages/api/src/files/code/index.ts +++ b/packages/api/src/files/code/index.ts @@ -3,4 +3,5 @@ export * from './destinations'; export * from './extract'; export * from './form'; export * from './identity'; +export * from './image'; export * from './preflight'; diff --git a/packages/api/src/utils/code.spec.ts b/packages/api/src/utils/code.spec.ts new file mode 100644 index 0000000000..be4d25bfa3 --- /dev/null +++ b/packages/api/src/utils/code.spec.ts @@ -0,0 +1,150 @@ +import { AxiosError, AxiosHeaders } from 'axios'; +import { getCodeApiRetryAfterMs, withCodeApiRateLimit, createCodeApiRateLimitBudget } from './code'; + +/** Builds the error axios raises for a Code API rate-limit response. */ +function rateLimited({ + status = 429, + headers = {}, + data = {}, +}: { + status?: number; + headers?: Record; + data?: unknown; +} = {}): AxiosError { + const error = new AxiosError('Request failed', 'ERR_BAD_REQUEST'); + error.response = { + status, + statusText: '', + headers: new AxiosHeaders(headers), + config: { headers: new AxiosHeaders() }, + data, + }; + return error; +} + +describe('getCodeApiRetryAfterMs', () => { + it('reads the standard Retry-After header', () => { + expect(getCodeApiRetryAfterMs(rateLimited({ headers: { 'retry-after': '17' } }))).toBe(17_000); + }); + + it('falls back to the response body when a proxy dropped the header', () => { + expect(getCodeApiRetryAfterMs(rateLimited({ data: { retry_after_seconds: 4 } }))).toBe(4_000); + }); + + it('prefers the header over the body', () => { + const error = rateLimited({ + headers: { 'retry-after': '2' }, + data: { retry_after_seconds: 30 }, + }); + expect(getCodeApiRetryAfterMs(error)).toBe(2_000); + }); + + it('returns null for a non-429 response', () => { + expect(getCodeApiRetryAfterMs(rateLimited({ status: 503 }))).toBeNull(); + }); + + it('returns null when neither channel carries a usable delay', () => { + expect(getCodeApiRetryAfterMs(rateLimited({ headers: { 'retry-after': 'soon' } }))).toBeNull(); + expect(getCodeApiRetryAfterMs(rateLimited({ data: 'rate_limited' }))).toBeNull(); + }); + + it('treats a blank header as absent rather than "retry immediately"', () => { + /* `Number('')` is 0, which would both skip the body fallback and read + * as a zero-second wait. */ + const error = rateLimited({ headers: { 'retry-after': '' }, data: { retry_after_seconds: 9 } }); + expect(getCodeApiRetryAfterMs(error)).toBe(9_000); + }); + + it('returns null for anything that is not an axios error', () => { + expect(getCodeApiRetryAfterMs(new Error('boom'))).toBeNull(); + expect(getCodeApiRetryAfterMs(undefined)).toBeNull(); + }); +}); + +describe('withCodeApiRateLimit', () => { + it('passes a successful attempt straight through', async () => { + const attempt = jest.fn(async () => 'ok'); + await expect(withCodeApiRateLimit({ attempt, label: 'reading' })).resolves.toBe('ok'); + expect(attempt).toHaveBeenCalledTimes(1); + }); + + it('waits out a rate limit the budget can absorb, then finishes', async () => { + /* A multi-request operation would otherwise abandon the work already + * done the moment a limiter window closes mid-flight. */ + const budget = createCodeApiRateLimitBudget(5_000); + const waits: number[] = []; + let attempts = 0; + const attempt = jest.fn(async () => { + attempts++; + if (attempts === 1) { + throw rateLimited({ headers: { 'retry-after': '1' } }); + } + return 'ok'; + }); + + await expect( + withCodeApiRateLimit({ attempt, label: 'reading', budget, onWait: (ms) => waits.push(ms) }), + ).resolves.toBe('ok'); + + expect(attempts).toBe(2); + expect(waits).toEqual([1_000]); + expect(budget.remainingMs).toBe(4_000); + }); + + it('names the limit when the wait exceeds the remaining budget', async () => { + const budget = createCodeApiRateLimitBudget(5_000); + const attempt = jest.fn(async () => { + throw rateLimited({ headers: { 'retry-after': '300' } }); + }); + + await expect( + withCodeApiRateLimit({ attempt, label: 'reading "x.png"', budget }), + ).rejects.toThrow(/rate limit reached while reading "x\.png" \(retry in 300s\)/); + expect(attempt).toHaveBeenCalledTimes(1); + expect(budget.remainingMs).toBe(5_000); + }); + + it('names the limit even when the response carries no usable delay', async () => { + /* Otherwise a bare "Request failed with status code 429" surfaces as an + * unspecified failure instead of one that clears on its own. */ + const attempt = jest.fn(async () => { + throw rateLimited(); + }); + + await expect( + withCodeApiRateLimit({ + attempt, + label: 'reading', + budget: createCodeApiRateLimitBudget(60_000), + }), + ).rejects.toThrow(/rate limit reached while reading\./); + }); + + it('charges at least a second so a zero delay cannot spin', async () => { + const budget = createCodeApiRateLimitBudget(2_500); + const waits: number[] = []; + let attempts = 0; + const attempt = jest.fn(async () => { + attempts++; + if (attempts < 3) { + throw rateLimited({ headers: { 'retry-after': '0' } }); + } + return 'ok'; + }); + + await expect( + withCodeApiRateLimit({ attempt, label: 'reading', budget, onWait: (ms) => waits.push(ms) }), + ).resolves.toBe('ok'); + expect(waits).toEqual([1_000, 1_000]); + expect(budget.remainingMs).toBe(500); + }); + + it('rethrows anything that is not a rate limit', async () => { + const attempt = jest.fn(async () => { + throw rateLimited({ status: 500 }); + }); + await expect(withCodeApiRateLimit({ attempt, label: 'reading' })).rejects.toThrow( + 'Request failed', + ); + }); +}); diff --git a/packages/api/src/utils/code.ts b/packages/api/src/utils/code.ts index 4293ef7f05..42931b5bca 100644 --- a/packages/api/src/utils/code.ts +++ b/packages/api/src/utils/code.ts @@ -1,5 +1,6 @@ import http from 'http'; import https from 'https'; +import { isAxiosError } from 'axios'; /** * Dedicated agents for code-server requests, preventing socket pool contamination. @@ -9,3 +10,89 @@ import https from 'https'; */ export const codeServerHttpAgent: http.Agent = new http.Agent({ keepAlive: false }); export const codeServerHttpsAgent: https.Agent = new https.Agent({ keepAlive: false }); + +/** + * Wait implied by a Code API 429, in milliseconds, or `null` when the error + * is not a rate-limit response. The service answers with both the standard + * `Retry-After` header and a `retry_after_seconds` body field; the header is + * preferred so a body the caller asked for as a stream/buffer never has to be + * parsed, and the body is the fallback for proxies that drop the header. + */ +export function getCodeApiRetryAfterMs(error: unknown): number | null { + if (!isAxiosError(error) || error.response?.status !== 429) { + return null; + } + const header = error.response.headers?.['retry-after']; + const raw = Array.isArray(header) ? header[0] : header; + /* `Number('')` is 0, which would read as "retry immediately" and skip the + * body fallback; treat a blank header as absent. */ + const headerSeconds = typeof raw === 'string' && raw.trim() === '' ? NaN : Number(raw); + if (Number.isFinite(headerSeconds) && headerSeconds >= 0) { + return headerSeconds * 1000; + } + const body = error.response.data; + if (typeof body === 'object' && body !== null && 'retry_after_seconds' in body) { + const bodySeconds = Number(body.retry_after_seconds); + if (Number.isFinite(bodySeconds) && bodySeconds >= 0) { + return bodySeconds * 1000; + } + } + return null; +} + +/** Total time one operation may spend waiting out Code API rate limits. + * Sized to cover a single limiter window without stalling a chat turn. */ +export const MAX_CODE_API_RATE_LIMIT_WAIT_MS = 20_000; + +/** Remaining wait allowance, shared by every request of one operation. */ +export interface CodeApiRateLimitBudget { + remainingMs: number; +} + +export function createCodeApiRateLimitBudget( + totalMs: number = MAX_CODE_API_RATE_LIMIT_WAIT_MS, +): CodeApiRateLimitBudget { + return { remainingMs: totalMs }; +} + +/** + * Runs a Code API request, waiting out a rate limit whenever the shared + * budget still covers the server's `Retry-After`. Multi-request operations + * (a windowed image read, a batch upload) would otherwise abandon the work + * already done the moment a limiter window closes mid-flight. A 429 the + * budget cannot absorb throws a named error rather than a bare + * "status code 429", which reads as an unspecified failure instead of one + * that clears on its own. + */ +export async function withCodeApiRateLimit(params: { + attempt: () => Promise; + label: string; + budget?: CodeApiRateLimitBudget; + onWait?: (waitMs: number) => void; +}): Promise { + const { attempt, label, budget, onWait } = params; + for (;;) { + try { + return await attempt(); + } catch (error) { + if (!isAxiosError(error) || error.response?.status !== 429) { + throw error; + } + const retryAfterMs = getCodeApiRetryAfterMs(error); + if (retryAfterMs == null) { + throw new Error(`Code API rate limit reached while ${label}.`); + } + /* A zero delay would spin; charge at least a second so the budget + * always drains and the loop terminates. */ + const waitMs = Math.max(retryAfterMs, 1000); + if (budget == null || waitMs > budget.remainingMs) { + throw new Error( + `Code API rate limit reached while ${label} (retry in ${Math.ceil(waitMs / 1000)}s).`, + ); + } + budget.remainingMs -= waitMs; + onWait?.(waitMs); + await new Promise((resolve) => setTimeout(resolve, waitMs)); + } + } +}