refactor + tests

Author: liwilliam2021

src/core/webview/__tests__/messageEnhancer.test.ts

@@ -0,0 +1,303 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
+import { MessageEnhancer } from "../messageEnhancer"
+import { ProviderSettings, ClineMessage } from "@roo-code/types"
+import { TelemetryService } from "@roo-code/telemetry"
+import * as singleCompletionHandlerModule from "../../../utils/single-completion-handler"
+import { ProviderSettingsManager } from "../../config/ProviderSettingsManager"
+
+// Mock dependencies
+vi.mock("../../../utils/single-completion-handler")
+vi.mock("@roo-code/telemetry")
+
+describe("MessageEnhancer", () => {
+	let mockProviderSettingsManager: ProviderSettingsManager
+	let mockSingleCompletionHandler: ReturnType<typeof vi.fn>
+
+	const mockApiConfiguration: ProviderSettings = {
+		apiProvider: "openai",
+		apiKey: "test-key",
+		apiModelId: "gpt-4",
+	}
+
+	const mockListApiConfigMeta = [
+		{ id: "config1", name: "Config 1" },
+		{ id: "config2", name: "Config 2" },
+	]
+
+	beforeEach(() => {
+		// Reset all mocks
+		vi.clearAllMocks()
+
+		// Mock provider settings manager
+		mockProviderSettingsManager = {
+			getProfile: vi.fn().mockResolvedValue({
+				name: "Enhancement Config",
+				apiProvider: "anthropic",
+				apiKey: "enhancement-key",
+				apiModelId: "claude-3",
+			}),
+		} as any
+
+		// Mock single completion handler
+		mockSingleCompletionHandler = vi.fn().mockResolvedValue("Enhanced prompt text")
+		vi.mocked(singleCompletionHandlerModule).singleCompletionHandler = mockSingleCompletionHandler
+
+		// Mock TelemetryService
+		vi.mocked(TelemetryService).hasInstance = vi.fn().mockReturnValue(true)
+		// Mock the instance getter
+		Object.defineProperty(TelemetryService, "instance", {
+			get: vi.fn().mockReturnValue({
+				capturePromptEnhanced: vi.fn(),
+			}),
+			configurable: true,
+		})
+	})
+
+	afterEach(() => {
+		vi.restoreAllMocks()
+	})
+
+	describe("enhanceMessage", () => {
+		it("should enhance a simple message successfully", async () => {
+			const result = await MessageEnhancer.enhanceMessage({
+				text: "Write a function to calculate fibonacci",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			expect(result.success).toBe(true)
+			expect(result.enhancedText).toBe("Enhanced prompt text")
+			expect(result.error).toBeUndefined()
+
+			// Verify single completion handler was called with correct prompt
+			expect(mockSingleCompletionHandler).toHaveBeenCalledWith(
+				mockApiConfiguration,
+				expect.stringContaining("Write a function to calculate fibonacci"),
+			)
+		})
+
+		it("should use enhancement API config when provided", async () => {
+			const result = await MessageEnhancer.enhanceMessage({
+				text: "Test prompt",
+				apiConfiguration: mockApiConfiguration,
+				customSupportPrompts: {},
+				listApiConfigMeta: mockListApiConfigMeta,
+				enhancementApiConfigId: "config2",
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			expect(result.success).toBe(true)
+			expect(mockProviderSettingsManager.getProfile).toHaveBeenCalledWith({ id: "config2" })
+
+			// Verify the enhancement config was used instead of default
+			const expectedConfig = {
+				apiProvider: "anthropic",
+				apiKey: "enhancement-key",
+				apiModelId: "claude-3",
+			}
+			expect(mockSingleCompletionHandler).toHaveBeenCalledWith(expectedConfig, expect.any(String))
+		})
+
+		it("should include task history when enabled", async () => {
+			const mockClineMessages: ClineMessage[] = [
+				{ type: "ask", text: "Create a React component", ts: 1000 },
+				{ type: "say", say: "text", text: "I'll create a React component for you", ts: 2000 },
+				{ type: "ask", text: "Add props to the component", ts: 3000 },
+				{ type: "say", say: "reasoning", text: "Using tool", ts: 4000 }, // Should be filtered out
+			]
+
+			const result = await MessageEnhancer.enhanceMessage({
+				text: "Improve the component",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				includeTaskHistoryInEnhance: true,
+				currentClineMessages: mockClineMessages,
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			expect(result.success).toBe(true)
+
+			// Verify the prompt includes task history
+			const calledPrompt = mockSingleCompletionHandler.mock.calls[0][1]
+			expect(calledPrompt).toContain("Improve the component")
+			expect(calledPrompt).toContain("previous conversation context")
+			expect(calledPrompt).toContain("User: Create a React component")
+			expect(calledPrompt).toContain("Assistant: I'll create a React component for you")
+			expect(calledPrompt).toContain("User: Add props to the component")
+			expect(calledPrompt).not.toContain("Using tool") // reasoning messages should be filtered
+		})
+
+		it("should limit task history to last 10 messages", async () => {
+			// Create 15 messages
+			const mockClineMessages: ClineMessage[] = Array.from({ length: 15 }, (_, i) => ({
+				type: i % 2 === 0 ? "ask" : "say",
+				say: i % 2 === 1 ? "text" : undefined,
+				text: `Message ${i + 1}`,
+				ts: i * 1000,
+			})) as ClineMessage[]
+
+			await MessageEnhancer.enhanceMessage({
+				text: "Test",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				includeTaskHistoryInEnhance: true,
+				currentClineMessages: mockClineMessages,
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			const calledPrompt = mockSingleCompletionHandler.mock.calls[0][1]
+
+			// Should include messages 6-15 (last 10)
+			expect(calledPrompt).toContain("Message 6")
+			expect(calledPrompt).toContain("Message 15")
+			expect(calledPrompt).not.toContain("Message 5")
+		})
+
+		it("should truncate long messages in task history", async () => {
+			const longText = "A".repeat(600) // 600 characters
+			const mockClineMessages: ClineMessage[] = [{ type: "ask", text: longText, ts: 1000 }]
+
+			await MessageEnhancer.enhanceMessage({
+				text: "Test",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				includeTaskHistoryInEnhance: true,
+				currentClineMessages: mockClineMessages,
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			const calledPrompt = mockSingleCompletionHandler.mock.calls[0][1]
+
+			// Should truncate to 500 chars + "..."
+			expect(calledPrompt).toContain("A".repeat(500) + "...")
+			expect(calledPrompt).not.toContain("A".repeat(501))
+		})
+
+		it("should use custom support prompts when provided", async () => {
+			const customSupportPrompts = {
+				ENHANCE: "Custom enhancement template: ${userInput}",
+			}
+
+			await MessageEnhancer.enhanceMessage({
+				text: "Test prompt",
+				apiConfiguration: mockApiConfiguration,
+				customSupportPrompts,
+				listApiConfigMeta: mockListApiConfigMeta,
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			const calledPrompt = mockSingleCompletionHandler.mock.calls[0][1]
+			expect(calledPrompt).toBe("Custom enhancement template: Test prompt")
+		})
+
+		it("should handle errors gracefully", async () => {
+			mockSingleCompletionHandler.mockRejectedValue(new Error("API error"))
+
+			const result = await MessageEnhancer.enhanceMessage({
+				text: "Test",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			expect(result.success).toBe(false)
+			expect(result.error).toBe("API error")
+			expect(result.enhancedText).toBeUndefined()
+		})
+
+		it("should handle non-Error exceptions", async () => {
+			mockSingleCompletionHandler.mockRejectedValue("String error")
+
+			const result = await MessageEnhancer.enhanceMessage({
+				text: "Test",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			expect(result.success).toBe(false)
+			expect(result.error).toBe("String error")
+		})
+
+		it("should fall back to default config if enhancement config is invalid", async () => {
+			mockProviderSettingsManager.getProfile = vi.fn().mockResolvedValue({
+				name: "Invalid Config",
+				// Missing apiProvider
+			})
+
+			await MessageEnhancer.enhanceMessage({
+				text: "Test",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				enhancementApiConfigId: "config2",
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			// Should use the default config
+			expect(mockSingleCompletionHandler).toHaveBeenCalledWith(mockApiConfiguration, expect.any(String))
+		})
+
+		it("should handle empty task history gracefully", async () => {
+			const result = await MessageEnhancer.enhanceMessage({
+				text: "Test",
+				apiConfiguration: mockApiConfiguration,
+				listApiConfigMeta: mockListApiConfigMeta,
+				includeTaskHistoryInEnhance: true,
+				currentClineMessages: [],
+				providerSettingsManager: mockProviderSettingsManager,
+			})
+
+			expect(result.success).toBe(true)
+
+			const calledPrompt = mockSingleCompletionHandler.mock.calls[0][1]
+			// Should not include task history section
+			expect(calledPrompt).not.toContain("previous conversation context")
+		})
+	})
+
+	describe("captureTelemetry", () => {
+		it("should capture telemetry when TelemetryService is available", () => {
+			const mockTaskId = "task-123"
+			MessageEnhancer.captureTelemetry(mockTaskId)
+
+			expect(TelemetryService.hasInstance).toHaveBeenCalled()
+			expect(TelemetryService.instance.capturePromptEnhanced).toHaveBeenCalledWith(mockTaskId)
+		})
+
+		it("should handle missing TelemetryService gracefully", () => {
+			vi.mocked(TelemetryService).hasInstance = vi.fn().mockReturnValue(false)
+
+			// Should not throw
+			expect(() => MessageEnhancer.captureTelemetry("task-123")).not.toThrow()
+		})
+
+		it("should work without task ID", () => {
+			MessageEnhancer.captureTelemetry()
+
+			expect(TelemetryService.instance.capturePromptEnhanced).toHaveBeenCalledWith(undefined)
+		})
+	})
+
+	describe("extractTaskHistory", () => {
+		it("should filter and format messages correctly", () => {
+			const messages: ClineMessage[] = [
+				{ type: "ask", text: "User message 1", ts: 1000 },
+				{ type: "say", say: "text", text: "Assistant message 1", ts: 2000 },
+				{ type: "say", say: "reasoning", text: "Tool use", ts: 3000 },
+				{ type: "ask", text: "", ts: 4000 }, // Empty text
+				{ type: "say", say: "text", text: undefined, ts: 5000 }, // No text
+				{ type: "ask", text: "User message 2", ts: 6000 },
+			]
+
+			// Access private method through any type assertion for testing
+			const history = (MessageEnhancer as any).extractTaskHistory(messages)
+
+			expect(history).toContain("User: User message 1")
+			expect(history).toContain("Assistant: Assistant message 1")
+			expect(history).toContain("User: User message 2")
+			expect(history).not.toContain("Tool use")
+			expect(history.split("\n").length).toBe(3) // Only 3 valid messages
+		})
+	})
+})