Merge pull request 'Make live snapshot revisions safe across workers and restarts' (#384) from timmy/383-multi-worker-live-revisions into main

README.md

@@ -209,10 +209,12 @@ below the route deadlines. Streaming diff reads share the same read capacity,
 while POST, PATCH, PUT, and DELETE requests are never coalesced.
 
 Live snapshots expose independent `context`, `events`, and `notifications`
-revisions. The browser sends its known revisions on later polls, so `/api/v1/live`
+revision tokens. Each bounded opaque token includes a per-process generation, so a token
-can omit unchanged section bodies while still returning current freshness and
+from another worker or from before a restart cannot suppress different content. The browser
-retry metadata. The client retains omitted data and only rebuilds or persists the
+sends its known tokens on later polls, so `/api/v1/live` can omit unchanged section bodies
-sections that changed.
+while still returning current freshness and retry metadata. The client retains omitted data
+and only rebuilds or persists the sections that changed. Malformed or oversized tokens are
+rejected before any upstream work.
 
 ## Offline mobile shell
 

frontend/context-poller.js

@@ -1,3 +1,15 @@
+function buildLiveRevisionQuery(revisions = {}) {
+  const params = new URLSearchParams();
+  const tokenPattern = /^[0-9a-f]{16}\.[0-9]{1,20}$/;
+  ['context', 'events', 'notifications'].forEach((section) => {
+    const revision = revisions[section];
+    if (typeof revision === 'string' && tokenPattern.test(revision)) {
+      params.set(section + '_revision', revision);
+    }
+  });
+  return params.toString();
+}
+
 function createContextPoller({
   fetchContext,
   onSnapshot,
@@ -123,6 +135,8 @@ function createContextPoller({
   };
 }
 
+createContextPoller.buildRevisionQuery = buildLiveRevisionQuery;
+
 if (typeof module !== 'undefined' && module.exports) {
   module.exports = createContextPoller;
 }

frontend/dashboard.js

@@ -324,11 +324,7 @@
   setClock(); setInterval(setClock, 1000);
 
   async function fetchLiveSnapshot(revisions = {}, { signal } = {}) {
-    const params = new URLSearchParams();
+    const query = createContextPoller.buildRevisionQuery(revisions);
-    Object.entries(revisions).forEach(([section, revision]) => {
-      if (Number.isInteger(revision) && revision >= 0) params.set(section + '_revision', revision);
-    });
-    const query = params.toString();
     const res = await fetch('api/v1/live' + (query ? '?' + query : ''), {
       headers: { Accept: 'application/json' },
       signal,

src/main.py

@@ -2,6 +2,7 @@ import asyncio
 import hmac
 import math
 import os
+import secrets
 import sqlite3
 import time
 from collections.abc import Awaitable, Coroutine
@@ -108,6 +109,7 @@ _live_snapshot_refreshing_sections: set[str] = set()
 _live_section_revisions: dict[str, int] = {
     section: 0 for section in LIVE_SNAPSHOT_SECTIONS
 }
+_live_revision_generation = secrets.token_hex(8)
 _read_notification_ids: set[int] = set()
 _authored_action_operations: dict[
     str, tuple[tuple[Any, ...], asyncio.Task, float]
@@ -1584,7 +1586,7 @@ def _live_snapshot_payload(
     *,
     stale: bool,
     revalidating: bool,
-    known_revisions: dict[str, int | None] | None = None,
+    known_revisions: dict[str, str | None] | None = None,
 ) -> dict:
     payload = dict(value)
     now = time.monotonic()
@@ -1617,9 +1619,13 @@ def _live_snapshot_payload(
         "retry_in_seconds": min(retries, default=0),
         "sections": section_freshness,
     }
-    payload["revisions"] = dict(_live_section_revisions)
+    revision_tokens = {
+        section: f"{_live_revision_generation}.{revision}"
+        for section, revision in _live_section_revisions.items()
+    }
+    payload["revisions"] = revision_tokens
     for section, known_revision in (known_revisions or {}).items():
-        if known_revision is None or known_revision != _live_section_revisions[section]:
+        if known_revision is None or known_revision != revision_tokens[section]:
             continue
         payload.pop(section, None)
         if section == "notifications":
@@ -1669,9 +1675,15 @@ def _without_read_notifications(snapshot: dict) -> dict:
 
 @app.get("/api/v1/live")
 async def live_snapshot(
-    context_revision: int | None = Query(default=None, ge=0),
+    context_revision: str | None = Query(
-    events_revision: int | None = Query(default=None, ge=0),
+        default=None, max_length=64, pattern=r"^[0-9a-f]{16}\.[0-9]{1,20}$"
-    notifications_revision: int | None = Query(default=None, ge=0),
+    ),
+    events_revision: str | None = Query(
+        default=None, max_length=64, pattern=r"^[0-9a-f]{16}\.[0-9]{1,20}$"
+    ),
+    notifications_revision: str | None = Query(
+        default=None, max_length=64, pattern=r"^[0-9a-f]{16}\.[0-9]{1,20}$"
+    ),
 ) -> JSONResponse:
     """Return a freshness-bounded snapshot and share identical upstream loads."""
     global _live_snapshot_task, _live_snapshot_value, _live_snapshot_created_at

tests/test_context_polling.py

@@ -116,6 +116,23 @@ const poller = createContextPoller({{
     }
 
 
+def test_live_revision_query_preserves_bounded_opaque_tokens():
+    script = f"""
+const createContextPoller = require({json.dumps(str(POLLER))});
+const query = createContextPoller.buildRevisionQuery({{
+  context: '0123456789abcdef.12',
+  events: 'fedcba9876543210.3',
+  notifications: 'not-a-token',
+  extra: '0123456789abcdef.1',
+}});
+process.stdout.write(JSON.stringify({{ query }}));
+"""
+
+    assert run_node(script) == {
+        "query": "context_revision=0123456789abcdef.12&events_revision=fedcba9876543210.3"
+    }
+
+
 def test_context_poller_aborts_a_stalled_request_and_recovers_on_schedule():
     script = f"""
 const createContextPoller = require({json.dumps(str(POLLER))});

tests/test_live_snapshot.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 
+import httpx
 import pytest
 
 from src import main
@@ -90,9 +91,9 @@ async def test_live_snapshot_fetches_user_once_and_updates_work_and_activity(mon
     assert result["sections"] == {
         "context": "fresh", "events": "fresh", "notifications": "fresh"
     }
-    assert result["revisions"] == {
+    assert set(result["revisions"]) == {"context", "events", "notifications"}
-        "context": 1, "events": 1, "notifications": 1
+    assert len(set(result["revisions"].values())) == 1
-    }
+    assert result["revisions"]["context"].endswith(".1")
 
 
 @pytest.mark.anyio
@@ -129,6 +130,56 @@ async def test_live_snapshot_omits_section_bodies_at_known_revisions(monkeypatch
     assert "freshness" in second
 
 
+@pytest.mark.anyio
+async def test_live_snapshot_token_from_another_worker_never_suppresses_content(monkeypatch):
+    async def user():
+        return {"id": 1, "login": "timmy"}
+
+    async def empty():
+        return []
+
+    monkeypatch.setattr(main, "current_user", user)
+    monkeypatch.setattr(main, "repos", empty)
+    monkeypatch.setattr(main, "issues", empty)
+    monkeypatch.setattr(main, "pull_requests", empty)
+    monkeypatch.setattr(main, "activity_events", lambda _user: empty())
+    monkeypatch.setattr(main, "notifications", empty)
+    monkeypatch.setattr(main, "_live_revision_generation", "worker-a")
+
+    worker_a = payload(await main.live_snapshot())
+    worker_a_token = worker_a["revisions"]["context"]
+
+    main._live_snapshot_value = None
+    main._live_snapshot_created_at = None
+    main._live_section_created_at = {
+        section: None for section in main.LIVE_SNAPSHOT_SECTIONS
+    }
+    main._live_section_revisions = {
+        section: 0 for section in main.LIVE_SNAPSHOT_SECTIONS
+    }
+    monkeypatch.setattr(main, "_live_revision_generation", "worker-b")
+
+    worker_b = payload(await main.live_snapshot(context_revision=worker_a_token))
+
+    assert worker_a_token == "worker-a.1"
+    assert worker_b["revisions"]["context"] == "worker-b.1"
+    assert worker_b["context"]["user"]["login"] == "timmy"
+
+
+@pytest.mark.anyio
+async def test_live_snapshot_rejects_malformed_revision_before_upstream_work(monkeypatch):
+    async def unexpected_user():
+        raise AssertionError("malformed revision reached upstream work")
+
+    monkeypatch.setattr(main, "current_user", unexpected_user)
+    transport = httpx.ASGITransport(app=main.app)
+
+    async with httpx.AsyncClient(transport=transport, base_url="https://test") as client:
+        response = await client.get("/api/v1/live?context_revision=not-a-token")
+
+    assert response.status_code == 422
+
+
 @pytest.mark.anyio
 async def test_live_snapshot_keeps_fresh_context_when_activity_fails(monkeypatch):
     async def user():

tests/test_shared_context_snapshot.py

@@ -23,7 +23,7 @@ async def test_one_live_snapshot_updates_work_and_activity_on_one_timer():
 
     assert "fetch('api/v1/live'" in html
     assert "onSnapshot: renderLiveSnapshot" in html
-    assert "section + '_revision'" in html
+    assert "createContextPoller.buildRevisionQuery(revisions)" in html
     assert "workChanged = contextChanged || notificationsChanged" in html
     assert "renderContextSnapshot(snapshot.context)" in html
     assert "paintEventStream(snapshot.events)" in html