Skip to content

Commit e296862

Browse files
RomneyDaRomneyDa
committed
fix: restore model prompt transform extraction
1 parent d76d23d commit e296862

3 files changed

Lines changed: 206 additions & 1 deletion

File tree

‎src/agents/embedded-agent-runner/run/attempt.llm-boundary.test.ts‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { describe, expect, it } from "vitest";
22
import {
3+
installModelPromptTransform,
34
insertRuntimeContextMessageForPrompt,
45
normalizeMessagesForLlmBoundary,
56
} from "./attempt.llm-boundary.js";
@@ -330,4 +331,72 @@ describe("normalizeMessagesForLlmBoundary", () => {
330331
expect(JSON.stringify(output)).not.toContain("matched secret prompt");
331332
expect(input[0]).toHaveProperty("__openclaw");
332333
});
334+
335+
it("replaces only the armed prompt with model prompt context", async () => {
336+
const messages = [
337+
{
338+
role: "user",
339+
content: [{ type: "text", text: "visible transcript prompt" }],
340+
timestamp: 1,
341+
},
342+
] as Parameters<typeof normalizeMessagesForLlmBoundary>[0];
343+
const captured: (typeof messages)[] = [];
344+
const session = {
345+
agent: {
346+
transformContext: async (nextMessages: typeof messages) => {
347+
captured.push(nextMessages);
348+
return nextMessages;
349+
},
350+
},
351+
};
352+
let armed = false;
353+
const cleanup = installModelPromptTransform({
354+
session,
355+
transcriptPrompt: "visible transcript prompt",
356+
modelPrompt: "private model prompt",
357+
prependContext: "before",
358+
appendContext: "after",
359+
shouldCapturePrompt: () => armed,
360+
});
361+
362+
const unarmed = await session.agent.transformContext(messages);
363+
armed = true;
364+
const armedResult = await session.agent.transformContext(messages);
365+
cleanup();
366+
const unarmedRecords = unarmed as Array<{ content?: unknown }>;
367+
const armedRecords = armedResult as Array<{ content?: unknown }>;
368+
369+
expect(unarmedRecords[0]?.content).toEqual([
370+
{ type: "text", text: "visible transcript prompt" },
371+
]);
372+
expect(armedRecords[0]?.content).toEqual([{ type: "text", text: "private model prompt" }]);
373+
expect(armedResult[0]).toHaveProperty(
374+
"__openclawTranscriptPromptText",
375+
"visible transcript prompt",
376+
);
377+
expect(captured).toHaveLength(2);
378+
expect(session.agent.transformContext).not.toBeUndefined();
379+
});
380+
381+
it("restores the original model prompt transform on cleanup", async () => {
382+
const originalTransform = async (
383+
messages: Parameters<typeof normalizeMessagesForLlmBoundary>[0],
384+
) => messages;
385+
const session = {
386+
agent: {
387+
transformContext: originalTransform,
388+
},
389+
};
390+
const cleanup = installModelPromptTransform({
391+
session,
392+
transcriptPrompt: "visible transcript prompt",
393+
prependContext: "before",
394+
shouldCapturePrompt: () => true,
395+
});
396+
397+
expect(session.agent.transformContext).not.toBe(originalTransform);
398+
cleanup();
399+
400+
expect(session.agent.transformContext).toBe(originalTransform);
401+
});
333402
});

‎src/agents/embedded-agent-runner/run/attempt.llm-boundary.ts‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { stripHistoricalRuntimeContextCustomMessages } from "../../internal-runt
33
import type { AgentMessage } from "../../runtime/index.js";
44
import { stripToolResultDetails } from "../../session-transcript-repair.js";
55
import { normalizeAssistantReplayContent } from "../replay-history.js";
6+
import { markTranscriptPromptText } from "../tool-result-context-guard.js";
67
import type { RuntimeContextCustomMessage } from "./runtime-context-prompt.js";
78

