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