feat(platform): Phase 3 — AI Budget & Cost Governance

- Scope expansion: BudgetScopeTypeSchema now includes 'org' + 'workspace'
- Cost dashboard: GET /ai-budgets/costs with groupBy (model/agent/day/scope)
  - Aggregates totalCostUsd, totalTokens, entryCount, breakdown
- Budget rollover: POST /ai-budgets/policies/:id/rollover
  - Computes previous period remaining, creates rollover doc, adjusts budget
  - GET /ai-budgets/policies/:id/rollovers for history
- Enforcement check: POST /ai-budgets/check (pre-flight, no spend recorded)
  - Model allowlist + threshold evaluation, returns verdict + reasons
- New types: CostDashboardQuerySchema, BudgetRolloverSchema
- New repo functions: listAllSpendEntries, createRollover, listRollovers
- New Cosmos container: ai_budget_rollovers
- 1,316 tests passing (9 new)
This commit is contained in:
saravanakumardb1 2026-03-20 03:26:23 -07:00
parent 84dc348687
commit 05acacd400
4 changed files with 478 additions and 1 deletions

View File

@ -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<BudgetRolloverDoc>('ai_budget_rollovers', '/productId');
}
export async function createRollover(doc: BudgetRolloverDoc): Promise<BudgetRolloverDoc> {
return rolloverCollection().create(doc);
}
export async function listRollovers(
productId: string,
policyId: string
): Promise<BudgetRolloverDoc[]> {
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<BudgetSpendEntryDoc[]> {
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;
});
}

View File

@ -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<string, unknown>) => 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');
});
});

View File

@ -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<string, { costUsd: number; tokensUsed: number; count: number }>();
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,
};
});
}

View File

@ -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<typeof ListBudgetPoliciesQuerySchema>;
export type ListBudgetAlertsQuery = z.infer<typeof ListBudgetAlertsQuerySchema>;
// ── 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<typeof CostDashboardQuerySchema>;
// ── 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<typeof BudgetRolloverSchema> & {
_ts?: number;
_etag?: string;
};