Translation dispatcher: single-item jobs, capability Hello, wanted-langs setting

pagerite/app.py

@@ -410,11 +410,11 @@ _render_gen = 0
 
 def _invalidate_pages() -> None:
     """Drop cached page bodies and bump the render generation (ETags);
-    any content change also pushes translation-job deltas to translators."""
+    any content change also re-runs translation dispatch."""
     global _render_gen
     _render_gen += 1
     _cached_body.cache_clear()
-    _schedule_translator_notify()
+    _schedule_translation_dispatch()
 
 
 @lru_cache(maxsize=128)
@@ -622,7 +622,8 @@ 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 and the translator service key (for the /_translate socket)."""
+    selectors, the translator service key and the wanted translation
+    languages (for the /_translate socket)."""
     return {
         "brand": data.brand,
         "brand_html": data.brand_html,
@@ -635,6 +636,7 @@ async def get_settings() -> dict:
         "transition": data.transition,
         "transitions": views._transition_names(),
         "translate_key": data.translate_key,
+        "translate_langs": sorted(data.translate_langs),
     }
 
 
@@ -646,6 +648,7 @@ class SettingsIn(BaseModel):
     custom_css: str
     brand_html: str = ""
     transition: str = "cube"
+    translate_langs: list[str] | None = None  # None keeps the current set
 
 
 @app.put("/_api/settings", status_code=204)
@@ -657,6 +660,12 @@ async def put_settings(settings: SettingsIn) -> None:
         data.theme = settings.theme
         data.custom_css = settings.custom_css
         data.transition = settings.transition
+        if settings.translate_langs is not None:
+            data.translate_langs = {
+                tag: True
+                for lang in settings.translate_langs
+                if (tag := i18n.base_tag(lang)) and tag != i18n.ORIGINAL_LANGUAGE
+            }
         _invalidate_pages()
 
 
@@ -973,28 +982,37 @@ async def delete_page(path: str) -> None:
         _invalidate_pages()
 
 
+# WebSocket API for external translation services (not under /_api: it is keyed
+# with Data.translate_key instead of the SSO forward-auth). The server is a
+# dispatcher: one single-item job at a time per connection, offered in the
+# intersection of the wanted languages (Data.translate_langs) and the
+# connection's announced capabilities. Results are matched to content by
+# chunk key alone.
 class _TranslatorState:
-    """One connected translator socket: announced languages and the
-    (lang, chunk-key) pairs already sent and still outstanding on it.
+    """One connected translator socket: the language codes it announced as
+    capabilities (Hello) and the (lang, chunk-key) job currently in flight
+    on it — one at a time, the next is sent only after its Result.
 
-    Per-connection only: a reconnecting client re-receives everything
-    pending for its languages (its Hello triggers a full Job push)."""
+    Per-connection only: in-flight lives solely here, so on disconnect the
+    item simply becomes pending again and is re-offered to any free capable
+    connection."""
 
-    def __init__(self, langs: set[str]) -> None:
-        self.langs = langs
-        self.outstanding: set[tuple[str, bytes]] = set()
+    def __init__(self, capable: set[str]) -> None:
+        self.capable = capable
+        self.inflight: tuple[str, bytes] | None = None
 
 
 #: Connected translator sockets and their per-connection state.
 _translator_clients: dict[WebSocket, _TranslatorState] = {}
 
 
-def _schedule_translator_notify() -> None:
-    """Schedule a delta Job push to connected translators, if any.
+def _schedule_translation_dispatch() -> None:
+    """Schedule a dispatch pass, if any translator is connected.
 
-    The hook is _invalidate_pages (sync, called inside transactions): the
-    task first runs once the current coroutine awaits again, i.e. after the
-    transaction has committed. No-op without a running loop (CLI use).
+    The content-change hook is _invalidate_pages (sync, called inside
+    transactions): the task first runs once the current coroutine awaits
+    again, i.e. after the transaction has committed. No-op without a
+    running loop (CLI use).
     """
     if not _translator_clients:
         return
