aevgarik/paperclip
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix budget auth and monthly spend rollups

Browse Source
This commit is contained in:
Dotta
2026-03-16 15:41:48 -05:00
parent 5f2c2ee0e2
commit 728d9729ed
7 changed files with 315 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
server/src/services/agents.ts
View File

@@ -1,5 +1,5 @@
import { createHash, randomBytes } from "node:crypto";
import { and, desc, eq, inArray, ne } from "drizzle-orm";
import { and, desc, eq, gte, inArray, lt, ne, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import {
agents,
@@ -8,6 +8,7 @@ import {
agentRuntimeState,
agentTaskSessions,
agentWakeupRequests,
costEvents,
heartbeatRunEvents,
heartbeatRuns,
} from "@paperclipai/db";
@@ -182,6 +183,15 @@ export function deduplicateAgentName(
}
export function agentService(db: Db) {
function currentUtcMonthWindow(now = new Date()) {
const year = now.getUTCFullYear();
const month = now.getUTCMonth();
return {
start: new Date(Date.UTC(year, month, 1, 0, 0, 0, 0)),
end: new Date(Date.UTC(year, month + 1, 1, 0, 0, 0, 0)),
};
}
function withUrlKey<T extends { id: string; name: string }>(row: T) {
return {
...row,
@@ -196,13 +206,47 @@ export function agentService(db: Db) {
});
}
async function getMonthlySpendByAgentIds(companyId: string, agentIds: string[]) {
if (agentIds.length === 0) return new Map<string, number>();
const { start, end } = currentUtcMonthWindow();
const rows = await db
.select({
agentId: costEvents.agentId,
spentMonthlyCents: sql<number>`coalesce(sum(${costEvents.costCents}), 0)::int`,
})
.from(costEvents)
.where(
and(
eq(costEvents.companyId, companyId),
inArray(costEvents.agentId, agentIds),
gte(costEvents.occurredAt, start),
lt(costEvents.occurredAt, end),
),
)
.groupBy(costEvents.agentId);
return new Map(rows.map((row) => [row.agentId, Number(row.spentMonthlyCents ?? 0)]));
}
async function hydrateAgentSpend<T extends { id: string; companyId: string; spentMonthlyCents: number }>(rows: T[]) {
const agentIds = rows.map((row) => row.id);
const companyId = rows[0]?.companyId;
if (!companyId || agentIds.length === 0) return rows;
const spendByAgentId = await getMonthlySpendByAgentIds(companyId, agentIds);
return rows.map((row) => ({
...row,
spentMonthlyCents: spendByAgentId.get(row.id) ?? 0,
}));
}
async function getById(id: string) {
const row = await db
.select()
.from(agents)
.where(eq(agents.id, id))
.then((rows) => rows[0] ?? null);
return row ? normalizeAgentRow(row) : null;
if (!row) return null;
const [hydrated] = await hydrateAgentSpend([row]);
return normalizeAgentRow(hydrated);
}
async function ensureManager(companyId: string, managerId: string) {
@@ -331,7 +375,8 @@ export function agentService(db: Db) {
conditions.push(ne(agents.status, "terminated"));
}
const rows = await db.select().from(agents).where(and(...conditions));
return rows.map(normalizeAgentRow);
const hydrated = await hydrateAgentSpend(rows);
return hydrated.map(normalizeAgentRow);
},
getById,

3
server/src/services/budgets.ts
View File

@@ -1,4 +1,4 @@
import { and, desc, eq, gte, inArray, lt, sql } from "drizzle-orm";
import { and, desc, eq, gte, inArray, lt, ne, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import {
agents,
@@ -360,6 +360,7 @@ export function budgetService(db: Db, hooks: BudgetServiceHooks = {}) {
eq(budgetIncidents.policyId, policy.id),
eq(budgetIncidents.windowStart, start),
eq(budgetIncidents.thresholdType, thresholdType),
ne(budgetIncidents.status, "dismissed"),
),
)
.then((rows) => rows[0] ?? null);

75
server/src/services/companies.ts
View File

@@ -1,4 +1,4 @@
import { eq, count } from "drizzle-orm";
import { and, count, eq, gte, inArray, lt, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import {
companies,
@@ -54,6 +54,49 @@ export function companyService(db: Db) {
};
}
function currentUtcMonthWindow(now = new Date()) {
const year = now.getUTCFullYear();
const month = now.getUTCMonth();
return {
start: new Date(Date.UTC(year, month, 1, 0, 0, 0, 0)),
end: new Date(Date.UTC(year, month + 1, 1, 0, 0, 0, 0)),
};
}
async function getMonthlySpendByCompanyIds(
companyIds: string[],
database: Pick<Db, "select"> = db,
) {
if (companyIds.length === 0) return new Map<string, number>();
const { start, end } = currentUtcMonthWindow();
const rows = await database
.select({
companyId: costEvents.companyId,
spentMonthlyCents: sql<number>`coalesce(sum(${costEvents.costCents}), 0)::int`,
})
.from(costEvents)
.where(
and(
inArray(costEvents.companyId, companyIds),
gte(costEvents.occurredAt, start),
lt(costEvents.occurredAt, end),
),
)
.groupBy(costEvents.companyId);
return new Map(rows.map((row) => [row.companyId, Number(row.spentMonthlyCents ?? 0)]));
}
async function hydrateCompanySpend<T extends { id: string; spentMonthlyCents: number }>(
rows: T[],
database: Pick<Db, "select"> = db,
) {
const spendByCompanyId = await getMonthlySpendByCompanyIds(rows.map((row) => row.id), database);
return rows.map((row) => ({
...row,
spentMonthlyCents: spendByCompanyId.get(row.id) ?? 0,
}));
}
function getCompanyQuery(database: Pick<Db, "select">) {
return database
.select(companySelection)
@@ -104,13 +147,20 @@ export function companyService(db: Db) {
}
return {
list: () =>
getCompanyQuery(db).then((rows) => rows.map((row) => enrichCompany(row))),
list: async () => {
const rows = await getCompanyQuery(db);
const hydrated = await hydrateCompanySpend(rows);
return hydrated.map((row) => enrichCompany(row));
},
getById: (id: string) =>
getCompanyQuery(db)
getById: async (id: string) => {
const row = await getCompanyQuery(db)
.where(eq(companies.id, id))
.then((rows) => (rows[0] ? enrichCompany(rows[0]) : null)),
.then((rows) => rows[0] ?? null);
if (!row) return null;
const [hydrated] = await hydrateCompanySpend([row], db);
return enrichCompany(hydrated);
},
create: async (data: typeof companies.$inferInsert) => {
const created = await createCompanyWithUniquePrefix(data);
@@ -118,7 +168,8 @@ export function companyService(db: Db) {
.where(eq(companies.id, created.id))
.then((rows) => rows[0] ?? null);
if (!row) throw notFound("Company not found after creation");
return enrichCompany(row);
const [hydrated] = await hydrateCompanySpend([row], db);
return enrichCompany(hydrated);
},
update: (
@@ -175,10 +226,12 @@ export function companyService(db: Db) {
await tx.delete(assets).where(eq(assets.id, existing.logoAssetId));
}
return enrichCompany({
const [hydrated] = await hydrateCompanySpend([{
...updated,
logoAssetId: logoAssetId === undefined ? existing.logoAssetId : logoAssetId,
});
}], tx);
return enrichCompany(hydrated);
}),
archive: (id: string) =>
@@ -193,7 +246,9 @@ export function companyService(db: Db) {
const row = await getCompanyQuery(tx)
.where(eq(companies.id, id))
.then((rows) => rows[0] ?? null);
return row ? enrichCompany(row) : null;
if (!row) return null;
const [hydrated] = await hydrateCompanySpend([row], tx);
return enrichCompany(hydrated);
}),
remove: (id: string) =>

42
server/src/services/costs.ts
View File

@@ -1,4 +1,4 @@
import { and, desc, eq, gte, isNotNull, lte, sql } from "drizzle-orm";
import { and, desc, eq, gte, isNotNull, lt, lte, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import { activityLog, agents, companies, costEvents, issues, projects } from "@paperclipai/db";
import { notFound, unprocessable } from "../errors.js";
@@ -12,6 +12,37 @@ export interface CostDateRange {
const METERED_BILLING_TYPE = "metered_api";
const SUBSCRIPTION_BILLING_TYPES = ["subscription_included", "subscription_overage"] as const;
function currentUtcMonthWindow(now = new Date()) {
const year = now.getUTCFullYear();
const month = now.getUTCMonth();
return {
start: new Date(Date.UTC(year, month, 1, 0, 0, 0, 0)),
end: new Date(Date.UTC(year, month + 1, 1, 0, 0, 0, 0)),
};
}
async function getMonthlySpendTotal(
db: Db,
scope: { companyId: string; agentId?: string | null },
) {
const { start, end } = currentUtcMonthWindow();
const conditions = [
eq(costEvents.companyId, scope.companyId),
gte(costEvents.occurredAt, start),
lt(costEvents.occurredAt, end),
];
if (scope.agentId) {
conditions.push(eq(costEvents.agentId, scope.agentId));
}
const [row] = await db
.select({
total: sql<number>`coalesce(sum(${costEvents.costCents}), 0)::int`,
})
.from(costEvents)
.where(and(...conditions));
return Number(row?.total ?? 0);
}
export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
const budgets = budgetService(db, budgetHooks);
return {
@@ -39,10 +70,15 @@ export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
.returning()
.then((rows) => rows[0]);
const [agentMonthSpend, companyMonthSpend] = await Promise.all([
getMonthlySpendTotal(db, { companyId, agentId: event.agentId }),
getMonthlySpendTotal(db, { companyId }),
]);
await db
.update(agents)
.set({
spentMonthlyCents: sql`${agents.spentMonthlyCents} + ${event.costCents}`,
spentMonthlyCents: agentMonthSpend,
updatedAt: new Date(),
})
.where(eq(agents.id, event.agentId));
@@ -50,7 +86,7 @@ export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
await db
.update(companies)
.set({
spentMonthlyCents: sql`${companies.spentMonthlyCents} + ${event.costCents}`,
spentMonthlyCents: companyMonthSpend,
updatedAt: new Date(),
})
.where(eq(companies.id, companyId));