Auth, Multi-Tenancy & Usage Quotas

# script_id: day_083_auth_multitenancy_quotas/api_keys
# API keys in 20 lines: issue, store only a hash, verify. Stdlib only: no install, no key.
import hashlib
import hmac
import secrets

def issue_key() -> tuple[str, str]:
    raw = "atk_" + secrets.token_urlsafe(32)                 # 32 random bytes: unguessable
    return raw, hashlib.sha256(raw.encode()).hexdigest()     # show raw ONCE; store only the hash

def verify(raw: str, stored_hash: str) -> bool:
    candidate = hashlib.sha256(raw.encode()).hexdigest()
    return hmac.compare_digest(candidate, stored_hash)       # constant-time comparison

raw, stored = issue_key()
print("give the customer (once):", raw[:10] + "…")           # never log a full key
print("store in the database:   ", stored[:16] + "…")
print("verify(correct):", verify(raw, stored))
print("verify(wrong):  ", verify(raw + "x", stored))

give the customer (once): atk_Qm3v9Z…
store in the database:    5f1c0d2e9a7b44c1…
verify(correct): True
verify(wrong):   False

# script_id: day_083_auth_multitenancy_quotas/auth
# Becomes atlas/api/auth.py in R5. Install: uv add "fastapi>=0.162"   No key; SQLite file only.
import hashlib
import hmac
import secrets
import sqlite3
import time
from contextlib import closing
from dataclasses import dataclass
from typing import Annotated, Literal

from fastapi import Depends, HTTPException, Request
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer

Plan = Literal["demo", "free", "team"]

@dataclass(frozen=True)
class Principal:
    """Who is calling and what they may do. Code below takes this, never the raw key."""
    tenant_id: str
    key_id: str
    plan: Plan
    scopes: frozenset[str]

def hash_key(raw: str) -> str:
    # fine for 256-bit random keys; human passwords need a slow hash (argon2/bcrypt)
    return hashlib.sha256(raw.encode()).hexdigest()

# script_id: day_083_auth_multitenancy_quotas/auth
class KeyStore:
    """API keys in SQLite. Raw format atk_<key_id>_<secret>; only the hash is stored."""
    def __init__(self, path: str = "keys.db") -> None:
        self.path = path
        with closing(sqlite3.connect(path)) as db, db:      # closing() closes; `db` commits
            db.execute("CREATE TABLE IF NOT EXISTS keys (key_id TEXT PRIMARY KEY, tenant_id TEXT,"
                       " plan TEXT, scopes TEXT, key_hash TEXT, created REAL, revoked REAL)")

    def issue(self, tenant_id: str, plan: Plan, scopes: set[str]) -> str:
        key_id = secrets.token_hex(4)                        # public half: safe in logs
        raw = f"atk_{key_id}_{secrets.token_urlsafe(32)}"
        with closing(sqlite3.connect(self.path)) as db, db:
            db.execute("INSERT INTO keys VALUES (?, ?, ?, ?, ?, ?, NULL)",
                       (key_id, tenant_id, plan, " ".join(sorted(scopes)), hash_key(raw), time.time()))
        return raw                                           # the only time anyone sees it

    def lookup(self, raw: str) -> Principal | None:
        parts = raw.split("_", 2)
        if len(parts) != 3 or parts[0] != "atk":
            return None
        with closing(sqlite3.connect(self.path)) as db:
            row = db.execute("SELECT tenant_id, plan, scopes, key_hash FROM keys"
                             " WHERE key_id = ? AND revoked IS NULL", (parts[1],)).fetchone()
        if row is None or not hmac.compare_digest(row[3], hash_key(raw)):
            return None
        return Principal(row[0], parts[1], row[1], frozenset(row[2].split()))

    def revoke(self, key_id: str) -> None:
        with closing(sqlite3.connect(self.path)) as db, db:
            db.execute("UPDATE keys SET revoked = ? WHERE key_id = ?", (time.time(), key_id))

