diff --git a/src/agent/modify.ts b/src/agent/modify.ts index 98794cc..f6b302b 100644 --- a/src/agent/modify.ts +++ b/src/agent/modify.ts @@ -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; + }; +} + +/** + * 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 { + 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; diff --git a/src/cli/App.tsx b/src/cli/App.tsx index 5efb444..4cdd308 100644 --- a/src/cli/App.tsx +++ b/src/cli/App.tsx @@ -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(), ); + const systemPromptRecompileByAgentRef = useRef( + new Map>(), + ); + const queuedSystemPromptRecompileByAgentRef = useRef(new Set()); 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]); }, }); diff --git a/src/cli/helpers/initCommand.ts b/src/cli/helpers/initCommand.ts index 64064e4..b840d35 100644 --- a/src/cli/helpers/initCommand.ts +++ b/src/cli/helpers/initCommand.ts @@ -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, ): Promise { if (hasActiveInitSubagent()) return false; if (!settingsManager.isMemfsEnabled(agentId)) return false; diff --git a/src/cli/helpers/memorySubagentCompletion.ts b/src/cli/helpers/memorySubagentCompletion.ts new file mode 100644 index 0000000..0ef282d --- /dev/null +++ b/src/cli/helpers/memorySubagentCompletion.ts @@ -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; + +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>; + recompileQueuedByAgent: Set; + updateInitProgress: ( + agentId: string, + update: Partial, + ) => 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 { + 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}`; +} diff --git a/src/tests/agent/recompile-system-prompt.test.ts b/src/tests/agent/recompile-system-prompt.test.ts new file mode 100644 index 0000000..c7945eb --- /dev/null +++ b/src/tests/agent/recompile-system-prompt.test.ts @@ -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) => + 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, + }); + }); +}); diff --git a/src/tests/cli/memory-subagent-recompile-wiring.test.ts b/src/tests/cli/memory-subagent-recompile-wiring.test.ts new file mode 100644 index 0000000..07dbbf7 --- /dev/null +++ b/src/tests/cli/memory-subagent-recompile-wiring.test.ts @@ -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() { + let resolve!: (value: T | PromiseLike) => void; + const promise = new Promise((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; + }> = []; + + 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, + }); + }, + }, + ); + + 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(); + const secondDeferred = createDeferred(); + recompileAgentSystemPromptMock + .mockImplementationOnce(() => firstDeferred.promise) + .mockImplementationOnce(() => secondDeferred.promise); + + const recompileByAgent = new Map>(); + const recompileQueuedByAgent = new Set(); + 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); + }); +}); diff --git a/src/tests/cli/task-notification-flush.test.ts b/src/tests/cli/task-notification-flush.test.ts index 3b6a403..4208756 100644 --- a/src/tests/cli/task-notification-flush.test.ts +++ b/src/tests/cli/task-notification-flush.test.ts @@ -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("); }); }); diff --git a/src/tests/tools/task-background-helper.test.ts b/src/tests/tools/task-background-helper.test.ts index 325cbb4..c15e73c 100644 --- a/src/tests/tools/task-background-helper.test.ts +++ b/src/tests/tools/task-background-helper.test.ts @@ -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"); diff --git a/src/tools/impl/Task.ts b/src/tools/impl/Task.ts index 8b239be..772a69b 100644 --- a/src/tools/impl/Task.ts +++ b/src/tools/impl/Task.ts @@ -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; /** * 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();