Agentic v2: flat tool args, conversation history, date-aware rehab math, 429 retry, stderr telemetry, auth fail-closed, Stop/export

README.md

@@ -12,14 +12,15 @@ client can still ask Canadian legal-research questions.
 
 ## How it works -- thin client
 
-This app holds **no copy of the legal corpus**. For each question it:
-
-1. calls the deployed **CanLex MCP server** (`canlex_search_legislation`) to
-   retrieve the cited source passages;
-2. sends those passages and the question to **Google Gemini Flash**, which
-   composes a grounded, cited answer following CanLex's own answering
-   instructions;
-3. displays the answer, with the retrieved passages shown for review.
+This app holds **no copy of the legal corpus**. For each question, Google
+Gemini runs an **agentic loop** over the deployed CanLex MCP server's seven
+tools -- searching, fetching sections, running the three-step IRPA s. 36
+screening chain, and checking case citations -- until it has enough material,
+then composes a grounded, cited answer following CanLex's own answering
+instructions. The answer streams live, with a per-question trace, a trust
+badge, and every retrieved passage shown for review. Follow-up questions keep
+the conversation's context; the answer and sources can be exported as a
+dated research memo.
 
 Because retrieval stays on the MCP server, a corpus or retrieval change is
 deployed once (to the MCP Space) and both the MCP connector and this website
@@ -39,10 +40,11 @@ Optional overrides:
 | Name | Default |
 |------|---------|
 | `CANLEX_MCP_URL` | `https://beemer0-canlex.hf.space/mcp` |
-| `CANLEX_GEMINI_MODEL` | `gemini-2.5-flash` |
+| `CANLEX_GEMINI_MODEL` | `gemini-2.5-flash` (this deployment sets `gemini-2.5-pro`) |
+| `CANLEX_THINKING_BUDGET` | `4096` |
 
-If `CANLEX_WEB_AUTH` is unset the app falls back to an insecure default login
-(`canlex` / `canlex`) and logs a warning -- set the secret before real use.
+If `CANLEX_WEB_AUTH` is unset the app **refuses to start** -- there is no
+default login. Set the secret first.
 
 ## Make the Space private
 

app.py

@@ -3,20 +3,23 @@
 
 A thin client that gives a non-Claude user roughly the same experience as Claude
 with the CanLex MCP server. For each question it opens one streamable-HTTP
-session against the deployed CanLex MCP, declares the four CanLex tools to
+session against the deployed CanLex MCP, declares the seven CanLex tools to
 Google Gemini, and lets the model agentically iterate -- searching, fetching
