Capstone R5, Part 1

atlas/
  config.py            CHANGED  D82 Settings + backend, index_version, phoenix_endpoint, cache_threshold
  services.py          CHANGED  D77 seam widened; reference adapters over R2-R4
  stubs.py             NEW      keyless, tenant-aware stand-ins
  api/deps.py          NEW      the composition root: Deps, build_deps()
  api/pipeline.py      NEW      the 12-stage request path
  api/app.py           CHANGED  create_app(settings): middleware, routes, readiness
  ui/main.py           NEW      D79 chat + API key + cards + feedback
scripts/seed_tenants.py, scripts/smoke_staging.py     NEW
tests/test_integration.py                            NEW

# script_id: day_089_capstone_r5_integrate/scripts/assemble_r5
# scripts/assemble_r5.py: write the carried files you don't have yet, from the course lessons.
# Run: uv run python scripts/assemble_r5.py [COURSE_DIR]   Stdlib only.
# COURSE_DIR defaults to the course checkout holding $ORBIT_DATA (<course>/datasets/orbit).
from __future__ import annotations

import os
import re
import sys
from pathlib import Path

CARRIED = {   # R5 file -> (day, that lesson's script_ids, stitched in order): §3's table as data
    "atlas/api/events.py": ("077", ["events"]), "atlas/jobs.py": ("078", ["atlas_jobs"]),
    "atlas/ui/app.py": ("079", ["atlas_ui"]), "atlas/ui/components.py": ("080", ["components"]),
    "scripts/smoke_test.py": ("082", ["smoke_test"]), "atlas/api/auth.py": ("083", ["auth"]),
    "atlas/tenancy.py": ("083", ["tenancy"]), "atlas/quotas.py": ("083", ["quotas"]),
    "atlas/resilience.py": ("084", ["atlas_resilience"]), "atlas/cache.py": ("085", ["atlas_cache"]),
    "atlas/model_router.py": ("085", ["atlas_router"]), "atlas/monitoring.py": ("086", ["atlas_monitoring"]),
    "atlas/feedback.py": ("086", ["feedback"]), "scripts/replay_traffic.py": ("086", ["replay_traffic"]),
    "atlas/flags.py": ("087", ["atlas_release"]), "atlas/ops.py": ("088", ["atlas_ops"]),
    "atlas/index_migration.py": ("088", ["index_migration", "openai_embed"]),
    "scripts/game_day.py": ("088", ["game_day"]),
    "atlas/tracing.py": ("042", ["atlas/tracing"]),
}
MODULES = {   # a lesson's flat import name -> its R5 module
    "auth": "atlas.api.auth", "tenancy": "atlas.tenancy", "quotas": "atlas.quotas",
    "atlas_resilience": "atlas.resilience", "atlas_cache": "atlas.cache", "first_semantic_cache": "atlas.cache",
    "atlas_router": "atlas.model_router", "atlas_monitoring": "atlas.monitoring", "atlas_release": "atlas.flags",
    "atlas_ops": "atlas.ops", "atlas_jobs": "atlas.jobs", "events": "atlas.api.events",
    "components": "atlas.ui.components", "smoke_test": "scripts.smoke_test", "feedback": "atlas.feedback",
    "index_migration": "atlas.index_migration", "atlas_ui": "atlas.ui.app",
}
# Day 77 kept its request/response models in `events`; R5 imports them from atlas.api.models
MODELS_SHIM = ("from atlas.api.events import (Answer, AskRequest, AskResponse, Citation,  # noqa: F401\n"
               "                              InvestigateRequest, TriageRequest)\n")
FENCE = re.compile(r"```python\n# script_id: day_\d{3}_\w+/(\S+)\n(.*?)```", re.S)

# script_id: day_089_capstone_r5_integrate/scripts/assemble_r5
def stitch(course: Path, day: str, ids: list[str]) -> str:
    lesson = next(course.glob(f"Phase_*/Day_{day}_*/*.md"))
    blocks: dict[str, list[str]] = {}
    for sid, code in FENCE.findall(lesson.read_text(encoding="utf-8")):
        blocks.setdefault(sid, []).append(code)
    code = "\n".join(part for sid in ids for part in blocks[sid])
    for flat, mod in MODULES.items():                 # `from auth import X` -> `from atlas.api.auth import X`
        code = re.sub(rf"^(\s*)from {flat} import", rf"\1from {mod} import", code, flags=re.M)
        code = re.sub(rf"^(\s*)import {flat}\b", rf"\1import {mod} as {flat}", code, flags=re.M)
    return code


def assemble(course: Path, root: Path) -> tuple[list[str], list[str]]:
    wrote, kept = [], []
    files = {path: (lambda d=day, i=ids: stitch(course, d, i)) for path, (day, ids) in CARRIED.items()}
    files["atlas/api/models.py"] = lambda: MODELS_SHIM
    for path, build in files.items():
        target = root / path
        if target.exists():                           # your own commit wins
            kept.append(path)
            continue
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(build(), encoding="utf-8")
        wrote.append(path)
    return wrote, kept


if __name__ == "__main__":
    default = Path(os.environ.get("ORBIT_DATA", "datasets/orbit")).resolve().parents[1]
    course = Path(sys.argv[1]) if len(sys.argv) > 1 else default
    wrote, kept = assemble(course, Path(__file__).resolve().parents[1])
    print(f"wrote {len(wrote)} carried files, kept {len(kept)} of yours")

wrote 20 carried files, kept 0 of yours        (on a clean copy; your own repo keeps more)

# script_id: day_089_capstone_r5_integrate/atlas/config
# atlas/config.py: Day 82's Settings plus the R5 fields. Stdlib only. Env: ATLAS_* (see .env.example)
# Keep the r4-reference constants block (MODEL, DATA_DIR, STATE_DIR, PRICES_PER_1M, ...) at the top
# of this file unchanged: R2-R4 modules import those names. Only Settings is shown here.
from __future__ import annotations

import os
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Literal


class ConfigError(RuntimeError):
    """Startup configuration is invalid; the message lists every problem found."""


def _split(value: str) -> tuple[str, ...]:
    return tuple(part.strip() for part in value.split(",") if part.strip())


@dataclass(frozen=True)
class Settings:
    env: Literal["dev", "staging", "prod"]
    release: str
    model: str
    fallback_models: tuple[str, ...]
    embed_model: str
    data_dir: Path
    redis_url: str
    database_url: str | None
    api_url: str
    allowed_origins: tuple[str, ...]
    backend: Literal["reference", "stub"] = "stub"   # new: which R2-R4 sit behind the API
    index_version: str = "v1"                       # new: part of every cache key (Day 88)
    phoenix_endpoint: str | None = None             # new: OTLP URL, e.g. http://localhost:6006/v1/traces; None = off
    cache_threshold: float = 0.95                   # new: from Day 85's threshold sweep
    state_dir: Path = Path(".atlas")                # keys.db, events.db
    release_file: Path = Path("release.toml")       # Day 87's pins and canaries

