diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 24e0fbd3c9..fbebc977db 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -610,6 +610,8 @@ Pricing estimates and gateway spend checks resolve latest-model aliases through An alias remains unavailable when its target is missing or has no usable pricing. The gateway rejects an alias request if provider routing selects a different target from the priced catalog target. The structured report remains available when generation fails or reaches a limit. +Summary charges accumulate across a report's generations: each generation keeps its settlement record, a pending update preserves the reconciled total, and retry becomes available only when every tracked generation's usage is final. +The settlement record also keeps the gateway's delivery outcome: a queued usage event keeps its generation open for the usage queue's delivery window, because queue acceptance precedes the committed row, while a failed accounting closes its generation at once. Apply database migrations before deploying the session-ingest, notification, gateway, and web changes. The CLI status producer must send occurrence timestamps and transition identities for reliable historical blocker counts. diff --git a/apps/web/src/lib/agent-briefing/summary.ts b/apps/web/src/lib/agent-briefing/summary.ts index d73e1908c0..d8b35c3036 100644 --- a/apps/web/src/lib/agent-briefing/summary.ts +++ b/apps/web/src/lib/agent-briefing/summary.ts @@ -30,6 +30,10 @@ import { type TextPricingRates, } from '@kilocode/web-shared/lib/ai-gateway/request-cost-limit'; import type { OpenRouterModelsResponse } from '@kilocode/web-shared/lib/organizations/organization-types'; +import { + BRIEFING_SETTLEMENT_HEADER, + briefingChargeAccountingSchema, +} from '@kilocode/web-shared/lib/agent-briefing-usage'; import { getBalanceAndOrgSettings } from '@kilocode/web-shared/lib/organizations/organization-usage'; import { resolveOrganizationMemberModelDecision } from '@kilocode/web-shared/lib/organizations/effective-model-access.server'; import { @@ -338,7 +342,9 @@ export async function estimateBriefingSummary( const model = input.model ?? setting.model; const maxSpendMicrodollars = input.maxSpendMicrodollars ?? setting.maxSpendMicrodollars; const maxDurationSeconds = input.maxDurationSeconds ?? setting.maxDurationSeconds; - const report = input.reportId ? await briefingReportRow(user.id, input.reportId) : null; + const report = input.reportId + ? await reconcileBriefingCharge(await briefingReportRow(user.id, input.reportId)) + : null; if (report && report.organization_id !== (input.organizationId ?? null)) throw new TRPCError({ code: 'BAD_REQUEST', message: 'Report context does not match estimate' }); const language = report?.language ?? (await getBriefingPreferences(user.id)).language; @@ -426,7 +432,7 @@ async function queueSummary( throw new TRPCError({ code: 'BAD_REQUEST', message: 'Not a one-off briefing' }); const setting = await briefingContextSettings(user.id, accessible.organization_id); const queued = await db.transaction(async tx => { - const [row] = await tx + const [locked] = await tx .select() .from(agent_briefing_reports) .where( @@ -436,7 +442,8 @@ async function queueSummary( ) ) .for('update'); - if (!row) throw new TRPCError({ code: 'NOT_FOUND', message: 'Briefing not found' }); + if (!locked) throw new TRPCError({ code: 'NOT_FOUND', message: 'Briefing not found' }); + const row = await reconcileLockedBriefingCharge(tx, locked); const summary = briefingSummarySchema.parse(row.summary); if (publishOneOff && row.published_at) return row; if (row.kind === 'one_off' && ['pending', 'running'].includes(summary.status)) return row; @@ -527,10 +534,11 @@ async function queueSummary( generation_started_at: null, generation_finished_at: null, generation_deadline_at: null, - next_charge_check_at: null, + next_charge_check_at: row.next_charge_check_at, updated_at: now, ...(publishOneOff ? { published_at: now } : {}), summary: { + ...row.summary, ...summary, status: 'pending', manualRequested: true, @@ -551,50 +559,138 @@ async function queueSummary( }); return serializeBriefing(await briefingReportRow(user.id, queued.id)); } -export async function reconcileBriefingCharge(row: BriefingReportRow): Promise { - if (!row.generation_id) return row; - const [charge] = await db +type BriefingTransaction = Parameters[0]>[0]; + +/** Margin past a generation's deadline before a receipt's missing row is final. */ +const CHARGE_REPAIR_GRACE_MS = 10 * 60_000; + +/** + * How long a queued usage delivery keeps its generation open. Queue acceptance + * precedes persistence, and Cloudflare Queues allows message retention up to 14 + * days. A queued event can therefore commit long after the inference request that + * produced it ended. + */ +const CHARGE_QUEUE_DELIVERY_WINDOW_MS = 14 * 86_400_000; + +function chargeAccounting(row: BriefingReportRow) { + return briefingChargeAccountingSchema.parse( + row.summary.chargeAccounting ?? { + baselineCostMicrodollars: z + .number() + .nonnegative() + .catch(0) + .parse(row.summary.priorCostMicrodollars), + generations: row.generation_id ? [{ id: row.generation_id, usageId: null }] : [], + } + ); +} + +async function reconcileLockedBriefingCharge( + tx: BriefingTransaction, + row: BriefingReportRow +): Promise { + const accounting = chargeAccounting(row); + if (accounting.generations.length === 0) return row; + const generationIds = accounting.generations.map(generation => generation.id); + const charges = await tx .select({ + generationId: microdollar_usage_metadata.session_id, cost: sql`coalesce(sum(${microdollar_usage.cost}), 0)::float8`, - count: sql`count(*)::integer`, + usageIds: sql`array_agg(${microdollar_usage.id}::text)`, }) .from(microdollar_usage) .innerJoin(microdollar_usage_metadata, eq(microdollar_usage_metadata.id, microdollar_usage.id)) .where( and( eq(microdollar_usage.kilo_user_id, row.kilo_user_id), - eq(microdollar_usage_metadata.session_id, row.generation_id), + sql`${microdollar_usage_metadata.session_id} = ANY(${sql.param(generationIds)}::text[])`, row.organization_id ? eq(microdollar_usage.organization_id, row.organization_id) : sql`${microdollar_usage.organization_id} IS NULL` ) - ); - const previous = z.number().nonnegative().catch(0).parse(row.summary.priorCostMicrodollars); - const total = previous + Number(charge?.cost ?? 0); - const chargeStatus = Number(charge?.count ?? 0) > 0 ? 'settled' : 'pending'; - const [updated] = await db + ) + .groupBy(microdollar_usage_metadata.session_id); + const chargesByGeneration = new Map(charges.map(charge => [charge.generationId, charge])); + const total = + accounting.baselineCostMicrodollars + + charges.reduce((cost, charge) => cost + Number(charge.cost), 0); + // A receipt names its delivery outcome: the ledger row is authoritative once it + // exists, while a queued delivery can still arrive until the queue's retention + // window closes. Only receipts without that marker use the request deadline. + const now = Date.now(); + const settlements = accounting.generations.map(generation => { + const charge = chargesByGeneration.get(generation.id); + if (generation.usageId !== null && charge?.usageIds.includes(generation.usageId) === true) + return 'settled' as const; + const final = charge === undefined ? ('unreachable' as const) : ('repaired' as const); + const deadline = + generation.deadlineAt ?? + (generation.id === row.generation_id ? row.generation_deadline_at : null); + const deadlineMs = deadline === null ? null : Date.parse(deadline); + if (generation.delivery === 'queued') { + // Receipts write the delivery time; a record without one falls back to the + // request deadline, which is inside the same request lifetime. + const anchorMs = + generation.deliveryAt === undefined ? deadlineMs : Date.parse(generation.deliveryAt); + if (anchorMs === null) return 'waiting' as const; + return now < anchorMs + CHARGE_QUEUE_DELIVERY_WINDOW_MS ? ('waiting' as const) : final; + } + if (deadlineMs === null) + return generation.settlementToken ? ('waiting' as const) : ('unreachable' as const); + return now >= deadlineMs + CHARGE_REPAIR_GRACE_MS ? final : ('waiting' as const); + }); + const settled = settlements.every( + settlement => settlement === 'settled' || settlement === 'repaired' + ); + const awaitingRows = settlements.includes('waiting'); + const chargeStatus = settled ? 'settled' : 'pending'; + const canRetry = + row.summary.status === 'failed' && + settled && + Number(chargesByGeneration.get(row.generation_id)?.cost ?? 0) > 0; + if ( + settled && + row.summary.chargeStatus === 'settled' && + row.summary.canRetry === canRetry && + z.number().nonnegative().catch(-1).parse(row.summary.costMicrodollars) === total && + row.next_charge_check_at === null + ) + return row; + const [updated] = await tx .update(agent_briefing_reports) .set({ - next_charge_check_at: - chargeStatus === 'settled' ? null : new Date(Date.now() + 60_000).toISOString(), - summary: sql`jsonb_set(${agent_briefing_reports.summary} || ${JSON.stringify({ costMicrodollars: total, chargeStatus })}::jsonb, '{canRetry}', to_jsonb((${agent_briefing_reports.summary}->>'status' = 'failed') AND ${total > previous}))`, + next_charge_check_at: awaitingRows ? new Date(now + 60_000).toISOString() : null, + summary: sql`${agent_briefing_reports.summary} || ${JSON.stringify({ costMicrodollars: total, chargeStatus, canRetry, chargeAccounting: accounting })}::jsonb`, }) - .where( - and( - eq(agent_briefing_reports.id, row.id), - eq(agent_briefing_reports.generation_id, row.generation_id) - ) - ) + .where(eq(agent_briefing_reports.id, row.id)) .returning(); return updated ?? row; } + +export async function reconcileBriefingCharge(row: BriefingReportRow): Promise { + return db.transaction(async tx => { + const [current] = await tx + .select() + .from(agent_briefing_reports) + .where( + and( + eq(agent_briefing_reports.id, row.id), + eq(agent_briefing_reports.kilo_user_id, row.kilo_user_id) + ) + ) + .for('update'); + return current ? reconcileLockedBriefingCharge(tx, current) : row; + }); +} export async function runBriefingSummary(reportId: string): Promise { const [candidate] = await db .select() .from(agent_briefing_reports) .where(eq(agent_briefing_reports.id, reportId)); if (!candidate || candidate.generation_id || candidate.summary.status !== 'pending') return; - const row = await briefingReportRow(candidate.kilo_user_id, reportId); + const row = await reconcileBriefingCharge( + await briefingReportRow(candidate.kilo_user_id, reportId) + ); if (row.summary.status !== 'pending') return; const setting = await briefingContextSettings(row.kilo_user_id, row.organization_id); const [user] = await db @@ -608,14 +704,19 @@ export async function runBriefingSummary(reportId: string): Promise { throw new Error('permission_unavailable'); if (row.kind === 'one_off' || row.summary.manualRequested === true) { const reviewed = summaryAuthorizationSchema.safeParse(row.summary.authorization); - const summary = briefingSummarySchema.parse(row.summary); + briefingSummarySchema.parse(row.summary); + const authorizedPrior = z + .number() + .nonnegative() + .catch(reviewed.success ? reviewed.data.priorCost : 0) + .parse(row.summary.priorCostMicrodollars); if ( !reviewed.success || reviewed.data.model !== setting.model || reviewed.data.maxSpendMicrodollars !== setting.maxSpendMicrodollars || reviewed.data.maxDurationSeconds !== setting.maxDurationSeconds || reviewed.data.revision !== row.revision || - reviewed.data.priorCost !== summary.costMicrodollars + reviewed.data.priorCost !== authorizedPrior ) throw new Error('permission_unavailable'); } @@ -632,25 +733,35 @@ export async function runBriefingSummary(reportId: string): Promise { const remainingMs = setting.maxDurationSeconds * 1000 - (Date.now() - Date.parse(startedAt)); if (remainingMs <= 0) throw new Error('time_limit'); generationId = `agent-briefing:${crypto.randomUUID()}`; + const settlementToken = crypto.randomUUID(); + const deadlineAt = new Date( + Date.parse(startedAt) + setting.maxDurationSeconds * 1000 + ).toISOString(); + const accounting = chargeAccounting(row); + accounting.generations.push({ + id: generationId, + settlementToken, + usageId: null, + deadlineAt, + }); const [claimed] = await db .update(agent_briefing_reports) .set({ generation_id: generationId, generation_started_at: startedAt, next_charge_check_at: startedAt, - generation_deadline_at: new Date( - Date.parse(startedAt) + setting.maxDurationSeconds * 1000 - ).toISOString(), + generation_deadline_at: deadlineAt, generation_revision: row.revision, generation_max_cost: setting.maxSpendMicrodollars, - summary: sql`${agent_briefing_reports.summary} || ${JSON.stringify({ status: 'running', chargeStatus: 'pending', model: setting.model, startedAt, estimatedCostMicrodollars: bound.estimatedCost })}::jsonb`, + summary: sql`${agent_briefing_reports.summary} || ${JSON.stringify({ status: 'running', chargeStatus: 'pending', model: setting.model, startedAt, estimatedCostMicrodollars: bound.estimatedCost, chargeAccounting: accounting })}::jsonb`, }) .where( and( eq(agent_briefing_reports.id, row.id), eq(agent_briefing_reports.revision, row.revision), sql`${agent_briefing_reports.generation_id} IS NULL`, - sql`${agent_briefing_reports.summary}->>'status' = 'pending'` + sql`${agent_briefing_reports.summary}->>'status' = 'pending'`, + sql`${agent_briefing_reports.summary}->'costMicrodollars' = ${JSON.stringify(row.summary.costMicrodollars)}::jsonb` ) ) .returning(); @@ -662,6 +773,7 @@ export async function runBriefingSummary(reportId: string): Promise { Authorization: `Bearer ${generateApiToken(user, { internalApiUse: true }, { expiresIn: setting.maxDurationSeconds + 60 })}`, 'X-Kilo-Session': generationId, 'X-Kilo-Request': generationId, + [BRIEFING_SETTLEMENT_HEADER]: settlementToken, 'X-Kilo-Max-Cost-Microdollars': String(setting.maxSpendMicrodollars), ...(row.organization_id ? { 'X-KiloCode-OrganizationId': row.organization_id } : {}), 'User-Agent': 'Kilo-Agent-Briefing', diff --git a/apps/web/src/routers/agent-briefing-router.one-off.test.ts b/apps/web/src/routers/agent-briefing-router.one-off.test.ts index e18eb35acf..f77bc4d2aa 100644 --- a/apps/web/src/routers/agent-briefing-router.one-off.test.ts +++ b/apps/web/src/routers/agent-briefing-router.one-off.test.ts @@ -89,6 +89,54 @@ beforeEach(async () => { }); afterEach(() => jest.restoreAllMocks()); +async function insertSummaryCharge(generationId: string, cost: number, id = crypto.randomUUID()) { + await db.insert(microdollar_usage).values({ + id, + kilo_user_id: user.id, + cost, + input_tokens: 1, + output_tokens: 1, + cache_write_tokens: 0, + cache_hit_tokens: 0, + }); + await db + .insert(microdollar_usage_metadata) + .values({ id, session_id: generationId, message_id: id }); + return id; +} + +async function prepareFailedSummary(receipt: 'missing' | 'uncommitted' | 'committed') { + const caller = agentBriefingRouter.createCaller({ user }); + const input = period(); + await seedActivity(input); + const prepared = await caller.prepareOneOff(input); + await caller.createOneOff({ + reportId: prepared.report.id, + estimateToken: prepared.estimate?.estimateToken, + }); + const generationId = `agent-briefing:${crypto.randomUUID()}`; + const usageId = crypto.randomUUID(); + const [report] = await db + .update(agent_briefing_reports) + .set({ + generation_id: generationId, + summary: { + ...prepared.report.summary, + status: 'failed', + chargeStatus: 'pending', + chargeAccounting: { + baselineCostMicrodollars: 100, + generations: [{ id: generationId, usageId: receipt === 'missing' ? null : usageId }], + }, + }, + }) + .where(eq(agent_briefing_reports.id, prepared.report.id)) + .returning(); + if (!report) throw new Error('Summary fixture was not inserted'); + await insertSummaryCharge(generationId, 1200, receipt === 'committed' ? usageId : undefined); + return { caller, report, generationId, usageId }; +} + type OneOffPeriod = { requestId: string; startAt: string; @@ -426,15 +474,23 @@ it('uses the existing billable retry flow without scheduled AI and cannot enqueu estimateToken: prepared.estimate?.estimateToken, }); const generationId = `agent-briefing:${crypto.randomUUID()}`; + const usageId = crypto.randomUUID(); const [failed] = await db .update(agent_briefing_reports) .set({ generation_id: generationId, - summary: { ...prepared.report.summary, status: 'failed', chargeStatus: 'pending' }, + summary: { + ...prepared.report.summary, + status: 'failed', + chargeStatus: 'pending', + chargeAccounting: { + baselineCostMicrodollars: 0, + generations: [{ id: generationId, usageId }], + }, + }, }) .where(eq(agent_briefing_reports.id, prepared.report.id)) .returning(); - const usageId = crypto.randomUUID(); await db.insert(microdollar_usage).values({ id: usageId, kilo_user_id: user.id, @@ -458,6 +514,10 @@ it('uses the existing billable retry flow without scheduled AI and cannot enqueu .where(eq(agent_briefing_reports.id, prepared.report.id)); expect(queued.summary.priorCostMicrodollars).toBe(1234); expect(queued.generation_id).toBeNull(); + expect(queued.summary.chargeAccounting).toEqual({ + baselineCostMicrodollars: 0, + generations: [{ id: generationId, usageId }], + }); }); it('enforces payer balance and organization access without creating preferences or AI jobs', async () => { @@ -883,3 +943,350 @@ it('pages large saved reports within the response budget without losing history' expect(page.nextCursor).toBeNull(); expect(seen.sort()).toEqual(ids.sort()); }, 30_000); + +it.each(['missing', 'uncommitted'] as const)( + 'keeps retry unavailable while the %s receipt is pending', + async receipt => { + const { caller, report } = await prepareFailedSummary(receipt); + await reconcileBriefingCharge(report); + const current = await caller.get({ reportId: report.id }); + expect(current.summary).toMatchObject({ + status: 'failed', + chargeStatus: 'pending', + costMicrodollars: 1300, + canRetry: false, + }); + const estimate = await caller.estimateSummary({ reportId: report.id }); + await expect( + caller.retrySummary({ reportId: report.id, estimateToken: estimate.estimateToken }) + ).rejects.toMatchObject({ code: 'CONFLICT' }); + } +); + +it('settles only after the receipt usage row is committed', async () => { + const { caller, report, generationId, usageId } = await prepareFailedSummary('uncommitted'); + expect((await reconcileBriefingCharge(report)).summary.chargeStatus).toBe('pending'); + await insertSummaryCharge(generationId, 300, usageId); + await reconcileBriefingCharge(report); + expect((await caller.get({ reportId: report.id })).summary).toMatchObject({ + status: 'failed', + chargeStatus: 'settled', + costMicrodollars: 1600, + canRetry: true, + }); +}); + +it('settles committed usage past the generation deadline without a receipt', async () => { + const { caller, report, generationId, usageId } = await prepareFailedSummary('uncommitted'); + const deadline = new Date(Date.now() - 20 * 60_000).toISOString(); + await db + .update(agent_briefing_reports) + .set({ + generation_deadline_at: deadline, + summary: { + ...report.summary, + chargeAccounting: { + baselineCostMicrodollars: 100, + generations: [{ id: generationId, usageId: null, deadlineAt: deadline }], + }, + }, + }) + .where(eq(agent_briefing_reports.id, report.id)); + await insertSummaryCharge(generationId, 300, usageId); + const [repaired] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + await reconcileBriefingCharge(repaired); + const current = await caller.get({ reportId: report.id }); + expect(current.summary).toMatchObject({ + status: 'failed', + chargeStatus: 'settled', + costMicrodollars: 1600, + canRetry: true, + }); + const [persisted] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(persisted.next_charge_check_at).toBeNull(); +}); + +it('keeps settlement pending within the generation deadline', async () => { + const { caller, report, generationId, usageId } = await prepareFailedSummary('uncommitted'); + const deadline = new Date(Date.now() + 5 * 60_000).toISOString(); + await db + .update(agent_briefing_reports) + .set({ + generation_deadline_at: deadline, + summary: { + ...report.summary, + chargeAccounting: { + baselineCostMicrodollars: 100, + generations: [{ id: generationId, usageId: null, deadlineAt: deadline }], + }, + }, + }) + .where(eq(agent_briefing_reports.id, report.id)); + await insertSummaryCharge(generationId, 300, usageId); + const [pendingRow] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + await reconcileBriefingCharge(pendingRow); + const current = await caller.get({ reportId: report.id }); + expect(current.summary).toMatchObject({ + chargeStatus: 'pending', + costMicrodollars: 1600, + canRetry: false, + }); + const [persisted] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(persisted.next_charge_check_at).not.toBeNull(); +}); +it('keeps queued usage scheduled past the inference grace until the sweep sees its commit', async () => { + const { caller, report, generationId, usageId } = await prepareFailedSummary('uncommitted'); + const deadline = new Date(Date.now() - 20 * 60_000).toISOString(); + await db + .update(agent_briefing_reports) + .set({ + generation_deadline_at: deadline, + next_charge_check_at: new Date(Date.now() - 60_000).toISOString(), + summary: { + ...report.summary, + chargeAccounting: { + baselineCostMicrodollars: 100, + generations: [ + { + id: generationId, + usageId, + deadlineAt: deadline, + deliveryAt: new Date().toISOString(), + delivery: 'queued', + }, + ], + }, + }, + }) + .where(eq(agent_briefing_reports.id, report.id)); + + await sweepAgentBriefings(); + const [waiting] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(waiting.next_charge_check_at).not.toBeNull(); + expect(waiting.summary).toMatchObject({ + chargeStatus: 'pending', + costMicrodollars: 1300, + canRetry: false, + }); + expect( + (await caller.list({ scope: 'personal' })).reports.find(item => item.id === report.id) + ).toMatchObject({ summary: { chargeStatus: 'pending', costMicrodollars: 1300 } }); + + await insertSummaryCharge(generationId, 300, usageId); + await db + .update(agent_briefing_reports) + .set({ next_charge_check_at: new Date(Date.now() - 60_000).toISOString() }) + .where(eq(agent_briefing_reports.id, report.id)); + await sweepAgentBriefings(); + const [settled] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(settled.next_charge_check_at).toBeNull(); + expect(settled.summary).toMatchObject({ + chargeStatus: 'settled', + costMicrodollars: 1600, + canRetry: true, + }); + expect( + (await caller.list({ scope: 'personal' })).reports.find(item => item.id === report.id) + ).toMatchObject({ summary: { chargeStatus: 'settled', costMicrodollars: 1600 } }); +}); + +it('ends queued reconciliation after the queue delivery window', async () => { + const { report, generationId, usageId } = await prepareFailedSummary('uncommitted'); + const deadline = new Date(Date.now() - 15 * 86_400_000).toISOString(); + const deliveryAt = new Date(Date.now() - 15 * 86_400_000).toISOString(); + await db + .update(agent_briefing_reports) + .set({ + generation_deadline_at: deadline, + next_charge_check_at: new Date(Date.now() - 60_000).toISOString(), + summary: { + ...report.summary, + chargeAccounting: { + baselineCostMicrodollars: 100, + generations: [ + { id: generationId, usageId, deadlineAt: deadline, deliveryAt, delivery: 'queued' }, + ], + }, + }, + }) + .where(eq(agent_briefing_reports.id, report.id)); + await sweepAgentBriefings(); + const [persisted] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(persisted.next_charge_check_at).toBeNull(); + expect(persisted.summary).toMatchObject({ chargeStatus: 'settled', costMicrodollars: 1300 }); +}); + +it('clears a leftover charge schedule on an already settled report', async () => { + const { caller, report } = await prepareFailedSummary('committed'); + await reconcileBriefingCharge(report); + expect((await caller.get({ reportId: report.id })).summary.chargeStatus).toBe('settled'); + await db + .update(agent_briefing_reports) + .set({ next_charge_check_at: new Date(Date.now() - 60_000).toISOString() }) + .where(eq(agent_briefing_reports.id, report.id)); + const [stale] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + await reconcileBriefingCharge(stale); + const [persisted] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(persisted.next_charge_check_at).toBeNull(); + expect(persisted.summary).toMatchObject({ + chargeStatus: 'settled', + costMicrodollars: 1300, + canRetry: true, + }); + expect((await caller.get({ reportId: report.id })).summary.costMicrodollars).toBe(1300); +}); + +it('keeps generations without receipt accounting pending', async () => { + const { caller, report, generationId } = await prepareFailedSummary('committed'); + await db + .update(agent_briefing_reports) + .set({ + summary: { + ...report.summary, + priorCostMicrodollars: 100, + chargeAccounting: undefined, + }, + }) + .where(eq(agent_briefing_reports.id, report.id)); + const [legacy] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + await reconcileBriefingCharge(legacy); + expect((await caller.get({ reportId: report.id })).summary).toMatchObject({ + chargeStatus: 'pending', + costMicrodollars: 1300, + canRetry: false, + }); + expect(legacy.generation_id).toBe(generationId); + const [persisted] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(persisted.next_charge_check_at).toBeNull(); +}); + +it.each(['queued', 'running'] as const)( + 'retains all generation charges while the retry is %s', + async status => { + const { caller, report, generationId, usageId } = await prepareFailedSummary('committed'); + await insertSummaryCharge(generationId, 300); + await reconcileBriefingCharge(report); + const estimate = await caller.estimateSummary({ reportId: report.id }); + await caller.retrySummary({ reportId: report.id, estimateToken: estimate.estimateToken }); + const [queued] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(queued.generation_id).toBeNull(); + expect(queued.summary).toMatchObject({ + status: 'pending', + costMicrodollars: 1600, + priorCostMicrodollars: 1600, + chargeAccounting: { + baselineCostMicrodollars: 100, + generations: [{ id: generationId, usageId }], + }, + }); + const nextGenerationId = `agent-briefing:${crypto.randomUUID()}`; + const nextUsageId = crypto.randomUUID(); + if (status === 'running') { + await db + .update(agent_briefing_reports) + .set({ + generation_id: nextGenerationId, + summary: { + ...queued.summary, + status: 'running', + chargeStatus: 'pending', + chargeAccounting: { + baselineCostMicrodollars: 100, + generations: [ + { id: generationId, usageId }, + { id: nextGenerationId, usageId: nextUsageId }, + ], + }, + }, + }) + .where(eq(agent_briefing_reports.id, report.id)); + await insertSummaryCharge(nextGenerationId, 700, nextUsageId); + } + await insertSummaryCharge(generationId, 200); + await reconcileBriefingCharge(report); + const current = await caller.get({ reportId: report.id }); + expect(current.summary).toMatchObject({ + status: status === 'running' ? 'running' : 'pending', + costMicrodollars: status === 'running' ? 2500 : 1800, + canRetry: false, + }); + const [persisted] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(persisted.generation_id).toBe(status === 'running' ? nextGenerationId : null); + expect(persisted.summary.chargeAccounting).toEqual({ + baselineCostMicrodollars: 100, + generations: + status === 'running' + ? [ + { id: generationId, usageId }, + { id: nextGenerationId, usageId: nextUsageId }, + ] + : [{ id: generationId, usageId }], + }); + } +); + +it('requires a fresh estimate when committed generation cost changes', async () => { + const { caller, report, generationId } = await prepareFailedSummary('committed'); + await reconcileBriefingCharge(report); + const estimate = await caller.estimateSummary({ reportId: report.id }); + await insertSummaryCharge(generationId, 300); + await expect( + caller.retrySummary({ reportId: report.id, estimateToken: estimate.estimateToken }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + const [failed] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(failed.generation_id).toBe(generationId); + expect(failed.summary.status).toBe('failed'); + const refreshed = await caller.estimateSummary({ reportId: report.id }); + await caller.retrySummary({ reportId: report.id, estimateToken: refreshed.estimateToken }); + const [queued] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, report.id)); + expect(queued.summary).toMatchObject({ + status: 'pending', + costMicrodollars: 1600, + priorCostMicrodollars: 1600, + }); +}); diff --git a/apps/web/src/routers/agent-briefing-router.test.ts b/apps/web/src/routers/agent-briefing-router.test.ts index c703017961..cc49f5bfb9 100644 --- a/apps/web/src/routers/agent-briefing-router.test.ts +++ b/apps/web/src/routers/agent-briefing-router.test.ts @@ -604,17 +604,25 @@ describe('agent briefing ownership and lifecycle', () => { it('reconciles delayed real ledger rows without clobbering summary state and enables only billable retry', async () => { const { report } = await seedReport(); const generationId = `agent-briefing:${crypto.randomUUID()}`; + const usageId = crypto.randomUUID(); const [failed] = await db .update(agent_briefing_reports) .set({ generation_id: generationId, - summary: { ...emptySummary, status: 'failed', chargeStatus: 'pending' }, + summary: { + ...emptySummary, + status: 'failed', + chargeStatus: 'pending', + chargeAccounting: { + baselineCostMicrodollars: 0, + generations: [{ id: generationId, usageId }], + }, + }, }) .where(eq(agent_briefing_reports.id, report.id)) .returning(); if (!failed) throw new Error('Failed fixture was not inserted'); expect((await reconcileBriefingCharge(failed)).summary.canRetry).toBe(false); - const usageId = crypto.randomUUID(); await db.insert(microdollar_usage).values({ id: usageId, kilo_user_id: user.id, diff --git a/packages/web-shared/src/lib/agent-briefing-usage.test.ts b/packages/web-shared/src/lib/agent-briefing-usage.test.ts new file mode 100644 index 0000000000..7b47f34262 --- /dev/null +++ b/packages/web-shared/src/lib/agent-briefing-usage.test.ts @@ -0,0 +1,230 @@ +import { randomUUID } from 'crypto'; +import { eq } from 'drizzle-orm'; +import { agent_briefing_reports } from '@kilocode/db/schema'; +import { db } from '@kilocode/web-shared/lib/drizzle'; +import { insertTestUser } from '@kilocode/web-shared/tests/helpers/user.helper'; +import { createTestOrganization } from '@kilocode/web-shared/tests/helpers/organization.helper'; +import { recordBriefingUsage } from './agent-briefing-usage'; +import type { UsageAccountingResult } from './ai-gateway/processUsage.types'; + +async function fixture(organizationId?: string) { + const user = await insertTestUser(); + const generationId = `agent-briefing:${randomUUID()}`; + const settlementToken = randomUUID(); + const previousGeneration = { + id: `agent-briefing:${randomUUID()}`, + settlementToken: randomUUID(), + usageId: randomUUID(), + }; + const summary = { + status: 'failed', + error: 'time_limit', + text: null, + costMicrodollars: 120, + chargeStatus: 'pending', + canRetry: false, + retained: { value: 'preserved' }, + chargeAccounting: { + baselineCostMicrodollars: 80, + retained: 'accounting-field', + generations: [previousGeneration, { id: generationId, settlementToken, usageId: null }], + }, + }; + const now = new Date().toISOString(); + const [report] = await db + .insert(agent_briefing_reports) + .values({ + kilo_user_id: user.id, + context_key: organizationId ?? 'personal', + organization_id: organizationId, + kind: 'one_off', + night_start_at: now, + cutoff_at: now, + morning_end_at: now, + evidence_hash: 'receipt-test', + facts: { retained: 'facts' }, + summary, + generation_id: generationId, + next_charge_check_at: now, + }) + .returning(); + if (!report) throw new Error('Report fixture was not inserted'); + return { + report, + summary, + input: { kiloUserId: user.id, organizationId, generationId, settlementToken }, + }; +} + +async function readReport(id: string) { + const [report] = await db + .select() + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, id)); + if (!report) throw new Error('Report fixture was not found'); + return report; +} + +function usage(status: UsageAccountingResult['status'] = 'recorded'): UsageAccountingResult { + return { status, usageId: randomUUID(), createdAt: new Date().toISOString() }; +} + +describe('recordBriefingUsage', () => { + it.each(['recorded', 'queued'] as const)( + 'stores the %s accounting identity while preserving report fields', + async status => { + const prepared = await fixture(); + const result = usage(status); + await recordBriefingUsage({ ...prepared.input, usage: result }); + const persisted = await readReport(prepared.report.id); + const expectedSummary = { + ...prepared.summary, + chargeAccounting: { + ...prepared.summary.chargeAccounting, + generations: prepared.summary.chargeAccounting.generations.map(entry => + entry.id === prepared.input.generationId + ? { + ...entry, + usageId: result.usageId, + delivery: status === 'queued' ? 'queued' : 'committed', + deliveryAt: expect.any(String), + } + : entry + ), + }, + }; + expect(persisted).toEqual({ ...prepared.report, summary: expectedSummary }); + } + ); + + it('retains the first accounting identity for repeated completion', async () => { + const prepared = await fixture(); + const first = usage('queued'); + await recordBriefingUsage({ ...prepared.input, usage: first }); + const persisted = await readReport(prepared.report.id); + await recordBriefingUsage({ ...prepared.input, usage: first }); + await recordBriefingUsage({ ...prepared.input, usage: usage() }); + expect(await readReport(prepared.report.id)).toEqual(persisted); + }); + + it('serializes concurrent completion identities', async () => { + const prepared = await fixture(); + const first = usage(); + const second = usage(); + await Promise.all([ + recordBriefingUsage({ ...prepared.input, usage: first }), + recordBriefingUsage({ ...prepared.input, usage: second }), + ]); + const persisted = await readReport(prepared.report.id); + const accounting = persisted.summary.chargeAccounting; + expect(accounting).toEqual({ + ...prepared.summary.chargeAccounting, + generations: [ + prepared.summary.chargeAccounting.generations[0], + { + id: prepared.input.generationId, + settlementToken: prepared.input.settlementToken, + usageId: expect.stringMatching(`${first.usageId}|${second.usageId}`), + delivery: 'committed', + deliveryAt: expect.any(String), + }, + ], + }); + }); + + it('requires the generation capability and report owner', async () => { + const prepared = await fixture(); + const otherUser = await insertTestUser(); + for (const changed of [ + { settlementToken: randomUUID() }, + { settlementToken: '' }, + { kiloUserId: otherUser.id }, + { generationId: `agent-briefing:${randomUUID()}` }, + ]) { + await recordBriefingUsage({ ...prepared.input, ...changed, usage: usage() }); + expect(await readReport(prepared.report.id)).toEqual(prepared.report); + } + }); + + it('requires exact personal and organization scope', async () => { + const owner = await insertTestUser(); + const organization = await createTestOrganization('Receipt scope', owner.id, 0); + const otherOrganization = await createTestOrganization('Other receipt scope', owner.id, 0); + const personal = await fixture(); + const scoped = await fixture(organization.id); + await recordBriefingUsage({ + ...personal.input, + organizationId: organization.id, + usage: usage(), + }); + await recordBriefingUsage({ ...scoped.input, organizationId: undefined, usage: usage() }); + await recordBriefingUsage({ + ...scoped.input, + organizationId: otherOrganization.id, + usage: usage(), + }); + expect(await readReport(personal.report.id)).toEqual(personal.report); + expect(await readReport(scoped.report.id)).toEqual(scoped.report); + const result = usage(); + await recordBriefingUsage({ ...scoped.input, usage: result }); + expect((await readReport(scoped.report.id)).summary).toMatchObject({ + chargeAccounting: { + generations: [expect.anything(), expect.objectContaining({ usageId: result.usageId })], + }, + }); + }); + + it('requires the current generation and matching accounting entry', async () => { + const prepared = await fixture(); + const previous = prepared.summary.chargeAccounting.generations[0]; + if (!previous) throw new Error('Previous generation fixture was not found'); + await recordBriefingUsage({ + ...prepared.input, + generationId: previous.id, + settlementToken: previous.settlementToken, + usage: usage(), + }); + expect(await readReport(prepared.report.id)).toEqual(prepared.report); + await db + .update(agent_briefing_reports) + .set({ generation_id: `agent-briefing:${randomUUID()}` }) + .where(eq(agent_briefing_reports.id, prepared.report.id)); + const changed = await readReport(prepared.report.id); + await recordBriefingUsage({ ...prepared.input, usage: usage() }); + expect(await readReport(prepared.report.id)).toEqual(changed); + }); + + it('leaves completion pending without an accounting outcome', async () => { + const prepared = await fixture(); + await recordBriefingUsage({ ...prepared.input, usage: null }); + expect(await readReport(prepared.report.id)).toEqual(prepared.report); + }); + + it('requires a valid generation accounting structure', async () => { + const prepared = await fixture(); + for (const chargeAccounting of [ + undefined, + { baselineCostMicrodollars: -1, generations: [] }, + { baselineCostMicrodollars: 0, generations: [] }, + { + baselineCostMicrodollars: 0, + generations: [{ id: prepared.input.generationId, usageId: null }], + }, + { + ...prepared.summary.chargeAccounting, + generations: [ + ...prepared.summary.chargeAccounting.generations, + { id: prepared.input.generationId, settlementToken: randomUUID(), usageId: null }, + ], + }, + ]) { + await db + .update(agent_briefing_reports) + .set({ summary: { ...prepared.summary, chargeAccounting } }) + .where(eq(agent_briefing_reports.id, prepared.report.id)); + const changed = await readReport(prepared.report.id); + await recordBriefingUsage({ ...prepared.input, usage: usage() }); + expect(await readReport(prepared.report.id)).toEqual(changed); + } + }); +}); diff --git a/packages/web-shared/src/lib/agent-briefing-usage.ts b/packages/web-shared/src/lib/agent-briefing-usage.ts new file mode 100644 index 0000000000..62f18eece1 --- /dev/null +++ b/packages/web-shared/src/lib/agent-briefing-usage.ts @@ -0,0 +1,87 @@ +import 'server-only'; +import * as z from 'zod'; +import { and, eq, isNull, sql } from 'drizzle-orm'; +import { agent_briefing_reports } from '@kilocode/db/schema'; +import { db } from '@kilocode/web-shared/lib/drizzle'; +import type { UsageAccountingResult } from '@kilocode/web-shared/lib/ai-gateway/processUsage.types'; + +export const BRIEFING_SETTLEMENT_HEADER = 'X-Kilo-Briefing-Settlement'; + +/** + * What the gateway reported for a generation's usage write. `committed` means the + * row was written before the response; `queued` means only queue acceptance is + * known. + */ +export const briefingChargeDeliverySchema = z.enum(['committed', 'queued']); + +export const briefingChargeAccountingSchema = z.object({ + baselineCostMicrodollars: z.number().nonnegative(), + generations: z.array( + z.object({ + id: z.string(), + settlementToken: z.string().optional(), + usageId: z.uuid().nullable(), + deadlineAt: z.string().datetime().optional(), + deliveryAt: z.string().datetime().optional(), + delivery: briefingChargeDeliverySchema.optional(), + }) + ), +}); + +export async function recordBriefingUsage(input: { + kiloUserId: string; + organizationId?: string; + generationId: string; + settlementToken: string; + usage: UsageAccountingResult | null; +}): Promise { + const usage = input.usage; + if (!usage || !input.settlementToken || !z.uuid().safeParse(usage.usageId).success) return; + const delivery = usage.status === 'queued' ? 'queued' : 'committed'; + // Queue acceptance precedes this receipt, so it anchors a full retention window. + const deliveryAt = new Date().toISOString(); + + // The gateway runs in sfo1 while db is the Frankfurt primary: this is the rare + // post-response write for a briefing request, and it takes a row lock so the + // receipt cannot land on a report whose generation has already moved on. + await db.transaction(async tx => { + const [report] = await tx + .select({ id: agent_briefing_reports.id, summary: agent_briefing_reports.summary }) + .from(agent_briefing_reports) + .where( + and( + eq(agent_briefing_reports.generation_id, input.generationId), + eq(agent_briefing_reports.kilo_user_id, input.kiloUserId), + input.organizationId === undefined + ? isNull(agent_briefing_reports.organization_id) + : eq(agent_briefing_reports.organization_id, input.organizationId) + ) + ) + .for('update'); + if (!report) return; + + const accounting = briefingChargeAccountingSchema.safeParse(report.summary.chargeAccounting); + if (!accounting.success) return; + const index = accounting.data.generations.findIndex(entry => entry.id === input.generationId); + const generation = accounting.data.generations[index]; + if ( + !generation || + generation.settlementToken !== input.settlementToken || + generation.usageId !== null || + accounting.data.generations.some( + (entry, otherIndex) => otherIndex !== index && entry.id === input.generationId + ) + ) { + return; + } + + const path = sql`ARRAY['chargeAccounting', 'generations', ${String(index)}]::text[]`; + const receipt = JSON.stringify({ usageId: usage.usageId, delivery, deliveryAt }); + await tx + .update(agent_briefing_reports) + .set({ + summary: sql`jsonb_set(${agent_briefing_reports.summary}, ${path}, coalesce(${agent_briefing_reports.summary} #> ${path}, '{}'::jsonb) || ${receipt}::jsonb)`, + }) + .where(eq(agent_briefing_reports.id, report.id)); + }); +} diff --git a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts index 698e58f1f9..6d63b6d23c 100644 --- a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts +++ b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts @@ -1,6 +1,11 @@ import { describe, it, expect, beforeEach } from '@jest/globals'; import { simHash64 } from '@kilocode/web-shared/lib/bouncer/simhash'; -import type { User } from '@kilocode/db/schema'; +import { randomUUID } from 'crypto'; +import { eq } from 'drizzle-orm'; +import { agent_briefing_reports, type User } from '@kilocode/db/schema'; +import { BRIEFING_SETTLEMENT_HEADER } from '@kilocode/web-shared/lib/agent-briefing-usage'; +import type { UsageAccountingResult } from '@kilocode/web-shared/lib/ai-gateway/processUsage.types'; +import { insertTestUser } from '@kilocode/web-shared/tests/helpers/user.helper'; import jwt from 'jsonwebtoken'; import { getUserFromAuth } from '@kilocode/web-shared/lib/user/server'; import { failureResult } from '@kilocode/web-shared/lib/maybe-result'; @@ -57,7 +62,7 @@ import { applyResolvedAutoModel } from '@kilocode/web-shared/lib/ai-gateway/auto import type * as AutoModelResolution from '@kilocode/web-shared/lib/ai-gateway/auto-model/resolution'; import { getDirectByokModel } from '@kilocode/web-shared/lib/ai-gateway/providers/direct-byok'; import { rewriteModelResponse } from '@kilocode/web-shared/lib/ai-gateway/rewriteModelResponse'; -import { readDb } from '@kilocode/web-shared/lib/drizzle'; +import { db, readDb } from '@kilocode/web-shared/lib/drizzle'; import { isFreeModelRateLimited } from '@kilocode/web-shared/lib/ai-gateway/free-model-rate-limit'; import { isCustomLlm } from '@kilocode/web-shared/lib/ai-gateway/custom-llm/custom-llm-catalog'; import { gemma_4_26b_a4b_it_free_model } from '@kilocode/web-shared/lib/ai-gateway/kilo-exclusive-models'; @@ -69,7 +74,7 @@ import { releaseDecideLease, type DecideResponse, } from '@kilocode/web-shared/lib/bouncer/client'; -import { NextRequest, NextResponse } from 'next/server'; +import { after, NextRequest, NextResponse } from 'next/server'; import { handleLlmProxyRequest } from './llm-proxy'; jest.mock('next/server', () => { @@ -98,7 +103,10 @@ jest.mock('@kilocode/web-shared/lib/creditTransactions', () => ({ payments_total_microdollars: 0, })), })); -jest.mock('@kilocode/web-shared/lib/drizzle', () => ({ readDb: {} })); +jest.mock('@kilocode/web-shared/lib/drizzle', () => ({ + ...jest.requireActual>('@kilocode/web-shared/lib/drizzle'), + readDb: {}, +})); jest.mock('@kilocode/web-shared/lib/ai-gateway/custom-llm/custom-llm-catalog', () => ({ isCustomLlm: jest.fn(async () => false), })); @@ -379,7 +387,7 @@ describe('POST /api/openrouter/v1/chat/completions bearer audiences', () => { response: upstreamJsonResponse({ id: 'chatcmpl-1', model: 'openai/gpt-4o', choices: [] }), }); mockedEmitApiMetricsForResponse.mockReturnValue(undefined); - mockedAccountForMicrodollarUsage.mockReturnValue(undefined); + mockedAccountForMicrodollarUsage.mockResolvedValue(null); }); it('serves a free model to a client that sends the anonymous sentinel', async () => { @@ -696,7 +704,7 @@ describe('POST /api/openrouter/v1/chat/completions request handling', () => { response: upstreamJsonResponse({ id: 'chatcmpl-1', model: 'openai/gpt-4o', choices: [] }), }); mockedEmitApiMetricsForResponse.mockReturnValue(undefined); - mockedAccountForMicrodollarUsage.mockReturnValue(undefined); + mockedAccountForMicrodollarUsage.mockResolvedValue(null); }); it('rejects providerOptions and directs clients to provider', async () => { @@ -1741,7 +1749,7 @@ describe('kilo-auto/efficient classifier billing', () => { }), }); mockedEmitApiMetricsForResponse.mockReturnValue(undefined); - mockedAccountForMicrodollarUsage.mockReturnValue(undefined); + mockedAccountForMicrodollarUsage.mockResolvedValue(null); mockedLogMicrodollarUsage.mockResolvedValue(null); mockedGetEffectiveModelDecision.mockResolvedValue({ allowed: true }); mockedCollectDeniedAutoRoutingModelIds.mockResolvedValue([]); @@ -2442,7 +2450,7 @@ describe('auto-routing shadow classifier', () => { response: upstreamJsonResponse({ id: 'chatcmpl-1', model: 'openai/gpt-4o', choices: [] }), }); mockedEmitApiMetricsForResponse.mockReturnValue(undefined); - mockedAccountForMicrodollarUsage.mockReturnValue(undefined); + mockedAccountForMicrodollarUsage.mockResolvedValue(null); mockedApplyResolvedAutoModel.mockImplementation(async (opts, request) => { if (opts.efficientDecision) await opts.efficientDecision(); request.body.model = 'openai/gpt-4o'; @@ -2529,7 +2537,7 @@ describe('custom routers and signed auto decisions', () => { response: upstreamJsonResponse({ id: 'chatcmpl-1', model: haiku, choices: [] }), }); mockedEmitApiMetricsForResponse.mockReturnValue(undefined); - mockedAccountForMicrodollarUsage.mockReturnValue(undefined); + mockedAccountForMicrodollarUsage.mockResolvedValue(null); mockedGetEffectiveModelDecision.mockResolvedValue({ allowed: true }); mockedDeniedAutoRoutingCandidates.mockResolvedValue([]); mockedFindCustomRouter.mockResolvedValue(personalRouter); @@ -2892,3 +2900,195 @@ describe('custom routers and signed auto decisions', () => { }); }); }); + +describe('briefing gateway accounting receipts', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockedGetBalanceAndOrgSettings.mockResolvedValue({ + balance: 1000, + settings: undefined, + plan: undefined, + }); + mockedGetProvider.mockResolvedValue({ + kind: 'provider', + provider, + userByok: null, + isFree: false, + bypassAccessCheck: false, + }); + mockedGetOpenRouterModels.mockResolvedValue(new Set()); + mockedIsValidOpenRouterModelId.mockResolvedValue(true); + mockedIsFreeModelRateLimited.mockResolvedValue(false); + mockedUpstreamRequest.mockResolvedValue({ + type: 'success', + response: upstreamJsonResponse({ id: 'chatcmpl-1', model: 'openai/gpt-4o', choices: [] }), + }); + mockedAccountForMicrodollarUsage.mockResolvedValue(null); + }); + + async function prepare() { + const user = await insertTestUser(); + mockedGetUserFromAuth.mockResolvedValue({ + user, + authFailedResponse: null, + organizationId: undefined, + }); + const generationId = `agent-briefing:${randomUUID()}`; + const settlementToken = randomUUID(); + const now = new Date().toISOString(); + const summary = { + status: 'failed', + chargeStatus: 'pending', + chargeAccounting: { + baselineCostMicrodollars: 0, + generations: [{ id: generationId, settlementToken, usageId: null }], + }, + }; + const [report] = await db + .insert(agent_briefing_reports) + .values({ + kilo_user_id: user.id, + context_key: 'personal', + kind: 'one_off', + night_start_at: now, + cutoff_at: now, + morning_end_at: now, + evidence_hash: 'gateway-receipt-test', + facts: {}, + summary, + generation_id: generationId, + }) + .returning(); + if (!report) throw new Error('Report fixture was not inserted'); + return { + report, + headers: { + [BRIEFING_SETTLEMENT_HEADER]: settlementToken, + 'x-kilocode-taskid': generationId, + 'x-kilo-request': generationId, + }, + usage: { + status: 'queued', + usageId: randomUUID(), + createdAt: now, + } satisfies UsageAccountingResult, + }; + } + + async function persistedSummary(reportId: string) { + const [report] = await db + .select({ summary: agent_briefing_reports.summary }) + .from(agent_briefing_reports) + .where(eq(agent_briefing_reports.id, reportId)); + if (!report) throw new Error('Report fixture was not found'); + return report.summary; + } + + async function finishAfterTasks() { + await Promise.all( + jest.mocked(after).mock.calls.map(([task]) => (typeof task === 'function' ? undefined : task)) + ); + } + + it('stores completion after accounting without changing the response', async () => { + const prepared = await prepare(); + const completion = Promise.withResolvers(); + mockedAccountForMicrodollarUsage.mockReturnValueOnce(completion.promise); + const response = await handleLlmProxyRequest(makeRequest(makeBody(), prepared.headers)); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + id: 'chatcmpl-1', + model: 'openai/gpt-4o', + choices: [], + }); + expect(await persistedSummary(prepared.report.id)).toEqual(prepared.report.summary); + completion.resolve(prepared.usage); + await finishAfterTasks(); + expect(await persistedSummary(prepared.report.id)).toMatchObject({ + chargeAccounting: { + generations: [expect.objectContaining({ usageId: prepared.usage.usageId })], + }, + }); + }); + + it('stores accounting completion after response handling fails', async () => { + const prepared = await prepare(); + const completion = Promise.withResolvers(); + mockedAccountForMicrodollarUsage.mockReturnValueOnce(completion.promise); + mockedRewriteModelResponse.mockRejectedValueOnce(new Error('Response unavailable')); + await expect(handleLlmProxyRequest(makeRequest(makeBody(), prepared.headers))).rejects.toThrow( + 'Response unavailable' + ); + expect(await persistedSummary(prepared.report.id)).toEqual(prepared.report.summary); + completion.resolve(prepared.usage); + await finishAfterTasks(); + expect(await persistedSummary(prepared.report.id)).toMatchObject({ + chargeAccounting: { + generations: [expect.objectContaining({ usageId: prepared.usage.usageId })], + }, + }); + }); + + it('leaves completion pending when accounting is unavailable', async () => { + const prepared = await prepare(); + mockedAccountForMicrodollarUsage.mockRejectedValueOnce(new Error('Accounting unavailable')); + const response = await handleLlmProxyRequest(makeRequest(makeBody(), prepared.headers)); + expect(response.status).toBe(200); + await finishAfterTasks(); + expect(await persistedSummary(prepared.report.id)).toEqual(prepared.report.summary); + }); + + it.each([ + 'without capability', + 'different request', + 'different session', + 'streaming', + 'anonymous', + 'auto routing', + ])('leaves accounting pending for requests %s', async variant => { + const prepared = await prepare(); + const headers = new Headers(prepared.headers); + const body = makeBody(); + if (variant === 'without capability') headers.delete(BRIEFING_SETTLEMENT_HEADER); + if (variant === 'different request') headers.set('x-kilo-request', randomUUID()); + if (variant === 'different session') { + headers.set('x-kilocode-taskid', 'another-session'); + headers.set('x-kilo-request', 'another-session'); + } + if (variant === 'anonymous') { + headers.set('authorization', 'Bearer anonymous'); + body.model = 'test/free-model:free'; + mockedGetUserFromAuth.mockResolvedValueOnce({ + user: null, + authFailedResponse: NextResponse.json(failureResult('Unauthorized'), { status: 401 }), + organizationId: undefined, + }); + mockedGetProvider.mockResolvedValueOnce({ + kind: 'provider', + provider, + userByok: null, + isFree: true, + bypassAccessCheck: false, + }); + } + if (variant === 'auto routing') { + body.model = 'kilo-auto/efficient'; + mockedGetEffectiveModelDecision.mockResolvedValue({ allowed: true }); + mockedCollectDeniedAutoRoutingModelIds.mockResolvedValue([]); + mockedCollectDataCollectionRequiredAutoRoutingModelIds.mockResolvedValue([]); + mockedApplyResolvedAutoModel.mockImplementationOnce(async (_options, request) => { + request.body.model = 'openai/gpt-4o'; + return { kind: 'ok', resolved: { model: 'openai/gpt-4o' } }; + }); + } + mockedAccountForMicrodollarUsage.mockResolvedValueOnce(prepared.usage); + const request = makeRequest( + { ...body, stream: variant === 'streaming' }, + Object.fromEntries(headers) + ); + const response = await handleLlmProxyRequest(request); + expect(response.status).toBe(200); + await finishAfterTasks(); + expect(await persistedSummary(prepared.report.id)).toEqual(prepared.report.summary); + }); +}); diff --git a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts index 2e3311ea44..49c423c837 100644 --- a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts +++ b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts @@ -6,6 +6,10 @@ import { toMicrodollars } from '@kilocode/web-shared/lib/microdollars'; import { extractPromptInfo } from '@kilocode/web-shared/lib/ai-gateway/extractPromptInfo'; import { determineFallbackFeature } from '@kilocode/web-shared/lib/ai-gateway/determineFallbackFeature'; import { validateFeatureHeader, FEATURE_HEADER } from '@kilocode/web-shared/lib/feature-detection'; +import { + BRIEFING_SETTLEMENT_HEADER, + recordBriefingUsage, +} from '@kilocode/web-shared/lib/agent-briefing-usage'; import type { OpenRouterChatCompletionRequest, GatewayResponsesRequest, @@ -136,7 +140,7 @@ import { hasMiddleOutTransform, } from '@kilocode/web-shared/lib/ai-gateway/providers/openrouter/request-helpers'; import { redactProviderHints } from '@kilocode/auto-routing-contracts'; -import { logExceptInTest } from '@kilocode/web-shared/lib/utils.server'; +import { errorExceptInTest, logExceptInTest } from '@kilocode/web-shared/lib/utils.server'; import { readDb } from '@kilocode/web-shared/lib/drizzle'; import { getOrganizationGroupPolicyContext } from '@kilocode/web-shared/lib/organizations/organization-group-policy-context.server'; import { @@ -1106,7 +1110,33 @@ export async function handleLlmProxyRequest( const clonedReponse = response.clone(); // reading from body is side-effectful - accountForMicrodollarUsage(clonedReponse, usageContext, openrouterRequestSpan); + const accounting = accountForMicrodollarUsage(clonedReponse, usageContext, openrouterRequestSpan); + const settlementToken = request.headers.get(BRIEFING_SETTLEMENT_HEADER); + if ( + settlementToken && + !isAnonymousContext(user) && + !usageContext.isStreaming && + autoModel === null && + usageContext.session_id?.startsWith('agent-briefing:') && + clientRequestId === usageContext.session_id + ) { + const generationId = usageContext.session_id; + after( + accounting + .then(usage => + recordBriefingUsage({ + kiloUserId: user.id, + organizationId, + generationId, + settlementToken, + usage, + }) + ) + .catch(() => { + errorExceptInTest('Briefing accounting receipt could not be stored'); + }) + ); + } const requestLogging = { user: maybeUser, diff --git a/packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts b/packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts index bc22845d2d..1130bd45bf 100644 --- a/packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts @@ -664,7 +664,9 @@ export function accountForMicrodollarUsage( ) { const logFileExtension = usageContext.isStreaming ? '.log.resp.sse' : '.log.resp.json'; debugSaveProxyResponseStream(clonedReponse, logFileExtension); - after(countAndStoreUsage(clonedReponse, usageContext, openrouterRequestSpan)); + const accounting = countAndStoreUsage(clonedReponse, usageContext, openrouterRequestSpan); + after(accounting); + return accounting; } function isObjectRecord(value: unknown): value is Record { diff --git a/packages/worker-utils/src/redact-headers.test.ts b/packages/worker-utils/src/redact-headers.test.ts index 92f69d6748..e43687df37 100644 --- a/packages/worker-utils/src/redact-headers.test.ts +++ b/packages/worker-utils/src/redact-headers.test.ts @@ -12,6 +12,7 @@ describe('redactSensitiveHeaders', () => { 'x-gitlab-token': 'glpat-secret', 'x-hub-signature': 'sha1=abc', 'x-hub-signature-256': 'sha256=def', + 'x-kilo-briefing-settlement': 'generation-capability', 'content-type': 'application/json', 'x-request-id': '123', }; @@ -27,6 +28,7 @@ describe('redactSensitiveHeaders', () => { 'x-gitlab-token': '[REDACTED]', 'x-hub-signature': '[REDACTED]', 'x-hub-signature-256': '[REDACTED]', + 'x-kilo-briefing-settlement': '[REDACTED]', 'content-type': 'application/json', 'x-request-id': '123', }); diff --git a/packages/worker-utils/src/redact-headers.ts b/packages/worker-utils/src/redact-headers.ts index a9463bf07a..ada86224be 100644 --- a/packages/worker-utils/src/redact-headers.ts +++ b/packages/worker-utils/src/redact-headers.ts @@ -4,6 +4,7 @@ const SENSITIVE_HEADERS = new Set([ 'cookie', 'set-cookie', 'x-api-key', + 'x-kilo-briefing-settlement', 'x-gitlab-token', 'x-hub-signature', 'x-hub-signature-256',