fix: citation numbering gaps, flag-click annotation reset, stale client directory

apps/backend/src/taxflow/services/agents/research.py

@@ -1053,6 +1053,51 @@ def _build_trace(
             trace["session"] = session
         return trace
 
+    def _renumber_citations(self, answer: str, citation_map: list[dict]) -> tuple[str, list[dict]]:
+        """Remap the answer's ``[N]`` markers to a contiguous 1..M sequence in
+        first-appearance order, instead of preserving the model's own raw
+        numbering.
+
+        The model is free to cite ``[1]``, ``[2]``, then jump straight to
+        ``[4]`` without ever citing ``[3]`` - nothing about ``citation_map``
+        requires every rendered source to be referenced. Left alone, that
+        produces exactly the bug two accountants independently hit in the
+        round-three audit: a visible source list that skips a number (e.g.
+        ``[1][2][4][5]``), which reads as broken even though the underlying
+        citations are fine. This closes the gap in both the rendered markdown
+        and the returned citations array so ``[1]`` in the text always lines
+        up with the first entry in ``citations``, ``[2]`` the second, etc.
+        """
+        first_seen: list[int] = []
+        seen = set()
+        for group in CITATION_PATTERN.findall(answer):
+            for n_str in group.split(","):
+                n = int(n_str.strip())
+                if 1 <= n <= len(citation_map) and n not in seen:
+                    seen.add(n)
+                    first_seen.append(n)
+
+        if not first_seen:
+            return answer, []
+
+        remap = {old: new for new, old in enumerate(first_seen, start=1)}
+
+        def _replace(match: re.Match) -> str:
+            new_numbers = [
+                remap[n]
+                for n_str in match.group(1).split(",")
+                if (n := int(n_str.strip())) in remap
+            ]
+            return f"[{', '.join(str(n) for n in new_numbers)}]" if new_numbers else ""
+
+        renumbered_answer = CITATION_PATTERN.sub(_replace, answer)
+        # Build citation entries from `first_seen` (the OLD numbers, in
+        # first-appearance order) against the ORIGINAL citation_map - not by
+        # re-parsing renumbered_answer, which now reads [1][2][3] and would
+        # double-remap back onto citation_map's own first entries.
+        citations = self._citation_entries(first_seen, citation_map)
+        return renumbered_answer, citations
+
     def _parse_citations(self, answer: str, citation_map: list[dict]) -> list[dict]:
         """Resolve the answer's ``[N]`` markers against the ``citation_map``
         returned by ``_build_context_string`` (one entry per rendered block, in
@@ -1069,8 +1114,14 @@ def _parse_citations(self, answer: str, citation_map: list[dict]) -> list[dict]:
             for group in CITATION_PATTERN.findall(answer)
             for n in group.split(",")
         }
+        return self._citation_entries(sorted(cited_numbers), citation_map)
+
+    def _citation_entries(self, numbers: list[int], citation_map: list[dict]) -> list[dict]:
+        """Build one citation dict per (valid) number in ``numbers``, in the
+        order given - the shared core of ``_parse_citations`` (ascending raw
+        order) and ``_renumber_citations`` (first-appearance order)."""
         citations = []
-        for n in sorted(cited_numbers):
+        for n in numbers:
             if 1 <= n <= len(citation_map):
                 entry = citation_map[n - 1]
                 chunk = entry["chunks"][0]
@@ -1394,7 +1445,7 @@ async def run(
         model = self._model_for(routed)
 
         answer, stats = await self._generate(question, context, model, steering=steering)
-        citations = self._parse_citations(answer, citation_map)
+        answer, citations = self._renumber_citations(answer, citation_map)
         confidence = self._estimate_confidence(answer, chunks, citations)
 
         # Task C5: bump usage_count for the CITED firm chunks (best-effort) and
@@ -1490,7 +1541,7 @@ async def regenerate_with_feedback(
             question, corrective_context, corrective_model, steering=steering,
             max_tokens=CORRECTIVE_MAX_TOKENS,
         )
-        citations = self._parse_citations(answer, citation_map)
+        answer, citations = self._renumber_citations(answer, citation_map)
         confidence = self._estimate_confidence(answer, chunks, citations)
         re_retrieval = (
             {"fired": True, "reason": "reviewer_flag"} if widened else {"fired": False}

apps/backend/tests/test_parent_expansion.py

@@ -205,6 +205,62 @@ def test_parse_citations_ignores_non_numeric_bracket_content(agent):
     assert [c["citation"] for c in citations] == ["ITAA 1997"]
 
 
+# --- citation renumbering (accountant audit round three, #3) ------------------
+# The model is free to cite [1], [2], then skip straight to [4] - nothing
+# requires every rendered source to be referenced. Two accountants
+# independently hit exactly this: a visible source list that skips a number.
+# _renumber_citations closes the gap so what's on screen is always 1..N.
+def test_renumber_citations_closes_gap_in_answer_text(agent):
+    chunks = [
+        _child("TR 2024/1", "http://a", "first"),
+        _child("TR 2024/2", "http://b", "second"),
+        _child("TR 2024/3", "http://c", "third"),
+        _child("TR 2024/4", "http://d", "fourth"),
+    ]
+    _, citation_map = agent._build_context_string(chunks)
+    answer, citations = agent._renumber_citations(
+        "First point [1]. Second point [2]. Fourth point [4].", citation_map
+    )
+    assert answer == "First point [1]. Second point [2]. Fourth point [3]."
+    assert [c["citation"] for c in citations] == ["TR 2024/1", "TR 2024/2", "TR 2024/4"]
+
+
+def test_renumber_citations_orders_by_first_appearance_not_raw_number(agent):
+    chunks = [
+        _child("TR 2024/1", "http://a", "first"),
+        _child("TR 2024/2", "http://b", "second"),
+        _child("TR 2024/3", "http://c", "third"),
+    ]
+    _, citation_map = agent._build_context_string(chunks)
+    answer, citations = agent._renumber_citations(
+        "Cites the third source first [3], then the first [1].", citation_map
+    )
+    assert answer == "Cites the third source first [1], then the first [2]."
+    assert [c["citation"] for c in citations] == ["TR 2024/3", "TR 2024/1"]
+
+
+def test_renumber_citations_handles_multi_number_bracket(agent):
+    chunks = [
+        _child("TR 2024/1", "http://a", "first"),
+        _child("TR 2024/2", "http://b", "second"),
+        _child("TR 2024/3", "http://c", "third"),
+    ]
+    _, citation_map = agent._build_context_string(chunks)
+    answer, citations = agent._renumber_citations(
+        "First point [1]. Third point [3]. Both again [1, 3].", citation_map
+    )
+    assert answer == "First point [1]. Third point [2]. Both again [1, 2]."
+    assert [c["citation"] for c in citations] == ["TR 2024/1", "TR 2024/3"]
+
+
+def test_renumber_citations_no_citations_is_a_noop(agent):
+    chunks = [_child("TR 2024/1", "http://a", "first")]
+    _, citation_map = agent._build_context_string(chunks)
+    answer, citations = agent._renumber_citations("No sources cited here.", citation_map)
+    assert answer == "No sources cited here."
+    assert citations == []
+
+
 # --- reliance posture flags (business audit P1) -------------------------------
 def test_parse_citations_flags_historical_and_engagement_memo(agent):
     """Historical/superseded and engagement-memo citations must carry their own

apps/dashboard/app/dashboard/workspace/page.tsx

@@ -1,5 +1,6 @@
 "use client";
 
+import { useState } from "react";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { WorkspaceClientsTable } from "@/components/WorkspaceClientsTable";
 import { DocumentTemplatesPanel } from "@/components/DocumentTemplatesPanel";
@@ -14,15 +15,31 @@ import DocumentsPage from "@/app/dashboard/documents/page";
 // instructions for every document type, in one place instead of only
 // reachable per-document-type from inside the Ask TaxFlow save flow).
 export default function WorkspacePage() {
+  // Radix keeps every TabsContent mounted once rendered, so
+  // WorkspaceClientsTable's own mount-time fetch only ever reflected
+  // whatever was true the first time this page loaded - asking a fresh
+  // question in Ask TaxFlow, then switching back to this already-mounted
+  // Clients tab, showed stale (often zero) counts (accountant audit round
+  // three, Priya/Michael). Tabs is controlled here so a fresh key is handed
+  // to the table each time the Clients tab is actually selected, forcing its
+  // fetch effect to re-run instead of trusting a fetch from whenever the
+  // page happened to first load.
+  const [clientsTabVisits, setClientsTabVisits] = useState(0);
+
   return (
-    <Tabs defaultValue="clients">
+    <Tabs
+      defaultValue="clients"
+      onValueChange={(value) => {
+        if (value === "clients") setClientsTabVisits((n) => n + 1);
+      }}
+    >
       <TabsList>
         <TabsTrigger value="clients">Clients</TabsTrigger>
         <TabsTrigger value="documents">Documents</TabsTrigger>
         <TabsTrigger value="templates">Templates</TabsTrigger>
       </TabsList>
       <TabsContent value="clients">
-        <WorkspaceClientsTable />
+        <WorkspaceClientsTable key={clientsTabVisits} />
       </TabsContent>
       <TabsContent value="documents">
         <DocumentsPage />

apps/dashboard/components/AnnotatableMarkdown.tsx

@@ -263,8 +263,20 @@ export const AnnotatableMarkdown = forwardRef<AnnotatableMarkdownHandle, {
     }));
   }, [annotations, serverHash]);
 
-  const visibleThreads = threads.filter((t) =>
-    showResolved ? t.root.resolved_at != null : t.root.resolved_at == null
+  // Memoized deliberately (not a plain .filter()): RecogitoLayer's
+  // setAnnotations effect keys off this array's identity (deps
+  // [anno, threads, verifyAnchors]) and calls Recogito with `replace: true`,
+  // which tears down and rebuilds the ENTIRE highlight layer. An unmemoized
+  // filter here produces a new array reference on every AnnotatableMarkdown
+  // render, so that effect - and the full highlight rebuild - fired far more
+  // often than `threads`/`showResolved` actually changed, including renders
+  // that happen to land mid-click. A rebuild racing the click's own
+  // select-and-open-popup sequence is what made clicking a flagged claim
+  // clear its underline instead of opening the detail popup (accountant
+  // audit round three, #1).
+  const visibleThreads = useMemo(
+    () => threads.filter((t) => (showResolved ? t.root.resolved_at != null : t.root.resolved_at == null)),
+    [threads, showResolved]
   );
   const openCount = threads.filter((t) => t.root.resolved_at == null).length;
   const resolvedCount = threads.length - openCount;