| 1 | import type { Env } from "../env.js"; |
| 2 | import type { Account, Directory, Repo } from "../lib/directory.js"; |
| 3 | import { HttpError, now } from "../lib/util.js"; |
| 4 | import { ModelKeys, PROVIDERS, type Provider } from "./keys.js"; |
| 5 | import { MAX_RUNNING_MAX, MAX_RUNNING_MIN } from "./project.js"; |
| 6 | import { RUNTIMES, runtimeLabel, runtimeNeeds, type RuntimeId } from "./runtimes.js"; |
| 7 | |
| 8 | /** |
| 9 | * Organizations' agent policies and model keys (phase 5b): what owners decide for agents in the organization's |
| 10 | * repositories, and whose key pays. Nothing changes for an organization that set nothing, and repositories owned |
| 11 | * by a person have no policy. |
| 12 | * |
| 13 | * Whose key pays is decided when a task starts (a review or coordinator: when its run starts) and kept on the task |
| 14 | * as `key_owner`: the person's id, or the organization's. The model proxy re-checks it on every request, uncached, |
| 15 | * before it decrypts the key (proxyPayer): only a member who is a person (never an agent account) spends the |
| 16 | * organization's key, and only in a repository the organization still owns, under its monthly allowance. |
| 17 | */ |
| 18 | |
| 19 | export type KeysPolicy = "personal" | "either" | "organization"; |
| 20 | export const KEYS_POLICIES: readonly KeysPolicy[] = ["personal", "either", "organization"]; |
| 21 | const RUNTIME_IDS = Object.keys(RUNTIMES) as RuntimeId[]; |
| 22 | const PROVIDER_IDS = Object.keys(PROVIDERS) as Provider[]; |
| 23 | |
| 24 | /** The largest monthly allowance per person on the organization's keys: 10 billion weighted tokens. */ |
| 25 | export const MONTHLY_TOKENS_MAX = 10_000_000_000; |
| 26 | export const MONTHLY_TOKENS_MIN = 10_000; |
| 27 | |
| 28 | export interface AgentPolicy { |
| 29 | keys: KeysPolicy; |
| 30 | runtimes: RuntimeId[]; |
| 31 | providers: Provider[]; |
| 32 | /** Most tasks at once per repository; each project's own setting is clamped to it. */ |
| 33 | maxRunning: number; |
| 34 | terminals: boolean; |
| 35 | /** Whether terminals may run agent tasks (only while terminals are allowed). */ |
| 36 | terminalTasks: boolean; |
| 37 | review: boolean; |
| 38 | /** Most weighted tokens each person may spend on the organization's keys per calendar month (UTC); null for no limit. */ |
| 39 | memberMonthlyTokens: number | null; |
| 40 | } |
| 41 | |
| 42 | export const DEFAULT_POLICY: AgentPolicy = { |
| 43 | keys: "personal", |
| 44 | runtimes: [...RUNTIME_IDS], |
| 45 | providers: [...PROVIDER_IDS], |
| 46 | maxRunning: MAX_RUNNING_MAX, |
| 47 | terminals: true, |
| 48 | terminalTasks: true, |
| 49 | review: true, |
| 50 | memberMonthlyTokens: null, |
| 51 | }; |
| 52 | |
| 53 | export interface StoredPolicy extends AgentPolicy { |
| 54 | updatedBy: string | null; |
| 55 | updatedAt: string | null; |
| 56 | } |
| 57 | |
| 58 | export async function orgPolicy(env: Env, orgId: string): Promise<StoredPolicy> { |
| 59 | const row = await env.DB.prepare("SELECT p.json, a.login AS by_login, p.updated_at FROM org_agent_policies p LEFT JOIN accounts a ON a.id = p.updated_by WHERE p.org_id = ?") |
| 60 | .bind(orgId) |
| 61 | .first<{ json: string; by_login: string | null; updated_at: string }>(); |
| 62 | if (!row) return { ...DEFAULT_POLICY, updatedBy: null, updatedAt: null }; |
| 63 | return { ...DEFAULT_POLICY, ...(JSON.parse(row.json) as Partial<AgentPolicy>), updatedBy: row.by_login, updatedAt: row.updated_at }; |
| 64 | } |
| 65 | |
| 66 | /** The policy of the organization that owns `repo`, or null for a person's repository. */ |
| 67 | export async function repoPolicy(env: Env, repo: Pick<Repo, "owner_id" | "owner_kind">): Promise<{ orgId: string; policy: AgentPolicy } | null> { |
| 68 | if (repo.owner_kind !== "org") return null; |
| 69 | return { orgId: repo.owner_id, policy: await orgPolicy(env, repo.owner_id) }; |
| 70 | } |
| 71 | |
| 72 | export interface PolicyInput { |
| 73 | keys?: unknown; |
| 74 | runtimes?: unknown; |
| 75 | providers?: unknown; |
| 76 | max_running?: unknown; |
| 77 | terminals?: unknown; |
| 78 | terminal_tasks?: unknown; |
| 79 | review?: unknown; |
| 80 | member_monthly_tokens?: unknown; |
| 81 | } |
| 82 | |
| 83 | const list = <T extends string>(v: unknown, all: readonly T[], name: string): T[] => { |
| 84 | if (!Array.isArray(v) || v.some((x) => typeof x !== "string" || !all.includes(x as T))) throw new HttpError(422, `\`${name}\` must be a list of: ${all.join(", ")}`); |
| 85 | return all.filter((x) => v.includes(x)); |
| 86 | }; |
| 87 | |
| 88 | const bool = (v: unknown, name: string): boolean => { |
| 89 | if (typeof v !== "boolean") throw new HttpError(422, `\`${name}\` must be true or false`); |
| 90 | return v; |
| 91 | }; |
| 92 | |
| 93 | /** A policy from the REST body or the settings form; fields left out keep `current`. */ |
| 94 | export function parsePolicy(input: PolicyInput, current: AgentPolicy): AgentPolicy { |
| 95 | const p: AgentPolicy = { ...current }; |
| 96 | if (input.keys !== undefined) { |
| 97 | if (typeof input.keys !== "string" || !KEYS_POLICIES.includes(input.keys as KeysPolicy)) throw new HttpError(422, `\`keys\` must be ${KEYS_POLICIES.join(", ")}`); |
| 98 | p.keys = input.keys as KeysPolicy; |
| 99 | } |
| 100 | if (input.runtimes !== undefined) p.runtimes = list(input.runtimes, RUNTIME_IDS, "runtimes"); |
| 101 | if (input.providers !== undefined) p.providers = list(input.providers, PROVIDER_IDS, "providers"); |
| 102 | if (input.max_running !== undefined) { |
| 103 | const n = Number(input.max_running); |
| 104 | if (input.max_running === "" || input.max_running === null || !Number.isInteger(n) || n < MAX_RUNNING_MIN || n > MAX_RUNNING_MAX) throw new HttpError(422, `\`max_running\` must be a whole number from ${MAX_RUNNING_MIN} to ${MAX_RUNNING_MAX}`); |
| 105 | p.maxRunning = n; |
| 106 | } |
| 107 | if (input.terminals !== undefined) p.terminals = bool(input.terminals, "terminals"); |
| 108 | if (input.terminal_tasks !== undefined) p.terminalTasks = bool(input.terminal_tasks, "terminal_tasks"); |
| 109 | if (input.review !== undefined) p.review = bool(input.review, "review"); |
| 110 | if (input.member_monthly_tokens !== undefined) { |
| 111 | const v = input.member_monthly_tokens; |
| 112 | if (v === null || v === "") p.memberMonthlyTokens = null; |
| 113 | else { |
| 114 | const n = Number(v); |
| 115 | if (!Number.isInteger(n) || n < MONTHLY_TOKENS_MIN || n > MONTHLY_TOKENS_MAX) throw new HttpError(422, `\`member_monthly_tokens\` must be null (no limit) or a whole number from ${MONTHLY_TOKENS_MIN} to ${MONTHLY_TOKENS_MAX}`); |
| 116 | p.memberMonthlyTokens = n; |
| 117 | } |
| 118 | } |
| 119 | if (!p.runtimes.length) throw new HttpError(422, "allow at least one runtime"); |
| 120 | if (!p.providers.length) throw new HttpError(422, "allow at least one provider"); |
| 121 | return p; |
| 122 | } |
| 123 | |
| 124 | export const policyJson = (p: StoredPolicy) => ({ |
| 125 | keys: p.keys, |
| 126 | runtimes: p.runtimes, |
| 127 | providers: p.providers, |
| 128 | max_running: p.maxRunning, |
| 129 | terminals: p.terminals, |
| 130 | terminal_tasks: p.terminalTasks, |
| 131 | review: p.review, |
| 132 | member_monthly_tokens: p.memberMonthlyTokens, |
| 133 | updated_by: p.updatedBy, |
| 134 | updated_at: p.updatedAt, |
| 135 | }); |
| 136 | |
| 137 | /** |
| 138 | * Save an organization's policy (the caller checked it's an owner, or a site admin), and clamp its repositories' |
| 139 | * "how many at once" now: their ProjectDOs keep the organization's limit, which every task start refreshes too. |
| 140 | */ |
| 141 | export async function savePolicy(env: Env, dir: Directory, org: Account, me: Account, p: AgentPolicy, wait: (p: Promise<unknown>) => void): Promise<StoredPolicy> { |
| 142 | const at = now(); |
| 143 | const json: AgentPolicy = { keys: p.keys, runtimes: p.runtimes, providers: p.providers, maxRunning: p.maxRunning, terminals: p.terminals, terminalTasks: p.terminalTasks, review: p.review, memberMonthlyTokens: p.memberMonthlyTokens }; |
| 144 | await env.DB.prepare("INSERT INTO org_agent_policies (org_id, json, updated_by, updated_at) VALUES (?, ?, ?, ?) ON CONFLICT (org_id) DO UPDATE SET json = excluded.json, updated_by = excluded.updated_by, updated_at = excluded.updated_at") |
| 145 | .bind(org.id, JSON.stringify(json), me.id, at) |
| 146 | .run(); |
| 147 | // A site administrator acting for the owners is recorded as such, as on the key audits. |
| 148 | const asSiteAdmin = !!me.site_admin && (await dir.orgRole(org.id, me.id)) !== "owner"; |
| 149 | wait(dir.audit({ actorId: me.id, actorLogin: me.login, action: "org.agent_policy.update", target: org.login, detail: { ...policyJson({ ...json, updatedBy: null, updatedAt: null }), updated_by: undefined, updated_at: undefined, ...(asSiteAdmin ? { site_admin: true } : {}) } })); |
| 150 | wait( |
| 151 | (async () => { |
| 152 | // Every repository the organization owns, a page at a time. |
| 153 | for (let after = ""; ; ) { |
| 154 | const page = (await env.DB.prepare("SELECT id FROM repos WHERE owner_id = ? AND id > ? ORDER BY id LIMIT 500").bind(org.id, after).all<{ id: string }>()).results; |
| 155 | for (const r of page) await setRepoLimit(env, r.id, p.maxRunning); |
| 156 | if (page.length < 500) break; |
| 157 | after = page[page.length - 1]!.id; |
| 158 | } |
| 159 | })(), |
| 160 | ); |
| 161 | return { ...json, updatedBy: me.login, updatedAt: at }; |
| 162 | } |
| 163 | |
| 164 | const setRepoLimit = (env: Env, repoId: string, n: number | null) => |
| 165 | env.PROJECT.get(env.PROJECT.idFromName(repoId)) |
| 166 | .setOrgMaxRunning(n) |
| 167 | .catch((e: unknown) => console.error("clamping a project's tasks at once failed", e)); |
| 168 | |
| 169 | /** After a transfer: the repository takes its new owner's limit on tasks at once, or none when a person owns it. */ |
| 170 | export async function refreshRepoLimit(env: Env, repo: Repo): Promise<void> { |
| 171 | await setRepoLimit(env, repo.id, (await repoPolicy(env, repo))?.policy.maxRunning ?? null); |
| 172 | } |
| 173 | |
| 174 | // ---------------------------------------------------------------- checks when work starts |
| 175 | |
| 176 | const orgName = async (dir: Directory, orgId: string) => (await dir.account(orgId))?.login ?? "the organization"; |
| 177 | |
| 178 | /** Why the policy refuses a runtime, or the providers a run would use; null when it allows them. */ |
| 179 | export function policyRefusal(org: string, p: AgentPolicy, runtime: RuntimeId, providers: Provider[]): string | null { |
| 180 | if (!p.runtimes.includes(runtime)) return `${org}'s agent policy doesn't allow ${RUNTIMES[runtime].name}`; |
| 181 | const off = providers.filter((x) => !p.providers.includes(x)); |
| 182 | if (off.length) return `${org}'s agent policy doesn't allow ${off.map((x) => PROVIDERS[x].name).join(" or ")}`; |
| 183 | return null; |
| 184 | } |
| 185 | |
| 186 | /** Whether `person` may spend `orgId`'s key: an active person (never an agent account) who is a member, admin or owner. */ |
| 187 | export async function spendsOrgKey(dir: Directory, orgId: string, person: Account | null): Promise<boolean> { |
| 188 | if (!person || person.kind !== "user" || person.status !== "active") return false; |
| 189 | return !!(await dir.orgRole(orgId, person.id)); |
| 190 | } |
| 191 | |
| 192 | export interface Payer { |
| 193 | /** Whose key pays: the person's id, or the organization's. */ |
| 194 | keyOwner: string; |
| 195 | providers: Provider[]; |
| 196 | /** The organization's limit on tasks at once, for its repositories. */ |
| 197 | orgMaxRunning: number | null; |
| 198 | } |
| 199 | |
| 200 | /** |
| 201 | * Whose key a run of `person`'s pays, and the providers it gets, under the repository's policy. `fixed` is the key |
| 202 | * owner a task already has (messages, retries and plan approval keep it); left out, it is decided now. Throws an |
| 203 | * HttpError saying what to do when nobody's key may pay. |
| 204 | */ |
| 205 | export async function payerFor(env: Env, dir: Directory, repo: Repo, person: Account, runtime: RuntimeId, via: Provider | null, fixed?: string | null): Promise<Payer> { |
| 206 | const keys = new ModelKeys(env); |
| 207 | const needs = runtimeNeeds(runtime, via); |
| 208 | const rp = await repoPolicy(env, repo); |
| 209 | const personal = async (): Promise<Payer> => { |
| 210 | const have = await keys.providers(person.id); |
| 211 | for (const p of needs) if (!have.has(p)) throw new HttpError(422, `${runtimeLabel(runtime, { via })} needs your ${PROVIDERS[p].name} API key: add it in Settings → Model access`); |
| 212 | let providers = runtime === "custom" ? [...have] : needs; |
| 213 | if (rp && runtime === "custom") providers = providers.filter((p) => rp.policy.providers.includes(p)); |
| 214 | if (rp && runtime === "custom" && !providers.length) throw new HttpError(403, `${repo.owner_login}'s agent policy allows none of the providers you have a key for`); |
| 215 | return { keyOwner: person.id, providers, orgMaxRunning: rp?.policy.maxRunning ?? null }; |
| 216 | }; |
| 217 | if (!rp) { |
| 218 | if (fixed && fixed !== person.id) throw new HttpError(403, "this task was paid for by an organization that no longer owns the repository: start a new task"); |
| 219 | return personal(); |
| 220 | } |
| 221 | const { orgId, policy } = rp; |
| 222 | const refusal = policyRefusal(repo.owner_login, policy, runtime, needs); |
| 223 | if (refusal) throw new HttpError(403, refusal); |
| 224 | const org = async (): Promise<Payer | null> => { |
| 225 | if (!(await spendsOrgKey(dir, orgId, person))) return null; |
| 226 | const have = [...(await keys.providers(orgId))].filter((p) => policy.providers.includes(p)); |
| 227 | const providers = runtime === "custom" ? have : needs.filter((p) => have.includes(p)); |
| 228 | if (!providers.length || (runtime !== "custom" && providers.length < needs.length)) return null; |
| 229 | return { keyOwner: orgId, providers, orgMaxRunning: policy.maxRunning }; |
| 230 | }; |
| 231 | if (fixed) { |
| 232 | // A task keeps whose key pays from its start: re-check that it still may. |
| 233 | if (fixed === person.id) { |
| 234 | if (policy.keys === "organization") throw new HttpError(403, `${repo.owner_login} now pays for agents in its repositories only with its own keys: start a new task`); |
| 235 | return personal(); |
| 236 | } |
| 237 | if (fixed !== orgId) throw new HttpError(403, `this task was paid for by an organization that no longer owns ${repo.owner_login}/${repo.name}: start a new task`); |
| 238 | if (policy.keys === "personal") throw new HttpError(403, `${repo.owner_login} no longer pays for agents with its keys: start a new task with yours`); |
| 239 | const o = await org(); |
| 240 | if (!o) throw new HttpError(403, (await spendsOrgKey(dir, orgId, person)) ? `${repo.owner_login} no longer has the ${needs.map((p) => PROVIDERS[p].name).join(" and ")} key this task uses` : `only members of ${repo.owner_login} can run agents on its keys`); |
| 241 | return o; |
| 242 | } |
| 243 | if (policy.keys === "personal") return personal(); |
| 244 | const o = await org(); |
| 245 | if (o) return o; |
| 246 | if (policy.keys === "organization") { |
| 247 | if (!(await spendsOrgKey(dir, orgId, person))) throw new HttpError(403, `${repo.owner_login} pays for agents in its repositories with its own keys, which only its members can use`); |
| 248 | throw new HttpError(422, `${repo.owner_login} pays for agents with its own keys, and has no ${needs.map((p) => PROVIDERS[p].name).join(" and ")} key: an owner can add it on the organization's Settings tab`); |
| 249 | } |
| 250 | return personal(); |
| 251 | } |
| 252 | |
| 253 | /** The policy's word on web terminals in `repo`: null when allowed. */ |
| 254 | export async function terminalRefusal(env: Env, repo: Repo, opts: { tasks?: boolean; runtime?: RuntimeId; person?: { dir: Directory; account: Account } }): Promise<string | null> { |
| 255 | const rp = await repoPolicy(env, repo); |
| 256 | if (!rp) return null; |
| 257 | const org = repo.owner_login; |
| 258 | if (!rp.policy.terminals) return `${org}'s agent policy doesn't allow web terminals`; |
| 259 | if (opts.tasks && !rp.policy.terminalTasks) return `${org}'s agent policy doesn't allow agent tasks in terminals`; |
| 260 | // A terminal's task runs on the person's own subscription, with its runtime's provider: the policy's runtimes and providers still hold. |
| 261 | if (opts.runtime) { |
| 262 | const refused = policyRefusal(org, rp.policy, opts.runtime, RUNTIMES[opts.runtime].needs); |
| 263 | if (refused) return refused; |
| 264 | } |
| 265 | // Under "only the organization's keys", only its members run agents in its repositories, terminals included. |
| 266 | if (opts.tasks && opts.person && rp.policy.keys === "organization" && !(await spendsOrgKey(opts.person.dir, repo.owner_id, opts.person.account))) |
| 267 | return `${org}'s agent policy lets only its members run agent tasks in its repositories`; |
| 268 | return null; |
| 269 | } |
| 270 | |
| 271 | /** The policy's word on agent review in `repo`: null when allowed. */ |
| 272 | export async function reviewRefusal(env: Env, repo: Repo): Promise<string | null> { |
| 273 | const rp = await repoPolicy(env, repo); |
| 274 | return rp && !rp.policy.review ? `${repo.owner_login}'s agent policy doesn't allow agent review` : null; |
| 275 | } |
| 276 | |
| 277 | // ---------------------------------------------------------------- the model proxy's check, every request |
| 278 | |
| 279 | /** |
| 280 | * Whether a run's request may spend `keyOwner`'s key now. Nothing is cached: a member removed, a repository |
| 281 | * transferred, a policy or key changed applies at the next request. For an organization's key it returns the |
| 282 | * person's monthly allowance, which the proxy books against. |
| 283 | */ |
| 284 | export async function proxyPayer( |
| 285 | env: Env, |
| 286 | dir: Directory, |
| 287 | run: { repo: Repo; personId: string; keyOwner: string; provider: Provider }, |
| 288 | ): Promise<{ org: { id: string; login: string; cap: number | null } | null } | { refuse: string }> { |
| 289 | const { repo, personId, keyOwner, provider } = run; |
| 290 | const rp = await repoPolicy(env, repo); |
| 291 | if (keyOwner === personId) { |
| 292 | if (!rp) return { org: null }; |
| 293 | if (rp.policy.keys === "organization") return { refuse: `${repo.owner_login} now pays for agents in its repositories only with its own keys` }; |
| 294 | if (!rp.policy.providers.includes(provider)) return { refuse: `${repo.owner_login}'s agent policy doesn't allow ${PROVIDERS[provider].name}` }; |
| 295 | return { org: null }; |
| 296 | } |
| 297 | if (!rp || rp.orgId !== keyOwner) return { refuse: `the organization whose key this task uses no longer owns ${repo.owner_login}/${repo.name}` }; |
| 298 | const { policy } = rp; |
| 299 | if (policy.keys === "personal") return { refuse: `${repo.owner_login} no longer pays for agents with its keys` }; |
| 300 | if (!policy.providers.includes(provider)) return { refuse: `${repo.owner_login}'s agent policy doesn't allow ${PROVIDERS[provider].name}` }; |
| 301 | const person = await dir.account(personId); |
| 302 | if (!(await spendsOrgKey(dir, keyOwner, person))) return { refuse: `${person?.login ?? "the person who started this task"} is no longer a member of ${repo.owner_login}, whose key this task uses` }; |
| 303 | return { org: { id: keyOwner, login: repo.owner_login, cap: policy.memberMonthlyTokens } }; |
| 304 | } |
| 305 | |
| 306 | /** The calendar month (UTC) an organization's allowance counts in, as `2026-10`. */ |
| 307 | export const spendMonth = (at = new Date()): string => at.toISOString().slice(0, 7); |
| 308 | |
| 309 | /** |
| 310 | * Book `tokens` (a request's estimate, weighted) against the person's allowance on the organization's keys this |
| 311 | * month, in one statement: false, booking nothing, when it would go over `cap`. Without a cap it always books, so |
| 312 | * the month's count is right if an owner sets one later. |
| 313 | */ |
| 314 | export async function bookSpend(env: Env, orgId: string, personId: string, tokens: number, cap: number | null): Promise<boolean> { |
| 315 | const month = spendMonth(); |
| 316 | if (cap === null) { |
| 317 | await env.DB.prepare("INSERT INTO org_key_spend (org_id, person_id, month, tokens) VALUES (?, ?, ?, ?) ON CONFLICT (org_id, person_id, month) DO UPDATE SET tokens = tokens + excluded.tokens").bind(orgId, personId, month, tokens).run(); |
| 318 | return true; |
| 319 | } |
| 320 | const r = await env.DB.prepare( |
| 321 | "INSERT INTO org_key_spend (org_id, person_id, month, tokens) SELECT ?1, ?2, ?3, ?4 WHERE ?4 <= ?5 AND ?5 > 0 ON CONFLICT (org_id, person_id, month) DO UPDATE SET tokens = tokens + excluded.tokens WHERE org_key_spend.tokens < ?5 AND org_key_spend.tokens + excluded.tokens <= ?5", |
| 322 | ) |
| 323 | .bind(orgId, personId, month, tokens, cap) |
| 324 | .run(); |
| 325 | return (r.meta.changes ?? 0) > 0; |
| 326 | } |
| 327 | |
| 328 | /** Settle a booking: `delta` is what the answer used minus what was booked (negative for a refund). */ |
| 329 | export async function settleSpend(env: Env, orgId: string, personId: string, delta: number, month = spendMonth()): Promise<void> { |
| 330 | if (!delta) return; |
| 331 | await env.DB.prepare("UPDATE org_key_spend SET tokens = MAX(tokens + ?, 0) WHERE org_id = ? AND person_id = ? AND month = ?").bind(delta, orgId, personId, month).run(); |
| 332 | } |
| 333 | |
| 334 | export interface OrgSpend { |
| 335 | personId: string; |
| 336 | login: string | null; |
| 337 | tokens: number; |
| 338 | } |
| 339 | |
| 340 | /** Each person's tokens on the organization's keys this month, most first. */ |
| 341 | export async function orgSpend(env: Env, orgId: string, month = spendMonth()): Promise<OrgSpend[]> { |
| 342 | const r = await env.DB.prepare("SELECT s.person_id, a.login, s.tokens FROM org_key_spend s LEFT JOIN accounts a ON a.id = s.person_id WHERE s.org_id = ? AND s.month = ? ORDER BY s.tokens DESC LIMIT 500") |
| 343 | .bind(orgId, month) |
| 344 | .all<{ person_id: string; login: string | null; tokens: number }>(); |
| 345 | return r.results.map((x) => ({ personId: x.person_id, login: x.login, tokens: x.tokens })); |
| 346 | } |