Blame · Line-by-line history

`agent-multiplayer.ts`

Each line is annotated with the commit that last touched it. Click any SHA to jump to that commit and see the surrounding change.

agent-multiplayer.tsBlame408 lines · 1 contributor

Normal view Raw

e75eddcClaude	1	/**
	2	* Agent multiplayer v1 — sessions, leases, budgets.
	3	*
	4	* Goal: when 10-100 AI agents push to the same repo, they don't step on
	5	* each other. Each agent gets a stable session with:
	6	* - a token (`agt_<hex>`) used as Bearer credentials,
	7	* - a `branch_namespace` prefix the git plumbing enforces on every
	8	* ref update,
	9	* - a daily spend budget (cents) that callers consult via
	10	* `chargeAgent()` before billing an action.
	11	*
	12	* Coordination on shared targets (issues, PRs, file paths, branches)
	13	* flows through `agent_leases` — a soft mutex with a TTL. The active
	14	* lease is enforced by a partial UNIQUE index on
	15	* (target_type, target_id) WHERE status='active'; conflicting INSERTs
	16	* raise a 23505 we translate to `null`.
	17	*
	18	* All helpers swallow DB errors and return null/false, matching the
	19	* graceful-degradation pattern used elsewhere on this codebase. Pure
	20	* format helpers (token generation, namespace normalisation) work
	21	* without a DB so unit tests can exercise them in isolation.
	22	*/
	23
	24	import { and, eq, lt, sql } from "drizzle-orm";
	25	import { db } from "../db";
	26	import { agentSessions, agentLeases } from "../db/schema";
	27	import type { AgentSession, AgentLease } from "../db/schema";
	28
	29	/** Plaintext token prefix — distinct from PAT (`glc_`) and OAuth (`glct_`). */
	30	export const AGENT_TOKEN_PREFIX = "agt_";
	31
	32	/** Default lease duration when callers don't supply one. */
	33	export const DEFAULT_LEASE_DURATION_MS = 5 * 60 * 1000;
	34
	35	/** Recognised target_type values. Keep in sync with the migration comment. */
	36	export const LEASE_TARGET_TYPES = [
	37	"issue",
	38	"pr",
	39	"file_path",
	40	"branch",
	41	] as const;
	42	export type LeaseTargetType = (typeof LEASE_TARGET_TYPES)[number];
	43
	44	// ---------------------------------------------------------------------------
	45	// Pure helpers (no DB)
	46	// ---------------------------------------------------------------------------
	47
	48	/** 32-byte hex token, `agt_`-prefixed. Mirrors the PAT generator shape. */
	49	export function generateAgentToken(): string {
	50	const bytes = crypto.getRandomValues(new Uint8Array(32));
	51	return (
	52	AGENT_TOKEN_PREFIX +
	53	Array.from(bytes)
	54	.map((b) => b.toString(16).padStart(2, "0"))
	55	.join("")
	56	);
	57	}
	58
	59	/** SHA-256 hex of the plaintext token — what we persist. */
	60	export async function hashAgentToken(token: string): Promise<string> {
	61	const data = new TextEncoder().encode(token);
	62	const digest = await crypto.subtle.digest("SHA-256", data);
	63	return Array.from(new Uint8Array(digest))
	64	.map((b) => b.toString(16).padStart(2, "0"))
	65	.join("");
	66	}
	67
	68	/**
	69	* Detect the agent token shape. Used by middleware to short-circuit
	70	* before hitting the DB on regular PAT/OAuth tokens.
	71	*/
	72	export function isAgentToken(token: string): boolean {
	73	return /^agt_[0-9a-f]{64}$/.test(token);
	74	}
	75
	76	/**
	77	* Normalise a caller-supplied branch namespace to a canonical
	78	* `agents/<name>/` form. We keep it lower-case, strip leading
	79	* `refs/heads/`, and force a trailing slash so `startsWith()` checks
	80	* elsewhere are unambiguous.
	81	*/
	82	export function normaliseBranchNamespace(name: string, raw?: string): string {
	83	let ns = (raw ?? `agents/${name}`).trim();
	84	if (ns.startsWith("refs/heads/")) ns = ns.slice("refs/heads/".length);
	85	ns = ns.replace(/^\/+\|\/+$/g, ""); // strip surrounding slashes
	86	if (!ns) ns = `agents/${name}`;
	87	return ns + "/";
	88	}
	89
	90	/**
	91	* Check whether a fully-qualified ref (e.g. `refs/heads/agents/claude-1/foo`)
	92	* sits inside the agent's allowed namespace. The caller is expected to pass
	93	* a fully-qualified ref OR a short branch name — both work.
	94	*/
	95	export function refIsInNamespace(ref: string, namespace: string): boolean {
	96	const short = ref.startsWith("refs/heads/")
	97	? ref.slice("refs/heads/".length)
	98	: ref;
	99	return short.startsWith(namespace);
	100	}
	101
	102	// ---------------------------------------------------------------------------
	103	// Session lifecycle
	104	// ---------------------------------------------------------------------------
	105
	106	export interface CreateAgentSessionInput {
	107	ownerUserId: string;
	108	name: string;
	109	repositoryId?: string \| null;
	110	branchNamespace?: string;
	111	budgetCentsPerDay?: number;
	112	}
	113
	114	export interface CreateAgentSessionResult {
	115	session: AgentSession;
	116	/** Plaintext token — returned exactly once, never persisted. */
	117	token: string;
	118	}
	119
	120	/**
	121	* Mint a new agent session. Returns the plaintext token alongside the
	122	* stored row so the caller can hand it to the agent. The token is only
	123	* recoverable here; we persist only its SHA-256 hash.
	124	*/
	125	export async function createAgentSession(
	126	input: CreateAgentSessionInput
	127	): Promise<CreateAgentSessionResult \| null> {
	128	const name = input.name.trim();
	129	if (!name) return null;
	130
	131	const token = generateAgentToken();
	132	const tokenHash = await hashAgentToken(token);
	133	const branchNamespace = normaliseBranchNamespace(name, input.branchNamespace);
	134	const budget =
	135	typeof input.budgetCentsPerDay === "number" && input.budgetCentsPerDay >= 0
	136	? Math.floor(input.budgetCentsPerDay)
	137	: 500;
	138
	139	try {
	140	const [row] = await db
	141	.insert(agentSessions)
	142	.values({
	143	name,
	144	ownerUserId: input.ownerUserId,
	145	repositoryId: input.repositoryId ?? null,
	146	tokenHash,
	147	branchNamespace,
	148	budgetCentsPerDay: budget,
	149	spentCentsToday: 0,
	150	})
	151	.returning();
	152	if (!row) return null;
	153	return { session: row, token };
	154	} catch {
	155	return null;
	156	}
	157	}
	158
	159	/**
	160	* Look up a session by plaintext token. Returns null if the token is
	161	* malformed, unknown, or the DB call throws. Side-effect: bumps
	162	* `last_active_at` on a hit (fire-and-forget).
	163	*/
	164	export async function authenticateAgent(
	165	token: string
	166	): Promise<AgentSession \| null> {
	167	if (!isAgentToken(token)) return null;
	168	try {
	169	const tokenHash = await hashAgentToken(token);
	170	const [row] = await db
	171	.select()
	172	.from(agentSessions)
	173	.where(eq(agentSessions.tokenHash, tokenHash))
	174	.limit(1);
	175	if (!row) return null;
	176	// Fire-and-forget — never block the request on the activity bump.
	177	db.update(agentSessions)
	178	.set({ lastActiveAt: new Date() })
	179	.where(eq(agentSessions.id, row.id))
	180	.catch(() => undefined);
	181	return row;
	182	} catch {
	183	return null;
	184	}
	185	}
	186
	187	export async function listAgentSessionsForOwner(
	188	ownerUserId: string
	189	): Promise<AgentSession[]> {
	190	try {
	191	return await db
	192	.select()
	193	.from(agentSessions)
	194	.where(eq(agentSessions.ownerUserId, ownerUserId));
	195	} catch {
	196	return [];
	197	}
	198	}
	199
	200	export async function revokeAgentSession(
	201	sessionId: string,
	202	ownerUserId: string
	203	): Promise<boolean> {
	204	try {
	205	const result = await db
	206	.delete(agentSessions)
	207	.where(
	208	and(
	209	eq(agentSessions.id, sessionId),
	210	eq(agentSessions.ownerUserId, ownerUserId)
	211	)
	212	)
	213	.returning({ id: agentSessions.id });
	214	return result.length > 0;
	215	} catch {
	216	return false;
	217	}
	218	}
	219
	220	// ---------------------------------------------------------------------------
	221	// Leases
	222	// ---------------------------------------------------------------------------
	223
	224	/**
	225	* Try to grab an exclusive lease on a target. Returns the lease row on
	226	* success, null when another agent already holds an active lease. The
	227	* uniqueness invariant is enforced by a partial UNIQUE index, so this
	228	* is race-safe even under concurrent INSERTs — the loser hits a 23505
	229	* we translate to null.
	230	*
	231	* Before INSERT we sweep any expired-but-still-active rows for this
	232	* target, flipping their status to 'expired' so a stale holder doesn't
	233	* block a fresh request.
	234	*/
	235	export async function acquireLease(
	236	agentSessionId: string,
	237	targetType: string,
	238	targetId: string,
	239	durationMs: number = DEFAULT_LEASE_DURATION_MS
	240	): Promise<AgentLease \| null> {
	241	if (!agentSessionId \|\| !targetType \|\| !targetId) return null;
	242	const now = new Date();
	243	const expiresAt = new Date(now.getTime() + Math.max(1, durationMs));
	244
	245	try {
	246	// Expire stale leases on this target so the unique index doesn't
	247	// wrongly block a new acquire.
	248	await db
	249	.update(agentLeases)
	250	.set({ status: "expired" })
	251	.where(
	252	and(
	253	eq(agentLeases.targetType, targetType),
	254	eq(agentLeases.targetId, targetId),
	255	eq(agentLeases.status, "active"),
	256	lt(agentLeases.expiresAt, now)
	257	)
	258	);
	259
	260	const [row] = await db
	261	.insert(agentLeases)
	262	.values({
	263	agentSessionId,
	264	targetType,
	265	targetId,
	266	acquiredAt: now,
	267	expiresAt,
	268	status: "active",
	269	})
	270	.returning();
	271	return row ?? null;
	272	} catch {
	273	// 23505 unique violation (another active lease) or any other DB error
	274	// — treat as "couldn't acquire". The caller should fall back.
	275	return null;
	276	}
	277	}
	278
	279	export async function releaseLease(leaseId: string): Promise<boolean> {
	280	if (!leaseId) return false;
	281	try {
	282	const result = await db
	283	.update(agentLeases)
	284	.set({ status: "released" })
	285	.where(
	286	and(eq(agentLeases.id, leaseId), eq(agentLeases.status, "active"))
	287	)
	288	.returning({ id: agentLeases.id });
	289	return result.length > 0;
	290	} catch {
	291	return false;
	292	}
	293	}
	294
	295	export async function listLeasesForAgent(
	296	agentSessionId: string
	297	): Promise<AgentLease[]> {
	298	try {
	299	return await db
	300	.select()
	301	.from(agentLeases)
	302	.where(eq(agentLeases.agentSessionId, agentSessionId));
	303	} catch {
	304	return [];
	305	}
	306	}
	307
	308	/**
	309	* Mark every active lease whose `expires_at` has passed as 'expired'.
	310	* Safe to call from the autopilot ticker.
	311	*/
	312	export async function expireStaleLeases(now: Date = new Date()): Promise<number> {
	313	try {
	314	const result = await db
	315	.update(agentLeases)
	316	.set({ status: "expired" })
	317	.where(
	318	and(eq(agentLeases.status, "active"), lt(agentLeases.expiresAt, now))
	319	)
	320	.returning({ id: agentLeases.id });
	321	return result.length;
	322	} catch {
	323	return 0;
	324	}
	325	}
	326
	327	// ---------------------------------------------------------------------------
	328	// Budgets
	329	// ---------------------------------------------------------------------------
	330
	331	/**
	332	* Atomically attempt to charge `cents` against an agent's daily budget.
	333	* Returns true when the charge was recorded, false when the agent is
	334	* over budget (caller should refuse the request).
	335	*
	336	* We use a single UPDATE with a WHERE clause that includes the budget
	337	* predicate, so we never race against a parallel charger: either the
	338	* row update succeeds (we landed under cap) or it doesn't (over).
	339	*/
	340	export async function chargeAgent(
	341	agentSessionId: string,
	342	cents: number
	343	): Promise<boolean> {
	344	if (!agentSessionId) return false;
	345	const amount = Math.max(0, Math.floor(cents));
	346	if (amount === 0) return true; // free actions still pass.
	347
	348	try {
	349	const result = await db
	350	.update(agentSessions)
	351	.set({
	352	spentCentsToday: sql`${agentSessions.spentCentsToday} + ${amount}`,
	353	})
	354	.where(
	355	and(
	356	eq(agentSessions.id, agentSessionId),
	357	// spent + amount must still fit under cap.
	358	sql`${agentSessions.spentCentsToday} + ${amount} <= ${agentSessions.budgetCentsPerDay}`
	359	)
	360	)
	361	.returning({ id: agentSessions.id });
	362	return result.length > 0;
	363	} catch {
	364	return false;
	365	}
	366	}
	367
	368	/** Convenience getter — returns 0 spent / 0 cap when the session is missing. */
	369	export async function getAgentUsage(agentSessionId: string): Promise<{
	370	spent: number;
	371	cap: number;
	372	remaining: number;
	373	}> {
	374	try {
	375	const [row] = await db
	376	.select({
	377	spent: agentSessions.spentCentsToday,
	378	cap: agentSessions.budgetCentsPerDay,
	379	})
	380	.from(agentSessions)
	381	.where(eq(agentSessions.id, agentSessionId))
	382	.limit(1);
	383	if (!row) return { spent: 0, cap: 0, remaining: 0 };
	384	return {
	385	spent: row.spent,
	386	cap: row.cap,
	387	remaining: Math.max(0, row.cap - row.spent),
	388	};
	389	} catch {
	390	return { spent: 0, cap: 0, remaining: 0 };
	391	}
	392	}
	393
	394	/**
	395	* Reset every agent's `spent_cents_today` to 0. Designed to run at the
	396	* UTC day boundary from the autopilot.
	397	*/
	398	export async function resetDailyBudgets(): Promise<number> {
	399	try {
	400	const result = await db
	401	.update(agentSessions)
	402	.set({ spentCentsToday: 0 })
	403	.returning({ id: agentSessions.id });
	404	return result.length;
	405	} catch {
	406	return 0;
	407	}
	408	}