Browse Source

feat: allow variable interpolation into the custom system prompt (#2863)

* feat: allow variable interpolation into the custom system prompt

* fix: allow the test to pass on windows by using the path module
Daniel 8 months ago
parent
commit
0f64849542

+ 131 - 0
src/core/prompts/sections/__tests__/custom-system-prompt.test.ts

@@ -0,0 +1,131 @@
+import path from "path"
+import { readFile } from "fs/promises"
+import { Mode } from "../../../../shared/modes" // Adjusted import path
+import { loadSystemPromptFile, PromptVariables } from "../custom-system-prompt"
+
+// Mock the fs/promises module
+jest.mock("fs/promises")
+
+// Cast the mocked readFile to the correct Jest mock type
+const mockedReadFile = readFile as jest.MockedFunction<typeof readFile>
+
+describe("loadSystemPromptFile", () => {
+	// Corrected PromptVariables type and added mockMode
+	const mockVariables: PromptVariables = {
+		workspace: "/path/to/workspace",
+	}
+	const mockCwd = "/mock/cwd"
+	const mockMode: Mode = "test" // Use Mode type, e.g., 'test'
+	// Corrected expected file path format
+	const expectedFilePath = path.join(mockCwd, ".roo", `system-prompt-${mockMode}`)
+
+	beforeEach(() => {
+		// Clear mocks before each test
+		mockedReadFile.mockClear()
+	})
+
+	it("should return an empty string if the file does not exist (ENOENT)", async () => {
+		const error: NodeJS.ErrnoException = new Error("File not found")
+		error.code = "ENOENT"
+		mockedReadFile.mockRejectedValue(error)
+
+		// Added mockMode argument
+		const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
+
+		expect(result).toBe("")
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+
+	// Updated test: should re-throw unexpected errors
+	it("should re-throw unexpected errors from readFile", async () => {
+		const expectedError = new Error("Some other error")
+		mockedReadFile.mockRejectedValue(expectedError)
+
+		// Assert that the promise rejects with the specific error
+		await expect(loadSystemPromptFile(mockCwd, mockMode, mockVariables)).rejects.toThrow(expectedError)
+
+		// Verify readFile was still called correctly
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+
+	it("should return an empty string if the file content is empty", async () => {
+		mockedReadFile.mockResolvedValue("")
+
+		// Added mockMode argument
+		const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
+
+		expect(result).toBe("")
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+
+	// Updated test to only check workspace interpolation
+	it("should correctly interpolate workspace variable", async () => {
+		const template = "Workspace is: {{workspace}}"
+		mockedReadFile.mockResolvedValue(template)
+
+		// Added mockMode argument
+		const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
+
+		expect(result).toBe("Workspace is: /path/to/workspace")
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+
+	// Updated test for multiple occurrences of workspace
+	it("should handle multiple occurrences of the workspace variable", async () => {
+		const template = "Path: {{workspace}}/{{workspace}}"
+		mockedReadFile.mockResolvedValue(template)
+
+		// Added mockMode argument
+		const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
+
+		expect(result).toBe("Path: /path/to/workspace//path/to/workspace")
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+
+	// Updated test for mixed used/unused
+	it("should handle mixed used workspace and unused variables", async () => {
+		const template = "Workspace: {{workspace}}, Unused: {{unusedVar}}, Another: {{another}}"
+		mockedReadFile.mockResolvedValue(template)
+
+		// Added mockMode argument
+		const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
+
+		// Unused variables should remain untouched
+		expect(result).toBe("Workspace: /path/to/workspace, Unused: {{unusedVar}}, Another: {{another}}")
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+
+	// Test remains valid, just needs the mode argument and updated template
+	it("should handle templates with placeholders not present in variables", async () => {
+		const template = "Workspace: {{workspace}}, Missing: {{missingPlaceholder}}"
+		mockedReadFile.mockResolvedValue(template)
+
+		// Added mockMode argument
+		const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
+
+		expect(result).toBe("Workspace: /path/to/workspace, Missing: {{missingPlaceholder}}")
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+
+	// Removed the test for extra keys as PromptVariables is simple now
+
+	// Test remains valid, just needs the mode argument
+	it("should handle template with no variables", async () => {
+		const template = "This is a static prompt."
+		mockedReadFile.mockResolvedValue(template)
+
+		// Added mockMode argument
+		const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
+
+		expect(result).toBe("This is a static prompt.")
+		expect(mockedReadFile).toHaveBeenCalledTimes(1)
+		expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
+	})
+})

+ 25 - 2
src/core/prompts/sections/custom-system-prompt.ts

@@ -3,6 +3,24 @@ import path from "path"
 import { Mode } from "../../../shared/modes"
 import { fileExistsAtPath } from "../../../utils/fs"
 
+export type PromptVariables = {
+	workspace?: string
+}
+
+function interpolatePromptContent(content: string, variables: PromptVariables): string {
+	let interpolatedContent = content
+	for (const key in variables) {
+		if (
+			Object.prototype.hasOwnProperty.call(variables, key) &&
+			variables[key as keyof PromptVariables] !== undefined
+		) {
+			const placeholder = new RegExp(`\\{\\{${key}\\}\\}`, "g")
+			interpolatedContent = interpolatedContent.replace(placeholder, variables[key as keyof PromptVariables]!)
+		}
+	}
+	return interpolatedContent
+}
+
 /**
  * Safely reads a file, returning an empty string if the file doesn't exist
  */
@@ -31,9 +49,14 @@ export function getSystemPromptFilePath(cwd: string, mode: Mode): string {
  * Loads custom system prompt from a file at .roo/system-prompt-[mode slug]
  * If the file doesn't exist, returns an empty string
  */
-export async function loadSystemPromptFile(cwd: string, mode: Mode): Promise<string> {
+export async function loadSystemPromptFile(cwd: string, mode: Mode, variables: PromptVariables): Promise<string> {
 	const filePath = getSystemPromptFilePath(cwd, mode)
-	return safeReadFile(filePath)
+	const rawContent = await safeReadFile(filePath)
+	if (!rawContent) {
+		return ""
+	}
+	const interpolatedContent = interpolatePromptContent(rawContent, variables)
+	return interpolatedContent
 }
 
 /**

+ 6 - 1
src/core/prompts/system.ts

@@ -9,6 +9,7 @@ import {
 	getModeBySlug,
 	getGroupName,
 } from "../../shared/modes"
+import { PromptVariables } from "./sections/custom-system-prompt"
 import { DiffStrategy } from "../../shared/tools"
 import { McpHub } from "../../services/mcp/McpHub"
 import { getToolDescriptionsForMode } from "./tools"
@@ -125,7 +126,10 @@ export const SYSTEM_PROMPT = async (
 	}
 
 	// Try to load custom system prompt from file
-	const fileCustomSystemPrompt = await loadSystemPromptFile(cwd, mode)
+	const variablesForPrompt: PromptVariables = {
+		workspace: cwd,
+	}
+	const fileCustomSystemPrompt = await loadSystemPromptFile(cwd, mode, variablesForPrompt)
 
 	// Check if it's a custom mode
 	const promptComponent = getPromptComponent(customModePrompts?.[mode])
@@ -143,6 +147,7 @@ export const SYSTEM_PROMPT = async (
 			mode,
 			{ language: language ?? formatLanguage(vscode.env.language), rooIgnoreInstructions },
 		)
+
 		// For file-based prompts, don't include the tool sections
 		return `${roleDefinition}