diff --git a/src/ccstatusline.ts b/src/ccstatusline.ts index af6a8853..09b7213c 100644 --- a/src/ccstatusline.ts +++ b/src/ccstatusline.ts @@ -40,6 +40,7 @@ import { getPackageVersion, getTerminalWidth } from './utils/terminal'; +import { isWidgetSubagentsEnabled } from './utils/token-subagents'; import { prefetchUsageDataIfNeeded } from './utils/usage-prefetch'; function hasSessionDurationInStatusJson(data: StatusJSON): boolean { @@ -108,6 +109,9 @@ async function renderMultipleLines(data: StatusJSON) { const speedWidgetTypes = new Set(['output-speed', 'input-speed', 'total-speed']); const hasSpeedItems = lines.some(line => line.some(item => speedWidgetTypes.has(item.type))); const hasCompactionWidget = lines.some(line => line.some(item => item.type === 'compaction-counter')); + const subagentTokenWidgetTypes = new Set(['tokens-input', 'tokens-output', 'tokens-cached', 'tokens-total']); + const needsSessionTokens = lines.some(line => line.some(item => item.type === 'tokens-session-total' + || (subagentTokenWidgetTypes.has(item.type) && isWidgetSubagentsEnabled(item)))); const hasThinkingEffortWidget = lines.some(line => line.some(item => item.type === 'thinking-effort')); const hasSessionNameWidget = lines.some(line => line.some(item => item.type === 'session-name')); const needsTranscriptThinkingEffort = hasThinkingEffortWidget @@ -126,6 +130,7 @@ async function renderMultipleLines(data: StatusJSON) { includeSessionDuration: hasSessionClock && !hasSessionDurationInStatusJson(data), includeSpeedMetrics: hasSpeedItems, includeSubagents: true, + includeSubagentTokens: needsSessionTokens, speedWindowSeconds: Array.from(requestedSpeedWindows), includeCompactionStats: hasCompactionWidget, includeThinkingEffort: needsTranscriptThinkingEffort, @@ -139,6 +144,7 @@ async function renderMultipleLines(data: StatusJSON) { ]); const tokenMetrics = transcriptAnalysis?.tokenMetrics ?? null; + const sessionTokenMetrics = transcriptAnalysis?.sessionTokenMetrics ?? null; const sessionDuration = transcriptAnalysis?.sessionDuration ?? null; const speedMetrics = transcriptAnalysis?.speedMetricsCollection?.sessionAverage ?? null; const windowedSpeedMetrics = transcriptAnalysis?.speedMetricsCollection?.windowed ?? null; @@ -156,6 +162,7 @@ async function renderMultipleLines(data: StatusJSON) { const context: RenderContext = { data, tokenMetrics, + sessionTokenMetrics, speedMetrics, windowedSpeedMetrics, usageData, diff --git a/src/types/RenderContext.ts b/src/types/RenderContext.ts index 9e0a0606..21983227 100644 --- a/src/types/RenderContext.ts +++ b/src/types/RenderContext.ts @@ -45,6 +45,7 @@ export interface CompactionData { export interface RenderContext { data?: StatusJSON; tokenMetrics?: TokenMetrics | null; + sessionTokenMetrics?: TokenMetrics | null; speedMetrics?: SpeedMetrics | null; windowedSpeedMetrics?: Record | null; usageData?: RenderUsageData | null; diff --git a/src/utils/__tests__/jsonl-metrics.test.ts b/src/utils/__tests__/jsonl-metrics.test.ts index 348ea432..b63375ce 100644 --- a/src/utils/__tests__/jsonl-metrics.test.ts +++ b/src/utils/__tests__/jsonl-metrics.test.ts @@ -29,6 +29,15 @@ async function getTokenMetrics(transcriptPath: string): Promise { return analysis.tokenMetrics; } +async function getSessionTokenMetrics(transcriptPath: string): Promise { + const analysis = await getTranscriptAnalysis(transcriptPath, { includeSubagentTokens: true }); + if (!analysis.sessionTokenMetrics) { + throw new Error('subagent-inclusive token metrics were requested but not collected'); + } + + return analysis.sessionTokenMetrics; +} + async function getSpeedMetricsCollection( transcriptPath: string, options: { includeSubagents?: boolean; windowSeconds?: number[] } = {} @@ -786,6 +795,122 @@ describe('jsonl transcript metrics', () => { }); }); + it('leaves the main token metrics alone and counts sub-agents only in the session totals', async () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), 'ccstatusline-jsonl-token-sub-')); + tempRoots.push(root); + const transcriptPath = path.join(root, 'main.jsonl'); + const subagentsDir = path.join(root, 'subagents'); + + fs.writeFileSync(transcriptPath, [ + makeUsageLine({ + timestamp: '2026-01-01T10:00:00.000Z', + input: 100, output: 50, cacheRead: 20, cacheCreate: 10 + }), + // Inline sidechain entry that ALSO lives in the separate file below. + makeUsageLine({ + timestamp: '2026-01-01T10:05:00.000Z', + input: 500, output: 60, cacheRead: 5, cacheCreate: 5, + isSidechain: true + }), + JSON.stringify({ type: 'progress', data: { agentId: 'x' } }) + ].join('\n')); + + fs.mkdirSync(subagentsDir, { recursive: true }); + fs.writeFileSync(path.join(subagentsDir, 'agent-x.jsonl'), `${makeUsageLine({ + timestamp: '2026-01-01T10:05:00.000Z', + input: 500, output: 60, cacheRead: 5, cacheCreate: 5, + isSidechain: true + })}\n`); + + // Unchanged behavior: the main metrics count the inline rows and never + // read the separate file. + const mainOnly = await getTokenMetrics(transcriptPath); + expect(mainOnly).toEqual({ + inputTokens: 600, + outputTokens: 110, + cachedTokens: 40, + cacheReadTokens: 25, + cacheCreationTokens: 15, + totalTokens: 750, + contextLength: 130 + }); + + // Session totals: the inline sidechain rows are dropped because the + // file represents them, and the file is added once. + const sessionMetrics = await getSessionTokenMetrics(transcriptPath); + expect(sessionMetrics).toEqual({ + inputTokens: 600, + outputTokens: 110, + cachedTokens: 40, + cacheReadTokens: 25, + cacheCreationTokens: 15, + totalTokens: 750, + contextLength: 130 + }); + }); + + it('keeps inline sidechain entries in the session totals when no sub-agent files exist', async () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), 'ccstatusline-jsonl-token-sub-')); + tempRoots.push(root); + const transcriptPath = path.join(root, 'main-no-files.jsonl'); + + fs.writeFileSync(transcriptPath, [ + makeUsageLine({ + timestamp: '2026-01-01T10:00:00.000Z', + input: 100, output: 50, cacheRead: 20, cacheCreate: 10 + }), + makeUsageLine({ + timestamp: '2026-01-01T10:05:00.000Z', + input: 500, output: 60, cacheRead: 5, cacheCreate: 5, + isSidechain: true + }) + ].join('\n')); + + // Older transcript format: nothing to subtract, so the inline rows stand. + const sessionMetrics = await getSessionTokenMetrics(transcriptPath); + expect(sessionMetrics).toEqual({ + inputTokens: 600, + outputTokens: 110, + cachedTokens: 40, + cacheReadTokens: 25, + cacheCreationTokens: 15, + totalTokens: 750, + contextLength: 130 + }); + }); + + it('sums every referenced sub-agent transcript and ignores unreferenced ones', async () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), 'ccstatusline-jsonl-token-sub-')); + tempRoots.push(root); + const transcriptPath = path.join(root, 'main-multi.jsonl'); + const subagentsDir = path.join(root, 'subagents'); + + fs.writeFileSync(transcriptPath, [ + makeUsageLine({ timestamp: '2026-01-01T10:00:00.000Z', input: 10, output: 5 }), + JSON.stringify({ type: 'progress', data: { agentId: 'a' } }), + JSON.stringify({ type: 'progress', data: { agentId: 'b' } }) + ].join('\n')); + + fs.mkdirSync(subagentsDir, { recursive: true }); + fs.writeFileSync(path.join(subagentsDir, 'agent-a.jsonl'), + `${makeUsageLine({ timestamp: '2026-01-01T10:01:00.000Z', input: 100, output: 200 })}\n`); + fs.writeFileSync(path.join(subagentsDir, 'agent-b.jsonl'), + `${makeUsageLine({ timestamp: '2026-01-01T10:02:00.000Z', input: 30, output: 40 })}\n`); + fs.writeFileSync(path.join(subagentsDir, 'agent-unreferenced.jsonl'), + `${makeUsageLine({ timestamp: '2026-01-01T10:03:00.000Z', input: 9999, output: 9999 })}\n`); + + const sessionMetrics = await getSessionTokenMetrics(transcriptPath); + expect(sessionMetrics).toEqual({ + inputTokens: 140, + outputTokens: 245, + cachedTokens: 0, + cacheReadTokens: 0, + cacheCreationTokens: 0, + totalTokens: 385, + contextLength: 10 + }); + }); + it('calculates speed metrics from user-to-assistant processing windows', async () => { const root = fs.mkdtempSync(path.join(os.tmpdir(), 'ccstatusline-jsonl-speed-')); tempRoots.push(root); diff --git a/src/utils/__tests__/token-subagents.test.ts b/src/utils/__tests__/token-subagents.test.ts new file mode 100644 index 00000000..ff96d5b8 --- /dev/null +++ b/src/utils/__tests__/token-subagents.test.ts @@ -0,0 +1,45 @@ +import { + describe, + expect, + it +} from 'vitest'; + +import type { WidgetItem } from '../../types/Widget'; +import { + SUBAGENTS_MARKER, + isWidgetSubagentsEnabled, + withWidgetSubagentsEnabled +} from '../token-subagents'; + +function makeItem(metadata?: Record): WidgetItem { + return { id: '1', type: 'tokens-input', metadata }; +} + +describe('token-subagents helper', () => { + it('defaults to disabled', () => { + expect(isWidgetSubagentsEnabled(makeItem())).toBe(false); + expect(isWidgetSubagentsEnabled(makeItem({}))).toBe(false); + }); + + it('reads the includeSubagents flag', () => { + expect(isWidgetSubagentsEnabled(makeItem({ includeSubagents: 'true' }))).toBe(true); + expect(isWidgetSubagentsEnabled(makeItem({ includeSubagents: 'false' }))).toBe(false); + }); + + it('enables and clears the flag immutably', () => { + const base = makeItem({ color: 'red' }); + + const enabled = withWidgetSubagentsEnabled(base, true); + expect(enabled).not.toBe(base); + expect(isWidgetSubagentsEnabled(enabled)).toBe(true); + expect(enabled.metadata?.color).toBe('red'); + + const disabled = withWidgetSubagentsEnabled(enabled, false); + expect(isWidgetSubagentsEnabled(disabled)).toBe(false); + expect(disabled.metadata?.includeSubagents).toBeUndefined(); + }); + + it('exposes the sigma marker', () => { + expect(SUBAGENTS_MARKER).toBe('Σ '); + }); +}); diff --git a/src/utils/jsonl-metrics.ts b/src/utils/jsonl-metrics.ts index 88f4c47d..ec0a39b9 100644 --- a/src/utils/jsonl-metrics.ts +++ b/src/utils/jsonl-metrics.ts @@ -38,6 +38,8 @@ export interface TranscriptAnalysisOptions { includeSessionDuration?: boolean; includeSpeedMetrics?: boolean; includeSubagents?: boolean; + /** Also fold sub-agent transcripts into `sessionTokenMetrics`. */ + includeSubagentTokens?: boolean; speedWindowSeconds?: number[]; includeCompactionStats?: boolean; includeThinkingEffort?: boolean; @@ -46,6 +48,8 @@ export interface TranscriptAnalysisOptions { export interface TranscriptAnalysis { tokenMetrics: TokenMetrics; + /** Token metrics including the session's sub-agents; only set when `includeSubagents` is requested. */ + sessionTokenMetrics?: TokenMetrics; sessionDuration: string | null; speedMetricsCollection: SpeedMetricsCollection | null; compactionData: CompactionData | null; @@ -57,6 +61,7 @@ interface TranscriptScanOptions extends TranscriptAnalysisOptions { includeToken interface TranscriptScanResult { tokenMetrics: TokenMetrics | null; + sessionTokenMetrics: TokenMetrics | null; sessionDuration: string | null; speedMetricsCollection: SpeedMetricsCollection | null; compactionData: CompactionData | null; @@ -86,6 +91,7 @@ interface TokenMetricEntry { stopReason: string | null | undefined; timestampMs: number | null; isMainChain: boolean; + isSidechain: boolean; } interface TokenMetricAccumulator { @@ -101,6 +107,10 @@ interface TokenMetricAccumulator { interface TokenMetricState { metrics: TokenMetricAccumulator; + // Sidechain rows counted inside `metrics`. Older transcripts inline + // sub-agent turns here instead of writing separate files; when both exist + // this is subtracted so a sub-agent turn is never counted twice. + sidechainMetrics: TokenMetricAccumulator; hasStopReasonField: boolean; lastUsageEntry: TokenMetricEntry | null; sawCompactBoundary: boolean; @@ -168,6 +178,7 @@ function accumulateTokenMetricEntry( function createTokenMetricState(): TokenMetricState { return { metrics: createTokenMetricAccumulator(), + sidechainMetrics: createTokenMetricAccumulator(), hasStopReasonField: false, lastUsageEntry: null, sawCompactBoundary: false, @@ -192,16 +203,21 @@ function collectTokenMetricRecord(state: TokenMetricState, data: TranscriptLine usage: parseUsageTokens(usage), stopReason: message.stop_reason, timestampMs, - isMainChain: data?.isSidechain !== true && !data?.isApiErrorMessage + isMainChain: data?.isSidechain !== true && !data?.isApiErrorMessage, + isSidechain: data?.isSidechain === true }; const hasStopReason = Object.prototype.hasOwnProperty.call(message, 'stop_reason'); if (hasStopReason && !state.hasStopReasonField) { state.hasStopReasonField = true; state.metrics = createTokenMetricAccumulator(); + state.sidechainMetrics = createTokenMetricAccumulator(); } if (!state.hasStopReasonField || entry.stopReason) { accumulateTokenMetricEntry(state.metrics, entry, !compactBoundary); + if (entry.isSidechain) { + accumulateTokenMetricEntry(state.sidechainMetrics, entry, !compactBoundary); + } } state.lastUsageEntry = entry; state.boundaryAfterLastUsage = compactBoundary; @@ -211,6 +227,9 @@ function collectTokenMetricRecord(state: TokenMetricState, data: TranscriptLine function finishTokenMetrics(state: TokenMetricState): TokenMetrics { if (state.hasStopReasonField && state.lastUsageEntry?.stopReason === null) { accumulateTokenMetricEntry(state.metrics, state.lastUsageEntry, !state.boundaryAfterLastUsage); + if (state.lastUsageEntry.isSidechain) { + accumulateTokenMetricEntry(state.sidechainMetrics, state.lastUsageEntry, !state.boundaryAfterLastUsage); + } } const contextLengthFromUsage = (usage: UsageTokens | null): number | null => usage @@ -232,6 +251,52 @@ function finishTokenMetrics(state: TokenMetricState): TokenMetrics { }; } +/** + * Folds sub-agent transcripts into the session's token counts. + * + * Sub-agent turns reach the metrics one of two ways: inlined into the main + * transcript as `isSidechain` rows (older format), or written to their own + * `subagents/agent-*.jsonl` files. When the files exist they are authoritative, + * so the inlined rows are taken back out first and a turn is never counted + * twice. `contextLength` stays a main-chain concept and is carried over as-is. + */ +function buildSessionTokenMetrics( + mainMetrics: TokenMetrics, + sidechainMetrics: TokenMetricAccumulator, + subagentMetrics: TokenMetrics[] +): TokenMetrics { + const hasSubagentFiles = subagentMetrics.length > 0; + let inputTokens = mainMetrics.inputTokens; + let outputTokens = mainMetrics.outputTokens; + let cacheReadTokens = mainMetrics.cacheReadTokens ?? 0; + let cacheCreationTokens = mainMetrics.cacheCreationTokens ?? 0; + + if (hasSubagentFiles) { + inputTokens -= sidechainMetrics.inputTokens; + outputTokens -= sidechainMetrics.outputTokens; + cacheReadTokens -= sidechainMetrics.cacheReadTokens; + cacheCreationTokens -= sidechainMetrics.cacheCreationTokens; + } + + for (const metrics of subagentMetrics) { + inputTokens += metrics.inputTokens; + outputTokens += metrics.outputTokens; + cacheReadTokens += metrics.cacheReadTokens ?? 0; + cacheCreationTokens += metrics.cacheCreationTokens ?? 0; + } + + const cachedTokens = cacheReadTokens + cacheCreationTokens; + return { + inputTokens, + outputTokens, + cachedTokens, + cacheReadTokens, + cacheCreationTokens, + totalTokens: inputTokens + outputTokens + cachedTokens, + contextLength: mainMetrics.contextLength + }; +} + function collectAgentIds(value: unknown, agentIds: Set) { if (!value || typeof value !== 'object') { return; @@ -364,14 +429,34 @@ function collectSpeedMetricRecord( }); } -async function collectSpeedMetricsFromFile(filePath: string, ignoreSidechain: boolean): Promise { - const state = createSpeedMetricCollector(); +interface SubagentScanResult { + tokenMetrics: TokenMetrics | null; + speedMetrics: CollectedSpeedMetrics | null; +} + +/** One read of a sub-agent transcript, feeding whichever collectors were asked for. */ +async function scanSubagentFile( + filePath: string, + options: { includeTokenMetrics: boolean; includeSpeedMetrics: boolean } +): Promise { + const tokenState = options.includeTokenMetrics ? createTokenMetricState() : null; + const speedState = options.includeSpeedMetrics ? createSpeedMetricCollector() : null; + for await (const line of iterateJsonlLines(filePath)) { const data = parseJsonlLine(line) as TranscriptLine | null; - collectSpeedMetricRecord(state, data, parseTimestampMs(data?.timestamp), ignoreSidechain); + const timestampMs = parseTimestampMs(data?.timestamp); + if (tokenState) { + collectTokenMetricRecord(tokenState, data, timestampMs); + } + if (speedState) { + collectSpeedMetricRecord(speedState, data, timestampMs, false); + } } - return state; + return { + tokenMetrics: tokenState ? finishTokenMetrics(tokenState) : null, + speedMetrics: speedState + }; } function mergeCollectedSpeedMetrics(parts: CollectedSpeedMetrics[]): CollectedSpeedMetrics { @@ -509,6 +594,9 @@ function normalizeSpeedWindows(windowSeconds: number[] | undefined): number[] { function createEmptyScanResult(options: TranscriptScanOptions, speedWindows: number[]): TranscriptScanResult { return { tokenMetrics: options.includeTokenMetrics ? createEmptyTokenMetrics() : null, + sessionTokenMetrics: options.includeTokenMetrics && options.includeSubagentTokens === true + ? createEmptyTokenMetrics() + : null, sessionDuration: null, speedMetricsCollection: options.includeSpeedMetrics ? { @@ -532,7 +620,9 @@ async function scanTranscript(transcriptPath: string, options: TranscriptScanOpt const tokenState = options.includeTokenMetrics ? createTokenMetricState() : null; const speedState = options.includeSpeedMetrics ? createSpeedMetricCollector() : null; const compactionData = options.includeCompactionStats ? createCompactionStats() : null; - const referencedAgentIds = options.includeSpeedMetrics && options.includeSubagents + const needsSubagentTokens = options.includeSubagentTokens === true && options.includeTokenMetrics === true; + const needsSubagentSpeed = options.includeSubagents === true && options.includeSpeedMetrics === true; + const referencedAgentIds = needsSubagentTokens || needsSubagentSpeed ? new Set() : null; let firstTimestampMs: number | null = null; @@ -575,29 +665,46 @@ async function scanTranscript(transcriptPath: string, options: TranscriptScanOpt } } - let speedMetricsCollection: SpeedMetricsCollection | null = null; - if (speedState) { - const collected: CollectedSpeedMetrics[] = [speedState]; - if (referencedAgentIds) { - const subagentPaths = getSubagentTranscriptPaths(transcriptPath, referencedAgentIds); - const subagentMetrics = await Promise.all(subagentPaths.map(async (subagentPath) => { + const subagentResults = referencedAgentIds + ? await Promise.all( + getSubagentTranscriptPaths(transcriptPath, referencedAgentIds).map(async (subagentPath) => { try { - return await collectSpeedMetricsFromFile(subagentPath, false); + return await scanSubagentFile(subagentPath, { + includeTokenMetrics: needsSubagentTokens, + includeSpeedMetrics: needsSubagentSpeed + }); } catch { return null; } - })); - for (const metrics of subagentMetrics) { - if (metrics) { - collected.push(metrics); - } + }) + ) + : []; + + const mainTokenMetrics = tokenState ? finishTokenMetrics(tokenState) : null; + const sessionTokenMetrics = tokenState && needsSubagentTokens + ? buildSessionTokenMetrics( + mainTokenMetrics ?? createEmptyTokenMetrics(), + tokenState.sidechainMetrics, + subagentResults + .map(result => result?.tokenMetrics ?? null) + .filter((metrics): metrics is TokenMetrics => metrics !== null) + ) + : null; + + let speedMetricsCollection: SpeedMetricsCollection | null = null; + if (speedState) { + const collected: CollectedSpeedMetrics[] = [speedState]; + for (const result of subagentResults) { + if (result?.speedMetrics) { + collected.push(result.speedMetrics); } } speedMetricsCollection = buildSpeedMetricsCollection(collected, speedWindows); } return { - tokenMetrics: tokenState ? finishTokenMetrics(tokenState) : null, + tokenMetrics: mainTokenMetrics, + sessionTokenMetrics, sessionDuration: options.includeSessionDuration ? formatSessionDuration(firstTimestampMs, lastTimestampMs) : null, @@ -673,6 +780,7 @@ export async function getTranscriptAnalysis( return { tokenMetrics: result.tokenMetrics ?? createEmptyTokenMetrics(), + ...result.sessionTokenMetrics ? { sessionTokenMetrics: result.sessionTokenMetrics } : {}, sessionDuration: result.sessionDuration, speedMetricsCollection: result.speedMetricsCollection, compactionData: result.compactionData, diff --git a/src/utils/token-subagents.ts b/src/utils/token-subagents.ts new file mode 100644 index 00000000..2a1225e1 --- /dev/null +++ b/src/utils/token-subagents.ts @@ -0,0 +1,37 @@ +import type { RenderContext } from '../types/RenderContext'; +import type { TokenMetrics } from '../types/TokenMetrics'; +import type { WidgetItem } from '../types/Widget'; +import { + isMetadataFlagEnabled, + removeMetadataKeys +} from '../widgets/shared/metadata'; + +export const SUBAGENTS_METADATA_KEY = 'includeSubagents'; +export const SUBAGENTS_MARKER = 'Σ '; + +export function isWidgetSubagentsEnabled(item: WidgetItem): boolean { + return isMetadataFlagEnabled(item, SUBAGENTS_METADATA_KEY); +} + +// Off is the absence of the key, so a widget left at its default stores nothing. +export function withWidgetSubagentsEnabled(item: WidgetItem, on: boolean): WidgetItem { + if (!on) { + return removeMetadataKeys(item, [SUBAGENTS_METADATA_KEY]); + } + + return { + ...item, + metadata: { + ...item.metadata, + [SUBAGENTS_METADATA_KEY]: 'true' + } + }; +} + +// Selects the subagent-inclusive metrics when the widget opts in, otherwise the +// main-only metrics. Returns null when the needed metrics are unavailable. +export function tokenMetricsForWidget(item: WidgetItem, context: RenderContext): TokenMetrics | null { + return isWidgetSubagentsEnabled(item) + ? (context.sessionTokenMetrics ?? null) + : (context.tokenMetrics ?? null); +} diff --git a/src/utils/widget-manifest.ts b/src/utils/widget-manifest.ts index cec52154..ef67d646 100644 --- a/src/utils/widget-manifest.ts +++ b/src/utils/widget-manifest.ts @@ -58,6 +58,7 @@ export const WIDGET_MANIFEST: WidgetManifestEntry[] = [ { type: 'tokens-output', create: () => new widgets.TokensOutputWidget() }, { type: 'tokens-cached', create: () => new widgets.TokensCachedWidget() }, { type: 'tokens-total', create: () => new widgets.TokensTotalWidget() }, + { type: 'tokens-session-total', create: () => new widgets.SessionTotalTokensWidget() }, { type: 'cache-hit-rate', create: () => new widgets.CacheHitRateWidget() }, { type: 'cache-read', create: () => new widgets.CacheReadWidget() }, { type: 'cache-write', create: () => new widgets.CacheWriteWidget() }, diff --git a/src/widgets/SessionTotalTokens.ts b/src/widgets/SessionTotalTokens.ts new file mode 100644 index 00000000..93ba4dc5 --- /dev/null +++ b/src/widgets/SessionTotalTokens.ts @@ -0,0 +1,99 @@ +import type { RenderContext } from '../types/RenderContext'; +import type { Settings } from '../types/Settings'; +import type { + CustomKeybind, + HideableState, + Widget, + WidgetEditorDisplay, + WidgetItem +} from '../types/Widget'; +import { resolveNumberFormat } from '../utils/number-format'; +import { formatTokens } from '../utils/renderer'; +import { SUBAGENTS_MARKER } from '../utils/token-subagents'; + +import { isHidden } from './shared/hideable'; +import { + isMetadataFlagEnabled, + removeMetadataKeys +} from './shared/metadata'; +import { formatRawOrLabeledValue } from './shared/raw-or-labeled'; + +const BREAKDOWN_METADATA_KEY = 'breakdown'; +const ZERO_HIDEABLE_STATE: HideableState = { key: 'zero', label: 'when token count is zero' }; + +function isBreakdownEnabled(item: WidgetItem): boolean { + return isMetadataFlagEnabled(item, BREAKDOWN_METADATA_KEY); +} + +export class SessionTotalTokensWidget implements Widget { + getDefaultColor(): string { return 'cyan'; } + getDescription(): string { return 'Shows total session tokens (input + output + cache) including sub-agents'; } + getDisplayName(): string { return 'Session Total Tokens'; } + getCategory(): string { return 'Tokens'; } + getEditorDisplay(item: WidgetItem): WidgetEditorDisplay { + return isBreakdownEnabled(item) + ? { displayText: this.getDisplayName(), modifierText: '[breakdown]' } + : { displayText: this.getDisplayName() }; + } + + getHideableStates(): HideableState[] { + return [ZERO_HIDEABLE_STATE]; + } + + render(item: WidgetItem, context: RenderContext, settings: Settings): string | null { + const format = resolveNumberFormat('token', item, settings); + const label = `${SUBAGENTS_MARKER}Total: `; + if (context.isPreview) { + const preview = formatRawOrLabeledValue(item, label, formatTokens(152000, format)); + return isBreakdownEnabled(item) && !item.rawValue + ? `${preview} (in ${formatTokens(90000, format)}/out ${formatTokens(40000, format)}/cache ${formatTokens(22000, format)})` + : preview; + } + + const metrics = context.sessionTokenMetrics; + if (!metrics) { + return null; + } + + if (metrics.totalTokens === 0 && isHidden(item, ZERO_HIDEABLE_STATE.key)) { + return null; + } + + const base = formatRawOrLabeledValue(item, label, formatTokens(metrics.totalTokens, format)); + if (isBreakdownEnabled(item) && !item.rawValue) { + return `${base} (in ${formatTokens(metrics.inputTokens, format)}/out ${formatTokens(metrics.outputTokens, format)}/cache ${formatTokens(metrics.cachedTokens, format)})`; + } + + return base; + } + + getCustomKeybinds(): CustomKeybind[] { + return [{ key: 'b', label: '(b)reakdown', action: 'toggle-breakdown' }]; + } + + handleEditorAction(action: string, item: WidgetItem): WidgetItem | null { + if (action !== 'toggle-breakdown') { + return null; + } + + if (isBreakdownEnabled(item)) { + return removeMetadataKeys(item, [BREAKDOWN_METADATA_KEY]); + } + + return { + ...item, + metadata: { + ...item.metadata, + [BREAKDOWN_METADATA_KEY]: 'true' + } + }; + } + + getNumericValue(context: RenderContext): number | null { + return context.sessionTokenMetrics?.totalTokens ?? null; + } + + supportsRawValue(): boolean { return true; } + supportsNumberFormat(): boolean { return true; } + supportsColors(item: WidgetItem): boolean { return true; } +} diff --git a/src/widgets/TokensCached.ts b/src/widgets/TokensCached.ts index 237ac27b..df528433 100644 --- a/src/widgets/TokensCached.ts +++ b/src/widgets/TokensCached.ts @@ -1,6 +1,7 @@ import type { RenderContext } from '../types/RenderContext'; import type { Settings } from '../types/Settings'; import type { + CustomKeybind, HideableState, Widget, WidgetEditorDisplay, @@ -8,6 +9,12 @@ import type { } from '../types/Widget'; import { resolveNumberFormat } from '../utils/number-format'; import { formatTokens } from '../utils/renderer'; +import { + SUBAGENTS_MARKER, + isWidgetSubagentsEnabled, + tokenMetricsForWidget, + withWidgetSubagentsEnabled +} from '../utils/token-subagents'; import { isHidden } from './shared/hideable'; import { formatRawOrLabeledValue } from './shared/raw-or-labeled'; @@ -20,7 +27,9 @@ export class TokensCachedWidget implements Widget { getDisplayName(): string { return 'Tokens Cached'; } getCategory(): string { return 'Tokens'; } getEditorDisplay(item: WidgetItem): WidgetEditorDisplay { - return { displayText: this.getDisplayName() }; + return isWidgetSubagentsEnabled(item) + ? { displayText: this.getDisplayName(), modifierText: '[+sub]' } + : { displayText: this.getDisplayName() }; } getHideableStates(): HideableState[] { @@ -29,19 +38,34 @@ export class TokensCachedWidget implements Widget { render(item: WidgetItem, context: RenderContext, settings: Settings): string | null { const format = resolveNumberFormat('token', item, settings); + const label = isWidgetSubagentsEnabled(item) ? `${SUBAGENTS_MARKER}Cached: ` : 'Cached: '; if (context.isPreview) { - return formatRawOrLabeledValue(item, 'Cached: ', formatTokens(12000, format)); + return formatRawOrLabeledValue(item, label, formatTokens(12000, format)); } - if (context.tokenMetrics) { - if (context.tokenMetrics.cachedTokens === 0 && isHidden(item, ZERO_HIDEABLE_STATE.key)) { + // Subagent-inclusive metrics when the widget opts in, main-only otherwise. + const metrics = tokenMetricsForWidget(item, context); + if (metrics) { + if (metrics.cachedTokens === 0 && isHidden(item, ZERO_HIDEABLE_STATE.key)) { return null; } - return formatRawOrLabeledValue(item, 'Cached: ', formatTokens(context.tokenMetrics.cachedTokens, format)); + return formatRawOrLabeledValue(item, label, formatTokens(metrics.cachedTokens, format)); } return null; } + getCustomKeybinds(): CustomKeybind[] { + return [{ key: 's', label: '(s)ubagents', action: 'toggle-subagents' }]; + } + + handleEditorAction(action: string, item: WidgetItem): WidgetItem | null { + if (action !== 'toggle-subagents') { + return null; + } + + return withWidgetSubagentsEnabled(item, !isWidgetSubagentsEnabled(item)); + } + supportsRawValue(): boolean { return true; } supportsColors(item: WidgetItem): boolean { return true; } supportsNumberFormat(): boolean { return true; } diff --git a/src/widgets/TokensInput.ts b/src/widgets/TokensInput.ts index 1266deac..7d162d6d 100644 --- a/src/widgets/TokensInput.ts +++ b/src/widgets/TokensInput.ts @@ -1,6 +1,7 @@ import type { RenderContext } from '../types/RenderContext'; import type { Settings } from '../types/Settings'; import type { + CustomKeybind, HideableState, Widget, WidgetEditorDisplay, @@ -9,6 +10,12 @@ import type { import { getContextWindowInputTotalTokens } from '../utils/context-window'; import { resolveNumberFormat } from '../utils/number-format'; import { formatTokens } from '../utils/renderer'; +import { + SUBAGENTS_MARKER, + isWidgetSubagentsEnabled, + tokenMetricsForWidget, + withWidgetSubagentsEnabled +} from '../utils/token-subagents'; import { isHidden } from './shared/hideable'; import { formatRawOrLabeledValue } from './shared/raw-or-labeled'; @@ -21,7 +28,9 @@ export class TokensInputWidget implements Widget { getDisplayName(): string { return 'Tokens Input'; } getCategory(): string { return 'Tokens'; } getEditorDisplay(item: WidgetItem): WidgetEditorDisplay { - return { displayText: this.getDisplayName() }; + return isWidgetSubagentsEnabled(item) + ? { displayText: this.getDisplayName(), modifierText: '[+sub]' } + : { displayText: this.getDisplayName() }; } getHideableStates(): HideableState[] { @@ -30,12 +39,16 @@ export class TokensInputWidget implements Widget { render(item: WidgetItem, context: RenderContext, settings: Settings): string | null { const format = resolveNumberFormat('token', item, settings); + const subagents = isWidgetSubagentsEnabled(item); + const label = subagents ? `${SUBAGENTS_MARKER}In: ` : 'In: '; if (context.isPreview) { - return formatRawOrLabeledValue(item, 'In: ', formatTokens(15200, format)); + return formatRawOrLabeledValue(item, label, formatTokens(15200, format)); } - const inputTotalTokens = context.tokenMetrics?.inputTokens - ?? getContextWindowInputTotalTokens(context.data) + // The status JSON's context_window is main-agent only, so it is only a + // fallback while this widget counts the main agent. + const inputTotalTokens = tokenMetricsForWidget(item, context)?.inputTokens + ?? (subagents ? null : getContextWindowInputTotalTokens(context.data)) ?? null; if (inputTotalTokens === null) { return null; @@ -45,7 +58,19 @@ export class TokensInputWidget implements Widget { return null; } - return formatRawOrLabeledValue(item, 'In: ', formatTokens(inputTotalTokens, format)); + return formatRawOrLabeledValue(item, label, formatTokens(inputTotalTokens, format)); + } + + getCustomKeybinds(): CustomKeybind[] { + return [{ key: 's', label: '(s)ubagents', action: 'toggle-subagents' }]; + } + + handleEditorAction(action: string, item: WidgetItem): WidgetItem | null { + if (action !== 'toggle-subagents') { + return null; + } + + return withWidgetSubagentsEnabled(item, !isWidgetSubagentsEnabled(item)); } supportsRawValue(): boolean { return true; } diff --git a/src/widgets/TokensOutput.ts b/src/widgets/TokensOutput.ts index 43dc46d4..897cd30d 100644 --- a/src/widgets/TokensOutput.ts +++ b/src/widgets/TokensOutput.ts @@ -1,6 +1,7 @@ import type { RenderContext } from '../types/RenderContext'; import type { Settings } from '../types/Settings'; import type { + CustomKeybind, HideableState, Widget, WidgetEditorDisplay, @@ -9,6 +10,12 @@ import type { import { getContextWindowOutputTotalTokens } from '../utils/context-window'; import { resolveNumberFormat } from '../utils/number-format'; import { formatTokens } from '../utils/renderer'; +import { + SUBAGENTS_MARKER, + isWidgetSubagentsEnabled, + tokenMetricsForWidget, + withWidgetSubagentsEnabled +} from '../utils/token-subagents'; import { isHidden } from './shared/hideable'; import { formatRawOrLabeledValue } from './shared/raw-or-labeled'; @@ -21,7 +28,9 @@ export class TokensOutputWidget implements Widget { getDisplayName(): string { return 'Tokens Output'; } getCategory(): string { return 'Tokens'; } getEditorDisplay(item: WidgetItem): WidgetEditorDisplay { - return { displayText: this.getDisplayName() }; + return isWidgetSubagentsEnabled(item) + ? { displayText: this.getDisplayName(), modifierText: '[+sub]' } + : { displayText: this.getDisplayName() }; } getHideableStates(): HideableState[] { @@ -30,12 +39,16 @@ export class TokensOutputWidget implements Widget { render(item: WidgetItem, context: RenderContext, settings: Settings): string | null { const format = resolveNumberFormat('token', item, settings); + const subagents = isWidgetSubagentsEnabled(item); + const label = subagents ? `${SUBAGENTS_MARKER}Out: ` : 'Out: '; if (context.isPreview) { - return formatRawOrLabeledValue(item, 'Out: ', formatTokens(3400, format)); + return formatRawOrLabeledValue(item, label, formatTokens(3400, format)); } - const outputTotalTokens = context.tokenMetrics?.outputTokens - ?? getContextWindowOutputTotalTokens(context.data) + // The status JSON's context_window is main-agent only, so it is only a + // fallback while this widget counts the main agent. + const outputTotalTokens = tokenMetricsForWidget(item, context)?.outputTokens + ?? (subagents ? null : getContextWindowOutputTotalTokens(context.data)) ?? null; if (outputTotalTokens === null) { return null; @@ -45,7 +58,19 @@ export class TokensOutputWidget implements Widget { return null; } - return formatRawOrLabeledValue(item, 'Out: ', formatTokens(outputTotalTokens, format)); + return formatRawOrLabeledValue(item, label, formatTokens(outputTotalTokens, format)); + } + + getCustomKeybinds(): CustomKeybind[] { + return [{ key: 's', label: '(s)ubagents', action: 'toggle-subagents' }]; + } + + handleEditorAction(action: string, item: WidgetItem): WidgetItem | null { + if (action !== 'toggle-subagents') { + return null; + } + + return withWidgetSubagentsEnabled(item, !isWidgetSubagentsEnabled(item)); } supportsRawValue(): boolean { return true; } diff --git a/src/widgets/TokensTotal.ts b/src/widgets/TokensTotal.ts index a780ad31..4d7c663d 100644 --- a/src/widgets/TokensTotal.ts +++ b/src/widgets/TokensTotal.ts @@ -1,6 +1,7 @@ import type { RenderContext } from '../types/RenderContext'; import type { Settings } from '../types/Settings'; import type { + CustomKeybind, HideableState, Widget, WidgetEditorDisplay, @@ -8,6 +9,12 @@ import type { } from '../types/Widget'; import { resolveNumberFormat } from '../utils/number-format'; import { formatTokens } from '../utils/renderer'; +import { + SUBAGENTS_MARKER, + isWidgetSubagentsEnabled, + tokenMetricsForWidget, + withWidgetSubagentsEnabled +} from '../utils/token-subagents'; import { isHidden } from './shared/hideable'; import { formatRawOrLabeledValue } from './shared/raw-or-labeled'; @@ -20,7 +27,9 @@ export class TokensTotalWidget implements Widget { getDisplayName(): string { return 'Tokens Total'; } getCategory(): string { return 'Tokens'; } getEditorDisplay(item: WidgetItem): WidgetEditorDisplay { - return { displayText: this.getDisplayName() }; + return isWidgetSubagentsEnabled(item) + ? { displayText: this.getDisplayName(), modifierText: '[+sub]' } + : { displayText: this.getDisplayName() }; } getHideableStates(): HideableState[] { @@ -29,19 +38,34 @@ export class TokensTotalWidget implements Widget { render(item: WidgetItem, context: RenderContext, settings: Settings): string | null { const format = resolveNumberFormat('token', item, settings); + const label = isWidgetSubagentsEnabled(item) ? `${SUBAGENTS_MARKER}Total: ` : 'Total: '; if (context.isPreview) { - return formatRawOrLabeledValue(item, 'Total: ', formatTokens(30600, format)); + return formatRawOrLabeledValue(item, label, formatTokens(30600, format)); } - if (context.tokenMetrics) { - if (context.tokenMetrics.totalTokens === 0 && isHidden(item, ZERO_HIDEABLE_STATE.key)) { + // Subagent-inclusive metrics when the widget opts in, main-only otherwise. + const metrics = tokenMetricsForWidget(item, context); + if (metrics) { + if (metrics.totalTokens === 0 && isHidden(item, ZERO_HIDEABLE_STATE.key)) { return null; } - return formatRawOrLabeledValue(item, 'Total: ', formatTokens(context.tokenMetrics.totalTokens, format)); + return formatRawOrLabeledValue(item, label, formatTokens(metrics.totalTokens, format)); } return null; } + getCustomKeybinds(): CustomKeybind[] { + return [{ key: 's', label: '(s)ubagents', action: 'toggle-subagents' }]; + } + + handleEditorAction(action: string, item: WidgetItem): WidgetItem | null { + if (action !== 'toggle-subagents') { + return null; + } + + return withWidgetSubagentsEnabled(item, !isWidgetSubagentsEnabled(item)); + } + supportsRawValue(): boolean { return true; } supportsColors(item: WidgetItem): boolean { return true; } supportsNumberFormat(): boolean { return true; } diff --git a/src/widgets/__tests__/session-total-tokens.test.ts b/src/widgets/__tests__/session-total-tokens.test.ts new file mode 100644 index 00000000..05548b2d --- /dev/null +++ b/src/widgets/__tests__/session-total-tokens.test.ts @@ -0,0 +1,74 @@ +import { + describe, + expect, + it +} from 'vitest'; + +import type { RenderContext } from '../../types/RenderContext'; +import type { Settings } from '../../types/Settings'; +import type { WidgetItem } from '../../types/Widget'; +// Statically load renderer first so the eager widget registry initializes before +// the widget module is dynamically imported (avoids a circular-import init order). +import '../../utils/renderer'; + +async function loadWidget() { + const { SessionTotalTokensWidget } = await import('../SessionTotalTokens'); + return SessionTotalTokensWidget; +} + +const settings = {} as Settings; +const sessionMetrics = { + inputTokens: 90000, outputTokens: 40000, cachedTokens: 22000, + totalTokens: 152000, contextLength: 50000 +}; + +function item(metadata?: Record, rawValue?: boolean): WidgetItem { + return { id: '1', type: 'tokens-session-total', metadata, rawValue }; +} + +function ctx(over: Partial = {}): RenderContext { + return { sessionTokenMetrics: sessionMetrics, ...over }; +} + +describe('SessionTotalTokens widget', () => { + it('renders the session grand total with the Σ marker', async () => { + const SessionTotalTokensWidget = await loadWidget(); + const w = new SessionTotalTokensWidget(); + expect(w.render(item(), ctx(), settings)).toBe('Σ Total: 152.0k'); + }); + + it('renders raw value without marker or label', async () => { + const SessionTotalTokensWidget = await loadWidget(); + const w = new SessionTotalTokensWidget(); + expect(w.render(item(undefined, true), ctx(), settings)).toBe('152.0k'); + }); + + it('appends a breakdown when enabled', async () => { + const SessionTotalTokensWidget = await loadWidget(); + const w = new SessionTotalTokensWidget(); + const out = w.render(item({ breakdown: 'true' }), ctx(), settings); + expect(out).toBe('Σ Total: 152.0k (in 90.0k/out 40.0k/cache 22.0k)'); + }); + + it('returns null when session metrics are unavailable', async () => { + const SessionTotalTokensWidget = await loadWidget(); + const w = new SessionTotalTokensWidget(); + expect(w.render(item(), ctx({ sessionTokenMetrics: null }), settings)).toBeNull(); + }); + + it('toggles the breakdown flag through handleEditorAction', async () => { + const SessionTotalTokensWidget = await loadWidget(); + const w = new SessionTotalTokensWidget(); + const on = w.handleEditorAction('toggle-breakdown', item()); + expect(on?.metadata?.breakdown).toBe('true'); + const off = w.handleEditorAction('toggle-breakdown', on ?? item()); + expect(off?.metadata?.breakdown).toBeUndefined(); + }); + + it('exposes the total as its numeric value', async () => { + const SessionTotalTokensWidget = await loadWidget(); + const w = new SessionTotalTokensWidget(); + expect(w.getNumericValue(ctx())).toBe(152000); + expect(w.getNumericValue(ctx({ sessionTokenMetrics: null }))).toBeNull(); + }); +}); diff --git a/src/widgets/__tests__/token-widgets-subagents.test.ts b/src/widgets/__tests__/token-widgets-subagents.test.ts new file mode 100644 index 00000000..743cc0de --- /dev/null +++ b/src/widgets/__tests__/token-widgets-subagents.test.ts @@ -0,0 +1,100 @@ +import { + describe, + expect, + it +} from 'vitest'; + +import type { RenderContext } from '../../types/RenderContext'; +import type { Settings } from '../../types/Settings'; +import type { WidgetItem } from '../../types/Widget'; +// Statically load renderer first so the eager widget registry initializes before +// any widget module is dynamically imported (avoids a circular-import init order). +import '../../utils/renderer'; + +async function loadWidgets() { + const [{ TokensInputWidget }, { TokensTotalWidget }] = await Promise.all([ + import('../TokensInput'), + import('../TokensTotal') + ]); + return { TokensInputWidget, TokensTotalWidget }; +} + +const settings = {} as Settings; + +function item(metadata?: Record, rawValue?: boolean): WidgetItem { + return { id: '1', type: 'tokens-input', metadata, rawValue }; +} + +const baseMetrics = { + inputTokens: 90000, outputTokens: 40000, cachedTokens: 22000, + totalTokens: 152000, contextLength: 50000 +}; +const sessionMetrics = { + inputTokens: 150000, outputTokens: 60000, cachedTokens: 30000, + totalTokens: 240000, contextLength: 50000 +}; + +function ctx(over: Partial = {}): RenderContext { + return { + tokenMetrics: baseMetrics, + sessionTokenMetrics: sessionMetrics, + ...over + }; +} + +describe('token widgets — subagents toggle', () => { + it('TokensInput uses main metrics and no marker when disabled', async () => { + const { TokensInputWidget } = await loadWidgets(); + const w = new TokensInputWidget(); + expect(w.render(item(), ctx(), settings)).toBe('In: 90.0k'); + }); + + it('TokensInput uses session metrics and Σ marker when enabled', async () => { + const { TokensInputWidget } = await loadWidgets(); + const w = new TokensInputWidget(); + const out = w.render(item({ includeSubagents: 'true' }), ctx(), settings); + expect(out).toBe('Σ In: 150.0k'); + }); + + it('TokensInput omits the marker in raw mode but still uses session metrics', async () => { + const { TokensInputWidget } = await loadWidgets(); + const w = new TokensInputWidget(); + const out = w.render(item({ includeSubagents: 'true' }, true), ctx(), settings); + expect(out).toBe('150.0k'); + }); + + it('TokensInput ignores the stdin context_window total when subagents are on', async () => { + const { TokensInputWidget } = await loadWidgets(); + const w = new TokensInputWidget(); + const out = w.render( + item({ includeSubagents: 'true' }), + ctx({ data: { context_window: { total_input_tokens: 12345 } } }), + settings + ); + // Must come from sessionTokenMetrics (150.0k), not the 12345 payload value. + expect(out).toBe('Σ In: 150.0k'); + }); + + it('TokensTotal switches to session total with marker when enabled', async () => { + const { TokensTotalWidget } = await loadWidgets(); + const w = new TokensTotalWidget(); + expect(w.render(item(), ctx(), settings)).toBe('Total: 152.0k'); + expect(w.render(item({ includeSubagents: 'true' }), ctx(), settings)).toBe('Σ Total: 240.0k'); + }); + + it('toggles the flag through handleEditorAction', async () => { + const { TokensTotalWidget } = await loadWidgets(); + const w = new TokensTotalWidget(); + const on = w.handleEditorAction('toggle-subagents', item()); + expect(on?.metadata?.includeSubagents).toBe('true'); + const off = w.handleEditorAction('toggle-subagents', on ?? item()); + expect(off?.metadata?.includeSubagents).toBeUndefined(); + }); + + it('reports the [+sub] modifier in the editor display when enabled', async () => { + const { TokensTotalWidget } = await loadWidgets(); + const w = new TokensTotalWidget(); + expect(w.getEditorDisplay(item()).modifierText).toBeUndefined(); + expect(w.getEditorDisplay(item({ includeSubagents: 'true' })).modifierText).toBe('[+sub]'); + }); +}); diff --git a/src/widgets/index.ts b/src/widgets/index.ts index c4113d0e..d240fa8f 100644 --- a/src/widgets/index.ts +++ b/src/widgets/index.ts @@ -30,6 +30,7 @@ export { TokensInputWidget } from './TokensInput'; export { TokensOutputWidget } from './TokensOutput'; export { TokensCachedWidget } from './TokensCached'; export { TokensTotalWidget } from './TokensTotal'; +export { SessionTotalTokensWidget } from './SessionTotalTokens'; export { CacheHitRateWidget } from './CacheHitRate'; export { CacheReadWidget } from './CacheRead'; export { CacheWriteWidget } from './CacheWrite';