Add a total token test

Fix echo model echoing the first message on subsequent responses

Author: timtmok

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[_messages.length - 2]; // Get the last user message, the last message is the context
 
 		if (typeof message.content === 'string') {
 			message.content = [{ type: 'text', text: message.content }];

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

@@ -215,4 +215,23 @@ test.describe('Positron Assistant Chat Tokens', { tag: [tags.WIN, tags.ASSISTANT
 		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
+		});
+	});
 });