@edgehero/pi-dispatch 0.1.0

package/src/budget.mjs

@@ -0,0 +1,179 @@
+/**
+ * CONST-BUDGET-BEFORE-TOKENS. The ordering IS the mechanism.
+ *
+ * Reserve a slot against the spend caps BEFORE a container starts. Check-after-spend means fifty
+ * junk triggers cost fifty jobs of real money before the cap engages -- the exact scenario the cap
+ * exists for. So: atomically INCR each active window's counter, and only start the container if every
+ * window is within its cap and outside its soft-hold band.
+ *
+ * INCR is atomic -- verified against a real Valkey under 20 parallel increments yielding exactly
+ * 20, no lost updates -- so this needs no lock even at concurrency 3. Three jobs racing at count 9
+ * with cap 10 reserve 10, 11, 12; exactly one proceeds, the others are refused, and no container
+ * starts over budget.
+ *
+ * Three windows bound spend on different horizons: a mandatory DAY cap, plus OPTIONAL WEEK and MONTH
+ * ceilings (REQ-SPEND-CAPS-MULTI-WINDOW). A window whose cap is null/absent is disabled -- not counted,
+ * not evaluated. On top of the hard caps a single SOFT-HOLD percentage refuses new starts once any
+ * window enters its band (e.g. >= 80% of that window's cap), a distinct operator-visible brake before
+ * the hard wall; in-flight containers are untouched because the reservation happens pre-container.
+ *
+ * `redis` is any ioredis-compatible client (BullMQ bundles ioredis). Injected so the logic is
+ * testable without a running server.
+ */
+
+/** UTC day key: `budget:YYYY-MM-DD`. */
+export function dayKey(now = new Date(), prefix = "budget") {
+	return `${prefix}:${now.toISOString().slice(0, 10)}`; // YYYY-MM-DD, UTC
+}
+
+/**
+ * UTC week key, bucketed by the Monday of the run's week: `budget:w:YYYY-MM-DD` (that Monday's date).
+ * Keying by an actual date sidesteps ISO week-number edge cases (week 53, year rollovers) while still
+ * mapping every run in one Mon-Sun week to one key. The `w:` sub-namespace never collides with the day
+ * key (`budget:YYYY-MM-DD`).
+ */
+export function weekKey(now = new Date(), prefix = "budget") {
+	const d = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate()));
+	const sinceMonday = (d.getUTCDay() + 6) % 7; // Mon=0 .. Sun=6
+	d.setUTCDate(d.getUTCDate() - sinceMonday);
+	return `${prefix}:w:${d.toISOString().slice(0, 10)}`;
+}
+
+/** UTC month key: `budget:m:YYYY-MM`. The `m:` sub-namespace never collides with the day/week keys. */
+export function monthKey(now = new Date(), prefix = "budget") {
+	return `${prefix}:m:${now.toISOString().slice(0, 7)}`; // YYYY-MM, UTC
+}
+
+/**
+ * UTC day key for the daily TOKEN counter: `budget:t:YYYY-MM-DD`. The `t:` sub-namespace never collides
+ * with the job-count day/week/month keys -- token spend and job count are separate ledgers.
+ */
+export function tokenDayKey(now = new Date(), prefix = "budget") {
+	return `${prefix}:t:${now.toISOString().slice(0, 10)}`; // YYYY-MM-DD, UTC
+}
+
+// Each window's TTL outlives its bucket with slack, so a stale counter is reclaimed shortly after the
+// window rolls over. Set once, on first reservation only (reserved === 1), so a busy window cannot push
+// its own expiry forward indefinitely.
+const DAY_TTL_SECONDS = 2 * 24 * 60 * 60; // outlives the UTC day
+const WEEK_TTL_SECONDS = 9 * 24 * 60 * 60; // outlives the Mon-Sun week
+const MONTH_TTL_SECONDS = 40 * 24 * 60 * 60; // outlives the longest month
+
+/**
+ * Classify one window's reserved count against its cap and the soft-hold band. PURE, and exported so the
+ * worker's `reserveBudget` and the admin panel compute the same state from the same inputs and cannot
+ * drift on where the band and the wall begin (the read-model already imports `dayKey` for the same reason).
+ *
+ *   reserved > cap                      -> "over"       (hard cap exceeded; also fails closed for cap <= 0)
+ *   reserved > floor(cap * pct / 100)   -> "soft-hold"  (inside the band; only when pct is set)
+ *   otherwise                           -> "ok"
+ *
+ * `softHoldPct` null/absent disables the band. Callers only classify ACTIVE windows (cap is a real number);
+ * a disabled window (cap null) is never passed here.
+ */
+export function windowState(reserved, cap, softHoldPct = null) {
+	if (reserved > cap) return "over";
+	if (Number.isInteger(softHoldPct) && reserved > Math.floor((cap * softHoldPct) / 100)) return "soft-hold";
+	return "ok";
+}
+
+/** The active windows for a `caps` set: day plus week/month only when their cap is set (not null/undefined). */
+function activeWindows(caps, now, keyPrefix) {
+	return [
+		{ name: "day", key: dayKey(now, keyPrefix), cap: caps?.day, ttl: DAY_TTL_SECONDS },
+		{ name: "week", key: weekKey(now, keyPrefix), cap: caps?.week, ttl: WEEK_TTL_SECONDS },
+		{ name: "month", key: monthKey(now, keyPrefix), cap: caps?.month, ttl: MONTH_TTL_SECONDS },
+	].filter((w) => w.cap !== null && w.cap !== undefined);
+}
+
+/**
+ * Atomically reserve one slot across every ACTIVE window. Returns
+ * `{ allowed, reason, blockedWindow, windows }` where `windows[name]` is `{ reserved, cap, state }` for an
+ * active window or `null` for a disabled one, and `reason` is the worst window state:
+ *   any "over"       -> "over-budget"
+ *   else any "soft-hold" -> "soft-hold"
+ *   else                 -> "ok"  (allowed)
+ * `blockedWindow` names the day>week>month-first window carrying that worst state, so the caller can say
+ * WHICH ceiling blocked. Both refusals are pre-container, so a refused job spends nothing; a soft-hold is
+ * refused exactly like over-budget, only under a distinct reason and an amber panel state.
+ *
+ * Every active window is INCR'd UNCONDITIONALLY -- the same "a refused reservation still counts, no
+ * give-back" invariant the daily cap has always had, now per window. This is deliberate: conditional
+ * increments would reintroduce the read-then-write race the atomic INCR design exists to avoid, and
+ * unconditional INCR keeps `releaseBudget` a clean symmetric DECR of the same windows. It fails closed --
+ * a burst that exhausts one window counts against the others too -- which is the money-safe direction.
+ *
+ * A cap of 0/negative on an ACTIVE window fails closed (every job "over") rather than meaning "unlimited";
+ * a DISABLED window (cap null/absent) is not counted or evaluated at all.
+ */
+export async function reserveBudget(redis, { caps, softHoldPct = null, now = new Date(), keyPrefix = "budget" } = {}) {
+	const windows = { day: null, week: null, month: null };
+	for (const w of activeWindows(caps, now, keyPrefix)) {
+		const reserved = Number(await redis.incr(w.key));
+		// Set the TTL only when the key is first created, so a long window cannot push its expiry forward.
+		if (reserved === 1) await redis.expire(w.key, w.ttl);
+		windows[w.name] = { reserved, cap: w.cap, state: windowState(reserved, w.cap, softHoldPct) };
+	}
+
+	const order = ["day", "week", "month"];
+	for (const [state, reason] of [
+		["over", "over-budget"],
+		["soft-hold", "soft-hold"],
+	]) {
+		const hit = order.find((name) => windows[name]?.state === state);
+		if (hit) return { allowed: false, reason, blockedWindow: hit, windows };
+	}
+	return { allowed: true, reason: "ok", blockedWindow: null, windows };
+}
+
+/**
+ * Give a reservation back in every window `reserveBudget` would have reserved. Used ONLY when the container
+ * never started because of an INFRA fault AFTER reserving (e.g. the docker daemon was unreachable) -- an
+ * infra failure that spent nothing should not permanently consume a slot. DECRs exactly the ACTIVE windows,
+ * mirroring `reserveBudget` so the refund cannot drift from the reservation. NOT used for a completed run
+ * (which really consumed its slot), an exit-1 infra retry (the container ran and spent), or an
+ * over-budget/soft-hold refusal (which never gives back its slot, deliberately).
+ */
+export async function releaseBudget(redis, { caps, now = new Date(), keyPrefix = "budget" } = {}) {
+	for (const w of activeWindows(caps, now, keyPrefix)) {
+		await redis.decr(w.key);
+	}
+}
+
+/**
+ * The daily TOKEN cap check -- read-only, no increment. Returns `{ allowed, reason, spent, cap }`.
+ *
+ * This is the deliberate ASYMMETRY to `reserveBudget` (issue #25 / OQ-010). A job's token cost is
+ * knowable only AFTER it runs, so a token cap cannot check-and-increment BEFORE the spend the way the
+ * job-COUNT cap does (CONST-BUDGET-BEFORE-TOKENS, which governs only the ordering of the job-count cap
+ * and stays unchanged). Instead the token count is a check-BEFORE read of PRIOR jobs' recorded spend
+ * (this function) paired with a record-AFTER `INCRBY` (`recordTokenSpend`). It can only stop the NEXT
+ * job once a day's accumulated spend has reached the cap -- a lagging backstop, not a before-the-spend
+ * cap; `maxTurns` remains the one proactive per-job lever. Being read-only it consumes nothing, so it
+ * safely precedes the job-count reservation and a refusal here spends no job-count slot.
+ *
+ * `cap` null/absent disables the cap (always allowed). A cap <= 0 fails closed (spent >= 0 always
+ * blocks), matching `reserveBudget`'s fail-closed posture on a nonsensical cap.
+ */
+export async function checkTokenCap(redis, { cap, now = new Date(), keyPrefix = "budget" } = {}) {
+	if (cap === null || cap === undefined) return { allowed: true, reason: "ok", spent: 0, cap: null };
+	const spent = Number(await redis.get(tokenDayKey(now, keyPrefix))) || 0;
+	if (spent >= cap) return { allowed: false, reason: "daily-token-cap", spent, cap };
+	return { allowed: true, reason: "ok", spent, cap };
+}
+
+/**
+ * Record a job's token spend against the daily counter -- the record-AFTER half of the lagging token cap
+ * (see `checkTokenCap`). `INCRBY` by the job's total tokens; set the day TTL once, on first write, so a
+ * busy day cannot push its own expiry forward (mirrors `reserveBudget`'s set-once TTL). Under concurrency
+ * the counter is best-effort: N in-flight jobs each passed `checkTokenCap` before any recorded, so the
+ * daily total can overshoot the cap by up to N per-job budgets -- expected and acceptable for a lagging
+ * backstop (OQ-010). Returns the new counter value. `tokens` is assumed a non-negative integer (the
+ * caller guards `> 0`), so the first-write TTL test `total === tokens` holds.
+ */
+export async function recordTokenSpend(redis, tokens, { now = new Date(), keyPrefix = "budget" } = {}) {
+	const key = tokenDayKey(now, keyPrefix);
+	const total = Number(await redis.incrby(key, tokens));
+	if (total === tokens) await redis.expire(key, DAY_TTL_SECONDS);
+	return total;
+}

