import asyncio from datetime import UTC, date, datetime, timedelta from test_conversations import ScriptedClient, StubFrontend, World, world from beaver_gateway.core.conversations import ConversationTexts from beaver_gateway.core.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) == "ночь" 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) == "возраст" 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) == "транскрипт" 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, "возраст") 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="Новый день {day}." ) 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, "ночь") 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("[сид: morning] master") assert "[инжект: ротация" in prompt assert "Новый день" in prompt assert "переехало из старого мастера: 1" 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] == "транскрипт"