# script_id: day_089_capstone_r5_integrate/atlas/config
    # (Settings continues) every knob from the environment, checked once, every problem reported
    @classmethod
    def from_env(cls, environ: Mapping[str, str] | None = None) -> Settings:
        e = os.environ if environ is None else environ
        env, backend = e.get("ATLAS_ENV", "dev"), e.get("ATLAS_BACKEND", "stub")
        problems: list[str] = []
        if env not in ("dev", "staging", "prod"):
            problems.append(f"ATLAS_ENV must be dev, staging or prod, not {env!r}")
        required = ["ATLAS_RELEASE", "ATLAS_REDIS_URL"] + (["OPENAI_API_KEY"] if backend == "reference" else [])
        if env != "dev" and (missing := [k for k in required if not e.get(k)]):
            problems.append("missing " + ", ".join(missing))
        if backend not in ("reference", "stub") or (env == "prod" and backend == "stub"):
            problems.append(f"ATLAS_BACKEND={backend!r} is not allowed in {env}")
        origins = _split(e.get("ATLAS_ALLOWED_ORIGINS", "http://localhost:7860"))
        if env != "dev" and "*" in origins:          # Day 82 blocked it in prod; R5 blocks staging too
            problems.append("ATLAS_ALLOWED_ORIGINS='*' is only allowed in dev")
        try:
            threshold = float(e.get("ATLAS_CACHE_THRESHOLD", "0.95"))
        except ValueError:
            threshold = -1.0
        if not 0 < threshold <= 1:
            problems.append("ATLAS_CACHE_THRESHOLD must be a number in (0, 1]")
        if problems:
            raise ConfigError(f"ATLAS_ENV={env}: " + "; ".join(problems))
        return cls(env=env, release=e.get("ATLAS_RELEASE", "dev"),  # type: ignore[arg-type]
                   model=e.get("ATLAS_MODEL", "gpt-6-luna"),
                   fallback_models=_split(e.get("ATLAS_FALLBACK_MODELS", "gpt-6.1-terra")),
                   embed_model=e.get("ATLAS_EMBED_MODEL", "text-embedding-3-small"),
                   data_dir=Path(e.get("ORBIT_DATA", "datasets/orbit")),
                   redis_url=e.get("ATLAS_REDIS_URL", "redis://localhost:6379/0"),
                   database_url=e.get("ATLAS_DATABASE_URL") or None,
                   api_url=e.get("ATLAS_API_URL", "http://localhost:8000"), allowed_origins=origins,
                   backend=backend, index_version=e.get("ATLAS_INDEX_VERSION", "v1"),  # type: ignore[arg-type]
                   phoenix_endpoint=e.get("ATLAS_PHOENIX_ENDPOINT") or None, cache_threshold=threshold,
                   state_dir=Path(e.get("ATLAS_STATE_DIR", ".atlas")),
                   release_file=Path(e.get("ATLAS_RELEASE_FILE", "release.toml")))

# script_id: day_089_capstone_r5_integrate/atlas/config
if __name__ == "__main__":
    s = Settings.from_env({"ATLAS_REDIS_URL": "fake://"})
    print(s.env, s.backend, s.model, s.fallback_models, s.index_version)
    try:
        Settings.from_env({"ATLAS_ENV": "prod", "ATLAS_RELEASE": "3f2c1ab", "ATLAS_REDIS_URL": "redis://r:6379/0"})
    except ConfigError as err:
        print(f"ConfigError: {err}")

dev stub gpt-6-luna ('gpt-6.1-terra',) v1
ConfigError: ATLAS_ENV=prod: ATLAS_BACKEND='stub' is not allowed in prod

# script_id: day_089_capstone_r5_integrate/atlas/api/deps
# atlas/api/deps.py: the composition root. Every long-lived object is built here, once, at startup.
# Install: uv add "openai>=3.46,<4" "redis>=8.1" fakeredis numpy   Env: see atlas/config.py
from __future__ import annotations

from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any

import numpy as np

from atlas.api.auth import KeyStore
from atlas.cache import SemanticCache, hashed_embed
from atlas.config import Settings
from atlas.flags import ReleaseWatcher
from atlas.jobs import get_redis as get_sync_redis     # Day 78's sync client: Day 88's levers use it
from atlas.model_router import STRONG
from atlas.monitoring import EventLog
from atlas.quotas import RateLimiter, TokenQuota, get_redis   # Day 83's async client
from atlas.resilience import (CircuitBreaker, DocChunk, FakeAsyncClient, FaultInjector, ResilientLLM,
                              RetryPolicy, load_chunks)
from atlas.services import AtlasServices, ParseShim, load_services

INTERACTIVE = RetryPolicy(max_attempts=3, max_wait_s=5.0)    # a person is waiting: fall back sooner


@dataclass
class Deps:
    settings: Settings
    redis: Any                          # async Redis: rate limits, token quotas, stream slots
    ops_redis: Any                      # sync Redis: kill switches and their audit log
    services: AtlasServices
    llms: dict[str, ResilientLLM]       # one pool per routable model; all share `breakers`
    breakers: dict[str, CircuitBreaker]
    cache: SemanticCache
    events: EventLog
    keys: KeyStore
    quotas: TokenQuota
    limiter: RateLimiter
    release: ReleaseWatcher
    chunks: list[DocChunk]              # degraded mode's keyword index (Day 84)
    client: Any = None                  # the real AsyncOpenAI, closed at shutdown

# script_id: day_089_capstone_r5_integrate/atlas/api/deps
def build_pools(client: Any, models: list[str], *, policy: RetryPolicy = INTERACTIVE
                ) -> tuple[dict[str, ResilientLLM], dict[str, CircuitBreaker]]:
    """Pool m tries m first, then the others. One breaker per model, shared by every pool."""
    pools = {m: ResilientLLM(client, [m, *(x for x in models if x != m)], policy=policy) for m in models}
    breakers = pools[models[0]].breakers
    for pool in pools.values():
        pool.breakers = breakers
    return pools, breakers


def _embedder(settings: Settings):
    if settings.backend == "stub":
        return hashed_embed                         # Day 85's offline stand-in: tests and demos only
    from atlas.index_migration import openai_embed  # Day 88's real embedder (needs a key)
    v = SimpleNamespace(embed_model=settings.embed_model, dims=256)   # small vectors: cache only

    def embed(texts: list[str]) -> np.ndarray:
        vecs = openai_embed(texts, v)
        return vecs / np.linalg.norm(vecs, axis=1, keepdims=True)    # the cache's dot product = cosine
    return embed


def build_deps(settings: Settings) -> Deps:
    models = list(dict.fromkeys([settings.model, *settings.fallback_models, STRONG]))
    client = None
    if settings.backend == "reference":
        from openai import AsyncOpenAI
        client = AsyncOpenAI(timeout=30.0, max_retries=0)    # ResilientLLM is the ONE retry layer
    pools, breakers = build_pools(ParseShim(client) if client else FakeAsyncClient(FaultInjector({})), models)
    redis = get_redis(settings.redis_url)
    settings.state_dir.mkdir(parents=True, exist_ok=True)
    return Deps(settings=settings, redis=redis, ops_redis=get_sync_redis(settings.redis_url),
                services=load_services(settings), llms=pools, breakers=breakers,
                cache=SemanticCache(_embedder(settings), threshold=settings.cache_threshold),
                events=EventLog(settings.state_dir / "events.db"),
                keys=KeyStore(str(settings.state_dir / "keys.db")),
                quotas=TokenQuota(redis), limiter=RateLimiter(redis),
                release=ReleaseWatcher(settings.release_file), chunks=load_chunks(settings.data_dir),
                client=client)

