mcp-server/src/legal_mcp/services/block_writer.py

@@ -19,7 +19,7 @@ from datetime import date
 from uuid import UUID
 
 from legal_mcp import config
-from legal_mcp.services import db, embeddings, claude_session
+from legal_mcp.services import db, embeddings, claude_session, audit
 from legal_mcp.services.lessons import get_content_checklist, get_methodology_summary
 
 logger = logging.getLogger(__name__)
@@ -305,7 +305,9 @@ async def write_block(
     # Template blocks
     if block_id in TEMPLATE_WRITERS:
         content = TEMPLATE_WRITERS[block_id](case, decision)
-        return _build_result(block_id, content, block_cfg)
+        r = _build_result(block_id, content, block_cfg)
+        r["sources"] = {"document_ids": [], "claim_ids": [], "case_law_ids": []}
+        return r
 
     # AI-generated blocks
     prompt_template = BLOCK_PROMPTS.get(block_id)
@@ -318,7 +320,7 @@ async def write_block(
     claims_context = await _build_claims_context(case_id)
     direction_context = _build_direction_context(decision)
     plans_context = await _build_plans_context(case_id)
-    precedents_context = await _build_precedents_context(case_id, block_id)
+    precedents_context, _precedent_case_law_ids = await _build_precedents_context(case_id, block_id)
     style_context = await _build_style_context()
     discussion_context = await _build_previous_blocks_context(case_id, decision)
     appraiser_facts_context = await _build_appraiser_facts_context(case_id)
@@ -391,7 +393,11 @@ async def write_block(
     timeout = claude_session.LONG_TIMEOUT if model_key == "opus" else claude_session.DEFAULT_TIMEOUT
     content = await claude_session.query(prompt, timeout=timeout)
 
-    return _build_result(block_id, content, block_cfg)
+    sources = await _collect_block_sources(case_id, block_id)
+    sources["case_law_ids"] = _precedent_case_law_ids
+    result = _build_result(block_id, content, block_cfg)
+    result["sources"] = sources
+    return result
 
 
 def _build_result(block_id: str, content: str, block_cfg: dict) -> dict:
@@ -408,6 +414,24 @@ def _build_result(block_id: str, content: str, block_cfg: dict) -> dict:
     }
 
 
+async def _collect_block_sources(case_id: UUID, block_id: str) -> dict:
+    """Deterministic source ids available to a block's generation (GAP-19).
+
+    document_ids: case documents matching the block's allowed doc-types.
+    claim_ids: extracted claims for the case. (case_law_ids are captured
+    separately from the precedent search inside write_block.)
+    """
+    allowed = _BLOCK_DOC_TYPES.get(block_id, [])
+    docs = await db.list_documents(case_id)
+    if allowed:
+        docs = [d for d in docs if d.get("doc_type") in allowed]
+    claims = await db.get_claims(case_id)
+    return {
+        "document_ids": [str(d["id"]) for d in docs],
+        "claim_ids": [str(c["id"]) for c in claims],
+    }
+
+
 # ── Context builders ──────────────────────────────────────────────
 
 def _build_case_context(case: dict, decision: dict | None) -> str:
@@ -668,9 +692,10 @@ async def _build_post_hearing_context(case_id: UUID) -> str:
     return "\n".join(lines)
 
 
-async def _build_precedents_context(case_id: UUID, block_id: str) -> str:
+async def _build_precedents_context(case_id: UUID, block_id: str) -> tuple[str, list[str]]:
     """Search for similar precedent paragraphs from other decisions and case law."""
     parts = []
+    case_law_ids: list[str] = []
     try:
         case = await db.get_case(case_id)
         case_number = case.get("case_number", "") if case else ""
@@ -694,7 +719,7 @@ async def _build_precedents_context(case_id: UUID, block_id: str) -> str:
         pool = await db.get_pool()
         async with pool.acquire() as conn:
             caselaw_rows = await conn.fetch(
-                """SELECT cl.case_number, cl.case_name, cl.court, cl.summary, cl.key_quote,
+                """SELECT cl.id, cl.case_number, cl.case_name, cl.court, cl.summary, cl.key_quote,
                           1 - (cle.embedding <=> $1) AS score
                    FROM case_law_embeddings cle
                    JOIN case_law cl ON cl.id = cle.case_law_id
@@ -703,6 +728,7 @@ async def _build_precedents_context(case_id: UUID, block_id: str) -> str:
                 query_emb,
             )
         for r in caselaw_rows[:3]:
+            case_law_ids.append(str(r["id"]))
             text = r["key_quote"] or r["summary"] or ""
             if text:
                 parts.append(
@@ -713,7 +739,7 @@ async def _build_precedents_context(case_id: UUID, block_id: str) -> str:
     except Exception as e:
         logger.warning("Failed to fetch precedents: %s", e)
 
-    return "\n\n".join(parts) if parts else "(אין תקדימים)"
+    return ("\n\n".join(parts) if parts else "(אין תקדימים)"), case_law_ids
 
 
 async def _build_style_context() -> str:
@@ -841,7 +867,7 @@ async def get_block_context(case_id: UUID, block_id: str, instructions: str = ""
     claims_context = await _build_claims_context(case_id)
     direction_context = _build_direction_context(decision)
     plans_context = await _build_plans_context(case_id)
-    precedents_context = await _build_precedents_context(case_id, block_id)
+    precedents_context, _ = await _build_precedents_context(case_id, block_id)
     style_context = await _build_style_context()
     discussion_context = await _build_previous_blocks_context(case_id, decision)
     appraiser_facts_context = await _build_appraiser_facts_context(case_id)
@@ -920,6 +946,7 @@ async def save_block_content(case_id: UUID, block_id: str, content: str) -> dict
     result["model_used"] = "claude-code"
 
     await store_block(UUID(decision["id"]), result)
+    await db.mark_blocks_stale(case_id, False)
 
     # Also write/update the draft file on disk
     await _update_draft_file(case_id, UUID(decision["id"]))
@@ -1049,4 +1076,15 @@ async def write_and_store_block(
 
     result = await write_block(case_id, block_id, instructions)
     await store_block(UUID(decision["id"]), result)
+    await audit.log_action_safe(
+        "write_block", case_id=case_id,
+        details={
+            "decision_id": str(decision["id"]),
+            "block_id": block_id,
+            "model_used": result.get("model_used"),
+            "generation_type": result.get("generation_type"),
+            "sources": result.get("sources", {}),
+        },
+    )
+    await db.mark_blocks_stale(case_id, False)
     return result