-sections, and looking up case-law citations -- until it decides it has enough
-material to compose a grounded answer.
+sections, running the s. 36 screening chain, and looking up case-law citations
+-- until it decides it has enough material to compose a grounded answer.
 
 All configuration comes from environment variables, set as Hugging Face Space
 secrets. Run locally with:  python app.py
 """
 import asyncio
+import datetime as _dt
 import json
 import os
 import queue
 import sys
+import tempfile
 import threading
+import time
 import urllib.error
 import urllib.request
 from datetime import timedelta
@@ -34,13 +37,24 @@ MCP_URL = os.environ.get(
     "CANLEX_MCP_URL", "https://beemer0-canlex.hf.space/mcp").strip()
 
 # Google Gemini -- the free-tier key is supplied as the GEMINI_API_KEY secret.
-GEMINI_MODEL = os.environ.get("CANLEX_GEMINI_MODEL", "gemini-2.5-pro").strip()
+# Default is flash (fast, generous free quota); the deployed Space overrides to
+# gemini-2.5-pro via the CANLEX_GEMINI_MODEL variable for answer quality.
+GEMINI_MODEL = os.environ.get("CANLEX_GEMINI_MODEL", "gemini-2.5-flash").strip()
 GEMINI_ENDPOINT = ("https://generativelanguage.googleapis.com/v1beta/models/"
                    f"{GEMINI_MODEL}:generateContent")
 
 MAX_OUTPUT_TOKENS = 8192   # generous -- covers Gemini 2.5 thinking plus the answer
 MAX_TOOL_ITERATIONS = 8    # loop guard for the agent
 REQUEST_TIMEOUT = 180      # seconds, applied to the MCP and Gemini calls
+# Bound Gemini 2.5's internal thinking so it cannot eat the whole output
+# budget on the free tier; env-overridable for experimentation.
+THINKING_BUDGET = int(os.environ.get("CANLEX_THINKING_BUDGET", "4096"))
+# A single tool result re-enters the prompt on EVERY later iteration, so an
+# unbounded top_k=20 search x 8 iterations balloons prompt tokens against the
+# free-tier TPM cap. Cap what is sent back to the model; the sources panel
+# still shows the full text.
+TOOL_OUTPUT_CAP = 15_000
+_GEMINI_RETRIES = 3        # bounded retry on 429/500/503, with backoff
 
 
 def _load_auth() -> list[tuple[str, str]]:
@@ -54,11 +68,13 @@ def _load_auth() -> list[tuple[str, str]]:
         if user and password:
             creds.append((user, password))
     if not creds:
-        print("WARNING: CANLEX_WEB_AUTH is not set; using the insecure default "
-              "login 'canlex' / 'canlex'. Set CANLEX_WEB_AUTH as a Space secret "
-              "(one 'username:password' per line) before sharing this app.",
-              file=sys.stderr)
-        creds = [("canlex", "canlex")]
+        # Fail closed: a guessable default on a public Space is a live door to
+        # the operator's Gemini quota. Refuse to start instead.
+        print("FATAL: CANLEX_WEB_AUTH is not set (or contains no valid "
+              "'username:password' line). Set it as a Space secret -- one "
+              "'username:password' per line -- and restart. Refusing to start "
+              "with default credentials.", file=sys.stderr)
+        sys.exit(1)
     return creds
 
 
@@ -67,10 +83,9 @@ AUTH = _load_auth()
 
 # --- Tool declarations (Gemini function-calling schema) -----------------------
 
-# The four CanLex MCP tools, declared so Gemini can call them. Three of the four
-# wrap their arguments inside a single 'params' object on the server side; the
-# Gemini schema is kept flat for the model's convenience and re-wrapped at the
-# MCP edge in _run_tool.
+# The seven CanLex MCP tools, declared so Gemini can call them. The MCP server
+# publishes flat argument schemas, so these declarations mirror the server
+# signatures one-to-one and _run_tool passes arguments straight through.
 TOOL_DECLARATIONS = [
     {
         "name": "canlex_search_legislation",
@@ -107,12 +122,17 @@ TOOL_DECLARATIONS = [
                 },
                 "doc_type": {
                     "type": "string",
+                    "enum": ["legislation", "memorandum", "agreement",
+                             "directive", "caselaw", "delegation",
+                             "benefits", "commentary"],
                     "description": (
                         "Optional. Restrict to one source type: 'legislation', "
                         "'memorandum' (CBSA D-Memoranda), 'agreement' "
                         "(collective agreements), 'directive' (NJC), "
-                        "'caselaw' (court and tribunal decisions), or "
-                        "'delegation' (IRPA/IRPR delegation and designation)."
+                        "'caselaw' (court and tribunal decisions), "
+                        "'delegation' (IRPA/IRPR delegation and designation), "
+                        "'benefits' (plan booklets), or 'commentary' "
+                        "(CanLex's curated US-record analysis)."
                     ),
                 },
             },
@@ -273,21 +293,23 @@ TOOL_DECLARATIONS = [
     },
 ]
 
-# Tools whose MCP signature wraps arguments under a single 'params' object.
-# canlex_list_acts takes none and is handled separately in _run_tool.
-_PARAMS_WRAPPED = {"canlex_search_legislation", "canlex_get_section",
-                   "canlex_case", "canlex_us_disposition",
-                   "canlex_us_equivalency", "canlex_rehabilitation"}
+_KNOWN_TOOLS = {d["name"] for d in TOOL_DECLARATIONS}
 
 
 # --- System prompt ------------------------------------------------------------
 
 SYSTEM_INSTRUCTION = """\