# script_id: day_089_capstone_r5_integrate/atlas/services
# atlas/services.py: the seam between the API and R2-R4 (replaces Day 77's AtlasServices).
# ATLAS_BACKEND=stub -> atlas/stubs.py; reference -> thin adapters over the r4-reference modules.
from __future__ import annotations

import asyncio
import sqlite3
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import dataclass
from typing import Any

import pydantic

from atlas.api.events import AgentEvent
from atlas.api.models import Answer
from atlas.config import Settings
from atlas.resilience import CallMeta


@dataclass
class AskOutcome:
    answer: Answer
    hits: list[str]              # retrieved doc ids: online checks compare citations with them
    meta: CallMeta               # which model answered, attempts, fallback (Day 84)
    input_tokens: int            # real usage from the response, not a length heuristic
    output_tokens: int


@dataclass(frozen=True)
class AtlasServices:
    """Day 77's seam, widened: ask() gets the tenant's groups, a model pool and a prompt version."""
    name: str                                        # "reference" | "stub", shown by /readyz
    ask: Callable[..., Awaitable[AskOutcome]]        # (question, *, groups, llm, effort, ...)
    investigate: Callable[[int], AsyncIterator[AgentEvent]]
    triage: Callable[[int], Awaitable[dict]]


def load_services(settings: Settings) -> AtlasServices:
    if settings.backend == "stub":
        from atlas.stubs import stub_services
        return stub_services()
    try:
        return reference_services(settings)
    except ImportError as exc:                       # fail at startup, never on a customer's request
        raise RuntimeError(f"ATLAS_BACKEND=reference needs the r4-reference modules ({exc}); "
                           "check out the tag, or set ATLAS_BACKEND=stub") from exc

# script_id: day_089_capstone_r5_integrate/atlas/services
async def adapt_graph_events(updates: AsyncIterator[Any]) -> AsyncIterator[AgentEvent]:
    """LangGraph stream_mode="updates" chunks, {node: state_update}, -> Day 77 AgentEvents."""
    usage: dict = {}
    async for chunk in updates:
        for node, update in chunk.items():
            if not isinstance(update, dict):          # "__interrupt__" carries a tuple
                continue
            usage = update.get("usage", usage)
            for m in update.get("messages", []):
                for call in getattr(m, "tool_calls", None) or []:
                    yield AgentEvent(type="tool_call", data={"name": call["name"], "args": call["args"]})
                if getattr(m, "type", "") == "tool":
                    yield AgentEvent(type="tool_result", data={"name": m.name, "chars": len(str(m.content))})
            if d := update.get("diagnosis"):          # R3's Diagnosis, as a dict
                for c in d.get("doc_citations", []):
                    yield AgentEvent(type="citation", data={"doc_id": c["doc_id"]})
                yield AgentEvent(type="final", data={"text": f"{d['summary']}\n\nProbable cause: "
                                                     f"{d['probable_cause']}", "usage": usage,
                                                     "needs_human": d["needs_human"]})
            else:
                yield AgentEvent(type="status", data={"step": node, "state": "done"})


class ParseShim:
    """ResilientLLM calls client.responses.create(); route it to .parse() so R2's schema is kept."""

    def __init__(self, client: Any) -> None:
        self._client, self.responses = client, self

    async def create(self, **kwargs: Any) -> Any:
        return await self._client.responses.parse(**kwargs)

# script_id: day_089_capstone_r5_integrate/atlas/services
def reference_services(settings: Settings) -> AtlasServices:
    """Adapters over R2 (Day 34), R3 (Day 61) and R4 (Day 76). Any missing module raises ImportError."""
    from langgraph.checkpoint.sqlite import SqliteSaver

    from atlas import rag
    from atlas.config import STATE_DIR                  # R4's constant, kept in config.py
    from atlas.graph import AgentConfig, build_investigator_graph
    from atlas.prompts.registry import load_prompt      # Day 21: returns a Prompt, so use .text; release.toml picks the version
    from atlas.schemas import Answer as R2Answer
    from atlas.tools import read_issue
    from atlas.triage import build_triage_graph

    index = rag.make_index()                            # ATLAS_INDEX=pg|memory, as in R2

    async def ask(question: str, *, groups: tuple[str, ...], llm: Any, effort: str,
                  max_output_tokens: int, prompt_version: str) -> AskOutcome:
        hits = await asyncio.to_thread(rag.retrieve, question, index=index, groups=groups)
        context, used = rag.build_context(hits)
        try:
            resp, meta = await llm.create(instructions=load_prompt("ask", prompt_version).text,
                                          input=f"{context}\n\nQuestion: {question}", text_format=R2Answer,
                                          reasoning={"effort": effort}, max_output_tokens=max_output_tokens,
                                          store=False)    # Day 73: the provider keeps nothing
        except pydantic.ValidationError:                # JSON cut off at the output cap (Day 16)
            idk = Answer(text=rag.IDK_TEXT, citations=[])
            return AskOutcome(idk, [], CallMeta(llm.models[0], 1, False, False, 0.0), 0, max_output_tokens)
        draft = resp.output_parsed or R2Answer(text=rag.IDK_TEXT, citations=[])
        good = [c for c in draft.citations if c not in rag.verify_citations(draft, used)]
        answer = Answer.model_validate({"text": draft.text if good else rag.IDK_TEXT,   # R2's rule
                                        "citations": [c.model_dump() for c in good]})
        u = resp.usage
        return AskOutcome(answer, [h.chunk.doc_id for h in used], meta,
                          u.input_tokens if u else 0, u.output_tokens if u else 0)

# script_id: day_089_capstone_r5_integrate/atlas/services
    # (reference_services continues)
    async def investigate(issue_id: int) -> AsyncIterator[AgentEvent]:
        cfg = AgentConfig()
        graph = build_investigator_graph(config=cfg)    # no checkpointer: a one-shot API run
        start = {"thread_id": f"api-{issue_id}", "issue": read_issue(issue_id).model_dump()}
        updates = graph.astream(start, {"recursion_limit": 4 * cfg.max_steps + 12}, stream_mode="updates")
        async for event in adapt_graph_events(updates):
            yield event

    def triage_sync(issue_id: int) -> dict:
        saver = SqliteSaver(sqlite3.connect(str(STATE_DIR / "triage.sqlite"), check_same_thread=False))
        graph = build_triage_graph(saver)               # same file and thread id as R4's cli.py
        cfg = {"configurable": {"thread_id": f"triage-{issue_id}"}, "recursion_limit": 30}
        if done := graph.get_state(cfg).values.get("proposal"):
            return done                                 # already proposed: same answer, $0
        return graph.invoke({"issue_id": issue_id}, cfg)["proposal"]   # pauses at approval

    async def triage(issue_id: int) -> dict:
        return await asyncio.to_thread(triage_sync, issue_id)

    return AtlasServices("reference", ask, investigate, triage)

