Fimeg/letta-code
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

39
src/agent/modify.ts
View File

@@ -247,6 +247,45 @@ export async function updateConversationLLMConfig(
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 {
success: boolean;
message: string;

108
src/cli/App.tsx
View File

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

5
src/cli/helpers/initCommand.ts
View File

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

123
src/cli/helpers/memorySubagentCompletion.ts Normal file
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}`;
}

31
src/tests/agent/recompile-system-prompt.test.ts Normal file
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,
});
});
});

140
src/tests/cli/memory-subagent-recompile-wiring.test.ts Normal file
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);
});
});

24
src/tests/cli/task-notification-flush.test.ts
View File

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

93
src/tests/tools/task-background-helper.test.ts
View File

@@ -181,6 +181,99 @@ describe("spawnBackgroundSubagentTask", () => {
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 () => {
const spawnSubagentImpl = mock(async () => {
throw new Error("subagent exploded");

36
src/tools/impl/Task.ts
View File

@@ -75,9 +75,13 @@ export interface SpawnBackgroundSubagentTaskArgs {
silentCompletion?: boolean;
/**
* 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.
* Production callers should not provide this.
@@ -248,6 +252,9 @@ export function spawnBackgroundSubagentTask(
backgroundTasks.set(taskId, bgTask);
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(
subagentType,
prompt,
@@ -258,7 +265,7 @@ export function spawnBackgroundSubagentTask(
existingConversationId,
maxTurns,
)
.then((result) => {
.then(async (result) => {
bgTask.status = result.success ? "completed" : "failed";
if (result.error) {
bgTask.error = result.error;
@@ -276,7 +283,13 @@ export function spawnBackgroundSubagentTask(
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) {
const subagentSnapshot = getSubagentSnapshotFn();
@@ -325,7 +338,7 @@ export function spawnBackgroundSubagentTask(
// Silently ignore hook errors
});
})
.catch((error) => {
.catch(async (error) => {
const errorMessage =
error instanceof Error ? error.message : String(error);
bgTask.status = "failed";
@@ -333,7 +346,18 @@ export function spawnBackgroundSubagentTask(
appendToOutputFile(outputFile, `[error] ${errorMessage}\n`);
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) {
const subagentSnapshot = getSubagentSnapshotFn();