diff --git a/packages/core/src/core/tokenLimits.test.ts b/packages/core/src/core/tokenLimits.test.ts index 872e3b1cab..a8c520b4fe 100644 --- a/packages/core/src/core/tokenLimits.test.ts +++ b/packages/core/src/core/tokenLimits.test.ts @@ -2,6 +2,7 @@ import { describe, it, expect } from 'vitest'; import { normalize, tokenLimit, + knownTokenLimit, DEFAULT_TOKEN_LIMIT, DEFAULT_OUTPUT_TOKEN_LIMIT, } from './tokenLimits.js'; @@ -234,6 +235,21 @@ describe('tokenLimit', () => { }); }); +describe('knownTokenLimit', () => { + it('returns a limit for known input models', () => { + expect(knownTokenLimit('qwen3-max')).toBe(262144); + expect(knownTokenLimit('gpt-5')).toBe(272000); + }); + + it('returns a limit for known output models', () => { + expect(knownTokenLimit('qwen3-max', 'output')).toBe(32768); + }); + + it('returns undefined for unknown models instead of the default fallback', () => { + expect(knownTokenLimit('unknown-model-v1.0')).toBeUndefined(); + }); +}); + describe('tokenLimit with output type', () => { describe('latest models output limits', () => { it('should return correct output limits for GPT-5.x', () => { diff --git a/packages/core/src/core/tokenLimits.ts b/packages/core/src/core/tokenLimits.ts index 108b80c73c..ff0966430a 100644 --- a/packages/core/src/core/tokenLimits.ts +++ b/packages/core/src/core/tokenLimits.ts @@ -184,6 +184,22 @@ const OUTPUT_PATTERNS: Array<[RegExp, TokenCount]> = [ [/^kimi-k2\.5/, LIMITS['32k']], ]; +function findTokenLimit( + model: Model, + type: TokenLimitType = 'input', +): TokenCount | undefined { + const norm = normalize(model); + const patterns = type === 'output' ? OUTPUT_PATTERNS : PATTERNS; + + for (const [regex, limit] of patterns) { + if (regex.test(norm)) { + return limit; + } + } + + return undefined; +} + /** * Check if a model has an explicitly defined output token limit. * This distinguishes between models with known limits in OUTPUT_PATTERNS @@ -197,6 +213,13 @@ export function hasExplicitOutputLimit(model: Model): boolean { return OUTPUT_PATTERNS.some(([regex]) => regex.test(norm)); } +export function knownTokenLimit( + model: Model, + type: TokenLimitType = 'input', +): TokenCount | undefined { + return findTokenLimit(model, type); +} + /** * Return the token limit for a model string based on the specified type. * @@ -216,17 +239,8 @@ export function tokenLimit( model: Model, type: TokenLimitType = 'input', ): TokenCount { - const norm = normalize(model); - - // Choose the appropriate patterns based on token type - const patterns = type === 'output' ? OUTPUT_PATTERNS : PATTERNS; - - for (const [regex, limit] of patterns) { - if (regex.test(norm)) { - return limit; - } - } - - // Return appropriate default based on token type - return type === 'output' ? DEFAULT_OUTPUT_TOKEN_LIMIT : DEFAULT_TOKEN_LIMIT; + return ( + knownTokenLimit(model, type) ?? + (type === 'output' ? DEFAULT_OUTPUT_TOKEN_LIMIT : DEFAULT_TOKEN_LIMIT) + ); } diff --git a/packages/vscode-ide-companion/src/utils/acpModelInfo.test.ts b/packages/vscode-ide-companion/src/utils/acpModelInfo.test.ts index d69d405659..03ff8ee8cf 100644 --- a/packages/vscode-ide-companion/src/utils/acpModelInfo.test.ts +++ b/packages/vscode-ide-companion/src/utils/acpModelInfo.test.ts @@ -136,6 +136,26 @@ describe('extractSessionModelState', () => { // The function should still return a state with empty availableModels expect(result?.availableModels).toHaveLength(0); }); + + it('derives contextLimit for known models when the ACP payload omits it', () => { + const result = extractSessionModelState({ + models: { + currentModelId: 'qwen3-max', + availableModels: [{ modelId: 'qwen3-max', name: 'Qwen3 Max' }], + }, + }); + + expect(result).toEqual({ + currentModelId: 'qwen3-max', + availableModels: [ + { + modelId: 'qwen3-max', + name: 'Qwen3 Max', + _meta: { contextLimit: 262144 }, + }, + ], + }); + }); }); describe('extractModelInfoFromNewSessionResult', () => { @@ -205,4 +225,36 @@ describe('extractModelInfoFromNewSessionResult', () => { expect(extractModelInfoFromNewSessionResult({})).toBeNull(); expect(extractModelInfoFromNewSessionResult(null)).toBeNull(); }); + + it('derives contextLimit for known models when the payload has null metadata', () => { + expect( + extractModelInfoFromNewSessionResult({ + model: { + name: 'Qwen3 Max', + modelId: 'qwen3-max', + _meta: null, + }, + }), + ).toEqual({ + name: 'Qwen3 Max', + modelId: 'qwen3-max', + _meta: { contextLimit: 262144 }, + }); + }); + + it('preserves null contextLimit for unknown models', () => { + expect( + extractModelInfoFromNewSessionResult({ + model: { + name: 'Unknown', + modelId: 'unknown-model-v1.0', + _meta: { contextLimit: null }, + }, + }), + ).toEqual({ + name: 'Unknown', + modelId: 'unknown-model-v1.0', + _meta: { contextLimit: null }, + }); + }); }); diff --git a/packages/vscode-ide-companion/src/utils/acpModelInfo.ts b/packages/vscode-ide-companion/src/utils/acpModelInfo.ts index d2c8b5e1b4..e9d92caf06 100644 --- a/packages/vscode-ide-companion/src/utils/acpModelInfo.ts +++ b/packages/vscode-ide-companion/src/utils/acpModelInfo.ts @@ -5,6 +5,7 @@ */ import type { ModelInfo } from '@agentclientprotocol/sdk'; +import { knownTokenLimit } from '@qwen-code/qwen-code-core/src/core/tokenLimits.js'; import type { ApprovalModeValue } from '../types/approvalModeValueTypes.js'; type AcpMeta = Record; @@ -19,6 +20,15 @@ const asMeta = (value: unknown): AcpMeta | null | undefined => { return undefined; }; +const getContextLimitFromMeta = ( + meta: AcpMeta | null | undefined, +): number | null | undefined => { + const metaLimit = meta?.['contextLimit']; + return typeof metaLimit === 'number' || metaLimit === null + ? metaLimit + : undefined; +}; + const normalizeModelInfo = (value: unknown): ModelInfo | null => { if (!value || typeof value !== 'object') { return null; @@ -48,10 +58,25 @@ const normalizeModelInfo = (value: unknown): ModelInfo | null => { // Back-compat: older implementations used `contextLimit` at the top-level. const legacyContextLimit = obj['contextLimit']; - const contextLimit = + const legacyLimit = typeof legacyContextLimit === 'number' || legacyContextLimit === null ? legacyContextLimit : undefined; + const metaLimit = getContextLimitFromMeta(metaFromWire); + const derivedLimit = knownTokenLimit(modelId || name); + + // Priority: legacy numeric > meta numeric > derived from known model > explicit null > undefined. + // An explicit `null` from the server means "limit intentionally unknown"; `undefined` means "not provided". + const contextLimit = + typeof legacyLimit === 'number' + ? legacyLimit + : typeof metaLimit === 'number' + ? metaLimit + : typeof derivedLimit === 'number' + ? derivedLimit + : legacyLimit === null || metaLimit === null + ? null + : undefined; let mergedMeta: AcpMeta | null | undefined = metaFromWire; if (typeof contextLimit !== 'undefined') { diff --git a/packages/vscode-ide-companion/src/webview/App.tsx b/packages/vscode-ide-companion/src/webview/App.tsx index 3f01f30e51..63b17000d9 100644 --- a/packages/vscode-ide-companion/src/webview/App.tsx +++ b/packages/vscode-ide-companion/src/webview/App.tsx @@ -52,11 +52,8 @@ import type { ApprovalModeValue } from '../types/approvalModeValueTypes.js'; import type { PlanEntry, UsageStatsPayload } from '../types/chatTypes.js'; import type { ModelInfo, AvailableCommand } from '@agentclientprotocol/sdk'; import type { Question } from '../types/acpTypes.js'; -import { - DEFAULT_TOKEN_LIMIT, - tokenLimit, -} from '@qwen-code/qwen-code-core/src/core/tokenLimits.js'; import { useImagePaste, type WebViewImageMessage } from './hooks/useImage.js'; +import { computeContextUsage } from './utils/contextUsage.js'; export const App: React.FC = () => { const vscode = useVSCode(); @@ -208,52 +205,7 @@ export const App: React.FC = () => { const completion = useCompletionTrigger(inputFieldRef, getCompletionItems); - const contextUsage = useMemo(() => { - if (!usageStats && !modelInfo) { - return null; - } - - const modelName = - modelInfo?.modelId && typeof modelInfo.modelId === 'string' - ? modelInfo.modelId - : modelInfo?.name && typeof modelInfo.name === 'string' - ? modelInfo.name - : undefined; - - // Note: In the webview context, the contextWindowSize is already reflected in - // modelInfo._meta.contextLimit which is computed on the extension side with the proper config. - // We only use tokenLimit as a fallback if metaLimit is not available. - const derivedLimit = - modelName && modelName.length > 0 - ? tokenLimit(modelName, 'input') - : undefined; - - const metaLimitRaw = modelInfo?._meta?.['contextLimit']; - const metaLimit = - typeof metaLimitRaw === 'number' || metaLimitRaw === null - ? metaLimitRaw - : undefined; - - const limit = - usageStats?.tokenLimit ?? - metaLimit ?? - derivedLimit ?? - DEFAULT_TOKEN_LIMIT; - - const used = usageStats?.usage?.promptTokens ?? 0; - if (typeof limit !== 'number' || limit <= 0 || used < 0) { - return null; - } - const percentLeft = Math.max( - 0, - Math.min(100, Math.round(((limit - used) / limit) * 100)), - ); - return { - percentLeft, - usedTokens: used, - tokenLimit: limit, - }; - }, [usageStats, modelInfo]); + const contextUsage = useMemo(() => computeContextUsage(usageStats, modelInfo), [usageStats, modelInfo]); // Track a lightweight signature of workspace files to detect content changes even when length is unchanged const workspaceFilesSignature = useMemo( diff --git a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts index ff873cfda4..de71c5aa2b 100644 --- a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts +++ b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts @@ -178,7 +178,7 @@ describe('SessionMessageHandler', () => { const handler = new SessionMessageHandler( agentManager as never, conversationStore as never, - null, + 'conversation-1', sendToWebView, ); @@ -186,8 +186,13 @@ describe('SessionMessageHandler', () => { type: 'newQwenSession', }); + expect(handler.getCurrentConversationId()).toBeNull(); expect(agentManager.createNewSession).toHaveBeenCalledWith('/workspace', { forceNew: true, }); + expect(sendToWebView).toHaveBeenCalledWith({ + type: 'conversationCleared', + data: {}, + }); }); }); diff --git a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts index 8fab3919de..b0a0179cd9 100644 --- a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts +++ b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts @@ -588,6 +588,7 @@ export class SessionMessageHandler extends BaseMessageHandler { const workingDir = workspaceFolder?.uri.fsPath || process.cwd(); await this.agentManager.createNewSession(workingDir, { forceNew: true }); + this.currentConversationId = null; this.sendToWebView({ type: 'conversationCleared', diff --git a/packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.test.ts b/packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.test.ts new file mode 100644 index 0000000000..d22a27716d --- /dev/null +++ b/packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.test.ts @@ -0,0 +1,49 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it, vi } from 'vitest'; +import { resetConversationState } from './useWebViewMessages.js'; + +describe('resetConversationState', () => { + it('clears retained usage stats when a conversation is reset', () => { + const clearMessages = vi.fn(); + const clearToolCalls = vi.fn(); + const setCurrentSessionId = vi.fn(); + const setCurrentSessionTitle = vi.fn(); + const setUsageStats = vi.fn(); + const clearImageResolutions = vi.fn(); + const postMessage = vi.fn(); + + resetConversationState({ + handlers: { + messageHandling: { + clearMessages, + }, + clearToolCalls, + sessionManagement: { + setCurrentSessionId, + setCurrentSessionTitle, + }, + setUsageStats, + }, + clearImageResolutions, + vscode: { + postMessage, + }, + }); + + expect(clearMessages).toHaveBeenCalled(); + expect(clearToolCalls).toHaveBeenCalled(); + expect(setCurrentSessionId).toHaveBeenCalledWith(null); + expect(clearImageResolutions).toHaveBeenCalled(); + expect(setUsageStats).toHaveBeenCalledWith(undefined); + expect(setCurrentSessionTitle).toHaveBeenCalledWith('Past Conversations'); + expect(postMessage).toHaveBeenCalledWith({ + type: 'updatePanelTitle', + data: { title: 'Qwen Code' }, + }); + }); +}); diff --git a/packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.ts b/packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.ts index 8d5eef683a..bbf97ad9d9 100644 --- a/packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.ts +++ b/packages/vscode-ide-companion/src/webview/hooks/useWebViewMessages.ts @@ -133,6 +133,41 @@ interface UseWebViewMessagesProps { setAvailableModels?: (models: ModelInfo[]) => void; } +type ConversationResetHandlers = { + messageHandling: Pick< + UseWebViewMessagesProps['messageHandling'], + 'clearMessages' + >; + clearToolCalls: UseWebViewMessagesProps['clearToolCalls']; + sessionManagement: Pick< + UseWebViewMessagesProps['sessionManagement'], + 'setCurrentSessionId' | 'setCurrentSessionTitle' + >; + setUsageStats?: UseWebViewMessagesProps['setUsageStats']; +}; + +export function resetConversationState({ + handlers, + clearImageResolutions, + vscode, +}: { + handlers: ConversationResetHandlers; + clearImageResolutions: () => void; + vscode: { postMessage: (message: unknown) => void }; +}) { + handlers.messageHandling.clearMessages(); + handlers.clearToolCalls(); + handlers.sessionManagement.setCurrentSessionId(null); + clearImageResolutions(); + handlers.setUsageStats?.(undefined); + handlers.sessionManagement.setCurrentSessionTitle('Past Conversations'); + // Reset the VS Code tab title to default label + vscode.postMessage({ + type: 'updatePanelTitle', + data: { title: 'Qwen Code' }, + }); +} + /** * WebView message handling Hook * Handles all messages from VSCode Extension uniformly @@ -914,17 +949,10 @@ export const useWebViewMessages = ({ break; case 'conversationCleared': - handlers.messageHandling.clearMessages(); - handlers.clearToolCalls(); - handlers.sessionManagement.setCurrentSessionId(null); - clearImageResolutions(); - handlers.sessionManagement.setCurrentSessionTitle( - 'Past Conversations', - ); - // Reset the VS Code tab title to default label - vscode.postMessage({ - type: 'updatePanelTitle', - data: { title: 'Qwen Code' }, + resetConversationState({ + handlers, + clearImageResolutions, + vscode, }); lastPlanSnapshotRef.current = null; break; diff --git a/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.test.ts b/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.test.ts index 6ba2d3bf86..0062e0439d 100644 --- a/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.test.ts +++ b/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.test.ts @@ -103,6 +103,8 @@ vi.mock('./MessageHandler.js', () => ({ setLoginHandler = vi.fn(); setPermissionHandler = vi.fn(); setAskUserQuestionHandler = vi.fn(); + setCurrentConversationId = vi.fn(); + getCurrentConversationId = vi.fn(() => null); setupFileWatchers = vi.fn(() => ({ dispose: vi.fn() })); appendStreamContent = vi.fn(); route = vi.fn(); @@ -303,6 +305,13 @@ describe('WebViewProvider.createNewSession', () => { }; } ).agentManager; + const messageHandler = ( + provider as unknown as { + messageHandler: { + setCurrentConversationId: ReturnType; + }; + } + ).messageHandler; await provider.createNewSession(); @@ -310,5 +319,6 @@ describe('WebViewProvider.createNewSession', () => { '/workspace-root', { forceNew: true }, ); + expect(messageHandler.setCurrentConversationId).toHaveBeenCalledWith(null); }); }); diff --git a/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.ts b/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.ts index 3e7d38a2b2..18d8a86721 100644 --- a/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.ts +++ b/packages/vscode-ide-companion/src/webview/providers/WebViewProvider.ts @@ -1615,6 +1615,7 @@ export class WebViewProvider { // Create new Qwen session via agent manager await this.agentManager.createNewSession(workingDir, { forceNew: true }); + this.messageHandler.setCurrentConversationId(null); // Clear current conversation UI this.sendMessageToWebView({ diff --git a/packages/vscode-ide-companion/src/webview/utils/contextUsage.test.ts b/packages/vscode-ide-companion/src/webview/utils/contextUsage.test.ts new file mode 100644 index 0000000000..d1e07ca185 --- /dev/null +++ b/packages/vscode-ide-companion/src/webview/utils/contextUsage.test.ts @@ -0,0 +1,87 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { computeContextUsage } from './contextUsage.js'; + +describe('computeContextUsage', () => { + it('returns null when there is no trusted token limit', () => { + expect( + computeContextUsage( + { + usage: { + promptTokens: 1234, + }, + }, + { + modelId: 'unknown-model', + name: 'Unknown Model', + }, + ), + ).toBeNull(); + }); + + it('prefers usageStats.tokenLimit over model metadata', () => { + expect( + computeContextUsage( + { + usage: { + promptTokens: 1000, + }, + tokenLimit: 4000, + }, + { + modelId: 'qwen3-max', + name: 'Qwen3 Max', + _meta: { contextLimit: 8000 }, + }, + ), + ).toEqual({ + percentLeft: 75, + usedTokens: 1000, + tokenLimit: 4000, + }); + }); + + it('falls back to model metadata when usageStats does not include a limit', () => { + expect( + computeContextUsage( + { + usage: { + promptTokens: 2000, + }, + }, + { + modelId: 'qwen3-max', + name: 'Qwen3 Max', + _meta: { contextLimit: 8000 }, + }, + ), + ).toEqual({ + percentLeft: 75, + usedTokens: 2000, + tokenLimit: 8000, + }); + }); + + it('uses inputTokens when promptTokens is unavailable', () => { + expect( + computeContextUsage( + { + usage: { + inputTokens: 3000, + }, + tokenLimit: 12000, + }, + null, + ), + ).toEqual({ + percentLeft: 75, + usedTokens: 3000, + tokenLimit: 12000, + }); + }); +}); diff --git a/packages/vscode-ide-companion/src/webview/utils/contextUsage.ts b/packages/vscode-ide-companion/src/webview/utils/contextUsage.ts new file mode 100644 index 0000000000..394cdb0361 --- /dev/null +++ b/packages/vscode-ide-companion/src/webview/utils/contextUsage.ts @@ -0,0 +1,46 @@ +/** + * @license + * Copyright 2025 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ModelInfo } from '@agentclientprotocol/sdk'; +import type { ContextUsage } from '@qwen-code/webui'; +import type { UsageStatsPayload } from '../../types/chatTypes.js'; + +export function computeContextUsage( + usageStats: UsageStatsPayload | null, + modelInfo: ModelInfo | null, +): ContextUsage | null { + if (!usageStats && !modelInfo) { + return null; + } + + const metaLimitRaw = modelInfo?._meta?.['contextLimit']; + const metaLimit = + typeof metaLimitRaw === 'number' || metaLimitRaw === null + ? metaLimitRaw + : undefined; + // Intentionally avoid DEFAULT_TOKEN_LIMIT here. The footer should disappear + // when neither ACP nor trusted model metadata provides a numeric limit. + const limit = usageStats?.tokenLimit ?? metaLimit; + // Prefer the ACP SDK's canonical inputTokens field and only fall back to the + // legacy promptTokens name for older payloads. + const used = + usageStats?.usage?.inputTokens ?? usageStats?.usage?.promptTokens ?? 0; + + if (typeof limit !== 'number' || limit <= 0 || used < 0) { + return null; + } + + const percentLeft = Math.max( + 0, + Math.min(100, Math.round(((limit - used) / limit) * 100)), + ); + + return { + percentLeft, + usedTokens: used, + tokenLimit: limit, + }; +}