Skip to content

fix(code-review): split cached token usage #4086

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alex-alecu wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(code-review): split cached token usage #4086

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fed3ff7
fix(code-review): split cached token usage
alex-alecu Jun 18, 2026
71b08f9
fix(code-review): bound review usage lookup
alex-alecu Jun 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions apps/web/src/app/(app)/code-reviews/[reviewId]/CodeReviewDetailClient.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { Button } from '@/components/ui/button';
import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
import { PageContainer } from '@/components/layouts/PageContainer';
import { CodeReviewStreamView } from '@/components/code-reviews/CodeReviewStreamView';
import { formatTokenCount } from '@/lib/code-reviews/summary/usage-footer';
import {
ExternalLink,
GitPullRequest,
Expand Down Expand Up @@ -55,6 +56,10 @@ type CodeReviewDetailClientProps = {
reviewId: string;
};

function formatAvailableTokenCount(count: number | null): string {
return count == null ? '—' : formatTokenCount(count);
}

export function CodeReviewDetailClient({ reviewId }: CodeReviewDetailClientProps) {
const trpc = useTRPC();
const queryClient = useQueryClient();
Expand Down Expand Up @@ -268,12 +273,15 @@ export function CodeReviewDetailClient({ reviewId }: CodeReviewDetailClientProps
<dd>${(review.total_cost_musd / 1_000_000).toFixed(4)}</dd>
</div>
)}
{(review.total_tokens_in != null || review.total_tokens_out != null) && (
{(data.tokenUsage.input != null ||
data.tokenUsage.output != null ||
data.tokenUsage.cached != null) && (
<div>
<dt className="text-muted-foreground">Tokens</dt>
<dd>
{review.total_tokens_in?.toLocaleString() ?? '—'} in /{' '}
{review.total_tokens_out?.toLocaleString() ?? '—'} out
<dd className="tabular-nums">
Input {formatAvailableTokenCount(data.tokenUsage.input)} / Output{' '}
{formatAvailableTokenCount(data.tokenUsage.output)} / Cached{' '}
{formatAvailableTokenCount(data.tokenUsage.cached)}
</dd>
</div>
)}
Expand Down
46 changes: 44 additions & 2 deletions apps/web/src/app/api/internal/code-review-status/[reviewId]/route.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1317,6 +1317,8 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
model: 'anthropic/claude-sonnet-4.6',
totalTokensIn: 100_001,
totalTokensOut: 0,
totalCachedTokens: 0,
totalUncachedTokens: 100_001,
totalCostMusd: 200_000,
},
},
Expand All @@ -1326,6 +1328,8 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
model: 'anthropic/claude-sonnet-4.6',
totalTokensIn: 60_000,
totalTokensOut: 40_000,
totalCachedTokens: 0,
totalUncachedTokens: 60_000,
totalCostMusd: 200_000,
},
},
Expand All @@ -1335,6 +1339,8 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
model: 'anthropic/claude-sonnet-4.6',
totalTokensIn: 100_001,
totalTokensOut: 0,
totalCachedTokens: 0,
totalUncachedTokens: 100_001,
totalCostMusd: 199_999,
},
},
Expand All @@ -1344,6 +1350,8 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
model: 'anthropic/claude-sonnet-4.6',
totalTokensIn: 99_999,
totalTokensOut: 0,
totalCachedTokens: 0,
totalUncachedTokens: 99_999,
totalCostMusd: 200_000,
},
},
Expand Down Expand Up @@ -1575,6 +1583,8 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
model: 'anthropic/claude-sonnet-4.6',
totalTokensIn: 99_999,
totalTokensOut: 0,
totalCachedTokens: 0,
totalUncachedTokens: 99_999,
totalCostMusd: 200_000,
});

Expand Down Expand Up @@ -1614,6 +1624,8 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
model: 'anthropic/claude-sonnet-4.6',
totalTokensIn: 99_999,
totalTokensOut: 0,
totalCachedTokens: 0,
totalUncachedTokens: 99_999,
totalCostMusd: 199_999,
},
},
Expand Down Expand Up @@ -2762,11 +2774,21 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
repository_review_instructions_used: true,
repository_review_instructions_ref: 'main',
repository_review_instructions_truncated: false,
cli_session_id: 'ses_review_with_cache',
model: 'anthropic/claude-sonnet-4.6',
total_tokens_in: 1000,
total_tokens_out: 200,
completed_at: '2025-01-01T00:10:00Z',
});
mockGetCodeReviewById.mockResolvedValue(review);
mockGetSessionUsageFromBilling.mockResolvedValue({
model: 'openai/gpt-4o',
totalTokensIn: 1000,
totalTokensOut: 200,
totalCachedTokens: 800,
totalUncachedTokens: 200,
totalCostMusd: 100,
});

await POST(makeRequest({ status: 'completed' }), makeParams(REVIEW_ID));

