Skip to content

Commit 5169041

Browse files
eshurakoveshurakov
authored
feat(cloud-agent-next): show context window usage (#3635)
* feat(cloud-agent-next): show context window usage * fix(cloud-agent-next): guard indexed context usage access
1 parent 03c5edc commit 5169041

10 files changed

Lines changed: 685 additions & 12 deletions

‎apps/web/src/components/cloud-agent-next/CloudChatPage.tsx‎

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ import {
3737
} from './terminal-tabs';
3838
import { isMessageStreaming } from './types';
3939
import { useOrganizationModels } from './hooks/useOrganizationModels';
40+
import { ContextUsageIndicator } from './ContextUsageIndicator';
41+
import { resolveContextWindow } from './model-context-lengths';
4042
import { useSlashCommandSets } from '@/hooks/useSlashCommandSets';
4143
import { useCelebrationSound } from '@/hooks/useCelebrationSound';
4244
import type { CloudAgentAttachments } from '@/lib/cloud-agent/constants';
@@ -216,6 +218,7 @@ export default function CloudChatPage({ organizationId }: CloudChatPageProps) {
216218
const dynamicMessages = useAtomValue(manager.atoms.dynamicMessages);
217219
const pendingMessages = useAtomValue(manager.atoms.pendingMessages);
218220
const totalCost = useAtomValue(manager.atoms.totalCost);
221+
const contextUsage = useAtomValue(manager.atoms.contextUsage);
219222
const getChildMessages = useAtomValue(manager.atoms.childMessages);
220223
const fetchedSessionData = useAtomValue(manager.atoms.fetchedSessionData);
221224

@@ -234,7 +237,9 @@ export default function CloudChatPage({ organizationId }: CloudChatPageProps) {
234237
}, [sessionId]);
235238

236239
// -- Organization models --------------------------------------------------
237-
const { modelOptions, isLoadingModels } = useOrganizationModels(organizationId);
240+
const { modelOptions, isLoadingModels, contextLengthByModelId } =
241+
useOrganizationModels(organizationId);
242+
const contextWindow = resolveContextWindow(contextUsage, contextLengthByModelId);
238243
const { availableCommands } = useSlashCommandSets();
239244

240245
// -- Sound effects --------------------------------------------------------
@@ -826,17 +831,30 @@ export default function CloudChatPage({ organizationId }: CloudChatPageProps) {
826831
},
827832
}}
828833
/>
829-
{sessionConfig?.repository && (
830-
<div className="text-muted-foreground flex items-center gap-1.5 px-[max(1rem,calc(50%_-_27rem))] pb-3 text-xs md:pb-4">
831-
<GitBranch className="h-3 w-3 shrink-0" />
832-
<span className="truncate">{sessionConfig.repository}</span>
833-
{fetchedSessionData?.gitBranch && (
834-
<>
835-
<span>·</span>
836-
<span className="truncate">
837-
{fetchedSessionData.gitBranch}
838-
</span>
839-
</>
834+
{(sessionConfig?.repository ||
835+
(contextUsage !== undefined && contextWindow !== undefined)) && (
836+
<div className="text-muted-foreground flex items-center gap-3 px-[max(1rem,calc(50%_-_27rem))] pb-3 text-xs md:pb-4">
837+
{sessionConfig?.repository && (
838+
<div className="flex min-w-0 items-center gap-1.5">
839+
<GitBranch className="h-3 w-3 shrink-0" />
840+
<span className="truncate">{sessionConfig.repository}</span>
841+
{fetchedSessionData?.gitBranch && (
842+
<>
843+
<span>·</span>
844+
<span className="truncate">
845+
{fetchedSessionData.gitBranch}
846+
</span>
847+
</>
848+
)}
849+
</div>
850+
)}
851+
{contextUsage !== undefined && contextWindow !== undefined && (
852+
<div className="ml-auto shrink-0">
853+
<ContextUsageIndicator
854+
contextTokens={contextUsage.contextTokens}
855+
contextWindow={contextWindow}
856+
/>
857+
</div>
840858
)}
841859
</div>
842860
)}

‎apps/web/src/components/cloud-agent-next/ContextUsageIndicator.tsx‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
'use client';
2+
3+
import React from 'react';
4+
import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip';
5+
import { calculateContextUsagePercentage } from '@/lib/cloud-agent-sdk/context-usage';
6+
7+
type ContextUsageIndicatorProps = {
8+
contextTokens?: number;
9+
contextWindow?: number;
10+
};
11+
12+
function formatTokenCount(tokens: number): string {
13+
return tokens.toLocaleString('en-US');
14+
}
15+
16+
function formatCompactTokenCount(tokens: number): string {
17+
if (tokens < 1_000) return formatTokenCount(tokens);
18+
return `${(tokens / 1_000).toFixed(1)}K`;
19+
}
20+
21+
export function formatContextUsageTooltip(contextTokens: number, contextWindow: number): string {
22+
return `${formatTokenCount(contextTokens)} / ${formatTokenCount(contextWindow)} tokens used`;
23+
}
24+
25+
export function ContextUsageIndicator({
26+
contextTokens,
27+
contextWindow,
28+
}: ContextUsageIndicatorProps) {
29+
if (contextTokens === undefined || contextWindow === undefined) return null;
30+
31+
const percentage = calculateContextUsagePercentage(contextTokens, contextWindow);
32+
if (percentage === undefined) return null;
33+
34+
const formattedContextTokens = formatTokenCount(contextTokens);
35+
const formattedContextWindow = formatTokenCount(contextWindow);
36+
37+
return (
38+
<Tooltip>
39+
<TooltipTrigger asChild>
40+
<button
41+
type="button"
42+
aria-label={`${percentage}% of context used. ${formattedContextTokens} of ${formattedContextWindow} tokens used.`}
43+
className="text-muted-foreground hover:text-foreground focus-visible:ring-ring focus-visible:ring-offset-background relative inline-flex h-8 min-w-8 shrink-0 items-center justify-center rounded-sm px-1 font-mono text-xs whitespace-nowrap tabular-nums transition-colors before:absolute before:-inset-1.5 before:content-[''] focus-visible:ring-2 focus-visible:ring-offset-2 focus-visible:outline-none"
44+
>
45+
{formatCompactTokenCount(contextTokens)} ({percentage}%)
46+
</button>
47+
</TooltipTrigger>
48+
<TooltipContent side="top">
49+
{formatContextUsageTooltip(contextTokens, contextWindow)}
50+
</TooltipContent>
51+
</Tooltip>
52+
);
53+
}

‎apps/web/src/components/cloud-agent-next/context-usage-indicator.test.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import React from 'react';
2+
import { renderToStaticMarkup } from 'react-dom/server';
3+
import { ContextUsageIndicator, formatContextUsageTooltip } from './ContextUsageIndicator';
4+
5+
function renderIndicator(contextTokens?: number, contextWindow?: number) {
6+
return renderToStaticMarkup(
7+
React.createElement(ContextUsageIndicator, { contextTokens, contextWindow })
8+
);
9+
}
10+
11+
describe('ContextUsageIndicator', () => {
12+
it('renders a visible integer percentage in a named button trigger', () => {
13+
const html = renderIndicator(32_418, 80_000);
14+
15+
expect(html.match(/<button/g)).toHaveLength(1);
16+
expect(html).toMatch(/<button[^>]*type="button"[^>]*>32.4K \(41%\)<\/button>/);
17+
expect(html).toContain('aria-label="41% of context used. 32,418 of 80,000 tokens used."');
18+
});
19+
20+
it('matches the CLI compact token-count label', () => {
21+
expect(renderIndicator(239_100, 1_000_000)).toContain('239.1K (24%)');
22+
});
23+
24+
it('preserves rendered percentages above one hundred', () => {
25+
expect(renderIndicator(101, 100)).toContain('101 (101%)');
26+
});
27+
28+
it.each([undefined, 0, -1])('omits markup for invalid context window %s', contextWindow => {
29+
expect(renderIndicator(32_418, contextWindow)).toBe('');
30+
});
31+
32+
it('does not announce streaming updates as live status changes', () => {
33+
const html = renderIndicator(32_418, 80_000);
34+
35+
expect(html).not.toContain('aria-live');
36+
expect(html).not.toContain('role="status"');
37+
});
38+
});
39+
40+
describe('formatContextUsageTooltip', () => {
41+
it('formats exact token counts with stable grouping', () => {
42+
expect(formatContextUsageTooltip(32_418, 80_000)).toBe('32,418 / 80,000 tokens used');
43+
});
44+
});

‎apps/web/src/components/cloud-agent-next/hooks/useOrganizationModels.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,15 @@ import { useOrganizationDefaults } from '@/app/api/organizations/hooks';
99
import { useModelSelectorList } from '@/app/api/openrouter/hooks';
1010
import type { ModelOption } from '@/components/shared/ModelCombobox';
1111
import { appendCloudAgentNextLocalTestModel } from '@/components/cloud-agent-next/model-preferences';
12+
import { buildContextLengthByModelId } from '@/components/cloud-agent-next/model-context-lengths';
1213

1314
type UseOrganizationModelsReturn = {
1415
/** Models formatted for the ModelCombobox component */
1516
modelOptions: ModelOption[];
1617
/** Whether models are still loading */
1718
isLoadingModels: boolean;
19+
/** Context windows keyed by exact catalog model ID */
20+
contextLengthByModelId: ReadonlyMap<string, number>;
1821
/** The organization's default model */
1922
defaultModel: string | undefined;
2023
};
@@ -44,9 +47,15 @@ export function useOrganizationModels(organizationId?: string): UseOrganizationM
4447
);
4548
}, [openRouterModels]);
4649

50+
const contextLengthByModelId = useMemo(
51+
() => buildContextLengthByModelId(openRouterModels?.data ?? []),
52+
[openRouterModels]
53+
);
54+
4755
return {
4856
modelOptions,
4957
isLoadingModels: isLoadingOpenRouter,
58+
contextLengthByModelId,
5059
defaultModel: defaultsData?.defaultModel,
5160
};
5261
}

