refactor(api): retire daily talk-time quotas and account auth
This commit is contained in:
@@ -1,229 +0,0 @@
|
||||
"""
|
||||
HF OAuth + per-request identity for the duration limiter.
|
||||
|
||||
Login uses Hugging Face's native Spaces OAuth via `huggingface_hub`
|
||||
(`attach_huggingface_oauth` / `parse_huggingface_oauth`). The OAuth env
|
||||
(`OAUTH_CLIENT_ID`, ...) is injected by the platform when the Space README sets
|
||||
`hf_oauth: true`, so this only activates on a deployed Space — locally and in
|
||||
direct mode there's no OAuth and the limiter treats everyone as anonymous.
|
||||
|
||||
Identity:
|
||||
- signed in -> tier 'pro' | 'free', keyed by hashed HF `sub`
|
||||
- anonymous -> tier 'anon', keyed by BOTH hashed client IP and a hashed
|
||||
signed-cookie id (OR-matched in the limiter)
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
|
||||
import limiter
|
||||
|
||||
logger = logging.getLogger("s2s.auth")
|
||||
|
||||
# huggingface_hub adds these routes when OAuth is attached. Centralised so a
|
||||
# version change is a one-line fix; the paths are handed to the client via
|
||||
# /api/me rather than hardcoded there.
|
||||
OAUTH_LOGIN_PATH = "/oauth/huggingface/login"
|
||||
OAUTH_LOGOUT_PATH = "/oauth/huggingface/logout"
|
||||
|
||||
ANON_COOKIE = "s2s_anon"
|
||||
_COOKIE_MAX_AGE = 60 * 60 * 24 * 30 # 30 days
|
||||
|
||||
|
||||
# Members of these orgs get unlimited usage (like PRO) out of the box. The
|
||||
# UNLIMITED_ORGS env adds to this set; it doesn't replace it.
|
||||
_DEFAULT_UNLIMITED_ORGS = {"cerebras", "huggingfacem4", "smolagents", "pollen-robotics"}
|
||||
|
||||
|
||||
def _unlimited_orgs() -> "set[str]":
|
||||
"""Org usernames whose members get unlimited usage (like PRO).
|
||||
|
||||
Defaults to {cerebras, HuggingFaceM4, smolagents}; the UNLIMITED_ORGS env
|
||||
(comma/space-separated, e.g. `UNLIMITED_ORGS=my-team`) adds more. Matched
|
||||
case-insensitively against the signed-in user's organisations."""
|
||||
raw = os.environ.get("UNLIMITED_ORGS", "")
|
||||
extra = {o.strip().lower() for o in raw.replace(",", " ").split() if o.strip()}
|
||||
return _DEFAULT_UNLIMITED_ORGS | extra
|
||||
|
||||
try:
|
||||
from huggingface_hub import attach_huggingface_oauth, parse_huggingface_oauth
|
||||
_OAUTH_IMPORTABLE = True
|
||||
except Exception as exc: # pragma: no cover - import guard
|
||||
logger.info("huggingface_hub OAuth unavailable (%s); sign-in disabled.", exc)
|
||||
_OAUTH_IMPORTABLE = False
|
||||
|
||||
# Set by attach(): True once OAuth is actually wired (importable + env present).
|
||||
oauth_enabled = False
|
||||
|
||||
|
||||
def attach(app) -> bool:
|
||||
"""Wire HF OAuth onto the app if it's importable and configured. Returns
|
||||
whether sign-in is available."""
|
||||
global oauth_enabled
|
||||
if not _OAUTH_IMPORTABLE or not os.environ.get("OAUTH_CLIENT_ID"):
|
||||
return False
|
||||
try:
|
||||
attach_huggingface_oauth(app)
|
||||
oauth_enabled = True
|
||||
logger.info("HF OAuth attached (sign-in enabled).")
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
logger.warning("Failed to attach HF OAuth: %r", exc)
|
||||
oauth_enabled = False
|
||||
return oauth_enabled
|
||||
|
||||
|
||||
def _field(obj, name, default=None):
|
||||
"""Read a field whether the user-info is an object or a dict."""
|
||||
if obj is None:
|
||||
return default
|
||||
if isinstance(obj, dict):
|
||||
return obj.get(name, default)
|
||||
return getattr(obj, name, default)
|
||||
|
||||
|
||||
# Surface what we detect (orgs, tier) in logs and on /api/me when set. Handy for
|
||||
# verifying org gating on the live Space without guessing.
|
||||
AUTH_DEBUG = bool(os.environ.get("AUTH_DEBUG"))
|
||||
|
||||
# whoami-v2 org lookups are cached for the process lifetime, keyed by token, so
|
||||
# /api/me + /api/session don't each hit the Hub.
|
||||
_orgs_cache: "dict[str, set[str]]" = {}
|
||||
|
||||
|
||||
def current_oauth(request):
|
||||
"""The parsed HF OAuth info (user_info + access_token), or None."""
|
||||
if not oauth_enabled:
|
||||
return None
|
||||
try:
|
||||
return parse_huggingface_oauth(request)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def current_user(request):
|
||||
"""The signed-in HF user-info, or None."""
|
||||
return _field(current_oauth(request), "user_info")
|
||||
|
||||
|
||||
def _user_org_names(user) -> "set[str]":
|
||||
"""The user's organisations from the OAuth userinfo, by username/name/id."""
|
||||
names = set()
|
||||
for org in _field(user, "orgs", []) or []:
|
||||
for key in ("preferred_username", "name", "sub"):
|
||||
val = _field(org, key)
|
||||
if val:
|
||||
names.add(str(val).lower())
|
||||
return names
|
||||
|
||||
|
||||
def _orgs_via_token(token: str) -> "set[str]":
|
||||
"""Fallback org lookup via the Hub `whoami-v2` API, using the user's OAuth
|
||||
access token. Covers the case where the userinfo claim omits `orgs`."""
|
||||
if not token:
|
||||
return set()
|
||||
if token in _orgs_cache:
|
||||
return _orgs_cache[token]
|
||||
names: "set[str]" = set()
|
||||
try:
|
||||
import httpx
|
||||
|
||||
resp = httpx.get(
|
||||
"https://huggingface.co/api/whoami-v2",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
timeout=5.0,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
for org in resp.json().get("orgs", []) or []:
|
||||
for key in ("name", "fullname"):
|
||||
val = org.get(key)
|
||||
if val:
|
||||
names.add(str(val).lower())
|
||||
except Exception as exc: # pragma: no cover - network/permission dependent
|
||||
logger.info("whoami-v2 org lookup failed: %r", exc)
|
||||
_orgs_cache[token] = names
|
||||
return names
|
||||
|
||||
|
||||
def _org_names(user, token=None, allow=None) -> "set[str]":
|
||||
"""The user's org usernames from the OAuth userinfo claim. If that doesn't
|
||||
already satisfy `allow`, fall back to the Hub `whoami-v2` API (the claim is
|
||||
often empty or partial), so membership is resolved either way."""
|
||||
names = _user_org_names(user)
|
||||
if token and (allow is None or not (allow & names)):
|
||||
names = names | _orgs_via_token(token)
|
||||
return names
|
||||
|
||||
|
||||
def resolve_tier(user, token=None) -> str:
|
||||
"""Tier for a signed-in user: 'pro' (paying), 'org' (allow-listed org
|
||||
member, unlimited), or 'free'. PRO wins over org if both apply."""
|
||||
if bool(_field(user, "is_pro", False)):
|
||||
return "pro"
|
||||
allow = _unlimited_orgs()
|
||||
names = _org_names(user, token, allow)
|
||||
tier = "org" if (allow & names) else "free"
|
||||
if AUTH_DEBUG:
|
||||
logger.info("tier=%s orgs=%s allow=%s", tier, sorted(names), sorted(allow))
|
||||
return tier
|
||||
|
||||
|
||||
def user_view(request) -> dict:
|
||||
"""Public profile for /api/me."""
|
||||
info = current_oauth(request)
|
||||
user = _field(info, "user_info")
|
||||
if not user:
|
||||
return {"loggedIn": False, "tier": "anon"}
|
||||
token = _field(info, "access_token")
|
||||
out = {
|
||||
"loggedIn": True,
|
||||
"username": _field(user, "preferred_username") or _field(user, "name") or "tú",
|
||||
"avatar": _field(user, "picture"),
|
||||
"tier": resolve_tier(user, token),
|
||||
}
|
||||
if AUTH_DEBUG:
|
||||
out["orgs"] = sorted(_org_names(user, token))
|
||||
return out
|
||||
|
||||
|
||||
def _client_ip(request) -> str:
|
||||
"""Real client IP. On HF the app sits behind a proxy, so the user's address
|
||||
is the first hop in X-Forwarded-For, not request.client.host."""
|
||||
xff = request.headers.get("x-forwarded-for", "")
|
||||
if xff:
|
||||
return xff.split(",")[0].strip()
|
||||
return request.client.host if request.client else "unknown"
|
||||
|
||||
|
||||
def resolve_identity(request):
|
||||
"""Resolve (tier, keys, set_cookie) for this request.
|
||||
|
||||
`keys` are the limiter usage_daily keys to debit (one for signed-in, two for
|
||||
anonymous). `set_cookie` is a signed value to Set-Cookie when we minted a new
|
||||
anonymous id, else None.
|
||||
"""
|
||||
info = current_oauth(request)
|
||||
user = _field(info, "user_info")
|
||||
if user:
|
||||
sub = _field(user, "sub") or _field(user, "preferred_username")
|
||||
token = _field(info, "access_token")
|
||||
return resolve_tier(user, token), [limiter.hash_key(f"sub:{sub}")], None
|
||||
|
||||
# Anonymous: key by IP and a signed cookie id, minting the cookie if absent.
|
||||
ip = _client_ip(request)
|
||||
cookie_id = limiter.verify_cookie(request.cookies.get(ANON_COOKIE, ""))
|
||||
set_cookie = None
|
||||
if not cookie_id:
|
||||
cookie_id = secrets.token_urlsafe(18)
|
||||
set_cookie = limiter.sign_cookie(cookie_id)
|
||||
keys = [limiter.hash_key(f"ip:{ip}"), limiter.hash_key(f"cookie:{cookie_id}")]
|
||||
return "anon", keys, set_cookie
|
||||
|
||||
|
||||
def set_anon_cookie(response, signed: str) -> None:
|
||||
# The Space runs inside an iframe on huggingface.co, so the cookie lives in a
|
||||
# cross-site context — it must be SameSite=None; Secure or the browser drops it.
|
||||
response.set_cookie(
|
||||
ANON_COOKIE, signed,
|
||||
max_age=_COOKIE_MAX_AGE, httponly=True, samesite="none", secure=True,
|
||||
)
|
||||
@@ -1,276 +0,0 @@
|
||||
"""
|
||||
Per-day talk-time budget for the speech-to-speech demo.
|
||||
|
||||
Our server isn't in the audio path (the browser dials the compute WebSocket
|
||||
directly), so it can't cut a live stream. What it *can* do is meter time with a
|
||||
server-clock, chunked reservation:
|
||||
|
||||
- At grant we reserve the first chunk (CHUNK_SEC) and debit it from the day's
|
||||
budget. A parallel grant therefore sees the budget already spent.
|
||||
- The client heartbeats; each heartbeat extends the reservation one chunk at a
|
||||
time until the daily budget runs out, then we report `expired` so the client
|
||||
tears down.
|
||||
- On a clean end (sendBeacon) we reconcile to the real elapsed time and refund
|
||||
the unused chunk. A crash (no end, no heartbeats) is reaped by a sweep and
|
||||
forfeits at most one chunk.
|
||||
|
||||
All time is the server's clock. Budgets are per UTC day; a new day is simply a
|
||||
new row (no explicit reset). Logged-in users are keyed by a hashed HF `sub`;
|
||||
anonymous users by BOTH a hashed IP and a hashed signed-cookie id, OR-matched
|
||||
(spent = max of the two) so clearing one identifier doesn't reset the budget.
|
||||
|
||||
Storage is SQLite at $USAGE_DB_PATH, else /data (persistent Spaces storage),
|
||||
else a /tmp fallback. On /tmp the budget is only per-uptime — flagged in logs.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger("s2s.limiter")
|
||||
|
||||
# ── Tunables (env-overridable) ───────────────────────────────────────────────
|
||||
ANON_SEC = int(os.environ.get("LIMIT_ANON_SEC", "300")) # 5 min/day, not signed in
|
||||
FREE_SEC = int(os.environ.get("LIMIT_FREE_SEC", "600")) # 10 min/day, signed in, no PRO
|
||||
CHUNK_SEC = int(os.environ.get("RESERVE_CHUNK_SEC", "10")) # reservation granularity
|
||||
HEARTBEAT_SEC = int(os.environ.get("HEARTBEAT_SEC", "5")) # advertised client cadence
|
||||
REAP_AFTER_SEC = int(os.environ.get("SESSION_REAP_SEC", "15")) # silence before sweep
|
||||
|
||||
# Stable across restarts or the hashed keys (and signed cookies) rotate and the
|
||||
# budget effectively resets. Set it as a Space secret. Falls back to a per-boot
|
||||
# random value (keys then only hold within one uptime).
|
||||
_HASH_SECRET = (os.environ.get("USAGE_HASH_SECRET", "").strip() or os.urandom(32).hex()).encode()
|
||||
|
||||
_lock = threading.Lock()
|
||||
_db_path: "Path | None" = None
|
||||
|
||||
|
||||
def budget_for(tier: str) -> "int | None":
|
||||
"""Daily second-budget for a tier, or None for unlimited.
|
||||
|
||||
Unlimited tiers: 'pro' (paying PRO members) and 'org' (members of an
|
||||
allow-listed organisation, see UNLIMITED_ORGS in auth.py)."""
|
||||
if tier in ("pro", "org"):
|
||||
return None
|
||||
if tier == "free":
|
||||
return FREE_SEC
|
||||
return ANON_SEC
|
||||
|
||||
|
||||
def hash_key(raw: str) -> str:
|
||||
"""HMAC a raw identifier (sub / ip / cookie id) into an opaque storage key."""
|
||||
digest = hmac.new(_HASH_SECRET, raw.encode("utf-8"), hashlib.sha256).hexdigest()
|
||||
return f"k_{digest}"
|
||||
|
||||
|
||||
def sign_cookie(value: str) -> str:
|
||||
"""`<id>.<sig>` so a forged anon-cookie id is rejected on read."""
|
||||
sig = hmac.new(_HASH_SECRET, value.encode("utf-8"), hashlib.sha256).hexdigest()[:32]
|
||||
return f"{value}.{sig}"
|
||||
|
||||
|
||||
def verify_cookie(signed: str) -> "str | None":
|
||||
"""Return the id if the signature checks out, else None."""
|
||||
if not signed or "." not in signed:
|
||||
return None
|
||||
value, _, sig = signed.rpartition(".")
|
||||
want = hmac.new(_HASH_SECRET, value.encode("utf-8"), hashlib.sha256).hexdigest()[:32]
|
||||
return value if hmac.compare_digest(sig, want) else None
|
||||
|
||||
|
||||
def _today() -> str:
|
||||
return datetime.now(timezone.utc).date().isoformat()
|
||||
|
||||
|
||||
def _resolve_db_path() -> Path:
|
||||
explicit = os.environ.get("USAGE_DB_PATH", "").strip()
|
||||
if explicit:
|
||||
return Path(explicit)
|
||||
data = Path("/data")
|
||||
if data.is_dir() and os.access(data, os.W_OK):
|
||||
return data / "s2s-usage.sqlite3"
|
||||
logger.warning("No persistent /data — usage budget falls back to /tmp (per-uptime only).")
|
||||
return Path(tempfile.gettempdir()) / "s2s-usage.sqlite3"
|
||||
|
||||
|
||||
def _connect() -> sqlite3.Connection:
|
||||
con = sqlite3.connect(_db_path, timeout=5.0)
|
||||
con.execute("PRAGMA journal_mode=WAL")
|
||||
con.execute("PRAGMA busy_timeout=5000")
|
||||
return con
|
||||
|
||||
|
||||
def init() -> None:
|
||||
"""Create the schema. Call once at startup."""
|
||||
global _db_path
|
||||
_db_path = _resolve_db_path()
|
||||
with _lock, _connect() as con:
|
||||
con.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS usage_daily (
|
||||
user_key TEXT NOT NULL,
|
||||
day TEXT NOT NULL,
|
||||
spent_sec INTEGER NOT NULL DEFAULT 0,
|
||||
updated_at INTEGER NOT NULL,
|
||||
PRIMARY KEY (user_key, day)
|
||||
)"""
|
||||
)
|
||||
con.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS sessions (
|
||||
session_id TEXT PRIMARY KEY,
|
||||
keys TEXT NOT NULL, -- comma-joined usage_daily keys to debit
|
||||
day TEXT NOT NULL,
|
||||
tier TEXT NOT NULL,
|
||||
grant_ts REAL NOT NULL,
|
||||
last_seen_ts REAL NOT NULL,
|
||||
reserved_sec INTEGER NOT NULL,
|
||||
ended INTEGER NOT NULL DEFAULT 0
|
||||
)"""
|
||||
)
|
||||
logger.info("Usage limiter ready at %s (anon=%ss free=%ss chunk=%ss)", _db_path, ANON_SEC, FREE_SEC, CHUNK_SEC)
|
||||
|
||||
|
||||
# ── Internal helpers (call under _lock) ───────────────────────────────────────
|
||||
|
||||
def _spent(con, key: str, day: str) -> int:
|
||||
row = con.execute(
|
||||
"SELECT spent_sec FROM usage_daily WHERE user_key=? AND day=?", (key, day)
|
||||
).fetchone()
|
||||
return int(row[0]) if row else 0
|
||||
|
||||
|
||||
def _spent_max(con, keys, day: str) -> int:
|
||||
"""OR-match: the most-spent identifier governs."""
|
||||
return max((_spent(con, k, day) for k in keys), default=0)
|
||||
|
||||
|
||||
def _add(con, keys, day: str, delta: int) -> None:
|
||||
now = int(time.time())
|
||||
for k in keys:
|
||||
cur = _spent(con, k, day)
|
||||
nxt = max(0, cur + delta)
|
||||
con.execute(
|
||||
"""INSERT INTO usage_daily (user_key, day, spent_sec, updated_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(user_key, day) DO UPDATE SET
|
||||
spent_sec = excluded.spent_sec, updated_at = excluded.updated_at""",
|
||||
(k, day, nxt, now),
|
||||
)
|
||||
|
||||
|
||||
# ── Public API ────────────────────────────────────────────────────────────────
|
||||
|
||||
def remaining(keys, tier: str) -> "int | None":
|
||||
"""Seconds left today for these keys (None = unlimited). No mutation."""
|
||||
budget = budget_for(tier)
|
||||
if budget is None:
|
||||
return None
|
||||
with _lock, _connect() as con:
|
||||
return max(0, budget - _spent_max(con, keys, _today()))
|
||||
|
||||
|
||||
def begin(session_id: str, keys, tier: str) -> int:
|
||||
"""Reserve the first chunk for a new session and record it. Returns the
|
||||
chunk reserved (0 if the budget is already exhausted — the first heartbeat
|
||||
will then expire it). PRO (unlimited) is never tracked; don't call it here."""
|
||||
day = _today()
|
||||
budget = budget_for(tier)
|
||||
now = time.time()
|
||||
with _lock, _connect() as con:
|
||||
avail = budget - _spent_max(con, keys, day) if budget is not None else CHUNK_SEC
|
||||
chunk = max(0, min(CHUNK_SEC, avail))
|
||||
if chunk:
|
||||
_add(con, keys, day, chunk)
|
||||
con.execute(
|
||||
"""INSERT OR REPLACE INTO sessions
|
||||
(session_id, keys, day, tier, grant_ts, last_seen_ts, reserved_sec, ended)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, 0)""",
|
||||
(session_id, ",".join(keys), day, tier, now, now, chunk),
|
||||
)
|
||||
return chunk
|
||||
|
||||
|
||||
def heartbeat(session_id: str) -> bool:
|
||||
"""Keep a session alive: extend the reservation toward `elapsed + 1 chunk`,
|
||||
debiting the budget chunk by chunk. Returns True while alive, False once the
|
||||
budget is spent (caller should tear down) or the session is unknown/ended."""
|
||||
now = time.time()
|
||||
with _lock, _connect() as con:
|
||||
row = con.execute(
|
||||
"SELECT keys, day, tier, grant_ts, reserved_sec, ended FROM sessions WHERE session_id=?",
|
||||
(session_id,),
|
||||
).fetchone()
|
||||
if not row or row[5]:
|
||||
return False
|
||||
keys = row[0].split(",")
|
||||
day, tier, grant_ts, reserved = row[1], row[2], row[3], int(row[4])
|
||||
budget = budget_for(tier)
|
||||
elapsed = now - grant_ts
|
||||
|
||||
# Grow the reservation one chunk at a time until it covers elapsed + a
|
||||
# one-chunk lookahead, or the budget runs dry.
|
||||
while reserved < elapsed + CHUNK_SEC:
|
||||
if budget is not None and _spent_max(con, keys, day) >= budget:
|
||||
break
|
||||
_add(con, keys, day, CHUNK_SEC)
|
||||
reserved += CHUNK_SEC
|
||||
|
||||
alive = reserved > elapsed # could we cover the time already elapsed?
|
||||
con.execute(
|
||||
"UPDATE sessions SET last_seen_ts=?, reserved_sec=?, ended=? WHERE session_id=?",
|
||||
(now, reserved, 0 if alive else 1, session_id),
|
||||
)
|
||||
if not alive:
|
||||
_reconcile(con, keys, day, grant_ts, reserved, end_ts=now)
|
||||
return alive
|
||||
|
||||
|
||||
def end(session_id: str) -> None:
|
||||
"""Clean teardown: reconcile to actual elapsed time and refund the unused
|
||||
reservation. Idempotent."""
|
||||
now = time.time()
|
||||
with _lock, _connect() as con:
|
||||
row = con.execute(
|
||||
"SELECT keys, day, grant_ts, reserved_sec, ended FROM sessions WHERE session_id=?",
|
||||
(session_id,),
|
||||
).fetchone()
|
||||
if not row or row[4]:
|
||||
return
|
||||
keys, day, grant_ts, reserved = row[0].split(","), row[1], row[2], int(row[3])
|
||||
_reconcile(con, keys, day, grant_ts, reserved, end_ts=now)
|
||||
con.execute("UPDATE sessions SET ended=1, last_seen_ts=? WHERE session_id=?", (now, session_id))
|
||||
|
||||
|
||||
def sweep() -> None:
|
||||
"""Reap sessions that went silent (crash / closed without a beacon): bill
|
||||
their elapsed time, refund the rest, mark ended. Forfeits ≤ one chunk."""
|
||||
now = time.time()
|
||||
cutoff = now - REAP_AFTER_SEC
|
||||
with _lock, _connect() as con:
|
||||
stale = con.execute(
|
||||
"SELECT session_id, keys, day, grant_ts, last_seen_ts, reserved_sec FROM sessions "
|
||||
"WHERE ended=0 AND last_seen_ts < ?",
|
||||
(cutoff,),
|
||||
).fetchall()
|
||||
for session_id, keys_s, day, grant_ts, last_seen, reserved in stale:
|
||||
_reconcile(con, keys_s.split(","), day, grant_ts, int(reserved), end_ts=last_seen)
|
||||
con.execute("UPDATE sessions SET ended=1 WHERE session_id=?", (session_id,))
|
||||
if stale:
|
||||
logger.debug("swept %d stale session(s)", len(stale))
|
||||
|
||||
|
||||
def _reconcile(con, keys, day: str, grant_ts: float, reserved: int, end_ts: float) -> None:
|
||||
"""Refund reserved-but-unused time. Bill elapsed rounded up to a chunk,
|
||||
capped at what was reserved."""
|
||||
elapsed = max(0.0, end_ts - grant_ts)
|
||||
billed = min(reserved, int(math.ceil(elapsed / CHUNK_SEC) * CHUNK_SEC))
|
||||
refund = reserved - billed
|
||||
if refund > 0:
|
||||
_add(con, keys, day, -refund)
|
||||
@@ -2,5 +2,3 @@ fastapi==0.115.6
|
||||
uvicorn[standard]==0.34.0
|
||||
httpx==0.28.1
|
||||
python-dotenv==1.2.2
|
||||
# [oauth] extra pulls authlib + itsdangerous for attach_huggingface_oauth.
|
||||
huggingface_hub[oauth]>=0.30,<1.0
|
||||
|
||||
+27
-256
@@ -1,43 +1,24 @@
|
||||
"""
|
||||
Tiny server for the speech-to-speech demo.
|
||||
API and WebSocket proxy for the speech-to-speech frontend.
|
||||
|
||||
The demo used to ship as a `sdk: static` Space, but the web-search tool needs a
|
||||
search key the browser must NOT see. A static Space has no runtime process, so it
|
||||
can't hold a secret the front-end uses. This server fixes that: it serves the
|
||||
unchanged front-end AND exposes a same-origin `/api/search` proxy that holds the
|
||||
Serper key server-side (see docs/adr/0001).
|
||||
|
||||
Everything lives in one container; the speech-to-speech backend stays a separate,
|
||||
load-balanced service the browser talks to over WebSocket as before. The load
|
||||
balancer's address is a secret too (like the Serper key): the browser never sees
|
||||
it. `/api/session` proxies the session handshake server-side so only the
|
||||
per-session compute URL the LB hands back (which the browser must dial) is exposed.
|
||||
|
||||
On the deployed Space the server also meters conversation time by HF login tier
|
||||
(anonymous / signed-in / PRO) — see `limiter.py` and `auth.py`. That whole feature
|
||||
is off unless BOTH `LOAD_BALANCER_URL` and `SPACE_ID` are set, so it runs only on
|
||||
the live Space, never locally (even with the LB exported for testing).
|
||||
|
||||
`SPEECH_TO_SPEECH_URL` overrides everything: when set, the LB logic above is
|
||||
disabled entirely (no session proxy, no queue, no metering, no sign-in) and the
|
||||
browser connects directly to that URL, shown read-only in Settings.
|
||||
Vite serves the React application and forwards backend requests here. Search,
|
||||
caller-profile, and tool credentials stay server-side. `/api/session` proxies
|
||||
load-balancer allocation so the browser receives only the per-session compute
|
||||
URL it must dial. A pinned direct URL or WebSocket proxy target disables the LB.
|
||||
|
||||
Endpoints:
|
||||
GET /api/config -> { search, lb, allowDirect, s2sUrl, auth }
|
||||
GET /api/me -> login + tier + remaining budget (LB mode only)
|
||||
GET /api/config -> frontend settings and available integrations
|
||||
POST /api/search -> { results, answer } Google via Serper.dev
|
||||
POST /api/session -> proxies <LB>/session: a grant, or a queue ticket
|
||||
GET /api/queue/{id} -> proxies <LB>/queue/{id}: position, or a grant on claim
|
||||
DELETE /api/queue/{id} -> leave the queue (explicit "Leave queue" button)
|
||||
POST /api/queue/end -> leave the queue (sendBeacon on teardown)
|
||||
POST /api/session/heartbeat-> extend the reservation; { expired }
|
||||
POST /api/session/end -> reconcile + refund (sendBeacon on teardown)
|
||||
/* -> static files (index.html, main.js, ...)
|
||||
WS /v1/realtime -> optional speech-service WebSocket proxy
|
||||
GET /docs/* -> shared project documentation
|
||||
|
||||
When every compute slot is busy the load balancer hands back a queue ticket
|
||||
instead of a grant; the browser polls /api/queue/{id} until it reaches the front
|
||||
and a slot frees. Waiting reserves nothing — the daily budget is only reserved at
|
||||
the moment a slot is actually claimed (a grant), never while queued.
|
||||
and a slot frees. Queue deadlines and pending-grant timeouts belong to the LB.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -58,9 +39,6 @@ from fastapi.responses import JSONResponse, RedirectResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel
|
||||
|
||||
import auth
|
||||
import limiter
|
||||
|
||||
logger = logging.getLogger("s2s")
|
||||
_logging_configured = False
|
||||
if not _logging_configured:
|
||||
@@ -1402,9 +1380,9 @@ NODE_RED_REASONING_URL = os.environ.get("S2S_REASONING_NODE_RED_URL", "").strip(
|
||||
# browser connects to it straight (no load balancer).
|
||||
LOAD_BALANCER_URL = os.environ.get("LOAD_BALANCER_URL", "").strip()
|
||||
# Direct s2s server URL pinned by the deploy. Takes priority over the load
|
||||
# balancer: when set, ALL LB logic is disabled (no /api/session proxy, no queue,
|
||||
# no limiter, no sign-in) and the browser connects to this URL directly. Unlike
|
||||
# the LB address it is NOT a secret — /api/config sends it to the client, which
|
||||
# balancer: when set, LB allocation and queuing are disabled, and the browser
|
||||
# connects to this URL directly. Unlike the LB address it is NOT a secret —
|
||||
# /api/config sends it to the client, which
|
||||
# shows it read-only in Settings.
|
||||
SPEECH_TO_SPEECH_URL = os.environ.get("SPEECH_TO_SPEECH_URL", "").strip()
|
||||
# Optional same-origin WebSocket relay target. When set, the browser dials this
|
||||
@@ -1413,13 +1391,6 @@ SPEECH_TO_SPEECH_URL = os.environ.get("SPEECH_TO_SPEECH_URL", "").strip()
|
||||
S2S_PROXY_TARGET = os.environ.get("S2S_PROXY_TARGET", "").strip()
|
||||
if SPEECH_TO_SPEECH_URL or S2S_PROXY_TARGET:
|
||||
LOAD_BALANCER_URL = ""
|
||||
# HF injects SPACE_ID ("owner/space") into every Space runtime; it's absent
|
||||
# locally and on a plain `docker run`. We meter conversation time ONLY on the
|
||||
# deployed Space — i.e. when BOTH the LB is configured AND we're on a Space.
|
||||
# Off-Space (local dev, even with the LB exported) the app still proxies the LB,
|
||||
# but nothing is metered: no budget, no reservations, no sign-in gating.
|
||||
SPACE_ID = os.environ.get("SPACE_ID", "").strip()
|
||||
LIMITER_ENABLED = bool(LOAD_BALANCER_URL) and bool(SPACE_ID)
|
||||
SERPER_URL = "https://google.serper.dev/search"
|
||||
# Cap results so the tool output stays small enough to feed back to the model.
|
||||
MAX_RESULTS = 5
|
||||
@@ -1427,10 +1398,6 @@ HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
app = FastAPI(title="s2s-demo")
|
||||
|
||||
# Wire HF OAuth before the app serves (no-op unless the OAuth env is present).
|
||||
# Sign-in only matters when we're metering (prod Space), so gate it on that.
|
||||
AUTH_ENABLED = LIMITER_ENABLED and auth.attach(app)
|
||||
|
||||
# Resolve the execution URL for the startup banner using the same normalization
|
||||
# as _apipilot_resolve_execute_url.
|
||||
def _resolve_exec_url_for_startup_banner():
|
||||
@@ -1444,7 +1411,7 @@ _STARTUP_APIPILOT_EXEC_URL = _resolve_exec_url_for_startup_banner()
|
||||
del _resolve_exec_url_for_startup_banner
|
||||
|
||||
logger.info(
|
||||
"[STARTUP] Other externals — reasoning=%s nodeRed=%s serper=%s lb=%s s2sProxy=%s s2sDirect=%s userAuth=%s reportsApi=%s injectReports=%s reportsHintMax=%d reportsInjectMode=%s injectToday=%s summaryApi=%s summaryText=%s loginApi=%s injectProfile=%s debugPrompt=%s startSessionMode=%s theme=%s auth=%s apipilotToolsUrl=%s apipilotExecuteUrlOverride=%s apipilotServiceToken=%s apipilotTtl=%ds apipilotTimeout=%.1fs apipilotReasoningPass2=%s apipilotApiKeyBypass=%s",
|
||||
"[STARTUP] Other externals — reasoning=%s nodeRed=%s serper=%s lb=%s s2sProxy=%s s2sDirect=%s userAuth=%s reportsApi=%s injectReports=%s reportsHintMax=%d reportsInjectMode=%s injectToday=%s summaryApi=%s summaryText=%s loginApi=%s injectProfile=%s debugPrompt=%s startSessionMode=%s theme=%s apipilotToolsUrl=%s apipilotExecuteUrlOverride=%s apipilotServiceToken=%s apipilotTtl=%ds apipilotTimeout=%.1fs apipilotReasoningPass2=%s apipilotApiKeyBypass=%s",
|
||||
S2S_REASONING_TYPE,
|
||||
NODE_RED_REASONING_URL or "-",
|
||||
bool(SERPER_KEY),
|
||||
@@ -1464,7 +1431,6 @@ logger.info(
|
||||
S2S_DEBUG_PROMPT,
|
||||
S2S_START_SESSION_MODE or "-",
|
||||
S2S_THEME or "-",
|
||||
AUTH_ENABLED,
|
||||
S2S_APIPILOT_TOOLS_URL or "-",
|
||||
(S2S_APIPILOT_TOOLS_EXECUTE_URL + f"(resolved:{_STARTUP_APIPILOT_EXEC_URL})") if S2S_APIPILOT_TOOLS_EXECUTE_URL else "-",
|
||||
"Y" if bool(S2S_APIPILOT_SERVICE_TOKEN) else "-",
|
||||
@@ -1474,10 +1440,6 @@ logger.info(
|
||||
("YES head=" + (S2S_APIPILOT_API_KEY[:6] + "..." if len(S2S_APIPILOT_API_KEY) > 6 else S2S_APIPILOT_API_KEY) + " (3 FORMS!)") if S2S_APIPILOT_API_KEY else "-",
|
||||
)
|
||||
del _STARTUP_APIPILOT_EXEC_URL
|
||||
logger.debug(
|
||||
"[STARTUP] Runtime config — SPACE_ID=%s LIMITER_ENABLED=%s MAX_RESULTS=%d",
|
||||
bool(SPACE_ID), LIMITER_ENABLED, MAX_RESULTS,
|
||||
)
|
||||
|
||||
|
||||
async def _report_server_debug(hypothesis_id: str, location: str, msg: str, data: dict | None = None) -> None:
|
||||
@@ -1652,22 +1614,6 @@ _NOISY_CLIENT_EVENT_TYPES = frozenset({
|
||||
"input_audio_buffer.speech_started",
|
||||
"input_audio_buffer.speech_stopped",
|
||||
})
|
||||
@app.on_event("startup")
|
||||
async def _startup():
|
||||
"""Stand up the usage DB and a periodic sweeper — metered (prod Space) only."""
|
||||
if not LIMITER_ENABLED:
|
||||
return
|
||||
limiter.init()
|
||||
asyncio.create_task(_sweeper())
|
||||
|
||||
|
||||
async def _sweeper():
|
||||
while True:
|
||||
await asyncio.sleep(limiter.REAP_AFTER_SEC)
|
||||
try:
|
||||
await asyncio.to_thread(limiter.sweep)
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
logger.warning("usage sweep failed: %r", exc)
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
@@ -3334,8 +3280,8 @@ def _effective_s2s_url(request: Request) -> str:
|
||||
@app.get("/api/config")
|
||||
async def config(request: Request):
|
||||
"""Client bootstrap: whether web search is available, whether the deploy runs
|
||||
behind a load balancer (so the browser uses the /api/session proxy + limiter),
|
||||
whether HF sign-in is available, and whether the user may instead set a direct
|
||||
behind a load balancer (so the browser uses the /api/session proxy),
|
||||
and whether the user may instead set a direct
|
||||
s2s server URL. The LB address itself is intentionally NOT included.
|
||||
|
||||
2026-09-24: new top-level `apipilot` flag. When True the deploy exposes the
|
||||
@@ -3351,7 +3297,6 @@ async def config(request: Request):
|
||||
"s2sUrl": _effective_s2s_url(request),
|
||||
"primaryPrompt": primary_prompt,
|
||||
"reasoningType": S2S_REASONING_TYPE,
|
||||
"auth": AUTH_ENABLED,
|
||||
"injectProfile": S2S_INJECT_PROFILE,
|
||||
"debugPrompt": S2S_DEBUG_PROMPT,
|
||||
"startSessionMode": S2S_START_SESSION_MODE,
|
||||
@@ -5029,40 +4974,6 @@ async def user_reports(phoneNumber: str | None = None):
|
||||
return JSONResponse(resp_body)
|
||||
|
||||
|
||||
@app.get("/api/me")
|
||||
async def me(request: Request):
|
||||
"""Login state, tier, and remaining daily budget. Only meaningful in LB mode;
|
||||
sets the anonymous tracking cookie when first seen."""
|
||||
client_ip = request.client.host if request.client else "?"
|
||||
logger.debug(
|
||||
"[INTERNAL] /api/me client=%s LIMITER=%s",
|
||||
client_ip, LIMITER_ENABLED,
|
||||
)
|
||||
if not LIMITER_ENABLED:
|
||||
return {"enabled": False}
|
||||
view = auth.user_view(request)
|
||||
tier, keys, set_cookie = auth.resolve_identity(request)
|
||||
unlimited = limiter.budget_for(tier) is None
|
||||
rem = None if unlimited else await asyncio.to_thread(limiter.remaining, keys, tier)
|
||||
logger.debug(
|
||||
"[INTERNAL] /api/me client=%s tier=%s unlimited=%s remaining=%s setCookie=%s",
|
||||
client_ip, tier, unlimited, rem, bool(set_cookie),
|
||||
)
|
||||
out = {
|
||||
"enabled": True,
|
||||
"auth": AUTH_ENABLED,
|
||||
**view,
|
||||
"remainingSec": rem,
|
||||
"limitSec": limiter.budget_for(tier),
|
||||
"loginUrl": auth.OAUTH_LOGIN_PATH if AUTH_ENABLED else None,
|
||||
"logoutUrl": auth.OAUTH_LOGOUT_PATH if AUTH_ENABLED else None,
|
||||
}
|
||||
resp = JSONResponse(out)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
|
||||
|
||||
@app.post("/api/search")
|
||||
async def search(req: SearchRequest):
|
||||
"""Proxy a Google search via Serper.dev. The key stays on the server unless
|
||||
@@ -5136,44 +5047,17 @@ async def search(req: SearchRequest):
|
||||
|
||||
@app.post("/api/session")
|
||||
async def session(request: Request):
|
||||
"""Proxy the session handshake to the load balancer, keeping its URL secret,
|
||||
and meter conversation time by tier.
|
||||
"""Relay a session grant or queue ticket, keeping the LB address server-side.
|
||||
|
||||
The browser POSTs here (same-origin); we resolve the caller's tier, refuse if
|
||||
today's budget is already spent (402), otherwise POST <LOAD_BALANCER_URL>/session
|
||||
and relay the JSON back. The LB body carries a per-session `connect_url`
|
||||
(compute host + short-lived token) the browser must dial directly — that one
|
||||
URL is unavoidably exposed, but the stable load-balancer address is not. On a
|
||||
successful grant we reserve the first time chunk against the day's budget."""
|
||||
The grant's `connect_url` contains the compute host and short-lived token the
|
||||
browser needs to connect directly to the allocated speech service."""
|
||||
client_ip = request.client.host if request.client else "?"
|
||||
if not LOAD_BALANCER_URL:
|
||||
# No LB configured — this deploy is direct-mode only; the browser should
|
||||
# never call this. 404 so it's indistinguishable from a missing route.
|
||||
raise HTTPException(status_code=404, detail="No encontrado.")
|
||||
|
||||
tier, keys, set_cookie = auth.resolve_identity(request)
|
||||
# Metering runs only on the deployed Space; off-Space the LB still proxies but
|
||||
# nothing is tracked. Within metering, unlimited tiers (pro, org) aren't either.
|
||||
tracked = LIMITER_ENABLED and limiter.budget_for(tier) is not None
|
||||
|
||||
logger.debug(
|
||||
"[EXT][LB] session: client=%s tier=%s tracked=%s LIMITER=%s",
|
||||
client_ip, tier, tracked, LIMITER_ENABLED,
|
||||
)
|
||||
|
||||
# Refuse before troubling the LB if the day's budget is already gone. Done
|
||||
# here (at enqueue) so we never put a user who can't talk into the queue.
|
||||
if tracked:
|
||||
rem = await asyncio.to_thread(limiter.remaining, keys, tier)
|
||||
logger.debug("[EXT][LB] session: client=%s budget remaining=%s", client_ip, rem)
|
||||
if rem is not None and rem <= 0:
|
||||
resp = JSONResponse(
|
||||
{"tier": tier, "reason": "limit", "remainingSec": 0}, status_code=402
|
||||
)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
|
||||
logger.debug("[EXT][LB] session: client=%s", client_ip)
|
||||
url = f"{LOAD_BALANCER_URL.rstrip('/')}/session"
|
||||
logger.debug("[EXT][LB] session POST %s (handshake)", url)
|
||||
try:
|
||||
@@ -5193,10 +5077,7 @@ async def session(request: Request):
|
||||
body = _safe_json(lb)
|
||||
if body.get("state") == "at_capacity":
|
||||
logger.debug("[EXT][LB] session: at_capacity — queue full")
|
||||
resp = JSONResponse({"state": "at_capacity"}, status_code=503)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
return JSONResponse({"state": "at_capacity"}, status_code=503)
|
||||
|
||||
if lb.status_code != 200:
|
||||
# The LB's error body may name the reason (e.g. capacity); it carries no
|
||||
@@ -5213,36 +5094,24 @@ async def session(request: Request):
|
||||
state, session_id, bool(connect_url), connect_url,
|
||||
)
|
||||
|
||||
# Busy pool: the LB queued us. Relay the ticket untouched — crucially with NO
|
||||
# reservation, so waiting in line never costs the day's budget.
|
||||
# Busy pool: the LB queued us. Relay the ticket for subsequent polling.
|
||||
if data.get("state") == "queued":
|
||||
data["tier"] = tier
|
||||
logger.debug("[EXT][LB] session: queued queueId=%s position=%s",
|
||||
data.get("queue_id") or data.get("queueId"),
|
||||
data.get("position"))
|
||||
resp = JSONResponse(data)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
|
||||
# A slot was free: reserve the first chunk now and return the grant.
|
||||
return await _finalize_grant(data, keys, tier, tracked, set_cookie)
|
||||
return JSONResponse(data)
|
||||
|
||||
|
||||
@app.get("/api/queue/{queue_id}")
|
||||
async def queue_status(queue_id: str, request: Request):
|
||||
"""Poll a waiting ticket: relay the position, or — when the head of the line
|
||||
claims a freed slot — reserve the budget now and return the grant. Re-checks the
|
||||
daily budget at claim, since a multi-minute wait could have spent it elsewhere."""
|
||||
"""Relay a waiting ticket's position or its newly claimed session grant."""
|
||||
client_ip = request.client.host if request.client else "?"
|
||||
if not LOAD_BALANCER_URL:
|
||||
raise HTTPException(status_code=404, detail="No encontrado.")
|
||||
|
||||
tier, keys, set_cookie = auth.resolve_identity(request)
|
||||
tracked = LIMITER_ENABLED and limiter.budget_for(tier) is not None
|
||||
logger.debug(
|
||||
"[EXT][LB] queue poll: client=%s queueId=%s tier=%s",
|
||||
client_ip, queue_id[:16], tier,
|
||||
"[EXT][LB] queue poll: client=%s queueId=%s",
|
||||
client_ip, queue_id[:16],
|
||||
)
|
||||
|
||||
url = f"{LOAD_BALANCER_URL.rstrip('/')}/queue/{queue_id}"
|
||||
@@ -5262,10 +5131,7 @@ async def queue_status(queue_id: str, request: Request):
|
||||
# Ticket unknown/expired (reaped after we stopped polling). Tell the client
|
||||
# to start over rather than spin.
|
||||
logger.debug("[EXT][LB] queue: queueId=%s EXPIRED (404)", queue_id[:16])
|
||||
resp = JSONResponse({"state": "expired"}, status_code=404)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
return JSONResponse({"state": "expired"}, status_code=404)
|
||||
|
||||
if lb.status_code != 200:
|
||||
logger.warning("Queue poll failed %s: %s", lb.status_code, lb.text[:300])
|
||||
@@ -5280,31 +5146,7 @@ async def queue_status(queue_id: str, request: Request):
|
||||
data.get("session_id") or data.get("sessionId"),
|
||||
)
|
||||
|
||||
if data.get("state") == "queued":
|
||||
data["tier"] = tier
|
||||
resp = JSONResponse(data)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
|
||||
# Claimed a slot. Re-check the budget: it may have been spent in another tab
|
||||
# during the wait. If so, refuse — the just-claimed slot is now a pending
|
||||
# session on the LB and its pending-timeout reaper reclaims it shortly.
|
||||
if tracked:
|
||||
rem = await asyncio.to_thread(limiter.remaining, keys, tier)
|
||||
logger.debug(
|
||||
"[EXT][LB] queue CLAIMED queueId=%s budget remaining=%s",
|
||||
queue_id[:16], rem,
|
||||
)
|
||||
if rem is not None and rem <= 0:
|
||||
resp = JSONResponse(
|
||||
{"tier": tier, "reason": "limit", "remainingSec": 0}, status_code=402
|
||||
)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
|
||||
return await _finalize_grant(data, keys, tier, tracked, set_cookie)
|
||||
return JSONResponse(data)
|
||||
|
||||
|
||||
@app.delete("/api/queue/{queue_id}")
|
||||
@@ -5330,35 +5172,6 @@ async def queue_end(request: Request):
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
async def _finalize_grant(data, keys, tier, tracked, set_cookie):
|
||||
"""Shared grant tail (fast path or queue claim): reserve the first chunk, attach
|
||||
the metering fields the client needs, and set the anon cookie."""
|
||||
remaining = None
|
||||
session_id = data.get("session_id")
|
||||
if tracked and session_id:
|
||||
logger.debug(
|
||||
"[EXT][LB] grant FINALIZE sessionId=%s tier=%s tracked=%s",
|
||||
session_id, tier, tracked,
|
||||
)
|
||||
await asyncio.to_thread(limiter.begin, session_id, keys, tier)
|
||||
remaining = await asyncio.to_thread(limiter.remaining, keys, tier)
|
||||
logger.debug(
|
||||
"[EXT][LB] grant RESERVED sessionId=%s remaining=%s",
|
||||
session_id, remaining,
|
||||
)
|
||||
|
||||
data.update({
|
||||
"tier": tier,
|
||||
"limited": tracked,
|
||||
"remainingSec": remaining,
|
||||
"heartbeatSec": limiter.HEARTBEAT_SEC,
|
||||
})
|
||||
resp = JSONResponse(data)
|
||||
if set_cookie:
|
||||
auth.set_anon_cookie(resp, set_cookie)
|
||||
return resp
|
||||
|
||||
|
||||
async def _lb_leave(queue_id: str) -> None:
|
||||
"""Best-effort: tell the LB to drop a waiting ticket."""
|
||||
url = f"{LOAD_BALANCER_URL.rstrip('/')}/queue/{queue_id}"
|
||||
@@ -6480,48 +6293,6 @@ async def _queue_id(request: Request) -> str:
|
||||
return (data or {}).get("queueId", "") if isinstance(data, dict) else ""
|
||||
|
||||
|
||||
async def _session_id(request: Request) -> str:
|
||||
"""Pull `sessionId` from a JSON body, tolerating sendBeacon's blob posts."""
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception:
|
||||
return ""
|
||||
return (data or {}).get("sessionId", "") if isinstance(data, dict) else ""
|
||||
|
||||
|
||||
@app.post("/api/session/heartbeat")
|
||||
async def session_heartbeat(request: Request):
|
||||
"""Extend the live reservation one chunk at a time. `expired` once the day's
|
||||
budget is spent — the client then tears down."""
|
||||
client_ip = request.client.host if request.client else "?"
|
||||
if not LIMITER_ENABLED:
|
||||
raise HTTPException(status_code=404, detail="No encontrado.")
|
||||
sid = await _session_id(request)
|
||||
alive = bool(sid) and await asyncio.to_thread(limiter.heartbeat, sid)
|
||||
logger.debug(
|
||||
"[INTERNAL] /api/session/heartbeat client=%s sessionId=%s alive=%s",
|
||||
client_ip, (sid or "")[:16], alive,
|
||||
)
|
||||
return {"expired": not alive}
|
||||
|
||||
|
||||
@app.post("/api/session/end")
|
||||
async def session_end(request: Request):
|
||||
"""Clean teardown: reconcile to real elapsed time and refund the unused
|
||||
chunk. Sent via navigator.sendBeacon, so it must succeed without a response."""
|
||||
client_ip = request.client.host if request.client else "?"
|
||||
if not LIMITER_ENABLED:
|
||||
raise HTTPException(status_code=404, detail="No encontrado.")
|
||||
sid = await _session_id(request)
|
||||
logger.debug(
|
||||
"[INTERNAL] /api/session/end client=%s sessionId=%s",
|
||||
client_ip, (sid or "")[:16],
|
||||
)
|
||||
if sid:
|
||||
await asyncio.to_thread(limiter.end, sid)
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def no_cache_static_assets(request: Request, call_next):
|
||||
response = await call_next(request)
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Backend contract checks with local HTTP fixtures and a loopback voice server.
|
||||
|
||||
Run from the repository root: .venv/bin/python -m unittest discover -s s2s-client/tests
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
from websockets.sync.server import serve
|
||||
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
# Never consume deployment credentials or contact services from a developer's .env.
|
||||
with patch.dict(os.environ, {}, clear=True), patch("dotenv.load_dotenv", return_value=False):
|
||||
import server
|
||||
|
||||
|
||||
GRANT = {
|
||||
"session_id": "local-session",
|
||||
"connect_url": "ws://voice.test/v1/realtime?token=fixture",
|
||||
"pending_timeout_s": 60,
|
||||
}
|
||||
TICKET = {"state": "queued", "queue_id": "local-ticket", "position": 2, "poll_interval_s": 1}
|
||||
ASYNC_CLIENT = httpx.AsyncClient
|
||||
|
||||
|
||||
class BackendContractTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.upstream = Mock(side_effect=AssertionError("Unexpected upstream HTTP request"))
|
||||
self.enterContext(patch.object(server, "LOAD_BALANCER_URL", "http://lb.test"))
|
||||
self.enterContext(patch.object(
|
||||
server.httpx, "AsyncClient",
|
||||
side_effect=lambda **kwargs: ASYNC_CLIENT(
|
||||
transport=httpx.MockTransport(self.upstream), **kwargs,
|
||||
),
|
||||
))
|
||||
self.client = self.enterContext(TestClient(server.app))
|
||||
|
||||
def respond(self, *responses):
|
||||
self.upstream.side_effect = iter(responses)
|
||||
|
||||
def test_bootstrap_and_removed_account_contract(self):
|
||||
response = self.client.get("/api/config")
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertTrue(response.json()["lb"])
|
||||
self.assertFalse(response.json()["allowDirect"])
|
||||
self.assertNotIn("auth", response.json())
|
||||
paths = self.client.get("/openapi.json").json()["paths"]
|
||||
for path in ("/api/me", "/api/session/heartbeat", "/api/session/end"):
|
||||
self.assertNotIn(path, paths)
|
||||
self.assertFalse(any(getattr(route, "path", "").startswith("/oauth") for route in server.app.routes))
|
||||
self.upstream.assert_not_called()
|
||||
|
||||
def test_session_grant_passes_through_without_extra_fields_or_cookies(self):
|
||||
self.respond(httpx.Response(200, json=GRANT))
|
||||
response = self.client.post("/api/session", json={})
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.json(), GRANT)
|
||||
self.assertNotIn("set-cookie", response.headers)
|
||||
request = self.upstream.call_args.args[0]
|
||||
self.assertEqual((request.method, str(request.url)), ("POST", "http://lb.test/session"))
|
||||
self.assertEqual(json.loads(request.content), {})
|
||||
|
||||
def test_queue_polling_returns_positions_and_then_grant(self):
|
||||
self.respond(
|
||||
httpx.Response(200, json=TICKET),
|
||||
httpx.Response(200, json={**TICKET, "position": 1}),
|
||||
httpx.Response(200, json=GRANT),
|
||||
)
|
||||
self.assertEqual(self.client.post("/api/session").json(), TICKET)
|
||||
self.assertEqual(self.client.get("/api/queue/local-ticket").json(), {**TICKET, "position": 1})
|
||||
response = self.client.get("/api/queue/local-ticket")
|
||||
self.assertEqual(response.json(), GRANT)
|
||||
self.assertNotIn("set-cookie", response.headers)
|
||||
self.assertEqual(
|
||||
[(call.args[0].method, call.args[0].url.path) for call in self.upstream.call_args_list],
|
||||
[("POST", "/session"), ("GET", "/queue/local-ticket"), ("GET", "/queue/local-ticket")],
|
||||
)
|
||||
|
||||
def test_capacity_and_expired_queue_responses_remain(self):
|
||||
self.respond(httpx.Response(503, json={"state": "at_capacity"}), httpx.Response(404))
|
||||
response = self.client.post("/api/session")
|
||||
self.assertEqual((response.status_code, response.json()), (503, {"state": "at_capacity"}))
|
||||
response = self.client.get("/api/queue/local-ticket")
|
||||
self.assertEqual((response.status_code, response.json()), (404, {"state": "expired"}))
|
||||
|
||||
def test_explicit_and_beacon_queue_cancellation(self):
|
||||
self.respond(httpx.Response(204), httpx.Response(204))
|
||||
self.assertEqual(self.client.delete("/api/queue/local-ticket").json(), {"ok": True})
|
||||
# sendBeacon posts a JSON blob, not a form or a session identifier.
|
||||
response = self.client.post(
|
||||
"/api/queue/end", content=json.dumps({"queueId": "local-ticket"}),
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
self.assertEqual(response.json(), {"ok": True})
|
||||
for call in self.upstream.call_args_list:
|
||||
self.assertEqual((call.args[0].method, call.args[0].url.path), ("DELETE", "/queue/local-ticket"))
|
||||
self.assertEqual(self.upstream.call_count, 2)
|
||||
|
||||
def test_empty_beacon_is_harmless_and_cancellation_is_best_effort(self):
|
||||
self.assertEqual(self.client.post("/api/queue/end", content="").json(), {"ok": True})
|
||||
self.upstream.assert_not_called()
|
||||
self.upstream.side_effect = httpx.ConnectError("Local fixture unavailable")
|
||||
self.assertEqual(self.client.delete("/api/queue/local-ticket").json(), {"ok": True})
|
||||
|
||||
def test_upstream_failures_are_still_reported(self):
|
||||
self.upstream.side_effect = httpx.ConnectError("Local fixture unavailable")
|
||||
self.assertEqual(self.client.post("/api/session").status_code, 502)
|
||||
self.assertEqual(self.client.get("/api/queue/local-ticket").status_code, 502)
|
||||
self.respond(httpx.Response(429), httpx.Response(500))
|
||||
self.assertEqual(self.client.post("/api/session").status_code, 502)
|
||||
self.assertEqual(self.client.get("/api/queue/local-ticket").status_code, 502)
|
||||
|
||||
def test_direct_mode_disables_lb_and_preserves_pinned_voice_url(self):
|
||||
with patch.object(server, "LOAD_BALANCER_URL", ""), patch.object(
|
||||
server, "SPEECH_TO_SPEECH_URL", GRANT["connect_url"],
|
||||
):
|
||||
config = self.client.get("/api/config").json()
|
||||
self.assertEqual(config["s2sUrl"], GRANT["connect_url"])
|
||||
self.assertFalse(config["lb"])
|
||||
self.assertTrue(config["allowDirect"])
|
||||
for method, path in (
|
||||
("POST", "/api/session"), ("GET", "/api/queue/local-ticket"),
|
||||
("DELETE", "/api/queue/local-ticket"), ("POST", "/api/queue/end"),
|
||||
):
|
||||
self.assertEqual(self.client.request(method, path).status_code, 404)
|
||||
self.upstream.assert_not_called()
|
||||
|
||||
def test_caller_profile_still_uses_upstream_login_token(self):
|
||||
profile = {"result": {"name": "Local Caller"}}
|
||||
self.respond(
|
||||
httpx.Response(200, json={"token": "fixture-token"}),
|
||||
httpx.Response(200, json=profile),
|
||||
)
|
||||
with patch.object(server, "S2S_LOGIN_BY_PHONE", "http://profile.test/login"), patch.object(
|
||||
server, "VOCERO_USER_PHONE_AUTH_URL", "http://profile.test/user",
|
||||
):
|
||||
response = self.client.get("/api/user/phone", params={"phoneNumber": "+56000000000"})
|
||||
self.assertEqual(response.json(), profile)
|
||||
login, lookup = [call.args[0] for call in self.upstream.call_args_list]
|
||||
self.assertEqual(json.loads(login.content), {"phoneNumber": "+56000000000"})
|
||||
self.assertEqual(lookup.headers["authorization"], "Bearer fixture-token")
|
||||
self.assertEqual(lookup.url.params["phoneNumber"], "+56000000000")
|
||||
|
||||
def test_caller_auth_headers_remain_independent(self):
|
||||
request = server.Request({
|
||||
"type": "http", "headers": [
|
||||
(b"authorization", b"Bearer fixture-token"), (b"x-demo-session", b"fixture-demo"),
|
||||
],
|
||||
})
|
||||
self.assertEqual(server._apipilot_parse_authorization(request), ("Bearer fixture-token", "fixture-demo"))
|
||||
|
||||
|
||||
class VoiceProxyTests(unittest.TestCase):
|
||||
def test_proxy_voice_frames_keepalive_and_client_disconnect(self):
|
||||
closed = threading.Event()
|
||||
received = []
|
||||
keepalive = threading.Event()
|
||||
|
||||
def voice(socket):
|
||||
try:
|
||||
socket.send(json.dumps({"type": "session.created"}))
|
||||
if socket.ping().wait(3):
|
||||
keepalive.set()
|
||||
for frame in socket:
|
||||
received.append(frame)
|
||||
socket.send(frame)
|
||||
finally:
|
||||
closed.set()
|
||||
|
||||
with serve(voice, "127.0.0.1", 0) as upstream:
|
||||
thread = threading.Thread(target=upstream.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
self.addCleanup(thread.join, 5)
|
||||
target = f"ws://127.0.0.1:{upstream.socket.getsockname()[1]}/v1/realtime"
|
||||
with patch.object(server, "S2S_PROXY_TARGET", target), TestClient(server.app) as client:
|
||||
config = client.get("/api/config", headers={"host": "localhost:19420"}).json()
|
||||
self.assertEqual(config["s2sUrl"], "ws://localhost:19420/v1/realtime")
|
||||
with client.websocket_connect("/v1/realtime") as socket:
|
||||
self.assertEqual(socket.receive_json(), {"type": "session.created"})
|
||||
update = {"type": "session.update", "session": {"instructions": "Local fixture"}}
|
||||
socket.send_json(update)
|
||||
self.assertEqual(socket.receive_json(), update)
|
||||
socket.send_bytes(b"\x00\x01\x02\x03")
|
||||
self.assertEqual(socket.receive_bytes(), b"\x00\x01\x02\x03")
|
||||
self.assertTrue(keepalive.wait(3))
|
||||
self.assertTrue(closed.wait(3), "Client disconnect must close the upstream WebSocket")
|
||||
self.assertEqual(len(received), 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user