import asyncio import tempfile from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any, cast import pytest from claude_agent_sdk import ( AssistantMessage, RateLimitEvent, RateLimitInfo, ResultMessage, TextBlock, ToolResultBlock, ToolUseBlock, UserMessage, ) from httpx import ASGITransport, AsyncClient from test_conversations import ScriptedClient, World from beaver_gateway.core.conversation_store import rewrite_messages from beaver_gateway.core.auth import TokenStore from beaver_gateway.core.registry import McpRegistry from beaver_gateway.core.transcript import build_entries from beaver_gateway.frontends.admin import AdminFrontend from beaver_gateway.frontends.admin.frontend import build_app as build_admin from beaver_gateway.frontends.api import ApiFrontend from beaver_gateway.frontends.api.frontend import build_app as build_api from beaver_gateway.frontends.base import Frontend, GatewayRuntime from beaver_gateway.frontends.root import build_root_app from beaver_gateway.storage.models import RateLimit, Usage TOKEN = "tok" HEADERS = {"Authorization": f"Bearer {TOKEN}"} class ToolClient(ScriptedClient): async def receive_response(self): yield AssistantMessage( content=[ToolUseBlock(id="tu_1", name="Bash", input={"command": "ls"})], model="m", ) yield AssistantMessage( content=[ToolUseBlock(id="tu_2", name="Read", input={"file_path": "x"})], model="m", parent_tool_use_id="tu_1", ) yield RateLimitEvent( rate_limit_info=RateLimitInfo( status="allowed_warning", resets_at=int(datetime.now(UTC).timestamp()) + 3600, rate_limit_type="five_hour", utilization=0.8, raw={"status": "allowed_warning"}, ), uuid="r", session_id=self.session_id, ) hold = ScriptedClient.hold if hold is not None: await hold.wait() yield UserMessage(content=[ToolResultBlock(tool_use_id="tu_1", content="done")]) yield AssistantMessage(content=[TextBlock(text="ok")], model="m") yield ResultMessage( subtype="success", duration_ms=1, duration_api_ms=1, is_error=False, num_turns=1, session_id=self.session_id, stop_reason="end_turn", total_cost_usd=0.5, usage={"input_tokens": 10, "output_tokens": 5}, model_usage=cast( "Any", { "claude-opus-5": { "inputTokens": 8, "outputTokens": 4, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0, "webSearchRequests": 1, "costUSD": 0.4, }, "claude-haiku-4-5": { "inputTokens": 2, "outputTokens": 1, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 0, "webSearchRequests": 0, "costUSD": 0.1, }, }, ), ) class Api: def __init__(self, world: World, memory_root: Path | None = None) -> None: self.world = world self.store = TokenStore(world.db, bootstrap={"t": TOKEN}) self.runtime = GatewayRuntime( agents=world.conversations._agents, # noqa: SLF001 mcps=McpRegistry([]), backends={"a": world.backend, "d": world.deep_backend}, token_store=self.store, db=world.db, admin_user="admin", admin_pass="secret", session_secret="s" * 32, frontends=(world.api, world.markdown), conversations=world.conversations, bus=world.bus, pool=world.pool, ) self.app = build_api(self.runtime, memory_root=memory_root) self.http = AsyncClient( transport=ASGITransport(app=self.app), base_url="http://api" ) def root(self, *frontends: Frontend) -> AsyncClient: return AsyncClient( transport=ASGITransport(app=build_root_app(frontends)), base_url="http://gw" ) async def get(self, path: str, params: dict[str, Any] | None = None) -> Any: res = await self.http.get(path, params=params, headers=HEADERS) assert res.status_code == 200, res.text return res.json() @pytest.fixture async def world() -> World: root = Path(tempfile.mkdtemp(prefix="beaver-api-")) w = await World(root).setup() yield w await w.conversations.stop() await w.pool.close_all() await w.db.dispose() async def seed_usage(world: World, rows: list[dict[str, Any]]) -> None: async with world.db.session() as session: for row in rows: session.add(Usage(**row)) await session.commit() async def test_usage_groups_by_agent_day_and_model(world: World) -> None: api = Api(world) now = datetime.now(UTC).replace(tzinfo=None) await seed_usage( world, [ { "ts": now - timedelta(hours=1), "agent_name": "a", "conversation_id": "c1", "model": "claude-opus-5", "input_tokens": 10, "output_tokens": 20, "cache_read_tokens": 100, "cache_creation_tokens": 5, "cost_usd": 0.5, "model_usage": { "claude-opus-5": { "inputTokens": 10, "outputTokens": 20, "cacheReadInputTokens": 100, "cacheCreationInputTokens": 5, "webSearchRequests": 2, "costUSD": 0.5, } }, }, { "ts": now - timedelta(days=2), "agent_name": "d", "conversation_id": "c2", "model": "claude-sonnet-5", "input_tokens": 1, "output_tokens": 2, "cost_usd": 0.1, }, ], ) by_agent = await api.get("/usage", {"group_by": "agent"}) assert [r["agent"] for r in by_agent["rows"]] == ["a"] assert by_agent["total"] == { "turns": 1, "input": 10, "output": 20, "cache_read": 100, "cache_creation": 5, "cost_usd": 0.5, "web_searches": 2, } since = (datetime.now(UTC) - timedelta(days=3)).isoformat() by_day = await api.get("/usage", {"group_by": "day", "since": since}) assert len(by_day["rows"]) == 2 assert by_day["rows"][0]["day"] < by_day["rows"][1]["day"] by_model = await api.get("/usage", {"group_by": "model", "since": since}) models = {r["model"]: r for r in by_model["rows"]} assert models["claude-opus-5"]["web_searches"] == 2 assert models["claude-sonnet-5"]["output"] == 2 by_conv = await api.get("/usage", {"group_by": "conversation", "since": since}) assert {r["conversation"] for r in by_conv["rows"]} == {"c1", "c2"} bad = await api.http.get("/usage", params={"group_by": "x"}, headers=HEADERS) assert bad.status_code == 400 async def test_limits_report_latest_window_with_gateway_spend(world: World) -> None: api = Api(world) now = datetime.now(UTC) async with world.db.session() as session: session.add( RateLimit( window="five_hour", status="allowed", utilization=0.2, resets_at=now + timedelta(hours=4), ) ) session.add( RateLimit( window="five_hour", status="allowed_warning", utilization=0.9, resets_at=now + timedelta(hours=4), ) ) session.add(RateLimit(window="seven_day", status="allowed", utilization=0.3)) await session.commit() await seed_usage( world, [ { "ts": (now - timedelta(minutes=30)).replace(tzinfo=None), "agent_name": "a", "model": "m", "output_tokens": 7, "cost_usd": 0.25, } ], ) out = await api.get("/limits") windows = {w["window"]: w for w in out["windows"]} assert windows["five_hour"]["utilization"] == 0.9 assert windows["five_hour"]["status"] == "allowed_warning" assert windows["five_hour"]["gateway"]["output"] == 7 assert windows["five_hour"]["gateway"]["cost_usd"] == 0.25 assert windows["seven_day"]["utilization"] == 0.3 assert [w["window"] for w in out["windows"]] == ["five_hour", "seven_day"] assert len(out["history"]) == 3 async def test_memory_tree_and_file(world: World) -> None: root = world.root / "zone" (root / "дни").mkdir(parents=True) (root / "дни" / "2026-08-27.md").write_text("# day", encoding="utf-8") (root / "состояние.md").write_text("# state", encoding="utf-8") (root / ".hidden").write_text("x", encoding="utf-8") api = Api(world, memory_root=root) tree = await api.get("/memory") names = [n["name"] for n in tree["tree"]] assert names == ["дни", "состояние.md"] assert tree["tree"][0]["children"][0]["path"] == "дни/2026-08-27.md" file = await api.get("/memory/file", {"path": "дни/2026-08-27.md"}) assert file["content"] == "# day" escape = await api.http.get( "/memory/file", params={"path": "../w.db"}, headers=HEADERS ) assert escape.status_code == 404 unset = Api(world) res = await unset.http.get("/memory", headers=HEADERS) assert res.status_code == 404 async def test_describe_snapshots_open_tools_and_records_rate_limit( world: World, ) -> None: world.backend._factory = ToolClient # noqa: SLF001 api = Api(world) async def record(event: Any) -> None: async with world.db.session() as session: session.add( Usage( agent_name=event.agent_name, conversation_id=event.conversation_id, model=event.model, cost_usd=event.usage.cost_usd, model_usage=event.usage.model_usage, ) ) await session.commit() world.backend._usage_sink = record # noqa: SLF001 conv = await world.conversations.create(kind="master", agent="a", origin="test") ScriptedClient.hold = asyncio.Event() await world.conversations.post(conv, "go") await asyncio.sleep(0.3) described = await api.get(f"/conversations/{conv.external_id}") turn = described["turn"] assert turn is not None and turn["id"] == described["running_turn"] tools = {t["tool_use_id"]: t for t in turn["tools"]} assert tools["tu_1"]["name"] == "Bash" and tools["tu_1"]["ended_at"] is None assert tools["tu_2"]["parent_tool_use_id"] == "tu_1" ScriptedClient.hold.set() await world.settle(conv, 1) described = await api.get(f"/conversations/{conv.external_id}") assert described["turn"] is None limits = await api.get("/limits") assert limits["windows"][0]["utilization"] == 0.8 assert limits["windows"][0]["gateway"]["cost_usd"] == 0.5 usage = await api.get("/usage", {"group_by": "model"}) assert {r["model"] for r in usage["rows"]} == {"claude-opus-5", "claude-haiku-4-5"} assert usage["total"]["web_searches"] == 1 conv = await world.conversations.get(conv.external_id) assert conv is not None and conv.session_id is not None await world.store.append( world.key(conv.session_id), build_entries( [ {"role": "user", "content": "go"}, {"role": "assistant", "content": [{"type": "text", "text": "ok"}]}, ], session_id=conv.session_id, cwd=str(world.root), model="m", ), ) history = await api.get(f"/conversations/{conv.external_id}/history") assert [m["role"] for m in history["messages"]] == ["user", "assistant"] entries = await api.get(f"/conversations/{conv.external_id}/entries") assert entries["total"] == len(entries["entries"]) == 2 page = await api.get( f"/conversations/{conv.external_id}/entries", {"limit": 1, "offset": 0} ) assert page["offset"] == 0 and len(page["entries"]) == 1 async def test_tokens_and_audit_need_admin_scope(world: World) -> None: api = Api(world) api.store.grant("api-only", "weak", scope="api") weak = {"Authorization": "Bearer weak"} assert (await api.http.get("/tokens", headers=weak)).status_code == 403 created = await api.http.post( "/tokens", json={"name": "cursor", "scope": "mcp"}, headers=HEADERS ) assert created.status_code == 201 plaintext = created.json()["plaintext"] await api.store.invalidate() identity = await api.store.verify(plaintext) assert identity is not None and identity.scope == "mcp" listed = await api.get("/tokens") assert [t["name"] for t in listed["tokens"]] == ["cursor"] token_id = listed["tokens"][0]["id"] revoked = await api.http.post(f"/tokens/{token_id}/revoke", headers=HEADERS) assert revoked.status_code == 200 assert (await api.get("/tokens"))["tokens"] == [] audit = await api.get("/audit") assert [r["kind"] for r in audit["records"]] == ["token_revoke", "token_create"] async def test_admin_login_hands_out_the_ui_bearer(world: World) -> None: api = Api(world) ui_dir = world.root / "build" ui_dir.mkdir() (ui_dir / "index.html").write_text("ui", encoding="utf-8") admin = AdminFrontend(ui_dir=ui_dir) admin._app = build_admin(api.runtime, token="ui-bearer", ui_dir=ui_dir) # noqa: SLF001 api.store.grant("admin-ui", "ui-bearer") api_fe = ApiFrontend() api_fe._app = api.app # noqa: SLF001 http = api.root(admin, api_fe) assert (await http.get("/admin/auth/session")).status_code == 401 bad = await http.post( "/admin/auth/login", json={"username": "admin", "password": "nope"} ) assert bad.status_code == 401 ok = await http.post( "/admin/auth/login", json={"username": "admin", "password": "secret"} ) assert ok.status_code == 200 session = (await http.get("/admin/auth/session")).json() assert session == {"user": "admin", "token": "ui-bearer"} spa = await http.get("/admin/conversations/abc") assert spa.status_code == 200 and spa.text == "ui" root = await http.get("/") assert root.status_code == 307 and root.headers["location"] == "/admin/" assert (await http.get("/admin")).status_code == 307 health = await http.get("/healthz") assert health.json()["frontends"] == ["/admin", "/api"] escape = await http.get("/admin/%2e%2e/pyproject.toml") assert escape.status_code == 200 and escape.text == "ui" me = await api.http.get("/agents", headers={"Authorization": "Bearer ui-bearer"}) assert me.status_code == 200 assert (await http.post("/admin/auth/logout")).status_code == 204 assert (await http.get("/admin/auth/session")).status_code == 401 async def test_pre_sdk_rows_show_canonical_history_and_a_title(world: World) -> None: api = Api(world) conv = await world.conversations.create(kind="deep", agent="d", origin="markdown") async with world.db.session() as session: await rewrite_messages( session, conversation_id=cast("int", conv.id), messages=[ {"role": "user", "content": "так смотри, план на май\nвторая строка"}, {"role": "assistant", "content": [{"type": "text", "text": "ок"}]}, ], ) rows = (await api.get("/conversations"))["conversations"] assert [r["title"] for r in rows if r["id"] == conv.external_id] == [ "так смотри, план на май" ] shown = await api.get(f"/conversations/{conv.external_id}") assert shown["title"] == "так смотри, план на май" and shown["session_id"] is None history = await api.get(f"/conversations/{conv.external_id}/history") assert [m["role"] for m in history["messages"]] == ["user", "assistant"] async def test_close_distills_a_deep_chat(world: World) -> None: from test_distill import DistillerClient, deep_chat, distiller config = distiller(world, DistillerClient) api = Api(world) master = await world.conversations.create(kind="master", agent="a", origin="test") chat = await deep_chat(world) res = await api.http.post( f"/conversations/{chat.external_id}/close", headers=HEADERS ) assert res.status_code == 200, res.text body = res.json() assert body["status"] == "closed" and body["error"] is None assert body["digest"] == str(config.dir / "2026-08-29 - тема.md") assert body["text"].count("\n") == 2 assert (await world.conversations.get(chat.external_id)).status == "closed" assert (await world.conversations.queue.recent(master.id))[0].origin == "выжимка" again = await api.http.post( f"/conversations/{chat.external_id}/close", headers=HEADERS ) assert again.status_code == 400 branch = await world.conversations.create(kind="branch", agent="a", origin="test") plain = await api.http.post( f"/conversations/{branch.external_id}/close", headers=HEADERS ) assert plain.status_code == 200 and plain.json()["status"] == "closed" async def test_bind_without_a_window_materializes_one(world: World) -> None: api = Api(world) chat = await world.conversations.create(kind="deep", agent="d", origin="test") world.markdown.materialized.clear() res = await api.http.post( f"/conversations/{chat.external_id}/bind", json={"frontend": "markdown"}, headers=HEADERS, ) assert res.status_code == 200, res.text assert world.markdown.materialized == [chat.external_id] assert res.json()["bindings"] == [ { "frontend": "markdown", "external_id": f"markdown:{chat.external_id}", "visible": True, } ] hidden = await api.http.post( f"/conversations/{chat.external_id}/bind", json={ "frontend": "markdown", "external_id": f"markdown:{chat.external_id}", "visible": False, }, headers=HEADERS, ) assert hidden.json()["bindings"][0]["visible"] is False no_window = await api.http.post( f"/conversations/{chat.external_id}/bind", json={"frontend": "api"}, headers=HEADERS, ) assert no_window.status_code == 400 unknown = await api.http.post( f"/conversations/{chat.external_id}/bind", json={"frontend": "nope"}, headers=HEADERS, ) assert unknown.status_code == 400