Fix document annotation live updates

ui/src/context/LiveUpdatesProvider.test.ts

@@ -139,6 +139,9 @@ describe("LiveUpdatesProvider issue invalidation", () => {
     expect(invalidations).toContainEqual({
       queryKey: queryKeys.issues.documentRevisions("issue-1", "plan"),
     });
+    expect(invalidations).toContainEqual({
+      queryKey: ["issues", "document-annotations", "issue-1", "plan"],
+    });
     expect(invalidations).toContainEqual({
       queryKey: queryKeys.issues.documents("PAP-9403"),
     });
@@ -148,6 +151,9 @@ describe("LiveUpdatesProvider issue invalidation", () => {
     expect(invalidations).toContainEqual({
       queryKey: queryKeys.issues.documentRevisions("PAP-9403", "plan"),
     });
+    expect(invalidations).toContainEqual({
+      queryKey: ["issues", "document-annotations", "PAP-9403", "plan"],
+    });
     expect(invalidations).not.toContainEqual({
       queryKey: queryKeys.issues.documents("issue-1"),
       refetchType: "inactive",
@@ -186,6 +192,48 @@ describe("LiveUpdatesProvider issue invalidation", () => {
     expect(invalidations).toContainEqual({
       queryKey: ["issues", "document-revisions", "issue-1"],
     });
+    expect(invalidations).toContainEqual({
+      queryKey: ["issues", "document-annotations", "issue-1"],
+    });
+  });
+
+  it("refreshes document annotation caches when annotation activity arrives", () => {
+    const invalidations: unknown[] = [];
+    const queryClient = {
+      invalidateQueries: (input: unknown) => {
+        invalidations.push(input);
+      },
+      getQueryData: () => undefined,
+    };
+
+    __liveUpdatesTestUtils.invalidateActivityQueries(
+      queryClient as never,
+      "company-1",
+      {
+        entityType: "issue",
+        entityId: "issue-1",
+        action: "issue.document_annotation_comment_added",
+        actorType: "user",
+        actorId: "user-2",
+        details: {
+          identifier: "PAP-9403",
+          documentKey: "plan",
+          threadId: "thread-1",
+          commentId: "comment-1",
+        },
+      },
+      { userId: "user-1", agentId: null },
+    );
+
+    expect(invalidations).toContainEqual({
+      queryKey: ["issues", "document-annotations", "issue-1", "plan"],
+    });
+    expect(invalidations).toContainEqual({
+      queryKey: ["issues", "document-annotations", "PAP-9403", "plan"],
+    });
+    expect(invalidations).not.toContainEqual({
+      queryKey: queryKeys.issues.documents("issue-1"),
+    });
   });
 
   it("keeps self-authored comment events from refetching the active issue tree", () => {

ui/src/context/LiveUpdatesProvider.tsx

@@ -415,6 +415,13 @@ const ISSUE_DOCUMENT_ACTIVITY_ACTIONS = new Set([
   "issue.document_restored",
   "issue.document_deleted",
 ]);
+const ISSUE_DOCUMENT_ANNOTATION_ACTIVITY_ACTIONS = new Set([
+  "issue.document_annotation_thread_created",
+  "issue.document_annotation_comment_added",
+  "issue.document_annotation_thread_resolved",
+  "issue.document_annotation_thread_reopened",
+  "issue.document_annotation_remapped",
+]);
 const AGENT_TOAST_STATUSES = new Set(["error"]);
 const RUN_TOAST_STATUSES = new Set(["failed", "timed_out", "cancelled"]);
 
@@ -710,6 +717,18 @@ function invalidateActivityQueries(
             queryClient.invalidateQueries({ queryKey: ["issues", "document-revisions", ref], ...invalidationOptions });
           }
         }
+        if (
+          action &&
+          (ISSUE_DOCUMENT_ACTIVITY_ACTIONS.has(action) || ISSUE_DOCUMENT_ANNOTATION_ACTIVITY_ACTIONS.has(action))
+        ) {
+          const documentKey = readString(details?.key) ?? readString(details?.documentKey);
+          queryClient.invalidateQueries({
+            queryKey: documentKey
+              ? ["issues", "document-annotations", ref, documentKey]
+              : ["issues", "document-annotations", ref],
+            ...invalidationOptions,
+          });
+        }
         if (action?.startsWith("issue.thread_interaction_")) {
           queryClient.invalidateQueries({ queryKey: queryKeys.issues.interactions(ref), ...invalidationOptions });
         }