@assistant-ui/core 0.3.19 → 0.3.21

package/src/tests/external-store-thread-runtime-core.test.ts

@@ -1282,6 +1282,37 @@ describe("ExternalStoreThreadRuntimeCore - deleteMessage via setMessages", () =>
     }
   });
 
+  it("notifies the eviction when the host resyncs synchronously", async () => {
+    let current = [
+      message("u1", "user", "one"),
+      message("a1", "assistant", "two"),
+      message("u2", "user", "three"),
+      message("a2", "assistant", "four"),
+    ];
+    const store = (): ExternalStoreAdapter =>
+      makeStore({
+        messages: current,
+        setMessages: (m: import("../types/message").ThreadMessage[]) => {
+          current = m;
+          runtime.__internal_setAdapter(store());
+        },
+      });
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      store(),
+    );
+
+    const observed: (readonly string[])[] = [];
+    runtime.subscribe(() => {
+      observed.push(runtime.getBranches("a2"));
+    });
+
+    await runtime.deleteMessage("u2");
+
+    expect(runtime.messages.map((m) => m.id)).toEqual(["u1", "a1", "a2"]);
+    expect(observed.at(-1)).toEqual(["a2"]);
+  });
+
   it("evicts the deleted message on the onDelete path too", async () => {
     let current = [
       message("u1", "user", "hi"),
@@ -1378,6 +1409,192 @@ describe("ExternalStoreThreadRuntimeCore - deleteMessage via setMessages", () =>
     );
   });
 
+  const deferredDeleteStore = (
+    initial: import("../types/message").ThreadMessage[],
+  ) => {
+    let current = initial;
+    let isRunning = false;
+    let confirmDelete!: () => void;
+    const onDelete = vi.fn(
+      (id: string) =>
+        new Promise<void>((resolve) => {
+          confirmDelete = () => {
+            current = current.filter((m) => m.id !== id);
+            resolve();
+          };
+        }),
+    );
+    const onReload = vi.fn(async (parentId: string | null) => {
+      isRunning = true;
+      current = current.slice(
+        0,
+        current.findIndex((m) => m.id === parentId) + 1,
+      );
+    });
+    const store = () =>
+      makeStore({
+        messages: current,
+        onDelete,
+        onReload,
+        isRunning,
+      });
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      store(),
+    );
+    return {
+      runtime,
+      confirmDelete: () => confirmDelete(),
+      rerender: () => {
+        current = [...current];
+        runtime.__internal_setAdapter(store());
+      },
+      setStoreMessages: (m: import("../types/message").ThreadMessage[]) => {
+        current = m;
+        isRunning = false;
+        runtime.__internal_setAdapter(store());
+      },
+    };
+  };
+
+  it("keeps the eviction when the host re-renders while onDelete is in flight", async () => {
+    const { runtime, confirmDelete, rerender } = deferredDeleteStore([
+      message("u1", "user", "one"),
+      message("a1", "assistant", "two"),
+      message("u2", "user", "three"),
+      message("a2", "assistant", "four"),
+    ]);
+
+    const deletion = runtime.deleteMessage("u2");
+    rerender();
+    confirmDelete();
+    await deletion;
+    rerender();
+
+    expect(runtime.messages.map((m) => m.id)).toEqual(["u1", "a1", "a2"]);
+    expect(runtime.getBranches("a2")).toEqual(["a2"]);
+  });
+
+  it("keeps the eviction when the thread reloads while onDelete is in flight", async () => {
+    const { runtime, confirmDelete, rerender } = deferredDeleteStore([
+      message("u1", "user", "one"),
+      message("a1", "assistant", "two"),
+      message("u2", "user", "three"),
+      message("a2", "assistant", "four"),
+    ]);
+
+    const deletion = runtime.deleteMessage("u1");
+    await runtime.startRun({ parentId: "u2", sourceId: "a2", runConfig: {} });
+    rerender();
+    confirmDelete();
+    await deletion;
+    rerender();
+
+    expect(runtime.getBranches("a1")).toEqual(["a1"]);
+  });
+
+  it("keeps the eviction of a reload's parent deleted while onDelete is in flight", async () => {
+    const { runtime, confirmDelete, rerender, setStoreMessages } =
+      deferredDeleteStore([
+        message("u1", "user", "one"),
+        message("a1", "assistant", "two"),
+      ]);
+
+    const deletion = runtime.deleteMessage("u1");
+    await runtime.startRun({ parentId: "u1", sourceId: "a1", runConfig: {} });
+    rerender();
+    confirmDelete();
+    await deletion;
+    rerender();
+    setStoreMessages([message("a1b", "assistant", "regenerated")]);
+
+    expect(runtime.messages.map((m) => m.id)).toEqual(["a1b"]);
+    expect(runtime.getBranches("a1b")).toEqual(["a1", "a1b"]);
+  });
+
+  it.each([
+    ["resolves", () => Promise.resolve()],
+    ["rejects", () => Promise.reject(new Error("already deleting"))],
+  ])(
+    "keeps the eviction while another onDelete for the id is in flight when one %s",
+    async (_, settleDuplicate) => {
+      let current = [
+        message("u1", "user", "one"),
+        message("a1", "assistant", "two"),
+        message("u2", "user", "three"),
+        message("a2", "assistant", "four"),
+      ];
+      let confirmDelete!: () => void;
+      let calls = 0;
+      const onDelete = vi.fn((id: string) =>
+        ++calls > 1
+          ? settleDuplicate()
+          : new Promise<void>((resolve) => {
+              confirmDelete = () => {
+                current = current.filter((m) => m.id !== id);
+                resolve();
+              };
+            }),
+      );
+      const store = () => makeStore({ messages: current, onDelete });
+      const rerender = () => {
+        current = [...current];
+        runtime.__internal_setAdapter(store());
+      };
+      const runtime = new ExternalStoreThreadRuntimeCore(
+        mockContextProvider,
+        store(),
+      );
+
+      const deletion = runtime.deleteMessage("u2");
+      await runtime.deleteMessage("u2").catch(() => {});
+      rerender();
+      confirmDelete();
+      await deletion;
+      rerender();
+
+      expect(runtime.messages.map((m) => m.id)).toEqual(["u1", "a1", "a2"]);
+      expect(runtime.getBranches("a2")).toEqual(["a2"]);
+    },
+  );
+
+  it("keeps the eviction when a second onDelete for the id rejects after the host dropped it", async () => {
+    let current = [
+      message("u1", "user", "one"),
+      message("a1", "assistant", "two"),
+    ];
+    const settlers: Array<() => void> = [];
+    const onDelete = vi.fn(
+      (id: string) =>
+        new Promise<void>((resolve, reject) => {
+          settlers.push(
+            settlers.length === 0
+              ? () => {
+                  current = current.filter((m) => m.id !== id);
+                  resolve();
+                }
+              : () => reject(new Error("already deleted")),
+          );
+        }),
+    );
+    const store = () => makeStore({ messages: current, onDelete });
+    const runtime = new ExternalStoreThreadRuntimeCore(
+      mockContextProvider,
+      store(),
+    );
+
+    const first = runtime.deleteMessage("u1");
+    const second = runtime.deleteMessage("u1");
+    settlers[0]!();
+    await first;
+    settlers[1]!();
+    await expect(second).rejects.toThrow("already deleted");
+    runtime.__internal_setAdapter(store());
+
+    expect(runtime.messages.map((m) => m.id)).toEqual(["a1"]);
+    expect(runtime.getBranches("a1")).toEqual(["a1"]);
+  });
+
   it("keeps a pending eviction across a branch switch swallowed mid-run", async () => {
     let current = [
       message("u1", "user", "hi"),