-You are CanLex Web, a Canadian legal-research assistant. A member of the public \
-has asked the legal question shown below through a web form. Answer it by \
-agentically using the CanLex tools to retrieve primary sources, then \
-compose a clear, well-organised answer grounded entirely in what those tools \
-return.
+You are CanLex Web, a Canadian legal-research assistant. Your reader is a \
+working border-services professional: lead with the operative statutory \
+language and the decision points, cite precisely, and skip lay simplification \
+-- though they may not be a lawyer, so define genuinely technical terms once. \
+Answer the question by agentically using the CanLex tools to retrieve primary \
+sources, then compose a clear, well-organised answer grounded entirely in \
+what those tools return.
+
+Today's date is {today}. Use it for any elapsed-time computation (e.g. years \
+since a sentence was completed, for rehabilitation) instead of assuming a \
+date.
 
 Tool-use guidance:
 - Start with canlex_search_legislation on the user's question. Read the \
@@ -332,20 +354,34 @@ class _AgentError(RuntimeError):
     """Surfaced to the UI; the message text is shown verbatim."""
 
 
-def _gemini_request_body(contents: list[dict]) -> dict:
+class _Stopped(RuntimeError):
+    """Raised inside the agent loop when the user pressed Stop."""
+
+
+# Per-question run stats for the structured stderr log line -- the app serves
+# one user and runs one question at a time, so a module-level dict is fine.
+_RUN_STATS: dict = {}
+
+
+def _gemini_request_body(contents: list[dict], tool_mode: str = "AUTO") -> dict:
     """The JSON body sent to Gemini -- identical between the streaming and the
     non-streaming endpoints. Tool declarations turn on function calling; the
     safety filters are relaxed because legal research routinely discusses
     crime, weapons and the like, and the high-threshold defaults spuriously
-    block legitimate legal text."""
+    block legitimate legal text. tool_mode='NONE' forces a text-only turn
+    (used for the terminal budget-exhausted answer). The system instruction
+    is stamped with today's date at request time -- rehabilitation math
+    depends on it."""
     return {
-        "systemInstruction": {"parts": [{"text": SYSTEM_INSTRUCTION}]},
+        "systemInstruction": {"parts": [{"text": SYSTEM_INSTRUCTION.format(
+            today=_dt.date.today().isoformat())}]},
         "contents": contents,
         "tools": [{"functionDeclarations": TOOL_DECLARATIONS}],
-        "toolConfig": {"functionCallingConfig": {"mode": "AUTO"}},
+        "toolConfig": {"functionCallingConfig": {"mode": tool_mode}},
         "generationConfig": {
             "temperature": 0.2,
             "maxOutputTokens": MAX_OUTPUT_TOKENS,
+            "thinkingConfig": {"thinkingBudget": THINKING_BUDGET},
         },
         "safetySettings": [
             {"category": c, "threshold": "BLOCK_ONLY_HIGH"}
@@ -365,39 +401,71 @@ _STREAM_ENDPOINT = GEMINI_ENDPOINT.replace(
     ":generateContent", ":streamGenerateContent") + "?alt=sse"
 
 
-async def _gemini_stream(api_key: str, contents: list[dict]):
+def _retry_delay(code: int, body_text: str, attempt: int) -> float:
+    """Backoff for a retryable Gemini error. A 429 body may carry a
+    google.rpc.RetryInfo detail with an explicit retryDelay ('17s'); honour
+    it when present, else exponential backoff."""
+    if code == 429:
+        try:
+            for detail in json.loads(body_text)["error"]["details"]:
+                delay = detail.get("retryDelay")
+                if delay:
+                    return min(float(str(delay).rstrip("s")), 60.0)
+        except Exception:
+            pass
+    return 2.0 * (2 ** attempt)
+
+
+async def _gemini_stream(api_key: str, contents: list[dict],
+                         tool_mode: str = "AUTO"):
     """Async generator over Gemini's streaming response.
 
     Yields dicts of one of three shapes:
       {"type": "text_delta", "text": str}        -- a partial answer fragment
       {"type": "function_call", "call": dict}    -- a complete tool call
       {"type": "finish", "reason": str|None,     -- end of stream; `parts` is
-       "parts": list[dict]}                          the whole assistant turn
+       "parts": list[dict], "usage": dict}           the whole assistant turn
+
+    429/500/503 responses are retried up to _GEMINI_RETRIES times with
+    backoff -- on the free tier, mid-question 429s are the most likely
+    failure mode, and an abort discards every tool result already gathered.
     """
-    body = _gemini_request_body(contents)
-    request = urllib.request.Request(
-        _STREAM_ENDPOINT,
-        data=json.dumps(body).encode("utf-8"),
-        headers={"Content-Type": "application/json",
-                 "x-goog-api-key": api_key,
-                 "Accept": "text/event-stream"},
-        method="POST",
-    )
-    try:
-        # `timeout` is a kwarg of urlopen; passing it positionally to
-        # asyncio.to_thread would forward it as `data` (POST body) and break
-        # the request.
-        response = await asyncio.to_thread(
-            lambda: urllib.request.urlopen(request, timeout=REQUEST_TIMEOUT))
-    except urllib.error.HTTPError as exc:
-        detail = await asyncio.to_thread(exc.read)
-        text = detail.decode("utf-8", "replace")[:600]
-        raise _AgentError(f"Gemini API returned HTTP {exc.code}: {text}") from None
-    except urllib.error.URLError as exc:
-        raise _AgentError(f"Could not reach the Gemini API: {exc.reason}") from None
+    body = _gemini_request_body(contents, tool_mode=tool_mode)
+    payload = json.dumps(body).encode("utf-8")
+    response = None
+    for attempt in range(_GEMINI_RETRIES + 1):
+        request = urllib.request.Request(
+            _STREAM_ENDPOINT,
+            data=payload,
+            headers={"Content-Type": "application/json",
+                     "x-goog-api-key": api_key,
+                     "Accept": "text/event-stream"},
+            method="POST",
+        )
+        try:
+            # `timeout` is a kwarg of urlopen; passing it positionally to
+            # asyncio.to_thread would forward it as `data` (POST body) and
+            # break the request.
+            response = await asyncio.to_thread(
+                lambda: urllib.request.urlopen(request, timeout=REQUEST_TIMEOUT))
+            break
+        except urllib.error.HTTPError as exc:
+            detail = await asyncio.to_thread(exc.read)
+            text = detail.decode("utf-8", "replace")[:600]
+            if exc.code in (429, 500, 503) and attempt < _GEMINI_RETRIES:
+                await asyncio.sleep(_retry_delay(exc.code, text, attempt))
+                continue
+            raise _AgentError(
+                f"Gemini API returned HTTP {exc.code}"
+                + (" after retries" if attempt else "") + f": {text}") from None
+        except urllib.error.URLError as exc:
+            raise _AgentError(
+                f"Could not reach the Gemini API: {exc.reason}") from None
 
     accumulated_parts: list[dict] = []
     finish_reason = None
+    usage: dict = {}
+    block_reason = None
     try:
         while True:
             raw = await asyncio.to_thread(response.readline)
@@ -410,6 +478,11 @@ async def _gemini_stream(api_key: str, contents: list[dict]):
                 chunk = json.loads(line[6:])
             except ValueError:
                 continue
+            if chunk.get("usageMetadata"):
+                usage = chunk["usageMetadata"]
+            fb = chunk.get("promptFeedback") or {}
+            if fb.get("blockReason"):
+                block_reason = fb["blockReason"]
             candidate = (chunk.get("candidates") or [{}])[0]
             for part in (candidate.get("content") or {}).get("parts") or []:
                 accumulated_parts.append(part)
@@ -425,22 +498,29 @@ async def _gemini_stream(api_key: str, contents: list[dict]):
                 finish_reason = candidate["finishReason"]
     finally:
         await asyncio.to_thread(response.close)
-    yield {"type": "finish", "reason": finish_reason, "parts": accumulated_parts}
+    if block_reason and not accumulated_parts:
+        # A safety block returns no candidates, only promptFeedback -- name it
+        # rather than surfacing the confusing 'empty answer' error downstream.
+        raise _AgentError(
+            f"Gemini blocked this request (blockReason: {block_reason}). "
+            "Rephrase the question; legal fact patterns occasionally trip "
+            "the safety filters despite the relaxed thresholds.")
+    yield {"type": "finish", "reason": finish_reason,
+           "parts": accumulated_parts, "usage": usage}
 
 
 async def _run_tool(session: ClientSession, name: str, args: dict) -> str:
     """Execute a Gemini function call against the MCP, returning text output."""
-    if name == "canlex_list_acts":
-        mcp_args: dict = {}
-    elif name in _PARAMS_WRAPPED:
-        # The MCP server's tools accept their schema as a single 'params' object.
-        mcp_args = {"params": args or {}}
-    else:
+    if name not in _KNOWN_TOOLS:
         return f"(unknown tool '{name}')"
+    t0 = time.monotonic()
     try:
-        result = await session.call_tool(name, mcp_args)
+        result = await session.call_tool(name, args or {})
     except Exception as exc:                        # MCP transport errors
         return f"(tool '{name}' failed: {type(exc).__name__}: {exc})"
+    finally:
+        _RUN_STATS.setdefault("tools", []).append(
+            (name, int((time.monotonic() - t0) * 1000)))
     text = "\n".join(
         block.text for block in result.content
         if getattr(block, "type", None) == "text" and getattr(block, "text", None)
@@ -512,11 +592,16 @@ def _format_sources(tool_log: list[tuple[str, dict, str]]) -> str:
     return "\n\n---\n\n".join(blocks)
 
 
-async def _agentic_answer(question: str):
+async def _agentic_answer(question: str,
+                          history: list[tuple[str, str]] | None = None,
+                          stop_event: threading.Event | None = None):
     """Run the Gemini-driven agentic loop against a single MCP session.
 
     Yields tuples of (status, answer_md, sources_md). The final yield carries
     the composed answer; earlier yields are progress updates the UI can show.