# script_id: day_089_capstone_r5_integrate/atlas/stubs
# atlas/stubs.py: keyless stand-ins for R2-R4 (Day 77's stubs, now tenant-aware). ATLAS_BACKEND=stub.
# No key, no Postgres. Generation still goes through the ResilientLLM pool it is handed (a
# FakeAsyncClient in stub mode), so breakers, fallbacks and degraded mode behave as in production.
from __future__ import annotations

import json
import re
from collections.abc import AsyncIterator

from fastapi import HTTPException

from atlas.api.events import AgentEvent
from atlas.api.models import Answer, Citation
from atlas.jobs import load_open_issues, rules_triage
from atlas.quotas import count_tokens
from atlas.resilience import ResilientLLM
from atlas.services import AskOutcome, AtlasServices
from atlas.tenancy import load_corpus, search
from atlas.ui.components import to_ui_event

WORD = re.compile(r"[a-z0-9_]{3,}")


def _cite_id(doc_id: str) -> str:
    return doc_id.removeprefix("docs/")             # public docs are cited by slug, as in R2


def _best_line(text: str, question: str) -> str:
    """The prose line sharing the most words with the question: a stand-in for a real answer."""
    want = set(WORD.findall(question.lower()))
    lines = [ln.strip() for ln in text.splitlines() if ln.strip() and ln.lstrip()[0] not in "-#|`>"]
    return max(lines, key=lambda ln: len(want & set(WORD.findall(ln.lower()))), default="")[:200]

# script_id: day_089_capstone_r5_integrate/atlas/stubs
def stub_services() -> AtlasServices:
    corpus, issues = load_corpus(), load_open_issues()   # Day 83's toy tenants, Day 78's open issues

    async def ask(question: str, *, groups: tuple[str, ...], llm: ResilientLLM, effort: str,
                  max_output_tokens: int, prompt_version: str) -> AskOutcome:
        docs = search(corpus, question, groups)          # the tenant filter runs BEFORE ranking
        prompt = "\n\n".join(d.text[:2000] for d in docs) + f"\n\nQuestion: {question}"
        resp, meta = await llm.create(input=prompt, max_output_tokens=max_output_tokens)
        cites = [Citation(doc_id=_cite_id(d.doc_id), quote=_best_line(d.text, question)) for d in docs[:2]]
        text = f"[{resp.model}, prompt {prompt_version}] " + (cites[0].quote if cites else "I don't know.")
        return AskOutcome(Answer(text=text, citations=cites), [_cite_id(d.doc_id) for d in docs], meta,
                          count_tokens(prompt), count_tokens(text))

    async def triage(issue_id: int) -> dict:
        if issue_id not in issues:
            raise HTTPException(404, "not found")
        return rules_triage(issues[issue_id]).model_dump()   # Day 78's rules: proposes, applies nothing

# script_id: day_089_capstone_r5_integrate/atlas/stubs
    # (stub_services continues)
    async def investigate(issue_id: int) -> AsyncIterator[AgentEvent]:
        issue = issues.get(issue_id)
        if issue is None:
            yield AgentEvent(type="error", data={"code": "not_found"})
            return
        words = set(WORD.findall(issue["title"].lower()))
        yield AgentEvent(type="tool_call", data={"name": "search_issues", "args": {"query": issue["title"]}})
        similar = sorted((i for i in issues.values() if i["id"] != issue_id), reverse=True,
                         key=lambda i: len(words & set(WORD.findall(i["title"].lower()))))[:2]
        yield AgentEvent(type="tool_result", data={"name": "search_issues", "chars": len(json.dumps(similar))})
        yield to_ui_event("DuplicateCandidates", json.dumps({"issue_id": issue_id, "candidates": [
            {"issue_id": i["id"], "title": i["title"][:120], "confidence": "low",
             "reason": "shared title words (stub)"} for i in similar]}))     # a Day 80 card
        for d in search(corpus, issue["title"], ("public",))[:2]:
            yield AgentEvent(type="citation", data={"doc_id": _cite_id(d.doc_id)})
        yield AgentEvent(type="final", data={"text": f"Stub diagnosis for #{issue_id}: {issue['title']}",
                                             "usage": {"input_tokens": count_tokens(issue["body"]),
                                                       "output_tokens": 60}})

    return AtlasServices("stub", ask, investigate, triage)

# script_id: day_089_capstone_r5_integrate/atlas/api/pipeline
# atlas/api/pipeline.py: ONE request path for every route, in a fixed order (see the diagram).
# No key needed with ATLAS_BACKEND=stub. Every stage reuses the day that built it.
from __future__ import annotations

import asyncio
import time
from collections.abc import AsyncIterator
from dataclasses import dataclass, field

import anyio
from fastapi import BackgroundTasks, HTTPException

from atlas.api.auth import Principal
from atlas.api.deps import Deps
from atlas.api.events import AgentEvent
from atlas.api.models import AskRequest, AskResponse, InvestigateRequest
from atlas.cache import cache_scope, detect_version
from atlas.flags import resolve_prompt
from atlas.model_router import choose_model, cost_usd
from atlas.monitoring import Answer as CheckedAnswer, RequestEvent, score_request
from atlas.ops import get_overrides, plan_request
from atlas.quotas import PLANS, PlanLimits, QuotaExceeded, Reservation, count_tokens
from atlas.resilience import AllModelsUnavailable, degraded_answer, keyword_search
from atlas.tenancy import tenant_key, visible_tenants
from atlas.tracing import step_span

CONTEXT_TOKENS = 1_500              # retrieved sources: Day 83's estimate
INVESTIGATE_TOKENS = 20_000         # flat reservation per agent run, settled to real usage
EFFORTS = ["none", "low", "medium", "high"]


class AskResult(AskResponse):
    """Day 77's response plus the flags that tests, dashboards and the UI read."""
    model: str
    variant: str = "control"
    cached: bool = False
    fallback: bool = False
    degraded: bool = False

# script_id: day_089_capstone_r5_integrate/atlas/api/pipeline
@dataclass
class Ticket:
    """What the guards hand to a handler. `ev` is filled in stage by stage and written once."""
    principal: Principal
    route: str
    limits: PlanLimits
    input_tokens: int
    ev: RequestEvent
    mode: str = "full"                                  # Day 88: "full" | "retrieval_only"
    force_fallback: bool = False
    background: BackgroundTasks | None = None
    reservation: Reservation | None = None
    stream_key: str | None = None
    check: tuple | None = None                          # (CheckedAnswer, hits, question) for Day 86
    started: float = field(default_factory=time.perf_counter)


def hour_key(p: Principal) -> str:                      # feeds Day 88's per-tenant hourly caps
    return tenant_key(p, "tokens_h", time.strftime("%Y%m%d%H", time.gmtime()))