package/src/cli.mjs

@@ -0,0 +1,208 @@
+#!/usr/bin/env node
+import { existsSync } from "node:fs";
+import { resolve } from "node:path";
+import { parseArgs } from "node:util";
+import { loadConfig } from "./config.mjs";
+import { EXIT_POLICY } from "./exit-code.mjs";
+import { gitDirty } from "./git-dirty.mjs";
+
+const USAGE = `pi-dispatch — run pi coding-agent flows on your own folders
+
+  pi-dispatch init         scaffold .env + triggers.json + pause-windows.json + pi-packages.json + subscriptions.json here
+  pi-dispatch doctor [--fix]  preflight Docker, Valkey, the job image, and your provider key; --fix offers to run each fix (y/N per action)
+  pi-dispatch up [--yes]   one consented pass: pull+tag the job image, start Valkey, init, doctor
+  pi-dispatch setup github mint GitHub App credentials in one browser click (App Manifest flow);
+                           every write shown first and individually consented — no --yes here
+                           (--webhook-url <URL> | --no-webhook) [--org <org>] [--name <appName>]
+  pi-dispatch import-pi    stage your host pi setup (models/skills/persona) into a global overlay
+                           [--no-extensions] [--with-packages] [--packages-file <path>]
+                           [--from <agentDir>] [--to <overlayDir>]
+
+  pi-dispatch run <folder> --task "<what to do>" [--flow <name>]
+                           [--provider <p>] [--model <m>] [--max-turns <n>] [--image <ref>] [--force]
+  pi-dispatch sandbox <jobId>
+                           re-open a finished run's sandbox as a shell — same image, same workspace,
+                           no credentials  [--publish <port>[:<containerPort>]] [--pin]
+  pi-dispatch sandbox --list
+                           what is still re-openable, for how long, and what is running now
+
+  pi-dispatch worker       drain the queue (run this in another terminal, or as a service)
+  pi-dispatch-receiver     webhook receiver for forge triggers — its own bin (see the GitHub section of the README)
+  pi-dispatch service <render|install|uninstall|status|start|stop|restart> [--receiver] [--user|--system] [--force]
+                           run the worker (or --receiver) as an OS service — the deploy/ templates
+                           rendered with this host's real paths, installed user-level;
+                           \`service restart --drain\` lets the in-flight job finish first
+  pi-dispatch pause        stop taking new jobs (durable; survives worker restart)
+  pi-dispatch resume       resume taking jobs
+  pi-dispatch status       show paused state + job counts
+
+Config comes from the environment (see .env.example); flags override it per run.`;
+
+export async function main(argv = process.argv.slice(2), env = process.env) {
+	const cmd = argv[0];
+
+	if (cmd === "init") {
+		const { runInit } = await import("./init.mjs");
+		return runInit(process.cwd());
+	}
+
+	if (cmd === "doctor") {
+		const { runDoctor } = await import("./doctor.mjs");
+		// `fix` rides in the deps position (runDoctor(env, depsOrOpts)) — one options bag, no third arg.
+		return runDoctor(env, { fix: argv.slice(1).includes("--fix") });
+	}
+
+	if (cmd === "up") {
+		const { runUp } = await import("./up.mjs");
+		return runUp(argv.slice(1), { env });
+	}
+
+	if (cmd === "setup") {
+		// Subcommand shape (`setup <forge>`) so future forges can land beside github without a new
+		// top-level verb; an unknown or missing target prints guidance rather than guessing.
+		if (argv[1] === "github") {
+			const { runGithubAppSetup } = await import("./github-app-setup.mjs");
+			return runGithubAppSetup(argv.slice(2), { env });
+		}
+		process.stdout.write(`pi-dispatch setup <target> — guided credential setup\n\n  targets: github\n\n  pi-dispatch setup github (--webhook-url <URL> | --no-webhook) [--org <org>] [--name <appName>]\n      mint GitHub App credentials via the App Manifest flow — one browser click returns the app id,\n      private key, and webhook secret; every write is shown first and individually consented\n`);
+		return argv[1] ? 1 : 0;
+	}
+
+	if (cmd === "import-pi") {
+		const { runImportPi } = await import("./import-pi.mjs");
+		return runImportPi(argv.slice(1), { env });
+	}
+
+	if (cmd === "sandbox") {
+		const { runSandbox } = await import("./sandbox-cli.mjs");
+		return runSandbox(argv.slice(1), { env });
+	}
+
+	if (cmd === "worker") {
+		const { startWorker } = await import("./start.mjs");
+		await startWorker(env);
+		return 0; // the worker keeps the process alive until SIGTERM
+	}
+
+	if (cmd === "service") {
+		const { runService } = await import("./service.mjs");
+		return runService(argv.slice(1), { env });
+	}
+
+	if (cmd === "run") {
+		const { values, positionals } = parseArgs({
+			args: argv.slice(1),
+			allowPositionals: true,
+			options: {
+				task: { type: "string" },
+				flow: { type: "string" },
+				provider: { type: "string" },
+				model: { type: "string" },
+				"max-turns": { type: "string" },
+				image: { type: "string" }, // the container image for this one job; blank/absent = PI_JOB_IMAGE
+				force: { type: "boolean", default: false },
+			},
+		});
+		const folder = positionals[0] && resolve(positionals[0]);
+		if (!folder || !existsSync(folder)) return fail(`folder not found: ${positionals[0] ?? "(none given)"}`);
+		if (!values.task) return fail("a --task is required");
+
+		// A local job edits the folder IN PLACE with no undo (SECURITY.md). Refuse a dirty working
+		// tree unless --force, so a bad run cannot mix with uncommitted work the operator can't
+		// cleanly separate. A non-git folder is caught later by prepare (v1 requires a git repo).
+		if (existsSync(`${folder}/.git`) && !values.force) {
+			const dirty = gitDirty(folder);
+			if (dirty === null) return fail(`${folder} is not a usable git repository`);
+			if (dirty) return fail(`${folder} has uncommitted changes. Commit or stash them, or pass --force.`);
+		}
+
+		const config = loadConfig(env);
+		const { parseConnection } = await import("./connection.mjs");
+		const { makeQueue, enqueueLocalJob } = await import("./queue.mjs");
+		// failFast: a one-shot enqueue must not hang forever if Valkey is down -- error clearly.
+		const queue = makeQueue(parseConnection(config.valkeyUrl, { failFast: true }));
+		try {
+			// Absent flags stay absent (undefined) so the value resolves at job start against the
+			// settings overlay/env, not a default frozen here (INT-CONFIG-OVERLAY-CONTRACT).
+			const jobId = await enqueueLocalJob(queue, {
+				folder,
+				task: values.task,
+				flow: values.flow,
+				provider: values.provider,
+				model: values.model,
+				maxTurns: values["max-turns"] ? Number(values["max-turns"]) : undefined,
+				// || not the raw value: `--image ""` must collapse to absent rather than becoming a falsy string that
+				// throws inside buildDockerRunArgs after a budget slot is reserved.
+				image: values.image || undefined,
+			});
+			process.stdout.write(`queued ${jobId} — folder ${folder}\nrun \`pi-dispatch worker\` to process it.\n`);
+		} catch (error) {
+			return fail(`could not reach Valkey at ${config.valkeyUrl} — is it running? (docker compose up)\n  ${error.message}`);
+		} finally {
+			await queue.close().catch(() => {});
+		}
+		return 0;
+	}
+
+	if (cmd === "pause" || cmd === "resume" || cmd === "status") {
+		// The kill switch reads ONLY VALKEY_URL, not the full loadConfig -- it must work even when
+		// GitHub auth is misconfigured, so an operator can always stop the queue.
+		const url = env.VALKEY_URL ?? "redis://127.0.0.1:6379";
+		const { parseConnection } = await import("./connection.mjs");
+		const { makeQueue } = await import("./queue.mjs");
+		// failFast: reach the kill switch in seconds when Valkey is down, not a hang. makeQueue pins
+		// the "pi-jobs" name -- a different name would pause an empty queue (silent no-op).
+		const queue = makeQueue(parseConnection(url, { failFast: true }));
+		try {
+			if (cmd === "pause") {
+				await queue.pause();
+				process.stdout.write("paused — worker will stop taking new jobs (jobs still enqueue)\n");
+			} else if (cmd === "resume") {
+				await queue.resume();
+				process.stdout.write("resumed\n");
+			} else {
+				// "paused" is included in the counts because jobs enqueued while paused land in the
+				// `paused` list, not `wait` -- omitting it would report backlog 0 in the exact state
+				// the pause switch creates. `pausedState` (the boolean) is named apart from the
+				// `paused` count `getJobCounts` returns, so the two do not collide in the output.
+				const pausedState = await queue.isPaused();
+				const counts = await queue.getJobCounts("waiting", "active", "paused", "delayed", "failed");
+				process.stdout.write(`${JSON.stringify({ pausedState, ...counts })}\n`);
+			}
+		} catch (error) {
+			return fail(`could not reach Valkey at ${url} — is it running? (docker compose up)\n  ${error.message}`);
+		} finally {
+			await queue.close().catch(() => {});
+		}
+		return 0;
+	}
+
+	process.stdout.write(`${USAGE}\n`);
+	return cmd ? 1 : 0;
+}
+
+function fail(message) {
+	process.stderr.write(`error: ${message}\n`);
+	return 1;
+}
+
+/**
+ * Exit code for an error that escaped main() as a rejection. A tagged config error (loadConfig's
+ * `piDispatchConfig`) is a determinate refusal -> EXIT_POLICY (2, never retried); anything else is
+ * infra -> 1 (retryable). Mirrors INT-RUNNER-EXIT-CODE-PROTOCOL for the CLI's own exit space.
+ */
+export function entryExitCode(err) {
+	return err?.piDispatchConfig ? EXIT_POLICY : 1;
+}
+
+// Entry point when run as a bin. Kept out of the exported main so tests can call main() directly.
+if (import.meta.url === `file://${process.argv[1]}` || process.argv[1]?.endsWith("cli.mjs")) {
+	main()
+		.then((code) => {
+			if (code) process.exitCode = code;
+		})
+		.catch((err) => {
+			process.stderr.write(`error: ${err.message}\n`);
+			process.exitCode = entryExitCode(err);
+		});
+}