diff --git a/src/main/database.ts b/src/main/database.ts index 7e2c8c8e..221615ca 100644 --- a/src/main/database.ts +++ b/src/main/database.ts @@ -14,6 +14,8 @@ import type { AgentRecord, ApprovalRecord, BrowserActionRecord, + BudgetIncidentRecord, + BudgetPolicyRecord, CommentRecord, CompanyMetrics, CompanyRecord, @@ -1369,9 +1371,12 @@ export class AppDatabase { ) .run(input.costUsd, ts, run.agentId); } + // Look up task's project_id for cost attribution + const taskProjectRow = this.db.prepare("SELECT project_id FROM tasks WHERE id = ?").get(run.taskId) as { project_id: string | null } | undefined; + const projectId = taskProjectRow?.project_id ?? null; this.db - .prepare("insert into cost_entries (id, company_id, run_id, task_id, agent_id, connector_id, amount_usd, unattributed, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)") - .run(randomUUID(), run.companyId, input.runId, run.taskId, run.agentId, run.connectorId, input.costUsd, input.unattributedCost ? 1 : 0, ts); + .prepare("insert into cost_entries (id, company_id, run_id, task_id, agent_id, connector_id, amount_usd, unattributed, project_id, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)") + .run(randomUUID(), run.companyId, input.runId, run.taskId, run.agentId, run.connectorId, input.costUsd, input.unattributedCost ? 1 : 0, projectId, ts); if (budgetHardStopped) { this.db .prepare("update agents set status = 'paused', updated_at = ? where id = ?") @@ -3793,6 +3798,31 @@ export class AppDatabase { this.db.exec("DROP TABLE IF EXISTS pending_wakes"); }, }, + { + version: 18, + name: "migrate_agent_budget_to_policies", + fn: () => { + const agents = this.readAll<{ id: string; companyId: string; budgetMonthlyUsd: number }>( + "SELECT id, company_id as companyId, budget_monthly_usd as budgetMonthlyUsd FROM agents WHERE budget_monthly_usd > 0", + ); + const ts = nowIso(); + for (const agent of agents) { + const existing = this.db + .prepare( + "SELECT id FROM budget_policies WHERE company_id = ? AND scope_type = 'agent' AND scope_id = ? AND metric = 'billed_cents' AND window_kind = 'calendar_month_utc'", + ) + .get(agent.companyId, agent.id); + if (!existing) { + this.db + .prepare( + "INSERT INTO budget_policies (id, company_id, scope_type, scope_id, metric, window_kind, amount, warn_percent, hard_stop_enabled, notify_enabled, is_active, created_at, updated_at) VALUES (?, ?, 'agent', ?, 'billed_cents', 'calendar_month_utc', ?, 80, 1, 1, 1, ?, ?)", + ) + .run(randomUUID(), agent.companyId, agent.id, Math.round(agent.budgetMonthlyUsd * 100), ts, ts); + } + } + log.info(`Budget migration: converted ${agents.length} agent budget(s) to policy rows`); + }, + }, ]; for (const migration of migrations) { @@ -4308,6 +4338,278 @@ export class AppDatabase { .run(taskType, ts, taskId); } + // --------------------------------------------------------------------------- + // Budget Policy CRUD + // --------------------------------------------------------------------------- + + saveBudgetPolicy(input: { + id?: string | null; + companyId: string; + scopeType: string; + scopeId: string; + metric?: string; + windowKind?: string; + amount: number; + warnPercent?: number; + hardStopEnabled?: boolean; + notifyEnabled?: boolean; + isActive?: boolean; + }): string { + const id = input.id || randomUUID(); + const ts = nowIso(); + this.db + .prepare( + `INSERT OR REPLACE INTO budget_policies + (id, company_id, scope_type, scope_id, metric, window_kind, amount, warn_percent, + hard_stop_enabled, notify_enabled, is_active, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .run( + id, + input.companyId, + input.scopeType, + input.scopeId, + input.metric ?? "billed_cents", + input.windowKind ?? "calendar_month_utc", + input.amount, + input.warnPercent ?? 80, + (input.hardStopEnabled ?? true) ? 1 : 0, + (input.notifyEnabled ?? true) ? 1 : 0, + (input.isActive ?? true) ? 1 : 0, + ts, + ts, + ); + return id; + } + + deleteBudgetPolicy(companyId: string, policyId: string): void { + this.transaction(() => { + this.db.prepare("DELETE FROM budget_incidents WHERE policy_id = ?").run(policyId); + this.db.prepare("DELETE FROM budget_policies WHERE id = ? AND company_id = ?").run(policyId, companyId); + }); + } + + listBudgetPolicies(companyId: string): BudgetPolicyRecord[] { + return this.readAll( + "SELECT * FROM budget_policies WHERE company_id = ? ORDER BY scope_type, created_at", + companyId, + ).map((r) => ({ + id: r.id, + companyId: r.company_id, + scopeType: r.scope_type, + scopeId: r.scope_id, + metric: r.metric, + windowKind: r.window_kind, + amount: r.amount, + warnPercent: r.warn_percent, + hardStopEnabled: !!r.hard_stop_enabled, + notifyEnabled: !!r.notify_enabled, + isActive: !!r.is_active, + createdAt: r.created_at, + updatedAt: r.updated_at, + })); + } + + getBudgetPolicy(policyId: string): BudgetPolicyRecord | null { + const r = this.db.prepare("SELECT * FROM budget_policies WHERE id = ?").get(policyId) as SqliteRow | undefined; + if (!r) return null; + return { + id: r.id, + companyId: r.company_id, + scopeType: r.scope_type, + scopeId: r.scope_id, + metric: r.metric, + windowKind: r.window_kind, + amount: r.amount, + warnPercent: r.warn_percent, + hardStopEnabled: !!r.hard_stop_enabled, + notifyEnabled: !!r.notify_enabled, + isActive: !!r.is_active, + createdAt: r.created_at, + updatedAt: r.updated_at, + }; + } + + getActivePoliciesForRun(companyId: string, agentId: string, projectId: string | null): BudgetPolicyRecord[] { + const sql = projectId + ? `SELECT * FROM budget_policies + WHERE company_id = ? AND is_active = 1 + AND ( + (scope_type = 'company' AND scope_id = ?) + OR (scope_type = 'agent' AND scope_id = ?) + OR (scope_type = 'project' AND scope_id = ?) + )` + : `SELECT * FROM budget_policies + WHERE company_id = ? AND is_active = 1 + AND ( + (scope_type = 'company' AND scope_id = ?) + OR (scope_type = 'agent' AND scope_id = ?) + )`; + const params: SQLInputValue[] = projectId + ? [companyId, companyId, agentId, projectId] + : [companyId, companyId, agentId]; + return this.readAll(sql, ...params).map((r) => ({ + id: r.id, + companyId: r.company_id, + scopeType: r.scope_type, + scopeId: r.scope_id, + metric: r.metric, + windowKind: r.window_kind, + amount: r.amount, + warnPercent: r.warn_percent, + hardStopEnabled: !!r.hard_stop_enabled, + notifyEnabled: !!r.notify_enabled, + isActive: !!r.is_active, + createdAt: r.created_at, + updatedAt: r.updated_at, + })); + } + + getCalendarMonthWindow(): { start: string; end: string } { + const now = new Date(); + const year = now.getUTCFullYear(); + const month = now.getUTCMonth(); + const start = new Date(Date.UTC(year, month, 1)).toISOString(); + const end = new Date(Date.UTC(year, month + 1, 1)).toISOString(); + return { start, end }; + } + + getSpentInWindow(policy: BudgetPolicyRecord): number { + const { start, end } = this.getCalendarMonthWindow(); + let sql: string; + let params: SQLInputValue[]; + if (policy.scopeType === "agent") { + sql = "SELECT COALESCE(SUM(amount_usd), 0) * 100 as total_cents FROM cost_entries WHERE company_id = ? AND created_at >= ? AND created_at < ? AND agent_id = ?"; + params = [policy.companyId, start, end, policy.scopeId]; + } else if (policy.scopeType === "project") { + sql = "SELECT COALESCE(SUM(amount_usd), 0) * 100 as total_cents FROM cost_entries WHERE company_id = ? AND created_at >= ? AND created_at < ? AND project_id = ?"; + params = [policy.companyId, start, end, policy.scopeId]; + } else { + sql = "SELECT COALESCE(SUM(amount_usd), 0) * 100 as total_cents FROM cost_entries WHERE company_id = ? AND created_at >= ? AND created_at < ?"; + params = [policy.companyId, start, end]; + } + const row = this.db.prepare(sql).get(...params) as { total_cents: number } | undefined; + return Math.round(row?.total_cents ?? 0); + } + + // --------------------------------------------------------------------------- + // Budget Incident CRUD + // --------------------------------------------------------------------------- + + upsertBudgetIncident( + policy: BudgetPolicyRecord, + thresholdType: "warn_threshold" | "hard_stop", + amountObserved: number, + ): string { + const { start, end } = this.getCalendarMonthWindow(); + const existing = this.db + .prepare( + "SELECT id FROM budget_incidents WHERE policy_id = ? AND window_start = ? AND threshold_type = ? AND status = 'open'", + ) + .get(policy.id, start, thresholdType) as { id: string } | undefined; + + const ts = nowIso(); + if (existing) { + this.db + .prepare("UPDATE budget_incidents SET amount_observed = ?, updated_at = ? WHERE id = ?") + .run(amountObserved, ts, existing.id); + return existing.id; + } + + const id = randomUUID(); + this.db + .prepare( + `INSERT INTO budget_incidents + (id, company_id, policy_id, scope_type, scope_id, metric, window_kind, + window_start, window_end, threshold_type, amount_limit, amount_observed, + status, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'open', ?, ?)`, + ) + .run( + id, + policy.companyId, + policy.id, + policy.scopeType, + policy.scopeId, + policy.metric, + policy.windowKind, + start, + end, + thresholdType, + policy.amount, + amountObserved, + ts, + ts, + ); + return id; + } + + listBudgetIncidents(companyId: string, statusFilter: string = "all"): BudgetIncidentRecord[] { + const sql = statusFilter === "all" + ? "SELECT * FROM budget_incidents WHERE company_id = ? ORDER BY created_at DESC" + : "SELECT * FROM budget_incidents WHERE company_id = ? AND status = ? ORDER BY created_at DESC"; + const rows = statusFilter === "all" + ? this.readAll(sql, companyId) + : this.readAll(sql, companyId, statusFilter); + return rows.map((r) => ({ + id: r.id, + companyId: r.company_id, + policyId: r.policy_id, + scopeType: r.scope_type, + scopeId: r.scope_id, + metric: r.metric, + windowKind: r.window_kind, + windowStart: r.window_start, + windowEnd: r.window_end, + thresholdType: r.threshold_type, + amountLimit: r.amount_limit, + amountObserved: r.amount_observed, + status: r.status, + resolvedAt: r.resolved_at, + createdAt: r.created_at, + updatedAt: r.updated_at, + })); + } + + resolveBudgetIncident(incidentId: string, companyId: string): void { + const ts = nowIso(); + this.db + .prepare( + "UPDATE budget_incidents SET status = 'resolved', resolved_at = ?, updated_at = ? WHERE id = ? AND company_id = ? AND status = 'open'", + ) + .run(ts, ts, incidentId, companyId); + } + + resolveIncidentsForPolicy(policyId: string, resolutionStatus: string): number { + const ts = nowIso(); + const result = this.db + .prepare( + "UPDATE budget_incidents SET status = ?, resolved_at = ?, updated_at = ? WHERE policy_id = ? AND status = 'open'", + ) + .run(resolutionStatus, ts, ts, policyId); + return Number(result.changes); + } + + resolveWarnIncidentsForPolicy(policyId: string, resolutionStatus: string): number { + const ts = nowIso(); + const result = this.db + .prepare( + "UPDATE budget_incidents SET status = ?, resolved_at = ?, updated_at = ? WHERE policy_id = ? AND status = 'open' AND threshold_type = 'warn_threshold'", + ) + .run(resolutionStatus, ts, ts, policyId); + return Number(result.changes); + } + + autoCloseExpiredIncidents(companyId: string): number { + const ts = nowIso(); + const result = this.db + .prepare( + "UPDATE budget_incidents SET status = 'auto_closed', resolved_at = ?, updated_at = ? WHERE company_id = ? AND status = 'open' AND window_end <= ?", + ) + .run(ts, ts, companyId, ts); + return Number(result.changes); + } + private openConnection(): SqliteDatabase { return new DatabaseSync(this.dbPath, { timeout: 5000 }); } diff --git a/src/preload/index.ts b/src/preload/index.ts index 1e139a5a..2f0a6dbc 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -95,6 +95,12 @@ const api: DesktopApi = { listAgentSkills: (companyId, agentId) => ipcRenderer.invoke(IPC_CHANNELS.listAgentSkills, { companyId, agentId }), listWakeupRequests: (input) => ipcRenderer.invoke(IPC_CHANNELS.listWakeupRequests, input), + saveBudgetPolicy: (input) => ipcRenderer.invoke(IPC_CHANNELS.saveBudgetPolicy, input), + deleteBudgetPolicy: (input) => ipcRenderer.invoke(IPC_CHANNELS.deleteBudgetPolicy, input), + listBudgetPolicies: (companyId) => ipcRenderer.invoke(IPC_CHANNELS.listBudgetPolicies, companyId), + listBudgetIncidents: (input) => ipcRenderer.invoke(IPC_CHANNELS.listBudgetIncidents, input), + resolveBudgetIncident: (input) => ipcRenderer.invoke(IPC_CHANNELS.resolveBudgetIncident, input), + getBudgetStatus: (companyId) => ipcRenderer.invoke(IPC_CHANNELS.getBudgetStatus, companyId), recoveryRetry: (taskId: string, agentId: string, companyId: string) => ipcRenderer.invoke("recovery:retry", { taskId, agentId, companyId }), recoveryDismiss: (taskId: string) => diff --git a/src/shared/contracts.ts b/src/shared/contracts.ts index 7087dddf..0bbe8242 100644 --- a/src/shared/contracts.ts +++ b/src/shared/contracts.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import type { AgentMessageRecord, AgentMetrics, BrowserActionRecord, CommentRecord, CompanyMetrics, DesktopEvent, InboxItem, OnboardingBootstrapResult, ProfileSnapshot, RunLogChunk, SkillImportResult, SkillRecord, SkillUpdateStatus, StandupReport } from "./types"; +import type { AgentMessageRecord, AgentMetrics, BrowserActionRecord, BudgetIncidentRecord, BudgetPolicyRecord, BudgetStatusSummary, CommentRecord, CompanyMetrics, DesktopEvent, InboxItem, OnboardingBootstrapResult, ProfileSnapshot, RunLogChunk, SkillImportResult, SkillRecord, SkillUpdateStatus, StandupReport } from "./types"; export const IPC_CHANNELS = { loadSnapshot: "agent-company:load-snapshot", @@ -83,6 +83,12 @@ export const IPC_CHANNELS = { listSkillAgents: "agent-company:list-skill-agents", listAgentSkills: "agent-company:list-agent-skills", listWakeupRequests: "agent-company:list-wakeup-requests", + saveBudgetPolicy: "agent-company:save-budget-policy", + deleteBudgetPolicy: "agent-company:delete-budget-policy", + listBudgetPolicies: "agent-company:list-budget-policies", + listBudgetIncidents: "agent-company:list-budget-incidents", + resolveBudgetIncident: "agent-company:resolve-budget-incident", + getBudgetStatus: "agent-company:get-budget-status", } as const; const optionalId = z.preprocess( @@ -570,6 +576,43 @@ export const listAgentSkillsSchema = z.object({ agentId: z.string().min(1), }); +export const budgetPolicyInputSchema = z.object({ + id: optionalId, + companyId: z.string().uuid(), + scopeType: z.enum(["company", "agent", "project"]), + scopeId: z.string().uuid(), + metric: z.literal("billed_cents").default("billed_cents"), + windowKind: z.literal("calendar_month_utc").default("calendar_month_utc"), + amount: z.number().int().min(1), + warnPercent: z.number().int().min(1).max(100).default(80), + hardStopEnabled: z.boolean().default(true), + notifyEnabled: z.boolean().default(true), + isActive: z.boolean().default(true), +}); + +export const listBudgetPoliciesSchema = z.object({ + companyId: z.string().uuid(), +}); + +export const deleteBudgetPolicySchema = z.object({ + companyId: z.string().uuid(), + policyId: z.string().uuid(), +}); + +export const listBudgetIncidentsSchema = z.object({ + companyId: z.string().uuid(), + status: z.enum(["open", "resolved", "auto_resolved_limit_increase", "auto_closed", "all"]).default("all"), +}); + +export const resolveBudgetIncidentSchema = z.object({ + companyId: z.string().uuid(), + incidentId: z.string().uuid(), +}); + +export const getBudgetStatusSchema = z.object({ + companyId: z.string().uuid(), +}); + export const wakeupQuerySchema = z.object({ companyId: z.string().uuid(), status: z.enum(["queued", "claimed", "coalesced", "deferred", "completed"]).optional(), @@ -678,6 +721,12 @@ export type DesktopApi = { listSkillAgents: (companyId: string, skillId: string) => Promise>; listAgentSkills: (companyId: string, agentId: string) => Promise>; listWakeupRequests: (input: z.infer) => Promise>; + saveBudgetPolicy: (input: z.infer) => Promise>; + deleteBudgetPolicy: (input: z.infer) => Promise>; + listBudgetPolicies: (companyId: string) => Promise>; + listBudgetIncidents: (input: z.infer) => Promise>; + resolveBudgetIncident: (input: z.infer) => Promise>; + getBudgetStatus: (companyId: string) => Promise>; recoveryRetry: (taskId: string, agentId: string, companyId: string) => Promise>; recoveryDismiss: (taskId: string) => Promise>; recoveryFail: (taskId: string) => Promise>; diff --git a/src/shared/types.ts b/src/shared/types.ts index 6c2d9d53..1792f6ad 100644 --- a/src/shared/types.ts +++ b/src/shared/types.ts @@ -269,6 +269,54 @@ export interface CostRecord { createdAt: string; } +export interface BudgetPolicyRecord { + id: string; + companyId: string; + scopeType: "company" | "agent" | "project"; + scopeId: string; + metric: string; + windowKind: string; + amount: number; // INTEGER cents + warnPercent: number; + hardStopEnabled: boolean; + notifyEnabled: boolean; + isActive: boolean; + createdAt: string; + updatedAt: string; +} + +export interface BudgetIncidentRecord { + id: string; + companyId: string; + policyId: string; + scopeType: string; + scopeId: string; + metric: string; + windowKind: string; + windowStart: string; + windowEnd: string; + thresholdType: "warn_threshold" | "hard_stop"; + amountLimit: number; + amountObserved: number; + status: "open" | "resolved" | "auto_resolved_limit_increase" | "auto_closed"; + resolvedAt: string | null; + createdAt: string; + updatedAt: string; +} + +export interface BudgetStatusSummary { + policies: Array<{ + policyId: string; + scopeType: string; + scopeId: string; + amount: number; + spent: number; + warnPercent: number; + hardStopEnabled: boolean; + hasOpenIncident: boolean; + }>; +} + export interface ActivityRecord { id: string; companyId: string;