import { beforeEach, describe, expect, it, vi } from "vitest"; import { CloudMessagePersistence } from "../CloudMessagePersistence"; import type { AssistantCloud } from "../AssistantCloud"; import { CloudAPIError } from "../AssistantCloudAPI"; function createMockCloud() { return { threads: { messages: { create: vi.fn(), list: vi.fn(), update: vi.fn(), delete: vi.fn().mockResolvedValue(undefined), }, }, } as unknown as AssistantCloud; } function createCloudMessages(count: number) { return Array.from({ length: count }, (_, index) => ({ id: `message-${index + 1}`, parent_id: null, height: index, created_at: "2025-01-01T00:00:00Z" as unknown as Date, updated_at: "2025-01-01T00:00:00Z" as unknown as Date, format: "aui/v0", content: {}, })); } describe("CloudMessagePersistence", () => { let cloud: AssistantCloud; let persistence: CloudMessagePersistence; beforeEach(() => { vi.restoreAllMocks(); cloud = createMockCloud(); persistence = new CloudMessagePersistence(cloud); }); it("waits for a pending append before deleting its remote ID", async () => { let resolveAppend!: (value: { message_id: string }) => void; vi.mocked(cloud.threads.messages.create).mockImplementationOnce( () => new Promise((resolve) => { resolveAppend = resolve; }), ); const append = persistence.append( "thread-1", "local-1", null, "aui/v0", {}, ); const deletion = persistence.delete("thread-1", ["local-1"]); expect(cloud.threads.messages.delete).not.toHaveBeenCalled(); resolveAppend({ message_id: "msg_1" }); await append; await deletion; expect(cloud.threads.messages.delete).toHaveBeenCalledWith("thread-1", [ "msg_1", ]); expect(persistence.isPersisted("local-1")).toBe(false); }); it("skips unmapped IDs and failed appends without a delete request", async () => { await persistence.delete("thread-1", ["missing"]); let rejectAppend!: (error: Error) => void; vi.mocked(cloud.threads.messages.create).mockImplementationOnce( () => new Promise((_, reject) => { rejectAppend = reject; }), ); const append = persistence.append( "thread-1", "local-1", null, "aui/v0", {}, ); const rejected = expect(append).rejects.toThrow("append failed"); const deletion = persistence.delete("thread-1", ["local-1"]); rejectAppend(new Error("append failed")); await rejected; await deletion; expect(cloud.threads.messages.delete).not.toHaveBeenCalled(); }); it("keeps a replacement mapping made during deletion", async () => { persistence.record("local-1", "msg_1"); let resolveDelete!: () => void; vi.mocked(cloud.threads.messages.delete).mockImplementationOnce( () => new Promise((resolve) => { resolveDelete = resolve; }), ); const deletion = persistence.delete("thread-1", ["local-1"]); await vi.waitFor(() => expect(cloud.threads.messages.delete).toHaveBeenCalledOnce(), ); persistence.record("local-1", "msg_2"); resolveDelete(); await deletion; expect(await persistence.getRemoteId("local-1")).toBe("msg_2"); }); it("treats a 404 as already deleted and clears the mapping", async () => { persistence.record("local-1", "msg_1"); vi.mocked(cloud.threads.messages.delete).mockRejectedValueOnce( new CloudAPIError("Messages not found", 404), ); await persistence.delete("thread-1", ["local-1"]); expect(persistence.isPersisted("local-1")).toBe(false); }); it("propagates a server error and retains the mapping", async () => { persistence.record("local-1", "msg_1"); const failure = new CloudAPIError("failed", 500); vi.mocked(cloud.threads.messages.delete).mockRejectedValueOnce(failure); await expect(persistence.delete("thread-1", ["local-1"])).rejects.toBe( failure, ); expect(persistence.isPersisted("local-1")).toBe(true); }); it("appends a message and maps local ID to remote ID", async () => { vi.mocked(cloud.threads.messages.create).mockResolvedValue({ message_id: "remote-1", }); await persistence.append("thread-1", "local-1", null, "aui/v0", { text: "hello", }); expect(persistence.isPersisted("local-1")).toBe(true); expect(await persistence.getRemoteId("local-1")).toBe("remote-1"); expect(persistence.getResolvedRemoteId("local-1")).toBe("remote-1"); }); it.each(["__proto__", "constructor", "toString"])( "supports prototype-named local ID %s", async (messageId) => { vi.mocked(cloud.threads.messages.create).mockResolvedValue({ message_id: `remote-${messageId}`, }); expect(persistence.isPersisted(messageId)).toBe(false); await persistence.append("thread-1", messageId, null, "aui/v0", { text: "hello", }); expect(persistence.isPersisted(messageId)).toBe(true); expect(await persistence.getRemoteId(messageId)).toBe( `remote-${messageId}`, ); }, ); it("uses the current client without losing ID mappings", async () => { const firstCloud = createMockCloud(); const secondCloud = createMockCloud(); let currentCloud = firstCloud; persistence = new CloudMessagePersistence(() => currentCloud); vi.mocked(firstCloud.threads.messages.create).mockResolvedValue({ message_id: "remote-parent", }); vi.mocked(secondCloud.threads.messages.create).mockResolvedValue({ message_id: "remote-child", }); await persistence.append("thread-1", "parent", null, "aui/v0", { text: "parent", }); currentCloud = secondCloud; await persistence.append("thread-1", "child", "parent", "aui/v0", { text: "child", }); expect(secondCloud.threads.messages.create).toHaveBeenCalledWith( "thread-1", expect.objectContaining({ parent_id: "remote-parent" }), ); expect(await persistence.getRemoteId("child")).toBe("remote-child"); }); it("resolves parent ID from a concurrent append", async () => { // Parent creation is delayed — the promise won't resolve immediately let resolveParent!: (v: { message_id: string }) => void; vi.mocked(cloud.threads.messages.create).mockImplementationOnce( () => new Promise((resolve) => { resolveParent = resolve; }), ); vi.mocked(cloud.threads.messages.create).mockResolvedValueOnce({ message_id: "remote-child", }); // Start parent append (doesn't resolve yet) const parentPromise = persistence.append( "thread-1", "parent", null, "aui/v0", { text: "parent" }, ); // Start child append — it will await the parent's promise for its remote ID const childPromise = persistence.append( "thread-1", "child", "parent", "aui/v0", { text: "child" }, ); // Now resolve the parent resolveParent({ message_id: "remote-parent" }); await parentPromise; await childPromise; // The child's create call should have used the parent's resolved remote ID const childCreateCall = vi.mocked(cloud.threads.messages.create).mock .calls[1]!; expect(childCreateCall[1]).toMatchObject({ parent_id: "remote-parent", }); }); it("deduplicates concurrent child appends while the parent is pending", async () => { let resolveParent!: (value: { message_id: string }) => void; vi.mocked(cloud.threads.messages.create) .mockImplementationOnce( () => new Promise((resolve) => { resolveParent = resolve; }), ) .mockResolvedValue({ message_id: "remote-child" }); const parent = persistence.append("thread-1", "parent", null, "aui/v0", { text: "parent", }); const firstChild = persistence.append( "thread-1", "child", "parent", "aui/v0", { text: "child" }, ); const secondChild = persistence.append( "thread-1", "child", "parent", "aui/v0", { text: "child" }, ); expect(persistence.isPersisted("child")).toBe(true); resolveParent({ message_id: "remote-parent" }); await Promise.all([parent, firstChild, secondChild]); expect(cloud.threads.messages.create).toHaveBeenCalledTimes(2); }); it("resolves remote IDs throughout a concurrent message chain", async () => { let resolveParent!: (value: { message_id: string }) => void; vi.mocked(cloud.threads.messages.create) .mockImplementationOnce( () => new Promise((resolve) => { resolveParent = resolve; }), ) .mockResolvedValueOnce({ message_id: "remote-child" }) .mockResolvedValueOnce({ message_id: "remote-grandchild" }); const parent = persistence.append("thread-1", "parent", null, "aui/v0", { text: "parent", }); const child = persistence.append("thread-1", "child", "parent", "aui/v0", { text: "child", }); const grandchild = persistence.append( "thread-1", "grandchild", "child", "aui/v0", { text: "grandchild" }, ); resolveParent({ message_id: "remote-parent" }); await Promise.all([parent, child, grandchild]); expect(cloud.threads.messages.create).toHaveBeenCalledWith("thread-1", { parent_id: "remote-parent", format: "aui/v0", content: { text: "child" }, }); expect(cloud.threads.messages.create).toHaveBeenCalledWith("thread-1", { parent_id: "remote-child", format: "aui/v0", content: { text: "grandchild" }, }); }); it("re-appends a message after its mapping has settled", async () => { vi.mocked(cloud.threads.messages.create) .mockResolvedValueOnce({ message_id: "remote-1" }) .mockResolvedValueOnce({ message_id: "remote-2" }); await persistence.append("thread-1", "local-1", null, "aui/v0", { text: "first", }); await persistence.append("thread-1", "local-1", null, "aui/v0", { text: "second", }); expect(cloud.threads.messages.create).toHaveBeenCalledTimes(2); expect(await persistence.getRemoteId("local-1")).toBe("remote-2"); }); it("loaded messages are marked as persisted and not re-created", async () => { vi.mocked(cloud.threads.messages.list).mockResolvedValue({ messages: [ { id: "msg-1", parent_id: null, height: 0, created_at: "2025-01-01T00:00:00Z" as unknown as Date, updated_at: "2025-01-01T00:00:00Z" as unknown as Date, format: "aui/v0", content: { text: "loaded" }, }, ], }); await persistence.load("thread-1"); expect(persistence.isPersisted("msg-1")).toBe(true); expect(await persistence.getRemoteId("msg-1")).toBe("msg-1"); }); it("preserves a pending append mapping when load returns the same ID", async () => { const messages = createCloudMessages(1); let resolveLoad!: (value: { messages: typeof messages }) => void; let resolveAppend!: (value: { message_id: string }) => void; vi.mocked(cloud.threads.messages.list).mockImplementationOnce( () => new Promise((resolve) => { resolveLoad = resolve; }), ); vi.mocked(cloud.threads.messages.create).mockImplementationOnce( () => new Promise((resolve) => { resolveAppend = resolve; }), ); const load = persistence.load("thread-1"); const append = persistence.append( "thread-1", "message-1", null, "aui/v0", {}, ); resolveLoad({ messages }); await load; resolveAppend({ message_id: "remote-2" }); await append; expect(await persistence.getRemoteId("message-1")).toBe("remote-2"); await persistence.update("thread-1", "message-1", "aui/v0", { text: "updated", }); expect(cloud.threads.messages.update).toHaveBeenCalledWith( "thread-1", "remote-2", { content: { text: "updated" } }, ); }); it("preserves a resolved append mapping when load returns the same ID", async () => { const messages = createCloudMessages(1); vi.mocked(cloud.threads.messages.create).mockResolvedValue({ message_id: "remote-1", }); vi.mocked(cloud.threads.messages.list).mockResolvedValue({ messages, }); await persistence.append("thread-1", "message-1", null, "aui/v0", {}); await persistence.load("thread-1"); expect(await persistence.getRemoteId("message-1")).toBe("remote-1"); }); it("keeps the loaded ID when a pending append for it fails", async () => { const messages = createCloudMessages(1); const failure = new Error("create failed"); let resolveLoad!: (value: { messages: typeof messages }) => void; let rejectAppend!: (error: Error) => void; vi.mocked(cloud.threads.messages.list).mockImplementationOnce( () => new Promise((resolve) => { resolveLoad = resolve; }), ); vi.mocked(cloud.threads.messages.create).mockImplementationOnce( () => new Promise((_resolve, reject) => { rejectAppend = reject; }), ); const load = persistence.load("thread-1"); const append = persistence.append( "thread-1", "message-1", null, "aui/v0", {}, ); resolveLoad({ messages }); await load; rejectAppend(failure); await expect(append).rejects.toBe(failure); expect(persistence.isPersisted("message-1")).toBe(true); expect(await persistence.getRemoteId("message-1")).toBe("message-1"); }); it("does not restore a loaded ID after reset when its pending append fails", async () => { const messages = createCloudMessages(1); const failure = new Error("create failed"); let rejectAppend!: (error: Error) => void; vi.mocked(cloud.threads.messages.list).mockResolvedValue({ messages }); vi.mocked(cloud.threads.messages.create).mockImplementationOnce( () => new Promise((_resolve, reject) => { rejectAppend = reject; }), ); const append = persistence.append( "thread-1", "message-1", null, "aui/v0", {}, ); await persistence.load("thread-1"); persistence.reset(); rejectAppend(failure); await expect(append).rejects.toBe(failure); expect(persistence.isPersisted("message-1")).toBe(false); }); it("does not restore IDs from a load that finishes after reset", async () => { const oldMessages = createCloudMessages(1); const newMessages = [{ ...oldMessages[0]!, id: "new-message" }]; let resolveOld!: (value: { messages: typeof oldMessages }) => void; vi.mocked(cloud.threads.messages.list) .mockImplementationOnce( () => new Promise((resolve) => { resolveOld = resolve; }), ) .mockResolvedValueOnce({ messages: newMessages }); const oldLoad = persistence.load("old-thread"); persistence.reset(); await persistence.load("new-thread"); resolveOld({ messages: oldMessages }); expect(await oldLoad).toEqual(oldMessages); expect(persistence.isPersisted("message-1")).toBe(false); expect(persistence.isPersisted("new-message")).toBe(true); }); it("keeps a new append mapping when an old load returns the same ID", async () => { const messages = createCloudMessages(1); let resolveOld!: (value: { messages: typeof messages }) => void; vi.mocked(cloud.threads.messages.list).mockImplementationOnce( () => new Promise((resolve) => { resolveOld = resolve; }), ); vi.mocked(cloud.threads.messages.create).mockResolvedValue({ message_id: "new-remote-id", }); const oldLoad = persistence.load("old-thread"); persistence.reset(); await persistence.append("new-thread", "message-1", null, "aui/v0", {}); resolveOld({ messages }); await oldLoad; expect(await persistence.getRemoteId("message-1")).toBe("new-remote-id"); }); it("does not populate any old IDs when reset happens between history pages", async () => { const messages = createCloudMessages(201); let resolveLast!: (value: { messages: typeof messages }) => void; vi.mocked(cloud.threads.messages.list) .mockResolvedValueOnce({ messages: messages.slice(0, 200) }) .mockImplementationOnce( () => new Promise((resolve) => { resolveLast = resolve; }), ); const load = persistence.load("old-thread"); await vi.waitFor(() => expect(cloud.threads.messages.list).toHaveBeenCalledTimes(2), ); persistence.reset(); resolveLast({ messages: messages.slice(200) }); expect(await load).toEqual(messages); expect( messages.some((message) => persistence.isPersisted(message.id)), ).toBe(false); }); it("propagates an old load failure without disturbing post-reset state", async () => { let rejectOld!: (error: Error) => void; vi.mocked(cloud.threads.messages.list).mockImplementationOnce( () => new Promise((_resolve, reject) => { rejectOld = reject; }), ); vi.mocked(cloud.threads.messages.create).mockResolvedValue({ message_id: "new-remote-id", }); const failure = new Error("history failed"); const oldLoad = persistence.load("old-thread"); const rejected = expect(oldLoad).rejects.toBe(failure); persistence.reset(); await persistence.append("new-thread", "new-message", null, "aui/v0", {}); rejectOld(failure); await rejected; expect(await persistence.getRemoteId("new-message")).toBe("new-remote-id"); }); it("maps prototype-named loaded messages", async () => { vi.mocked(cloud.threads.messages.list).mockResolvedValue({ messages: [ { id: "__proto__", parent_id: null, height: 0, created_at: "2025-01-01T00:00:00Z" as unknown as Date, updated_at: "2025-01-01T00:00:00Z" as unknown as Date, format: "aui/v0", content: { text: "loaded" }, }, ], }); await persistence.load("thread-1"); expect(persistence.isPersisted("__proto__")).toBe(true); expect(await persistence.getRemoteId("__proto__")).toBe("__proto__"); }); it("follows the message cursor across pages in server order", async () => { const rows = createCloudMessages(450); vi.mocked(cloud.threads.messages.list).mockImplementation( async (_threadId, query) => { const start = query?.after ? rows.findIndex((row) => row.id === query.after) + 1 : 0; return { messages: rows.slice(start, start + 200) }; }, ); const messages = await persistence.load("thread-1", "aui/v0"); expect(messages.map((message) => message.id)).toEqual( rows.map((row) => row.id), ); expect(cloud.threads.messages.list).toHaveBeenCalledTimes(3); expect(cloud.threads.messages.list).toHaveBeenNthCalledWith(1, "thread-1", { format: "aui/v0", limit: 200, }); expect(cloud.threads.messages.list).toHaveBeenNthCalledWith(3, "thread-1", { format: "aui/v0", limit: 200, after: "message-400", }); }); it("stops without duplicates when a later cursor stops resolving", async () => { const rows = createCloudMessages(450); vi.mocked(cloud.threads.messages.list).mockImplementation( async (_threadId, query) => { // The third request replays page one, the way the endpoint answers a // cursor whose message no longer resolves. const start = query?.after && query.after !== "message-400" ? rows.findIndex((row) => row.id === query.after) + 1 : 0; return { messages: rows.slice(start, start + 200) }; }, ); const messages = await persistence.load("thread-1", "aui/v0"); expect(messages.map((message) => message.id)).toEqual( rows.slice(0, 400).map((row) => row.id), ); expect(cloud.threads.messages.list).toHaveBeenCalledTimes(3); }); it("stops when a page does not advance the cursor", async () => { vi.mocked(cloud.threads.messages.list).mockResolvedValue({ messages: createCloudMessages(200), }); const messages = await persistence.load("thread-1", "aui/v0"); expect(messages).toHaveLength(200); expect(cloud.threads.messages.list).toHaveBeenCalledTimes(2); }); it("updates an already-persisted message", async () => { vi.mocked(cloud.threads.messages.create).mockResolvedValue({ message_id: "remote-1", }); vi.mocked(cloud.threads.messages.update).mockResolvedValue(undefined); await persistence.append("thread-1", "local-1", null, "aui/v0", { text: "original", }); await persistence.update("thread-1", "local-1", "aui/v0", { text: "updated", }); expect(cloud.threads.messages.update).toHaveBeenCalledWith( "thread-1", "remote-1", { content: { text: "updated" } }, ); }); it("warns and skips update when no remote id is mapped", async () => { const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); await persistence.update("thread-1", "unmapped-1", "aui/v0", { text: "x", }); expect(cloud.threads.messages.update).not.toHaveBeenCalled(); expect(warn).toHaveBeenCalledWith( "Skipping update for message unmapped-1: no remote id is mapped.", ); }); it("cleans up ID mapping on append failure", async () => { vi.mocked(cloud.threads.messages.create).mockRejectedValue( new Error("network error"), ); await expect( persistence.append("thread-1", "local-1", null, "aui/v0", { text: "fail", }), ).rejects.toThrow("network error"); expect(persistence.isPersisted("local-1")).toBe(false); }); it("does not restore an append mapping after reset", async () => { let resolveAppend!: (value: { message_id: string }) => void; vi.mocked(cloud.threads.messages.create).mockImplementationOnce( () => new Promise((resolve) => { resolveAppend = resolve; }), ); const append = persistence.append( "thread-1", "local-1", null, "aui/v0", {}, ); persistence.reset(); resolveAppend({ message_id: "remote-1" }); await expect(append).resolves.toBeUndefined(); expect(persistence.isPersisted("local-1")).toBe(false); expect(await persistence.getRemoteId("local-1")).toBeUndefined(); }); it("reset clears all ID mappings", async () => { vi.mocked(cloud.threads.messages.create).mockResolvedValue({ message_id: "remote-1", }); await persistence.append("thread-1", "local-1", null, "aui/v0", { text: "hello", }); expect(persistence.isPersisted("local-1")).toBe(true); persistence.reset(); expect(persistence.isPersisted("local-1")).toBe(false); }); });