89
export function normalizeMessagesForLlmBoundary(messages: AgentMessage[]): AgentMessage[] {
@@ -103,6 +104,141 @@ export function insertRuntimeContextMessageForPrompt(params: {
103104
];
104105
}
105106

107+
function replaceLastUserTextPrompt(params: {
108+
messages: AgentMessage[];
109+
shouldCapture?: (message: AgentMessage) => boolean;
110+
transcriptText?: string;
111+
replace: (text: string) => string | undefined;
112+
}): AgentMessage[] {
113+
const userIndex = params.messages.findLastIndex((message) => message.role === "user");
114+
if (userIndex === -1) {
115+
return params.messages;
116+
}
117+
const message = params.messages[userIndex];
118+
if (!message || message.role !== "user") {
119+
return params.messages;
120+
}
121+
if (params.shouldCapture && !params.shouldCapture(message)) {
122+
return params.messages;
123+
}
124+
const content = (message as { content?: unknown }).content;
125+
if (typeof content === "string") {
126+
const replacement = params.replace(content);
127+
if (replacement === undefined) {
128+
return params.messages;
129+
}
130+
const next = params.messages.slice();
131+
next[userIndex] = { ...message, content: replacement } as AgentMessage;
132+
if (params.transcriptText !== undefined) {
133+
markTranscriptPromptText(next[userIndex], params.transcriptText);
134+
}
135+
return next;
136+
}
137+
if (!Array.isArray(content)) {
138+
return params.messages;
139+
}
140+
let replaced = false;
141+
const nextContent = content.map((block) => {
142+
if (replaced || !block || typeof block !== "object") {
143+
return block;
144+
}
145+
const textBlock = block as { type?: unknown; text?: unknown };
146+
if (textBlock.type !== "text" || typeof textBlock.text !== "string") {
147+
return block;
148+
}
149+
const replacement = params.replace(textBlock.text);
150+
if (replacement === undefined) {
151+
return block;
152+
}
153+
replaced = true;
154+
return Object.assign({}, block, { text: replacement });
155+
});
156+
if (!replaced) {
157+
return params.messages;
158+
}
159+
const next = params.messages.slice();
160+
next[userIndex] = { ...message, content: nextContent } as AgentMessage;
161+
if (params.transcriptText !== undefined) {
162+
markTranscriptPromptText(next[userIndex], params.transcriptText);
163+
}
164+
return next;
165+
}
166+
167+
function composeModelPromptContext(params: {
168+
prompt: string;
169+
prependContext?: string;
170+
appendContext?: string;
171+
}): string {
172+
return [params.prependContext, params.prompt, params.appendContext]
173+
.filter((value): value is string => Boolean(value?.trim()))
174+
.join("\n\n");
175+
}
176+
177+
export function installModelPromptTransform(params: {
178+
session: {
179+
agent: {
180+
transformContext?: (
181+
messages: AgentMessage[],
182+
signal?: AbortSignal,
183+
) => Promise<AgentMessage[]>;
184+
};
185+
};
186+
transcriptPrompt: string;
187+
modelPrompt?: string;
188+
prependContext?: string;
189+
appendContext?: string;
190+
shouldCapturePrompt: () => boolean;
191+
}): () => void {
192+
const modelPrompt = params.modelPrompt;
193+
const hasPromptContext =
194+
Boolean(params.prependContext?.trim()) || Boolean(params.appendContext?.trim());
195+
if ((!modelPrompt?.trim() || modelPrompt === params.transcriptPrompt) && !hasPromptContext) {
196+
return () => undefined;
197+
}
198+
const agent = params.session.agent;
199+
const originalTransformContext = agent.transformContext;
200+
let targetPromptTimestamp: number | undefined;
201+
agent.transformContext = async (messages, signal) => {
202+
const promptMessages = replaceLastUserTextPrompt({
203+
messages,
204+
transcriptText: params.transcriptPrompt,
205+
shouldCapture: (message) => {
206+
const timestamp = (message as { timestamp?: unknown }).timestamp;
207+
if (targetPromptTimestamp !== undefined) {
208+
return timestamp === targetPromptTimestamp;
209+
}
210+
if (!params.shouldCapturePrompt()) {
211+
return false;
212+
}
213+
if (typeof timestamp === "number") {
214+
targetPromptTimestamp = timestamp;
215+
}
216+
return true;
217+
},
218+
replace: (text) => {
219+
if (modelPrompt?.trim() && text === params.transcriptPrompt) {
220+
return modelPrompt;
221+
}
222+
if (!hasPromptContext) {
223+
return undefined;
224+
}
225+
const replacement = composeModelPromptContext({
226+
prompt: text,
227+
prependContext: params.prependContext,
228+
appendContext: params.appendContext,
229+
});
230+
return replacement === text ? undefined : replacement;
231+
},
232+
});
233+
return originalTransformContext
234+
? await originalTransformContext.call(agent, promptMessages, signal)
235+
: promptMessages;
236+
};
237+
return () => {
238+
agent.transformContext = originalTransformContext;
239+
};
240+
}
241+
106242
function stripHistoricalInboundMetadataFromUserMessages(messages: AgentMessage[]): AgentMessage[] {
107243
const activeUserMessageIndex = findActiveUserMessageIndex(messages);
108244
let changed = false;

‎src/agents/embedded-agent-runner/run/attempt.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -262,7 +262,6 @@ import {
262262
import {
263263
installContextEngineLoopHook,
264264
installToolResultContextGuard,
265-
markTranscriptPromptText,
266265
} from "../tool-result-context-guard.js";
267266
import {
268267
resolveLiveToolResultMaxChars,
@@ -312,6 +311,7 @@ import {
312311
runAttemptContextEngineBootstrap,
313312
} from "./attempt.context-engine-helpers.js";
314313
import {
314+
installModelPromptTransform,
315315
installRuntimeContextMessageForPrompt,
316316
normalizeMessagesForCurrentPromptBoundary,
317317
normalizeMessagesForLlmBoundary,

0 commit comments

Comments
 (0)