fix: address PR review feedback on Mathpix failure paths

Adds timeouts to all Mathpix HTTP calls, polls the conversion status
endpoint instead of treating every non-200 as "not ready", surfaces
Mathpix's in-band error bodies instead of raising a bare KeyError, warns
instead of silently skipping a failed figure download, and returns the
markdown as a string rather than writing it into out_dir (which would
otherwise collide with the user's chosen output file once the wizard
command wires this up). Also documents that PDFs are sent to a
third-party OCR service and that Mathpix offers a training opt-out.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: m-messer

in2lambda/wizard/mathpix.py

@@ -3,11 +3,17 @@
 Needs ``MATHPIX_APP_ID`` and ``MATHPIX_API_KEY`` in the environment (a ``.env``
 file is honoured by the wizard). Figures referenced by the returned markdown are
 downloaded next to it so the ``Markdown`` filter can pick them up.
+
+The PDF is uploaded to Mathpix, a third-party OCR service, for processing.
+Instructors converting student work should be told their PDFs leave the
+local machine. Mathpix also offers an opt-out from using submitted data to
+improve its models; see https://mathpix.com/privacy for how to enable it.
 """
 
 import os
 import re
 import time
+import warnings
 from pathlib import Path
 
 import requests
@@ -35,65 +41,93 @@ def pdf_to_markdown(
     out_dir: str,
     poll_interval: float = 5.0,
     max_polls: int = 60,
-) -> Path:
-    """Convert ``pdf_path`` to markdown, writing it and its figures under ``out_dir``.
+    timeout: float = 30.0,
+) -> str:
+    """Convert ``pdf_path`` to markdown, downloading its figures under ``out_dir``.
 
     Args:
         pdf_path: Path to the source PDF.
-        out_dir: Directory to write ``<stem>.md`` and a ``media/`` folder into.
+        out_dir: Directory to write a ``media/`` folder of figures into.
         poll_interval: Seconds to wait between Mathpix "is it ready yet" polls.
         max_polls: How many times to poll before giving up.
+        timeout: Seconds to wait for each individual HTTP request.
 
     Returns:
-        The path to the written markdown file. Figures are saved in
-        ``<out_dir>/media/`` and referenced from the markdown as
-        ``./media/<name>``.
+        The converted markdown, with figures saved in ``<out_dir>/media/`` and
+        referenced from the markdown as ``./media/<name>``. The caller is
+        responsible for writing the markdown out wherever it belongs.
 
     Raises:
-        RuntimeError: if credentials are missing or Mathpix does not finish in time.
+        RuntimeError: if credentials are missing, Mathpix rejects the PDF or
+            fails to convert it, or the conversion does not finish in time.
     """
     headers = _headers()
     out = Path(out_dir)
     (out / "media").mkdir(parents=True, exist_ok=True)
 
     with open(pdf_path, "rb") as pdf:
         response = requests.post(
-            MATHPIX_PDF_ENDPOINT, headers=headers, files={"file": pdf}
+            MATHPIX_PDF_ENDPOINT,
+            headers=headers,
+            files={"file": pdf},
+            timeout=timeout,
         )
     response.raise_for_status()
-    pdf_id = response.json()["pdf_id"]
+    data = response.json()
+    if "error" in data:
+        raise RuntimeError(f"Mathpix rejected the PDF: {data['error']}")
+    pdf_id = data["pdf_id"]
 
-    markdown = _poll_for_markdown(pdf_id, headers, poll_interval, max_polls)
-    markdown = _localise_figures(markdown, out)
-
-    md_path = out / f"{Path(pdf_path).stem}.md"
-    md_path.write_text(markdown, encoding="utf-8")
-    return md_path
+    markdown = _poll_for_markdown(pdf_id, headers, poll_interval, max_polls, timeout)
+    return _localise_figures(markdown, out, timeout)
 
 
 def _poll_for_markdown(
-    pdf_id: str, headers: dict, poll_interval: float, max_polls: int
+    pdf_id: str,
+    headers: dict,
+    poll_interval: float,
+    max_polls: int,
+    timeout: float,
 ) -> str:
-    """Poll Mathpix until the ``.md`` render of ``pdf_id`` is ready."""
-    url = f"{MATHPIX_PDF_ENDPOINT}/{pdf_id}.md"
+    """Poll Mathpix until ``pdf_id`` finishes converting, then return its markdown."""
+    status_url = f"{MATHPIX_PDF_ENDPOINT}/{pdf_id}"
     for _ in range(max_polls):
-        response = requests.get(url, headers=headers)
-        if response.status_code == 200:
-            return response.text
+        response = requests.get(status_url, headers=headers, timeout=timeout)
+        response.raise_for_status()
+        data = response.json()
+        status = data.get("status")
+        if status == "completed":
+            break
+        if status == "error":
+            raise RuntimeError(
+                f"Mathpix failed to convert {pdf_id}: {data.get('error', 'unknown error')}"
+            )
         time.sleep(poll_interval)
-    raise RuntimeError(f"Mathpix did not finish converting {pdf_id} in time.")
+    else:
+        raise RuntimeError(f"Mathpix did not finish converting {pdf_id} in time.")
+
+    md_response = requests.get(
+        f"{MATHPIX_PDF_ENDPOINT}/{pdf_id}.md", headers=headers, timeout=timeout
+    )
+    md_response.raise_for_status()
+    return md_response.text
 
 
-def _localise_figures(markdown: str, out_dir: Path) -> str:
+def _localise_figures(markdown: str, out_dir: Path, timeout: float) -> str:
     """Download remote figures into ``out_dir/media`` and repoint the markdown at them."""
     markdown = markdown.replace("![]", "![pictureTag]")
 
     for idx, url in enumerate(dict.fromkeys(_REMOTE_IMAGE.findall(markdown))):
         basename = os.path.basename(url).split("?")[0] or f"figure_{idx}.png"
         local_name = f"{idx}_{basename}"
 
-        image = requests.get(url)
+        image = requests.get(url, timeout=timeout)
         if image.status_code != 200:
+            warnings.warn(
+                f"Mathpix figure download failed for {url} "
+                f"(status {image.status_code}); markdown will reference a "
+                f"missing file: ./media/{local_name}"
+            )
             continue
 
         (out_dir / "media" / local_name).write_bytes(image.content)

tests/test_mathpix.py

@@ -19,62 +19,116 @@ def _pdf(tmp_path):
     return pdf
 
 
-def test_pdf_to_markdown_writes_md_and_localises_figures(tmp_path):
+def _post(pdf_id="abc123", error=None):
+    post = MagicMock(status_code=200)
+    post.json.return_value = {"error": error} if error else {"pdf_id": pdf_id}
+    return post
+
+
+def _status(status, error=None):
+    body = {"status": status}
+    if error:
+        body["error"] = error
+    resp = MagicMock(status_code=200)
+    resp.json.return_value = body
+    return resp
+
+
+def test_pdf_to_markdown_returns_markdown_and_localises_figures(tmp_path):
     pdf = _pdf(tmp_path)
     out_dir = tmp_path / "out"
 
-    post = MagicMock(status_code=200)
-    post.json.return_value = {"pdf_id": "abc123"}
+    completed = _status("completed")
     md = MagicMock(
         status_code=200,
         text="# Heading\n\n![](https://cdn.mathpix.com/x/fig.png?width=8) done\n",
     )
     image = MagicMock(status_code=200, content=b"PNGBYTES")
 
     with patch("in2lambda.wizard.mathpix.requests") as req:
-        req.post.return_value = post
-        req.get.side_effect = [md, image]
-        md_path = pdf_to_markdown(str(pdf), str(out_dir), poll_interval=0.0)
+        req.post.return_value = _post()
+        req.get.side_effect = [completed, md, image]
+        markdown = pdf_to_markdown(str(pdf), str(out_dir), poll_interval=0.0)
 