# script_id: day_083_auth_multitenancy_quotas/auth
_bearer = HTTPBearer(auto_error=False)

def require_principal(request: Request,
                      creds: Annotated[HTTPAuthorizationCredentials | None, Depends(_bearer)]) -> Principal:
    store: KeyStore = request.app.state.keys
    principal = store.lookup(creds.credentials) if creds else None
    if principal is None:
        raise HTTPException(401, "invalid or missing API key", headers={"WWW-Authenticate": "Bearer"})
    request.state.principal = principal          # logs/traces get tenant_id + key_id, never the key
    return principal

def require_scope(scope: str):
    """Dependency factory: Depends(require_scope("triage:write"))."""
    def check(p: Annotated[Principal, Depends(require_principal)]) -> Principal:
        if scope not in p.scopes:
            raise HTTPException(403, f"key lacks scope {scope!r}")
        return p
    return check

if __name__ == "__main__":
    import os
    import tempfile
    store = KeyStore(os.path.join(tempfile.mkdtemp(), "keys.db"))
    raw = store.issue("acme", "team", {"ask", "reports"})
    p = store.lookup(raw)
    print(p)
    store.revoke(p.key_id)
    print("after revoke:", store.lookup(raw))

Principal(tenant_id='acme', key_id='9f3a1c2e', plan='team', scopes=frozenset({'ask', 'reports'}))
after revoke: None

# script_id: day_083_auth_multitenancy_quotas/tenancy
# Becomes atlas/tenancy.py in R5. Install: uv add "fastapi>=0.162"   Env: ORBIT_DATA (docs/*.md)
import os
import re
from dataclasses import dataclass
from pathlib import Path

from fastapi import HTTPException

from auth import Principal

PUBLIC_TENANT = "public"                     # Orbit's own docs: every tenant may read them
DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))

def tenant_key(p: Principal, *parts: str) -> str:
    """The ONE way to build a storage key: t:<tenant>:... Grep for 't:' to find every bypass."""
    return ":".join(["t", p.tenant_id, *parts])

def visible_tenants(p: Principal) -> tuple[str, ...]:
    # demo keys read public docs only; company keys read public docs plus their own
    return (PUBLIC_TENANT,) if p.plan == "demo" else (PUBLIC_TENANT, p.tenant_id)

def assert_owner(p: Principal, owner_tenant: str) -> None:
    if owner_tenant != p.tenant_id:
        # 404, not 403: a 403 confirms the id exists, which leaks another customer's activity
        raise HTTPException(404, "not found")

# script_id: day_083_auth_multitenancy_quotas/tenancy
@dataclass(frozen=True)
class Doc:
    doc_id: str
    tenant: str                              # every row carries its owner
    text: str

TENANT_NOTES = [                             # inline toy data: two customers' private runbooks
    Doc("acme/runbook", "acme", "Acme runbook: we use lease_seconds = 900 for the nightly ETL queue."),
    Doc("zephyr/runbook", "zephyr", "Zephyr runbook: lease_seconds = 120, heartbeat_seconds = 30."),
]

def load_corpus() -> list[Doc]:
    docs = sorted((DATA_DIR / "docs").glob("*.md"))
    return [Doc(f"docs/{p.stem}", PUBLIC_TENANT, p.read_text(encoding="utf-8")) for p in docs] + TENANT_NOTES

def _score(question: str, text: str) -> float:
    # toy keyword overlap, normalised by length; Day 27's hybrid search in real life
    q, words = set(re.findall(r"\w+", question.lower())), re.findall(r"\w+", text.lower())
    return len(q & set(words)) / (1 + len(words)) ** 0.5

def search(corpus: list[Doc], question: str, tenants: tuple[str, ...], k: int = 3) -> list[Doc]:
    allowed = [d for d in corpus if d.tenant in tenants]       # filter FIRST...
    return sorted(allowed, key=lambda d: _score(question, d.text), reverse=True)[:k]   # ...then rank

