fix(labs): stop silently dropping WASM DesignLedger saves (#1985)

save() previously fire-and-forgot the WASM persistence task via
asyncio.create_task(), so IndexedDB failures were swallowed and
students could lose progress with no indication anything went wrong.

- save_async() now raises on IndexedDB failure instead of print()-ing
- save() attaches a done-callback to the background task and records
  failures on last_save_error / save_pending
- added asave() and flush() for callers that can await
- added regression tests in tests/test_state.py

Author: aadityansha06

mlsysim/mlsysim/labs/state.py

@@ -26,8 +26,25 @@ def __init__(self):
         self.config_dir = Path.home() / ".mlsys"
         self.file_path = self.config_dir / "ledger.json"
         self._state = LedgerState()
+        self._pending_save_task = None
+        self._last_save_error: Optional[str] = None
         self.load()
 
+    @property
+    def last_save_error(self) -> Optional[str]:
+        """Error message from the most recent failed background save, if any.
+
+        Callers in WASM environments should check this after calling
+        ``save()`` (e.g. on the next reactive cell run) since ``save()``
+        cannot block on the result inside a synchronous notebook cell.
+        """
+        return self._last_save_error
+
+    @property
+    def save_pending(self) -> bool:
+        """True while a WASM background save has not finished yet."""
+        return self._pending_save_task is not None and not self._pending_save_task.done()
+
     @property
     def is_wasm(self) -> bool:
         """Detect if we are running in a browser environment (Pyodide)."""
@@ -110,53 +127,62 @@ async def load_async(self) -> LedgerState:
         return self._state
 
     async def save_async(self):
-        """Async save for WASM environments using IndexedDB."""
+        """Async save for WASM environments using IndexedDB.
+
+        Unlike the previous implementation, this raises on failure instead
+        of printing to the console and returning ``None``. Callers that can
+        await (e.g. :meth:`asave`, tests) will see the real exception.
+        Callers that cannot await (the synchronous :meth:`save`) capture the
+        exception via a done-callback on the background task instead.
+        """
         if not self.is_wasm:
-            return
-        try:
-            import json
-            from pyodide.code import run_js
-            from js import globalThis
-            
-            state_json = json.dumps(asdict(self._state))
-            globalThis.__mlsys_temp_state = state_json
-            
-            js_code = """
-            (async () => {
-                return new Promise((resolve, reject) => {
-                    const request = indexedDB.open("mlsys_ledger_db", 1);
-                    request.onupgradeneeded = (e) => {
-                        const db = e.target.result;
-                        if (!db.objectStoreNames.contains("ledger")) {
-                            db.createObjectStore("ledger");
-                        }
-                    };
-                    request.onsuccess = (e) => {
-                        const db = e.target.result;
-                        try {
-                            const tx = db.transaction("ledger", "readwrite");
-                            const store = tx.objectStore("ledger");
-                            const putReq = store.put(globalThis.__mlsys_temp_state, "mlsys_design_ledger");
-                            putReq.onsuccess = () => resolve(true);
-                            putReq.onerror = () => resolve(false);
-                        } catch (err) {
-                            resolve(false);
-                        }
-                    };
-                    request.onerror = () => resolve(false);
-                });
-            })()
-            """
-            await run_js(js_code)
-        except Exception as e:
-            print(f"Failed to save to IndexedDB: {e}")
+            return True
 
-    def save(self, track: str = None, step: int = None, design: dict = None, chapter: int = None):
-        """Persists the design decisions to storage.
+        import json
+        from pyodide.code import run_js
+        from js import globalThis
 
-        ``chapter`` is kept as a compatibility alias for existing Co-Labs,
-        while ``step`` is the newer generic ledger key.
+        state_json = json.dumps(asdict(self._state))
+        globalThis.__mlsys_temp_state = state_json
+
+        js_code = """
+        (async () => {
+            return new Promise((resolve, reject) => {
+                const request = indexedDB.open("mlsys_ledger_db", 1);
+                request.onupgradeneeded = (e) => {
+                    const db = e.target.result;
+                    if (!db.objectStoreNames.contains("ledger")) {
+                        db.createObjectStore("ledger");
+                    }
+                };
+                request.onsuccess = (e) => {
+                    const db = e.target.result;
+                    try {
+                        const tx = db.transaction("ledger", "readwrite");
+                        const store = tx.objectStore("ledger");
+                        const putReq = store.put(globalThis.__mlsys_temp_state, "mlsys_design_ledger");
+                        putReq.onsuccess = () => resolve(true);
+                        putReq.onerror = () => reject(new Error(
+                            "IndexedDB put failed: " + (putReq.error ? putReq.error.message : "unknown error")
+                        ));
+                    } catch (err) {
+                        reject(new Error("IndexedDB transaction failed: " + err.message));
+                    }
+                };
+                request.onerror = () => reject(new Error(
+                    "indexedDB.open failed: " + (request.error ? request.error.message : "unknown error")
+                    + " (storage may be disabled, full, or unavailable in private browsing)"
+                ));
+            });
+        })()
         """
+        # If the JS promise rejects, Pyodide raises the corresponding
+        # exception here rather than us having to poll a boolean result.
+        await run_js(js_code)
+        return True
+
+    def _apply_pending_state(self, track, step, design, chapter):
+        """Shared bookkeeping used by both save() and asave()."""
         if track:
             self._state.track = track
 
@@ -169,14 +195,77 @@ def save(self, track: str = None, step: int = None, design: dict = None, chapter
         if design is not None:
             self._state.history[step_id] = design
 
+    def _on_save_task_done(self, task):
+        """Done-callback for the background WASM save task.
+
+        This is what actually fixes the silent-failure bug: instead of the
+        exception from `save_async()` disappearing into an unobserved task,
+        we record it and log it loudly so it's impossible to miss in the
+        browser console, and expose it via `last_save_error` /
+        `save_pending` so the notebook UI can surface it to the student.
+        """
+        try:
+            task.result()
+            self._last_save_error = None
+        except Exception as e:
+            self._last_save_error = str(e)
+            message = f"[DesignLedger] SAVE FAILED - progress was NOT persisted: {e}"
+            try:
+                from js import console
+                console.error(message)
+            except Exception:
+                print(message, file=sys.stderr)
+
+    def save(self, track: str = None, step: int = None, design: dict = None, chapter: int = None):
+        """Persists the design decisions to storage.
+
+        ``chapter`` is kept as a compatibility alias for existing Co-Labs,
+        while ``step`` is the newer generic ledger key.
+
+        In WASM this schedules a background save (marimo cells are
+        synchronous, so we can't await here), but unlike the previous
+        implementation the resulting task's outcome is observed: failures
+        are logged loudly and recorded on ``self.last_save_error`` instead
+        of vanishing silently. Callers that need a hard persistence
+        guarantee (tests, async code) should use :meth:`asave` instead.
+        """
+        self._apply_pending_state(track, step, design, chapter)
+
         if self.is_wasm:
             import asyncio
-            asyncio.create_task(self.save_async())
+            task = asyncio.ensure_future(self.save_async())
+            task.add_done_callback(self._on_save_task_done)
+            self._pending_save_task = task
+        else:
+            self.config_dir.mkdir(exist_ok=True)
+            with open(self.file_path, 'w') as f:
+                json.dump(asdict(self._state), f, indent=2)
+
+    async def asave(self, track: str = None, step: int = None, design: dict = None, chapter: int = None):
+        """Async variant of :meth:`save` that awaits persistence.
+
+        Guarantees the data is written (or raises) before returning. Use
+        this from async contexts (tests, async marimo cells) whenever you
+        need certainty that a save actually succeeded.
+        """
+        self._apply_pending_state(track, step, design, chapter)
+
+        if self.is_wasm:
+            await self.save_async()
+            self._last_save_error = None
         else:
             self.config_dir.mkdir(exist_ok=True)
             with open(self.file_path, 'w') as f:
                 json.dump(asdict(self._state), f, indent=2)
 
+    async def flush(self):
+        """Await any in-flight background save scheduled by `save()`.
+
+        Raises the underlying exception if that save failed.
+        """
+        if self._pending_save_task is not None:
+            await self._pending_save_task
+
     def get_design(self, step_id: int) -> Optional[Dict[str, Any]]:
         """Retrieves the design dictionary for a specific step."""
         return self._state.history.get(step_id)

mlsysim/tests/test_state.py

@@ -0,0 +1,99 @@
+"""Tests for DesignLedger persistence (mlsysim/labs/state.py).
+
+Covers the WASM background-save failure path fixed in #1985: previously
+`save()` used `asyncio.create_task(...)` fire-and-forget, so IndexedDB
+failures inside `save_async()` were silently swallowed and never surfaced
+to the caller. See:
+https://github.com/harvard-edge/cs249r_book/issues/1985
+"""
+import asyncio
+
+import pytest
+
+from mlsysim.labs.state import DesignLedger
+import mlsysim.labs.state as state_mod
+
+
+def _ledger_with_home(monkeypatch, tmp_path):
+    monkeypatch.setattr(state_mod.Path, "home", lambda: tmp_path)
+    return DesignLedger()
+
+
+def _force_wasm(monkeypatch, ledger):
+    monkeypatch.setattr(type(ledger), "is_wasm", property(lambda self: True))
+
+
+def test_save_and_load_roundtrip_native(tmp_path, monkeypatch):
+    """Non-WASM path: save() should persist synchronously to disk."""
+    ledger = _ledger_with_home(monkeypatch, tmp_path)
+    ledger.save(track="edge", step=1, design={"foo": "bar"})
+
+    reloaded = _ledger_with_home(monkeypatch, tmp_path)
+    assert reloaded.get_track() == "edge"
+    assert reloaded.get_design(1) == {"foo": "bar"}
+
+
+def test_wasm_save_success_clears_error(tmp_path, monkeypatch):
+    """A successful WASM background save should leave last_save_error unset."""
+    ledger = _ledger_with_home(monkeypatch, tmp_path)
+    _force_wasm(monkeypatch, ledger)
+
+    async def fake_save_async(self):
+        return True
+
+    monkeypatch.setattr(DesignLedger, "save_async", fake_save_async)
+
+    async def run():
+        ledger.save(step=1, design={"ok": True})
+        await ledger.flush()
+
+    asyncio.run(run())
+    assert ledger.last_save_error is None
+    assert ledger.save_pending is False
+
+
+def test_wasm_save_failure_is_captured_not_silent(tmp_path, monkeypatch):
+    """Regression test for #1985.
+
+    Previously: an IndexedDB failure inside save_async() was swallowed by
+    a fire-and-forget asyncio.create_task(), so save() returned as if it
+    had succeeded and students never learned their progress was lost.
+
+    Now: the failure must be captured on `last_save_error`, and
+    `save_pending` must go back to False once the background task settles.
+    """
+    ledger = _ledger_with_home(monkeypatch, tmp_path)
+    _force_wasm(monkeypatch, ledger)
+
+    async def failing_save_async(self):
+        raise RuntimeError("indexedDB.open failed: QuotaExceededError")
+
+    monkeypatch.setattr(DesignLedger, "save_async", failing_save_async)
+
+    async def run():
+        ledger.save(step=1, design={"will_fail": True})
+        with pytest.raises(RuntimeError):
+            await ledger.flush()
+
+    asyncio.run(run())
+
+    assert ledger.save_pending is False
+    assert ledger.last_save_error is not None
+    assert "QuotaExceededError" in ledger.last_save_error
+
+
+def test_asave_raises_on_failure_directly(tmp_path, monkeypatch):
+    """asave() should propagate the exception directly to an awaiting caller."""
+    ledger = _ledger_with_home(monkeypatch, tmp_path)
+    _force_wasm(monkeypatch, ledger)
+
+    async def failing_save_async(self):
+        raise RuntimeError("storage disabled")
+
+    monkeypatch.setattr(DesignLedger, "save_async", failing_save_async)
+
+    async def run():
+        with pytest.raises(RuntimeError, match="storage disabled"):
+            await ledger.asave(step=2, design={"x": 1})
+
+    asyncio.run(run())