Files
beaver-gateway/tests/test_conversations.py
T

746 lines
27 KiB
Python

import asyncio
import tempfile
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
import pytest
from claude_agent_sdk import (
AssistantMessage,
InMemorySessionStore,
ResultMessage,
StreamEvent,
TextBlock,
ToolResultBlock,
ToolUseBlock,
UserMessage,
project_key_for_directory,
)
from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions
from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend
from beaver_gateway.core.bus import EventBus
from beaver_gateway.core.conversations import Conversations, ConversationTexts, parse_at
from beaver_gateway.core.registry import AgentRegistry
from beaver_gateway.core.sessions import SessionPool
from beaver_gateway.core.transcript import (
build_entries,
close_open_tool_uses,
open_tool_uses,
render_messages,
strip_tool_entries,
window_entries,
)
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.storage import Database
from beaver_gateway.storage.models import (
Conversation,
ConversationBinding,
InjectQueueItem,
)
class StubFrontend(Frontend):
def __init__(
self,
name: str,
kinds: tuple[str, ...],
agents: dict[str, str] | None = None,
*,
home: bool = False,
) -> None:
self.name = name
self.kinds = kinds
self.agents = agents or {}
self.home = home
self.materialized: list[str] = []
self.conversations = None
def configure(self, runtime) -> None:
pass
async def serve(self) -> None:
pass
def agent_for(self, kind: str) -> str | None:
return self.agents.get(kind)
async def materialize(self, conv: Conversation) -> ConversationBinding | None:
if not self.home:
return None
self.materialized.append(conv.external_id)
return await self.conversations.bind(
conv, frontend=self.name, external_id=f"{self.name}:{conv.external_id}"
)
class ScriptedClient:
instances: list["ScriptedClient"] = []
hold: asyncio.Event | None = None
def __init__(self, options: Any) -> None:
self.options = options
self.prompts: list[str] = []
self.session_id = options.resume or str(uuid.uuid4())
self.interrupted = False
self.connected = False
ScriptedClient.instances.append(self)
async def connect(self) -> None:
self.connected = True
async def query(self, prompt: str) -> None:
self.prompts.append(prompt)
async def receive_response(self):
prompt = self.prompts[-1]
yield StreamEvent(
uuid="u",
session_id=self.session_id,
event={"type": "message_start", "message": {}},
)
yield StreamEvent(
uuid="u",
session_id=self.session_id,
event={
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
)
yield StreamEvent(
uuid="u",
session_id=self.session_id,
event={
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": f"ok:{prompt}"},
},
)
yield StreamEvent(
uuid="u",
session_id=self.session_id,
event={"type": "content_block_stop", "index": 0},
)
yield AssistantMessage(content=[TextBlock(text=f"ok:{prompt}")], model="m")
hold = ScriptedClient.hold
if hold is not None and not self.interrupted:
await hold.wait()
cut = self.interrupted
self.interrupted = False
yield ResultMessage(
subtype="error_during_execution" if cut else "success",
duration_ms=1,
duration_api_ms=1,
is_error=cut,
num_turns=1,
session_id=self.session_id,
stop_reason="end_turn",
total_cost_usd=0.0,
usage={"input_tokens": 1, "output_tokens": 1},
)
async def interrupt(self) -> None:
self.interrupted = True
if ScriptedClient.hold is not None:
ScriptedClient.hold.set()
async def disconnect(self) -> None:
self.connected = False
class World:
def __init__(self, root: Path) -> None:
ScriptedClient.instances.clear()
ScriptedClient.hold = None
self.root = root
self.db = Database(f"sqlite:///{root / 'w.db'}")
self.store = InMemorySessionStore()
self.agent = ClaudeAgent(
name="a",
model="m",
system_prompt="hi",
cwd=root,
kinds=("master", "branch", "job", "fork"),
gateway_tools=("say",),
options=ClaudeOptions(effort="low"),
)
self.deep_agent = ClaudeAgent(
name="d", model="m", system_prompt="deep", cwd=root
)
self.markdown = StubFrontend("markdown", ("deep",), {"deep": "d"}, home=True)
self.api = StubFrontend("api", ("master", "branch", "deep", "job"))
self.pool = SessionPool(rss_limit=1 << 40, max_live=100)
self.backend, self.deep_backend = (
ClaudeSdkBackend(
agent=agent,
mcp_internal_urls={},
session_store=self.store,
client_factory=ScriptedClient,
work_dir=root / "work",
pool=self.pool,
)
for agent in (self.agent, self.deep_agent)
)
self.bus = EventBus()
self.conversations = Conversations(
db=self.db,
agents=AgentRegistry([self.agent, self.deep_agent]),
backends={"a": self.backend, "d": self.deep_backend},
bus=self.bus,
pool=self.pool,
store=self.store,
texts=ConversationTexts(),
frontends=[self.api, self.markdown],
idle_interval=3600,
)
self.markdown.conversations = self.conversations
self.api.conversations = self.conversations
async def setup(self) -> "World":
await self.db.create_all()
return self
def key(self, session_id: str) -> dict[str, str]:
return {
"project_key": project_key_for_directory(str(self.root)),
"session_id": session_id,
}
async def statuses(self, conv) -> list[tuple[str, str]]:
rows = await self.conversations.queue.recent(conv.id)
return [(r.priority, r.status) for r in reversed(rows)]
async def settle(self, conv, expected: int, timeout: float = 5.0) -> None:
deadline = asyncio.get_running_loop().time() + timeout
while asyncio.get_running_loop().time() < deadline:
rows = await self.conversations.queue.recent(conv.id)
if sum(1 for r in rows if r.status == "done") >= expected:
return
await asyncio.sleep(0.02)
msg = f"queue did not settle: {await self.statuses(conv)}"
raise AssertionError(msg)
@pytest.fixture
async def world() -> World:
root = Path(tempfile.mkdtemp(prefix="beaver-conv-"))
w = await World(root).setup()
yield w
await w.conversations.stop()
await w.pool.close_all()
await w.db.dispose()
async def test_two_messages_run_one_at_a_time_in_order(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
ScriptedClient.hold = asyncio.Event()
await world.conversations.post(conv, "first")
await asyncio.sleep(0.2)
await world.conversations.post(conv, "second")
await asyncio.sleep(0.2)
assert await world.statuses(conv) == [("user", "running"), ("user", "queued")]
live = world.pool.get(conv.external_id)
assert live is not None and live.busy
assert (await world.conversations.get(conv.external_id)).running_turn is not None
ScriptedClient.hold.set()
await world.settle(conv, 2)
assert len(ScriptedClient.instances) == 1
assert ScriptedClient.instances[0].prompts == ["first", "second"]
row = await world.conversations.get(conv.external_id)
assert row.running_turn is None
assert row.session_id == ScriptedClient.instances[0].session_id
async def test_urgent_interrupts_and_goes_first(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
ScriptedClient.hold = asyncio.Event()
await world.conversations.post(conv, "first")
await asyncio.sleep(0.2)
await world.conversations.post(conv, "second")
await world.conversations.inject(conv, "ALERT", urgency="urgent", origin="крон")
await world.settle(conv, 2)
client = ScriptedClient.instances[0]
assert await world.statuses(conv) == [
("user", "interrupted"),
("user", "done"),
("urgent", "done"),
]
assert await world.statuses(conv) == [
("user", "interrupted"),
("user", "done"),
("urgent", "done"),
]
assert client.prompts[0] == "first"
assert client.prompts[1].startswith("[инжект: крон")
assert "прервал предыдущий тёрн" in client.prompts[1]
assert client.prompts[1].endswith("ALERT")
assert client.prompts[2] == "second"
async def test_inject_header_is_configurable(world: World) -> None:
from beaver_gateway.core.conversations import ConversationTexts
world.conversations._texts = ConversationTexts( # noqa: SLF001
inject_header=lambda ctx: (
f"[от {ctx.origin}, {ctx.priority}, cut={ctx.interrupted_turn}]"
)
)
conv = await world.conversations.create(kind="master", agent="a", origin="test")
await world.conversations.inject(conv, "hi", urgency="urgent", origin="panel")
await world.settle(conv, 1)
assert ScriptedClient.instances[0].prompts[0] == "[от panel, urgent, cut=False]\nhi"
async def test_normal_injects_ride_with_the_next_user_message(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
await world.conversations.inject(
conv, "vault changed", urgency="normal", origin="watch"
)
await asyncio.sleep(0.2)
assert await world.statuses(conv) == [("normal", "queued")]
await world.conversations.post(conv, "hello")
await world.settle(conv, 2)
prompts = ScriptedClient.instances[0].prompts
assert len(prompts) == 1
assert prompts[0].startswith("hello\n\n[инжекты")
assert "- [watch] vault changed" in prompts[0]
async def test_inject_turn_reply_is_not_routed(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
seen: list[dict[str, Any]] = []
async def collect() -> None:
async for event in world.bus.stream(conversation_id=conv.external_id):
seen.append(event)
task = asyncio.create_task(collect())
world.conversations._normal_window = 0.05
await world.conversations.inject(conv, "tick", urgency="normal", origin="крон")
await world.settle(conv, 1)
await asyncio.sleep(0.05)
task.cancel()
types = [e["type"] for e in seen]
assert "turn.start" in types and "turn.end" in types
assert "reply" not in types
assert all(e.get("origin") == "inject" for e in seen if e["type"] == "turn.start")
async def test_spawn_seeds_first_user_message(world: World) -> None:
conv = await world.conversations.spawn(
kind="branch", agent="a", seed="brief", text="do X", title="t"
)
await world.settle(conv, 1)
prompt = ScriptedClient.instances[0].prompts[0]
assert prompt.startswith("[сид: brief] branch «t», ")
assert prompt.endswith("\n\ndo X")
assert ScriptedClient.instances[0].options.system_prompt == "hi"
async def test_fork_leaves_original_untouched(world: World) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": "one"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": "x"}],
},
{"role": "assistant", "content": "done one"},
{"role": "user", "content": "two"},
{"role": "assistant", "content": "done two"},
{"role": "user", "content": "three"},
{"role": "assistant", "content": "done three"},
]
entries = build_entries(history, session_id=sid, cwd=str(world.root), model="m")
await world.store.append(world.key(sid), entries)
conv = await world.conversations.create(
kind="branch", agent="a", origin="test", session_id=sid
)
before = [dict(e) for e in await world.store.load(world.key(sid))]
result = await world.conversations.fork(
conv, "summarize", window=2, strip_tools=True
)
after = await world.store.load(world.key(sid))
assert after == before
assert result.text == "ok:summarize"
assert result.conversation.status == "closed"
assert result.conversation.parent_id == conv.id
fork_entries = await world.store.load(world.key(result.conversation.session_id))
assert not {e["uuid"] for e in fork_entries} & {e["uuid"] for e in before}
texts = [e["message"]["content"] for e in fork_entries]
assert texts[0] == "two"
assert len(fork_entries) == 4
assert fork_entries[0]["parentUuid"] is None
assert all(
b["type"] == "text"
for e in fork_entries[1:]
for b in e["message"]["content"]
if isinstance(e["message"]["content"], list)
)
client = ScriptedClient.instances[0]
assert client.options.resume == result.conversation.session_id
assert client.options.mcp_servers == {}
assert world.pool.get(result.conversation.external_id) is None
async def test_copy_seed_forks_parent_with_window(world: World) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": f"q{i}"}
if i % 2 == 0
else {"role": "assistant", "content": f"a{i}"}
for i in range(8)
]
await world.store.append(
world.key(sid),
build_entries(history, session_id=sid, cwd=str(world.root), model="m"),
)
parent = await world.conversations.create(
kind="master", agent="a", origin="test", session_id=sid
)
child = await world.conversations.spawn(
kind="branch", agent="a", seed="copy", parent=parent, window=1
)
assert child.session_id and child.session_id != sid
copied = await world.store.load(world.key(child.session_id))
assert [e["message"]["content"] for e in copied] == [
"q6",
[{"type": "text", "text": "a7"}],
]
assert (await world.conversations.get(child.external_id)).flags["seed"] == "copy"
await world.conversations.post(child, "go")
await world.settle(child, 1)
assert ScriptedClient.instances[0].options.resume == child.session_id
prompt = ScriptedClient.instances[0].prompts[0]
assert prompt.startswith("[сид: copy] branch, ")
assert "последние 1 тёрнов" in prompt and prompt.endswith("\n\ngo")
assert (await world.conversations.get(child.external_id)).flags["seed"] is None
async def test_merge_injects_summary_into_parent(world: World) -> None:
sid = str(uuid.uuid4())
await world.store.append(
world.key(sid),
build_entries(
[{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}],
session_id=sid,
cwd=str(world.root),
model="m",
),
)
master = await world.conversations.create(kind="master", agent="a", origin="test")
branch = await world.conversations.create(
kind="branch", agent="a", origin="test", parent=master, session_id=sid
)
result = await world.conversations.merge(branch)
assert result.text.startswith("ok:")
assert (await world.conversations.get(branch.external_id)).status == "merged"
assert await world.statuses(master) == [("normal", "queued")]
item = (await world.conversations.queue.recent(master.id))[0]
assert item.origin == "слив" and item.text == result.text
async def test_recover_closes_open_tool_use_and_injects_interrupted(
world: World,
) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": "go"},
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "t9",
"name": "Bash",
"input": {"command": "sleep"},
}
],
},
]
await world.store.append(
world.key(sid),
build_entries(history, session_id=sid, cwd=str(world.root), model="m"),
)
conv = await world.conversations.create(
kind="master", agent="a", origin="test", session_id=sid
)
async with world.db.session() as session:
row = await session.get(type(conv), conv.id)
row.running_turn = "turn_dead"
session.add(row)
session.add(
InjectQueueItem(
conversation_id=conv.id,
priority="user",
origin="user",
text="go",
status="running",
turn_id="turn_dead",
)
)
await session.commit()
cut = await world.conversations.recover()
assert [c.external_id for c in cut] == [conv.external_id]
entries = await world.store.load(world.key(sid))
assert not open_tool_uses(entries)
tail = entries[-1]
assert tail["type"] == "user"
assert tail["message"]["content"][0] == {
"type": "tool_result",
"tool_use_id": "t9",
"content": "прервано",
"is_error": True,
}
assert tail["parentUuid"] == entries[-2]["uuid"]
assert (await world.conversations.get(conv.external_id)).running_turn is None
assert await world.statuses(conv) == [("user", "interrupted"), ("normal", "queued")]
note = (await world.conversations.queue.recent(conv.id))[0]
assert (
"turn_dead" in note.text
and "оборван" in note.text
and "1 незакрытых" in note.text
)
await asyncio.sleep(0.2)
assert ScriptedClient.instances == []
async def test_graceful_stop_keeps_running_turn_for_recover(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
ScriptedClient.hold = asyncio.Event()
await world.conversations.post(conv, "long")
await asyncio.sleep(0.2)
assert (await world.conversations.get(conv.external_id)).running_turn is not None
await world.conversations.stop()
row = await world.conversations.get(conv.external_id)
assert row.running_turn is not None
assert row.session_id == ScriptedClient.instances[0].session_id
cut = await world.conversations.recover()
assert [c.external_id for c in cut] == [conv.external_id]
assert (await world.conversations.get(conv.external_id)).running_turn is None
assert await world.statuses(conv) == [("user", "interrupted"), ("normal", "queued")]
ScriptedClient.hold = None
async def test_read_and_bindings(world: World) -> None:
sid = str(uuid.uuid4())
history = [
{"role": "user", "content": "q1"},
{
"role": "assistant",
"content": [
{"type": "text", "text": "a1"},
{"type": "tool_use", "id": "t", "name": "Read", "input": {}},
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t", "content": "x"}],
},
{"role": "assistant", "content": "a1b"},
{"role": "user", "content": "q2"},
{"role": "assistant", "content": "a2"},
]
await world.store.append(
world.key(sid),
build_entries(history, session_id=sid, cwd=str(world.root), model="m"),
)
conv = await world.conversations.create(
kind="deep", agent="d", origin="test", session_id=sid
)
assert (
await world.conversations.read(conv)
== "user:\nq1\n\nassistant:\na1\n\na1b\n(tools: Read)\n\nuser:\nq2\n\nassistant:\na2"
)
assert (
await world.conversations.read(conv, window=1) == "user:\nq2\n\nassistant:\na2"
)
await world.conversations.bind(conv, frontend="markdown", external_id="a.md")
await world.conversations.bind(conv, frontend="markdown", external_id="b.md")
bindings = await world.conversations.bindings(conv)
assert [(b.external_id, b.visible) for b in bindings] == [
("a.md", False),
("b.md", True),
]
found = await world.conversations.find_bound(
frontend="markdown", external_id="b.md"
)
assert found is not None and found.id == conv.id
assert (
await world.conversations.find_bound(frontend="markdown", external_id="a.md")
is None
)
async def test_schedule_without_scheduler_and_parse_at(world: World) -> None:
conv = await world.conversations.create(kind="master", agent="a", origin="test")
with pytest.raises(RuntimeError, match="no scheduler"):
await world.conversations.schedule(conv, "+15m", "push X")
assert await world.conversations.schedules(conv) == []
assert parse_at("2026-09-01T10:00:00+02:00") == datetime(
2026, 9, 1, 8, 0, tzinfo=UTC
)
delta = (parse_at("+15m") - datetime.now(UTC)).total_seconds()
assert 14 * 60 < delta <= 15 * 60
with pytest.raises(ValueError, match="Invalid isoformat"):
parse_at("tomorrow")
def test_transcript_helpers() -> None:
history = [
{"role": "user", "content": "q1"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "hm", "signature": "s"},
{"type": "tool_use", "id": "t", "name": "Read", "input": {}},
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t", "content": "x"}],
},
{"role": "assistant", "content": "a1"},
{"role": "user", "content": "q2"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "t2", "name": "Bash", "input": {}}],
},
]
entries = build_entries(history, session_id="s", cwd="/v", model="m")
assert [b["id"] for _, b in open_tool_uses(entries)] == ["t2"]
fixes = close_open_tool_uses(entries)
assert len(fixes) == 1 and fixes[0]["sessionId"] == "s" and fixes[0]["cwd"] == "/v"
assert not open_tool_uses([*entries, *fixes])
stripped = strip_tool_entries(entries)
assert [e["message"]["content"] for e in stripped] == [
"q1",
[{"type": "text", "text": "a1"}],
"q2",
]
assert (
stripped[0]["parentUuid"] is None
and stripped[1]["parentUuid"] == stripped[0]["uuid"]
)
windowed = window_entries(entries, window=1)
assert windowed[0]["message"]["content"] == "q2" and len(windowed) == 2
assert (
render_messages(
[{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}]
)
== "user:\nq\n\nassistant:\na"
)
async def test_agent_kinds_gate_create(world: World) -> None:
with pytest.raises(ValueError, match="does not serve kind 'deep'"):
await world.conversations.create(kind="deep", agent="a", origin="test")
with pytest.raises(ValueError, match="does not serve kind 'master'"):
await world.conversations.create(kind="master", agent="d", origin="test")
with pytest.raises(ValueError, match="'d' does not serve kind 'master'"):
await world.conversations.spawn(kind="master", seed="clean", agent="d")
conv = await world.conversations.create(kind="deep", agent="d", origin="test")
assert conv.agent_name == "d"
async def test_frontend_kinds_gate_bind(world: World) -> None:
master = await world.conversations.create(kind="master", agent="a", origin="test")
with pytest.raises(ValueError, match="'markdown' does not show kind 'master'"):
await world.conversations.bind(master, frontend="markdown", external_id="m.md")
with pytest.raises(ValueError, match="unknown frontend 'telegram'"):
await world.conversations.bind(master, frontend="telegram", external_id="1")
await world.conversations.bind(master, frontend="api", external_id="x")
async def test_spawn_defaults_agent_and_materializes(world: World) -> None:
master = await world.conversations.create(kind="master", agent="a", origin="test")
deep = await world.conversations.spawn(
kind="deep", seed="brief", text="dig", title="t", parent=master, origin="mcp"
)
assert deep.agent_name == "d"
assert world.markdown.materialized == [deep.external_id]
bindings = await world.conversations.bindings(deep)
assert [(b.frontend, b.external_id) for b in bindings] == [
("markdown", f"markdown:{deep.external_id}")
]
branch = await world.conversations.spawn(kind="branch", seed="clean", parent=master)
assert branch.agent_name == "a"
assert world.markdown.materialized == [deep.external_id]
with pytest.raises(ValueError, match="no default agent for kind 'job'"):
await world.conversations.spawn(kind="job", seed="clean")
await world.settle(deep, 1)
def test_agent_kinds_follow_prompts(tmp_path: Path) -> None:
from beaver_gateway.agents.claude import Prompts
plain = ClaudeAgent(name="p", model="m", system_prompt="hi", cwd=tmp_path)
assert plain.kinds == ("deep",)
dispatcher = ClaudeAgent(
name="x", model="m", cwd=tmp_path, prompts=Prompts(master=("a.md",), fork=())
)
assert dispatcher.kinds == ("master", "fork")
assert dispatcher.prompt_for("master") == ("a.md",)
assert dispatcher.prompt_for("deep") is None
with pytest.raises(ValueError, match=r"serves \['deep'\] without a prompt"):
ClaudeAgent(
name="y",
model="m",
cwd=tmp_path,
kinds=("master", "deep"),
prompts=Prompts(master=()),
)
async def test_observer_publishes_tool_results_for_the_panel(world: World) -> None:
seen: list[dict[str, Any]] = []
conv = await world.conversations.create(kind="master", agent="a", origin="test")
async def collect() -> None:
async for event in world.bus.stream(conversation_id=conv.external_id):
seen.append(event)
task = asyncio.create_task(collect())
await asyncio.sleep(0)
runner = world.conversations._runner(conv.id)
runner.turn_id = "t1"
observe = world.conversations._observer(conv, runner, "t1", "user")
observe(
AssistantMessage(
content=[ToolUseBlock(id="toolu_1", name="Bash", input={"command": "ls"})],
model="m",
parent_tool_use_id="toolu_0",
)
)
observe(
UserMessage(
content=[
ToolResultBlock(tool_use_id="toolu_1", content="a\nb", is_error=False)
],
parent_tool_use_id="toolu_0",
)
)
observe(UserMessage(content="plain text, no tool result"))
await asyncio.sleep(0.05)
task.cancel()
assert [e["type"] for e in seen] == ["tool", "tool.result"]
tool, result = seen
assert tool["tool_use_id"] == "toolu_1" and tool["parent_tool_use_id"] == "toolu_0"
snapshot = runner.snapshot()
assert snapshot is not None and snapshot["tools"][0]["content"] == "a\nb"
assert result["tool_use_id"] == "toolu_1"
assert result["parent_tool_use_id"] == "toolu_0"
assert result["turn_id"] == "t1" and result["is_error"] is False
assert result["content"] == "a\nb"