async def guard(p: Principal, route: str, text: str, deps: Deps, *, request_id: str,
                background: BackgroundTasks | None = None) -> Ticket:
    """Stages 4-6: kill switches, input cap, rate limit. Cheapest first; any of them can say no."""
    with step_span("guard", "pipeline", **{"atlas.tenant": p.tenant_id, "atlas.route": route}):
        used = int(await deps.redis.get(hour_key(p)) or 0)
        ov = get_overrides(deps.ops_redis)              # sync, but read at most every 5 s (Day 88)
        plan = plan_request(route, p.tenant_id, ov, tokens_this_hour=used)
        if plan.status == 202:                          # queue mode: R5.0 has no queued ask path
            raise HTTPException(503, plan.message, headers={"Retry-After": "30"})
        if plan.status >= 400:
            raise HTTPException(plan.status, plan.message)
        limits, n_in = PLANS[p.plan], count_tokens(text)
        if n_in > limits.max_input_tokens:
            raise QuotaExceeded("input_too_large")      # 413 before Redis or a model is touched
        rate = await deps.limiter.hit(tenant_key(p, "rpm"), limits.rpm)
        if not rate.allowed:
            raise QuotaExceeded("rate_limited", rate.retry_after_s)   # 429 + Retry-After
    ev = RequestEvent(request_id=request_id, ts=time.time(), tenant_id=p.tenant_id, route=route,
                      model="none", status=200, latency_ms=0.0)
    return Ticket(p, route, limits, n_in, ev, plan.mode, plan.force_fallback, background)

# script_id: day_089_capstone_r5_integrate/atlas/api/pipeline
async def reserve_stream(t: Ticket, deps: Deps) -> None:
    """Streams must say no BEFORE the 200 goes out, so they reserve inside the route dependency."""
    out = t.limits.max_output_tokens
    est = INVESTIGATE_TOKENS if t.route == "investigate" else t.input_tokens + CONTEXT_TOKENS + out
    t.reservation = await deps.quotas.reserve(t.principal, est)
    try:
        t.stream_key = await deps.quotas.acquire_stream(t.principal)
    except QuotaExceeded:
        await deps.quotas.settle(t.reservation, 0)      # hand the tokens back
        raise


async def answer(t: Ticket, question: str, deps: Deps) -> AskResult:
    """Stages 7-10: cache, reserve, route, resilient call. Fills t.ev as it goes."""
    p, ev, out = t.principal, t.ev, t.limits.max_output_tokens
    prompt_v, ev.variant = resolve_prompt("ask", p, deps.release.get(), user_id=p.key_id)   # Day 87
    route = choose_model(question, plan=p.plan)                                              # Day 85
    fallbacks = deps.settings.fallback_models
    model = fallbacks[0] if t.force_fallback and fallbacks else route.model                 # Day 88 lever
    scope = cache_scope(p, model=model, prompt_version=prompt_v, index_version=deps.settings.index_version)
    scope += f":orbit={detect_version(question)}"           # as Day 85's cached_ask() does
    with step_span("cache", "pipeline"):
        hit = await asyncio.to_thread(deps.cache.lookup, scope, question)   # embedding may be a network call
    if hit:
        ev.model, ev.cached, ev.citations = model, True, len(hit.answer.citations)
        return AskResult(answer=hit.answer, usage={}, request_id=ev.request_id, model=model,
                         variant=ev.variant, cached=True)
    if t.reservation is None:                               # JSON routes reserve only on a miss
        t.reservation = await deps.quotas.reserve(p, t.input_tokens + CONTEXT_TOKENS + out)
    try:
        if t.mode == "retrieval_only":                      # Day 88's lever takes the outage path
            raise AllModelsUnavailable([], 0)
        with step_span("generate", "llm", **{"atlas.model": model}):
            o = await deps.services.ask(question, groups=visible_tenants(p), prompt_version=prompt_v,
                                        llm=deps.llms.get(model, deps.llms[deps.settings.model]),
                                        effort=min(route.effort, t.limits.effort, key=EFFORTS.index),
                                        max_output_tokens=out)
    except AllModelsUnavailable:                            # Day 84: cited links, no generation
        d = degraded_answer(question, keyword_search(question, deps.chunks))
        ev.fallback = ev.degraded = True
        ev.citations = len(d.citations)
        return AskResult(answer=d.model_dump(include={"text", "citations"}), usage={}, model="none",
                         request_id=ev.request_id, variant=ev.variant, fallback=True, degraded=True)

# script_id: day_089_capstone_r5_integrate/atlas/api/pipeline
    # (answer continues: a model answered)
    ev.model, ev.fallback, ev.citations = o.meta.model, o.meta.fallback, len(o.answer.citations)
    ev.input_tokens, ev.output_tokens = o.input_tokens, o.output_tokens
    ev.cost_usd = cost_usd(o.meta.model, o.input_tokens, o.output_tokens)   # real usage x Day 85 prices
    t.check = (CheckedAnswer(o.answer.text, [c.doc_id for c in o.answer.citations]), o.hits, question)
    if o.answer.citations and not o.meta.fallback:          # cache only healthy, cited answers
        await asyncio.to_thread(deps.cache.store, scope, question, o.answer)
    usage = {"input_tokens": o.input_tokens, "output_tokens": o.output_tokens, "cost_usd": ev.cost_usd}
    return AskResult(answer=o.answer, usage=usage, request_id=ev.request_id, model=o.meta.model,
                     variant=ev.variant, fallback=o.meta.fallback)


async def finish(t: Ticket, deps: Deps, *, status: int) -> None:
    """Stages 11-12. Always runs (call it from `finally`): settle tokens, then write ONE event row."""
    used = t.ev.input_tokens + t.ev.output_tokens
    with anyio.CancelScope(shield=True):                    # a cancelled request still settles (Day 83)
        if t.reservation:
            await deps.quotas.settle(t.reservation, used)
            await deps.redis.incrby(hour_key(t.principal), used)
            await deps.redis.expire(hour_key(t.principal), 7_200)
        if t.stream_key:
            await deps.quotas.release_stream(t.stream_key)
    t.ev.status, t.ev.latency_ms = status, round((time.perf_counter() - t.started) * 1000, 1)
    deps.events.record(t.ev)                                # metrics only: no question or answer text
    if t.check and t.background is not None and status == 200:
        t.background.add_task(score_request, deps.events, t.ev, *t.check)   # Day 86, after the response


async def handle_ask(t: Ticket, req: AskRequest, deps: Deps) -> AskResult:
    status = 500
    try:
        result = await answer(t, req.question, deps)
        status = 200
        return result
    except HTTPException as exc:
        status = exc.status_code
        raise
    except QuotaExceeded:
        status = 429
        raise
    finally:
        await finish(t, deps, status=status)