# script_id: day_083_auth_multitenancy_quotas/tenancy
if __name__ == "__main__":
    corpus = load_corpus()
    zephyr = Principal("zephyr", "k_z", "team", frozenset({"ask"}))
    q = "what lease_seconds do we use?"
    everyone = tuple({d.tenant for d in corpus})
    top = search(corpus, q, everyone)                          # the launch-test bug: no filter
    print("BROKEN, no filter:      ", [d.doc_id for d in top])
    print("BROKEN, filter after k: ", [d.doc_id for d in top if d.tenant in visible_tenants(zephyr)])
    print("FIXED, filter first:    ", [d.doc_id for d in search(corpus, q, visible_tenants(zephyr))])
    try:
        assert_owner(zephyr, owner_tenant="acme")              # zephyr fetches an Acme job id
    except HTTPException as e:
        print("foreign job id ->", e.status_code, e.detail)

BROKEN, no filter:       ['acme/runbook', 'zephyr/runbook', 'docs/broker-sqs']
BROKEN, filter after k:  ['zephyr/runbook', 'docs/broker-sqs']
FIXED, filter first:     ['zephyr/runbook', 'docs/broker-sqs', 'docs/scaling-workers']
foreign job id -> 404 not found

# script_id: day_083_auth_multitenancy_quotas/quotas
# Becomes atlas/quotas.py in R5. Install: uv add "redis>=8.1" tiktoken   (tests/demos: uv add fakeredis)
# Env: ATLAS_REDIS_URL (default redis://localhost:6379/0; "fake://" = in-process fakeredis, no Docker).
# Services: Redis 7+ (EXPIRE ... NX), e.g. the redis:8 compose.yaml from Day 78.
# Redis 6 rejects nx= ("wrong number of arguments for 'expire'"): drop nx=True there; re-setting the TTL is harmless.
import os
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Literal

import tiktoken
from redis.asyncio import Redis

from auth import Plan, Principal
from tenancy import tenant_key

REDIS_URL = os.environ.get("ATLAS_REDIS_URL", "redis://localhost:6379/0")
_enc = tiktoken.get_encoding("o200k_base")

def count_tokens(text: str) -> int:          # compact copy of Day 4: approximate for GPT-6 models
    return len(_enc.encode(text))

def get_redis(url: str = REDIS_URL) -> Redis:
    """Async client: these calls run on every request, so they mustn't block the event loop."""
    if url.startswith("fake://"):
        import fakeredis
        return fakeredis.FakeAsyncRedis(decode_responses=True)    # fresh in-process store
    return Redis.from_url(url, decode_responses=True)

@dataclass(frozen=True)
class PlanLimits:
    rpm: int                  # requests per minute, per tenant
    tokens_per_day: int       # input + output, per tenant, per UTC day
    max_input_tokens: int     # one question
    max_output_tokens: int    # clamp; reasoning tokens count against it too
    max_streams: int          # concurrent SSE streams
    effort: str               # reasoning-effort ceiling for the plan

PLANS: dict[Plan, PlanLimits] = {             # illustrative: derive yours from cost per tenant
    "demo": PlanLimits(10, 50_000, 500, 300, 1, "none"),
    "free": PlanLimits(60, 200_000, 2_000, 800, 2, "low"),
    "team": PlanLimits(300, 2_000_000, 8_000, 2_000, 3, "medium"),
}

# script_id: day_083_auth_multitenancy_quotas/quotas
Code = Literal["rate_limited", "quota_exceeded", "too_many_streams", "input_too_large"]

class QuotaExceeded(Exception):
    def __init__(self, code: Code, retry_after_s: int = 0) -> None:
        super().__init__(code)
        self.code, self.retry_after_s = code, retry_after_s

@dataclass(frozen=True)
class Decision:
    allowed: bool
    remaining: int
    retry_after_s: int        # integer SECONDS until the window resets, never a timestamp

