diff --git a/src/lib/inference/local-adapter-lifecycle.test.ts b/src/lib/inference/local-adapter-lifecycle.test.ts index 64f36c0ab8c..a45e8321bf5 100644 --- a/src/lib/inference/local-adapter-lifecycle.test.ts +++ b/src/lib/inference/local-adapter-lifecycle.test.ts @@ -12,6 +12,7 @@ import { ensureLocalAdapterStateDir, isLocalAdapterProcess, killLocalAdapterPid, + LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES, loadLocalAdapterPid, localAdapterTokenHash, persistLocalAdapterPid, @@ -155,6 +156,113 @@ describe("local adapter lifecycle", () => { }), ).resolves.toBe(false); }); + + it("fails closed when a chunked health response exceeds the memory budget", async () => { + const expectedTokenHash = localAdapterTokenHash("secret-token"); + const payload = Buffer.from( + JSON.stringify({ + tokenHash: expectedTokenHash, + padding: " ".repeat(LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES), + }), + ); + const server = http.createServer((_req, res) => { + res.writeHead(200, { "Content-Type": "application/json" }); + res.write(payload.subarray(0, LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES)); + res.end(payload.subarray(LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES)); + }); + const port = await listen(server); + + await expect( + probeLocalAdapterHealth({ + host: "127.0.0.1", + port, + expectedTokenHash, + }), + ).resolves.toBe(false); + }); + + it("accepts a valid health response at the memory budget", async () => { + const expectedTokenHash = localAdapterTokenHash("secret-token"); + const emptyPayload = JSON.stringify({ tokenHash: expectedTokenHash, padding: "" }); + const payload = Buffer.from( + JSON.stringify({ + tokenHash: expectedTokenHash, + padding: " ".repeat( + LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES - Buffer.byteLength(emptyPayload), + ), + }), + ); + expect(payload).toHaveLength(LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES); + const server = http.createServer((_req, res) => { + res.writeHead(200, { + "Content-Length": String(payload.length), + "Content-Type": "application/json", + }); + res.end(payload); + }); + const port = await listen(server); + + await expect( + probeLocalAdapterHealth({ + host: "127.0.0.1", + port, + expectedTokenHash, + }), + ).resolves.toBe(true); + }); + + it("fails closed when a health response closes before completion", async () => { + const expectedTokenHash = localAdapterTokenHash("secret-token"); + const server = http.createServer((_req, res) => { + res.writeHead(200, { + "Content-Length": "128", + "Content-Type": "application/json", + }); + res.write('{"tokenHash":"'); + res.socket?.destroy(); + }); + const port = await listen(server); + + await expect( + probeLocalAdapterHealth({ + host: "127.0.0.1", + port, + expectedTokenHash, + }), + ).resolves.toBe(false); + }); + + it("destroys a declared health response above the memory budget before buffering", async () => { + const expectedTokenHash = localAdapterTokenHash("secret-token"); + const declaredBytes = LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES + 1; + const destroySpy = vi.spyOn(http.IncomingMessage.prototype, "destroy"); + const server = http.createServer((_req, res) => { + res.writeHead(200, { + "Content-Length": String(declaredBytes), + "Content-Type": "application/json", + }); + res.end(); + }); + const port = await listen(server); + + await expect( + probeLocalAdapterHealth({ + host: "127.0.0.1", + port, + expectedTokenHash, + }), + ).resolves.toBe(false); + expect( + destroySpy.mock.calls.some((args, index) => { + const response = destroySpy.mock.contexts[index] as http.IncomingMessage; + return ( + response.statusCode === 200 && + response.headers["content-length"] === String(declaredBytes) && + args.length === 0 + ); + }), + ).toBe(true); + }); }); describe("ensureLocalAdapterStateDir", () => { diff --git a/src/lib/inference/local-adapter-lifecycle.ts b/src/lib/inference/local-adapter-lifecycle.ts index 65a8ed9891f..2565dc56fd5 100644 --- a/src/lib/inference/local-adapter-lifecycle.ts +++ b/src/lib/inference/local-adapter-lifecycle.ts @@ -28,6 +28,7 @@ export type RunFn = ( export type LocalAdapterProcessMatcher = string | RegExp | ((commandLine: string) => boolean); export const DEFAULT_LOCAL_ADAPTER_STATE_DIR = nemoclawStateRoot(os.homedir(), GATEWAY_PORT); +export const LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES = 64 * 1024; export function ensureLocalAdapterStateDir(stateDir = DEFAULT_LOCAL_ADAPTER_STATE_DIR): void { rejectSymlinksOnPath(stateDir); @@ -191,6 +192,12 @@ export function probeLocalAdapterHealth(options: { tokenHashField?: string; }): Promise { return new Promise((resolve) => { + let settled = false; + const settle = (healthy: boolean) => { + if (settled) return; + settled = true; + resolve(healthy); + }; const req = http.request( { hostname: options.host, @@ -200,31 +207,59 @@ export function probeLocalAdapterHealth(options: { timeout: options.timeoutMs || 1000, }, (res) => { + const declaredBytes = Number(res.headers["content-length"]); + if ( + Number.isFinite(declaredBytes) && + declaredBytes > LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES + ) { + res.destroy(); + settle(false); + return; + } + const chunks: Buffer[] = []; - res.on("data", (chunk) => chunks.push(Buffer.from(chunk))); + let receivedBytes = 0; + res.on("data", (chunk) => { + if (settled) return; + const chunkBytes = Buffer.isBuffer(chunk) ? chunk.length : Buffer.byteLength(chunk); + if (receivedBytes + chunkBytes > LOCAL_ADAPTER_HEALTH_MAX_RESPONSE_BYTES) { + chunks.length = 0; + res.destroy(); + settle(false); + return; + } + receivedBytes += chunkBytes; + const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); + chunks.push(buffer); + }); + res.on("close", () => { + if (!res.complete) settle(false); + }); + res.on("error", () => settle(false)); res.on("end", () => { + if (settled) return; if (res.statusCode !== 200) { - resolve(false); + settle(false); return; } if (!options.expectedTokenHash) { - resolve(true); + settle(true); return; } try { const body = JSON.parse(Buffer.concat(chunks).toString("utf8")) as JsonObject; - resolve(body[options.tokenHashField || "tokenHash"] === options.expectedTokenHash); + settle(body[options.tokenHashField || "tokenHash"] === options.expectedTokenHash); } catch { - resolve(false); + settle(false); } }); }, ); req.on("timeout", () => { req.destroy(); - resolve(false); + settle(false); }); - req.on("error", () => resolve(false)); + req.on("error", () => settle(false)); req.end(); }); }