@assistant-ui/react 0.10.44 → 0.10.46

package/src/runtimes/utils/MessageRepository.tsx

@@ -414,8 +414,7 @@ export class MessageRepository {
    */
   resetHead(messageId: string | null) {
     if (messageId === null) {
-      this.head = null;
-      this._messages.dirty();
+      this.clear();
       return;
     }
 
@@ -425,6 +424,22 @@ export class MessageRepository {
         "MessageRepository(resetHead): Branch not found. This is likely an internal bug in assistant-ui.",
       );
 
+    if (message.children.length > 0) {
+      const deleteDescendants = (msg: RepositoryMessage) => {
+        for (const childId of msg.children) {
+          const childMessage = this.messages.get(childId);
+          if (childMessage) {
+            deleteDescendants(childMessage);
+            this.messages.delete(childId);
+          }
+        }
+      };
+      deleteDescendants(message);
+
+      message.children = [];
+      message.next = null;
+    }
+
     this.head = message;
     for (
       let current: RepositoryMessage | null = message;

package/src/tests/MessageRepository.test.ts

@@ -308,6 +308,51 @@ describe("MessageRepository", () => {
       expect(messages.map((m) => m.id)).toEqual(["parent-id"]);
     });
 
+    /**
+     * Tests that resetting head to a message with children removes those children.
+     * All descendants should be deleted from the repository.
+     */
+    it("should remove children when resetting head to a message with children", () => {
+      const parent = createTestMessage({ id: "parent-id" });
+      const child = createTestMessage({ id: "child-id" });
+      const grandchild1 = createTestMessage({ id: "grandchild1-id" });
+      const grandchild2 = createTestMessage({ id: "grandchild2-id" });
+      const greatGrandchild = createTestMessage({ id: "greatgrandchild-id" });
+
+      // Build tree: parent -> child -> grandchild1
+      //                            \-> grandchild2 -> greatGrandchild
+      repository.addOrUpdateMessage(null, parent);
+      repository.addOrUpdateMessage("parent-id", child);
+      repository.addOrUpdateMessage("child-id", grandchild1);
+      repository.addOrUpdateMessage("child-id", grandchild2);
+      repository.addOrUpdateMessage("grandchild2-id", greatGrandchild);
+
+      // Reset to child (which has children)
+      repository.resetHead("child-id");
+
+      // Head should be child
+      expect(repository.headId).toBe("child-id");
+
+      // Messages should only include parent and child
+      const messages = repository.getMessages();
+      expect(messages.map((m) => m.id)).toEqual(["parent-id", "child-id"]);
+
+      // Verify children are removed from repository
+      expect(() => repository.getMessage("grandchild1-id")).toThrow(
+        /Message not found/,
+      );
+      expect(() => repository.getMessage("grandchild2-id")).toThrow(
+        /Message not found/,
+      );
+      expect(() => repository.getMessage("greatgrandchild-id")).toThrow(
+        /Message not found/,
+      );
+
+      // Verify branches are empty for the child
+      const branches = repository.getBranches("child-id");
+      expect(branches).toEqual(["child-id"]);
+    });
+
     /**
      * Tests resetting the head to null.
      * This should clear the active branch completely.