class RateLimiter:
    """Fixed window: one counter per (key, clock minute), shared by every replica."""
    def __init__(self, redis: Redis) -> None:
        self.r = redis

    async def hit(self, key: str, limit: int, window_s: int = 60) -> Decision:
        now = int(time.time())
        bucket = f"{key}:{now // window_s}"
        async with self.r.pipeline(transaction=True) as pipe:     # MULTI/EXEC: both or neither
            pipe.incr(bucket)
            pipe.expire(bucket, window_s, nx=True)                # TTL set once, by the first hit
            count, _ = await pipe.execute()
        if count > limit:
            return Decision(False, 0, window_s - now % window_s)
        return Decision(True, limit - count, 0)

# script_id: day_083_auth_multitenancy_quotas/quotas
@dataclass(frozen=True)
class Reservation:
    key: str
    reserved: int

def _day_key(p: Principal) -> str:
    return tenant_key(p, "tokens", datetime.now(UTC).strftime("%Y-%m-%d"))

def _seconds_to_utc_midnight() -> int:
    now = datetime.now(UTC)
    midnight = (now + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)
    return int((midnight - now).total_seconds()) + 1

class TokenQuota:
    """Daily token budget per tenant: reserve the worst case, settle to real usage."""
    def __init__(self, redis: Redis) -> None:
        self.r = redis

    async def reserve(self, p: Principal, est: int) -> Reservation:
        key = _day_key(p)
        used = await self.r.incrby(key, est)             # atomic across replicas
        await self.r.expire(key, 2 * 86_400, nx=True)    # day keys clean themselves up
        if used > PLANS[p.plan].tokens_per_day:
            await self.r.decrby(key, est)                # rejected calls cost nothing
            raise QuotaExceeded("quota_exceeded", _seconds_to_utc_midnight())
        return Reservation(key, est)

    async def settle(self, res: Reservation, actual: int) -> None:
        if actual != res.reserved:
            await self.r.incrby(res.key, actual - res.reserved)

    async def remaining(self, p: Principal) -> int:
        return PLANS[p.plan].tokens_per_day - int(await self.r.get(_day_key(p)) or 0)

# script_id: day_083_auth_multitenancy_quotas/quotas
    async def acquire_stream(self, p: Principal) -> str:
        key = tenant_key(p, "streams")
        n = await self.r.incr(key)
        await self.r.expire(key, 300)                    # > the longest stream deadline
        if n > PLANS[p.plan].max_streams:
            await self.r.decr(key)
            raise QuotaExceeded("too_many_streams", retry_after_s=5)
        return key

    async def release_stream(self, key: str) -> None:
        if await self.r.decr(key) < 0:                   # key expired mid-stream: don't go negative
            await self.r.set(key, 0)

    @asynccontextmanager
    async def stream_slot(self, p: Principal):
        """For workers and WebSockets. HTTP routes acquire in a dependency instead (see below)."""
        key = await self.acquire_stream(p)
        try:
            yield
        finally:
            await self.release_stream(key)

# script_id: day_083_auth_multitenancy_quotas/quota_demo
# Run: python quota_demo.py   (fakeredis by default; ATLAS_REDIS_URL=redis://... for real Redis)
import asyncio
import os

from auth import Principal
from quotas import PLANS, QuotaExceeded, RateLimiter, TokenQuota, count_tokens, get_redis
from tenancy import tenant_key