+    `history` is the session's prior (question, answer) pairs -- threaded into
+    the model's context so follow-ups ('same facts but Texas') work without
+    re-typing, at the cost of only the final answers, not old tool outputs.
     """
     api_key = os.environ.get("GEMINI_API_KEY", "").strip()
     if not api_key:
@@ -534,9 +619,11 @@ async def _agentic_answer(question: str):
         async with ClientSession(read, write) as session:
             await session.initialize()
 
-            contents: list[dict] = [
-                {"role": "user", "parts": [{"text": question}]}
-            ]
+            contents: list[dict] = []
+            for past_q, past_a in (history or []):
+                contents.append({"role": "user", "parts": [{"text": past_q}]})
+                contents.append({"role": "model", "parts": [{"text": past_a}]})
+            contents.append({"role": "user", "parts": [{"text": question}]})
             tool_log: list[tuple[str, dict, str]] = []
             trace: list[str] = []
 
@@ -553,7 +640,12 @@ async def _agentic_answer(question: str):
                     lines.append("- _Thinking..._")
                 return "\n".join(lines) if lines else ""
 
+            def _check_stop():
+                if stop_event is not None and stop_event.is_set():
+                    raise _Stopped()
+
             for step in range(MAX_TOOL_ITERATIONS):
+                _check_stop()
                 yield status_md(), answer_buf, sources_md
 
                 # Stream Gemini's next turn. Stream text deltas to the answer
@@ -566,6 +658,7 @@ async def _agentic_answer(question: str):
                 optimistic = True
 
                 async for chunk in _gemini_stream(api_key, contents):
+                    _check_stop()
                     if chunk["type"] == "text_delta" and optimistic:
                         turn_text += chunk["text"]
                         yield (status_md(),
@@ -585,6 +678,10 @@ async def _agentic_answer(question: str):
                         # Capture any text-only finish reason so the caller can
                         # surface a useful error for an empty answer.
                         finish_reason = chunk.get("reason")
+                        _RUN_STATS["finish"] = finish_reason
+                        _RUN_STATS["usage"] = (chunk.get("usage")
+                                               or _RUN_STATS.get("usage", {}))
+                        _RUN_STATS["iterations"] = step + 1
 
                 contents.append({"role": "model", "parts": turn_parts})
 
@@ -610,25 +707,37 @@ async def _agentic_answer(question: str):
                         snippet = snippet[:137].rstrip() + "..."
                     trace.append(f"_{snippet}_")
 
-                # Execute every function call in this turn, then send the
-                # responses back as a single 'user' message.
-                function_responses = []
+                # Execute every function call in this turn concurrently (a
+                # parallel-call turn is common in the screening flow), then
+                # send the responses back as a single 'user' message.
                 for call in turn_calls:
+                    trace.append(_summarize_call(call.get("name", ""),
+                                                 call.get("args") or {}))
+                yield status_md(), answer_buf, sources_md
+
+                outputs = await asyncio.gather(*(
+                    _run_tool(session, call.get("name", ""),
+                              call.get("args") or {})
+                    for call in turn_calls))
+                function_responses = []
+                for call, output in zip(turn_calls, outputs):
                     name = call.get("name", "")
                     args = call.get("args") or {}
-                    label = _summarize_call(name, args)
-                    trace.append(label)
-                    yield status_md(), answer_buf, sources_md
-
-                    output = await _run_tool(session, name, args)
                     tool_log.append((name, args, output))
-                    sources_md = _format_sources(tool_log)
+                    # The model gets a capped copy; the sources panel keeps
+                    # the full text.
+                    sent = output
+                    if len(sent) > TOOL_OUTPUT_CAP:
+                        sent = (sent[:TOOL_OUTPUT_CAP]
+                                + "\n\n(truncated for context; the full text "
+                                "was shown to the user)")
                     function_responses.append({
                         "functionResponse": {
                             "name": name,
-                            "response": {"output": output},
+                            "response": {"output": sent},
                         }
                     })
+                sources_md = _format_sources(tool_log)
                 contents.append({"role": "user", "parts": function_responses})
 
             # Loop budget exhausted -- ask Gemini for a final answer without
@@ -639,8 +748,11 @@ async def _agentic_answer(question: str):
                 "the best answer you can from the material gathered so far, "
                 "without calling further tools. If the material is "
                 "insufficient, say so plainly."}]})
+            # tool_mode NONE: the model CANNOT waste the terminal turn on
+            # another function call the loop would discard.
             turn_text = ""
-            async for chunk in _gemini_stream(api_key, contents):
+            async for chunk in _gemini_stream(api_key, contents,
+                                              tool_mode="NONE"):
                 if chunk["type"] == "text_delta":
                     turn_text += chunk["text"]
                     yield (status_md(thinking=False),
@@ -661,7 +773,30 @@ ANSWER_PLACEHOLDER = "*Your answer will appear here.*"
 _SENTINEL = object()
 
 
-def answer(question: str):
+def _log_run(question: str, started: float, outcome: str):
+    """One structured stderr line per question -- HF captures container
+    stderr, so this is the whole postmortem story for 'it failed yesterday'
+    plus free-tier quota tracking."""
+    usage = _RUN_STATS.get("usage") or {}
+    try:
+        print(json.dumps({
+            "ts": _dt.datetime.now().isoformat(timespec="seconds"),
+            "question": question[:120],
+            "outcome": outcome,
+            "duration_s": round(time.monotonic() - started, 1),
+            "iterations": _RUN_STATS.get("iterations"),
+            "finish_reason": _RUN_STATS.get("finish"),
+            "tools": _RUN_STATS.get("tools", []),
+            "prompt_tokens": usage.get("promptTokenCount"),
+            "response_tokens": usage.get("candidatesTokenCount"),
+            "thinking_tokens": usage.get("thoughtsTokenCount"),
+            "model": GEMINI_MODEL,
+        }, ensure_ascii=False), file=sys.stderr, flush=True)
+    except Exception:
+        pass
+
+
+def answer(question: str, history: list | None):
     """Generator wrapping the async agent for Gradio's progressive UI.
 
     The async work runs on a dedicated worker thread with its own event loop
