Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 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
1 change: 1 addition & 0 deletions packages/types/src/global-settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ export const globalSettingsSchema = z.object({
customModePrompts: customModePromptsSchema.optional(),
customSupportPrompts: customSupportPromptsSchema.optional(),
enhancementApiConfigId: z.string().optional(),
includeTaskHistoryInEnhance: z.boolean().optional(),
historyPreviewCollapsed: z.boolean().optional(),
profileThresholds: z.record(z.string(), z.number()).optional(),
hasOpenedModeSelector: z.boolean().optional(),
Expand Down
4 changes: 4 additions & 0 deletions src/core/webview/ClineProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1511,6 +1511,7 @@ export class ClineProvider
followupAutoApproveTimeoutMs,
includeDiagnosticMessages,
maxDiagnosticMessages,
includeTaskHistoryInEnhance,
} = await this.getState()

const telemetryKey = process.env.POSTHOG_API_KEY
Expand Down Expand Up @@ -1635,6 +1636,7 @@ export class ClineProvider
followupAutoApproveTimeoutMs: followupAutoApproveTimeoutMs ?? 60000,
includeDiagnosticMessages: includeDiagnosticMessages ?? true,
maxDiagnosticMessages: maxDiagnosticMessages ?? 50,
includeTaskHistoryInEnhance: includeTaskHistoryInEnhance ?? false,
}
}

Expand Down Expand Up @@ -1805,6 +1807,8 @@ export class ClineProvider
// Add diagnostic message settings
includeDiagnosticMessages: stateValues.includeDiagnosticMessages ?? true,
maxDiagnosticMessages: stateValues.maxDiagnosticMessages ?? 50,
// Add includeTaskHistoryInEnhance setting
includeTaskHistoryInEnhance: stateValues.includeTaskHistoryInEnhance ?? false,
}
}

Expand Down
303 changes: 303 additions & 0 deletions src/core/webview/__tests__/messageEnhancer.test.ts
Original file line number Diff line number Diff line change
@@ -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
})
})
})
Loading
Loading