diff --git a/ui/src/components/SidebarAgents.test.tsx b/ui/src/components/SidebarAgents.test.tsx index 3ee1aeea3b8..89941d58f8f 100644 --- a/ui/src/components/SidebarAgents.test.tsx +++ b/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, }); }); diff --git a/ui/src/components/SidebarProjects.test.tsx b/ui/src/components/SidebarProjects.test.tsx index 0289a143494..abe3ac11d64 100644 --- a/ui/src/components/SidebarProjects.test.tsx +++ b/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)); diff --git a/ui/src/components/SidebarStarredProjects.test.tsx b/ui/src/components/SidebarStarredProjects.test.tsx index fcce78b7183..3c4ed424f38 100644 --- a/ui/src/components/SidebarStarredProjects.test.tsx +++ b/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)); diff --git a/ui/src/hooks/useResourceMemberships.ts b/ui/src/hooks/useResourceMemberships.ts index c20d06cd55d..db0d2a8c935 100644 --- a/ui/src/hooks/useResourceMemberships.ts +++ b/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; + 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 : []; }