feat: Recompile system prompt after memory subagents finish (#1310)

This commit is contained in:
Devansh Jain
2026-03-09 16:50:40 -07:00
committed by GitHub
parent f5d6f095a6
commit 89d6ed2c87
9 changed files with 547 additions and 52 deletions

View File

@@ -247,6 +247,45 @@ export async function updateConversationLLMConfig(
return client.conversations.update(conversationId, payload); return client.conversations.update(conversationId, payload);
} }
export interface RecompileAgentSystemPromptOptions {
dryRun?: boolean;
updateTimestamp?: boolean;
}
interface AgentSystemPromptRecompileClient {
agents: {
recompile: (
agentId: string,
params: {
dry_run?: boolean;
update_timestamp?: boolean;
},
) => Promise<string>;
};
}
/**
* Recompile an agent's system prompt after memory writes so server-side prompt
* state picks up the latest memory content.
*
* @param agentId - The agent ID to recompile
* @param options - Optional dry-run/timestamp controls
* @param clientOverride - Optional injected client for tests
* @returns The compiled system prompt returned by the API
*/
export async function recompileAgentSystemPrompt(
agentId: string,
options: RecompileAgentSystemPromptOptions = {},
clientOverride?: AgentSystemPromptRecompileClient,
): Promise<string> {
const client = clientOverride ?? (await getClient());
return client.agents.recompile(agentId, {
dry_run: options.dryRun,
update_timestamp: options.updateTimestamp,
});
}
export interface SystemPromptUpdateResult { export interface SystemPromptUpdateResult {
success: boolean; success: boolean;
message: string; message: string;

View File

@@ -232,6 +232,7 @@ import {
type ReflectionSettings, type ReflectionSettings,
reflectionSettingsToLegacyMode, reflectionSettingsToLegacyMode,
} from "./helpers/memoryReminder"; } from "./helpers/memoryReminder";
import { handleMemorySubagentCompletion } from "./helpers/memorySubagentCompletion";
import { import {
type QueuedMessage, type QueuedMessage,
setMessageQueueAdder, setMessageQueueAdder,
@@ -1746,6 +1747,10 @@ export default function App({
const initProgressByAgentRef = useRef( const initProgressByAgentRef = useRef(
new Map<string, { shallowCompleted: boolean; deepFired: boolean }>(), new Map<string, { shallowCompleted: boolean; deepFired: boolean }>(),
); );
const systemPromptRecompileByAgentRef = useRef(
new Map<string, Promise<void>>(),
);
const queuedSystemPromptRecompileByAgentRef = useRef(new Set<string>());
const updateInitProgress = ( const updateInitProgress = (
forAgentId: string, forAgentId: string,
update: Partial<{ shallowCompleted: boolean; deepFired: boolean }>, update: Partial<{ shallowCompleted: boolean; deepFired: boolean }>,
@@ -9305,13 +9310,24 @@ export default function App({
prompt: initPrompt, prompt: initPrompt,
description: "Initializing memory", description: "Initializing memory",
silentCompletion: true, silentCompletion: true,
onComplete: ({ success, error }) => { onComplete: async ({ success, error }) => {
if (success) { const msg = await handleMemorySubagentCompletion(
updateInitProgress(agentId, { deepFired: true }); {
} agentId,
const msg = success subagentType: "init",
? "Built a memory palace of you. Visit it with /palace." initDepth: "deep",
: `Memory initialization failed: ${error || "Unknown error"}`; success,
error,
},
{
recompileByAgent: systemPromptRecompileByAgentRef.current,
recompileQueuedByAgent:
queuedSystemPromptRecompileByAgentRef.current,
updateInitProgress,
logRecompileFailure: (message) =>
debugWarn("memory", message),
},
);
appendTaskNotificationEvents([msg]); appendTaskNotificationEvents([msg]);
}, },
}); });
@@ -9479,15 +9495,29 @@ export default function App({
// attempt (e.g. another /init subagent in flight) preserves the entry for retry. // attempt (e.g. another /init subagent in flight) preserves the entry for retry.
if (autoInitPendingAgentIdsRef.current.has(agentId) && !isSystemOnly) { if (autoInitPendingAgentIdsRef.current.has(agentId) && !isSystemOnly) {
try { try {
const fired = await fireAutoInit(agentId, ({ success, error }) => { const fired = await fireAutoInit(
if (success) { agentId,
updateInitProgress(agentId, { shallowCompleted: true }); async ({ success, error }) => {
} const msg = await handleMemorySubagentCompletion(
const msg = success {
? "Built a memory palace of you. Visit it with /palace." agentId,
: `Memory initialization failed: ${error || "Unknown error"}`; subagentType: "init",
appendTaskNotificationEvents([msg]); initDepth: "shallow",
}); success,
error,
},
{
recompileByAgent: systemPromptRecompileByAgentRef.current,
recompileQueuedByAgent:
queuedSystemPromptRecompileByAgentRef.current,
updateInitProgress,
logRecompileFailure: (message) =>
debugWarn("memory", message),
},
);
appendTaskNotificationEvents([msg]);
},
);
if (fired) { if (fired) {
autoInitPendingAgentIdsRef.current.delete(agentId); autoInitPendingAgentIdsRef.current.delete(agentId);
sharedReminderStateRef.current.pendingAutoInitReminder = true; sharedReminderStateRef.current.pendingAutoInitReminder = true;
@@ -9596,10 +9626,23 @@ ${SYSTEM_REMINDER_CLOSE}
prompt: AUTO_REFLECTION_PROMPT, prompt: AUTO_REFLECTION_PROMPT,
description: AUTO_REFLECTION_DESCRIPTION, description: AUTO_REFLECTION_DESCRIPTION,
silentCompletion: true, silentCompletion: true,
onComplete: ({ success, error }) => { onComplete: async ({ success, error }) => {
const msg = success const msg = await handleMemorySubagentCompletion(
? "Reflected on /palace, the halls remember more now." {
: `Tried to reflect, but got lost in the palace: ${error}`; agentId,
subagentType: "reflection",
success,
error,
},
{
recompileByAgent: systemPromptRecompileByAgentRef.current,
recompileQueuedByAgent:
queuedSystemPromptRecompileByAgentRef.current,
updateInitProgress,
logRecompileFailure: (message) =>
debugWarn("memory", message),
},
);
appendTaskNotificationEvents([msg]); appendTaskNotificationEvents([msg]);
}, },
}); });
@@ -9638,13 +9681,24 @@ ${SYSTEM_REMINDER_CLOSE}
prompt: initPrompt, prompt: initPrompt,
description: "Deep memory initialization", description: "Deep memory initialization",
silentCompletion: true, silentCompletion: true,
onComplete: ({ success, error }) => { onComplete: async ({ success, error }) => {
if (success) { const msg = await handleMemorySubagentCompletion(
updateInitProgress(agentId, { deepFired: true }); {
} agentId,
const msg = success subagentType: "init",
? "Built a memory palace of you. Visit it with /palace." initDepth: "deep",
: `Deep memory initialization failed: ${error || "Unknown error"}`; success,
error,
},
{
recompileByAgent: systemPromptRecompileByAgentRef.current,
recompileQueuedByAgent:
queuedSystemPromptRecompileByAgentRef.current,
updateInitProgress,
logRecompileFailure: (message) =>
debugWarn("memory", message),
},
);
appendTaskNotificationEvents([msg]); appendTaskNotificationEvents([msg]);
}, },
}); });

View File

@@ -139,7 +139,10 @@ Instructions:
*/ */
export async function fireAutoInit( export async function fireAutoInit(
agentId: string, agentId: string,
onComplete: (result: { success: boolean; error?: string }) => void, onComplete: (result: {
success: boolean;
error?: string;
}) => void | Promise<void>,
): Promise<boolean> { ): Promise<boolean> {
if (hasActiveInitSubagent()) return false; if (hasActiveInitSubagent()) return false;
if (!settingsManager.isMemfsEnabled(agentId)) return false; if (!settingsManager.isMemfsEnabled(agentId)) return false;

View File

@@ -0,0 +1,123 @@
import {
type RecompileAgentSystemPromptOptions,
recompileAgentSystemPrompt,
} from "../../agent/modify";
export type MemorySubagentType = "init" | "reflection";
export type MemoryInitDepth = "shallow" | "deep";
export interface MemoryInitProgressUpdate {
shallowCompleted: boolean;
deepFired: boolean;
}
type RecompileAgentSystemPromptFn = (
agentId: string,
options?: RecompileAgentSystemPromptOptions,
) => Promise<string>;
export type MemorySubagentCompletionArgs =
| {
agentId: string;
subagentType: "init";
initDepth: MemoryInitDepth;
success: boolean;
error?: string;
}
| {
agentId: string;
subagentType: "reflection";
initDepth?: never;
success: boolean;
error?: string;
};
export interface MemorySubagentCompletionDeps {
recompileByAgent: Map<string, Promise<void>>;
recompileQueuedByAgent: Set<string>;
updateInitProgress: (
agentId: string,
update: Partial<MemoryInitProgressUpdate>,
) => void;
logRecompileFailure?: (message: string) => void;
recompileAgentSystemPromptImpl?: RecompileAgentSystemPromptFn;
}
/**
* Finalize a memory-writing subagent by updating init progress, recompiling the
* parent agent's system prompt, and returning the user-facing completion text.
*/
export async function handleMemorySubagentCompletion(
args: MemorySubagentCompletionArgs,
deps: MemorySubagentCompletionDeps,
): Promise<string> {
const { agentId, subagentType, initDepth, success, error } = args;
const recompileAgentSystemPromptFn =
deps.recompileAgentSystemPromptImpl ?? recompileAgentSystemPrompt;
let recompileError: string | null = null;
if (success) {
if (subagentType === "init") {
deps.updateInitProgress(
agentId,
initDepth === "shallow"
? { shallowCompleted: true }
: { deepFired: true },
);
}
try {
let inFlight = deps.recompileByAgent.get(agentId);
if (!inFlight) {
inFlight = (async () => {
do {
deps.recompileQueuedByAgent.delete(agentId);
await recompileAgentSystemPromptFn(agentId, {
updateTimestamp: true,
});
} while (deps.recompileQueuedByAgent.has(agentId));
})().finally(() => {
// Cleanup runs only after the shared promise settles, so every
// concurrent caller awaits the same full recompile lifecycle.
deps.recompileQueuedByAgent.delete(agentId);
deps.recompileByAgent.delete(agentId);
});
deps.recompileByAgent.set(agentId, inFlight);
} else {
deps.recompileQueuedByAgent.add(agentId);
}
await inFlight;
} catch (recompileFailure) {
recompileError =
recompileFailure instanceof Error
? recompileFailure.message
: String(recompileFailure);
deps.logRecompileFailure?.(
`Failed to recompile system prompt after ${subagentType} subagent for ${agentId}: ${recompileError}`,
);
}
}
if (!success) {
const normalizedError = error || "Unknown error";
if (subagentType === "reflection") {
return `Tried to reflect, but got lost in the palace: ${normalizedError}`;
}
return initDepth === "deep"
? `Deep memory initialization failed: ${normalizedError}`
: `Memory initialization failed: ${normalizedError}`;
}
const baseMessage =
subagentType === "reflection"
? "Reflected on /palace, the halls remember more now."
: "Built a memory palace of you. Visit it with /palace.";
if (!recompileError) {
return baseMessage;
}
return `${baseMessage} System prompt recompilation failed: ${recompileError}`;
}

View File

@@ -0,0 +1,31 @@
import { describe, expect, mock, test } from "bun:test";
import { recompileAgentSystemPrompt } from "../../agent/modify";
describe("recompileAgentSystemPrompt", () => {
test("calls the Letta agent recompile endpoint with mapped params", async () => {
const agentsRecompileMock = mock(
(_agentId: string, _params?: Record<string, unknown>) =>
Promise.resolve("compiled-system-prompt"),
);
const client = {
agents: {
recompile: agentsRecompileMock,
},
};
const compiledPrompt = await recompileAgentSystemPrompt(
"agent-123",
{
updateTimestamp: true,
dryRun: true,
},
client,
);
expect(compiledPrompt).toBe("compiled-system-prompt");
expect(agentsRecompileMock).toHaveBeenCalledWith("agent-123", {
dry_run: true,
update_timestamp: true,
});
});
});

View File

@@ -0,0 +1,140 @@
import { beforeEach, describe, expect, mock, test } from "bun:test";
import type { RecompileAgentSystemPromptOptions } from "../../agent/modify";
import { handleMemorySubagentCompletion } from "../../cli/helpers/memorySubagentCompletion";
const recompileAgentSystemPromptMock = mock(
(_agentId: string, _opts?: RecompileAgentSystemPromptOptions) =>
Promise.resolve("compiled-system-prompt"),
);
function createDeferred<T>() {
let resolve!: (value: T | PromiseLike<T>) => void;
const promise = new Promise<T>((res) => {
resolve = res;
});
return { promise, resolve };
}
describe("memory subagent recompile handling", () => {
beforeEach(() => {
recompileAgentSystemPromptMock.mockReset();
recompileAgentSystemPromptMock.mockImplementation(
(_agentId: string, _opts?: RecompileAgentSystemPromptOptions) =>
Promise.resolve("compiled-system-prompt"),
);
});
test("updates init progress and recompiles after successful shallow init", async () => {
const progressUpdates: Array<{
agentId: string;
update: Record<string, boolean>;
}> = [];
const message = await handleMemorySubagentCompletion(
{
agentId: "agent-init-1",
subagentType: "init",
initDepth: "shallow",
success: true,
},
{
recompileByAgent: new Map(),
recompileQueuedByAgent: new Set(),
recompileAgentSystemPromptImpl: recompileAgentSystemPromptMock,
updateInitProgress: (agentId, update) => {
progressUpdates.push({
agentId,
update: update as Record<string, boolean>,
});
},
},
);
expect(message).toBe(
"Built a memory palace of you. Visit it with /palace.",
);
expect(progressUpdates).toEqual([
{
agentId: "agent-init-1",
update: { shallowCompleted: true },
},
]);
expect(recompileAgentSystemPromptMock).toHaveBeenCalledWith(
"agent-init-1",
{
updateTimestamp: true,
},
);
});
test("queues a trailing recompile when later completions land mid-flight", async () => {
const firstDeferred = createDeferred<string>();
const secondDeferred = createDeferred<string>();
recompileAgentSystemPromptMock
.mockImplementationOnce(() => firstDeferred.promise)
.mockImplementationOnce(() => secondDeferred.promise);
const recompileByAgent = new Map<string, Promise<void>>();
const recompileQueuedByAgent = new Set<string>();
const deps = {
recompileByAgent,
recompileQueuedByAgent,
recompileAgentSystemPromptImpl: recompileAgentSystemPromptMock,
updateInitProgress: () => {},
};
const first = handleMemorySubagentCompletion(
{
agentId: "agent-shared",
subagentType: "reflection",
success: true,
},
deps,
);
const second = handleMemorySubagentCompletion(
{
agentId: "agent-shared",
subagentType: "reflection",
success: true,
},
deps,
);
const third = handleMemorySubagentCompletion(
{
agentId: "agent-shared",
subagentType: "reflection",
success: true,
},
deps,
);
expect(recompileAgentSystemPromptMock).toHaveBeenCalledTimes(1);
expect(recompileByAgent.has("agent-shared")).toBe(true);
expect(recompileQueuedByAgent.has("agent-shared")).toBe(true);
firstDeferred.resolve("compiled-system-prompt");
await Promise.resolve();
expect(recompileAgentSystemPromptMock).toHaveBeenCalledTimes(2);
expect(recompileByAgent.has("agent-shared")).toBe(true);
secondDeferred.resolve("compiled-system-prompt");
const [firstMessage, secondMessage, thirdMessage] = await Promise.all([
first,
second,
third,
]);
expect(firstMessage).toBe(
"Reflected on /palace, the halls remember more now.",
);
expect(secondMessage).toBe(
"Reflected on /palace, the halls remember more now.",
);
expect(thirdMessage).toBe(
"Reflected on /palace, the halls remember more now.",
);
expect(recompileByAgent.size).toBe(0);
expect(recompileQueuedByAgent.size).toBe(0);
});
});

View File

@@ -143,17 +143,9 @@ describe("background onComplete → flush wiring in App.tsx", () => {
const onCompleteIdx = source.indexOf("onComplete:", initBlock); const onCompleteIdx = source.indexOf("onComplete:", initBlock);
expect(onCompleteIdx).toBeGreaterThan(-1); expect(onCompleteIdx).toBeGreaterThan(-1);
// The appendTaskNotificationEvents call must be within the onComplete body const onCompleteWindow = source.slice(onCompleteIdx, onCompleteIdx + 900);
// (before the closing of spawnBackgroundSubagentTask). expect(onCompleteWindow).toContain("await handleMemorySubagentCompletion(");
// Use "});" with leading whitespace to match the spawn call's closing, expect(onCompleteWindow).toContain("appendTaskNotificationEvents(");
// not inner statements like updateInitProgress(...});
const callIdx = source.indexOf(
"appendTaskNotificationEvents(",
onCompleteIdx,
);
const blockEnd = source.indexOf(" });", onCompleteIdx);
expect(callIdx).toBeGreaterThan(onCompleteIdx);
expect(callIdx).toBeLessThan(blockEnd);
}); });
test("reflection onComplete calls appendTaskNotificationEvents", () => { test("reflection onComplete calls appendTaskNotificationEvents", () => {
@@ -165,12 +157,8 @@ describe("background onComplete → flush wiring in App.tsx", () => {
const onCompleteIdx = source.indexOf("onComplete:", reflectionBlock); const onCompleteIdx = source.indexOf("onComplete:", reflectionBlock);
expect(onCompleteIdx).toBeGreaterThan(-1); expect(onCompleteIdx).toBeGreaterThan(-1);
const callIdx = source.indexOf( const onCompleteWindow = source.slice(onCompleteIdx, onCompleteIdx + 700);
"appendTaskNotificationEvents(", expect(onCompleteWindow).toContain("await handleMemorySubagentCompletion(");
onCompleteIdx, expect(onCompleteWindow).toContain("appendTaskNotificationEvents(");
);
const blockEnd = source.indexOf("});", onCompleteIdx);
expect(callIdx).toBeGreaterThan(onCompleteIdx);
expect(callIdx).toBeLessThan(blockEnd);
}); });
}); });

View File

@@ -181,6 +181,99 @@ describe("spawnBackgroundSubagentTask", () => {
expect(runSubagentStopHooksImpl).toHaveBeenCalledTimes(1); expect(runSubagentStopHooksImpl).toHaveBeenCalledTimes(1);
}); });
test("awaits async onComplete before queue notification and hooks", async () => {
const callOrder: string[] = [];
const spawnSubagentImpl = mock(async () => ({
agentId: "agent-ordered",
conversationId: "default",
report: "reflection done",
success: true,
totalTokens: 22,
}));
const addToMessageQueueOrdered = mock(
(_msg: { kind: "user" | "task_notification"; text: string }) => {
callOrder.push("queue");
},
);
const runSubagentStopHooksOrdered = mock(async () => {
callOrder.push("hooks");
return {
blocked: false,
errored: false,
feedback: [],
results: [],
};
});
const onComplete = mock(async () => {
callOrder.push("onComplete:start");
await new Promise((resolve) => setTimeout(resolve, 20));
callOrder.push("onComplete:end");
});
spawnBackgroundSubagentTask({
subagentType: "reflection",
prompt: "Reflect",
description: "Reflect on memory",
onComplete,
deps: {
spawnSubagentImpl,
addToMessageQueueImpl: addToMessageQueueOrdered,
formatTaskNotificationImpl,
runSubagentStopHooksImpl: runSubagentStopHooksOrdered,
generateSubagentIdImpl,
registerSubagentImpl,
completeSubagentImpl,
getSubagentSnapshotImpl,
},
});
await new Promise((resolve) => setTimeout(resolve, 60));
expect(callOrder).toEqual([
"onComplete:start",
"onComplete:end",
"queue",
"hooks",
]);
});
test("continues queue notification and hooks when onComplete throws", async () => {
const spawnSubagentImpl = mock(async () => ({
agentId: "agent-oncomplete-error",
conversationId: "default",
report: "reflection done",
success: true,
totalTokens: 19,
}));
const onComplete = mock(async () => {
throw new Error("callback exploded");
});
const launched = spawnBackgroundSubagentTask({
subagentType: "reflection",
prompt: "Reflect",
description: "Reflect on memory",
onComplete,
deps: {
spawnSubagentImpl,
addToMessageQueueImpl,
formatTaskNotificationImpl,
runSubagentStopHooksImpl,
generateSubagentIdImpl,
registerSubagentImpl,
completeSubagentImpl,
getSubagentSnapshotImpl,
},
});
await new Promise((resolve) => setTimeout(resolve, 20));
expect(queueMessages).toHaveLength(1);
expect(runSubagentStopHooksImpl).toHaveBeenCalledTimes(1);
const outputContent = readFileSync(launched.outputFile, "utf-8");
expect(outputContent).toContain("[onComplete error] callback exploded");
});
test("marks background task failed and emits notification on error", async () => { test("marks background task failed and emits notification on error", async () => {
const spawnSubagentImpl = mock(async () => { const spawnSubagentImpl = mock(async () => {
throw new Error("subagent exploded"); throw new Error("subagent exploded");

View File

@@ -75,9 +75,13 @@ export interface SpawnBackgroundSubagentTaskArgs {
silentCompletion?: boolean; silentCompletion?: boolean;
/** /**
* Called after the subagent finishes (success or failure). * Called after the subagent finishes (success or failure).
* Runs regardless of `silentCompletion`. * Runs regardless of `silentCompletion` and is awaited before
* completion notifications/hooks continue.
*/ */
onComplete?: (result: { success: boolean; error?: string }) => void; onComplete?: (result: {
success: boolean;
error?: string;
}) => void | Promise<void>;
/** /**
* Optional dependency overrides for tests. * Optional dependency overrides for tests.
* Production callers should not provide this. * Production callers should not provide this.
@@ -248,6 +252,9 @@ export function spawnBackgroundSubagentTask(
backgroundTasks.set(taskId, bgTask); backgroundTasks.set(taskId, bgTask);
writeTaskTranscriptStart(outputFile, description, subagentType); writeTaskTranscriptStart(outputFile, description, subagentType);
// Intentionally fire-and-forget: background tasks own their lifecycle and
// capture failures in task state/transcripts instead of surfacing a promise
// back to the caller.
spawnSubagentFn( spawnSubagentFn(
subagentType, subagentType,
prompt, prompt,
@@ -258,7 +265,7 @@ export function spawnBackgroundSubagentTask(
existingConversationId, existingConversationId,
maxTurns, maxTurns,
) )
.then((result) => { .then(async (result) => {
bgTask.status = result.success ? "completed" : "failed"; bgTask.status = result.success ? "completed" : "failed";
if (result.error) { if (result.error) {
bgTask.error = result.error; bgTask.error = result.error;
@@ -276,7 +283,13 @@ export function spawnBackgroundSubagentTask(
totalTokens: result.totalTokens, totalTokens: result.totalTokens,
}); });
onComplete?.({ success: result.success, error: result.error }); try {
await onComplete?.({ success: result.success, error: result.error });
} catch (error) {
const errorMessage =
error instanceof Error ? error.message : String(error);
appendToOutputFile(outputFile, `[onComplete error] ${errorMessage}\n`);
}
if (!silentCompletion) { if (!silentCompletion) {
const subagentSnapshot = getSubagentSnapshotFn(); const subagentSnapshot = getSubagentSnapshotFn();
@@ -325,7 +338,7 @@ export function spawnBackgroundSubagentTask(
// Silently ignore hook errors // Silently ignore hook errors
}); });
}) })
.catch((error) => { .catch(async (error) => {
const errorMessage = const errorMessage =
error instanceof Error ? error.message : String(error); error instanceof Error ? error.message : String(error);
bgTask.status = "failed"; bgTask.status = "failed";
@@ -333,7 +346,18 @@ export function spawnBackgroundSubagentTask(
appendToOutputFile(outputFile, `[error] ${errorMessage}\n`); appendToOutputFile(outputFile, `[error] ${errorMessage}\n`);
completeSubagentFn(subagentId, { success: false, error: errorMessage }); completeSubagentFn(subagentId, { success: false, error: errorMessage });
onComplete?.({ success: false, error: errorMessage }); try {
await onComplete?.({ success: false, error: errorMessage });
} catch (onCompleteError) {
const callbackMessage =
onCompleteError instanceof Error
? onCompleteError.message
: String(onCompleteError);
appendToOutputFile(
outputFile,
`[onComplete error] ${callbackMessage}\n`,
);
}
if (!silentCompletion) { if (!silentCompletion) {
const subagentSnapshot = getSubagentSnapshotFn(); const subagentSnapshot = getSubagentSnapshotFn();