@@ -1002,28 +1020,49 @@ def _schedule_translator_notify() -> None:
         asyncio.get_running_loop()
     except RuntimeError:
         return
-    asyncio.create_task(_notify_translators())
+    asyncio.create_task(_dispatch_translations())
 
 
-async def _notify_translators() -> None:
-    """Push newly-pending items to every connected translator, deltas only:
-    items never sent on that connection and still untranslated."""
+async def _dispatch_translations() -> None:
+    """Offer one pending item to every free capable connection.
+
+    Runs on every relevant event: Hello, Result, disconnect and content
+    change (via _invalidate_pages). A connection with no wanted ∩ capable
+    overlap simply stays idle. Pending is derived from the trans store
+    (translate.pending_items) minus the items in flight on any connection.
+    """
+    wanted = {
+        tag
+        for lang in data.translate_langs
+        if (tag := i18n.base_tag(lang)) and tag != i18n.ORIGINAL_LANGUAGE
+    }
+    if not wanted:
+        return
     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
-            ]
-            if not items:
+        if state.inflight is not None:
             continue
-            state.outstanding |= {(lang, item.key) for item in items}
+        langs = wanted & state.capable
+        if not langs:
+            continue
+        inflight = {s.inflight for s in _translator_clients.values() if s.inflight}
+        job = None
+        for lang in sorted(langs):
+            for item in translate.pending_items(data, lang):
+                if (lang, item.key) not in inflight:
+                    job = translate.Job(
+                        lang=lang, key=item.key, text=item.text,
+                        path=item.path, kind=item.kind,
+                    )
+                    break
+            if job is not None:
+                break
+        if job is None:
+            continue
+        state.inflight = (job.lang, job.key)  # before the await: no double-assign
         try:
-                job = translate.Job(lang=lang, items=items)
             await ws.send_text(msgspec.json.encode(job).decode())
         except Exception:  # send failed: the receive loop cleans up
             _translator_clients.pop(ws, None)
-                break
 
 
 @app.websocket("/_translate/{clientkey}")
@@ -1037,12 +1076,13 @@ async def translate_ws(ws: WebSocket, clientkey: str) -> None:
     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.
+    opens with Hello(langs) announcing its CAPABILITIES — the language
+    codes its model can produce (normalized to base subtags; "en"/empty
+    dropped). The dispatcher sends one Job(lang, key, text, path, kind)
+    at a time and waits for the matching Result(lang, key, text) before
+    offering the next. A Result without an in-flight job or with a
+    different (lang, key), a duplicate Hello, 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
@@ -1061,38 +1101,36 @@ async def translate_ws(ws: WebSocket, clientkey: str) -> None:
                 if state is not None:  # one Hello per connection
                     await ws.close(code=1002)
                     return
