diff --git a/server/src/__tests__/heartbeat-ledger-billing-code.test.ts b/server/src/__tests__/heartbeat-ledger-billing-code.test.ts new file mode 100644 index 0000000000..9b9a9fc4c7 --- /dev/null +++ b/server/src/__tests__/heartbeat-ledger-billing-code.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it, vi } from "vitest"; +import type { heartbeatRuns } from "@paperclipai/db"; +import { resolveLedgerScopeForRun } from "../services/heartbeat.ts"; + +type IssueRow = { id: string; projectId: string | null; billingCode: string | null }; + +/** + * Minimal Drizzle stand-in for the single lookup `resolveLedgerScopeForRun` + * performs: `db.select({...}).from(issues).where(...)` resolved as a thenable. + */ +type LedgerDb = Parameters[0]; + +function makeDb(rows: IssueRow[]) { + const where = vi.fn(() => Promise.resolve(rows)); + const from = vi.fn(() => ({ where })); + const select = vi.fn(() => ({ from })); + return { db: { select } as unknown as LedgerDb, select, from, where }; +} + +function makeRun(contextSnapshot: Record) { + return { id: "run-1", contextSnapshot } as unknown as typeof heartbeatRuns.$inferSelect; +} + +describe("resolveLedgerScopeForRun billing code propagation", () => { + it("carries the issue's billing code onto the ledger scope", async () => { + const { db } = makeDb([{ id: "issue-1", projectId: "project-1", billingCode: "ACME-42" }]); + + const scope = await resolveLedgerScopeForRun(db, "company-1", makeRun({ + issueId: "issue-1", + projectId: "context-project", + })); + + expect(scope).toEqual({ + issueId: "issue-1", + projectId: "project-1", + billingCode: "ACME-42", + }); + }); + + it("resolves a null billing code when the issue has none set", async () => { + const { db } = makeDb([{ id: "issue-1", projectId: "project-1", billingCode: null }]); + + const scope = await resolveLedgerScopeForRun(db, "company-1", makeRun({ + issueId: "issue-1", + projectId: "context-project", + })); + + expect(scope.billingCode).toBeNull(); + expect(scope.issueId).toBe("issue-1"); + }); + + it("resolves a null billing code without querying when the run has no issue in context", async () => { + const { db, select } = makeDb([]); + + const scope = await resolveLedgerScopeForRun(db, "company-1", makeRun({ + projectId: "context-project", + })); + + expect(scope).toEqual({ + issueId: null, + projectId: "context-project", + billingCode: null, + }); + expect(select).not.toHaveBeenCalled(); + }); + + it("resolves a null billing code when the issue is not visible to the company", async () => { + const { db } = makeDb([]); + + const scope = await resolveLedgerScopeForRun(db, "other-company", makeRun({ + issueId: "issue-1", + projectId: "context-project", + })); + + expect(scope).toEqual({ + issueId: null, + projectId: "context-project", + billingCode: null, + }); + }); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 2953969e22..6fbe84b699 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -2667,7 +2667,7 @@ export function resolveLedgerCostStatus(input: { return input.costUsd == null && hasTokenUsage ? "unpriced" : "reported"; } -async function resolveLedgerScopeForRun( +export async function resolveLedgerScopeForRun( db: Db, companyId: string, run: typeof heartbeatRuns.$inferSelect, @@ -2680,6 +2680,7 @@ async function resolveLedgerScopeForRun( return { issueId: null, projectId: contextProjectId, + billingCode: null, }; } @@ -2687,6 +2688,7 @@ async function resolveLedgerScopeForRun( .select({ id: issues.id, projectId: issues.projectId, + billingCode: issues.billingCode, }) .from(issues) .where(and(eq(issues.id, contextIssueId), eq(issues.companyId, companyId))) @@ -2695,6 +2697,7 @@ async function resolveLedgerScopeForRun( return { issueId: issue?.id ?? null, projectId: issue?.projectId ?? contextProjectId, + billingCode: issue?.billingCode ?? null, }; } @@ -11653,6 +11656,7 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) agentId: agent.id, issueId: ledgerScope.issueId, projectId: ledgerScope.projectId, + billingCode: ledgerScope.billingCode, provider, biller, billingType,