205 lines
7.9 KiB
Python
205 lines
7.9 KiB
Python
import asyncio
|
|
from datetime import UTC, date, datetime, timedelta
|
|
|
|
from test_conversations import ScriptedClient, StubFrontend, World, world
|
|
|
|
from beaver_gateway.conversations.service import ConversationTexts
|
|
from beaver_gateway.conversations.rotation import (
|
|
HandoutContext,
|
|
Rotation,
|
|
RotationPolicy,
|
|
)
|
|
from beaver_gateway.storage.models import Conversation, Usage
|
|
|
|
__all__ = ["world"]
|
|
|
|
POLICY = RotationPolicy(tz="Europe/Warsaw")
|
|
|
|
|
|
def master(
|
|
*, created_ago: timedelta, silence: timedelta, now: datetime
|
|
) -> Conversation:
|
|
return Conversation(
|
|
frontend="test",
|
|
external_id="m",
|
|
agent_name="a",
|
|
kind="master",
|
|
created_at=now - created_ago,
|
|
last_user_activity_at=now - silence,
|
|
)
|
|
|
|
|
|
def test_night_rule_needs_silence_and_a_master_from_before_four() -> None:
|
|
now = datetime(2026, 8, 29, 2, 30, tzinfo=UTC)
|
|
quiet = master(created_ago=timedelta(hours=20), silence=timedelta(hours=4), now=now)
|
|
assert POLICY.reason(quiet, now=now, context_tokens=0) == "night"
|
|
active = master(
|
|
created_ago=timedelta(hours=20), silence=timedelta(minutes=5), now=now
|
|
)
|
|
assert POLICY.reason(active, now=now, context_tokens=0) is None
|
|
fresh = master(
|
|
created_ago=timedelta(minutes=20), silence=timedelta(hours=4), now=now
|
|
)
|
|
assert POLICY.reason(fresh, now=now, context_tokens=0) is None
|
|
early = datetime(2026, 8, 29, 1, 30, tzinfo=UTC)
|
|
before = master(
|
|
created_ago=timedelta(hours=4), silence=timedelta(hours=4), now=early
|
|
)
|
|
assert POLICY.reason(before, now=early, context_tokens=0) is None
|
|
|
|
|
|
def test_age_and_context_rules_need_thirty_minutes_of_silence() -> None:
|
|
now = datetime(2026, 8, 29, 12, 0, tzinfo=UTC)
|
|
old = master(
|
|
created_ago=timedelta(hours=37), silence=timedelta(minutes=31), now=now
|
|
)
|
|
assert POLICY.reason(old, now=now, context_tokens=0) == "age"
|
|
busy = master(
|
|
created_ago=timedelta(hours=37), silence=timedelta(minutes=5), now=now
|
|
)
|
|
assert POLICY.reason(busy, now=now, context_tokens=0) is None
|
|
big = master(created_ago=timedelta(hours=2), silence=timedelta(minutes=31), now=now)
|
|
assert POLICY.reason(big, now=now, context_tokens=90_000) == "context"
|
|
assert POLICY.reason(big, now=now, context_tokens=70_000) is None
|
|
|
|
|
|
async def age(world: World, conv: Conversation, created: datetime) -> Conversation:
|
|
async def apply(row: Conversation) -> None:
|
|
row.created_at = created
|
|
row.last_user_activity_at = created + timedelta(hours=1)
|
|
|
|
return await world.conversations._update(conv, apply) # noqa: SLF001
|
|
|
|
|
|
async def test_rotation_does_not_touch_a_master_mid_turn(world: World) -> None:
|
|
conv = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
ScriptedClient.hold = asyncio.Event()
|
|
await world.conversations.post(conv, "working")
|
|
await asyncio.sleep(0.2)
|
|
rotation = Rotation(world.conversations, POLICY)
|
|
assert await rotation.rotate(conv, "age") is None
|
|
assert (await world.conversations.get(conv.external_id)).status == "open"
|
|
assert len(await world.conversations.find(kind="master")) == 1
|
|
ScriptedClient.hold.set()
|
|
await world.settle(conv, 1)
|
|
|
|
|
|
async def test_rotation_order_handout_close_marks_moves_and_new_day(
|
|
world: World,
|
|
) -> None:
|
|
marked: list[str] = []
|
|
handouts: list[HandoutContext] = []
|
|
|
|
class MarkingFrontend(StubFrontend):
|
|
async def mark_closed(self, conv: Conversation) -> bool:
|
|
marked.append(conv.external_id)
|
|
return True
|
|
|
|
tg = MarkingFrontend("tg", ("master", "branch"), home=True)
|
|
tg.conversations = world.conversations
|
|
world.conversations._frontends = [tg, world.api] # noqa: SLF001
|
|
|
|
def handout(ctx: HandoutContext) -> str:
|
|
handouts.append(ctx)
|
|
return f"напиши хендаут за {ctx.day}"
|
|
|
|
world.conversations._texts = ConversationTexts( # noqa: SLF001
|
|
handout=handout, new_day=lambda ctx: f"Новый день {ctx.day} ({ctx.reason})."
|
|
)
|
|
old = await world.conversations.spawn(kind="master", agent="a", seed="clean")
|
|
old = await age(world, old, datetime(2026, 8, 27, 9, 0, tzinfo=UTC))
|
|
await world.conversations.post(old, "hi")
|
|
await world.settle(old, 1)
|
|
old = await world.conversations.get(old.external_id)
|
|
merged = await world.conversations.spawn(
|
|
kind="branch", parent=old, seed="brief", text="done"
|
|
)
|
|
await world.settle(merged, 1)
|
|
await world.conversations.set_status(merged, "merged")
|
|
live = await world.conversations.spawn(
|
|
kind="branch", parent=old, seed="brief", text="still going"
|
|
)
|
|
await world.settle(live, 1)
|
|
await world.conversations.inject(old, "later", urgency="normal", origin="крон")
|
|
old_client = ScriptedClient.instances[0]
|
|
|
|
new = await Rotation(world.conversations, POLICY).rotate(old, "night")
|
|
assert new is not None and new.kind == "master"
|
|
assert handouts[0].day == date(2026, 8, 27)
|
|
assert old_client.prompts[-1] == "напиши хендаут за 2026-08-27"
|
|
closed = await world.conversations.get(old.external_id)
|
|
assert closed.status == "closed"
|
|
assert world.pool.get(old.external_id) is None
|
|
assert marked == [merged.external_id]
|
|
assert (await world.conversations.get(live.external_id)).parent_id == new.id
|
|
bound = await world.conversations.find_bound(
|
|
frontend="tg", external_id=f"tg:{new.external_id}"
|
|
)
|
|
assert bound is not None and bound.id == new.id
|
|
old_bindings = await world.conversations.bindings(closed)
|
|
assert all(b.visible for b in old_bindings)
|
|
assert [i.text for i in await world.conversations.queue.pending(old.id)] == []
|
|
await world.settle(new, 1)
|
|
new_client = ScriptedClient.instances[-1]
|
|
prompt = new_client.prompts[0]
|
|
assert prompt.startswith("[seed: morning] master")
|
|
assert "[inject: rotation" in prompt
|
|
assert "Новый день 20" in prompt and "(night)" in prompt
|
|
assert "1 queued injects moved over" in prompt
|
|
moved = await world.conversations.queue.pending(new.id)
|
|
assert [(i.priority, i.text) for i in moved] == [("normal", "later")]
|
|
assert (await world.conversations.find(kind="master", status="open")) == [
|
|
await world.conversations.get(new.external_id)
|
|
]
|
|
|
|
|
|
async def test_due_uses_last_usage_row_for_context_size(world: World) -> None:
|
|
conv = await world.conversations.create(kind="master", agent="a", origin="test")
|
|
await age(world, conv, datetime.now(UTC) - timedelta(hours=2))
|
|
rotation = Rotation(world.conversations, RotationPolicy(max_context_tokens=10))
|
|
assert await rotation.due() == []
|
|
async with world.db.session() as session:
|
|
session.add(
|
|
Usage(
|
|
agent_name="a",
|
|
conversation_id=conv.external_id,
|
|
model="m",
|
|
input_tokens=1,
|
|
cache_read_tokens=2,
|
|
cache_creation_tokens=3,
|
|
output_tokens=100,
|
|
)
|
|
)
|
|
await session.commit()
|
|
assert await world.conversations.context_tokens(conv) == 6
|
|
assert await rotation.due() == []
|
|
rotation = Rotation(world.conversations, RotationPolicy(max_context_tokens=5))
|
|
(pair,) = await rotation.due()
|
|
assert pair[0].id == conv.id and pair[1] == "context"
|
|
|
|
|
|
def test_context_of_prefers_last_call_and_averages_old_rows() -> None:
|
|
from beaver_gateway.conversations.service import context_of
|
|
from beaver_gateway.storage.models import Usage
|
|
|
|
fresh = Usage(
|
|
agent_name="a",
|
|
model="m",
|
|
input_tokens=12,
|
|
cache_read_tokens=142_984,
|
|
cache_creation_tokens=29_584,
|
|
num_turns=6,
|
|
context_tokens=29_000,
|
|
)
|
|
assert context_of(fresh) == 29_000
|
|
old = Usage(
|
|
agent_name="a",
|
|
model="m",
|
|
input_tokens=12,
|
|
cache_read_tokens=142_984,
|
|
cache_creation_tokens=29_584,
|
|
num_turns=6,
|
|
)
|
|
assert context_of(old) == 28_763
|
|
assert context_of(None) == 0
|