‎apps/web/src/components/cloud-agent-next/model-context-lengths.test.ts‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import type { ContextUsage } from '@/lib/cloud-agent-sdk/context-usage';
2+
import { buildContextLengthByModelId, resolveContextWindow } from './model-context-lengths';
3+
4+
const contextUsage = {
5+
contextTokens: 32_418,
6+
providerID: 'kilo',
7+
modelID: 'anthropic/claude-sonnet-4',
8+
} satisfies ContextUsage;
9+
10+
describe('buildContextLengthByModelId', () => {
11+
it('maps exact catalog model ids without aliases', () => {
12+
const lengths = buildContextLengthByModelId([
13+
{ id: 'anthropic/claude-sonnet-4', context_length: 200_000 },
14+
{ id: 'kilo-auto/free', context_length: 114_688 },
15+
{ id: 'fake-deterministic', context_length: 200_000 },
16+
{ id: 'kilo/preview-allowed-by-policy', context_length: 80_000 },
17+
]);
18+
19+
expect(lengths).toEqual(
20+
new Map([
21+
['anthropic/claude-sonnet-4', 200_000],
22+
['kilo-auto/free', 114_688],
23+
['fake-deterministic', 200_000],
24+
['kilo/preview-allowed-by-policy', 80_000],
25+
])
26+
);
27+
expect(lengths.has('preview-allowed-by-policy')).toBe(false);
28+
});
29+
30+
it('omits invalid context lengths', () => {
31+
expect(
32+
buildContextLengthByModelId([
33+
{ id: 'zero', context_length: 0 },
34+
{ id: 'negative', context_length: -1 },
35+
{ id: 'missing', context_length: undefined },
36+
{ id: 'nan', context_length: Number.NaN },
37+
{ id: 'infinity', context_length: Number.POSITIVE_INFINITY },
38+
])
39+
).toEqual(new Map());
40+
});
41+
42+
it('retains agreeing duplicate exact ids', () => {
43+
expect(
44+
buildContextLengthByModelId([
45+
{ id: 'anthropic/claude-sonnet-4', context_length: 200_000 },
46+
{ id: 'anthropic/claude-sonnet-4', context_length: 200_000 },
47+
])
48+
).toEqual(new Map([['anthropic/claude-sonnet-4', 200_000]]));
49+
});
50+
51+
it('permanently omits a conflicting duplicate exact id', () => {
52+
expect(
53+
buildContextLengthByModelId([
54+
{ id: 'anthropic/claude-sonnet-4', context_length: 200_000 },
55+
{ id: 'anthropic/claude-sonnet-4', context_length: 80_000 },
56+
{ id: 'anthropic/claude-sonnet-4', context_length: 200_000 },
57+
])
58+
).toEqual(new Map());
59+
});
60+
});
61+
62+
describe('resolveContextWindow', () => {
63+
it('resolves a known kilo response by exact emitted model id', () => {
64+
expect(resolveContextWindow(contextUsage, new Map([[contextUsage.modelID, 200_000]]))).toBe(
65+
200_000
66+
);
67+
});
68+
69+
it('returns undefined for missing usage or a non-kilo provider', () => {
70+
expect(
71+
resolveContextWindow(undefined, new Map([[contextUsage.modelID, 200_000]]))
72+
).toBeUndefined();
73+
expect(
74+
resolveContextWindow(
75+
{ ...contextUsage, providerID: 'anthropic' },
76+
new Map([[contextUsage.modelID, 200_000]])
77+
)
78+
).toBeUndefined();
79+
});
80+
81+
it('returns undefined for missing or non-positive capacities', () => {
82+
expect(resolveContextWindow(contextUsage, new Map())).toBeUndefined();
83+
expect(
84+
resolveContextWindow(contextUsage, new Map([[contextUsage.modelID, 0]]))
85+
).toBeUndefined();
86+
expect(
87+
resolveContextWindow(contextUsage, new Map([[contextUsage.modelID, -1]]))
88+
).toBeUndefined();
89+
});
90+
91+
it('returns undefined rather than guessing a stripped experiment id alias', () => {
92+
expect(
93+
resolveContextWindow(
94+
{ ...contextUsage, modelID: 'preview-allowed-by-policy' },
95+
new Map([['kilo/preview-allowed-by-policy', 80_000]])
96+
)
97+
).toBeUndefined();
98+
});
99+
});

