webmail-client/plugins/webmail_plugin_labels/webmail_plugin_labels/router.py

"""HTTP surface for the labels plugin, mounted at /api/plugins/labels (SPEC §11).

State is the namespaced KV store: key ``<message_id>`` -> that message's labels,
and key ``_labels`` -> the user's known label set. The KV store is scoped to
(plugin, user_id) by construction, so users never see each other's labels.
"""

from __future__ import annotations

from pathlib import Path

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

from app.plugins.sdk import CurrentUserId, require_csrf, sdk

router = APIRouter()

_UI = Path(__file__).resolve().parent / "ui"
_SDK = sdk("labels")
_ALL_KEY = "_labels"
_BUNDLES = {"actions.js", "sidebar.js"}


class LabelsIn(BaseModel):
    labels: list[str]


class IdsIn(BaseModel):
    ids: list[str]


class NameIn(BaseModel):
    name: str


async def _remember(kv: object, labels: list[str]) -> None:
    """Merge ``labels`` into the user's known-label set."""
    known = set(await kv.get(_ALL_KEY, []))  # type: ignore[attr-defined]
    if not labels or set(labels) <= known:
        return
    known.update(labels)
    await kv.set(_ALL_KEY, sorted(known))  # type: ignore[attr-defined]


@router.get("/message/{message_id}/labels", response_model=list[str])
async def get_labels(message_id: str, user_id: CurrentUserId) -> list[str]:
    return await _SDK.kv(user_id).get(message_id, [])


@router.put("/message/{message_id}/labels", dependencies=[Depends(require_csrf)])
async def set_labels(message_id: str, body: LabelsIn, user_id: CurrentUserId) -> dict[str, object]:
    labels = sorted({name.strip() for name in body.labels if name.strip()})
    kv = _SDK.kv(user_id)
    await kv.set(message_id, labels)
    await _remember(kv, labels)
    return {"status": "ok", "labels": labels}


@router.post("/messages/labels")
async def batch_message_labels(body: IdsIn, user_id: CurrentUserId) -> dict[str, list[str]]:
    """Labels for a set of messages (the reading pane aggregates the thread union)."""
    kv = _SDK.kv(user_id)
    return {mid: await kv.get(mid, []) for mid in body.ids}


@router.put("/thread/{thread_id}/labels", dependencies=[Depends(require_csrf)])
async def set_thread_labels(
    thread_id: str, body: LabelsIn, user_id: CurrentUserId
) -> dict[str, object]:
    """The thread's aggregated labels (union of its messages), written by the reading
    pane so the message list — which has no message ids — can display them."""
    labels = sorted({name.strip() for name in body.labels if name.strip()})
    kv = _SDK.kv(user_id)
    await kv.set(f"t:{thread_id}", labels)
    await _remember(kv, labels)
    return {"status": "ok", "labels": labels}


@router.post("/threads/labels")
async def batch_thread_labels(body: IdsIn, user_id: CurrentUserId) -> dict[str, list[str]]:
    """Aggregated labels for a set of threads (list-view chips). Reconciles first so
    the chips reflect live messages — stale ones from deleted mail are dropped."""
    live = await _reconcile(user_id)
    return {tid: live[tid] for tid in body.ids if live.get(tid)}


_GMAIL_SUFFIX = " [Gmail]"


async def _reconcile(user_id: object) -> dict[str, list[str]]:
    """Self-heal label state from the per-message keys (the source of truth) and
    return the live ``thread_id -> labels`` map.

    Per-message keys (``<message_id>``) are authoritative; the ``t:<thread_id>``
    aggregates and the ``_labels`` set are derived. Messages get deleted (account
    removal) and re-threaded (JWZ merges) out from under us, which orphans thread
    aggregates — so rebuild them here: drop labels of deleted messages, regroup the
    survivors by their *current* thread, rewrite changed/stale ``t:`` keys, and
    rebuild ``_labels`` from labels on live messages. User-created labels (no
    ``… [Gmail]`` suffix) are preserved even when unused; orphaned Gmail labels drop."""
    kv = _SDK.kv(user_id)
    items = await kv.items()
    msg_keys = {
        k: v
        for k, v in items.items()
        if k != _ALL_KEY and not k.startswith("t:") and isinstance(v, list)
    }
    homes = await _SDK.messages.threads_of(user_id, list(msg_keys))

    thread_labels: dict[str, set[str]] = {}
    live_labels: set[str] = set()
    for message_id, labels in msg_keys.items():
        if message_id not in homes:  # message deleted -> drop its orphaned key
            await kv.delete(message_id)
            continue
        live_labels.update(labels)
        thread_id = homes[message_id]
        if thread_id is not None:
            thread_labels.setdefault(thread_id, set()).update(labels)

    # Rewrite only the thread aggregates that changed; delete ones with no messages.
    old_thread_keys = {k[2:] for k in items if k.startswith("t:")}
    result: dict[str, list[str]] = {}
    for thread_id, names in thread_labels.items():
        value = sorted(names)
        result[thread_id] = value
        if items.get(f"t:{thread_id}") != value:
            await kv.set(f"t:{thread_id}", value)
    for thread_id in old_thread_keys - set(thread_labels):
        await kv.delete(f"t:{thread_id}")

    known = set(items.get(_ALL_KEY, []) or [])
    new_known = {n for n in known if not n.endswith(_GMAIL_SUFFIX)} | live_labels
    if new_known != known:
        await kv.set(_ALL_KEY, sorted(new_known))
    return result


@router.get("/labels", response_model=list[str])
async def all_labels(user_id: CurrentUserId) -> list[str]:
    await _reconcile(user_id)  # self-heal from live messages, then return the known set
    return await _SDK.kv(user_id).get(_ALL_KEY, [])


@router.get("/labels/{name}/threads", response_model=list[str])
async def label_threads(name: str, user_id: CurrentUserId) -> list[str]:
    """Thread ids whose aggregate labels include ``name`` — powers the label view in
    the mail list. Reconciles first, so it reflects live messages' current threads."""
    live = await _reconcile(user_id)
    return [tid for tid, labels in live.items() if name in labels]


@router.post("/labels", dependencies=[Depends(require_csrf)])
async def create_label(body: NameIn, user_id: CurrentUserId) -> list[str]:
    name = body.name.strip()
    kv = _SDK.kv(user_id)
    if name:
        await _remember(kv, [name])
    return await kv.get(_ALL_KEY, [])


@router.delete("/labels/{name}", dependencies=[Depends(require_csrf)])
async def delete_label(name: str, user_id: CurrentUserId) -> list[str]:
    """Delete a label everywhere: drop it from the known set and unhook it from every
    message (``<message_id>``) and thread (``t:<thread_id>``) it was attached to."""
    kv = _SDK.kv(user_id)
    for key, labels in (await kv.items()).items():
        if key == _ALL_KEY or not isinstance(labels, list) or name not in labels:
            continue
        await kv.set(key, [n for n in labels if n != name])
    known = [n for n in await kv.get(_ALL_KEY, []) if n != name]
    await kv.set(_ALL_KEY, known)
    return known


@router.get("/ui/{filename}")
async def bundle(filename: str) -> FileResponse:
    if filename not in _BUNDLES:
        raise HTTPException(status_code=404)
    return FileResponse(_UI / filename, media_type="text/javascript")