async def main() -> None:
    r = get_redis(os.environ.get("ATLAS_REDIS_URL", "fake://"))
    zephyr = Principal("zephyr", "k_z", "free", frozenset({"ask"}))
    limiter = RateLimiter(r)
    hits = [await limiter.hit(tenant_key(zephyr, "rpm"), PLANS["free"].rpm) for _ in range(70)]
    no = [d for d in hits if not d.allowed]
    print(f"70 calls: {70 - len(no)} allowed, {len(no)} rejected, Retry-After: {no[0].retry_after_s} s")

    demo, limits = Principal("demo", "k_d", "demo", frozenset({"ask"})), PLANS["demo"]
    quota = TokenQuota(r)
    for question, real_out in [("How do I retry a failed job?", 120), ("What is a lease?", 80)]:
        n_in = count_tokens(question)
        res = await quota.reserve(demo, n_in + limits.max_output_tokens)    # worst case
        await quota.settle(res, n_in + real_out)                            # after the call
        print(f"reserved {res.reserved}, settled {n_in + real_out}, left {await quota.remaining(demo)}")

    huge = "Why does my worker stop? " * 2_000                              # 50,000 characters
    print(f"50k-char question: {count_tokens(huge)} tokens > max_input {limits.max_input_tokens}")
    try:
        while True:
            await quota.reserve(demo, 5_000)
    except QuotaExceeded as e:
        print(f"budget exhausted -> {e.code}, Retry-After: {e.retry_after_s} s")
    await r.aclose()

asyncio.run(main())

70 calls: 60 allowed, 10 rejected, Retry-After: 17 s
reserved 308, settled 128, left 49872
reserved 305, settled 85, left 49787
50k-char question: 12001 tokens > max_input 500
budget exhausted -> quota_exceeded, Retry-After: 30113 s

# script_id: day_083_auth_multitenancy_quotas/app
# Compact stand-in for Day 77's app with auth + quotas. Becomes atlas/api/app.py wiring in R5.
# Install: uv add "fastapi>=0.162" uvicorn "redis>=8.1" tiktoken "openai>=3.46"  (+ fakeredis)
# Env: ATLAS_REDIS_URL, ATLAS_KEYS_DB (default keys.db), ATLAS_LLM=fake|openai, OPENAI_API_KEY
# Run: uv run python app.py   (port 8000)
import os
import uuid
from collections.abc import AsyncIterable, AsyncIterator, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Annotated, Protocol

import anyio
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.sse import EventSourceResponse, ServerSentEvent
from pydantic import BaseModel, Field

from auth import KeyStore, Principal, require_scope
from quotas import PLANS, QuotaExceeded, RateLimiter, Reservation, TokenQuota, count_tokens, get_redis
from tenancy import Doc, assert_owner, load_corpus, search, tenant_key, visible_tenants

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
Event = tuple[str, str | int]                 # ("delta", text) ... then ("usage", total_tokens)

class LLM(Protocol):
    def stream(self, question: str, docs: list[Doc], max_output_tokens: int,
               effort: str) -> AsyncIterator[Event]: ...

# script_id: day_083_auth_multitenancy_quotas/app
class FakeLLM:
    """No key, fixed output; records the caps it was given so tests can check them."""
    def __init__(self) -> None:
        self.calls: list[dict] = []

    async def stream(self, question: str, docs: list[Doc], max_output_tokens: int,
                     effort: str) -> AsyncIterator[Event]:
        self.calls.append({"max_output_tokens": max_output_tokens, "effort": effort})
        text = f"[{docs[0].doc_id}] {docs[0].text[:80]}" if docs else "No sources found."
        for word in text.split(" "):
            yield "delta", word + " "
        yield "usage", count_tokens(question) + sum(count_tokens(d.text[:2000]) for d in docs) + count_tokens(text)

class OpenAILLM:
    def __init__(self) -> None:
        from openai import AsyncOpenAI
        self.client = AsyncOpenAI(timeout=60.0, max_retries=2)

    async def stream(self, question: str, docs: list[Doc], max_output_tokens: int,
                     effort: str) -> AsyncIterator[Event]:
        sources = "\n\n".join(f"[{d.doc_id}] {d.text[:2000]}" for d in docs)   # caps context cost
        events = await self.client.responses.create(
            model=MODEL, instructions="Answer only from the sources and cite [doc_id].",
            input=f"Sources:\n{sources}\n\nQuestion: {question}", stream=True, store=False,
            max_output_tokens=max_output_tokens, reasoning={"effort": effort})
        async for ev in events:
            if ev.type == "response.output_text.delta":
                yield "delta", ev.delta
            elif ev.type == "response.completed" and ev.response.usage:
                yield "usage", ev.response.usage.total_tokens