‎apps/web/src/components/cloud-agent-next/model-context-lengths.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
import type { ContextUsage } from '@/lib/cloud-agent-sdk/context-usage';
2+
3+
type ModelContextLength = {
4+
id: string;
5+
context_length?: number | null;
6+
};
7+
8+
export function buildContextLengthByModelId(
9+
models: readonly ModelContextLength[]
10+
): ReadonlyMap<string, number> {
11+
const contextLengthByModelId = new Map<string, number>();
12+
const conflictingModelIds = new Set<string>();
13+
14+
for (const model of models) {
15+
const contextLength = model.context_length;
16+
if (contextLength === undefined || contextLength === null) continue;
17+
if (!Number.isFinite(contextLength) || contextLength <= 0) continue;
18+
if (conflictingModelIds.has(model.id)) continue;
19+
20+
const existingContextLength = contextLengthByModelId.get(model.id);
21+
if (existingContextLength === undefined) {
22+
contextLengthByModelId.set(model.id, contextLength);
23+
continue;
24+
}
25+
26+
if (existingContextLength !== contextLength) {
27+
contextLengthByModelId.delete(model.id);
28+
conflictingModelIds.add(model.id);
29+
}
30+
}
31+
32+
return contextLengthByModelId;
33+
}
34+
35+
export function resolveContextWindow(
36+
contextUsage: ContextUsage | undefined,
37+
contextLengthByModelId: ReadonlyMap<string, number>
38+
): number | undefined {
39+
if (contextUsage?.providerID !== 'kilo') return undefined;
40+
41+
const contextWindow = contextLengthByModelId.get(contextUsage.modelID);
42+
if (contextWindow === undefined || !Number.isFinite(contextWindow) || contextWindow <= 0) {
43+
return undefined;
44+
}
45+
46+
return contextWindow;
47+
}

0 commit comments

Comments
 (0)