Add chat session token usage

extensions/positron-assistant/package.nls.json

@@ -23,5 +23,5 @@
 	"configuration.inlineCompletionExcludes.description": "A list of [glob patterns](https://aka.ms/vscode-glob-patterns) to exclude from inline completions.",
 	"configuration.gitIntegration.description": "Enable Positron Assistant git integration.",
 	"configuration.getTableSummary.description": "Enable Positron Assistant get table summary tool.",
-	"configuration.showTokenUsage.description": "Show token usage in the chat view for supported providers for each message and response including prompts. Check with your provider for detailed usage."
+	"configuration.showTokenUsage.description": "Show token usage in the chat view for supported providers for the session and each message and response including prompts. Check with your provider for detailed usage."
 }

extensions/positron-assistant/src/models.ts

@@ -117,7 +117,7 @@ class EchoLanguageModel implements positron.ai.LanguageModelChatProvider {
 		token: vscode.CancellationToken
 	): Promise<any> {
 		const _messages = toAIMessage(messages);
-		const message = _messages[0];
+		const message = _messages.length > 1 ? _messages[_messages.length - 2] : _messages[0]; // Get the last user message, the last message is the context
 
 		if (typeof message.content === 'string') {
 			message.content = [{ type: 'text', text: message.content }];

src/vs/workbench/contrib/chat/browser/chatInputPart.ts

@@ -107,6 +107,8 @@ import { IModePickerDelegate, ModePickerActionItem } from './modelPicker/modePic
 import { ChatRuntimeSessionContext } from './contrib/chatRuntimeSessionContext.js';
 import { RuntimeSessionContextAttachmentWidget } from './attachments/runtimeSessionContextAttachment.js';
 import { RuntimeSessionAttachmentWidget } from './chatRuntimeAttachmentWidget.js';
+// eslint-disable-next-line no-duplicate-imports
+import { isResponseVM } from '../common/chatViewModel.js';
 // --- End Positron ---
 
 const $ = dom.$;
@@ -246,6 +248,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 
 	private chatEditingSessionWidgetContainer!: HTMLElement;
 
+	// --- Start Positron ---
+	private tokenUsageContainer!: HTMLElement;
+	// --- End Positron ---
+
 	private _inputPartHeight: number;
 	get inputPartHeight() {
 		return this._inputPartHeight;
@@ -1041,6 +1047,9 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 			elements = dom.h('.interactive-input-part', [
 				dom.h('.interactive-input-and-edit-session', [
 					dom.h('.chat-editing-session@chatEditingSessionWidgetContainer'),
+					// --- Start Positron ---
+					dom.h('.chat-token-usage-status@tokenUsageContainer'),
+					// --- End Positron ---
 					dom.h('.interactive-input-and-side-toolbar@inputAndSideToolbar', [
 						dom.h('.chat-input-container@inputContainer', [
 							dom.h('.chat-editor-container@editorContainer'),
@@ -1059,6 +1068,9 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 			elements = dom.h('.interactive-input-part', [
 				dom.h('.interactive-input-followups@followupsContainer'),
 				dom.h('.chat-editing-session@chatEditingSessionWidgetContainer'),
+				// --- Start Positron ---
+				dom.h('.chat-token-usage-status@tokenUsageContainer'),
+				// --- End Positron ---
 				dom.h('.interactive-input-and-side-toolbar@inputAndSideToolbar', [
 					dom.h('.chat-input-container@inputContainer', [
 						dom.h('.chat-attachments-container@attachmentsContainer', [
@@ -1087,6 +1099,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		const toolbarsContainer = elements.inputToolbars;
 		const attachmentToolbarContainer = elements.attachmentToolbar;
 		this.chatEditingSessionWidgetContainer = elements.chatEditingSessionWidgetContainer;
+		// --- Start Positron ---
+		this.tokenUsageContainer = elements.tokenUsageContainer;
+		this.tokenUsageContainer.style.display = 'none'; // Initially hidden
+		// --- End Positron ---
 		if (this.options.enableImplicitContext) {
 			this._implicitContext = this._register(
 				this.instantiationService.createInstance(ChatImplicitContext),
@@ -1732,6 +1748,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		this._followupsHeight = data.followupsHeight;
 		this._editSessionWidgetHeight = data.chatEditingStateHeight;
 
+		// --- Start Positron ---
+		this._inputPartHeight += this.tokenUsageHeight;
+		// --- End Positron ---
+
 		const initialEditorScrollWidth = this._inputEditor.getScrollWidth();
 		const newEditorWidth = width - data.inputPartHorizontalPadding - data.editorBorder - data.inputPartHorizontalPaddingInside - data.toolbarsWidth - data.sideToolbarWidth;
 		const newDimension = { width: newEditorWidth, height: inputEditorHeight };
@@ -1766,6 +1786,9 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 			toolbarsHeight: this.options.renderStyle === 'compact' ? 0 : 22,
 			chatEditingStateHeight: this.chatEditingSessionWidgetContainer.offsetHeight,
 			sideToolbarWidth: this.inputSideToolbarContainer ? dom.getTotalWidth(this.inputSideToolbarContainer) + 4 /*gap*/ : 0,
+			// --- Start Positron ---
+			tokenUsageHeight: this.tokenUsageHeight,
+			// --- End Positron ---
 		};
 	}
 
@@ -1781,6 +1804,69 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		const inputHistory = [...this.history];
 		this.historyService.saveHistory(this.location, inputHistory);
 	}
+
+	// --- Start Positron ---
+	/**
+	 * Calculate the total token usage from a view model's items
+	 */
+	private calculateTotalTokenUsage(viewModel: any): { totalInputTokens: number; totalOutputTokens: number } | undefined {
+		if (!viewModel) {
+			return undefined;
+		}
+
+		let totalInputTokens = 0;
+		let totalOutputTokens = 0;
+		let hasAnyTokenUsage = false;
+
+		for (const item of viewModel.getItems()) {
+			if (isResponseVM(item) && item.tokenUsage && item.isComplete) {
+				totalInputTokens += item.tokenUsage.inputTokens;
+				totalOutputTokens += item.tokenUsage.outputTokens;
+				hasAnyTokenUsage = true;
+			}
+		}
+
+		return hasAnyTokenUsage ? { totalInputTokens, totalOutputTokens } : undefined;
+	}
+
+	/**
+	 * Update the token usage status display
+	 */
+	updateTokenUsageDisplay(viewModel: any): void {
+		if (!this.tokenUsageContainer) {
+			return;
+		}
+
+		const previousDisplay = this.tokenUsageContainer.style.display;
+		const showTokens = this.configurationService.getValue<boolean>('positron.assistant.showTokenUsage.enable');
+		if (!showTokens) {
+			this.tokenUsageContainer.style.display = 'none';
+		} else {
+			const totalTokens = this.calculateTotalTokenUsage(viewModel);
+			if (totalTokens && totalTokens.totalInputTokens > 0 && totalTokens.totalOutputTokens > 0) {
+				dom.clearNode(this.tokenUsageContainer);
+				this.tokenUsageContainer.appendChild(
+					dom.$('.token-usage-total', undefined,
+						localize('totalTokenUsage', "Total tokens: ↑{0} ↓{1}", totalTokens.totalInputTokens, totalTokens.totalOutputTokens)
+					)
+				);
+				this.tokenUsageContainer.style.display = 'block';
+			} else {
+				this.tokenUsageContainer.style.display = 'none';
+			}
+		}
+
+		// Fire height change event if visibility changed
+		if (previousDisplay !== this.tokenUsageContainer.style.display) {
+			this._onDidChangeHeight.fire();
+		}
+	}
+
+	get tokenUsageHeight(): number {
+		return (this.tokenUsageContainer && this.tokenUsageContainer.style.display !== 'none')
+			? this.tokenUsageContainer.offsetHeight : 0;
+	}
+	// --- End Positron ---
 }
 
 const historyKeyFn = (entry: IChatHistoryEntry) => JSON.stringify({ ...entry, state: { ...entry.state, chatMode: undefined } });

src/vs/workbench/contrib/chat/browser/chatListRenderer.ts

@@ -635,14 +635,19 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 			experimentalTokenUsage = experimentalTokenUsage.concat(approximateTokenCount);
 		}
 
+		const tokenUsageElements = templateData.value.getElementsByClassName('token-usage');
 		if (element.tokenUsage && element.isComplete && showTokens && experimentalTokenUsage.includes(element.tokenUsage.provider)) {
-			const tokenUsageElements = templateData.value.getElementsByClassName('token-usage');
 			const tokenUsageText = localize('tokenUsage', "Tokens: ↑{0} ↓{1}", element.tokenUsage.inputTokens, element.tokenUsage.outputTokens);
 			if (tokenUsageElements.length > 0) {
 				tokenUsageElements[0].textContent = tokenUsageText;
 			} else {
 				templateData.value.appendChild(dom.$('.token-usage', undefined, tokenUsageText));
 			}
+		} else {
+			// Remove token usage elements if they exist and should not be shown
+			while (tokenUsageElements.length > 0) {
+				tokenUsageElements[0].remove();
+			}
 		}
 		// --- End Positron ---
 

src/vs/workbench/contrib/chat/browser/chatWidget.ts

@@ -375,6 +375,9 @@ export class ChatWidget extends Disposable implements IChatWidget {
 			} else if (e.affectsConfiguration(ChatConfiguration.EditRequests)) {
 				this.settingChangeCounter++;
 				this.onDidChangeItems();
+			} else if (e.affectsConfiguration('positron.assistant.showTokenUsage.enable') || e.affectsConfiguration('positron.assistant.approximateTokenCount')) {
+				this.settingChangeCounter++;
+				this.onDidChangeItems();
 			}
 		}));
 
@@ -764,6 +767,11 @@ export class ChatWidget extends Disposable implements IChatWidget {
 			}
 
 			this.renderFollowups();
+
+			// --- Start Positron ---
+			// Update token usage display when items change
+			this.input.updateTokenUsageDisplay(this.viewModel);
+			// --- End Positron ---
 		}
 	}
 
@@ -988,6 +996,10 @@ Type \`/\` to use predefined commands such as \`/help\`.`,
 				// Do it after a timeout because the container is not visible yet (it should be but offsetHeight returns 0 here)
 				if (this._visible) {
 					this.onDidChangeItems(true);
+					// --- Start Positron ---
+					// Update token usage display when widget becomes visible
+					this.input.updateTokenUsageDisplay(this.viewModel);
+					// --- End Positron ---
 				}
 			}, 0));
 		} else if (wasVisible) {

src/vs/workbench/contrib/chat/browser/media/positronChat.css

@@ -7,3 +7,15 @@
 	width: 14px;
 	height: 14px;
 }
+
+.chat-token-usage-status {
+	padding-top: 2px;
+	font-size: 12px;
+	color: var(--vscode-descriptionForeground);
+	border-top: 1px solid var(--vscode-panel-border);
+	flex-shrink: 0;
+}
+
+.chat-token-usage-status .token-usage-total {
+	display: inline-block;
+}

test/e2e/pages/positronAssistant.ts

@@ -186,6 +186,19 @@ export class Assistant {
 		await expect(this.code.driver.page.locator('.token-usage')).not.toBeVisible();
 	}
 
+	async verifyTotalTokenUsageVisible() {
+		await expect(this.code.driver.page.locator('.token-usage-total')).toBeVisible();
+		await expect(this.code.driver.page.locator('.token-usage-total')).toHaveText(/Total tokens: ↑\d+ ↓\d+/);
+	}
+
+	async verifyNumberOfVisibleResponses(expectedCount: number, checkTokenUsage: boolean = false) {
+		const responses = this.code.driver.page.locator('.interactive-response');
+		await expect(responses).toHaveCount(expectedCount);
+		if (checkTokenUsage) {
+			this.code.driver.page.locator('.token-usage').nth(expectedCount - 1).waitFor({ state: 'visible' });
+		}
+	}
+
 	async getTokenUsage() {
 		const tokenUsageElement = this.code.driver.page.locator('.token-usage');
 		await expect(tokenUsageElement).toBeVisible();
@@ -198,4 +211,22 @@ export class Assistant {
 			outputTokens: outputMatch ? parseInt(outputMatch[1], 10) : 0
 		};
 	}
+
+	async getTotalTokenUsage() {
+		const totalTokenUsageElement = this.code.driver.page.locator('.token-usage-total');
+		await expect(totalTokenUsageElement).toBeVisible();
+		const text = await totalTokenUsageElement.textContent();
+		console.log('Total Token Usage Text:', text);
+		expect(text).not.toBeNull();
+		const totalMatch = text ? text.match(/Total tokens: ↑(\d+) ↓(\d+)/) : null;
+		return {
+			inputTokens: totalMatch ? parseInt(totalMatch[1], 10) : 0,
+			outputTokens: totalMatch ? parseInt(totalMatch[2], 10) : 0
+		};
+	}
+
+	async waitForReadyToSend(timeout: number = 5000) {
+		await this.code.driver.page.waitForSelector('.chat-input-toolbars .codicon-send', { timeout });
+		await this.code.driver.page.waitForSelector('.detail-container .detail:has-text("Working")', { state: 'hidden', timeout });
+	}
 }

test/e2e/tests/positron-assistant/positron-assistant.test.ts

@@ -104,6 +104,7 @@ test.describe('Positron Assistant Chat Editing', { tag: [tags.WIN, tags.ASSISTAN
 	test.beforeAll('Enable Assistant', async function ({ app }) {
 		await app.workbench.assistant.openPositronAssistantChat();
 		await app.workbench.quickaccess.runCommand('positron-assistant.configureModels');
+
 		await app.workbench.assistant.selectModelProvider('echo');
 		await app.workbench.assistant.clickSignInButton();
 		await app.workbench.assistant.clickCloseButton();
@@ -158,13 +159,12 @@ test.describe('Positron Assistant Chat Tokens', { tag: [tags.WIN, tags.ASSISTANT
 		await app.workbench.assistant.selectModelProvider('echo');
 		await app.workbench.assistant.clickSignInButton();
 		await app.workbench.assistant.clickCloseButton();
-
-		await settings.set({ 'positron.assistant.showTokenUsage.enable': true });
-		await settings.set({ 'positron.assistant.approximateTokenCount': ['echo'] });
 	});
 
-	test.beforeEach('Clear chat', async function ({ app }) {
+	test.beforeEach('Clear chat', async function ({ app, settings }) {
+		await settings.set({ 'positron.assistant.showTokenUsage.enable': true });
 		await app.workbench.assistant.clickNewChatButton();
+		await settings.set({ 'positron.assistant.approximateTokenCount': ['echo'] });
 	});
 
 	test.afterAll('Sign out of Assistant', async function ({ app }) {
@@ -198,4 +198,40 @@ test.describe('Positron Assistant Chat Tokens', { tag: [tags.WIN, tags.ASSISTANT
 
 		expect(await app.workbench.assistant.verifyTokenUsageNotVisible());
 	});
+
+	test('Token usage updates when settings change', async function ({ app, settings }) {
+		await app.workbench.assistant.enterChatMessage('What is the meaning of life?');
+		await app.workbench.assistant.verifyTokenUsageVisible();
+
+		await settings.set({ 'positron.assistant.approximateTokenCount': [] });
+		expect(await app.workbench.assistant.verifyTokenUsageNotVisible());
+
+		await settings.set({ 'positron.assistant.approximateTokenCount': ['echo'] });
+		await app.workbench.assistant.verifyTokenUsageVisible();
+
+		await settings.set({ 'positron.assistant.showTokenUsage.enable': false });
+		expect(await app.workbench.assistant.verifyTokenUsageNotVisible());
+
+		await settings.set({ 'positron.assistant.showTokenUsage.enable': true });
+		await app.workbench.assistant.verifyTokenUsageVisible();
+	});
+
+	test('Total token usage is displayed in chat header', async function ({ app }) {
+		const message1 = 'What is the meaning of life?';
+		const message2 = 'Forty-two';
+
+		await app.workbench.assistant.enterChatMessage(message1);
+		await app.workbench.assistant.waitForReadyToSend();
+		await app.workbench.assistant.enterChatMessage(message2);
+
+		await app.workbench.assistant.waitForReadyToSend();
+		await app.workbench.assistant.verifyNumberOfVisibleResponses(2, true);
+
+		const totalTokens = await app.workbench.assistant.getTotalTokenUsage();
+		expect(totalTokens).toBeDefined();
+		expect(totalTokens).toMatchObject({
+			inputTokens: message1.length + message2.length,
+			outputTokens: message1.length + message2.length
+		});
+	});
 });