| 1 | import assert from "node:assert/strict"; |
| 2 | import { test } from "node:test"; |
| 3 | |
| 4 | import { |
| 5 | type ComputeEntitlements, |
| 6 | type ComputeKind, |
| 7 | type PlanMemory, |
| 8 | ComputeGate, |
| 9 | MODEL_ESTIMATE_MICROS, |
| 10 | WAITING_PREFIX, |
| 11 | actualMicros, |
| 12 | agentEstimateMicros, |
| 13 | alwaysPasses, |
| 14 | embeddingEstimateMicros, |
| 15 | entitlementsKeptSeconds, |
| 16 | isWaiting, |
| 17 | issueCapReached, |
| 18 | localRefusal, |
| 19 | lowerCap, |
| 20 | onBillingError, |
| 21 | readEntitlements, |
| 22 | refusalCode, |
| 23 | refusalMessage, |
| 24 | runBudgetUsd, |
| 25 | runMinutes, |
| 26 | sandboxEstimateMicros, |
| 27 | slotFree, |
| 28 | waitingMessage, |
| 29 | } from "../../../packages/contracts/src/compute.ts"; |
| 30 | import { BUILD_HOSTS, buildHosts, withPlanLimits } from "./egress.ts"; |
| 31 | import { WAITING, handleMention } from "./mentions.ts"; |
| 32 | |
| 33 | const repo = { namespace: "acme", name: "web" }; |
| 34 | |
| 35 | function ent(change: Partial<ComputeEntitlements> = {}): ComputeEntitlements { |
| 36 | return { |
| 37 | plan: "paid", |
| 38 | compute: true, |
| 39 | trialMicrosLeft: 0, |
| 40 | trialVerified: true, |
| 41 | firstMonth: false, |
| 42 | maxConcurrentAgents: 10, |
| 43 | maxRunMinutes: 0, |
| 44 | runCapMicros: 2_000_000, |
| 45 | issueCapMicros: 10_000_000, |
| 46 | ceilingMicros: 100_000_000, |
| 47 | exposureMicros: 0, |
| 48 | paused: null, |
| 49 | ...change, |
| 50 | }; |
| 51 | } |
| 52 | |
| 53 | /** A billing service that answers as told, and records what it was asked. */ |
| 54 | function billing(answers: { |
| 55 | entitlements?: unknown; |
| 56 | reserve?: unknown; |
| 57 | fail?: Set<string>; |
| 58 | }) { |
| 59 | const asked: { method: string; body: Record<string, unknown> }[] = []; |
| 60 | const binding = { |
| 61 | async fetch(url: string, init?: RequestInit): Promise<Response> { |
| 62 | const method = url.split("/rpc/")[1]; |
| 63 | asked.push({ method, body: JSON.parse(String(init?.body ?? "{}")) }); |
| 64 | if (answers.fail?.has(method)) return new Response("down", { status: 503 }); |
| 65 | if (method === "entitlements") return Response.json(answers.entitlements ?? null); |
| 66 | if (method === "reserve") return Response.json(answers.reserve ?? { ok: true, value: { id: "rsv_1", paidBy: "credit" } }); |
| 67 | if (method === "settle") return Response.json({ ok: true, value: true }); |
| 68 | if (method === "prices") return Response.json({ prices: [{ meter: "sandbox_second", costMicros: 20 }] }); |
| 69 | return new Response("not found", { status: 404 }); |
| 70 | }, |
| 71 | }; |
| 72 | return { binding, asked }; |
| 73 | } |
| 74 | |
| 75 | /** A plan memory that remembers nothing unless told. */ |
| 76 | function memory(plan: ComputeEntitlements["plan"] | null = null): PlanMemory { |
| 77 | return { get: async () => plan, put: async () => undefined }; |
| 78 | } |
| 79 | |
| 80 | const quiet = () => undefined; |
| 81 | |
| 82 | function wire(e: ComputeEntitlements) { |
| 83 | // As billing serialises it: camelCase. |
| 84 | return { workspace: "acme", ...e }; |
| 85 | } |
| 86 | |
| 87 | const request = (kind: ComputeKind, isPublic = false) => ({ |
| 88 | workspace: "acme", |
| 89 | repo, |
| 90 | public: isPublic, |
| 91 | kind, |
| 92 | estimateMicros: 100_000, |
| 93 | }); |
| 94 | |
| 95 | // ---- Gating decisions --------------------------------------------------------------- |
| 96 | |
| 97 | test("a paid workspace's start is reserved and goes ahead", async () => { |
| 98 | const { binding, asked } = billing({ entitlements: wire(ent()) }); |
| 99 | const admitted = await new ComputeGate(binding, memory(), quiet).admit(request("agent")); |
| 100 | assert.equal(admitted.ok, true); |
| 101 | assert.deepEqual(admitted.ok && admitted.reservation, { id: "rsv_1", paidBy: "credit" }); |
| 102 | const reserve = asked.find((call) => call.method === "reserve")!; |
| 103 | assert.deepEqual(reserve.body, { workspace: "acme", repo, public: false, kind: "agent", estimateMicros: 100_000 }); |
| 104 | }); |
| 105 | |
| 106 | test("billing's refusal is shown in its own words", async () => { |
| 107 | for (const code of ["not_paid", "trial_used", "limit", "oss_pool_empty"] as const) { |
| 108 | const { binding } = billing({ |
| 109 | entitlements: wire(ent({ plan: "free", compute: false })), |
| 110 | reserve: { ok: false, error: { code, message: `Refused: ${code}. /acme/-/billing` } }, |
| 111 | }); |
| 112 | const admitted = await new ComputeGate(binding, memory(), quiet).admit(request("check", true)); |
| 113 | assert.equal(admitted.ok, false, code); |
| 114 | assert.equal(!admitted.ok && admitted.code, code); |
| 115 | assert.equal(!admitted.ok && admitted.message, `Refused: ${code}. /acme/-/billing`); |
| 116 | } |
| 117 | }); |
| 118 | |
| 119 | test("a refusal without words gets g1t's, with what to do and where", async () => { |
| 120 | const { binding } = billing({ |
| 121 | entitlements: wire(ent({ plan: "free", compute: false })), |
| 122 | reserve: { ok: false, error: { code: "not_paid", message: "" } }, |
| 123 | }); |
| 124 | const admitted = await new ComputeGate(binding, memory(), quiet).admit(request("agent")); |
| 125 | assert.equal(!admitted.ok && admitted.message, refusalMessage("not_paid", "acme", "agent")); |
| 126 | assert.match(refusalMessage("not_paid", "acme", "agent"), /^Agents need a paid workspace\. Start the \$20 plan or try it with \$5 of free usage after a card check: \/acme\/-\/billing$/); |
| 127 | }); |
| 128 | |
| 129 | test("a paused workspace starts nothing, and billing is not asked to reserve", async () => { |
| 130 | const { binding, asked } = billing({ entitlements: wire(ent({ paused: "A spend spike is waiting for an owner." })) }); |
| 131 | const admitted = await new ComputeGate(binding, memory(), quiet).admit(request("workflow")); |
| 132 | assert.equal(!admitted.ok && admitted.code, "paused"); |
| 133 | assert.match(!admitted.ok ? admitted.message : "", /A spend spike is waiting for an owner/); |
| 134 | assert.equal(asked.some((call) => call.method === "reserve"), false); |
| 135 | }); |
| 136 | |
| 137 | test("an answer that holds compute back is kept seconds, one that lets it through half a minute", () => { |
| 138 | assert.equal(entitlementsKeptSeconds(ent()), 30); |
| 139 | assert.equal(entitlementsKeptSeconds(ent({ paused: "A spend spike is waiting for an owner." })), 3); |
| 140 | assert.equal(entitlementsKeptSeconds(ent({ plan: "free", compute: false })), 3); |
| 141 | // A free workspace with compute (its trial) runs, so it is kept as long. |
| 142 | assert.equal(entitlementsKeptSeconds(ent({ plan: "free", compute: true })), 30); |
| 143 | }); |
| 144 | |
| 145 | test("internal and enterprise plans pass even when billing refuses", async () => { |
| 146 | for (const plan of ["internal", "enterprise"] as const) { |
| 147 | const { binding } = billing({ |
| 148 | entitlements: wire(ent({ plan })), |
| 149 | reserve: { ok: false, error: { code: "limit", message: "Over." } }, |
| 150 | }); |
| 151 | const admitted = await new ComputeGate(binding, memory(), quiet).admit(request("agent")); |
| 152 | assert.equal(admitted.ok, true, plan); |
| 153 | } |
| 154 | assert.equal(alwaysPasses("internal") && alwaysPasses("enterprise") && !alwaysPasses("paid"), true); |
| 155 | }); |
| 156 | |
| 157 | test("a pause from billing holds for every plan: g1t's own budget and its daily breaker", async () => { |
| 158 | for (const plan of ["internal", "enterprise", "paid"] as const) { |
| 159 | const message = "flagon-io's monthly budget for g1t's own agents is used up ($150.00 of $150.00 this month at cost), so new runs wait."; |
| 160 | const { binding, asked } = billing({ entitlements: wire(ent({ plan })), reserve: { ok: false, error: { code: "paused", message } } }); |
| 161 | const admitted = await new ComputeGate(binding, memory(), quiet).admit({ ...request("agent"), hostedModel: true }); |
| 162 | assert.equal(!admitted.ok && admitted.code, "paused", plan); |
| 163 | assert.equal(!admitted.ok && admitted.message, message); |
| 164 | // Billing hears whether the run is on g1t's hosted models. |
| 165 | assert.equal(asked.find((call) => call.method === "reserve")!.body.hostedModel, true); |
| 166 | } |
| 167 | // Billing's pause reason reads cleanly inside g1t's sentence. |
| 168 | assert.equal( |
| 169 | refusalMessage("paused", "acme", "agent", "Its budget is used up."), |
| 170 | "g1t paused compute for this workspace: Its budget is used up. Contact support@g1t.sh to have it looked at.", |
| 171 | ); |
| 172 | }); |
| 173 | |
| 174 | test("billing down: free workspaces fail closed, paying ones go on", async () => { |
| 175 | const down = new Set(["entitlements", "reserve"]); |
| 176 | // Nothing known about the workspace: treated as free. |
| 177 | let gate = new ComputeGate(billing({ fail: down }).binding, memory(null), quiet); |
| 178 | let admitted = await gate.admit(request("agent")); |
| 179 | assert.equal(!admitted.ok && admitted.code, "billing_unavailable"); |
| 180 | // Last seen paid: goes ahead, without a reservation. |
| 181 | gate = new ComputeGate(billing({ fail: down }).binding, memory("paid"), quiet); |
| 182 | admitted = await gate.admit(request("agent")); |
| 183 | assert.deepEqual(admitted, { ok: true, reservation: null, entitlements: null }); |
| 184 | // Entitlements answer, reserve errors: the plan decides. |
| 185 | gate = new ComputeGate(billing({ entitlements: wire(ent({ plan: "free", compute: false, trialMicrosLeft: 5_000_000 })), fail: new Set(["reserve"]) }).binding, memory(), quiet); |
| 186 | admitted = await gate.admit(request("agent")); |
| 187 | assert.equal(admitted.ok, false); |
| 188 | gate = new ComputeGate(billing({ entitlements: wire(ent()), fail: new Set(["reserve"]) }).binding, memory(), quiet); |
| 189 | assert.equal((await gate.admit(request("agent"))).ok, true); |
| 190 | assert.equal(onBillingError("free"), "refuse"); |
| 191 | assert.equal(onBillingError(null), "refuse"); |
| 192 | assert.equal(onBillingError("paid"), "allow"); |
| 193 | }); |
| 194 | |
| 195 | test("an answer that is not a refusal counts as billing being down", async () => { |
| 196 | const { binding } = billing({ |
| 197 | entitlements: wire(ent({ plan: "free", compute: false })), |
| 198 | reserve: { ok: false, error: { code: "invalid", message: "Unknown method" } }, |
| 199 | }); |
| 200 | const admitted = await new ComputeGate(binding, memory(), quiet).admit(request("agent")); |
| 201 | assert.equal(!admitted.ok && admitted.code, "billing_unavailable"); |
| 202 | assert.equal(refusalCode({ code: "payment_required" }), "not_paid"); |
| 203 | assert.equal(refusalCode({ code: "conflict", reason: "trial_used" }), "trial_used"); |
| 204 | assert.equal(refusalCode({ code: "invalid" }), null); |
| 205 | }); |
| 206 | |
| 207 | test("entitlements read in either case, and not at all without a plan", () => { |
| 208 | const snake = readEntitlements({ plan: "free", trial_micros_left: 5, trial_verified: true, max_concurrent_agents: 2, paused: "" }); |
| 209 | assert.equal(snake?.trialMicrosLeft, 5); |
| 210 | assert.equal(snake?.maxConcurrentAgents, 2); |
| 211 | assert.equal(snake?.paused, null); |
| 212 | assert.equal(readEntitlements({ team: true })?.plan, undefined); |
| 213 | assert.equal(readEntitlements({ ok: true, value: { plan: "paid" } })?.plan, "paid"); |
| 214 | }); |
| 215 | |
| 216 | test("before they try: what a free workspace can start, by kind and repository", () => { |
| 217 | const free = ent({ plan: "free", compute: false, trialVerified: false }); |
| 218 | assert.equal(localRefusal(free, "agent", false), "not_paid"); |
| 219 | assert.equal(localRefusal(free, "check", true), "not_paid", "the pool needs a card check"); |
| 220 | const verified = ent({ plan: "free", compute: false, trialVerified: true, trialMicrosLeft: 0 }); |
| 221 | // The open-source path: public repositories' checks, workflows and queue. |
| 222 | for (const kind of ["check", "workflow", "queue"] as const) assert.equal(localRefusal(verified, kind, true), null, kind); |
| 223 | for (const kind of ["agent", "deploy", "embedding"] as const) assert.equal(localRefusal(verified, kind, true), "trial_used", kind); |
| 224 | assert.equal(localRefusal(verified, "check", false), "trial_used"); |
| 225 | const trial = ent({ plan: "free", compute: true, trialVerified: true, trialMicrosLeft: 4_000_000 }); |
| 226 | assert.equal(localRefusal(trial, "agent", false), null); |
| 227 | assert.equal(localRefusal(ent(), "deploy", false), null); |
| 228 | assert.equal(localRefusal(ent({ paused: "Held." }), "agent", false), "paused"); |
| 229 | }); |
| 230 | |
| 231 | // ---- Estimates ------------------------------------------------------------------------ |
| 232 | |
| 233 | test("an agent's estimate is its model's average plus its sandbox for its time cap", () => { |
| 234 | assert.equal(sandboxEstimateMicros(60, 25), 90_000); |
| 235 | assert.equal(agentEstimateMicros("implement", 90, 25), 100_000 + 135_000); |
| 236 | assert.equal(agentEstimateMicros("review", 30, 25), 70_000 + 45_000); |
| 237 | assert.equal(agentEstimateMicros("plan", 30, 20), 100_000 + 36_000); |
| 238 | // The workspace's own provider pays for the model. |
| 239 | assert.equal(agentEstimateMicros("implement", 90, 25, true), 135_000); |
| 240 | assert.equal(MODEL_ESTIMATE_MICROS.implement, 100_000); |
| 241 | assert.equal(embeddingEstimateMicros(1_000_000), 67_000); |
| 242 | assert.equal(sandboxEstimateMicros(-5, 25), 0); |
| 243 | }); |
| 244 | |
| 245 | test("what work cost: its seconds, plus its model in dollars", () => { |
| 246 | assert.equal(actualMicros(600, 20), 12_000); |
| 247 | assert.equal(actualMicros(600, 20, 0.094), 12_000 + 94_000); |
| 248 | assert.equal(actualMicros(10, 20, Number.NaN), 200); |
| 249 | }); |
| 250 | |
| 251 | test("the gate reads the sandbox price from the price book", async () => { |
| 252 | const gate = new ComputeGate(billing({}).binding, memory(), quiet); |
| 253 | assert.equal(await gate.microsPerSecond(), 20); |
| 254 | const down = new ComputeGate(billing({ fail: new Set(["prices"]) }).binding, memory(), quiet); |
| 255 | assert.equal(await down.microsPerSecond(), 25); |
| 256 | }); |
| 257 | |
| 258 | // ---- Caps ------------------------------------------------------------------------------- |
| 259 | |
| 260 | test("caps are the lower of the guardrails' and the plan's", () => { |
| 261 | assert.equal(lowerCap(90, 60), 60); |
| 262 | assert.equal(lowerCap(30, 60), 30); |
| 263 | assert.equal(lowerCap(null, 60), 60); |
| 264 | assert.equal(lowerCap(0, 0), null); |
| 265 | assert.equal(runMinutes(90, ent({ maxRunMinutes: 60 })), 60); |
| 266 | assert.equal(runMinutes(45, ent({ maxRunMinutes: 0 })), 45); |
| 267 | assert.equal(runBudgetUsd(5, ent({ runCapMicros: 2_000_000 })), 2); |
| 268 | assert.equal(runBudgetUsd(1, ent({ runCapMicros: 2_000_000 })), 1); |
| 269 | assert.equal(runBudgetUsd(null, ent({ runCapMicros: 0 })), null); |
| 270 | const guard = { policy: { budgetUsd: 5 } as never, minutes: 90 }; |
| 271 | const limited = withPlanLimits(guard as never, { minutes: 60, budgetUsd: 2 }); |
| 272 | assert.equal(limited.minutes, 60); |
| 273 | assert.equal((limited.policy as { budgetUsd: number }).budgetUsd, 2); |
| 274 | const unlimited = withPlanLimits(guard as never, { minutes: null, budgetUsd: null }); |
| 275 | assert.equal(unlimited.minutes, 90); |
| 276 | assert.equal((unlimited.policy as { budgetUsd: number }).budgetUsd, 5); |
| 277 | }); |
| 278 | |
| 279 | test("agents at once: a run waits for a free slot past the plan's cap", () => { |
| 280 | const first = ent({ firstMonth: true, maxConcurrentAgents: 2 }); |
| 281 | assert.equal(slotFree(1, first), true); |
| 282 | assert.equal(slotFree(2, first), false); |
| 283 | assert.equal(slotFree(50, ent({ plan: "internal", maxConcurrentAgents: 2 })), true); |
| 284 | assert.equal(slotFree(50, ent({ maxConcurrentAgents: 0 })), true); |
| 285 | assert.equal(slotFree(50, null), true); |
| 286 | const said = waitingMessage(2); |
| 287 | assert.equal(isWaiting(said), true); |
| 288 | assert.match(said, /runs 2 agents at a time/); |
| 289 | assert.equal(isWaiting("Agents need a paid workspace."), false); |
| 290 | // mentions.ts keeps its own copy, to have no runtime imports. |
| 291 | assert.equal(WAITING, WAITING_PREFIX); |
| 292 | }); |
| 293 | |
| 294 | test("an issue's agents stop at its cap, with a way to raise it", () => { |
| 295 | assert.equal(issueCapReached(9_999_999, ent(), 12), null); |
| 296 | const capped = issueCapReached(10_000_000, ent(), 12); |
| 297 | assert.match(capped ?? "", /#12 have spent \$10\.00, its cap of \$10\.00/); |
| 298 | assert.equal(issueCapReached(50_000_000, ent({ issueCapMicros: 0 }), 12), null); |
| 299 | assert.equal(issueCapReached(50_000_000, ent({ plan: "internal" }), 12), null); |
| 300 | assert.match(refusalMessage("issue_cap", "acme", "agent", capped), /raise the cap per issue: \/acme\/-\/billing#caps$/); |
| 301 | }); |
| 302 | |
| 303 | // ---- Builds' network --------------------------------------------------------------------- |
| 304 | |
| 305 | test("builds reach registries and git hosts, and no mining pool", () => { |
| 306 | for (const host of ["registry.npmjs.org", "codeload.github.com", "nodejs.org", "crates.io"]) { |
| 307 | assert.ok(BUILD_HOSTS.includes(host), host); |
| 308 | } |
| 309 | assert.ok(buildHosts("deploy").includes("api.cloudflare.com")); |
| 310 | assert.ok(!buildHosts("actions").includes("api.cloudflare.com")); |
| 311 | for (const host of buildHosts("deploy")) { |
| 312 | assert.doesNotMatch(host, /pool|xmr|monero|nicehash|^\*/, host); |
| 313 | } |
| 314 | }); |
| 315 | |
| 316 | // ---- Mentions -------------------------------------------------------------------------------- |
| 317 | |
| 318 | test("a mention whose run waits for a slot says so, and is not a failure", async () => { |
| 319 | const replies: string[] = []; |
| 320 | const recorded: string[] = []; |
| 321 | const job = { |
| 322 | commentId: "cmt_1", |
| 323 | actor: { id: "u1", username: "ana" }, |
| 324 | repo, |
| 325 | number: 3, |
| 326 | member: true, |
| 327 | pull: null, |
| 328 | intent: "work", |
| 329 | issueOpen: true, |
| 330 | workingPull: null, |
| 331 | body: "@g1t take this", |
| 332 | defaultBranch: "main", |
| 333 | } as never; |
| 334 | await handleMention(job, { |
| 335 | mentions: { replyMention: async (_id: string, text: string) => (replies.push(text), true) } as never, |
| 336 | refusal: async () => null, |
| 337 | assign: async () => ({ ok: false, error: { code: "conflict", message: waitingMessage(2) } }), |
| 338 | revise: async () => null, |
| 339 | review: async () => ({ ok: true, value: true }), |
| 340 | answer: async () => ({ ok: true, value: true }), |
| 341 | message: async () => ({ ok: true, value: true }), |
| 342 | record: async (_job, why) => void recorded.push(why), |
| 343 | }); |
| 344 | assert.equal(recorded.length, 0); |
| 345 | assert.equal(replies.length, 1); |
| 346 | assert.match(replies[0], /^@ana, Waiting for a free slot/); |
| 347 | }); |