Skip to content
125 linesCodeBlameRaw
1/**
2 * The workspace's say over all its agents together
3 * (docs.g1t.sh/guides/agent-budgets/): one monthly budget across every
4 * agent, the budget a new agent starts with, the cap a session starts with,
5 * the budget each person's agents share, and alerts at 75, 90 and 100% of
6 * the monthly budget. Owners set it. Pure
7 * apart from the statements it builds, so the rules are tested on their
8 * own.
9 */
10import type { AgentPolicy } from "@g1t/contracts";
11
12import { dayKey, monthKey } from "./budget.ts";
13
14/** A session's cap when nobody set one: $2. */
15export const DEFAULT_SESSION_MICROS = 2_000_000;
16/** The thresholds owners hear about, in % of the workspace's agent budget. */
17export const ALERTS = [75, 90, 100] as const;
18
19export type PolicyRow = AgentPolicy & {
20 /** Every agent's spend this month. */
21 spent: number;
22 /** The highest alert already sent this month, or 0. */
23 alerted: number;
24};
25
26export const DEFAULT_POLICY: AgentPolicy = {
27 monthly_micros: null,
28 default_agent_monthly_micros: null,
29 default_session_micros: DEFAULT_SESSION_MICROS,
30 person_monthly_micros: null,
31 members_create_agents: true,
32};
33
34/** The workspace's policy and this month's spend across its agents. */
35export async function readPolicy(db: D1Database, workspaceId: string, month: string): Promise<PolicyRow> {
36 const [policy, spend] = await Promise.all([
37 db
38 .prepare("SELECT monthly_micros, default_agent_monthly_micros, default_session_micros, person_monthly_micros, members_create_agents FROM agent_policies WHERE workspace_id = ?")
39 .bind(workspaceId)
40 .first<Omit<AgentPolicy, "members_create_agents"> & { members_create_agents: number | null }>(),
41 db
42 .prepare("SELECT micros, alerted FROM workspace_agent_spend WHERE workspace_id = ? AND period = ?")
43 .bind(workspaceId, month)
44 .first<{ micros: number; alerted: number }>(),
45 ]);
46 return {
47 monthly_micros: positive(policy?.monthly_micros),
48 default_agent_monthly_micros: positive(policy?.default_agent_monthly_micros),
49 default_session_micros: positive(policy?.default_session_micros) ?? DEFAULT_SESSION_MICROS,
50 person_monthly_micros: positive(policy?.person_monthly_micros),
51 // On unless an owner turned it off: no row, or a row from before the setting, is on.
52 members_create_agents: policy?.members_create_agents !== 0,
53 spent: spend?.micros ?? 0,
54 alerted: spend?.alerted ?? 0,
55 };
56}
57
58function positive(value: number | null | undefined): number | null {
59 return typeof value === "number" && Number.isFinite(value) && value > 0 ? Math.floor(value) : null;
60}
61
62/** Why no agent may start more work this month, as people are told, or null. */
63export function policyBlock(policy: Pick<PolicyRow, "monthly_micros" | "spent">): string | null {
64 if (policy.monthly_micros && policy.spent >= policy.monthly_micros) {
65 return "This workspace's agents have used their budget for the month. An owner can raise it under Agents → Budget.";
66 }
67 return null;
68}
69
70/** The alert a workspace has newly crossed, or null: the highest threshold at or under its share, past what was sent. */
71export function alertDue(policy: Pick<PolicyRow, "monthly_micros" | "spent" | "alerted">): number | null {
72 if (!policy.monthly_micros) return null;
73 const share = (policy.spent * 100) / policy.monthly_micros;
74 const crossed = ALERTS.filter((level) => share >= level).at(-1) ?? null;
75 return crossed && crossed > policy.alerted ? crossed : null;
76}
77
78/** Adds a charge to the workspace's agent spend, this month and today. */
79export function workspaceSpendStatements(db: D1Database, workspaceId: string, micros: number, now: Date): D1PreparedStatement[] {
80 return [monthKey(now), dayKey(now)].map((period) =>
81 db
82 .prepare(
83 `INSERT INTO workspace_agent_spend (workspace_id, period, micros) VALUES (?1, ?2, ?3)
84 ON CONFLICT (workspace_id, period) DO UPDATE SET micros = micros + ?3`,
85 )
86 .bind(workspaceId, period, Math.max(0, Math.ceil(micros))),
87 );
88}
89
90/** Marks an alert as sent, once: true when this call is the one that sent it. */
91export async function markAlerted(db: D1Database, workspaceId: string, month: string, level: number): Promise<boolean> {
92 const changed = await db
93 .prepare("UPDATE workspace_agent_spend SET alerted = ? WHERE workspace_id = ? AND period = ? AND alerted < ?")
94 .bind(level, workspaceId, month, level)
95 .run();
96 return changed.meta.changes > 0;
97}
98
99/**
100 * A policy as owners change it, checked: budgets are whole micro-dollars
101 * or null (none); a session's default cap is between 10 cents and $500.
102 */
103export function checkPolicy(current: AgentPolicy, changes: Partial<AgentPolicy>): { ok: true; value: AgentPolicy } | { ok: false; message: string } {
104 const next = { ...current };
105 for (const key of ["monthly_micros", "default_agent_monthly_micros", "person_monthly_micros"] as const) {
106 if (!(key in changes)) continue;
107 const value = changes[key];
108 if (value === null || value === 0) next[key] = null;
109 else if (typeof value === "number" && Number.isFinite(value) && value > 0 && value <= 1_000_000_000_000) next[key] = Math.floor(value);
110 else return { ok: false, message: "A budget is a positive amount, or none." };
111 }
112 if ("default_session_micros" in changes) {
113 const value = changes.default_session_micros;
114 if (typeof value !== "number" || !Number.isFinite(value) || value < 100_000 || value > 500_000_000) {
115 return { ok: false, message: "A session's cap is between $0.10 and $500." };
116 }
117 next.default_session_micros = Math.floor(value);
118 }
119 if ("members_create_agents" in changes) {
120 if (typeof changes.members_create_agents !== "boolean") return { ok: false, message: "Whether members create personal agents is true or false." };
121 next.members_create_agents = changes.members_create_agents;
122 }
123 return { ok: true, value: next };
124}
125