Fix Work Memories UX

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

server/src/__tests__/memory-routes.test.ts

@@ -27,6 +27,7 @@ const mockMemoryService = vi.hoisted(() => ({
   review: vi.fn(),
   sweepRetention: vi.fn(),
   listRecords: vi.fn(),
+  countRecords: vi.fn(),
   getRecord: vi.fn(),
   listOperations: vi.fn(),
   listExtractionJobs: vi.fn(),
@@ -245,6 +246,35 @@ describe("memory routes", () => {
     expect(mockLogActivity).toHaveBeenCalledOnce();
   });
 
+  it("routes count-only record queries through memory service", async () => {
+    mockMemoryService.countRecords.mockResolvedValue({ count: 152 });
+    const app = createApp({
+      type: "board",
+      userId: "board-user",
+      source: "session",
+      companyIds: [companyA],
+      isInstanceAdmin: false,
+    });
+
+    const res = await request(app)
+      .get(`/api/companies/${companyA}/memory/records`)
+      .query({ count: "only", reviewState: "pending" })
+      .set("Origin", "http://localhost:3100");
+
+    expect(res.status).toBe(200);
+    expect(res.body).toEqual({ count: 152 });
+    expect(mockMemoryService.countRecords).toHaveBeenCalledWith(
+      companyA,
+      expect.objectContaining({
+        count: "only",
+        reviewState: "pending",
+        includeRevoked: false,
+      }),
+      expect.objectContaining({ actorType: "user", userId: "board-user" }),
+    );
+    expect(mockMemoryService.listRecords).not.toHaveBeenCalled();
+  });
+
   it("sets project memory overrides through the owning project company", async () => {
     const projectId = "77777777-7777-4777-8777-777777777777";
     mockMemoryService.setProjectOverride.mockResolvedValue({

server/src/routes/memory.ts

@@ -268,6 +268,10 @@ export function memoryRoutes(
     const companyId = req.params.companyId as string;
     assertCompanyAccess(req, companyId);
     const parsed = memoryListRecordsQuerySchema.parse(req.query);
+    if (parsed.count === "only") {
+      res.json(await memory.countRecords(companyId, parsed, actorInfoFromReq(req)));
+      return;
+    }
     res.json(await memory.listRecords(companyId, parsed, actorInfoFromReq(req)));
   });
 

server/src/services/memory.ts

@@ -40,6 +40,7 @@ import type {
   MemoryQuery,
   MemoryQueryResult,
   MemoryRecord,
+  MemoryRecordCountResult,
   MemoryRefreshJob,
   MemoryRefreshJobResult,
   MemoryRefreshJobSourceCounts,
@@ -1121,7 +1122,7 @@ export function memoryService(
     });
   }
 
-  async function listLocalBasic(companyId: string, filters: MemoryListRecordsQuery, actor: ActorInfo) {
+  function buildLocalBasicConditions(companyId: string, filters: MemoryListRecordsQuery, actor: ActorInfo) {
     const conditions = [eq(memoryLocalRecords.companyId, companyId)];
     if (filters.bindingId) conditions.push(eq(memoryLocalRecords.bindingId, filters.bindingId));
     if (filters.providerKey) conditions.push(eq(memoryLocalRecords.providerKey, filters.providerKey));
@@ -1181,6 +1182,11 @@ export function memoryService(
       );
     }
 
+    return conditions;
+  }
+
+  async function listLocalBasic(companyId: string, filters: MemoryListRecordsQuery, actor: ActorInfo) {
+    const conditions = buildLocalBasicConditions(companyId, filters, actor);
     const rows = await db
       .select()
       .from(memoryLocalRecords)
@@ -1190,6 +1196,19 @@ export function memoryService(
     return rows.map((row) => mapRecord(row));
   }
 
+  async function countLocalBasic(
+    companyId: string,
+    filters: MemoryListRecordsQuery,
+    actor: ActorInfo,
+  ): Promise<MemoryRecordCountResult> {
+    const conditions = buildLocalBasicConditions(companyId, filters, actor);
+    const [row] = await db
+      .select({ count: sql<number>`count(*)::int` })
+      .from(memoryLocalRecords)
+      .where(and(...conditions));
+    return { count: Number(row?.count ?? 0) };
+  }
+
   async function logOperation(input: {
     id?: string;
     companyId: string;
@@ -2675,6 +2694,9 @@ export function memoryService(
     listRecords: async (companyId: string, filters: MemoryListRecordsQuery, actor: ActorInfo) =>
       listLocalBasic(companyId, filters, actor),
 
+    countRecords: async (companyId: string, filters: MemoryListRecordsQuery, actor: ActorInfo) =>
+      countLocalBasic(companyId, filters, actor),
+
     getRecord: async (companyId: string, recordId: string, actor?: ActorInfo) => {
       const row = await db
         .select()