Skip to content

Commit 1875137

Browse files
committed
feat: allow variable interpolation into the custom system prompt
1 parent f06567d commit 1875137

File tree

3 files changed

+162
-3
lines changed

3 files changed

+162
-3
lines changed
Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
// src/core/prompts/sections/__tests__/custom-system-prompt.test.ts
2+
import { readFile } from "fs/promises"
3+
import { Mode } from "../../../../shared/modes" // Adjusted import path
4+
import { loadSystemPromptFile, PromptVariables } from "../custom-system-prompt"
5+
6+
// Mock the fs/promises module
7+
jest.mock("fs/promises")
8+
9+
// Cast the mocked readFile to the correct Jest mock type
10+
const mockedReadFile = readFile as jest.MockedFunction<typeof readFile>
11+
12+
describe("loadSystemPromptFile", () => {
13+
// Corrected PromptVariables type and added mockMode
14+
const mockVariables: PromptVariables = {
15+
workspace: "/path/to/workspace",
16+
}
17+
const mockCwd = "/mock/cwd"
18+
const mockMode: Mode = "test" // Use Mode type, e.g., 'test'
19+
// Corrected expected file path format
20+
const expectedFilePath = `${mockCwd}/.roo/system-prompt-${mockMode}`
21+
22+
beforeEach(() => {
23+
// Clear mocks before each test
24+
mockedReadFile.mockClear()
25+
})
26+
27+
it("should return an empty string if the file does not exist (ENOENT)", async () => {
28+
const error: NodeJS.ErrnoException = new Error("File not found")
29+
error.code = "ENOENT"
30+
mockedReadFile.mockRejectedValue(error)
31+
32+
// Added mockMode argument
33+
const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
34+
35+
expect(result).toBe("")
36+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
37+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
38+
})
39+
40+
// Updated test: should re-throw unexpected errors
41+
it("should re-throw unexpected errors from readFile", async () => {
42+
const expectedError = new Error("Some other error")
43+
mockedReadFile.mockRejectedValue(expectedError)
44+
45+
// Assert that the promise rejects with the specific error
46+
await expect(loadSystemPromptFile(mockCwd, mockMode, mockVariables)).rejects.toThrow(expectedError)
47+
48+
// Verify readFile was still called correctly
49+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
50+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
51+
})
52+
53+
it("should return an empty string if the file content is empty", async () => {
54+
mockedReadFile.mockResolvedValue("")
55+
56+
// Added mockMode argument
57+
const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
58+
59+
expect(result).toBe("")
60+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
61+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
62+
})
63+
64+
// Updated test to only check workspace interpolation
65+
it("should correctly interpolate workspace variable", async () => {
66+
const template = "Workspace is: {{workspace}}"
67+
mockedReadFile.mockResolvedValue(template)
68+
69+
// Added mockMode argument
70+
const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
71+
72+
expect(result).toBe("Workspace is: /path/to/workspace")
73+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
74+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
75+
})
76+
77+
// Updated test for multiple occurrences of workspace
78+
it("should handle multiple occurrences of the workspace variable", async () => {
79+
const template = "Path: {{workspace}}/{{workspace}}"
80+
mockedReadFile.mockResolvedValue(template)
81+
82+
// Added mockMode argument
83+
const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
84+
85+
expect(result).toBe("Path: /path/to/workspace//path/to/workspace")
86+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
87+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
88+
})
89+
90+
// Updated test for mixed used/unused
91+
it("should handle mixed used workspace and unused variables", async () => {
92+
const template = "Workspace: {{workspace}}, Unused: {{unusedVar}}, Another: {{another}}"
93+
mockedReadFile.mockResolvedValue(template)
94+
95+
// Added mockMode argument
96+
const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
97+
98+
// Unused variables should remain untouched
99+
expect(result).toBe("Workspace: /path/to/workspace, Unused: {{unusedVar}}, Another: {{another}}")
100+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
101+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
102+
})
103+
104+
// Test remains valid, just needs the mode argument and updated template
105+
it("should handle templates with placeholders not present in variables", async () => {
106+
const template = "Workspace: {{workspace}}, Missing: {{missingPlaceholder}}"
107+
mockedReadFile.mockResolvedValue(template)
108+
109+
// Added mockMode argument
110+
const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
111+
112+
expect(result).toBe("Workspace: /path/to/workspace, Missing: {{missingPlaceholder}}")
113+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
114+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
115+
})
116+
117+
// Removed the test for extra keys as PromptVariables is simple now
118+
119+
// Test remains valid, just needs the mode argument
120+
it("should handle template with no variables", async () => {
121+
const template = "This is a static prompt."
122+
mockedReadFile.mockResolvedValue(template)
123+
124+
// Added mockMode argument
125+
const result = await loadSystemPromptFile(mockCwd, mockMode, mockVariables)
126+
127+
expect(result).toBe("This is a static prompt.")
128+
expect(mockedReadFile).toHaveBeenCalledTimes(1)
129+
expect(mockedReadFile).toHaveBeenCalledWith(expectedFilePath, "utf-8")
130+
})
131+
})

