diff --git a/services/platform-service/src/modules/ai-budgets/repository.ts b/services/platform-service/src/modules/ai-budgets/repository.ts index 5e44f955..61b113ee 100644 --- a/services/platform-service/src/modules/ai-budgets/repository.ts +++ b/services/platform-service/src/modules/ai-budgets/repository.ts @@ -3,6 +3,7 @@ import { getCollection } from '../../lib/datastore.js'; import type { BudgetAlertDoc, BudgetPolicyDoc, + BudgetRolloverDoc, BudgetSpendEntryDoc, ListBudgetAlertsQuery, ListBudgetPoliciesQuery, @@ -98,3 +99,54 @@ export async function listAlerts( limit: query.limit, }); } + +// ── Budget Rollover ───────────────────────────────────────── + +function rolloverCollection() { + return getCollection('ai_budget_rollovers', '/productId'); +} + +export async function createRollover(doc: BudgetRolloverDoc): Promise { + return rolloverCollection().create(doc); +} + +export async function listRollovers( + productId: string, + policyId: string +): Promise { + return rolloverCollection().findMany({ + filter: { productId, policyId }, + sort: { createdAt: -1 }, + limit: 50, + }); +} + +// ── Cost Dashboard ────────────────────────────────────────── + +export async function listAllSpendEntries( + productId: string, + options: { + scopeType?: string; + scopeId?: string; + since?: string; + until?: string; + limit?: number; + } = {} +): Promise { + const entries = await spendCollection().findMany({ + filter: { + productId, + ...(options.scopeType ? { scopeType: options.scopeType } : {}), + ...(options.scopeId ? { scopeId: options.scopeId } : {}), + }, + sort: { recordedAt: -1 }, + limit: options.limit ?? 500, + }); + + // Filter by time range in-memory + return entries.filter(e => { + if (options.since && e.recordedAt < options.since) return false; + if (options.until && e.recordedAt > options.until) return false; + return true; + }); +} diff --git a/services/platform-service/src/modules/ai-budgets/routes.test.ts b/services/platform-service/src/modules/ai-budgets/routes.test.ts index 8b1c0720..03bb370b 100644 --- a/services/platform-service/src/modules/ai-budgets/routes.test.ts +++ b/services/platform-service/src/modules/ai-budgets/routes.test.ts @@ -10,6 +10,9 @@ const repoMock = { listSpendEntries: vi.fn(), createAlerts: vi.fn(), listAlerts: vi.fn(), + listAllSpendEntries: vi.fn(), + createRollover: vi.fn(), + listRollovers: vi.fn(), }; vi.mock('./repository.js', () => repoMock); @@ -130,4 +133,198 @@ describe('aiBudgetRoutes', () => { }) ); }); + + // ── Scope expansion ───────────────────────────────────── + + it('POST /ai-budgets/policies supports org scope', async () => { + repoMock.createPolicy.mockResolvedValue({ id: 'aibudget_2', scopeType: 'org' }); + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + + const res = await app.inject({ + method: 'POST', + url: '/api/ai-budgets/policies', + payload: { scopeType: 'org', scopeId: 'org_1', name: 'Org Budget', budgetUsd: 500 }, + }); + + expect(res.statusCode).toBe(200); + expect(repoMock.createPolicy).toHaveBeenCalledWith( + expect.objectContaining({ scopeType: 'org', scopeId: 'org_1' }) + ); + }); + + it('POST /ai-budgets/policies supports workspace scope', async () => { + repoMock.createPolicy.mockResolvedValue({ id: 'aibudget_3', scopeType: 'workspace' }); + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + + const res = await app.inject({ + method: 'POST', + url: '/api/ai-budgets/policies', + payload: { scopeType: 'workspace', scopeId: 'ws_1', name: 'WS Budget', budgetUsd: 100 }, + }); + + expect(res.statusCode).toBe(200); + expect(repoMock.createPolicy).toHaveBeenCalledWith( + expect.objectContaining({ scopeType: 'workspace' }) + ); + }); + + // ── Cost Dashboard ────────────────────────────────────── + + it('GET /ai-budgets/costs returns grouped cost breakdown', async () => { + repoMock.listAllSpendEntries.mockResolvedValue([ + { costUsd: 5, tokensUsed: 100, model: 'gpt-4o', recordedAt: '2026-03-15T00:00:00Z' }, + { costUsd: 3, tokensUsed: 80, model: 'gpt-4o', recordedAt: '2026-03-15T01:00:00Z' }, + { costUsd: 1, tokensUsed: 50, model: 'gpt-4o-mini', recordedAt: '2026-03-15T02:00:00Z' }, + ]); + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + + const res = await app.inject({ method: 'GET', url: '/api/ai-budgets/costs?groupBy=model' }); + expect(res.statusCode).toBe(200); + const body = res.json(); + expect(body.totalCostUsd).toBe(9); + expect(body.totalTokens).toBe(230); + expect(body.breakdown).toHaveLength(2); + expect(body.breakdown[0].key).toBe('gpt-4o'); + }); + + it('GET /ai-budgets/costs groups by day', async () => { + repoMock.listAllSpendEntries.mockResolvedValue([ + { costUsd: 5, tokensUsed: 100, recordedAt: '2026-03-15T00:00:00Z' }, + { costUsd: 3, tokensUsed: 80, recordedAt: '2026-03-16T01:00:00Z' }, + ]); + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + + const res = await app.inject({ method: 'GET', url: '/api/ai-budgets/costs?groupBy=day' }); + expect(res.statusCode).toBe(200); + expect(res.json().breakdown).toHaveLength(2); + }); + + // ── Budget Rollover ───────────────────────────────────── + + it('POST /ai-budgets/policies/:id/rollover creates rollover and adjusts budget', async () => { + repoMock.getPolicy.mockResolvedValue({ + id: 'aibudget_1', + productId: 'lysnrai', + period: 'monthly', + budgetUsd: 100, + metadata: {}, + }); + repoMock.listSpendEntries.mockResolvedValue([ + { costUsd: 60, recordedAt: '2026-02-10T00:00:00Z' }, + ]); + repoMock.createRollover.mockImplementation(async (doc: Record) => doc); + repoMock.updatePolicy.mockResolvedValue({ budgetUsd: 140 }); + + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + const res = await app.inject({ + method: 'POST', + url: '/api/ai-budgets/policies/aibudget_1/rollover', + }); + + expect(res.statusCode).toBe(200); + expect(repoMock.createRollover).toHaveBeenCalled(); + expect(repoMock.updatePolicy).toHaveBeenCalledWith( + 'aibudget_1', + 'lysnrai', + expect.objectContaining({ budgetUsd: expect.any(Number) }) + ); + }); + + it('GET /ai-budgets/policies/:id/rollovers returns rollover history', async () => { + repoMock.listRollovers.mockResolvedValue([ + { id: 'airollover_1', policyId: 'aibudget_1', rolledOverUsd: 40 }, + ]); + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + + const res = await app.inject({ + method: 'GET', + url: '/api/ai-budgets/policies/aibudget_1/rollovers', + }); + expect(res.statusCode).toBe(200); + expect(res.json()).toHaveLength(1); + }); + + // ── Enforcement Check ─────────────────────────────────── + + it('POST /ai-budgets/check returns allow when under budget', async () => { + repoMock.listPolicies.mockResolvedValue([ + { + id: 'aibudget_1', + status: 'active', + period: 'monthly', + budgetUsd: 100, + softThreshold: 0.8, + hardThreshold: 1, + modelAllowlist: [], + name: 'Test', + }, + ]); + repoMock.listSpendEntries.mockResolvedValue([ + { costUsd: 20, recordedAt: '2026-03-15T00:00:00Z' }, + ]); + + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + const res = await app.inject({ + method: 'POST', + url: '/api/ai-budgets/check', + payload: { scopeType: 'agent', scopeId: 'agt_1', estimatedCostUsd: 5 }, + }); + + expect(res.statusCode).toBe(200); + expect(res.json().verdict).toBe('allow'); + }); + + it('POST /ai-budgets/check returns block when model not allowed', async () => { + repoMock.listPolicies.mockResolvedValue([ + { + id: 'aibudget_1', + status: 'active', + period: 'monthly', + budgetUsd: 100, + softThreshold: 0.8, + hardThreshold: 1, + modelAllowlist: ['gpt-4o-mini'], + name: 'Test', + }, + ]); + + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + const res = await app.inject({ + method: 'POST', + url: '/api/ai-budgets/check', + payload: { scopeType: 'agent', scopeId: 'agt_1', estimatedCostUsd: 5, model: 'gpt-4.1' }, + }); + + expect(res.statusCode).toBe(200); + expect(res.json().verdict).toBe('block'); + expect(res.json().reasons[0]).toContain('not in allowlist'); + }); + + it('POST /ai-budgets/check returns warn when approaching soft threshold', async () => { + repoMock.listPolicies.mockResolvedValue([ + { + id: 'aibudget_1', + status: 'active', + period: 'monthly', + budgetUsd: 100, + softThreshold: 0.8, + hardThreshold: 1, + modelAllowlist: [], + name: 'Test', + }, + ]); + repoMock.listSpendEntries.mockResolvedValue([ + { costUsd: 75, recordedAt: '2026-03-15T00:00:00Z' }, + ]); + + const app = await buildApp({ sub: 'admin_1', productId: 'lysnrai', role: 'admin' }); + const res = await app.inject({ + method: 'POST', + url: '/api/ai-budgets/check', + payload: { scopeType: 'agent', scopeId: 'agt_1', estimatedCostUsd: 10 }, + }); + + expect(res.statusCode).toBe(200); + expect(res.json().verdict).toBe('warn'); + }); }); diff --git a/services/platform-service/src/modules/ai-budgets/routes.ts b/services/platform-service/src/modules/ai-budgets/routes.ts index 439e9c09..eeeffd83 100644 --- a/services/platform-service/src/modules/ai-budgets/routes.ts +++ b/services/platform-service/src/modules/ai-budgets/routes.ts @@ -4,7 +4,9 @@ import { BadRequestError, ForbiddenError } from '../../lib/errors.js'; import { BudgetAlertDoc, BudgetPolicyDoc, + BudgetRolloverDoc, BudgetSpendEntryDoc, + CostDashboardQuerySchema, CreateBudgetPolicySchema, ListBudgetAlertsQuerySchema, ListBudgetPoliciesQuerySchema, @@ -248,4 +250,197 @@ export async function aiBudgetRoutes(app: FastifyInstance) { alerts, }; }); + + // ── Cost Dashboard ────────────────────────────────────── + + app.get('/ai-budgets/costs', async req => { + const access = requireAdmin(req); + const parsed = CostDashboardQuerySchema.safeParse(req.query); + if (!parsed.success) { + validationError(parsed.error.issues.map(issue => issue.message).join('; ')); + } + + const entries = await repo.listAllSpendEntries(access.productId, { + scopeType: parsed.data.scopeType, + scopeId: parsed.data.scopeId, + since: parsed.data.since, + until: parsed.data.until, + limit: parsed.data.limit, + }); + + const totalCostUsd = entries.reduce((sum, e) => sum + e.costUsd, 0); + const totalTokens = entries.reduce((sum, e) => sum + e.tokensUsed, 0); + + // Group by requested dimension + const groups = new Map(); + for (const entry of entries) { + let key: string; + switch (parsed.data.groupBy) { + case 'model': + key = entry.model ?? 'unknown'; + break; + case 'agent': + key = entry.agentId ?? 'unknown'; + break; + case 'day': + key = entry.recordedAt.slice(0, 10); + break; + case 'scope': + key = `${entry.scopeType}:${entry.scopeId}`; + break; + default: + key = 'all'; + } + const existing = groups.get(key) ?? { costUsd: 0, tokensUsed: 0, count: 0 }; + existing.costUsd += entry.costUsd; + existing.tokensUsed += entry.tokensUsed; + existing.count++; + groups.set(key, existing); + } + + const breakdown = Array.from(groups.entries()) + .map(([key, data]) => ({ key, ...data })) + .sort((a, b) => b.costUsd - a.costUsd); + + return { + totalCostUsd: Math.round(totalCostUsd * 10000) / 10000, + totalTokens, + entryCount: entries.length, + groupBy: parsed.data.groupBy, + breakdown, + }; + }); + + // ── Budget Rollover ───────────────────────────────────── + + app.post('/ai-budgets/policies/:id/rollover', async req => { + const access = requireAdmin(req); + const { id } = req.params as { id: string }; + const policy = await repo.getPolicy(id, access.productId); + + const now = new Date(); + const currentPeriodStart = toIsoDayBoundary(periodStart(now, policy.period)); + + // Calculate previous period + const prev = new Date(now); + if (policy.period === 'monthly') { + prev.setUTCMonth(prev.getUTCMonth() - 1); + } else { + prev.setUTCDate(prev.getUTCDate() - 1); + } + const prevPeriodStart = toIsoDayBoundary(periodStart(prev, policy.period)); + + // Get spend from previous period + const prevEntries = await repo.listSpendEntries(access.productId, { + policyId: id, + since: prevPeriodStart, + }); + // Only count entries in previous period (before current period) + const prevPeriodEntries = prevEntries.filter(e => e.recordedAt < currentPeriodStart); + const spentUsd = prevPeriodEntries.reduce((sum, e) => sum + e.costUsd, 0); + const remainingUsd = Math.max(0, policy.budgetUsd - spentUsd); + const rolledOverUsd = remainingUsd; // Full rollover — could add a cap later + + const rollover: BudgetRolloverDoc = { + id: `airollover_${randomUUID()}`, + productId: access.productId, + policyId: id, + fromPeriod: prevPeriodStart, + toPeriod: currentPeriodStart, + spentUsd, + budgetUsd: policy.budgetUsd, + remainingUsd, + rolledOverUsd, + createdAt: now.toISOString(), + }; + + const saved = await repo.createRollover(rollover); + + // Temporarily increase budget for current period + await repo.updatePolicy(id, access.productId, { + budgetUsd: policy.budgetUsd + rolledOverUsd, + metadata: { + ...(policy.metadata ?? {}), + lastRollover: saved.id, + originalBudgetUsd: policy.budgetUsd, + }, + }); + + return { + rollover: saved, + newBudgetUsd: policy.budgetUsd + rolledOverUsd, + }; + }); + + app.get('/ai-budgets/policies/:id/rollovers', async req => { + const access = requireAdmin(req); + const { id } = req.params as { id: string }; + return repo.listRollovers(access.productId, id); + }); + + // ── Budget Enforcement Check ──────────────────────────── + // Lightweight endpoint for pre-flight budget checks (no spend recorded) + + app.post('/ai-budgets/check', async req => { + const access = requireAdmin(req); + const body = req.body as { + scopeType: string; + scopeId: string; + model?: string; + estimatedCostUsd: number; + }; + + if (!body.scopeType || !body.scopeId || body.estimatedCostUsd === undefined) { + validationError('scopeType, scopeId, and estimatedCostUsd are required'); + } + + const policies = await repo.listPolicies(access.productId, { + scopeType: body.scopeType as 'product' | 'agent' | 'org' | 'workspace', + scopeId: body.scopeId, + status: 'active', + limit: 100, + }); + + const now = new Date(); + let verdict: 'allow' | 'warn' | 'block' = 'allow'; + const reasons: string[] = []; + + for (const policy of policies) { + // Model allowlist check + if (body.model && policy.modelAllowlist.length > 0) { + if (!policy.modelAllowlist.includes(body.model)) { + verdict = 'block'; + reasons.push(`Model '${body.model}' not in allowlist for '${policy.name}'`); + continue; + } + } + + const since = toIsoDayBoundary(periodStart(now, policy.period)); + const entries = await repo.listSpendEntries(access.productId, { + policyId: policy.id, + since, + }); + const currentSpend = entries.reduce((sum, e) => sum + e.costUsd, 0); + const projectedSpend = currentSpend + body.estimatedCostUsd; + const percentUsed = policy.budgetUsd > 0 ? projectedSpend / policy.budgetUsd : 0; + + if (percentUsed >= policy.hardThreshold) { + verdict = 'block'; + reasons.push( + `Would exceed hard threshold (${Math.round(percentUsed * 100)}%) for '${policy.name}'` + ); + } else if (percentUsed >= policy.softThreshold && verdict !== 'block') { + verdict = 'warn'; + reasons.push( + `Would exceed soft threshold (${Math.round(percentUsed * 100)}%) for '${policy.name}'` + ); + } + } + + return { + verdict, + reasons, + policiesEvaluated: policies.length, + }; + }); } diff --git a/services/platform-service/src/modules/ai-budgets/types.ts b/services/platform-service/src/modules/ai-budgets/types.ts index 6fa77478..c0490ee6 100644 --- a/services/platform-service/src/modules/ai-budgets/types.ts +++ b/services/platform-service/src/modules/ai-budgets/types.ts @@ -1,6 +1,6 @@ import { z } from 'zod'; -export const BudgetScopeTypeSchema = z.enum(['product', 'agent']); +export const BudgetScopeTypeSchema = z.enum(['product', 'agent', 'org', 'workspace']); export const BudgetPolicyStatusSchema = z.enum(['active', 'paused', 'archived']); export const BudgetPeriodSchema = z.enum(['daily', 'monthly']); export const BudgetAlertSeveritySchema = z.enum(['warn', 'block']); @@ -122,3 +122,36 @@ export const ListBudgetAlertsQuerySchema = z.object({ export type ListBudgetPoliciesQuery = z.infer; export type ListBudgetAlertsQuery = z.infer; + +// ── Cost Dashboard ────────────────────────────────────────── + +export const CostDashboardQuerySchema = z.object({ + scopeType: BudgetScopeTypeSchema.optional(), + scopeId: z.string().optional(), + since: z.string().optional(), + until: z.string().optional(), + groupBy: z.enum(['model', 'agent', 'day', 'scope']).default('model'), + limit: z.coerce.number().min(1).max(500).default(100), +}); + +export type CostDashboardQuery = z.infer; + +// ── Budget Rollover ───────────────────────────────────────── + +export const BudgetRolloverSchema = z.object({ + id: z.string().min(1), + productId: z.string().min(1), + policyId: z.string().min(1), + fromPeriod: z.string().min(1), + toPeriod: z.string().min(1), + spentUsd: z.number().min(0), + budgetUsd: z.number().min(0), + remainingUsd: z.number(), + rolledOverUsd: z.number().min(0), + createdAt: z.string(), +}); + +export type BudgetRolloverDoc = z.infer & { + _ts?: number; + _etag?: string; +};