Skip to content

Commit fb715a4

Browse files
rajbosrajbos
authored
Merge pull request #593 from rajbos/rajbos/add-token-count-diagnostics
feat: add token count column to diagnostics session files tab
2 parents f1a5e97 + f8f1ada commit fb715a4

3 files changed

Lines changed: 72 additions & 24 deletions

File tree

‎vscode-extension/src/extension.ts‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3138,12 +3138,14 @@ class CopilotTokenTracker implements vscode.Disposable {
31383138
lastInteraction = stat.mtime.toISOString();
31393139
}
31403140

3141-
// Reconstruct SessionFileDetails from cache
3141+
// Reconstruct SessionFileDetails from cache.
3142+
// Prefer actualTokens (real API count) when available; fall back to estimated tokens.
31423143
const details: SessionFileDetails = {
31433144
file: sessionFile,
31443145
size: cached.size || stat.size,
31453146
modified: stat.mtime.toISOString(),
31463147
interactions: cached.interactions,
3148+
tokens: cached.actualTokens || cached.tokens || 0,
31473149
contextReferences: cached.usageAnalysis.contextReferences,
31483150
firstInteraction: cached.firstInteraction || null,
31493151
lastInteraction: lastInteraction,
@@ -3170,6 +3172,10 @@ class CopilotTokenTracker implements vscode.Disposable {
31703172
// Get existing cache entry if available
31713173
const existingCache = this.getCachedSessionData(sessionFile);
31723174

3175+
// Enrich details with token count from existing cache (populated by main stats calculation).
3176+
// Prefer actualTokens (real API count) when available; fall back to estimated tokens.
3177+
details.tokens = existingCache?.actualTokens || existingCache?.tokens || 0;
3178+
31733179
// Create or update cache entry
31743180
const cacheEntry: SessionFileCache = {
31753181
tokens: existingCache?.tokens || 0,
@@ -4135,13 +4141,25 @@ class CopilotTokenTracker implements vscode.Disposable {
41354141
let cliSessionModel = 'gpt-4o';
41364142
let cliSessionEffort: string | undefined;
41374143

4138-
// Pre-scan for session.start to extract default model and effort
4144+
// Pre-scan for model and effort:
4145+
// 1. session.start.data.selectedModel (older CLI format)
4146+
// 2. First tool.execution_complete.data.model (newer CLI format — session.start has no selectedModel)
4147+
let cliModelFound = false;
41394148
for (const line of lines) {
41404149
try {
41414150
const ev = JSON.parse(line);
41424151
if (ev.type === 'session.start' && ev.data) {
4143-
if (typeof ev.data.selectedModel === 'string') { cliSessionModel = ev.data.selectedModel; }
4152+
if (typeof ev.data.selectedModel === 'string') {
4153+
cliSessionModel = ev.data.selectedModel;
4154+
cliModelFound = true;
4155+
}
41444156
if (typeof ev.data.reasoningEffort === 'string') { cliSessionEffort = ev.data.reasoningEffort; }
4157+
if (cliModelFound) { break; }
4158+
// No model in session.start — continue scanning for tool.execution_complete
4159+
}
4160+
// Newer format: model stored per tool call result
4161+
if (ev.type === 'tool.execution_complete' && typeof ev.data?.model === 'string') {
4162+
cliSessionModel = ev.data.model;
41454163
break;
41464164
}
41474165
} catch { /* skip */ }

‎vscode-extension/src/types.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,7 @@ export interface SessionFileDetails {
314314
size: number;
315315
modified: string;
316316
interactions: number;
317+
tokens?: number; // estimated token count for the session
317318
contextReferences: ContextReferenceUsage;
318319
firstInteraction: string | null;
319320
lastInteraction: string | null;

‎vscode-extension/src/webview/diagnostics/main.ts‎

Lines changed: 50 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ type SessionFileDetails = {
2424
size: number;
2525
modified: string;
2626
interactions: number;
27+
tokens?: number;
2728
contextReferences: ContextReferenceUsage;
2829
firstInteraction: string | null;
2930
lastInteraction: string | null;
@@ -102,7 +103,7 @@ const vscode = acquireVsCodeApi<DiagnosticsViewState>();
102103
const initialData = window.__INITIAL_DIAGNOSTICS__;
103104

104105
// Sorting and filtering state
105-
let currentSortColumn: "lastInteraction" = "lastInteraction";
106+
let currentSortColumn: "lastInteraction" | "size" | "tokens" | "interactions" | "contextRefs" = "lastInteraction";
106107
let currentSortDirection: "asc" | "desc" = "desc";
107108
let currentEditorFilter: string | null = null; // null = show all
108109
let currentContextRefFilter: keyof ContextReferenceUsage | null = null; // null = show all
@@ -188,6 +189,14 @@ function sanitizeNumber(value: number | undefined | null): string {
188189
return Math.floor(n).toString();
189190
}
190191

192+
function formatTokenCount(value: number | undefined | null): string {
193+
const n = Number(value ?? 0);
194+
if (!Number.isFinite(n) || n === 0) { return "0"; }
195+
if (n >= 1_000_000) { return `${(n / 1_000_000).toFixed(1)}M`; }
196+
if (n >= 1_000) { return `${(n / 1_000).toFixed(1)}K`; }
197+
return Math.floor(n).toString();
198+
}
199+
191200
/**
192201
* Build a DOM element showing all candidate paths the extension considers,
193202
* with their existence status. Helps users understand why data may be missing.
@@ -409,28 +418,38 @@ function getEditorIcon(editor: string): string {
409418

410419
function sortSessionFiles(files: SessionFileDetails[]): SessionFileDetails[] {
411420
return [...files].sort((a, b) => {
412-
const aVal = a.lastInteraction;
413-
const bVal = b.lastInteraction;
414-
415-
// Handle null values - push them to the end
416-
if (!aVal && !bVal) {
421+
let aNum: number;
422+
let bNum: number;
423+
424+
if (currentSortColumn === "lastInteraction") {
425+
const aVal = a.lastInteraction;
426+
const bVal = b.lastInteraction;
427+
if (!aVal && !bVal) { return 0; }
428+
if (!aVal) { return 1; }
429+
if (!bVal) { return -1; }
430+
aNum = new Date(aVal).getTime();
431+
bNum = new Date(bVal).getTime();
432+
} else if (currentSortColumn === "size") {
433+
aNum = a.size || 0;
434+
bNum = b.size || 0;
435+
} else if (currentSortColumn === "tokens") {
436+
aNum = a.tokens || 0;
437+
bNum = b.tokens || 0;
438+
} else if (currentSortColumn === "interactions") {
439+
aNum = a.interactions || 0;
440+
bNum = b.interactions || 0;
441+
} else if (currentSortColumn === "contextRefs") {
442+
aNum = getTotalContextRefs(a.contextReferences);
443+
bNum = getTotalContextRefs(b.contextReferences);
444+
} else {
417445
return 0;
418446
}
419-
if (!aVal) {
420-
return 1;
421-
}
422-
if (!bVal) {
423-
return -1;
424-
}
425447

426-
const aTime = new Date(aVal).getTime();
427-
const bTime = new Date(bVal).getTime();
428-
429-
return currentSortDirection === "desc" ? bTime - aTime : aTime - bTime;
448+
return currentSortDirection === "desc" ? bNum - aNum : aNum - bNum;
430449
});
431450
}
432451

433-
function getSortIndicator(column: "lastInteraction"): string {
452+
function getSortIndicator(column: typeof currentSortColumn): string {
434453
if (currentSortColumn !== column) {
435454
return "";
436455
}
@@ -509,6 +528,10 @@ function renderSessionTable(
509528
(sum, sf) => sum + Number(sf.interactions || 0),
510529
0,
511530
);
531+
const totalTokens = filteredFiles.reduce(
532+
(sum, sf) => sum + Number(sf.tokens || 0),
533+
0,
534+
);
512535
const totalContextRefs = filteredFiles.reduce(
513536
(sum, sf) => sum + getTotalContextRefs(sf.contextReferences),
514537
0,
@@ -581,6 +604,10 @@ function renderSessionTable(
581604
<div class="summary-label">💬 Interactions</div>
582605
<div class="summary-value">${totalInteractions}</div>
583606
</div>
607+
<div class="summary-card">
608+
<div class="summary-label">🪙 Tokens</div>
609+
<div class="summary-value" title="${totalTokens.toLocaleString()} tokens">${formatTokenCount(totalTokens)}</div>
610+
</div>
584611
<div class="summary-card">
585612
<div class="summary-label">🔗 Context References</div>
586613
<div class="summary-value">${safeText(totalContextRefs)}</div>
@@ -617,9 +644,10 @@ function renderSessionTable(
617644
<th>Editor</th>
618645
<th>Title</th>
619646
<th>Repository</th>
620-
<th>Size</th>
621-
<th>Interactions</th>
622-
<th>Context Refs</th>
647+
<th class="sortable" data-sort="size">Size${getSortIndicator("size")}</th>
648+
<th class="sortable" data-sort="tokens">Tokens${getSortIndicator("tokens")}</th>
649+
<th class="sortable" data-sort="interactions">Interactions${getSortIndicator("interactions")}</th>
650+
<th class="sortable" data-sort="contextRefs">Context Refs${getSortIndicator("contextRefs")}</th>
623651
<th class="sortable" data-sort="lastInteraction">Last Interaction${getSortIndicator("lastInteraction")}</th>
624652
<th>Actions</th>
625653
</tr>
@@ -636,6 +664,7 @@ function renderSessionTable(
636664
</td>
637665
<td class="repository-cell" title="${sf.repository ? escapeHtml(sf.repository) : "No repository detected"}">${sf.repository ? escapeHtml(getRepoDisplayName(sf.repository)) : '<span style="color: #666;">—</span>'}</td>
638666
<td>${formatFileSize(sf.size)}</td>
667+
<td title="${Number(sf.tokens || 0).toLocaleString()} tokens">${formatTokenCount(sf.tokens)}</td>
639668
<td>${sanitizeNumber(sf.interactions)}</td>
640669
<td title="${escapeHtml(getContextRefsSummary(sf.contextReferences))}">${sanitizeNumber(getTotalContextRefs(sf.contextReferences))}</td>
641670
<td>${formatDate(sf.lastInteraction)}</td>
@@ -1556,7 +1585,7 @@ function renderLayout(data: DiagnosticsData): void {
15561585
header.addEventListener("click", () => {
15571586
const sortColumn = (header as HTMLElement).getAttribute(
15581587
"data-sort",
1559-
) as "lastInteraction";
1588+
) as typeof currentSortColumn;
15601589
if (sortColumn) {
15611590
// Toggle direction if same column, otherwise default to desc
15621591
if (currentSortColumn === sortColumn) {

0 commit comments

Comments
 (0)