Blame for relay.py

#!/usr/bin/env python3
"""
aichat 訊息中繼（REST）。自架，不經過任何第三方轉發站。

    POST /api/aichat/put                      送一則訊息
    GET  /api/aichat/poll?channel=&since=&wait= 長輪詢取回後續
    GET  /api/aichat/peek?channel=&limit=     不阻塞地抓最近幾則（測試／除錯用）
    GET  /api/aichat/health                   存活與統計

為什麼要有這個
--------------
原本訊息是經過 ntfy.sh 之類的**公開 anyone-can-read 轉發站**：內文明文、topic 誰都能猜，
而且免費站台會限速、會把跑太多的來源 IP 直接封掉且不會自己恢復。訊息走別人的機器，
也就等於把對話交出去。

現在自己架一個：資料只落在這個 Space（SQLite），走 HTTPS，需要 HF token 才能進。
「不會漏掉」也不再靠對方的善意：

* 每則訊息有單調遞增的 ``id``。接收端記住最後收到的 ``id``，重連帶 ``since=<id>``，
  伺服器精確地從那裡接續 —— 不是「重送我沒收到的」，是「從沒斷過的那裡接」。
* 送出端（bridge / 網頁）自己寫磁碟 outbox 才送出，失敗只是延後。
* 就算游標太舊被修剪，bridge 端還有整篇快照（history.json）可以補齊。

## 路由
挂在 app.py 的 FastAPI 上（``app.include_router(relay.router)``），不動 app.py 既有的
``/``、``/files``、``/download``、``/ws``。前綴 ``/api/aichat`` 跟那些都不撞。

## 授權
Space 本身是 private，HF 的閘門已經擋掉沒 token 的人。這個 Space 的 AICHAT_TOKEN 有設定
時，還要再對一次 ``X-Aichat-Token``（雙層）。
"""

from __future__ import annotations

import asyncio
import json
import os
import re
import sqlite3
import time
from pathlib import Path

from fastapi import APIRouter, Header, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

router = APIRouter(prefix="/api/aichat", tags=["aichat"])

VERSION = 2

# ---------------------------------------------------------------- CORS
# 網頁在 aichatv（static Space）、中繼在 testzerogpu2，兩者不同 host → 瀏覽器一定
# 會先送 OPTIONS 預檢。這裡放行指定來源，**不是**萬用 *（萬用 * 不能帶
# Authorization 標頭，實際等於不能用）。
ALLOW_ORIGINS = [
    o.strip() for o in os.environ.get(
        "AICHAT_ORIGINS",
        "https://hellosun-aichatv.static.hf.space,"
        "https://huggingface.co,"
        "http://localhost:8000,http://127.0.0.1:8000"
    ).split(",") if o.strip()
]


def install_cors(app) -> None:
    app.add_middleware(
        CORSMiddleware,
        allow_origins=ALLOW_ORIGINS,
        allow_credentials=False,
        allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
        allow_headers=["Authorization", "Content-Type", "X-Aichat-Token"],
        max_age=600,
    )

# ---------------------------------------------------------------- 設定
DB_PATH = Path(os.environ.get("AICHAT_DB", "/home/user/data/aichat_relay.db"))
SHARED_TOKEN = os.environ.get("AICHAT_TOKEN", "").strip()

KEEP_PER_CHANNEL = int(os.environ.get("AICHAT_KEEP", "2000"))   # 每個 channel 留幾則
CHANNEL_MAX = 64
CHANNEL_RE = re.compile(r"^[A-Za-z0-9._~-]{1,%d}$" % CHANNEL_MAX)
PAYLOAD_MAX = 200_000          # 單則最多 200 KB（附件切片後仍在範圍內）
MAX_WAIT = 25                  # 長輪詢最多吊幾秒：HF 的 proxy 逾時前一定要回


# ---------------------------------------------------------------- 儲存
def _connect() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    # Space 重啟後檔案還在（重建才會沒），所以 WAL + 自動建表
    con = sqlite3.connect(DB_PATH, timeout=10, check_same_thread=False)
    con.execute("PRAGMA journal_mode=WAL")
    con.execute("PRAGMA synchronous=NORMAL")
    con.execute(
        "CREATE TABLE IF NOT EXISTS msgs ("
        " id INTEGER PRIMARY KEY AUTOINCREMENT,"
        " channel TEXT NOT NULL,"
        " payload TEXT NOT NULL,"
        " ts REAL NOT NULL)"
    )
    con.execute("CREATE INDEX IF NOT EXISTS msgs_ch ON msgs(channel, id)")
    return con


_con: sqlite3.Connection | None = None
_lock = asyncio.Lock()
# channel → 一組等待中的未來；有新訊息就把它們叫醒
_waiters: dict[str, set[asyncio.Future]] = {}


def db() -> sqlite3.Connection:
    global _con
    if _con is None:
        _con = _connect()
    return _con


def check_channel(c: str) -> str:
    if not c or not CHANNEL_RE.match(c):
        raise HTTPException(400, f"channel 只能是 1..{CHANNEL_MAX} 個字元，且只限 A-Za-z0-9._~-")
    return c


def prune(channel: str) -> None:
    con = db()
    row = con.execute("SELECT COUNT(*) FROM msgs WHERE channel=?", (channel,)).fetchone()
    n = row[0] if row else 0
    if n > KEEP_PER_CHANNEL:
        con.execute(
            "DELETE FROM msgs WHERE channel=? AND id <= ("
            "  SELECT id FROM msgs WHERE channel=? ORDER BY id DESC LIMIT 1 OFFSET ?)",
            (channel, channel, KEEP_PER_CHANNEL - 1),
        )
        con.commit()


