Translator API as authed WebSocket /_translate/{key} with delta job push

pagerite/app.py

@@ -13,12 +13,12 @@ walking the tree (``resolve``), moves are slot detach/attach
 """
 
 import asyncio
-import base64
 import gzip
 import ipaddress
 import mimetypes
 import os
 import re
+import secrets
 import shutil
 import socket
 import tempfile
@@ -48,9 +48,9 @@ from mediapreview import dispatch
 from pydantic import BaseModel
 from zstandard import ZstdCompressor
 
-from pagerite import analytics, i18n, seed, views
+from pagerite import analytics, i18n, seed, translate, views
 from pagerite.__main__ import DEVMODE
-from pagerite.chunks import chunk_key, needs_translation, store_chunks
+from pagerite.chunks import store_chunks
 from pagerite.data import (
     Data,
     Node,
@@ -287,6 +287,10 @@ def _seed(data: Data) -> None:
 async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
     """Open the database (migrations run inside kanta.open), load assets, load GeoIP."""
     await kanta.open()
+    # The translator service key is persisted; generated once on first boot.
+    if not data.translate_key:
+        with kanta.transaction("generate translate key"):
+            data.translate_key = secrets.token_urlsafe(24)
     await asyncio.to_thread(file_store.load)
     await frontend.load()
     # Decompress/open the DB-IP MMDB once at startup.  Lookups are then
@@ -405,10 +409,12 @@ _render_gen = 0
 
 
 def _invalidate_pages() -> None:
-    """Drop cached page bodies and bump the render generation (ETags)."""
+    """Drop cached page bodies and bump the render generation (ETags);
+    any content change also pushes translation-job deltas to translators."""
     global _render_gen
     _render_gen += 1
     _cached_body.cache_clear()
+    _schedule_translator_notify()
 
 
 @lru_cache(maxsize=128)
@@ -616,7 +622,7 @@ async def update_structure(op: StructureOp) -> None:
 async def get_settings() -> dict:
     """Site-wide settings (brand, theme, custom CSS and favicon URL), plus
     the themes, banner designs and user fonts available on disk for the
-    selectors."""
+    selectors and the translator service key (for the /_translate socket)."""
     return {
         "brand": data.brand,
         "brand_html": data.brand_html,
@@ -628,6 +634,7 @@ async def get_settings() -> dict:
         "fonts": views._user_fonts(),
         "transition": data.transition,
         "transitions": views._transition_names(),
+        "translate_key": data.translate_key,
     }
 
 
@@ -966,122 +973,126 @@ async def delete_page(path: str) -> None:
         _invalidate_pages()
 
 
