From a7468eaa3f2166be339eaeeabe33d7b7617c39b2 Mon Sep 17 00:00:00 2001 From: Saravana Achu Mac Date: Fri, 3 Apr 2026 13:48:52 -0700 Subject: [PATCH] feat(cowork-service): add budget policy proxy routes --- .../src/modules/usage/routes.test.ts | 95 +++++++++++ .../src/modules/usage/routes.ts | 154 ++++++++++++++++-- .../cowork-service/src/modules/usage/types.ts | 7 + 3 files changed, 245 insertions(+), 11 deletions(-) diff --git a/services/cowork-service/src/modules/usage/routes.test.ts b/services/cowork-service/src/modules/usage/routes.test.ts index 8d7bdbc8..5b508962 100644 --- a/services/cowork-service/src/modules/usage/routes.test.ts +++ b/services/cowork-service/src/modules/usage/routes.test.ts @@ -108,4 +108,99 @@ describe('usage proxy routes', () => { expect(res.statusCode).toBe(502); }); }); + + describe('GET /api/usage/budget-policy', () => { + it('loads active daily and monthly product policies', async () => { + mockFetch.mockResolvedValue({ + ok: true, + json: async () => [ + { id: 'day_1', period: 'daily', budgetUsd: 5, modelAllowlist: ['gpt-4o-mini'] }, + { id: 'month_1', period: 'monthly', budgetUsd: 50, modelAllowlist: ['gpt-4o-mini'] }, + ], + }); + + const res = await app.inject({ + method: 'GET', + url: '/api/usage/budget-policy', + headers: { authorization: 'Bearer test-token' }, + }); + + expect(res.statusCode).toBe(200); + const body = JSON.parse(res.payload); + expect(body.dailyBudgetUsd).toBe(5); + expect(body.monthlyBudgetUsd).toBe(50); + expect((mockFetch.mock.calls[0][1].headers as Record).authorization).toBe( + 'Bearer test-token' + ); + }); + }); + + describe('PUT /api/usage/budget-policy', () => { + it('creates daily and monthly policies when missing', async () => { + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: async () => [], + }) + .mockResolvedValueOnce({ + ok: true, + json: async () => ({ id: 'day_1', period: 'daily', budgetUsd: 10 }), + }) + .mockResolvedValueOnce({ + ok: true, + json: async () => ({ id: 'month_1', period: 'monthly', budgetUsd: 200 }), + }); + + const res = await app.inject({ + method: 'PUT', + url: '/api/usage/budget-policy', + headers: { authorization: 'Bearer admin-token' }, + payload: { + dailyBudgetUsd: 10, + monthlyBudgetUsd: 200, + modelAllowlist: ['gpt-4o-mini', 'claude-sonnet-4-20250514'], + }, + }); + + expect(res.statusCode).toBe(200); + const body = JSON.parse(res.payload); + expect(body.dailyBudgetUsd).toBe(10); + expect(body.monthlyBudgetUsd).toBe(200); + expect(mockFetch.mock.calls[1][1].method).toBe('POST'); + expect(mockFetch.mock.calls[2][1].method).toBe('POST'); + }); + + it('updates existing policies when present', async () => { + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: async () => [ + { id: 'day_1', period: 'daily', budgetUsd: 5 }, + { id: 'month_1', period: 'monthly', budgetUsd: 50 }, + ], + }) + .mockResolvedValueOnce({ + ok: true, + json: async () => ({ id: 'day_1', period: 'daily', budgetUsd: 15 }), + }) + .mockResolvedValueOnce({ + ok: true, + json: async () => ({ id: 'month_1', period: 'monthly', budgetUsd: 150 }), + }); + + const res = await app.inject({ + method: 'PUT', + url: '/api/usage/budget-policy', + payload: { + dailyBudgetUsd: 15, + monthlyBudgetUsd: 150, + modelAllowlist: ['gpt-4o-mini'], + }, + }); + + expect(res.statusCode).toBe(200); + expect(mockFetch.mock.calls[1][1].method).toBe('PATCH'); + expect(mockFetch.mock.calls[2][1].method).toBe('PATCH'); + }); + }); }); diff --git a/services/cowork-service/src/modules/usage/routes.ts b/services/cowork-service/src/modules/usage/routes.ts index bc4e8008..e2ddc682 100644 --- a/services/cowork-service/src/modules/usage/routes.ts +++ b/services/cowork-service/src/modules/usage/routes.ts @@ -9,7 +9,90 @@ import type { FastifyInstance } from 'fastify'; import { config } from '../../lib/config.js'; import { PRODUCT_ID } from '../../lib/product-config.js'; import { getUserId } from '../../lib/request-context.js'; -import { UsageSummaryQuerySchema, CheckLimitsSchema } from './types.js'; +import { + BudgetPolicyPayloadSchema, + CheckLimitsSchema, + UsageSummaryQuerySchema, + type BudgetPolicyPayload, +} from './types.js'; + +type PlatformBudgetPolicy = { + id: string; + period: 'daily' | 'monthly'; + budgetUsd: number; + modelAllowlist?: string[]; + scopeType: string; + scopeId: string; + status?: string; +}; + +function buildProxyHeaders( + req: import('fastify').FastifyRequest, + contentType = false +): Record { + const headers: Record = { + 'x-product-id': PRODUCT_ID, + 'x-request-id': req.id, + }; + if (contentType) headers['Content-Type'] = 'application/json'; + const auth = req.headers.authorization; + if (typeof auth === 'string') headers.authorization = auth; + return headers; +} + +async function listBudgetPolicies( + platformUrl: string, + req: import('fastify').FastifyRequest +): Promise { + const params = new URLSearchParams({ + scopeType: 'product', + scopeId: PRODUCT_ID, + status: 'active', + limit: '20', + }); + const res = await fetch(`${platformUrl}/ai-budgets/policies?${params.toString()}`, { + headers: buildProxyHeaders(req), + }); + if (!res.ok) { + throw new Error(`Platform returned ${res.status}`); + } + return (await res.json()) as PlatformBudgetPolicy[]; +} + +async function upsertPolicy( + platformUrl: string, + req: import('fastify').FastifyRequest, + existing: PlatformBudgetPolicy | undefined, + period: 'daily' | 'monthly', + payload: BudgetPolicyPayload +): Promise { + const budgetUsd = period === 'daily' ? payload.dailyBudgetUsd : payload.monthlyBudgetUsd; + const body = { + scopeType: 'product', + scopeId: PRODUCT_ID, + name: `Claw Cowork ${period === 'daily' ? 'Daily' : 'Monthly'} Budget`, + period, + budgetUsd, + softThreshold: 0.8, + hardThreshold: 1, + modelAllowlist: payload.modelAllowlist, + }; + + const targetUrl = existing + ? `${platformUrl}/ai-budgets/policies/${encodeURIComponent(existing.id)}` + : `${platformUrl}/ai-budgets/policies`; + const method = existing ? 'PATCH' : 'POST'; + + const res = await fetch(targetUrl, { + method, + headers: buildProxyHeaders(req, true), + body: JSON.stringify(body), + }); + if (!res.ok) { + throw new Error(`Platform returned ${res.status}`); + } + return (await res.json()) as PlatformBudgetPolicy; +} export async function usageRoutes(app: FastifyInstance) { const platformUrl = config.PLATFORM_SERVICE_URL; @@ -29,11 +112,8 @@ export async function usageRoutes(app: FastifyInstance) { const res = await fetch( `${platformUrl}/usage/summary?userId=${encodeURIComponent(userId)}&days=${days}`, { - headers: { - 'x-product-id': PRODUCT_ID, - 'x-request-id': req.id, - }, - }, + headers: buildProxyHeaders(req), + } ); if (!res.ok) { reply.code(res.status); @@ -61,11 +141,7 @@ export async function usageRoutes(app: FastifyInstance) { try { const res = await fetch(`${platformUrl}/usage/check-limits`, { method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'x-product-id': PRODUCT_ID, - 'x-request-id': req.id, - }, + headers: buildProxyHeaders(req, true), body: JSON.stringify({ userId, plan }), }); if (!res.ok) { @@ -79,4 +155,60 @@ export async function usageRoutes(app: FastifyInstance) { return { error: 'Platform-service unavailable' }; } }); + + app.get('/api/usage/budget-policy', async (req, reply) => { + try { + const policies = await listBudgetPolicies(platformUrl, req); + const daily = policies.find(policy => policy.period === 'daily'); + const monthly = policies.find(policy => policy.period === 'monthly'); + + return { + dailyBudgetUsd: daily?.budgetUsd ?? null, + monthlyBudgetUsd: monthly?.budgetUsd ?? null, + modelAllowlist: monthly?.modelAllowlist ?? daily?.modelAllowlist ?? [], + policies, + }; + } catch (err) { + req.log.warn({ err }, 'Failed to load budget policy from platform-service'); + reply.code(502); + return { error: 'Platform-service unavailable' }; + } + }); + + app.put('/api/usage/budget-policy', async (req, reply) => { + const parsed = BudgetPolicyPayloadSchema.safeParse(req.body); + if (!parsed.success) { + reply.code(400); + return { error: 'Invalid body', details: parsed.error.issues }; + } + + try { + const existing = await listBudgetPolicies(platformUrl, req); + const daily = await upsertPolicy( + platformUrl, + req, + existing.find(policy => policy.period === 'daily'), + 'daily', + parsed.data + ); + const monthly = await upsertPolicy( + platformUrl, + req, + existing.find(policy => policy.period === 'monthly'), + 'monthly', + parsed.data + ); + + return { + dailyBudgetUsd: daily.budgetUsd, + monthlyBudgetUsd: monthly.budgetUsd, + modelAllowlist: monthly.modelAllowlist ?? daily.modelAllowlist ?? [], + policies: [daily, monthly], + }; + } catch (err) { + req.log.warn({ err }, 'Failed to save budget policy to platform-service'); + reply.code(502); + return { error: 'Platform-service unavailable' }; + } + }); } diff --git a/services/cowork-service/src/modules/usage/types.ts b/services/cowork-service/src/modules/usage/types.ts index 581da3e6..63bd07c4 100644 --- a/services/cowork-service/src/modules/usage/types.ts +++ b/services/cowork-service/src/modules/usage/types.ts @@ -12,5 +12,12 @@ export const CheckLimitsSchema = z.object({ plan: z.string().default('free'), }); +export const BudgetPolicyPayloadSchema = z.object({ + dailyBudgetUsd: z.number().min(0.01), + monthlyBudgetUsd: z.number().min(0.01), + modelAllowlist: z.array(z.string().min(1)).default([]), +}); + export type UsageSummaryQuery = z.infer; export type CheckLimitsInput = z.infer; +export type BudgetPolicyPayload = z.infer;