Merge pull request #128 from jmjava/cursor/wizard-state-segments-2ccd

Fail closed when wizard state segments is not a mapping

Author: cursor[bot]

milestones/README.md

@@ -5,10 +5,13 @@ repositories that install `docgen` and maintain their own demo bundle. The
 library no longer ships an in-repo dogfood; consumers are the integration test
 of record.
 
-**Active:** **[timing-inner-lists.md](timing-inner-lists.md)** —
-`timing.json` `words` / `segments` must be JSON arrays of objects.
+**Active:** **[wizard-state-segments.md](wizard-state-segments.md)** —
+wizard `.docgen-state.json` `segments` must be a mapping of objects.
 
 **Shipped:**
+- **[timing-inner-lists.md](timing-inner-lists.md)** —
+  `timing.json` `words` / `segments` must be JSON arrays of objects
+  (#127).
 - **[timing-stem-objects.md](timing-stem-objects.md)** —
   `timing.json` per-stem values must be JSON objects (not lists/scalars)
   (#126).

milestones/timing-inner-lists.md

@@ -1,6 +1,6 @@
 # Milestone: timing.json words/segments must be object arrays
 
-**Status:** Active  
+**Status:** Shipped  
 **PR:** [#127](https://github.com/jmjava/documentation-generator/pull/127)  
 **Depends on:** `milestones/timing-stem-objects.md` (PR #126)
 

milestones/wizard-state-segments.md

@@ -0,0 +1,38 @@
+# Milestone: wizard state segments must be objects
+
+**Status:** Active  
+**PR:** [#128](https://github.com/jmjava/documentation-generator/pull/128)  
+**Depends on:** `milestones/timing-inner-lists.md` (PR #127)
+
+## Problem
+
+``load_state`` accepts any JSON object. ``GET /api/segments`` then does
+``state.get("segments", {}).get(seg_id, {})``. A list, string, or
+per-id non-object under ``segments`` raises ``AttributeError`` (500
+traceback). ``POST /api/state`` writes that payload unchanged.
+
+Corrupt JSON / a non-object root still reset to ``{"segments": {}}``
+so a broken file does not brick the wizard.
+
+## Goal
+
+When ``segments`` is present, it must be a JSON object whose values
+are objects. ``POST /api/state`` rejects a non-object body and invalid
+``segments`` with 400. ``GET`` on a corrupt-typed file returns 500 with
+the parse error (not ``AttributeError``).
+
+## Done when
+
+- [x] ``load_state`` raises ``WizardError`` on list / scalar ``segments``
+- [x] per-id non-object rows raise
+- [x] missing / null ``segments`` still means ``{}``
+- [x] corrupt JSON still resets to empty (existing contract)
+- [x] ``POST /api/state`` rejects a list body and list ``segments``
+- [x] `ruff check src/ tests/`
+- [x] `pytest tests/` (741 passed, 1 skipped)
+- [x] `docgen benchmark` (no clock change; meets baseline)
+
+## Out of scope
+
+- Changing the corrupt-JSON → empty reset
+- Wizard ``except Exception`` around ``narration_topic_label``

src/docgen/wizard.py

@@ -12,6 +12,31 @@
 STATE_FILENAME = ".docgen-state.json"
 
 
+class WizardError(RuntimeError):
+    """Raised when wizard persisted state or API payload is invalid."""
+
+
+def _json_kind(value: Any) -> str:
+    return "null" if value is None else type(value).__name__
+
+
+def require_state_segments(data: dict[str, Any], *, label: str) -> dict[str, Any]:
+    """Return the ``segments`` mapping, or raise if it is present but not objects."""
+    segs = data.get("segments")
+    if segs is None:
+        return {}
+    if not isinstance(segs, dict):
+        raise WizardError(
+            f"{label} segments must be a JSON object, not {_json_kind(segs)}"
+        )
+    for sid, row in segs.items():
+        if not isinstance(row, dict):
+            raise WizardError(
+                f"{label} segments[{sid!r}] must be a JSON object, not {_json_kind(row)}"
+            )
+    return segs
+
+
 def session_payload(config: Any | None) -> dict[str, Any]:
     """GUI session: frozen shell vs pip CLI, and whether a bundle is attached."""
     from docgen.resources import is_frozen
@@ -229,7 +254,11 @@ def load_state(base_dir: Path) -> dict[str, Any]:
         data = json.loads(p.read_text(encoding="utf-8"))
     except (json.JSONDecodeError, OSError, UnicodeDecodeError):
         return {"segments": {}}
-    return data if isinstance(data, dict) else {"segments": {}}
+    if not isinstance(data, dict):
+        return {"segments": {}}
+    out = dict(data)
+    out["segments"] = require_state_segments(data, label=".docgen-state.json")
+    return out
 
 
 def save_state(base_dir: Path, state: dict[str, Any]) -> None:
@@ -648,13 +677,26 @@ def api_generate_narration():
     def api_get_state():
         cfg = _cfg()
         base = cfg.base_dir if cfg else Path.cwd()
-        return jsonify(load_state(base))
+        try:
+            return jsonify(load_state(base))
+        except WizardError as exc:
+            return jsonify({"error": str(exc)}), 500
 
     @app.route("/api/state", methods=["POST"])
     def api_set_state():
         cfg = _cfg()
         base = cfg.base_dir if cfg else Path.cwd()
-        state = request.json or {}
+        raw = request.get_json(silent=True)
+        if raw is None:
+            raw = {}
+        if not isinstance(raw, dict):
+            return jsonify({"error": "state must be a JSON object"}), 400
+        try:
+            segs = require_state_segments(raw, label="state")
+        except WizardError as exc:
+            return jsonify({"error": str(exc)}), 400
+        state = dict(raw)
+        state["segments"] = segs
         save_state(base, state)
         return jsonify({"ok": True})
 
@@ -668,7 +710,10 @@ def api_segments():
         from docgen.yaml_generate import read_hint_focus_paths
 
         base = cfg.base_dir
-        state = load_state(base)
+        try:
+            state = load_state(base)
+        except WizardError as exc:
+            return jsonify({"error": str(exc)}), 500
         result = []
         for seg_id in cfg.segments_all:
             seg_name = cfg.resolve_segment_name(seg_id)

tests/test_wizard.py

@@ -66,6 +66,98 @@ def test_load_state_corrupt_json_returns_empty(tmp_path):
     assert load_state(tmp_path) == {"segments": {}}
 
 
+def test_load_state_rejects_non_object_segments(tmp_path):
+    import json
+
+    from docgen.wizard import WizardError, load_state
+
+    (tmp_path / ".docgen-state.json").write_text(
+        json.dumps({"segments": ["01"]}), encoding="utf-8"
+    )
+    with pytest.raises(WizardError, match=r"\.docgen-state.json segments must be a JSON object, not list"):
+        load_state(tmp_path)
+
+
+def test_load_state_rejects_non_object_segment_row(tmp_path):
+    import json
+
+    from docgen.wizard import WizardError, load_state
+
+    (tmp_path / ".docgen-state.json").write_text(
+        json.dumps({"segments": {"01": "draft"}}), encoding="utf-8"
+    )
+    with pytest.raises(
+        WizardError,
+        match=r"\.docgen-state.json segments\['01'\] must be a JSON object, not str",
+    ):
+        load_state(tmp_path)
+
+
+def test_load_state_null_segments_is_empty(tmp_path):
+    import json
+
+    from docgen.wizard import load_state
+
+    (tmp_path / ".docgen-state.json").write_text(
+        json.dumps({"segments": None, "extra": 1}), encoding="utf-8"
+    )
+    assert load_state(tmp_path) == {"segments": {}, "extra": 1}
+
+
+def _wizard_client(tmp_path):
+    from docgen.config import Config
+    from docgen.wizard import create_app
+
+    yaml_path = tmp_path / "docgen.yaml"
+    yaml_path.write_text(
+        "repo_root: .\nsegments:\n  default: ['01']\n  all: ['01']\n",
+        encoding="utf-8",
+    )
+    cfg = Config.from_yaml(yaml_path)
+    return create_app(cfg).test_client(), cfg
+
+
+def test_api_segments_rejects_list_state_segments(tmp_path):
+    import json
+
+    client, cfg = _wizard_client(tmp_path)
+    (cfg.base_dir / ".docgen-state.json").write_text(
+        json.dumps({"segments": []}), encoding="utf-8"
+    )
+    res = client.get("/api/segments")
+    assert res.status_code == 500
+    assert "segments must be a JSON object" in res.get_json()["error"]
+
+
+def test_api_state_post_rejects_list_body(tmp_path):
+    client, _cfg = _wizard_client(tmp_path)
+    res = client.post("/api/state", json=["not", "an", "object"])
+    assert res.status_code == 400
+    assert res.get_json()["error"] == "state must be a JSON object"
+
+
+def test_api_state_post_rejects_list_segments(tmp_path):
+    client, cfg = _wizard_client(tmp_path)
+    res = client.post("/api/state", json={"segments": ["01"]})
+    assert res.status_code == 400
+    assert "segments must be a JSON object" in res.get_json()["error"]
+    assert not (cfg.base_dir / ".docgen-state.json").exists()
+
+
+def test_api_state_roundtrip_object_segments(tmp_path):
+    client, cfg = _wizard_client(tmp_path)
+    payload = {"segments": {"01": {"status": "ready", "revision_notes": "n"}}}
+    res = client.post("/api/state", json=payload)
+    assert res.status_code == 200
+    got = client.get("/api/state")
+    assert got.status_code == 200
+    assert got.get_json()["segments"]["01"]["status"] == "ready"
+    segs = client.get("/api/segments")
+    assert segs.status_code == 200
+    assert segs.get_json()["segments"][0]["status"] == "ready"
+    assert (cfg.base_dir / ".docgen-state.json").is_file()
+
+
 def test_api_file_rejects_prefix_escape(tmp_path):
     from docgen.config import Config
     from docgen.wizard import create_app