From e3074c266a2ddee764dea82e795183636fc1d22a Mon Sep 17 00:00:00 2001 From: h Date: Fri, 28 Aug 2026 03:08:30 +0200 Subject: [PATCH] feat(core,backends,frontends,storage): conversations, inject queue, session pool, gateway mcp tools, api frontend --- pyproject.toml | 1 + src/beaver_gateway/agents/claude.py | 20 +- src/beaver_gateway/backends/claude_sdk.py | 301 +++-- src/beaver_gateway/cli.py | 48 +- src/beaver_gateway/config_loader.py | 2 + src/beaver_gateway/core/auth.py | 3 +- src/beaver_gateway/core/bus.py | 69 ++ src/beaver_gateway/core/conversation_store.py | 6 +- src/beaver_gateway/core/conversations.py | 1024 +++++++++++++++++ src/beaver_gateway/core/gateway_tools.py | 169 +++ src/beaver_gateway/core/injects.py | 130 +++ src/beaver_gateway/core/registry.py | 3 + src/beaver_gateway/core/sessions.py | 257 +++++ src/beaver_gateway/core/transcript.py | 230 +++- src/beaver_gateway/core/turn_capture.py | 2 + src/beaver_gateway/frontends/_sse.py | 63 + src/beaver_gateway/frontends/api/__init__.py | 5 + src/beaver_gateway/frontends/api/frontend.py | 491 ++++++++ src/beaver_gateway/frontends/base.py | 8 +- .../frontends/markdown/frontend.py | 278 ++--- .../frontends/markdown/renderer.py | 24 +- src/beaver_gateway/storage/__init__.py | 15 +- src/beaver_gateway/storage/db.py | 26 +- src/beaver_gateway/storage/models.py | 106 +- tests/test_claude_sdk_backend.py | 3 + tests/test_conversations.py | 546 +++++++++ tests/test_db_migrate.py | 28 + uv.lock | 30 + 28 files changed, 3543 insertions(+), 345 deletions(-) create mode 100644 src/beaver_gateway/core/bus.py create mode 100644 src/beaver_gateway/core/conversations.py create mode 100644 src/beaver_gateway/core/gateway_tools.py create mode 100644 src/beaver_gateway/core/injects.py create mode 100644 src/beaver_gateway/core/sessions.py create mode 100644 src/beaver_gateway/frontends/_sse.py create mode 100644 src/beaver_gateway/frontends/api/__init__.py create mode 100644 src/beaver_gateway/frontends/api/frontend.py create mode 100644 tests/test_conversations.py diff --git a/pyproject.toml b/pyproject.toml index a3c2129..f77068d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "greenlet>=3.5.0", "itsdangerous>=2.2.0", "jinja2>=3.1.6", + "psutil>=7.2.2", "psycopg[binary]>=3.3.4", "pydantic>=2.13.4", "pydantic-settings>=2.14.1", diff --git a/src/beaver_gateway/agents/claude.py b/src/beaver_gateway/agents/claude.py index 38f969f..9604a46 100644 --- a/src/beaver_gateway/agents/claude.py +++ b/src/beaver_gateway/agents/claude.py @@ -42,15 +42,27 @@ class ClaudeOptions(BaseModel): """Extra inherited variable names to let through the env whitelist.""" max_turns: int | None = None - idle_session_ttl: float = 1800.0 - """Seconds a live session may sit unused before it is closed.""" - - session_store_flush: str = "batched" + session_store_flush: str = "eager" + """``eager`` mirrors every transcript frame as it lands, so a gateway + killed mid-turn loses at most the frame in flight.""" class ClaudeAgent(BaseAgent): cwd: Path system_prompt: str = "" prompt_sources: tuple[PromptSource, ...] = () + prompt_sources_by_kind: Mapping[str, tuple[PromptSource, ...]] = Field( + default_factory=dict + ) + """Per conversation kind (``master``/``branch``/``deep``/``job``/``fork``) + assembly; falls back to ``prompt_sources``. Constant per kind (§3.12).""" + skill_sets: tuple[Path, ...] = () + gateway_tools: tuple[str, ...] = () + """Gateway tools exposed in-process (``read_conversation``, ``spawn``, + ``say``, ``schedule``, ``inject``); empty = no gateway MCP server.""" + options: ClaudeOptions = Field(default_factory=ClaudeOptions) + + def prompt_for(self, kind: str) -> tuple[PromptSource, ...]: + return self.prompt_sources_by_kind.get(kind, self.prompt_sources) diff --git a/src/beaver_gateway/backends/claude_sdk.py b/src/beaver_gateway/backends/claude_sdk.py index 4589dcf..97227c1 100644 --- a/src/beaver_gateway/backends/claude_sdk.py +++ b/src/beaver_gateway/backends/claude_sdk.py @@ -2,11 +2,19 @@ One :class:`ClaudeSdkBackend` per :class:`ClaudeAgent`. A live session is one ``ClaudeSDKClient`` (one claude subprocess) and runs one turn at a -time. Sessions are keyed by ``conversation_id`` when the frontend passes -one (markdown chats) or by a text-only fingerprint of ``messages[:-1]`` +time; the sessions of every agent live in one shared +:class:`~beaver_gateway.core.sessions.SessionPool` that owns TTL and +memory-pressure eviction. Sessions are keyed by ``conversation_id`` when +the caller passes one or by a text-only fingerprint of ``messages[:-1]`` for stateless callers (``/v1/messages``). Without a live session the -adapter resumes ``session_id`` from the session store, or seeds the -incoming history into the store via ``core/transcript`` and resumes that. +adapter resumes ``session_id`` from the session store (after closing any +``tool_use`` left open by a crash), or seeds the incoming history into the +store via ``core/transcript`` and resumes that. + +Per-turn ``**options`` beyond the protocol's: ``kind`` (conversation kind, +picks the prompt assembly and the pool TTL), ``pinned`` (never evicted), +``tools=False`` (no MCP at all - forks and jobs), ``observer`` (callback +receiving every raw SDK message, subagent ones included), ``turn_id``. Events on the wire are the Anthropic ``MessageStreamEvent`` family: one ``message_start``/``message_stop`` envelope per turn, block indices @@ -36,13 +44,14 @@ import uuid from collections.abc import Mapping from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING, Any, Protocol, Self, cast +from typing import TYPE_CHECKING, Any, Self, cast import claude_agent_sdk from claude_agent_sdk import ( AssistantMessage, ClaudeAgentOptions, ClaudeSDKClient, + MirrorErrorMessage, ResultMessage, StreamEvent, TextBlock, @@ -68,14 +77,15 @@ from beaver_gateway.core.events import ( build_thinking_delta, build_tool_use_block_start, ) -from beaver_gateway.core.transcript import build_entries +from beaver_gateway.core.sessions import Session, SessionClient, SessionPool +from beaver_gateway.core.transcript import build_entries, close_open_tool_uses from beaver_gateway.core.turn_capture import TurnCapture, TurnUsage if TYPE_CHECKING: from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence from anthropic.types import MessageParam - from claude_agent_sdk import SessionStore + from claude_agent_sdk import McpSdkServerConfig, SessionStore from beaver_gateway.agents.base import BaseAgent from beaver_gateway.agents.claude import ClaudeAgent @@ -88,6 +98,7 @@ __all__ = [ "ClaudeSdkBackend", "RunnerConfig", "SessionClient", + "ToolServerFactory", "UsageSink", "fingerprint", ] @@ -120,7 +131,6 @@ ENV_KEEP: tuple[str, ...] = ( ) ENV_KEEP_PREFIXES: tuple[str, ...] = ("CLAUDE_", "ANTHROPIC_", "DISABLE_") -_REAP_INTERVAL = 60.0 _STOP_REASONS: dict[str, StopReason] = { "end_turn": "end_turn", "tool_use": "tool_use", @@ -131,15 +141,10 @@ _STOP_REASONS: dict[str, StopReason] = { } -class SessionClient(Protocol): - async def connect(self) -> None: ... - async def query(self, prompt: str) -> None: ... - def receive_response(self) -> AsyncIterator[Any]: ... - async def disconnect(self) -> None: ... - - ClientFactory = "Callable[[ClaudeAgentOptions], SessionClient]" UsageSink = "Callable[[UsageEvent], Awaitable[None]]" +ToolServerFactory = "Callable[[str, str], McpSdkServerConfig | None]" +"""``(conversation_key, kind) -> in-process MCP server config`` or ``None``.""" @dataclass(frozen=True, slots=True) @@ -158,16 +163,6 @@ class UsageEvent: usage: TurnUsage -@dataclass -class _Live: - client: SessionClient - session_id: str | None - resumed: bool - lock: asyncio.Lock = field(default_factory=asyncio.Lock) - last_used: float = field(default_factory=time.monotonic) - turns: int = 0 - - class _RunnerClient(ClaudeSDKClient): """``ClaudeSDKClient`` that hands the materialized resume dir to the runner uid.""" @@ -194,6 +189,8 @@ class ClaudeSdkBackend: usage_sink: Callable[[UsageEvent], Awaitable[None]] | None = None, client_factory: Callable[[ClaudeAgentOptions], SessionClient] | None = None, work_dir: Path | None = None, + pool: SessionPool | None = None, + tool_server: Callable[[str, str], McpSdkServerConfig | None] | None = None, ) -> None: self._agent = agent self._store = session_store @@ -203,8 +200,8 @@ class ClaudeSdkBackend: self._work_dir = work_dir or Path(tempfile.gettempdir()) / "beaver-claude" self._servers = _mcp_servers(agent, mcp_internal_urls) self._mcp_disallowed = _mcp_disallowed(agent, mcp_tool_names or {}) - self._sessions: dict[str, _Live] = {} - self._reaper: asyncio.Task[None] | None = None + self._pool = pool if pool is not None else SessionPool() + self._tool_server = tool_server self._uid, self._gid = _resolve_ids(self._runner.user) self._wrapper: Path | None = None @@ -212,34 +209,56 @@ class ClaudeSdkBackend: def agent(self) -> ClaudeAgent: return self._agent + @property + def pool(self) -> SessionPool: + return self._pool + @property def sessions(self) -> dict[str, dict[str, Any]]: - now = time.monotonic() return { - key: { - "session_id": live.session_id, - "idle_seconds": now - live.last_used, - "turns": live.turns, - "busy": live.lock.locked(), - } - for key, live in self._sessions.items() + row["key"]: row + for row in self._pool.snapshot() + if row["agent"] == self._agent.name } async def __aenter__(self) -> Self: - self._reaper = asyncio.create_task(self._reap_loop()) return self async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: await self.aclose() async def aclose(self) -> None: - if self._reaper is not None: - self._reaper.cancel() - with contextlib.suppress(BaseException): - await self._reaper - self._reaper = None - for key in list(self._sessions): - await self._close(key) + await self._pool.close_all(agent=self._agent.name) + + async def close(self, key: str) -> None: + await self._pool.close(key) + + async def interrupt(self, key: str) -> bool: + live = self._pool.get(key) + if live is None or not live.busy: + return False + live.interrupt_requested = True + await live.client.interrupt() + return True + + def live(self, key: str) -> Session | None: + return self._pool.get(key) + + async def repair_session(self, session_id: str) -> int: + """Close ``tool_use`` blocks a crash left without a result; count added.""" + key = self._store_key(session_id) + entries = await self._store.load(cast("Any", key)) + if not entries: + return 0 + fixes = close_open_tool_uses(cast("list[Mapping[str, Any]]", entries)) + if fixes: + await self._store.append(cast("Any", key), cast("Any", fixes)) + _log.warning( + "session %s: closed %d open tool_use with synthetic results", + session_id, + len(fixes), + ) + return len(fixes) async def complete( self, @@ -250,6 +269,11 @@ class ClaudeSdkBackend: conversation_id: str | None = None, session_id: str | None = None, capture: TurnCapture | None = None, + kind: str = "deep", + pinned: bool = False, + tools: bool = True, + observer: Callable[[Any], None] | None = None, + turn_id: str | None = None, **options: Any, # noqa: ARG002 - per-request knobs are not supported ) -> AsyncIterator[MessageStreamEvent]: if agent.name != self._agent.name: @@ -262,37 +286,70 @@ class ClaudeSdkBackend: prompt = _prompt_text(history[-1].get("content")) prior = history[:-1] key = conversation_id or fingerprint(prior) - live = await self._acquire(key, session_id=session_id, history=prior) + spec = _SessionSpec(kind=kind, pinned=pinned, tools=tools) + live = await self._acquire(key, session_id=session_id, history=prior, spec=spec) message_id = f"msg_{uuid.uuid4().hex}" yield build_message_start(message_id=message_id, model=self._agent.model) async with live.lock: + live.running_turn = turn_id or message_id live.last_used = time.monotonic() try: - turn = await self._run_turn(live, prompt) + turn = await self._run_turn(live, prompt, observer=observer) except Exception: - if live.resumed and live.turns == 0: - _log.exception( - "resume of %s failed, reseeding from history", live.session_id - ) - await self._close(key) - live = await self._acquire(key, session_id=None, history=prior) - async with live.lock: - turn = await self._run_turn(live, prompt) - else: + live.running_turn = None + await self._pool.close(key) + if not (live.resumed and live.turns == 0): raise + _log.exception( + "resume of %s failed, reseeding from history", live.session_id + ) + live = await self._acquire( + key, session_id=None, history=prior, spec=spec + ) + async with live.lock: + live.running_turn = turn_id or message_id + turn = await self._run_turn(live, prompt, observer=observer) for event in turn.events: yield event live.turns += 1 live.last_used = time.monotonic() + live.running_turn = None + usage, interrupted = await self._after_turn( + live, + turn, + conversation_id=conversation_id, + history=history, + capture=capture, + ) + if turn.result is not None and turn.result.is_error and not interrupted: + msg = f"claude: {turn.result.result or turn.result.subtype}" + raise RuntimeError(msg) + yield build_message_delta( + stop_reason=turn.stop_reason, usage=_wire_usage(usage) + ) + yield build_message_stop() + + async def _after_turn( + self, + live: Session, + turn: _Turn, + *, + conversation_id: str | None, + history: list[dict[str, Any]], + capture: TurnCapture | None, + ) -> tuple[TurnUsage, bool]: if turn.result is not None and turn.result.session_id: live.session_id = turn.result.session_id if conversation_id is None: - self._rekey(key, fingerprint([*history, *turn.synthesized])) + self._rekey(live.key, fingerprint([*history, *turn.synthesized])) usage = _usage_of(turn.result) + interrupted = live.interrupt_requested + live.interrupt_requested = False if capture is not None: capture.synthesized_messages = turn.synthesized capture.session_id = live.session_id capture.usage = usage + capture.interrupted = interrupted if self._usage_sink is not None: await self._usage_sink( UsageEvent( @@ -304,15 +361,15 @@ class ClaudeSdkBackend: usage=usage, ) ) - if turn.result is not None and turn.result.is_error: - msg = f"claude: {turn.result.result or turn.result.subtype}" - raise RuntimeError(msg) - yield build_message_delta( - stop_reason=turn.stop_reason, usage=_wire_usage(usage) - ) - yield build_message_stop() + return usage, interrupted - async def _run_turn(self, live: _Live, prompt: str) -> _Turn: + async def _run_turn( + self, + live: Session, + prompt: str, + *, + observer: Callable[[Any], None] | None = None, + ) -> _Turn: streaming = self._agent.options.include_partial_messages turn = _Turn() raw: list[Any] = [] @@ -320,6 +377,16 @@ class ClaudeSdkBackend: offset = 0 await live.client.query(prompt) async for message in live.client.receive_response(): + if observer is not None: + observer(message) + if isinstance(message, MirrorErrorMessage): + live.dirty = True + _log.error( + "session %s: mirror error, marked dirty: %s", + live.session_id, + message.error, + ) + continue if getattr(message, "parent_tool_use_id", None) is not None: continue if isinstance(message, StreamEvent): @@ -357,17 +424,30 @@ class ClaudeSdkBackend: return turn async def _acquire( - self, key: str, *, session_id: str | None, history: list[dict[str, Any]] - ) -> _Live: - live = self._sessions.get(key) + self, + key: str, + *, + session_id: str | None, + history: list[dict[str, Any]], + spec: _SessionSpec, + ) -> Session: + live = self._pool.get(key) if live is not None: return live resume = session_id - if resume is None and history: + if resume is not None: + await self.repair_session(resume) + elif history: resume = await self._seed(history) - live = await self._spawn(resume) - self._sessions[key] = live - return live + await self._pool.make_room() + live = await self._spawn(resume, key=key, spec=spec) + return self._pool.add(live) + + def _store_key(self, session_id: str) -> dict[str, str]: + return { + "project_key": project_key_for_directory(str(self._agent.cwd)), + "session_id": session_id, + } async def _seed(self, history: list[dict[str, Any]]) -> str: session_id = str(uuid.uuid4()) @@ -378,11 +458,9 @@ class ClaudeSdkBackend: model=self._agent.model, permission_mode=self._agent.options.permission_mode, ) - key = { - "project_key": project_key_for_directory(str(self._agent.cwd)), - "session_id": session_id, - } - await self._store.append(cast("Any", key), cast("Any", entries)) + await self._store.append( + cast("Any", self._store_key(session_id)), cast("Any", entries) + ) _log.info( "seeded session %s with %d entries from %d messages", session_id, @@ -391,22 +469,36 @@ class ClaudeSdkBackend: ) return session_id - async def _spawn(self, resume: str | None) -> _Live: - options = self._build_options(resume) + async def _spawn( + self, resume: str | None, *, key: str, spec: _SessionSpec + ) -> Session: + options = self._build_options(resume, key=key, spec=spec) client = self._factory(options) await client.connect() _log.info( - "spawned claude: agent=%s resume=%s user=%s", + "spawned claude: agent=%s kind=%s resume=%s tools=%s user=%s", self._agent.name, + spec.kind, resume, + spec.tools, self._runner.user, ) - return _Live(client=client, session_id=resume, resumed=resume is not None) + return Session( + key=key, + agent=self._agent.name, + kind=spec.kind, + client=client, + session_id=resume, + resumed=resume is not None, + pinned=spec.pinned, + ) def _default_factory(self, options: ClaudeAgentOptions) -> SessionClient: return _RunnerClient(options, uid=self._uid) - def _build_options(self, resume: str | None) -> ClaudeAgentOptions: + def _build_options( + self, resume: str | None, *, key: str, spec: _SessionSpec + ) -> ClaudeAgentOptions: agent = self._agent opt = agent.options env = dict(opt.env) @@ -414,18 +506,25 @@ class ClaudeSdkBackend: env["HOME"] = str(self._runner.home) env.setdefault("CLAUDE_CONFIG_DIR", str(self._runner.home / ".claude")) plugins = self._plugins() + sources = agent.prompt_for(spec.kind) system_prompt = ( - prompt_assembly.assemble(agent.prompt_sources) - if agent.prompt_sources - else agent.system_prompt + prompt_assembly.assemble(sources) if sources else agent.system_prompt ) + servers: dict[str, Any] = dict(self._servers) if spec.tools else {} + gateway = ( + self._tool_server(key, spec.kind) + if spec.tools and self._tool_server is not None and agent.gateway_tools + else None + ) + if gateway is not None: + servers[str(gateway["name"])] = gateway return ClaudeAgentOptions( model=agent.model or None, effort=cast("Any", opt.effort), system_prompt=system_prompt, setting_sources=[], strict_mcp_config=True, - mcp_servers=cast("Any", self._servers), + mcp_servers=cast("Any", servers), permission_mode=cast("Any", opt.permission_mode), tools=list(opt.tools) if opt.tools is not None else None, disallowed_tools=[*opt.disallowed_tools, *self._mcp_disallowed], @@ -473,8 +572,8 @@ class ClaudeSdkBackend: target=json.dumps(target), keep=json.dumps(keep), prefixes=json.dumps(list(ENV_KEEP_PREFIXES)), - uid=json.dumps(self._uid), - gid=json.dumps(self._gid), + uid=repr(self._uid), + gid=repr(self._gid), ) digest = hashlib.sha256(script.encode("utf-8")).hexdigest()[:12] path = self._work_dir / f"claude-exec-{digest}.py" @@ -487,30 +586,16 @@ class ClaudeSdkBackend: return path def _rekey(self, old: str, new: str) -> None: - live = self._sessions.pop(old, None) - if live is None: - return - stale = self._sessions.pop(new, None) - self._sessions[new] = live - if stale is not None and stale is not live: + stale = self._pool.rekey(old, new) + if stale is not None: asyncio.get_running_loop().create_task(_disconnect(stale)) - async def _close(self, key: str) -> None: - live = self._sessions.pop(key, None) - if live is not None: - await _disconnect(live) - async def _reap_loop(self) -> None: - ttl = self._agent.options.idle_session_ttl - while True: - await asyncio.sleep(_REAP_INTERVAL) - if ttl <= 0: - continue - now = time.monotonic() - for key, live in list(self._sessions.items()): - if not live.lock.locked() and now - live.last_used > ttl: - _log.info("closing idle session %s (%s)", live.session_id, key) - await self._close(key) +@dataclass(frozen=True, slots=True) +class _SessionSpec: + kind: str + pinned: bool + tools: bool @dataclass @@ -541,7 +626,7 @@ os.execve(TARGET, [TARGET, *sys.argv[1:]], env) """ -async def _disconnect(live: _Live) -> None: +async def _disconnect(live: Session) -> None: try: await live.client.disconnect() except Exception: # noqa: BLE001 diff --git a/src/beaver_gateway/cli.py b/src/beaver_gateway/cli.py index 6010293..e89ef24 100644 --- a/src/beaver_gateway/cli.py +++ b/src/beaver_gateway/cli.py @@ -22,6 +22,7 @@ from __future__ import annotations import asyncio import contextlib +import functools import logging import signal from contextlib import AsyncExitStack @@ -43,13 +44,18 @@ from beaver_gateway.backends.claude_sdk import ( ) 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.gateway_tools import build_tool_server from beaver_gateway.core.registry import AgentRegistry, McpRegistry +from beaver_gateway.core.sessions import SessionPool from beaver_gateway.frontends.base import GatewayRuntime from beaver_gateway.mcp.internal_app import build_internal_app from beaver_gateway.settings import Settings from beaver_gateway.storage import Database, PostgresSessionStore, Usage, append_usage if TYPE_CHECKING: + from claude_agent_sdk import McpSdkServerConfig from fastmcp import FastMCP from fastmcp.tools.base import Tool as FastMCPTool from starlette.applications import Starlette @@ -140,15 +146,32 @@ async def _async_main() -> None: # surfaces at startup instead of mid-conversation. mcp_tools = await _prefetch_mcp_tools(mcp_servers) + pool = SessionPool() + bus = EventBus() + late = _LateConversations() + session_store = PostgresSessionStore(db) backends: dict[str, Backend] = await _build_backends( settings=settings, agents=agents, stack=stack, db=db, + session_store=session_store, mcp_internal_urls=internal_urls, mcp_servers=mcp_servers, mcp_tools=mcp_tools, + pool=pool, + late=late, ) + conversations = Conversations( + db=db, + agents=agents, + backends=backends, + bus=bus, + pool=pool, + store=session_store, + texts=gateway.texts, + ) + late.conversations = conversations runtime = GatewayRuntime( agents=agents, @@ -161,6 +184,9 @@ async def _async_main() -> None: admin_pass=settings.admin_pass, session_secret=settings.session_secret, frontends=tuple(gateway.frontends), + conversations=conversations, + bus=bus, + pool=pool, ) for fe in gateway.frontends: @@ -186,7 +212,10 @@ async def _async_main() -> None: # it in this path and exit cleanly (Phase 0 DoD). return + await conversations.start() + stack.push_async_callback(conversations.stop) 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)) for fe in gateway.frontends: @@ -252,15 +281,31 @@ async def _serve_internal_mcp(app: Starlette, *, settings: Settings) -> None: await server.serve() +class _LateConversations: + """Backends need a tool-server factory before the service that backs it exists.""" + + conversations: Conversations | None = None + + def server( + self, key: str, _kind: str, names: tuple[str, ...] + ) -> McpSdkServerConfig | None: + if self.conversations is None or not names: + return None + return build_tool_server(self.conversations, conversation_key=key, names=names) + + async def _build_backends( *, settings: Settings, agents: AgentRegistry, stack: AsyncExitStack, db: Database, + session_store: PostgresSessionStore, mcp_internal_urls: dict[str, str], mcp_servers: dict[str, FastMCP], mcp_tools: dict[str, list[FastMCPTool]], + pool: SessionPool, + late: _LateConversations, ) -> dict[str, Backend]: """Construct one backend per agent name. @@ -285,7 +330,6 @@ async def _build_backends( for a in raycast_agents: backends[a.name] = raycast_backend - session_store = PostgresSessionStore(db) runner = RunnerConfig(user=settings.claude_runner_user, home=settings.claude_home) mcp_tool_names = { name: [t.name for t in tools] for name, tools in mcp_tools.items() @@ -321,6 +365,8 @@ async def _build_backends( mcp_tool_names=mcp_tool_names, runner=runner, usage_sink=record_usage, + pool=pool, + tool_server=functools.partial(late.server, names=a.gateway_tools), ) await stack.enter_async_context(adapter) backends[a.name] = adapter diff --git a/src/beaver_gateway/config_loader.py b/src/beaver_gateway/config_loader.py index 87876e8..102a97b 100644 --- a/src/beaver_gateway/config_loader.py +++ b/src/beaver_gateway/config_loader.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any 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.frontends.base import Frontend from beaver_gateway.mcp.types import HttpMcp, McpServer, PythonToolMcp, StdioMcp @@ -37,6 +38,7 @@ _PUBLIC_NAMES: dict[str, Any] = { "McpServer": McpServer, "ExposedMcp": ExposedMcp, "Gateway": Gateway, + "ConversationTexts": ConversationTexts, } _McpInstance = StdioMcp | HttpMcp | PythonToolMcp diff --git a/src/beaver_gateway/core/auth.py b/src/beaver_gateway/core/auth.py index f3f9908..b1824a7 100644 --- a/src/beaver_gateway/core/auth.py +++ b/src/beaver_gateway/core/auth.py @@ -60,12 +60,13 @@ class TokenStoreError(ValueError): """Malformed ``BOOTSTRAP_TOKENS`` value or duplicate token.""" -VALID_SCOPES: frozenset[str] = frozenset({"*", "messages", "mcp", "admin"}) +VALID_SCOPES: frozenset[str] = frozenset({"*", "messages", "mcp", "admin", "api"}) """The scopes a ``Token.scope`` may hold (Phase 4.3 admin UI enforces). * ``*`` — wildcard, may use any frontend * ``messages`` — Anthropic Messages frontend only * ``mcp`` — MCP server frontend only +* ``api`` — conversations API (``/api``) and its SSE * ``admin`` — reserved for programmatic admin access; the AdminFrontend itself authenticates via session cookies, not bearer tokens, so this scope is unused today and kept for forward compatibility. diff --git a/src/beaver_gateway/core/bus.py b/src/beaver_gateway/core/bus.py new file mode 100644 index 0000000..e2fe37a --- /dev/null +++ b/src/beaver_gateway/core/bus.py @@ -0,0 +1,69 @@ +"""In-process event bus: what the gateway does, as a stream frontends can tap. + +Events are plain dicts ``{"type", "seq", "ts", ...}``. Publishers never +block; a subscriber that falls behind loses its oldest events rather than +stalling the turn that produced them. ``/api/events`` serialises the +stream as SSE, the panel builds its subagent tree from ``stream`` events +carrying ``parent_tool_use_id``. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import itertools +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +__all__ = ["Event", "EventBus"] + +Event = dict[str, Any] + + +class EventBus: + def __init__(self, *, maxsize: int = 2000) -> None: + self._maxsize = maxsize + self._subscribers: set[asyncio.Queue[Event]] = set() + self._seq = itertools.count(1) + + def publish(self, type_: str, **data: Any) -> Event: + event: Event = { + "type": type_, + "seq": next(self._seq), + "ts": datetime.now(UTC).isoformat(timespec="milliseconds"), + **data, + } + for queue in list(self._subscribers): + if queue.full(): + with contextlib.suppress(asyncio.QueueEmpty): + queue.get_nowait() + queue.put_nowait(event) + return event + + @contextlib.asynccontextmanager + async def subscribe(self) -> AsyncIterator[asyncio.Queue[Event]]: + queue: asyncio.Queue[Event] = asyncio.Queue(maxsize=self._maxsize) + self._subscribers.add(queue) + try: + yield queue + finally: + self._subscribers.discard(queue) + + async def stream( + self, *, conversation_id: str | None = None + ) -> AsyncIterator[Event]: + async with self.subscribe() as queue: + while True: + event = await queue.get() + if ( + conversation_id is None + or event.get("conversation_id") == conversation_id + ): + yield event + + @property + def subscribers(self) -> int: + return len(self._subscribers) diff --git a/src/beaver_gateway/core/conversation_store.py b/src/beaver_gateway/core/conversation_store.py index 4cd0e04..abe44c0 100644 --- a/src/beaver_gateway/core/conversation_store.py +++ b/src/beaver_gateway/core/conversation_store.py @@ -496,12 +496,14 @@ def _walk_prefix( s.name for s in inc.structure if isinstance(s, ToolSegment) ) inc_text_count = sum(1 for s in inc.structure if isinstance(s, TextSegment)) - if inc_skeleton != st.skeleton: + # Files rendered without tool callouts (§3.10) carry no skeleton: + # prose alone decides whether the turn matched. + if inc_skeleton and inc_skeleton != st.skeleton: return spliced_groups, i if inc.text == st.spoken_text: spliced_groups.append(list(st.messages)) continue - if inc_text_count != st.text_segment_count: + if inc_skeleton and inc_text_count != st.text_segment_count: return spliced_groups, i spliced = _splice_assistant_group(stored_group=st, incoming=inc) if spliced is None: diff --git a/src/beaver_gateway/core/conversations.py b/src/beaver_gateway/core/conversations.py new file mode 100644 index 0000000..8e046a0 --- /dev/null +++ b/src/beaver_gateway/core/conversations.py @@ -0,0 +1,1024 @@ +"""Conversations, threads and the one place a turn runs (§3.1, §3.4, §8.2-8.3). + +The service owns the rows (:class:`Conversation`, :class:`ConversationBinding`), +the per-conversation queue workers, the fork/spawn/read/inject/say/schedule +API that frontends, jobs and the in-process MCP tools all call, and the +restart recovery: a conversation with ``running_turn`` set at startup gets +its open ``tool_use`` closed in the transcript and an inject saying the +turn was cut. + +Routing is deterministic (§0.1): a turn started by a user message streams +back to whoever asked; a turn started by an inject streams nowhere - its +events still hit the bus flagged ``origin="inject"`` so the panel can show +activity, but the only way it speaks is ``say``. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import logging +import re +import uuid +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Any, cast + +from claude_agent_sdk import ( + AssistantMessage, + ResultMessage, + StreamEvent, + ToolUseBlock, + fork_session_via_store, + project_key_for_directory, +) +from sqlmodel import col, select + +from beaver_gateway.core.injects import InjectQueue, inject_header +from beaver_gateway.core.transcript import ( + messages_from_entries, + render_messages, + strip_tool_entries, + window_entries, +) +from beaver_gateway.core.turn_capture import TurnCapture +from beaver_gateway.frontends._accumulate import StreamAccumulator +from beaver_gateway.storage.models import ( + Conversation, + ConversationBinding, + InjectQueueItem, + Schedule, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Awaitable, Callable, Sequence + + from claude_agent_sdk import SessionStore + + 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.events import MessageStreamEvent + from beaver_gateway.core.injects import Priority + from beaver_gateway.core.registry import AgentRegistry + from beaver_gateway.core.sessions import SessionPool + from beaver_gateway.storage.db import Database + +__all__ = [ + "KINDS", + "SEEDS", + "ConversationTexts", + "Conversations", + "ForkResult", + "SeedContext", +] + +_log = logging.getLogger("beaver_gateway.core.conversations") + +KINDS = ("master", "branch", "deep", "job", "fork") +SEEDS = ("clean", "morning", "copy", "brief") +_STATUSES = ("open", "merged", "closed", "archived") +_DEFAULT_MERGE_PROMPT = ( + "Ветка закрывается. Напиши слив для мастера: что решили, что сделано, " + "что не сделано и почему, открытые вопросы. Идентификаторы и ссылки - " + "дословно. Коротко, прошедшее время." +) +_RELATIVE = re.compile(r"^\+(\d+)\s*([smhd])$") +_UNITS = {"s": 1, "m": 60, "h": 3600, "d": 86400} + + +@dataclass(frozen=True, slots=True) +class SeedContext: + kind: str + seed: str + agent: str + parent: Conversation | None + text: str | None + title: str | None + + +@dataclass(frozen=True, slots=True) +class ConversationTexts: + """Texts the gateway cannot invent for a setup. + + What a merge asks for and what a seed says; ``seed`` may return the body + for any seed mode (the morning handout lives in the vault the gateway + knows nothing about). + """ + + merge_prompt: str = _DEFAULT_MERGE_PROMPT + interrupted: str = "прервано" + seed: Callable[[SeedContext], Awaitable[str | None] | str | None] | None = None + + +@dataclass(frozen=True, slots=True) +class ForkResult: + conversation: Conversation + text: str + capture: TurnCapture + + +@dataclass +class _Runner: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + wake: asyncio.Event = field(default_factory=asyncio.Event) + task: asyncio.Task[None] | None = None + turn_id: str | None = None + + +class Conversations: + def __init__( + self, + *, + db: Database, + agents: AgentRegistry, + backends: dict[str, Any], + bus: EventBus, + pool: SessionPool, + store: SessionStore, + texts: ConversationTexts | None = None, + normal_window: float = 3600.0, + idle_days: Sequence[int] = (2,), + idle_interval: float = 3600.0, + ) -> None: + self._db = db + self._agents = agents + self._backends = backends + self._bus = bus + self._pool = pool + self._store = store + self._texts = texts or ConversationTexts() + self._normal_window = normal_window + self._idle_days = tuple(sorted(idle_days)) + self._idle_interval = idle_interval + self._queue = InjectQueue(db) + self._runners: dict[int, _Runner] = {} + self._tasks: set[asyncio.Task[None]] = set() + self._idle_task: asyncio.Task[None] | None = None + + @property + def queue(self) -> InjectQueue: + return self._queue + + @property + def bus(self) -> EventBus: + return self._bus + + @property + def pool(self) -> SessionPool: + return self._pool + + # ---- rows ---------------------------------------------------------- + + async def create( + self, + *, + kind: str, + agent: str, + parent: Conversation | None = None, + title: str | None = None, + origin: str = "api", + session_id: str | None = None, + flags: dict[str, Any] | None = None, + ) -> Conversation: + if kind not in KINDS: + msg = f"unknown conversation kind {kind!r}" + raise ValueError(msg) + self._claude_agent(agent) + now = datetime.now(UTC) + row = Conversation( + frontend=origin, + external_id=str(uuid.uuid4()), + agent_name=agent, + kind=kind, + parent_id=parent.id if parent is not None else None, + title=title, + session_id=session_id, + flags=dict(flags or {}), + last_activity_at=now, + ) + async with self._db.session() as session: + session.add(row) + await session.commit() + await session.refresh(row) + self._bus.publish("conversation.created", **self.public(row)) + return row + + async def get(self, public_id: str) -> Conversation | None: + async with self._db.session() as session: + result = await session.exec( + select(Conversation).where(Conversation.external_id == public_id) + ) + return result.first() + + async def get_row(self, row_id: int) -> Conversation | None: + async with self._db.session() as session: + return await session.get(Conversation, row_id) + + async def find( + self, + *, + status: str | None = None, + kind: str | None = None, + parent: Conversation | None = None, + limit: int = 200, + ) -> list[Conversation]: + stmt = select(Conversation).order_by(col(Conversation.id).desc()).limit(limit) + if status is not None: + stmt = stmt.where(Conversation.status == status) + if kind is not None: + stmt = stmt.where(Conversation.kind == kind) + if parent is not None: + stmt = stmt.where(Conversation.parent_id == parent.id) + async with self._db.session() as session: + return list((await session.exec(stmt)).all()) + + async def bindings(self, conv: Conversation) -> list[ConversationBinding]: + async with self._db.session() as session: + result = await session.exec( + select(ConversationBinding) + .where(ConversationBinding.conversation_id == conv.id) + .order_by(col(ConversationBinding.id)) + ) + return list(result.all()) + + async def bind( + self, + conv: Conversation, + *, + frontend: str, + external_id: str, + visible: bool = True, + ) -> ConversationBinding: + async with self._db.session() as session: + existing = list( + ( + await session.exec( + select(ConversationBinding).where( + ConversationBinding.conversation_id == conv.id, + ConversationBinding.frontend == frontend, + ) + ) + ).all() + ) + row = next((b for b in existing if b.external_id == external_id), None) + if visible: + for other in existing: + if other is not row and other.visible: + other.visible = False + session.add(other) + if row is None: + row = ConversationBinding( + conversation_id=cast("int", conv.id), + frontend=frontend, + external_id=external_id, + visible=visible, + ) + else: + row.visible = visible + session.add(row) + await session.commit() + await session.refresh(row) + self._bus.publish( + "conversation.bound", + conversation_id=conv.external_id, + frontend=frontend, + external_id=external_id, + visible=visible, + ) + return row + + async def find_bound( + self, *, frontend: str, external_id: str + ) -> Conversation | None: + async with self._db.session() as session: + result = await session.exec( + select(Conversation) + .join( + ConversationBinding, + col(ConversationBinding.conversation_id) == col(Conversation.id), + ) + .where( + ConversationBinding.frontend == frontend, + ConversationBinding.external_id == external_id, + col(ConversationBinding.visible).is_(True), + ) + .order_by(col(Conversation.id).desc()) + ) + return result.first() + + async def set_flags( + self, conv: Conversation, flags: dict[str, Any] + ) -> Conversation: + async def apply(row: Conversation) -> None: + row.flags = {**row.flags, **flags} + + return await self._update(conv, apply) + + async def set_status(self, conv: Conversation, status: str) -> Conversation: + if status not in _STATUSES: + msg = f"unknown status {status!r}" + raise ValueError(msg) + + async def apply(row: Conversation) -> None: + row.status = status + + return await self._update(conv, apply) + + async def set_title(self, conv: Conversation, title: str) -> Conversation: + async def apply(row: Conversation) -> None: + row.title = title + + return await self._update(conv, apply) + + async def touch_user(self, conv: Conversation) -> Conversation: + async def apply(row: Conversation) -> None: + row.last_user_activity_at = datetime.now(UTC) + + return await self._update(conv, apply) + + async def _update( + self, conv: Conversation, apply: Callable[[Conversation], Awaitable[None]] + ) -> Conversation: + async with self._db.session() as session: + row = await session.get(Conversation, conv.id) + if row is None: + msg = f"conversation {conv.external_id} vanished" + raise LookupError(msg) + await apply(row) + row.updated_at = datetime.now(UTC) + session.add(row) + await session.commit() + await session.refresh(row) + self._bus.publish("conversation.updated", **self.public(row)) + return row + + def public(self, conv: Conversation) -> dict[str, Any]: + return { + "id": conv.external_id, + "kind": conv.kind, + "agent": conv.agent_name, + "title": conv.title, + "status": conv.status, + "parent_row": conv.parent_id, + "session_id": conv.session_id, + "running_turn": conv.running_turn, + "pending_question": conv.pending_question, + "flags": conv.flags, + "origin": conv.frontend, + "created_at": _iso(conv.created_at), + "last_user_activity_at": _iso(conv.last_user_activity_at), + "last_activity_at": _iso(conv.last_activity_at), + } + + async def describe(self, conv: Conversation) -> dict[str, Any]: + out = self.public(conv) + parent = await self.get_row(conv.parent_id) if conv.parent_id else None + out["parent"] = parent.external_id if parent is not None else None + out["bindings"] = [ + {"frontend": b.frontend, "external_id": b.external_id, "visible": b.visible} + for b in await self.bindings(conv) + ] + live = self._pool.get(conv.external_id) + out["live"] = live is not None + out["busy"] = live.busy if live is not None else False + return out + + # ---- §3.1 api ------------------------------------------------------ + + async def spawn( + self, + *, + kind: str, + agent: str, + seed: str = "clean", + parent: Conversation | None = None, + text: str | None = None, + title: str | None = None, + window: int | None = None, + origin: str = "api", + ) -> Conversation: + if seed not in SEEDS: + msg = f"unknown seed {seed!r}" + raise ValueError(msg) + if seed == "brief" and not text: + msg = "seed=brief needs text" + raise ValueError(msg) + session_id: str | None = None + if seed == "copy": + if parent is None or parent.session_id is None: + msg = "seed=copy needs a parent with a session" + raise ValueError(msg) + session_id = await self._copy_session( + parent, window=window, strip_tools=False + ) + conv = await self.create( + kind=kind, + agent=agent, + parent=parent, + title=title, + origin=origin, + session_id=session_id, + ) + prompt = await self._seed_text( + SeedContext( + kind=kind, seed=seed, agent=agent, parent=parent, text=text, title=title + ), + window=window, + ) + await self._queue.push( + conversation_id=cast("int", conv.id), + priority="user", + origin=f"сид:{seed}", + text=prompt, + ) + self._ensure_worker(cast("int", conv.id)) + return conv + + async def fork( + self, + conv: Conversation, + prompt: str, + *, + window: int | None = None, + strip_tools: bool = False, + title: str | None = None, + ) -> ForkResult: + session_id = await self._copy_session( + conv, window=window, strip_tools=strip_tools + ) + child = await self.create( + kind="fork", + agent=conv.agent_name, + parent=conv, + title=title or f"fork: {conv.title or conv.external_id}", + origin="system", + session_id=session_id, + ) + try: + text, capture = await self.run_text_turn( + child, prompt, origin="fork", tools=False + ) + finally: + await self._backend(conv.agent_name).close(child.external_id) + child = await self.set_status(child, "closed") + return ForkResult(conversation=child, text=text, capture=capture) + + async def read(self, conv: Conversation, *, window: int | None = None) -> str: + if conv.session_id is None: + return "" + entries = await self._store.load(cast("Any", self._store_key(conv))) + if not entries: + return "" + return render_messages( + messages_from_entries(cast("Any", entries)), window=window + ) + + async def inject( + self, + conv: Conversation, + text: str, + *, + urgency: Priority = "normal", + origin: str = "system", + ) -> InjectQueueItem: + item = await self._queue.push( + conversation_id=cast("int", conv.id), + priority=urgency, + origin=origin, + text=text, + ) + self._bus.publish( + "inject.queued", + conversation_id=conv.external_id, + item=item.id, + priority=urgency, + origin=origin, + ) + if urgency == "urgent": + backend = self._backend(conv.agent_name) + if await backend.interrupt(conv.external_id): + _log.info( + "conversation %s: interrupted for urgent inject", conv.external_id + ) + self._ensure_worker(cast("int", conv.id)) + return item + + async def post( + self, conv: Conversation, text: str, *, origin: str = "user" + ) -> InjectQueueItem: + item = await self._queue.push( + conversation_id=cast("int", conv.id), + priority="user", + origin=origin, + text=text, + ) + await self.touch_user(conv) + self._bus.publish( + "message.queued", + conversation_id=conv.external_id, + item=item.id, + origin=origin, + ) + self._ensure_worker(cast("int", conv.id)) + return item + + async def say(self, conv: Conversation, text: str) -> dict[str, Any]: + runner = self._runners.get(cast("int", conv.id)) + _log.info("say[%s]: %s", conv.external_id, text[:200]) + return self._bus.publish( + "say", + conversation_id=conv.external_id, + text=text, + turn_id=runner.turn_id if runner is not None else None, + ) + + async def merge(self, conv: Conversation) -> ForkResult: + if conv.parent_id is None: + msg = "merge needs a parent conversation" + raise ValueError(msg) + parent = await self.get_row(conv.parent_id) + if parent is None: + msg = "parent conversation vanished" + raise LookupError(msg) + result = await self.fork( + conv, + self._texts.merge_prompt, + title=f"слив: {conv.title or conv.external_id}", + ) + if result.text.strip(): + await self.inject(parent, result.text, urgency="normal", origin="слив") + await self.set_status(conv, "merged") + self._bus.publish( + "conversation.merged", + conversation_id=conv.external_id, + parent=parent.external_id, + fork=result.conversation.external_id, + ) + 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) + self._bus.publish( + "schedule.created", + conversation_id=conv.external_id, + schedule=row.id, + execute_at=_iso(row.execute_at), + ) + return row + + 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( + self, + conv: Conversation, + *, + messages: Sequence[Any], + origin: str, + capture: TurnCapture | None = None, + session_id: str | None = None, + use_session: bool = True, + tools: bool = True, + turn_id: str | None = None, + ) -> AsyncIterator[MessageStreamEvent]: + """Run one turn on ``conv`` under its lock; the only path to the backend. + + ``use_session=False`` withholds the stored ``session_id`` so the + backend reseeds from ``messages`` (markdown file diverged). + """ + row_id = cast("int", conv.id) + runner = self._runner(row_id) + backend = self._backend(conv.agent_name) + turn_id = turn_id or f"turn_{uuid.uuid4().hex[:12]}" + capture = capture or TurnCapture() + resume = session_id if session_id is not None else conv.session_id + async with runner.lock: + runner.turn_id = turn_id + await self._mark_running(conv, turn_id) + self._bus.publish( + "turn.start", + conversation_id=conv.external_id, + turn_id=turn_id, + origin=origin, + ) + stop = "error" + try: + events = backend.complete( + agent=self._claude_agent(conv.agent_name), + messages=messages, + conversation_id=conv.external_id, + session_id=resume if use_session else None, + capture=capture, + kind=conv.kind, + pinned=conv.kind == "master", + tools=tools, + observer=self._observer(conv.external_id, turn_id, origin), + turn_id=turn_id, + ) + async for event in events: + yield event + stop = "interrupted" if capture.interrupted else "end_turn" + finally: + runner.turn_id = None + await self._mark_done(conv, capture) + self._bus.publish( + "turn.end", + conversation_id=conv.external_id, + turn_id=turn_id, + origin=origin, + stop=stop, + usage=_usage_dict(capture), + ) + + async def run_text_turn( + self, + conv: Conversation, + text: str, + *, + origin: str, + tools: bool = True, + turn_id: str | None = None, + ) -> tuple[str, TurnCapture]: + capture = TurnCapture() + acc = StreamAccumulator() + agent = self._claude_agent(conv.agent_name) + async for event in self.turn( + conv, + messages=[{"role": "user", "content": text}], + origin=origin, + capture=capture, + tools=tools, + turn_id=turn_id, + ): + acc.feed(event) + message = acc.finalize(model=agent.model) + reply = "\n\n".join( + getattr(b, "text", "") + for b in message.content + if getattr(b, "type", "") == "text" + ).strip() + return reply, capture + + # ---- lifecycle ----------------------------------------------------- + + async def start(self) -> None: + await self.recover() + for row_id in await self._queue.conversations_with_pending(): + self._ensure_worker(row_id) + self._idle_task = asyncio.create_task(self._idle_loop()) + + async def stop(self) -> None: + tasks = list(self._tasks) + if self._idle_task is not None: + tasks.append(self._idle_task) + for task in tasks: + task.cancel() + for task in tasks: + with contextlib.suppress(BaseException): + await task + self._tasks.clear() + self._idle_task = None + + async def recover(self) -> list[Conversation]: + """Restart path (§8.1): repair transcripts of turns cut mid-flight.""" + async with self._db.session() as session: + result = await session.exec( + select(Conversation).where(col(Conversation.running_turn).is_not(None)) + ) + cut = list(result.all()) + for conv in cut: + fixed = 0 + if conv.session_id is not None: + backend = self._backend(conv.agent_name) + try: + fixed = await backend.repair_session(conv.session_id) + except Exception: # noqa: BLE001 + _log.exception("repair of %s failed", conv.session_id) + turn_id = conv.running_turn + + async def clear(row: Conversation) -> None: + row.running_turn = None + + await self._update(conv, clear) + note = f"тёрн {turn_id} оборван рестартом gateway" + if fixed: + note += ( + f"; {fixed} незакрытых тул-коллов получили tool_result " + f"«{self._texts.interrupted}»" + ) + await self.inject(conv, note, urgency="normal", origin="система") + _log.warning("conversation %s: %s", conv.external_id, note) + interrupted = await self._queue.interrupted() + for item in interrupted: + _log.warning( + "queue item #%s (%s) was running at restart; marked interrupted", + item.id, + item.priority, + ) + return cut + + # ---- internals ----------------------------------------------------- + + def _backend(self, agent: str) -> ClaudeSdkBackend: + backend = self._backends.get(agent) + if backend is None or not hasattr(backend, "repair_session"): + msg = f"agent {agent!r} has no Claude SDK backend" + raise LookupError(msg) + return cast("ClaudeSdkBackend", backend) + + def _claude_agent(self, name: str) -> ClaudeAgent: + agent = self._agents.get(name) + if agent is None or not hasattr(agent, "cwd"): + msg = f"unknown Claude agent {name!r}" + raise LookupError(msg) + return cast("ClaudeAgent", agent) + + def _store_key(self, conv: Conversation) -> dict[str, str]: + agent = self._claude_agent(conv.agent_name) + return { + "project_key": project_key_for_directory(str(agent.cwd)), + "session_id": cast("str", conv.session_id), + } + + async def _copy_session( + self, conv: Conversation, *, window: int | None, strip_tools: bool + ) -> str: + if conv.session_id is None: + msg = f"conversation {conv.external_id} has no session to copy" + raise ValueError(msg) + live = self._pool.get(conv.external_id) + if live is not None and live.dirty: + msg = f"conversation {conv.external_id} has a mirror gap; not forking" + raise RuntimeError(msg) + agent = self._claude_agent(conv.agent_name) + forked = await fork_session_via_store( + self._store, conv.session_id, directory=str(agent.cwd) + ) + if window is not None or strip_tools: + key = { + "project_key": project_key_for_directory(str(agent.cwd)), + "session_id": forked.session_id, + } + entries = await self._store.load(cast("Any", key)) or [] + trimmed = window_entries(cast("Any", entries), window=window) + if strip_tools: + trimmed = strip_tool_entries(trimmed) + await self._store.delete(cast("Any", key)) + await self._store.append(cast("Any", key), cast("Any", trimmed)) + _log.info( + "forked session %s -> %s (window=%s, strip_tools=%s)", + conv.session_id, + forked.session_id, + window, + strip_tools, + ) + return forked.session_id + + async def _seed_text(self, ctx: SeedContext, *, window: int | None) -> str: + stamp = datetime.now(UTC).astimezone().strftime("%Y-%m-%d %H:%M") + title = f" «{ctx.title}»" if ctx.title else "" + head = f"[сид: {ctx.seed}] {ctx.kind}{title}, {stamp}." + body: str | None = None + if self._texts.seed is not None: + produced: Any = self._texts.seed(ctx) + if inspect.isawaitable(produced): + produced = await produced + body = cast("str | None", produced) + if body is None: + if ctx.seed == "brief": + body = ctx.text + elif ctx.seed == "copy": + scope = f"последние {window} тёрнов" if window else "вся история" + body = f"История родителя скопирована ({scope}); продолжай в ней." + elif ctx.seed == "morning": + body = "Хендаут не приехал." + return f"{head}\n\n{body}" if body else head + + def _runner(self, row_id: int) -> _Runner: + runner = self._runners.get(row_id) + if runner is None: + runner = _Runner() + self._runners[row_id] = runner + return runner + + def _ensure_worker(self, row_id: int) -> None: + runner = self._runner(row_id) + runner.wake.set() + if runner.task is None or runner.task.done(): + runner.task = asyncio.create_task(self._worker(row_id)) + self._tasks.add(runner.task) + runner.task.add_done_callback(self._tasks.discard) + + async def _worker(self, row_id: int) -> None: + runner = self._runner(row_id) + while True: + items = await self._queue.pending(row_id) + batch, wait = self._pick(items) + if batch is None: + runner.wake.clear() + if wait is None: + return + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(runner.wake.wait(), timeout=wait) + continue + conv = await self.get_row(row_id) + if conv is None: + await self._queue.finish(batch, status="failed") + return + await self._run_batch(conv, batch) + + def _pick( + self, items: list[InjectQueueItem] + ) -> tuple[list[InjectQueueItem] | None, float | None]: + if not items: + return None, None + normals = [i for i in items if i.priority == "normal"] + head = items[0] + if head.priority == "urgent": + return [head], None + if head.priority == "user": + return [head, *normals], None + age = (datetime.now(UTC) - _aware(head.created_at)).total_seconds() + if age >= self._normal_window: + return normals, None + return None, max(self._normal_window - age, 1.0) + + async def _run_batch( + self, conv: Conversation, batch: list[InjectQueueItem] + ) -> None: + head = batch[0] + turn_id = f"turn_{uuid.uuid4().hex[:12]}" + await self._queue.start(batch, turn_id) + if head.priority == "user": + origin = "user" + prompt = head.text + if len(batch) > 1: + prompt += "\n\n" + _bundle(batch[1:]) + else: + origin = "inject" + prompt = "\n\n".join(f"{inject_header(i.origin)}\n{i.text}" for i in batch) + try: + text, capture = await self.run_text_turn( + conv, prompt, origin=origin, turn_id=turn_id + ) + except Exception: # noqa: BLE001 + _log.exception("turn %s on %s failed", turn_id, conv.external_id) + await self._queue.finish(batch, status="failed") + return + await self._queue.finish( + batch, status="interrupted" if capture.interrupted else "done" + ) + if origin == "user": + self._bus.publish( + "reply", + conversation_id=conv.external_id, + turn_id=turn_id, + item=head.id, + text=text, + ) + + def _observer( + self, conversation_id: str, turn_id: str, origin: str + ) -> Callable[[Any], None]: + def observe(message: Any) -> None: + parent = getattr(message, "parent_tool_use_id", None) + if isinstance(message, StreamEvent): + self._bus.publish( + "stream", + conversation_id=conversation_id, + turn_id=turn_id, + origin=origin, + parent_tool_use_id=parent, + event=message.event, + ) + elif isinstance(message, AssistantMessage): + for block in message.content: + if isinstance(block, ToolUseBlock): + self._bus.publish( + "tool", + conversation_id=conversation_id, + turn_id=turn_id, + origin=origin, + parent_tool_use_id=parent, + tool_use_id=block.id, + name=block.name, + input=block.input, + ) + elif isinstance(message, ResultMessage) and parent is None: + self._bus.publish( + "result", + conversation_id=conversation_id, + turn_id=turn_id, + origin=origin, + subtype=message.subtype, + is_error=message.is_error, + num_turns=message.num_turns, + ) + + return observe + + async def _mark_running(self, conv: Conversation, turn_id: str) -> None: + async def apply(row: Conversation) -> None: + row.running_turn = turn_id + row.last_activity_at = datetime.now(UTC) + + await self._update(conv, apply) + + async def _mark_done(self, conv: Conversation, capture: TurnCapture) -> None: + async def apply(row: Conversation) -> None: + row.running_turn = None + row.last_activity_at = datetime.now(UTC) + if capture.session_id is not None: + row.session_id = capture.session_id + + await self._update(conv, apply) + + async def _idle_loop(self) -> None: + while True: + try: + await self._emit_idle() + except Exception: # noqa: BLE001 + _log.exception("idle watcher failed") + await asyncio.sleep(self._idle_interval) + + async def _emit_idle(self) -> None: + if not self._idle_days: + return + now = datetime.now(UTC) + for conv in await self.find(status="open", limit=10_000): + last = _aware(conv.last_activity_at or conv.created_at) + days = int((now - last).total_seconds() // 86400) + due = [d for d in self._idle_days if days >= d] + if not due: + continue + notified = int(conv.flags.get("idle_notified", 0) or 0) + if due[-1] <= notified: + continue + await self.set_flags(conv, {"idle_notified": due[-1]}) + bindings = await self.bindings(conv) + self._bus.publish( + "conversation.idle", + conversation_id=conv.external_id, + kind=conv.kind, + agent=conv.agent_name, + days=due[-1], + bindings=[ + {"frontend": b.frontend, "external_id": b.external_id} + for b in bindings + if b.visible + ], + ) + + +def parse_at(at: str) -> datetime: + raw = at.strip() + match = _RELATIVE.match(raw.replace(" ", "")) + if match: + amount, unit = match.groups() + return datetime.now(UTC) + timedelta(seconds=int(amount) * _UNITS[unit]) + parsed = datetime.fromisoformat(raw) + if parsed.tzinfo is None: + parsed = parsed.astimezone() + return parsed.astimezone(UTC) + + +def _bundle(items: Sequence[InjectQueueItem]) -> str: + lines = [f"[инжекты, накопившиеся с {_iso(items[0].created_at)}; это не Бобёр]"] # noqa: RUF001 + lines.extend(f"- [{i.origin}] {i.text}" for i in items) + return "\n".join(lines) + + +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 + + +def _usage_dict(capture: TurnCapture) -> dict[str, Any] | None: + usage = capture.usage + if usage is None: + return None + return { + "input": usage.input_tokens, + "output": usage.output_tokens, + "cache_read": usage.cache_read_tokens, + "cache_creation": usage.cache_creation_tokens, + "cost_usd": usage.cost_usd, + "duration_ms": usage.duration_ms, + } diff --git a/src/beaver_gateway/core/gateway_tools.py b/src/beaver_gateway/core/gateway_tools.py new file mode 100644 index 0000000..6e45c37 --- /dev/null +++ b/src/beaver_gateway/core/gateway_tools.py @@ -0,0 +1,169 @@ +"""In-process MCP server with the gateway's own tools (§3.1, §3.2). + +One server per live session so every tool knows which conversation is +calling; ``alwaysLoad`` keeps the tools out of tool search. Which names a +session gets comes from ``ClaudeAgent.gateway_tools``. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, cast + +from claude_agent_sdk import create_sdk_mcp_server, tool + +if TYPE_CHECKING: + from collections.abc import Iterable + + from claude_agent_sdk import McpSdkServerConfig, SdkMcpTool + + from beaver_gateway.core.conversations import Conversations + +__all__ = ["SERVER_NAME", "TOOL_NAMES", "build_tool_server"] + +_log = logging.getLogger("beaver_gateway.core.gateway_tools") + +SERVER_NAME = "gateway" +TOOL_NAMES = ("read_conversation", "spawn", "say", "schedule", "inject") + + +def build_tool_server( + conversations: Conversations, *, conversation_key: str, names: Iterable[str] +) -> McpSdkServerConfig | None: + wanted = set(names) + unknown = wanted - set(TOOL_NAMES) + if unknown: + msg = f"unknown gateway tools: {sorted(unknown)}" + raise ValueError(msg) + tools = [t for t in _tools(conversations, conversation_key) if t.name in wanted] + if not tools: + return None + server = create_sdk_mcp_server(SERVER_NAME, tools=tools) + return cast("McpSdkServerConfig", {**server, "alwaysLoad": True}) + + +def _tools(conversations: Conversations, key: str) -> list[SdkMcpTool[Any]]: + async def current() -> Any: + conv = await conversations.get(key) + if conv is None: + msg = f"conversation {key} not found" + raise LookupError(msg) + return conv + + @tool( + "read_conversation", + "Read another conversation (a branch, the master, a deep chat) as plain " + "text. `window` limits it to the last N user turns.", + { + "type": "object", + "properties": { + "id": {"type": "string", "description": "conversation id"}, + "window": {"type": "integer", "minimum": 1}, + }, + "required": ["id"], + }, + ) + async def read_conversation(args: dict[str, Any]) -> dict[str, Any]: + conv = await conversations.get(str(args["id"])) + if conv is None: + return _error(f"conversation {args['id']} not found") + text = await conversations.read(conv, window=args.get("window")) + return _text(text or "(empty)") + + @tool( + "spawn", + "Open a new conversation of the given kind (branch = your own thread, " + "deep = a long research chat, job = a headless task). `seed` is how it " + "starts: clean (nothing), morning (handout), copy (copy of this " + "conversation, last `window` turns), brief (your `text`). Returns the id.", + { + "type": "object", + "properties": { + "kind": {"type": "string", "enum": ["branch", "deep", "job"]}, + "seed": { + "type": "string", + "enum": ["clean", "morning", "copy", "brief"], + "default": "clean", + }, + "text": {"type": "string", "description": "brief for seed=brief"}, + "title": {"type": "string"}, + "window": {"type": "integer", "minimum": 1}, + }, + "required": ["kind"], + }, + ) + async def spawn(args: dict[str, Any]) -> dict[str, Any]: + parent = await current() + child = await conversations.spawn( + kind=str(args["kind"]), + agent=parent.agent_name, + seed=str(args.get("seed") or "clean"), + parent=parent, + text=args.get("text"), + title=args.get("title"), + window=args.get("window"), + origin="mcp", + ) + return _text(f"spawned {child.kind} {child.external_id}") + + @tool( + "say", + "Say something to the human in the frontend this conversation is bound " + "to. The only way an inject-started turn can speak; silence is simply " + "not calling it.", + {"text": str}, + ) + async def say(args: dict[str, Any]) -> dict[str, Any]: + conv = await current() + await conversations.say(conv, str(args["text"])) + return _text("ok") + + @tool( + "schedule", + "Promise yourself an inject later: `at` is `+15m`, `+2h`, `+1d` or an " + "ISO datetime; `text` arrives in this conversation at that time.", + {"at": str, "text": str}, + ) + 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()}") + + @tool( + "inject", + "Put a system-origin message into another conversation's queue.", + { + "type": "object", + "properties": { + "conversation": {"type": "string"}, + "text": {"type": "string"}, + "urgency": { + "type": "string", + "enum": ["normal", "urgent"], + "default": "normal", + }, + }, + "required": ["conversation", "text"], + }, + ) + async def inject(args: dict[str, Any]) -> dict[str, Any]: + target = await conversations.get(str(args["conversation"])) + if target is None: + return _error(f"conversation {args['conversation']} not found") + item = await conversations.inject( + target, + str(args["text"]), + urgency=cast("Any", args.get("urgency") or "normal"), + origin="агент", + ) + return _text(f"queued #{item.id}") + + return [read_conversation, spawn, say, schedule, inject] + + +def _text(text: str) -> dict[str, Any]: + return {"content": [{"type": "text", "text": text}]} + + +def _error(text: str) -> dict[str, Any]: + return {"content": [{"type": "text", "text": text}], "is_error": True} diff --git a/src/beaver_gateway/core/injects.py b/src/beaver_gateway/core/injects.py new file mode 100644 index 0000000..1ba5f9f --- /dev/null +++ b/src/beaver_gateway/core/injects.py @@ -0,0 +1,130 @@ +"""Persisted per-conversation queue with priorities ``urgent > user > normal`` (§3.4). + +One ``ClaudeSDKClient`` runs one turn at a time, so ordering has to happen +before the client: the rows here are the queue, ``core/conversations`` +runs one worker per conversation over them. A row that is still +``running`` when the gateway starts was cut by a restart; it is flagged +``interrupted`` and never re-run. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Literal + +from sqlmodel import col, select + +from beaver_gateway.storage.models import InjectQueueItem + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + + from beaver_gateway.storage.db import Database + +__all__ = ["PRIORITY_RANK", "InjectQueue", "Priority", "inject_header"] + +Priority = Literal["urgent", "user", "normal"] +PRIORITY_RANK: dict[str, int] = {"urgent": 0, "user": 1, "normal": 2} + + +def inject_header(origin: str) -> str: + return f"[инжект: {origin} - это не Бобёр, отвечать не нужно, голос не обязателен]" + + +class InjectQueue: + def __init__(self, db: Database) -> None: + self._db = db + + async def push( + self, *, conversation_id: int, priority: Priority, origin: str, text: str + ) -> InjectQueueItem: + if priority not in PRIORITY_RANK: + msg = f"unknown priority {priority!r}" + raise ValueError(msg) + row = InjectQueueItem( + conversation_id=conversation_id, priority=priority, origin=origin, text=text + ) + async with self._db.session() as session: + session.add(row) + await session.commit() + await session.refresh(row) + return row + + async def pending(self, conversation_id: int) -> list[InjectQueueItem]: + async with self._db.session() as session: + result = await session.exec( + select(InjectQueueItem) + .where( + InjectQueueItem.conversation_id == conversation_id, + InjectQueueItem.status == "queued", + ) + .order_by(col(InjectQueueItem.created_at), col(InjectQueueItem.id)) + ) + rows = list(result.all()) + rows.sort(key=lambda r: (PRIORITY_RANK.get(r.priority, 9), r.created_at)) + return rows + + async def conversations_with_pending(self) -> list[int]: + async with self._db.session() as session: + result = await session.exec( + select(InjectQueueItem.conversation_id) + .where(InjectQueueItem.status == "queued") + .distinct() + ) + return list(result.all()) + + async def start(self, items: Iterable[InjectQueueItem], turn_id: str) -> None: + await self._mark(items, status="running", turn_id=turn_id, delivered=True) + + async def finish( + self, items: Iterable[InjectQueueItem], status: str = "done" + ) -> None: + await self._mark(items, status=status) + + async def interrupted(self) -> Sequence[InjectQueueItem]: + async with self._db.session() as session: + result = await session.exec( + select(InjectQueueItem).where(InjectQueueItem.status == "running") + ) + rows = list(result.all()) + for row in rows: + row.status = "interrupted" + session.add(row) + await session.commit() + return rows + + async def recent( + self, conversation_id: int, *, limit: int = 50 + ) -> list[InjectQueueItem]: + async with self._db.session() as session: + result = await session.exec( + select(InjectQueueItem) + .where(InjectQueueItem.conversation_id == conversation_id) + .order_by(col(InjectQueueItem.id).desc()) + .limit(limit) + ) + return list(result.all()) + + async def _mark( + self, + items: Iterable[InjectQueueItem], + *, + status: str, + turn_id: str | None = None, + delivered: bool = False, + ) -> None: + ids = [i.id for i in items if i.id is not None] + if not ids: + return + async with self._db.session() as session: + result = await session.exec( + select(InjectQueueItem).where(col(InjectQueueItem.id).in_(ids)) + ) + for row in result.all(): + row.status = status + if turn_id is not None: + row.turn_id = turn_id + if delivered: + row.delivered_at = datetime.now(UTC) + session.add(row) + await session.commit() diff --git a/src/beaver_gateway/core/registry.py b/src/beaver_gateway/core/registry.py index d45c365..a7704a7 100644 --- a/src/beaver_gateway/core/registry.py +++ b/src/beaver_gateway/core/registry.py @@ -16,6 +16,7 @@ if TYPE_CHECKING: from collections.abc import Iterable, Iterator from beaver_gateway.agents.base import BaseAgent + from beaver_gateway.core.conversations import ConversationTexts from beaver_gateway.frontends.base import Frontend from beaver_gateway.mcp.types import McpServerT @@ -81,3 +82,5 @@ class Gateway: agents: list[BaseAgent] = field(default_factory=list) mcps: list[McpServerT] = field(default_factory=list) frontends: list[Frontend] = field(default_factory=list) + texts: ConversationTexts | None = None + """Merge prompt and seed bodies for ``core/conversations`` (§8.2-8.3).""" diff --git a/src/beaver_gateway/core/sessions.py b/src/beaver_gateway/core/sessions.py new file mode 100644 index 0000000..9914a00 --- /dev/null +++ b/src/beaver_gateway/core/sessions.py @@ -0,0 +1,257 @@ +"""Pool of live Agent SDK sessions across every Claude agent (§3.2). + +One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess). +The pool owns the two decisions the adapters used to make on their own: +when a session is closed for idleness (TTL by conversation kind) and +which one goes when memory runs out (measured RSS of the subprocess tree +against the cgroup limit, ``max_live`` where there is no limit). Eviction +only ever picks ``idle && !running_turn && !pending_question`` sessions +that are neither pinned (the master) nor ``dirty`` (mirror gap not yet +repaired); forks and jobs go first. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any, Protocol + +import psutil + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterator, Mapping + +__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"] + +_log = logging.getLogger("beaver_gateway.core.sessions") + +DEFAULT_TTL: Mapping[str, float | None] = { + "master": None, + "branch": 7200.0, + "deep": 1800.0, + "job": 0.0, + "fork": 0.0, +} +"""Idle seconds before a session is closed; ``None`` = never (pinned kinds).""" + +_EVICT_ORDER = {"fork": 0, "job": 0, "deep": 1, "branch": 2, "master": 3} +_RSS_HEADROOM = 0.8 + + +class SessionClient(Protocol): + async def connect(self) -> None: ... + async def query(self, prompt: str) -> None: ... + def receive_response(self) -> AsyncIterator[Any]: ... + async def interrupt(self) -> None: ... + async def disconnect(self) -> None: ... + + +@dataclass +class Session: + key: str + agent: str + kind: str + client: SessionClient + session_id: str | None + resumed: bool + pinned: bool = False + dirty: bool = False + running_turn: str | None = None + pending_question: bool = False + interrupt_requested: bool = False + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + last_used: float = field(default_factory=time.monotonic) + created_at: float = field(default_factory=time.monotonic) + turns: int = 0 + + @property + def busy(self) -> bool: + return self.lock.locked() or self.running_turn is not None + + @property + def evictable(self) -> bool: + return not (self.pinned or self.dirty or self.busy or self.pending_question) + + @property + def pid(self) -> int | None: + transport = getattr(self.client, "_transport", None) + process = getattr(transport, "_process", None) + pid = getattr(process, "pid", None) + return pid if isinstance(pid, int) else None + + +class SessionPool: + def __init__( + self, + *, + rss_limit: int | None = None, + max_live: int = 8, + ttl: Mapping[str, float | None] = DEFAULT_TTL, + reap_interval: float = 60.0, + ) -> None: + self._sessions: dict[str, Session] = {} + self._rss_limit = rss_limit if rss_limit is not None else cgroup_limit() + self._max_live = max_live + self._ttl = dict(ttl) + self._reap_interval = reap_interval + + def __len__(self) -> int: + return len(self._sessions) + + def __iter__(self) -> Iterator[Session]: + return iter(list(self._sessions.values())) + + def __contains__(self, key: object) -> bool: + return key in self._sessions + + @property + def rss_limit(self) -> int | None: + return self._rss_limit + + def get(self, key: str) -> Session | None: + return self._sessions.get(key) + + def add(self, session: Session) -> Session: + self._sessions[session.key] = session + return session + + def pop(self, key: str) -> Session | None: + return self._sessions.pop(key, None) + + def rekey(self, old: str, new: str) -> Session | None: + session = self._sessions.pop(old, None) + if session is None: + return None + stale = self._sessions.pop(new, None) + session.key = new + self._sessions[new] = session + return stale if stale is not session else None + + def ttl_for(self, kind: str) -> float | None: + return self._ttl.get(kind, self._ttl.get("deep")) + + def rss(self) -> int: + try: + children = psutil.Process(os.getpid()).children(recursive=True) + except psutil.Error: + return 0 + total = 0 + for child in children: + with contextlib.suppress(psutil.Error): + total += child.memory_info().rss + return total + + @staticmethod + def rss_of(session: Session) -> int | None: + pid = session.pid + if pid is None: + return None + try: + process = psutil.Process(pid) + return process.memory_info().rss + sum( + c.memory_info().rss for c in process.children(recursive=True) + ) + except psutil.Error: + return None + + def over_limit(self) -> bool: + if self._rss_limit is not None: + return self.rss() > self._rss_limit * _RSS_HEADROOM + return len(self._sessions) >= self._max_live + + def victims(self) -> list[Session]: + candidates = [s for s in self._sessions.values() if s.evictable] + candidates.sort(key=lambda s: (_EVICT_ORDER.get(s.kind, 1), s.last_used)) + return candidates + + async def make_room(self) -> int: + closed = 0 + while self.over_limit(): + victims = self.victims() + if not victims: + _log.warning( + "session pool over limit (%d live, rss=%d) but nothing evictable", + len(self._sessions), + self.rss(), + ) + break + await self.close(victims[0].key) + closed += 1 + return closed + + async def close(self, key: str) -> None: + session = self._sessions.pop(key, None) + if session is None: + return + _log.info("closing session %s (%s, %s)", session.session_id, session.kind, key) + try: + await session.client.disconnect() + except Exception: # noqa: BLE001 + _log.exception("disconnect failed for session %s", session.session_id) + + async def close_all(self, *, agent: str | None = None) -> None: + for session in list(self._sessions.values()): + if agent is None or session.agent == agent: + await self.close(session.key) + + async def reap_once(self) -> int: + now = time.monotonic() + closed = 0 + for session in list(self._sessions.values()): + ttl = self.ttl_for(session.kind) + if ttl is None or not session.evictable: + continue + if now - session.last_used > ttl: + await self.close(session.key) + closed += 1 + return closed + + async def reap_loop(self) -> None: + while True: + await asyncio.sleep(self._reap_interval) + try: + await self.reap_once() + await self.make_room() + except Exception: # noqa: BLE001 + _log.exception("session reaper failed") + + def snapshot(self) -> list[dict[str, Any]]: + now = time.monotonic() + return [ + { + "key": s.key, + "agent": s.agent, + "kind": s.kind, + "session_id": s.session_id, + "pid": s.pid, + "rss": self.rss_of(s), + "idle_seconds": round(now - s.last_used, 1), + "age_seconds": round(now - s.created_at, 1), + "turns": s.turns, + "busy": s.busy, + "running_turn": s.running_turn, + "pending_question": s.pending_question, + "pinned": s.pinned, + "dirty": s.dirty, + } + for s in self._sessions.values() + ] + + +def cgroup_limit() -> int | None: + for path in ( + "/sys/fs/cgroup/memory.max", + "/sys/fs/cgroup/memory/memory.limit_in_bytes", + ): + try: + raw = Path(path).read_text(encoding="ascii").strip() + except OSError: + continue + if raw.isdigit() and int(raw) < 1 << 60: + return int(raw) + return None diff --git a/src/beaver_gateway/core/transcript.py b/src/beaver_gateway/core/transcript.py index e039744..1529417 100644 --- a/src/beaver_gateway/core/transcript.py +++ b/src/beaver_gateway/core/transcript.py @@ -27,7 +27,17 @@ except ImportError: # pragma: no cover if TYPE_CHECKING: from collections.abc import Iterable, Mapping -__all__ = ["CLI_VERSION", "build_entries", "messages_from_entries"] +__all__ = [ + "CLI_VERSION", + "build_entries", + "close_open_tool_uses", + "messages_from_entries", + "open_tool_uses", + "prompt_count", + "render_messages", + "strip_tool_entries", + "window_entries", +] CLI_VERSION = _cli_version _ENTRYPOINT = "sdk-py" @@ -275,3 +285,221 @@ def _zero_usage() -> dict[str, Any]: "iterations": [], "speed": "standard", } + + +# ---- repair, windows, projections --------------------------------------- + +_PROMPT_TYPES = ("user", "assistant") +_INTERRUPTED = "прервано" + + +def open_tool_uses( + entries: Iterable[Mapping[str, Any]], +) -> list[tuple[Mapping[str, Any], dict[str, Any]]]: + """``(assistant entry, tool_use block)`` pairs that never got a ``tool_result``.""" + closed: set[str] = set() + uses: list[tuple[Mapping[str, Any], dict[str, Any]]] = [] + for entry in entries: + content = _entry_content(entry) + if entry.get("type") == "user": + closed.update( + str(b.get("tool_use_id", "")) + for b in content + if b.get("type") == "tool_result" + ) + elif entry.get("type") == "assistant": + uses.extend((entry, b) for b in content if b.get("type") == "tool_use") + return [ + (owner, block) + for owner, block in uses + if str(block.get("id", "")) not in closed + ] + + +def close_open_tool_uses( + entries: list[Mapping[str, Any]], *, text: str = _INTERRUPTED +) -> list[dict[str, Any]]: + """Synthetic ``tool_result`` entries on the leaf, one per open ``tool_use``. + + Appending the result to the session store gives the next ``resume`` a + transcript the CLI accepts: an assistant message ending in ``tool_use`` + without its result is rejected by the API on the next call. + """ + pending = open_tool_uses(entries) + if not pending: + return [] + leaf = next( + ( + e + for e in reversed(entries) + if e.get("type") in _PROMPT_TYPES and e.get("uuid") + ), + None, + ) + parent = str(leaf["uuid"]) if leaf is not None else None + stamp = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S.000Z") + out: list[dict[str, Any]] = [] + for owner, block in pending: + uid = _new_uuid() + result = { + "type": "tool_result", + "tool_use_id": block.get("id", ""), + "content": text, + "is_error": True, + } + out.append( + { + "parentUuid": parent, + "promptId": owner.get("promptId"), + "type": "user", + "message": {"role": "user", "content": [result]}, + "uuid": uid, + "timestamp": stamp, + "toolUseResult": text, + "sourceToolAssistantUUID": owner.get("uuid"), + **{k: owner[k] for k in _COMMON_KEYS if k in owner}, + } + ) + parent = uid + return out + + +_COMMON_KEYS = ( + "isSidechain", + "userType", + "entrypoint", + "cwd", + "sessionId", + "version", + "gitBranch", +) + + +def window_entries( + entries: Iterable[Mapping[str, Any]], *, window: int | None = None +) -> list[dict[str, Any]]: + """``user``/``assistant`` entries only, last ``window`` prompts, relinked.""" + kept = [ + dict(e) + for e in entries + if e.get("type") in _PROMPT_TYPES and isinstance(e.get("uuid"), str) + ] + if window is not None and window > 0: + starts = [i for i, e in enumerate(kept) if _is_prompt_entry(e)] + if len(starts) > window: + kept = kept[starts[-window] :] + return _relink(kept) + + +def strip_tool_entries(entries: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: + """Drop ``tool_result`` entries and ``tool_use``/``thinking`` blocks; relink.""" + out: list[dict[str, Any]] = [] + for raw in entries: + entry = dict(raw) + message = entry.get("message") + if not isinstance(message, dict): + out.append(entry) + continue + content = message.get("content") + if isinstance(content, list): + blocks = [ + b + for b in content + if isinstance(b, dict) + and b.get("type") not in ("tool_use", "tool_result", "thinking") + ] + if not blocks: + continue + entry["message"] = {**message, "content": blocks} + out.append(entry) + return _relink(out) + + +def prompt_count(entries: Iterable[Mapping[str, Any]]) -> int: + return sum(1 for e in entries if _is_prompt_entry(e)) + + +def render_messages( + messages: Iterable[Mapping[str, Any]], *, window: int | None = None +) -> str: + """Plain-text projection for ``read_conversation``: ``user:``/``assistant:`` turns. + + Tool calls collapse to a one-line summary per assistant turn; tool + results and thinking are dropped. + """ + turns: list[str] = [] + tools: list[str] = [] + current: list[str] = [] + + def flush() -> None: + if not current and not tools: + return + body = "\n\n".join(current).strip() + if tools: + body = (body + "\n" if body else "") + "(tools: " + ", ".join(tools) + ")" + turns.append("assistant:\n" + body) + current.clear() + tools.clear() + + for message in messages: + role = message.get("role") + content = message.get("content") + if role == "user": + if _tool_results(content): + continue + flush() + turns.append("user:\n" + _text_of_content(content)) + continue + for block in _assistant_blocks(content): + if block.get("type") == "text" and block.get("text"): + current.append(str(block["text"])) + elif block.get("type") == "tool_use": + tools.append(str(block.get("name", ""))) + flush() + if window is not None and window > 0: + starts = [i for i, t in enumerate(turns) if t.startswith("user:")] + if len(starts) > window: + turns = turns[starts[-window] :] + return "\n\n".join(turns) + + +def _entry_content(entry: Mapping[str, Any]) -> list[dict[str, Any]]: + message = entry.get("message") + if not isinstance(message, dict): + return [] + content = message.get("content") + if not isinstance(content, list): + return [] + return [b for b in content if isinstance(b, dict)] + + +def _is_prompt_entry(entry: Mapping[str, Any]) -> bool: + if entry.get("type") != "user": + return False + message = entry.get("message") + if not isinstance(message, dict): + return False + content = message.get("content") + if isinstance(content, str): + return True + return not _tool_results(content) + + +def _relink(entries: list[dict[str, Any]]) -> list[dict[str, Any]]: + parent: str | None = None + for entry in entries: + entry["parentUuid"] = parent + parent = entry.get("uuid") + return entries + + +def _text_of_content(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + return "\n\n".join( + str(b.get("text", "")) + for b in content + if isinstance(b, dict) and b.get("type") == "text" and b.get("text") + ) + return "" diff --git a/src/beaver_gateway/core/turn_capture.py b/src/beaver_gateway/core/turn_capture.py index fb3c497..5526b13 100644 --- a/src/beaver_gateway/core/turn_capture.py +++ b/src/beaver_gateway/core/turn_capture.py @@ -35,3 +35,5 @@ class TurnCapture: """Backend session that ran the turn; persist and pass back as ``session_id``.""" usage: TurnUsage | None = None + interrupted: bool = False + """The turn was cut by ``interrupt()`` (urgent inject), not by an error.""" diff --git a/src/beaver_gateway/frontends/_sse.py b/src/beaver_gateway/frontends/_sse.py new file mode 100644 index 0000000..66f7756 --- /dev/null +++ b/src/beaver_gateway/frontends/_sse.py @@ -0,0 +1,63 @@ +"""Server-sent events helpers shared by the markdown and API frontends.""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +__all__ = ["HEARTBEAT_INTERVAL", "SSE_HEADERS", "events_with_heartbeat", "sse_pack"] + +HEARTBEAT_INTERVAL = 15.0 +"""Seconds of backend silence before a comment frame keeps the socket warm.""" + +SSE_HEADERS = { + "Cache-Control": "no-cache, no-transform", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", +} + +KEEPALIVE = b": keepalive\n\n" + + +async def events_with_heartbeat( + events: AsyncIterator[Any], interval: float = HEARTBEAT_INTERVAL +) -> AsyncIterator[Any]: + """Pass ``events`` through, yielding ``None`` after ``interval`` seconds of silence. + + One in-flight ``__anext__`` task is reused across timeouts: a second + consumer on the same async generator raises ``RuntimeError``. + Cancellation of the outer scope cancels that task instead of leaving + it dangling. + """ + src = events.__aiter__() + next_task: asyncio.Task[Any] | None = None + try: + while True: + if next_task is None: + next_task = asyncio.ensure_future(src.__anext__()) + done, _pending = await asyncio.wait({next_task}, timeout=interval) + if not done: + yield None + continue + task = next_task + next_task = None + try: + result = task.result() + except StopAsyncIteration: + return + yield result + finally: + if next_task is not None and not next_task.done(): + next_task.cancel() + with contextlib.suppress(BaseException): + await next_task + + +def sse_pack(event: str, data: dict[str, Any]) -> bytes: + body = json.dumps(data, ensure_ascii=False) + return f"event: {event}\ndata: {body}\n\n".encode() diff --git a/src/beaver_gateway/frontends/api/__init__.py b/src/beaver_gateway/frontends/api/__init__.py new file mode 100644 index 0000000..d568772 --- /dev/null +++ b/src/beaver_gateway/frontends/api/__init__.py @@ -0,0 +1,5 @@ +"""``ApiFrontend`` - the conversations API and event stream (§3.9).""" + +from beaver_gateway.frontends.api.frontend import ApiFrontend + +__all__ = ["ApiFrontend"] diff --git a/src/beaver_gateway/frontends/api/frontend.py b/src/beaver_gateway/frontends/api/frontend.py new file mode 100644 index 0000000..e21c9a5 --- /dev/null +++ b/src/beaver_gateway/frontends/api/frontend.py @@ -0,0 +1,491 @@ +"""``ApiFrontend`` - ``/api/conversations``, SSE events, sessions, usage (§3.9). + +Bearer scope ``api``. Every write goes through ``core/conversations``; the +frontend only shapes JSON. ``/api/events`` and +``/api/conversations/{id}/events`` replay the gateway bus as SSE with the +same keepalive the markdown frontend uses, so a proxy never sees a +silent socket. +""" + +from __future__ import annotations + +import json +import logging +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Any, cast + +from fastapi import FastAPI, HTTPException, Request, status +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse, StreamingResponse +from sqlalchemy import func +from sqlalchemy import select as sa_select +from sqlmodel import col + +from beaver_gateway.core import audit +from beaver_gateway.core.conversations import KINDS, SEEDS +from beaver_gateway.frontends._auth import require_token +from beaver_gateway.frontends._sse import ( + KEEPALIVE, + SSE_HEADERS, + events_with_heartbeat, + sse_pack, +) +from beaver_gateway.frontends.base import Frontend +from beaver_gateway.storage.models import Usage + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from beaver_gateway.core.conversations import Conversations + from beaver_gateway.frontends.base import GatewayRuntime + from beaver_gateway.storage.models import Conversation + +_log = logging.getLogger("beaver_gateway.frontends.api") + +__all__ = ["ApiFrontend"] + +SCOPE = "api" + + +class ApiFrontend(Frontend): + def __init__( + self, + *, + host: str = "0.0.0.0", # noqa: S104 + port: int = 8004, + public_base_url: str | None = None, + ) -> None: + self.host = host + self.port = port + self.public_base_url = public_base_url.rstrip("/") if public_base_url else None + self._app: FastAPI | None = None + + def configure(self, runtime: GatewayRuntime) -> None: + if runtime.conversations is None or runtime.bus is None: + msg = "ApiFrontend needs runtime.conversations and runtime.bus" + raise RuntimeError(msg) + self._app = _build_app(runtime) + + async def serve(self) -> None: + import uvicorn + + if self._app is None: + msg = "configure() must be called before serve()" + raise RuntimeError(msg) + server = uvicorn.Server( + uvicorn.Config(self._app, host=self.host, port=self.port, log_level="info") + ) + await server.serve() + + +def _build_app(runtime: GatewayRuntime) -> FastAPI: # noqa: PLR0915 + app = FastAPI(title="beaver-gateway / API") + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=False, + allow_methods=["*"], + allow_headers=["*"], + ) + conversations = cast("Conversations", runtime.conversations) + + async def body_of(request: Request) -> dict[str, Any]: + if not await request.body(): + return {} + try: + data = await request.json() + except json.JSONDecodeError as exc: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, f"invalid JSON: {exc}" + ) from exc + if not isinstance(data, dict): + raise HTTPException(status.HTTP_400_BAD_REQUEST, "body must be an object") + return data + + async def conv_of(public_id: str) -> Conversation: + conv = await conversations.get(public_id) + if conv is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, f"unknown conversation {public_id}" + ) + return conv + + def text_of(data: dict[str, Any], key: str = "text") -> str: + text = data.get(key) + if not isinstance(text, str) or not text.strip(): + raise HTTPException(status.HTTP_400_BAD_REQUEST, f"missing `{key}`") + return text + + def int_or_none(data: dict[str, Any], key: str) -> int | None: + value = data.get(key) + if value is None: + return None + if not isinstance(value, int) or value < 1: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, f"`{key}` must be a positive int" + ) + return value + + @app.get("/healthz") + async def healthz() -> dict[str, str]: + return {"status": "ok"} + + @app.get("/api/conversations") + async def list_conversations(request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + q = request.query_params + rows = await conversations.find( + status=q.get("status"), kind=q.get("kind"), limit=int(q.get("limit", "200")) + ) + return {"conversations": [conversations.public(r) for r in rows]} + + @app.post("/api/conversations", status_code=status.HTTP_201_CREATED) + async def create_conversation(request: Request) -> dict[str, Any]: + token = await require_token(request, runtime, scope=SCOPE) + data = await body_of(request) + kind = str(data.get("kind") or "deep") + agent = data.get("agent") + if kind not in KINDS or kind == "fork": + raise HTTPException( + status.HTTP_400_BAD_REQUEST, f"kind must be one of {KINDS[:-1]}" + ) + if not isinstance(agent, str) or agent not in runtime.agents: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, "unknown or missing `agent`" + ) + seed = str(data.get("seed") or "clean") + if seed not in SEEDS: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, f"seed must be one of {SEEDS}" + ) + parent = await conv_of(str(data["parent"])) if data.get("parent") else None + try: + conv = await conversations.spawn( + kind=kind, + agent=agent, + seed=seed, + parent=parent, + text=data.get("text"), + title=data.get("title"), + window=int_or_none(data, "window"), + origin="api", + ) + except ValueError as exc: + raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc + await audit.log( + runtime, + actor=f"token:{token}", + kind="api_spawn", + agent_name=agent, + conversation=conv.external_id, + seed=seed, + ) + return await conversations.describe(conv) + + @app.get("/api/conversations/{public_id}") + async def get_conversation(public_id: str, request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + out = await conversations.describe(conv) + out["queue"] = [ + { + "id": i.id, + "priority": i.priority, + "origin": i.origin, + "status": i.status, + "created_at": i.created_at.isoformat(), + "text": i.text[:200], + } + for i in await conversations.queue.recent(cast("int", conv.id), limit=20) + ] + return out + + @app.get("/api/conversations/{public_id}/messages") + async def get_messages(public_id: str, request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + raw = request.query_params.get("window") + window = int(raw) if raw and raw.isdigit() else None + return { + "id": conv.external_id, + "text": await conversations.read(conv, window=window), + } + + @app.post( + "/api/conversations/{public_id}/messages", status_code=status.HTTP_202_ACCEPTED + ) + async def post_message(public_id: str, request: Request) -> dict[str, Any]: + token = await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + data = await body_of(request) + item = await conversations.post( + conv, text_of(data), origin=str(data.get("origin") or "user") + ) + await audit.log( + runtime, + actor=f"token:{token}", + kind="api_message", + agent_name=conv.agent_name, + conversation=conv.external_id, + ) + return {"id": conv.external_id, "item": item.id, "status": item.status} + + @app.post( + "/api/conversations/{public_id}/inject", status_code=status.HTTP_202_ACCEPTED + ) + async def post_inject(public_id: str, request: Request) -> dict[str, Any]: + token = await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + data = await body_of(request) + urgency = str(data.get("urgency") or "normal") + if urgency not in ("normal", "urgent"): + raise HTTPException( + status.HTTP_400_BAD_REQUEST, "urgency must be normal|urgent" + ) + item = await conversations.inject( + conv, + text_of(data), + urgency=cast("Any", urgency), + origin=str(data.get("origin") or "api"), + ) + await audit.log( + runtime, + actor=f"token:{token}", + kind="api_inject", + agent_name=conv.agent_name, + conversation=conv.external_id, + urgency=urgency, + ) + return {"id": conv.external_id, "item": item.id, "priority": item.priority} + + @app.post("/api/conversations/{public_id}/say") + async def post_say(public_id: str, request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + return await conversations.say(conv, text_of(await body_of(request))) + + @app.post( + "/api/conversations/{public_id}/branch", status_code=status.HTTP_201_CREATED + ) + async def post_branch(public_id: str, request: Request) -> dict[str, Any]: + token = await require_token(request, runtime, scope=SCOPE) + parent = await conv_of(public_id) + data = await body_of(request) + seed = str(data.get("seed_mode") or data.get("seed") or "morning") + if seed not in SEEDS: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, f"seed must be one of {SEEDS}" + ) + agent = str(data.get("agent") or parent.agent_name) + try: + child = await conversations.spawn( + kind="branch", + agent=agent, + seed=seed, + parent=parent, + text=data.get("text"), + title=data.get("title"), + window=int_or_none(data, "window"), + origin="api", + ) + except (ValueError, LookupError) as exc: + raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc + await audit.log( + runtime, + actor=f"token:{token}", + kind="api_branch", + agent_name=agent, + conversation=child.external_id, + parent=parent.external_id, + seed=seed, + ) + return await conversations.describe(child) + + @app.post("/api/conversations/{public_id}/merge") + async def post_merge(public_id: str, request: Request) -> dict[str, Any]: + token = await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + try: + result = await conversations.merge(conv) + except (ValueError, LookupError) as exc: + raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc + await audit.log( + runtime, + actor=f"token:{token}", + kind="api_merge", + agent_name=conv.agent_name, + conversation=conv.external_id, + ) + return { + "id": conv.external_id, + "status": "merged", + "fork": result.conversation.external_id, + "text": result.text, + } + + @app.post("/api/conversations/{public_id}/fork") + async def post_fork(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: + result = await conversations.fork( + conv, + text_of(data, "prompt"), + window=int_or_none(data, "window"), + strip_tools=bool(data.get("strip_tools", False)), + ) + except (ValueError, RuntimeError) as exc: + raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc + return { + "id": conv.external_id, + "fork": result.conversation.external_id, + "text": result.text, + } + + @app.post("/api/conversations/{public_id}/bind") + async def post_bind(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) + frontend = text_of(data, "frontend") + external_id = text_of(data, "external_id") + await conversations.bind( + conv, + frontend=frontend, + external_id=external_id, + visible=bool(data.get("visible", True)), + ) + return await conversations.describe(conv) + + @app.patch("/api/conversations/{public_id}/flags") + async def patch_flags(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) + return conversations.public(await conversations.set_flags(conv, data)) + + @app.patch("/api/conversations/{public_id}") + async def patch_conversation(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: + if isinstance(data.get("status"), str): + conv = await conversations.set_status(conv, data["status"]) + if isinstance(data.get("title"), str): + conv = await conversations.set_title(conv, data["title"]) + except ValueError as exc: + raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc + return conversations.public(conv) + + @app.get("/api/conversations/{public_id}/events") + async def conversation_events(public_id: str, request: Request) -> Any: + await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + return _sse(runtime, conversation_id=conv.external_id) + + @app.get("/api/events") + async def all_events(request: Request) -> Any: + await require_token(request, runtime, scope=SCOPE) + return _sse(runtime, conversation_id=None) + + @app.get("/api/sessions") + async def sessions(request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + pool = runtime.pool + return { + "rss": pool.rss() if pool is not None else None, + "rss_limit": pool.rss_limit if pool is not None else None, + "sessions": pool.snapshot() if pool is not None else [], + } + + @app.get("/api/schedules") + async def schedules(request: Request) -> dict[str, Any]: + 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": s.execute_at.isoformat(), + "text": s.text, + "delivered_at": s.delivered_at.isoformat() + if s.delivered_at + else None, + } + for s in await conversations.schedules(conv) + ] + } + + @app.get("/api/usage") + async def usage(request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + hours = float(request.query_params.get("hours", "24")) + since = datetime.now(UTC) - timedelta(hours=hours) + columns = ( + func.count(), + func.coalesce(func.sum(Usage.input_tokens), 0), + func.coalesce(func.sum(Usage.output_tokens), 0), + func.coalesce(func.sum(Usage.cache_read_tokens), 0), + func.coalesce(func.sum(Usage.cache_creation_tokens), 0), + func.coalesce(func.sum(Usage.cost_usd), 0.0), + ) + async with runtime.db.session() as session: + by_agent = ( + await session.execute( # ty: ignore[deprecated] + sa_select(col(Usage.agent_name), *columns) + .where(col(Usage.ts) >= since.replace(tzinfo=None)) + .group_by(col(Usage.agent_name)) + ) + ).all() + by_conversation = ( + await session.execute( # ty: ignore[deprecated] + sa_select(col(Usage.conversation_id), *columns) + .where(col(Usage.ts) >= since.replace(tzinfo=None)) + .group_by(col(Usage.conversation_id)) + ) + ).all() + return { + "since": since.isoformat(timespec="seconds"), + "by_agent": [_usage_row("agent", r) for r in by_agent], + "by_conversation": [_usage_row("conversation", r) for r in by_conversation], + } + + @app.exception_handler(HTTPException) + async def http_error(_request: Request, exc: HTTPException) -> JSONResponse: + return JSONResponse( + status_code=exc.status_code, + content={"error": exc.detail}, + headers=exc.headers, + ) + + return app + + +def _usage_row(label: str, row: Any) -> dict[str, Any]: + key, turns, inp, out, cache_read, cache_creation, cost = row + return { + label: key, + "turns": turns, + "input": inp, + "output": out, + "cache_read": cache_read, + "cache_creation": cache_creation, + "cost_usd": round(float(cost or 0.0), 4), + } + + +def _sse(runtime: GatewayRuntime, *, conversation_id: str | None) -> StreamingResponse: + async def gen() -> AsyncIterator[bytes]: + stream = runtime.bus.stream(conversation_id=conversation_id) + yield sse_pack("hello", {"type": "hello", "conversation_id": conversation_id}) + async for event in events_with_heartbeat(stream): + if event is None: + yield KEEPALIVE + continue + yield sse_pack(str(event["type"]), event) + + return StreamingResponse(gen(), media_type="text/event-stream", headers=SSE_HEADERS) diff --git a/src/beaver_gateway/frontends/base.py b/src/beaver_gateway/frontends/base.py index 9babb7d..c1686fd 100644 --- a/src/beaver_gateway/frontends/base.py +++ b/src/beaver_gateway/frontends/base.py @@ -12,7 +12,7 @@ from __future__ import annotations from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Mapping, Sequence @@ -77,6 +77,12 @@ class GatewayRuntime: # to keep ``anthropic.types`` out of the runtime import graph for # this base module. turn_log_handlers: list[TurnLogHandler] = field(default_factory=list) + # M1b: conversations service, event bus and the shared session pool. + # ``Any`` for the same import-graph reason as above; ``None`` only in + # tests that build a runtime without them. + conversations: Any = None + bus: Any = None + pool: Any = None class Frontend(ABC): diff --git a/src/beaver_gateway/frontends/markdown/frontend.py b/src/beaver_gateway/frontends/markdown/frontend.py index 14ea497..501acb5 100644 --- a/src/beaver_gateway/frontends/markdown/frontend.py +++ b/src/beaver_gateway/frontends/markdown/frontend.py @@ -13,7 +13,13 @@ Concurrency model: an in-memory ``set[Path]`` of files currently in flight. Two concurrent requests for the same file → the second gets 409. The set is single-process (one gateway instance) — that's by design; the markdown frontend is the only writer in its vault from -the gateway side. +the gateway side. The turn itself runs through ``core/conversations`` +(one turn per conversation, ``running_turn`` in the DB), so a message +posted to the same conversation via ``/api`` waits its turn. + +A chat file is a ``deep`` conversation bound as +``(markdown, )``; frontmatter carries only ``agent`` +and ``conversation_id`` (§3.10), tool calls are never rendered. Cross-frontend logging: when ``log_all_chats=True``, ``configure()`` registers a handler on ``runtime.turn_log_handlers`` so every other @@ -24,7 +30,6 @@ shape. from __future__ import annotations -import asyncio import contextlib import json import logging @@ -33,7 +38,7 @@ import tempfile import time from collections.abc import AsyncIterator from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import aiofile from anthropic.types import RawContentBlockStopEvent @@ -44,29 +49,28 @@ from fastapi.responses import JSONResponse, StreamingResponse from beaver_gateway.core import audit from beaver_gateway.core.conversation_store import ( diff_and_fork, - load_conversation, load_messages, - mint_conversation, rewrite_messages, - set_session_id, ) from beaver_gateway.core.turn_capture import TurnCapture from beaver_gateway.core.turn_record import TurnRecord from beaver_gateway.frontends._accumulate import StreamAccumulator from beaver_gateway.frontends._auth import require_token +from beaver_gateway.frontends._sse import ( + KEEPALIVE, + SSE_HEADERS, + events_with_heartbeat, + sse_pack, +) from beaver_gateway.frontends.base import Frontend from beaver_gateway.frontends.markdown import parser, renderer -from beaver_gateway.frontends.markdown.crossfront import ( - CrossFrontendLogger, - fingerprint_messages, -) +from beaver_gateway.frontends.markdown.crossfront import CrossFrontendLogger if TYPE_CHECKING: from collections.abc import AsyncIterator, Callable - from anthropic.types import MessageParam - from beaver_gateway.frontends.base import GatewayRuntime + from beaver_gateway.storage.models import Conversation _log = logging.getLogger("beaver_gateway.frontends.markdown") @@ -89,12 +93,7 @@ _STREAM_FLUSH_DEBOUNCE = 0.4 # disk round-trip). _SSE_FLUSH_DEBOUNCE = 0.1 -# Interval between SSE comment-frames sent when the backend is silent -# (e.g. claude is mid-thinking on a large context). The Obsidian plugin -# and any intermediate proxies will hold the connection open as long as -# bytes keep flowing; a comment-frame is the cheapest legal SSE keepalive. -# Set well under typical proxy/client idle timeouts (60s). -_SSE_HEARTBEAT_INTERVAL = 15.0 +FRONTEND = "markdown" class MarkdownFrontend(Frontend): @@ -134,6 +133,9 @@ class MarkdownFrontend(Frontend): self._crossfront: CrossFrontendLogger | None = None def configure(self, runtime: GatewayRuntime) -> None: + if runtime.conversations is None: + msg = "MarkdownFrontend needs runtime.conversations" + raise RuntimeError(msg) self._runtime = runtime self.vault_path.mkdir(parents=True, exist_ok=True) if self.log_all_chats: @@ -286,17 +288,7 @@ class MarkdownFrontend(Frontend): self._busy.discard(file_path) return StreamingResponse( - gen(), - media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache, no-transform", - "Connection": "keep-alive", - # nginx default-buffers SSE bodies; this header tells - # both nginx and uvicorn-behind-proxy to flush as we - # write. Harmless if the deployment has no reverse - # proxy in front. - "X-Accel-Buffering": "no", - }, + gen(), media_type="text/event-stream", headers=SSE_HEADERS ) return app @@ -384,16 +376,19 @@ class MarkdownFrontend(Frontend): # stored history, and feed the aligned messages to the backend # - see ``core/conversation_store.py`` for the full rationale. conv, conv_external_id, stored_msgs = await self._resolve_conversation( - runtime=runtime, metadata=parsed.metadata, agent_name=agent.name + runtime=runtime, + metadata=parsed.metadata, + agent_name=agent.name, + file_path=file_path, ) outcome = diff_and_fork(stored=stored_msgs, incoming=parsed.turns) capture = TurnCapture() - events = backend.complete( - agent=agent, + events = runtime.conversations.turn( + conv, messages=outcome.messages, - system=None, + origin="user", capture=capture, - **_session_options(conv, outcome.divergence_index), + use_session=outcome.divergence_index is None, ) try: message = await self._stream_to_file( @@ -422,7 +417,7 @@ class MarkdownFrontend(Frontend): await self._persist_canonical_history( runtime=runtime, - conversation_id=conv.id, + conversation_id=cast("int", conv.id), persist_messages=outcome.persist_messages, new_user_text=parsed.turns[-1].text, capture=capture, @@ -487,7 +482,7 @@ class MarkdownFrontend(Frontend): elif content_override is None: file_text = await _read_or_empty(file_path) else: - yield _sse_pack( + yield sse_pack( "error", { "status_code": status.HTTP_400_BAD_REQUEST, @@ -503,7 +498,7 @@ class MarkdownFrontend(Frontend): default=self.default_agent, ) if not agent_name: - yield _sse_pack( + yield sse_pack( "error", { "status_code": status.HTTP_400_BAD_REQUEST, @@ -516,7 +511,7 @@ class MarkdownFrontend(Frontend): return if not parsed.messages: - yield _sse_pack( + yield sse_pack( "done", { "status": "nothing_to_do", @@ -527,7 +522,7 @@ class MarkdownFrontend(Frontend): return if parser.last_role(parsed.messages) == "assistant": - yield _sse_pack( + yield sse_pack( "done", { "status": "nothing_to_do", @@ -539,7 +534,7 @@ class MarkdownFrontend(Frontend): agent = runtime.agents.get(agent_name) if agent is None: - yield _sse_pack( + yield sse_pack( "error", { "status_code": status.HTTP_404_NOT_FOUND, @@ -549,7 +544,7 @@ class MarkdownFrontend(Frontend): return backend = runtime.backends.get(agent.name) if backend is None: - yield _sse_pack( + yield sse_pack( "error", { "status_code": status.HTTP_503_SERVICE_UNAVAILABLE, @@ -575,7 +570,10 @@ class MarkdownFrontend(Frontend): ) conv, conv_external_id, stored_msgs = await self._resolve_conversation( - runtime=runtime, metadata=parsed.metadata, agent_name=agent.name + runtime=runtime, + metadata=parsed.metadata, + agent_name=agent.name, + file_path=file_path, ) _log.info( "chat/stream: file=%s conv_external_id=%s conv_id=%d " @@ -602,12 +600,12 @@ class MarkdownFrontend(Frontend): agent.name, conv.session_id, ) - events = backend.complete( - agent=agent, + events = runtime.conversations.turn( + conv, messages=outcome.messages, - system=None, + origin="user", capture=capture, - **_session_options(conv, outcome.divergence_index), + use_session=outcome.divergence_index is None, ) acc = StreamAccumulator() @@ -624,13 +622,9 @@ class MarkdownFrontend(Frontend): return _reattach_frontmatter(parsed.metadata, new_body) try: - async for ev in _events_with_heartbeat(events): + async for ev in events_with_heartbeat(events): if ev is None: - # Backend is quiet (claude mid-thinking, MCP slow, - # whatever). SSE comment-frame keeps the TCP socket - # warm so the plugin / uvicorn / any reverse proxy - # doesn't time the request out before we finish. - yield b": keepalive\n\n" + yield KEEPALIVE continue acc.feed(ev) now = time.monotonic() @@ -643,7 +637,7 @@ class MarkdownFrontend(Frontend): # render to the same prefix as before they closed # (we don't surface the tool-call args in markdown). if payload is not None and payload != last_payload: - yield _sse_pack("delta", {"new_content": payload}) + yield sse_pack("delta", {"new_content": payload}) last_payload = payload last_flush = now except Exception as exc: # noqa: BLE001 — wire any backend failure as an SSE error frame @@ -662,7 +656,7 @@ class MarkdownFrontend(Frontend): await _write_atomic( file_path, _reattach_frontmatter(parsed.metadata, new_body) ) - yield _sse_pack( + yield sse_pack( "error", { "status_code": status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -684,7 +678,7 @@ class MarkdownFrontend(Frontend): await self._persist_canonical_history( runtime=runtime, - conversation_id=conv.id, + conversation_id=cast("int", conv.id), persist_messages=outcome.persist_messages, new_user_text=parsed.turns[-1].text, capture=capture, @@ -704,7 +698,7 @@ class MarkdownFrontend(Frontend): except Exception: # noqa: BLE001 _log.exception("turn_log_handler raised; continuing") - yield _sse_pack( + yield sse_pack( "done", { "status": "ok", @@ -795,71 +789,53 @@ class MarkdownFrontend(Frontend): rendered = renderer.render_assistant_message(message) new_body = renderer.append_to_body(parsed.body, rendered) new_body = renderer.append_to_body(new_body, renderer.USER_SCAFFOLD) - # Recompute fingerprint so a future cross-frontend hit on this - # same conversation can find it. Stored as hex string in - # frontmatter — only the markdown frontend reads it. - assistant_param: MessageParam = { - "role": "assistant", - "content": _flatten_assistant_text(message), - } - updated_messages: list[MessageParam] = [*parsed.messages, assistant_param] updated_metadata = dict(parsed.metadata) + updated_metadata.pop("fingerprint", None) updated_metadata["agent"] = agent_name updated_metadata["conversation_id"] = conv_external_id - updated_metadata["fingerprint"] = fingerprint_messages(updated_messages) new_content = _reattach_frontmatter(updated_metadata, new_body) if write_disk: await _write_atomic(file_path, new_content) return new_content async def _resolve_conversation( - self, *, runtime: GatewayRuntime, metadata: dict[str, Any], agent_name: str - ) -> tuple[Any, str, list[dict[str, Any]]]: - """Resolve the conversation row + stored messages for this request. + self, + *, + runtime: GatewayRuntime, + metadata: dict[str, Any], + agent_name: str, + file_path: Path, + ) -> tuple[Conversation, str, list[dict[str, Any]]]: + """Resolve the ``deep`` conversation for this file + its stored messages. - Looks up by frontmatter ``conversation_id``, mints a new row if - missing, and returns ``(conv, external_id, stored_messages)``. - ``conv.id`` is guaranteed non-None because both - ``load_conversation`` (after refresh on a committed row) and - ``mint_conversation`` (post-commit refresh) populate it. We - coerce with a runtime check so the rest of the handler can - treat it as ``int``. + Frontmatter ``conversation_id`` wins; a file that lost it is found + by its visible ``(markdown, path)`` binding; otherwise a new + conversation is created. The binding follows the file: a moved + chat re-binds to its new path on the next turn. """ + conversations = runtime.conversations + rel = file_path.relative_to(self.vault_path).as_posix() raw = metadata.get("conversation_id") - lookup_id = raw if isinstance(raw, str) and raw else None + conv = await conversations.get(raw) if isinstance(raw, str) and raw else None + if conv is None: + conv = await conversations.find_bound(frontend=FRONTEND, external_id=rel) + if conv is None: + conv = await conversations.create( + kind="deep", agent=agent_name, origin=FRONTEND, title=file_path.stem + ) + _log.info("minted conversation %s for %s", conv.external_id, rel) + bound = [ + b + for b in await conversations.bindings(conv) + if b.frontend == FRONTEND and b.visible and b.external_id == rel + ] + if not bound: + await conversations.bind(conv, frontend=FRONTEND, external_id=rel) + await conversations.touch_user(conv) + if conv.id is None: + msg = "conversation row missing primary key after commit" + raise RuntimeError(msg) async with runtime.db.session() as session: - conv = None - if lookup_id is not None: - conv = await load_conversation( - session, frontend="markdown", external_id=lookup_id - ) - if conv is None: - _log.info( - "_resolve_conversation: frontmatter conv_id=%s " - "not found in DB, will mint new", - lookup_id, - ) - else: - _log.info( - "_resolve_conversation: LOADED existing conv " - "id=%d external_id=%s", - conv.id or -1, - conv.external_id, - ) - if conv is None: - conv = await mint_conversation( - session, frontend="markdown", agent_name=agent_name - ) - _log.info( - "_resolve_conversation: MINTED new conv " - "id=%d external_id=%s agent=%s", - conv.id or -1, - conv.external_id, - agent_name, - ) - if conv.id is None: - msg = "conversation row missing primary key after commit" - raise RuntimeError(msg) stored = await load_messages(session, conversation_id=conv.id) return conv, conv.external_id, stored @@ -894,12 +870,6 @@ class MarkdownFrontend(Frontend): await rewrite_messages( session, conversation_id=conversation_id, messages=canonical ) - if capture.session_id is not None: - await set_session_id( - session, - conversation_id=conversation_id, - session_id=capture.session_id, - ) _log.info( "_persist_canonical_history: conv_id=%d DB committed", conversation_id ) @@ -927,71 +897,6 @@ class MarkdownFrontend(Frontend): # ---- module-level utilities ---------------------------------------------- -def _session_options(conv: Any, divergence_index: int | None) -> dict[str, Any]: - """Backend options that pin the turn to the conversation's live session. - - A divergence means the file's history no longer matches what the - session saw, so the stored ``session_id`` is withheld and the backend - seeds a fresh one from the aligned messages. - """ - return { - "conversation_id": conv.external_id, - "session_id": conv.session_id if divergence_index is None else None, - } - - -async def _events_with_heartbeat( - events: AsyncIterator[Any], interval: float = _SSE_HEARTBEAT_INTERVAL -) -> AsyncIterator[Any]: - """Wrap an async event stream with idle-time heartbeat markers. - - Yields ``None`` every ``interval`` seconds during silence; real - events pass through unchanged. When the wrapped iterator is - exhausted, this generator returns. Cancellation propagates: if the - outer scope is cancelled we cancel the pending ``__anext__`` task - instead of leaving it dangling. - """ - src = events.__aiter__() - next_task: asyncio.Task[Any] | None = None - try: - while True: - # Reuse the in-flight task across timeouts. Spawning a fresh - # ``__anext__()`` while the previous one is still pending - # puts two consumers on the same async generator — that - # raises ``RuntimeError: anext(): asynchronous generator is - # already running``. - if next_task is None: - next_task = asyncio.ensure_future(src.__anext__()) - done, _pending = await asyncio.wait({next_task}, timeout=interval) - if not done: - yield None - continue - task = next_task - next_task = None - try: - result = task.result() - except StopAsyncIteration: - return - yield result - finally: - if next_task is not None and not next_task.done(): - next_task.cancel() - with contextlib.suppress(BaseException): - await next_task - - -def _sse_pack(event: str, data: dict[str, Any]) -> bytes: - r"""Format one Server-Sent Event frame. - - Uses named events (``event: ``) so the plugin can dispatch on - type without parsing JSON discriminators. ``ensure_ascii=False`` so - multibyte content rides through verbatim instead of becoming - ``\uXXXX`` blobs that bloat the wire. - """ - body = json.dumps(data, ensure_ascii=False) - return f"event: {event}\ndata: {body}\n\n".encode() - - async def _read_or_empty(path: Path) -> str: """Return file contents, or empty string if the file doesn't exist.""" # ``path.exists()`` here is a metadata stat — microseconds — and @@ -1083,21 +988,6 @@ def _fallback_synthesized(message: Any) -> list[dict[str, Any]]: return [{"role": "assistant", "content": content}] -def _flatten_assistant_text(message: Any) -> str: - """Pull all text blocks from an assistant ``Message`` and join them. - - Used when we need the assistant content as a plain string for - fingerprinting / equality with a parser-shaped history (parser - already drops thinking + tool_use from assistant turns). - """ - chunks = [ - getattr(block, "text", "") or "" - for block in getattr(message, "content", ()) - if getattr(block, "type", None) == "text" - ] - return "\n\n".join(c for c in chunks if c) - - def _render_error_block(exc: BaseException) -> str: """Render a backend failure as an Assistant turn with a ``[!error]-`` callout.""" msg = str(exc) or exc.__class__.__name__ diff --git a/src/beaver_gateway/frontends/markdown/renderer.py b/src/beaver_gateway/frontends/markdown/renderer.py index 2c3e765..983eb0c 100644 --- a/src/beaver_gateway/frontends/markdown/renderer.py +++ b/src/beaver_gateway/frontends/markdown/renderer.py @@ -9,11 +9,10 @@ from other frontends). from __future__ import annotations -import json import re from typing import TYPE_CHECKING, Any, cast -from anthropic.types import Message, TextBlock, ThinkingBlock, ToolUseBlock +from anthropic.types import Message, TextBlock, ThinkingBlock if TYPE_CHECKING: from collections.abc import Iterable @@ -61,9 +60,8 @@ def render_assistant_message(message: Message) -> str: * ``ThinkingBlock`` → ``> [!thinking]-`` collapsed callout * ``TextBlock`` → plain text (the spoken answer) - * ``ToolUseBlock`` → ``> [!tool]- `` callout with the ``input`` - JSON quoted inside. Tool *results* are not persisted — see - module docstring on ``parser.py`` for why. + * ``ToolUseBlock`` → nothing (§3.10: tool calls never reach the file; + "what the agent is doing" is the activity panel fed by SSE) Blank lines separate adjacent blocks; trailing newline guarantees the next ``---`` / ``### User:`` marker lands on its own line. @@ -147,10 +145,7 @@ def _render_block(block: object) -> Iterable[str]: if isinstance(block, ThinkingBlock): yield from _render_thinking(block.thinking or "") return - if isinstance(block, ToolUseBlock): - yield from _render_tool_use(block) - return - # Unknown block type — skip silently rather than corrupting the file. + # Tool-use blocks and unknown block types never reach the file. def _render_thinking(text: str) -> Iterable[str]: @@ -159,17 +154,6 @@ def _render_thinking(text: str) -> Iterable[str]: yield f"> {line}" if line else ">" -def _render_tool_use(block: ToolUseBlock) -> Iterable[str]: - title = summarize_tool_input(block.name, block.input) - yield f"> [!tool]- {title}" - yield "> **input:**" - yield "> ```json" - pretty = json.dumps(block.input, indent=2, ensure_ascii=False, sort_keys=True) - for line in pretty.splitlines(): - yield f"> {line}" if line else ">" - yield "> ```" - - def adaptive_fence(content: str) -> str: """Return a backtick fence at least one longer than the longest run in ``content``. diff --git a/src/beaver_gateway/storage/__init__.py b/src/beaver_gateway/storage/__init__.py index 561ee73..8093662 100644 --- a/src/beaver_gateway/storage/__init__.py +++ b/src/beaver_gateway/storage/__init__.py @@ -16,13 +16,26 @@ from beaver_gateway.storage.db import ( revoke_token, touch_token, ) -from beaver_gateway.storage.models import AuditLog, Token, TranscriptEntry, Usage +from beaver_gateway.storage.models import ( + AuditLog, + Conversation, + ConversationBinding, + InjectQueueItem, + Schedule, + Token, + TranscriptEntry, + Usage, +) from beaver_gateway.storage.session_store import PostgresSessionStore __all__ = [ "AuditLog", + "Conversation", + "ConversationBinding", "Database", + "InjectQueueItem", "PostgresSessionStore", + "Schedule", "Token", "TranscriptEntry", "Usage", diff --git a/src/beaver_gateway/storage/db.py b/src/beaver_gateway/storage/db.py index 331f22e..b51dcfe 100644 --- a/src/beaver_gateway/storage/db.py +++ b/src/beaver_gateway/storage/db.py @@ -105,9 +105,29 @@ def _add_missing_columns(conn: Connection) -> None: if column.name in existing: continue kind = column.type.compile(conn.dialect) - conn.execute( - text(f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}") - ) + ddl = f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}" + default = _column_default(column) + if default is not None: + ddl += f" DEFAULT {default}" + conn.execute(text(ddl)) + + +def _column_default(column: Any) -> str | None: + """Literal for ``ADD COLUMN ... DEFAULT`` so old rows get the model default.""" + server_default = getattr(column, "server_default", None) + if server_default is not None: + return str(server_default.arg.text) + default = getattr(column, "default", None) + if default is None or default.is_callable: + return None + value = default.arg + if isinstance(value, bool): + return "TRUE" if value else "FALSE" + if isinstance(value, int | float): + return repr(value) + if isinstance(value, str): + return "'" + value.replace("'", "''") + "'" + return None # ---- Token CRUD --------------------------------------------------------- diff --git a/src/beaver_gateway/storage/models.py b/src/beaver_gateway/storage/models.py index 9c3a680..b2c3a55 100644 --- a/src/beaver_gateway/storage/models.py +++ b/src/beaver_gateway/storage/models.py @@ -1,6 +1,6 @@ """SQLModel tables. -Four tables, all flat, no FK relationships modelled (``actor`` and +Flat tables, no FK relationships modelled (``actor`` and ``agent_name`` are stored as strings — joining audit→token by name is fine at this volume; we'll introduce FKs when the admin UI actually demands them). @@ -74,15 +74,16 @@ class AuditLog(SQLModel, table=True): class Conversation(SQLModel, table=True): - """One chat thread, scoped to a frontend. + """One conversation (§3.1): a master thread, a branch, a deep chat or a job. - ``external_id`` is the identifier the frontend uses to find this - thread again on the next request — for the markdown frontend it's a - uuid we mint and persist into the file's frontmatter, for the - anthropic frontend it'd be the same metadata.conversation_id the - client passes. Unique per ``(frontend, external_id)`` because two - frontends sharing a uuid is fine; the same frontend reusing one is - a bug. + ``external_id`` is the public id (uuid) every frontend, the API and the + usage table refer to; ``frontend`` names the frontend that created the + row (``markdown``, ``api``, ``mcp``, ``system``). Where a conversation + is *visible* lives in :class:`ConversationBinding`. + + ``running_turn`` survives a restart: a non-null value at startup means + the gateway died mid-turn and the transcript needs its open ``tool_use`` + closed before the session is resumed (``core/conversations``). """ __tablename__ = "conversations" @@ -95,8 +96,92 @@ class Conversation(SQLModel, table=True): external_id: str = Field(index=True) agent_name: str = Field(index=True) session_id: str | None = Field(default=None, index=True) + kind: str = Field(default="deep", index=True) + parent_id: int | None = Field(default=None, index=True) + title: str | None = Field(default=None) + status: str = Field(default="open", index=True) + running_turn: str | None = Field(default=None) + pending_question: bool = Field(default=False) + flags: dict[str, Any] = Field( + default_factory=dict, + sa_column=Column( + JSON().with_variant(JSONB(), "postgresql"), + nullable=False, + server_default=text("'{}'"), + ), + ) created_at: datetime = Field(default_factory=_utcnow) updated_at: datetime = Field(default_factory=_utcnow) + last_user_activity_at: datetime | None = Field(default=None) + last_activity_at: datetime | None = Field(default=None) + + +class ConversationBinding(SQLModel, table=True): + """Where a conversation shows up: a Telegram topic id, a vault-relative path. + + Invariant (§3.1): at most one *visible* binding per frontend per + conversation - enforced by the partial unique index. The same external + id may point at several conversations over time (a renamed topic gets a + new branch), only one of them visible. + """ + + __tablename__ = "conversation_bindings" + __table_args__ = ( + Index( + "uq_binding_visible", + "conversation_id", + "frontend", + unique=True, + postgresql_where=text("visible"), + sqlite_where=text("visible"), + ), + Index("ix_binding_lookup", "frontend", "external_id"), + ) + + id: int | None = Field(default=None, primary_key=True) + conversation_id: int = Field(index=True) + frontend: str + external_id: str + visible: bool = Field(default=True) + created_at: datetime = Field(default_factory=_utcnow) + + +class InjectQueueItem(SQLModel, table=True): + """Persisted per-conversation queue (§3.4), priorities ``urgent > user > normal``. + + ``status`` walks ``queued -> running -> done``; a row still ``running`` + at startup was cut by a restart and becomes ``interrupted`` - it is + never re-run, the agent gets an inject saying so instead. + """ + + __tablename__ = "inject_queue" + + id: int | None = Field(default=None, primary_key=True) + conversation_id: int = Field(index=True) + priority: str = Field(index=True) + origin: str = Field(default="system") + text: str + status: str = Field(default="queued", index=True) + turn_id: str | None = Field(default=None) + created_at: datetime = Field(default_factory=_utcnow) + 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 ConversationMessage(SQLModel, table=True): @@ -193,7 +278,10 @@ class Usage(SQLModel, table=True): __all__ = [ "AuditLog", "Conversation", + "ConversationBinding", "ConversationMessage", + "InjectQueueItem", + "Schedule", "Token", "TranscriptEntry", "Usage", diff --git a/tests/test_claude_sdk_backend.py b/tests/test_claude_sdk_backend.py index 1744d39..3e498c6 100644 --- a/tests/test_claude_sdk_backend.py +++ b/tests/test_claude_sdk_backend.py @@ -93,6 +93,9 @@ class FakeClient: async def query(self, prompt: str) -> None: self.prompts.append(prompt) + async def interrupt(self) -> None: + self.interrupted = True + async def receive_response(self): start = {"type": "message_start", "message": {}} yield StreamEvent(uuid="u", session_id="s", event=start) diff --git a/tests/test_conversations.py b/tests/test_conversations.py new file mode 100644 index 0000000..4cf86d5 --- /dev/null +++ b/tests/test_conversations.py @@ -0,0 +1,546 @@ +import asyncio +import tempfile +import uuid +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import pytest +from claude_agent_sdk import ( + AssistantMessage, + InMemorySessionStore, + ResultMessage, + StreamEvent, + TextBlock, + project_key_for_directory, +) + +from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions +from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend +from beaver_gateway.core.bus import EventBus +from beaver_gateway.core.conversations import Conversations, ConversationTexts, parse_at +from beaver_gateway.core.registry import AgentRegistry +from beaver_gateway.core.sessions import SessionPool +from beaver_gateway.core.transcript import ( + build_entries, + close_open_tool_uses, + open_tool_uses, + render_messages, + strip_tool_entries, + window_entries, +) +from beaver_gateway.storage import Database +from beaver_gateway.storage.models import InjectQueueItem + + +class ScriptedClient: + instances: list["ScriptedClient"] = [] + hold: asyncio.Event | None = None + + def __init__(self, options: Any) -> None: + self.options = options + self.prompts: list[str] = [] + self.session_id = options.resume or str(uuid.uuid4()) + self.interrupted = False + self.connected = False + ScriptedClient.instances.append(self) + + async def connect(self) -> None: + self.connected = True + + async def query(self, prompt: str) -> None: + self.prompts.append(prompt) + + async def receive_response(self): + prompt = self.prompts[-1] + yield StreamEvent( + uuid="u", session_id="s", event={"type": "message_start", "message": {}} + ) + yield StreamEvent( + uuid="u", + session_id="s", + event={ + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ) + yield StreamEvent( + uuid="u", + session_id="s", + event={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": f"ok:{prompt}"}, + }, + ) + yield StreamEvent( + uuid="u", session_id="s", event={"type": "content_block_stop", "index": 0} + ) + yield AssistantMessage(content=[TextBlock(text=f"ok:{prompt}")], model="m") + hold = ScriptedClient.hold + if hold is not None and not self.interrupted: + await hold.wait() + cut = self.interrupted + self.interrupted = False + yield ResultMessage( + subtype="error_during_execution" if cut else "success", + duration_ms=1, + duration_api_ms=1, + is_error=cut, + num_turns=1, + session_id=self.session_id, + stop_reason="end_turn", + total_cost_usd=0.0, + usage={"input_tokens": 1, "output_tokens": 1}, + ) + + async def interrupt(self) -> None: + self.interrupted = True + if ScriptedClient.hold is not None: + ScriptedClient.hold.set() + + async def disconnect(self) -> None: + self.connected = False + + +class World: + def __init__(self, root: Path) -> None: + ScriptedClient.instances.clear() + ScriptedClient.hold = None + self.root = root + self.db = Database(f"sqlite:///{root / 'w.db'}") + self.store = InMemorySessionStore() + self.agent = ClaudeAgent( + name="a", + model="m", + system_prompt="hi", + cwd=root, + gateway_tools=("say",), + options=ClaudeOptions(effort="low"), + ) + self.pool = SessionPool(rss_limit=1 << 40, max_live=100) + self.backend = ClaudeSdkBackend( + agent=self.agent, + mcp_internal_urls={}, + session_store=self.store, + client_factory=ScriptedClient, + work_dir=root / "work", + pool=self.pool, + ) + self.bus = EventBus() + self.conversations = Conversations( + db=self.db, + agents=AgentRegistry([self.agent]), + backends={"a": self.backend}, + bus=self.bus, + pool=self.pool, + store=self.store, + texts=ConversationTexts(), + idle_interval=3600, + ) + + async def setup(self) -> "World": + await self.db.create_all() + return self + + def key(self, session_id: str) -> dict[str, str]: + return { + "project_key": project_key_for_directory(str(self.root)), + "session_id": session_id, + } + + async def statuses(self, conv) -> list[tuple[str, str]]: + rows = await self.conversations.queue.recent(conv.id) + return [(r.priority, r.status) for r in reversed(rows)] + + async def settle(self, conv, expected: int, timeout: float = 5.0) -> None: + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + rows = await self.conversations.queue.recent(conv.id) + if sum(1 for r in rows if r.status == "done") >= expected: + return + await asyncio.sleep(0.02) + msg = f"queue did not settle: {await self.statuses(conv)}" + raise AssertionError(msg) + + +@pytest.fixture +async def world() -> World: + root = Path(tempfile.mkdtemp(prefix="beaver-conv-")) + w = await World(root).setup() + yield w + await w.conversations.stop() + await w.pool.close_all() + await w.db.dispose() + + +async def test_two_messages_run_one_at_a_time_in_order(world: World) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + ScriptedClient.hold = asyncio.Event() + await world.conversations.post(conv, "first") + await asyncio.sleep(0.2) + await world.conversations.post(conv, "second") + await asyncio.sleep(0.2) + assert await world.statuses(conv) == [("user", "running"), ("user", "queued")] + live = world.pool.get(conv.external_id) + assert live is not None and live.busy + assert (await world.conversations.get(conv.external_id)).running_turn is not None + ScriptedClient.hold.set() + await world.settle(conv, 2) + assert len(ScriptedClient.instances) == 1 + assert ScriptedClient.instances[0].prompts == ["first", "second"] + row = await world.conversations.get(conv.external_id) + assert row.running_turn is None + assert row.session_id == ScriptedClient.instances[0].session_id + + +async def test_urgent_interrupts_and_goes_first(world: World) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + ScriptedClient.hold = asyncio.Event() + await world.conversations.post(conv, "first") + await asyncio.sleep(0.2) + await world.conversations.post(conv, "second") + await world.conversations.inject(conv, "ALERT", urgency="urgent", origin="крон") + await world.settle(conv, 2) + client = ScriptedClient.instances[0] + assert await world.statuses(conv) == [ + ("user", "interrupted"), + ("user", "done"), + ("urgent", "done"), + ] + assert await world.statuses(conv) == [ + ("user", "interrupted"), + ("user", "done"), + ("urgent", "done"), + ] + assert client.prompts[0] == "first" + assert client.prompts[1].startswith("[инжект: крон") + assert client.prompts[1].endswith("ALERT") + assert client.prompts[2] == "second" + + +async def test_normal_injects_ride_with_the_next_user_message(world: World) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + await world.conversations.inject( + conv, "vault changed", urgency="normal", origin="watch" + ) + await asyncio.sleep(0.2) + assert await world.statuses(conv) == [("normal", "queued")] + await world.conversations.post(conv, "hello") + await world.settle(conv, 2) + prompts = ScriptedClient.instances[0].prompts + assert len(prompts) == 1 + assert prompts[0].startswith("hello\n\n[инжекты") + assert "- [watch] vault changed" in prompts[0] + + +async def test_inject_turn_reply_is_not_routed(world: World) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + seen: list[dict[str, Any]] = [] + + async def collect() -> None: + async for event in world.bus.stream(conversation_id=conv.external_id): + seen.append(event) + + task = asyncio.create_task(collect()) + world.conversations._normal_window = 0.05 + await world.conversations.inject(conv, "tick", urgency="normal", origin="крон") + await world.settle(conv, 1) + await asyncio.sleep(0.05) + task.cancel() + types = [e["type"] for e in seen] + assert "turn.start" in types and "turn.end" in types + assert "reply" not in types + assert all(e.get("origin") == "inject" for e in seen if e["type"] == "turn.start") + + +async def test_spawn_seeds_first_user_message(world: World) -> None: + conv = await world.conversations.spawn( + kind="branch", agent="a", seed="brief", text="do X", title="t" + ) + await world.settle(conv, 1) + prompt = ScriptedClient.instances[0].prompts[0] + assert prompt.startswith("[сид: brief] branch «t», ") + assert prompt.endswith("\n\ndo X") + assert ScriptedClient.instances[0].options.system_prompt == "hi" + + +async def test_fork_leaves_original_untouched(world: World) -> None: + sid = str(uuid.uuid4()) + history = [ + {"role": "user", "content": "one"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "x"}], + }, + {"role": "assistant", "content": "done one"}, + {"role": "user", "content": "two"}, + {"role": "assistant", "content": "done two"}, + {"role": "user", "content": "three"}, + {"role": "assistant", "content": "done three"}, + ] + entries = build_entries(history, session_id=sid, cwd=str(world.root), model="m") + await world.store.append(world.key(sid), entries) + conv = await world.conversations.create( + kind="branch", agent="a", origin="test", session_id=sid + ) + before = [dict(e) for e in await world.store.load(world.key(sid))] + + result = await world.conversations.fork( + conv, "summarize", window=2, strip_tools=True + ) + + after = await world.store.load(world.key(sid)) + assert after == before + assert result.text == "ok:summarize" + assert result.conversation.status == "closed" + assert result.conversation.parent_id == conv.id + fork_entries = await world.store.load(world.key(result.conversation.session_id)) + assert not {e["uuid"] for e in fork_entries} & {e["uuid"] for e in before} + texts = [e["message"]["content"] for e in fork_entries] + assert texts[0] == "two" + assert len(fork_entries) == 4 + assert fork_entries[0]["parentUuid"] is None + assert all( + b["type"] == "text" + for e in fork_entries[1:] + for b in e["message"]["content"] + if isinstance(e["message"]["content"], list) + ) + client = ScriptedClient.instances[0] + assert client.options.resume == result.conversation.session_id + assert client.options.mcp_servers == {} + assert world.pool.get(result.conversation.external_id) is None + + +async def test_copy_seed_forks_parent_with_window(world: World) -> None: + sid = str(uuid.uuid4()) + history = [ + {"role": "user", "content": f"q{i}"} + if i % 2 == 0 + else {"role": "assistant", "content": f"a{i}"} + for i in range(8) + ] + await world.store.append( + world.key(sid), + build_entries(history, session_id=sid, cwd=str(world.root), model="m"), + ) + parent = await world.conversations.create( + kind="master", agent="a", origin="test", session_id=sid + ) + child = await world.conversations.spawn( + kind="branch", agent="a", seed="copy", parent=parent, window=1 + ) + assert child.session_id and child.session_id != sid + copied = await world.store.load(world.key(child.session_id)) + assert [e["message"]["content"] for e in copied] == [ + "q6", + [{"type": "text", "text": "a7"}], + ] + await world.settle(child, 1) + assert ScriptedClient.instances[0].options.resume == child.session_id + + +async def test_merge_injects_summary_into_parent(world: World) -> None: + sid = str(uuid.uuid4()) + await world.store.append( + world.key(sid), + build_entries( + [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}], + session_id=sid, + cwd=str(world.root), + model="m", + ), + ) + master = await world.conversations.create(kind="master", agent="a", origin="test") + branch = await world.conversations.create( + kind="branch", agent="a", origin="test", parent=master, session_id=sid + ) + result = await world.conversations.merge(branch) + assert result.text.startswith("ok:") + assert (await world.conversations.get(branch.external_id)).status == "merged" + assert await world.statuses(master) == [("normal", "queued")] + item = (await world.conversations.queue.recent(master.id))[0] + assert item.origin == "слив" and item.text == result.text + + +async def test_recover_closes_open_tool_use_and_injects_interrupted( + world: World, +) -> None: + sid = str(uuid.uuid4()) + history = [ + {"role": "user", "content": "go"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "t9", + "name": "Bash", + "input": {"command": "sleep"}, + } + ], + }, + ] + await world.store.append( + world.key(sid), + build_entries(history, session_id=sid, cwd=str(world.root), model="m"), + ) + conv = await world.conversations.create( + kind="master", agent="a", origin="test", session_id=sid + ) + async with world.db.session() as session: + row = await session.get(type(conv), conv.id) + row.running_turn = "turn_dead" + session.add(row) + session.add( + InjectQueueItem( + conversation_id=conv.id, + priority="user", + origin="user", + text="go", + status="running", + turn_id="turn_dead", + ) + ) + await session.commit() + + cut = await world.conversations.recover() + + assert [c.external_id for c in cut] == [conv.external_id] + entries = await world.store.load(world.key(sid)) + assert not open_tool_uses(entries) + tail = entries[-1] + assert tail["type"] == "user" + assert tail["message"]["content"][0] == { + "type": "tool_result", + "tool_use_id": "t9", + "content": "прервано", + "is_error": True, + } + assert tail["parentUuid"] == entries[-2]["uuid"] + assert (await world.conversations.get(conv.external_id)).running_turn is None + assert await world.statuses(conv) == [("user", "interrupted"), ("normal", "queued")] + note = (await world.conversations.queue.recent(conv.id))[0] + assert ( + "turn_dead" in note.text + and "оборван" in note.text + and "1 незакрытых" in note.text + ) + await asyncio.sleep(0.2) + assert ScriptedClient.instances == [] + + +async def test_read_and_bindings(world: World) -> None: + sid = str(uuid.uuid4()) + history = [ + {"role": "user", "content": "q1"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "a1"}, + {"type": "tool_use", "id": "t", "name": "Read", "input": {}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t", "content": "x"}], + }, + {"role": "assistant", "content": "a1b"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ] + await world.store.append( + world.key(sid), + build_entries(history, session_id=sid, cwd=str(world.root), model="m"), + ) + conv = await world.conversations.create( + kind="deep", agent="a", origin="test", session_id=sid + ) + assert ( + await world.conversations.read(conv) + == "user:\nq1\n\nassistant:\na1\n\na1b\n(tools: Read)\n\nuser:\nq2\n\nassistant:\na2" + ) + assert ( + await world.conversations.read(conv, window=1) == "user:\nq2\n\nassistant:\na2" + ) + + await world.conversations.bind(conv, frontend="markdown", external_id="a.md") + await world.conversations.bind(conv, frontend="markdown", external_id="b.md") + bindings = await world.conversations.bindings(conv) + assert [(b.external_id, b.visible) for b in bindings] == [ + ("a.md", False), + ("b.md", True), + ] + found = await world.conversations.find_bound( + frontend="markdown", external_id="b.md" + ) + assert found is not None and found.id == conv.id + assert ( + await world.conversations.find_bound(frontend="markdown", external_id="a.md") + is None + ) + + +async def test_schedule_rows_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"] + assert parse_at("2026-09-01T10:00:00+02:00") == datetime( + 2026, 9, 1, 8, 0, tzinfo=UTC + ) + with pytest.raises(ValueError, match="Invalid isoformat"): + parse_at("tomorrow") + + +def test_transcript_helpers() -> None: + history = [ + {"role": "user", "content": "q1"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "hm", "signature": "s"}, + {"type": "tool_use", "id": "t", "name": "Read", "input": {}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t", "content": "x"}], + }, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t2", "name": "Bash", "input": {}}], + }, + ] + entries = build_entries(history, session_id="s", cwd="/v", model="m") + assert [b["id"] for _, b in open_tool_uses(entries)] == ["t2"] + fixes = close_open_tool_uses(entries) + assert len(fixes) == 1 and fixes[0]["sessionId"] == "s" and fixes[0]["cwd"] == "/v" + assert not open_tool_uses([*entries, *fixes]) + stripped = strip_tool_entries(entries) + assert [e["message"]["content"] for e in stripped] == [ + "q1", + [{"type": "text", "text": "a1"}], + "q2", + ] + assert ( + stripped[0]["parentUuid"] is None + and stripped[1]["parentUuid"] == stripped[0]["uuid"] + ) + windowed = window_entries(entries, window=1) + assert windowed[0]["message"]["content"] == "q2" and len(windowed) == 2 + assert ( + render_messages( + [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}] + ) + == "user:\nq\n\nassistant:\na" + ) diff --git a/tests/test_db_migrate.py b/tests/test_db_migrate.py index 99580ab..826971f 100644 --- a/tests/test_db_migrate.py +++ b/tests/test_db_migrate.py @@ -31,3 +31,31 @@ async def test_create_all_adds_missing_columns() -> None: session.add(conv) await session.commit() await db.dispose() + + +async def test_create_all_backfills_defaults_for_old_rows() -> None: + path = Path(tempfile.mkdtemp(prefix="beaver-migrate-")) / "old.db" + raw = sqlite3.connect(path) + raw.execute( + "CREATE TABLE conversations (id INTEGER PRIMARY KEY, frontend VARCHAR NOT NULL, " + "external_id VARCHAR NOT NULL, agent_name VARCHAR NOT NULL, session_id VARCHAR, " + "created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)" + ) + raw.execute( + "INSERT INTO conversations VALUES (1, 'markdown', 'x', 'a', NULL, '2026-01-01', '2026-01-01')" + ) + raw.commit() + raw.close() + + db = Database(f"sqlite:///{path}") + await db.create_all() + async with db.session() as session: + conv = (await session.exec(select(Conversation))).one() + assert (conv.kind, conv.status, conv.pending_question, conv.flags) == ( + "deep", + "open", + False, + {}, + ) + assert conv.running_turn is None + await db.dispose() diff --git a/uv.lock b/uv.lock index 4c736e4..cc9fb2c 100644 --- a/uv.lock +++ b/uv.lock @@ -273,6 +273,7 @@ dependencies = [ { name = "greenlet" }, { name = "itsdangerous" }, { name = "jinja2" }, + { name = "psutil" }, { name = "psycopg", extra = ["binary"] }, { name = "pydantic" }, { name = "pydantic-settings" }, @@ -313,6 +314,7 @@ requires-dist = [ { name = "greenlet", specifier = ">=3.5.0" }, { name = "itsdangerous", specifier = ">=2.2.0" }, { name = "jinja2", specifier = ">=3.1.6" }, + { name = "psutil", specifier = ">=7.2.2" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.3.4" }, { name = "pydantic", specifier = ">=2.13.4" }, { name = "pydantic-settings", specifier = ">=2.14.1" }, @@ -1386,6 +1388,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/ed/1cdcab6ba3d6ab7feca11fc14f0eeea80755bb53ef4e892079f31b10a25f/propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe", size = 14036, upload-time = "2026-05-08T21:02:10.673Z" }, ] +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, + { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, + { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" }, + { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" }, + { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" }, + { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" }, + { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + [[package]] name = "psycopg" version = "3.3.4"