Files
beaver-gateway/src/beaver_gateway/frontends/markdown/frontend.py
T

874 lines
30 KiB
Python

"""``MarkdownFrontend`` — chat-via-markdown-files frontend.
``POST /chat`` (and SSE ``/chat/stream``) parses the vault file, runs
the agent on the trailing user turn, and appends the reply.
"""
from __future__ import annotations
import json
import logging
import time
from collections.abc import AsyncIterator
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from anthropic.types import RawContentBlockStopEvent
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from beaver_gateway.backends.capture import TurnCapture
from beaver_gateway.frontends.accumulate import StreamAccumulator
from beaver_gateway.frontends.base import Frontend
from beaver_gateway.frontends.bearer import require_token
from beaver_gateway.frontends.markdown import parser, renderer
from beaver_gateway.frontends.markdown.crossfront import CrossFrontendLogger
from beaver_gateway.frontends.markdown.files import (
read_or_empty,
reattach_frontmatter,
write_atomic,
)
from beaver_gateway.frontends.markdown.history import (
diff_and_fork,
load_messages,
rewrite_messages,
)
from beaver_gateway.frontends.markdown.mirror import FRONTEND, ChatMirror
from beaver_gateway.frontends.sse import (
KEEPALIVE,
SSE_HEADERS,
events_with_heartbeat,
sse_pack,
)
from beaver_gateway.frontends.turn_record import TurnRecord
from beaver_gateway.security import audit
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Callable
from beaver_gateway.conversations.kinds import Kind
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.storage.models import Conversation, ConversationBinding
_log = logging.getLogger("beaver_gateway.frontends.markdown")
__all__ = ["MarkdownFrontend"]
_STREAM_FLUSH_DEBOUNCE = 0.4
_SSE_FLUSH_DEBOUNCE = 0.1
class MarkdownFrontend(Frontend):
"""FastAPI app behind ``POST /chat`` driven by Obsidian-vault files.
``_busy`` tracks in-flight files; check-and-add must stay atomic (no
``await`` between them) so concurrent requests reliably lose to 409.
"""
name = FRONTEND
kinds = ("deep",)
path = "/md"
def __init__(
self,
*,
vault_path: Path | str,
default_agent: str | None = None,
log_all_chats: bool = False,
logged_subdir: str = "_logs",
chat_path: Callable[[str, str, Path], Path] | None = None,
) -> None:
"""Configure the vault-backed frontend.
``chat_path``, if given, overrides where new chat files are created;
``logged_subdir`` holds cross-frontend logs when ``log_all_chats`` is set.
"""
self.vault_path = Path(vault_path).expanduser().resolve()
self.default_agent = default_agent
self.log_all_chats = log_all_chats
self.logged_subdir = logged_subdir
self.chat_path = chat_path
self._runtime: GatewayRuntime | None = None
self._app: FastAPI | None = None
self._busy: set[Path] = set()
self._crossfront: CrossFrontendLogger | None = None
self._mirror: ChatMirror | 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)
self._mirror = ChatMirror(
vault_path=self.vault_path,
runtime=runtime,
logged_subdir=self.logged_subdir,
chat_path=self.chat_path,
)
if self.log_all_chats:
self._crossfront = CrossFrontendLogger(
vault_path=self.vault_path,
logged_subdir=self.logged_subdir,
chat_path=self.chat_path,
)
self._crossfront.warm_index()
runtime.turn_log_handlers.append(self._crossfront.handle)
self._app = self._build_app(runtime)
@property
def mirror(self) -> ChatMirror:
if self._mirror is None:
msg = "configure() must be called first"
raise RuntimeError(msg)
return self._mirror
def agent_for(self, kind: Kind) -> str | None:
return self.default_agent if kind == "deep" else None
async def materialize(self, conv: Conversation) -> ConversationBinding | None:
return await self.mirror.materialize(conv)
def app(self) -> FastAPI | None:
return self._app
async def serve(self) -> None:
await self.mirror.run()
def _build_app(self, runtime: GatewayRuntime) -> FastAPI:
"""CORS is wide open here since auth is bearer-token, not cookie-based."""
app = FastAPI(title="beaver-gateway / Markdown")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=False,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/healthz")
async def healthz() -> dict[str, str]:
return {"status": "ok"}
@app.get("/agents")
async def list_agents(request: Request) -> dict[str, Any]:
await require_token(request, runtime, scope="messages")
return {"agents": [{"name": a.name} for a in runtime.agents]}
@app.post("/chat")
async def chat(request: Request) -> Any:
token_name = await require_token(request, runtime, scope="messages")
try:
body = await request.json()
except json.JSONDecodeError as exc:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"invalid JSON: {exc}"
) from exc
filename = body.get("filename")
if not isinstance(filename, str) or not filename.strip():
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "missing or non-string `filename`"
)
content_override = body.get("content")
agent_override = body.get("agent")
if agent_override is not None and not isinstance(agent_override, str):
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "`agent` must be a string"
)
file_path = self._resolve_path(filename)
if file_path in self._busy:
return JSONResponse(
status_code=status.HTTP_409_CONFLICT,
content={"status": "in_progress", "filename": filename},
)
self._busy.add(file_path)
try:
return await self._handle_chat(
runtime=runtime,
token_name=token_name,
filename=filename,
file_path=file_path,
content_override=content_override,
agent_override=agent_override,
)
finally:
self._busy.discard(file_path)
@app.post("/chat/stream")
async def chat_stream(request: Request) -> Any:
"""SSE variant of ``/chat``: ``delta`` events, one disk write at the end.
The 409-conflict response is still plain JSON — the stream hasn't started.
"""
token_name = await require_token(request, runtime, scope="messages")
try:
body = await request.json()
except json.JSONDecodeError as exc:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"invalid JSON: {exc}"
) from exc
filename = body.get("filename")
if not isinstance(filename, str) or not filename.strip():
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "missing or non-string `filename`"
)
content_override = body.get("content")
agent_override = body.get("agent")
if agent_override is not None and not isinstance(agent_override, str):
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "`agent` must be a string"
)
file_path = self._resolve_path(filename)
if file_path in self._busy:
return JSONResponse(
status_code=status.HTTP_409_CONFLICT,
content={"status": "in_progress", "filename": filename},
)
self._busy.add(file_path)
async def gen() -> AsyncIterator[bytes]:
try:
async for chunk in self._handle_chat_streaming(
runtime=runtime,
token_name=token_name,
filename=filename,
file_path=file_path,
content_override=content_override,
agent_override=agent_override,
):
yield chunk
finally:
self._busy.discard(file_path)
return StreamingResponse(
gen(), media_type="text/event-stream", headers=SSE_HEADERS
)
return app
async def _handle_chat(
self,
*,
runtime: GatewayRuntime,
token_name: str,
filename: str,
file_path: Path,
content_override: Any,
agent_override: str | None,
) -> Any:
"""Non-streaming ``/chat`` handler: parse, align, run, persist."""
write_disk = content_override is None
if isinstance(content_override, str):
file_text = content_override
elif content_override is None:
file_text = await read_or_empty(file_path)
else:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, "`content` must be a string when present"
)
parsed = parser.parse(file_text)
agent_name = parser.resolve_agent(
metadata=parsed.metadata,
request_override=agent_override,
default=self.default_agent,
)
if not agent_name:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
"no agent specified: pass `agent`, set frontmatter, "
"or configure `default_agent`",
)
if not parsed.messages:
return {
"status": "nothing_to_do",
"reason": "empty file",
"new_content": file_text,
}
if parser.last_role(parsed.messages) == "assistant":
return {
"status": "nothing_to_do",
"reason": "last turn is assistant",
"new_content": file_text,
}
agent = runtime.agents.get(agent_name)
if agent is None:
raise HTTPException(
status.HTTP_404_NOT_FOUND, f"unknown agent: {agent_name!r}"
)
backend = runtime.backends.get(agent.name)
if backend is None:
raise HTTPException(
status.HTTP_503_SERVICE_UNAVAILABLE,
f"no backend configured for agent {agent.name!r}",
)
_log.info(
"chat: actor=%s agent=%s file=%s msgs=%d",
token_name,
agent.name,
filename,
len(parsed.messages),
)
await audit.log(
runtime,
actor=f"token:{token_name}",
kind="markdown_chat",
agent_name=agent.name,
filename=filename,
msgs=len(parsed.messages),
)
conv, conv_external_id, stored_msgs = await self._resolve_conversation(
runtime=runtime,
metadata=parsed.metadata,
agent_name=agent.name,
file_path=file_path,
first_user_text=parsed.turns[0].text,
)
outcome = diff_and_fork(stored=stored_msgs, incoming=parsed.turns)
capture = TurnCapture()
events = runtime.conversations.turn(
conv,
messages=outcome.messages,
origin="user",
capture=capture,
use_session=outcome.reuse_session,
)
try:
message = await self._stream_to_file(
events=events,
file_path=file_path,
parsed=parsed,
model=agent.model or agent.name,
filename=filename,
write_disk=write_disk,
)
except HTTPException:
raise
except Exception as exc:
raise HTTPException(
status.HTTP_500_INTERNAL_SERVER_ERROR, f"backend error: {exc}"
) from exc
new_content = await self._write_assistant_reply(
file_path=file_path,
parsed=parsed,
message=message,
agent_name=agent.name,
conv_external_id=conv_external_id,
write_disk=write_disk,
)
await self._persist_canonical_history(
runtime=runtime,
conversation_id=cast("int", conv.id),
persist_messages=outcome.persist_messages,
new_user_text=parsed.turns[-1].text,
capture=capture,
message=message,
)
record = TurnRecord(
agent_name=agent.name,
input_messages=list(parsed.messages),
output_message=message,
system=None,
source="markdown",
)
for handler in runtime.turn_log_handlers:
try:
await handler(record)
except Exception: # noqa: BLE001
_log.exception("turn_log_handler raised; continuing")
return {
"status": "ok",
"turns_appended": 1,
"agent": agent.name,
"new_content": new_content,
}
async def _handle_chat_streaming( # noqa: PLR0915
self,
*,
runtime: GatewayRuntime,
token_name: str,
filename: str,
file_path: Path,
content_override: Any,
agent_override: str | None,
) -> AsyncIterator[bytes]:
"""SSE counterpart of :meth:`_handle_chat`.
Errors surface as ``error`` frames (HTTP is already 200 by then); disk
is only written once at the end, avoiding a write race with Obsidian Sync.
"""
write_disk = content_override is None
if isinstance(content_override, str):
file_text = content_override
elif content_override is None:
file_text = await read_or_empty(file_path)
else:
yield sse_pack(
"error",
{
"status_code": status.HTTP_400_BAD_REQUEST,
"detail": "`content` must be a string when present",
},
)
return
parsed = parser.parse(file_text)
agent_name = parser.resolve_agent(
metadata=parsed.metadata,
request_override=agent_override,
default=self.default_agent,
)
if not agent_name:
yield sse_pack(
"error",
{
"status_code": status.HTTP_400_BAD_REQUEST,
"detail": (
"no agent specified: pass `agent`, set frontmatter, "
"or configure `default_agent`"
),
},
)
return
if not parsed.messages:
yield sse_pack(
"done",
{
"status": "nothing_to_do",
"reason": "empty file",
"new_content": file_text,
},
)
return
if parser.last_role(parsed.messages) == "assistant":
yield sse_pack(
"done",
{
"status": "nothing_to_do",
"reason": "last turn is assistant",
"new_content": file_text,
},
)
return
agent = runtime.agents.get(agent_name)
if agent is None:
yield sse_pack(
"error",
{
"status_code": status.HTTP_404_NOT_FOUND,
"detail": f"unknown agent: {agent_name!r}",
},
)
return
backend = runtime.backends.get(agent.name)
if backend is None:
yield sse_pack(
"error",
{
"status_code": status.HTTP_503_SERVICE_UNAVAILABLE,
"detail": f"no backend configured for agent {agent.name!r}",
},
)
return
_log.info(
"chat/stream: actor=%s agent=%s file=%s msgs=%d",
token_name,
agent.name,
filename,
len(parsed.messages),
)
await audit.log(
runtime,
actor=f"token:{token_name}",
kind="markdown_chat_stream",
agent_name=agent.name,
filename=filename,
msgs=len(parsed.messages),
)
try:
conv, conv_external_id, stored_msgs = await self._resolve_conversation(
runtime=runtime,
metadata=parsed.metadata,
agent_name=agent.name,
file_path=file_path,
first_user_text=parsed.turns[0].text,
)
except HTTPException as exc:
yield sse_pack(
"error", {"status_code": exc.status_code, "detail": exc.detail}
)
return
_log.info(
"chat/stream: file=%s conv_external_id=%s conv_id=%d "
"stored_msgs=%d incoming_turns=%d",
filename,
conv_external_id,
conv.id or -1,
len(stored_msgs),
len(parsed.turns),
)
outcome = diff_and_fork(stored=stored_msgs, incoming=parsed.turns)
_log.info(
"chat/stream: file=%s diff_and_fork divergence=%s "
"backend_msgs=%d persist_msgs=%d",
filename,
outcome.divergence_index,
len(outcome.messages),
len(outcome.persist_messages),
)
capture = TurnCapture()
_log.info(
"chat/stream: file=%s calling backend.complete agent=%s session=%s",
filename,
agent.name,
conv.session_id,
)
events = runtime.conversations.turn(
conv,
messages=outcome.messages,
origin="user",
capture=capture,
use_session=outcome.reuse_session,
)
acc = StreamAccumulator()
model = agent.model or agent.name
last_flush = time.monotonic()
last_payload: str | None = None
def snapshot() -> str | None:
partial = acc.finalize(model=model)
if not partial.content:
return None
rendered = renderer.render_assistant_message(partial)
new_body = renderer.append_to_body(parsed.body, rendered)
return reattach_frontmatter(parsed.metadata, new_body)
try:
async for ev in events_with_heartbeat(events):
if ev is None:
yield KEEPALIVE
continue
acc.feed(ev)
now = time.monotonic()
if (
isinstance(ev, RawContentBlockStopEvent)
or (now - last_flush) >= _SSE_FLUSH_DEBOUNCE
):
payload = snapshot()
if payload is not None and payload != last_payload:
yield sse_pack("delta", {"new_content": payload})
last_payload = payload
last_flush = now
except Exception as exc: # noqa: BLE001
_log.exception("backend failed for %s", filename)
partial = acc.finalize(model=model)
new_body = parsed.body
if partial.content:
new_body = renderer.append_to_body(
new_body, renderer.render_assistant_message(partial)
)
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
if write_disk:
await write_atomic(
file_path, reattach_frontmatter(parsed.metadata, new_body)
)
yield sse_pack(
"error",
{
"status_code": status.HTTP_500_INTERNAL_SERVER_ERROR,
"detail": f"backend error: {exc}",
},
)
return
message = acc.finalize(model=model)
new_content = await self._write_assistant_reply(
file_path=file_path,
parsed=parsed,
message=message,
agent_name=agent.name,
conv_external_id=conv_external_id,
write_disk=write_disk,
)
await self._persist_canonical_history(
runtime=runtime,
conversation_id=cast("int", conv.id),
persist_messages=outcome.persist_messages,
new_user_text=parsed.turns[-1].text,
capture=capture,
message=message,
)
record = TurnRecord(
agent_name=agent.name,
input_messages=list(parsed.messages),
output_message=message,
system=None,
source="markdown",
)
for handler in runtime.turn_log_handlers:
try:
await handler(record)
except Exception: # noqa: BLE001
_log.exception("turn_log_handler raised; continuing")
yield sse_pack(
"done",
{
"status": "ok",
"turns_appended": 1,
"agent": agent.name,
"new_content": new_content,
},
)
async def _stream_to_file(
self,
*,
events: Any,
file_path: Path,
parsed: parser.ParsedFile,
model: str,
filename: str,
write_disk: bool = True,
) -> Any:
"""Drain ``events`` into a ``Message``, flushing partials to disk.
Flushes on each block boundary and on a debounce timer; on backend
failure it still flushes a partial + error callout, then re-raises.
"""
acc = StreamAccumulator()
async def flush_partial() -> None:
if not write_disk:
return
partial = acc.finalize(model=model)
if not partial.content:
return
rendered = renderer.render_assistant_message(partial)
new_body = renderer.append_to_body(parsed.body, rendered)
await write_atomic(
file_path, reattach_frontmatter(parsed.metadata, new_body)
)
try:
last_flush = time.monotonic()
async for ev in events:
acc.feed(ev)
now = time.monotonic()
if (
isinstance(ev, RawContentBlockStopEvent)
or (now - last_flush) >= _STREAM_FLUSH_DEBOUNCE
):
await flush_partial()
last_flush = now
except Exception as exc:
_log.exception("backend failed for %s", filename)
partial = acc.finalize(model=model)
new_body = parsed.body
if partial.content:
new_body = renderer.append_to_body(
new_body, renderer.render_assistant_message(partial)
)
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
if write_disk:
await write_atomic(
file_path, reattach_frontmatter(parsed.metadata, new_body)
)
raise
return acc.finalize(model=model)
async def _write_assistant_reply(
self,
*,
file_path: Path,
parsed: parser.ParsedFile,
message: Any,
agent_name: str,
conv_external_id: str,
write_disk: bool = True,
) -> str:
"""Render the assistant turn, refresh frontmatter, write if we own the file."""
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)
updated_metadata = dict(parsed.metadata)
updated_metadata.pop("fingerprint", None)
updated_metadata["agent"] = agent_name
updated_metadata["conversation_id"] = conv_external_id
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,
file_path: Path,
first_user_text: str = "",
) -> tuple[Conversation, str, list[dict[str, Any]]]:
"""Resolve the ``deep`` conversation for this file + its stored messages.
Precedence: frontmatter ``conversation_id`` > existing path binding >
an adopted unbound conversation with the same first prompt > a new one.
"""
conversations = runtime.conversations
rel = file_path.relative_to(self.vault_path).as_posix()
raw = metadata.get("conversation_id")
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.adopt(
kind="deep", first_user_text=first_user_text
)
if conv is not None:
_log.info("adopted conversation %s for %s", conv.external_id, rel)
if conv is None:
try:
conv = await conversations.create(
kind="deep", agent=agent_name, origin=FRONTEND, title=file_path.stem
)
except ValueError as exc:
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
_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:
stored = await load_messages(session, conversation_id=conv.id)
return conv, conv.external_id, stored
async def _persist_canonical_history(
self,
*,
runtime: GatewayRuntime,
conversation_id: int,
persist_messages: list[dict[str, Any]],
new_user_text: str,
capture: TurnCapture,
message: Any,
) -> None:
"""Stamp the DB with the post-turn canonical Anthropic-shape history.
Combines prior state + new user prompt + the backend's synthesized
cycle, falling back to text-only if ``capture`` is empty.
"""
new_user_msg = {"role": "user", "content": new_user_text}
synthesized = capture.synthesized_messages or _fallback_synthesized(message)
canonical = [*persist_messages, new_user_msg, *synthesized]
_log.info(
"_persist_canonical_history: conv_id=%d writing %d msgs "
"(prior=%d + new_user + synth=%d)",
conversation_id,
len(canonical),
len(persist_messages),
len(synthesized),
)
async with runtime.db.session() as session:
await rewrite_messages(
session, conversation_id=conversation_id, messages=canonical
)
_log.info(
"_persist_canonical_history: conv_id=%d DB committed", conversation_id
)
def _resolve_path(self, filename: str) -> Path:
"""Resolve ``filename`` under the vault; reject escapes.
Leading slashes are stripped first — ``Path.__truediv__`` would
otherwise discard ``vault_path`` for an absolute ``filename``.
"""
rel = filename.lstrip("/")
if not rel.endswith(".md"):
rel = rel + ".md"
candidate = (self.vault_path / rel).resolve()
try:
candidate.relative_to(self.vault_path)
except ValueError as exc:
raise HTTPException(
status.HTTP_400_BAD_REQUEST, f"filename escapes vault: {filename!r}"
) from exc
return candidate
def _fallback_synthesized(message: Any) -> list[dict[str, Any]]:
"""Build a single-assistant ``synthesized_messages`` list from a raw ``Message``.
Fallback for backends that don't populate :class:`TurnCapture` (anthropic
HTTP, raycast, …); the reply lands as one canonical-block message.
"""
content: list[dict[str, Any]] = []
for block in getattr(message, "content", ()):
btype = getattr(block, "type", None)
if btype == "text":
content.append({"type": "text", "text": getattr(block, "text", "") or ""})
elif btype == "tool_use":
content.append(
{
"type": "tool_use",
"id": getattr(block, "id", ""),
"name": getattr(block, "name", ""),
"input": getattr(block, "input", {}),
}
)
elif btype == "thinking":
content.append(
{
"type": "thinking",
"thinking": getattr(block, "thinking", "") or "",
"signature": getattr(block, "signature", "") or "",
}
)
if not content:
return []
return [{"role": "assistant", "content": content}]
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__
safe = msg.replace("\n", " ").strip()
return f"### Assistant:\n\n> [!error]-\n> {safe}\n"