# script_id: day_083_auth_multitenancy_quotas/app
CONTEXT_TOKENS = 1_500        # retrieved sources: 3 docs x ~2,000 chars, truncated in the prompt

class AskRequest(BaseModel):
    question: str = Field(min_length=1, max_length=100_000)   # coarse body guard; plans cap tokens
    max_output_tokens: int | None = Field(default=None, ge=1)

@dataclass
class Admission:
    principal: Principal
    reservation: Reservation
    max_output_tokens: int
    stream_key: str | None = None

def admit(scope: str, *, stream: bool = False) -> Callable:
    """Every check that can say no, cheapest first, BEFORE the handler (and before SSE sends 200)."""
    async def _admit(request: Request, body: AskRequest,
                     p: Annotated[Principal, Depends(require_scope(scope))]) -> Admission:
        s, limits = request.app.state, PLANS[p.plan]
        n_in = count_tokens(body.question)
        if n_in > limits.max_input_tokens:                            # free: no Redis, no model
            raise QuotaExceeded("input_too_large")
        request.state.rate = await s.limiter.hit(tenant_key(p, "rpm"), limits.rpm)
        if not request.state.rate.allowed:
            raise QuotaExceeded("rate_limited", request.state.rate.retry_after_s)
        out = min(body.max_output_tokens or limits.max_output_tokens, limits.max_output_tokens)
        adm = Admission(p, await s.quota.reserve(p, n_in + CONTEXT_TOKENS + out), out)
        if stream:
            try:
                adm.stream_key = await s.quota.acquire_stream(p)
            except QuotaExceeded:
                await s.quota.settle(adm.reservation, 0)                # hand the tokens back
                raise
        return adm
    return _admit

# script_id: day_083_auth_multitenancy_quotas/app
async def _run(request: Request, adm: Admission, question: str) -> AsyncIterator[Event]:
    s, p = request.app.state, adm.principal
    docs = search(s.corpus, question, visible_tenants(p))           # tenant filter before ranking
    used = adm.reservation.reserved             # no usage yet (crash, cancel): keep the estimate
    try:
        async for kind, value in s.llm.stream(question, docs, adm.max_output_tokens,
                                              PLANS[p.plan].effort):
            if kind == "usage":
                used = int(value)
            yield kind, value
        yield "sources", ",".join(d.doc_id for d in docs)
    finally:
        with anyio.CancelScope(shield=True):    # cleanup must finish even when cancelled
            await s.quota.settle(adm.reservation, used)
            if adm.stream_key:
                await s.quota.release_stream(adm.stream_key)

def create_app(keys: KeyStore | None = None, llm: LLM | None = None,
               redis_factory: Callable = get_redis) -> FastAPI:
    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
        app.state.keys = keys or KeyStore(os.environ.get("ATLAS_KEYS_DB", "keys.db"))
        app.state.llm = llm or (OpenAILLM() if os.environ.get("ATLAS_LLM") == "openai" else FakeLLM())
        app.state.redis = redis_factory()
        app.state.limiter, app.state.quota = RateLimiter(app.state.redis), TokenQuota(app.state.redis)
        app.state.corpus = load_corpus()
        yield
        await app.state.redis.aclose()

    app = FastAPI(title="Atlas: auth + quotas", lifespan=lifespan)

    @app.exception_handler(QuotaExceeded)
    async def on_quota(request: Request, exc: QuotaExceeded) -> JSONResponse:
        headers = {"Retry-After": str(exc.retry_after_s)} if exc.retry_after_s else {}
        return JSONResponse({"error": exc.code, "retry_after_s": exc.retry_after_s},
                            status_code=413 if exc.code == "input_too_large" else 429, headers=headers)

    add_routes(app)
    return app