-def _check_translate_lang(lang: str) -> str:
+class _TranslatorState:
-    """Normalize a ``/_api/translate`` language parameter; translations never
+    """One connected translator socket: announced languages and the
-    target the original language (its text lives in the chunk store, and an
+    (lang, chunk-key) pairs already sent and still outstanding on it.
-    "en" langs key would advertise a bogus ?lang=en alternate)."""
+
-    tag = i18n.base_tag(lang)
+    Per-connection only: a reconnecting client re-receives everything
-    if not tag or tag == i18n.ORIGINAL_LANGUAGE:
+    pending for its languages (its Hello triggers a full Job push)."""
-        raise HTTPException(400, "bad target language")
+
-    return tag
+    def __init__(self, langs: set[str]) -> None:
+        self.langs = langs
+        self.outstanding: set[tuple[str, bytes]] = set()
 
 
-@app.get("/_api/translate/{lang}")
+#: Connected translator sockets and their per-connection state.
-async def translate_pending(lang: str) -> list[dict]:
+_translator_clients: dict[WebSocket, _TranslatorState] = {}
-    """Pending translation items for ``lang`` (translation service API).
 
-    Every page node (published or not) contributes its title and each of
+
-    its chunks that needs translation (``needs_translation``), is not
+def _schedule_translator_notify() -> None:
-    editor-flagged no-translate (``node.no_trans``) and has no ``trans``
+    """Schedule a delta Job push to connected translators, if any.
-    entry for ``lang`` yet. Each item is
+
-    ``{"key", "text", "path", "kind"}``: ``key`` is the base64 of the
+    The hook is _invalidate_pages (sync, called inside transactions): the
-    9-byte chunk hash — the handle the service translates against and
+    task first runs once the current coroutine awaits again, i.e. after the
-    POSTs back; ``text`` the original Markdown (or title); ``path`` the
+    transaction has committed. No-op without a running loop (CLI use).
-    article it came from (no leading slash); ``kind`` "chunk" or "title".
-    Items are deduped by key: content-addressed text (shared paragraphs,
-    repeated titles) is translated once, whichever page it first came from.
     """
-    lang = _check_translate_lang(lang)
+    if not _translator_clients:
-    items = []
-    seen = set()
-
-    def emit(key: bytes, text: str, path: str, kind: str) -> None:
-        if key in seen or lang in data.trans.get(key, {}):
         return
-        seen.add(key)
-        items.append({
-            "key": base64.b64encode(key).decode(),
-            "text": text,
-            "path": path,
-            "kind": kind,
-        })
-
-    def walk(nodes: dict[str, Node], prefix: str) -> None:
-        for slug, node in sorted_nodes(nodes):
-            path = f"{prefix}/{slug}" if prefix else slug
-            if node.chunks is not None:
-                if node.title:
-                    emit(chunk_key(node.title), node.title, path, "title")
-                for h in node.chunks:
-                    text = data.chunks.get(h)
-                    if (
-                        text is not None
-                        and h not in node.no_trans
-                        and needs_translation(text)
-                    ):
-                        emit(h, text, path, "chunk")
-            walk(node.children, path)
-
-    walk(data.menu, "")
-    return items
-
-
-class TranslationItemIn(BaseModel):
-    """One translated fragment submitted by the translation service."""
-
-    key: str  # base64 of the 9-byte chunk hash, as issued by the GET
-    text: str  # the translated Markdown (or title)
-
-
-class TranslationsIn(BaseModel):
-    """Batch of translations for one language."""
-
-    items: list[TranslationItemIn]
-
-
-@app.post("/_api/translate/{lang}")
-async def translate_submit(lang: str, body: TranslationsIn) -> dict:
-    """Store a batch of machine translations for ``lang``, one transaction.
-
-    Keys are the base64 chunk hashes issued by the GET; each text lands in
-    ``trans[key][lang]``. Unknown keys are stored anyway (unreferenced
-    hashes are never read, and the content may simply have moved on since
-    the GET); duplicates within a batch overwrite, last wins; a malformed
-    base64 key rejects the whole batch with 400. Every article that gained
-    at least one entry gets ``node.langs[lang]`` set (the availability
-    index), and cached pages are invalidated. Returns the stored count and
-    the paths of the articles that gained the language.
-    """
-    lang = _check_translate_lang(lang)
     try:
-        entries = [
+        asyncio.get_running_loop()
-            (base64.b64decode(item.key, validate=True), item.text)
+    except RuntimeError:
-            for item in body.items
+        return
+    asyncio.create_task(_notify_translators())
+
+
+async def _notify_translators() -> None:
+    """Push newly-pending items to every connected translator, deltas only:
+    items never sent on that connection and still untranslated."""
+    for ws, state in list(_translator_clients.items()):
+        for lang in sorted(state.langs):
+            items = [
+                item
+                for item in translate.pending_items(data, lang)
+                if (lang, item.key) not in state.outstanding
             ]
-    except ValueError:
+            if not items:
-        raise HTTPException(400, "malformed base64 key") from None
+                continue
-    with kanta.transaction("submit translations", extra=lang):
+            state.outstanding |= {(lang, item.key) for item in items}
-        stored = set()
+            try:
-        for key, text in entries:
+                job = translate.Job(lang=lang, items=items)
-            data.trans.setdefault(key, {})[lang] = text
+                await ws.send_text(msgspec.json.encode(job).decode())
-            stored.add(key)
+            except Exception:  # send failed: the receive loop cleans up
-        pages = []
+                _translator_clients.pop(ws, None)
+                break
 
-        def walk(nodes: dict[str, Node], prefix: str) -> None:
-            for slug, node in sorted_nodes(nodes):
-                path = f"{prefix}/{slug}" if prefix else slug
-                if node.chunks is not None:
-                    keys = set(node.chunks)
-                    if node.title:
-                        keys.add(chunk_key(node.title))
-                    if keys & stored:
-                        node.langs[lang] = True
-                        pages.append(path)
-                walk(node.children, path)
 
-        walk(data.menu, "")
+@app.websocket("/_translate/{clientkey}")
+async def translate_ws(ws: WebSocket, clientkey: str) -> None:
+    """Translator service channel (docs/localization.md).
+
+    Deliberately NOT under /_api/: the external forward-auth is skipped;
+    the server-generated client key in the path is the access control
+    (``Data.translate_key``, generated at startup, shown in the admin's
+    /_api/settings). A wrong/empty key rejects the handshake — closing
+    before accept makes Starlette answer HTTP 403.
+
+    Protocol (JSON frames, msgspec structs in translate.py): the client
+    opens with Hello(langs) (normalized to base subtags; "en"/empty
+    dropped); the server pushes Job(lang, items) — on connect everything
+    pending per language, afterwards deltas on content change — and the
+    client answers with Result(lang, items) batches. A Result for an
+    unannounced or invalid language, or any malformed frame, closes the
+    socket with a protocol error.
+    """
+    if not data.translate_key or clientkey != data.translate_key:
+        await ws.close(code=1008)  # policy violation; pre-accept = HTTP 403
+        return
+    await ws.accept()
+    state: _TranslatorState | None = None
+    try:
+        while True:
+            raw = await ws.receive_text()
+            try:
+                msg = msgspec.json.decode(raw.encode(), type=translate.ClientMsg)
+            except msgspec.DecodeError:
+                await ws.close(code=1002)  # protocol error
+                return
+            if isinstance(msg, translate.Hello):
+                if state is not None:  # one Hello per connection
+                    await ws.close(code=1002)
+                    return
+                langs = {
+                    tag
+                    for lang in msg.langs
+                    if (tag := i18n.base_tag(lang)) and tag != i18n.ORIGINAL_LANGUAGE
+                }
+                state = _TranslatorState(langs)
+                _translator_clients[ws] = state
+                for lang in sorted(langs):
+                    items = translate.pending_items(data, lang)
+                    if not items:
+                        continue
+                    state.outstanding |= {(lang, item.key) for item in items}
+                    job = translate.Job(lang=lang, items=items)
+                    await ws.send_text(msgspec.json.encode(job).decode())
+            else:  # translate.Result
+                lang = i18n.base_tag(msg.lang)
+                if (
+                    state is None  # results before Hello
+                    or not lang
+                    or lang == i18n.ORIGINAL_LANGUAGE
+                    or lang not in state.langs  # not announced in Hello
+                ):
+                    await ws.close(code=1002)
+                    return
+                with kanta.transaction("translator results", extra=lang):
+                    translate.store_results(data, lang, msg.items)
                     _invalidate_pages()
-    return {"stored": len(entries), "pages": pages}
+                state.outstanding -= {(lang, item.key) for item in msg.items}
+    except WebSocketDisconnect:
+        pass
+    finally:
+        _translator_clients.pop(ws, None)
 
 
 _SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$")