fix(ui): sync resource membership document stars

Co-Authored-By: Paperclip <noreply@paperclip.ing>

ui/src/components/SidebarAgents.test.tsx

@@ -230,6 +230,8 @@ describe("SidebarAgents", () => {
     memberships = {
       projectMemberships: {},
       agentMemberships: {},
+      starredDocumentIds: [],
+      documentStarredAt: {},
       updatedAt: null,
     };
     mockResourceMembershipsApi.listMine.mockImplementation(() => Promise.resolve(memberships));
@@ -368,8 +370,10 @@ describe("SidebarAgents", () => {
       agentMemberships: {},
       starredProjectIds: [],
       starredAgentIds: ["agent-b"],
+      starredDocumentIds: [],
       projectStarredAt: {},
       agentStarredAt: {},
+      documentStarredAt: {},
       updatedAt: new Date(),
     };
 
@@ -417,8 +421,10 @@ describe("SidebarAgents", () => {
       agentMemberships: { "agent-b": "joined" },
       starredProjectIds: [],
       starredAgentIds: ["agent-b"],
+      starredDocumentIds: [],
       projectStarredAt: {},
       agentStarredAt: {},
+      documentStarredAt: {},
       updatedAt: new Date(),
     };
     mockResourceMembershipsApi.updateAgent.mockRejectedValue(new Error("nope"));
@@ -544,6 +550,8 @@ describe("SidebarAgents", () => {
       resolveMemberships({
         projectMemberships: {},
         agentMemberships: { "agent-1": "left" },
+        starredDocumentIds: [],
+        documentStarredAt: {},
         updatedAt: null,
       });
     });

ui/src/components/SidebarProjects.test.tsx

@@ -297,6 +297,8 @@ describe("SidebarProjects", () => {
     memberships = {
       projectMemberships: {},
       agentMemberships: {},
+      starredDocumentIds: [],
+      documentStarredAt: {},
       updatedAt: null,
     };
     mockResourceMembershipsApi.listMine.mockImplementation(() => Promise.resolve(memberships));

ui/src/components/SidebarStarredProjects.test.tsx

@@ -135,8 +135,10 @@ describe("SidebarStarredProjects", () => {
       agentMemberships: {},
       starredProjectIds: [],
       starredAgentIds: [],
+      starredDocumentIds: [],
       projectStarredAt: {},
       agentStarredAt: {},
+      documentStarredAt: {},
       updatedAt: null,
     };
     mockResourceMembershipsApi.listMine.mockImplementation(() => Promise.resolve(memberships));

ui/src/hooks/useResourceMemberships.ts

@@ -9,7 +9,7 @@ import { useToastActions } from "../context/ToastContext";
 import { queryKeys } from "../lib/queryKeys";
 
 type MutationVariables = {
-  resourceType: ResourceMembershipResourceType;
+  resourceType: JoinableResourceType;
   resourceId: string;
   resourceName: string;
   /** Join / leave transition. Omit to only change the starred flag. */
@@ -18,19 +18,23 @@ type MutationVariables = {
   starred?: boolean;
 };
 
+type JoinableResourceType = Exclude<ResourceMembershipResourceType, "document">;
+
 function emptyMemberships(): ResourceMemberships {
   return {
     projectMemberships: {},
     agentMemberships: {},
     starredProjectIds: [],
     starredAgentIds: [],
+    starredDocumentIds: [],
     projectStarredAt: {},
     agentStarredAt: {},
+    documentStarredAt: {},
     updatedAt: null,
   };
 }
 
-function starKeys(resourceType: ResourceMembershipResourceType) {
+function starKeys(resourceType: JoinableResourceType) {
   return resourceType === "project"
     ? { ids: "starredProjectIds", at: "projectStarredAt", state: "projectMemberships" }
     : { ids: "starredAgentIds", at: "agentStarredAt", state: "agentMemberships" };
@@ -45,7 +49,7 @@ function starKeys(resourceType: ResourceMembershipResourceType) {
  */
 function applyMembershipChange(
   current: ResourceMemberships | undefined,
-  resourceType: ResourceMembershipResourceType,
+  resourceType: JoinableResourceType,
   resourceId: string,
   change: { state?: ResourceMembershipState; starred?: boolean },
 ): ResourceMemberships {
@@ -96,7 +100,7 @@ function applyMembershipChange(
 
 export function resourceMembershipState(
   memberships: ResourceMemberships | undefined,
-  resourceType: ResourceMembershipResourceType,
+  resourceType: JoinableResourceType,
   resourceId: string,
 ): ResourceMembershipState {
   const state = resourceType === "project"
@@ -113,7 +117,9 @@ export function isStarred(
 ): boolean {
   const ids = resourceType === "project"
     ? memberships?.starredProjectIds
-    : memberships?.starredAgentIds;
+    : resourceType === "agent"
+      ? memberships?.starredAgentIds
+      : memberships?.starredDocumentIds;
   return Array.isArray(ids) && ids.includes(resourceId);
 }
 
@@ -124,7 +130,9 @@ export function starredResourceIds(
 ): string[] {
   const ids = resourceType === "project"
     ? memberships?.starredProjectIds
-    : memberships?.starredAgentIds;
+    : resourceType === "agent"
+      ? memberships?.starredAgentIds
+      : memberships?.starredDocumentIds;
   return Array.isArray(ids) ? ids : [];
 }