Skip to content

Commit 7e81de7

Browse files
authored
fix(code-reviews): show model and tokens in review summary for v2 reviews (#978)
## Summary [PR #407](#407) added model + token info to the PR review summary comment (the "Reviewed by claude-sonnet-4.6 · 12,345 tokens" footer). But it only works for v1 (SSE-based) reviews. All reviews now run on v2 (cloud-agent-next), which is callback-based — it never streams SSE events, so usage data is never collected and the `model`, `total_tokens_in`, `total_tokens_out` columns stay null. **The fix:** when the `code_reviews` record has no usage data, we query the billing tables (`microdollar_usage` + `microdollar_usage_metadata`) by `cli_session_id`. The billing system already tracks every LLM call with model, tokens, and cost — we just aggregate it per session. We also back-fill the `code_reviews` record so future reads don't repeat the aggregation. ## Verification - [x] `pnpm typecheck` — no new errors (pre-existing errors in kiloclaw only) - [x] `pnpm test usage-footer` — 10/10 pass - [x] Read through the billing schema to confirm `microdollar_usage_metadata.session_id` matches `code_reviews.cli_session_id` ## Visual Changes N/A ## Reviewer Notes - The billing query groups by model and picks the one with the most tokens (the primary review model). This handles sessions that use multiple models (e.g. a cheaper model for sub-tasks). - The back-fill write to `code_reviews` is fire-and-forget (`.catch()`) — if it fails, the footer still shows correctly; we just won't cache the result. - Long-term, cloud-agent-next could include usage data in its `ExecutionCallbackPayload`, but that's a bigger change. This fix works today with no changes outside the Next.js app.
2 parents 7905afd + ad71c5e commit 7e81de7

2 files changed

Lines changed: 125 additions & 9 deletions

File tree

src/app/api/internal/code-review-status/[reviewId]/route.ts

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,12 @@
1818

1919
import type { NextRequest } from 'next/server';
2020
import { NextResponse } from 'next/server';
21-
import { updateCodeReviewStatus, getCodeReviewById } from '@/lib/code-reviews/db/code-reviews';
21+
import {
22+
updateCodeReviewStatus,
23+
updateCodeReviewUsage,
24+
getCodeReviewById,
25+
getSessionUsageFromBilling,
26+
} from '@/lib/code-reviews/db/code-reviews';
2227
import { tryDispatchPendingReviews } from '@/lib/code-reviews/dispatch/dispatch-pending-reviews';
2328
import { getBotUserId } from '@/lib/bot-users/bot-user-service';
2429
import { logExceptInTest, errorExceptInTest } from '@/lib/utils.server';
@@ -112,24 +117,58 @@ function normalizePayload(raw: StatusUpdatePayload): {
112117

113118
/**
114119
* Read a review's usage data, polling with exponential backoff if not yet available.
115-
* Handles the race between the orchestrator's usage report and the cloud agent's completion callback.
120+
*
121+
* For v1 (SSE) reviews the orchestrator reports usage before the completion
122+
* callback fires, so a short poll handles the race. For v2 (cloud-agent-next)
123+
* reviews the orchestrator never reports usage — we fall back to aggregating
124+
* from the billing tables (microdollar_usage) keyed by cli_session_id.
125+
*
126+
* When the billing fallback is used we also back-fill the code_reviews record
127+
* so subsequent reads (e.g. the admin panel) don't need the aggregation again.
116128
*/
117129
async function getReviewUsageData(reviewId: string) {
118130
const MAX_RETRIES = 3;
119131
const BASE_DELAY_MS = 200;
120132

121133
let review = await getCodeReviewById(reviewId);
122134

135+
// Short poll: usage may arrive from the orchestrator just before the callback
123136
for (let attempt = 0; attempt < MAX_RETRIES && review && !review.model; attempt++) {
124137
await new Promise(resolve => setTimeout(resolve, BASE_DELAY_MS * 2 ** attempt));
125138
review = await getCodeReviewById(reviewId);
126139
}
127140

128-
return {
129-
model: review?.model ?? null,
130-
tokensIn: review?.total_tokens_in ?? null,
131-
tokensOut: review?.total_tokens_out ?? null,
132-
};
141+
if (review?.model) {
142+
return {
143+
model: review.model,
144+
tokensIn: review.total_tokens_in ?? null,
145+
tokensOut: review.total_tokens_out ?? null,
146+
};
147+
}
148+
149+
// Fallback: aggregate from billing tables (covers v2 / cloud-agent-next reviews)
150+
if (review?.cli_session_id) {
151+
const billing = await getSessionUsageFromBilling(review.cli_session_id);
152+
if (billing) {
153+
// Back-fill the code_reviews record so we don't repeat this aggregation
154+
updateCodeReviewUsage(reviewId, {
155+
model: billing.model,
156+
totalTokensIn: billing.totalTokensIn,
157+
totalTokensOut: billing.totalTokensOut,
158+
totalCostMusd: billing.totalCostMusd,
159+
}).catch(err => {
160+
logExceptInTest('[code-review-status] Failed to back-fill usage from billing', err);
161+
});
162+
163+
return {
164+
model: billing.model,
165+
tokensIn: billing.totalTokensIn,
166+
tokensOut: billing.totalTokensOut,
167+
};
168+
}
169+
}
170+
171+
return { model: null, tokensIn: null, tokensOut: null };
133172
}
134173

135174
/**

src/lib/code-reviews/db/code-reviews.ts

Lines changed: 79 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,12 @@
66
*/
77

88
import { db } from '@/lib/drizzle';
9-
import { cloud_agent_code_reviews } from '@kilocode/db/schema';
10-
import { eq, and, desc, count, ne, inArray } from 'drizzle-orm';
9+
import {
10+
cloud_agent_code_reviews,
11+
microdollar_usage,
12+
microdollar_usage_metadata,
13+
} from '@kilocode/db/schema';
14+
import { eq, and, desc, count, ne, inArray, sql, sum } from 'drizzle-orm';
1115
import { captureException } from '@sentry/nextjs';
1216
import type { CreateReviewParams, CodeReviewStatus, ListReviewsParams, Owner } from '../core';
1317
import type { CloudAgentCodeReview } from '@kilocode/db/schema';
@@ -479,3 +483,76 @@ export async function userOwnsReview(reviewId: string, userId: string): Promise<
479483
throw error;
480484
}
481485
}
486+
487+
/**
488+
* Result of aggregating billing usage for a session.
489+
*/
490+
export type SessionUsageSummary = {
491+
model: string;
492+
totalTokensIn: number;
493+
totalTokensOut: number;
494+
totalCostMusd: number;
495+
};
496+
497+
/**
498+
* Aggregates LLM usage from the billing tables for a given kilo session ID.
499+
*
500+
* This is the fallback path for v2 (cloud-agent-next) reviews where the
501+
* orchestrator does not accumulate usage from SSE events. The billing
502+
* system (processUsage → microdollar_usage) already records per-request
503+
* usage keyed by session_id, so we aggregate here.
504+
*
505+
* Uses two queries:
506+
* 1. Session-wide totals (tokens + cost across all models)
507+
* 2. The model with the most tokens (the primary review model name)
508+
*
509+
* This avoids undercounting when a session uses more than one model.
510+
*/
511+
export async function getSessionUsageFromBilling(
512+
cliSessionId: string
513+
): Promise<SessionUsageSummary | null> {
514+
try {
515+
const sessionFilter = eq(microdollar_usage_metadata.session_id, cliSessionId);
516+
const joinCondition = eq(microdollar_usage.id, microdollar_usage_metadata.id);
517+
518+
// 1. Session-wide totals (all models combined)
519+
const [totals] = await db
520+
.select({
521+
totalTokensIn: sum(microdollar_usage.input_tokens).mapWith(Number),
522+
totalTokensOut: sum(microdollar_usage.output_tokens).mapWith(Number),
523+
totalCostMusd: sum(microdollar_usage.cost).mapWith(Number),
524+
})
525+
.from(microdollar_usage)
526+
.innerJoin(microdollar_usage_metadata, joinCondition)
527+
.where(sessionFilter);
528+
529+
if (totals?.totalTokensIn == null) return null;
530+
531+
// 2. Pick the model with the most tokens (the primary review model)
532+
const [topModel] = await db
533+
.select({ model: microdollar_usage.model })
534+
.from(microdollar_usage)
535+
.innerJoin(microdollar_usage_metadata, joinCondition)
536+
.where(sessionFilter)
537+
.groupBy(microdollar_usage.model)
538+
.orderBy(
539+
sql`sum(${microdollar_usage.input_tokens} + ${microdollar_usage.output_tokens}) desc`
540+
)
541+
.limit(1);
542+
543+
if (!topModel?.model) return null;
544+
545+
return {
546+
model: topModel.model,
547+
totalTokensIn: totals.totalTokensIn,
548+
totalTokensOut: totals.totalTokensOut ?? 0,
549+
totalCostMusd: totals.totalCostMusd ?? 0,
550+
};
551+
} catch (error) {
552+
captureException(error, {
553+
tags: { operation: 'getSessionUsageFromBilling' },
554+
extra: { cliSessionId },
555+
});
556+
return null;
557+
}
558+
}

0 commit comments

Comments
 (0)