287 lines
11 KiB
Python
287 lines
11 KiB
Python
import asyncio
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import cast
|
|
|
|
import frontmatter
|
|
import httpx
|
|
import pytest
|
|
|
|
from beaver_gateway.core.auth import TokenStore
|
|
from beaver_gateway.core.conversation_store import load_messages, rewrite_messages
|
|
from beaver_gateway.core.gateway_tools import _tools
|
|
from beaver_gateway.core.registry import McpRegistry
|
|
from beaver_gateway.frontends.anthropic import AnthropicMessagesFrontend
|
|
from beaver_gateway.frontends.api import ApiFrontend
|
|
from beaver_gateway.frontends.base import GatewayRuntime
|
|
from beaver_gateway.frontends.markdown import MarkdownFrontend
|
|
from beaver_gateway.frontends.root import build_root_app
|
|
from test_conversations import ScriptedClient, World
|
|
|
|
AUTH = {"Authorization": "Bearer tok"}
|
|
|
|
|
|
class Stack:
|
|
def __init__(self, world: World) -> None:
|
|
self.world = world
|
|
self.vault = world.root / "vault"
|
|
self.api = ApiFrontend(master_agent="a")
|
|
self.markdown = MarkdownFrontend(vault_path=self.vault, default_agent="d")
|
|
self.anthropic = AnthropicMessagesFrontend()
|
|
frontends = [self.api, self.anthropic, self.markdown]
|
|
world.conversations._frontends = frontends
|
|
self.runtime = GatewayRuntime(
|
|
agents=world.conversations._agents,
|
|
mcps=McpRegistry([]),
|
|
backends=world.conversations._backends,
|
|
token_store=TokenStore(bootstrap={"t": "tok"}),
|
|
db=world.db,
|
|
conversations=world.conversations,
|
|
bus=world.bus,
|
|
pool=world.pool,
|
|
frontends=tuple(frontends),
|
|
)
|
|
for fe in frontends:
|
|
fe.configure(self.runtime)
|
|
self.mirror = asyncio.create_task(self.markdown.mirror.run())
|
|
self.root = build_root_app(frontends)
|
|
|
|
def client(self, fe) -> httpx.AsyncClient:
|
|
return httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=self.root), base_url=f"http://t{fe.path}"
|
|
)
|
|
|
|
async def close(self) -> None:
|
|
self.mirror.cancel()
|
|
|
|
async def file_of(self, conv_id: str) -> Path:
|
|
conv = await self.world.conversations.get(conv_id)
|
|
path = await self.markdown.mirror.bound_path(conv)
|
|
assert path is not None
|
|
return path
|
|
|
|
async def wait_file(self, path: Path, needle: str, timeout: float = 5.0) -> str:
|
|
deadline = asyncio.get_running_loop().time() + timeout
|
|
while asyncio.get_running_loop().time() < deadline:
|
|
if path.exists() and needle in path.read_text(encoding="utf-8"):
|
|
return path.read_text(encoding="utf-8")
|
|
await asyncio.sleep(0.02)
|
|
msg = f"{needle!r} never showed up in {path}"
|
|
raise AssertionError(msg)
|
|
|
|
|
|
@pytest.fixture
|
|
async def stack() -> Stack:
|
|
root = Path(tempfile.mkdtemp(prefix="beaver-routing-"))
|
|
world = await World(root).setup()
|
|
await world.conversations.start()
|
|
s = Stack(world)
|
|
yield s
|
|
await s.close()
|
|
await world.conversations.stop()
|
|
await world.pool.close_all()
|
|
await world.db.dispose()
|
|
|
|
|
|
async def test_api_rejects_deep_with_dispatcher(stack: Stack) -> None:
|
|
async with stack.client(stack.api) as c:
|
|
r = await c.post(
|
|
"/conversations", json={"kind": "deep", "agent": "a"}, headers=AUTH
|
|
)
|
|
assert r.status_code == 400
|
|
assert "does not serve kind 'deep'" in r.json()["error"]
|
|
r = await c.post("/conversations", json={"kind": "job"}, headers=AUTH)
|
|
assert r.status_code == 400
|
|
assert "no default agent" in r.json()["error"]
|
|
r = await c.get("/agents", headers=AUTH)
|
|
agents = {a["name"]: a["kinds"] for a in r.json()["agents"]}
|
|
assert agents == {"a": ["master", "branch", "job", "fork"], "d": ["deep"]}
|
|
homes = {f["name"]: f["default_agents"] for f in r.json()["frontends"]}
|
|
assert homes == {
|
|
"api": {"master": "a"},
|
|
"anthropic": {},
|
|
"markdown": {"deep": "d"},
|
|
}
|
|
|
|
|
|
async def test_api_spawn_deep_lands_in_vault(stack: Stack) -> None:
|
|
async with stack.client(stack.api) as c:
|
|
r = await c.post("/conversations", json={"kind": "master"}, headers=AUTH)
|
|
assert r.status_code == 201 and r.json()["agent"] == "a"
|
|
master = r.json()["id"]
|
|
r = await c.post(
|
|
"/conversations",
|
|
json={"kind": "deep", "seed": "brief", "text": "dig", "title": "Тема"},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 201, r.text
|
|
body = r.json()
|
|
assert body["agent"] == "d"
|
|
assert [b["frontend"] for b in body["bindings"]] == ["markdown"]
|
|
rel = body["bindings"][0]["external_id"]
|
|
assert rel.endswith("_Тема.md") and rel.startswith("_logs/d/")
|
|
r = await c.post(
|
|
f"/conversations/{master}/bind",
|
|
json={"frontend": "markdown", "external_id": "x.md"},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 400
|
|
path = stack.vault / rel
|
|
text = await stack.wait_file(path, "ok:[сид: brief]")
|
|
post = frontmatter.loads(text)
|
|
assert post.metadata == {"agent": "d", "conversation_id": body["id"]}
|
|
assert post.content.startswith("### User:\n\n[сид: brief] deep «Тема», ")
|
|
assert post.content.rstrip().endswith("### User:")
|
|
|
|
|
|
async def test_markdown_chat_rejects_dispatcher(stack: Stack) -> None:
|
|
async with stack.client(stack.markdown) as c:
|
|
r = await c.post(
|
|
"/chat",
|
|
json={"filename": "x.md", "agent": "a", "content": "### User:\n\nhi\n"},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 400
|
|
assert "does not serve kind 'deep'" in r.json()["detail"]
|
|
|
|
|
|
async def test_spawn_tool_reports_bad_pair(stack: Stack) -> None:
|
|
master = await stack.world.conversations.create(
|
|
kind="master", agent="a", origin="test"
|
|
)
|
|
spawn = next(
|
|
t
|
|
for t in _tools(stack.world.conversations, master.external_id)
|
|
if t.name == "spawn"
|
|
)
|
|
out = await spawn.handler({"kind": "deep", "agent": "a"})
|
|
assert (
|
|
out.get("is_error")
|
|
and "does not serve kind 'deep'" in out["content"][0]["text"]
|
|
)
|
|
out = await spawn.handler({"kind": "deep", "seed": "brief", "text": "go"})
|
|
assert not out.get("is_error")
|
|
conv_id = out["content"][0]["text"].split()[-1]
|
|
path = await stack.file_of(conv_id)
|
|
assert path.exists()
|
|
|
|
|
|
async def test_anthropic_turns_become_one_deep_conversation(stack: Stack) -> None:
|
|
async with stack.client(stack.anthropic) as c:
|
|
r = await c.post(
|
|
"/v1/messages",
|
|
json={"model": "a", "messages": [{"role": "user", "content": "hi"}]},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 400
|
|
first = [{"role": "user", "content": "first question"}]
|
|
r = await c.post(
|
|
"/v1/messages", json={"model": "d", "messages": first}, headers=AUTH
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
reply = r.json()["content"][0]["text"]
|
|
assert reply == "ok:first question"
|
|
convs = await stack.world.conversations.find(kind="deep")
|
|
assert len(convs) == 1
|
|
conv = convs[0]
|
|
assert conv.title == "first question" and conv.frontend == "anthropic"
|
|
path = await stack.file_of(conv.external_id)
|
|
text = await stack.wait_file(path, "ok:first question")
|
|
assert frontmatter.loads(text).metadata == {
|
|
"agent": "d",
|
|
"conversation_id": conv.external_id,
|
|
}
|
|
|
|
second = [
|
|
*first,
|
|
{"role": "assistant", "content": reply},
|
|
{"role": "user", "content": "second"},
|
|
]
|
|
r = await c.post(
|
|
"/v1/messages",
|
|
json={"model": "d", "messages": second, "stream": True},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 200
|
|
assert "ok:second" in r.text
|
|
assert len(await stack.world.conversations.find(kind="deep")) == 1
|
|
assert len(ScriptedClient.instances) == 1
|
|
assert ScriptedClient.instances[0].prompts == ["first question", "second"]
|
|
text = await stack.wait_file(path, "ok:second")
|
|
assert text.count("### User:") == 3 and text.count("### Assistant:") == 2
|
|
bindings = await stack.world.conversations.bindings(conv)
|
|
assert [b.frontend for b in bindings if b.visible] == ["markdown", "anthropic"]
|
|
async with stack.world.db.session() as session:
|
|
stored = await load_messages(session, conversation_id=conv.id)
|
|
assert [m["role"] for m in stored] == ["user", "assistant", "user", "assistant"]
|
|
assert stored[-1]["content"] == [{"type": "text", "text": "ok:second"}]
|
|
|
|
r = await c.post(
|
|
"/v1/messages",
|
|
json={
|
|
"model": "a",
|
|
"messages": [
|
|
*second,
|
|
{"role": "assistant", "content": "ok:second"},
|
|
{"role": "user", "content": "x"},
|
|
],
|
|
},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 400 and "runs on 'd'" in r.json()["detail"]
|
|
|
|
|
|
async def test_markdown_edited_reply_reseeds_the_session(stack: Stack) -> None:
|
|
async with stack.client(stack.markdown) as c:
|
|
r = await c.post(
|
|
"/chat",
|
|
json={"filename": "e.md", "agent": "d", "content": "### User:\n\nanimal\n"},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
content = r.json()["new_content"]
|
|
assert "ok:animal" in content and len(ScriptedClient.instances) == 1
|
|
edited = content.replace("ok:animal", "penguin") + "\nwhich one?\n"
|
|
r = await c.post(
|
|
"/chat",
|
|
json={"filename": "e.md", "agent": "d", "content": edited},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
assert len(ScriptedClient.instances) == 2
|
|
first, second = ScriptedClient.instances
|
|
assert not first.connected and second.prompts == ["which one?"]
|
|
assert second.options.resume not in (None, first.session_id)
|
|
assert "penguin" in repr(vars(stack.world.store))
|
|
async with stack.world.db.session() as session:
|
|
stored = await load_messages(session, conversation_id=1)
|
|
assert stored[1]["content"] == [{"type": "text", "text": "penguin"}]
|
|
|
|
|
|
async def test_markdown_adopts_a_pre_sdk_conversation(stack: Stack) -> None:
|
|
world = stack.world
|
|
old = 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", old.id),
|
|
messages=[
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "ok:hello"},
|
|
],
|
|
)
|
|
content = "### User:\n\nhello\n\n---\n\n### Assistant:\n\nok:hello\n\n---\n\n### User:\n\nnext\n"
|
|
async with stack.client(stack.markdown) as c:
|
|
r = await c.post(
|
|
"/chat",
|
|
json={"filename": "old.md", "agent": "d", "content": content},
|
|
headers=AUTH,
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
assert f"conversation_id: {old.external_id}" in r.json()["new_content"]
|
|
assert len(await world.conversations.find(kind="deep")) == 1
|
|
bindings = await world.conversations.bindings(old)
|
|
assert [(b.frontend, b.external_id) for b in bindings] == [("markdown", "old.md")]
|
|
assert ScriptedClient.instances[-1].prompts == ["next"]
|
|
assert "ok:hello" in repr(vars(world.store))
|