fix: persist orphaned user leaf repair

serkonyc · serkonyc · commit e42bceefd658 · 2026-05-04T01:21:40.000+03:00
diff --git a/extensions/discord/src/monitor/message-handler.process.ts b/extensions/discord/src/monitor/message-handler.process.ts
@@ -645,10 +645,10 @@ export async function processDiscordMessage(
                   ? (payload) => draftPreview.updateFromPartial(payload.text)
                   : undefined,
                 onAssistantMessageStart: draftPreview.draftStream
-                  ? draftPreview.handleAssistantMessageBoundary
+                  ? () => draftPreview.handleAssistantMessageBoundary()
                   : undefined,
                 onReasoningEnd: draftPreview.draftStream
-                  ? draftPreview.handleAssistantMessageBoundary
+                  ? () => draftPreview.handleAssistantMessageBoundary()
                   : undefined,
                 onModelSelected,
                 suppressDefaultToolProgressMessages:
diff --git a/src/agents/pi-embedded-runner/run/attempt.test.ts b/src/agents/pi-embedded-runner/run/attempt.test.ts
@@ -14,6 +14,7 @@ import {
   isPrimaryBootstrapRun,
   mergeOrphanedTrailingUserPrompt,
   normalizeMessagesForLlmBoundary,
+  removeSessionManagerLeafEntry,
   prependSystemPromptAddition,
   remapInjectedContextFilesToWorkspace,
   resetEmbeddedAgentBaseStreamFnCacheForTest,
@@ -479,6 +480,54 @@ describe("remapInjectedContextFilesToWorkspace", () => {
   });
 });
 
+describe("removeSessionManagerLeafEntry", () => {
+  it("removes the persisted leaf entry before rebuilding context", () => {
+    const orphan = { type: "message", id: "orphan-user", parentId: "assistant-1" };
+    const fileEntries = [
+      { type: "session", id: "session-1" },
+      { type: "message", id: "user-1", parentId: null },
+      { type: "message", id: "assistant-1", parentId: "user-1" },
+      orphan,
+    ];
+    const manager = {
+      fileEntries,
+      byId: new Map(fileEntries.map((entry) => [entry.id, entry])),
+      leafId: "orphan-user" as string | null,
+      branch: vi.fn(),
+      resetLeaf: vi.fn(),
+      _rewriteFile: vi.fn(),
+    };
+
+    removeSessionManagerLeafEntry({ sessionManager: manager, leafEntry: orphan });
+
+    expect(manager.fileEntries.map((entry) => entry.id)).toEqual([
+      "session-1",
+      "user-1",
+      "assistant-1",
+    ]);
+    expect(manager.byId.has("orphan-user")).toBe(false);
+    expect(manager.leafId).toBe("assistant-1");
+    expect(manager._rewriteFile).toHaveBeenCalledTimes(1);
+    expect(manager.branch).not.toHaveBeenCalled();
+    expect(manager.resetLeaf).not.toHaveBeenCalled();
+  });
+
+  it("falls back to in-memory branching when the manager cannot rewrite", () => {
+    const manager = {
+      branch: vi.fn(),
+      resetLeaf: vi.fn(),
+    };
+
+    removeSessionManagerLeafEntry({
+      sessionManager: manager,
+      leafEntry: { id: "orphan-user", parentId: "assistant-1" },
+    });
+
+    expect(manager.branch).toHaveBeenCalledWith("assistant-1");
+    expect(manager.resetLeaf).not.toHaveBeenCalled();
+  });
+});
+
 describe("shouldWarnOnOrphanedUserRepair", () => {
   it("warns for user and manual runs", () => {
     expect(shouldWarnOnOrphanedUserRepair("user")).toBe(true);
diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -563,6 +563,44 @@ function isMidTurnPrecheckAssistantError(message: AgentMessage | undefined): boo
   return record.stopReason === "error" && record.errorMessage === MID_TURN_PRECHECK_ERROR_MESSAGE;
 }
 
+export function removeSessionManagerLeafEntry(params: {
+  sessionManager: {
+    branch: (branchFromId: string) => void;
+    resetLeaf: () => void;
+  };
+  leafEntry: { id?: string; parentId?: string | null };
+}): void {
+  const mutableSessionManager = params.sessionManager as unknown as {
+    fileEntries?: Array<{ id?: string; type?: string }>;
+    byId?: Map<string, unknown>;
+    leafId?: string | null;
+    _rewriteFile?: () => void;
+  };
+  const leafId = params.leafEntry.id;
+  const fileEntries = mutableSessionManager.fileEntries;
+  const leafIndex = leafId ? fileEntries?.findIndex((entry) => entry.id === leafId) : -1;
+
+  if (
+    leafId &&
+    fileEntries &&
+    leafIndex !== undefined &&
+    leafIndex >= 0 &&
+    typeof mutableSessionManager._rewriteFile === "function"
+  ) {
+    fileEntries.splice(leafIndex, 1);
+    mutableSessionManager.byId?.delete(leafId);
+    mutableSessionManager.leafId = params.leafEntry.parentId ?? null;
+    mutableSessionManager._rewriteFile();
+    return;
+  }
+
+  if (params.leafEntry.parentId) {
+    params.sessionManager.branch(params.leafEntry.parentId);
+  } else {
+    params.sessionManager.resetLeaf();
+  }
+}
+
 function removeTrailingMidTurnPrecheckAssistantError(params: {
   activeSession: { agent: { state: { messages: AgentMessage[] } } };
   sessionManager: ReturnType<typeof guardSessionManager>;
@@ -2738,11 +2776,7 @@ export async function runEmbeddedAttempt(
           });
           effectivePrompt = orphanPromptMerge.prompt;
           if (orphanPromptMerge.removeLeaf) {
-            if (leafEntry.parentId) {
-              sessionManager.branch(leafEntry.parentId);
-            } else {
-              sessionManager.resetLeaf();
-            }
+            removeSessionManagerLeafEntry({ sessionManager, leafEntry });
             const sessionContext = sessionManager.buildSessionContext();
             activeSession.agent.state.messages = sessionContext.messages;
           }