@@ -670,19 +805,33 @@ def answer(question: str):
     previous loop.run_until_complete-per-anext pattern created a fresh task
     on every yield, which tripped anyio's cancel-scope check inside the MCP
     streamable-HTTP client ('Attempted to exit cancel scope in a different
-    task than it was entered in')."""
+    task than it was entered in').
+
+    Yields (status, answer_md, sources_md, history). On error, the partial
+    answer already streamed stays visible under the error notice instead of
+    vanishing. If Gradio cancels this generator (Stop button, disconnect),
+    the stop event tells the worker to abandon the run instead of silently
+    burning quota to completion."""
     question = (question or "").strip()
+    history = list(history or [])
     if not question:
-        yield "Please enter a legal question above.", ANSWER_PLACEHOLDER, ""
+        yield ("Please enter a legal question above.", ANSWER_PLACEHOLDER,
+               "", history)
         return
 
+    _RUN_STATS.clear()
+    started = time.monotonic()
     events: queue.Queue = queue.Queue()
+    stop_event = threading.Event()
 
     def worker():
         async def run():
             try:
-                async for tup in _agentic_answer(question):
+                async for tup in _agentic_answer(question, history,
+                                                 stop_event):
                     events.put(("yield", tup))
+            except _Stopped:
+                events.put(("stopped", None))
             except _AgentError as exc:
                 events.put(("agent_error", exc))
             except Exception as exc:                # network blip, MCP transport
@@ -697,31 +846,58 @@ def answer(question: str):
 
     threading.Thread(target=worker, daemon=True).start()
 
-    while True:
-        kind, *payload = events.get()
-        if kind is _SENTINEL:
-            return
-        if kind == "yield":
-            yield payload[0]
-        elif kind == "agent_error":
-            yield (f"**{payload[0]}**", ANSWER_PLACEHOLDER, "")
-        elif kind == "error":
-            exc = payload[0]
-            # Unwrap ExceptionGroup (from anyio TaskGroups in the MCP client)
-            # so the user sees the actual root cause, not the wrapper.
-            lines = []
-            def _walk(e, depth=0):
-                indent = "  " * depth
-                lines.append(f"{indent}- `{type(e).__name__}: {e}`")
-                inner = getattr(e, "exceptions", None)
-                if inner:
-                    for sub in inner:
-                        _walk(sub, depth + 1)
-            _walk(exc)
-            yield ("**Could not complete the request.**\n\n"
-                   + "\n".join(lines) +
-                   "\n\nThe MCP service may be waking from sleep -- "
-                   "try again in a moment.", ANSWER_PLACEHOLDER, "")
+    last_answer, last_sources = "", ""
+    outcome = "ok"
+    try:
+        while True:
+            kind, *payload = events.get()
+            if kind is _SENTINEL:
+                if outcome == "ok" and last_answer and \
+                        last_answer != ANSWER_PLACEHOLDER:
+                    history.append((question, last_answer))
+                return
+            if kind == "yield":
+                status, ans, src = payload[0]
+                last_answer, last_sources = ans or last_answer, src or last_sources
+                yield status, ans, src, history
+            elif kind == "stopped":
+                outcome = "stopped"
+                yield ("_Stopped._", last_answer or ANSWER_PLACEHOLDER,
+                       last_sources, history)
+            elif kind == "agent_error":
+                outcome = "agent_error"
+                kept = ""
+                if last_answer and last_answer != ANSWER_PLACEHOLDER:
+                    kept = ("\n\n---\n\n_Partial answer before the error "
+                            "(retry to complete):_\n\n" + last_answer)
+                yield (f"**{payload[0]}**",
+                       (last_answer or ANSWER_PLACEHOLDER) if not kept
+                       else f"**{payload[0]}**{kept}",
+                       last_sources, history)
+            elif kind == "error":
+                outcome = "error"
+                exc = payload[0]
+                # Unwrap ExceptionGroup (from anyio TaskGroups in the MCP
+                # client) so the user sees the actual root cause.
+                lines = []
+                def _walk(e, depth=0):
+                    indent = "  " * depth
+                    lines.append(f"{indent}- `{type(e).__name__}: {e}`")
+                    inner = getattr(e, "exceptions", None)
+                    if inner:
+                        for sub in inner:
+                            _walk(sub, depth + 1)
+                _walk(exc)
+                notice = ("**Could not complete the request.**\n\n"
+                          + "\n".join(lines) +
+                          "\n\nThe MCP service may be waking from sleep -- "
+                          "try again in a moment.")
+                body = last_answer if last_answer and \
+                    last_answer != ANSWER_PLACEHOLDER else ANSWER_PLACEHOLDER
+                yield notice, body, last_sources, history
+    finally:
+        stop_event.set()          # Stop button / disconnect: end the worker
+        _log_run(question, started, outcome)
 
 
 # --- UI -----------------------------------------------------------------------
@@ -734,16 +910,20 @@ labour or related federal law**. CanLex finds the governing statutory
 provisions, D-Memoranda, collective-agreement terms and leading court
 decisions, then composes an answer that cites them.
 
-The CanLex corpus contains 31 federal Acts and regulations -- including the
-Immigration and Refugee Protection Act, the Customs Act and the Criminal Code
--- alongside the CBSA D-Memoranda, the FB (Border Services) collective
-agreement, the National Joint Council directives, leading decisions of the
-Supreme Court, the Federal Courts and the federal labour and immigration
-tribunals, and the IRPA/IRPR instruments of delegation and designation.
+The CanLex corpus contains dozens of federal Acts and regulations with their
+Schedules -- including the Immigration and Refugee Protection Act, the Customs
+Act, the Criminal Code and the controlled-substance and firearms-classification
+schedules -- alongside the CBSA D-Memoranda, the FB (Border Services)
+collective agreement, the National Joint Council directives, leading decisions
+of the Supreme Court, the Federal Courts and the federal labour and
+immigration tribunals, the IRPA/IRPR instruments of delegation, and CanLex's
+curated US-record screening commentary. (Ask "what does the corpus contain?"
+for the live inventory.)
 
 The assistant iterates over the corpus -- searching, fetching sections and
 looking up case-law citations -- before composing a grounded answer. A complex
-question may take 30 seconds or more.
+question may take 30 seconds or more. Follow-up questions keep the
+conversation's context; Clear starts fresh.
 
 Legal information, not legal advice -- always verify against the primary sources.
 """
@@ -766,21 +946,56 @@ with gr.Blocks(title="CanLex", analytics_enabled=False) as demo:
     )
     with gr.Row():
         submit = gr.Button("Ask CanLex", variant="primary")
