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