# ---------------------------------------------------------------- 授權
def authorize(x_aichat_token: str | None) -> None:
    """Space 已經是 private（HF 擋掉沒 token 的）。這裡是第二層，設定了才啟用。"""
    if not SHARED_TOKEN:
        return
    if not x_aichat_token or x_aichat_token != SHARED_TOKEN:
        raise HTTPException(401, "X-Aichat-Token 不對")


# ---------------------------------------------------------------- 路由
@router.get("/health")
async def health() -> JSONResponse:
    con = db()
    total = con.execute("SELECT COUNT(*) FROM msgs").fetchone()[0]
    chans = con.execute("SELECT COUNT(DISTINCT channel) FROM msgs").fetchone()[0]
    oldest = con.execute("SELECT MIN(id) FROM msgs").fetchone()[0]
    newest = con.execute("SELECT MAX(id) FROM msgs").fetchone()[0]
    return JSONResponse({
        "ok": True, "version": VERSION,
        "hub": "self-hosted (no third party)",
        "messages": total, "channels": chans,
        "oldest_id": oldest, "newest_id": newest,
        "shared_token": bool(SHARED_TOKEN),
        "waiters": sum(len(v) for v in _waiters.values()),
    })


@router.post("/put")
async def put(request: Request, x_aichat_token: str | None = Header(None)) -> JSONResponse:
    authorize(x_aichat_token)
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(400, "body 不是合法 JSON")
    channel = check_channel(str(body.get("channel") or ""))
    payload = body.get("payload")
    if not isinstance(payload, str):
        raise HTTPException(400, "payload 必須是字串")
    if len(payload.encode("utf-8")) > PAYLOAD_MAX:
        raise HTTPException(413, f"payload 超過 {PAYLOAD_MAX} bytes")

    con = db()
    cur = con.execute("INSERT INTO msgs(channel, payload, ts) VALUES(?,?,?)",
                      (channel, payload, time.time()))
    con.commit()
    mid = cur.lastrowid
    prune(channel)

    # 叫醒在長輪詢的人
    for fut in list(_waiters.get(channel, ())):
        if not fut.done():
            fut.set_result(True)
    return JSONResponse({"ok": True, "id": mid, "ts": time.time()})


@router.get("/peek")
async def peek(channel: str = Query(...), limit: int = Query(50, ge=1, le=500),
               x_aichat_token: str | None = Header(None)) -> JSONResponse:
    """不阻塞：抓最近 limit 則（給測試與除錯用）。"""
    authorize(x_aichat_token)
    channel = check_channel(channel)
    con = db()
    rows = con.execute(
        "SELECT id, payload, ts FROM msgs WHERE channel=? ORDER BY id DESC LIMIT ?",
        (channel, limit)).fetchall()
    items = [{"id": r[0], "payload": r[1], "ts": r[2]} for r in reversed(rows)]
    return JSONResponse({"items": items, "cursor": items[-1]["id"] if items else 0})


@router.get("/poll")
async def poll(channel: str = Query(...),
               since: int = Query(0, ge=0),
               wait: int = Query(20, ge=0, le=MAX_WAIT),
               x_aichat_token: str | None = Header(None)) -> JSONResponse:
    """長輪詢：等到有新訊息或逾時。回傳的 cursor 請記住，下次帶回來。"""
    authorize(x_aichat_token)
    channel = check_channel(channel)
    wait = max(0, min(wait, MAX_WAIT))

    def fetch() -> list:
        con = db()
        rows = con.execute(
            "SELECT id, payload, ts FROM msgs WHERE channel=? AND id>? ORDER BY id ASC LIMIT 500",
            (channel, since)).fetchall()
        return [{"id": r[0], "payload": r[1], "ts": r[2]} for r in rows]

    items = fetch()
    if items or wait == 0:
        return JSONResponse({"items": items,
                             "cursor": items[-1]["id"] if items else since})

    # 沒有東西 → 掛一個 future，有人送訊息時會被 set_result 叫醒
    fut: asyncio.Future = asyncio.get_running_loop().create_future()
    _waiters.setdefault(channel, set()).add(fut)
    try:
        await asyncio.wait_for(fut, timeout=wait)
    except asyncio.TimeoutError:
        pass
    finally:
        _waiters.get(channel, set()).discard(fut)

    items = fetch()
    return JSONResponse({"items": items,
                         "cursor": items[-1]["id"] if items else since})


@router.delete("/channel/{channel}")
async def wipe(channel: str, x_aichat_token: str | None = Header(None)) -> JSONResponse:
    """清掉某個 channel 的全部訊息（測試重置用）。"""
    authorize(x_aichat_token)
    channel = check_channel(channel)
    con = db()
    con.execute("DELETE FROM msgs WHERE channel=?", (channel,))
    con.commit()
    return JSONResponse({"ok": True, "channel": channel})


@router.on_event("startup")
async def _startup() -> None:
    try:
        db()
    except Exception as e:                      # 磁碟壞掉不該讓整個 Space 掛掉
        print(f"[relay] 資料庫開不起來：{e!r}")
    print(f"[relay] v{VERSION} 就緒  db={DB_PATH}  共享 token={'有' if SHARED_TOKEN else '無'}")