423 lines
16 KiB
Python
423 lines
16 KiB
Python
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("<html>ui</html>", 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 == "<html>ui</html>"
|
|
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 == "<html>ui</html>"
|
|
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"]
|