# script_id: day_089_capstone_r5_integrate/atlas/api/pipeline
async def handle_stream(t: Ticket, req: AskRequest | InvestigateRequest,
                        deps: Deps) -> AsyncIterator[AgentEvent]:
    """Stages 7-12 for SSE routes. Auth, guards and the reservation already ran in the dependency."""
    status = 500
    try:
        if isinstance(req, AskRequest):
            r = await answer(t, req.question, deps)
            if t.check:
                yield AgentEvent(type="status", data={"step": "retrieved", "doc_ids": t.check[1]})
            yield AgentEvent(type="token", data={"text": r.answer.text})   # one chunk: see Go Deeper
            for c in r.answer.citations:
                yield AgentEvent(type="citation", data={"doc_id": c.doc_id})
            yield AgentEvent(type="final", data=r.model_dump(
                include={"request_id", "model", "cached", "fallback", "degraded"}))
        else:
            async for e in deps.services.investigate(req.issue_id):
                if e.type == "final":                       # the run's real usage, priced like ask
                    u, t.ev.model = e.data.get("usage") or {}, deps.settings.model   # R3 runs on ATLAS_MODEL
                    t.ev.input_tokens, t.ev.output_tokens = u.get("input_tokens", 0), u.get("output_tokens", 0)
                    t.ev.cost_usd = cost_usd(t.ev.model, t.ev.input_tokens, t.ev.output_tokens)
                    e = e.model_copy(update={"data": {**e.data, "request_id": t.ev.request_id}})
                elif e.type == "citation":
                    t.ev.citations += 1
                yield e
        status = 200
    except (GeneratorExit, asyncio.CancelledError):
        status = 499                                        # the client left: not an outage
        raise
    finally:
        await finish(t, deps, status=status)

# script_id: day_089_capstone_r5_integrate/atlas/api/app
# atlas/api/app.py: create_app(settings) wires middleware, routes and the composition root.
# Run: ATLAS_BACKEND=stub ATLAS_REDIS_URL=fake:// uv run uvicorn atlas.api.app:create_app --factory --port 8000
# Install: uv add "fastapi>=0.162" uvicorn "redis>=8.1" fakeredis   Services: Redis (or fake://)
# No `from __future__ import annotations` here: FastAPI must see the Depends() inside closures.
import asyncio
import re
import uuid
from collections.abc import AsyncIterable, AsyncIterator
from contextlib import asynccontextmanager
from typing import Annotated

from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.sse import EventSourceResponse, ServerSentEvent

MAX_BODY_BYTES = 64_000                  # JSON bodies only; AskRequest caps a question at 2,000 chars
_RID = re.compile(r"[A-Za-z0-9-]{8,64}")


class RequestContext:
    """Give every request an id before anything else runs (a caller's valid X-Request-Id is kept)."""
    def __init__(self, app) -> None:
        self.app = app

    async def __call__(self, scope, receive, send) -> None:
        if scope["type"] == "http":
            sent = dict(scope["headers"]).get(b"x-request-id", b"").decode()
            scope.setdefault("state", {})["request_id"] = sent if _RID.fullmatch(sent) else uuid.uuid4().hex
        await self.app(scope, receive, send)


class BodyLimit:
    """413 for big bodies, 411 when a POST has no Content-Length. Runs before any JSON parsing."""
    def __init__(self, app, max_bytes: int = MAX_BODY_BYTES) -> None:
        self.app, self.max_bytes = app, max_bytes

    async def __call__(self, scope, receive, send) -> None:
        if scope["type"] == "http" and scope["method"] in ("POST", "PUT", "PATCH"):
            raw = dict(scope["headers"]).get(b"content-length")
            status = 411 if raw is None else 413 if not raw.isdigit() or int(raw) > self.max_bytes else 0
            if status:
                body = {"error": "length_required" if status == 411 else "body_too_large"}
                await JSONResponse(body, status_code=status)(scope, receive, send)
                return
        await self.app(scope, receive, send)

# script_id: day_089_capstone_r5_integrate/atlas/api/app
from atlas.api.auth import Principal, require_principal, require_scope
from atlas.api.deps import Deps, build_deps
from atlas.api.pipeline import Ticket, guard, reserve_stream
from atlas.config import Settings
from atlas.feedback import feedback_router
from atlas.quotas import QuotaExceeded
from atlas.tracing import setup_tracing


async def on_quota(request: Request, exc: QuotaExceeded) -> JSONResponse:    # Day 83's mapping
    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)


def admitted(route: str, *, stream: bool = False):
    """Dependency factory: auth, guards and (for streams) the reservation, BEFORE any byte is sent."""
    async def dep(request: Request, background: BackgroundTasks,
                  p: Annotated[Principal, Depends(require_scope(route))]) -> Ticket:
        deps: Deps = request.app.state.deps
        raw = (await request.body()).decode()           # cap the raw JSON: covers every body shape
        t = await guard(p, route, raw, deps, request_id=request.state.request_id, background=background)
        if stream:
            await reserve_stream(t, deps)
        return t
    return dep


def create_app(settings: Settings | None = None) -> FastAPI:
    settings = settings or Settings.from_env()          # ConfigError lists every problem, at boot

    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
        if settings.phoenix_endpoint:
            setup_tracing("atlas", endpoint=settings.phoenix_endpoint)   # Day 42 (endpoint is keyword-only)
        app.state.deps = deps = build_deps(settings)
        app.state.keys = deps.keys                      # Day 83's require_principal reads this
        app.include_router(feedback_router(deps.events, require_principal))   # Day 86 needs the log
        yield
        await deps.redis.aclose()
        if deps.client is not None:
            await deps.client.close()

    app = FastAPI(title="Atlas", version=settings.release, lifespan=lifespan)
    app.add_exception_handler(QuotaExceeded, on_quota)
    app.add_middleware(BodyLimit)                       # added first = innermost
    app.add_middleware(CORSMiddleware, allow_origins=list(settings.allowed_origins),
                       allow_methods=["GET", "POST"], allow_headers=["Authorization", "Content-Type"])
    app.add_middleware(RequestContext)                  # added last = outermost: runs first
    add_routes(app)
    return app

# script_id: day_089_capstone_r5_integrate/atlas/api/app
from atlas.api.events import sse_stream
from atlas.api.models import AskRequest, InvestigateRequest, TriageRequest
from atlas.api.pipeline import AskResult, finish, handle_ask, handle_stream
from atlas.ops import get_overrides


def add_routes(app: FastAPI) -> None:
    @app.post("/v1/ask")
    async def ask(body: AskRequest, request: Request,
                  t: Annotated[Ticket, Depends(admitted("ask"))]) -> AskResult:
        return await handle_ask(t, body, request.app.state.deps)

    @app.post("/v1/ask/stream", response_class=EventSourceResponse)
    async def ask_stream(body: AskRequest, request: Request,
                         t: Annotated[Ticket, Depends(admitted("ask", stream=True))]
                         ) -> AsyncIterable[ServerSentEvent]:
        async for sse in sse_stream(handle_stream(t, body, request.app.state.deps), request, timeout_s=60):
            yield sse

    @app.post("/v1/investigate/stream", response_class=EventSourceResponse)
    async def investigate(body: InvestigateRequest, request: Request,
                          t: Annotated[Ticket, Depends(admitted("investigate", stream=True))]
                          ) -> AsyncIterable[ServerSentEvent]:
        async for sse in sse_stream(handle_stream(t, body, request.app.state.deps), request, timeout_s=120):
            yield sse

    @app.post("/v1/triage")
    async def triage(body: TriageRequest, request: Request,
                     t: Annotated[Ticket, Depends(admitted("triage"))]) -> dict:
        status = 500
        try:
            proposal = await request.app.state.deps.services.triage(body.issue_id)
            status = 200
            return proposal                             # proposes only: a maintainer approves (Day 76)
        except HTTPException as exc:
            status = exc.status_code
            raise
        finally:
            await finish(t, request.app.state.deps, status=status)

