420 lines
15 KiB
Python
420 lines
15 KiB
Python
import asyncio
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from claude_agent_sdk import (
|
|
AssistantMessage,
|
|
ResultMessage,
|
|
TextBlock,
|
|
ToolResultBlock,
|
|
ToolUseBlock,
|
|
UserMessage,
|
|
project_key_for_directory,
|
|
)
|
|
from test_conversations import ScriptedClient, World, world
|
|
|
|
from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions
|
|
from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend
|
|
from beaver_gateway.conversations.service import ConversationTexts
|
|
from beaver_gateway.conversations.distill import (
|
|
Distiller,
|
|
DistillContext,
|
|
LineCap,
|
|
check_digest,
|
|
trim_summary,
|
|
)
|
|
from beaver_gateway.conversations.tools import _tools, build_tool_server
|
|
from beaver_gateway.app import AgentRegistry
|
|
from beaver_gateway.jobs.scheduler import Job, JobRun, Scheduler
|
|
from beaver_gateway.backends.transcript import build_entries
|
|
from beaver_gateway.storage.models import Conversation
|
|
|
|
__all__ = ["world"]
|
|
|
|
DIGEST = """---
|
|
type: digest
|
|
source: "[[{chat}]]"
|
|
date: 2026-08-29
|
|
---
|
|
# тема
|
|
|
|
## контекст
|
|
зачем открывали
|
|
|
|
## решили
|
|
- одно
|
|
"""
|
|
|
|
|
|
class DistillerClient(ScriptedClient):
|
|
"""The distiller: writes a digest file when told to, answers in N lines."""
|
|
|
|
digest_dir: Path | None = None
|
|
lines = 3
|
|
write = True
|
|
frontmatter = DIGEST
|
|
|
|
async def receive_response(self):
|
|
prompt = self.prompts[-1]
|
|
written: list[ToolUseBlock] = []
|
|
if self.write and self.digest_dir is not None and "write no file" not in prompt:
|
|
chat = prompt.split("«", 1)[1].split("»", 1)[0] if "«" in prompt else "чат"
|
|
path = self.digest_dir / "2026-08-29 - тема.md"
|
|
path.write_text(self.frontmatter.format(chat=chat), encoding="utf-8")
|
|
written.append(
|
|
ToolUseBlock(id="w1", name="Write", input={"file_path": str(path)})
|
|
)
|
|
for block in written:
|
|
yield AssistantMessage(content=[block], model="m")
|
|
yield UserMessage(content=[ToolResultBlock(tool_use_id=block.id)])
|
|
text = "\n".join(f"строка {i} про {prompt[:20]!r}" for i in range(self.lines))
|
|
yield AssistantMessage(content=[TextBlock(text=text)], 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.0,
|
|
usage={"input_tokens": 1, "output_tokens": 1},
|
|
)
|
|
|
|
|
|
class ClosingClient(ScriptedClient):
|
|
"""A deep chat that calls ``close_chat`` inside its reply."""
|
|
|
|
conversations: Any = None
|
|
key = ""
|
|
|
|
async def receive_response(self):
|
|
prompt = self.prompts[-1]
|
|
if "обсудили" in prompt:
|
|
assert "gateway" in self.options.mcp_servers
|
|
tools = _tools(ClosingClient.conversations, ClosingClient.key)
|
|
close = next(t for t in tools if t.name == "close_chat")
|
|
result = await close.handler({})
|
|
assert "closes after this reply" in result["content"][0]["text"]
|
|
yield AssistantMessage(
|
|
content=[
|
|
ToolUseBlock(id="c1", name="mcp__gateway__close_chat", input={})
|
|
],
|
|
model="m",
|
|
)
|
|
yield AssistantMessage(content=[TextBlock(text=f"ok:{prompt}")], 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.0,
|
|
usage={"input_tokens": 1, "output_tokens": 1},
|
|
)
|
|
|
|
|
|
class CapClient(ScriptedClient):
|
|
"""A job that rewrites the capped file with ``lines`` lines."""
|
|
|
|
target: Path | None = None
|
|
lines = 70
|
|
|
|
async def receive_response(self):
|
|
assert self.target is not None
|
|
self.target.write_text(
|
|
"\n".join(f"- строка {i}" for i in range(self.lines)), encoding="utf-8"
|
|
)
|
|
CapClient.lines = 10
|
|
async for event in super().receive_response():
|
|
yield event
|
|
|
|
|
|
def distiller(world: World, client: type[ScriptedClient]) -> Distiller:
|
|
vault = world.root / "vault"
|
|
digests = vault / "мета" / "бобер" / "выжимки"
|
|
digests.mkdir(parents=True)
|
|
agent = ClaudeAgent(
|
|
name="x",
|
|
model="m",
|
|
system_prompt="distill",
|
|
cwd=vault,
|
|
kinds=("fork", "job"),
|
|
options=ClaudeOptions(
|
|
effort="medium", tools=("Read", "Write"), include_partial_messages=False
|
|
),
|
|
)
|
|
backend = ClaudeSdkBackend(
|
|
agent=agent,
|
|
mcp_internal_urls={},
|
|
session_store=world.store,
|
|
client_factory=client,
|
|
work_dir=world.root / "work",
|
|
pool=world.pool,
|
|
)
|
|
world.conversations._agents = AgentRegistry( # noqa: SLF001
|
|
[world.agent, world.deep_agent, agent]
|
|
)
|
|
world.conversations._backends["x"] = backend # noqa: SLF001
|
|
config = Distiller(agent="x", dir=digests, index=digests.parent / "индекс.md")
|
|
world.conversations._distiller = config # noqa: SLF001
|
|
DistillerClient.digest_dir = digests
|
|
DistillerClient.lines = 3
|
|
DistillerClient.write = True
|
|
DistillerClient.frontmatter = DIGEST
|
|
return config
|
|
|
|
|
|
async def deep_chat(world: World, name: str = "2026-08-20 - тема чата") -> Conversation:
|
|
sid = str(uuid.uuid4())
|
|
await world.store.append(
|
|
world.key(sid),
|
|
build_entries(
|
|
[
|
|
{"role": "user", "content": "q"},
|
|
{
|
|
"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": "a"},
|
|
],
|
|
session_id=sid,
|
|
cwd=str(world.root),
|
|
model="m",
|
|
),
|
|
)
|
|
conv = await world.conversations.create(
|
|
kind="deep", agent="d", origin="markdown", session_id=sid
|
|
)
|
|
await world.conversations.bind(
|
|
conv, frontend="markdown", external_id=f"2026-08/{name}.md"
|
|
)
|
|
return conv
|
|
|
|
|
|
async def test_distill_writes_the_digest_indexes_it_and_merges_short(
|
|
world: World,
|
|
) -> None:
|
|
config = distiller(world, DistillerClient)
|
|
prompts: list[DistillContext] = []
|
|
|
|
def distill_prompt(ctx: DistillContext) -> str:
|
|
prompts.append(ctx)
|
|
return f"Чат «{ctx.chat_name}» закрыт ({ctx.reason})."
|
|
|
|
world.conversations._texts = ConversationTexts(distill=distill_prompt) # noqa: SLF001
|
|
master = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
chat = await deep_chat(world)
|
|
|
|
result = await world.conversations.distill(chat, reason="api")
|
|
|
|
assert prompts[0].chat_name == "2026-08-20 - тема чата"
|
|
assert prompts[0].memory is True
|
|
assert result.digest is not None and result.error is None
|
|
assert result.digest.path == config.dir / "2026-08-29 - тема.md"
|
|
assert result.digest.source == "[[2026-08-20 - тема чата]]"
|
|
index = config.index.read_text(encoding="utf-8")
|
|
assert index.startswith("# index")
|
|
assert "- 2026-08-29 [[2026-08-20 - тема чата]] → [[2026-08-29 - тема]]" in index
|
|
assert result.text.count("\n") == 2 and not result.trimmed
|
|
closed = await world.conversations.get(chat.external_id)
|
|
assert closed.status == "closed"
|
|
assert closed.flags["digest"] == str(result.digest.path)
|
|
assert closed.flags["closed_reason"] == "api"
|
|
fork = await world.conversations.get(result.fork.external_id)
|
|
assert fork.kind == "fork" and fork.agent_name == "x" and fork.status == "closed"
|
|
items = await world.conversations.queue.recent(master.id)
|
|
assert items[0].origin == "digest"
|
|
assert items[0].text.startswith(
|
|
"Deep chat [[2026-08-20 - тема чата]] closed, digest [[2026-08-29 - тема]]."
|
|
)
|
|
assert items[0].text.endswith(result.text)
|
|
forked = ScriptedClient.instances[-1]
|
|
assert forked.options.resume is not None
|
|
vault_key = {
|
|
"project_key": project_key_for_directory(str(config.dir.parents[2])),
|
|
"session_id": forked.options.resume,
|
|
}
|
|
entries = await world.store.load(vault_key)
|
|
assert entries
|
|
assert all(
|
|
not isinstance(e.get("message", {}).get("content"), list)
|
|
or all(b.get("type") == "text" for b in e["message"]["content"])
|
|
for e in entries
|
|
)
|
|
|
|
|
|
async def test_memory_off_merges_without_a_file(world: World) -> None:
|
|
config = distiller(world, DistillerClient)
|
|
master = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
chat = await deep_chat(world)
|
|
await world.conversations.set_flags(chat, {"memory": False})
|
|
|
|
result = await world.conversations.distill(chat, reason="api")
|
|
|
|
assert result.digest is None and result.error is None
|
|
assert list(config.dir.iterdir()) == []
|
|
assert not config.index.exists()
|
|
assert (await world.conversations.get(chat.external_id)).status == "closed"
|
|
items = await world.conversations.queue.recent(master.id)
|
|
assert len(items) == 1 and items[0].text.endswith(result.text)
|
|
assert ", digest" not in items[0].text
|
|
assert "write no file" in ScriptedClient.instances[-1].prompts[0]
|
|
|
|
|
|
async def test_bad_frontmatter_and_long_merge_are_reported(world: World) -> None:
|
|
distiller(world, DistillerClient)
|
|
DistillerClient.frontmatter = "---\ntype: заметка\n---\n# тема\n"
|
|
DistillerClient.lines = 9
|
|
await world.conversations.create(kind="master", agent="a", origin="test")
|
|
chat = await deep_chat(world)
|
|
|
|
result = await world.conversations.distill(chat, reason="idle 2d")
|
|
|
|
assert result.digest is None
|
|
assert result.error is not None and "`type`" in result.error
|
|
assert result.trimmed and result.text.count("\n") == 4
|
|
closed = await world.conversations.get(chat.external_id)
|
|
assert closed.status == "closed"
|
|
assert closed.flags["digest_error"] == result.error
|
|
|
|
|
|
def test_check_digest_rejects_what_is_not_a_digest(tmp_path: Path) -> None:
|
|
config = Distiller(agent="x", dir=tmp_path, index=tmp_path / "i.md")
|
|
path = tmp_path / "d.md"
|
|
path.write_text("---\ntype: digest\nsource: ''\ndate: 2026-08-29\n---\nx\n")
|
|
assert check_digest(path, config) == "`source` is empty"
|
|
path.write_text("---\ntype: digest\nsource: '[[a]]'\ndate: вчера\n---\nx\n")
|
|
assert "`date`" in check_digest(path, config)
|
|
path.write_text("---\ntype: digest\nsource: '[[a]]'\ndate: 2026-08-29\n---\n\n")
|
|
assert check_digest(path, config) == "empty body"
|
|
path.write_text("---\ntype: digest\nsource: '[[a]]'\ndate: 2026-08-29\n---\nx\n")
|
|
digest = check_digest(path, config)
|
|
assert not isinstance(digest, str) and digest.date.isoformat() == "2026-08-29"
|
|
assert trim_summary("a\n\nb\nc\nd\ne\nf") == ("a\nb\nc\nd\ne", True)
|
|
|
|
|
|
async def test_close_chat_tool_closes_after_the_reply(world: World) -> None:
|
|
distiller(world, DistillerClient)
|
|
world.deep_agent = ClaudeAgent(
|
|
name="d",
|
|
model="m",
|
|
system_prompt="deep",
|
|
cwd=world.root,
|
|
gateway_tools=("close_chat",),
|
|
)
|
|
world.conversations._agents = AgentRegistry( # noqa: SLF001
|
|
[world.agent, world.deep_agent, world.conversations._agents.get("x")] # noqa: SLF001
|
|
)
|
|
world.conversations._backends["d"] = ClaudeSdkBackend( # noqa: SLF001
|
|
agent=world.deep_agent,
|
|
mcp_internal_urls={},
|
|
session_store=world.store,
|
|
client_factory=ClosingClient,
|
|
work_dir=world.root / "work",
|
|
pool=world.pool,
|
|
tool_server=lambda key, _kind: build_tool_server(
|
|
world.conversations, conversation_key=key, names=("close_chat",)
|
|
),
|
|
)
|
|
master = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
chat = await deep_chat(world)
|
|
ClosingClient.conversations = world.conversations
|
|
ClosingClient.key = chat.external_id
|
|
|
|
await world.conversations.post(chat, "ок, обсудили")
|
|
await world.settle(chat, 1)
|
|
for _ in range(100):
|
|
row = await world.conversations.get(chat.external_id)
|
|
if row.status == "closed":
|
|
break
|
|
await asyncio.sleep(0.05)
|
|
else:
|
|
raise AssertionError("chat did not close")
|
|
|
|
assert row.flags["closed_reason"] == "close_chat"
|
|
assert row.flags["digest"] is not None
|
|
items = await world.conversations.queue.recent(master.id)
|
|
assert items[0].origin == "digest"
|
|
|
|
|
|
async def test_idle_picks_quiet_chats_after_launch_at_most_limit(world: World) -> None:
|
|
distiller(world, DistillerClient)
|
|
now = datetime.now(UTC)
|
|
launch = now - timedelta(days=10)
|
|
|
|
async def aged(name: str, days: float) -> Conversation:
|
|
conv = await deep_chat(world, name)
|
|
|
|
async def apply(row: Conversation) -> None:
|
|
row.last_activity_at = now - timedelta(days=days)
|
|
|
|
return await world.conversations._update(conv, apply) # noqa: SLF001
|
|
|
|
old = await aged("old", 30)
|
|
a = await aged("a", 5)
|
|
b = await aged("b", 4)
|
|
c = await aged("c", 3)
|
|
d = await aged("d", 2.5)
|
|
fresh = await aged("fresh", 1)
|
|
no_session = await world.conversations.create(kind="deep", agent="d", origin="t")
|
|
|
|
idle = await world.conversations.idle(kind="deep", days=2, since=launch)
|
|
assert [x.external_id for x in idle] == [
|
|
a.external_id,
|
|
b.external_id,
|
|
c.external_id,
|
|
d.external_id,
|
|
]
|
|
assert old.external_id not in {x.external_id for x in idle}
|
|
assert fresh.external_id not in {x.external_id for x in idle}
|
|
assert no_session.external_id not in {x.external_id for x in idle}
|
|
|
|
scheduler = Scheduler(conversations=world.conversations)
|
|
run = JobRun(Job("закрытие", lambda _run: asyncio.sleep(0)), "cron", {}, scheduler)
|
|
closed = await run.close_idle(kind="deep", days=2, limit=3, since=launch)
|
|
assert [r.conversation.external_id for r in closed] == [
|
|
a.external_id,
|
|
b.external_id,
|
|
c.external_id,
|
|
]
|
|
assert (await world.conversations.get(d.external_id)).status == "open"
|
|
assert len(await world.conversations.idle(kind="deep", days=2, since=launch)) == 1
|
|
|
|
|
|
async def test_line_cap_bounces_a_long_rewrite_and_asks_to_shorten(
|
|
world: World,
|
|
) -> None:
|
|
distiller(world, CapClient)
|
|
state = world.root / "vault" / "состояние.md"
|
|
state.write_text("# состояние\n- было так\n", encoding="utf-8")
|
|
CapClient.target = state
|
|
CapClient.lines = 70
|
|
scheduler = Scheduler(conversations=world.conversations)
|
|
run = JobRun(Job("память", lambda _run: asyncio.sleep(0)), "cron", {}, scheduler)
|
|
|
|
job = await run.spawn_job(
|
|
agent="x", text="перепиши", line_cap=LineCap(state, max_lines=60)
|
|
)
|
|
await world.settle(job, 2)
|
|
|
|
client = ScriptedClient.instances[-1]
|
|
assert len(client.prompts) == 2
|
|
assert "70 lines against a cap of 60" in client.prompts[1]
|
|
assert "[inject: cap" in client.prompts[1]
|
|
assert state.read_text(encoding="utf-8").count("\n") == 10 - 1
|
|
row = await world.conversations.get(job.external_id)
|
|
assert row.flags["line_cap_attempts"] == 1
|