-        clear = gr.Button("Clear")
+        stop = gr.Button("Stop", variant="stop")
+        clear = gr.Button("Clear conversation")
 
     gr.Examples(examples=EXAMPLE_QUESTIONS, inputs=question, label="Examples")
 
     # Three panels: a progress trace (also used to surface errors), the final
-    # composed answer, and the raw tool outputs the agent gathered.
+    # composed answer, and the raw tool outputs the agent gathered. History
+    # carries the session's prior Q&A pairs so follow-ups have context.
+    history_state = gr.State([])
     progress_md = gr.Markdown(value="")
-    answer_md = gr.Markdown(value=ANSWER_PLACEHOLDER)
+    answer_md = gr.Markdown(value=ANSWER_PLACEHOLDER, show_copy_button=True)
     with gr.Accordion("Retrieved source passages (every tool call)", open=False):
-        sources_md = gr.Markdown()
-
-    submit.click(answer, [question], [progress_md, answer_md, sources_md])
-    question.submit(answer, [question], [progress_md, answer_md, sources_md])
-    clear.click(lambda: ("", "", ANSWER_PLACEHOLDER, ""), None,
-                [question, progress_md, answer_md, sources_md])
+        sources_md = gr.Markdown(show_copy_button=True)
+    export_btn = gr.Button("Export research memo (.md)", size="sm")
+    export_file = gr.File(label="Research memo", visible=False)
+
+    def export_memo(q, ans, src, hist):
+        """Write the session's research to a dated Markdown memo -- an
+        audit-trail artifact for a case file."""
+        stamp = _dt.datetime.now().strftime("%Y-%m-%d %H:%M")
+        parts = [f"# CanLex research memo — {stamp}", ""]
+        for past_q, past_a in (hist or []):
+            parts += [f"## Q: {past_q}", "", past_a, ""]
+        if q and ans and ans != ANSWER_PLACEHOLDER and not any(
+                q == pq for pq, _ in (hist or [])):
+            parts += [f"## Q: {q}", "", ans, ""]
+        if src:
+            parts += ["---", "", "## Retrieved source passages", "", src]
+        parts += ["", "---", "_Generated by CanLex Web. Legal information, "
+                  "not legal advice; verify against the primary sources._"]
+        path = os.path.join(tempfile.gettempdir(),
+                            f"canlex-memo-{_dt.date.today().isoformat()}.md")
+        with open(path, "w", encoding="utf-8") as fh:
+            fh.write("\n".join(parts))
+        return gr.File(value=path, visible=True)
+
+    ask_evt = submit.click(answer, [question, history_state],
+                           [progress_md, answer_md, sources_md, history_state])
+    enter_evt = question.submit(answer, [question, history_state],
+                                [progress_md, answer_md, sources_md,
+                                 history_state])
+    stop.click(None, cancels=[ask_evt, enter_evt])
+    export_btn.click(export_memo,
+                     [question, answer_md, sources_md, history_state],
+                     [export_file])
+    clear.click(lambda: ("", "", ANSWER_PLACEHOLDER, "", [],
+                         gr.File(visible=False)),
+                None,
+                [question, progress_md, answer_md, sources_md, history_state,
+                 export_file])
 
 
 if __name__ == "__main__":