961 lines
36 KiB
Python
961 lines
36 KiB
Python
import asyncio
|
|
import tempfile
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import pytest
|
|
from claude_agent_sdk import (
|
|
TaskNotificationMessage,
|
|
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.events.bus import EventBus
|
|
from beaver_gateway.conversations.service import (
|
|
Conversations,
|
|
ConversationTexts,
|
|
parse_at,
|
|
)
|
|
from beaver_gateway.conversations.tools import SAY_IN_USER_TURN, _tools
|
|
from beaver_gateway.app import AgentRegistry
|
|
from beaver_gateway.backends.sessions import SessionPool
|
|
from beaver_gateway.backends.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.contents: list[list[dict[str, Any]]] = []
|
|
self.session_id = options.resume or str(uuid.uuid4())
|
|
self.interrupted = False
|
|
self.connected = False
|
|
self.asked = asyncio.Event()
|
|
self.extra: list[Any] = []
|
|
ScriptedClient.instances.append(self)
|
|
|
|
async def connect(self) -> None:
|
|
self.connected = True
|
|
|
|
async def query(self, prompt) -> None:
|
|
if not isinstance(prompt, str):
|
|
messages = [m async for m in prompt]
|
|
content = messages[-1]["message"]["content"]
|
|
self.contents.append(content)
|
|
prompt = "\n".join(
|
|
b.get("text", "") for b in content if b.get("type") == "text"
|
|
)
|
|
self.prompts.append(prompt)
|
|
self.asked.set()
|
|
|
|
async def receive_messages(self):
|
|
"""One response per prompt, forever; ``extra`` is what the CLI sends
|
|
on its own between prompts (a subagent reporting back)."""
|
|
served = 0
|
|
while True:
|
|
while len(self.prompts) <= served:
|
|
self.asked.clear()
|
|
if self.extra:
|
|
for message in self.extra:
|
|
if isinstance(message, asyncio.Event):
|
|
await message.wait()
|
|
continue
|
|
yield message
|
|
self.extra = []
|
|
await self.asked.wait()
|
|
served += 1
|
|
async for message in self.receive_response():
|
|
yield message
|
|
|
|
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_say_is_refused_in_user_turns_and_works_in_inject_turns(
|
|
world: World,
|
|
) -> None:
|
|
conv = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
say = next(
|
|
t for t in _tools(world.conversations, conv.external_id) if t.name == "say"
|
|
)
|
|
said: list[str] = []
|
|
|
|
async def watch() -> None:
|
|
async for event in world.bus.stream(conversation_id=conv.external_id):
|
|
if event["type"] == "say":
|
|
said.append(str(event["text"]))
|
|
|
|
watcher = asyncio.create_task(watch())
|
|
ScriptedClient.hold = asyncio.Event()
|
|
await world.conversations.post(conv, "hi")
|
|
await asyncio.sleep(0.2)
|
|
assert world.conversations.turn_origin(conv) == "user"
|
|
out = await say.handler({"text": "dup"})
|
|
assert out.get("is_error") and out["content"][0]["text"] == SAY_IN_USER_TURN
|
|
ScriptedClient.hold.set()
|
|
await world.settle(conv, 1)
|
|
assert world.conversations.turn_origin(conv) is None
|
|
ScriptedClient.hold = asyncio.Event()
|
|
await world.conversations.inject(conv, "tick", urgency="urgent", origin="крон")
|
|
await asyncio.sleep(0.2)
|
|
assert world.conversations.turn_origin(conv) == "inject"
|
|
out = await say.handler({"text": "psst"})
|
|
assert not out.get("is_error")
|
|
ScriptedClient.hold.set()
|
|
await world.settle(conv, 2)
|
|
await asyncio.sleep(0.1)
|
|
watcher.cancel()
|
|
assert said == ["psst"]
|
|
|
|
|
|
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("[inject: крон")
|
|
assert "cut the previous turn" 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.conversations.service 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[injects")
|
|
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("[seed: 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("[seed: copy] branch, ")
|
|
assert "last 1 turns" 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 == "merge" 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": "interrupted",
|
|
"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 "cut by a gateway restart" in note.text
|
|
and "1 open tool calls" 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
|
|
)
|
|
assert parse_at("2026-09-01T10:00", ZoneInfo("Europe/Warsaw")) == 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"
|
|
|
|
|
|
async def test_wake_injects_start_a_turn_and_take_normals_along(world: World) -> None:
|
|
conv = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
await world.conversations.inject(conv, "digest", urgency="normal", origin="watch")
|
|
await asyncio.sleep(0.2)
|
|
assert await world.statuses(conv) == [("normal", "queued")]
|
|
await world.conversations.inject(
|
|
conv, "reminder", urgency="wake", origin="schedule"
|
|
)
|
|
await world.settle(conv, 2)
|
|
prompts = ScriptedClient.instances[0].prompts
|
|
assert len(prompts) == 1
|
|
assert prompts[0].startswith("[inject: schedule")
|
|
assert "reminder" in prompts[0]
|
|
assert "[inject: watch" in prompts[0]
|
|
assert "digest" in prompts[0]
|
|
assert await world.statuses(conv) == [("normal", "done"), ("wake", "done")]
|
|
|
|
|
|
async def test_wake_does_not_cut_a_running_turn(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.inject(conv, "later", urgency="wake", origin="schedule")
|
|
await asyncio.sleep(0.2)
|
|
assert await world.statuses(conv) == [("user", "running"), ("wake", "queued")]
|
|
assert not ScriptedClient.instances[0].interrupted
|
|
assert await world.conversations.busy(conv)
|
|
ScriptedClient.hold.set()
|
|
await world.settle(conv, 2)
|
|
assert await world.statuses(conv) == [("user", "done"), ("wake", "done")]
|
|
assert ScriptedClient.instances[0].prompts[1].endswith("later")
|
|
|
|
|
|
async def test_resolve_aliases_master_and_parent(world: World) -> None:
|
|
master = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
branch = await world.conversations.create(
|
|
kind="branch", agent="a", parent=master, origin="test"
|
|
)
|
|
resolve = world.conversations.resolve
|
|
assert (await resolve("master")).id == master.id
|
|
assert (await resolve(" master ")).id == master.id
|
|
assert (await resolve("parent", origin=branch)).id == master.id
|
|
assert await resolve("parent", origin=master) is None
|
|
assert await resolve("parent") is None
|
|
assert (await resolve(branch.external_id)).id == branch.id
|
|
assert await resolve("nope") is None
|
|
await world.conversations.set_status(master, "closed")
|
|
assert await resolve("master") is None
|
|
|
|
|
|
async def test_inject_into_a_closed_master_lands_in_the_open_one(world: World) -> None:
|
|
old = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
old = await world.conversations.set_status(old, "closed")
|
|
new = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
item = await world.conversations.inject(old, "late", urgency="normal", origin="x")
|
|
assert item.conversation_id == new.id
|
|
assert await world.statuses(old) == []
|
|
assert await world.statuses(new) == [("normal", "queued")]
|
|
|
|
|
|
def _self_started(session_id: str, text: str) -> list[Any]:
|
|
"""What the CLI sends by itself once a background subagent reports back."""
|
|
return [
|
|
AssistantMessage(
|
|
content=[ToolUseBlock(id="tu_dig", name="Grep", input={"pattern": "au"})],
|
|
model="m",
|
|
parent_tool_use_id="tu_bg",
|
|
),
|
|
TaskNotificationMessage(
|
|
subtype="task_notification",
|
|
data={},
|
|
task_id="t1",
|
|
status="completed",
|
|
output_file="",
|
|
summary='Agent "dig" finished',
|
|
uuid="n1",
|
|
session_id=session_id,
|
|
),
|
|
AssistantMessage(content=[TextBlock(text=text)], model="m"),
|
|
ResultMessage(
|
|
subtype="success",
|
|
duration_ms=1,
|
|
duration_api_ms=1,
|
|
is_error=False,
|
|
num_turns=1,
|
|
session_id=session_id,
|
|
stop_reason="end_turn",
|
|
total_cost_usd=0.0,
|
|
usage={"input_tokens": 1, "output_tokens": 1},
|
|
origin={"kind": "task-notification"},
|
|
),
|
|
]
|
|
|
|
|
|
async def test_subagent_report_between_turns_is_its_own_turn(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())
|
|
await world.conversations.post(conv, "go")
|
|
await world.settle(conv, 1)
|
|
client = ScriptedClient.instances[0]
|
|
client.extra = _self_started(client.session_id, "the agent found gold")
|
|
client.asked.set()
|
|
await asyncio.sleep(0.3)
|
|
task.cancel()
|
|
replies = [e for e in seen if e["type"] == "reply"]
|
|
assert [r["text"] for r in replies] == ["ok:go", "the agent found gold"]
|
|
assert replies[1]["item_origin"] == "task"
|
|
starts = [e for e in seen if e["type"] == "turn.start"]
|
|
assert [s["origin"] for s in starts] == ["user", "task"]
|
|
assert starts[1]["text"] == 'Agent "dig" finished (completed)'
|
|
assert len(client.prompts) == 1
|
|
first_end = next(i for i, e in enumerate(seen) if e["type"] == "turn.end")
|
|
later = [e for e in seen[first_end + 1 :] if e["type"] in ("stream", "tool")]
|
|
assert later and all(e["turn_id"] == starts[0]["turn_id"] for e in later[:1])
|
|
assert all(e["turn_id"] == starts[1]["turn_id"] for e in later[1:])
|
|
assert (await world.conversations.get(conv.external_id)).running_turn is None
|
|
|
|
|
|
async def test_subagent_result_inside_a_turn_does_not_end_it(world: World) -> None:
|
|
conv = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
await world.conversations.post(conv, "one")
|
|
await world.settle(conv, 1)
|
|
client = ScriptedClient.instances[0]
|
|
client.extra = _self_started(client.session_id, "late report")
|
|
await world.conversations.post(conv, "two")
|
|
await world.settle(conv, 2)
|
|
assert client.prompts[-1].startswith("two")
|
|
rows = await world.conversations.queue.recent(conv.id)
|
|
assert [r.status for r in rows] == ["done", "done"]
|