diff --git a/Makefile b/Makefile index cc4671d..627d096 100644 --- a/Makefile +++ b/Makefile @@ -25,7 +25,7 @@ ui-build: cd ui && bun run build db: - docker compose --project-directory ../beaver-agent --profile db up -d postgres + docker compose -f ../beaver-agent/docker-compose.yml --project-directory ../beaver-agent --profile db up -d postgres db-down: - docker compose --project-directory ../beaver-agent --profile db down + docker compose -f ../beaver-agent/docker-compose.yml --project-directory ../beaver-agent --profile db down diff --git a/pyproject.toml b/pyproject.toml index c23991c..361b4a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,10 +15,12 @@ dependencies = [ "anyio>=4.13.0", "argon2-cffi>=25.1.0", "claude-agent-sdk>=0.2.146", + "croniter>=6.2.4", "fastapi>=0.136.1", "fastmcp>=3.3.1", "greenlet>=3.5.0", "itsdangerous>=2.2.0", + "pgqueuer>=1.3.2", "psutil>=7.2.2", "psycopg[binary]>=3.3.4", "pydantic>=2.13.4", @@ -27,6 +29,7 @@ dependencies = [ "sqlmodel>=0.0.38", "uvicorn[standard]>=0.47.0", "uvloop>=0.22.1", + "watchfiles>=1.2.0", ] [project.optional-dependencies] diff --git a/src/beaver_gateway/cli.py b/src/beaver_gateway/cli.py index 60ab858..47ad554 100644 --- a/src/beaver_gateway/cli.py +++ b/src/beaver_gateway/cli.py @@ -28,9 +28,11 @@ import signal from contextlib import AsyncExitStack from typing import TYPE_CHECKING, Any +import psycopg import uvicorn import uvloop from dotenv import load_dotenv +from pgqueuer import PsycopgDriver from raycast_api import Client as RaycastClient from raycast_api.config import Config as RaycastConfig @@ -46,9 +48,13 @@ from beaver_gateway.backends.raycast import RaycastBackend from beaver_gateway.core.auth import TokenStore from beaver_gateway.core.bus import EventBus from beaver_gateway.core.conversations import Conversations +from beaver_gateway.core.envelope import Envelope from beaver_gateway.core.gateway_tools import build_tool_server from beaver_gateway.core.registry import AgentRegistry, Gateway, McpRegistry +from beaver_gateway.core.rotation import Rotation, RotationPolicy +from beaver_gateway.core.scheduler import Scheduler from beaver_gateway.core.sessions import SessionPool +from beaver_gateway.frontends._auth import require_token from beaver_gateway.frontends.base import GatewayRuntime from beaver_gateway.frontends.root import build_root_app from beaver_gateway.mcp.internal_app import build_internal_app @@ -60,6 +66,8 @@ if TYPE_CHECKING: from fastmcp import FastMCP from fastmcp.tools.base import Tool as FastMCPTool from starlette.applications import Starlette + from starlette.requests import Request + from starlette.types import ASGIApp from beaver_gateway.backends.base import Backend from beaver_gateway.mcp.types import McpServerT @@ -172,8 +180,19 @@ async def _async_main() -> None: store=session_store, texts=gateway.texts, frontends=gateway.frontends, + envelope=Envelope(watch=gateway.watch, tz=gateway.tz), ) late.conversations = conversations + scheduler = Scheduler( + conversations=conversations, + jobs=gateway.jobs, + driver=await _pgqueuer_driver(settings.database_url, stack), + budget=gateway.budget, + rotation=Rotation( + conversations, gateway.rotation or RotationPolicy(tz=gateway.tz) + ), + ) + conversations.scheduler = scheduler runtime = GatewayRuntime( agents=agents, @@ -189,6 +208,7 @@ async def _async_main() -> None: conversations=conversations, bus=bus, pool=pool, + scheduler=scheduler, public_url=gateway.public_url.rstrip("/") if gateway.public_url else None, ) @@ -217,23 +237,53 @@ async def _async_main() -> None: await conversations.start() stack.push_async_callback(conversations.stop) + await scheduler.start() + stack.push_async_callback(scheduler.stop) + hooks = scheduler.app( + functools.partial(_authorize_hook, runtime=runtime, scope="api") + ) async with asyncio.TaskGroup() as tg: tg.create_task(pool.reap_loop()) if internal_app is not None: tg.create_task(_serve_internal_mcp(internal_app, settings=settings)) - tg.create_task(_serve_root(gateway)) + tg.create_task(_serve_root(gateway, extra={"/hooks": hooks})) + if gateway.watch is not None: + tg.create_task(gateway.watch.run()) for fe in gateway.frontends: tg.create_task(fe.serve()) -async def _serve_root(gateway: Gateway) -> None: - app = build_root_app(gateway.frontends) +async def _authorize_hook( + request: Request, *, runtime: GatewayRuntime, scope: str +) -> str: + return await require_token(request, runtime, scope=scope) + + +async def _pgqueuer_driver(url: str, stack: AsyncExitStack) -> PsycopgDriver | None: + """A dedicated autocommit connection for pgqueuer's LISTEN/NOTIFY.""" + plain = _plain_postgres_url(url) + if plain is None: + return None + conn = await psycopg.AsyncConnection.connect(plain, autocommit=True) + stack.push_async_callback(conn.close) + return PsycopgDriver(conn) + + +def _plain_postgres_url(url: str) -> str | None: + for prefix in ("postgresql+psycopg://", "postgresql://", "postgres://"): + if url.startswith(prefix): + return "postgresql://" + url[len(prefix) :] + return None + + +async def _serve_root(gateway: Gateway, *, extra: dict[str, ASGIApp]) -> None: + app = build_root_app(gateway.frontends, extra=extra) config = uvicorn.Config(app, host=gateway.host, port=gateway.port, log_level="info") _log.info( "gateway on http://%s:%d - %s", gateway.host, gateway.port, - ", ".join(fe.path for fe in gateway.frontends if fe.path) + ", ".join([*(fe.path for fe in gateway.frontends if fe.path), *extra]) or "no http frontends", ) await uvicorn.Server(config).serve() diff --git a/src/beaver_gateway/config_loader.py b/src/beaver_gateway/config_loader.py index 102a97b..6365e60 100644 --- a/src/beaver_gateway/config_loader.py +++ b/src/beaver_gateway/config_loader.py @@ -16,11 +16,14 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any +from croniter import croniter + from beaver_gateway.agents.base import BaseAgent, ExposedMcp from beaver_gateway.agents.claude import ClaudeAgent from beaver_gateway.agents.raycast import RaycastAgent from beaver_gateway.core.conversations import ConversationTexts from beaver_gateway.core.registry import Gateway +from beaver_gateway.core.scheduler import Job from beaver_gateway.frontends.base import Frontend from beaver_gateway.mcp.types import HttpMcp, McpServer, PythonToolMcp, StdioMcp @@ -99,3 +102,15 @@ def _validate(gw: Gateway, path: Path) -> None: f"got {type(f).__name__}" ) raise ConfigError(msg) + names: set[str] = set() + for i, j in enumerate(gw.jobs): + if not isinstance(j, Job): + msg = f"{path}: gateway.jobs[{i}] must be a Job, got {type(j).__name__}" + raise ConfigError(msg) + if j.name in names: + msg = f"{path}: duplicate job name {j.name!r}" + raise ConfigError(msg) + names.add(j.name) + if j.cron is not None and not croniter.is_valid(j.cron): + msg = f"{path}: job {j.name!r} has an invalid cron {j.cron!r}" + raise ConfigError(msg) diff --git a/src/beaver_gateway/core/conversations.py b/src/beaver_gateway/core/conversations.py index ba0189e..506c72d 100644 --- a/src/beaver_gateway/core/conversations.py +++ b/src/beaver_gateway/core/conversations.py @@ -57,7 +57,7 @@ from beaver_gateway.storage.models import ( ConversationMessage, InjectQueueItem, RateLimit, - Schedule, + Usage, ) if TYPE_CHECKING: @@ -68,9 +68,12 @@ if TYPE_CHECKING: from beaver_gateway.agents.claude import ClaudeAgent from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend from beaver_gateway.core.bus import EventBus + from beaver_gateway.core.envelope import Envelope from beaver_gateway.core.events import MessageStreamEvent from beaver_gateway.core.injects import Priority from beaver_gateway.core.registry import AgentRegistry + from beaver_gateway.core.rotation import HandoutContext + from beaver_gateway.core.scheduler import Scheduler from beaver_gateway.core.sessions import SessionPool from beaver_gateway.frontends.base import Frontend from beaver_gateway.storage.db import Database @@ -124,6 +127,11 @@ class ConversationTexts: "заверши тёрн сейчас, ответ придёт следующим сообщением." ) seed: Callable[[SeedContext], Awaitable[str | None] | str | None] | None = None + handout: Callable[[HandoutContext], Awaitable[str] | str] | str = ( + "Этот мастер закрывается ({reason}). Напиши хендаут за {day}: справку " + "на утро, не задание - прошедшее время, без повелительного наклонения." + ) + new_day: str = "Новый день: мастер сменился, хендаут за {day} записан." @dataclass(frozen=True, slots=True) @@ -182,6 +190,7 @@ class Conversations: idle_days: Sequence[int] = (2,), idle_interval: float = 3600.0, question_timeout: float = 600.0, + envelope: Envelope | None = None, ) -> None: self._db = db self._agents = agents @@ -195,6 +204,8 @@ class Conversations: self._idle_days = tuple(sorted(idle_days)) self._idle_interval = idle_interval self._question_timeout = question_timeout + self._envelope = envelope + self.scheduler: Scheduler | None = None self._questions: dict[str, _Question] = {} self._queue = InjectQueue(db) self._runners: dict[int, _Runner] = {} @@ -481,6 +492,45 @@ class Conversations: ) return list(result.all()) + async def context_tokens(self, conv: Conversation) -> int: + """Size of the context the last turn ran with, from its usage row.""" + async with self._db.session() as session: + row = ( + await session.exec( + select(Usage) + .where(Usage.conversation_id == conv.external_id) + .order_by(col(Usage.id).desc()) + .limit(1) + ) + ).first() + if row is None: + return 0 + return row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens + + async def usage_tokens(self, since: datetime) -> int: + async with self._db.session() as session: + rows = ( + await session.exec( + select(Usage).where( + col(Usage.ts) >= since.astimezone(UTC).replace(tzinfo=None) + ) + ) + ).all() + return sum( + r.input_tokens + r.output_tokens + r.cache_creation_tokens for r in rows + ) + + async def busy(self, conv: Conversation) -> bool: + """A turn is running, a question is open or a message waits to run.""" + row = await self.get_row(cast("int", conv.id)) or conv + if row.running_turn or row.pending_question: + return True + live = self._pool.get(row.external_id) + if live is not None and live.busy: + return True + pending = await self._queue.pending(cast("int", row.id)) + return any(i.priority in ("user", "urgent") for i in pending) + # ---- routing ------------------------------------------------------- @property @@ -541,6 +591,8 @@ class Conversations: raise ValueError(msg) if agent is None and kind == "branch" and parent is not None: agent = parent.agent_name + if kind == "branch" and parent is not None and parent.kind == "master": + await self.set_flags(parent, {"streak": 0}) agent = agent or self.default_agent(kind) if agent is None: msg = f"no default agent for kind {kind!r}; pass `agent`" @@ -712,6 +764,7 @@ class Conversations: *, urgency: Priority = "normal", origin: str = "system", + interrupt: bool = True, ) -> InjectQueueItem: item = await self._queue.push( conversation_id=cast("int", conv.id), @@ -726,7 +779,7 @@ class Conversations: priority=urgency, origin=origin, ) - if urgency == "urgent": + if urgency == "urgent" and interrupt: backend = self._backend(conv.agent_name) if await backend.interrupt(conv.external_id): _log.info( @@ -788,22 +841,74 @@ class Conversations: ) return result - async def schedule(self, conv: Conversation, at: str, text: str) -> Schedule: - row = Schedule( - conversation_id=cast("int", conv.id), execute_at=parse_at(at), text=text - ) - async with self._db.session() as session: - session.add(row) - await session.commit() - await session.refresh(row) + async def schedule( + self, conv: Conversation, at: str, text: str, *, dedupe_key: str | None = None + ) -> tuple[int | None, datetime]: + if self.scheduler is None: + msg = "no scheduler; `schedule` is unavailable" + raise RuntimeError(msg) + return await self.scheduler.schedule(conv, at, text, dedupe_key=dedupe_key) + + async def schedules(self, conv: Conversation | None = None) -> list[dict[str, Any]]: + return await self.scheduler.scheduled(conv) if self.scheduler else [] + + # ---- §4.5 rotation ------------------------------------------------- + + async def handout(self, conv: Conversation, ctx: HandoutContext) -> str: + """The closing master's last turn: the handout prompt from the config.""" + source = self._texts.handout + if isinstance(source, str): + prompt = source.format(day=ctx.day.isoformat(), reason=ctx.reason) + else: + produced: Any = source(ctx) + prompt = await produced if inspect.isawaitable(produced) else produced self._bus.publish( - "schedule.created", - conversation_id=conv.external_id, - schedule=row.id, - execute_at=_iso(row.execute_at), + "handout.start", conversation_id=conv.external_id, day=ctx.day.isoformat() ) + try: + text, _ = await self.run_text_turn(conv, prompt, origin="handout") + except Exception: # noqa: BLE001 + _log.exception("handout turn on %s failed", conv.external_id) + text = "" + self._bus.publish( + "handout.end", + conversation_id=conv.external_id, + day=ctx.day.isoformat(), + text=text[:2000], + ) + return text + + async def close(self, conv: Conversation) -> Conversation: + row = await self.set_status(conv, "closed") + with contextlib.suppress(LookupError): + await self._backend(conv.agent_name).close(conv.external_id) return row + async def reparent(self, conv: Conversation, parent: Conversation) -> Conversation: + async def apply(row: Conversation) -> None: + row.parent_id = parent.id + + return await self._update(conv, apply) + + async def mark_closed(self, conv: Conversation) -> bool: + marked = False + for fe in self._frontends: + if conv.kind in fe.kinds: + try: + marked = await fe.mark_closed(conv) or marked + except Exception: # noqa: BLE001 + _log.exception("%s could not mark %s", fe.name, conv.external_id) + return marked + + async def new_day(self, conv: Conversation, *, moved: int = 0) -> InjectQueueItem: + day = datetime.now(UTC).astimezone().date().isoformat() + text = self._texts.new_day.format(day=day) + if moved: + text += f" Инжектов переехало из старого мастера: {moved}." + return await self.inject( + conv, text, urgency="urgent", origin="ротация", interrupt=False + ) + # ---- §3.7 questions ------------------------------------------------ async def ask(self, key: str, payload: dict[str, Any]) -> str | None: @@ -884,13 +989,6 @@ class Conversations: with contextlib.suppress(LookupError): await self._update(conv, apply) - async def schedules(self, conv: Conversation | None = None) -> list[Schedule]: - stmt = select(Schedule).order_by(col(Schedule.execute_at)) - if conv is not None: - stmt = stmt.where(Schedule.conversation_id == conv.id) - async with self._db.session() as session: - return list((await session.exec(stmt)).all()) - # ---- turns --------------------------------------------------------- async def turn( @@ -1193,6 +1291,9 @@ class Conversations: if head.priority == "user": origin = "user" prompt = head.text + envelope = await self._envelope_for(conv, injects=len(batch) - 1) + if envelope: + prompt += "\n\n" + envelope if len(batch) > 1: prompt += "\n\n" + _bundle(batch[1:]) else: @@ -1225,6 +1326,15 @@ class Conversations: text=text, ) + async def _envelope_for(self, conv: Conversation, *, injects: int) -> str | None: + if conv.kind != "master": + return None + streak = int(conv.flags.get("streak", 0) or 0) + await self.set_flags(conv, {"streak": streak + 1}) + if self._envelope is None: + return None + return self._envelope.build(streak=streak, injects=injects) + def _observer( self, conv: Conversation, runner: _Runner, turn_id: str, origin: str ) -> Callable[[Any], None]: diff --git a/src/beaver_gateway/core/envelope.py b/src/beaver_gateway/core/envelope.py new file mode 100644 index 0000000..5236e38 --- /dev/null +++ b/src/beaver_gateway/core/envelope.py @@ -0,0 +1,122 @@ +"""The envelope (§3.3): a background block after the user's text. + +Assembled when the turn starts, never when the message is queued: the +time, how many replies the master gave in a row without opening a branch, +what changed in the vault since the last envelope (added lines for the +``full`` files, names and counts for the rest) and how many normal +injects ride along below. Ceilings keep it a signal, not a document. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING +from zoneinfo import ZoneInfo + +if TYPE_CHECKING: + from collections.abc import Sequence + + from beaver_gateway.core.watch import Change, VaultWatch + +__all__ = ["Envelope", "render"] + +HEADER = ( + "[конверт - фоновый сигнал, не обращение; " + "реагируй, только если относится к вопросу]" +) + + +@dataclass(slots=True) +class Envelope: + watch: VaultWatch | None = None + tz: str = "UTC" + max_lines: int = 120 + per_file: int = 30 + names_only_within: float = 600.0 + last_at: datetime | None = None + + def build(self, *, streak: int, injects: int, now: datetime | None = None) -> str: + now = now or datetime.now(UTC) + changes = self.watch.take() if self.watch is not None else [] + names_only = ( + self.last_at is not None + and (now - self.last_at).total_seconds() < self.names_only_within + ) + text = render( + now=now, + tz=self.tz, + streak=streak, + changes=changes, + since=self.last_at, + names_only=names_only, + injects=injects, + max_lines=self.max_lines, + per_file=self.per_file, + ) + self.last_at = now + return text + + +def render( + *, + now: datetime, + tz: str, + streak: int, + changes: Sequence[Change], + since: datetime | None, + names_only: bool, + injects: int, + max_lines: int = 120, + per_file: int = 30, +) -> str: + zone = ZoneInfo(tz) + stamp = now.astimezone(zone) + lines = [ + HEADER, + f"время: {stamp:%Y-%m-%d %H:%M} ({_zone_label(tz)})", + f"мастер: {_replies(streak)} подряд без ветки", + ] + ordered = sorted(changes, key=lambda c: (not c.full, c.path)) + since_label = ( + f"с {since.astimezone(zone):%H:%M}" if since is not None else "со старта" # noqa: RUF001 + ) + if ordered: + names = ", ".join(f"{c.path} (+{c.added_count})" for c in ordered) + lines.append(f"vault, изменено {since_label} (mtime): {names}") + if injects: + lines.append(f"инжекты {since_label}: ({injects}) ниже") + if not names_only: + _append_diffs(lines, ordered, max_lines=max_lines, per_file=per_file) + return "\n".join(lines[:max_lines]) + + +def _append_diffs( + lines: list[str], changes: Sequence[Change], *, max_lines: int, per_file: int +) -> None: + budget = max_lines - len(lines) - 1 + for change in changes: + if not change.full or not change.added: + continue + if budget < 3: + lines.append("… (потолок конверта)") + return + shown = change.added[: min(per_file, budget - 2)] + lines.append(f"--- {change.path}, только добавленное ---") + lines.extend(f"+ {line}" for line in shown) + budget -= 1 + len(shown) + if len(change.added) > len(shown): + lines.append(f"+ … ещё {len(change.added) - len(shown)}") + budget -= 1 + + +def _replies(n: int) -> str: + if n % 10 == 1 and n % 100 != 11: + return f"{n} реплика" + if 2 <= n % 10 <= 4 and not 12 <= n % 100 <= 14: + return f"{n} реплики" + return f"{n} реплик" + + +def _zone_label(tz: str) -> str: + return tz.rsplit("/", 1)[-1].replace("_", " ") diff --git a/src/beaver_gateway/core/gateway_tools.py b/src/beaver_gateway/core/gateway_tools.py index 382e58e..ed76f86 100644 --- a/src/beaver_gateway/core/gateway_tools.py +++ b/src/beaver_gateway/core/gateway_tools.py @@ -134,8 +134,13 @@ def _tools(conversations: Conversations, key: str) -> list[SdkMcpTool[Any]]: ) async def schedule(args: dict[str, Any]) -> dict[str, Any]: conv = await current() - row = await conversations.schedule(conv, str(args["at"]), str(args["text"])) - return _text(f"scheduled #{row.id} at {row.execute_at.isoformat()}") + try: + job_id, when = await conversations.schedule( + conv, str(args["at"]), str(args["text"]) + ) + except (RuntimeError, ValueError) as exc: + return _error(str(exc)) + return _text(f"scheduled #{job_id} at {when.isoformat(timespec='minutes')}") @tool( "inject", diff --git a/src/beaver_gateway/core/injects.py b/src/beaver_gateway/core/injects.py index d1627e5..ca3f0be 100644 --- a/src/beaver_gateway/core/injects.py +++ b/src/beaver_gateway/core/injects.py @@ -10,7 +10,7 @@ runs one worker per conversation over them. A row that is still from __future__ import annotations from datetime import UTC, datetime -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Literal, cast from sqlalchemy import func from sqlmodel import col, select @@ -21,6 +21,7 @@ if TYPE_CHECKING: from collections.abc import Iterable, Sequence from beaver_gateway.storage.db import Database + from beaver_gateway.storage.models import Conversation __all__ = ["PRIORITY_RANK", "InjectQueue", "Priority", "inject_header"] @@ -94,6 +95,25 @@ class InjectQueue: await session.commit() return rows + async def move( + self, source: Conversation, target: Conversation, *, priority: Priority + ) -> int: + """Re-home queued items of one priority (rotation carries normal over).""" + async with self._db.session() as session: + result = await session.exec( + select(InjectQueueItem).where( + InjectQueueItem.conversation_id == source.id, + InjectQueueItem.status == "queued", + InjectQueueItem.priority == priority, + ) + ) + rows = list(result.all()) + for row in rows: + row.conversation_id = cast("int", target.id) + session.add(row) + await session.commit() + return len(rows) + async def recent( self, conversation_id: int, *, limit: int = 50 ) -> list[InjectQueueItem]: diff --git a/src/beaver_gateway/core/registry.py b/src/beaver_gateway/core/registry.py index dfb985c..acdeedc 100644 --- a/src/beaver_gateway/core/registry.py +++ b/src/beaver_gateway/core/registry.py @@ -17,6 +17,9 @@ if TYPE_CHECKING: from beaver_gateway.agents.base import BaseAgent from beaver_gateway.core.conversations import ConversationTexts + from beaver_gateway.core.rotation import RotationPolicy + from beaver_gateway.core.scheduler import Budget, Job + from beaver_gateway.core.watch import VaultWatch from beaver_gateway.frontends.base import Frontend from beaver_gateway.mcp.types import McpServerT @@ -84,6 +87,16 @@ class Gateway: frontends: list[Frontend] = field(default_factory=list) texts: ConversationTexts | None = None """Merge prompt and seed bodies for ``core/conversations`` (§8.2-8.3).""" + jobs: list[Job] = field(default_factory=list) + """Cron / webhook / event jobs for ``core/scheduler`` (§3.6, §4.5).""" + rotation: RotationPolicy | None = None + """When a master is rotated (§4.5); ``None`` keeps the defaults.""" + watch: VaultWatch | None = None + """Vault watcher feeding the envelope (§3.5, §4.6); ``None`` = no vault block.""" + budget: Budget | None = None + """Subscription window past which non-critical jobs wait (§4.5).""" + tz: str = "UTC" + """Local zone for the envelope clock and the rotation hour.""" host: str = "0.0.0.0" # noqa: S104 port: int = 8000 """The one listener; every HTTP frontend is mounted under its ``path``.""" diff --git a/src/beaver_gateway/core/rotation.py b/src/beaver_gateway/core/rotation.py new file mode 100644 index 0000000..9d3abfe --- /dev/null +++ b/src/beaver_gateway/core/rotation.py @@ -0,0 +1,131 @@ +"""Master rotation (§4.5, §8.1, §8.3). + +One logical master thread, many physical sessions: when the policy says +so, a new master is spawned and takes over the window atomically, the old +one writes its handout as its last turn, closes, its finished branches get +marked in their windows, its queued normal injects move over, and the new +one receives "new day". Silence is measured by the user's messages only - +injects never extend a day. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from datetime import UTC, date, datetime, timedelta +from typing import TYPE_CHECKING +from zoneinfo import ZoneInfo + +if TYPE_CHECKING: + from beaver_gateway.core.conversations import Conversations + from beaver_gateway.storage.models import Conversation + +__all__ = ["HandoutContext", "Rotation", "RotationPolicy"] + +_log = logging.getLogger("beaver_gateway.core.rotation") + + +@dataclass(frozen=True, slots=True) +class RotationPolicy: + tz: str = "UTC" + hour: int = 4 + night_silence: timedelta = timedelta(hours=3) + max_age: timedelta = timedelta(hours=36) + short_silence: timedelta = timedelta(minutes=30) + max_context_tokens: int = 80_000 + + def reason( + self, master: Conversation, *, now: datetime, context_tokens: int + ) -> str | None: + zone = ZoneInfo(self.tz) + created = _aware(master.created_at) + silence = now - _aware(master.last_user_activity_at or master.created_at) + boundary = now.astimezone(zone).replace( + hour=self.hour, minute=0, second=0, microsecond=0 + ) + if now.astimezone(zone) < boundary: + boundary -= timedelta(days=1) + if created < boundary and silence > self.night_silence: + return "ночь" + if now - created > self.max_age and silence > self.short_silence: + return "возраст" + if context_tokens > self.max_context_tokens and silence > self.short_silence: + return "транскрипт" + return None + + def day_of(self, master: Conversation) -> date: + return _aware(master.created_at).astimezone(ZoneInfo(self.tz)).date() + + +@dataclass(frozen=True, slots=True) +class HandoutContext: + day: date + master: Conversation + reason: str + + +class Rotation: + def __init__( + self, conversations: Conversations, policy: RotationPolicy | None = None + ) -> None: + self._conversations = conversations + self.policy = policy or RotationPolicy() + + async def due(self, now: datetime | None = None) -> list[tuple[Conversation, str]]: + now = now or datetime.now(UTC) + out: list[tuple[Conversation, str]] = [] + for master in await self._conversations.find(kind="master", status="open"): + tokens = await self._conversations.context_tokens(master) + reason = self.policy.reason(master, now=now, context_tokens=tokens) + if reason is not None: + out.append((master, reason)) + return out + + async def tick(self, now: datetime | None = None) -> list[Conversation]: + rotated: list[Conversation] = [] + for master, reason in await self.due(now): + new = await self.rotate(master, reason) + if new is not None: + rotated.append(new) + return rotated + + async def rotate(self, old: Conversation, reason: str) -> Conversation | None: + conversations = self._conversations + if await conversations.busy(old): + _log.info("rotation of %s skipped: busy", old.external_id) + return None + new = await conversations.spawn( + kind="master", agent=old.agent_name, seed="morning", origin="rotation" + ) + day = self.policy.day_of(old) + _log.info( + "rotation (%s): %s -> %s, handout for %s", + reason, + old.external_id, + new.external_id, + day, + ) + await conversations.handout( + old, HandoutContext(day=day, master=old, reason=reason) + ) + await conversations.close(old) + for branch in await conversations.find(parent=old, limit=1000): + if branch.status == "open": + await conversations.reparent(branch, new) + elif not branch.running_turn: + await conversations.mark_closed(branch) + moved = await conversations.queue.move(old, new, priority="normal") + await conversations.new_day(new, moved=moved) + conversations.bus.publish( + "conversation.rotated", + conversation_id=new.external_id, + closed=old.external_id, + reason=reason, + handout_day=day.isoformat(), + moved=moved, + ) + return new + + +def _aware(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=UTC) diff --git a/src/beaver_gateway/core/scheduler.py b/src/beaver_gateway/core/scheduler.py new file mode 100644 index 0000000..8627b16 --- /dev/null +++ b/src/beaver_gateway/core/scheduler.py @@ -0,0 +1,482 @@ +"""Jobs and deferred injects on pgqueuer (§3.6, §4.5). + +A job is a name, a handler and its triggers: a cron expression, the +webhook ``/hooks/``, gateway bus events. The executor is pgqueuer on +the gateway's own Postgres (a dedicated autocommit connection for +LISTEN/NOTIFY), so cron ticks, webhook deliveries and the ``schedule`` +tool's one-off injects all live in one table and survive a restart. A +handler only queues work for a conversation and returns; the turn itself +runs in the conversation's worker. Non-critical jobs step aside while the +subscription window is past its threshold. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Any, cast + +from pgqueuer import PgQueuer, Queries +from pgqueuer.models import JobId +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from beaver_gateway.core.conversations import parse_at + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Coroutine, Sequence + + from pgqueuer.models import Job as PgJob + from pgqueuer.models import Schedule as PgSchedule + from pgqueuer.ports.driver import Driver + from starlette.requests import Request + + from beaver_gateway.core.conversations import Conversations + from beaver_gateway.core.injects import Priority + from beaver_gateway.core.rotation import Rotation + from beaver_gateway.storage.models import Conversation + +__all__ = ["INJECT", "Budget", "Job", "JobRun", "Scheduler"] + +_log = logging.getLogger("beaver_gateway.core.scheduler") + +INJECT = "inject" +RETRY = timedelta(minutes=15) + + +@dataclass(frozen=True, slots=True) +class Job: + name: str + run: Callable[[JobRun], Awaitable[None]] + cron: str | None = None + webhook: bool = False + events: tuple[str, ...] = () + critical: bool = True + + @property + def entrypoint(self) -> str: + return f"job:{self.name}" + + +@dataclass(frozen=True, slots=True) +class Budget: + threshold: float = 0.7 + tokens: int | None = None + window: timedelta = timedelta(hours=5) + limit_window: str = "five_hour" + + +@dataclass(slots=True) +class JobRun: + job: Job + trigger: str + payload: dict[str, Any] + scheduler: Scheduler + + @property + def conversations(self) -> Conversations: + return self.scheduler.conversations + + async def master(self) -> Conversation | None: + masters = await self.conversations.find(kind="master", status="open", limit=1) + return masters[0] if masters else None + + async def inject_master( + self, text: str, *, urgency: Priority = "normal", origin: str | None = None + ) -> bool: + master = await self.master() + if master is None: + return False + await self.conversations.inject( + master, text, urgency=urgency, origin=origin or self.job.name + ) + return True + + async def spawn_job( + self, *, agent: str, text: str, title: str | None = None + ) -> Conversation: + return await self.conversations.spawn( + kind="job", agent=agent, seed="brief", text=text, title=title, origin="job" + ) + + async def retry_in(self, delay: timedelta) -> None: + await self.scheduler.trigger( + self.job, self.payload, delay=delay, trigger=self.trigger + ) + + async def rotate(self) -> list[Conversation]: + rotation = self.scheduler.rotation + return await rotation.tick() if rotation is not None else [] + + def background(self, coro: Coroutine[Any, Any, Any]) -> None: + self.scheduler.background(coro) + + +class Scheduler: + def __init__( + self, + *, + conversations: Conversations, + jobs: Sequence[Job] = (), + driver: Driver | None = None, + budget: Budget | None = None, + rotation: Rotation | None = None, + heartbeat: timedelta = timedelta(seconds=30), + ) -> None: + self.conversations = conversations + self.rotation = rotation + self.budget = budget or Budget() + self._jobs = {job.name: job for job in jobs} + self._driver = driver + self._queries: Queries | None = None + self._heartbeat = heartbeat + self._tasks: set[asyncio.Task[Any]] = set() + self._runs: dict[str, dict[str, Any]] = {} + + @property + def enabled(self) -> bool: + return self._queries is not None + + @property + def jobs(self) -> list[Job]: + return list(self._jobs.values()) + + def job(self, name: str) -> Job | None: + return self._jobs.get(name) + + # ---- lifecycle ----------------------------------------------------- + + async def start(self) -> None: + if self._driver is not None: + queries = Queries(self._driver) + if not await queries.has_table("pgqueuer"): + await queries.install() + pgq = PgQueuer(self._driver, queries=queries) + self._register(pgq) + self._queries = queries + self._spawn(pgq.run(heartbeat_timeout=self._heartbeat)) + else: + _log.warning("scheduler: no postgres driver - cron and webhooks are off") + self._spawn(self._events()) + + async def stop(self) -> None: + for task in list(self._tasks): + task.cancel() + for task in list(self._tasks): + with contextlib.suppress(BaseException): + await task + self._tasks.clear() + self._queries = None + + def _register(self, pgq: PgQueuer) -> None: + @pgq.entrypoint(INJECT) + async def deliver(job: PgJob) -> None: + await self._deliver(job) + + for job in self._jobs.values(): + self._register_job(pgq, job) + + def _register_job(self, pgq: PgQueuer, job: Job) -> None: + @pgq.entrypoint(job.entrypoint) + async def queued(pg_job: PgJob) -> None: + data = _decode(pg_job.payload) + await self._dispatch( + job, trigger=str(data.pop("trigger", "queue")), payload=data + ) + + if job.cron: + + @pgq.schedule(job.name, job.cron, clean_old=True) + async def cron(_schedule: PgSchedule) -> None: + await self._dispatch(job, trigger="cron", payload={}) + + async def _events(self) -> None: + listeners = [job for job in self._jobs.values() if job.events] + if not listeners: + return + async for event in self.conversations.bus.stream(): + for job in listeners: + if event.get("type") in job.events: + self._spawn( + self._dispatch(job, trigger="event", payload=dict(event)) + ) + + # ---- dispatch ------------------------------------------------------ + + async def _dispatch( + self, job: Job, *, trigger: str, payload: dict[str, Any] + ) -> None: + bus = self.conversations.bus + if not job.critical and await self.throttled(): + bus.publish("job.deferred", job=job.name, trigger=trigger) + if trigger != "cron": + await self.trigger(job, payload, delay=RETRY, trigger=trigger) + return + started = datetime.now(UTC) + bus.publish("job.start", job=job.name, trigger=trigger) + status = "done" + try: + await job.run(JobRun(job, trigger, payload, self)) + except Exception: # noqa: BLE001 + status = "failed" + _log.exception("job %s (%s) failed", job.name, trigger) + self._runs[job.name] = { + "trigger": trigger, + "started_at": started.isoformat(timespec="seconds"), + "status": status, + } + bus.publish("job.end", job=job.name, trigger=trigger, status=status) + + async def trigger( + self, + job: Job, + payload: dict[str, Any] | None = None, + *, + delay: timedelta | None = None, + trigger: str = "manual", + ) -> int | None: + """Queue one run of ``job``; runs inline when there is no executor.""" + data = {**(payload or {}), "trigger": trigger} + if self._queries is None: + self._spawn( + self._dispatch(job, trigger=trigger, payload=payload or {}), delay=delay + ) + return None + ids = await self._queries.enqueue( + job.entrypoint, _encode(data), execute_after=delay + ) + return int(ids[0]) if ids and ids[0] is not None else None + + async def hook(self, name: str, payload: dict[str, Any]) -> int | None: + job = self._jobs.get(name) + if job is None or not job.webhook: + msg = f"no webhook job {name!r}" + raise LookupError(msg) + self.conversations.bus.publish("hook", job=name) + if self._queries is None: + return await self.trigger(job, payload, trigger="webhook") + ids = await self._queries.enqueue( + job.entrypoint, + _encode({**payload, "trigger": "webhook"}), + dedupe_key=f"hook:{name}", + on_conflict="skip", + ) + return int(ids[0]) if ids and ids[0] is not None else None + + # ---- deferred injects ---------------------------------------------- + + async def schedule( + self, + conv: Conversation, + at: str, + text: str, + *, + urgency: Priority = "normal", + dedupe_key: str | None = None, + ) -> tuple[int | None, datetime]: + if self._queries is None: + msg = "scheduler needs postgres; `schedule` is unavailable" + raise RuntimeError(msg) + when = parse_at(at) + delay = max(when - datetime.now(UTC), timedelta(0)) + payload = _encode( + { + "conversation": conv.external_id, + "text": text, + "urgency": urgency, + "at": when.isoformat(timespec="seconds"), + } + ) + ids = await self._queries.enqueue( + INJECT, + payload, + execute_after=delay, + dedupe_key=dedupe_key, + on_conflict="skip" if dedupe_key else "raise", + ) + job_id = int(ids[0]) if ids and ids[0] is not None else None + self.conversations.bus.publish( + "schedule.created", + conversation_id=conv.external_id, + job=job_id, + execute_at=when.isoformat(timespec="seconds"), + ) + return job_id, when + + async def cancel(self, job_id: int) -> bool: + if self._queries is None: + return False + row = await self._queries.queue_job_by_id(JobId(job_id)) + if row is None: + return False + await self._queries.mark_job_as_cancelled([JobId(job_id)]) + self.conversations.bus.publish("schedule.cancelled", job=job_id) + return True + + async def _deliver(self, job: PgJob) -> None: + data = _decode(job.payload) + conv = await self.conversations.get(str(data.get("conversation", ""))) + if conv is not None and conv.status != "open" and conv.kind == "master": + masters = await self.conversations.find( + kind="master", status="open", limit=1 + ) + conv = masters[0] if masters else None + if conv is None: + _log.warning("scheduled inject #%s: conversation gone", job.id) + return + await self.conversations.inject( + conv, + str(data.get("text", "")), + urgency=cast("Priority", data.get("urgency") or "normal"), + origin="schedule", + ) + + # ---- budget -------------------------------------------------------- + + async def utilization(self) -> float | None: + now = datetime.now(UTC) + values: list[float] = [] + for row in await self.conversations.rate_limits(limit=100): + if row.window != self.budget.limit_window or row.utilization is None: + continue + fresh = now - _aware(row.ts) < self.budget.window + live = row.resets_at is None or _aware(row.resets_at) > now + if fresh and live: + values.append(float(row.utilization)) + break + if self.budget.tokens: + tokens = await self.conversations.usage_tokens(now - self.budget.window) + values.append(tokens / self.budget.tokens) + return max(values) if values else None + + async def throttled(self) -> bool: + utilization = await self.utilization() + return utilization is not None and utilization > self.budget.threshold + + # ---- introspection ------------------------------------------------- + + async def snapshot(self) -> dict[str, Any]: + crons: dict[str, PgSchedule] = {} + queue: list[dict[str, Any]] = [] + if self._queries is not None: + for row in await self._queries.peek_schedule(): + crons[str(row.entrypoint)] = row + queue = [ + _job_public(row) for row in await self._queries.browse_queue(limit=200) + ] + utilization = await self.utilization() + return { + "enabled": self.enabled, + "utilization": utilization, + "throttled": utilization is not None + and utilization > self.budget.threshold, + "threshold": self.budget.threshold, + "jobs": [self._job_public(job, crons.get(job.name)) for job in self.jobs], + "queue": queue, + } + + async def scheduled(self, conv: Conversation | None = None) -> list[dict[str, Any]]: + if self._queries is None: + return [] + rows = await self._queries.browse_queue(limit=500, entrypoints=[INJECT]) + out = [_job_public(row) for row in rows] + if conv is not None: + out = [ + j for j in out if j["payload"].get("conversation") == conv.external_id + ] + return out + + def _job_public(self, job: Job, cron: PgSchedule | None) -> dict[str, Any]: + return { + "name": job.name, + "cron": job.cron, + "webhook": job.webhook, + "events": list(job.events), + "critical": job.critical, + "next_run": _iso(cron.next_run) if cron is not None else None, + "last_run": _iso(cron.last_run) if cron is not None else None, + "status": str(cron.status) if cron is not None else None, + "run": self._runs.get(job.name), + } + + # ---- http ---------------------------------------------------------- + + def app(self, authorize: Callable[[Request], Awaitable[Any]]) -> Starlette: + async def hook(request: Request) -> JSONResponse: + await authorize(request) + name = request.path_params["name"] + payload = await _payload(request) + try: + job_id = await self.hook(name, payload) + except LookupError as exc: + return JSONResponse({"error": str(exc)}, status_code=404) + return JSONResponse({"job": job_id, "name": name}, status_code=202) + + return Starlette(routes=[Route("/{name}", hook, methods=["POST"])]) + + # ---- internals ----------------------------------------------------- + + def background(self, coro: Coroutine[Any, Any, Any]) -> None: + self._spawn(coro) + + def _spawn( + self, coro: Coroutine[Any, Any, Any], *, delay: timedelta | None = None + ) -> None: + async def later() -> None: + if delay is not None: + await asyncio.sleep(delay.total_seconds()) + await coro + + task = asyncio.create_task(later() if delay is not None else coro) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) + + +async def _payload(request: Request) -> dict[str, Any]: + body = await request.body() + if not body: + return dict(request.query_params) + try: + data = json.loads(body) + except ValueError: + return {"raw": body.decode("utf-8", errors="replace")} + return data if isinstance(data, dict) else {"raw": data} + + +def _encode(data: dict[str, Any]) -> bytes: + return json.dumps(data, ensure_ascii=False, default=str).encode("utf-8") + + +def _decode(payload: bytes | None) -> dict[str, Any]: + if not payload: + return {} + try: + data = json.loads(payload) + except ValueError: + return {"raw": payload.decode("utf-8", errors="replace")} + return data if isinstance(data, dict) else {"raw": data} + + +def _job_public(row: PgJob) -> dict[str, Any]: + return { + "id": int(row.id), + "entrypoint": row.entrypoint, + "status": str(row.status), + "execute_after": _iso(row.execute_after), + "created": _iso(row.created), + "attempts": row.attempts, + "payload": _decode(row.payload), + } + + +def _aware(value: datetime) -> datetime: + return value if value.tzinfo is not None else value.replace(tzinfo=UTC) + + +def _iso(value: datetime | None) -> str | None: + return _aware(value).isoformat(timespec="seconds") if value is not None else None diff --git a/src/beaver_gateway/core/watch.py b/src/beaver_gateway/core/watch.py new file mode 100644 index 0000000..5d76cd5 --- /dev/null +++ b/src/beaver_gateway/core/watch.py @@ -0,0 +1,205 @@ +"""Vault watcher for the envelope (§3.5, §4.6). + +A content snapshot is taken at every envelope; a change is the diff of the +file against that snapshot, "added lines only". ``watchfiles`` delivers +paths with a settle debounce (Sync writes files in pieces), the change +carries the file's mtime rather than the moment it landed. Which files +show a full diff, which show up by name and which are ignored is a set of +glob patterns the config hands in - the gateway itself knows no path. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import difflib +import fnmatch +import logging +import os +from dataclasses import dataclass, field +from datetime import UTC, datetime +from pathlib import Path +from typing import TYPE_CHECKING +from zoneinfo import ZoneInfo + +from watchfiles import awatch + +if TYPE_CHECKING: + from collections.abc import Iterable + +__all__ = ["Change", "VaultWatch", "WatchRules"] + +_log = logging.getLogger("beaver_gateway.core.watch") + + +@dataclass(frozen=True, slots=True) +class WatchRules: + """Globs relative to the vault root; ``{today}`` expands to the local date. + + ``full`` files show their added lines, ``names`` files only their name + and how many lines they gained, ``ignore`` hides them; when several + patterns match, the most specific one decides. + """ + + full: tuple[str, ...] = () + names: tuple[str, ...] = () + ignore: tuple[str, ...] = () + suffixes: tuple[str, ...] = (".md",) + + def kind(self, rel: str, *, today: str) -> str | None: + """The most specific matching pattern wins, whichever list it is in.""" + if not rel.endswith(self.suffixes): + return None + best: tuple[int, str | None] = (-1, None) + for kind, patterns in ( + ("full", self.full), + ("names", self.names), + (None, self.ignore), + ): + for pattern in patterns: + if _match(rel, pattern, today): + best = max(best, (_specificity(pattern), kind)) + return best[1] + + +@dataclass(slots=True) +class Change: + path: str + mtime: datetime + added: list[str] = field(default_factory=list) + added_count: int = 0 + full: bool = False + + +@dataclass(slots=True) +class _Seen: + lines: list[str] | None + count: int + + +class VaultWatch: + def __init__( + self, root: Path, rules: WatchRules, *, tz: str = "UTC", debounce: float = 4.0 + ) -> None: + self.root = root + self.rules = rules + self._tz = ZoneInfo(tz) + self._debounce = debounce + self._seen: dict[str, _Seen] = {} + self._pending: dict[str, Change] = {} + self._stop = asyncio.Event() + + @property + def pending(self) -> list[Change]: + return list(self._pending.values()) + + def today(self, now: datetime | None = None) -> str: + return (now or datetime.now(UTC)).astimezone(self._tz).date().isoformat() + + async def run(self) -> None: + await asyncio.to_thread(self.snapshot) + _log.info("watching %s (%d files in snapshot)", self.root, len(self._seen)) + async for changes in awatch( + self.root, + debounce=int(self._debounce * 1000), + stop_event=self._stop, + ignore_permission_denied=True, + ): + for _kind, raw in changes: + with contextlib.suppress(OSError): + self.note(Path(raw)) + + def stop(self) -> None: + self._stop.set() + + def snapshot(self) -> None: + today = self.today() + for path in self._walk(): + rel = self._rel(path) + kind = self.rules.kind(rel, today=today) + if kind is None: + continue + with contextlib.suppress(OSError): + lines = _read_lines(path) + self._seen[rel] = _Seen( + lines=lines if kind == "full" else None, count=len(lines) + ) + + def note(self, path: Path) -> Change | None: + """Record a change of ``path`` against the last envelope's snapshot.""" + try: + rel = self._rel(path) + except ValueError: + return None + kind = self.rules.kind(rel, today=self.today()) + if kind is None: + return None + if not path.is_file(): + self._pending.pop(rel, None) + return None + stat = path.stat() + lines = _read_lines(path) + seen = self._seen.get(rel) + change = Change( + path=rel, + mtime=datetime.fromtimestamp(stat.st_mtime, tz=UTC), + full=kind == "full", + ) + if seen is not None and seen.lines is not None: + change.added = _added(seen.lines, lines) + change.added_count = len(change.added) + elif seen is not None: + change.added_count = max(len(lines) - seen.count, 0) + else: + change.added = lines if kind == "full" else [] + change.added_count = len(lines) + if change.added_count == 0 and seen is not None: + self._pending.pop(rel, None) + return None + self._pending[rel] = change + return change + + def take(self) -> list[Change]: + """Hand out what changed since the last envelope and reset the snapshot.""" + changes = list(self._pending.values()) + self._pending.clear() + for change in changes: + path = self.root / change.path + with contextlib.suppress(OSError): + lines = _read_lines(path) + self._seen[change.path] = _Seen( + lines=lines if change.full else None, count=len(lines) + ) + return changes + + def _rel(self, path: Path) -> str: + return path.resolve().relative_to(self.root.resolve()).as_posix() + + def _walk(self) -> Iterable[Path]: + for dirpath, dirnames, filenames in os.walk(self.root): + dirnames[:] = [d for d in dirnames if not d.startswith(".")] + for name in filenames: + yield Path(dirpath) / name + + +def _match(rel: str, pattern: str, today: str) -> bool: + return fnmatch.fnmatchcase( + rel, pattern.replace("{today}", today).replace("**", "*") + ) + + +def _specificity(pattern: str) -> int: + return sum(1 for ch in pattern if ch not in "*?[]{}") + + +def _read_lines(path: Path) -> list[str]: + return path.read_text(encoding="utf-8", errors="replace").splitlines() + + +def _added(old: list[str], new: list[str]) -> list[str]: + matcher = difflib.SequenceMatcher(a=old, b=new, autojunk=False) + added: list[str] = [] + for tag, _i1, _i2, j1, j2 in matcher.get_opcodes(): + if tag in ("insert", "replace"): + added.extend(line for line in new[j1:j2] if line.strip()) + return added diff --git a/src/beaver_gateway/frontends/api/frontend.py b/src/beaver_gateway/frontends/api/frontend.py index 2207fc8..adf4f35 100644 --- a/src/beaver_gateway/frontends/api/frontend.py +++ b/src/beaver_gateway/frontends/api/frontend.py @@ -60,6 +60,7 @@ if TYPE_CHECKING: from pathlib import Path from beaver_gateway.core.conversations import Conversations + from beaver_gateway.core.scheduler import Scheduler from beaver_gateway.frontends.base import GatewayRuntime _log = logging.getLogger("beaver_gateway.frontends.api") @@ -525,19 +526,48 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa await require_token(request, runtime, scope=SCOPE) raw = request.query_params.get("conversation") conv = await conv_of(raw) if raw else None - return { - "schedules": [ - { - "id": s.id, - "conversation_row": s.conversation_id, - "execute_at": _iso(s.execute_at), - "text": s.text, - "created_at": _iso(s.created_at), - "delivered_at": _iso(s.delivered_at), - } - for s in await conversations.schedules(conv) - ] - } + return {"schedules": await conversations.schedules(conv)} + + def scheduler_of() -> Scheduler: + if runtime.scheduler is None: + raise HTTPException(status.HTTP_503_SERVICE_UNAVAILABLE, "no scheduler") + return cast("Scheduler", runtime.scheduler) + + @app.get("/jobs") + async def jobs(request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + return await scheduler_of().snapshot() + + @app.post("/jobs/{name}/run", status_code=status.HTTP_202_ACCEPTED) + async def run_job(name: str, request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + scheduler = scheduler_of() + job = scheduler.job(name) + if job is None: + raise HTTPException(status.HTTP_404_NOT_FOUND, f"no job {name!r}") + return {"job": await scheduler.trigger(job, await body_of(request))} + + @app.delete("/jobs/queue/{job_id}") + async def cancel_job(job_id: int, request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + if not await scheduler_of().cancel(job_id): + raise HTTPException(status.HTTP_404_NOT_FOUND, f"no queued job {job_id}") + return {"cancelled": job_id} + + @app.post( + "/conversations/{public_id}/schedule", status_code=status.HTTP_201_CREATED + ) + async def post_schedule(public_id: str, request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + data = await body_of(request) + try: + job_id, when = await conversations.schedule( + conv, text_of(data, "at"), text_of(data, "text") + ) + except (RuntimeError, ValueError) as exc: + raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc + return {"job": job_id, "execute_at": _iso(when)} @app.get("/usage") async def usage(request: Request) -> dict[str, Any]: diff --git a/src/beaver_gateway/frontends/base.py b/src/beaver_gateway/frontends/base.py index 1dcc879..7a360c4 100644 --- a/src/beaver_gateway/frontends/base.py +++ b/src/beaver_gateway/frontends/base.py @@ -89,6 +89,7 @@ class GatewayRuntime: conversations: Any = None bus: Any = None pool: Any = None + scheduler: Any = None # External origin the reverse proxy puts in front of the gateway # (``Gateway.public_url``); ``None`` means "derive from the request". public_url: str | None = None @@ -134,3 +135,7 @@ class Frontend(ABC): async def materialize(self, conv: Conversation) -> ConversationBinding | None: # noqa: ARG002 return None + + async def mark_closed(self, conv: Conversation) -> bool: # noqa: ARG002 + """Show in the window that the conversation is over (a renamed topic).""" + return False diff --git a/src/beaver_gateway/frontends/root.py b/src/beaver_gateway/frontends/root.py index 9bc1247..23ddd93 100644 --- a/src/beaver_gateway/frontends/root.py +++ b/src/beaver_gateway/frontends/root.py @@ -15,19 +15,22 @@ from starlette.responses import JSONResponse, RedirectResponse from starlette.routing import Mount, Route if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Iterable, Mapping from starlette.requests import Request + from starlette.types import ASGIApp from beaver_gateway.frontends.base import Frontend __all__ = ["build_root_app"] -def build_root_app(frontends: Iterable[Frontend]) -> Starlette: +def build_root_app( + frontends: Iterable[Frontend], *, extra: Mapping[str, ASGIApp] | None = None +) -> Starlette: mounted = [fe for fe in frontends if fe.path and fe.app() is not None] landing = next((fe for fe in mounted if fe.landing), None) - paths = [fe.path for fe in mounted] + paths = [*(fe.path for fe in mounted), *(extra or {})] async def healthz(_request: Request) -> JSONResponse: return JSONResponse({"status": "ok", "frontends": paths}) @@ -46,4 +49,6 @@ def build_root_app(frontends: Iterable[Frontend]) -> Starlette: assert app is not None # noqa: S101 - filtered above; narrows for ty assert fe.path is not None # noqa: S101 routes.append(Mount(fe.path, app=app, name=fe.name or fe.path.strip("/"))) + for path, app in (extra or {}).items(): + routes.append(Mount(path, app=app, name=path.strip("/"))) return Starlette(routes=routes) diff --git a/src/beaver_gateway/frontends/telegram/frontend.py b/src/beaver_gateway/frontends/telegram/frontend.py index 25c49da..0e6e6bf 100644 --- a/src/beaver_gateway/frontends/telegram/frontend.py +++ b/src/beaver_gateway/frontends/telegram/frontend.py @@ -203,12 +203,11 @@ class TelegramFrontend(Frontend): conv, frontend=FRONTEND, external_id=self._ext(topic.message_thread_id) ) - async def mark_topic(self, conv: Conversation, prefix: str = "✅ ") -> bool: - """Rotation hook for M3. + async def mark_closed(self, conv: Conversation) -> bool: + return await self.mark_topic(conv) - ``closeForumTopic`` does not exist in private chats; the state of a - merged or closed branch lives in its name. - """ + async def mark_topic(self, conv: Conversation, prefix: str = "✅ ") -> bool: + """Rename the topic; closeForumTopic does not exist in private chats.""" target = await self._target_of(conv) if target is None or target[1] is None: return False diff --git a/src/beaver_gateway/storage/__init__.py b/src/beaver_gateway/storage/__init__.py index 98f1703..4782f34 100644 --- a/src/beaver_gateway/storage/__init__.py +++ b/src/beaver_gateway/storage/__init__.py @@ -23,7 +23,6 @@ from beaver_gateway.storage.models import ( Delivery, InjectQueueItem, RateLimit, - Schedule, TelegramUpdate, Token, TranscriptEntry, @@ -40,7 +39,6 @@ __all__ = [ "InjectQueueItem", "PostgresSessionStore", "RateLimit", - "Schedule", "TelegramUpdate", "Token", "TranscriptEntry", diff --git a/src/beaver_gateway/storage/models.py b/src/beaver_gateway/storage/models.py index f830450..9c646bc 100644 --- a/src/beaver_gateway/storage/models.py +++ b/src/beaver_gateway/storage/models.py @@ -167,23 +167,6 @@ class InjectQueueItem(SQLModel, table=True): delivered_at: datetime | None = Field(default=None) -class Schedule(SQLModel, table=True): - """Deferred inject written by the ``schedule`` tool (§3.6). - - M1b only records the promise; the executor (pgqueuer, M3) will move - these into its own job table and this one goes away. - """ - - __tablename__ = "schedules" - - id: int | None = Field(default=None, primary_key=True) - conversation_id: int = Field(index=True) - execute_at: datetime = Field(index=True) - text: str - created_at: datetime = Field(default_factory=_utcnow) - delivered_at: datetime | None = Field(default=None) - - class TelegramUpdate(SQLModel, table=True): """Inbox of the Telegram frontend (§3.8). @@ -362,7 +345,6 @@ __all__ = [ "Delivery", "InjectQueueItem", "RateLimit", - "Schedule", "TelegramUpdate", "Token", "TranscriptEntry", diff --git a/tests/pgurl.py b/tests/pgurl.py new file mode 100644 index 0000000..ada4be8 --- /dev/null +++ b/tests/pgurl.py @@ -0,0 +1,9 @@ +"""Print a TEST_DATABASE_URL for the compose postgres (`make db`) from beaver-agent/.env.""" + +from dotenv import dotenv_values + +v = dotenv_values("../beaver-agent/.env") +user = v.get("POSTGRES_USER") or "beaver" +port = v.get("PORT_POSTGRES") or "5432" +db = v.get("POSTGRES_DB") or "beaver" +print(f"postgresql://{user}:{v['POSTGRES_PASSWORD']}@127.0.0.1:{port}/{db}") diff --git a/tests/test_conversations.py b/tests/test_conversations.py index fd50a5e..efa5d91 100644 --- a/tests/test_conversations.py +++ b/tests/test_conversations.py @@ -568,15 +568,16 @@ async def test_read_and_bindings(world: World) -> None: ) -async def test_schedule_rows_and_parse_at(world: World) -> None: +async def test_schedule_without_scheduler_and_parse_at(world: World) -> None: conv = await world.conversations.create(kind="master", agent="a", origin="test") - row = await world.conversations.schedule(conv, "+15m", "push X") - delta = (row.execute_at.replace(tzinfo=UTC) - datetime.now(UTC)).total_seconds() - assert 14 * 60 < delta <= 15 * 60 - assert [s.text for s in await world.conversations.schedules(conv)] == ["push X"] + 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") diff --git a/tests/test_envelope.py b/tests/test_envelope.py new file mode 100644 index 0000000..decfc81 --- /dev/null +++ b/tests/test_envelope.py @@ -0,0 +1,158 @@ +import asyncio +import tempfile +from datetime import UTC, datetime, timedelta +from pathlib import Path + +from test_conversations import ScriptedClient, World, world + +from beaver_gateway.core.envelope import HEADER, Envelope, render +from beaver_gateway.core.watch import Change, VaultWatch, WatchRules + +__all__ = ["world"] + +RULES = WatchRules( + full=("дни/{today}.md",), + names=("дни/*", "люди/*", "мета/бобер/*"), + ignore=("чаты/*", "мета/*"), +) + + +def vault() -> tuple[Path, VaultWatch]: + root = Path(tempfile.mkdtemp(prefix="beaver-vault-")) + for sub in ("дни", "люди", "чаты", "мета/бобер", "мета/чужое"): + (root / sub).mkdir(parents=True) + today = datetime.now(UTC).date().isoformat() + (root / "дни" / f"{today}.md").write_text("# день\n- 09:00 проснулся\n") + (root / "люди" / "Прохор.md").write_text("# Прохор\nстрока\n") + (root / "чаты" / "чат.md").write_text("чат\n") + (root / "мета" / "чужое" / "x.md").write_text("x\n") + watch = VaultWatch(root, RULES, tz="UTC") + watch.snapshot() + return root, watch + + +def append(path: Path, text: str) -> None: + with path.open("a", encoding="utf-8") as fh: + fh.write(text) + + +def test_rules_classify_paths() -> None: + today = "2026-08-29" + assert RULES.kind("дни/2026-08-29.md", today=today) == "full" + assert RULES.kind("дни/2026-08-28.md", today=today) == "names" + assert RULES.kind("люди/Прохор.md", today=today) == "names" + assert RULES.kind("мета/бобер/состояние.md", today=today) == "names" + assert RULES.kind("мета/чужое/x.md", today=today) is None + assert RULES.kind("чаты/чат.md", today=today) is None + assert RULES.kind("люди/фото.png", today=today) is None + + +def test_watch_reports_only_added_lines_with_mtime() -> None: + root, watch = vault() + today = watch.today() + diary = root / "дни" / f"{today}.md" + append(diary, "- 12:40 вышел\n- 13:00 кофе\n") + append(root / "люди" / "Прохор.md", "ещё\n") + append(root / "чаты" / "чат.md", "ignored\n") + for path in (diary, root / "люди" / "Прохор.md", root / "чаты" / "чат.md"): + watch.note(path) + changes = {c.path: c for c in watch.take()} + assert set(changes) == {f"дни/{today}.md", "люди/Прохор.md"} + assert changes[f"дни/{today}.md"].added == ["- 12:40 вышел", "- 13:00 кофе"] + assert changes[f"дни/{today}.md"].full + assert changes["люди/Прохор.md"].added_count == 1 + assert changes["люди/Прохор.md"].added == [] + assert abs( + changes[f"дни/{today}.md"].mtime + - datetime.fromtimestamp(diary.stat().st_mtime, tz=UTC) + ) < timedelta(seconds=1) + assert watch.take() == [] + append(diary, "- 14:00 снова\n") + watch.note(diary) + (only,) = watch.take() + assert only.added == ["- 14:00 снова"] + + +def test_envelope_respects_ceilings_and_names_only_window() -> None: + root, watch = vault() + today = watch.today() + diary = root / "дни" / f"{today}.md" + append(diary, "".join(f"- строка {i}\n" for i in range(200))) + for name in ("Петя", "Маша"): + (root / "люди" / f"{name}.md").write_text("новый\n" * 50) + watch.note(root / "люди" / f"{name}.md") + watch.note(diary) + envelope = Envelope(watch=watch, tz="Europe/Warsaw") + text = envelope.build(streak=7, injects=2) + lines = text.splitlines() + assert lines[0] == HEADER + assert "(Warsaw)" in lines[1] + assert lines[2] == "мастер: 7 реплик подряд без ветки" + assert f"дни/{today}.md (+200)" in lines[3] + assert "люди/Петя.md (+50)" in lines[3] + assert "инжекты со старта: (2) ниже" in lines[4] + assert sum(1 for line in lines if line.startswith("+ ")) == 31 + assert "+ … ещё 170" in lines + assert len(lines) <= 120 + append(diary, "- ещё одна\n") + watch.note(diary) + second = envelope.build(streak=8, injects=0) + assert f"дни/{today}.md (+1)" in second + assert "+ - ещё одна" not in second + assert "инжекты" not in second + + +def test_render_hits_total_ceiling() -> None: + now = datetime(2026, 8, 26, 11, 4, tzinfo=UTC) + changes = [ + Change( + path=f"дни/{i}.md", + mtime=now, + added=[f"l{j}" for j in range(30)], + added_count=30, + full=True, + ) + for i in range(10) + ] + text = render( + now=now, + tz="Europe/Warsaw", + streak=1, + changes=changes, + since=now - timedelta(hours=1), + names_only=False, + injects=0, + ) + lines = text.splitlines() + assert len(lines) <= 120 + assert "… (потолок конверта)" in lines + assert lines[1] == "время: 2026-08-26 13:04 (Warsaw)" + assert lines[2] == "мастер: 1 реплика подряд без ветки" + + +async def test_master_turn_gets_envelope_after_text_and_before_injects( + world: World, +) -> None: + root, watch = vault() + world.conversations._envelope = Envelope(watch=watch, tz="UTC") # noqa: SLF001 + master = await world.conversations.create(kind="master", agent="a", origin="test") + append(root / "люди" / "Прохор.md", "новое\n") + watch.note(root / "люди" / "Прохор.md") + await world.conversations.inject(master, "tick", urgency="normal", origin="крон") + await world.conversations.post(master, "hello") + await world.settle(master, 2) + prompt = ScriptedClient.instances[0].prompts[0] + head, _, rest = prompt.partition("\n\n") + assert head == "hello" + assert rest.startswith(HEADER) + assert "люди/Прохор.md (+1)" in rest + assert "мастер: 0 реплик подряд без ветки" in rest + assert rest.index("[инжекты") > rest.index(HEADER) + assert (await world.conversations.get(master.external_id)).flags["streak"] == 1 + branch = await world.conversations.spawn( + kind="branch", parent=master, seed="brief", text="do X" + ) + await world.settle(branch, 1) + assert (await world.conversations.get(master.external_id)).flags["streak"] == 0 + assert "[конверт" not in ScriptedClient.instances[-1].prompts[0] + await asyncio.sleep(0) diff --git a/tests/test_rotation.py b/tests/test_rotation.py new file mode 100644 index 0000000..9809fcd --- /dev/null +++ b/tests/test_rotation.py @@ -0,0 +1,174 @@ +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] == "транскрипт" diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 0000000..b9dd4e8 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,196 @@ +import asyncio +import os +import tempfile +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any + +import psycopg +import pytest +from httpx import ASGITransport, AsyncClient +from pgqueuer import PsycopgDriver, Queries +from test_conversations import World + +from beaver_gateway.core.scheduler import Budget, Job, JobRun, Scheduler +from beaver_gateway.storage.models import RateLimit + +DATABASE_URL = os.environ.get("TEST_DATABASE_URL") + +pytestmark = pytest.mark.skipif( + not DATABASE_URL, reason="TEST_DATABASE_URL (postgres) is not set" +) + + +class Pg: + def __init__(self) -> None: + self.connections: list[psycopg.AsyncConnection] = [] + + async def driver(self) -> PsycopgDriver: + conn = await psycopg.AsyncConnection.connect(str(DATABASE_URL), autocommit=True) + self.connections.append(conn) + return PsycopgDriver(conn) + + async def close(self) -> None: + for conn in self.connections: + await conn.close() + + +@pytest.fixture +async def pg() -> Pg: + handle = Pg() + queries = Queries(await handle.driver()) + if await queries.has_table("pgqueuer"): + await queries.uninstall() + await queries.install() + yield handle + await queries.uninstall() + await handle.close() + + +@pytest.fixture +async def world() -> World: + root = Path(tempfile.mkdtemp(prefix="beaver-sched-")) + w = await World(root).setup() + yield w + await w.conversations.stop() + await w.pool.close_all() + await w.db.dispose() + + +async def until(pred, timeout: float = 5.0) -> Any: + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + value = pred() + if value: + return value + await asyncio.sleep(0.05) + msg = "condition never happened" + raise AssertionError(msg) + + +async def test_schedule_survives_a_restart(world: World, pg: Pg) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + first = Scheduler(conversations=world.conversations, driver=await pg.driver()) + world.conversations.scheduler = first + await first.start() + job_id, when = await world.conversations.schedule(conv, "+1s", "push X") + assert job_id is not None + assert 0 < (when - datetime.now(UTC)).total_seconds() <= 1 + queued = await world.conversations.schedules(conv) + assert [q["payload"]["text"] for q in queued] == ["push X"] + await first.stop() + + second = Scheduler(conversations=world.conversations, driver=await pg.driver()) + world.conversations.scheduler = second + await second.start() + world.conversations._normal_window = 0.05 # noqa: SLF001 + await until(lambda: len(ScriptedClient_prompts(world)) == 1) + prompt = ScriptedClient_prompts(world)[0] + assert prompt.startswith("[инжект: schedule") + assert prompt.endswith("push X") + assert await world.conversations.schedules(conv) == [] + await second.stop() + + +def ScriptedClient_prompts(world: World) -> list[str]: # noqa: N802 + from test_conversations import ScriptedClient + + return [p for c in ScriptedClient.instances for p in c.prompts] + + +async def test_cancel_removes_a_pending_inject(world: World, pg: Pg) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + scheduler = Scheduler(conversations=world.conversations, driver=await pg.driver()) + world.conversations.scheduler = scheduler + await scheduler.start() + job_id, _ = await scheduler.schedule(conv, "+1h", "never") + assert job_id is not None + assert await scheduler.cancel(job_id) + assert not await scheduler.cancel(job_id) + await until(lambda: True) + assert await scheduler.scheduled(conv) == [] + await scheduler.stop() + + +async def test_webhook_runs_the_job_with_its_payload(world: World, pg: Pg) -> None: + seen: list[dict[str, Any]] = [] + + async def ping(run: JobRun) -> None: + seen.append({"trigger": run.trigger, **run.payload}) + + scheduler = Scheduler( + conversations=world.conversations, + jobs=[Job("ping", ping, webhook=True), Job("quiet", ping)], + driver=await pg.driver(), + ) + await scheduler.start() + + async def authorize(_request: Any) -> str: + return "test" + + transport = ASGITransport(app=scheduler.app(authorize)) + async with AsyncClient(transport=transport, base_url="http://hooks") as client: + accepted = await client.post("/ping", json={"stack": "x"}) + assert accepted.status_code == 202 + assert accepted.json()["name"] == "ping" + missing = await client.post("/quiet", json={}) + assert missing.status_code == 404 + await until(lambda: seen) + assert seen == [{"trigger": "webhook", "stack": "x"}] + snapshot = await scheduler.snapshot() + assert [j["name"] for j in snapshot["jobs"]] == ["ping", "quiet"] + assert snapshot["enabled"] and snapshot["jobs"][0]["webhook"] + await scheduler.stop() + + +async def test_non_critical_jobs_wait_while_the_window_is_hot( + world: World, pg: Pg +) -> None: + ran: list[str] = [] + + async def job(run: JobRun) -> None: + ran.append(run.job.name) + + async with world.db.session() as session: + session.add( + RateLimit( + window="five_hour", + status="allowed_warning", + utilization=0.9, + resets_at=datetime.now(UTC) + timedelta(hours=2), + ) + ) + await session.commit() + scheduler = Scheduler( + conversations=world.conversations, + jobs=[ + Job("vibegram", job, cron="*/10 * * * *", critical=False), + Job("rotation", job, cron="0 * * * *"), + ], + driver=await pg.driver(), + budget=Budget(threshold=0.7), + ) + await scheduler.start() + events: list[dict[str, Any]] = [] + + async def collect() -> None: + async for event in world.bus.stream(): + events.append(event) + + task = asyncio.create_task(collect()) + await scheduler._dispatch( # noqa: SLF001 + scheduler.job("vibegram"), trigger="cron", payload={} + ) + await scheduler._dispatch( # noqa: SLF001 + scheduler.job("rotation"), trigger="cron", payload={} + ) + await asyncio.sleep(0.05) + task.cancel() + assert ran == ["rotation"] + assert [e["job"] for e in events if e["type"] == "job.deferred"] == ["vibegram"] + snapshot = await scheduler.snapshot() + assert snapshot["throttled"] and snapshot["utilization"] == 0.9 + crons = {j["name"]: j["cron"] for j in snapshot["jobs"]} + assert crons == {"vibegram": "*/10 * * * *", "rotation": "0 * * * *"} + await until(lambda: all(j["next_run"] for j in snapshot["jobs"]) or True) + await scheduler.stop() diff --git a/ui/src/lib/api/client.ts b/ui/src/lib/api/client.ts index 236aad4..bbdc7ab 100644 --- a/ui/src/lib/api/client.ts +++ b/ui/src/lib/api/client.ts @@ -7,10 +7,10 @@ import type { ConversationSummary, EntriesPage, HistoryMessage, + JobsResponse, LimitsResponse, MemoryFile, MemoryTree, - Schedule, SessionsResponse, TokenRow, UsageGroup, @@ -275,8 +275,16 @@ export class ApiClient { return this.get("/api/sessions"); } - schedules(): Promise<{ schedules: Schedule[] }> { - return this.get("/api/schedules"); + jobs(): Promise { + return this.get("/api/jobs"); + } + + runJob(name: string): Promise<{ job: number | null }> { + return this.post(`/api/jobs/${encodeURIComponent(name)}/run`); + } + + cancelJob(id: number): Promise<{ cancelled: number }> { + return this.request("DELETE", `/api/jobs/queue/${id}`); } usage(params: { diff --git a/ui/src/lib/api/types.ts b/ui/src/lib/api/types.ts index 9d92c11..4de9cce 100644 --- a/ui/src/lib/api/types.ts +++ b/ui/src/lib/api/types.ts @@ -260,13 +260,35 @@ export interface AuditPage { records: AuditRecord[]; } -export interface Schedule { - conversation_row: number; - created_at: string; - delivered_at: string | null; - execute_at: string; +export interface QueuedJob { + attempts: number; + created: string | null; + entrypoint: string; + execute_after: string | null; id: number; - text: string; + payload: Record; + status: string; +} + +export interface JobInfo { + critical: boolean; + cron: string | null; + events: string[]; + last_run: string | null; + name: string; + next_run: string | null; + run: { trigger: string; started_at: string; status: string } | null; + status: string | null; + webhook: boolean; +} + +export interface JobsResponse { + enabled: boolean; + jobs: JobInfo[]; + queue: QueuedJob[]; + threshold: number; + throttled: boolean; + utilization: number | null; } export interface TurnUsage { diff --git a/ui/src/routes/jobs/+page.svelte b/ui/src/routes/jobs/+page.svelte index aeb7f8d..94b6b83 100644 --- a/ui/src/routes/jobs/+page.svelte +++ b/ui/src/routes/jobs/+page.svelte @@ -1,16 +1,23 @@ Jobs · Beaver
- {#if failure} - {:else if schedules === null} + {:else if data === null} - {:else if schedules.length === 0} -

No deferred injects.

{:else} + {#if !data.enabled} + + {/if} + {#if data.throttled} +

+ Subscription window at {fmtPct(data.utilization)} (threshold + {fmtPct( + data.threshold + )}): non-critical jobs are deferred. +

+ {/if} +

- Deferred injects + Jobs

-
    - {#each [...pending, ...delivered] as s (s.id)} -
  • - - {s.delivered_at ? "delivered" : fmtRelative(s.execute_at)} - - {clip(s.text, 160)} - - {fmtDateTime(s.execute_at)} - -
  • - {/each} -
+ {#if data.jobs.length === 0} +

No jobs in config.

+ {:else} +
    + {#each data.jobs as job (job.name)} +
  • + + {job.name} + {#if !job.critical} + · soft + {/if} + + + {triggers(job)} + + + {job.next_run ? `next ${fmtRelative(job.next_run)}` : ""} + + + {#if job.run} + + + {fmtRelative(job.run.started_at)} + · {job.run.trigger} + + {:else if job.last_run} + + ran {fmtRelative(job.last_run)} + + {/if} + + +
  • + {/each} +
+ {/if} +
+ +
+

+ Queue +

+ {#if data.queue.length === 0} +

+ Nothing queued: no deferred injects, no pending webhooks. +

+ {:else} +
    + {#each [...picked, ...pending] as job (job.id)} +
  • + + {job.status === "queued" + ? fmtRelative(job.execute_after) + : job.status} + + {job.entrypoint} + + {clip(describe(job), PAYLOAD_MAX)} + + + {fmtDateTime(job.execute_after)} + + {#if job.status === "queued"} + + {:else} + + {/if} +
  • + {/each} +
+ {/if}
{/if}
diff --git a/uv.lock b/uv.lock index 8414421..cf1c21a 100644 --- a/uv.lock +++ b/uv.lock @@ -295,10 +295,12 @@ dependencies = [ { name = "anyio" }, { name = "argon2-cffi" }, { name = "claude-agent-sdk" }, + { name = "croniter" }, { name = "fastapi" }, { name = "fastmcp" }, { name = "greenlet" }, { name = "itsdangerous" }, + { name = "pgqueuer" }, { name = "psutil" }, { name = "psycopg", extra = ["binary"] }, { name = "pydantic" }, @@ -307,6 +309,7 @@ dependencies = [ { name = "sqlmodel" }, { name = "uvicorn", extra = ["standard"] }, { name = "uvloop" }, + { name = "watchfiles" }, ] [package.optional-dependencies] @@ -336,10 +339,12 @@ requires-dist = [ { name = "anyio", specifier = ">=4.13.0" }, { name = "argon2-cffi", specifier = ">=25.1.0" }, { name = "claude-agent-sdk", specifier = ">=0.2.146" }, + { name = "croniter", specifier = ">=6.2.4" }, { name = "fastapi", specifier = ">=0.136.1" }, { name = "fastmcp", specifier = ">=3.3.1" }, { name = "greenlet", specifier = ">=3.5.0" }, { name = "itsdangerous", specifier = ">=2.2.0" }, + { name = "pgqueuer", specifier = ">=1.3.2" }, { name = "psutil", specifier = ">=7.2.2" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.3.4" }, { name = "pydantic", specifier = ">=2.13.4" }, @@ -350,6 +355,7 @@ requires-dist = [ { name = "sqlmodel", specifier = ">=0.0.38" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.47.0" }, { name = "uvloop", specifier = ">=0.22.1" }, + { name = "watchfiles", specifier = ">=1.2.0" }, ] provides-extras = ["local", "prod"] @@ -482,6 +488,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "croniter" +version = "6.2.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/37/57/2e2a65aee2a70483cb28e2b7e15a072d00a523207593b44400d4717bb100/croniter-6.2.4.tar.gz", hash = "sha256:fc124f751b1b04805c2a04b061898b436b45ab2320b045e1e052ea895de65189", size = 166267, upload-time = "2026-07-10T09:52:59.955Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cd/ba/d678e5bd329646ca51d3c92addbc77804e86d21f4b6b6a027218e6abb010/croniter-6.2.4-py3-none-any.whl", hash = "sha256:8ef3d544107a5c05a150a2d78f8bf5a8eb9c5c4d93405a736b824109574e3f4d", size = 46677, upload-time = "2026-07-10T09:52:58.425Z" }, +] + [[package]] name = "cryptography" version = "48.0.0" @@ -1264,6 +1282,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/96/5a770e5c461462575474468e5af931cff9de036e7c2b4fea23c1c58d2cbe/pathable-0.5.0-py3-none-any.whl", hash = "sha256:646e3d09491a6351a0c82632a09c02cdf70a252e73196b36d8a15ba0a114f0a6", size = 16867, upload-time = "2026-02-20T08:46:59.536Z" }, ] +[[package]] +name = "pgqueuer" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "croniter" }, + { name = "pydantic" }, + { name = "pydantic-settings" }, + { name = "tabulate" }, + { name = "typer" }, + { name = "uvloop", marker = "sys_platform != 'win32' or (extra == 'extra-14-beaver-gateway-local' and extra == 'extra-14-beaver-gateway-prod')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/33/10/6af18eede8a31c6d016564fd815d0cf1aaa69b8e584108a604b88f1ea56c/pgqueuer-1.3.2.tar.gz", hash = "sha256:11f4c67d9dfc343c02d8bb2e362733a0d44272f6c7e846630fe198f688fa3ed1", size = 505605, upload-time = "2026-07-27T16:58:00.864Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/16/1e501ee36fb8e6bdde6a3a4dd9b10d1b45c548e02f9c1f11e9ffc62ee922/pgqueuer-1.3.2-py3-none-any.whl", hash = "sha256:0823fcb089bb3ffd0580518e91b53888c9be742b42e706f9860b6a533f32bc44", size = 146747, upload-time = "2026-07-27T16:57:59.153Z" }, +] + [[package]] name = "platformdirs" version = "4.9.6" @@ -1617,6 +1653,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" }, ] +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + [[package]] name = "python-dotenv" version = "1.2.2" @@ -1887,6 +1935,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl", hash = "sha256:0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137", size = 15554, upload-time = "2025-11-23T19:02:51.545Z" }, ] +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + [[package]] name = "sniffio" version = "1.3.1" @@ -1974,6 +2040,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0b/c9/584bc9651441b4ba60cc4d557d8a547b5aff901af35bda3a4ee30c819b82/starlette-1.0.0-py3-none-any.whl", hash = "sha256:d3ec55e0bb321692d275455ddfd3df75fff145d009685eb40dc91fc66b03d38b", size = 72651, upload-time = "2026-03-22T18:29:45.111Z" }, ] +[[package]] +name = "tabulate" +version = "0.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/58/8c37dea7bbf769b20d58e7ace7e5edfe65b849442b00ffcdd56be88697c6/tabulate-0.10.0.tar.gz", hash = "sha256:e2cfde8f79420f6deeffdeda9aaec3b6bc5abce947655d17ac662b126e48a60d", size = 91754, upload-time = "2026-03-04T18:55:34.402Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/55/db07de81b5c630da5cbf5c7df646580ca26dfaefa593667fc6f2fe016d2e/tabulate-0.10.0-py3-none-any.whl", hash = "sha256:f0b0622e567335c8fabaaa659f1b33bcb6ddfe2e496071b743aa113f8774f2d3", size = 39814, upload-time = "2026-03-04T18:55:31.284Z" }, +] + [[package]] name = "ty" version = "0.0.37" @@ -1999,6 +2074,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1a/ed/5ec4b501479bc5dad55467e2fe72e797cb9c178468c0d1a514536872ebc5/ty-0.0.37-py3-none-win_arm64.whl", hash = "sha256:6c3c2b997f68c71e14242b96d48cba3c086439556af02bb4613aa458950d5c23", size = 10958817, upload-time = "2026-05-16T05:57:08.907Z" }, ] +[[package]] +name = "typer" +version = "0.27.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-14-beaver-gateway-local' and extra == 'extra-14-beaver-gateway-prod')" }, + { name = "rich" }, + { name = "shellingham" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/16/f7/57713ba479fd405eb76de31404b2c744c289e336b2d999511ebf51e496f7/typer-0.27.2.tar.gz", hash = "sha256:269b7eb9d3c202ca84b4bc9618cb04ebb43d3d4d1e567e4c768607232c05f945", size = 204045, upload-time = "2026-08-28T10:26:55.046Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/bf/205d0004930ede8f542fb58f601526fccf4ae7626075ca1e6c4de5d3d652/typer-0.27.2-py3-none-any.whl", hash = "sha256:b3a5fc4342d5fc8fda8fc3010b1cf117e9249aab7fae800c2eff62fd3842d97d", size = 123130, upload-time = "2026-08-28T10:26:53.752Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0"