| 1 | // A single JSON-RPC message larger than the host's stdout budget kills the |
| 2 | // stdio transport and takes every tool with it (Claude Code disconnects at |
| 3 | // 16MB; a full-screen 5K PNG base64s to ~29MB). These tests drive the real |
| 4 | // server over stdio and assert that an over-budget raster degrades to its |
| 5 | // text receipt while the connection keeps serving. |
| 6 | import { test } from "node:test"; |
| 7 | import assert from "node:assert/strict"; |
| 8 | import { spawn } from "node:child_process"; |
| 9 | import fs from "node:fs"; |
| 10 | import os from "node:os"; |
| 11 | import path from "node:path"; |
| 12 | import url from "node:url"; |
| 13 | |
| 14 | const __dirname = path.dirname(url.fileURLToPath(import.meta.url)); |
| 15 | const ROOT = path.resolve(__dirname, ".."); |
| 16 | |
| 17 | async function session(t, extraEnv = {}) { |
| 18 | const stateDir = fs.mkdtempSync(path.join(os.tmpdir(), "cu-budget-state-")); |
| 19 | const recDir = fs.mkdtempSync(path.join(os.tmpdir(), "cu-budget-rec-")); |
| 20 | const server = spawn(process.execPath, [path.join(ROOT, "mcp", "server.mjs")], { |
| 21 | env: { |
| 22 | ...process.env, |
| 23 | CODEWHALE_CU_APP: "off", |
| 24 | CODEWHALE_CU_APP_WARM: "off", |
| 25 | CODEWHALE_CU_STATE_DIR: stateDir, |
| 26 | CODEWHALE_CU_RECORDINGS_DIR: recDir, |
| 27 | CODEWHALE_CU_TEST_BACKEND: path.join(__dirname, "fixtures", "fake-backend.mjs"), |
| 28 | FAKE_BACKEND_CALLS: path.join(stateDir, "calls.jsonl"), |
| 29 | FAKE_BACKEND_CONTROL: path.join(stateDir, "control.json"), |
| 30 | ...extraEnv, |
| 31 | }, |
| 32 | stdio: ["pipe", "pipe", "pipe"], |
| 33 | }); |
| 34 | t.after(() => { |
| 35 | server.kill("SIGTERM"); |
| 36 | for (const dir of [stateDir, recDir]) { try { fs.rmSync(dir, { recursive: true, force: true }); } catch {} } |
| 37 | }); |
| 38 | |
| 39 | const pending = new Map(); |
| 40 | let nextId = 1; |
| 41 | let buf = ""; |
| 42 | server.stdout.setEncoding("utf8"); |
| 43 | server.stdout.on("data", (chunk) => { |
| 44 | buf += chunk; |
| 45 | let i; |
| 46 | while ((i = buf.indexOf("\n")) !== -1) { |
| 47 | const line = buf.slice(0, i).trim(); |
| 48 | buf = buf.slice(i + 1); |
| 49 | if (!line) continue; |
| 50 | try { |
| 51 | const msg = JSON.parse(line); |
| 52 | if (msg.id && pending.has(msg.id)) { pending.get(msg.id)(msg); pending.delete(msg.id); } |
| 53 | } catch {} |
| 54 | } |
| 55 | }); |
| 56 | |
| 57 | const rpc = (method, params) => new Promise((resolve, reject) => { |
| 58 | const id = nextId++; |
| 59 | const timer = setTimeout(() => { pending.delete(id); reject(new Error(`timeout: ${method}`)); }, 30_000); |
| 60 | pending.set(id, (msg) => { clearTimeout(timer); resolve(msg); }); |
| 61 | server.stdin.write(JSON.stringify({ jsonrpc: "2.0", id, method, params }) + "\n"); |
| 62 | }); |
| 63 | |
| 64 | const init = await rpc("initialize", { protocolVersion: "2025-06-18" }); |
| 65 | assert.equal(init.result.serverInfo.name, "codewhale-cu"); |
| 66 | |
| 67 | return { |
| 68 | server, |
| 69 | async call(name, args = {}) { |
| 70 | const res = await rpc("tools/call", { name, arguments: args }); |
| 71 | assert.ok(res.result, `${name}: protocol error ${JSON.stringify(res.error ?? {})}`); |
| 72 | return { content: res.result.content, receipt: JSON.parse(res.result.content[0].text) }; |
| 73 | }, |
| 74 | }; |
| 75 | } |
| 76 | |
| 77 | test("an over-budget raster degrades to its receipt and keeps the transport alive", async (t) => { |
| 78 | // The fixture's 1x1 PNG base64s to ~92 bytes, so a 64-byte budget puts any |
| 79 | // capture over the line without writing a huge file. |
| 80 | const cu = await session(t, { CODEWHALE_CU_MAX_IMAGE_BYTES: "64" }); |
| 81 | |
| 82 | const shot = await cu.call("screenshot"); |
| 83 | assert.equal(shot.receipt.ok, true, "an over-budget capture still succeeds"); |
| 84 | assert.equal(shot.content.length, 1, "the image block is dropped, not truncated"); |
| 85 | assert.equal(shot.content.every((c) => c.type !== "image"), true); |
| 86 | |
| 87 | const omitted = shot.receipt.image_omitted; |
| 88 | assert.ok(omitted, "the receipt must say why no image came back"); |
| 89 | assert.equal(omitted.reason, "raster_too_large"); |
| 90 | assert.equal(omitted.limit_bytes, 64); |
| 91 | assert.ok(omitted.encoded_bytes > omitted.limit_bytes); |
| 92 | assert.match(omitted.note, /zoom|region/i); |
| 93 | // The path and geometry survive, so the caller can still act on the capture. |
| 94 | assert.ok(shot.receipt.file); |
| 95 | assert.deepEqual(shot.receipt.pixels, { w: 1600, h: 1200 }); |
| 96 | |
| 97 | // The whole point: the connection is still usable afterwards. |
| 98 | assert.equal(server_alive(cu.server), true); |
| 99 | const after = await cu.call("get_app_state", {}); |
| 100 | assert.equal(after.receipt.ok, true, "every other tool survives an over-budget capture"); |
| 101 | |
| 102 | // Geometry stays bound, so coordinate targets still resolve off the receipt. |
| 103 | const click = await cu.call("left_click", { target: { type: "coordinate", x: 400, y: 300 } }); |
| 104 | assert.equal(click.receipt.ok, true, JSON.stringify(click.receipt.error)); |
| 105 | }); |
| 106 | |
| 107 | test("a raster within budget is still inlined as an image block", async (t) => { |
| 108 | const cu = await session(t); |
| 109 | const shot = await cu.call("screenshot"); |
| 110 | assert.equal(shot.receipt.ok, true); |
| 111 | assert.equal(shot.receipt.image_omitted, undefined); |
| 112 | const image = shot.content.find((c) => c.type === "image"); |
| 113 | assert.ok(image, "normal captures must keep inlining the raster"); |
| 114 | assert.equal(image.mimeType, "image/png"); |
| 115 | assert.ok(Buffer.from(image.data, "base64").length > 0); |
| 116 | }); |
| 117 | |
| 118 | function server_alive(child) { |
| 119 | return child.exitCode === null && child.signalCode === null; |
| 120 | } |
| 121 |