# script_id: day_083_auth_multitenancy_quotas/app
def add_routes(app: FastAPI) -> None:
    @app.post("/v1/ask")
    async def ask(body: AskRequest, request: Request,
                  adm: Annotated[Admission, Depends(admit("ask"))]) -> dict:
        parts, sources = [], ""
        async for kind, value in _run(request, adm, body.question):
            parts += [value] if kind == "delta" else []
            sources = value if kind == "sources" else sources
        return {"answer": "".join(parts).strip(), "sources": sources.split(",")}

    @app.post("/v1/ask/stream", response_class=EventSourceResponse)
    async def ask_stream(body: AskRequest, request: Request,
                         adm: Annotated[Admission, Depends(admit("ask", stream=True))]
                         ) -> AsyncIterable[ServerSentEvent]:
        async for kind, value in _run(request, adm, body.question):
            if kind != "usage":                  # usage is for the meter, not the client
                yield ServerSentEvent(raw_data=str(value), event=kind)

    @app.post("/v1/reports", status_code=202)
    async def create_report(request: Request,
                            p: Annotated[Principal, Depends(require_scope("reports"))]) -> dict:
        job_id = uuid.uuid4().hex                # Day 78 enqueues the real job; we record the owner
        await request.app.state.redis.hset(f"job:{job_id}",
                                           mapping={"tenant": p.tenant_id, "status": "queued"})
        return {"job_id": job_id}

    @app.get("/v1/reports/{job_id}")
    async def get_report(job_id: str, request: Request,
                         p: Annotated[Principal, Depends(require_scope("reports"))]) -> dict:
        job = await request.app.state.redis.hgetall(f"job:{job_id}")
        if not job:
            raise HTTPException(404, "not found")
        assert_owner(p, job["tenant"])           # same 404 as "never existed"
        return {"job_id": job_id, "status": job["status"]}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(create_app(), port=8000)

# script_id: day_083_auth_multitenancy_quotas/walkthrough
# Run: python walkthrough.py   (in-process: TestClient + fakeredis + FakeLLM; no key, no Docker)
import tempfile
from pathlib import Path

from fastapi.testclient import TestClient

from app import create_app
from auth import KeyStore
from quotas import get_redis

keys = KeyStore(str(Path(tempfile.mkdtemp()) / "keys.db"))
acme = keys.issue("acme", "team", {"ask", "reports"})
zephyr = keys.issue("zephyr", "team", {"ask", "reports"})
demo = keys.issue("demo", "demo", {"ask"})
def h(key: str) -> dict:
    return {"Authorization": f"Bearer {key}"}
Q = {"question": "what lease_seconds do we use?"}

with TestClient(create_app(keys, redis_factory=lambda: get_redis("fake://"))) as c:
    r = c.post("/v1/ask", json=Q)
    print("no key       ", r.status_code, r.json(), r.headers.get("www-authenticate"))
    for name, key in [("acme", acme), ("zephyr", zephyr)]:
        print(f"{name:<13}", 200, c.post("/v1/ask", json=Q, headers=h(key)).json()["sources"])
    r = c.post("/v1/ask", json={"question": "Why does my worker stop? " * 2_000}, headers=h(demo))
    print("demo, 50k    ", r.status_code, r.json())
    rs = [c.post("/v1/ask", json=Q, headers=h(demo)) for _ in range(11)]
    print("demo x11     ", [x.status_code for x in rs][-3:], rs[-1].json(), "Retry-After:", rs[-1].headers.get("retry-after"))
    print("demo reports ", c.post("/v1/reports", headers=h(demo)).status_code)
    job = c.post("/v1/reports", headers=h(acme)).json()["job_id"]
    print("acme job     ", c.get(f"/v1/reports/{job}", headers=h(acme)).status_code,
          "| zephyr reads it:", c.get(f"/v1/reports/{job}", headers=h(zephyr)).status_code)