src/core/prompts/sections/custom-system-prompt.ts

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,24 @@ import path from "path"
33
import { Mode } from "../../../shared/modes"
44
import { fileExistsAtPath } from "../../../utils/fs"
55

6+
export type PromptVariables = {
7+
workspace?: string
8+
}
9+
10+
function interpolatePromptContent(content: string, variables: PromptVariables): string {
11+
let interpolatedContent = content
12+
for (const key in variables) {
13+
if (
14+
Object.prototype.hasOwnProperty.call(variables, key) &&
15+
variables[key as keyof PromptVariables] !== undefined
16+
) {
17+
const placeholder = new RegExp(`\\{\\{${key}\\}\\}`, "g")
18+
interpolatedContent = interpolatedContent.replace(placeholder, variables[key as keyof PromptVariables]!)
19+
}
20+
}
21+
return interpolatedContent
22+
}
23+
624
/**
725
* Safely reads a file, returning an empty string if the file doesn't exist
826
*/
@@ -31,9 +49,14 @@ export function getSystemPromptFilePath(cwd: string, mode: Mode): string {
3149
* Loads custom system prompt from a file at .roo/system-prompt-[mode slug]
3250
* If the file doesn't exist, returns an empty string
3351
*/
34-
export async function loadSystemPromptFile(cwd: string, mode: Mode): Promise<string> {
52+
export async function loadSystemPromptFile(cwd: string, mode: Mode, variables: PromptVariables): Promise<string> {
3553
const filePath = getSystemPromptFilePath(cwd, mode)
36-
return safeReadFile(filePath)
54+
const rawContent = await safeReadFile(filePath)
55+
if (!rawContent) {
56+
return ""
57+
}
58+
const interpolatedContent = interpolatePromptContent(rawContent, variables)
59+
return interpolatedContent
3760
}
3861

3962
/**

src/core/prompts/system.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
getModeBySlug,
1010
getGroupName,
1111
} from "../../shared/modes"
12+
import { PromptVariables } from "./sections/custom-system-prompt"
1213
import { DiffStrategy } from "../../shared/tools"
1314
import { McpHub } from "../../services/mcp/McpHub"
1415
import { getToolDescriptionsForMode } from "./tools"
@@ -125,7 +126,10 @@ export const SYSTEM_PROMPT = async (
125126
}
126127

127128
// Try to load custom system prompt from file
128-
const fileCustomSystemPrompt = await loadSystemPromptFile(cwd, mode)
129+
const variablesForPrompt: PromptVariables = {
130+
workspace: cwd,
131+
}
132+
const fileCustomSystemPrompt = await loadSystemPromptFile(cwd, mode, variablesForPrompt)
129133

130134
// Check if it's a custom mode
131135
const promptComponent = getPromptComponent(customModePrompts?.[mode])
@@ -143,6 +147,7 @@ export const SYSTEM_PROMPT = async (
143147
mode,
144148
{ language: language ?? formatLanguage(vscode.env.language), rooIgnoreInstructions },
145149
)
150+
146151
// For file-based prompts, don't include the tool sections
147152
return `${roleDefinition}
148153

0 commit comments

Comments
 (0)