# script_id: day_089_capstone_r5_integrate/atlas/api/app
    # (add_routes continues)
    @app.get("/healthz")
    async def healthz(request: Request) -> dict:       # liveness + "what is deployed?" (Day 82)
        return {"status": "ok", "release": request.app.state.deps.settings.release}

    @app.get("/readyz")
    async def readyz(request: Request) -> dict:        # local dependencies only, NEVER the model
        deps: Deps = request.app.state.deps
        ov = get_overrides(deps.ops_redis)
        checks = {"redis": await _redis_ok(deps.redis),
                  "postgres": await asyncio.to_thread(_postgres_ok, deps.settings.database_url),
                  "index": ov.index_version in (None, deps.settings.index_version) and bool(deps.chunks)}
        if not all(checks.values()):
            raise HTTPException(503, detail=checks)
        return {"status": "ready", "backend": deps.services.name, "checks": checks}


async def _redis_ok(redis) -> bool:
    try:
        return bool(await redis.ping())
    except Exception:                                   # readiness reports problems; it never raises
        return False


def _postgres_ok(url: str | None) -> bool:
    if url is None:
        return True                                     # not configured: stub backend or memory index
    try:
        import psycopg
        with psycopg.connect(url, connect_timeout=2) as conn:
            conn.execute("SELECT 1")
        return True
    except Exception:
        return False

# script_id: day_089_capstone_r5_integrate/atlas/ui/main
# atlas/ui/main.py: Day 79's chat, connected to R5: per-session API key, Day 80 cards, Day 86 feedback.
# Install: uv add "gradio>=6.31" httpx2   Env: ATLAS_API_URL (http://api:8000 in compose)
# Run: uv run python -m atlas.ui.main   (Day 79's atlas/ui/app.py stays as is; this file adds the wiring)
from __future__ import annotations

import os
from collections.abc import Iterator

import gradio as gr
import httpx2

from atlas.ui.app import AgentEvent, _route, iter_sse, to_messages
from atlas.ui.components import render_html

API_URL = os.environ.get("ATLAS_API_URL", "http://localhost:8000")   # from env, never hard-coded


def chat_fn(message: str, history: list, api_key: str) -> Iterator[tuple[list[gr.ChatMessage], str]]:
    path, body = _route(message)                    # "#209" investigates; anything else asks
    events, cards, request_id = [], [], ""
    headers = {"Authorization": f"Bearer {api_key.strip()}"}
    try:
        with httpx2.stream("POST", API_URL + path, json=body, headers=headers,
                           timeout=httpx2.Timeout(10.0, read=130.0)) as r:   # > the 120 s agent deadline
            r.raise_for_status()
            for e in iter_sse(r.iter_lines()):
                if e.type == "component":           # validated again, then filled into a fixed template
                    cards.append(gr.ChatMessage(content=render_html(e.data["name"], e.data["props"])))
                    continue
                request_id = e.data.get("request_id", request_id)
                events.append(e)
                yield cards + to_messages(events), request_id
    except httpx2.HTTPStatusError as exc:            # 401 bad key, 429 slow down, 503 paused
        events.append(AgentEvent(type="error", data={"code": exc.response.status_code}))
    except httpx2.HTTPError as exc:
        events.append(AgentEvent(type="error", data={"code": type(exc).__name__}))
    yield cards + to_messages(events), request_id

# script_id: day_089_capstone_r5_integrate/atlas/ui/main
def send_feedback(rating: str, api_key: str, request_id: str) -> str:
    if not request_id:
        return "Ask something first."
    r = httpx2.post(f"{API_URL}/v1/feedback", json={"request_id": request_id, "rating": rating},
                    headers={"Authorization": f"Bearer {api_key.strip()}"}, timeout=10.0)
    return "Thanks, recorded." if r.status_code == 204 else f"Feedback failed ({r.status_code})."


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Atlas") as demo:
        api_key = gr.Textbox(label="Atlas API key", type="password", placeholder="atk_...")
        last_id = gr.State("")                      # the latest answer's request_id, per session
        gr.ChatInterface(fn=chat_fn, additional_inputs=[api_key], additional_outputs=[last_id],
                         chatbot=gr.Chatbot(height=520, sanitize_html=True, render_markdown=True),
                         description="Ask about Orbit, or type #209 to investigate an issue.")
        with gr.Row():
            helpful, wrong = gr.Button("Helpful"), gr.Button("Wrong")
        note = gr.Markdown()
        helpful.click(lambda key, rid: send_feedback("up", key, rid), [api_key, last_id], note)
        wrong.click(lambda key, rid: send_feedback("down", key, rid), [api_key, last_id], note)
    return demo


if __name__ == "__main__":
    build_ui().launch(server_name=os.environ.get("ATLAS_UI_HOST", "127.0.0.1"), server_port=7860)

# script_id: day_089_capstone_r5_integrate/scripts/seed_tenants
# scripts/seed_tenants.py: create the four staging tenants and their API keys.
# Run: uv run python -m scripts.seed_tenants   Env: ATLAS_STATE_DIR (where keys.db lives)
# Raw keys go ONCE into a 0600 file; only key ids are printed, so keys never reach logs.
from __future__ import annotations

import json
import os
from pathlib import Path

from atlas.api.auth import KeyStore

TENANTS = {   # tenant -> (plan, scopes); limits per plan come from Day 83's PLANS
    "orbit-maintainers": ("team", {"ask", "investigate", "triage"}),
    "acme": ("free", {"ask", "investigate"}),
    "zephyr": ("free", {"ask", "investigate"}),
    "demo": ("demo", {"ask"}),
}


def seed_tenants(store: KeyStore) -> dict[str, str]:
    """Issue one key per tenant and return tenant -> raw key (shown nowhere else, ever)."""
    return {tenant: store.issue(tenant, plan, scopes) for tenant, (plan, scopes) in TENANTS.items()}


if __name__ == "__main__":
    state = Path(os.environ.get("ATLAS_STATE_DIR", ".atlas"))
    state.mkdir(parents=True, exist_ok=True)
    out = state / "staging_keys.json"
    keys = seed_tenants(KeyStore(str(state / "keys.db")))
    fd = os.open(out, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)    # owner-only, like an SSH key
    with os.fdopen(fd, "w") as f:
        json.dump(keys, f)
    for tenant, raw in keys.items():
        print(f"{tenant:18} key_id={raw.split('_')[1]}")
    print(f"raw keys -> {out} (mode 600): copy them to your secret store, then delete the file")

# script_id: day_089_capstone_r5_integrate/tests/test_integration
# tests/test_integration.py: 8 end-to-end tests at the HTTP boundary. No key, no Docker.
# Install: uv add --dev pytest fakeredis   Run: uv run pytest tests/test_integration.py -q
from __future__ import annotations

import os
from types import SimpleNamespace

import pytest
from fastapi.testclient import TestClient

from atlas import quotas
from atlas.api.app import create_app
from atlas.api.deps import build_pools
from atlas.config import Settings
from atlas.flags import RELEASE_TOML
from atlas.resilience import FakeAsyncClient, FaultInjector, RetryPolicy
from scripts.seed_tenants import seed_tenants