-    assert md_path == out_dir / "paper.md"
-    text = md_path.read_text()
-    assert "![pictureTag](./media/0_fig.png)" in text
+    assert "![pictureTag](./media/0_fig.png)" in markdown
+    assert not (out_dir / "paper.md").exists()
     assert (out_dir / "media" / "0_fig.png").read_bytes() == b"PNGBYTES"
 
 
 def test_pdf_to_markdown_polls_until_ready(tmp_path):
     pdf = _pdf(tmp_path)
 
-    post = MagicMock(status_code=200)
-    post.json.return_value = {"pdf_id": "abc123"}
-    not_ready = MagicMock(status_code=202)
-    ready = MagicMock(status_code=200, text="# Only text, no figures\n")
+    processing = _status("processing")
+    completed = _status("completed")
+    md = MagicMock(status_code=200, text="# Only text, no figures\n")
 
     with patch("in2lambda.wizard.mathpix.requests") as req:
-        req.post.return_value = post
-        req.get.side_effect = [not_ready, not_ready, ready]
-        md_path = pdf_to_markdown(
+        req.post.return_value = _post()
+        req.get.side_effect = [processing, processing, completed, md]
+        markdown = pdf_to_markdown(
             str(pdf), str(tmp_path / "out"), poll_interval=0.0, max_polls=5
         )
 
-    assert md_path.read_text().startswith("# Only text")
+    assert markdown.startswith("# Only text")
 
 
 def test_pdf_to_markdown_times_out(tmp_path):
     pdf = _pdf(tmp_path)
 
-    post = MagicMock(status_code=200)
-    post.json.return_value = {"pdf_id": "abc123"}
-
     with patch("in2lambda.wizard.mathpix.requests") as req:
-        req.post.return_value = post
-        req.get.return_value = MagicMock(status_code=202)
+        req.post.return_value = _post()
+        req.get.return_value = _status("processing")
         with pytest.raises(RuntimeError, match="did not finish"):
             pdf_to_markdown(
                 str(pdf), str(tmp_path / "out"), poll_interval=0.0, max_polls=3
             )
 
 
+def test_pdf_to_markdown_raises_on_rejected_upload(tmp_path):
+    pdf = _pdf(tmp_path)
+
+    with patch("in2lambda.wizard.mathpix.requests") as req:
+        req.post.return_value = _post(error="Invalid file type")
+        with pytest.raises(RuntimeError, match="Mathpix rejected the PDF"):
+            pdf_to_markdown(str(pdf), str(tmp_path / "out"))
+
+
+def test_pdf_to_markdown_raises_immediately_on_conversion_error(tmp_path):
+    pdf = _pdf(tmp_path)
+
+    with patch("in2lambda.wizard.mathpix.requests") as req:
+        req.post.return_value = _post()
+        req.get.return_value = _status("error", error="conversion failed")
+        with pytest.raises(RuntimeError, match="conversion failed"):
+            pdf_to_markdown(
+                str(pdf), str(tmp_path / "out"), poll_interval=0.0, max_polls=60
+            )
+
+    # Only the single status poll should have happened, not all 60.
+    assert req.get.call_count == 1
+
+
+def test_pdf_to_markdown_warns_on_failed_figure_download(tmp_path):
+    pdf = _pdf(tmp_path)
+    out_dir = tmp_path / "out"
+
+    completed = _status("completed")
+    md = MagicMock(
+        status_code=200,
+        text="![](https://cdn.mathpix.com/x/fig.png) done\n",
+    )
+    image = MagicMock(status_code=404, content=b"")
+
+    with patch("in2lambda.wizard.mathpix.requests") as req:
+        req.post.return_value = _post()
+        req.get.side_effect = [completed, md, image]
+        with pytest.warns(UserWarning, match="figure download failed"):
+            markdown = pdf_to_markdown(str(pdf), str(out_dir), poll_interval=0.0)
+
+    assert "https://cdn.mathpix.com/x/fig.png" in markdown
+    assert not (out_dir / "media" / "0_fig.png").exists()
+
+
 def test_missing_credentials_raise(tmp_path, monkeypatch):
     monkeypatch.delenv("MATHPIX_APP_ID", raising=False)
     monkeypatch.delenv("MATHPIX_API_KEY", raising=False)