Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
18 changes: 13 additions & 5 deletions server/claude-sdk.js
Original file line number Diff line number Diff line change
Expand Up @@ -279,7 +279,7 @@ function transformMessage(sdkMessage) {
* @param {Object} resultMessage - SDK result message
* @returns {Object|null} Token budget object or null
*/
function extractTokenBudget(resultMessage) {
function extractTokenBudget(resultMessage, selectedModel) {
if (resultMessage.type !== 'result' || !resultMessage.modelUsage) {
return null;
}
Expand All @@ -302,9 +302,17 @@ function extractTokenBudget(resultMessage) {
// Total used = input + output + cache tokens
const totalUsed = inputTokens + outputTokens + cacheReadTokens + cacheCreationTokens;

// Use configured context window budget from environment (default 160000)
// This is the user's budget limit, not the model's context window
const contextWindow = parseInt(process.env.CONTEXT_WINDOW) || 160000;
// Determine context window: env override > model usage data > selected model > default
let contextWindow = parseInt(process.env.CONTEXT_WINDOW) || 0;
if (!contextWindow) {
Comment on lines +305 to +307
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.

⚠️ Potential issue | 🟡 Minor

Validate CONTEXT_WINDOW as a positive integer before applying override.

At Line 306, parseInt(process.env.CONTEXT_WINDOW) || 0 treats negative values as valid and can emit an invalid budget total.

Suggested fix
-  let contextWindow = parseInt(process.env.CONTEXT_WINDOW) || 0;
+  const envContextWindow = Number.parseInt(process.env.CONTEXT_WINDOW ?? '', 10);
+  let contextWindow = Number.isFinite(envContextWindow) && envContextWindow > 0 ? envContextWindow : 0;
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Determine context window: env override > model usage data > selected model > default
let contextWindow = parseInt(process.env.CONTEXT_WINDOW) || 0;
if (!contextWindow) {
// Determine context window: env override > model usage data > selected model > default
const envContextWindow = Number.parseInt(process.env.CONTEXT_WINDOW ?? '', 10);
let contextWindow = Number.isFinite(envContextWindow) && envContextWindow > 0 ? envContextWindow : 0;
if (!contextWindow) {
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@server/claude-sdk.js` around lines 305 - 307, The override for CONTEXT_WINDOW
currently uses parseInt(process.env.CONTEXT_WINDOW) || 0 which accepts negative
or non-integer values; update the logic that sets contextWindow so it parses the
env var, verifies it's a finite positive integer (>0) (e.g., use
Number.isInteger(parsed) && parsed > 0) and only applies it when valid,
otherwise fall back to 0 or the existing model-derived value; modify the code
where contextWindow is assigned (the parseInt(...) usage and the subsequent if
(!contextWindow) check) to perform this validation and fall back safely.

if (modelData.contextWindow && modelData.contextWindow > 0) {
contextWindow = modelData.contextWindow;
} else if (selectedModel && selectedModel.toLowerCase().includes('1m')) {
contextWindow = 1000000;
} else {
contextWindow = 200000;
}
}

// Token calc logged via token-budget WS event

Expand Down Expand Up @@ -664,7 +672,7 @@ async function queryClaudeSDK(command, options = {}, ws) {
if (models.length > 0) {
// Model info available in result message
}
const tokenBudgetData = extractTokenBudget(message);
const tokenBudgetData = extractTokenBudget(message, sdkOptions.model);
if (tokenBudgetData) {
ws.send(createNormalizedMessage({ kind: 'status', text: 'token_budget', tokenBudget: tokenBudgetData, sessionId: capturedSessionId || sessionId || null, provider: 'claude' }));
}
Expand Down
1 change: 1 addition & 0 deletions shared/modelConstants.js
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ export const CLAUDE_MODELS = {
{ value: "haiku", label: "Haiku" },
{ value: "opusplan", label: "Opus Plan" },
{ value: "sonnet[1m]", label: "Sonnet [1M]" },
{ value: "opus[1m]", label: "Opus [1M]" },
],

DEFAULT: "sonnet",
Expand Down
15 changes: 12 additions & 3 deletions src/components/chat/hooks/useChatSessionState.ts
Original file line number Diff line number Diff line change
Expand Up @@ -548,20 +548,29 @@ export function useChatSessionState({
const sessionProvider = selectedSession.__provider || 'claude';
if (sessionProvider !== 'claude') return;

// Fetch initial token usage but don't overwrite a higher value from WebSocket
let cancelled = false;
const fetchInitialTokenUsage = async () => {
try {
const url = `/api/projects/${selectedProject.name}/sessions/${selectedSession.id}/token-usage`;
const response = await authenticatedFetch(url);
if (response.ok) {
setTokenBudget(await response.json());
} else {
if (response.ok && !cancelled) {
const data = await response.json();
setTokenBudget((prev: Record<string, unknown> | null) => {
if (prev && typeof prev.total === 'number' && prev.total > (data.total || 0)) {
return prev;
}
return data;
});
} else if (!cancelled) {
setTokenBudget(null);
}
Comment on lines +565 to 567
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.

⚠️ Potential issue | 🟡 Minor

Avoid clearing existing token budget on transient fetch failures.

At Line 565-Line 567, setting tokenBudget to null can erase a valid WebSocket-provided budget and cause UI regression.

Suggested fix
-        } else if (!cancelled) {
-          setTokenBudget(null);
-        }
+        } else if (!cancelled) {
+          // Preserve existing budget (especially realtime WS updates) on HTTP fetch failures.
+          setTokenBudget((prev: Record<string, unknown> | null) => prev);
+        }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
} else if (!cancelled) {
setTokenBudget(null);
}
} else if (!cancelled) {
// Preserve existing budget (especially realtime WS updates) on HTTP fetch failures.
setTokenBudget((prev: Record<string, unknown> | null) => prev);
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/components/chat/hooks/useChatSessionState.ts` around lines 565 - 567, The
code is clearing an existing valid WebSocket-provided token budget on transient
fetch failures by calling setTokenBudget(null) in the fetch error path; update
the error/else path inside useChatSessionState so it does not overwrite a
current tokenBudget on transient failures—either remove the setTokenBudget(null)
call or guard it so you only clear tokenBudget when you have a definitive “no
budget” response (e.g., only setTokenBudget(null) when the fetch returned a
confirmed null budget), and leave tokenBudget untouched when the fetch failed or
was cancelled (check cancelled and the fetched value variable before calling
setTokenBudget).

} catch (error) {
console.error('Failed to fetch initial token usage:', error);
}
};
fetchInitialTokenUsage();
return () => { cancelled = true; };
}, [selectedProject, selectedSession?.id, selectedSession?.__provider]);

const visibleMessages = useMemo(() => {
Expand Down
1 change: 1 addition & 0 deletions src/components/chat/view/ChatInterface.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,7 @@ function ChatInterface({
thinkingMode={thinkingMode}
setThinkingMode={setThinkingMode}
tokenBudget={tokenBudget}
selectedModel={claudeModel}
slashCommandsCount={slashCommandsCount}
onToggleCommandMenu={handleToggleCommandMenu}
hasInput={Boolean(input.trim())}
Expand Down
3 changes: 3 additions & 0 deletions src/components/chat/view/subcomponents/ChatComposer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ interface ChatComposerProps {
thinkingMode: string;
setThinkingMode: Dispatch<SetStateAction<string>>;
tokenBudget: { used?: number; total?: number } | null;
selectedModel?: string;
slashCommandsCount: number;
onToggleCommandMenu: () => void;
hasInput: boolean;
Expand Down Expand Up @@ -105,6 +106,7 @@ export default function ChatComposer({
thinkingMode,
setThinkingMode,
tokenBudget,
selectedModel,
slashCommandsCount,
onToggleCommandMenu,
hasInput,
Expand Down Expand Up @@ -192,6 +194,7 @@ export default function ChatComposer({
thinkingMode={thinkingMode}
setThinkingMode={setThinkingMode}
tokenBudget={tokenBudget}
selectedModel={selectedModel}
slashCommandsCount={slashCommandsCount}
onToggleCommandMenu={onToggleCommandMenu}
hasInput={hasInput}
Expand Down
4 changes: 3 additions & 1 deletion src/components/chat/view/subcomponents/ChatInputControls.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ interface ChatInputControlsProps {
thinkingMode: string;
setThinkingMode: React.Dispatch<React.SetStateAction<string>>;
tokenBudget: { used?: number; total?: number } | null;
selectedModel?: string;
slashCommandsCount: number;
onToggleCommandMenu: () => void;
hasInput: boolean;
Expand All @@ -27,6 +28,7 @@ export default function ChatInputControls({
thinkingMode,
setThinkingMode,
tokenBudget,
selectedModel,
slashCommandsCount,
onToggleCommandMenu,
hasInput,
Expand Down Expand Up @@ -78,7 +80,7 @@ export default function ChatInputControls({
<ThinkingModeSelector selectedMode={thinkingMode} onModeChange={setThinkingMode} onClose={() => {}} className="" />
)}

<TokenUsagePie used={tokenBudget?.used || 0} total={tokenBudget?.total || parseInt(import.meta.env.VITE_CONTEXT_WINDOW) || 160000} />
<TokenUsagePie used={tokenBudget?.used || 0} total={tokenBudget?.total || (selectedModel && selectedModel.includes('1m') ? 1000000 : 200000)} />

<button
type="button"
Expand Down