-                langs = {
+                state = _TranslatorState({
                     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())
+                _schedule_translation_dispatch()
             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
+                    or state.inflight is None  # no job in flight
+                    or (lang, msg.key) != state.inflight  # wrong job
                 ):
                     await ws.close(code=1002)
                     return
                 with kanta.transaction("translator results", extra=lang):
-                    translate.store_results(data, lang, msg.items)
-                    _invalidate_pages()
-                state.outstanding -= {(lang, item.key) for item in msg.items}
+                    paths = translate.store_results(
+                        data, lang, [translate.TransResult(key=msg.key, text=msg.text)]
+                    )
+                    _invalidate_pages()  # schedules the next dispatch
+                state.inflight = None
+                if paths:
+                    print(f"[{lang}] now available for {len(paths)} page(s): {', '.join(sorted(paths))}")
     except WebSocketDisconnect:
         pass
     finally:
-        _translator_clients.pop(ws, None)
+        if _translator_clients.pop(ws, None) is not None:
+            # The in-flight item (if any) is pending again; offer it around.
+            _schedule_translation_dispatch()
 
 
 _SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$")

pagerite/data.py

@@ -105,6 +105,11 @@ class Data(msgspec.Struct):
     #: the external forward-auth does not cover that route). Generated
     #: lazily at startup when empty (see lifespan in app.py).
     translate_key: str = ""
+    #: Wanted target languages for the translator service (presence-keys,
+    #: value always True). The dispatcher offers jobs only in the
+    #: intersection of these and a connection's announced capabilities.
+    #: Read/set via /_api/settings (no editing UI yet).
+    translate_langs: dict[str, bool] = {}
     #: All original-language page text, content-addressed:
     #: chunk_key (9 bytes; base64 at the JSON level) -> Markdown chunk.
     #: Shared by every article.

pagerite/translate.py

@@ -5,8 +5,9 @@ The external machine-translation service connects over WebSocket
 the tagged msgspec structs below (``bytes`` fields ride as base64 — no
 manual encoding anywhere). This module holds the message structs plus the
 two computations shared by the WS handler: which fragments are pending for
-a language (``pending_items``) and storing a batch of results
-(``store_results``).
+a language (``pending_items``) and storing a result (``store_results``).
+The dispatcher itself (one job at a time per connection, wanted ∩ capable
+language matching, requeue on disconnect) lives in app.py.
 """
 
 import msgspec
@@ -16,7 +17,9 @@ from pagerite.data import Data, Node, sorted_nodes
 
 
 class Hello(msgspec.Struct, tag="hello"):
-    """Client greeting on connect: the target languages it handles."""
+    """Client greeting on connect: the language codes its model CAN produce
+    (capabilities). The server offers jobs only in the intersection with
+    the wanted target languages (``Data.translate_langs``)."""
 
     langs: list[str]
 
@@ -31,24 +34,33 @@ class TransItem(msgspec.Struct):
 
 
 class Job(msgspec.Struct, tag="job"):
-    """Server push: pending items for one language."""
+    """Server push: ONE fragment to translate.
+
+    Exactly one job is in flight per connection — the next is sent only
+    after this one's Result. Clients wanting parallelism open multiple
+    connections."""
 
     lang: str
-    items: list[TransItem]
+    key: bytes  #: 9-byte chunk hash (base64 in the JSON frame)
+    text: str
+    path: str  #: article it came from ("" = front page), no leading slash
+    kind: str  #: "chunk" | "title"
 
 
 class TransResult(msgspec.Struct):
-    """One translated fragment."""
+    """One translated fragment (storage level, see store_results)."""
 
     key: bytes
     text: str
 
 
 class Result(msgspec.Struct, tag="result"):
-    """Client reply: a batch of translations for one language."""
+    """Client reply: the translation of the connection's current Job
+    (must match its lang and key exactly)."""
 
     lang: str
-    items: list[TransResult]
+    key: bytes
+    text: str
 
 
 #: Union of the client -> server frames (the "type" tag selects).
@@ -94,13 +106,13 @@ def pending_items(data: Data, lang: str) -> list[TransItem]:
 
 
 def store_results(data: Data, lang: str, items: list[TransResult]) -> list[str]:
-    """Store a batch of machine translations for ``lang``; return the paths
-    of the articles that gained at least one entry.
+    """Store machine translations for ``lang``; return the paths of the
+    articles that gained at least one entry.
 
     Pure data operations: the caller wraps this in a kanta transaction and
     invalidates pages. Unknown keys are stored anyway (unreferenced hashes
     are never read, and the content may simply have moved on since the job
-    was pushed); duplicates within a batch overwrite, last wins. Every
+    was pushed); re-storing an existing key overwrites, last wins. Every
     article that gained an entry gets ``node.langs[lang]`` set (the
     availability index, docs/migrate.md) — because chunks are
     content-addressed, that includes pages merely sharing a fragment.