diff --git a/mcp-server/src/legal_mcp/services/case_citation_verification.py b/mcp-server/src/legal_mcp/services/case_citation_verification.py index 88a1084..bbd3fc3 100644 --- a/mcp-server/src/legal_mcp/services/case_citation_verification.py +++ b/mcp-server/src/legal_mcp/services/case_citation_verification.py @@ -20,6 +20,7 @@ from __future__ import annotations import asyncio import logging +import time from uuid import UUID from legal_mcp.services import ( @@ -74,7 +75,92 @@ def _passes_floor(hit: dict) -> bool: _RETRIEVAL_BUDGET_S = 22.0 -async def build_view(case_number: str) -> dict: +#: Cached views, keyed by case number → (fingerprint, expires_at|None, view). +#: Assembling a view costs 20+ seconds of Voyage embeds and vector searches, and +#: the result changes only when its inputs do. Keyed by a FINGERPRINT of those +#: inputs rather than a guessed TTL: a chair who verifies a precedent must see it +#: reflected on the next load, not after a timer expires. +_view_cache: dict[str, tuple[str, float | None, dict]] = {} + +#: One lock per case so two concurrent loads of the same page don't both compute. +#: Reload-during-a-slow-load is how the original outage compounded into Postgres +#: deadlocks; the second reader now waits for the first and gets its result. +_view_locks: dict[str, asyncio.Lock] = {} + +#: A partial view (retrieval budget exhausted) is still worth caching — otherwise +#: every reload re-pays 22 seconds — but only briefly, so the missing suggestions +#: get another chance while the corpus is quieter. +_PARTIAL_TTL_S = 90.0 + + +async def _inputs_fingerprint(case_id: UUID) -> str: + """Cheap signature of everything the view is derived from. + + Changes when: arguments are re-aggregated, the chair attaches or verifies a + precedent, or the corpus grows (new rulings change the suggestions). One + round-trip — the point is to be far cheaper than the 20s it guards. + """ + pool = await db.get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + """ + SELECT + (SELECT count(*)::text || ':' || coalesce(max(updated_at)::text, '-') + FROM legal_arguments WHERE case_id = $1) AS args, + (SELECT count(*)::text || ':' || coalesce(max(updated_at)::text, '-') + FROM case_precedents WHERE case_id = $1) AS attached, + (SELECT count(*)::text || ':' || coalesce(max(created_at)::text, '-') + FROM case_law WHERE searchable = true) AS corpus + """, + case_id, + ) + return f"{row['args']}|{row['attached']}|{row['corpus']}" + + +def invalidate(case_number: str = "") -> None: + """Drop cached views — one case, or all when called with no argument. + + Callers that mutate the inputs may use this for immediacy, but correctness + does not depend on it: the fingerprint check catches any change on the next + read regardless of whether anyone remembered to invalidate (G1 — the guard + lives at the source of truth, not in every writer's memory). + """ + if case_number: + _view_cache.pop(case_number, None) + else: + _view_cache.clear() + + +async def build_view(case_number: str, *, use_cache: bool = True) -> dict: + """Assemble the citation-verification view, served from cache when unchanged.""" + if not use_cache: + return await _build_view_uncached(case_number) + + lock = _view_locks.setdefault(case_number, asyncio.Lock()) + async with lock: + case = await db.get_case_by_number(case_number) + if not case: + return {"status": "case_not_found", "case_number": case_number, + "arguments": []} + cid = case["id"] + fp = await _inputs_fingerprint(UUID(cid) if isinstance(cid, str) else cid) + + hit = _view_cache.get(case_number) + if hit is not None: + cached_fp, expires_at, view = hit + fresh = cached_fp == fp and (expires_at is None or time.monotonic() < expires_at) + if fresh: + return {**view, "cached": True} + + view = await _build_view_uncached(case_number) + if view.get("status") == "ok": + expires_at = (None if view.get("retrieval_complete", True) + else time.monotonic() + _PARTIAL_TTL_S) + _view_cache[case_number] = (fp, expires_at, view) + return {**view, "cached": False} + + +async def _build_view_uncached(case_number: str) -> dict: case = await db.get_case_by_number(case_number) if not case: return {"status": "case_not_found", "case_number": case_number, "arguments": []} diff --git a/mcp-server/tests/test_citation_view_cache.py b/mcp-server/tests/test_citation_view_cache.py new file mode 100644 index 0000000..26db495 --- /dev/null +++ b/mcp-server/tests/test_citation_view_cache.py @@ -0,0 +1,125 @@ +"""The citation view is cached by input fingerprint, not by a guessed TTL. + +Assembling it costs 20+ seconds of Voyage embeds and vector searches. A TTL +would force a choice between staleness and cost: a chair who verifies a +precedent must see it on the very next load, and a timer cannot promise that. +Keying on a fingerprint of the inputs (arguments · attachments · corpus) makes +freshness a property of the data rather than of the clock. +""" + +import asyncio + +import pytest + +from legal_mcp.services import case_citation_verification as ccv + + +@pytest.fixture(autouse=True) +def _clean_cache(): + ccv.invalidate() + yield + ccv.invalidate() + + +def _stub(monkeypatch, fingerprint: str, calls: list, *, complete: bool = True): + """Point build_view at a fake case, a controllable fingerprint, and a + counted assembler, so these tests never touch the corpus.""" + async def fake_case(_cn): + return {"id": "00000000-0000-0000-0000-000000000001"} + + async def fake_fp(_cid): + return fingerprint + + async def fake_build(cn): + calls.append(cn) + return {"status": "ok", "case_number": cn, "arguments": [], + "retrieval_complete": complete, "summary": {}} + + monkeypatch.setattr(ccv.db, "get_case_by_number", fake_case) + monkeypatch.setattr(ccv, "_inputs_fingerprint", fake_fp) + monkeypatch.setattr(ccv, "_build_view_uncached", fake_build) + + +@pytest.mark.asyncio +async def test_second_load_is_served_from_cache(monkeypatch): + calls: list = [] + _stub(monkeypatch, "fp-1", calls) + first = await ccv.build_view("1069-04-26") + second = await ccv.build_view("1069-04-26") + assert first["cached"] is False + assert second["cached"] is True + assert len(calls) == 1, "the expensive assembly must run once" + + +@pytest.mark.asyncio +async def test_changed_inputs_invalidate_without_anyone_calling_invalidate(monkeypatch): + """The chair verifies a precedent → fingerprint moves → next load rebuilds. + Correctness must not depend on a writer remembering to clear the cache.""" + calls: list = [] + _stub(monkeypatch, "fp-before", calls) + await ccv.build_view("1069-04-26") + + _stub(monkeypatch, "fp-after", calls) # e.g. a new case_precedents row + again = await ccv.build_view("1069-04-26") + assert again["cached"] is False + assert len(calls) == 2 + + +@pytest.mark.asyncio +async def test_concurrent_loads_do_not_stampede(monkeypatch): + """Reload-during-a-slow-load is how the original outage compounded into + Postgres deadlocks. The later readers must wait, not pile on.""" + calls: list = [] + + async def fake_case(_cn): + return {"id": "00000000-0000-0000-0000-000000000001"} + + async def fake_fp(_cid): + return "fp-1" + + async def slow_build(cn): + calls.append(cn) + await asyncio.sleep(0.05) + return {"status": "ok", "case_number": cn, "arguments": [], + "retrieval_complete": True, "summary": {}} + + monkeypatch.setattr(ccv.db, "get_case_by_number", fake_case) + monkeypatch.setattr(ccv, "_inputs_fingerprint", fake_fp) + monkeypatch.setattr(ccv, "_build_view_uncached", slow_build) + + await asyncio.gather(*(ccv.build_view("1069-04-26") for _ in range(6))) + assert len(calls) == 1, "six concurrent loads, one assembly" + + +@pytest.mark.asyncio +async def test_partial_view_expires_so_missing_suggestions_get_retried(monkeypatch): + """A budget-exhausted view is cached to stop reload storms, but it must not + become permanent — the arguments with no suggestions deserve another try.""" + calls: list = [] + _stub(monkeypatch, "fp-1", calls, complete=False) + await ccv.build_view("1069-04-26") + _fp, expires_at, _view = ccv._view_cache["1069-04-26"] + assert expires_at is not None, "a partial view must carry an expiry" + + monkeypatch.setattr(ccv.time, "monotonic", lambda: expires_at + 1) + after = await ccv.build_view("1069-04-26") + assert after["cached"] is False + assert len(calls) == 2 + + +@pytest.mark.asyncio +async def test_complete_view_does_not_expire_on_a_timer(monkeypatch): + calls: list = [] + _stub(monkeypatch, "fp-1", calls, complete=True) + await ccv.build_view("1069-04-26") + _fp, expires_at, _view = ccv._view_cache["1069-04-26"] + assert expires_at is None, "freshness comes from the fingerprint, not a clock" + + +@pytest.mark.asyncio +async def test_use_cache_false_always_rebuilds(monkeypatch): + calls: list = [] + _stub(monkeypatch, "fp-1", calls) + await ccv.build_view("1069-04-26") + await ccv.build_view("1069-04-26", use_cache=False) + assert len(calls) == 2