diff --git a/ui/src/hooks/useDocumentStarMutation.test.tsx b/ui/src/hooks/useDocumentStarMutation.test.tsx index 0d5cdbfa8a0..b9c3b02a70d 100644 --- a/ui/src/hooks/useDocumentStarMutation.test.tsx +++ b/ui/src/hooks/useDocumentStarMutation.test.tsx @@ -149,6 +149,32 @@ describe("useDocumentStarMutation", () => { expect(starredIds(queryClient)).toContain("doc-1"); }); + it("invalidates the Artifacts list on settle so an unstarred doc leaves the Starred tab", async () => { + membershipsApiMock.listMine.mockResolvedValue({ + ...emptyMemberships(), + starredDocumentIds: ["doc-1"], + documentStarredAt: { "doc-1": new Date().toISOString() }, + }); + membershipsApiMock.updateDocument.mockResolvedValue({ + resourceType: "document", + resourceId: "doc-1", + state: "left", + starredAt: null, + updatedAt: new Date(), + }); + + let handle: HarnessHandle | null = null; + const queryClient = await render((h) => { handle = h; }); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + await act(async () => { + handle!.mutate({ documentId: "doc-1", documentName: "Launch Brief", starred: false }); + await flush(); + }); + + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ["artifacts"] }); + }); + it("rolls back the optimistic unstar and toasts when the server rejects", async () => { membershipsApiMock.listMine.mockResolvedValue({ ...emptyMemberships(), diff --git a/ui/src/hooks/useResourceMemberships.ts b/ui/src/hooks/useResourceMemberships.ts index 2531ce7c71d..0f8bca1a218 100644 --- a/ui/src/hooks/useResourceMemberships.ts +++ b/ui/src/hooks/useResourceMemberships.ts @@ -259,16 +259,22 @@ export function useDocumentStarMutation(companyId: string | null | undefined) { onMutate: async (variables) => { await queryClient.cancelQueries({ queryKey }); const previous = queryClient.getQueryData(queryKey); + // Capture only this document's prior star so a rollback reverts just this + // change and leaves any concurrent star mutation's optimistic state intact. + const previousStarred = isStarred(previous, "document", variables.documentId); queryClient.setQueryData( queryKey, applyDocumentStarChange(previous, variables.documentId, variables.starred), ); - return { previous }; + return { previousStarred }; }, onError: (error, variables, context) => { - if (context?.previous) { - queryClient.setQueryData(queryKey, context.previous); - } + // Targeted rollback: re-apply this document's prior star on top of the + // current cache, so we don't clobber a newer concurrent update. + queryClient.setQueryData( + queryKey, + (current) => applyDocumentStarChange(current, variables.documentId, context?.previousStarred ?? false), + ); pushToast({ title: `Couldn't ${variables.starred ? "star" : "unstar"} ${variables.documentName}.`, body: error instanceof Error ? error.message : "Try again.", @@ -284,6 +290,10 @@ export function useDocumentStarMutation(companyId: string | null | undefined) { }, onSettled: () => { queryClient.invalidateQueries({ queryKey }); + // The Artifacts "Starred" tab is a separate list query keyed by ["artifacts", ...], + // so refresh it too — otherwise an unstarred document lingers in the grid + // until the next natural refetch. + queryClient.invalidateQueries({ queryKey: ["artifacts"] }); }, }); }