FAST = RetryPolicy(max_attempts=2, base_s=0.01, cap_s=0.02)     # tests shouldn't sleep for real


@pytest.fixture
def api(tmp_path, monkeypatch):
    (tmp_path / "release.toml").write_text(RELEASE_TOML)         # Day 87: prompt.ask v3 -> v4 at 10%
    settings = Settings.from_env({"ATLAS_BACKEND": "stub", "ATLAS_REDIS_URL": "fake://",
                                  "ATLAS_STATE_DIR": str(tmp_path),
                                  "ATLAS_RELEASE_FILE": str(tmp_path / "release.toml"),
                                  "ORBIT_DATA": os.environ.get("ORBIT_DATA", "datasets/orbit")})
    monkeypatch.setattr(quotas, "time", SimpleNamespace(time=lambda: 1_800_000_000.0))  # one rate window
    with TestClient(create_app(settings)) as client:
        client.keys = seed_tenants(client.app.state.deps.keys)
        yield client


def ask(c: TestClient, tenant: str | None, q: str = "How do retries back off?", path: str = "/v1/ask"):
    headers = {"Authorization": f"Bearer {c.keys[tenant]}"} if tenant else {}
    return c.post(path, json={"question": q}, headers=headers)

# script_id: day_089_capstone_r5_integrate/tests/test_integration
def test_1_no_key_is_401(api):
    assert ask(api, None).status_code == 401
    assert api.post("/v1/ask", json={"question": "x"}, headers={"Authorization": "Bearer atk_x_y"}).status_code == 401


def test_2_tenants_are_isolated(api):
    q = "What lease_seconds do we use for the nightly ETL queue?"
    mine = ask(api, "acme", q).json()
    theirs = ask(api, "zephyr", q).json()
    assert any(c["doc_id"] == "acme/runbook" for c in mine["answer"]["citations"])
    assert all(not c["doc_id"].startswith("acme/") for c in theirs["answer"]["citations"])
    assert "900" not in theirs["answer"]["text"]                  # Acme's private value
    r = api.post("/v1/feedback", json={"request_id": mine["request_id"], "rating": "down"},
                 headers={"Authorization": f"Bearer {api.keys['zephyr']}"})
    assert r.status_code == 404                                   # a foreign id looks like no id


def test_3_rate_limit_61st_request(api):
    codes = [ask(api, "acme").status_code for _ in range(61)]     # free plan: 60 requests/minute
    assert codes[:60] == [200] * 60 and codes[60] == 429
    assert int(ask(api, "acme").headers["Retry-After"]) > 0


def test_4_cache_hits_are_per_tenant(api):
    q = "How do I fix ORB-1101 errors?"
    assert ask(api, "acme", q).json()["cached"] is False
    assert ask(api, "acme", q).json()["cached"] is True
    assert ask(api, "zephyr", q).json()["cached"] is False        # same words, other tenant: miss

# script_id: day_089_capstone_r5_integrate/tests/test_integration
def test_5_fallback_then_degraded(api):
    deps = api.app.state.deps
    deps.llms, deps.breakers = build_pools(FakeAsyncClient(FaultInjector({}, down={"gpt-6-luna"})),
                                           list(deps.llms), policy=FAST)
    body = ask(api, "acme", "How do workers send heartbeats?").json()
    assert body["fallback"] is True and body["model"] == "gpt-6.1-terra"
    deps.llms, deps.breakers = build_pools(FakeAsyncClient(FaultInjector({}, down=set(deps.llms))),
                                           list(deps.llms), policy=FAST)
    r = ask(api, "acme", "How does the dead letter queue work?")
    assert r.status_code == 200 and r.json()["degraded"] is True   # an answer, never a 500


def test_6_sse_has_citations_and_ends_with_done(api):
    r = ask(api, "acme", path="/v1/ask/stream")
    names = [line.removeprefix("event: ") for line in r.text.splitlines() if line.startswith("event: ")]
    assert r.status_code == 200 and "citation" in names and names[-1] == "done"


def test_7_canary_bucket_is_stable(api):
    for q in ["How do retries back off?", "What is a lease?", "How do I pause a queue?"]:
        assert ask(api, "zephyr", q).status_code == 200
    variants = {e.variant for e in api.app.state.deps.events.window(5, tenant_id="zephyr")}
    assert len(variants) == 1                                     # one user, one variant, every time


def test_8_one_event_row_per_request(api):
    ids = [ask(api, "zephyr", q).json()["request_id"] for q in ["What is Orbit?", "What is a flow?"]]
    ids.append(ask(api, "zephyr", "What is Orbit?").json()["request_id"])   # a cache hit
    rows = api.app.state.deps.events.window(5, tenant_id="zephyr")
    assert sorted(e.request_id for e in rows) == sorted(ids)
    assert [e.cost_usd > 0 for e in sorted(rows, key=lambda e: e.ts)] == [True, True, False]

$ uv run pytest tests/test_integration.py -q
........                                                                 [100%]
8 passed in 1.2s            (your timing will differ)

# file: compose.yaml  (changes from Day 82: R5 commands, state volume, Postgres; api/ui healthchecks unchanged)
x-atlas: &atlas
  image: ${ATLAS_IMAGE:-atlas}:${ATLAS_TAG:-dev}
  env_file: ${ATLAS_ENV_FILE:-.env}
  environment:
    ATLAS_ENV: ${ATLAS_ENV:-staging}
    ATLAS_BACKEND: ${ATLAS_BACKEND:-stub}               # reference once R2-R4 + Postgres are in place
    ATLAS_REDIS_URL: redis://redis:6379/0
    ATLAS_DATABASE_URL: postgresql://atlas:atlas@postgres:5432/atlas
    ATLAS_API_URL: http://api:8000
    ATLAS_STATE_DIR: /app/state
    GRADIO_SERVER_NAME: 0.0.0.0
  volumes: [atlas-state:/app/state]
  restart: unless-stopped

services:
  api:
    <<: *atlas
    command: ["uvicorn", "atlas.api.app:create_app", "--factory", "--host", "0.0.0.0", "--port", "8000",
              "--timeout-graceful-shutdown", "130"]
    depends_on:
      redis: { condition: service_healthy }
      postgres: { condition: service_healthy }
  ui:
    <<: *atlas
    command: ["python", "-m", "atlas.ui.main"]
    environment: { ATLAS_UI_HOST: 0.0.0.0, ATLAS_API_URL: http://api:8000 }
  postgres:
    image: pgvector/pgvector:0.8.7-pg18-trixie
    environment: { POSTGRES_USER: atlas, POSTGRES_PASSWORD: atlas, POSTGRES_DB: atlas }
    healthcheck: { test: ["CMD", "pg_isready", "-U", "atlas"], interval: 5s, timeout: 3s, retries: 10 }
    volumes: [pg-data:/var/lib/postgresql]              # PG 18 images moved the data path
  # redis and worker: unchanged from Day 82

volumes: { atlas-state: {}, pg-data: {}, redis-data: {}, atlas-logs: {} }