Expand All @@ -2783,8 +2805,23 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
null,
{ maxBodyCharacters: 65_536, reservedCharacters: 8 }
);
expect(mockGetSessionUsageFromBilling).toHaveBeenCalledWith(
'ses_review_with_cache',
'2025-01-01T00:00:00Z',
'2025-01-01T00:10:00Z'
);
expect(mockUpdateCodeReviewUsage).toHaveBeenCalledWith(REVIEW_ID, {
totalTokensIn: 1000,
totalTokensOut: 200,
totalCostMusd: 100,
});
expect(mockAppendReviewSummaryFooter).toHaveBeenCalledWith('existing body', {
usage: { model: 'anthropic/claude-sonnet-4.6', tokensIn: 1000, tokensOut: 200 },
usage: {
model: 'anthropic/claude-sonnet-4.6',
tokensIn: 200,
tokensOut: 200,
cachedTokens: 800,
},
reviewGuidance: { used: true, ref: 'main', truncated: false },
});
expect(mockUpdateKiloReviewComment).toHaveBeenCalledWith(
Expand Down Expand Up @@ -2820,7 +2857,12 @@ describe('POST /api/internal/code-review-status/[reviewId]', () => {
'https://gitlab.com'
);
expect(mockAppendReviewSummaryFooter).toHaveBeenCalledWith('existing note body', {
usage: { model: 'anthropic/claude-sonnet-4.6', tokensIn: 1000, tokensOut: 200 },
usage: {
model: 'anthropic/claude-sonnet-4.6',
tokensIn: 1000,
tokensOut: 200,
cachedTokens: null,
},
reviewGuidance: { used: true, ref: 'main', truncated: true },
});
expect(mockUpdateKiloReviewNote).toHaveBeenCalledWith(
Expand Down
49 changes: 28 additions & 21 deletions apps/web/src/app/api/internal/code-review-status/[reviewId]/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -573,8 +573,9 @@ const BILLING_NOTICE_BODY = `${BILLING_NOTICE_MARKER}
* For v2 (cloud-agent-next) the orchestrator never writes usage — we
* skip the poll and go straight to the billing tables.
*
* When the billing fallback is used we also back-fill the code_reviews
* record so later reads (e.g. admin panel) don't repeat the aggregation.
* Billing usage is preferred when a CLI session is available because the
* persisted review columns do not include cache details. Raw cache-inclusive
* input, output, and cost totals are still back-filled for later reads.
*/
async function getReviewUsageData(reviewId: string) {
let review = await getCodeReviewById(reviewId);
Expand All @@ -591,21 +592,15 @@ async function getReviewUsageData(reviewId: string) {
}
}

if (review?.model) {
return {
model: review.model,
tokensIn: review.total_tokens_in ?? null,
tokensOut: review.total_tokens_out ?? null,
};
}

// Fallback: aggregate from billing tables (covers v2 / cloud-agent-next reviews)
if (review?.cli_session_id && review.created_at) {
const billing = await getSessionUsageFromBilling(review.cli_session_id, review.created_at);
const billing = await getSessionUsageFromBilling(
review.cli_session_id,
review.created_at,
review.completed_at ?? undefined
);
if (billing) {
// Back-fill the code_reviews record so we don't repeat this aggregation
updateCodeReviewUsage(reviewId, {
model: billing.model,
...(review.model == null ? { model: billing.model } : {}),
totalTokensIn: billing.totalTokensIn,
totalTokensOut: billing.totalTokensOut,
totalCostMusd: billing.totalCostMusd,
Expand All @@ -614,14 +609,24 @@ async function getReviewUsageData(reviewId: string) {
});

return {
model: billing.model,
tokensIn: billing.totalTokensIn,
model: review.model ?? billing.model,
tokensIn: billing.totalUncachedTokens,
tokensOut: billing.totalTokensOut,
cachedTokens: billing.totalCachedTokens,
};
}
}

return { model: null, tokensIn: null, tokensOut: null };
if (review?.model) {
return {
model: review.model,
tokensIn: review.total_tokens_in ?? null,

@kilo-code-bot kilo-code-bot Bot Jun 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: "Input" token count has inconsistent semantics between the two return paths of getReviewUsageData.

The billing path (line 609) returns tokensIn: billing.totalUncachedTokens (cache-exclusive), so the footer/details render Input 200 · Cached 800. This fallback returns tokensIn: review.total_tokens_in, which is cache-inclusive: for v2 reviews it is back-filled from billing.totalTokensIn (the raw input_tokens sum, line 600), and for v1 reviews it is the SSE-accumulated prompt_tokens (cache-inclusive per OpenRouter convention, see processUsage.ts:755).

Consequently two reviews with identical actual usage render very differently depending on whether billing rows are still available: a new review shows Input 200 · Cached 800 while an older one shows Input 1000 · Cached —. Since the stated goal of this PR is to stop cached work from looking like new work, the fallback silently re-introduces that confusion for old reviews. Consider normalizing the fallback (e.g. only show total_tokens_in as a combined input when no cached breakdown exists, or relabel it Input (incl. cached)), so the "Input" label means the same thing across reviews.

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

@RSO RSO Jun 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds like something you want to fix?

tokensOut: review.total_tokens_out ?? null,
cachedTokens: null,
};
}

return { model: null, tokensIn: null, tokensOut: null, cachedTokens: null };
}

function getReviewGuidanceFooterData(review: CloudAgentCodeReview) {
Expand Down Expand Up @@ -1455,10 +1460,11 @@ export async function POST(

// Summary history and footer (completed only)
if (status === 'completed') {
const { model, tokensIn, tokensOut } = await getReviewUsageData(reviewId);
const { model, tokensIn, tokensOut, cachedTokens } =
await getReviewUsageData(reviewId);
const usage =
model && tokensIn != null && tokensOut != null
? { model, tokensIn, tokensOut }
? { model, tokensIn, tokensOut, cachedTokens }
: undefined;
const reviewGuidance = getReviewGuidanceFooterData(review);
const summaryFooter = { usage, reviewGuidance };
Expand Down Expand Up @@ -1548,10 +1554,11 @@ export async function POST(

// Summary history and footer (completed only)
if (status === 'completed') {
const { model, tokensIn, tokensOut } = await getReviewUsageData(reviewId);
const { model, tokensIn, tokensOut, cachedTokens } =
await getReviewUsageData(reviewId);
const usage =
model && tokensIn != null && tokensOut != null
? { model, tokensIn, tokensOut }
? { model, tokensIn, tokensOut, cachedTokens }
: undefined;
const reviewGuidance = getReviewGuidanceFooterData(review);

Expand Down
76 changes: 75 additions & 1 deletion apps/web/src/lib/code-reviews/db/code-reviews.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,11 @@ import {
cloud_agent_code_review_attempts,
cloud_agent_code_reviews,
kilocode_users,
microdollar_usage,
microdollar_usage_metadata,
platform_integrations,
} from '@kilocode/db/schema';
import { eq } from 'drizzle-orm';
import { eq, inArray } from 'drizzle-orm';
import { insertTestUser } from '@/tests/helpers/user.helper';
import type { User } from '@kilocode/db/schema';
import {
Expand All @@ -14,6 +16,7 @@ import {
createCodeReviewAttempt,
createInfraRetryAttemptIfMissing,
getCodeReviewAttemptForReview,
getSessionUsageFromBilling,
listCodeReviewAttempts,
updateCodeReviewAttemptForCallback,
findPreviousCompletedReview,
Expand Down Expand Up @@ -618,3 +621,74 @@ describe('findPreviousCompletedReview', () => {
).rejects.toThrow('not found');
});
});

describe('getSessionUsageFromBilling', () => {
const usageIds: string[] = [];

afterEach(async () => {
if (usageIds.length === 0) return;

await db
.delete(microdollar_usage_metadata)
.where(inArray(microdollar_usage_metadata.id, usageIds));
await db.delete(microdollar_usage).where(inArray(microdollar_usage.id, usageIds));
usageIds.length = 0;
});

it('excludes later usage when a completed review session is reused', async () => {
const sessionId = `ses_usage_window_${crypto.randomUUID()}`;
const firstUsageId = crypto.randomUUID();
const laterUsageId = crypto.randomUUID();
usageIds.push(firstUsageId, laterUsageId);

await db.insert(microdollar_usage).values([
{
id: firstUsageId,
kilo_user_id: 'code-review-usage-test',
cost: 100,
input_tokens: 1000,
output_tokens: 100,
cache_write_tokens: 100,
cache_hit_tokens: 600,
created_at: '2026-06-18T10:00:00.000Z',
model: 'anthropic/claude-sonnet-4.6',
},
{
id: laterUsageId,
kilo_user_id: 'code-review-usage-test',
cost: 200,
input_tokens: 2000,
output_tokens: 200,
cache_write_tokens: 200,
cache_hit_tokens: 1200,
created_at: '2026-06-18T12:00:00.000Z',
model: 'openai/gpt-4o',
},
]);
await db.insert(microdollar_usage_metadata).values([
{
id: firstUsageId,
message_id: `msg_${firstUsageId}`,
session_id: sessionId,
created_at: '2026-06-18T10:00:00.000Z',
},
{
id: laterUsageId,
message_id: `msg_${laterUsageId}`,
session_id: sessionId,
created_at: '2026-06-18T12:00:00.000Z',
},
]);

await expect(
getSessionUsageFromBilling(sessionId, '2026-06-18T09:00:00.000Z', '2026-06-18T11:00:00.000Z')
).resolves.toEqual({
model: 'anthropic/claude-sonnet-4.6',
totalTokensIn: 1000,
totalTokensOut: 100,
totalCachedTokens: 700,
totalUncachedTokens: 300,
totalCostMusd: 100,
});
});
});
Loading