Skip to content

Commit 4305358

Browse files
authored
fix: prepend agent prompt in serve mode (#8393)
* fix: prepend agent prompt in serve mode * fix: prompt prepending and tests * fix: lint and format prepend prompt
1 parent 03fda47 commit 4305358

File tree

4 files changed

+213
-14
lines changed

4 files changed

+213
-14
lines changed

extensions/cli/src/commands/chat.ts

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import { gracefulExit } from "../util/exit.js";
3030
import { formatAnthropicError, formatError } from "../util/formatError.js";
3131
import { logger } from "../util/logger.js";
3232
import { question } from "../util/prompt.js";
33+
import { prependPrompt } from "../util/promptProcessor.js";
3334
import {
3435
calculateContextUsagePercentage,
3536
countChatHistoryTokens,
@@ -481,9 +482,11 @@ async function runHeadlessMode(
481482
const agentFileState = await serviceContainer.get<AgentFileServiceState>(
482483
SERVICE_NAMES.AGENT_FILE,
483484
);
484-
const initialPrompt =
485-
`${agentFileState?.agentFile?.prompt ?? ""}\n\n${prompt ?? ""}`.trim() ||
486-
undefined;
485+
486+
const initialPrompt = prependPrompt(
487+
agentFileState?.agentFile?.prompt,
488+
prompt,
489+
);
487490
const initialUserInput = await processAndCombinePrompts(
488491
options.prompt,
489492
initialPrompt,
@@ -556,9 +559,10 @@ export async function chat(prompt?: string, options: ChatOptions = {}) {
556559
SERVICE_NAMES.AGENT_FILE,
557560
);
558561

559-
const initialPrompt =
560-
`${agentFileState?.agentFile?.prompt ?? ""}\n\n${prompt ?? ""}`.trim() ||
561-
undefined;
562+
const initialPrompt = prependPrompt(
563+
agentFileState?.agentFile?.prompt,
564+
prompt,
565+
);
562566

563567
// Start TUI with skipOnboarding since we already handled it
564568
const tuiOptions: any = {

extensions/cli/src/commands/serve.ts

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import express, { Request, Response } from "express";
44

55
import { ToolPermissionServiceState } from "src/services/ToolPermissionService.js";
66
import { posthogService } from "src/telemetry/posthogService.js";
7+
import { prependPrompt } from "src/util/promptProcessor.js";
78

89
import { getAccessToken, getAssistantSlug } from "../auth/workos.js";
910
import { runEnvironmentInstallSafe } from "../environment/environmentHandler.js";
@@ -16,6 +17,7 @@ import {
1617
services,
1718
} from "../services/index.js";
1819
import {
20+
AgentFileServiceState,
1921
AuthServiceState,
2022
ConfigServiceState,
2123
ModelServiceState,
@@ -73,11 +75,13 @@ export async function serve(prompt?: string, options: ServeOptions = {}) {
7375
});
7476

7577
// Get initialized services from the service container
76-
const [configState, modelState, permissionsState] = await Promise.all([
77-
getService<ConfigServiceState>(SERVICE_NAMES.CONFIG),
78-
getService<ModelServiceState>(SERVICE_NAMES.MODEL),
79-
getService<ToolPermissionServiceState>(SERVICE_NAMES.TOOL_PERMISSIONS),
80-
]);
78+
const [configState, modelState, permissionsState, agentFileState] =
79+
await Promise.all([
80+
getService<ConfigServiceState>(SERVICE_NAMES.CONFIG),
81+
getService<ModelServiceState>(SERVICE_NAMES.MODEL),
82+
getService<ToolPermissionServiceState>(SERVICE_NAMES.TOOL_PERMISSIONS),
83+
getService<AgentFileServiceState>(SERVICE_NAMES.AGENT_FILE),
84+
]);
8185

8286
if (!configState.config || !modelState.llmApi || !modelState.model) {
8387
throw new Error("Failed to initialize required services");
@@ -389,9 +393,13 @@ export async function serve(prompt?: string, options: ServeOptions = {}) {
389393
runEnvironmentInstallSafe();
390394

391395
// If initial prompt provided, queue it for processing
392-
if (actualPrompt) {
396+
const initialPrompt = prependPrompt(
397+
agentFileState?.agentFile?.prompt,
398+
actualPrompt,
399+
);
400+
if (initialPrompt) {
393401
console.log(chalk.dim("\nProcessing initial prompt..."));
394-
await messageQueue.enqueueMessage(actualPrompt);
402+
await messageQueue.enqueueMessage(initialPrompt);
395403
processMessages(state, llmApi);
396404
}
397405
});

extensions/cli/src/util/promptProcessor.test.ts

Lines changed: 177 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
22

33
import { processRule } from "../hubLoader.js";
44

5-
import { processAndCombinePrompts } from "./promptProcessor.js";
5+
import { prependPrompt, processAndCombinePrompts } from "./promptProcessor.js";
66

77
// Mock the args module
88
vi.mock("../hubLoader.js", () => ({
@@ -125,3 +125,179 @@ describe("promptProcessor", () => {
125125
});
126126
});
127127
});
128+
129+
describe("prependPrompt", () => {
130+
describe("normal cases", () => {
131+
it("should combine two non-empty strings with double newlines", () => {
132+
const result = prependPrompt("Hello", "World");
133+
expect(result).toBe("Hello\n\nWorld");
134+
});
135+
136+
it("should combine multi-line strings correctly", () => {
137+
const prepend = "Line 1\nLine 2";
138+
const original = "Line 3\nLine 4";
139+
const result = prependPrompt(prepend, original);
140+
expect(result).toBe("Line 1\nLine 2\n\nLine 3\nLine 4");
141+
});
142+
143+
it("should handle strings with special characters", () => {
144+
const prepend = "Special: @#$%^&*()";
145+
const original = "Unicode: 🚀 ∑ ∆";
146+
const result = prependPrompt(prepend, original);
147+
expect(result).toBe("Special: @#$%^&*()\n\nUnicode: 🚀 ∑ ∆");
148+
});
149+
});
150+
151+
describe("undefined parameter cases", () => {
152+
it("should handle undefined prepend with defined original", () => {
153+
const result = prependPrompt(undefined, "original");
154+
expect(result).toBe("original");
155+
});
156+
157+
it("should handle defined prepend with undefined original", () => {
158+
const result = prependPrompt("prepend", undefined);
159+
expect(result).toBe("prepend");
160+
});
161+
162+
it("should return undefined when both parameters are undefined", () => {
163+
const result = prependPrompt(undefined, undefined);
164+
expect(result).toBeUndefined();
165+
});
166+
});
167+
168+
describe("empty string cases", () => {
169+
it("should handle empty prepend with non-empty original", () => {
170+
const result = prependPrompt("", "original");
171+
expect(result).toBe("original");
172+
});
173+
174+
it("should handle non-empty prepend with empty original", () => {
175+
const result = prependPrompt("prepend", "");
176+
expect(result).toBe("prepend");
177+
});
178+
179+
it("should return undefined when both parameters are empty strings", () => {
180+
const result = prependPrompt("", "");
181+
expect(result).toBeUndefined();
182+
});
183+
184+
it("should handle mixed undefined and empty string", () => {
185+
const result1 = prependPrompt(undefined, "");
186+
const result2 = prependPrompt("", undefined);
187+
expect(result1).toBeUndefined();
188+
expect(result2).toBeUndefined();
189+
});
190+
});
191+
192+
describe("whitespace cases", () => {
193+
it("should trim leading and trailing whitespace", () => {
194+
const result = prependPrompt(" hello ", " world ");
195+
expect(result).toBe("hello\n\nworld");
196+
});
197+
198+
it("should return undefined when result is only whitespace", () => {
199+
const result = prependPrompt(" ", " ");
200+
expect(result).toBeUndefined();
201+
});
202+
203+
it("should handle newlines in parameters", () => {
204+
const prepend = "\nhello\n";
205+
const original = "\nworld\n";
206+
const result = prependPrompt(prepend, original);
207+
expect(result).toBe("hello\n\nworld");
208+
});
209+
210+
it("should handle mixed whitespace characters", () => {
211+
const prepend = "\t hello \t";
212+
const original = "\r\n world \r\n";
213+
const result = prependPrompt(prepend, original);
214+
expect(result).toBe("hello\n\nworld");
215+
});
216+
217+
it("should preserve internal whitespace and newlines", () => {
218+
const prepend = "hello\n world";
219+
const original = "foo\t\tbar";
220+
const result = prependPrompt(prepend, original);
221+
expect(result).toBe("hello\n world\n\nfoo\t\tbar");
222+
});
223+
});
224+
225+
describe("edge cases", () => {
226+
it("should handle very long strings", () => {
227+
const longString = "a".repeat(10000);
228+
const result = prependPrompt(longString, "short");
229+
expect(result).toBe(`${longString}\n\nshort`);
230+
expect(result?.length).toBe(10000 + 2 + 5); // long + \n\n + short
231+
});
232+
233+
it("should handle strings with only newlines", () => {
234+
const result = prependPrompt("\n\n\n", "\n\n");
235+
expect(result).toBeUndefined();
236+
});
237+
238+
it("should handle one parameter with content and other with only whitespace", () => {
239+
const result1 = prependPrompt("content", " \n ");
240+
const result2 = prependPrompt(" \t ", "content");
241+
expect(result1).toBe("content");
242+
expect(result2).toBe("content");
243+
});
244+
});
245+
246+
describe("boundary conditions", () => {
247+
it("should return undefined for whitespace-only result after trim", () => {
248+
const cases = [
249+
["", ""],
250+
[" ", ""],
251+
["", " "],
252+
[" ", " "],
253+
["\n", "\n"],
254+
["\t", "\r"],
255+
[undefined, ""],
256+
["", undefined],
257+
[undefined, undefined],
258+
[" \n ", " \n "],
259+
] as const;
260+
261+
cases.forEach(([prepend, original]) => {
262+
const result = prependPrompt(prepend, original);
263+
expect(result).toBeUndefined();
264+
});
265+
});
266+
267+
it("should return trimmed string for non-empty result", () => {
268+
const cases = [
269+
["a", "", "a"],
270+
["", "b", "b"],
271+
["a", "b", "a\n\nb"],
272+
[" a ", "", "a"],
273+
["", " b ", "b"],
274+
[" a ", " b ", "a\n\nb"],
275+
[undefined, "content", "content"],
276+
["content", undefined, "content"],
277+
] as const;
278+
279+
cases.forEach(([prepend, original, expected]) => {
280+
const result = prependPrompt(prepend, original);
281+
expect(result).toBe(expected);
282+
});
283+
});
284+
});
285+
286+
describe("return type consistency", () => {
287+
it("should always return string or undefined", () => {
288+
const testCases = [
289+
["hello", "world"],
290+
[undefined, "world"],
291+
["hello", undefined],
292+
[undefined, undefined],
293+
["", ""],
294+
[" ", " "],
295+
] as const;
296+
297+
testCases.forEach(([prepend, original]) => {
298+
const result = prependPrompt(prepend, original);
299+
expect(typeof result === "string" || result === undefined).toBe(true);
300+
});
301+
});
302+
});
303+
});

extensions/cli/src/util/promptProcessor.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,3 +36,14 @@ export async function processAndCombinePrompts(
3636
? `${combinedPrompts}\n\n${initialPrompt}`
3737
: combinedPrompts;
3838
}
39+
40+
// Merges two prompts with new lines between them, handling undefined
41+
export function prependPrompt(
42+
prepend: string | undefined,
43+
original: string | undefined,
44+
) {
45+
return (
46+
`${(prepend ?? "").trim()}\n\n${(original ?? "").trim()}`.trim() ||
47+
undefined
48+
);
49+
}

0 commit comments

Comments
 (0)