no key        401 {'detail': 'invalid or missing API key'} Bearer
acme          200 ['acme/runbook', 'docs/broker-sqs', 'docs/scaling-workers']
zephyr        200 ['zephyr/runbook', 'docs/broker-sqs', 'docs/scaling-workers']
demo, 50k     413 {'error': 'input_too_large', 'retry_after_s': 0}
demo x11      [200, 200, 429] {'error': 'rate_limited', 'retry_after_s': 17} Retry-After: 17
demo reports  403
acme job      200 | zephyr reads it: 404

# script_id: day_083_auth_multitenancy_quotas/test_isolation
# Install: uv add --dev pytest fakeredis   Run: uv run pytest test_isolation.py -q   (no key, no Docker)
import tempfile
from pathlib import Path

import fakeredis
import pytest
from fastapi.testclient import TestClient

from app import FakeLLM, create_app
from auth import KeyStore

ASK = {"question": "what lease_seconds do we use?"}

def hdr(key: str) -> dict:
    return {"Authorization": f"Bearer {key}"}

@pytest.fixture
def env():
    keys = KeyStore(str(Path(tempfile.mkdtemp()) / "keys.db"))
    raw = {t: keys.issue(t, "team", {"ask", "reports"}) for t in ("acme", "zephyr")}
    raw["demo"] = keys.issue("demo", "demo", {"ask"})
    server, llm = fakeredis.FakeServer(), FakeLLM()
    app = create_app(keys, llm, lambda: fakeredis.FakeAsyncRedis(server=server, decode_responses=True))
    with TestClient(app) as client:
        yield client, raw, llm, fakeredis.FakeRedis(server=server, decode_responses=True)

@pytest.mark.parametrize("headers", [{}, {"Authorization": "Bearer atk_nope_x"}, {"Authorization": "Basic x"}])
def test_missing_or_bad_key_is_401(env, headers):
    r = env[0].post("/v1/ask", json=ASK, headers=headers)
    assert r.status_code == 401 and r.headers["www-authenticate"] == "Bearer"

# script_id: day_083_auth_multitenancy_quotas/test_isolation
@pytest.mark.parametrize("me,other", [("acme", "zephyr"), ("zephyr", "acme")])
@pytest.mark.parametrize("path", ["/v1/ask", "/v1/ask/stream"])
def test_answers_never_include_other_tenants(env, me, other, path):
    client, raw, *_ = env
    body = client.post(path, json=ASK, headers=hdr(raw[me])).text
    assert f"{me}/runbook" in body and f"{other}/" not in body

@pytest.mark.parametrize("owner,intruder", [("acme", "zephyr"), ("zephyr", "acme")])
def test_foreign_job_is_404_like_a_missing_one(env, owner, intruder):
    client, raw, *_ = env
    job = client.post("/v1/reports", headers=hdr(raw[owner])).json()["job_id"]
    assert client.get(f"/v1/reports/{job}", headers=hdr(raw[owner])).status_code == 200
    foreign = client.get(f"/v1/reports/{job}", headers=hdr(raw[intruder]))
    missing = client.get("/v1/reports/does-not-exist", headers=hdr(raw[intruder]))
    assert foreign.status_code == missing.status_code == 404 and foreign.json() == missing.json()

def test_demo_has_no_reports_scope(env):
    assert env[0].post("/v1/reports", headers=hdr(env[1]["demo"])).status_code == 403

def test_every_redis_key_is_tenant_scoped(env):
    client, raw, _, r = env
    client.post("/v1/ask", json=ASK, headers=hdr(raw["acme"]))
    client.post("/v1/reports", headers=hdr(raw["zephyr"]))
    assert r.keys() and all(k.startswith(("t:", "job:")) for k in r.keys())

if __name__ == "__main__":
    raise SystemExit(pytest.main([__file__, "-q"]))

...........                                                              [100%]
11 passed in 0.4s