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

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<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,
+    });
+  });
+});

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<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);
+  });
+});

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(");
   });
 });

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");