Compare commits
91
Commits
9b1127acff
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
944b68e275 | ||
|
|
7d812c7962 | ||
|
|
8f26263d4a | ||
|
|
7b36cb2156 | ||
|
|
2a47c6771d | ||
|
|
b0e5aaf599 | ||
|
|
0fec15fb3b | ||
|
|
712c24de4e | ||
|
|
9875f08016 | ||
|
|
65564d8541 | ||
|
|
1d1285de46 | ||
|
|
120985deab | ||
|
|
13873f857b | ||
|
|
bbce425eeb | ||
|
|
e58153628c | ||
|
|
c120953b02 | ||
|
|
5eb9f435ec | ||
|
|
e79e9fc149 | ||
|
|
60782d6276 | ||
|
|
421248b0ad | ||
|
|
8019c76e53 | ||
|
|
2afd3bc3c1 | ||
|
|
5b851c11d4 | ||
|
|
2d903c7c5f | ||
|
|
ef7474f636 | ||
|
|
fe1c894b93 | ||
|
|
2815956f5b | ||
|
|
c46ba026ba | ||
|
|
81d6560bf7 | ||
|
|
f35980ba85 | ||
|
|
c6401f1114 | ||
|
|
025f2dbc3f | ||
|
|
85b14e2c2f | ||
|
|
7a8ee0f200 | ||
|
|
2d34d488d7 | ||
|
|
0e07e8409d | ||
|
|
3915ab48f9 | ||
|
|
9aaddbed75 | ||
|
|
253bde1b11 | ||
|
|
61b5bda31a | ||
|
|
e2d611dd81 | ||
|
|
bac40ff268 | ||
|
|
39f2bc05d9 | ||
|
|
76d6f81789 | ||
|
|
8f9ca2800c | ||
|
|
f9bf51badf | ||
|
|
90063a7cfe | ||
|
|
a8491330fa | ||
|
|
064379ca90 | ||
|
|
e90963a6c3 | ||
|
|
064d855382 | ||
|
|
5d67cfcc9e | ||
|
|
b4853b9fd0 | ||
|
|
2d12391176 | ||
|
|
60dc7db925 | ||
|
|
be114e6eaf | ||
|
|
4a59c59bda | ||
|
|
72639c4b90 | ||
|
|
1786406b31 | ||
|
|
12728b2eb4 | ||
|
|
bc4233977d | ||
|
|
fd69eae05d | ||
|
|
d3a3e7b1f1 | ||
|
|
772db11c7f | ||
|
|
4fb020c788 | ||
|
|
932a7a80af | ||
|
|
3ca5843868 | ||
|
|
8ea4cea3ef | ||
|
|
b6c5a80326 | ||
|
|
98796a82c6 | ||
|
|
61562e947d | ||
|
|
609af9fe67 | ||
|
|
e469e7e174 | ||
|
|
c88726db9e | ||
|
|
38986cb82e | ||
|
|
43407016d0 | ||
|
|
7240e0cdac | ||
|
|
7b0c61c18e | ||
|
|
0110cbd1cc | ||
|
|
9663fbf956 | ||
|
|
5177392f46 | ||
|
|
5fc58e3bc7 | ||
|
|
5d2ce68f6b | ||
|
|
70349e5ff4 | ||
|
|
827fa0977b | ||
|
|
e3074c266a | ||
|
|
ab52fdc2b8 | ||
|
|
7424d52f88 | ||
|
|
b3a584a362 | ||
|
|
a2a80ab150 | ||
|
|
d33039f497 |
No files matched your search
+3
-1
@@ -6,9 +6,11 @@
|
|||||||
- **NEVER write comments** - code should be self-documenting
|
- **NEVER write comments** - code should be self-documenting
|
||||||
|
|
||||||
## Layout
|
## Layout
|
||||||
- `src/beaver_gateway` - python, unopinionated primitives (no vault paths, no agent names)
|
- `src/beaver_gateway` - python, unopinionated primitives (no vault paths, no agent names, no Russian): one package per capability - `agents`, `backends`, `conversations`, `jobs`, `frontends`, `mcp`, `vault`, `security`, `events`, `storage`; `config.py` loads a setup, `app.py` builds the runtime, `cli.py` is the entrypoint
|
||||||
|
- every string shown to a model or a user lives in a `*Texts` dataclass with English defaults (`conversations/texts.py`, `frontends/telegram/texts.py`); a setup overrides them
|
||||||
- `ui` - SvelteKit SPA (admin + Obsidian panel), served by the gateway from `ui/build`
|
- `ui` - SvelteKit SPA (admin + Obsidian panel), served by the gateway from `ui/build`
|
||||||
- `docs/MISSION.md`, `docs/PRD.md`, `docs/PLAN.md`, `docs/PROGRESS.md` - keep PROGRESS.md current
|
- `docs/MISSION.md`, `docs/PRD.md`, `docs/PLAN.md`, `docs/PROGRESS.md` - keep PROGRESS.md current
|
||||||
|
- docstrings: one line per module saying what it is for, one line per public configuration field; no section references, phases or history
|
||||||
|
|
||||||
## Checking commands
|
## Checking commands
|
||||||
After writing code, always run (or `make check`):
|
After writing code, always run (or `make check`):
|
||||||
|
|||||||
@@ -28,3 +28,7 @@ RAYCAST_BEARER=
|
|||||||
# токен от @BotFather, свой user id - @userinfobot
|
# токен от @BotFather, свой user id - @userinfobot
|
||||||
TELEGRAM_BOT_TOKEN=
|
TELEGRAM_BOT_TOKEN=
|
||||||
TELEGRAM_USER_ID=
|
TELEGRAM_USER_ID=
|
||||||
|
|
||||||
|
# Откуда тулза send_file имеет право брать файлы (через ':').
|
||||||
|
# Папка входящих вложений разрешена всегда, её дописывает сам фронтенд.
|
||||||
|
SEND_FILE_ROOTS=/mnt/hole/shared
|
||||||
+26
-3
@@ -45,17 +45,36 @@ RUN apt-get update \
|
|||||||
&& rm -rf /var/lib/apt/lists/*
|
&& rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
# Bun native binary (glibc): `bunx` runs the stdio MCP servers declared in
|
# Bun native binary (glibc): `bunx` runs the stdio MCP servers declared in
|
||||||
# the user's config. The claude CLI itself ships inside the
|
# the user's config.
|
||||||
# claude-agent-sdk wheel (`_bundled/claude`), nothing to install here.
|
|
||||||
COPY --from=oven/bun:1-slim /usr/local/bin/bun /usr/local/bin/bun
|
COPY --from=oven/bun:1-slim /usr/local/bin/bun /usr/local/bin/bun
|
||||||
RUN ln -s /usr/local/bin/bun /usr/local/bin/bunx
|
RUN ln -s /usr/local/bin/bun /usr/local/bin/bunx
|
||||||
|
|
||||||
|
# The claude-agent-sdk wheel bundles a CLI, but it trails the npm release:
|
||||||
|
# 0.2.157 carries 2.1.277, which does not know claude-opus-5-5 and writes an
|
||||||
|
# [claude-code:unrecognized_model] warning per turn. Install the npm build and
|
||||||
|
# point BEAVER_CLAUDE_CLI at it; the adapter prefers it over the bundled one.
|
||||||
|
#
|
||||||
|
# Only the glibc platform package is pulled: the meta package also drags in the
|
||||||
|
# musl build, and this image is glibc for the reason stated at the top. The
|
||||||
|
# tree lands under /opt world-readable because the model runs as beaver-runner,
|
||||||
|
# not as root.
|
||||||
|
ARG CLAUDE_CODE_VERSION=2.1.280
|
||||||
|
RUN mkdir -p /opt/claude-code \
|
||||||
|
&& cd /opt/claude-code \
|
||||||
|
&& echo '{"name":"claude-cli-holder","private":true}' > package.json \
|
||||||
|
&& bun add @anthropic-ai/claude-code-linux-x64@${CLAUDE_CODE_VERSION} \
|
||||||
|
&& ln -sf /opt/claude-code/node_modules/@anthropic-ai/claude-code-linux-x64/claude \
|
||||||
|
/usr/local/bin/claude-cli \
|
||||||
|
&& chmod -R a+rX /opt/claude-code \
|
||||||
|
&& /usr/local/bin/claude-cli --version
|
||||||
|
ENV BEAVER_CLAUDE_CLI=/usr/local/bin/claude-cli
|
||||||
|
|
||||||
# The model process runs as `beaver-runner`, not as the gateway: the
|
# The model process runs as `beaver-runner`, not as the gateway: the
|
||||||
# adapter spawns claude under this uid with a whitelisted environment.
|
# adapter spawns claude under this uid with a whitelisted environment.
|
||||||
# `acl` lets the entrypoint grant it write access to the vault sub-mounts
|
# `acl` lets the entrypoint grant it write access to the vault sub-mounts
|
||||||
# without chowning files that Obsidian Sync keeps rewriting as root.
|
# without chowning files that Obsidian Sync keeps rewriting as root.
|
||||||
RUN apt-get update \
|
RUN apt-get update \
|
||||||
&& apt-get install -y --no-install-recommends acl \
|
&& apt-get install -y --no-install-recommends acl sudo \
|
||||||
&& rm -rf /var/lib/apt/lists/* \
|
&& rm -rf /var/lib/apt/lists/* \
|
||||||
&& useradd --system --uid 1001 --create-home --shell /usr/sbin/nologin beaver-runner
|
&& useradd --system --uid 1001 --create-home --shell /usr/sbin/nologin beaver-runner
|
||||||
|
|
||||||
@@ -64,6 +83,10 @@ ENV PATH=/app/.venv/bin:$PATH
|
|||||||
COPY --from=builder /app/.venv /app/.venv
|
COPY --from=builder /app/.venv /app/.venv
|
||||||
COPY --from=builder /app /app
|
COPY --from=builder /app /app
|
||||||
COPY --from=ui /ui/build /app/ui/build
|
COPY --from=ui /ui/build /app/ui/build
|
||||||
|
COPY --chmod=0755 docker/secret /usr/local/bin/secret
|
||||||
|
COPY --chmod=0755 docker/beaver-secret /usr/local/libexec/beaver-secret
|
||||||
|
COPY --chmod=0440 docker/sudoers /etc/sudoers.d/beaver-secret
|
||||||
|
RUN visudo -cf /etc/sudoers.d/beaver-secret
|
||||||
WORKDIR /app
|
WORKDIR /app
|
||||||
|
|
||||||
ENTRYPOINT ["python", "-m"]
|
ENTRYPOINT ["python", "-m"]
|
||||||
|
|||||||
@@ -1,2 +1,79 @@
|
|||||||
# Beaver Gateway
|
# beaver-gateway
|
||||||
[](https://sladge.net)
|
|
||||||
|
A gateway for personal agents built on the Claude Agent SDK: conversations
|
||||||
|
with a lifecycle, windows to talk through, jobs, memory hooks, and a policy
|
||||||
|
boundary - without an opinion about what the agent is for. A setup is one
|
||||||
|
`config.py` (or a package next to it) that assembles a `Gateway`; this repo
|
||||||
|
holds the primitives, [beaver-agent](https://git.kotikot.com/beaver/beaver-agent)
|
||||||
|
is a full setup built on them.
|
||||||
|
|
||||||
|

|
||||||
|
|
||||||
|
## What it gives a setup
|
||||||
|
|
||||||
|
- **agents/** - what an agent is: a Claude agent with a prompt per
|
||||||
|
conversation kind, skill sets per kind, gateway tools, a `PreToolUse`
|
||||||
|
policy; a Raycast agent for cheap fast answers.
|
||||||
|
- **backends/** - how a turn runs: the Claude Agent SDK subprocess with a
|
||||||
|
session pool, transcripts and session stores; the Raycast wire.
|
||||||
|
- **conversations/** - the model of a conversation: kinds (master, branch,
|
||||||
|
deep, job, fork), rows and windows, the inject queue with priorities,
|
||||||
|
seeds, turns, questions, closing through a distiller, master rotation,
|
||||||
|
the envelope, the in-process gateway tools.
|
||||||
|
- **jobs/** - cron, webhook and event jobs on pgqueuer, deferred injects,
|
||||||
|
the subscription budget.
|
||||||
|
- **frontends/** - the windows: Telegram (master = General, topic = branch,
|
||||||
|
files out of the agent through `send_file`),
|
||||||
|
markdown files in a vault, `/api` + the admin SPA, an Anthropic-compatible
|
||||||
|
`/anthropic/v1/messages`, MCP re-exposure at `/mcp/<name>`, and
|
||||||
|
`WebhookFrontend` for a window a setup declares itself - its own request
|
||||||
|
and response schema, its own agent, its own token scope
|
||||||
|
(`docs/FRONTEND-PLUGINS.md`).
|
||||||
|
- **mcp/** - MCP servers a setup declares (stdio, http, python tools),
|
||||||
|
aggregated in-process and handed to agents by name.
|
||||||
|
- **vault/** - watching a directory of notes for the envelope.
|
||||||
|
- **security/** - bearer tokens, the audit log, credential redaction.
|
||||||
|
- **keystore/** - secrets the model runs commands with: values and files
|
||||||
|
sealed with `SECRETS_KEY` (AES-256-GCM) in the gateway's database, set
|
||||||
|
write-only from the admin, handed out only through `secret run NAME --
|
||||||
|
cmd` (one `sudo` rule for `beaver-runner`, a root-only socket, files in
|
||||||
|
`/dev/shm`, values masked in the command's output).
|
||||||
|
- **events/** - the in-process bus every frontend taps, and the stream
|
||||||
|
event protocol backends emit.
|
||||||
|
- **storage/** - SQLModel tables and the Postgres session store.
|
||||||
|
|
||||||
|
`config.py` defines `Gateway` and loads a setup; `app.py` builds the runtime
|
||||||
|
from it; `cli.py` is the entrypoint. The admin SPA in `ui/` is the browser
|
||||||
|
window on all of it: conversations by day, the thread with the tree of tool
|
||||||
|
calls, the context of the last turn, memory, jobs, usage, tokens, secrets and the audit.
|
||||||
|
|
||||||
|

|
||||||
|
|
||||||
|
 Every string the gateway puts in front of
|
||||||
|
a model or a user has an English default and is overridable
|
||||||
|
(`ConversationTexts`, `TelegramTexts`).
|
||||||
|
|
||||||
|
## Run
|
||||||
|
|
||||||
|
```sh
|
||||||
|
cd examples
|
||||||
|
CLAUDE_CODE_OAUTH_TOKEN=$(claude setup-token) docker compose up --build
|
||||||
|
```
|
||||||
|
|
||||||
|
`examples/config.py` is the smallest setup that starts: one agent, one
|
||||||
|
python tool, one cron job, the built-in frontends on port 8000; the compose
|
||||||
|
file next to it adds Postgres. Without Docker: `uv sync`, build the SPA with
|
||||||
|
`cd ui && bun install && bun run build`, then run `uv run beaver-gateway` with
|
||||||
|
`DATABASE_URL` and `CONFIG_PATH` set. A setup's own secrets come from `.env`
|
||||||
|
next to its config (`load_dotenv` runs before the config is executed).
|
||||||
|
|
||||||
|
## Develop
|
||||||
|
|
||||||
|
```sh
|
||||||
|
make check # ruff format --check, ruff check, ty, pytest
|
||||||
|
make fix
|
||||||
|
```
|
||||||
|
|
||||||
|
No comments in code: modules carry a one-line docstring saying what they are
|
||||||
|
for, public configuration fields carry a one-line docstring saying what they
|
||||||
|
mean; everything else is the code.
|
||||||
Executable
+2
@@ -0,0 +1,2 @@
|
|||||||
|
#!/bin/sh
|
||||||
|
exec /usr/local/bin/python3 -I /app/src/beaver_gateway/keystore/cli.py --root "$@"
|
||||||
Executable
+2
@@ -0,0 +1,2 @@
|
|||||||
|
#!/bin/sh
|
||||||
|
exec /usr/local/bin/python3 -I /app/src/beaver_gateway/keystore/cli.py "$@"
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
Defaults:beaver-runner !syslog
|
||||||
|
beaver-runner ALL=(root) NOPASSWD: /usr/local/libexec/beaver-secret
|
||||||
+35
-175
@@ -1,204 +1,64 @@
|
|||||||
# Sample user config — grows alongside the implementation phases.
|
"""The smallest setup: one agent, one tool, one job, the built-in frontends.
|
||||||
#
|
|
||||||
# Loader (beaver_gateway/config_loader.py) execs this file with
|
The loader execs this file with ``ClaudeAgent``, ``RaycastAgent``,
|
||||||
# ClaudeAgent, RaycastAgent, McpServer, ExposedMcp, Gateway already
|
``McpServer``, ``ExposedMcp`` and ``Gateway`` already bound; the imports
|
||||||
# bound — so importing them is optional. We import explicitly here so
|
below are for type checkers. A real setup grows into a package next to
|
||||||
# IDEs and type-checkers see real symbols instead of free variables.
|
this file (see beaver-agent).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
import tempfile
|
import tempfile
|
||||||
from datetime import date
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from beaver_gateway.agents.base import ExposedMcp
|
from beaver_gateway.agents.base import ExposedMcp
|
||||||
from beaver_gateway.agents.claude import ClaudeAgent
|
from beaver_gateway.agents.claude import ClaudeAgent, Prompts
|
||||||
from beaver_gateway.agents.raycast import RaycastAgent, RemoteTool, UserPreferences
|
from beaver_gateway.config import Gateway
|
||||||
from beaver_gateway.core.registry import Gateway
|
|
||||||
from beaver_gateway.core.turn_record import slugify
|
|
||||||
from beaver_gateway.frontends.admin import AdminFrontend
|
from beaver_gateway.frontends.admin import AdminFrontend
|
||||||
from beaver_gateway.frontends.anthropic import AnthropicMessagesFrontend
|
from beaver_gateway.frontends.anthropic import AnthropicMessagesFrontend
|
||||||
|
from beaver_gateway.frontends.api import ApiFrontend
|
||||||
from beaver_gateway.frontends.markdown import MarkdownFrontend
|
from beaver_gateway.frontends.markdown import MarkdownFrontend
|
||||||
from beaver_gateway.frontends.mcp_server import McpServerFrontend
|
from beaver_gateway.frontends.mcp_server import McpServerFrontend
|
||||||
|
from beaver_gateway.jobs.scheduler import Job, JobRun
|
||||||
from beaver_gateway.mcp.types import McpServer
|
from beaver_gateway.mcp.types import McpServer
|
||||||
|
|
||||||
|
log = logging.getLogger("example")
|
||||||
|
|
||||||
def chat_path(title: str, agent: str, vault: Path) -> Path: # noqa: ARG001
|
VAULT = Path(tempfile.mkdtemp(prefix="beaver-vault-")).resolve()
|
||||||
"""Where a new chat file lands in the vault.
|
PROMPT = VAULT / "assistant.md"
|
||||||
|
PROMPT.write_text(
|
||||||
Called by ``MarkdownFrontend`` for every conversation that needs a
|
"You are a concise assistant. When asked the time, call `current_time`.\n"
|
||||||
file: a ``deep`` chat spawned by the dispatcher, a ``/v1/messages``
|
)
|
||||||
chat (``title`` = first user message), or with ``log_all_chats=True``
|
|
||||||
the archive of a stateless agent's turns. Return value can be
|
|
||||||
absolute or relative; relative paths are anchored under ``vault``.
|
|
||||||
|
|
||||||
Layout below: ``<vault>/<YYYY-MM>/<YYYY-MM-DD>_<topic>.md``.
|
|
||||||
"""
|
|
||||||
today = date.today()
|
|
||||||
return vault / f"{today:%Y-%m}" / f"{today:%Y-%m-%d}_{slugify(title, maxlen=40)}.md"
|
|
||||||
|
|
||||||
|
|
||||||
def current_time() -> str:
|
def current_time() -> str:
|
||||||
"""Return the current local time as an ISO-8601 string.
|
return datetime.now(UTC).astimezone().isoformat()
|
||||||
|
|
||||||
Trivial demo tool for the Phase 2.1 internal MCP aggregator —
|
|
||||||
confirms a ``python_tool`` namespace is reachable on
|
|
||||||
``http://127.0.0.1:<INTERNAL_MCP_PORT>/mcp/time``.
|
|
||||||
"""
|
|
||||||
from datetime import datetime
|
|
||||||
|
|
||||||
return datetime.now().astimezone().isoformat()
|
async def heartbeat(run: JobRun) -> None:
|
||||||
|
master = await run.master()
|
||||||
|
log.info("heartbeat: master=%s", master.external_id if master else None)
|
||||||
|
|
||||||
|
|
||||||
gateway = Gateway(
|
gateway = Gateway(
|
||||||
agents=[
|
agents=[
|
||||||
# Phase 2.2 — ClaudeCodeBackendAdapter routes this agent's
|
|
||||||
# ``/v1/messages`` calls through ``claude-code-api``. The
|
|
||||||
# ``time`` MCP gets exposed as ``mcp__time__current_time`` to
|
|
||||||
# the subscription claude session via
|
|
||||||
# ``BackendOptions.mcp_servers`` pointing at the internal
|
|
||||||
# aggregator on ``127.0.0.1:INTERNAL_MCP_PORT/mcp/time/``.
|
|
||||||
#
|
|
||||||
# Fresh empty tempdir (not a hardcoded ``/tmp``) for two
|
|
||||||
# reasons: claude-code-api derives the JSONL project-key from
|
|
||||||
# ``cwd``, but claude itself writes the JSONL using the cwd's
|
|
||||||
# realpath — on macOS ``/tmp`` and ``/var/folders/...`` are
|
|
||||||
# both ``/private/...`` symlinks, so unresolved cwds make
|
|
||||||
# ``JsonlWatcher`` time out waiting on the wrong path. The
|
|
||||||
# explicit ``.resolve()`` collapses the symlink before claude
|
|
||||||
# ever sees the dir, and ``mkdtemp`` guarantees the directory
|
|
||||||
# is empty so claude does not pick up leftover files.
|
|
||||||
ClaudeAgent(
|
ClaudeAgent(
|
||||||
name="stub",
|
name="assistant",
|
||||||
model="claude-sonnet-4-6",
|
model="claude-sonnet-5",
|
||||||
# ``system_prompt`` is appended to claude's built-in agent
|
cwd=VAULT,
|
||||||
# prompt (via ``--append-system-prompt``) — so it adds the
|
prompts=Prompts(master=(PROMPT,), branch=(PROMPT,), deep=(PROMPT,)),
|
||||||
# agent's identity on top of claude-code's baseline tool
|
gateway_tools=("spawn", "say", "schedule"),
|
||||||
# knowledge, rather than replacing it. Same shape as the
|
|
||||||
# RaycastAgent's ``system_prompt → additional_system_instructions``
|
|
||||||
# mapping. For full ``BackendOptions`` knobs (timeouts,
|
|
||||||
# extra_args, history mode, etc.) import ``ClaudeCodeOptions``
|
|
||||||
# and pass ``options=ClaudeCodeOptions(...)``.
|
|
||||||
system_prompt=(
|
|
||||||
"You are a stub agent used to validate the Phase 0 skeleton.\n"
|
|
||||||
"If asked the current time, call the `current_time`"
|
|
||||||
" MCP tool instead of guessing."
|
|
||||||
),
|
|
||||||
cwd=Path(tempfile.mkdtemp(prefix="beaver-stub-cwd-")).resolve(),
|
|
||||||
expose_mcps=(ExposedMcp(name="time"),),
|
expose_mcps=(ExposedMcp(name="time"),),
|
||||||
),
|
)
|
||||||
# Phase 1.2 — a RaycastAgent the AnthropicMessagesFrontend will
|
|
||||||
# route via RaycastBackend. Phase 1.5 added the per-agent knobs
|
|
||||||
# (`temperature` / `additional_system_instructions` / etc.) —
|
|
||||||
# only `model`, `system_prompt`, and at least one of the others
|
|
||||||
# is mandatory.
|
|
||||||
RaycastAgent(
|
|
||||||
name="research",
|
|
||||||
model="Gemini 3.1 Flash Lite",
|
|
||||||
system_prompt=(
|
|
||||||
"You are a research assistant. "
|
|
||||||
"Reply in the user's language. Cite URLs when you use web search."
|
|
||||||
),
|
|
||||||
temperature=0.5,
|
|
||||||
available_native_tools=(RemoteTool.WEB_SEARCH, RemoteTool.READ_PAGE),
|
|
||||||
# Lambda so today's date is rebuilt on every request while
|
|
||||||
# locale/timezone stay pinned. ``True`` would give the same
|
|
||||||
# fresh-date behaviour but would also auto-pick host locale
|
|
||||||
# and timezone (``en-US`` / system tz), which isn't what we
|
|
||||||
# want here.
|
|
||||||
user_preferences=lambda: UserPreferences(
|
|
||||||
locale="en-GB",
|
|
||||||
timezone="Europe/Berlin",
|
|
||||||
current_date=date.today().isoformat(), # noqa: DTZ011 — local date is intended
|
|
||||||
),
|
|
||||||
),
|
|
||||||
],
|
|
||||||
mcps=[
|
|
||||||
# Phase 2.1 — bundle of plain Python callables exposed as one
|
|
||||||
# FastMCP namespace. The internal aggregator mounts it under
|
|
||||||
# ``/mcp/time`` on ``127.0.0.1:INTERNAL_MCP_PORT``; Phase 2.2's
|
|
||||||
# ClaudeCode adapter forwards that URL into
|
|
||||||
# ``BackendOptions.mcp_servers``. Phase 3's ``McpServerFrontend``
|
|
||||||
# reverse-proxies the same internal URL out to external clients.
|
|
||||||
McpServer.python_tool(name="time", tools=[current_time])
|
|
||||||
# Phase 3 — illustrates the ``lenient`` flag. Real-world stdio MCPs
|
|
||||||
# sometimes print "Processing..." or other chatter to stdout before
|
|
||||||
# their actual JSON-RPC frames; the default mcp client forwards
|
|
||||||
# those parse failures downstream as warnings (visible in
|
|
||||||
# Cursor/Cline). With ``lenient=True`` we silently drop non-JSON
|
|
||||||
# lines, so downstream UIs see clean JSON-RPC only. The command
|
|
||||||
# below is just a placeholder — replace with whatever stdio MCP
|
|
||||||
# you actually want gateway to ingest (e.g. an obsidian-mcp).
|
|
||||||
#
|
|
||||||
# Commented out by default: example users won't have the binary
|
|
||||||
# installed and an unreachable command makes ``docker compose up``
|
|
||||||
# surface a confusing "command not found" line at first request.
|
|
||||||
# Uncomment after pointing ``command`` at a real stdio MCP.
|
|
||||||
#
|
|
||||||
# McpServer.stdio(
|
|
||||||
# name="obsidian",
|
|
||||||
# command=["uvx", "mcp-obsidian"],
|
|
||||||
# env={"OBSIDIAN_API_KEY": "..."},
|
|
||||||
# lenient=True,
|
|
||||||
# ),
|
|
||||||
],
|
],
|
||||||
|
mcps=[McpServer.python_tool(name="time", tools=[current_time])],
|
||||||
frontends=[
|
frontends=[
|
||||||
# Phase 1.4 — expose the agents as `model=<name>` on an
|
|
||||||
# Anthropic-compatible Messages endpoint. Auth comes from
|
|
||||||
# `BOOTSTRAP_TOKENS` in the env (`name1:value1,name2:value2`).
|
|
||||||
#
|
|
||||||
# Every HTTP frontend is mounted under its own path on the one
|
|
||||||
# gateway port (`Gateway.port`, 8000 here): `/anthropic/v1/messages`.
|
|
||||||
# Behind a reverse proxy set `Gateway(public_url="https://domain.com")`
|
|
||||||
# so advertised endpoints use the outside origin; the proxy just
|
|
||||||
# forwards everything to the gateway, no prefix stripping.
|
|
||||||
AnthropicMessagesFrontend(),
|
AnthropicMessagesFrontend(),
|
||||||
# Phase 3 — re-exposes every declared `McpServer` outside the
|
|
||||||
# gateway with bearer auth + audit log. Each namespace lives at
|
|
||||||
# `/mcp/<name>/`; a flat bundle is published at `/mcp/all/`.
|
|
||||||
# Discovery page (HTML, auth-gated) at `/mcp/` with copy-pastable
|
|
||||||
# Cursor / Claude Desktop snippets. Auth re-uses `BOOTSTRAP_TOKENS`.
|
|
||||||
McpServerFrontend(),
|
McpServerFrontend(),
|
||||||
# Phase 4.3 — browser admin UI. Creds come from
|
|
||||||
# `ADMIN_USER`/`ADMIN_PASS`; the session cookie is signed with
|
|
||||||
# `SESSION_SECRET`. Use it to mint tokens (Argon2-hashed in
|
|
||||||
# the DB), revoke them, and watch the audit log. Scope is
|
|
||||||
# enforced on the bearer frontends: tokens minted with scope
|
|
||||||
# `messages` only work on `/v1/messages`; `mcp` only on
|
|
||||||
# `/mcp/<name>`; `*` works everywhere. Served at `/admin/`; `/`
|
|
||||||
# redirects there.
|
|
||||||
AdminFrontend(),
|
AdminFrontend(),
|
||||||
# Obsidian-vault chat frontend. Each `.md` is one conversation
|
ApiFrontend(master_agent="assistant", branch_agent="assistant"),
|
||||||
# (User/Assistant turn pairs). The Obsidian companion plugin
|
MarkdownFrontend(vault_path=VAULT / "chats", default_agent="assistant"),
|
||||||
# POSTs `{filename, content?}` to `/chat` — the frontend reads
|
|
||||||
# the file, runs the agent if the last turn is `user`, and
|
|
||||||
# appends the assistant reply back. With `log_all_chats=True`
|
|
||||||
# *every* turn (from Anthropic Messages too) is mirrored into
|
|
||||||
# `{vault}/_logs/<agent>/` so the vault is the central archive.
|
|
||||||
#
|
|
||||||
# `vault_path` here points at a per-restart tempdir so the
|
|
||||||
# example boots cleanly; in real deployments mount the
|
|
||||||
# Obsidian-sync container's vault volume to a stable path and
|
|
||||||
# pass that instead.
|
|
||||||
# Mounted at `/md` (`/md/chat`, `/md/chat/stream`).
|
|
||||||
MarkdownFrontend(
|
|
||||||
# Point at the dedicated chats subdir of your real Obsidian
|
|
||||||
# vault — the gateway has no idea (and no need) about other
|
|
||||||
# notes outside it. Path resolution / vault-escape checks
|
|
||||||
# are anchored here, so absolute-path attempts (and ``..``
|
|
||||||
# tricks) can't reach notes alongside it.
|
|
||||||
#
|
|
||||||
# vault_path=Path("/Users/me/Obsidian/Personal/chats"),
|
|
||||||
#
|
|
||||||
# Per-restart tempdir kept here so the example boots even
|
|
||||||
# without a real vault on the host.
|
|
||||||
vault_path=Path(tempfile.mkdtemp(prefix="beaver-vault-")).resolve(),
|
|
||||||
default_agent="research",
|
|
||||||
log_all_chats=True,
|
|
||||||
# ``chat_path`` (optional) overrides the default
|
|
||||||
# ``{vault}/_logs/<agent>/<date>_<slug>.md`` layout. Heads up:
|
|
||||||
# any custom path forces ``warm_index`` to scan the entire
|
|
||||||
# vault on startup so the fingerprint→file map of archived
|
|
||||||
# stateless chats survives a restart no matter where you put
|
|
||||||
# files.
|
|
||||||
chat_path=chat_path,
|
|
||||||
),
|
|
||||||
],
|
],
|
||||||
|
jobs=[Job("heartbeat", heartbeat, cron="0 * * * *")],
|
||||||
|
tz="UTC",
|
||||||
)
|
)
|
||||||
+27
-22
@@ -1,31 +1,36 @@
|
|||||||
# Minimal compose for the Phase 0 DoD: build the image, mount a
|
|
||||||
# single-file user config, watch the process print its summary line
|
|
||||||
# and exit cleanly.
|
|
||||||
#
|
|
||||||
# Real deployments add a postgres service, raycast-config.json bind,
|
|
||||||
# ~/.config/claude bind, etc. — Phase 5 lives in a separate example
|
|
||||||
# repo (see PLAN §5.1).
|
|
||||||
|
|
||||||
services:
|
services:
|
||||||
|
postgres:
|
||||||
|
image: postgres:16-alpine
|
||||||
|
environment:
|
||||||
|
POSTGRES_USER: beaver
|
||||||
|
POSTGRES_PASSWORD: beaver
|
||||||
|
POSTGRES_DB: beaver
|
||||||
|
volumes:
|
||||||
|
- postgres-data:/var/lib/postgresql/data
|
||||||
|
healthcheck:
|
||||||
|
test: ["CMD-SHELL", "pg_isready -U beaver -d beaver"]
|
||||||
|
interval: 5s
|
||||||
|
timeout: 3s
|
||||||
|
retries: 20
|
||||||
|
|
||||||
gateway:
|
gateway:
|
||||||
build:
|
build:
|
||||||
context: ..
|
context: ..
|
||||||
dockerfile: Dockerfile
|
depends_on:
|
||||||
|
postgres:
|
||||||
|
condition: service_healthy
|
||||||
environment:
|
environment:
|
||||||
DATABASE_URL: "sqlite:///./gateway.db"
|
CONFIG_PATH: /config/config.py
|
||||||
ADMIN_USER: "admin"
|
DATABASE_URL: postgresql+psycopg://beaver:beaver@postgres:5432/beaver
|
||||||
ADMIN_PASS: "change-me"
|
ADMIN_USER: admin
|
||||||
SESSION_SECRET: "dev-secret-change-me"
|
ADMIN_PASS: change-me
|
||||||
# Phase 1.4 — comma-separated `name:value` pairs. The frontend
|
SESSION_SECRET: change-me-too
|
||||||
# 401s every request until at least one token is set.
|
BOOTSTRAP_TOKENS: dev:change-me:*
|
||||||
BOOTSTRAP_TOKENS: "cursor:dev-token-change-me"
|
CLAUDE_CODE_OAUTH_TOKEN: ${CLAUDE_CODE_OAUTH_TOKEN}
|
||||||
# RAYCAST_BEARER / RAYCAST_DEVICE_ID + a raycast.json bind go
|
|
||||||
# here once you wire up an actual RaycastAgent. The example
|
|
||||||
# config.py declares one, so set these (or remove the agent)
|
|
||||||
# before exposing port 8000.
|
|
||||||
ports:
|
ports:
|
||||||
# The one gateway port: /anthropic, /mcp, /admin, /md under it
|
|
||||||
# (change ADMIN_USER/ADMIN_PASS/SESSION_SECRET before exposing).
|
|
||||||
- "8000:8000"
|
- "8000:8000"
|
||||||
volumes:
|
volumes:
|
||||||
- ./config.py:/config/config.py:ro
|
- ./config.py:/config/config.py:ro
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
postgres-data:
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 138 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 81 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 122 KiB |
+2
-1
@@ -14,8 +14,9 @@ dependencies = [
|
|||||||
"anthropic>=0.103.0",
|
"anthropic>=0.103.0",
|
||||||
"anyio>=4.13.0",
|
"anyio>=4.13.0",
|
||||||
"argon2-cffi>=25.1.0",
|
"argon2-cffi>=25.1.0",
|
||||||
"claude-agent-sdk>=0.2.150",
|
"claude-agent-sdk>=0.2.157",
|
||||||
"croniter>=6.2.4",
|
"croniter>=6.2.4",
|
||||||
|
"cryptography>=48.0.0",
|
||||||
"fastapi>=0.136.1",
|
"fastapi>=0.136.1",
|
||||||
"fastmcp>=3.3.1",
|
"fastmcp>=3.3.1",
|
||||||
"greenlet>=3.5.0",
|
"greenlet>=3.5.0",
|
||||||
|
|||||||
@@ -17,9 +17,8 @@ from pydantic import BaseModel, ConfigDict
|
|||||||
class ExposedMcp:
|
class ExposedMcp:
|
||||||
"""Reference to an ``McpServer`` (by name) exposed to a single agent.
|
"""Reference to an ``McpServer`` (by name) exposed to a single agent.
|
||||||
|
|
||||||
``tools`` is an allowlist of tool names (``None`` = every tool);
|
``tools`` allowlists tool names (``None`` = all); ``deny`` removes
|
||||||
``deny`` is a tuple of ``fnmatch`` patterns removed on top of that,
|
``fnmatch`` patterns on top, e.g. ``("delete_*",)``.
|
||||||
e.g. ``("delete_*",)``.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
name: str
|
name: str
|
||||||
@@ -37,3 +36,8 @@ class BaseAgent(BaseModel):
|
|||||||
system_prompt: str
|
system_prompt: str
|
||||||
expose_mcps: tuple[ExposedMcp, ...] = ()
|
expose_mcps: tuple[ExposedMcp, ...] = ()
|
||||||
accept_client_tools: bool = False
|
accept_client_tools: bool = False
|
||||||
|
|
||||||
|
@property
|
||||||
|
def thinking_is_progress(self) -> bool:
|
||||||
|
"""Whether non-empty thinking is text for the human rather than reasoning."""
|
||||||
|
return False
|
||||||
@@ -1,31 +1,35 @@
|
|||||||
"""Claude agent definition, backed by the Claude Agent SDK.
|
"""Claude agent definition, backed by the Claude Agent SDK.
|
||||||
|
|
||||||
The system prompt is ``system_prompt`` verbatim or, per conversation kind,
|
``system_prompt`` is used verbatim, or per-kind granules from ``prompts``
|
||||||
the granules named in ``prompts`` assembled at every session spawn (see
|
are assembled via ``agents/prompts.py``. ``skill_sets`` are
|
||||||
``core/prompt.py``). ``skill_sets`` are directories of ``<skill>/SKILL.md``
|
``<skill>/SKILL.md`` directories, each becoming a local SDK plugin.
|
||||||
folders; each becomes a local SDK plugin. Nothing from disk is loaded
|
|
||||||
otherwise: the adapter runs with ``setting_sources=[]``.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from collections.abc import Mapping # noqa: TC003 - pydantic runtime
|
from collections.abc import Mapping # noqa: TC003
|
||||||
from pathlib import Path # noqa: TC003 - pydantic runtime
|
from pathlib import Path # noqa: TC003
|
||||||
|
|
||||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||||
|
|
||||||
from beaver_gateway.agents.base import BaseAgent
|
from beaver_gateway.agents.base import BaseAgent
|
||||||
from beaver_gateway.core.kinds import KINDS, Kind
|
from beaver_gateway.agents.policy import PolicyRule # noqa: TC001
|
||||||
from beaver_gateway.core.policy import PolicyRule # noqa: TC001 - pydantic runtime
|
from beaver_gateway.agents.prompts import PromptSource # noqa: TC001
|
||||||
from beaver_gateway.core.prompt import PromptSource # noqa: TC001 - pydantic runtime
|
from beaver_gateway.conversations.kinds import KINDS, Kind
|
||||||
|
|
||||||
__all__ = ["ClaudeAgent", "ClaudeOptions", "Prompts"]
|
__all__ = ["PROGRESS_MODELS", "ClaudeAgent", "ClaudeOptions", "Prompts", "SkillSets"]
|
||||||
|
|
||||||
|
PROGRESS_MODELS = frozenset({"claude-opus-5-5", "claude-fable-5-1"})
|
||||||
|
"""Models whose non-empty thinking is the text written between tool calls."""
|
||||||
|
|
||||||
|
|
||||||
class ClaudeOptions(BaseModel):
|
class ClaudeOptions(BaseModel):
|
||||||
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
|
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
|
||||||
|
|
||||||
effort: str | None = None
|
effort: str | None = None
|
||||||
|
thinking: Mapping[str, str] | None = None
|
||||||
|
"""Thinking config for the SDK; ``None`` keeps the CLI default."""
|
||||||
|
|
||||||
include_partial_messages: bool = True
|
include_partial_messages: bool = True
|
||||||
"""Token-level deltas on the wire; off means one delta per finished block."""
|
"""Token-level deltas on the wire; off means one delta per finished block."""
|
||||||
|
|
||||||
@@ -42,16 +46,17 @@ class ClaudeOptions(BaseModel):
|
|||||||
"""Extra inherited variable names to let through the env whitelist."""
|
"""Extra inherited variable names to let through the env whitelist."""
|
||||||
|
|
||||||
max_turns: int | None = None
|
max_turns: int | None = None
|
||||||
|
max_buffer_size: int = 16 * 1024 * 1024
|
||||||
session_store_flush: str = "eager"
|
session_store_flush: str = "eager"
|
||||||
"""``eager`` mirrors every transcript frame as it lands, so a gateway
|
"""``eager`` mirrors every transcript frame as it lands, so a gateway
|
||||||
killed mid-turn loses at most the frame in flight."""
|
killed mid-turn loses at most the frame in flight."""
|
||||||
|
|
||||||
|
|
||||||
class Prompts(BaseModel):
|
class Prompts(BaseModel):
|
||||||
"""Prompt assembly per conversation kind (§3.12): the granules, in order.
|
"""Prompt assembly per conversation kind: the granules, in order.
|
||||||
|
|
||||||
A kind left ``None`` is not served by the agent; ``ClaudeAgent.kinds``
|
A kind left ``None`` is not served; ``ClaudeAgent.kinds`` follows from
|
||||||
follows from the kinds set here.
|
the kinds set here.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
model_config = ConfigDict(frozen=True)
|
model_config = ConfigDict(frozen=True)
|
||||||
@@ -70,6 +75,21 @@ class Prompts(BaseModel):
|
|||||||
return tuple(k for k in KINDS if self.for_kind(k) is not None)
|
return tuple(k for k in KINDS if self.for_kind(k) is not None)
|
||||||
|
|
||||||
|
|
||||||
|
class SkillSets(BaseModel):
|
||||||
|
"""Skill-set directories per conversation kind; ``None`` means none."""
|
||||||
|
|
||||||
|
model_config = ConfigDict(frozen=True)
|
||||||
|
|
||||||
|
master: tuple[Path, ...] | None = None
|
||||||
|
branch: tuple[Path, ...] | None = None
|
||||||
|
deep: tuple[Path, ...] | None = None
|
||||||
|
job: tuple[Path, ...] | None = None
|
||||||
|
fork: tuple[Path, ...] | None = None
|
||||||
|
|
||||||
|
def for_kind(self, kind: Kind) -> tuple[Path, ...]:
|
||||||
|
return getattr(self, kind) or ()
|
||||||
|
|
||||||
|
|
||||||
class ClaudeAgent(BaseAgent):
|
class ClaudeAgent(BaseAgent):
|
||||||
cwd: Path
|
cwd: Path
|
||||||
system_prompt: str = ""
|
system_prompt: str = ""
|
||||||
@@ -77,19 +97,21 @@ class ClaudeAgent(BaseAgent):
|
|||||||
|
|
||||||
prompts: Prompts = Field(default_factory=Prompts)
|
prompts: Prompts = Field(default_factory=Prompts)
|
||||||
kinds: tuple[Kind, ...] = ()
|
kinds: tuple[Kind, ...] = ()
|
||||||
"""Conversation kinds this agent serves; ``create``/``spawn`` reject the
|
"""Conversation kinds this agent serves; ``create``/``spawn`` reject the rest.
|
||||||
rest. Defaults to the kinds ``prompts`` covers, or ``("deep",)`` for a
|
Defaults to what ``prompts`` covers, or ``("deep",)`` for a verbatim prompt."""
|
||||||
verbatim ``system_prompt``."""
|
|
||||||
|
skill_sets: tuple[Path, ...] | SkillSets = ()
|
||||||
|
"""Skill-set directories, each a local plugin: one tuple for every kind,
|
||||||
|
or a ``SkillSets`` giving each kind its own (``None`` = no skills)."""
|
||||||
|
|
||||||
skill_sets: tuple[Path, ...] = ()
|
|
||||||
gateway_tools: tuple[str, ...] = ()
|
gateway_tools: tuple[str, ...] = ()
|
||||||
"""Gateway tools exposed in-process (``read_conversation``, ``spawn``,
|
"""Gateway tools exposed in-process (``read_conversation``, ``spawn``,
|
||||||
``say``, ``schedule``, ``inject``); empty = no gateway MCP server."""
|
``say``, ``schedule``, ``inject``); empty = no gateway MCP server."""
|
||||||
|
|
||||||
options: ClaudeOptions = Field(default_factory=ClaudeOptions)
|
options: ClaudeOptions = Field(default_factory=ClaudeOptions)
|
||||||
policy: tuple[PolicyRule, ...] = ()
|
policy: tuple[PolicyRule, ...] = ()
|
||||||
"""``PreToolUse`` rules (§3.7), run in order on every tool call; the
|
"""``PreToolUse`` rules, run in order on every tool call; the first
|
||||||
first ``Deny`` is what the model reads back. See ``core/policy``."""
|
``Deny`` is what the model reads back. See ``agents/policy.py``."""
|
||||||
|
|
||||||
@model_validator(mode="after")
|
@model_validator(mode="after")
|
||||||
def _kinds_follow_prompts(self) -> ClaudeAgent:
|
def _kinds_follow_prompts(self) -> ClaudeAgent:
|
||||||
@@ -103,8 +125,19 @@ class ClaudeAgent(BaseAgent):
|
|||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
@property
|
||||||
|
def thinking_is_progress(self) -> bool:
|
||||||
|
"""A progress model, unless ``display: summarized`` brings reasoning."""
|
||||||
|
summarized = (self.options.thinking or {}).get("display") == "summarized"
|
||||||
|
return self.model in PROGRESS_MODELS and not summarized
|
||||||
|
|
||||||
def prompt_for(self, kind: Kind) -> tuple[PromptSource, ...] | None:
|
def prompt_for(self, kind: Kind) -> tuple[PromptSource, ...] | None:
|
||||||
return self.prompts.for_kind(kind)
|
return self.prompts.for_kind(kind)
|
||||||
|
|
||||||
|
def skills_for(self, kind: Kind) -> tuple[Path, ...]:
|
||||||
|
if isinstance(self.skill_sets, SkillSets):
|
||||||
|
return self.skill_sets.for_kind(kind)
|
||||||
|
return self.skill_sets
|
||||||
|
|
||||||
def serves(self, kind: str) -> bool:
|
def serves(self, kind: str) -> bool:
|
||||||
return kind in self.kinds
|
return kind in self.kinds
|
||||||
@@ -1,18 +1,8 @@
|
|||||||
"""PreToolUse policy - the boundary without permission prompts (§3.7).
|
"""PreToolUse policy - the boundary without permission prompts.
|
||||||
|
|
||||||
``bypassPermissions`` everywhere; what a model may do is decided by
|
A rule is a callable ``(ToolCall) -> Deny | None`` declared per agent
|
||||||
mounts, ``disallowed_tools`` and the rules here. A rule is a callable
|
(``ClaudeAgent.policy``); the first :class:`Deny` becomes the tool result
|
||||||
``(ToolCall) -> Deny | None`` declared per agent (``ClaudeAgent.policy``);
|
the model reads back. A rule that raises is a deny too: fails closed.
|
||||||
the SDK backend registers one in-process ``PreToolUse`` hook that runs
|
|
||||||
the rules in order and turns the first :class:`Deny` into a hook deny
|
|
||||||
whose reason the model reads as the tool result. Rules never see
|
|
||||||
secrets and never prompt - they only say no, with a reason.
|
|
||||||
|
|
||||||
A rule that raises is a deny too: the boundary fails closed, the
|
|
||||||
traceback lands in the log.
|
|
||||||
|
|
||||||
Every tool call - allowed or denied - is reported to the audit sink the
|
|
||||||
backend was given, so the admin audit page shows what the model touched.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -90,11 +80,14 @@ class ToolCall:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def resolve(self, raw: str) -> Path:
|
def resolve(self, raw: str) -> Path:
|
||||||
|
"""Absolute path for ``raw``, normalised lexically.
|
||||||
|
|
||||||
|
Not via ``Path.resolve``, which would follow symlinks on the host,
|
||||||
|
not the model's view.
|
||||||
|
"""
|
||||||
path = Path(raw).expanduser()
|
path = Path(raw).expanduser()
|
||||||
if not path.is_absolute():
|
if not path.is_absolute():
|
||||||
path = self.cwd / path
|
path = self.cwd / path
|
||||||
# ``resolve`` would follow symlinks on the gateway host, which may
|
|
||||||
# not be the model's view; normalise lexically instead.
|
|
||||||
return Path(*_normalize(path.parts))
|
return Path(*_normalize(path.parts))
|
||||||
|
|
||||||
|
|
||||||
@@ -136,7 +129,7 @@ async def evaluate(
|
|||||||
verdict = rule(call)
|
verdict = rule(call)
|
||||||
if inspect.isawaitable(verdict):
|
if inspect.isawaitable(verdict):
|
||||||
verdict = await verdict
|
verdict = await verdict
|
||||||
except Exception: # noqa: BLE001 - a broken rule must fail closed
|
except Exception: # noqa: BLE001
|
||||||
name = getattr(rule, "__name__", repr(rule))
|
name = getattr(rule, "__name__", repr(rule))
|
||||||
_log.exception("policy rule %s failed on %s", name, call.tool)
|
_log.exception("policy rule %s failed on %s", name, call.tool)
|
||||||
return Deny(reason=f"policy rule {name} failed; the call is refused")
|
return Deny(reason=f"policy rule {name} failed; the call is refused")
|
||||||
@@ -1,12 +1,7 @@
|
|||||||
"""System prompt assembly from a list of source files.
|
"""System prompt assembly from a list of source files.
|
||||||
|
|
||||||
The gateway holds no prompt text: an agent names its granules (paths from
|
:func:`assemble` concatenates an agent's granules in order. A source is a
|
||||||
``config.py``) and :func:`assemble` concatenates them in that order, so the
|
path, or a ``(tag, path)`` pair whose content is wrapped in ``<tag>...</tag>``.
|
||||||
result is byte-for-byte identical for every session of the same agent as
|
|
||||||
long as the files are. A source is a path, or a ``(tag, path)`` pair whose
|
|
||||||
content is wrapped in ``<tag>...</tag>`` - the markup lives here, the vault
|
|
||||||
keeps plain markdown. Each granule's hash is logged at assembly so a
|
|
||||||
drifted prompt can be traced to the file that changed.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -21,12 +16,13 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
__all__ = ["PromptSource", "assemble"]
|
__all__ = ["PromptSource", "assemble"]
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.core.prompt")
|
_log = logging.getLogger("beaver_gateway.agents.prompts")
|
||||||
|
|
||||||
PromptSource = str | Path | tuple[str, str | Path]
|
PromptSource = str | Path | tuple[str, str | Path]
|
||||||
|
|
||||||
|
|
||||||
def assemble(sources: Iterable[PromptSource]) -> str:
|
def assemble(sources: Iterable[PromptSource]) -> str:
|
||||||
|
"""Concatenate ``sources`` in order into one system prompt."""
|
||||||
parts: list[str] = []
|
parts: list[str] = []
|
||||||
for source in sources:
|
for source in sources:
|
||||||
tag, raw = source if isinstance(source, tuple) else (None, source)
|
tag, raw = source if isinstance(source, tuple) else (None, source)
|
||||||
@@ -1,31 +1,12 @@
|
|||||||
"""Raycast agent definition.
|
"""Raycast agent definition.
|
||||||
|
|
||||||
Field set is the union of ``raycast_api.ChatAPI.stream`` parameters that
|
Field set is the union of ``raycast_api.ChatAPI.stream`` parameters useful
|
||||||
make sense as **per-agent defaults**. Per-request values (currently:
|
as per-agent defaults; per-request values override them where both exist.
|
||||||
``temperature``) win when both are set; the rest fall back to whatever
|
|
||||||
the agent declared, then to Raycast's own defaults.
|
|
||||||
|
|
||||||
``BaseAgent.system_prompt`` maps onto Raycast's wire field
|
|
||||||
``additional_system_instructions`` (the slot the real client uses for
|
|
||||||
*content*); the wire field ``system_instructions`` stays at the Raycast
|
|
||||||
source default — ``"markdown"`` for ``AI_CHAT``, ``"plain"`` otherwise.
|
|
||||||
We don't expose that wire dichotomy to the user — they get one
|
|
||||||
conceptual "system prompt".
|
|
||||||
|
|
||||||
Excluded on purpose:
|
|
||||||
|
|
||||||
* ``buffer_id``/``message_id``/``current_date`` — per-call ephemeral
|
|
||||||
* ``provider`` override — escape hatch for non-catalog models, no clear
|
|
||||||
use case yet (revisit in PRD §14 when discovery lands)
|
|
||||||
* ``locale`` — process-wide via ``Settings.raycast_locale`` because we
|
|
||||||
only spin up one ``raycast_api.Client`` per gateway
|
|
||||||
* ``system_instructions`` (wire) — that's a format marker, not content;
|
|
||||||
the SDK fills it from the source default and we let it
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from raycast_api import ( # noqa: F401 — UserPreferences re-exported for user configs
|
from raycast_api import ( # noqa: F401
|
||||||
RemoteTool,
|
RemoteTool,
|
||||||
Source,
|
Source,
|
||||||
UserPreferences,
|
UserPreferences,
|
||||||
@@ -36,44 +17,21 @@ from beaver_gateway.agents.base import BaseAgent
|
|||||||
|
|
||||||
|
|
||||||
class RaycastAgent(BaseAgent):
|
class RaycastAgent(BaseAgent):
|
||||||
"""Agent backed by ``raycast-api``.
|
"""Agent backed by ``raycast-api``."""
|
||||||
|
|
||||||
``available_native_tools`` is the closed set of Raycast's server-side
|
|
||||||
"remote tools" (``web_search``, ``search_images``, ``read_page``) —
|
|
||||||
typed as ``RemoteTool`` so config-time IDE completion lists exactly
|
|
||||||
the three valid values. Pydantic also coerces string literals, so
|
|
||||||
``("web_search", "read_page")`` keeps working unchanged.
|
|
||||||
|
|
||||||
``user_preferences`` toggles the auto-generated ``<user-preferences>``
|
|
||||||
block Raycast prepends to ``additional_system_instructions``:
|
|
||||||
|
|
||||||
* ``True`` (default) → auto from host locale/timezone/today, rebuilt
|
|
||||||
every request so the date stays fresh;
|
|
||||||
* ``False`` → omit the block entirely;
|
|
||||||
* ``UserPreferences(...)`` instance → used verbatim (frozen at the
|
|
||||||
time the agent was loaded, so the date won't auto-update);
|
|
||||||
* ``Callable[[], UserPreferencesArg]`` → re-invoked on every request.
|
|
||||||
Use this for the common case "fresh date but custom
|
|
||||||
locale/timezone": ``user_preferences=lambda:
|
|
||||||
UserPreferences(locale="ru-RU", timezone="Europe/Berlin",
|
|
||||||
current_date=date.today().isoformat())``. Callables may nest
|
|
||||||
(a lambda returning a lambda…) but there's no real reason to.
|
|
||||||
|
|
||||||
The library uses this block for date/locale-aware formatting, not
|
|
||||||
for personalisation/memory — those are out of scope upstream.
|
|
||||||
|
|
||||||
``reasoning_effort`` values vary by model: GPT-5 takes
|
|
||||||
``"minimal"|"low"|"medium"|"high"``; Anthropic exposes nothing here
|
|
||||||
(Claude reasoning lives in a separate ``…-reasoning`` model variant
|
|
||||||
in the catalog). Unknown effort for the chosen model is ignored
|
|
||||||
server-side, so we stay loose as ``str | None``.
|
|
||||||
"""
|
|
||||||
|
|
||||||
streaming: bool = True
|
streaming: bool = True
|
||||||
available_native_tools: tuple[RemoteTool, ...] = ()
|
available_native_tools: tuple[RemoteTool, ...] = ()
|
||||||
|
"""Raycast's server-side "remote tools": ``web_search``, ``search_images``,
|
||||||
|
``read_page``."""
|
||||||
|
|
||||||
source: Source = Source.AI_CHAT
|
source: Source = Source.AI_CHAT
|
||||||
|
|
||||||
temperature: float | None = None
|
temperature: float | None = None
|
||||||
reasoning_effort: str | None = None
|
reasoning_effort: str | None = None
|
||||||
|
"""Model-specific (e.g. GPT-5: ``"minimal"``/``"low"``/``"medium"``/``"high"``);
|
||||||
|
ignored server-side if the model doesn't support it."""
|
||||||
|
|
||||||
tool_choice: str | None = None
|
tool_choice: str | None = None
|
||||||
user_preferences: UserPreferencesArg = True
|
user_preferences: UserPreferencesArg = True
|
||||||
|
"""Auto-fills the ``<user-preferences>`` block from host locale/timezone/date;
|
||||||
|
``False`` omits it, or pass ``UserPreferences(...)`` / a zero-arg callable."""
|
||||||
@@ -0,0 +1,484 @@
|
|||||||
|
"""Build the runtime from a ``Gateway`` and run every part of it until shutdown."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import functools
|
||||||
|
import logging
|
||||||
|
from contextlib import AsyncExitStack
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
import psycopg
|
||||||
|
import uvicorn
|
||||||
|
from pgqueuer import PsycopgDriver
|
||||||
|
from raycast_api import Client as RaycastClient
|
||||||
|
from raycast_api.config import Config as RaycastConfig
|
||||||
|
|
||||||
|
from beaver_gateway.agents.claude import ClaudeAgent
|
||||||
|
from beaver_gateway.agents.raycast import RaycastAgent
|
||||||
|
from beaver_gateway.backends.claude_sdk import (
|
||||||
|
ClaudeSdkBackend,
|
||||||
|
RunnerConfig,
|
||||||
|
UsageEvent,
|
||||||
|
)
|
||||||
|
from beaver_gateway.backends.raycast import RaycastBackend
|
||||||
|
from beaver_gateway.backends.sessions import SessionPool
|
||||||
|
from beaver_gateway.conversations.envelope import Envelope
|
||||||
|
from beaver_gateway.conversations.rotation import Rotation, RotationPolicy
|
||||||
|
from beaver_gateway.conversations.service import Conversations
|
||||||
|
from beaver_gateway.conversations.tools import build_tool_server
|
||||||
|
from beaver_gateway.events.bus import EventBus
|
||||||
|
from beaver_gateway.frontends.base import GatewayRuntime
|
||||||
|
from beaver_gateway.frontends.bearer import require_token
|
||||||
|
from beaver_gateway.frontends.root import build_root_app
|
||||||
|
from beaver_gateway.jobs.scheduler import Scheduler
|
||||||
|
from beaver_gateway.keystore import SecretBox, SecretStore
|
||||||
|
from beaver_gateway.keystore import serve as serve_secrets
|
||||||
|
from beaver_gateway.mcp.internal_app import build_internal_app
|
||||||
|
from beaver_gateway.security.auth import TokenStore, scopes_with
|
||||||
|
from beaver_gateway.storage import (
|
||||||
|
Database,
|
||||||
|
PostgresSessionStore,
|
||||||
|
Usage,
|
||||||
|
append_audit,
|
||||||
|
append_usage,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Iterable, Iterator, Sequence
|
||||||
|
|
||||||
|
from claude_agent_sdk import McpSdkServerConfig
|
||||||
|
from fastmcp import FastMCP
|
||||||
|
from fastmcp.tools.base import Tool as FastMCPTool
|
||||||
|
from starlette.applications import Starlette
|
||||||
|
from starlette.requests import Request
|
||||||
|
from starlette.types import ASGIApp
|
||||||
|
|
||||||
|
from beaver_gateway.agents.base import BaseAgent
|
||||||
|
from beaver_gateway.agents.policy import ToolAudit
|
||||||
|
from beaver_gateway.backends.base import Backend
|
||||||
|
from beaver_gateway.config import Gateway
|
||||||
|
from beaver_gateway.mcp.types import McpServerT
|
||||||
|
from beaver_gateway.settings import Settings
|
||||||
|
|
||||||
|
__all__ = ["AgentRegistry", "McpRegistry", "run"]
|
||||||
|
|
||||||
|
_log = logging.getLogger("beaver_gateway.app")
|
||||||
|
|
||||||
|
|
||||||
|
class AgentRegistry:
|
||||||
|
def __init__(self, agents: Iterable[BaseAgent]) -> None:
|
||||||
|
self._agents: dict[str, BaseAgent] = {}
|
||||||
|
for a in agents:
|
||||||
|
if a.name in self._agents:
|
||||||
|
msg = f"duplicate agent name: {a.name!r}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
self._agents[a.name] = a
|
||||||
|
|
||||||
|
def __getitem__(self, name: str) -> BaseAgent:
|
||||||
|
return self._agents[name]
|
||||||
|
|
||||||
|
def get(self, name: str) -> BaseAgent | None:
|
||||||
|
return self._agents.get(name)
|
||||||
|
|
||||||
|
def __iter__(self) -> Iterator[BaseAgent]:
|
||||||
|
return iter(self._agents.values())
|
||||||
|
|
||||||
|
def __len__(self) -> int:
|
||||||
|
return len(self._agents)
|
||||||
|
|
||||||
|
def __contains__(self, name: object) -> bool:
|
||||||
|
return name in self._agents
|
||||||
|
|
||||||
|
|
||||||
|
class McpRegistry:
|
||||||
|
def __init__(self, mcps: Iterable[McpServerT]) -> None:
|
||||||
|
self._mcps: dict[str, McpServerT] = {}
|
||||||
|
for m in mcps:
|
||||||
|
if m.name in self._mcps:
|
||||||
|
msg = f"duplicate mcp name: {m.name!r}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
self._mcps[m.name] = m
|
||||||
|
|
||||||
|
def __getitem__(self, name: str) -> McpServerT:
|
||||||
|
return self._mcps[name]
|
||||||
|
|
||||||
|
def get(self, name: str) -> McpServerT | None:
|
||||||
|
return self._mcps.get(name)
|
||||||
|
|
||||||
|
def __iter__(self) -> Iterator[McpServerT]:
|
||||||
|
return iter(self._mcps.values())
|
||||||
|
|
||||||
|
def __len__(self) -> int:
|
||||||
|
return len(self._mcps)
|
||||||
|
|
||||||
|
def __contains__(self, name: object) -> bool:
|
||||||
|
return name in self._mcps
|
||||||
|
|
||||||
|
|
||||||
|
async def run(gateway: Gateway, settings: Settings) -> None:
|
||||||
|
agents = AgentRegistry(gateway.agents)
|
||||||
|
mcps = McpRegistry(gateway.mcps)
|
||||||
|
db = Database(settings.database_url)
|
||||||
|
await db.create_all()
|
||||||
|
secrets = SecretStore(db, _secret_box(settings))
|
||||||
|
token_store = TokenStore(
|
||||||
|
db,
|
||||||
|
bootstrap=TokenStore.parse_bootstrap(settings.bootstrap_tokens),
|
||||||
|
bootstrap_scopes=TokenStore.parse_bootstrap_scopes(settings.bootstrap_tokens),
|
||||||
|
)
|
||||||
|
|
||||||
|
async with AsyncExitStack() as stack:
|
||||||
|
stack.push_async_callback(db.dispose)
|
||||||
|
await token_store.start()
|
||||||
|
stack.push_async_callback(token_store.stop)
|
||||||
|
internal_app, internal_urls, mcp_servers = _build_internal_mcp(
|
||||||
|
gateway.mcps, settings=settings
|
||||||
|
)
|
||||||
|
mcp_tools = await _prefetch_mcp_tools(mcp_servers)
|
||||||
|
|
||||||
|
pool = SessionPool()
|
||||||
|
bus = EventBus()
|
||||||
|
late = _LateConversations()
|
||||||
|
session_store = PostgresSessionStore(db)
|
||||||
|
backends = 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,
|
||||||
|
frontends=gateway.frontends,
|
||||||
|
envelope=Envelope(
|
||||||
|
watch=gateway.watch, tz=gateway.tz, recall=gateway.recall
|
||||||
|
),
|
||||||
|
distiller=gateway.distiller,
|
||||||
|
user_sink=gateway.user_sink,
|
||||||
|
)
|
||||||
|
late.conversations = conversations
|
||||||
|
scheduler = Scheduler(
|
||||||
|
conversations=conversations,
|
||||||
|
jobs=gateway.jobs,
|
||||||
|
driver=await _pgqueuer_driver(settings.database_url, stack),
|
||||||
|
budget=gateway.budget,
|
||||||
|
rotation=Rotation(
|
||||||
|
conversations, gateway.rotation or RotationPolicy(tz=gateway.tz)
|
||||||
|
),
|
||||||
|
tz=gateway.tz,
|
||||||
|
)
|
||||||
|
conversations.scheduler = scheduler
|
||||||
|
|
||||||
|
runtime = GatewayRuntime(
|
||||||
|
agents=agents,
|
||||||
|
mcps=mcps,
|
||||||
|
backends=backends,
|
||||||
|
token_store=token_store,
|
||||||
|
db=db,
|
||||||
|
mcp_internal_urls=internal_urls,
|
||||||
|
admin_user=settings.admin_user,
|
||||||
|
admin_pass=settings.admin_pass,
|
||||||
|
session_secret=settings.session_secret,
|
||||||
|
frontends=tuple(gateway.frontends),
|
||||||
|
conversations=conversations,
|
||||||
|
bus=bus,
|
||||||
|
pool=pool,
|
||||||
|
scheduler=scheduler,
|
||||||
|
public_url=gateway.public_url.rstrip("/") if gateway.public_url else None,
|
||||||
|
scopes=scopes_with(fe.scope for fe in gateway.frontends),
|
||||||
|
send_file_roots=_roots(settings.send_file_roots),
|
||||||
|
secrets=secrets,
|
||||||
|
)
|
||||||
|
for fe in gateway.frontends:
|
||||||
|
fe.configure(runtime)
|
||||||
|
|
||||||
|
_log.info(
|
||||||
|
"beaver-gateway: loaded %d agents, %d mcps, %d frontends",
|
||||||
|
len(agents),
|
||||||
|
len(mcps),
|
||||||
|
len(gateway.frontends),
|
||||||
|
)
|
||||||
|
if not gateway.frontends:
|
||||||
|
return
|
||||||
|
|
||||||
|
await conversations.start()
|
||||||
|
stack.push_async_callback(conversations.stop)
|
||||||
|
await scheduler.start()
|
||||||
|
stack.push_async_callback(scheduler.stop)
|
||||||
|
hooks = scheduler.app(
|
||||||
|
functools.partial(_authorize_hook, runtime=runtime, scope="api")
|
||||||
|
)
|
||||||
|
async with asyncio.TaskGroup() as tg:
|
||||||
|
tg.create_task(pool.reap_loop())
|
||||||
|
if internal_app is not None:
|
||||||
|
tg.create_task(_serve_internal_mcp(internal_app, settings=settings))
|
||||||
|
tg.create_task(_serve_root(gateway, extra={"/hooks": hooks}))
|
||||||
|
if gateway.watch is not None:
|
||||||
|
tg.create_task(gateway.watch.run())
|
||||||
|
if secrets.enabled:
|
||||||
|
tg.create_task(_serve_secrets(secrets, settings=settings, db=db))
|
||||||
|
for fe in gateway.frontends:
|
||||||
|
tg.create_task(fe.serve())
|
||||||
|
|
||||||
|
|
||||||
|
def _secret_box(settings: Settings) -> SecretBox | None:
|
||||||
|
if not settings.secrets_key:
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
return SecretBox.from_hex(settings.secrets_key)
|
||||||
|
except ValueError as exc:
|
||||||
|
_log.error("secrets are off: %s", exc)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
async def _serve_secrets(
|
||||||
|
store: SecretStore, *, settings: Settings, db: Database
|
||||||
|
) -> None:
|
||||||
|
async def used(names: Sequence[str]) -> None:
|
||||||
|
try:
|
||||||
|
async with db.session() as session:
|
||||||
|
await append_audit(
|
||||||
|
session,
|
||||||
|
actor="runner",
|
||||||
|
kind="secret_use",
|
||||||
|
detail={"names": list(names)},
|
||||||
|
)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("secret audit write failed")
|
||||||
|
|
||||||
|
try:
|
||||||
|
await serve_secrets(store, settings.secrets_socket, on_use=used)
|
||||||
|
except OSError:
|
||||||
|
_log.exception("secrets socket %s did not open", settings.secrets_socket)
|
||||||
|
|
||||||
|
|
||||||
|
async def _authorize_hook(
|
||||||
|
request: Request, *, runtime: GatewayRuntime, scope: str
|
||||||
|
) -> str:
|
||||||
|
return await require_token(request, runtime, scope=scope)
|
||||||
|
|
||||||
|
|
||||||
|
async def _pgqueuer_driver(url: str, stack: AsyncExitStack) -> PsycopgDriver | None:
|
||||||
|
plain = _plain_postgres_url(url)
|
||||||
|
if plain is None:
|
||||||
|
return None
|
||||||
|
conn = await psycopg.AsyncConnection.connect(plain, autocommit=True)
|
||||||
|
stack.push_async_callback(conn.close)
|
||||||
|
return PsycopgDriver(conn)
|
||||||
|
|
||||||
|
|
||||||
|
def _plain_postgres_url(url: str) -> str | None:
|
||||||
|
for prefix in ("postgresql+psycopg://", "postgresql://", "postgres://"):
|
||||||
|
if url.startswith(prefix):
|
||||||
|
return "postgresql://" + url[len(prefix) :]
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
async def _serve_root(gateway: Gateway, *, extra: dict[str, ASGIApp]) -> None:
|
||||||
|
app = build_root_app(gateway.frontends, extra=extra)
|
||||||
|
config = uvicorn.Config(app, host=gateway.host, port=gateway.port, log_level="info")
|
||||||
|
_log.info(
|
||||||
|
"gateway on http://%s:%d - %s",
|
||||||
|
gateway.host,
|
||||||
|
gateway.port,
|
||||||
|
", ".join([*(fe.path for fe in gateway.frontends if fe.path), *extra])
|
||||||
|
or "no http frontends",
|
||||||
|
)
|
||||||
|
await uvicorn.Server(config).serve()
|
||||||
|
|
||||||
|
|
||||||
|
def _build_internal_mcp(
|
||||||
|
mcps: list[McpServerT], *, settings: Settings
|
||||||
|
) -> tuple[Starlette | None, dict[str, str], dict[str, FastMCP]]:
|
||||||
|
if not mcps:
|
||||||
|
return None, {}, {}
|
||||||
|
return build_internal_app(mcps, host="127.0.0.1", port=settings.internal_mcp_port)
|
||||||
|
|
||||||
|
|
||||||
|
async def _prefetch_mcp_tools(
|
||||||
|
servers: dict[str, FastMCP],
|
||||||
|
) -> dict[str, list[FastMCPTool]]:
|
||||||
|
out: dict[str, list[FastMCPTool]] = {}
|
||||||
|
for name, server in servers.items():
|
||||||
|
try:
|
||||||
|
out[name] = list(await server.list_tools())
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("failed to list tools for MCP %r, skipping", name)
|
||||||
|
out[name] = []
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
async def _serve_internal_mcp(app: Starlette, *, settings: Settings) -> None:
|
||||||
|
config = uvicorn.Config(
|
||||||
|
app,
|
||||||
|
host="127.0.0.1",
|
||||||
|
port=settings.internal_mcp_port,
|
||||||
|
log_level="warning",
|
||||||
|
loop="uvloop",
|
||||||
|
)
|
||||||
|
_log.info(
|
||||||
|
"internal MCP aggregator on http://127.0.0.1:%d/mcp/<name>",
|
||||||
|
settings.internal_mcp_port,
|
||||||
|
)
|
||||||
|
await uvicorn.Server(config).serve()
|
||||||
|
|
||||||
|
|
||||||
|
def _roots(value: str) -> tuple[Path, ...]:
|
||||||
|
return tuple(Path(p) for p in value.split(":") if p.strip())
|
||||||
|
|
||||||
|
|
||||||
|
class _LateConversations:
|
||||||
|
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 ask(self, key: str, payload: dict[str, Any]) -> str:
|
||||||
|
if self.conversations is None:
|
||||||
|
msg = "conversations service is not up yet"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
answer = await self.conversations.ask(key, payload)
|
||||||
|
return self.conversations.answer_text(answer)
|
||||||
|
|
||||||
|
|
||||||
|
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]:
|
||||||
|
backends: dict[str, Backend] = {}
|
||||||
|
|
||||||
|
raycast_agents = [a for a in agents if isinstance(a, RaycastAgent)]
|
||||||
|
if raycast_agents:
|
||||||
|
client = await _try_open_raycast_client(settings, stack)
|
||||||
|
if client is not None:
|
||||||
|
raycast_backend = RaycastBackend(
|
||||||
|
client, mcp_servers=mcp_servers, mcp_tools=mcp_tools
|
||||||
|
)
|
||||||
|
for a in raycast_agents:
|
||||||
|
backends[a.name] = raycast_backend
|
||||||
|
|
||||||
|
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()
|
||||||
|
}
|
||||||
|
|
||||||
|
async def record_usage(event: UsageEvent) -> None:
|
||||||
|
row = Usage(
|
||||||
|
agent_name=event.agent_name,
|
||||||
|
conversation_id=event.conversation_id,
|
||||||
|
session_id=event.session_id,
|
||||||
|
model=event.model,
|
||||||
|
effort=event.effort,
|
||||||
|
input_tokens=event.usage.input_tokens,
|
||||||
|
output_tokens=event.usage.output_tokens,
|
||||||
|
cache_read_tokens=event.usage.cache_read_tokens,
|
||||||
|
cache_creation_tokens=event.usage.cache_creation_tokens,
|
||||||
|
context_tokens=event.usage.context_tokens,
|
||||||
|
cost_usd=event.usage.cost_usd,
|
||||||
|
duration_ms=event.usage.duration_ms,
|
||||||
|
num_turns=event.usage.num_turns,
|
||||||
|
model_usage=event.usage.model_usage,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
async with db.session() as session:
|
||||||
|
await append_usage(session, row)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("usage write failed for %s", event.agent_name)
|
||||||
|
|
||||||
|
async def record_tool(event: ToolAudit) -> None:
|
||||||
|
detail = {
|
||||||
|
"conversation": event.conversation,
|
||||||
|
"kind": event.kind,
|
||||||
|
"tool": event.tool,
|
||||||
|
"decision": event.decision,
|
||||||
|
"reason": event.reason,
|
||||||
|
"brief": event.brief,
|
||||||
|
}
|
||||||
|
try:
|
||||||
|
async with db.session() as session:
|
||||||
|
await append_audit(
|
||||||
|
session,
|
||||||
|
actor=f"agent:{event.agent}",
|
||||||
|
kind="tool_call",
|
||||||
|
agent_name=event.agent,
|
||||||
|
detail=detail,
|
||||||
|
)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("tool audit write failed for %s", event.agent)
|
||||||
|
|
||||||
|
for a in agents:
|
||||||
|
if isinstance(a, ClaudeAgent):
|
||||||
|
adapter = ClaudeSdkBackend(
|
||||||
|
agent=a,
|
||||||
|
mcp_internal_urls=mcp_internal_urls,
|
||||||
|
session_store=session_store,
|
||||||
|
mcp_tool_names=mcp_tool_names,
|
||||||
|
runner=runner,
|
||||||
|
usage_sink=record_usage,
|
||||||
|
pool=pool,
|
||||||
|
tool_server=functools.partial(late.server, names=a.gateway_tools),
|
||||||
|
asker=late.ask,
|
||||||
|
audit_sink=record_tool,
|
||||||
|
)
|
||||||
|
await stack.enter_async_context(adapter)
|
||||||
|
backends[a.name] = adapter
|
||||||
|
|
||||||
|
return backends
|
||||||
|
|
||||||
|
|
||||||
|
async def _try_open_raycast_client(
|
||||||
|
settings: Settings, stack: AsyncExitStack
|
||||||
|
) -> RaycastClient | None:
|
||||||
|
if not settings.raycast_bearer:
|
||||||
|
_log.warning(
|
||||||
|
"RaycastAgent present but RAYCAST_BEARER is unset; those agents 503"
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
if not settings.raycast_device_id:
|
||||||
|
_log.warning(
|
||||||
|
"RaycastAgent present but RAYCAST_DEVICE_ID is unset, those agents 503 "
|
||||||
|
"(generate once with `python -c 'import secrets; "
|
||||||
|
"print(secrets.token_hex(32))'`)"
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
if not settings.raycast_config_path.exists():
|
||||||
|
_log.warning(
|
||||||
|
"RaycastAgent present but %s is missing, those agents 503",
|
||||||
|
settings.raycast_config_path,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
config = RaycastConfig.load(settings.raycast_config_path)
|
||||||
|
client = RaycastClient(
|
||||||
|
config=config,
|
||||||
|
bearer_token=settings.raycast_bearer,
|
||||||
|
device_id=settings.raycast_device_id,
|
||||||
|
locale=settings.raycast_locale,
|
||||||
|
)
|
||||||
|
return await stack.enter_async_context(client)
|
||||||
@@ -1,6 +1,6 @@
|
|||||||
"""Backend adapters.
|
"""Backend adapters.
|
||||||
|
|
||||||
Each backend wraps a provider SDK (``raycast-api``, ``claude-agent-sdk``)
|
Each backend wraps a provider SDK (``raycast-api``, ``claude-agent-sdk``)
|
||||||
and yields the unified :class:`~beaver_gateway.core.events.MessageStreamEvent`
|
and yields the unified :class:`~beaver_gateway.events.stream.MessageStreamEvent`
|
||||||
family. The Anthropic-style frontend serialises events straight to SSE.
|
family. The Anthropic-style frontend serialises events straight to SSE.
|
||||||
"""
|
"""
|
||||||
@@ -1,19 +1,8 @@
|
|||||||
"""Backend protocol.
|
"""Backend protocol.
|
||||||
|
|
||||||
A backend turns an Anthropic-style turn (``messages`` + agent definition)
|
A backend turns an Anthropic-style turn (``messages`` + agent definition)
|
||||||
into a stream of :class:`~beaver_gateway.core.events.MessageStreamEvent`
|
into a stream of :class:`~beaver_gateway.events.stream.MessageStreamEvent`
|
||||||
records. The frontend serializes whatever comes out straight to SSE, so
|
records; provider quirks are translated here, not in the frontend.
|
||||||
backends are the only place where provider quirks are translated.
|
|
||||||
|
|
||||||
Implementations are plain :class:`typing.Protocol` conformers - no ABC
|
|
||||||
subclassing - to keep them swappable in tests with bare async generators.
|
|
||||||
|
|
||||||
``**options`` is the one extension point. Known keys, all optional and
|
|
||||||
ignored by backends that don't keep state: ``conversation_id`` (stable id
|
|
||||||
the backend may pin a live session to), ``session_id`` (backend session to
|
|
||||||
resume when nothing is live), ``capture`` (a
|
|
||||||
:class:`~beaver_gateway.core.turn_capture.TurnCapture` the backend fills
|
|
||||||
after the stream closes).
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -26,7 +15,7 @@ if TYPE_CHECKING:
|
|||||||
from anthropic.types import MessageParam
|
from anthropic.types import MessageParam
|
||||||
|
|
||||||
from beaver_gateway.agents.base import BaseAgent
|
from beaver_gateway.agents.base import BaseAgent
|
||||||
from beaver_gateway.core.events import MessageStreamEvent
|
from beaver_gateway.events.stream import MessageStreamEvent
|
||||||
|
|
||||||
|
|
||||||
class Backend(Protocol):
|
class Backend(Protocol):
|
||||||
@@ -40,5 +29,13 @@ class Backend(Protocol):
|
|||||||
system: str | None = None,
|
system: str | None = None,
|
||||||
**options: Any,
|
**options: Any,
|
||||||
) -> AsyncIterator[MessageStreamEvent]:
|
) -> AsyncIterator[MessageStreamEvent]:
|
||||||
"""Yield Anthropic stream events for one turn against ``agent``."""
|
"""Yield Anthropic stream events for one turn against ``agent``.
|
||||||
|
|
||||||
|
``**options`` keys are optional and ignored by backends that don't
|
||||||
|
keep state: ``conversation_id`` (stable id the backend may pin a
|
||||||
|
live session to), ``session_id`` (backend session to resume when
|
||||||
|
nothing is live), ``capture`` (a
|
||||||
|
:class:`~beaver_gateway.backends.capture.TurnCapture` the backend
|
||||||
|
fills after the stream closes).
|
||||||
|
"""
|
||||||
...
|
...
|
||||||
@@ -1,10 +1,7 @@
|
|||||||
"""Side channel for what a backend learned during one turn.
|
"""Side channel for what a backend learned during one turn.
|
||||||
|
|
||||||
Frontends pass ``capture=TurnCapture()`` through ``Backend.complete``'s
|
Frontends pass ``capture=TurnCapture()`` through ``Backend.complete``'s
|
||||||
``**options``. Backends that keep state per conversation (the Claude SDK
|
``**options``; only backends that keep state per conversation fill it in.
|
||||||
adapter) fill it in after the stream closes; backends that don't
|
|
||||||
(anthropic, raycast) drop the kwarg and the frontend falls back to a
|
|
||||||
text-only history.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -1,29 +1,7 @@
|
|||||||
"""Claude Agent SDK backend adapter.
|
"""Claude Agent SDK backend adapter.
|
||||||
|
|
||||||
One :class:`ClaudeSdkBackend` per :class:`ClaudeAgent`. A live session is
|
One :class:`ClaudeSdkBackend` per :class:`ClaudeAgent`; its sessions live in
|
||||||
one ``ClaudeSDKClient`` (one claude subprocess) and runs one turn at a
|
a shared :class:`~beaver_gateway.backends.sessions.SessionPool`.
|
||||||
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 (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
|
|
||||||
rebased across the API calls claude makes inside the turn.
|
|
||||||
|
|
||||||
Process isolation: claude is spawned through a small exec wrapper that
|
|
||||||
drops every inherited environment variable outside a whitelist and, when
|
|
||||||
``RunnerConfig.user`` is set, switches to that uid before exec (done in the
|
|
||||||
wrapper rather than via ``subprocess(user=...)``, which uvloop rejects).
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -42,6 +20,7 @@ import tempfile
|
|||||||
import time
|
import time
|
||||||
import uuid
|
import uuid
|
||||||
import warnings
|
import warnings
|
||||||
|
from collections import deque
|
||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -59,6 +38,7 @@ from claude_agent_sdk import (
|
|||||||
PermissionResultDeny,
|
PermissionResultDeny,
|
||||||
ResultMessage,
|
ResultMessage,
|
||||||
StreamEvent,
|
StreamEvent,
|
||||||
|
TaskNotificationMessage,
|
||||||
TextBlock,
|
TextBlock,
|
||||||
ThinkingBlock,
|
ThinkingBlock,
|
||||||
ToolResultBlock,
|
ToolResultBlock,
|
||||||
@@ -67,9 +47,23 @@ from claude_agent_sdk import (
|
|||||||
project_key_for_directory,
|
project_key_for_directory,
|
||||||
)
|
)
|
||||||
|
|
||||||
from beaver_gateway.core import policy as policy_mod
|
from beaver_gateway.agents import policy as policy_mod
|
||||||
from beaver_gateway.core import prompt as prompt_assembly
|
from beaver_gateway.agents import prompts as prompt_assembly
|
||||||
from beaver_gateway.core.events import (
|
from beaver_gateway.backends.capture import TurnCapture, TurnUsage
|
||||||
|
from beaver_gateway.backends.sessions import (
|
||||||
|
STDERR_LINES,
|
||||||
|
Session,
|
||||||
|
SessionClient,
|
||||||
|
SessionPool,
|
||||||
|
)
|
||||||
|
from beaver_gateway.backends.transcript import (
|
||||||
|
build_entries,
|
||||||
|
close_open_tool_uses,
|
||||||
|
fingerprint,
|
||||||
|
text_of,
|
||||||
|
)
|
||||||
|
from beaver_gateway.conversations.kinds import as_kind
|
||||||
|
from beaver_gateway.events.stream import (
|
||||||
StopReason,
|
StopReason,
|
||||||
build_content_block_stop,
|
build_content_block_stop,
|
||||||
build_input_json_delta,
|
build_input_json_delta,
|
||||||
@@ -83,18 +77,17 @@ from beaver_gateway.core.events import (
|
|||||||
build_thinking_delta,
|
build_thinking_delta,
|
||||||
build_tool_use_block_start,
|
build_tool_use_block_start,
|
||||||
)
|
)
|
||||||
from beaver_gateway.core.kinds import as_kind
|
from beaver_gateway.security.redact import redact
|
||||||
from beaver_gateway.core.sessions import Session, SessionClient, SessionPool
|
|
||||||
from beaver_gateway.core.transcript import (
|
|
||||||
build_entries,
|
|
||||||
close_open_tool_uses,
|
|
||||||
fingerprint,
|
|
||||||
text_of,
|
|
||||||
)
|
|
||||||
from beaver_gateway.core.turn_capture import TurnCapture, TurnUsage
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
|
from collections.abc import (
|
||||||
|
AsyncIterator,
|
||||||
|
Awaitable,
|
||||||
|
Callable,
|
||||||
|
Coroutine,
|
||||||
|
Iterable,
|
||||||
|
Sequence,
|
||||||
|
)
|
||||||
|
|
||||||
from anthropic.types import MessageParam
|
from anthropic.types import MessageParam
|
||||||
from claude_agent_sdk import (
|
from claude_agent_sdk import (
|
||||||
@@ -107,18 +100,18 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
from beaver_gateway.agents.base import BaseAgent
|
from beaver_gateway.agents.base import BaseAgent
|
||||||
from beaver_gateway.agents.claude import ClaudeAgent
|
from beaver_gateway.agents.claude import ClaudeAgent
|
||||||
from beaver_gateway.core.events import MessageStreamEvent
|
from beaver_gateway.conversations.kinds import Kind
|
||||||
|
from beaver_gateway.events.stream import MessageStreamEvent
|
||||||
|
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.backends.claude_sdk")
|
_log = logging.getLogger("beaver_gateway.backends.claude_sdk")
|
||||||
|
|
||||||
# §3.7: in bypass the callback only ever sees AskUserQuestion, and that is
|
|
||||||
# exactly the one we want - the SDK's warning about the rest is noise here.
|
|
||||||
warnings.filterwarnings("ignore", category=CanUseToolShadowedWarning)
|
warnings.filterwarnings("ignore", category=CanUseToolShadowedWarning)
|
||||||
|
|
||||||
ASK_TOOL = "AskUserQuestion"
|
ASK_TOOL = "AskUserQuestion"
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
|
"ORIGIN_HEADER",
|
||||||
"AuditSink",
|
"AuditSink",
|
||||||
"ClaudeSdkBackend",
|
"ClaudeSdkBackend",
|
||||||
"RunnerConfig",
|
"RunnerConfig",
|
||||||
@@ -128,6 +121,9 @@ __all__ = [
|
|||||||
"fingerprint",
|
"fingerprint",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
ORIGIN_HEADER = "X-Conversation-Origin"
|
||||||
|
"""Rides on every exposed MCP request; the gateway proxy passes it upstream."""
|
||||||
|
|
||||||
ENV_KEEP: tuple[str, ...] = (
|
ENV_KEEP: tuple[str, ...] = (
|
||||||
"PATH",
|
"PATH",
|
||||||
"HOME",
|
"HOME",
|
||||||
@@ -175,8 +171,7 @@ AuditSink = "Callable[[policy_mod.ToolAudit], Awaitable[None]]"
|
|||||||
Asker = "Callable[[str, dict[str, Any]], Awaitable[str]]"
|
Asker = "Callable[[str, dict[str, Any]], Awaitable[str]]"
|
||||||
"""``(conversation_key, AskUserQuestion input) -> text the model reads as the
|
"""``(conversation_key, AskUserQuestion input) -> text the model reads as the
|
||||||
tool result``. The only channel an answer has in bypass mode is
|
tool result``. The only channel an answer has in bypass mode is
|
||||||
``PermissionResultDeny.message`` (spike S1, s05): ``updated_input`` never
|
``PermissionResultDeny.message``: ``updated_input`` never reaches the model."""
|
||||||
reaches the model."""
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
@@ -225,8 +220,16 @@ class ClaudeSdkBackend:
|
|||||||
tool_server: Callable[[str, str], McpSdkServerConfig | None] | None = None,
|
tool_server: Callable[[str, str], McpSdkServerConfig | None] | None = None,
|
||||||
asker: Callable[[str, dict[str, Any]], Awaitable[str]] | None = None,
|
asker: Callable[[str, dict[str, Any]], Awaitable[str]] | None = None,
|
||||||
audit_sink: Callable[[policy_mod.ToolAudit], Awaitable[None]] | None = None,
|
audit_sink: Callable[[policy_mod.ToolAudit], Awaitable[None]] | None = None,
|
||||||
|
waker: Callable[[str, str], Coroutine[Any, Any, None]] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
|
"""One backend per Claude agent.
|
||||||
|
|
||||||
|
``waker(key, note)`` is called when a session starts a turn by itself
|
||||||
|
(a background task reported back); the caller then attaches.
|
||||||
|
"""
|
||||||
self._agent = agent
|
self._agent = agent
|
||||||
|
self.waker = waker
|
||||||
|
self._tasks: set[asyncio.Task[None]] = set()
|
||||||
self._audit_sink = audit_sink
|
self._audit_sink = audit_sink
|
||||||
self._store = session_store
|
self._store = session_store
|
||||||
self._runner = runner or RunnerConfig()
|
self._runner = runner or RunnerConfig()
|
||||||
@@ -270,23 +273,67 @@ class ClaudeSdkBackend:
|
|||||||
await self._pool.close(key)
|
await self._pool.close(key)
|
||||||
|
|
||||||
async def interrupt(self, key: str) -> bool:
|
async def interrupt(self, key: str) -> bool:
|
||||||
live = self._pool.get(key)
|
return await self._pool.interrupt(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:
|
def live(self, key: str) -> Session | None:
|
||||||
return self._pool.get(key)
|
return self._pool.get(key)
|
||||||
|
|
||||||
async def repair_session(self, session_id: str) -> int:
|
def has_pending(self, key: str) -> bool:
|
||||||
|
"""A turn the session started on its own is waiting for ``attach``."""
|
||||||
|
live = self._pool.get(key)
|
||||||
|
return live is not None and bool(live.pending)
|
||||||
|
|
||||||
|
def pending_note(self, key: str) -> str | None:
|
||||||
|
live = self._pool.get(key)
|
||||||
|
return live.note if live is not None else None
|
||||||
|
|
||||||
|
async def attach(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
key: str,
|
||||||
|
capture: TurnCapture | None = None,
|
||||||
|
observer: Callable[[Any], None] | None = None,
|
||||||
|
turn_id: str | None = None,
|
||||||
|
) -> AsyncIterator[MessageStreamEvent]:
|
||||||
|
"""Consume the turn a live session started by itself as a normal turn.
|
||||||
|
|
||||||
|
The CLI runs one when a background task (a subagent) reports back:
|
||||||
|
no prompt goes out, the events are what the model said in reply to
|
||||||
|
the report, the turn ends on that result.
|
||||||
|
"""
|
||||||
|
live = self._pool.get(key)
|
||||||
|
if live is None or not live.pending:
|
||||||
|
return
|
||||||
|
message_id = f"msg_{uuid.uuid4().hex}"
|
||||||
|
yield build_message_start(message_id=message_id, model=self._agent.model)
|
||||||
|
turn = _Turn()
|
||||||
|
async with live.lock:
|
||||||
|
live.running_turn = turn_id or message_id
|
||||||
|
live.last_used = time.monotonic()
|
||||||
|
async for event in self._run_turn(live, None, turn, observer, capture):
|
||||||
|
yield event
|
||||||
|
live.turns += 1
|
||||||
|
live.last_used = time.monotonic()
|
||||||
|
live.running_turn = None
|
||||||
|
usage, _ = await self._after_turn(
|
||||||
|
live, turn, conversation_id=key, history=[], capture=capture
|
||||||
|
)
|
||||||
|
yield build_message_delta(
|
||||||
|
stop_reason=turn.stop_reason, usage=_wire_usage(usage)
|
||||||
|
)
|
||||||
|
yield build_message_stop()
|
||||||
|
|
||||||
|
async def repair_session(
|
||||||
|
self, session_id: str, *, text: str = "interrupted"
|
||||||
|
) -> int:
|
||||||
"""Close ``tool_use`` blocks a crash left without a result; count added."""
|
"""Close ``tool_use`` blocks a crash left without a result; count added."""
|
||||||
key = self._store_key(session_id)
|
key = self._store_key(session_id)
|
||||||
entries = await self._store.load(cast("Any", key))
|
entries = await self._store.load(cast("Any", key))
|
||||||
if not entries:
|
if not entries:
|
||||||
return 0
|
return 0
|
||||||
fixes = close_open_tool_uses(cast("list[Mapping[str, Any]]", entries))
|
fixes = close_open_tool_uses(
|
||||||
|
cast("list[Mapping[str, Any]]", entries), text=text
|
||||||
|
)
|
||||||
if fixes:
|
if fixes:
|
||||||
await self._store.append(cast("Any", key), cast("Any", fixes))
|
await self._store.append(cast("Any", key), cast("Any", fixes))
|
||||||
_log.warning(
|
_log.warning(
|
||||||
@@ -301,7 +348,7 @@ class ClaudeSdkBackend:
|
|||||||
*,
|
*,
|
||||||
agent: BaseAgent,
|
agent: BaseAgent,
|
||||||
messages: Iterable[MessageParam],
|
messages: Iterable[MessageParam],
|
||||||
system: str | None = None, # noqa: ARG002 - the agent owns its prompt
|
system: str | None = None, # noqa: ARG002
|
||||||
conversation_id: str | None = None,
|
conversation_id: str | None = None,
|
||||||
session_id: str | None = None,
|
session_id: str | None = None,
|
||||||
reseed: bool = False,
|
reseed: bool = False,
|
||||||
@@ -311,8 +358,21 @@ class ClaudeSdkBackend:
|
|||||||
tools: bool = True,
|
tools: bool = True,
|
||||||
observer: Callable[[Any], None] | None = None,
|
observer: Callable[[Any], None] | None = None,
|
||||||
turn_id: str | None = None,
|
turn_id: str | None = None,
|
||||||
**options: Any, # noqa: ARG002 - per-request knobs are not supported
|
**options: Any, # noqa: ARG002
|
||||||
) -> AsyncIterator[MessageStreamEvent]:
|
) -> AsyncIterator[MessageStreamEvent]:
|
||||||
|
"""Run one turn against ``agent``, yielding Anthropic stream events.
|
||||||
|
|
||||||
|
Sessions are keyed by ``conversation_id`` when given, else a
|
||||||
|
text-only fingerprint of ``messages[:-1]``. Without a live session,
|
||||||
|
resumes ``session_id`` (closing any ``tool_use`` a crash left open),
|
||||||
|
or seeds ``messages[:-1]`` via ``backends/transcript.py`` and resumes
|
||||||
|
that. ``kind`` picks the prompt assembly and pool TTL; ``pinned``
|
||||||
|
sessions are never evicted; ``tools=False`` mounts no MCP servers
|
||||||
|
(forks and jobs); ``observer`` sees every raw SDK message.
|
||||||
|
|
||||||
|
A resumed session that dies before any event reached the caller is
|
||||||
|
reseeded from history and retried once.
|
||||||
|
"""
|
||||||
if agent.name != self._agent.name:
|
if agent.name != self._agent.name:
|
||||||
msg = f"backend bound to {self._agent.name!r}, got {agent.name!r}"
|
msg = f"backend bound to {self._agent.name!r}, got {agent.name!r}"
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
@@ -320,7 +380,7 @@ class ClaudeSdkBackend:
|
|||||||
if not history or history[-1].get("role") != "user":
|
if not history or history[-1].get("role") != "user":
|
||||||
msg = "the last message must be a user turn"
|
msg = "the last message must be a user turn"
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
prompt = _prompt_text(history[-1].get("content"))
|
prompt = _prompt(history[-1].get("content"))
|
||||||
prior = history[:-1]
|
prior = history[:-1]
|
||||||
key = conversation_id or fingerprint(prior)
|
key = conversation_id or fingerprint(prior)
|
||||||
spec = _SessionSpec(kind=kind, pinned=pinned, tools=tools)
|
spec = _SessionSpec(kind=kind, pinned=pinned, tools=tools)
|
||||||
@@ -334,8 +394,6 @@ class ClaudeSdkBackend:
|
|||||||
live.running_turn = turn_id or message_id
|
live.running_turn = turn_id or message_id
|
||||||
live.last_used = time.monotonic()
|
live.last_used = time.monotonic()
|
||||||
try:
|
try:
|
||||||
# Events go out as the CLI produces them: the frontends
|
|
||||||
# stream text and thinking live, the turn is not buffered.
|
|
||||||
async for event in self._run_turn(
|
async for event in self._run_turn(
|
||||||
live, prompt, turn, observer, capture
|
live, prompt, turn, observer, capture
|
||||||
):
|
):
|
||||||
@@ -343,8 +401,6 @@ class ClaudeSdkBackend:
|
|||||||
except Exception:
|
except Exception:
|
||||||
live.running_turn = None
|
live.running_turn = None
|
||||||
await self._pool.close(key)
|
await self._pool.close(key)
|
||||||
# A dead resume can be reseeded from history, but only
|
|
||||||
# while nothing of this turn has reached the caller yet.
|
|
||||||
if not (live.resumed and live.turns == 0) or turn.events:
|
if not (live.resumed and live.turns == 0) or turn.events:
|
||||||
raise
|
raise
|
||||||
_log.exception(
|
_log.exception(
|
||||||
@@ -415,70 +471,71 @@ class ClaudeSdkBackend:
|
|||||||
async def _run_turn(
|
async def _run_turn(
|
||||||
self,
|
self,
|
||||||
live: Session,
|
live: Session,
|
||||||
prompt: str,
|
prompt: str | list[dict[str, Any]] | None,
|
||||||
turn: _Turn,
|
turn: _Turn,
|
||||||
observer: Callable[[Any], None] | None = None,
|
observer: Callable[[Any], None] | None = None,
|
||||||
capture: TurnCapture | None = None,
|
capture: TurnCapture | None = None,
|
||||||
) -> AsyncIterator[MessageStreamEvent]:
|
) -> AsyncIterator[MessageStreamEvent]:
|
||||||
"""Run one prompt, yielding wire events as they arrive.
|
"""Run one prompt, or attach to a self-started one, yielding events.
|
||||||
|
|
||||||
``turn`` is filled in place (result, synthesized history, count of
|
``turn`` is filled in place (result, synthesized history, count of
|
||||||
events already yielded) so the caller can finish bookkeeping - and
|
events already yielded) so the caller can finish bookkeeping - and
|
||||||
decide whether a retry is still possible - after a failure. The
|
decide whether a retry is still possible - after a failure. The
|
||||||
session id lands in ``capture`` with the first frame, so a turn cut
|
session id lands in ``capture`` with the first frame, so a turn cut
|
||||||
by a restart still leaves a resumable session behind.
|
by a restart still leaves a resumable session behind.
|
||||||
|
|
||||||
|
A result whose ``origin`` is not the human's ends a turn the CLI ran
|
||||||
|
on its own; inside a prompted turn it is passed over, the events it
|
||||||
|
came with stay in this turn rather than getting lost.
|
||||||
"""
|
"""
|
||||||
streaming = self._agent.options.include_partial_messages
|
|
||||||
raw: list[Any] = []
|
raw: list[Any] = []
|
||||||
next_index = 0
|
cursor = _Cursor()
|
||||||
offset = 0
|
live.observer = observer
|
||||||
await live.client.query(prompt)
|
live.attached = True
|
||||||
async for message in live.client.receive_response():
|
for queued in live.pending:
|
||||||
if observer is not None:
|
live.inbox.put_nowait(queued)
|
||||||
observer(message)
|
live.pending.clear()
|
||||||
session_id = getattr(message, "session_id", None)
|
try:
|
||||||
if isinstance(session_id, str) and session_id and live.session_id is None:
|
if prompt is not None:
|
||||||
live.session_id = session_id
|
await live.client.query(
|
||||||
if capture is not None:
|
prompt if isinstance(prompt, str) else _stream(prompt)
|
||||||
capture.session_id = session_id
|
|
||||||
if isinstance(message, MirrorErrorMessage):
|
|
||||||
live.dirty = True
|
|
||||||
_log.error(
|
|
||||||
"session %s: mirror error, marked dirty: %s",
|
|
||||||
live.session_id,
|
|
||||||
message.error,
|
|
||||||
)
|
)
|
||||||
continue
|
while True:
|
||||||
if getattr(message, "parent_tool_use_id", None) is not None:
|
message = await live.inbox.get()
|
||||||
continue
|
if message is _END:
|
||||||
if isinstance(message, StreamEvent):
|
if prompt is not None:
|
||||||
event = message.event
|
msg = f"claude exited before answering ({live.session_id})"
|
||||||
if event.get("type") == "message_start":
|
raise RuntimeError(msg)
|
||||||
offset = next_index
|
break
|
||||||
turn.context_tokens = _context_of(event.get("message"))
|
if isinstance(message, BaseException):
|
||||||
|
raise message
|
||||||
|
if observer is not None:
|
||||||
|
observer(message)
|
||||||
|
if not self._note(live, message, capture):
|
||||||
continue
|
continue
|
||||||
index = event.get("index")
|
if isinstance(message, ResultMessage):
|
||||||
if isinstance(index, int):
|
if prompt is not None and _self_started(message):
|
||||||
next_index = max(next_index, offset + index + 1)
|
_log.info(
|
||||||
if streaming:
|
"session %s: a self-started turn ended inside a "
|
||||||
for out in _emit_stream_event(event, offset + index):
|
"prompted one, folded in",
|
||||||
turn.events += 1
|
live.session_id,
|
||||||
yield out
|
)
|
||||||
elif isinstance(message, AssistantMessage):
|
continue
|
||||||
raw.append(message)
|
turn.result = message
|
||||||
if not streaming:
|
turn.stop_reason = _STOP_REASONS.get(
|
||||||
for block in message.content:
|
message.stop_reason or "", "end_turn"
|
||||||
for out in _emit_block(block, next_index):
|
)
|
||||||
turn.events += 1
|
break
|
||||||
yield out
|
for out in self._emit(message, turn, cursor, raw):
|
||||||
next_index += 1
|
turn.events += 1
|
||||||
elif isinstance(message, UserMessage):
|
yield out
|
||||||
raw.append(message)
|
if turn.events == 0:
|
||||||
elif isinstance(message, ResultMessage):
|
for out in _emit_whole(raw, cursor):
|
||||||
turn.result = message
|
turn.events += 1
|
||||||
turn.stop_reason = _STOP_REASONS.get(
|
yield out
|
||||||
message.stop_reason or "", "end_turn"
|
finally:
|
||||||
)
|
live.attached = False
|
||||||
|
live.woken = False
|
||||||
turn.synthesized = synthesize_turn_messages(raw)
|
turn.synthesized = synthesize_turn_messages(raw)
|
||||||
_log.info(
|
_log.info(
|
||||||
"turn: agent=%s session=%s events=%d synthesized=%d stop=%s",
|
"turn: agent=%s session=%s events=%d synthesized=%d stop=%s",
|
||||||
@@ -489,6 +546,88 @@ class ClaudeSdkBackend:
|
|||||||
turn.stop_reason,
|
turn.stop_reason,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _note(live: Session, message: Any, capture: TurnCapture | None) -> bool:
|
||||||
|
"""Session bookkeeping for one message; False if not the main thread's."""
|
||||||
|
session_id = getattr(message, "session_id", None)
|
||||||
|
if isinstance(session_id, str) and session_id and live.session_id is None:
|
||||||
|
live.session_id = session_id
|
||||||
|
if capture is not None:
|
||||||
|
capture.session_id = session_id
|
||||||
|
if isinstance(message, MirrorErrorMessage):
|
||||||
|
live.dirty = True
|
||||||
|
_log.error(
|
||||||
|
"session %s: mirror error, marked dirty: %s",
|
||||||
|
live.session_id,
|
||||||
|
message.error,
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
if isinstance(message, TaskNotificationMessage):
|
||||||
|
live.note = _note_of(message)
|
||||||
|
return False
|
||||||
|
return getattr(message, "parent_tool_use_id", None) is None
|
||||||
|
|
||||||
|
def _emit(
|
||||||
|
self, message: Any, turn: _Turn, cursor: _Cursor, raw: list[Any]
|
||||||
|
) -> Iterable[MessageStreamEvent]:
|
||||||
|
streaming = self._agent.options.include_partial_messages
|
||||||
|
if isinstance(message, StreamEvent):
|
||||||
|
event = message.event
|
||||||
|
if event.get("type") == "message_start":
|
||||||
|
cursor.offset = cursor.next_index
|
||||||
|
turn.context_tokens = _context_of(event.get("message"))
|
||||||
|
return
|
||||||
|
index = event.get("index")
|
||||||
|
if isinstance(index, int):
|
||||||
|
cursor.next_index = max(cursor.next_index, cursor.offset + index + 1)
|
||||||
|
if streaming:
|
||||||
|
yield from _emit_stream_event(event, cursor.offset + index)
|
||||||
|
elif isinstance(message, AssistantMessage):
|
||||||
|
raw.append(message)
|
||||||
|
if not streaming:
|
||||||
|
for block in message.content:
|
||||||
|
yield from _emit_block(block, cursor.next_index)
|
||||||
|
cursor.next_index += 1
|
||||||
|
elif isinstance(message, UserMessage):
|
||||||
|
raw.append(message)
|
||||||
|
|
||||||
|
async def _read(self, live: Session) -> None:
|
||||||
|
"""Read the client for the session's whole life.
|
||||||
|
|
||||||
|
During a turn every message goes to the turn's inbox. Between turns
|
||||||
|
a subagent still running is shown to the last observer, and the
|
||||||
|
first frame of the main thread means the CLI started a turn by
|
||||||
|
itself: it is buffered in ``pending`` and the waker is told once.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
async for message in live.client.receive_messages():
|
||||||
|
if live.attached:
|
||||||
|
live.inbox.put_nowait(message)
|
||||||
|
continue
|
||||||
|
if isinstance(message, TaskNotificationMessage):
|
||||||
|
live.note = _note_of(message)
|
||||||
|
if _main_thread(message):
|
||||||
|
live.pending.append(message)
|
||||||
|
if not live.woken:
|
||||||
|
live.woken = True
|
||||||
|
self._wake(live)
|
||||||
|
elif live.observer is not None:
|
||||||
|
live.observer(message)
|
||||||
|
live.inbox.put_nowait(_END)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except Exception as exc: # noqa: BLE001
|
||||||
|
_log.error("session %s: reader died: %s", live.session_id, exc)
|
||||||
|
live.dirty = True
|
||||||
|
live.inbox.put_nowait(exc)
|
||||||
|
|
||||||
|
def _wake(self, live: Session) -> None:
|
||||||
|
if self.waker is None:
|
||||||
|
return
|
||||||
|
task = asyncio.create_task(self.waker(live.key, live.note or ""))
|
||||||
|
self._tasks.add(task)
|
||||||
|
task.add_done_callback(self._tasks.discard)
|
||||||
|
|
||||||
async def _acquire(
|
async def _acquire(
|
||||||
self,
|
self,
|
||||||
key: str,
|
key: str,
|
||||||
@@ -541,7 +680,8 @@ class ClaudeSdkBackend:
|
|||||||
async def _spawn(
|
async def _spawn(
|
||||||
self, resume: str | None, *, key: str, spec: _SessionSpec
|
self, resume: str | None, *, key: str, spec: _SessionSpec
|
||||||
) -> Session:
|
) -> Session:
|
||||||
options = self._build_options(resume, key=key, spec=spec)
|
stderr: deque[str] = deque(maxlen=STDERR_LINES)
|
||||||
|
options = self._build_options(resume, key=key, spec=spec, stderr=stderr)
|
||||||
client = self._factory(options)
|
client = self._factory(options)
|
||||||
await client.connect()
|
await client.connect()
|
||||||
_log.info(
|
_log.info(
|
||||||
@@ -552,7 +692,7 @@ class ClaudeSdkBackend:
|
|||||||
spec.tools,
|
spec.tools,
|
||||||
self._runner.user,
|
self._runner.user,
|
||||||
)
|
)
|
||||||
return Session(
|
live = Session(
|
||||||
key=key,
|
key=key,
|
||||||
agent=self._agent.name,
|
agent=self._agent.name,
|
||||||
kind=spec.kind,
|
kind=spec.kind,
|
||||||
@@ -560,13 +700,21 @@ class ClaudeSdkBackend:
|
|||||||
session_id=resume,
|
session_id=resume,
|
||||||
resumed=resume is not None,
|
resumed=resume is not None,
|
||||||
pinned=spec.pinned,
|
pinned=spec.pinned,
|
||||||
|
stderr=stderr,
|
||||||
)
|
)
|
||||||
|
live.reader = asyncio.create_task(self._read(live))
|
||||||
|
return live
|
||||||
|
|
||||||
def _default_factory(self, options: ClaudeAgentOptions) -> SessionClient:
|
def _default_factory(self, options: ClaudeAgentOptions) -> SessionClient:
|
||||||
return _RunnerClient(options, uid=self._uid)
|
return _RunnerClient(options, uid=self._uid)
|
||||||
|
|
||||||
def _build_options(
|
def _build_options(
|
||||||
self, resume: str | None, *, key: str, spec: _SessionSpec
|
self,
|
||||||
|
resume: str | None,
|
||||||
|
*,
|
||||||
|
key: str,
|
||||||
|
spec: _SessionSpec,
|
||||||
|
stderr: deque[str] | None = None,
|
||||||
) -> ClaudeAgentOptions:
|
) -> ClaudeAgentOptions:
|
||||||
agent = self._agent
|
agent = self._agent
|
||||||
opt = agent.options
|
opt = agent.options
|
||||||
@@ -574,12 +722,20 @@ class ClaudeSdkBackend:
|
|||||||
if self._runner.home is not None:
|
if self._runner.home is not None:
|
||||||
env["HOME"] = str(self._runner.home)
|
env["HOME"] = str(self._runner.home)
|
||||||
env.setdefault("CLAUDE_CONFIG_DIR", str(self._runner.home / ".claude"))
|
env.setdefault("CLAUDE_CONFIG_DIR", str(self._runner.home / ".claude"))
|
||||||
plugins = self._plugins()
|
kind = as_kind(spec.kind)
|
||||||
sources = agent.prompt_for(as_kind(spec.kind))
|
plugins = self._plugins(kind)
|
||||||
|
sources = agent.prompt_for(kind)
|
||||||
system_prompt = (
|
system_prompt = (
|
||||||
prompt_assembly.assemble(sources) if 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 {}
|
servers: dict[str, Any] = (
|
||||||
|
{
|
||||||
|
name: {**server, "headers": {ORIGIN_HEADER: key}}
|
||||||
|
for name, server in self._servers.items()
|
||||||
|
}
|
||||||
|
if spec.tools
|
||||||
|
else {}
|
||||||
|
)
|
||||||
gateway = (
|
gateway = (
|
||||||
self._tool_server(key, spec.kind)
|
self._tool_server(key, spec.kind)
|
||||||
if spec.tools and self._tool_server is not None and agent.gateway_tools
|
if spec.tools and self._tool_server is not None and agent.gateway_tools
|
||||||
@@ -590,6 +746,7 @@ class ClaudeSdkBackend:
|
|||||||
return ClaudeAgentOptions(
|
return ClaudeAgentOptions(
|
||||||
model=agent.model or None,
|
model=agent.model or None,
|
||||||
effort=cast("Any", opt.effort),
|
effort=cast("Any", opt.effort),
|
||||||
|
thinking=cast("Any", dict(opt.thinking) if opt.thinking else None),
|
||||||
system_prompt=system_prompt,
|
system_prompt=system_prompt,
|
||||||
setting_sources=[],
|
setting_sources=[],
|
||||||
strict_mcp_config=True,
|
strict_mcp_config=True,
|
||||||
@@ -610,16 +767,26 @@ class ClaudeSdkBackend:
|
|||||||
plugins=cast("Any", plugins),
|
plugins=cast("Any", plugins),
|
||||||
skills="all" if plugins else None,
|
skills="all" if plugins else None,
|
||||||
max_turns=opt.max_turns,
|
max_turns=opt.max_turns,
|
||||||
stderr=lambda line: _log.warning(
|
max_buffer_size=opt.max_buffer_size,
|
||||||
"claude[%s]: %s", agent.name, line.rstrip()
|
stderr=lambda line: self._stderr(line.rstrip(), stderr),
|
||||||
),
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _stderr(self, line: str, tail: deque[str] | None) -> None:
|
||||||
|
_log.warning("claude[%s]: %s", self._agent.name, line)
|
||||||
|
if tail is not None:
|
||||||
|
tail.append(redact(line))
|
||||||
|
|
||||||
def _can_use_tool(
|
def _can_use_tool(
|
||||||
self, key: str
|
self, key: str
|
||||||
) -> Callable[
|
) -> Callable[
|
||||||
[str, dict[str, Any], ToolPermissionContext], Awaitable[PermissionResult]
|
[str, dict[str, Any], ToolPermissionContext], Awaitable[PermissionResult]
|
||||||
]:
|
]:
|
||||||
|
"""Build the ``can_use_tool`` callback for one session.
|
||||||
|
|
||||||
|
In bypass mode the SDK only ever calls this back for
|
||||||
|
``AskUserQuestion``; the warning about the rest is suppressed at
|
||||||
|
import time.
|
||||||
|
"""
|
||||||
asker = self._asker
|
asker = self._asker
|
||||||
|
|
||||||
async def can_use_tool(
|
async def can_use_tool(
|
||||||
@@ -642,7 +809,7 @@ class ClaudeSdkBackend:
|
|||||||
def _hooks(
|
def _hooks(
|
||||||
self, key: str, spec: _SessionSpec
|
self, key: str, spec: _SessionSpec
|
||||||
) -> dict[HookEvent, list[HookMatcher]] | None:
|
) -> dict[HookEvent, list[HookMatcher]] | None:
|
||||||
"""§3.7: one in-process ``PreToolUse`` hook - policy rules, then audit."""
|
"""One in-process ``PreToolUse`` hook: policy rules, then audit."""
|
||||||
agent = self._agent
|
agent = self._agent
|
||||||
if not agent.policy and self._audit_sink is None:
|
if not agent.policy and self._audit_sink is None:
|
||||||
return None
|
return None
|
||||||
@@ -683,10 +850,10 @@ class ClaudeSdkBackend:
|
|||||||
|
|
||||||
return {"PreToolUse": [HookMatcher(hooks=[cast("Any", pre_tool_use)])]}
|
return {"PreToolUse": [HookMatcher(hooks=[cast("Any", pre_tool_use)])]}
|
||||||
|
|
||||||
def _plugins(self) -> list[dict[str, str]]:
|
def _plugins(self, kind: Kind) -> list[dict[str, str]]:
|
||||||
plugins: list[dict[str, str]] = []
|
plugins: list[dict[str, str]] = []
|
||||||
root = self._work_dir / "plugins" / self._agent.name
|
root = self._work_dir / "plugins" / self._agent.name / kind
|
||||||
for raw in sorted(self._agent.skill_sets, key=str):
|
for raw in sorted(self._agent.skills_for(kind), key=str):
|
||||||
source = Path(str(raw))
|
source = Path(str(raw))
|
||||||
name = source.name
|
name = source.name
|
||||||
target = root / name
|
target = root / name
|
||||||
@@ -702,6 +869,12 @@ class ClaudeSdkBackend:
|
|||||||
return plugins
|
return plugins
|
||||||
|
|
||||||
def _exec_wrapper(self, *, extra_keep: tuple[str, ...]) -> Path:
|
def _exec_wrapper(self, *, extra_keep: tuple[str, ...]) -> Path:
|
||||||
|
"""Path to a generated wrapper that execs claude with a whitelisted env.
|
||||||
|
|
||||||
|
Drops every inherited variable outside the whitelist and, when
|
||||||
|
``RunnerConfig.user`` is set, switches to that uid before exec - done
|
||||||
|
here rather than via ``subprocess(user=...)``, which uvloop rejects.
|
||||||
|
"""
|
||||||
if self._wrapper is not None:
|
if self._wrapper is not None:
|
||||||
return self._wrapper
|
return self._wrapper
|
||||||
keep = sorted({*ENV_KEEP, *self._agent.options.env_keep, *extra_keep})
|
keep = sorted({*ENV_KEEP, *self._agent.options.env_keep, *extra_keep})
|
||||||
@@ -737,6 +910,15 @@ class _SessionSpec:
|
|||||||
tools: bool
|
tools: bool
|
||||||
|
|
||||||
|
|
||||||
|
_END = object()
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _Cursor:
|
||||||
|
next_index: int = 0
|
||||||
|
offset: int = 0
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class _Turn:
|
class _Turn:
|
||||||
events: int = 0
|
events: int = 0
|
||||||
@@ -776,7 +958,17 @@ async def _disconnect(live: Session) -> None:
|
|||||||
_log.exception("disconnect failed for session %s", live.session_id)
|
_log.exception("disconnect failed for session %s", live.session_id)
|
||||||
|
|
||||||
|
|
||||||
|
CLI_ENV = "BEAVER_CLAUDE_CLI"
|
||||||
|
"""Path to a ``claude`` CLI that wins over the one bundled in the wheel."""
|
||||||
|
|
||||||
|
|
||||||
def _claude_binary() -> str:
|
def _claude_binary() -> str:
|
||||||
|
override = os.environ.get(CLI_ENV)
|
||||||
|
if override:
|
||||||
|
if not Path(override).is_file():
|
||||||
|
msg = f"{CLI_ENV} points at {override!r}, which is not a file"
|
||||||
|
raise FileNotFoundError(msg)
|
||||||
|
return override
|
||||||
bundled = Path(claude_agent_sdk.__file__).parent / "_bundled" / "claude"
|
bundled = Path(claude_agent_sdk.__file__).parent / "_bundled" / "claude"
|
||||||
if bundled.is_file():
|
if bundled.is_file():
|
||||||
return str(bundled)
|
return str(bundled)
|
||||||
@@ -843,14 +1035,54 @@ def _mcp_disallowed(
|
|||||||
return out
|
return out
|
||||||
|
|
||||||
|
|
||||||
def _prompt_text(content: Any) -> str:
|
def _prompt(content: Any) -> str | list[dict[str, Any]]:
|
||||||
|
"""Plain text when the turn is text only; the block list when images ride along."""
|
||||||
text = text_of(content)
|
text = text_of(content)
|
||||||
if not text:
|
if not text:
|
||||||
msg = "user message has no text content"
|
msg = "user message has no text content"
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
|
if isinstance(content, list) and any(
|
||||||
|
isinstance(b, dict) and b.get("type") != "text" for b in content
|
||||||
|
):
|
||||||
|
return [dict(b) for b in content if isinstance(b, dict)]
|
||||||
return text
|
return text
|
||||||
|
|
||||||
|
|
||||||
|
async def _stream(content: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]:
|
||||||
|
yield {
|
||||||
|
"type": "user",
|
||||||
|
"message": {"role": "user", "content": content},
|
||||||
|
"parent_tool_use_id": None,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _emit_whole(raw: Iterable[Any], cursor: _Cursor) -> Iterable[MessageStreamEvent]:
|
||||||
|
"""Blocks of a turn that streamed no deltas, e.g. one reported whole."""
|
||||||
|
for message in raw:
|
||||||
|
if isinstance(message, AssistantMessage):
|
||||||
|
for block in message.content:
|
||||||
|
yield from _emit_block(block, cursor.next_index)
|
||||||
|
cursor.next_index += 1
|
||||||
|
|
||||||
|
|
||||||
|
def _self_started(result: ResultMessage) -> bool:
|
||||||
|
origin = result.origin
|
||||||
|
return origin is not None and origin.get("kind") != "human"
|
||||||
|
|
||||||
|
|
||||||
|
def _main_thread(message: Any) -> bool:
|
||||||
|
if getattr(message, "parent_tool_use_id", None) is not None:
|
||||||
|
return False
|
||||||
|
return isinstance(
|
||||||
|
message, (StreamEvent, AssistantMessage, UserMessage, ResultMessage)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _note_of(message: TaskNotificationMessage) -> str:
|
||||||
|
summary = message.summary or message.data.get("description") or message.task_id
|
||||||
|
return f"{summary} ({message.status})"
|
||||||
|
|
||||||
|
|
||||||
def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]:
|
def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]:
|
||||||
out: list[dict[str, Any]] = []
|
out: list[dict[str, Any]] = []
|
||||||
for message in raw:
|
for message in raw:
|
||||||
|
|||||||
@@ -0,0 +1,118 @@
|
|||||||
|
"""What runs under the gateway: the process tree with memory and CPU, and the kill."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import contextlib
|
||||||
|
import os
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import psutil
|
||||||
|
|
||||||
|
from beaver_gateway.security.redact import redact
|
||||||
|
|
||||||
|
__all__ = ["Processes", "subtree"]
|
||||||
|
|
||||||
|
_VALUE_HIDDEN = frozenset(
|
||||||
|
{
|
||||||
|
"--system-prompt",
|
||||||
|
"--append-system-prompt",
|
||||||
|
"--mcp-config",
|
||||||
|
"--settings",
|
||||||
|
"--agents",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
_ARG_CLIP = 80
|
||||||
|
_COMMAND_CLIP = 240
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(slots=True)
|
||||||
|
class Processes:
|
||||||
|
root: int = field(default_factory=os.getpid)
|
||||||
|
_cpu_seen: dict[int, tuple[float, float]] = field(default_factory=dict)
|
||||||
|
|
||||||
|
def tree_parents_first(self) -> list[dict[str, Any]]:
|
||||||
|
try:
|
||||||
|
found = psutil.Process(self.root).children(recursive=True)
|
||||||
|
except psutil.Error:
|
||||||
|
return []
|
||||||
|
now = time.monotonic()
|
||||||
|
rows: dict[int, dict[str, Any]] = {}
|
||||||
|
for process in found:
|
||||||
|
with contextlib.suppress(psutil.Error), process.oneshot():
|
||||||
|
spent = process.cpu_times()
|
||||||
|
rows[process.pid] = {
|
||||||
|
"pid": process.pid,
|
||||||
|
"parent": process.ppid(),
|
||||||
|
"name": process.name(),
|
||||||
|
"command": _command(process.cmdline()),
|
||||||
|
"rss": process.memory_info().rss,
|
||||||
|
"cpu": self._rate(process.pid, spent.user + spent.system, now),
|
||||||
|
"age_seconds": round(time.time() - process.create_time(), 1),
|
||||||
|
}
|
||||||
|
self._cpu_seen = {p: seen for p, seen in self._cpu_seen.items() if p in rows}
|
||||||
|
children: dict[int, list[int]] = {}
|
||||||
|
for pid, row in rows.items():
|
||||||
|
parent = row["parent"] if row["parent"] in rows else self.root
|
||||||
|
children.setdefault(parent, []).append(pid)
|
||||||
|
ordered: list[dict[str, Any]] = []
|
||||||
|
|
||||||
|
def walk(parent: int, depth: int) -> None:
|
||||||
|
for pid in sorted(children.get(parent, ())):
|
||||||
|
ordered.append({**rows[pid], "depth": depth})
|
||||||
|
walk(pid, depth + 1)
|
||||||
|
|
||||||
|
walk(self.root, 0)
|
||||||
|
return ordered
|
||||||
|
|
||||||
|
def kill_own_descendant(self, pid: int) -> bool:
|
||||||
|
try:
|
||||||
|
owned = psutil.Process(self.root).children(recursive=True)
|
||||||
|
target = next((p for p in owned if p.pid == pid), None)
|
||||||
|
if target is None:
|
||||||
|
return False
|
||||||
|
doomed = [*target.children(recursive=True), target]
|
||||||
|
except psutil.Error:
|
||||||
|
return False
|
||||||
|
for process in doomed:
|
||||||
|
with contextlib.suppress(psutil.Error):
|
||||||
|
process.kill()
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _rate(self, pid: int, spent: float, now: float) -> float | None:
|
||||||
|
before = self._cpu_seen.get(pid)
|
||||||
|
self._cpu_seen[pid] = (spent, now)
|
||||||
|
if before is None or now <= before[1]:
|
||||||
|
return None
|
||||||
|
return round(100 * max(spent - before[0], 0.0) / (now - before[1]), 1)
|
||||||
|
|
||||||
|
|
||||||
|
def subtree(rows: list[dict[str, Any]], pid: int | None) -> list[dict[str, Any]]:
|
||||||
|
start = next((i for i, row in enumerate(rows) if row["pid"] == pid), None)
|
||||||
|
if start is None:
|
||||||
|
return []
|
||||||
|
base = rows[start]["depth"]
|
||||||
|
out = [{**rows[start], "depth": 0}]
|
||||||
|
for row in rows[start + 1 :]:
|
||||||
|
if row["depth"] <= base:
|
||||||
|
break
|
||||||
|
out.append({**row, "depth": row["depth"] - base})
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def _command(argv: list[str]) -> str:
|
||||||
|
if not argv:
|
||||||
|
return ""
|
||||||
|
parts = [Path(argv[0]).name]
|
||||||
|
hide = False
|
||||||
|
for arg in argv[1:]:
|
||||||
|
if hide:
|
||||||
|
parts.append("…")
|
||||||
|
hide = False
|
||||||
|
continue
|
||||||
|
hide = arg in _VALUE_HIDDEN
|
||||||
|
parts.append(arg if len(arg) <= _ARG_CLIP else f"{arg[:_ARG_CLIP]}…")
|
||||||
|
command = redact(" ".join(parts))
|
||||||
|
return command if len(command) <= _COMMAND_CLIP else f"{command[:_COMMAND_CLIP]}…"
|
||||||
@@ -1,32 +1,8 @@
|
|||||||
"""Raycast backend adapter.
|
"""Raycast backend adapter.
|
||||||
|
|
||||||
Translates between Anthropic's ``/v1/messages`` wire vocabulary (incoming
|
Translates between Anthropic's ``/v1/messages`` wire vocabulary and the
|
||||||
``MessageParam`` history, outgoing ``MessageStreamEvent`` SSE) and the
|
``raycast-api`` SDK, splicing MCP tools into Raycast's ``tool_calls`` and
|
||||||
``raycast-api`` SDK (``Message`` history, ``ChatStreamChunk`` SSE).
|
routing them back in-process where the agent declares ``expose_mcps``.
|
||||||
|
|
||||||
Two halves live here:
|
|
||||||
|
|
||||||
* :func:`_to_raycast_messages` — pure conversion of an Anthropic message
|
|
||||||
list into ``list[raycast_api.Message]``. ``tool_result`` blocks carry no
|
|
||||||
tool name in Anthropic; we recover it by remembering each ``tool_use``
|
|
||||||
id we saw upstream.
|
|
||||||
* :meth:`RaycastBackend.complete` — opens a ``client.chat.stream`` and
|
|
||||||
walks chunks through a tiny block-state machine. The state machine
|
|
||||||
exists only because Raycast streams ``tool_calls`` in three phases
|
|
||||||
(open with id+name, deltas with empty id, final summary with the full
|
|
||||||
``arguments``) — Anthropic wants one ``content_block_start`` → deltas
|
|
||||||
→ ``content_block_stop`` per block, so we de-duplicate the final
|
|
||||||
summary against the per-delta increments already emitted.
|
|
||||||
|
|
||||||
MCP wiring: Raycast has no native MCP concept, so when an agent declares
|
|
||||||
``expose_mcps`` we splice each MCP's tools into the wire request as
|
|
||||||
``Tool.local(name=f"{mcp}__{tool}", ...)`` and run a gateway-internal
|
|
||||||
loop. Every time the model emits a tool_call for one of those local
|
|
||||||
tools we route it back to the underlying MCP in-process, append the
|
|
||||||
result as a ``tool`` message, and re-issue the stream — all inside one
|
|
||||||
Anthropic envelope (one ``message_start`` … one ``message_stop``). The
|
|
||||||
tool_use blocks DO surface to the caller (mirrors what
|
|
||||||
``ClaudeSdkBackend`` does), but tool_results stay internal.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -42,7 +18,7 @@ from raycast_api import Message as RaycastMessage
|
|||||||
from raycast_api import RemoteTool, Tool, ToolCall
|
from raycast_api import RemoteTool, Tool, ToolCall
|
||||||
|
|
||||||
from beaver_gateway.agents.raycast import RaycastAgent
|
from beaver_gateway.agents.raycast import RaycastAgent
|
||||||
from beaver_gateway.core.events import (
|
from beaver_gateway.events.stream import (
|
||||||
StopReason,
|
StopReason,
|
||||||
build_content_block_stop,
|
build_content_block_stop,
|
||||||
build_input_json_delta,
|
build_input_json_delta,
|
||||||
@@ -65,7 +41,7 @@ if TYPE_CHECKING:
|
|||||||
from raycast_api import ChatStreamChunk, Client
|
from raycast_api import ChatStreamChunk, Client
|
||||||
|
|
||||||
from beaver_gateway.agents.base import BaseAgent
|
from beaver_gateway.agents.base import BaseAgent
|
||||||
from beaver_gateway.core.events import MessageStreamEvent
|
from beaver_gateway.events.stream import MessageStreamEvent
|
||||||
else:
|
else:
|
||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
|
|
||||||
@@ -76,11 +52,9 @@ __all__ = ["RaycastBackend"]
|
|||||||
_log = logging.getLogger("beaver_gateway.backends.raycast")
|
_log = logging.getLogger("beaver_gateway.backends.raycast")
|
||||||
|
|
||||||
|
|
||||||
# Cap on consecutive tool-call turns inside one Anthropic envelope.
|
|
||||||
# Real conversations rarely chain more than a handful; the limit only
|
|
||||||
# fires on a model that loops, and surfaces as a clean ``end_turn``
|
|
||||||
# with an error tool_result rather than a hang.
|
|
||||||
_MAX_TOOL_TURNS = 20
|
_MAX_TOOL_TURNS = 20
|
||||||
|
"""Cap on consecutive tool-call turns inside one Anthropic envelope; a model
|
||||||
|
that loops hits this and gets a clean ``end_turn`` instead of a hang."""
|
||||||
|
|
||||||
|
|
||||||
_RAYCAST_TO_ANTHROPIC_STOP: dict[str, StopReason] = {
|
_RAYCAST_TO_ANTHROPIC_STOP: dict[str, StopReason] = {
|
||||||
@@ -392,7 +366,7 @@ class _BlockState:
|
|||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.index: int = -1
|
self.index: int = -1
|
||||||
self.kind: str | None = None # "text" | "thinking" | "tool_use" | None
|
self.kind: str | None = None
|
||||||
self.tool_id_to_block: dict[str, int] = {}
|
self.tool_id_to_block: dict[str, int] = {}
|
||||||
self.tool_idx_to_id: dict[int, str] = {}
|
self.tool_idx_to_id: dict[int, str] = {}
|
||||||
|
|
||||||
@@ -428,11 +402,14 @@ class RaycastBackend:
|
|||||||
self._client = client
|
self._client = client
|
||||||
self._mcp_servers: Mapping[str, FastMCP] = mcp_servers or {}
|
self._mcp_servers: Mapping[str, FastMCP] = mcp_servers or {}
|
||||||
self._mcp_tools: Mapping[str, list[FastMCPTool]] = mcp_tools or {}
|
self._mcp_tools: Mapping[str, list[FastMCPTool]] = mcp_tools or {}
|
||||||
# Cached per-agent catalog — agents are immutable, so a single
|
|
||||||
# render at first use covers the gateway's lifetime.
|
|
||||||
self._agent_catalog: dict[str, _AgentToolCatalog] = {}
|
self._agent_catalog: dict[str, _AgentToolCatalog] = {}
|
||||||
|
|
||||||
def _catalog_for(self, agent: RaycastAgent) -> _AgentToolCatalog:
|
def _catalog_for(self, agent: RaycastAgent) -> _AgentToolCatalog:
|
||||||
|
"""Build (or reuse) the tool catalog for ``agent``.
|
||||||
|
|
||||||
|
Agents are immutable, so a single render at first use covers the
|
||||||
|
gateway's lifetime.
|
||||||
|
"""
|
||||||
cached = self._agent_catalog.get(agent.name)
|
cached = self._agent_catalog.get(agent.name)
|
||||||
if cached is not None:
|
if cached is not None:
|
||||||
return cached
|
return cached
|
||||||
@@ -448,31 +425,26 @@ class RaycastBackend:
|
|||||||
system: str | None = None,
|
system: str | None = None,
|
||||||
**options: Any,
|
**options: Any,
|
||||||
) -> AsyncIterator[MessageStreamEvent]:
|
) -> AsyncIterator[MessageStreamEvent]:
|
||||||
|
"""Run one turn against ``agent``, yielding Anthropic stream events.
|
||||||
|
|
||||||
|
On the wire Raycast uses ``system_instructions`` as a format marker
|
||||||
|
(``"markdown"``/``"plain"``, filled from the source default when we
|
||||||
|
pass ``None``) and ``additional_system_instructions`` for the actual
|
||||||
|
prompt content; ``system`` (or ``agent.system_prompt``) flows into
|
||||||
|
the latter.
|
||||||
|
"""
|
||||||
if not isinstance(agent, RaycastAgent):
|
if not isinstance(agent, RaycastAgent):
|
||||||
msg = f"RaycastBackend requires RaycastAgent, got {type(agent).__name__}"
|
msg = f"RaycastBackend requires RaycastAgent, got {type(agent).__name__}"
|
||||||
raise TypeError(msg)
|
raise TypeError(msg)
|
||||||
|
|
||||||
raycast_messages = _to_raycast_messages(messages)
|
raycast_messages = _to_raycast_messages(messages)
|
||||||
catalog = self._catalog_for(agent)
|
catalog = self._catalog_for(agent)
|
||||||
# Native remote tools + spliced MCP locals. ``None`` keeps the
|
|
||||||
# SDK from sending a ``tools`` field at all when neither is
|
|
||||||
# declared.
|
|
||||||
tools_arg: list[Tool | RemoteTool | str] | None = (
|
tools_arg: list[Tool | RemoteTool | str] | None = (
|
||||||
list(catalog.tools) if catalog.tools else None
|
list(catalog.tools) if catalog.tools else None
|
||||||
)
|
)
|
||||||
|
|
||||||
# On the wire Raycast uses ``system_instructions`` as a format
|
|
||||||
# marker (``"markdown"`` for AI_CHAT, ``"plain"`` otherwise —
|
|
||||||
# filled in by the SDK from the source default when we pass
|
|
||||||
# ``None``) and ``additional_system_instructions`` as the actual
|
|
||||||
# prompt content. So our ``system_prompt`` (or the per-request
|
|
||||||
# Anthropic ``system``, if present) flows into the *additional*
|
|
||||||
# slot. The SDK still prepends ``<user-preferences>`` to whatever
|
|
||||||
# we hand it via ``_build_preamble``.
|
|
||||||
prompt_content = system if system is not None else agent.system_prompt
|
prompt_content = system if system is not None else agent.system_prompt
|
||||||
|
|
||||||
# Per-request options win over agent defaults; agent defaults
|
|
||||||
# win over Raycast SDK defaults. ``None`` means "fall back".
|
|
||||||
async for event in self._stream(
|
async for event in self._stream(
|
||||||
agent=agent,
|
agent=agent,
|
||||||
raycast_messages=raycast_messages,
|
raycast_messages=raycast_messages,
|
||||||
@@ -499,12 +471,16 @@ class RaycastBackend:
|
|||||||
reasoning_effort: str | None,
|
reasoning_effort: str | None,
|
||||||
tool_choice: str | None,
|
tool_choice: str | None,
|
||||||
) -> AsyncIterator[MessageStreamEvent]:
|
) -> AsyncIterator[MessageStreamEvent]:
|
||||||
|
"""Yield one Anthropic envelope spanning one or more ``chat.stream`` calls.
|
||||||
|
|
||||||
|
A tool_call not in ``mcp_routing`` is left to bubble out as a
|
||||||
|
regular ``tool_use`` block instead, for the caller to answer with a
|
||||||
|
``tool_result`` the usual way.
|
||||||
|
"""
|
||||||
message_id = f"msg_{uuid.uuid4().hex}"
|
message_id = f"msg_{uuid.uuid4().hex}"
|
||||||
yield build_message_start(message_id=message_id, model=agent.model)
|
yield build_message_start(message_id=message_id, model=agent.model)
|
||||||
|
|
||||||
state = _BlockState()
|
state = _BlockState()
|
||||||
# ``working_messages`` is the rolling history fed back to
|
|
||||||
# Raycast as we resolve MCP tool_calls turn by turn.
|
|
||||||
working_messages = list(raycast_messages)
|
working_messages = list(raycast_messages)
|
||||||
last_usage: dict[str, int] | None = None
|
last_usage: dict[str, int] | None = None
|
||||||
last_finish: str | None = None
|
last_finish: str | None = None
|
||||||
@@ -517,9 +493,6 @@ class RaycastBackend:
|
|||||||
model=agent.model,
|
model=agent.model,
|
||||||
messages=working_messages,
|
messages=working_messages,
|
||||||
source=agent.source,
|
source=agent.source,
|
||||||
# ``system_instructions=None`` → SDK substitutes the
|
|
||||||
# source default (``"markdown"`` / ``"plain"``). Real
|
|
||||||
# prompt goes into ``additional_system_instructions``.
|
|
||||||
additional_system_instructions=prompt_content,
|
additional_system_instructions=prompt_content,
|
||||||
user_preferences=agent.user_preferences,
|
user_preferences=agent.user_preferences,
|
||||||
tools=tools,
|
tools=tools,
|
||||||
@@ -540,27 +513,16 @@ class RaycastBackend:
|
|||||||
if acc.usage:
|
if acc.usage:
|
||||||
last_usage = acc.usage
|
last_usage = acc.usage
|
||||||
|
|
||||||
# Figure out which (if any) tool_calls land on our MCP
|
|
||||||
# routing table. Anything not in the table is left to bubble
|
|
||||||
# out of the envelope as a regular tool_use block — the
|
|
||||||
# Anthropic caller can then respond with a tool_result the
|
|
||||||
# usual way, and the next ``complete`` invocation will
|
|
||||||
# carry it back in.
|
|
||||||
pending_mcp_calls = [
|
pending_mcp_calls = [
|
||||||
tid for tid in acc.tool_order if acc.tool_names.get(tid) in mcp_routing
|
tid for tid in acc.tool_order if acc.tool_names.get(tid) in mcp_routing
|
||||||
]
|
]
|
||||||
if not pending_mcp_calls:
|
if not pending_mcp_calls:
|
||||||
break
|
break
|
||||||
|
|
||||||
# Close whatever block is still open before we step into
|
|
||||||
# tool execution — the next turn's chunks start a fresh
|
|
||||||
# block sequence.
|
|
||||||
if state.kind is not None:
|
if state.kind is not None:
|
||||||
yield build_content_block_stop(state.index)
|
yield build_content_block_stop(state.index)
|
||||||
state.kind = None
|
state.kind = None
|
||||||
|
|
||||||
# Echo the assistant turn so Raycast sees its own reply +
|
|
||||||
# tool_calls in subsequent context.
|
|
||||||
assistant_text = "".join(acc.text_parts)
|
assistant_text = "".join(acc.text_parts)
|
||||||
assistant_tool_calls = [
|
assistant_tool_calls = [
|
||||||
ToolCall(
|
ToolCall(
|
||||||
@@ -576,10 +538,6 @@ class RaycastBackend:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
# Dispatch each MCP-routed call and append a ``tool``
|
|
||||||
# message. Calls not in the routing table get a placeholder
|
|
||||||
# error so the model can correct itself rather than the
|
|
||||||
# gateway hanging the conversation.
|
|
||||||
for tid in acc.tool_order:
|
for tid in acc.tool_order:
|
||||||
tool_name = acc.tool_names.get(tid, "")
|
tool_name = acc.tool_names.get(tid, "")
|
||||||
args_str = "".join(acc.tool_args.get(tid, [])) or "{}"
|
args_str = "".join(acc.tool_args.get(tid, [])) or "{}"
|
||||||
@@ -588,16 +546,12 @@ class RaycastBackend:
|
|||||||
tool_name=tool_name, args_json=args_str, mcp_routing=mcp_routing
|
tool_name=tool_name, args_json=args_str, mcp_routing=mcp_routing
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# Non-MCP tool — shouldn't really happen because we
|
|
||||||
# haven't surfaced any other locals, but defend.
|
|
||||||
result_text = f"Tool {tool_name!r} is not handled by the gateway."
|
result_text = f"Tool {tool_name!r} is not handled by the gateway."
|
||||||
working_messages.append(
|
working_messages.append(
|
||||||
RaycastMessage.tool(
|
RaycastMessage.tool(
|
||||||
tool_call_id=tid, name=tool_name, result=result_text
|
tool_call_id=tid, name=tool_name, result=result_text
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
# Loop: re-stream with the new history.
|
|
||||||
else:
|
else:
|
||||||
_log.warning(
|
_log.warning(
|
||||||
"raycast tool-call loop hit %d-turn cap for agent %r; "
|
"raycast tool-call loop hit %d-turn cap for agent %r; "
|
||||||
@@ -606,7 +560,6 @@ class RaycastBackend:
|
|||||||
agent.name,
|
agent.name,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Close whatever block is still open before the final delta.
|
|
||||||
if state.kind is not None:
|
if state.kind is not None:
|
||||||
yield build_content_block_stop(state.index)
|
yield build_content_block_stop(state.index)
|
||||||
state.kind = None
|
state.kind = None
|
||||||
@@ -716,6 +669,12 @@ class RaycastBackend:
|
|||||||
the final-summary chunk's arguments string is dropped because the
|
the final-summary chunk's arguments string is dropped because the
|
||||||
deltas already streamed it.
|
deltas already streamed it.
|
||||||
|
|
||||||
|
Providers disagree on how ``arguments`` arrives: streaming ones
|
||||||
|
(GPT) send deltas across chunks and then restate the full string in
|
||||||
|
a final summary chunk (skip it, already streamed); non-streaming
|
||||||
|
ones (Gemini) send it only in that final summary (emit it once,
|
||||||
|
since nothing streamed it first).
|
||||||
|
|
||||||
Side-effects on ``acc`` mirror what gets emitted to the wire so
|
Side-effects on ``acc`` mirror what gets emitted to the wire so
|
||||||
the gateway can rebuild a full ``ToolCall`` for the next Raycast
|
the gateway can rebuild a full ``ToolCall`` for the next Raycast
|
||||||
turn (it needs the joined ``arguments`` JSON string, which the
|
turn (it needs the joined ``arguments`` JSON string, which the
|
||||||
@@ -728,9 +687,6 @@ class RaycastBackend:
|
|||||||
raw_tc = raw_tcs[i] if i < len(raw_tcs) else {}
|
raw_tc = raw_tcs[i] if i < len(raw_tcs) else {}
|
||||||
idx_field = raw_tc.get("index") if isinstance(raw_tc, dict) else None
|
idx_field = raw_tc.get("index") if isinstance(raw_tc, dict) else None
|
||||||
|
|
||||||
# Resolve this entry to a tool-id key, mirroring
|
|
||||||
# `raycast_api.ChatResult._merge_tool_calls`. Phase 1 carries
|
|
||||||
# id+index, phase 2 only index, phase 3 only id.
|
|
||||||
tool_id: str | None = None
|
tool_id: str | None = None
|
||||||
if tc.id:
|
if tc.id:
|
||||||
tool_id = tc.id
|
tool_id = tc.id
|
||||||
@@ -743,7 +699,6 @@ class RaycastBackend:
|
|||||||
|
|
||||||
block_idx = state.tool_id_to_block.get(tool_id)
|
block_idx = state.tool_id_to_block.get(tool_id)
|
||||||
if block_idx is None:
|
if block_idx is None:
|
||||||
# New tool_use block. Close any open text/thinking block first.
|
|
||||||
if state.kind is not None:
|
if state.kind is not None:
|
||||||
events.append(build_content_block_stop(state.index))
|
events.append(build_content_block_stop(state.index))
|
||||||
state.index += 1
|
state.index += 1
|
||||||
@@ -756,22 +711,11 @@ class RaycastBackend:
|
|||||||
block_idx, tool_use_id=tool_id, name=tc.name or ""
|
block_idx, tool_use_id=tool_id, name=tc.name or ""
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
# Streaming providers (GPT) deliver arguments as deltas
|
|
||||||
# across chunks and re-state the full string in a final
|
|
||||||
# summary chunk; non-streaming-args providers (Gemini)
|
|
||||||
# only send args in the final summary. Emit args in
|
|
||||||
# both cases when this is the first appearance — final
|
|
||||||
# summary then IS the full string.
|
|
||||||
if tc.arguments:
|
if tc.arguments:
|
||||||
events.append(build_input_json_delta(block_idx, tc.arguments))
|
events.append(build_input_json_delta(block_idx, tc.arguments))
|
||||||
acc.add_tool_args(tool_id, tc.arguments)
|
acc.add_tool_args(tool_id, tc.arguments)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Existing block. Final summary chunks restate the full args
|
|
||||||
# string; if we already streamed deltas, that restatement is
|
|
||||||
# a duplicate (skip). If we streamed nothing (the streaming
|
|
||||||
# provider didn't send mid-arg deltas — Gemini path again),
|
|
||||||
# the summary IS the args — emit it once.
|
|
||||||
if is_final_summary:
|
if is_final_summary:
|
||||||
if tc.arguments and not acc.tool_args.get(tool_id):
|
if tc.arguments and not acc.tool_args.get(tool_id):
|
||||||
events.append(build_input_json_delta(block_idx, tc.arguments))
|
events.append(build_input_json_delta(block_idx, tc.arguments))
|
||||||
|
|||||||
@@ -1,13 +1,9 @@
|
|||||||
"""Pool of live Agent SDK sessions across every Claude agent (§3.2).
|
"""Pool of live Agent SDK sessions across every Claude agent.
|
||||||
|
|
||||||
One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess).
|
One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess). The
|
||||||
The pool owns the two decisions the adapters used to make on their own:
|
pool decides when a session is closed for idleness (TTL by conversation
|
||||||
when a session is closed for idleness (TTL by conversation kind) and
|
kind) and which one goes when memory runs out (RSS of the subprocess tree
|
||||||
which one goes when memory runs out (measured RSS of the subprocess tree
|
against the cgroup limit, or ``max_live`` with no limit).
|
||||||
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
|
from __future__ import annotations
|
||||||
@@ -17,18 +13,27 @@ import contextlib
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
|
from collections import deque
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Any, Protocol
|
from typing import TYPE_CHECKING, Any, Protocol
|
||||||
|
|
||||||
import psutil
|
import psutil
|
||||||
|
|
||||||
|
from beaver_gateway.backends.processes import Processes, subtree
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
from collections.abc import (
|
||||||
|
AsyncIterable,
|
||||||
|
AsyncIterator,
|
||||||
|
Callable,
|
||||||
|
Iterator,
|
||||||
|
Mapping,
|
||||||
|
)
|
||||||
|
|
||||||
__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"]
|
__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"]
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.core.sessions")
|
_log = logging.getLogger("beaver_gateway.backends.sessions")
|
||||||
|
|
||||||
DEFAULT_TTL: Mapping[str, float | None] = {
|
DEFAULT_TTL: Mapping[str, float | None] = {
|
||||||
"master": None,
|
"master": None,
|
||||||
@@ -41,12 +46,13 @@ DEFAULT_TTL: Mapping[str, float | None] = {
|
|||||||
|
|
||||||
_EVICT_ORDER = {"fork": 0, "job": 0, "deep": 1, "branch": 2, "master": 3}
|
_EVICT_ORDER = {"fork": 0, "job": 0, "deep": 1, "branch": 2, "master": 3}
|
||||||
_RSS_HEADROOM = 0.8
|
_RSS_HEADROOM = 0.8
|
||||||
|
STDERR_LINES = 200
|
||||||
|
|
||||||
|
|
||||||
class SessionClient(Protocol):
|
class SessionClient(Protocol):
|
||||||
async def connect(self) -> None: ...
|
async def connect(self) -> None: ...
|
||||||
async def query(self, prompt: str) -> None: ...
|
async def query(self, prompt: str | AsyncIterable[dict[str, Any]]) -> None: ...
|
||||||
def receive_response(self) -> AsyncIterator[Any]: ...
|
def receive_messages(self) -> AsyncIterator[Any]: ...
|
||||||
async def interrupt(self) -> None: ...
|
async def interrupt(self) -> None: ...
|
||||||
async def disconnect(self) -> None: ...
|
async def disconnect(self) -> None: ...
|
||||||
|
|
||||||
@@ -69,7 +75,22 @@ class Session:
|
|||||||
created_at: float = field(default_factory=time.monotonic)
|
created_at: float = field(default_factory=time.monotonic)
|
||||||
turns: int = 0
|
turns: int = 0
|
||||||
state: dict[str, Any] = field(default_factory=dict)
|
state: dict[str, Any] = field(default_factory=dict)
|
||||||
"""Scratch for policy rules (``core/policy``); dies with the process."""
|
"""Scratch for policy rules (``agents/policy.py``); dies with the process."""
|
||||||
|
|
||||||
|
reader: asyncio.Task[None] | None = None
|
||||||
|
"""Reads the client for the session's whole life, not only during a turn."""
|
||||||
|
inbox: asyncio.Queue[Any] = field(default_factory=asyncio.Queue)
|
||||||
|
"""Where the reader puts messages while a turn is attached."""
|
||||||
|
pending: list[Any] = field(default_factory=list)
|
||||||
|
"""A turn the CLI started on its own (a background task reported back),
|
||||||
|
buffered until a consumer attaches."""
|
||||||
|
attached: bool = False
|
||||||
|
woken: bool = False
|
||||||
|
note: str | None = None
|
||||||
|
"""Summary of the last task notification, the pending turn's prompt."""
|
||||||
|
observer: Callable[[Any], None] | None = None
|
||||||
|
"""The last turn's observer; sees subagent traffic between turns."""
|
||||||
|
stderr: deque[str] = field(default_factory=lambda: deque(maxlen=STDERR_LINES))
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def busy(self) -> bool:
|
def busy(self) -> bool:
|
||||||
@@ -77,6 +98,10 @@ class Session:
|
|||||||
|
|
||||||
@property
|
@property
|
||||||
def evictable(self) -> bool:
|
def evictable(self) -> bool:
|
||||||
|
"""Never evict a pinned (master) session or a ``dirty`` one.
|
||||||
|
|
||||||
|
``dirty`` means a mirror gap not yet repaired.
|
||||||
|
"""
|
||||||
return not (self.pinned or self.dirty or self.busy or self.pending_question)
|
return not (self.pinned or self.dirty or self.busy or self.pending_question)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
@@ -101,6 +126,7 @@ class SessionPool:
|
|||||||
self._max_live = max_live
|
self._max_live = max_live
|
||||||
self._ttl = dict(ttl)
|
self._ttl = dict(ttl)
|
||||||
self._reap_interval = reap_interval
|
self._reap_interval = reap_interval
|
||||||
|
self._processes = Processes()
|
||||||
|
|
||||||
def __len__(self) -> int:
|
def __len__(self) -> int:
|
||||||
return len(self._sessions)
|
return len(self._sessions)
|
||||||
@@ -148,25 +174,13 @@ class SessionPool:
|
|||||||
total += child.memory_info().rss
|
total += child.memory_info().rss
|
||||||
return total
|
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:
|
def over_limit(self) -> bool:
|
||||||
if self._rss_limit is not None:
|
if self._rss_limit is not None:
|
||||||
return self.rss() > self._rss_limit * _RSS_HEADROOM
|
return self.rss() > self._rss_limit * _RSS_HEADROOM
|
||||||
return len(self._sessions) >= self._max_live
|
return len(self._sessions) >= self._max_live
|
||||||
|
|
||||||
def victims(self) -> list[Session]:
|
def victims(self) -> list[Session]:
|
||||||
|
"""Evictable sessions, forks/jobs first, then oldest idle."""
|
||||||
candidates = [s for s in self._sessions.values() if s.evictable]
|
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))
|
candidates.sort(key=lambda s: (_EVICT_ORDER.get(s.kind, 1), s.last_used))
|
||||||
return candidates
|
return candidates
|
||||||
@@ -191,11 +205,23 @@ class SessionPool:
|
|||||||
if session is None:
|
if session is None:
|
||||||
return
|
return
|
||||||
_log.info("closing session %s (%s, %s)", session.session_id, session.kind, key)
|
_log.info("closing session %s (%s, %s)", session.session_id, session.kind, key)
|
||||||
|
if session.reader is not None and session.reader is not asyncio.current_task():
|
||||||
|
session.reader.cancel()
|
||||||
|
if session.attached:
|
||||||
|
session.inbox.put_nowait(RuntimeError(f"session {key} was closed"))
|
||||||
try:
|
try:
|
||||||
await session.client.disconnect()
|
await session.client.disconnect()
|
||||||
except Exception: # noqa: BLE001
|
except Exception: # noqa: BLE001
|
||||||
_log.exception("disconnect failed for session %s", session.session_id)
|
_log.exception("disconnect failed for session %s", session.session_id)
|
||||||
|
|
||||||
|
async def interrupt(self, key: str) -> bool:
|
||||||
|
session = self._sessions.get(key)
|
||||||
|
if session is None or not session.busy:
|
||||||
|
return False
|
||||||
|
session.interrupt_requested = True
|
||||||
|
await session.client.interrupt()
|
||||||
|
return True
|
||||||
|
|
||||||
async def close_all(self, *, agent: str | None = None) -> None:
|
async def close_all(self, *, agent: str | None = None) -> None:
|
||||||
for session in list(self._sessions.values()):
|
for session in list(self._sessions.values()):
|
||||||
if agent is None or session.agent == agent:
|
if agent is None or session.agent == agent:
|
||||||
@@ -223,26 +249,57 @@ class SessionPool:
|
|||||||
_log.exception("session reaper failed")
|
_log.exception("session reaper failed")
|
||||||
|
|
||||||
def snapshot(self) -> list[dict[str, Any]]:
|
def snapshot(self) -> list[dict[str, Any]]:
|
||||||
|
return self.census()["sessions"]
|
||||||
|
|
||||||
|
def census(self) -> dict[str, Any]:
|
||||||
now = time.monotonic()
|
now = time.monotonic()
|
||||||
return [
|
rows = self._processes.tree_parents_first()
|
||||||
{
|
owned: set[int] = set()
|
||||||
"key": s.key,
|
sessions: list[dict[str, Any]] = []
|
||||||
"agent": s.agent,
|
for s in self._sessions.values():
|
||||||
"kind": s.kind,
|
tree = subtree(rows, s.pid)
|
||||||
"session_id": s.session_id,
|
owned.update(row["pid"] for row in tree)
|
||||||
"pid": s.pid,
|
sessions.append(
|
||||||
"rss": self.rss_of(s),
|
{
|
||||||
"idle_seconds": round(now - s.last_used, 1),
|
"key": s.key,
|
||||||
"age_seconds": round(now - s.created_at, 1),
|
"agent": s.agent,
|
||||||
"turns": s.turns,
|
"kind": s.kind,
|
||||||
"busy": s.busy,
|
"session_id": s.session_id,
|
||||||
"running_turn": s.running_turn,
|
"pid": s.pid,
|
||||||
"pending_question": s.pending_question,
|
"rss": sum(row["rss"] for row in tree) if tree else None,
|
||||||
"pinned": s.pinned,
|
"cpu": _cpu(tree),
|
||||||
"dirty": s.dirty,
|
"idle_seconds": round(now - s.last_used, 1),
|
||||||
}
|
"age_seconds": round(now - s.created_at, 1),
|
||||||
for s in self._sessions.values()
|
"turns": s.turns,
|
||||||
]
|
"busy": s.busy,
|
||||||
|
"running_turn": s.running_turn,
|
||||||
|
"pending": len(s.pending),
|
||||||
|
"pending_question": s.pending_question,
|
||||||
|
"pinned": s.pinned,
|
||||||
|
"dirty": s.dirty,
|
||||||
|
"processes": tree,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"rss": sum(row["rss"] for row in rows),
|
||||||
|
"rss_limit": self._rss_limit,
|
||||||
|
"cpu": _cpu(rows),
|
||||||
|
"sessions": sessions,
|
||||||
|
"others": [row for row in rows if row["pid"] not in owned],
|
||||||
|
}
|
||||||
|
|
||||||
|
async def kill(self, pid: int) -> bool:
|
||||||
|
if not self._processes.kill_own_descendant(pid):
|
||||||
|
return False
|
||||||
|
for session in list(self._sessions.values()):
|
||||||
|
if session.pid == pid:
|
||||||
|
await self.close(session.key)
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def _cpu(rows: list[dict[str, Any]]) -> float | None:
|
||||||
|
rates = [row["cpu"] for row in rows if row["cpu"] is not None]
|
||||||
|
return round(sum(rates), 1) if rates else None
|
||||||
|
|
||||||
|
|
||||||
def cgroup_limit() -> int | None:
|
def cgroup_limit() -> int | None:
|
||||||
+101
-24
@@ -1,16 +1,7 @@
|
|||||||
"""Anthropic messages <-> Agent SDK transcript entries.
|
"""Anthropic messages <-> Agent SDK transcript entries.
|
||||||
|
|
||||||
:func:`build_entries` renders a message list as the entries the Claude CLI
|
:func:`build_entries` renders a message list as entries the Claude CLI
|
||||||
itself writes (reference: ``t/spike_sdk/entries_reference.json``, CLI
|
itself writes; :func:`messages_from_entries` is the projection back.
|
||||||
2.1.248 via ``import_session_to_store``): one ``user`` entry per prompt,
|
|
||||||
one ``assistant`` entry per content block sharing a message id, one
|
|
||||||
``user`` entry per ``tool_result`` parented on the matching ``tool_use``
|
|
||||||
entry. Only ``user``/``assistant`` entries are produced - no attachments,
|
|
||||||
titles or queue markers. Appending the result to a session store and
|
|
||||||
resuming that session id seeds an external history into the SDK.
|
|
||||||
|
|
||||||
:func:`messages_from_entries` is the projection back, used by tests and by
|
|
||||||
anything that needs Anthropic-shape history out of a mirrored transcript.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -19,27 +10,33 @@ import hashlib
|
|||||||
import uuid as _uuid
|
import uuid as _uuid
|
||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from claude_agent_sdk._cli_version import __cli_version__ as _cli_version
|
from claude_agent_sdk._cli_version import __cli_version__
|
||||||
|
|
||||||
|
_cli_version = str(__cli_version__)
|
||||||
except ImportError: # pragma: no cover
|
except ImportError: # pragma: no cover
|
||||||
_cli_version = "2.1.248"
|
_cli_version = "2.1.248"
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Iterable
|
from collections.abc import Iterable, Sequence
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"CLI_VERSION",
|
"CLI_VERSION",
|
||||||
"build_entries",
|
"build_entries",
|
||||||
"close_open_tool_uses",
|
"close_open_tool_uses",
|
||||||
"fingerprint",
|
"fingerprint",
|
||||||
|
"last_message_before",
|
||||||
"messages_from_entries",
|
"messages_from_entries",
|
||||||
"open_tool_uses",
|
"open_tool_uses",
|
||||||
"prompt_count",
|
"prompt_count",
|
||||||
"render_messages",
|
"render_messages",
|
||||||
|
"reply_text",
|
||||||
|
"spoken_text",
|
||||||
"strip_tool_entries",
|
"strip_tool_entries",
|
||||||
"text_of",
|
"text_of",
|
||||||
|
"turn_starts",
|
||||||
"window_entries",
|
"window_entries",
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -59,6 +56,14 @@ def build_entries(
|
|||||||
permission_mode: str = "bypassPermissions",
|
permission_mode: str = "bypassPermissions",
|
||||||
now: datetime | None = None,
|
now: datetime | None = None,
|
||||||
) -> list[dict[str, Any]]:
|
) -> list[dict[str, Any]]:
|
||||||
|
"""Render ``messages`` as entries a session store can resume from.
|
||||||
|
|
||||||
|
One ``user`` entry per prompt, one ``assistant`` entry per content
|
||||||
|
block sharing a message id, one ``user`` entry per ``tool_result``
|
||||||
|
parented on the matching ``tool_use``. Appending the result to a
|
||||||
|
session store and resuming that session id seeds external history
|
||||||
|
into the SDK.
|
||||||
|
"""
|
||||||
stamp = (now or datetime.now(UTC)).strftime("%Y-%m-%dT%H:%M:%S.") + (
|
stamp = (now or datetime.now(UTC)).strftime("%Y-%m-%dT%H:%M:%S.") + (
|
||||||
f"{(now or datetime.now(UTC)).microsecond // 1000:03d}Z"
|
f"{(now or datetime.now(UTC)).microsecond // 1000:03d}Z"
|
||||||
)
|
)
|
||||||
@@ -160,7 +165,10 @@ def build_entries(
|
|||||||
return entries
|
return entries
|
||||||
|
|
||||||
|
|
||||||
def messages_from_entries(entries: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]:
|
def messages_from_entries(
|
||||||
|
entries: Iterable[Mapping[str, Any]], *, stamped: bool = False
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
"""Messages as the model saw them; ``stamped`` adds each one's entry id and time."""
|
||||||
out: list[dict[str, Any]] = []
|
out: list[dict[str, Any]] = []
|
||||||
last_message_id: str | None = None
|
last_message_id: str | None = None
|
||||||
for entry in entries:
|
for entry in entries:
|
||||||
@@ -179,7 +187,14 @@ def messages_from_entries(entries: Iterable[Mapping[str, Any]]) -> list[dict[str
|
|||||||
):
|
):
|
||||||
out[-1]["content"].extend(results)
|
out[-1]["content"].extend(results)
|
||||||
else:
|
else:
|
||||||
out.append({"role": "user", "content": _user_content(content)})
|
message_out: dict[str, Any] = {
|
||||||
|
"role": "user",
|
||||||
|
"content": _user_content(content),
|
||||||
|
}
|
||||||
|
if stamped:
|
||||||
|
message_out["id"] = entry.get("uuid")
|
||||||
|
message_out["ts"] = entry.get("timestamp")
|
||||||
|
out.append(message_out)
|
||||||
last_message_id = None
|
last_message_id = None
|
||||||
continue
|
continue
|
||||||
blocks = _assistant_blocks(content)
|
blocks = _assistant_blocks(content)
|
||||||
@@ -192,7 +207,11 @@ def messages_from_entries(entries: Iterable[Mapping[str, Any]]) -> list[dict[str
|
|||||||
):
|
):
|
||||||
out[-1]["content"].extend(blocks)
|
out[-1]["content"].extend(blocks)
|
||||||
else:
|
else:
|
||||||
out.append({"role": "assistant", "content": blocks})
|
message_out: dict[str, Any] = {"role": "assistant", "content": blocks}
|
||||||
|
if stamped:
|
||||||
|
message_out["id"] = entry.get("uuid")
|
||||||
|
message_out["ts"] = entry.get("timestamp")
|
||||||
|
out.append(message_out)
|
||||||
last_message_id = message_id
|
last_message_id = message_id
|
||||||
return out
|
return out
|
||||||
|
|
||||||
@@ -291,10 +310,8 @@ def _zero_usage() -> dict[str, Any]:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
# ---- repair, windows, projections ---------------------------------------
|
|
||||||
|
|
||||||
_PROMPT_TYPES = ("user", "assistant")
|
_PROMPT_TYPES = ("user", "assistant")
|
||||||
_INTERRUPTED = "прервано"
|
_INTERRUPTED = "interrupted"
|
||||||
|
|
||||||
|
|
||||||
def open_tool_uses(
|
def open_tool_uses(
|
||||||
@@ -423,13 +440,72 @@ def prompt_count(entries: Iterable[Mapping[str, Any]]) -> int:
|
|||||||
return sum(1 for e in entries if _is_prompt_entry(e))
|
return sum(1 for e in entries if _is_prompt_entry(e))
|
||||||
|
|
||||||
|
|
||||||
|
def turn_starts(entries: Iterable[Mapping[str, Any]]) -> list[int]:
|
||||||
|
"""Where each turn opens: the index of every prompt the SDK was handed.
|
||||||
|
|
||||||
|
The CLI stamps only those with ``promptSource``; what it writes as a
|
||||||
|
user on its own (an interruption note, a skill body) stays inside the
|
||||||
|
turn it belongs to.
|
||||||
|
"""
|
||||||
|
return [
|
||||||
|
index
|
||||||
|
for index, entry in enumerate(entries)
|
||||||
|
if entry.get("promptSource") and _is_prompt_entry(entry)
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def last_message_before(entries: Sequence[Mapping[str, Any]], index: int) -> str | None:
|
||||||
|
"""Uuid of the last ``user``/``assistant`` entry ahead of ``index``."""
|
||||||
|
for entry in reversed(entries[:index]):
|
||||||
|
uuid = entry.get("uuid")
|
||||||
|
if (
|
||||||
|
entry.get("type") in _PROMPT_TYPES
|
||||||
|
and isinstance(uuid, str)
|
||||||
|
and not entry.get("isSidechain")
|
||||||
|
):
|
||||||
|
return uuid
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def spoken_text(block: Any, *, thinking_is_progress: bool) -> str:
|
||||||
|
"""What a content block (dict or model) says to the human.
|
||||||
|
|
||||||
|
Text always; thinking only when the agent's thinking is progress
|
||||||
|
written for the human, never when it is reasoning.
|
||||||
|
"""
|
||||||
|
if isinstance(block, Mapping):
|
||||||
|
fields = cast("Mapping[str, Any]", block)
|
||||||
|
kind, text, thinking = (
|
||||||
|
fields.get("type"),
|
||||||
|
fields.get("text"),
|
||||||
|
fields.get("thinking"),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
kind = getattr(block, "type", None)
|
||||||
|
text, thinking = getattr(block, "text", None), getattr(block, "thinking", None)
|
||||||
|
if kind == "text":
|
||||||
|
return str(text or "").strip()
|
||||||
|
if kind == "thinking" and thinking_is_progress:
|
||||||
|
return str(thinking or "").strip()
|
||||||
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
def reply_text(content: Iterable[Any], *, thinking_is_progress: bool) -> str:
|
||||||
|
"""The reply as the human reads it: spoken blocks in order, a blank line apart."""
|
||||||
|
parts = (spoken_text(b, thinking_is_progress=thinking_is_progress) for b in content)
|
||||||
|
return "\n\n".join(p for p in parts if p)
|
||||||
|
|
||||||
|
|
||||||
def render_messages(
|
def render_messages(
|
||||||
messages: Iterable[Mapping[str, Any]], *, window: int | None = None
|
messages: Iterable[Mapping[str, Any]],
|
||||||
|
*,
|
||||||
|
thinking_is_progress: bool,
|
||||||
|
window: int | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Plain-text projection for ``read_conversation``: ``user:``/``assistant:`` turns.
|
"""Plain-text projection for ``read_conversation``: ``user:``/``assistant:`` turns.
|
||||||
|
|
||||||
Tool calls collapse to a one-line summary per assistant turn; tool
|
Tool calls collapse to a one-line summary per assistant turn; tool
|
||||||
results and thinking are dropped.
|
results and reasoning are dropped, progress thinking reads as text.
|
||||||
"""
|
"""
|
||||||
turns: list[str] = []
|
turns: list[str] = []
|
||||||
tools: list[str] = []
|
tools: list[str] = []
|
||||||
@@ -455,8 +531,9 @@ def render_messages(
|
|||||||
turns.append("user:\n" + _text_of_content(content))
|
turns.append("user:\n" + _text_of_content(content))
|
||||||
continue
|
continue
|
||||||
for block in _assistant_blocks(content):
|
for block in _assistant_blocks(content):
|
||||||
if block.get("type") == "text" and block.get("text"):
|
text = spoken_text(block, thinking_is_progress=thinking_is_progress)
|
||||||
current.append(str(block["text"]))
|
if text:
|
||||||
|
current.append(text)
|
||||||
elif block.get("type") == "tool_use":
|
elif block.get("type") == "tool_use":
|
||||||
tools.append(str(block.get("name", "")))
|
tools.append(str(block.get("name", "")))
|
||||||
flush()
|
flush()
|
||||||
+13
-504
@@ -1,532 +1,41 @@
|
|||||||
"""Process entrypoint.
|
"""Process entrypoint: logging, signals, ``.env``, the config, then ``app.run``."""
|
||||||
|
|
||||||
Phase 1.4 — async ``main``: install uvloop, load the user config, build
|
|
||||||
registries + per-agent backends (only ``RaycastBackend`` so far), wire
|
|
||||||
each frontend with a ``GatewayRuntime``, and run all
|
|
||||||
``frontend.serve()`` coroutines concurrently. Without any frontends we
|
|
||||||
still print the Phase 0 DoD line and exit cleanly so the bare skeleton
|
|
||||||
keeps working.
|
|
||||||
|
|
||||||
Phase 2.1 — when the user declares any ``McpServer``, we additionally
|
|
||||||
build the internal MCP aggregator app and run it on
|
|
||||||
``127.0.0.1:INTERNAL_MCP_PORT`` as another task inside the same
|
|
||||||
TaskGroup. URLs are surfaced through ``GatewayRuntime.mcp_internal_urls``
|
|
||||||
so Phase 2.2's ClaudeCode adapter can find them.
|
|
||||||
|
|
||||||
Phase 3 — the same aggregator backs the external ``McpServerFrontend``;
|
|
||||||
``cli`` doesn't have to know that the frontend reverse-proxies into it,
|
|
||||||
it just keeps the aggregator running for anyone who needs it.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import contextlib
|
import contextlib
|
||||||
import functools
|
|
||||||
import logging
|
import logging
|
||||||
import signal
|
import signal
|
||||||
from contextlib import AsyncExitStack
|
|
||||||
from typing import TYPE_CHECKING, Any
|
|
||||||
|
|
||||||
import psycopg
|
|
||||||
import uvicorn
|
|
||||||
import uvloop
|
import uvloop
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from pgqueuer import PsycopgDriver
|
|
||||||
from raycast_api import Client as RaycastClient
|
|
||||||
from raycast_api.config import Config as RaycastConfig
|
|
||||||
|
|
||||||
from beaver_gateway import config_loader
|
from beaver_gateway import app, config
|
||||||
from beaver_gateway.agents.claude import ClaudeAgent
|
from beaver_gateway.security.redact import install as install_redaction
|
||||||
from beaver_gateway.agents.raycast import RaycastAgent
|
from beaver_gateway.security.redact import load_secrets as load_secrets_to_mask
|
||||||
from beaver_gateway.backends.claude_sdk import (
|
|
||||||
ClaudeSdkBackend,
|
|
||||||
RunnerConfig,
|
|
||||||
UsageEvent,
|
|
||||||
)
|
|
||||||
from beaver_gateway.backends.raycast import RaycastBackend
|
|
||||||
from beaver_gateway.core.auth import TokenStore
|
|
||||||
from beaver_gateway.core.bus import EventBus
|
|
||||||
from beaver_gateway.core.conversations import Conversations
|
|
||||||
from beaver_gateway.core.envelope import Envelope
|
|
||||||
from beaver_gateway.core.gateway_tools import build_tool_server
|
|
||||||
from beaver_gateway.core.redact import install as install_redaction
|
|
||||||
from beaver_gateway.core.redact import load_secrets as load_secrets_to_mask
|
|
||||||
from beaver_gateway.core.registry import AgentRegistry, Gateway, McpRegistry
|
|
||||||
from beaver_gateway.core.rotation import Rotation, RotationPolicy
|
|
||||||
from beaver_gateway.core.scheduler import Scheduler
|
|
||||||
from beaver_gateway.core.sessions import SessionPool
|
|
||||||
from beaver_gateway.frontends._auth import require_token
|
|
||||||
from beaver_gateway.frontends.base import GatewayRuntime
|
|
||||||
from beaver_gateway.frontends.root import build_root_app
|
|
||||||
from beaver_gateway.mcp.internal_app import build_internal_app
|
|
||||||
from beaver_gateway.settings import Settings
|
from beaver_gateway.settings import Settings
|
||||||
from beaver_gateway.storage import (
|
|
||||||
Database,
|
|
||||||
PostgresSessionStore,
|
|
||||||
Usage,
|
|
||||||
append_audit,
|
|
||||||
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
|
|
||||||
from starlette.requests import Request
|
|
||||||
from starlette.types import ASGIApp
|
|
||||||
|
|
||||||
from beaver_gateway.backends.base import Backend
|
|
||||||
from beaver_gateway.core.policy import ToolAudit
|
|
||||||
from beaver_gateway.mcp.types import McpServerT
|
|
||||||
|
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.cli")
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
"""Sync wrapper: uvloop loop factory + asyncio.run."""
|
|
||||||
logging.basicConfig(
|
logging.basicConfig(
|
||||||
level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
|
level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
|
||||||
)
|
)
|
||||||
install_redaction()
|
install_redaction()
|
||||||
_install_sigterm_handler()
|
_sigterm_as_interrupt()
|
||||||
asyncio.run(_async_main(), loop_factory=uvloop.new_event_loop)
|
asyncio.run(_run(), loop_factory=uvloop.new_event_loop)
|
||||||
|
|
||||||
|
|
||||||
def _install_sigterm_handler() -> None:
|
async def _run() -> None:
|
||||||
"""Turn SIGTERM into a normal interpreter exit.
|
load_dotenv(override=False)
|
||||||
|
load_secrets_to_mask()
|
||||||
|
settings = Settings() # ty: ignore[missing-argument]
|
||||||
|
gateway = config.load(settings.config_path)
|
||||||
|
await app.run(gateway, settings)
|
||||||
|
|
||||||
Python's default SIGTERM disposition kills the process outright, so
|
|
||||||
neither ``AsyncExitStack`` unwinding nor ``atexit`` hooks run. That
|
|
||||||
matters because every ``claude`` we spawn lives in its own session
|
|
||||||
(ptyprocess calls ``setsid``), which makes it immune to the signal
|
|
||||||
that took us down — a hard SIGTERM leaves one orphaned CLI per live
|
|
||||||
session, each holding hundreds of MB. Raising ``KeyboardInterrupt``
|
|
||||||
instead routes ``docker stop`` / ``systemctl stop`` through the same
|
|
||||||
shutdown path as Ctrl-C, which does reap them.
|
|
||||||
|
|
||||||
Only installed when we own the main thread's signal handlers; under
|
|
||||||
an embedding host that isn't ours to take.
|
|
||||||
"""
|
|
||||||
|
|
||||||
|
def _sigterm_as_interrupt() -> None:
|
||||||
def _raise_interrupt(_signum: int, _frame: object) -> None:
|
def _raise_interrupt(_signum: int, _frame: object) -> None:
|
||||||
raise KeyboardInterrupt
|
raise KeyboardInterrupt
|
||||||
|
|
||||||
with contextlib.suppress(ValueError, OSError):
|
with contextlib.suppress(ValueError, OSError):
|
||||||
signal.signal(signal.SIGTERM, _raise_interrupt)
|
signal.signal(signal.SIGTERM, _raise_interrupt)
|
||||||
|
|
||||||
|
|
||||||
async def _async_main() -> None:
|
|
||||||
# Populate ``os.environ`` from ``.env`` before anything else so the
|
|
||||||
# user's ``config.py`` can read its own secrets via ``os.environ[...]``
|
|
||||||
# (Firefly PAT, third-party MCP creds, etc.). ``Settings`` already
|
|
||||||
# reads ``.env`` independently via pydantic-settings, but that path
|
|
||||||
# populates Settings fields, not the process environment.
|
|
||||||
# ``override=False``: real env vars (Docker, systemd) win over .env.
|
|
||||||
load_dotenv(override=False)
|
|
||||||
# Only now does the process environment hold the credentials the
|
|
||||||
# redactor masks literally (in Docker they arrive via ``env_file``).
|
|
||||||
load_secrets_to_mask()
|
|
||||||
|
|
||||||
settings = Settings() # ty: ignore[missing-argument]
|
|
||||||
gateway = config_loader.load(settings.config_path)
|
|
||||||
|
|
||||||
agents = AgentRegistry(gateway.agents)
|
|
||||||
mcps = McpRegistry(gateway.mcps)
|
|
||||||
|
|
||||||
# Phase 4.1 — open the async DB and run create_all once. Engine
|
|
||||||
# pool is process-wide; ``dispose()`` after the TaskGroup unwinds.
|
|
||||||
db = Database(settings.database_url)
|
|
||||||
await db.create_all()
|
|
||||||
|
|
||||||
# Phase 4.2 — TokenStore now reads from the DB (in-memory cache
|
|
||||||
# primed at start, TTL-refreshed, last_used_at flushed by a
|
|
||||||
# background task). BOOTSTRAP_TOKENS layers on top so first-run /
|
|
||||||
# examples still work without DB writes.
|
|
||||||
token_store = TokenStore(
|
|
||||||
db,
|
|
||||||
bootstrap=TokenStore.parse_bootstrap(settings.bootstrap_tokens),
|
|
||||||
bootstrap_scopes=TokenStore.parse_bootstrap_scopes(settings.bootstrap_tokens),
|
|
||||||
)
|
|
||||||
|
|
||||||
async with AsyncExitStack() as stack:
|
|
||||||
stack.push_async_callback(db.dispose)
|
|
||||||
await token_store.start()
|
|
||||||
stack.push_async_callback(token_store.stop)
|
|
||||||
# Internal MCP URLs must exist before we construct any
|
|
||||||
# ClaudeSdkBackend - adapters bake the URLs into their
|
|
||||||
# ``mcp_servers`` at construction time. The
|
|
||||||
# ``mcp_servers`` map is used by the Raycast backend, which
|
|
||||||
# needs in-process ``list_tools`` / ``call_tool`` access (the
|
|
||||||
# Raycast wire has no native MCP concept).
|
|
||||||
internal_app, internal_urls, mcp_servers = _build_internal_mcp(
|
|
||||||
gateway.mcps, settings=settings
|
|
||||||
)
|
|
||||||
|
|
||||||
# Prefetch tool catalogs for every MCP so RaycastAgent requests
|
|
||||||
# don't pay a per-turn list_tools roundtrip and so a broken MCP
|
|
||||||
# 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,
|
|
||||||
frontends=gateway.frontends,
|
|
||||||
envelope=Envelope(
|
|
||||||
watch=gateway.watch, tz=gateway.tz, recall=gateway.recall
|
|
||||||
),
|
|
||||||
distiller=gateway.distiller,
|
|
||||||
user_sink=gateway.user_sink,
|
|
||||||
)
|
|
||||||
late.conversations = conversations
|
|
||||||
scheduler = Scheduler(
|
|
||||||
conversations=conversations,
|
|
||||||
jobs=gateway.jobs,
|
|
||||||
driver=await _pgqueuer_driver(settings.database_url, stack),
|
|
||||||
budget=gateway.budget,
|
|
||||||
rotation=Rotation(
|
|
||||||
conversations, gateway.rotation or RotationPolicy(tz=gateway.tz)
|
|
||||||
),
|
|
||||||
)
|
|
||||||
conversations.scheduler = scheduler
|
|
||||||
|
|
||||||
runtime = GatewayRuntime(
|
|
||||||
agents=agents,
|
|
||||||
mcps=mcps,
|
|
||||||
backends=backends,
|
|
||||||
token_store=token_store,
|
|
||||||
db=db,
|
|
||||||
mcp_internal_urls=internal_urls,
|
|
||||||
admin_user=settings.admin_user,
|
|
||||||
admin_pass=settings.admin_pass,
|
|
||||||
session_secret=settings.session_secret,
|
|
||||||
frontends=tuple(gateway.frontends),
|
|
||||||
conversations=conversations,
|
|
||||||
bus=bus,
|
|
||||||
pool=pool,
|
|
||||||
scheduler=scheduler,
|
|
||||||
public_url=gateway.public_url.rstrip("/") if gateway.public_url else None,
|
|
||||||
)
|
|
||||||
|
|
||||||
for fe in gateway.frontends:
|
|
||||||
fe.configure(runtime)
|
|
||||||
|
|
||||||
_log.info(
|
|
||||||
"beaver-gateway: loaded %d agents, %d mcps, %d frontends",
|
|
||||||
len(agents),
|
|
||||||
len(mcps),
|
|
||||||
len(gateway.frontends),
|
|
||||||
)
|
|
||||||
|
|
||||||
# Keep the Phase 0 DoD line on stdout for grep-friendly smoke
|
|
||||||
# tests, in addition to the structured log line above.
|
|
||||||
print(
|
|
||||||
f"beaver-gateway: loaded {len(agents)} agents, "
|
|
||||||
f"{len(mcps)} mcps, {len(gateway.frontends)} frontends"
|
|
||||||
)
|
|
||||||
|
|
||||||
if not gateway.frontends:
|
|
||||||
# No external listeners → nothing to serve. The internal
|
|
||||||
# MCP app has no consumer on its own, so we skip running
|
|
||||||
# it in this path and exit cleanly (Phase 0 DoD).
|
|
||||||
return
|
|
||||||
|
|
||||||
await conversations.start()
|
|
||||||
stack.push_async_callback(conversations.stop)
|
|
||||||
await scheduler.start()
|
|
||||||
stack.push_async_callback(scheduler.stop)
|
|
||||||
hooks = scheduler.app(
|
|
||||||
functools.partial(_authorize_hook, runtime=runtime, scope="api")
|
|
||||||
)
|
|
||||||
async with asyncio.TaskGroup() as tg:
|
|
||||||
tg.create_task(pool.reap_loop())
|
|
||||||
if internal_app is not None:
|
|
||||||
tg.create_task(_serve_internal_mcp(internal_app, settings=settings))
|
|
||||||
tg.create_task(_serve_root(gateway, extra={"/hooks": hooks}))
|
|
||||||
if gateway.watch is not None:
|
|
||||||
tg.create_task(gateway.watch.run())
|
|
||||||
for fe in gateway.frontends:
|
|
||||||
tg.create_task(fe.serve())
|
|
||||||
|
|
||||||
|
|
||||||
async def _authorize_hook(
|
|
||||||
request: Request, *, runtime: GatewayRuntime, scope: str
|
|
||||||
) -> str:
|
|
||||||
return await require_token(request, runtime, scope=scope)
|
|
||||||
|
|
||||||
|
|
||||||
async def _pgqueuer_driver(url: str, stack: AsyncExitStack) -> PsycopgDriver | None:
|
|
||||||
"""A dedicated autocommit connection for pgqueuer's LISTEN/NOTIFY."""
|
|
||||||
plain = _plain_postgres_url(url)
|
|
||||||
if plain is None:
|
|
||||||
return None
|
|
||||||
conn = await psycopg.AsyncConnection.connect(plain, autocommit=True)
|
|
||||||
stack.push_async_callback(conn.close)
|
|
||||||
return PsycopgDriver(conn)
|
|
||||||
|
|
||||||
|
|
||||||
def _plain_postgres_url(url: str) -> str | None:
|
|
||||||
for prefix in ("postgresql+psycopg://", "postgresql://", "postgres://"):
|
|
||||||
if url.startswith(prefix):
|
|
||||||
return "postgresql://" + url[len(prefix) :]
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
async def _serve_root(gateway: Gateway, *, extra: dict[str, ASGIApp]) -> None:
|
|
||||||
app = build_root_app(gateway.frontends, extra=extra)
|
|
||||||
config = uvicorn.Config(app, host=gateway.host, port=gateway.port, log_level="info")
|
|
||||||
_log.info(
|
|
||||||
"gateway on http://%s:%d - %s",
|
|
||||||
gateway.host,
|
|
||||||
gateway.port,
|
|
||||||
", ".join([*(fe.path for fe in gateway.frontends if fe.path), *extra])
|
|
||||||
or "no http frontends",
|
|
||||||
)
|
|
||||||
await uvicorn.Server(config).serve()
|
|
||||||
|
|
||||||
|
|
||||||
def _build_internal_mcp(
|
|
||||||
mcps: list[McpServerT], *, settings: Settings
|
|
||||||
) -> tuple[Starlette | None, dict[str, str], dict[str, FastMCP]]:
|
|
||||||
"""Build the aggregator app + URL map + server map, or empty equivalents.
|
|
||||||
|
|
||||||
The URL map is always handed out (frontends may still introspect
|
|
||||||
``runtime.mcp_internal_urls`` even if nothing is configured); the
|
|
||||||
app is ``None`` when there are no MCPs to mount, so the caller
|
|
||||||
skips the uvicorn task entirely. The server map is the in-process
|
|
||||||
handle the Raycast backend needs to splice MCP tools into its
|
|
||||||
requests — empty when no MCPs are configured.
|
|
||||||
"""
|
|
||||||
if not mcps:
|
|
||||||
return None, {}, {}
|
|
||||||
return build_internal_app(mcps, host="127.0.0.1", port=settings.internal_mcp_port)
|
|
||||||
|
|
||||||
|
|
||||||
async def _prefetch_mcp_tools(
|
|
||||||
servers: dict[str, FastMCP],
|
|
||||||
) -> dict[str, list[FastMCPTool]]:
|
|
||||||
"""Eagerly enumerate tools per MCP so the Raycast loop has a static catalog.
|
|
||||||
|
|
||||||
Each underlying proxy is allowed to fail independently — a broken
|
|
||||||
MCP shouldn't take down the whole gateway. The result has one entry
|
|
||||||
per MCP that responded; agents that ``expose_mcps`` a missing entry
|
|
||||||
will simply expose no tools from it (logged once per request).
|
|
||||||
"""
|
|
||||||
out: dict[str, list[FastMCPTool]] = {}
|
|
||||||
for name, server in servers.items():
|
|
||||||
try:
|
|
||||||
out[name] = list(await server.list_tools())
|
|
||||||
except Exception: # noqa: BLE001 — proxy can raise any transport error; we degrade per-MCP rather than fail the whole gateway
|
|
||||||
_log.exception("failed to list tools for MCP %r — skipping", name)
|
|
||||||
out[name] = []
|
|
||||||
return out
|
|
||||||
|
|
||||||
|
|
||||||
async def _serve_internal_mcp(app: Starlette, *, settings: Settings) -> None:
|
|
||||||
"""Run the internal MCP aggregator on loopback.
|
|
||||||
|
|
||||||
Bound to ``127.0.0.1`` (never EXPOSE'd) — only the in-process
|
|
||||||
ClaudeCode subprocess reaches it. Logged at ``warning`` level so
|
|
||||||
we don't drown the gateway's own logs in per-request noise.
|
|
||||||
"""
|
|
||||||
config = uvicorn.Config(
|
|
||||||
app,
|
|
||||||
host="127.0.0.1",
|
|
||||||
port=settings.internal_mcp_port,
|
|
||||||
log_level="warning",
|
|
||||||
loop="uvloop",
|
|
||||||
)
|
|
||||||
server = uvicorn.Server(config)
|
|
||||||
_log.info(
|
|
||||||
"internal MCP aggregator on http://127.0.0.1:%d/mcp/<name>",
|
|
||||||
settings.internal_mcp_port,
|
|
||||||
)
|
|
||||||
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 ask(self, key: str, payload: dict[str, Any]) -> str:
|
|
||||||
if self.conversations is None:
|
|
||||||
msg = "conversations service is not up yet"
|
|
||||||
raise RuntimeError(msg)
|
|
||||||
answer = await self.conversations.ask(key, payload)
|
|
||||||
return self.conversations.answer_text(answer)
|
|
||||||
|
|
||||||
|
|
||||||
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.
|
|
||||||
|
|
||||||
The Raycast ``Client`` is shared across every ``RaycastAgent``
|
|
||||||
(bearer + device-id are process-wide), so we open it lazily - only
|
|
||||||
when at least one ``RaycastAgent`` is present - and close it via
|
|
||||||
the caller's exit stack.
|
|
||||||
|
|
||||||
Each :class:`ClaudeAgent` gets its own :class:`ClaudeSdkBackend`
|
|
||||||
(own live-session pool, own prompt and MCP set); all of them share
|
|
||||||
the session store and the usage sink.
|
|
||||||
"""
|
|
||||||
backends: dict[str, Backend] = {}
|
|
||||||
|
|
||||||
raycast_agents = [a for a in agents if isinstance(a, RaycastAgent)]
|
|
||||||
if raycast_agents:
|
|
||||||
client = await _try_open_raycast_client(settings, stack)
|
|
||||||
if client is not None:
|
|
||||||
raycast_backend = RaycastBackend(
|
|
||||||
client, mcp_servers=mcp_servers, mcp_tools=mcp_tools
|
|
||||||
)
|
|
||||||
for a in raycast_agents:
|
|
||||||
backends[a.name] = raycast_backend
|
|
||||||
|
|
||||||
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()
|
|
||||||
}
|
|
||||||
|
|
||||||
async def record_usage(event: UsageEvent) -> None:
|
|
||||||
row = Usage(
|
|
||||||
agent_name=event.agent_name,
|
|
||||||
conversation_id=event.conversation_id,
|
|
||||||
session_id=event.session_id,
|
|
||||||
model=event.model,
|
|
||||||
effort=event.effort,
|
|
||||||
input_tokens=event.usage.input_tokens,
|
|
||||||
output_tokens=event.usage.output_tokens,
|
|
||||||
cache_read_tokens=event.usage.cache_read_tokens,
|
|
||||||
cache_creation_tokens=event.usage.cache_creation_tokens,
|
|
||||||
context_tokens=event.usage.context_tokens,
|
|
||||||
cost_usd=event.usage.cost_usd,
|
|
||||||
duration_ms=event.usage.duration_ms,
|
|
||||||
num_turns=event.usage.num_turns,
|
|
||||||
model_usage=event.usage.model_usage,
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
async with db.session() as session:
|
|
||||||
await append_usage(session, row)
|
|
||||||
except Exception: # noqa: BLE001
|
|
||||||
_log.exception("usage write failed for %s", event.agent_name)
|
|
||||||
|
|
||||||
async def record_tool(event: ToolAudit) -> None:
|
|
||||||
detail = {
|
|
||||||
"conversation": event.conversation,
|
|
||||||
"kind": event.kind,
|
|
||||||
"tool": event.tool,
|
|
||||||
"decision": event.decision,
|
|
||||||
"reason": event.reason,
|
|
||||||
"brief": event.brief,
|
|
||||||
}
|
|
||||||
try:
|
|
||||||
async with db.session() as session:
|
|
||||||
await append_audit(
|
|
||||||
session,
|
|
||||||
actor=f"agent:{event.agent}",
|
|
||||||
kind="tool_call",
|
|
||||||
agent_name=event.agent,
|
|
||||||
detail=detail,
|
|
||||||
)
|
|
||||||
except Exception: # noqa: BLE001
|
|
||||||
_log.exception("tool audit write failed for %s", event.agent)
|
|
||||||
|
|
||||||
for a in agents:
|
|
||||||
if isinstance(a, ClaudeAgent):
|
|
||||||
adapter = ClaudeSdkBackend(
|
|
||||||
agent=a,
|
|
||||||
mcp_internal_urls=mcp_internal_urls,
|
|
||||||
session_store=session_store,
|
|
||||||
mcp_tool_names=mcp_tool_names,
|
|
||||||
runner=runner,
|
|
||||||
usage_sink=record_usage,
|
|
||||||
pool=pool,
|
|
||||||
tool_server=functools.partial(late.server, names=a.gateway_tools),
|
|
||||||
asker=late.ask,
|
|
||||||
audit_sink=record_tool,
|
|
||||||
)
|
|
||||||
await stack.enter_async_context(adapter)
|
|
||||||
backends[a.name] = adapter
|
|
||||||
|
|
||||||
return backends
|
|
||||||
|
|
||||||
|
|
||||||
async def _try_open_raycast_client(
|
|
||||||
settings: Settings, stack: AsyncExitStack
|
|
||||||
) -> RaycastClient | None:
|
|
||||||
"""Open a ``raycast_api.Client`` if creds are available, else warn + skip.
|
|
||||||
|
|
||||||
Skipping is gentler than failing startup: the user can bring up the
|
|
||||||
gateway, list their agents, and still hit a ``ClaudeAgent``; affected
|
|
||||||
``RaycastAgent`` instances 503 with a specific message at request time.
|
|
||||||
Missing ``raycast.json`` falls into the same bucket — first-run
|
|
||||||
users won't have it yet.
|
|
||||||
"""
|
|
||||||
if not settings.raycast_bearer:
|
|
||||||
_log.warning(
|
|
||||||
"RaycastAgent present but RAYCAST_BEARER is unset — those agents will 503"
|
|
||||||
)
|
|
||||||
return None
|
|
||||||
if not settings.raycast_device_id:
|
|
||||||
_log.warning(
|
|
||||||
"RaycastAgent present but RAYCAST_DEVICE_ID is unset — those agents "
|
|
||||||
"will 503 (generate once with `python -c 'import secrets; "
|
|
||||||
"print(secrets.token_hex(32))'`)"
|
|
||||||
)
|
|
||||||
return None
|
|
||||||
if not settings.raycast_config_path.exists():
|
|
||||||
_log.warning(
|
|
||||||
"RaycastAgent present but %s is missing — those agents will 503",
|
|
||||||
settings.raycast_config_path,
|
|
||||||
)
|
|
||||||
return None
|
|
||||||
|
|
||||||
config = RaycastConfig.load(settings.raycast_config_path)
|
|
||||||
client = RaycastClient(
|
|
||||||
config=config,
|
|
||||||
bearer_token=settings.raycast_bearer,
|
|
||||||
device_id=settings.raycast_device_id,
|
|
||||||
locale=settings.raycast_locale,
|
|
||||||
)
|
|
||||||
return await stack.enter_async_context(client)
|
|
||||||
@@ -1,20 +1,10 @@
|
|||||||
"""Load the user's ``/config/config.py``.
|
"""``Gateway`` - the one object a setup's ``config.py`` assembles - and its loader."""
|
||||||
|
|
||||||
The config file is regular Python. We ``exec`` it in a namespace seeded
|
|
||||||
with the public surface the user is expected to use (``ClaudeAgent``,
|
|
||||||
``RaycastAgent``, ``McpServer``, ``ExposedMcp``, ``Gateway``). The file
|
|
||||||
must assign a top-level ``gateway = Gateway(...)``.
|
|
||||||
|
|
||||||
Element-level validation already happens at construction time — agents
|
|
||||||
and ``McpServer`` factories are pydantic models that reject garbage.
|
|
||||||
What we *can't* validate at construction is the ``Gateway`` container
|
|
||||||
itself (deliberately a plain dataclass per PRD), so we type-check its
|
|
||||||
contents here before handing it back.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
import sys
|
import sys
|
||||||
|
from dataclasses import dataclass, field
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from croniter import croniter
|
from croniter import croniter
|
||||||
@@ -22,18 +12,60 @@ from croniter import croniter
|
|||||||
from beaver_gateway.agents.base import BaseAgent, ExposedMcp
|
from beaver_gateway.agents.base import BaseAgent, ExposedMcp
|
||||||
from beaver_gateway.agents.claude import ClaudeAgent
|
from beaver_gateway.agents.claude import ClaudeAgent
|
||||||
from beaver_gateway.agents.raycast import RaycastAgent
|
from beaver_gateway.agents.raycast import RaycastAgent
|
||||||
from beaver_gateway.core.conversations import ConversationTexts
|
from beaver_gateway.conversations.texts import ConversationTexts
|
||||||
from beaver_gateway.core.registry import Gateway
|
|
||||||
from beaver_gateway.core.scheduler import Job
|
|
||||||
from beaver_gateway.frontends.base import Frontend
|
from beaver_gateway.frontends.base import Frontend
|
||||||
|
from beaver_gateway.jobs.scheduler import Job
|
||||||
from beaver_gateway.mcp.types import HttpMcp, McpServer, PythonToolMcp, StdioMcp
|
from beaver_gateway.mcp.types import HttpMcp, McpServer, PythonToolMcp, StdioMcp
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.distill import Distiller
|
||||||
|
from beaver_gateway.conversations.envelope import RecallContext
|
||||||
|
from beaver_gateway.conversations.rotation import RotationPolicy
|
||||||
|
from beaver_gateway.conversations.texts import UserSaid
|
||||||
|
from beaver_gateway.jobs.scheduler import Budget
|
||||||
|
from beaver_gateway.mcp.types import McpServerT
|
||||||
|
from beaver_gateway.vault.watch import VaultWatch
|
||||||
|
|
||||||
|
__all__ = ["ConfigError", "Gateway", "load"]
|
||||||
|
|
||||||
|
_SCOPE = re.compile(r"[a-z][a-z0-9_-]*")
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(slots=True)
|
||||||
|
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
|
||||||
|
"""Every string the gateway says to a model; ``None`` = the English defaults."""
|
||||||
|
jobs: list[Job] = field(default_factory=list)
|
||||||
|
"""Cron, webhook and event jobs run by ``jobs.scheduler``."""
|
||||||
|
rotation: RotationPolicy | None = None
|
||||||
|
"""When a master is replaced by a fresh one; ``None`` keeps the defaults."""
|
||||||
|
watch: VaultWatch | None = None
|
||||||
|
"""Directory watcher feeding the envelope; ``None`` = no change block."""
|
||||||
|
recall: Callable[[RecallContext], str | None] | None = None
|
||||||
|
"""Envelope lookup on the user's text: pointers the setup derives from its files."""
|
||||||
|
user_sink: Callable[[UserSaid], Awaitable[None] | None] | None = None
|
||||||
|
"""Sees every user message entering a master or branch turn."""
|
||||||
|
budget: Budget | None = None
|
||||||
|
"""Subscription utilisation past which non-critical jobs wait."""
|
||||||
|
distiller: Distiller | None = None
|
||||||
|
"""Who closes deep chats and where digests and the index live."""
|
||||||
|
tz: str = "UTC"
|
||||||
|
"""Zone for the envelope clock, cron expressions and the rotation hour."""
|
||||||
|
host: str = "0.0.0.0" # noqa: S104
|
||||||
|
port: int = 8000
|
||||||
|
"""The one listener; every HTTP frontend is mounted under its ``path``."""
|
||||||
|
public_url: str | None = None
|
||||||
|
"""Origin the reverse proxy shows the world; ``None`` derives it per request."""
|
||||||
|
|
||||||
|
|
||||||
class ConfigError(Exception):
|
class ConfigError(Exception):
|
||||||
"""User config file is missing, unreadable, or structurally wrong."""
|
pass
|
||||||
|
|
||||||
|
|
||||||
_PUBLIC_NAMES: dict[str, Any] = {
|
_PUBLIC_NAMES: dict[str, Any] = {
|
||||||
@@ -49,7 +81,7 @@ _McpInstance = StdioMcp | HttpMcp | PythonToolMcp
|
|||||||
|
|
||||||
|
|
||||||
def load(path: Path) -> Gateway:
|
def load(path: Path) -> Gateway:
|
||||||
"""Execute ``path`` and return its top-level ``gateway`` object."""
|
"""Execute a setup's ``config.py`` and return its top-level ``gateway``."""
|
||||||
try:
|
try:
|
||||||
source = path.read_text(encoding="utf-8")
|
source = path.read_text(encoding="utf-8")
|
||||||
except FileNotFoundError as exc:
|
except FileNotFoundError as exc:
|
||||||
@@ -60,12 +92,11 @@ def load(path: Path) -> Gateway:
|
|||||||
raise ConfigError(msg) from exc
|
raise ConfigError(msg) from exc
|
||||||
|
|
||||||
code = compile(source, str(path), "exec")
|
code = compile(source, str(path), "exec")
|
||||||
# Siblings of the config (``policy.py``, ``mcps/``) import by name.
|
|
||||||
parent = str(path.resolve().parent)
|
parent = str(path.resolve().parent)
|
||||||
if parent not in sys.path:
|
if parent not in sys.path:
|
||||||
sys.path.insert(0, parent)
|
sys.path.insert(0, parent)
|
||||||
namespace: dict[str, Any] = {"__file__": str(path), **_PUBLIC_NAMES}
|
namespace: dict[str, Any] = {"__file__": str(path), **_PUBLIC_NAMES}
|
||||||
exec(code, namespace) # noqa: S102 - exec'ing user config is the feature
|
exec(code, namespace) # noqa: S102
|
||||||
|
|
||||||
try:
|
try:
|
||||||
gw = namespace["gateway"]
|
gw = namespace["gateway"]
|
||||||
@@ -107,6 +138,12 @@ def _validate(gw: Gateway, path: Path) -> None:
|
|||||||
f"got {type(f).__name__}"
|
f"got {type(f).__name__}"
|
||||||
)
|
)
|
||||||
raise ConfigError(msg)
|
raise ConfigError(msg)
|
||||||
|
if f.scope is not None and not _SCOPE.fullmatch(f.scope):
|
||||||
|
msg = (
|
||||||
|
f"{path}: gateway.frontends[{i}] has scope {f.scope!r}; "
|
||||||
|
f"a frontend scope is a lowercase name like 'voice', never '*'"
|
||||||
|
)
|
||||||
|
raise ConfigError(msg)
|
||||||
names: set[str] = set()
|
names: set[str] = set()
|
||||||
for i, j in enumerate(gw.jobs):
|
for i, j in enumerate(gw.jobs):
|
||||||
if not isinstance(j, Job):
|
if not isinstance(j, Job):
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
"""Conversations: rows, queue, seeds, turns, questions, closing, rotation, envelope."""
|
||||||
@@ -0,0 +1,322 @@
|
|||||||
|
"""Ending conversations: the distiller, the line cap, the master handover."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import contextlib
|
||||||
|
import inspect
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import UTC, datetime, timedelta
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.distill import (
|
||||||
|
Digest,
|
||||||
|
DistillContext,
|
||||||
|
LineCap,
|
||||||
|
append_index,
|
||||||
|
check_digest,
|
||||||
|
find_digest,
|
||||||
|
index_line,
|
||||||
|
trim_summary,
|
||||||
|
written_paths,
|
||||||
|
)
|
||||||
|
from beaver_gateway.conversations.questions import Questions
|
||||||
|
from beaver_gateway.conversations.state import aware
|
||||||
|
from beaver_gateway.conversations.texts import NewDayContext
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from beaver_gateway.conversations.rotation import HandoutContext
|
||||||
|
from beaver_gateway.storage.models import Conversation, InjectQueueItem
|
||||||
|
|
||||||
|
__all__ = ["Closing", "DistillResult"]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
CLOSE_WAIT = 0.25
|
||||||
|
CLOSE_TRIES = 40
|
||||||
|
CAP_TRIES = 3
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class DistillResult:
|
||||||
|
conversation: Conversation
|
||||||
|
fork: Conversation
|
||||||
|
text: str
|
||||||
|
digest: Digest | None
|
||||||
|
error: str | None
|
||||||
|
trimmed: bool
|
||||||
|
|
||||||
|
|
||||||
|
class Closing(Questions):
|
||||||
|
async def close(self, conv: Conversation) -> Conversation:
|
||||||
|
row = await self.set_status(conv, "closed")
|
||||||
|
with contextlib.suppress(LookupError):
|
||||||
|
await self._backend(conv.agent_name).close(conv.external_id)
|
||||||
|
return row
|
||||||
|
|
||||||
|
async def request_close(self, conv: Conversation) -> Conversation:
|
||||||
|
"""``close_chat`` from inside a turn: the chat closes once the turn ends."""
|
||||||
|
if conv.kind != "deep":
|
||||||
|
msg = f"only deep chats close this way, {conv.external_id} is {conv.kind}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
return await self.set_flags(conv, {"close_requested": True})
|
||||||
|
|
||||||
|
async def idle(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
kind: str,
|
||||||
|
days: int,
|
||||||
|
since: datetime | None = None,
|
||||||
|
limit: int | None = None,
|
||||||
|
) -> list[Conversation]:
|
||||||
|
"""Open conversations of ``kind`` with a session, quiet for ``days``."""
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
cutoff = now - timedelta(days=days)
|
||||||
|
out: list[tuple[datetime, Conversation]] = []
|
||||||
|
for conv in await self.find(status="open", kind=kind, limit=10_000):
|
||||||
|
if conv.session_id is None:
|
||||||
|
continue
|
||||||
|
last = aware(conv.last_activity_at or conv.created_at)
|
||||||
|
if last > cutoff or (since is not None and last < since):
|
||||||
|
continue
|
||||||
|
out.append((last, conv))
|
||||||
|
out.sort(key=lambda pair: pair[0])
|
||||||
|
rows = [conv for _, conv in out]
|
||||||
|
return rows[:limit] if limit is not None else rows
|
||||||
|
|
||||||
|
async def distill(
|
||||||
|
self, conv: Conversation, *, reason: str = "api"
|
||||||
|
) -> DistillResult:
|
||||||
|
"""Fork under the distiller: digest checked and indexed, merge to the master."""
|
||||||
|
if self._distiller is None:
|
||||||
|
msg = "no distiller configured (Gateway(distiller=...))"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
if conv.kind != "deep":
|
||||||
|
msg = f"only deep chats are distilled, {conv.external_id} is {conv.kind}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
row = await self.get_row(cast("int", conv.id)) or conv
|
||||||
|
if row.status != "open":
|
||||||
|
msg = f"conversation {row.external_id} is {row.status}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
if await self.busy(row):
|
||||||
|
msg = f"conversation {row.external_id} is busy"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
memory = bool(row.flags.get("memory", True))
|
||||||
|
chat_name = await self.chat_name(row)
|
||||||
|
ctx = DistillContext(
|
||||||
|
conversation=row,
|
||||||
|
title=await self.implied_title(row),
|
||||||
|
source=await self.window_of(row),
|
||||||
|
chat_name=chat_name,
|
||||||
|
memory=memory,
|
||||||
|
reason=reason,
|
||||||
|
day=datetime.now(UTC).astimezone().date(),
|
||||||
|
)
|
||||||
|
prompt = await self._distill_prompt(ctx)
|
||||||
|
started = datetime.now(UTC)
|
||||||
|
self._bus.publish(
|
||||||
|
"distill.start",
|
||||||
|
conversation_id=row.external_id,
|
||||||
|
reason=reason,
|
||||||
|
memory=memory,
|
||||||
|
)
|
||||||
|
result = await self.fork(
|
||||||
|
row,
|
||||||
|
prompt,
|
||||||
|
strip_tools=True,
|
||||||
|
agent=self._distiller.agent,
|
||||||
|
title=f"digest: {chat_name}",
|
||||||
|
)
|
||||||
|
text, trimmed = trim_summary(result.text)
|
||||||
|
digest: Digest | None = None
|
||||||
|
error: str | None = None
|
||||||
|
if memory:
|
||||||
|
written = written_paths(result.capture.synthesized_messages)
|
||||||
|
path = find_digest(self._distiller, since=started, written=written)
|
||||||
|
if path is None:
|
||||||
|
error = self._texts.digest_missing
|
||||||
|
else:
|
||||||
|
checked = check_digest(path, self._distiller)
|
||||||
|
if isinstance(checked, str):
|
||||||
|
error = f"{path.name}: {checked}"
|
||||||
|
else:
|
||||||
|
digest = checked
|
||||||
|
append_index(self._distiller, index_line(digest, chat_name))
|
||||||
|
if error is not None:
|
||||||
|
_log.warning("distill of %s: %s", row.external_id, error)
|
||||||
|
master = await self.open_master()
|
||||||
|
if master is not None and text:
|
||||||
|
note = self._texts.closed.format(
|
||||||
|
chat=chat_name,
|
||||||
|
digest=(
|
||||||
|
self._texts.closed_digest.format(digest=digest.path.stem)
|
||||||
|
if digest
|
||||||
|
else ""
|
||||||
|
),
|
||||||
|
text=text,
|
||||||
|
)
|
||||||
|
await self.inject(master, note, urgency="normal", origin="digest")
|
||||||
|
await self.close(row)
|
||||||
|
row = await self.set_flags(
|
||||||
|
row,
|
||||||
|
{
|
||||||
|
"close_requested": None,
|
||||||
|
"closed_reason": reason,
|
||||||
|
"digest": str(digest.path) if digest else None,
|
||||||
|
"digest_error": error,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
self._bus.publish(
|
||||||
|
"conversation.distilled",
|
||||||
|
conversation_id=row.external_id,
|
||||||
|
fork=result.conversation.external_id,
|
||||||
|
reason=reason,
|
||||||
|
memory=memory,
|
||||||
|
digest=str(digest.path) if digest else None,
|
||||||
|
error=error,
|
||||||
|
text=text,
|
||||||
|
trimmed=trimmed,
|
||||||
|
master=master.external_id if master is not None else None,
|
||||||
|
)
|
||||||
|
return DistillResult(
|
||||||
|
conversation=row,
|
||||||
|
fork=result.conversation,
|
||||||
|
text=text,
|
||||||
|
digest=digest,
|
||||||
|
error=error,
|
||||||
|
trimmed=trimmed,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _distill_prompt(self, ctx: DistillContext) -> str:
|
||||||
|
source = self._texts.distill
|
||||||
|
if source is None:
|
||||||
|
template = (
|
||||||
|
self._texts.distill_prompt
|
||||||
|
if ctx.memory
|
||||||
|
else self._texts.distill_prompt_no_memory
|
||||||
|
)
|
||||||
|
return template.format(
|
||||||
|
chat=ctx.chat_name, reason=ctx.reason, day=ctx.day.isoformat()
|
||||||
|
)
|
||||||
|
produced: Any = source(ctx)
|
||||||
|
return await produced if inspect.isawaitable(produced) else produced
|
||||||
|
|
||||||
|
async def _close_after_turn(self, conv: Conversation) -> None:
|
||||||
|
for _ in range(CLOSE_TRIES):
|
||||||
|
if await self.busy(conv):
|
||||||
|
await asyncio.sleep(CLOSE_WAIT)
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
await self.distill(conv, reason="close_chat")
|
||||||
|
except RuntimeError as exc:
|
||||||
|
_log.info("closing %s: %s, retrying", conv.external_id, exc)
|
||||||
|
await asyncio.sleep(CLOSE_WAIT)
|
||||||
|
continue
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("closing %s after its turn failed", conv.external_id)
|
||||||
|
return
|
||||||
|
_log.warning("closing %s: still busy, giving up", conv.external_id)
|
||||||
|
|
||||||
|
async def before_turn(self, conv: Conversation) -> str | None:
|
||||||
|
cap = LineCap.from_flags(conv.flags.get("line_cap"))
|
||||||
|
if cap is None:
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
return cap.path.read_text(encoding="utf-8") if cap.path.exists() else ""
|
||||||
|
except OSError:
|
||||||
|
_log.exception("line cap: cannot read %s", cap.path)
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def after_turn(self, conv: Conversation, before: str | None) -> None:
|
||||||
|
row = await self.get_row(cast("int", conv.id))
|
||||||
|
if row is None:
|
||||||
|
return
|
||||||
|
if row.kind == "deep" and row.flags.get("close_requested"):
|
||||||
|
self._track(asyncio.create_task(self._close_after_turn(row)))
|
||||||
|
cap = LineCap.from_flags(row.flags.get("line_cap"))
|
||||||
|
if cap is not None and before is not None:
|
||||||
|
await self._enforce_cap(row, cap, before)
|
||||||
|
|
||||||
|
async def _enforce_cap(self, conv: Conversation, cap: LineCap, before: str) -> None:
|
||||||
|
if not cap.path.exists():
|
||||||
|
return
|
||||||
|
after = cap.path.read_text(encoding="utf-8")
|
||||||
|
lines = sum(1 for line in after.splitlines() if line.strip())
|
||||||
|
if lines <= cap.max_lines:
|
||||||
|
return
|
||||||
|
if before:
|
||||||
|
cap.path.write_text(before, encoding="utf-8")
|
||||||
|
else:
|
||||||
|
cap.path.unlink()
|
||||||
|
attempts = int(conv.flags.get("line_cap_attempts", 0) or 0) + 1
|
||||||
|
await self.set_flags(conv, {"line_cap_attempts": attempts})
|
||||||
|
self._bus.publish(
|
||||||
|
"line_cap.bounced",
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
path=str(cap.path),
|
||||||
|
lines=lines,
|
||||||
|
max_lines=cap.max_lines,
|
||||||
|
attempt=attempts,
|
||||||
|
)
|
||||||
|
_log.warning(
|
||||||
|
"line cap: %s came back with %d lines (cap %d), restored; attempt %d",
|
||||||
|
cap.path,
|
||||||
|
lines,
|
||||||
|
cap.max_lines,
|
||||||
|
attempts,
|
||||||
|
)
|
||||||
|
if attempts > CAP_TRIES:
|
||||||
|
return
|
||||||
|
await self.inject(
|
||||||
|
conv,
|
||||||
|
self._texts.too_long.format(
|
||||||
|
name=cap.path.name, lines=lines, max_lines=cap.max_lines
|
||||||
|
),
|
||||||
|
urgency="urgent",
|
||||||
|
origin="cap",
|
||||||
|
interrupt=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def handout(self, conv: Conversation, ctx: HandoutContext) -> str:
|
||||||
|
"""The closing master's last turn."""
|
||||||
|
source = self._texts.handout
|
||||||
|
if isinstance(source, str):
|
||||||
|
prompt = source.format(day=ctx.day.isoformat(), reason=ctx.reason)
|
||||||
|
else:
|
||||||
|
produced: Any = source(ctx)
|
||||||
|
prompt = await produced if inspect.isawaitable(produced) else produced
|
||||||
|
self._bus.publish(
|
||||||
|
"handout.start", conversation_id=conv.external_id, day=ctx.day.isoformat()
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
text, _ = await self.run_text_turn(conv, prompt, origin="handout")
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("handout turn on %s failed", conv.external_id)
|
||||||
|
text = ""
|
||||||
|
self._bus.publish(
|
||||||
|
"handout.end",
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
day=ctx.day.isoformat(),
|
||||||
|
text=text[:2000],
|
||||||
|
)
|
||||||
|
return text
|
||||||
|
|
||||||
|
async def new_day(
|
||||||
|
self, conv: Conversation, *, reason: str = "night", moved: int = 0
|
||||||
|
) -> InjectQueueItem:
|
||||||
|
"""The new master's first inject."""
|
||||||
|
ctx = NewDayContext(
|
||||||
|
day=datetime.now(UTC).astimezone().date(), reason=reason, moved=moved
|
||||||
|
)
|
||||||
|
source = self._texts.new_day
|
||||||
|
if isinstance(source, str):
|
||||||
|
text = source.format(day=ctx.day.isoformat(), reason=ctx.reason)
|
||||||
|
else:
|
||||||
|
produced: Any = source(ctx)
|
||||||
|
text = await produced if inspect.isawaitable(produced) else produced
|
||||||
|
if moved:
|
||||||
|
text += self._texts.moved_injects.format(moved=moved)
|
||||||
|
return await self.inject(
|
||||||
|
conv, text, urgency="urgent", origin="rotation", interrupt=False
|
||||||
|
)
|
||||||
@@ -0,0 +1,246 @@
|
|||||||
|
"""What a conversation's model holds: prompt granules, skills, tools, files touched."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from collections import Counter
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Iterable, Mapping
|
||||||
|
|
||||||
|
from beaver_gateway.agents.claude import ClaudeAgent
|
||||||
|
from beaver_gateway.conversations.kinds import Kind
|
||||||
|
|
||||||
|
__all__ = ["bash_paths", "compose", "files_touched", "granules", "skill_sets"]
|
||||||
|
|
||||||
|
CHARS_PER_TOKEN = 3.2
|
||||||
|
READ_TOOLS = frozenset({"Read", "NotebookRead"})
|
||||||
|
WRITE_TOOLS = frozenset({"Write", "Edit", "MultiEdit", "NotebookEdit"})
|
||||||
|
SEARCH_TOOLS = frozenset({"Glob", "Grep"})
|
||||||
|
PATH_KEYS = ("file_path", "notebook_path", "path", "file", "filename")
|
||||||
|
_FRONTMATTER = re.compile(r"^---\s*\n(.*?)\n---", re.DOTALL)
|
||||||
|
# A shell command names files as quoted strings or bare words; only notes
|
||||||
|
# and plain text files count, code and binaries are the agent's own business.
|
||||||
|
_SHELL_TOKEN = re.compile(r'"([^"\n]+)"|\'([^\'\n]+)\'|([^\s"\'`;|&()<>]+)')
|
||||||
|
_NOTE_SUFFIXES = (".md", ".canvas", ".txt", ".csv", ".json", ".yaml", ".yml")
|
||||||
|
_SHELL_CD = re.compile(r'(?:^|&&|;|\|\|)\s*cd\s+("([^"]+)"|\'([^\']+)\'|(\S+))')
|
||||||
|
_SHELL_WRITES = re.compile(
|
||||||
|
r"write_text|write_bytes|>>?\s*[\"\']?[^\s&|]|\btee\b|\bsed\s+-i|\bmv\b|\brm\b"
|
||||||
|
r"|\btouch\b|\bcp\b|open\([^)]*[\"\'][wa]"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def bash_paths(command: str, cwd: Path | None) -> tuple[list[str], Path | None, bool]:
|
||||||
|
"""Files a shell command names, where it leaves the shell, whether it writes.
|
||||||
|
|
||||||
|
Relative names resolve against ``cwd``, or against the directory of a
|
||||||
|
leading ``cd`` in the same command. Absolute names stay as they are.
|
||||||
|
"""
|
||||||
|
here = cwd
|
||||||
|
for match in _SHELL_CD.finditer(command):
|
||||||
|
target = match.group(2) or match.group(3) or match.group(4) or ""
|
||||||
|
target = target.strip()
|
||||||
|
if not target:
|
||||||
|
continue
|
||||||
|
here = Path(target) if target.startswith("/") else (here or Path()) / target
|
||||||
|
found: list[str] = []
|
||||||
|
for match in _SHELL_TOKEN.finditer(command):
|
||||||
|
raw = (match.group(1) or match.group(2) or match.group(3) or "").strip()
|
||||||
|
if not raw.lower().endswith(_NOTE_SUFFIXES) or "*" in raw:
|
||||||
|
continue
|
||||||
|
path = Path(raw) if raw.startswith("/") else (here or Path()) / raw
|
||||||
|
text = str(path)
|
||||||
|
if text not in found:
|
||||||
|
found.append(text)
|
||||||
|
return found, here, bool(_SHELL_WRITES.search(command))
|
||||||
|
|
||||||
|
|
||||||
|
def _tokens(chars: int) -> int:
|
||||||
|
return round(chars / CHARS_PER_TOKEN)
|
||||||
|
|
||||||
|
|
||||||
|
def _relative(path: str, cwd: Path | None) -> str:
|
||||||
|
if cwd is None:
|
||||||
|
return path
|
||||||
|
try:
|
||||||
|
return str(Path(path).relative_to(cwd))
|
||||||
|
except ValueError:
|
||||||
|
return path
|
||||||
|
|
||||||
|
|
||||||
|
def _tool_uses(entries: Iterable[Mapping[str, Any]]) -> Iterable[tuple[str, Any, str]]:
|
||||||
|
for entry in entries:
|
||||||
|
message = entry.get("message")
|
||||||
|
if entry.get("type") != "assistant" or not isinstance(message, dict):
|
||||||
|
continue
|
||||||
|
content = message.get("content")
|
||||||
|
if not isinstance(content, list):
|
||||||
|
continue
|
||||||
|
stamp = str(entry.get("timestamp") or "")
|
||||||
|
for block in content:
|
||||||
|
if isinstance(block, dict) and block.get("type") == "tool_use":
|
||||||
|
yield str(block.get("name") or "?"), block.get("input"), stamp
|
||||||
|
|
||||||
|
|
||||||
|
def files_touched(
|
||||||
|
entries: Iterable[Mapping[str, Any]], *, cwd: Path | None = None
|
||||||
|
) -> tuple[list[dict[str, Any]], dict[str, int]]:
|
||||||
|
"""Files named in tool inputs, with how they were touched, plus tool counts."""
|
||||||
|
files: dict[str, dict[str, Any]] = {}
|
||||||
|
counts: Counter[str] = Counter()
|
||||||
|
# The shell keeps its directory between calls; follow it.
|
||||||
|
shell_cwd = cwd
|
||||||
|
|
||||||
|
def touch(raw: str, how: str, stamp: str) -> None:
|
||||||
|
key = _relative(raw, cwd)
|
||||||
|
row = files.setdefault(
|
||||||
|
key, {"path": key, "reads": 0, "writes": 0, "other": 0, "last_at": stamp}
|
||||||
|
)
|
||||||
|
row[how] += 1
|
||||||
|
row["last_at"] = max(row["last_at"], stamp)
|
||||||
|
|
||||||
|
for name, tool_input, stamp in _tool_uses(entries):
|
||||||
|
counts[name] += 1
|
||||||
|
if not isinstance(tool_input, dict):
|
||||||
|
continue
|
||||||
|
if name == "Bash" and isinstance(tool_input.get("command"), str):
|
||||||
|
paths, shell_cwd, writes = bash_paths(tool_input["command"], shell_cwd)
|
||||||
|
for raw in paths:
|
||||||
|
touch(raw, "writes" if writes else "reads", stamp)
|
||||||
|
continue
|
||||||
|
raw = next(
|
||||||
|
(v for k in PATH_KEYS if isinstance(v := tool_input.get(k), str) and v),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
if raw is None or (name in SEARCH_TOOLS and "file_path" not in tool_input):
|
||||||
|
continue
|
||||||
|
if name in READ_TOOLS:
|
||||||
|
touch(raw, "reads", stamp)
|
||||||
|
elif name in WRITE_TOOLS:
|
||||||
|
touch(raw, "writes", stamp)
|
||||||
|
else:
|
||||||
|
touch(raw, "other", stamp)
|
||||||
|
ordered = sorted(files.values(), key=lambda r: r["last_at"], reverse=True)
|
||||||
|
return ordered, dict(counts)
|
||||||
|
|
||||||
|
|
||||||
|
def granules(agent: ClaudeAgent, kind: Kind) -> dict[str, Any]:
|
||||||
|
"""The system prompt as its sources: tag, path, size, token estimate."""
|
||||||
|
sources = agent.prompt_for(kind) if agent.serves(kind) else None
|
||||||
|
rows: list[dict[str, Any]] = []
|
||||||
|
if sources is None:
|
||||||
|
chars = len(agent.system_prompt)
|
||||||
|
if chars:
|
||||||
|
rows.append(
|
||||||
|
{
|
||||||
|
"tag": None,
|
||||||
|
"path": None,
|
||||||
|
"bytes": chars,
|
||||||
|
"tokens_est": _tokens(chars),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
for source in sources:
|
||||||
|
tag, raw = source if isinstance(source, tuple) else (None, source)
|
||||||
|
path = Path(raw)
|
||||||
|
try:
|
||||||
|
chars = len(path.read_text(encoding="utf-8"))
|
||||||
|
except OSError:
|
||||||
|
chars = 0
|
||||||
|
rows.append(
|
||||||
|
{
|
||||||
|
"tag": tag,
|
||||||
|
"path": _relative(str(path), agent.cwd),
|
||||||
|
"bytes": chars,
|
||||||
|
"tokens_est": _tokens(chars),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
total = sum(r["bytes"] for r in rows)
|
||||||
|
return {"bytes": total, "tokens_est": _tokens(total), "granules": rows}
|
||||||
|
|
||||||
|
|
||||||
|
def _skill_meta(skill_md: Path) -> dict[str, str]:
|
||||||
|
try:
|
||||||
|
text = skill_md.read_text(encoding="utf-8")
|
||||||
|
except OSError:
|
||||||
|
return {}
|
||||||
|
match = _FRONTMATTER.match(text)
|
||||||
|
meta: dict[str, str] = {}
|
||||||
|
if match:
|
||||||
|
for line in match.group(1).splitlines():
|
||||||
|
key, sep, value = line.partition(":")
|
||||||
|
if sep:
|
||||||
|
meta[key.strip()] = value.strip()
|
||||||
|
return meta
|
||||||
|
|
||||||
|
|
||||||
|
def skill_sets(agent: ClaudeAgent, kind: Kind) -> list[dict[str, Any]]:
|
||||||
|
"""Skill-set directories the kind loads, with each skill's name and description."""
|
||||||
|
out: list[dict[str, Any]] = []
|
||||||
|
for directory in agent.skills_for(kind):
|
||||||
|
skills = []
|
||||||
|
try:
|
||||||
|
children = sorted(p for p in directory.iterdir() if p.is_dir())
|
||||||
|
except OSError:
|
||||||
|
children = []
|
||||||
|
for child in children:
|
||||||
|
skill_md = child / "SKILL.md"
|
||||||
|
if skill_md.is_file():
|
||||||
|
meta = _skill_meta(skill_md)
|
||||||
|
skills.append(
|
||||||
|
{
|
||||||
|
"name": meta.get("name") or child.name,
|
||||||
|
"description": meta.get("description", ""),
|
||||||
|
"path": _relative(str(child), agent.cwd),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
out.append(
|
||||||
|
{
|
||||||
|
"set": directory.name,
|
||||||
|
"path": _relative(str(directory), agent.cwd),
|
||||||
|
"skills": skills,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def compose(
|
||||||
|
agent: ClaudeAgent,
|
||||||
|
kind: Kind,
|
||||||
|
entries: Iterable[Mapping[str, Any]],
|
||||||
|
*,
|
||||||
|
context_tokens: int,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
files, counts = files_touched(entries, cwd=agent.cwd)
|
||||||
|
prompt = granules(agent, kind)
|
||||||
|
skills = skill_sets(agent, kind)
|
||||||
|
skills_chars = sum(
|
||||||
|
len(s["name"]) + len(s["description"])
|
||||||
|
for group in skills
|
||||||
|
for s in group["skills"]
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"agent": {
|
||||||
|
"name": agent.name,
|
||||||
|
"model": agent.model,
|
||||||
|
"effort": agent.options.effort,
|
||||||
|
},
|
||||||
|
"kind": kind,
|
||||||
|
"context_tokens": context_tokens,
|
||||||
|
"prompt": prompt,
|
||||||
|
"skills": skills,
|
||||||
|
"skills_tokens_est": _tokens(skills_chars),
|
||||||
|
"tools": {
|
||||||
|
"allowed": list(agent.options.tools) if agent.options.tools else None,
|
||||||
|
"disallowed": list(agent.options.disallowed_tools),
|
||||||
|
"gateway": list(agent.gateway_tools),
|
||||||
|
"mcps": [m.name for m in agent.expose_mcps],
|
||||||
|
},
|
||||||
|
"files": files,
|
||||||
|
"tool_counts": counts,
|
||||||
|
"history_tokens_est": max(
|
||||||
|
0, context_tokens - prompt["tokens_est"] - _tokens(skills_chars)
|
||||||
|
),
|
||||||
|
}
|
||||||
@@ -1,13 +1,8 @@
|
|||||||
"""Closing a deep chat: the digest, the index, the file line cap (§6.4, §8.4).
|
"""Closing a deep chat: the digest, the index, the file line cap.
|
||||||
|
|
||||||
The gateway knows no path by itself (§0.8): ``Distiller`` from ``config.py``
|
``Distiller`` says where digests and the index live; the gateway checks the
|
||||||
names the agent, says where digests land and where the index lives, and
|
digest's frontmatter and caps the merge text. ``LineCap`` is the same idea
|
||||||
the distiller writes the file on its own. What the gateway does is check that a file
|
for a file a job rewrites: too long, and it is bounced back unchanged.
|
||||||
with a valid frontmatter appeared under ``Distiller.dir`` during the fork
|
|
||||||
turn, put one line into the index, and cap the merge text at
|
|
||||||
``SUMMARY_LINES``. ``LineCap`` is the same idea for a file a job rewrites
|
|
||||||
(``состояние.md``): a result longer than the cap is bounced - the file
|
|
||||||
goes back to what it was and the job is told to shorten.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -39,7 +34,7 @@ __all__ = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
SUMMARY_LINES = 5
|
SUMMARY_LINES = 5
|
||||||
"""A merge into the master is at most this many lines (§6.4)."""
|
"""A merge into the master is at most this many lines."""
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
@@ -49,10 +44,10 @@ class Distiller:
|
|||||||
agent: str
|
agent: str
|
||||||
dir: Path
|
dir: Path
|
||||||
index: Path
|
index: Path
|
||||||
type: str = "выжимка"
|
type: str = "digest"
|
||||||
"""Value the ``type`` frontmatter key must carry."""
|
"""Value the ``type`` frontmatter key must carry."""
|
||||||
|
|
||||||
index_header: str = "# индекс\n\nстрока на выжимку: чат → его выжимка.\n" # noqa: RUF001
|
index_header: str = "# index\n\none line per digest: chat → its digest.\n"
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
@@ -150,13 +145,13 @@ def check_digest(path: Path, digests: Distiller) -> Digest | str:
|
|||||||
try:
|
try:
|
||||||
post = frontmatter.load(str(path))
|
post = frontmatter.load(str(path))
|
||||||
except (OSError, ValueError) as exc:
|
except (OSError, ValueError) as exc:
|
||||||
return f"не читается: {exc}"
|
return f"unreadable: {exc}"
|
||||||
meta = post.metadata
|
meta = post.metadata
|
||||||
if meta.get("type") != digests.type:
|
if meta.get("type") != digests.type:
|
||||||
return f"`type` должен быть `{digests.type}`, не {meta.get('type')!r}"
|
return f"`type` must be `{digests.type}`, not {meta.get('type')!r}"
|
||||||
source = meta.get("source")
|
source = meta.get("source")
|
||||||
if not isinstance(source, str) or not source.strip():
|
if not isinstance(source, str) or not source.strip():
|
||||||
return "`source` пустой"
|
return "`source` is empty"
|
||||||
when = meta.get("date")
|
when = meta.get("date")
|
||||||
if isinstance(when, datetime):
|
if isinstance(when, datetime):
|
||||||
when = when.date()
|
when = when.date()
|
||||||
@@ -164,11 +159,11 @@ def check_digest(path: Path, digests: Distiller) -> Digest | str:
|
|||||||
try:
|
try:
|
||||||
when = date.fromisoformat(when.strip())
|
when = date.fromisoformat(when.strip())
|
||||||
except ValueError:
|
except ValueError:
|
||||||
return f"`date` не дата: {when!r}"
|
return f"`date` is not a date: {when!r}"
|
||||||
if not isinstance(when, date):
|
if not isinstance(when, date):
|
||||||
return "`date` отсутствует"
|
return "`date` is missing"
|
||||||
if not post.content.strip():
|
if not post.content.strip():
|
||||||
return "тело пустое"
|
return "empty body"
|
||||||
return Digest(path=path, source=source.strip(), date=when)
|
return Digest(path=path, source=source.strip(), date=when)
|
||||||
|
|
||||||
|
|
||||||
+50
-37
@@ -1,42 +1,32 @@
|
|||||||
"""The envelope (§3.3): a background block after the user's text.
|
"""The envelope: a background block under the user's text - clock, changes, recall."""
|
||||||
|
|
||||||
Assembled when the turn starts, never when the message is queued: the
|
|
||||||
time and what changed in the vault since the last envelope (added lines
|
|
||||||
for the ``full`` files, names and counts for the rest). Ceilings keep it
|
|
||||||
a signal, not a document; the injects that ride along are bundled below
|
|
||||||
it by the queue, with their own header.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass, field
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
from zoneinfo import ZoneInfo
|
from zoneinfo import ZoneInfo
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.texts import EnvelopeTexts
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Callable, Sequence
|
from collections.abc import Callable, Sequence
|
||||||
|
|
||||||
from beaver_gateway.core.watch import Change, VaultWatch
|
from beaver_gateway.vault.watch import Change, VaultWatch
|
||||||
|
|
||||||
__all__ = ["Envelope", "RecallContext", "render"]
|
__all__ = ["Envelope", "RecallContext", "render"]
|
||||||
|
|
||||||
_log = logging.getLogger(__name__)
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
HEADER = (
|
|
||||||
"[конверт - фоновый сигнал, не обращение; "
|
|
||||||
"реагируй, только если относится к вопросу]"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
class RecallContext:
|
class RecallContext:
|
||||||
"""What the setup's ``recall`` hook sees: the user's text and where it landed."""
|
|
||||||
|
|
||||||
text: str
|
text: str
|
||||||
kind: str
|
kind: str
|
||||||
now: datetime
|
now: datetime
|
||||||
|
agent: str = ""
|
||||||
|
"""Agent whose turn it is; a setup with several tells them apart by it."""
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
@@ -48,13 +38,16 @@ class Envelope:
|
|||||||
names_only_within: float = 600.0
|
names_only_within: float = 600.0
|
||||||
last_at: datetime | None = None
|
last_at: datetime | None = None
|
||||||
recall: Callable[[RecallContext], str | None] | None = None
|
recall: Callable[[RecallContext], str | None] | None = None
|
||||||
"""Setup-side lookup run on the user's text at turn start: pointers into
|
"""Setup-side lookup on the user's text; its lines go under the change block."""
|
||||||
the vault (a person's card, the agent's own notes, due dates) that the
|
texts: EnvelopeTexts = field(default_factory=EnvelopeTexts)
|
||||||
gateway cannot know the paths of. Its lines go under the vault block;
|
|
||||||
a failure is logged and the envelope goes out without them."""
|
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
self, *, now: datetime | None = None, text: str = "", kind: str = "master"
|
self,
|
||||||
|
*,
|
||||||
|
now: datetime | None = None,
|
||||||
|
text: str = "",
|
||||||
|
kind: str = "master",
|
||||||
|
agent: str = "",
|
||||||
) -> str:
|
) -> str:
|
||||||
now = now or datetime.now(UTC)
|
now = now or datetime.now(UTC)
|
||||||
changes = self.watch.take() if self.watch is not None else []
|
changes = self.watch.take() if self.watch is not None else []
|
||||||
@@ -70,23 +63,29 @@ class Envelope:
|
|||||||
names_only=names_only,
|
names_only=names_only,
|
||||||
max_lines=self.max_lines,
|
max_lines=self.max_lines,
|
||||||
per_file=self.per_file,
|
per_file=self.per_file,
|
||||||
|
texts=self.texts,
|
||||||
)
|
)
|
||||||
self.last_at = now
|
self.last_at = now
|
||||||
block = self.recall_block(text=text, kind=kind, now=now)
|
block = self.recall_block(text=text, kind=kind, now=now, agent=agent)
|
||||||
return f"{out}\n{block}" if block else out
|
return f"{out}\n{block}" if block else out
|
||||||
|
|
||||||
def recall_only(
|
def recall_only(
|
||||||
self, *, text: str, kind: str, now: datetime | None = None
|
self, *, text: str, kind: str, now: datetime | None = None, agent: str = ""
|
||||||
) -> str | None:
|
) -> str | None:
|
||||||
"""The recall lines under the header, without the vault diff (branches)."""
|
block = self.recall_block(
|
||||||
block = self.recall_block(text=text, kind=kind, now=now or datetime.now(UTC))
|
text=text, kind=kind, now=now or datetime.now(UTC), agent=agent
|
||||||
return f"{HEADER}\n{block}" if block else None
|
)
|
||||||
|
return f"{self.texts.header}\n{block}" if block else None
|
||||||
|
|
||||||
def recall_block(self, *, text: str, kind: str, now: datetime) -> str | None:
|
def recall_block(
|
||||||
|
self, *, text: str, kind: str, now: datetime, agent: str = ""
|
||||||
|
) -> str | None:
|
||||||
if self.recall is None or not text.strip():
|
if self.recall is None or not text.strip():
|
||||||
return None
|
return None
|
||||||
try:
|
try:
|
||||||
block = self.recall(RecallContext(text=text, kind=kind, now=now))
|
block = self.recall(
|
||||||
|
RecallContext(text=text, kind=kind, now=now, agent=agent)
|
||||||
|
)
|
||||||
except Exception: # noqa: BLE001
|
except Exception: # noqa: BLE001
|
||||||
_log.exception("recall hook failed")
|
_log.exception("recall hook failed")
|
||||||
return None
|
return None
|
||||||
@@ -102,38 +101,52 @@ def render(
|
|||||||
names_only: bool,
|
names_only: bool,
|
||||||
max_lines: int = 120,
|
max_lines: int = 120,
|
||||||
per_file: int = 30,
|
per_file: int = 30,
|
||||||
|
texts: EnvelopeTexts | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
|
texts = texts or EnvelopeTexts()
|
||||||
zone = ZoneInfo(tz)
|
zone = ZoneInfo(tz)
|
||||||
stamp = now.astimezone(zone)
|
stamp = now.astimezone(zone)
|
||||||
lines = [HEADER, f"время: {stamp:%Y-%m-%d %H:%M} ({_zone_label(tz)})"]
|
lines = [
|
||||||
|
texts.header,
|
||||||
|
texts.time.format(stamp=f"{stamp:%Y-%m-%d %H:%M}", zone=_zone_label(tz)),
|
||||||
|
]
|
||||||
ordered = sorted(changes, key=lambda c: (not c.full, c.path))
|
ordered = sorted(changes, key=lambda c: (not c.full, c.path))
|
||||||
since_label = (
|
since_label = (
|
||||||
f"с {since.astimezone(zone):%H:%M}" if since is not None else "со старта" # noqa: RUF001
|
texts.since.format(time=f"{since.astimezone(zone):%H:%M}")
|
||||||
|
if since is not None
|
||||||
|
else texts.since_start
|
||||||
)
|
)
|
||||||
if ordered:
|
if ordered:
|
||||||
names = ", ".join(f"{c.path} (+{c.added_count})" for c in ordered)
|
names = ", ".join(f"{c.path} (+{c.added_count})" for c in ordered)
|
||||||
lines.append(f"vault, изменено {since_label}: {names}")
|
lines.append(texts.changed.format(since=since_label, names=names))
|
||||||
if not names_only:
|
if not names_only:
|
||||||
_append_diffs(lines, ordered, max_lines=max_lines, per_file=per_file)
|
_append_diffs(
|
||||||
|
lines, ordered, max_lines=max_lines, per_file=per_file, texts=texts
|
||||||
|
)
|
||||||
return "\n".join(lines[:max_lines])
|
return "\n".join(lines[:max_lines])
|
||||||
|
|
||||||
|
|
||||||
def _append_diffs(
|
def _append_diffs(
|
||||||
lines: list[str], changes: Sequence[Change], *, max_lines: int, per_file: int
|
lines: list[str],
|
||||||
|
changes: Sequence[Change],
|
||||||
|
*,
|
||||||
|
max_lines: int,
|
||||||
|
per_file: int,
|
||||||
|
texts: EnvelopeTexts,
|
||||||
) -> None:
|
) -> None:
|
||||||
budget = max_lines - len(lines) - 1
|
budget = max_lines - len(lines) - 1
|
||||||
for change in changes:
|
for change in changes:
|
||||||
if not change.full or not change.added:
|
if not change.full or not change.added:
|
||||||
continue
|
continue
|
||||||
if budget < 3:
|
if budget < 3:
|
||||||
lines.append("… (потолок конверта)")
|
lines.append(texts.truncated)
|
||||||
return
|
return
|
||||||
shown = change.added[: min(per_file, budget - 2)]
|
shown = change.added[: min(per_file, budget - 2)]
|
||||||
lines.append(f"--- {change.path}, только добавленное ---")
|
lines.append(texts.file_header.format(path=change.path))
|
||||||
lines.extend(f"+ {line}" for line in shown)
|
lines.extend(f"+ {line}" for line in shown)
|
||||||
budget -= 1 + len(shown)
|
budget -= 1 + len(shown)
|
||||||
if len(change.added) > len(shown):
|
if len(change.added) > len(shown):
|
||||||
lines.append(f"+ … ещё {len(change.added) - len(shown)}")
|
lines.append(texts.more_lines.format(count=len(change.added) - len(shown)))
|
||||||
budget -= 1
|
budget -= 1
|
||||||
|
|
||||||
|
|
||||||
@@ -1,20 +1,15 @@
|
|||||||
"""Persisted per-conversation queue, ``urgent > user > wake > normal`` (§3.4).
|
"""Persisted per-conversation queue: ``urgent > user > wake > normal``.
|
||||||
|
|
||||||
One ``ClaudeSDKClient`` runs one turn at a time, so ordering has to happen
|
conversations/service.py runs one worker per conversation over these rows.
|
||||||
before the client: the rows here are the queue, ``core/conversations``
|
A row still ``running`` when the gateway starts was cut by a restart and is
|
||||||
runs one worker per conversation over them. ``urgent`` cuts a running
|
flagged ``interrupted``, never re-run.
|
||||||
turn, ``user`` is the human, ``wake`` starts a turn as soon as the
|
|
||||||
conversation is idle and takes the queued normals with it, ``normal``
|
|
||||||
waits for the batching window or rides with the next turn. 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 __future__ import annotations
|
||||||
|
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from typing import TYPE_CHECKING, Literal, cast
|
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||||
|
|
||||||
from sqlalchemy import func
|
from sqlalchemy import func
|
||||||
from sqlmodel import col, select
|
from sqlmodel import col, select
|
||||||
@@ -42,8 +37,8 @@ URGENCY: tuple[Priority, ...] = ("normal", "wake", "urgent")
|
|||||||
"""What the API and the tools accept for ``urgency``: every priority but ``user``."""
|
"""What the API and the tools accept for ``urgency``: every priority but ``user``."""
|
||||||
|
|
||||||
INTERRUPTED_TURN = (
|
INTERRUPTED_TURN = (
|
||||||
"[этот инжект прервал предыдущий тёрн: «Request interrupted» выше - "
|
"[this inject cut the previous turn: the «Request interrupted» above is "
|
||||||
"прерывание, не отказ от тулзы]"
|
"an interruption, not a refused tool call]"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -58,9 +53,7 @@ class InjectContext:
|
|||||||
|
|
||||||
def inject_header(ctx: InjectContext) -> str:
|
def inject_header(ctx: InjectContext) -> str:
|
||||||
"""Default framing; a setup overrides it via ``ConversationTexts.inject_header``."""
|
"""Default framing; a setup overrides it via ``ConversationTexts.inject_header``."""
|
||||||
head = (
|
head = f"[inject: {ctx.origin} - not the user, no reply needed]"
|
||||||
f"[инжект: {ctx.origin} - это не Бобёр, отвечать не нужно, голос не обязателен]"
|
|
||||||
)
|
|
||||||
return f"{head}\n{INTERRUPTED_TURN}" if ctx.interrupted_turn else head
|
return f"{head}\n{INTERRUPTED_TURN}" if ctx.interrupted_turn else head
|
||||||
|
|
||||||
|
|
||||||
@@ -77,13 +70,23 @@ class InjectQueue:
|
|||||||
self._db = db
|
self._db = db
|
||||||
|
|
||||||
async def push(
|
async def push(
|
||||||
self, *, conversation_id: int, priority: Priority, origin: str, text: str
|
self,
|
||||||
|
*,
|
||||||
|
conversation_id: int,
|
||||||
|
priority: Priority,
|
||||||
|
origin: str,
|
||||||
|
text: str,
|
||||||
|
attachments: list[dict[str, Any]] | None = None,
|
||||||
) -> InjectQueueItem:
|
) -> InjectQueueItem:
|
||||||
if priority not in PRIORITY_RANK:
|
if priority not in PRIORITY_RANK:
|
||||||
msg = f"unknown priority {priority!r}"
|
msg = f"unknown priority {priority!r}"
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
row = InjectQueueItem(
|
row = InjectQueueItem(
|
||||||
conversation_id=conversation_id, priority=priority, origin=origin, text=text
|
conversation_id=conversation_id,
|
||||||
|
priority=priority,
|
||||||
|
origin=origin,
|
||||||
|
text=text,
|
||||||
|
attachments=attachments or None,
|
||||||
)
|
)
|
||||||
async with self._db.session() as session:
|
async with self._db.session() as session:
|
||||||
session.add(row)
|
session.add(row)
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
"""Conversation kinds (§3.1) as one closed type for agents, frontends, service."""
|
"""Conversation kinds as one closed type for agents, frontends, service."""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
@@ -0,0 +1,156 @@
|
|||||||
|
"""Putting words into a conversation: a message, an inject, ``say``, a file."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.turns import Turns
|
||||||
|
from beaver_gateway.frontends.files import SendFileError
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.injects import Priority
|
||||||
|
from beaver_gateway.storage.models import Conversation, InjectQueueItem
|
||||||
|
|
||||||
|
__all__ = ["Messaging"]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class Messaging(Turns):
|
||||||
|
async def post(
|
||||||
|
self,
|
||||||
|
conv: Conversation,
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
origin: str = "user",
|
||||||
|
attachments: list[dict[str, Any]] | None = None,
|
||||||
|
) -> InjectQueueItem:
|
||||||
|
item = await self._queue.push(
|
||||||
|
conversation_id=cast("int", conv.id),
|
||||||
|
priority="user",
|
||||||
|
origin=origin,
|
||||||
|
text=text,
|
||||||
|
attachments=attachments,
|
||||||
|
)
|
||||||
|
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 inject(
|
||||||
|
self,
|
||||||
|
conv: Conversation,
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
urgency: Priority = "normal",
|
||||||
|
origin: str = "system",
|
||||||
|
interrupt: bool = True,
|
||||||
|
) -> InjectQueueItem:
|
||||||
|
if conv.kind == "master" and conv.status != "open":
|
||||||
|
live = await self.open_master()
|
||||||
|
if live is None:
|
||||||
|
_log.error(
|
||||||
|
"inject (%s) for closed master %s: no open master, it stays there",
|
||||||
|
origin,
|
||||||
|
conv.external_id,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
_log.info(
|
||||||
|
"inject (%s) for closed master %s goes to %s",
|
||||||
|
origin,
|
||||||
|
conv.external_id,
|
||||||
|
live.external_id,
|
||||||
|
)
|
||||||
|
conv = live
|
||||||
|
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" and interrupt:
|
||||||
|
backend = self._backend(conv.agent_name)
|
||||||
|
if await backend.interrupt(conv.external_id):
|
||||||
|
_log.info(
|
||||||
|
"conversation %s: interrupted for urgent inject", conv.external_id
|
||||||
|
)
|
||||||
|
await self._queue.mark_interrupting(item)
|
||||||
|
self._ensure_worker(cast("int", conv.id))
|
||||||
|
return item
|
||||||
|
|
||||||
|
async def interrupt(self, conv: Conversation) -> bool:
|
||||||
|
"""Cut the running turn; the reply so far still lands."""
|
||||||
|
try:
|
||||||
|
backend = self._backend(conv.agent_name)
|
||||||
|
except LookupError:
|
||||||
|
return False
|
||||||
|
return await backend.interrupt(conv.external_id)
|
||||||
|
|
||||||
|
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 send_file(
|
||||||
|
self, conv: Conversation, path: str, *, caption: str = "", method: str = "auto"
|
||||||
|
) -> str:
|
||||||
|
"""Hand a file to the human: this conversation's window, else the master's."""
|
||||||
|
for target in (conv, await self.open_master()):
|
||||||
|
if target is None:
|
||||||
|
continue
|
||||||
|
for fe in self._frontends:
|
||||||
|
if target.kind in fe.kinds and fe.sends_files:
|
||||||
|
note = await fe.send_file(
|
||||||
|
target, path, caption=caption, method=method
|
||||||
|
)
|
||||||
|
_log.info("send_file[%s]: %s", target.external_id, note)
|
||||||
|
return note
|
||||||
|
msg = (
|
||||||
|
"nowhere to send: no frontend that sends files shows this "
|
||||||
|
"conversation, and there is no open master"
|
||||||
|
)
|
||||||
|
raise SendFileError(msg)
|
||||||
|
|
||||||
|
def file_note(self) -> str:
|
||||||
|
"""What the ``send_file`` tool tells the model about roots and size."""
|
||||||
|
sender = next((f for f in self._frontends if f.sends_files), None)
|
||||||
|
return sender.file_note() if sender else "no frontend here sends files"
|
||||||
|
|
||||||
|
async def schedule(
|
||||||
|
self,
|
||||||
|
conv: Conversation,
|
||||||
|
at: str,
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
urgency: Priority = "wake",
|
||||||
|
dedupe_key: str | None = None,
|
||||||
|
) -> tuple[int | None, datetime]:
|
||||||
|
if self.scheduler is None:
|
||||||
|
msg = "no scheduler; `schedule` is unavailable"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
return await self.scheduler.schedule(
|
||||||
|
conv, at, text, urgency=urgency, dedupe_key=dedupe_key
|
||||||
|
)
|
||||||
|
|
||||||
|
async def schedules(self, conv: Conversation | None = None) -> list[dict[str, Any]]:
|
||||||
|
return await self.scheduler.scheduled(conv) if self.scheduler else []
|
||||||
@@ -0,0 +1,85 @@
|
|||||||
|
"""``AskUserQuestion`` from inside a turn: answered or timed out."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import contextlib
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.rewinding import Rewinding
|
||||||
|
from beaver_gateway.conversations.state import Question
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = ["Questions"]
|
||||||
|
|
||||||
|
|
||||||
|
class Questions(Rewinding):
|
||||||
|
async def ask(self, key: str, payload: dict[str, Any]) -> str | None:
|
||||||
|
"""Show the question and wait; ``None`` when nobody answered in time."""
|
||||||
|
conv = await self.get(key)
|
||||||
|
if conv is None:
|
||||||
|
return None
|
||||||
|
runner = self._runners.get(cast("int", conv.id))
|
||||||
|
question_id = f"q_{uuid4().hex[:10]}"
|
||||||
|
pending = Question(
|
||||||
|
conversation_id=key,
|
||||||
|
turn_id=runner.turn_id if runner is not None else None,
|
||||||
|
questions=list(payload.get("questions") or []),
|
||||||
|
answer=asyncio.get_running_loop().create_future(),
|
||||||
|
)
|
||||||
|
self._questions[question_id] = pending
|
||||||
|
await self._set_pending_question(conv, value=True)
|
||||||
|
self._bus.publish(
|
||||||
|
"question",
|
||||||
|
conversation_id=key,
|
||||||
|
turn_id=pending.turn_id,
|
||||||
|
question_id=question_id,
|
||||||
|
questions=pending.questions,
|
||||||
|
timeout=self._question_timeout,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
async with asyncio.timeout(self._question_timeout):
|
||||||
|
answer = await pending.answer
|
||||||
|
except TimeoutError:
|
||||||
|
self._bus.publish(
|
||||||
|
"question.timeout",
|
||||||
|
conversation_id=key,
|
||||||
|
turn_id=pending.turn_id,
|
||||||
|
question_id=question_id,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
finally:
|
||||||
|
self._questions.pop(question_id, None)
|
||||||
|
await self._set_pending_question(conv, value=False)
|
||||||
|
self._bus.publish(
|
||||||
|
"question.answered",
|
||||||
|
conversation_id=key,
|
||||||
|
turn_id=pending.turn_id,
|
||||||
|
question_id=question_id,
|
||||||
|
answer=answer,
|
||||||
|
)
|
||||||
|
return answer
|
||||||
|
|
||||||
|
def answer(self, question_id: str, answer: str) -> bool:
|
||||||
|
pending = self._questions.get(question_id)
|
||||||
|
if pending is None or pending.answer.done():
|
||||||
|
return False
|
||||||
|
pending.answer.set_result(answer)
|
||||||
|
return True
|
||||||
|
|
||||||
|
def answer_text(self, answer: str | None) -> str:
|
||||||
|
if answer is None:
|
||||||
|
return self._texts.unanswered.format(
|
||||||
|
minutes=round(self._question_timeout / 60)
|
||||||
|
)
|
||||||
|
return self._texts.answered.format(answer=answer)
|
||||||
|
|
||||||
|
async def _set_pending_question(self, conv: Conversation, *, value: bool) -> None:
|
||||||
|
async def apply(row: Conversation) -> None:
|
||||||
|
row.pending_question = value
|
||||||
|
|
||||||
|
with contextlib.suppress(LookupError):
|
||||||
|
await self._update(conv, apply)
|
||||||
@@ -0,0 +1,114 @@
|
|||||||
|
"""Taking turns back: the thread moves to a fork that ends before them."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
|
from claude_agent_sdk import fork_session_via_store
|
||||||
|
|
||||||
|
from beaver_gateway.backends.transcript import (
|
||||||
|
last_message_before,
|
||||||
|
messages_from_entries,
|
||||||
|
reply_text,
|
||||||
|
text_of,
|
||||||
|
turn_starts,
|
||||||
|
)
|
||||||
|
from beaver_gateway.conversations.spawning import Spawning
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Sequence
|
||||||
|
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = ["Rewinding", "Rewound"]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
PREVIEW = 280
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class Rewound:
|
||||||
|
conversation: Conversation
|
||||||
|
turns: int
|
||||||
|
previous_session: str
|
||||||
|
|
||||||
|
|
||||||
|
class Rewinding(Spawning):
|
||||||
|
async def turns(
|
||||||
|
self, conv: Conversation, *, limit: int = 10
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
"""The latest turns that can be taken back, oldest first; the first stays."""
|
||||||
|
entries = await self.entries(conv)
|
||||||
|
starts = turn_starts(entries)
|
||||||
|
ends = [*starts[1:], len(entries)]
|
||||||
|
spans = list(zip(starts, ends, strict=True))[1:]
|
||||||
|
agent = self._agents.get(conv.agent_name)
|
||||||
|
progress = agent is not None and agent.thinking_is_progress
|
||||||
|
return [
|
||||||
|
self._turn(entries[start:end], thinking_is_progress=progress)
|
||||||
|
for start, end in spans[-limit:]
|
||||||
|
]
|
||||||
|
|
||||||
|
def _turn(
|
||||||
|
self, entries: Sequence[Any], *, thinking_is_progress: bool
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
prompt, *rest = messages_from_entries(entries)
|
||||||
|
replies = (
|
||||||
|
reply_text(m["content"], thinking_is_progress=thinking_is_progress)
|
||||||
|
for m in reversed(rest)
|
||||||
|
if m["role"] == "assistant"
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"id": entries[0].get("uuid"),
|
||||||
|
"ts": entries[0].get("timestamp"),
|
||||||
|
"text": text_of(self._unwrapped(prompt)["content"])[:PREVIEW],
|
||||||
|
"reply": next((text for text in replies if text), "")[:PREVIEW],
|
||||||
|
"failed": any(e.get("isApiErrorMessage") for e in entries),
|
||||||
|
}
|
||||||
|
|
||||||
|
async def rewind(self, conv: Conversation, turn: str) -> Rewound:
|
||||||
|
"""Take ``turn`` and every turn after it out of what the model sees.
|
||||||
|
|
||||||
|
Nothing is deleted: the conversation moves to a fork of its session
|
||||||
|
that ends on the message before ``turn``, the cut session stays in
|
||||||
|
the store under its old id.
|
||||||
|
"""
|
||||||
|
row_id = cast("int", conv.id)
|
||||||
|
runner = self._runner(row_id)
|
||||||
|
if runner.lock.locked():
|
||||||
|
msg = f"conversation {conv.external_id} is mid-turn; interrupt it first"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
async with runner.lock:
|
||||||
|
row = await self.get_row(row_id) or conv
|
||||||
|
entries = await self.entries(row)
|
||||||
|
starts = turn_starts(entries)
|
||||||
|
cut = next((i for i in starts[1:] if entries[i].get("uuid") == turn), None)
|
||||||
|
keep = last_message_before(entries, cut) if cut is not None else None
|
||||||
|
if cut is None or keep is None:
|
||||||
|
msg = f"conversation {row.external_id} has no turn {turn} to take back"
|
||||||
|
raise LookupError(msg)
|
||||||
|
previous = cast("str", row.session_id)
|
||||||
|
forked = await fork_session_via_store(
|
||||||
|
self._store,
|
||||||
|
previous,
|
||||||
|
directory=str(self._claude_agent(row.agent_name).cwd),
|
||||||
|
up_to_message_id=keep,
|
||||||
|
)
|
||||||
|
await self._backend(row.agent_name).close(row.external_id)
|
||||||
|
|
||||||
|
async def apply(moved: Conversation) -> None:
|
||||||
|
moved.session_id = forked.session_id
|
||||||
|
|
||||||
|
row = await self._update(row, apply)
|
||||||
|
turns = sum(1 for start in starts if start >= cut)
|
||||||
|
_log.info(
|
||||||
|
"rewound %s by %d turns: session %s -> %s",
|
||||||
|
row.external_id,
|
||||||
|
turns,
|
||||||
|
previous,
|
||||||
|
forked.session_id,
|
||||||
|
)
|
||||||
|
return Rewound(conversation=row, turns=turns, previous_session=previous)
|
||||||
+20
-14
@@ -1,11 +1,7 @@
|
|||||||
"""Master rotation (§4.5, §8.1, §8.3).
|
"""Master rotation: one logical master thread, many physical sessions.
|
||||||
|
|
||||||
One logical master thread, many physical sessions: when the policy says
|
When the policy says so, a new master takes over the window; the old one
|
||||||
so, a new master is spawned and takes over the window atomically, the old
|
hands out, closes, and its branches and queued normal injects move over.
|
||||||
one writes its handout as its last turn, closes, its finished branches get
|
|
||||||
marked in their windows, its queued normal injects move over, and the new
|
|
||||||
one receives "new day". Silence is measured by the user's messages only -
|
|
||||||
injects never extend a day.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -17,12 +13,12 @@ from typing import TYPE_CHECKING
|
|||||||
from zoneinfo import ZoneInfo
|
from zoneinfo import ZoneInfo
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from beaver_gateway.core.conversations import Conversations
|
from beaver_gateway.conversations.service import Conversations
|
||||||
from beaver_gateway.storage.models import Conversation
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
__all__ = ["HandoutContext", "Rotation", "RotationPolicy"]
|
__all__ = ["HandoutContext", "Rotation", "RotationPolicy"]
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.core.rotation")
|
_log = logging.getLogger("beaver_gateway.conversations.rotation")
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
@@ -46,11 +42,11 @@ class RotationPolicy:
|
|||||||
if now.astimezone(zone) < boundary:
|
if now.astimezone(zone) < boundary:
|
||||||
boundary -= timedelta(days=1)
|
boundary -= timedelta(days=1)
|
||||||
if created < boundary and silence > self.night_silence:
|
if created < boundary and silence > self.night_silence:
|
||||||
return "ночь"
|
return "night"
|
||||||
if now - created > self.max_age and silence > self.short_silence:
|
if now - created > self.max_age and silence > self.short_silence:
|
||||||
return "возраст"
|
return "age"
|
||||||
if context_tokens > self.max_context_tokens and silence > self.short_silence:
|
if context_tokens > self.max_context_tokens and silence > self.short_silence:
|
||||||
return "транскрипт"
|
return "context"
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def day_of(self, master: Conversation) -> date:
|
def day_of(self, master: Conversation) -> date:
|
||||||
@@ -81,9 +77,19 @@ class Rotation:
|
|||||||
out.append((master, reason))
|
out.append((master, reason))
|
||||||
return out
|
return out
|
||||||
|
|
||||||
async def tick(self, now: datetime | None = None) -> list[Conversation]:
|
async def tick(
|
||||||
|
self, now: datetime | None = None, *, force: str | None = None
|
||||||
|
) -> list[Conversation]:
|
||||||
|
"""Rotate the masters that are due; ``force`` rotates every open master."""
|
||||||
rotated: list[Conversation] = []
|
rotated: list[Conversation] = []
|
||||||
for master, reason in await self.due(now):
|
if force:
|
||||||
|
due = [
|
||||||
|
(m, force)
|
||||||
|
for m in await self._conversations.find(kind="master", status="open")
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
due = await self.due(now)
|
||||||
|
for master, reason in due:
|
||||||
new = await self.rotate(master, reason)
|
new = await self.rotate(master, reason)
|
||||||
if new is not None:
|
if new is not None:
|
||||||
rotated.append(new)
|
rotated.append(new)
|
||||||
@@ -0,0 +1,594 @@
|
|||||||
|
"""The conversation rows: create, find, bind to windows, flags, status, history."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from sqlalchemy import String, func
|
||||||
|
from sqlalchemy import cast as sa_cast
|
||||||
|
from sqlmodel import col, select
|
||||||
|
|
||||||
|
from beaver_gateway.backends.transcript import (
|
||||||
|
messages_from_entries,
|
||||||
|
render_messages,
|
||||||
|
text_of,
|
||||||
|
)
|
||||||
|
from beaver_gateway.conversations.kinds import KINDS, Kind
|
||||||
|
from beaver_gateway.conversations.state import State, iso
|
||||||
|
from beaver_gateway.conversations.texts import split_service
|
||||||
|
from beaver_gateway.frontends.markdown.history import load_messages
|
||||||
|
from beaver_gateway.storage.models import (
|
||||||
|
Conversation,
|
||||||
|
ConversationBinding,
|
||||||
|
ConversationMessage,
|
||||||
|
RateLimit,
|
||||||
|
TranscriptEntry,
|
||||||
|
Usage,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Awaitable, Callable, Iterable
|
||||||
|
|
||||||
|
from beaver_gateway.frontends.base import Frontend
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"MASTER_ALIAS",
|
||||||
|
"PARENT_ALIAS",
|
||||||
|
"STATUSES",
|
||||||
|
"TITLE_MAX",
|
||||||
|
"Rows",
|
||||||
|
"context_of",
|
||||||
|
"implied_title",
|
||||||
|
]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
MASTER_ALIAS = "master"
|
||||||
|
PARENT_ALIAS = "parent"
|
||||||
|
STATUSES = ("open", "merged", "closed", "archived")
|
||||||
|
TITLE_MAX = 80
|
||||||
|
|
||||||
|
|
||||||
|
class Rows(State):
|
||||||
|
async def create(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
kind: Kind,
|
||||||
|
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)
|
||||||
|
if not self._claude_agent(agent).serves(kind):
|
||||||
|
msg = f"agent {agent!r} does not serve kind {kind!r}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
row = Conversation(
|
||||||
|
frontend=origin,
|
||||||
|
external_id=str(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 resolve(
|
||||||
|
self, key: str, *, origin: Conversation | None = None
|
||||||
|
) -> Conversation | None:
|
||||||
|
"""By public id, or ``master`` / ``parent`` relative to ``origin``."""
|
||||||
|
key = key.strip()
|
||||||
|
if key == MASTER_ALIAS:
|
||||||
|
return await self.open_master()
|
||||||
|
if key == PARENT_ALIAS:
|
||||||
|
if origin is None or origin.parent_id is None:
|
||||||
|
return None
|
||||||
|
return await self.get_row(origin.parent_id)
|
||||||
|
return await self.get(key)
|
||||||
|
|
||||||
|
async def open_master(self) -> Conversation | None:
|
||||||
|
masters = await self.find(kind="master", status="open", limit=1)
|
||||||
|
return masters[0] if masters else None
|
||||||
|
|
||||||
|
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 search(self, query: str, *, limit: int = 20) -> list[Conversation]:
|
||||||
|
needle = f"%{query}%"
|
||||||
|
by_title = select(Conversation.id).where(col(Conversation.title).ilike(needle))
|
||||||
|
by_first = select(ConversationMessage.conversation_id).where(
|
||||||
|
ConversationMessage.seq == 0,
|
||||||
|
col(ConversationMessage.content_json).ilike(needle),
|
||||||
|
)
|
||||||
|
by_entry = (
|
||||||
|
select(Conversation.id)
|
||||||
|
.join(
|
||||||
|
TranscriptEntry,
|
||||||
|
col(TranscriptEntry.session_id) == col(Conversation.session_id),
|
||||||
|
)
|
||||||
|
.where(sa_cast(col(TranscriptEntry.entry), String).ilike(needle))
|
||||||
|
)
|
||||||
|
stmt = (
|
||||||
|
select(Conversation)
|
||||||
|
.where(col(Conversation.id).in_(by_title.union(by_first, by_entry)))
|
||||||
|
.order_by(col(Conversation.id).desc())
|
||||||
|
.limit(limit)
|
||||||
|
)
|
||||||
|
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:
|
||||||
|
if conv.kind not in self.frontend(frontend).kinds:
|
||||||
|
msg = f"frontend {frontend!r} does not show kind {conv.kind!r}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
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)
|
||||||
|
same_window = await session.exec(
|
||||||
|
select(ConversationBinding).where(
|
||||||
|
ConversationBinding.frontend == frontend,
|
||||||
|
ConversationBinding.external_id == external_id,
|
||||||
|
ConversationBinding.conversation_id != conv.id,
|
||||||
|
col(ConversationBinding.visible).is_(True),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
for other in same_window.all():
|
||||||
|
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 last_binding(
|
||||||
|
self, *, frontend: str, kind: str
|
||||||
|
) -> ConversationBinding | None:
|
||||||
|
"""The window ``frontend`` last used for ``kind``; outlives a rotation."""
|
||||||
|
async with self._db.session() as session:
|
||||||
|
result = await session.exec(
|
||||||
|
select(ConversationBinding)
|
||||||
|
.join(
|
||||||
|
Conversation,
|
||||||
|
col(Conversation.id) == col(ConversationBinding.conversation_id),
|
||||||
|
)
|
||||||
|
.where(
|
||||||
|
ConversationBinding.frontend == frontend, Conversation.kind == kind
|
||||||
|
)
|
||||||
|
.order_by(col(ConversationBinding.id).desc())
|
||||||
|
)
|
||||||
|
return result.first()
|
||||||
|
|
||||||
|
async def window_of(self, conv: Conversation) -> str | None:
|
||||||
|
for binding in await self.bindings(conv):
|
||||||
|
if binding.visible:
|
||||||
|
return binding.external_id
|
||||||
|
return None
|
||||||
|
|
||||||
|
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 reparent(self, conv: Conversation, parent: Conversation) -> Conversation:
|
||||||
|
async def apply(row: Conversation) -> None:
|
||||||
|
row.parent_id = parent.id
|
||||||
|
|
||||||
|
return await self._update(conv, apply)
|
||||||
|
|
||||||
|
async def _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)
|
||||||
|
out["title"] = await self.implied_title(conv)
|
||||||
|
out["context_tokens"] = await self.context_tokens(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)
|
||||||
|
agent = self._agents.get(conv.agent_name)
|
||||||
|
out["thinking_is_progress"] = agent is not None and agent.thinking_is_progress
|
||||||
|
out["live"] = live is not None
|
||||||
|
out["busy"] = live.busy if live is not None else False
|
||||||
|
runner = self._runners.get(cast("int", conv.id))
|
||||||
|
out["turn"] = runner.snapshot() if runner is not None else None
|
||||||
|
pending = self.pending_question(conv.external_id)
|
||||||
|
out["question"] = (
|
||||||
|
{"id": pending[0], "questions": pending[1]} if pending else None
|
||||||
|
)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def pending_question(self, key: str) -> tuple[str, list[dict[str, Any]]] | None:
|
||||||
|
for question_id, pending in self._questions.items():
|
||||||
|
if pending.conversation_id == key and not pending.answer.done():
|
||||||
|
return question_id, pending.questions
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def rate_limits(self, *, limit: int = 100) -> list[RateLimit]:
|
||||||
|
async with self._db.session() as session:
|
||||||
|
result = await session.exec(
|
||||||
|
select(RateLimit).order_by(col(RateLimit.id).desc()).limit(limit)
|
||||||
|
)
|
||||||
|
return list(result.all())
|
||||||
|
|
||||||
|
async def context_tokens(self, conv: Conversation) -> int:
|
||||||
|
async with self._db.session() as session:
|
||||||
|
row = (
|
||||||
|
await session.exec(
|
||||||
|
select(Usage)
|
||||||
|
.where(Usage.conversation_id == conv.external_id)
|
||||||
|
.order_by(col(Usage.id).desc())
|
||||||
|
.limit(1)
|
||||||
|
)
|
||||||
|
).first()
|
||||||
|
return context_of(row)
|
||||||
|
|
||||||
|
async def context_tokens_by_id(self, ids: Iterable[str]) -> dict[str, int]:
|
||||||
|
wanted = list(ids)
|
||||||
|
if not wanted:
|
||||||
|
return {}
|
||||||
|
latest = (
|
||||||
|
select(func.max(Usage.id))
|
||||||
|
.where(col(Usage.conversation_id).in_(wanted))
|
||||||
|
.group_by(Usage.conversation_id)
|
||||||
|
)
|
||||||
|
async with self._db.session() as session:
|
||||||
|
rows = (
|
||||||
|
await session.exec(select(Usage).where(col(Usage.id).in_(latest)))
|
||||||
|
).all()
|
||||||
|
return {
|
||||||
|
cast("str", row.conversation_id): context_of(row)
|
||||||
|
for row in rows
|
||||||
|
if row.conversation_id
|
||||||
|
}
|
||||||
|
|
||||||
|
async def usage_tokens(self, since: datetime) -> int:
|
||||||
|
async with self._db.session() as session:
|
||||||
|
rows = (
|
||||||
|
await session.exec(
|
||||||
|
select(Usage).where(
|
||||||
|
col(Usage.ts) >= since.astimezone(UTC).replace(tzinfo=None)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
).all()
|
||||||
|
return sum(
|
||||||
|
r.input_tokens + r.output_tokens + r.cache_creation_tokens for r in rows
|
||||||
|
)
|
||||||
|
|
||||||
|
async def busy(self, conv: Conversation) -> bool:
|
||||||
|
row = await self.get_row(cast("int", conv.id)) or conv
|
||||||
|
if row.running_turn or row.pending_question:
|
||||||
|
return True
|
||||||
|
live = self._pool.get(row.external_id)
|
||||||
|
if live is not None and live.busy:
|
||||||
|
return True
|
||||||
|
pending = await self._queue.pending(cast("int", row.id))
|
||||||
|
return any(i.priority in ("user", "urgent", "wake") for i in pending)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def frontends(self) -> list[Frontend]:
|
||||||
|
return list(self._frontends)
|
||||||
|
|
||||||
|
def frontend(self, name: str) -> Frontend:
|
||||||
|
for fe in self._frontends:
|
||||||
|
if fe.name == name:
|
||||||
|
return fe
|
||||||
|
msg = f"unknown frontend {name!r}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
|
||||||
|
def default_agent(self, kind: Kind) -> str | None:
|
||||||
|
for fe in self._frontends:
|
||||||
|
if kind in fe.kinds and (agent := fe.agent_for(kind)):
|
||||||
|
return agent
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def materialize(self, conv: Conversation) -> ConversationBinding | None:
|
||||||
|
for fe in self._frontends:
|
||||||
|
if conv.kind not in fe.kinds:
|
||||||
|
continue
|
||||||
|
binding = await fe.materialize(conv)
|
||||||
|
if binding is not None:
|
||||||
|
return binding
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def mark_closed(self, conv: Conversation) -> bool:
|
||||||
|
marked = False
|
||||||
|
for fe in self._frontends:
|
||||||
|
if conv.kind in fe.kinds:
|
||||||
|
try:
|
||||||
|
marked = await fe.mark_closed(conv) or marked
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("%s could not mark %s", fe.name, conv.external_id)
|
||||||
|
return marked
|
||||||
|
|
||||||
|
async def read(self, conv: Conversation, *, window: int | None = None) -> str:
|
||||||
|
return render_messages(
|
||||||
|
await self.history(conv),
|
||||||
|
thinking_is_progress=self._claude_agent(
|
||||||
|
conv.agent_name
|
||||||
|
).thinking_is_progress,
|
||||||
|
window=window,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def history(self, conv: Conversation) -> list[dict[str, Any]]:
|
||||||
|
if conv.session_id is None:
|
||||||
|
async with self._db.session() as session:
|
||||||
|
return await load_messages(
|
||||||
|
session, conversation_id=cast("int", conv.id)
|
||||||
|
)
|
||||||
|
messages = messages_from_entries(
|
||||||
|
cast("Any", await self.entries(conv)), stamped=True
|
||||||
|
)
|
||||||
|
return [self._unwrapped(message) for message in messages]
|
||||||
|
|
||||||
|
def _unwrapped(self, message: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
if message["role"] != "user":
|
||||||
|
return message
|
||||||
|
content = message["content"]
|
||||||
|
if isinstance(content, str):
|
||||||
|
text, service = split_service(content, self._texts)
|
||||||
|
return (
|
||||||
|
message
|
||||||
|
if service is None
|
||||||
|
else {**message, "content": text, "service": service}
|
||||||
|
)
|
||||||
|
for index, block in enumerate(content):
|
||||||
|
if block.get("type") != "text":
|
||||||
|
continue
|
||||||
|
text, service = split_service(str(block.get("text") or ""), self._texts)
|
||||||
|
if service is not None:
|
||||||
|
blocks = [*content]
|
||||||
|
blocks[index] = {**block, "text": text}
|
||||||
|
return {**message, "content": blocks, "service": service}
|
||||||
|
return message
|
||||||
|
|
||||||
|
async def entries(self, conv: Conversation, *, subpath: str = "") -> list[Any]:
|
||||||
|
if conv.session_id is None:
|
||||||
|
return []
|
||||||
|
key = {**self._store_key(conv), "subpath": subpath}
|
||||||
|
return list(await self._store.load(cast("Any", key)) or [])
|
||||||
|
|
||||||
|
async def subpaths(self, conv: Conversation) -> list[str]:
|
||||||
|
if conv.session_id is None:
|
||||||
|
return []
|
||||||
|
return list(await self._store.list_subkeys(cast("Any", self._store_key(conv))))
|
||||||
|
|
||||||
|
async def first_user_texts(self, ids: Iterable[int]) -> dict[int, str]:
|
||||||
|
wanted = list(ids)
|
||||||
|
if not wanted:
|
||||||
|
return {}
|
||||||
|
async with self._db.session() as session:
|
||||||
|
rows = (
|
||||||
|
await session.exec(
|
||||||
|
select(ConversationMessage).where(
|
||||||
|
col(ConversationMessage.conversation_id).in_(wanted),
|
||||||
|
ConversationMessage.seq == 0,
|
||||||
|
ConversationMessage.role == "user",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
).all()
|
||||||
|
return {
|
||||||
|
r.conversation_id: text_of(json.loads(r.content_json)).strip() for r in rows
|
||||||
|
}
|
||||||
|
|
||||||
|
async def implied_title(self, conv: Conversation) -> str | None:
|
||||||
|
if conv.title:
|
||||||
|
return conv.title
|
||||||
|
text = (await self.first_user_texts([cast("int", conv.id)])).get(
|
||||||
|
cast("int", conv.id)
|
||||||
|
)
|
||||||
|
return implied_title(text)
|
||||||
|
|
||||||
|
async def chat_name(self, conv: Conversation) -> str:
|
||||||
|
"""What a ``[[wikilink]]`` to the chat says: the file's stem when it has one."""
|
||||||
|
window = await self.window_of(conv)
|
||||||
|
if window and window.endswith(".md"):
|
||||||
|
return window.rsplit("/", 1)[-1][: -len(".md")]
|
||||||
|
return await self.implied_title(conv) or conv.external_id
|
||||||
|
|
||||||
|
async def adopt(self, *, kind: Kind, first_user_text: str) -> Conversation | None:
|
||||||
|
"""The one unbound, session-less conversation whose history starts here."""
|
||||||
|
text = first_user_text.strip()
|
||||||
|
if not text:
|
||||||
|
return None
|
||||||
|
bound = select(ConversationBinding.conversation_id).where(
|
||||||
|
col(ConversationBinding.visible).is_(True)
|
||||||
|
)
|
||||||
|
async with self._db.session() as session:
|
||||||
|
rows = (
|
||||||
|
await session.exec(
|
||||||
|
select(Conversation).where(
|
||||||
|
Conversation.kind == kind,
|
||||||
|
Conversation.status == "open",
|
||||||
|
col(Conversation.session_id).is_(None),
|
||||||
|
col(Conversation.id).not_in(bound),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
).all()
|
||||||
|
firsts = await self.first_user_texts(cast("int", r.id) for r in rows)
|
||||||
|
hits = [r for r in rows if firsts.get(cast("int", r.id)) == text]
|
||||||
|
return hits[0] if len(hits) == 1 else None
|
||||||
|
|
||||||
|
|
||||||
|
def implied_title(text: str | None) -> str | None:
|
||||||
|
if not text:
|
||||||
|
return None
|
||||||
|
line = text.strip().splitlines()[0].strip()
|
||||||
|
return line if len(line) <= TITLE_MAX else line[: TITLE_MAX - 1] + "…"
|
||||||
|
|
||||||
|
|
||||||
|
def context_of(row: Usage | None) -> int:
|
||||||
|
"""The last API call's input, or the per-call average for older rows."""
|
||||||
|
if row is None:
|
||||||
|
return 0
|
||||||
|
if row.context_tokens:
|
||||||
|
return row.context_tokens
|
||||||
|
total = row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens
|
||||||
|
return round(total / max(row.num_turns or 1, 1))
|
||||||
@@ -0,0 +1,71 @@
|
|||||||
|
"""How a new conversation starts: the seed rendered into its first prompt."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import inspect
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.kinds import as_kind
|
||||||
|
from beaver_gateway.conversations.rows import Rows
|
||||||
|
from beaver_gateway.conversations.texts import SeedContext
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = ["SEEDS", "Seeds"]
|
||||||
|
|
||||||
|
SEEDS = ("clean", "morning", "copy", "brief")
|
||||||
|
|
||||||
|
|
||||||
|
class Seeds(Rows):
|
||||||
|
async def pending_seed(self, conv: Conversation) -> str | None:
|
||||||
|
"""A seed nobody has spoken after yet: rendered now, spent once."""
|
||||||
|
seed = conv.flags.get("seed")
|
||||||
|
if not seed:
|
||||||
|
return None
|
||||||
|
parent = await self.get_row(conv.parent_id) if conv.parent_id else None
|
||||||
|
ctx = SeedContext(
|
||||||
|
kind=as_kind(conv.kind),
|
||||||
|
seed=str(seed),
|
||||||
|
agent=conv.agent_name,
|
||||||
|
parent=parent,
|
||||||
|
text=None,
|
||||||
|
title=conv.title,
|
||||||
|
)
|
||||||
|
window = conv.flags.get("seed_window")
|
||||||
|
text = await self.seed_text(
|
||||||
|
ctx, window=window if isinstance(window, int) else None
|
||||||
|
)
|
||||||
|
await self.set_flags(conv, {"seed": None, "seed_window": None})
|
||||||
|
return text
|
||||||
|
|
||||||
|
async def seed_text(self, ctx: SeedContext, *, window: int | None) -> str:
|
||||||
|
texts = self._texts
|
||||||
|
stamp = datetime.now(UTC).astimezone().strftime("%Y-%m-%d %H:%M")
|
||||||
|
head = texts.seed_head.format(
|
||||||
|
seed=ctx.seed,
|
||||||
|
kind=ctx.kind,
|
||||||
|
title=f" «{ctx.title}»" if ctx.title else "",
|
||||||
|
stamp=stamp,
|
||||||
|
)
|
||||||
|
body: str | None = None
|
||||||
|
if texts.seed is not None:
|
||||||
|
produced: Any = 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 = (
|
||||||
|
texts.seed_copy_window.format(window=window)
|
||||||
|
if window
|
||||||
|
else texts.seed_copy_all
|
||||||
|
)
|
||||||
|
body = texts.seed_copy.format(scope=scope)
|
||||||
|
elif ctx.seed == "morning":
|
||||||
|
body = texts.seed_morning_missing
|
||||||
|
parts = [head, body, ctx.text if ctx.seed != "brief" else None]
|
||||||
|
return "\n\n".join(p for p in parts if p)
|
||||||
@@ -0,0 +1,158 @@
|
|||||||
|
"""``Conversations`` - the service every frontend, job and gateway tool talks to.
|
||||||
|
|
||||||
|
Built as layers, one file each: rows → seeds → turns → messaging →
|
||||||
|
spawning → rewinding → questions → closing; this file adds start, stop and restart
|
||||||
|
recovery. A turn started by a user message streams back to whoever asked;
|
||||||
|
a turn started by an inject streams nowhere and can only speak via ``say``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import contextlib
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
from datetime import UTC, datetime, timedelta, tzinfo
|
||||||
|
|
||||||
|
from sqlmodel import col, select
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.closing import Closing, DistillResult
|
||||||
|
from beaver_gateway.conversations.kinds import KINDS
|
||||||
|
from beaver_gateway.conversations.rows import context_of, implied_title
|
||||||
|
from beaver_gateway.conversations.seeds import SEEDS
|
||||||
|
from beaver_gateway.conversations.spawning import ForkResult
|
||||||
|
from beaver_gateway.conversations.state import aware
|
||||||
|
from beaver_gateway.conversations.texts import (
|
||||||
|
ConversationTexts,
|
||||||
|
NewDayContext,
|
||||||
|
SeedContext,
|
||||||
|
UserSaid,
|
||||||
|
)
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"KINDS",
|
||||||
|
"SEEDS",
|
||||||
|
"ConversationTexts",
|
||||||
|
"Conversations",
|
||||||
|
"DistillResult",
|
||||||
|
"ForkResult",
|
||||||
|
"NewDayContext",
|
||||||
|
"SeedContext",
|
||||||
|
"UserSaid",
|
||||||
|
"context_of",
|
||||||
|
"implied_title",
|
||||||
|
"parse_at",
|
||||||
|
]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_RELATIVE = re.compile(r"^\+(\d+)\s*([smhd])$")
|
||||||
|
_UNITS = {"s": 1, "m": 60, "h": 3600, "d": 86400}
|
||||||
|
|
||||||
|
|
||||||
|
class Conversations(Closing):
|
||||||
|
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]:
|
||||||
|
"""Repair the transcripts of turns a restart cut and tell each conversation."""
|
||||||
|
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, text=self._texts.interrupted
|
||||||
|
)
|
||||||
|
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
|
||||||
|
row.pending_question = False
|
||||||
|
|
||||||
|
await self._update(conv, clear)
|
||||||
|
note = self._texts.cut_by_restart.format(turn_id=turn_id)
|
||||||
|
if fixed:
|
||||||
|
note += self._texts.repaired_tools.format(
|
||||||
|
fixed=fixed, interrupted=self._texts.interrupted
|
||||||
|
)
|
||||||
|
await self.inject(conv, note, urgency="normal", origin="system")
|
||||||
|
_log.warning("conversation %s: %s", conv.external_id, note)
|
||||||
|
for item in await self._queue.interrupted():
|
||||||
|
_log.warning(
|
||||||
|
"queue item #%s (%s) was running at restart; marked interrupted",
|
||||||
|
item.id,
|
||||||
|
item.priority,
|
||||||
|
)
|
||||||
|
return cut
|
||||||
|
|
||||||
|
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, tz: tzinfo = UTC) -> 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.replace(tzinfo=tz)
|
||||||
|
return parsed.astimezone(UTC)
|
||||||
@@ -0,0 +1,202 @@
|
|||||||
|
"""New conversations from old ones: spawn with a seed, fork a copy, merge back."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
|
from claude_agent_sdk import fork_session_via_store, project_key_for_directory
|
||||||
|
|
||||||
|
from beaver_gateway.backends.transcript import strip_tool_entries, window_entries
|
||||||
|
from beaver_gateway.conversations.messaging import Messaging
|
||||||
|
from beaver_gateway.conversations.seeds import SEEDS
|
||||||
|
from beaver_gateway.conversations.texts import SeedContext
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from beaver_gateway.backends.capture import TurnCapture
|
||||||
|
from beaver_gateway.conversations.kinds import Kind
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = ["ForkResult", "Spawning"]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class ForkResult:
|
||||||
|
conversation: Conversation
|
||||||
|
text: str
|
||||||
|
capture: TurnCapture
|
||||||
|
|
||||||
|
|
||||||
|
class Spawning(Messaging):
|
||||||
|
async def spawn(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
kind: Kind,
|
||||||
|
agent: str | None = None,
|
||||||
|
seed: str = "clean",
|
||||||
|
parent: Conversation | None = None,
|
||||||
|
text: str | None = None,
|
||||||
|
title: str | None = None,
|
||||||
|
window: int | None = None,
|
||||||
|
origin: str = "api",
|
||||||
|
binding: tuple[str, str] | None = None,
|
||||||
|
flags: dict[str, Any] | None = None,
|
||||||
|
attachments: list[dict[str, Any]] | None = None,
|
||||||
|
) -> Conversation:
|
||||||
|
"""Create a conversation in a window and queue its seed.
|
||||||
|
|
||||||
|
``binding`` reuses a window that already exists instead of asking the
|
||||||
|
home frontend for one. Without ``text`` the seed waits in ``flags``
|
||||||
|
and opens the first turn, so a fresh window costs nothing until
|
||||||
|
someone speaks.
|
||||||
|
"""
|
||||||
|
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)
|
||||||
|
if agent is None and kind == "branch" and parent is not None:
|
||||||
|
agent = parent.agent_name
|
||||||
|
agent = agent or self.default_agent(kind)
|
||||||
|
if agent is None:
|
||||||
|
msg = f"no default agent for kind {kind!r}; pass `agent`"
|
||||||
|
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,
|
||||||
|
flags=flags,
|
||||||
|
)
|
||||||
|
if binding is not None:
|
||||||
|
await self.bind(conv, frontend=binding[0], external_id=binding[1])
|
||||||
|
else:
|
||||||
|
await self.materialize(conv)
|
||||||
|
ctx = SeedContext(
|
||||||
|
kind=kind, seed=seed, agent=agent, parent=parent, text=text, title=title
|
||||||
|
)
|
||||||
|
if text is None:
|
||||||
|
return await self.set_flags(conv, {"seed": seed, "seed_window": window})
|
||||||
|
await self._queue.push(
|
||||||
|
conversation_id=cast("int", conv.id),
|
||||||
|
priority="user",
|
||||||
|
origin=f"seed:{seed}" if seed == "brief" else origin,
|
||||||
|
text=await self.seed_text(ctx, window=window),
|
||||||
|
attachments=attachments,
|
||||||
|
)
|
||||||
|
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,
|
||||||
|
agent: str | None = None,
|
||||||
|
) -> ForkResult:
|
||||||
|
"""Copy the history into a one-off session, run ``prompt`` on it, close it."""
|
||||||
|
agent = agent or conv.agent_name
|
||||||
|
session_id = await self._copy_session(
|
||||||
|
conv, window=window, strip_tools=strip_tools, agent=agent
|
||||||
|
)
|
||||||
|
child = await self.create(
|
||||||
|
kind="fork",
|
||||||
|
agent=agent,
|
||||||
|
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(agent).close(child.external_id)
|
||||||
|
child = await self.set_status(child, "closed")
|
||||||
|
return ForkResult(conversation=child, text=text, capture=capture)
|
||||||
|
|
||||||
|
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"merge: {conv.title or conv.external_id}",
|
||||||
|
)
|
||||||
|
if result.text.strip():
|
||||||
|
await self.inject(parent, result.text, urgency="normal", origin="merge")
|
||||||
|
await self.mark_closed(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 _copy_session(
|
||||||
|
self,
|
||||||
|
conv: Conversation,
|
||||||
|
*,
|
||||||
|
window: int | None,
|
||||||
|
strip_tools: bool,
|
||||||
|
agent: str | None = None,
|
||||||
|
) -> 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)
|
||||||
|
source = self._claude_agent(conv.agent_name)
|
||||||
|
target = self._claude_agent(agent) if agent else source
|
||||||
|
forked = await fork_session_via_store(
|
||||||
|
self._store, conv.session_id, directory=str(source.cwd)
|
||||||
|
)
|
||||||
|
source_key = {
|
||||||
|
"project_key": project_key_for_directory(str(source.cwd)),
|
||||||
|
"session_id": forked.session_id,
|
||||||
|
}
|
||||||
|
target_key = {
|
||||||
|
"project_key": project_key_for_directory(str(target.cwd)),
|
||||||
|
"session_id": forked.session_id,
|
||||||
|
}
|
||||||
|
if window is not None or strip_tools or target_key != source_key:
|
||||||
|
entries = await self._store.load(cast("Any", source_key)) or []
|
||||||
|
trimmed = window_entries(cast("Any", entries), window=window)
|
||||||
|
if strip_tools:
|
||||||
|
trimmed = strip_tool_entries(trimmed)
|
||||||
|
await self._store.delete(cast("Any", source_key))
|
||||||
|
await self._store.append(cast("Any", target_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
|
||||||
@@ -0,0 +1,191 @@
|
|||||||
|
"""What every layer of the conversations service shares: wiring and lookups."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
|
from claude_agent_sdk import project_key_for_directory
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.injects import InjectQueue
|
||||||
|
from beaver_gateway.conversations.texts import ConversationTexts
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Awaitable, Callable, Sequence
|
||||||
|
|
||||||
|
from claude_agent_sdk import SessionStore
|
||||||
|
|
||||||
|
from beaver_gateway.agents.claude import ClaudeAgent
|
||||||
|
from beaver_gateway.app import AgentRegistry
|
||||||
|
from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend
|
||||||
|
from beaver_gateway.backends.sessions import SessionPool
|
||||||
|
from beaver_gateway.conversations.distill import Distiller
|
||||||
|
from beaver_gateway.conversations.envelope import Envelope
|
||||||
|
from beaver_gateway.conversations.texts import UserSaid
|
||||||
|
from beaver_gateway.events.bus import EventBus
|
||||||
|
from beaver_gateway.frontends.base import Frontend
|
||||||
|
from beaver_gateway.jobs.scheduler import Scheduler
|
||||||
|
from beaver_gateway.storage.db import Database
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = ["Question", "Runner", "State", "aware", "iso"]
|
||||||
|
|
||||||
|
|
||||||
|
@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
|
||||||
|
origin: str | None = None
|
||||||
|
text: str | None = None
|
||||||
|
started_at: datetime | None = None
|
||||||
|
tools: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||||
|
|
||||||
|
def snapshot(self) -> dict[str, Any] | None:
|
||||||
|
if self.turn_id is None:
|
||||||
|
return None
|
||||||
|
return {
|
||||||
|
"id": self.turn_id,
|
||||||
|
"origin": self.origin,
|
||||||
|
"text": self.text,
|
||||||
|
"started_at": iso(self.started_at),
|
||||||
|
"tools": list(self.tools.values()),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class Question:
|
||||||
|
conversation_id: str
|
||||||
|
turn_id: str | None
|
||||||
|
questions: list[dict[str, Any]]
|
||||||
|
answer: asyncio.Future[str]
|
||||||
|
|
||||||
|
|
||||||
|
class State:
|
||||||
|
_db: Database
|
||||||
|
_agents: AgentRegistry
|
||||||
|
_backends: dict[str, Any]
|
||||||
|
_bus: EventBus
|
||||||
|
_pool: SessionPool
|
||||||
|
_store: SessionStore
|
||||||
|
_texts: ConversationTexts
|
||||||
|
_frontends: list[Frontend]
|
||||||
|
_normal_window: float
|
||||||
|
_idle_days: tuple[int, ...]
|
||||||
|
_idle_interval: float
|
||||||
|
_question_timeout: float
|
||||||
|
_envelope: Envelope | None
|
||||||
|
_distiller: Distiller | None
|
||||||
|
_user_sink: Callable[[UserSaid], Awaitable[None] | None] | None
|
||||||
|
_queue: InjectQueue
|
||||||
|
_runners: dict[int, Runner]
|
||||||
|
_questions: dict[str, Question]
|
||||||
|
_tasks: set[asyncio.Task[None]]
|
||||||
|
_idle_task: asyncio.Task[None] | None
|
||||||
|
scheduler: Scheduler | None
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
db: Database,
|
||||||
|
agents: AgentRegistry,
|
||||||
|
backends: dict[str, Any],
|
||||||
|
bus: EventBus,
|
||||||
|
pool: SessionPool,
|
||||||
|
store: SessionStore,
|
||||||
|
texts: ConversationTexts | None = None,
|
||||||
|
frontends: Sequence[Frontend] = (),
|
||||||
|
normal_window: float = 3600.0,
|
||||||
|
idle_days: Sequence[int] = (2,),
|
||||||
|
idle_interval: float = 3600.0,
|
||||||
|
question_timeout: float = 600.0,
|
||||||
|
envelope: Envelope | None = None,
|
||||||
|
distiller: Distiller | None = None,
|
||||||
|
user_sink: Callable[[UserSaid], Awaitable[None] | None] | None = None,
|
||||||
|
) -> 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._frontends = [f for f in frontends if f.name]
|
||||||
|
self._normal_window = normal_window
|
||||||
|
self._idle_days = tuple(sorted(idle_days))
|
||||||
|
self._idle_interval = idle_interval
|
||||||
|
self._question_timeout = question_timeout
|
||||||
|
self._envelope = envelope
|
||||||
|
self._distiller = distiller
|
||||||
|
self._user_sink = user_sink
|
||||||
|
self._queue = InjectQueue(db)
|
||||||
|
self._runners = {}
|
||||||
|
self._questions = {}
|
||||||
|
self._tasks = set()
|
||||||
|
self._idle_task = None
|
||||||
|
self.scheduler = None
|
||||||
|
for backend in backends.values():
|
||||||
|
if hasattr(backend, "waker"):
|
||||||
|
backend.waker = self._woken
|
||||||
|
|
||||||
|
async def _woken(self, key: str, note: str) -> None:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
@property
|
||||||
|
def db(self) -> Database:
|
||||||
|
return self._db
|
||||||
|
|
||||||
|
@property
|
||||||
|
def queue(self) -> InjectQueue:
|
||||||
|
return self._queue
|
||||||
|
|
||||||
|
@property
|
||||||
|
def bus(self) -> EventBus:
|
||||||
|
return self._bus
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pool(self) -> SessionPool:
|
||||||
|
return self._pool
|
||||||
|
|
||||||
|
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),
|
||||||
|
}
|
||||||
|
|
||||||
|
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 _track(self, task: asyncio.Task[None]) -> None:
|
||||||
|
self._tasks.add(task)
|
||||||
|
task.add_done_callback(self._tasks.discard)
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
@@ -0,0 +1,140 @@
|
|||||||
|
"""Every string the gateway puts in front of a model; English defaults, overridable."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from beaver_gateway.conversations import injects
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
|
from datetime import date, datetime
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.distill import DistillContext
|
||||||
|
from beaver_gateway.conversations.kinds import Kind
|
||||||
|
from beaver_gateway.conversations.rotation import HandoutContext
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"ConversationTexts",
|
||||||
|
"EnvelopeTexts",
|
||||||
|
"NewDayContext",
|
||||||
|
"SeedContext",
|
||||||
|
"UserSaid",
|
||||||
|
"split_service",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class SeedContext:
|
||||||
|
kind: Kind
|
||||||
|
seed: str
|
||||||
|
agent: str
|
||||||
|
parent: Conversation | None
|
||||||
|
text: str | None
|
||||||
|
title: str | None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class NewDayContext:
|
||||||
|
"""``reason`` is ``night``, ``age`` or ``context``; only ``night`` is a new day."""
|
||||||
|
|
||||||
|
day: date
|
||||||
|
reason: str
|
||||||
|
moved: int
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class UserSaid:
|
||||||
|
conversation_id: str
|
||||||
|
kind: str
|
||||||
|
title: str | None
|
||||||
|
text: str
|
||||||
|
at: datetime
|
||||||
|
agent: str = ""
|
||||||
|
"""Agent whose turn it is; a setup with several tells them apart by it."""
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class EnvelopeTexts:
|
||||||
|
header: str = (
|
||||||
|
"[envelope - background signal, not a message; react only if it bears "
|
||||||
|
"on the question]"
|
||||||
|
)
|
||||||
|
time: str = "time: {stamp} ({zone})"
|
||||||
|
changed: str = "vault, changed {since}: {names}"
|
||||||
|
since: str = "since {time}"
|
||||||
|
since_start: str = "since start"
|
||||||
|
truncated: str = "… (envelope cap)"
|
||||||
|
file_header: str = "--- {path}, added lines only ---"
|
||||||
|
more_lines: str = "+ … {count} more"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class ConversationTexts:
|
||||||
|
merge_prompt: str = (
|
||||||
|
"This branch is closing. Write a merge note for the master: what was "
|
||||||
|
"decided, what was done, what was not and why, open questions. "
|
||||||
|
"Identifiers and links verbatim. Brief, past tense."
|
||||||
|
)
|
||||||
|
inject_header: Callable[[injects.InjectContext], str] = injects.inject_header
|
||||||
|
bundle_header: str = "[injects accumulated since {since}; not the user]"
|
||||||
|
task_prompt: str = "[subagent reported back: {note}]"
|
||||||
|
"""Stands in for the user's words in mirrors of a turn a subagent started."""
|
||||||
|
interrupted: str = "interrupted"
|
||||||
|
answered: str = "The user answered: {answer}"
|
||||||
|
unanswered: str = (
|
||||||
|
"The user did not answer within {minutes} min. The question was shown "
|
||||||
|
"to them as text; finish the turn now, the answer comes as the next "
|
||||||
|
"message."
|
||||||
|
)
|
||||||
|
seed: Callable[[SeedContext], Awaitable[str | None] | str | None] | None = None
|
||||||
|
"""Body of a seed by mode; ``None`` from it falls back to the defaults below."""
|
||||||
|
seed_head: str = "[seed: {seed}] {kind}{title}, {stamp}."
|
||||||
|
seed_copy: str = "The parent's history is copied ({scope}); continue in it."
|
||||||
|
seed_copy_window: str = "last {window} turns"
|
||||||
|
seed_copy_all: str = "whole history"
|
||||||
|
seed_morning_missing: str = "No handout arrived."
|
||||||
|
handout: Callable[[HandoutContext], Awaitable[str] | str] | str = (
|
||||||
|
"This master is closing ({reason}). Write the handout for {day}: a "
|
||||||
|
"briefing for the morning, not a task list - past tense, no imperatives."
|
||||||
|
)
|
||||||
|
new_day: Callable[[NewDayContext], Awaitable[str] | str] | str = (
|
||||||
|
"The master was replaced ({reason}); the handout for {day} is written."
|
||||||
|
)
|
||||||
|
moved_injects: str = " {moved} queued injects moved over from the old master."
|
||||||
|
distill: Callable[[DistillContext], Awaitable[str] | str] | None = None
|
||||||
|
"""The distiller fork's first message; ``None`` uses the two templates below."""
|
||||||
|
distill_prompt: str = (
|
||||||
|
"Deep chat «{chat}» is closed ({reason}), today is {day}. Write the "
|
||||||
|
"digest as a file and the merge as your reply: up to 5 lines, third "
|
||||||
|
"person."
|
||||||
|
)
|
||||||
|
distill_prompt_no_memory: str = (
|
||||||
|
"Deep chat «{chat}» is closed ({reason}), today is {day}. Memory is off "
|
||||||
|
"for it: write no file, only the merge as your reply - up to 5 lines, "
|
||||||
|
"third person."
|
||||||
|
)
|
||||||
|
closed: str = "Deep chat [[{chat}]] closed{digest}.\n{text}"
|
||||||
|
closed_digest: str = ", digest [[{digest}]]"
|
||||||
|
digest_missing: str = "the digest file did not appear"
|
||||||
|
too_long: str = (
|
||||||
|
"`{name}`: {lines} lines against a cap of {max_lines}. The write was "
|
||||||
|
"rejected and the file restored. Shorten and rewrite."
|
||||||
|
)
|
||||||
|
cut_by_restart: str = "turn {turn_id} was cut by a gateway restart"
|
||||||
|
repaired_tools: str = (
|
||||||
|
"; {fixed} open tool calls received tool_result «{interrupted}»"
|
||||||
|
)
|
||||||
|
envelope: EnvelopeTexts = field(default_factory=EnvelopeTexts)
|
||||||
|
|
||||||
|
|
||||||
|
def split_service(text: str, texts: ConversationTexts) -> tuple[str, str | None]:
|
||||||
|
bundle = re.escape(texts.bundle_header).replace(r"\{since\}", ".+")
|
||||||
|
heads = f"{re.escape(texts.envelope.header)}|{bundle}"
|
||||||
|
match = re.search(rf"(?:\A|\n\n)(?=(?:{heads})$)", text, re.MULTILINE)
|
||||||
|
if match is None:
|
||||||
|
return text, None
|
||||||
|
return text[: match.start()], text[match.end() :]
|
||||||
+58
-16
@@ -1,8 +1,7 @@
|
|||||||
"""In-process MCP server with the gateway's own tools (§3.1, §3.2).
|
"""In-process MCP server with the gateway's own tools.
|
||||||
|
|
||||||
One server per live session so every tool knows which conversation is
|
One server per live session so every tool knows which conversation is
|
||||||
calling; ``alwaysLoad`` keeps the tools out of tool search. Which names a
|
calling. Which names a session gets comes from ``ClaudeAgent.gateway_tools``.
|
||||||
session gets comes from ``ClaudeAgent.gateway_tools``.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -13,9 +12,10 @@ from typing import TYPE_CHECKING, Any, cast
|
|||||||
|
|
||||||
from claude_agent_sdk import create_sdk_mcp_server, tool
|
from claude_agent_sdk import create_sdk_mcp_server, tool
|
||||||
|
|
||||||
from beaver_gateway.core.injects import URGENCY
|
from beaver_gateway.conversations.injects import URGENCY
|
||||||
from beaver_gateway.core.kinds import as_kind
|
from beaver_gateway.conversations.kinds import as_kind
|
||||||
from beaver_gateway.core.redact import redact_data
|
from beaver_gateway.frontends.files import AUTO, METHODS, SendFileError
|
||||||
|
from beaver_gateway.security.redact import redact_data
|
||||||
|
|
||||||
URGENCY_HELP = (
|
URGENCY_HELP = (
|
||||||
"normal waits for the hourly batch or rides with the next turn, wake "
|
"normal waits for the hourly batch or rides with the next turn, wake "
|
||||||
@@ -31,11 +31,11 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
from claude_agent_sdk import McpSdkServerConfig, SdkMcpTool
|
from claude_agent_sdk import McpSdkServerConfig, SdkMcpTool
|
||||||
|
|
||||||
from beaver_gateway.core.conversations import Conversations
|
from beaver_gateway.conversations.service import Conversations
|
||||||
|
|
||||||
__all__ = ["SERVER_NAME", "TOOL_NAMES", "build_tool_server"]
|
__all__ = ["SERVER_NAME", "TOOL_NAMES", "build_tool_server"]
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.core.gateway_tools")
|
_log = logging.getLogger("beaver_gateway.conversations.tools")
|
||||||
|
|
||||||
SERVER_NAME = "gateway"
|
SERVER_NAME = "gateway"
|
||||||
SAY_IN_USER_TURN = (
|
SAY_IN_USER_TURN = (
|
||||||
@@ -43,7 +43,23 @@ SAY_IN_USER_TURN = (
|
|||||||
"reply text reaches them by itself - put what you wanted to say into the "
|
"reply text reaches them by itself - put what you wanted to say into the "
|
||||||
"reply instead of repeating it here"
|
"reply instead of repeating it here"
|
||||||
)
|
)
|
||||||
TOOL_NAMES = ("read_conversation", "spawn", "say", "schedule", "inject", "close_chat")
|
TOOL_NAMES = (
|
||||||
|
"read_conversation",
|
||||||
|
"spawn",
|
||||||
|
"say",
|
||||||
|
"schedule",
|
||||||
|
"inject",
|
||||||
|
"close_chat",
|
||||||
|
"send_file",
|
||||||
|
)
|
||||||
|
SEND_FILE_HELP = (
|
||||||
|
"Send a file from disk to the human, into this conversation's window (the "
|
||||||
|
"master's if it has none): {note}. `method` is how Telegram renders it; "
|
||||||
|
"`auto` reads the extension - image is a photo, mp3/flac/m4a audio, ogg "
|
||||||
|
"voice, mp4/mov video, anything else a document. Photo and video are "
|
||||||
|
"recompressed, so pass `document` when the bytes must arrive untouched. "
|
||||||
|
"`caption` takes the same markdown as your replies."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def build_tool_server(
|
def build_tool_server(
|
||||||
@@ -68,11 +84,10 @@ def build_tool_server(
|
|||||||
def _redacting(spec: SdkMcpTool[Any]) -> SdkMcpTool[Any]:
|
def _redacting(spec: SdkMcpTool[Any]) -> SdkMcpTool[Any]:
|
||||||
"""Put a tool's result through the same mask as every other MCP.
|
"""Put a tool's result through the same mask as every other MCP.
|
||||||
|
|
||||||
These tools are mounted in-process by the SDK, so they bypass the
|
These tools are mounted in-process, bypassing the ``FastMCP`` middleware
|
||||||
``FastMCP`` middleware in :mod:`beaver_gateway.mcp.redacting` and
|
in :mod:`beaver_gateway.mcp.redacting`, so the filter is attached here
|
||||||
need the filter attached here instead. ``read_conversation`` is the
|
instead. ``read_conversation`` earns it: a replayed transcript can still
|
||||||
one that earns it: it replays a transcript, and a transcript written
|
hold a credential written before redaction existed.
|
||||||
before any of this existed can still hold a credential.
|
|
||||||
"""
|
"""
|
||||||
inner = spec.handler
|
inner = spec.handler
|
||||||
|
|
||||||
@@ -230,10 +245,37 @@ def _tools(conversations: Conversations, key: str) -> list[SdkMcpTool[Any]]:
|
|||||||
conv,
|
conv,
|
||||||
str(args["text"]),
|
str(args["text"]),
|
||||||
urgency=cast("Any", args.get("urgency") or "normal"),
|
urgency=cast("Any", args.get("urgency") or "normal"),
|
||||||
origin="агент",
|
origin="agent",
|
||||||
)
|
)
|
||||||
return _text(f"queued #{item.id}")
|
return _text(f"queued #{item.id}")
|
||||||
|
|
||||||
|
@tool(
|
||||||
|
"send_file",
|
||||||
|
SEND_FILE_HELP.format(note=conversations.file_note()),
|
||||||
|
{
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "absolute path on disk"},
|
||||||
|
"caption": {"type": "string"},
|
||||||
|
"method": {"type": "string", "enum": [AUTO, *METHODS], "default": AUTO},
|
||||||
|
},
|
||||||
|
"required": ["path"],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
async def send_file(args: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
conv = await current()
|
||||||
|
try:
|
||||||
|
return _text(
|
||||||
|
await conversations.send_file(
|
||||||
|
conv,
|
||||||
|
str(args["path"]),
|
||||||
|
caption=str(args.get("caption") or ""),
|
||||||
|
method=str(args.get("method") or AUTO),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
except SendFileError as exc:
|
||||||
|
return _error(str(exc))
|
||||||
|
|
||||||
@tool(
|
@tool(
|
||||||
"close_chat",
|
"close_chat",
|
||||||
"Close this deep chat once the current reply is finished: the "
|
"Close this deep chat once the current reply is finished: the "
|
||||||
@@ -250,7 +292,7 @@ def _tools(conversations: Conversations, key: str) -> list[SdkMcpTool[Any]]:
|
|||||||
return _error(str(exc))
|
return _error(str(exc))
|
||||||
return _text("ok: the chat closes after this reply")
|
return _text("ok: the chat closes after this reply")
|
||||||
|
|
||||||
return [read_conversation, spawn, say, schedule, inject, close_chat]
|
return [read_conversation, spawn, say, schedule, inject, send_file, close_chat]
|
||||||
|
|
||||||
|
|
||||||
def _text(text: str) -> dict[str, Any]:
|
def _text(text: str) -> dict[str, Any]:
|
||||||
@@ -0,0 +1,552 @@
|
|||||||
|
"""Running turns: one worker per conversation over its queue, the backend call."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import contextlib
|
||||||
|
import inspect
|
||||||
|
import logging
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from claude_agent_sdk import (
|
||||||
|
AssistantMessage,
|
||||||
|
RateLimitEvent,
|
||||||
|
ResultMessage,
|
||||||
|
StreamEvent,
|
||||||
|
ToolResultBlock,
|
||||||
|
ToolUseBlock,
|
||||||
|
UserMessage,
|
||||||
|
)
|
||||||
|
|
||||||
|
from beaver_gateway.backends.capture import TurnCapture
|
||||||
|
from beaver_gateway.backends.transcript import reply_text, text_of
|
||||||
|
from beaver_gateway.conversations import injects
|
||||||
|
from beaver_gateway.conversations.seeds import Seeds
|
||||||
|
from beaver_gateway.conversations.state import Runner, aware, iso
|
||||||
|
from beaver_gateway.conversations.texts import UserSaid
|
||||||
|
from beaver_gateway.frontends.accumulate import StreamAccumulator
|
||||||
|
from beaver_gateway.storage.models import Conversation, InjectQueueItem, RateLimit
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import AsyncIterator, Callable, Sequence
|
||||||
|
|
||||||
|
from beaver_gateway.events.stream import MessageStreamEvent
|
||||||
|
|
||||||
|
__all__ = ["Turns", "content_of", "image_block"]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class Turns(Seeds):
|
||||||
|
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,
|
||||||
|
item_origin: str | None = None,
|
||||||
|
shown: str | None = None,
|
||||||
|
attach: bool = False,
|
||||||
|
) -> AsyncIterator[MessageStreamEvent]:
|
||||||
|
"""Run one turn under the conversation's lock; the only path to the backend.
|
||||||
|
|
||||||
|
``shown`` is what frontends display as the turn's prompt (the user's
|
||||||
|
words, not the envelope). ``attach`` sends nothing: the turn is one
|
||||||
|
the live session started by itself, a subagent reporting back.
|
||||||
|
"""
|
||||||
|
row_id = cast("int", conv.id)
|
||||||
|
runner = self._runner(row_id)
|
||||||
|
backend = self._backend(conv.agent_name)
|
||||||
|
turn_id = turn_id or f"turn_{uuid4().hex[:12]}"
|
||||||
|
capture = capture or TurnCapture()
|
||||||
|
async with runner.lock:
|
||||||
|
runner.turn_id = turn_id
|
||||||
|
runner.origin = origin
|
||||||
|
runner.text = shown if shown is not None else _prompt_preview(messages)
|
||||||
|
runner.started_at = datetime.now(UTC)
|
||||||
|
runner.tools = {}
|
||||||
|
row = await self._mark_running(conv, turn_id)
|
||||||
|
resume = session_id if session_id is not None else row.session_id
|
||||||
|
before = await self.before_turn(conv)
|
||||||
|
self._bus.publish(
|
||||||
|
"turn.start",
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
turn_id=turn_id,
|
||||||
|
origin=origin,
|
||||||
|
item_origin=item_origin,
|
||||||
|
text=runner.text,
|
||||||
|
)
|
||||||
|
stop = "error"
|
||||||
|
cut = False
|
||||||
|
try:
|
||||||
|
observer = self._observer(conv, runner, turn_id, origin)
|
||||||
|
events = (
|
||||||
|
backend.attach(
|
||||||
|
key=conv.external_id,
|
||||||
|
capture=capture,
|
||||||
|
observer=observer,
|
||||||
|
turn_id=turn_id,
|
||||||
|
)
|
||||||
|
if attach
|
||||||
|
else backend.complete(
|
||||||
|
agent=self._claude_agent(conv.agent_name),
|
||||||
|
messages=messages,
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
session_id=resume if use_session else None,
|
||||||
|
reseed=not use_session,
|
||||||
|
capture=capture,
|
||||||
|
kind=conv.kind,
|
||||||
|
pinned=conv.kind == "master",
|
||||||
|
tools=tools,
|
||||||
|
observer=observer,
|
||||||
|
turn_id=turn_id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
async for event in events:
|
||||||
|
yield event
|
||||||
|
stop = "interrupted" if capture.interrupted else "end_turn"
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
cut = True
|
||||||
|
raise
|
||||||
|
finally:
|
||||||
|
runner.turn_id = None
|
||||||
|
await self._mark_done(conv, capture, cut=cut)
|
||||||
|
if stop != "error":
|
||||||
|
try:
|
||||||
|
await self.after_turn(conv, before)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("after-turn hook on %s failed", conv.external_id)
|
||||||
|
self._bus.publish(
|
||||||
|
"turn.end",
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
turn_id=turn_id,
|
||||||
|
origin=origin,
|
||||||
|
item_origin=item_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,
|
||||||
|
item_origin: str | None = None,
|
||||||
|
attachments: Sequence[dict[str, Any]] | None = None,
|
||||||
|
shown: str | None = None,
|
||||||
|
attach: bool = False,
|
||||||
|
) -> 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": content_of(text, attachments)}],
|
||||||
|
origin=origin,
|
||||||
|
capture=capture,
|
||||||
|
tools=tools,
|
||||||
|
turn_id=turn_id,
|
||||||
|
item_origin=item_origin,
|
||||||
|
shown=shown,
|
||||||
|
attach=attach,
|
||||||
|
):
|
||||||
|
acc.feed(event)
|
||||||
|
message = acc.finalize(model=agent.model)
|
||||||
|
return reply_text(
|
||||||
|
message.content, thinking_is_progress=agent.thinking_is_progress
|
||||||
|
), capture
|
||||||
|
|
||||||
|
async def before_turn(self, conv: Conversation) -> str | None: # noqa: ARG002
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def after_turn(self, conv: Conversation, before: str | None) -> None: # noqa: ARG002
|
||||||
|
return
|
||||||
|
|
||||||
|
def turn_origin(self, conv: Conversation) -> str | None:
|
||||||
|
runner = self._runners.get(cast("int", conv.id))
|
||||||
|
return runner.origin if runner is not None and runner.turn_id else None
|
||||||
|
|
||||||
|
async def _woken(self, key: str, note: str) -> None:
|
||||||
|
"""The live session of ``key`` started a turn by itself; run it next."""
|
||||||
|
conv = await self.get(key)
|
||||||
|
if conv is None:
|
||||||
|
_log.warning("self-started turn on unknown conversation %s: %s", key, note)
|
||||||
|
return
|
||||||
|
_log.info("conversation %s woke by itself: %s", key, note)
|
||||||
|
self._ensure_worker(cast("int", conv.id))
|
||||||
|
|
||||||
|
async def _run_pending(self, conv: Conversation) -> None:
|
||||||
|
turn_id = f"turn_{uuid4().hex[:12]}"
|
||||||
|
note = self._backend(conv.agent_name).pending_note(conv.external_id) or ""
|
||||||
|
try:
|
||||||
|
text, _ = await self.run_text_turn(
|
||||||
|
conv,
|
||||||
|
note,
|
||||||
|
origin="task",
|
||||||
|
turn_id=turn_id,
|
||||||
|
item_origin="task",
|
||||||
|
shown=note,
|
||||||
|
attach=True,
|
||||||
|
)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception(
|
||||||
|
"self-started turn %s on %s failed", turn_id, conv.external_id
|
||||||
|
)
|
||||||
|
return
|
||||||
|
if not text.strip():
|
||||||
|
return
|
||||||
|
self._bus.publish(
|
||||||
|
"reply",
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
turn_id=turn_id,
|
||||||
|
item=None,
|
||||||
|
item_origin="task",
|
||||||
|
source="task",
|
||||||
|
prompt=self._texts.task_prompt.format(note=note),
|
||||||
|
user_text=None,
|
||||||
|
text=text,
|
||||||
|
)
|
||||||
|
|
||||||
|
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._track(runner.task)
|
||||||
|
|
||||||
|
async def _worker(self, row_id: int) -> None:
|
||||||
|
runner = self._runner(row_id)
|
||||||
|
while True:
|
||||||
|
conv = await self.get_row(row_id)
|
||||||
|
if conv is not None and self._has_pending(conv):
|
||||||
|
await self._run_pending(conv)
|
||||||
|
continue
|
||||||
|
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
|
||||||
|
head = items[0]
|
||||||
|
if head.priority == "urgent":
|
||||||
|
return [head], None
|
||||||
|
tail = [i for i in items if i is not head and i.priority in ("wake", "normal")]
|
||||||
|
if head.priority in ("user", "wake"):
|
||||||
|
return [head, *tail], None
|
||||||
|
age = (datetime.now(UTC) - aware(head.created_at)).total_seconds()
|
||||||
|
if age >= self._normal_window:
|
||||||
|
return [head, *tail], 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_{uuid4().hex[:12]}"
|
||||||
|
await self._queue.start(batch, turn_id)
|
||||||
|
if head.priority == "user":
|
||||||
|
origin = "user"
|
||||||
|
prompt = head.text
|
||||||
|
await self._note_user(conv, head.text)
|
||||||
|
envelope = self._envelope_for(conv, head.text)
|
||||||
|
if envelope:
|
||||||
|
prompt += "\n\n" + envelope
|
||||||
|
if len(batch) > 1:
|
||||||
|
prompt += "\n\n" + self._bundle(batch[1:])
|
||||||
|
else:
|
||||||
|
origin = "inject"
|
||||||
|
prompt = "\n\n".join(
|
||||||
|
f"{self._texts.inject_header(injects.context_of(i))}\n{i.text}"
|
||||||
|
for i in batch
|
||||||
|
)
|
||||||
|
seed = await self.pending_seed(conv)
|
||||||
|
if seed:
|
||||||
|
prompt = f"{seed}\n\n{prompt}"
|
||||||
|
try:
|
||||||
|
text, capture = await self.run_text_turn(
|
||||||
|
conv,
|
||||||
|
prompt,
|
||||||
|
origin=origin,
|
||||||
|
turn_id=turn_id,
|
||||||
|
item_origin=head.origin,
|
||||||
|
attachments=head.attachments if origin == "user" else None,
|
||||||
|
shown=head.text if origin == "user" else None,
|
||||||
|
)
|
||||||
|
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,
|
||||||
|
item_origin=head.origin,
|
||||||
|
source="queue",
|
||||||
|
prompt=prompt,
|
||||||
|
user_text=head.text,
|
||||||
|
text=text,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _has_pending(self, conv: Conversation) -> bool:
|
||||||
|
try:
|
||||||
|
return self._backend(conv.agent_name).has_pending(conv.external_id)
|
||||||
|
except LookupError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _bundle(self, items: Sequence[InjectQueueItem]) -> str:
|
||||||
|
lines = [self._texts.bundle_header.format(since=iso(items[0].created_at))]
|
||||||
|
lines.extend(f"- [{i.origin}] {i.text}" for i in items)
|
||||||
|
return "\n".join(lines)
|
||||||
|
|
||||||
|
def _envelope_for(self, conv: Conversation, text: str = "") -> str | None:
|
||||||
|
if self._envelope is None:
|
||||||
|
return None
|
||||||
|
if conv.kind == "master":
|
||||||
|
return self._envelope.build(text=text, kind="master", agent=conv.agent_name)
|
||||||
|
if conv.kind == "branch":
|
||||||
|
return self._envelope.recall_only(
|
||||||
|
text=text, kind="branch", agent=conv.agent_name
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def _note_user(self, conv: Conversation, text: str) -> None:
|
||||||
|
if self._user_sink is None:
|
||||||
|
return
|
||||||
|
message = UserSaid(
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
kind=conv.kind,
|
||||||
|
title=conv.title,
|
||||||
|
text=text,
|
||||||
|
at=datetime.now(UTC),
|
||||||
|
agent=conv.agent_name,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
result = self._user_sink(message)
|
||||||
|
if inspect.isawaitable(result):
|
||||||
|
await result
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("user sink failed for %s", conv.external_id)
|
||||||
|
|
||||||
|
def _observer(
|
||||||
|
self, conv: Conversation, runner: Runner, turn_id: str, origin: str
|
||||||
|
) -> Callable[[Any], None]:
|
||||||
|
conversation_id = conv.external_id
|
||||||
|
|
||||||
|
def observe(message: Any) -> None:
|
||||||
|
parent = getattr(message, "parent_tool_use_id", None)
|
||||||
|
if isinstance(message, RateLimitEvent):
|
||||||
|
self._observe_rate_limit(conv, message)
|
||||||
|
elif 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):
|
||||||
|
event = 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,
|
||||||
|
)
|
||||||
|
runner.tools[block.id] = {
|
||||||
|
"tool_use_id": block.id,
|
||||||
|
"name": block.name,
|
||||||
|
"input": block.input,
|
||||||
|
"parent_tool_use_id": parent,
|
||||||
|
"started_at": event["ts"],
|
||||||
|
"ended_at": None,
|
||||||
|
"is_error": None,
|
||||||
|
"content": None,
|
||||||
|
}
|
||||||
|
elif isinstance(message, UserMessage):
|
||||||
|
blocks = message.content if isinstance(message.content, list) else ()
|
||||||
|
for block in blocks:
|
||||||
|
if isinstance(block, ToolResultBlock):
|
||||||
|
event = self._bus.publish(
|
||||||
|
"tool.result",
|
||||||
|
conversation_id=conversation_id,
|
||||||
|
turn_id=turn_id,
|
||||||
|
origin=origin,
|
||||||
|
parent_tool_use_id=parent,
|
||||||
|
tool_use_id=block.tool_use_id,
|
||||||
|
is_error=bool(block.is_error),
|
||||||
|
content=_result_preview(block.content),
|
||||||
|
)
|
||||||
|
node = runner.tools.get(block.tool_use_id)
|
||||||
|
if node is not None:
|
||||||
|
node["ended_at"] = event["ts"]
|
||||||
|
node["is_error"] = event["is_error"]
|
||||||
|
node["content"] = event["content"]
|
||||||
|
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
|
||||||
|
|
||||||
|
def _observe_rate_limit(self, conv: Conversation, message: RateLimitEvent) -> None:
|
||||||
|
info = message.rate_limit_info
|
||||||
|
row = RateLimit(
|
||||||
|
window=info.rate_limit_type or "unknown",
|
||||||
|
status=info.status,
|
||||||
|
utilization=info.utilization,
|
||||||
|
resets_at=_from_unix(info.resets_at),
|
||||||
|
overage_status=info.overage_status,
|
||||||
|
overage_resets_at=_from_unix(info.overage_resets_at),
|
||||||
|
agent_name=conv.agent_name,
|
||||||
|
session_id=message.session_id,
|
||||||
|
raw=dict(info.raw),
|
||||||
|
)
|
||||||
|
self._bus.publish(
|
||||||
|
"rate_limit",
|
||||||
|
conversation_id=conv.external_id,
|
||||||
|
window=row.window,
|
||||||
|
status=row.status,
|
||||||
|
utilization=row.utilization,
|
||||||
|
resets_at=iso(row.resets_at),
|
||||||
|
overage_status=row.overage_status,
|
||||||
|
)
|
||||||
|
self._track(asyncio.create_task(self._record_rate_limit(row)))
|
||||||
|
|
||||||
|
async def _record_rate_limit(self, row: RateLimit) -> None:
|
||||||
|
try:
|
||||||
|
async with self._db.session() as session:
|
||||||
|
session.add(row)
|
||||||
|
await session.commit()
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("rate limit write failed")
|
||||||
|
|
||||||
|
async def _mark_running(self, conv: Conversation, turn_id: str) -> Conversation:
|
||||||
|
async def apply(row: Conversation) -> None:
|
||||||
|
row.running_turn = turn_id
|
||||||
|
row.last_activity_at = datetime.now(UTC)
|
||||||
|
|
||||||
|
return await self._update(conv, apply)
|
||||||
|
|
||||||
|
async def _mark_done(
|
||||||
|
self, conv: Conversation, capture: TurnCapture, *, cut: bool = False
|
||||||
|
) -> None:
|
||||||
|
async def apply(row: Conversation) -> None:
|
||||||
|
if not cut:
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
IMAGE_MAX_BYTES = 5 * 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
|
def content_of(
|
||||||
|
text: str, attachments: Sequence[dict[str, Any]] | None
|
||||||
|
) -> str | list[dict[str, Any]]:
|
||||||
|
"""Plain text, or text plus the images that ride with it."""
|
||||||
|
blocks = [b for a in attachments or () if (b := image_block(a)) is not None]
|
||||||
|
if not blocks:
|
||||||
|
return text
|
||||||
|
return [{"type": "text", "text": text}, *blocks]
|
||||||
|
|
||||||
|
|
||||||
|
def image_block(attachment: dict[str, Any]) -> dict[str, Any] | None:
|
||||||
|
media_type = str(attachment.get("media_type") or "")
|
||||||
|
path = Path(str(attachment.get("path") or ""))
|
||||||
|
if not media_type.startswith("image/") or not path.is_file():
|
||||||
|
return None
|
||||||
|
if path.stat().st_size > IMAGE_MAX_BYTES:
|
||||||
|
return None
|
||||||
|
data = base64.b64encode(path.read_bytes()).decode("ascii")
|
||||||
|
return {
|
||||||
|
"type": "image",
|
||||||
|
"source": {"type": "base64", "media_type": media_type, "data": data},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _prompt_preview(messages: Sequence[Any], limit: int = 400) -> str | None:
|
||||||
|
if not messages:
|
||||||
|
return None
|
||||||
|
text = text_of(messages[-1].get("content"))
|
||||||
|
return text[:limit] if text else None
|
||||||
|
|
||||||
|
|
||||||
|
def _from_unix(value: int | None) -> datetime | None:
|
||||||
|
return datetime.fromtimestamp(value, tz=UTC) if value is not None else None
|
||||||
|
|
||||||
|
|
||||||
|
def _result_preview(
|
||||||
|
content: str | list[dict[str, Any]] | None, limit: int = 400
|
||||||
|
) -> str:
|
||||||
|
if content is None:
|
||||||
|
return ""
|
||||||
|
text = (
|
||||||
|
content
|
||||||
|
if isinstance(content, str)
|
||||||
|
else "\n".join(
|
||||||
|
str(part.get("text", ""))
|
||||||
|
for part in content
|
||||||
|
if isinstance(part, dict) and part.get("type") == "text"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return text if len(text) <= limit else text[:limit] + "…"
|
||||||
|
|
||||||
|
|
||||||
|
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,
|
||||||
|
}
|
||||||
@@ -1,7 +0,0 @@
|
|||||||
"""Cross-cutting machinery: registries, event protocol, auth, sessions."""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from beaver_gateway.core.registry import AgentRegistry, Gateway, McpRegistry
|
|
||||||
|
|
||||||
__all__ = ["AgentRegistry", "Gateway", "McpRegistry"]
|
|
||||||
@@ -1,94 +0,0 @@
|
|||||||
"""Single entry point for writing :class:`AuditLog` rows.
|
|
||||||
|
|
||||||
Every frontend ends up needing the same three-line pattern — open a DB
|
|
||||||
session, append a row, swallow failures so the user-visible request
|
|
||||||
still succeeds. Phase 4.3 inlined that pattern in the admin frontend
|
|
||||||
under a private ``_audit()`` helper; Phase 4.4 lifts it here so the
|
|
||||||
Messages and MCP frontends can call the same function and so the
|
|
||||||
swallow-and-log policy lives in one place.
|
|
||||||
|
|
||||||
The contract:
|
|
||||||
|
|
||||||
* ``log(runtime, actor=..., kind=...)`` is fire-and-forget. It awaits
|
|
||||||
the DB write (so callers can ``await`` it before responding and get
|
|
||||||
ordering), but never raises — if the audit insert fails, the function
|
|
||||||
emits an ``exception`` log line and returns.
|
|
||||||
* ``actor`` is a free-form string. By convention: ``"token:<name>"``
|
|
||||||
for bearer-authenticated traffic, ``"admin:<user>"`` for admin-UI
|
|
||||||
actions, ``"anon"`` for failed-auth paths we still want to record.
|
|
||||||
* ``kind`` is a short tag — see :data:`KNOWN_KINDS` for the set the
|
|
||||||
current frontends emit; new tags don't need a code change here, the
|
|
||||||
column is free-form.
|
|
||||||
* ``**detail`` is JSON-serialised by :func:`append_audit`. Keep it
|
|
||||||
small: paths, methods, status codes — not request bodies. Anything
|
|
||||||
passed here lands in ``AuditLog.detail_json`` verbatim.
|
|
||||||
|
|
||||||
Why a thin wrapper rather than ``append_audit`` directly: callers want
|
|
||||||
"write if you can, otherwise carry on", and pulling the try/except into
|
|
||||||
every frontend was already starting to drift (admin had it, bearer
|
|
||||||
frontends would have copy-pasted). One module, one policy.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import logging
|
|
||||||
from typing import TYPE_CHECKING, Any
|
|
||||||
|
|
||||||
from beaver_gateway.storage import append_audit
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from beaver_gateway.frontends.base import GatewayRuntime
|
|
||||||
|
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.audit")
|
|
||||||
|
|
||||||
|
|
||||||
# Tags currently emitted by the gateway. The set is informational —
|
|
||||||
# ``AuditLog.kind`` is free-form so new code can introduce new tags
|
|
||||||
# without touching this list — but listing them here gives the admin UI
|
|
||||||
# and any downstream log consumers one canonical reference.
|
|
||||||
KNOWN_KINDS: frozenset[str] = frozenset(
|
|
||||||
{
|
|
||||||
"messages", # POST /v1/messages accepted
|
|
||||||
"mcp_call", # /mcp/<ns>/... proxied
|
|
||||||
"tool_call", # a model's tool call seen by the PreToolUse hook
|
|
||||||
"login_ok",
|
|
||||||
"login_failed",
|
|
||||||
"logout",
|
|
||||||
"token_create",
|
|
||||||
"token_revoke",
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
async def log(
|
|
||||||
runtime: GatewayRuntime,
|
|
||||||
*,
|
|
||||||
actor: str,
|
|
||||||
kind: str,
|
|
||||||
agent_name: str | None = None,
|
|
||||||
**detail: Any,
|
|
||||||
) -> None:
|
|
||||||
"""Best-effort audit insert. Never raises.
|
|
||||||
|
|
||||||
Opens its own short-lived :class:`AsyncSession` so callers don't
|
|
||||||
have to thread one through. If the DB hiccups (table missing,
|
|
||||||
disk full, connection drop), we log and move on — the audit trail
|
|
||||||
is observability, not a hard precondition for serving the request.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
async with runtime.db.session() as session:
|
|
||||||
await append_audit(
|
|
||||||
session,
|
|
||||||
actor=actor,
|
|
||||||
kind=kind,
|
|
||||||
agent_name=agent_name,
|
|
||||||
detail=detail or None,
|
|
||||||
)
|
|
||||||
except Exception: # noqa: BLE001
|
|
||||||
_log.exception(
|
|
||||||
"audit write failed: actor=%s kind=%s agent=%s", actor, kind, agent_name
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
__all__ = ["KNOWN_KINDS", "log"]
|
|
||||||
File diff suppressed because it is too large.
Load diff
@@ -1,117 +0,0 @@
|
|||||||
"""Agent / MCP registries + the user-facing ``Gateway`` collector.
|
|
||||||
|
|
||||||
The user's ``/config/config.py`` ends with::
|
|
||||||
|
|
||||||
gateway = Gateway(agents=[...], mcps=[...], frontends=[...])
|
|
||||||
|
|
||||||
``cli.main`` picks that object up and builds the registries.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from dataclasses import dataclass, field
|
|
||||||
from typing import TYPE_CHECKING
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from collections.abc import Awaitable, Callable, Iterable, Iterator
|
|
||||||
|
|
||||||
from beaver_gateway.agents.base import BaseAgent
|
|
||||||
from beaver_gateway.core.conversations import ConversationTexts, UserSaid
|
|
||||||
from beaver_gateway.core.distill import Distiller
|
|
||||||
from beaver_gateway.core.envelope import RecallContext
|
|
||||||
from beaver_gateway.core.rotation import RotationPolicy
|
|
||||||
from beaver_gateway.core.scheduler import Budget, Job
|
|
||||||
from beaver_gateway.core.watch import VaultWatch
|
|
||||||
from beaver_gateway.frontends.base import Frontend
|
|
||||||
from beaver_gateway.mcp.types import McpServerT
|
|
||||||
|
|
||||||
|
|
||||||
class AgentRegistry:
|
|
||||||
"""Name → agent lookup with duplicate detection."""
|
|
||||||
|
|
||||||
def __init__(self, agents: Iterable[BaseAgent]) -> None:
|
|
||||||
self._agents: dict[str, BaseAgent] = {}
|
|
||||||
for a in agents:
|
|
||||||
if a.name in self._agents:
|
|
||||||
msg = f"duplicate agent name: {a.name!r}"
|
|
||||||
raise ValueError(msg)
|
|
||||||
self._agents[a.name] = a
|
|
||||||
|
|
||||||
def __getitem__(self, name: str) -> BaseAgent:
|
|
||||||
return self._agents[name]
|
|
||||||
|
|
||||||
def get(self, name: str) -> BaseAgent | None:
|
|
||||||
return self._agents.get(name)
|
|
||||||
|
|
||||||
def __iter__(self) -> Iterator[BaseAgent]:
|
|
||||||
return iter(self._agents.values())
|
|
||||||
|
|
||||||
def __len__(self) -> int:
|
|
||||||
return len(self._agents)
|
|
||||||
|
|
||||||
def __contains__(self, name: object) -> bool:
|
|
||||||
return name in self._agents
|
|
||||||
|
|
||||||
|
|
||||||
class McpRegistry:
|
|
||||||
"""Name → MCP server lookup with duplicate detection."""
|
|
||||||
|
|
||||||
def __init__(self, mcps: Iterable[McpServerT]) -> None:
|
|
||||||
self._mcps: dict[str, McpServerT] = {}
|
|
||||||
for m in mcps:
|
|
||||||
if m.name in self._mcps:
|
|
||||||
msg = f"duplicate mcp name: {m.name!r}"
|
|
||||||
raise ValueError(msg)
|
|
||||||
self._mcps[m.name] = m
|
|
||||||
|
|
||||||
def __getitem__(self, name: str) -> McpServerT:
|
|
||||||
return self._mcps[name]
|
|
||||||
|
|
||||||
def get(self, name: str) -> McpServerT | None:
|
|
||||||
return self._mcps.get(name)
|
|
||||||
|
|
||||||
def __iter__(self) -> Iterator[McpServerT]:
|
|
||||||
return iter(self._mcps.values())
|
|
||||||
|
|
||||||
def __len__(self) -> int:
|
|
||||||
return len(self._mcps)
|
|
||||||
|
|
||||||
def __contains__(self, name: object) -> bool:
|
|
||||||
return name in self._mcps
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
|
||||||
class Gateway:
|
|
||||||
"""Top-level object the user assembles in ``/config/config.py``."""
|
|
||||||
|
|
||||||
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)."""
|
|
||||||
jobs: list[Job] = field(default_factory=list)
|
|
||||||
"""Cron / webhook / event jobs for ``core/scheduler`` (§3.6, §4.5)."""
|
|
||||||
rotation: RotationPolicy | None = None
|
|
||||||
"""When a master is rotated (§4.5); ``None`` keeps the defaults."""
|
|
||||||
watch: VaultWatch | None = None
|
|
||||||
"""Vault watcher feeding the envelope (§3.5, §4.6); ``None`` = no vault block."""
|
|
||||||
recall: Callable[[RecallContext], str | None] | None = None
|
|
||||||
"""Envelope lookup on the user's text: pointers into the vault (cards,
|
|
||||||
the agent's notes, due dates) the gateway knows no paths for (§3.3)."""
|
|
||||||
user_sink: Callable[[UserSaid], Awaitable[None] | None] | None = None
|
|
||||||
"""Sees every user message as it enters a master or branch turn - the
|
|
||||||
setup's own grep-able log of what the user said, outside the transcript."""
|
|
||||||
budget: Budget | None = None
|
|
||||||
"""Subscription window past which non-critical jobs wait (§4.5)."""
|
|
||||||
distiller: Distiller | None = None
|
|
||||||
"""Who closes deep chats and where the digests and the index live (§8.4)."""
|
|
||||||
tz: str = "UTC"
|
|
||||||
"""Local zone for the envelope clock and the rotation hour."""
|
|
||||||
host: str = "0.0.0.0" # noqa: S104
|
|
||||||
port: int = 8000
|
|
||||||
"""The one listener; every HTTP frontend is mounted under its ``path``."""
|
|
||||||
public_url: str | None = None
|
|
||||||
"""Origin the reverse proxy shows the world (``https://b.example.com``).
|
|
||||||
|
|
||||||
Advertised endpoints and MCP discovery are built on it; ``None``
|
|
||||||
derives the origin from each request."""
|
|
||||||
@@ -1,133 +0,0 @@
|
|||||||
"""Cross-frontend turn record.
|
|
||||||
|
|
||||||
Frontends that finish a turn (the Anthropic Messages frontend, the
|
|
||||||
markdown frontend) emit a :class:`TurnRecord` to every handler in
|
|
||||||
``GatewayRuntime.turn_log_handlers``. The markdown frontend uses this
|
|
||||||
to persist chats from other frontends into the Obsidian vault — see
|
|
||||||
``frontends/markdown/crossfront.py``.
|
|
||||||
|
|
||||||
Kept tiny on purpose: it carries the structured-enough payload a logger
|
|
||||||
needs (which agent ran, input history, the assembled assistant reply)
|
|
||||||
and nothing else. The full event stream is gone by the time handlers
|
|
||||||
run — if a future consumer needs deltas it would subscribe at a lower
|
|
||||||
level, not here.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import re
|
|
||||||
from dataclasses import dataclass
|
|
||||||
from typing import TYPE_CHECKING, Any
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from anthropic.types import Message, MessageParam
|
|
||||||
|
|
||||||
|
|
||||||
__all__ = ["TurnRecord", "slugify"]
|
|
||||||
|
|
||||||
|
|
||||||
# Filesystem-safe slug: collapse anything that isn't a word char or
|
|
||||||
# space/hyphen to a hyphen, then squash runs of separators. Aimed at
|
|
||||||
# letting users build filenames from ``record.first_user_text`` without
|
|
||||||
# hand-rolling sanitization in every config.
|
|
||||||
_SLUG_BAD_RE = re.compile(r"[^\w\s\-]+", flags=re.UNICODE)
|
|
||||||
_SLUG_SEP_RE = re.compile(r"[\s\-]+", flags=re.UNICODE)
|
|
||||||
|
|
||||||
|
|
||||||
def slugify(text: str, *, maxlen: int = 40) -> str:
|
|
||||||
"""Sanitize ``text`` for use as a filename fragment.
|
|
||||||
|
|
||||||
Strips punctuation, collapses whitespace/hyphens into single ``-``,
|
|
||||||
and truncates to ``maxlen``. Returns ``"untitled"`` for empty input.
|
|
||||||
Unicode letters are preserved (Obsidian handles them fine; macOS
|
|
||||||
and modern Linux fs's too).
|
|
||||||
"""
|
|
||||||
cleaned = _SLUG_BAD_RE.sub(" ", text).strip()
|
|
||||||
cleaned = _SLUG_SEP_RE.sub("-", cleaned).strip("-")
|
|
||||||
if not cleaned:
|
|
||||||
return "untitled"
|
|
||||||
if len(cleaned) > maxlen:
|
|
||||||
cleaned = cleaned[:maxlen].rstrip("-")
|
|
||||||
return cleaned or "untitled"
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
|
||||||
class TurnRecord:
|
|
||||||
"""One completed turn, as seen by a frontend.
|
|
||||||
|
|
||||||
``input_messages`` is the conversation history sent to the backend
|
|
||||||
(everything *before* the assistant reply). ``output_message`` is the
|
|
||||||
finalized assistant ``Message`` (post-accumulation, with all content
|
|
||||||
blocks attached). ``system`` is the per-request system prompt if any
|
|
||||||
— agents own their own ``system_prompt``, this is the override the
|
|
||||||
caller passed; most handlers can ignore it.
|
|
||||||
|
|
||||||
``source`` names the frontend that produced the record so the
|
|
||||||
cross-frontend logger can avoid logging its own turns (markdown
|
|
||||||
frontend writing the file would otherwise also receive its own
|
|
||||||
broadcast and double-write).
|
|
||||||
"""
|
|
||||||
|
|
||||||
agent_name: str
|
|
||||||
input_messages: list[MessageParam]
|
|
||||||
output_message: Message
|
|
||||||
system: str | None = None
|
|
||||||
source: str = "anthropic"
|
|
||||||
|
|
||||||
@property
|
|
||||||
def first_user_text(self) -> str:
|
|
||||||
"""Plain text of the *earliest* user turn in this conversation.
|
|
||||||
|
|
||||||
Useful for naming new files by topic. Empty string if the input
|
|
||||||
history somehow has no user turn (shouldn't happen — turns are
|
|
||||||
broadcast only after a user → assistant cycle).
|
|
||||||
"""
|
|
||||||
for msg in self.input_messages:
|
|
||||||
if msg.get("role") == "user":
|
|
||||||
return _text_of(msg.get("content", ""))
|
|
||||||
return ""
|
|
||||||
|
|
||||||
@property
|
|
||||||
def last_user_text(self) -> str:
|
|
||||||
"""Plain text of the most recent user turn (the trigger)."""
|
|
||||||
for msg in reversed(self.input_messages):
|
|
||||||
if msg.get("role") == "user":
|
|
||||||
return _text_of(msg.get("content", ""))
|
|
||||||
return ""
|
|
||||||
|
|
||||||
@property
|
|
||||||
def assistant_text(self) -> str:
|
|
||||||
"""Plain text of the assistant reply, with thinking/tool blocks dropped."""
|
|
||||||
chunks = [
|
|
||||||
getattr(block, "text", "") or ""
|
|
||||||
for block in getattr(self.output_message, "content", ())
|
|
||||||
if getattr(block, "type", None) == "text"
|
|
||||||
]
|
|
||||||
return "\n\n".join(c for c in chunks if c)
|
|
||||||
|
|
||||||
|
|
||||||
def _text_of(content: object) -> str:
|
|
||||||
"""Flatten an Anthropic ``MessageParam.content`` field to a plain string.
|
|
||||||
|
|
||||||
Handles both shapes the SDK accepts: raw string, or a list of block
|
|
||||||
dicts (we pull ``text`` blocks only). Anything we don't recognize is
|
|
||||||
silently skipped — these helpers exist for naming files, not for
|
|
||||||
faithful content reconstruction.
|
|
||||||
"""
|
|
||||||
if isinstance(content, str):
|
|
||||||
return content
|
|
||||||
if isinstance(content, list):
|
|
||||||
parts: list[str] = []
|
|
||||||
for blk in content:
|
|
||||||
if not isinstance(blk, dict):
|
|
||||||
continue
|
|
||||||
# ``MessageParam.content`` is typed as a union of typed-dicts
|
|
||||||
# per Anthropic SDK; we only care about plain ``text`` blocks
|
|
||||||
# and look them up via ``Any`` to dodge the keyed-typed-dict
|
|
||||||
# variance gymnastics (``ty`` won't let an open dict alias a
|
|
||||||
# closed-keyed one).
|
|
||||||
d: Any = blk
|
|
||||||
if d.get("type") == "text":
|
|
||||||
parts.append(str(d.get("text", "")))
|
|
||||||
return "\n\n".join(p for p in parts if p)
|
|
||||||
return ""
|
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
"""The in-process event bus and the stream event protocol backends emit."""
|
||||||
@@ -1,10 +1,7 @@
|
|||||||
"""In-process event bus: what the gateway does, as a stream frontends can tap.
|
"""In-process event bus: what the gateway does, as a stream frontends can tap.
|
||||||
|
|
||||||
Events are plain dicts ``{"type", "seq", "ts", ...}``. Publishers never
|
Events are plain dicts. Publishers never block; a subscriber that falls
|
||||||
block; a subscriber that falls behind loses its oldest events rather than
|
behind loses its oldest events rather than stalling the publisher.
|
||||||
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
|
from __future__ import annotations
|
||||||
@@ -60,10 +57,16 @@ class EventBus:
|
|||||||
event = await queue.get()
|
event = await queue.get()
|
||||||
if (
|
if (
|
||||||
conversation_id is None
|
conversation_id is None
|
||||||
or event.get("conversation_id") == conversation_id
|
or _conversation_of(event) == conversation_id
|
||||||
):
|
):
|
||||||
yield event
|
yield event
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def subscribers(self) -> int:
|
def subscribers(self) -> int:
|
||||||
return len(self._subscribers)
|
return len(self._subscribers)
|
||||||
|
|
||||||
|
|
||||||
|
def _conversation_of(event: Event) -> Any:
|
||||||
|
if event["type"].startswith("conversation."):
|
||||||
|
return event.get("id")
|
||||||
|
return event.get("conversation_id")
|
||||||
@@ -1,11 +1,7 @@
|
|||||||
"""Unified streaming-event protocol shared by every backend adapter.
|
"""Streaming-event protocol shared by every backend adapter.
|
||||||
|
|
||||||
We piggyback on :mod:`anthropic.types` rather than reinvent the wire format:
|
Aliases over :mod:`anthropic.types` plus ``build_*`` helpers that keep
|
||||||
``AnthropicMessagesFrontend`` ultimately serializes whatever a backend yields
|
verbose constructors out of adapters.
|
||||||
straight to SSE via :py:meth:`pydantic.BaseModel.model_dump_json`, so events
|
|
||||||
must be valid Anthropic ``message_stream`` records. The aliases below give
|
|
||||||
the rest of the codebase one import path; the ``build_*`` helpers keep the
|
|
||||||
verbose constructors out of every adapter.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
+14
-18
@@ -1,11 +1,7 @@
|
|||||||
"""Collapse an Anthropic event stream into one ``Message``.
|
"""Collapse an Anthropic event stream into one ``Message``.
|
||||||
|
|
||||||
Extracted from ``AnthropicMessagesFrontend`` so the markdown frontend
|
Mirrors the Anthropic SDK's own accumulator: folds content-block deltas
|
||||||
can run the same accumulation logic when it wants the finalized turn
|
into finalized blocks, indexed by their ``content_block`` index.
|
||||||
rather than raw SSE chunks. Mirrors the Anthropic SDK's own accumulator:
|
|
||||||
walks events, builds block dicts indexed by their ``content_block``
|
|
||||||
index, folds text / thinking deltas in, buffers ``input_json_delta``
|
|
||||||
chunks until the block closes (then JSON-parses them once).
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -33,7 +29,7 @@ from anthropic.types import (
|
|||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator
|
from collections.abc import AsyncIterator
|
||||||
|
|
||||||
from beaver_gateway.core.events import MessageStreamEvent, StopReason
|
from beaver_gateway.events.stream import MessageStreamEvent, StopReason
|
||||||
|
|
||||||
|
|
||||||
__all__ = ["StreamAccumulator", "accumulate"]
|
__all__ = ["StreamAccumulator", "accumulate"]
|
||||||
@@ -42,9 +38,7 @@ __all__ = ["StreamAccumulator", "accumulate"]
|
|||||||
class StreamAccumulator:
|
class StreamAccumulator:
|
||||||
"""Folds a stream of events into one ``Message``, incrementally.
|
"""Folds a stream of events into one ``Message``, incrementally.
|
||||||
|
|
||||||
Use when you need to *both* forward events somewhere (SSE) *and*
|
Call :meth:`feed` for each event, then :meth:`finalize` once.
|
||||||
keep a finalized ``Message`` for post-stream work (audit, logging
|
|
||||||
to disk). Call :meth:`feed` for each event, :meth:`finalize` once.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
__slots__ = (
|
__slots__ = (
|
||||||
@@ -67,10 +61,11 @@ class StreamAccumulator:
|
|||||||
self._stop_sequence: str | None = None
|
self._stop_sequence: str | None = None
|
||||||
|
|
||||||
def feed(self, ev: MessageStreamEvent) -> None:
|
def feed(self, ev: MessageStreamEvent) -> None:
|
||||||
# isinstance, not ``ev.type == "..."``: ty narrows on the
|
"""Fold one stream event into the in-progress message state.
|
||||||
# discriminator only via the class, and the raw event union
|
|
||||||
# carries its own discriminators (``Raw*Event``) the SDK
|
Uses ``isinstance`` rather than ``ev.type == ...`` so ty narrows
|
||||||
# already promises.
|
the event type from the class, not a string comparison.
|
||||||
|
"""
|
||||||
if isinstance(ev, RawMessageStartEvent):
|
if isinstance(ev, RawMessageStartEvent):
|
||||||
self._message_id = ev.message.id
|
self._message_id = ev.message.id
|
||||||
self._role = ev.message.role
|
self._role = ev.message.role
|
||||||
@@ -109,6 +104,11 @@ class StreamAccumulator:
|
|||||||
)
|
)
|
||||||
|
|
||||||
def finalize(self, *, model: str) -> Message:
|
def finalize(self, *, model: str) -> Message:
|
||||||
|
"""Build the finalized ``Message`` from accumulated block state.
|
||||||
|
|
||||||
|
``role`` is always ``"assistant"`` at the wire level; the cast
|
||||||
|
avoids a runtime check ty would otherwise require.
|
||||||
|
"""
|
||||||
content: list[Any] = []
|
content: list[Any] = []
|
||||||
for idx in sorted(self._blocks):
|
for idx in sorted(self._blocks):
|
||||||
bd = self._blocks[idx]
|
bd = self._blocks[idx]
|
||||||
@@ -120,10 +120,6 @@ class StreamAccumulator:
|
|||||||
elif btype == "thinking":
|
elif btype == "thinking":
|
||||||
content.append(ThinkingBlock.model_validate(bd))
|
content.append(ThinkingBlock.model_validate(bd))
|
||||||
|
|
||||||
# ``role`` is always ``"assistant"`` at the wire level — we
|
|
||||||
# initialised the field to that and only overwrite from a
|
|
||||||
# ``RawMessageStartEvent`` which itself carries the same literal.
|
|
||||||
# The cast keeps both type-checkers happy without a runtime check.
|
|
||||||
return Message(
|
return Message(
|
||||||
id=self._message_id or "msg_unknown",
|
id=self._message_id or "msg_unknown",
|
||||||
type="message",
|
type="message",
|
||||||
@@ -1,15 +1,7 @@
|
|||||||
"""Admin console: serves the ``ui/`` SPA and signs the operator in.
|
"""Admin console: serves the ``ui/`` SPA and signs the operator in.
|
||||||
|
|
||||||
Everything the console shows comes from ``/api/*`` (``ApiFrontend``)
|
Login/logout/session live under ``/admin/auth``; a process-lifetime
|
||||||
with a bearer. This app, mounted at ``/admin``, owns three JSON routes
|
admin bearer (scope ``*``) is minted at startup and rotates on restart.
|
||||||
under ``/admin/auth`` - login (``ADMIN_USER`` / ``ADMIN_PASS`` from env,
|
|
||||||
session cookie signed with ``SESSION_SECRET``, 8 h), logout, and
|
|
||||||
``session``, which hands a signed-in browser the process-lifetime admin
|
|
||||||
bearer - and the static build under ``/admin/``.
|
|
||||||
|
|
||||||
The bearer is minted at startup, registered in the token store with
|
|
||||||
scope ``*`` and never persisted; a gateway restart rotates it, and the
|
|
||||||
SPA refetches ``session`` on a 401.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -27,8 +19,8 @@ import itsdangerous
|
|||||||
from fastapi import FastAPI, HTTPException, Request, status
|
from fastapi import FastAPI, HTTPException, Request, status
|
||||||
from fastapi.responses import FileResponse, JSONResponse, Response
|
from fastapi.responses import FileResponse, JSONResponse, Response
|
||||||
|
|
||||||
from beaver_gateway.core import audit
|
|
||||||
from beaver_gateway.frontends.base import Frontend
|
from beaver_gateway.frontends.base import Frontend
|
||||||
|
from beaver_gateway.security import audit
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from beaver_gateway.frontends.base import GatewayRuntime
|
from beaver_gateway.frontends.base import GatewayRuntime
|
||||||
|
|||||||
@@ -1,17 +1,8 @@
|
|||||||
"""``POST /v1/messages`` frontend.
|
"""``POST /v1/messages`` frontend.
|
||||||
|
|
||||||
Exposes the gateway as an Anthropic-compatible Messages endpoint, so any
|
Exposes the gateway as an Anthropic-compatible Messages endpoint, so any
|
||||||
client that already speaks Anthropic (Cursor, Cline, the official SDK,
|
client that already speaks Anthropic can reach a configured agent by
|
||||||
``curl``) can hit a configured agent by passing its name as ``model``.
|
passing its name as ``model``.
|
||||||
|
|
||||||
A Claude agent behind this endpoint is a ``deep`` conversation: the client
|
|
||||||
knows nothing about our ids, so the text fingerprint of the history it
|
|
||||||
sends is the ``(anthropic, fingerprint)`` binding of the conversation,
|
|
||||||
rebound after every turn to the fingerprint the next request will carry.
|
|
||||||
A history nobody has seen becomes a new conversation, materialized by the
|
|
||||||
home frontend of ``deep`` (the vault file), and every reply is published
|
|
||||||
on the bus so that file follows the chat. Other agents (Raycast) stay
|
|
||||||
stateless and are only archived through ``turn_log_handlers``.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -25,21 +16,22 @@ from fastapi import FastAPI, HTTPException, Request, status
|
|||||||
from fastapi.responses import JSONResponse, StreamingResponse
|
from fastapi.responses import JSONResponse, StreamingResponse
|
||||||
|
|
||||||
from beaver_gateway.agents.claude import ClaudeAgent
|
from beaver_gateway.agents.claude import ClaudeAgent
|
||||||
from beaver_gateway.core import audit
|
from beaver_gateway.backends.capture import TurnCapture
|
||||||
from beaver_gateway.core.transcript import fingerprint, text_of
|
from beaver_gateway.backends.transcript import fingerprint, reply_text, text_of
|
||||||
from beaver_gateway.core.turn_capture import TurnCapture
|
from beaver_gateway.frontends.accumulate import StreamAccumulator
|
||||||
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.base import Frontend
|
from beaver_gateway.frontends.base import Frontend
|
||||||
|
from beaver_gateway.frontends.bearer import require_token
|
||||||
|
from beaver_gateway.frontends.turn_record import TurnRecord
|
||||||
|
from beaver_gateway.security import audit
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||||
|
|
||||||
from anthropic.types import Message, MessageParam
|
from anthropic.types import Message, MessageParam
|
||||||
|
|
||||||
from beaver_gateway.core.conversations import Conversations
|
from beaver_gateway.agents.base import BaseAgent
|
||||||
from beaver_gateway.core.events import MessageStreamEvent
|
from beaver_gateway.conversations.service import Conversations
|
||||||
|
from beaver_gateway.events.stream import MessageStreamEvent
|
||||||
from beaver_gateway.frontends.base import GatewayRuntime
|
from beaver_gateway.frontends.base import GatewayRuntime
|
||||||
from beaver_gateway.storage.models import Conversation
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
@@ -167,6 +159,7 @@ class AnthropicMessagesFrontend(Frontend):
|
|||||||
capture=capture,
|
capture=capture,
|
||||||
turn_id=turn_id,
|
turn_id=turn_id,
|
||||||
history=history,
|
history=history,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
options: dict[str, Any] = {}
|
options: dict[str, Any] = {}
|
||||||
@@ -179,7 +172,7 @@ class AnthropicMessagesFrontend(Frontend):
|
|||||||
async def after(message: Message) -> None:
|
async def after(message: Message) -> None:
|
||||||
await _broadcast_turn(
|
await _broadcast_turn(
|
||||||
runtime,
|
runtime,
|
||||||
agent_name=agent.name,
|
agent=agent,
|
||||||
input_messages=messages,
|
input_messages=messages,
|
||||||
output_message=message,
|
output_message=message,
|
||||||
system=system_str,
|
system=system_str,
|
||||||
@@ -244,12 +237,9 @@ async def _finish(
|
|||||||
capture: TurnCapture,
|
capture: TurnCapture,
|
||||||
turn_id: str,
|
turn_id: str,
|
||||||
history: list[dict[str, str]] | None,
|
history: list[dict[str, str]] | None,
|
||||||
|
thinking_is_progress: bool,
|
||||||
) -> None:
|
) -> None:
|
||||||
reply = "\n\n".join(
|
reply = reply_text(message.content, thinking_is_progress=thinking_is_progress)
|
||||||
getattr(b, "text", "")
|
|
||||||
for b in message.content
|
|
||||||
if getattr(b, "type", "") == "text"
|
|
||||||
).strip()
|
|
||||||
tail = capture.synthesized_messages or [
|
tail = capture.synthesized_messages or [
|
||||||
{"role": "assistant", "content": [{"type": "text", "text": reply}]}
|
{"role": "assistant", "content": [{"type": "text", "text": reply}]}
|
||||||
]
|
]
|
||||||
@@ -275,10 +265,8 @@ async def _sse(
|
|||||||
) -> AsyncIterator[bytes]:
|
) -> AsyncIterator[bytes]:
|
||||||
r"""Serialize an event stream to SSE, then hand the assembled ``Message`` on.
|
r"""Serialize an event stream to SSE, then hand the assembled ``Message`` on.
|
||||||
|
|
||||||
Each event becomes ``event: <type>\ndata: <json>\n\n`` - the shape
|
Wire format is ``event: <type>\ndata: <json>\n\n``; mid-stream errors
|
||||||
the Anthropic SDK's SSE decoder expects. Errors mid-stream are
|
become a synthetic ``error`` event instead of a hung connection.
|
||||||
swallowed into a synthetic ``error`` event so the client sees the
|
|
||||||
failure rather than a hung connection.
|
|
||||||
"""
|
"""
|
||||||
acc = StreamAccumulator()
|
acc = StreamAccumulator()
|
||||||
try:
|
try:
|
||||||
@@ -302,7 +290,7 @@ async def _sse(
|
|||||||
async def _broadcast_turn(
|
async def _broadcast_turn(
|
||||||
runtime: GatewayRuntime,
|
runtime: GatewayRuntime,
|
||||||
*,
|
*,
|
||||||
agent_name: str,
|
agent: BaseAgent,
|
||||||
input_messages: list[MessageParam],
|
input_messages: list[MessageParam],
|
||||||
output_message: Message,
|
output_message: Message,
|
||||||
system: str | None,
|
system: str | None,
|
||||||
@@ -310,9 +298,10 @@ async def _broadcast_turn(
|
|||||||
if not runtime.turn_log_handlers:
|
if not runtime.turn_log_handlers:
|
||||||
return
|
return
|
||||||
record = TurnRecord(
|
record = TurnRecord(
|
||||||
agent_name=agent_name,
|
agent_name=agent.name,
|
||||||
input_messages=list(input_messages),
|
input_messages=list(input_messages),
|
||||||
output_message=output_message,
|
output_message=output_message,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
system=system,
|
system=system,
|
||||||
source=FRONTEND,
|
source=FRONTEND,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
"""``ApiFrontend`` - the conversations API and event stream (§3.9)."""
|
"""``ApiFrontend`` - the conversations API and event stream."""
|
||||||
|
|
||||||
from beaver_gateway.frontends.api.frontend import ApiFrontend
|
from beaver_gateway.frontends.api.frontend import ApiFrontend
|
||||||
|
|
||||||
|
|||||||
@@ -1,22 +1,13 @@
|
|||||||
"""``ApiFrontend`` - ``/api/*``: conversations, SSE, sessions, usage, limits (§3.9).
|
"""``ApiFrontend`` - ``/api/*``: conversations, SSE, sessions, usage, limits.
|
||||||
|
|
||||||
Bearer scope ``api``; token and audit management need ``admin``. Every
|
Bearer scope ``api``; tokens, secrets and audit need ``admin``. Every
|
||||||
write goes through ``core/conversations``; the frontend only shapes JSON.
|
write goes through conversations/service.py; 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.
|
|
||||||
|
|
||||||
Usage figures come from the ``usage`` table (one row per turn; API-price
|
|
||||||
``cost_usd`` and per-model ``model_usage`` are per-turn deltas of the SDK's
|
|
||||||
cumulative ``ResultMessage`` counters, see ``storage.append_usage``);
|
|
||||||
subscription quotas come from ``rate_limits``
|
|
||||||
(``RateLimitEvent``). The quota covers the whole subscription, so
|
|
||||||
``/api/limits`` puts the gateway's own spend for the window next to it
|
|
||||||
for calibration by eye.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import binascii
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import secrets
|
import secrets
|
||||||
@@ -29,20 +20,25 @@ from fastapi.middleware.cors import CORSMiddleware
|
|||||||
from fastapi.responses import JSONResponse, StreamingResponse
|
from fastapi.responses import JSONResponse, StreamingResponse
|
||||||
from sqlmodel import col, select
|
from sqlmodel import col, select
|
||||||
|
|
||||||
from beaver_gateway.core import audit
|
from beaver_gateway.agents.claude import ClaudeAgent
|
||||||
from beaver_gateway.core.auth import VALID_SCOPES, hash_token
|
from beaver_gateway.agents.prompts import assemble
|
||||||
from beaver_gateway.core.conversations import SEEDS, implied_title
|
from beaver_gateway.backends.transcript import prompt_count
|
||||||
from beaver_gateway.core.injects import URGENCY
|
from beaver_gateway.conversations.context import compose
|
||||||
from beaver_gateway.core.kinds import Kind, as_kind
|
from beaver_gateway.conversations.injects import URGENCY
|
||||||
from beaver_gateway.frontends._auth import require_token
|
from beaver_gateway.conversations.kinds import Kind, as_kind
|
||||||
from beaver_gateway.frontends._sse import (
|
from beaver_gateway.conversations.service import SEEDS, implied_title
|
||||||
|
from beaver_gateway.frontends.base import Frontend
|
||||||
|
from beaver_gateway.frontends.bearer import require_token
|
||||||
|
from beaver_gateway.frontends.sse import (
|
||||||
KEEPALIVE,
|
KEEPALIVE,
|
||||||
SSE_HEADERS,
|
SSE_HEADERS,
|
||||||
events_with_heartbeat,
|
events_with_heartbeat,
|
||||||
sse_pack,
|
sse_pack,
|
||||||
)
|
)
|
||||||
from beaver_gateway.frontends._urls import frontend_url
|
from beaver_gateway.frontends.urls import frontend_url
|
||||||
from beaver_gateway.frontends.base import Frontend
|
from beaver_gateway.keystore import SecretsDisabledError
|
||||||
|
from beaver_gateway.security import audit
|
||||||
|
from beaver_gateway.security.auth import hash_token
|
||||||
from beaver_gateway.storage import (
|
from beaver_gateway.storage import (
|
||||||
create_token,
|
create_token,
|
||||||
list_audit_records,
|
list_audit_records,
|
||||||
@@ -53,17 +49,20 @@ from beaver_gateway.storage.models import (
|
|||||||
Conversation,
|
Conversation,
|
||||||
InjectQueueItem,
|
InjectQueueItem,
|
||||||
RateLimit,
|
RateLimit,
|
||||||
|
Secret,
|
||||||
Token,
|
Token,
|
||||||
Usage,
|
Usage,
|
||||||
)
|
)
|
||||||
|
from beaver_gateway.vault.links import LinkIndex
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Iterable, Sequence
|
from collections.abc import AsyncIterator, Iterable, Sequence
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from beaver_gateway.core.conversations import Conversations
|
from beaver_gateway.backends.sessions import Session, SessionPool
|
||||||
from beaver_gateway.core.scheduler import Scheduler
|
from beaver_gateway.conversations.service import Conversations
|
||||||
from beaver_gateway.frontends.base import GatewayRuntime
|
from beaver_gateway.frontends.base import GatewayRuntime
|
||||||
|
from beaver_gateway.jobs.scheduler import Job, Scheduler
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.frontends.api")
|
_log = logging.getLogger("beaver_gateway.frontends.api")
|
||||||
|
|
||||||
@@ -81,6 +80,12 @@ WINDOWS: dict[str, timedelta] = {
|
|||||||
MEMORY_MAX_DEPTH = 12
|
MEMORY_MAX_DEPTH = 12
|
||||||
MEMORY_MAX_FILE = 2_000_000
|
MEMORY_MAX_FILE = 2_000_000
|
||||||
MEMORY_MAX_ENTRIES = 5000
|
MEMORY_MAX_ENTRIES = 5000
|
||||||
|
SEARCH_MIN = 2
|
||||||
|
GRAPH_LIMIT = 120
|
||||||
|
SEARCH_LIMIT = 20
|
||||||
|
TURNS_LIMIT = 10
|
||||||
|
TURNS_MAX = 50
|
||||||
|
SNIPPET_MAX = 160
|
||||||
|
|
||||||
|
|
||||||
class ApiFrontend(Frontend):
|
class ApiFrontend(Frontend):
|
||||||
@@ -96,12 +101,14 @@ class ApiFrontend(Frontend):
|
|||||||
deep_agent: str | None = None,
|
deep_agent: str | None = None,
|
||||||
job_agent: str | None = None,
|
job_agent: str | None = None,
|
||||||
memory_root: Path | None = None,
|
memory_root: Path | None = None,
|
||||||
|
vault_root: Path | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.master_agent = master_agent
|
self.master_agent = master_agent
|
||||||
self.branch_agent = branch_agent
|
self.branch_agent = branch_agent
|
||||||
self.deep_agent = deep_agent
|
self.deep_agent = deep_agent
|
||||||
self.job_agent = job_agent
|
self.job_agent = job_agent
|
||||||
self.memory_root = memory_root.resolve() if memory_root is not None else None
|
self.memory_root = memory_root.resolve() if memory_root is not None else None
|
||||||
|
self.vault_root = vault_root.resolve() if vault_root is not None else None
|
||||||
self._app: FastAPI | None = None
|
self._app: FastAPI | None = None
|
||||||
|
|
||||||
def agent_for(self, kind: Kind) -> str | None:
|
def agent_for(self, kind: Kind) -> str | None:
|
||||||
@@ -116,14 +123,22 @@ class ApiFrontend(Frontend):
|
|||||||
if runtime.conversations is None or runtime.bus is None:
|
if runtime.conversations is None or runtime.bus is None:
|
||||||
msg = "ApiFrontend needs runtime.conversations and runtime.bus"
|
msg = "ApiFrontend needs runtime.conversations and runtime.bus"
|
||||||
raise RuntimeError(msg)
|
raise RuntimeError(msg)
|
||||||
self._app = build_app(runtime, memory_root=self.memory_root)
|
self._app = build_app(
|
||||||
|
runtime, memory_root=self.memory_root, vault_root=self.vault_root
|
||||||
|
)
|
||||||
|
|
||||||
def app(self) -> FastAPI | None:
|
def app(self) -> FastAPI | None:
|
||||||
return self._app
|
return self._app
|
||||||
|
|
||||||
|
|
||||||
def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> FastAPI: # noqa: PLR0915
|
def build_app( # noqa: PLR0915
|
||||||
|
runtime: GatewayRuntime,
|
||||||
|
*,
|
||||||
|
memory_root: Path | None = None,
|
||||||
|
vault_root: Path | None = None,
|
||||||
|
) -> FastAPI:
|
||||||
app = FastAPI(title="beaver-gateway / API")
|
app = FastAPI(title="beaver-gateway / API")
|
||||||
|
links = LinkIndex(vault_root) if vault_root is not None else None
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
allow_origins=["*"],
|
allow_origins=["*"],
|
||||||
@@ -209,6 +224,7 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
k: fe.agent_for(k) for k in fe.kinds if fe.agent_for(k)
|
k: fe.agent_for(k) for k in fe.kinds if fe.agent_for(k)
|
||||||
},
|
},
|
||||||
"path": fe.path,
|
"path": fe.path,
|
||||||
|
"scope": fe.scope,
|
||||||
"url": frontend_url(request, runtime, fe),
|
"url": frontend_url(request, runtime, fe),
|
||||||
}
|
}
|
||||||
for fe in runtime.frontends
|
for fe in runtime.frontends
|
||||||
@@ -229,17 +245,49 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
firsts = await conversations.first_user_texts(
|
firsts = await conversations.first_user_texts(
|
||||||
cast("int", r.id) for r in rows if not r.title
|
cast("int", r.id) for r in rows if not r.title
|
||||||
)
|
)
|
||||||
|
contexts = await conversations.context_tokens_by_id(r.external_id for r in rows)
|
||||||
|
public_ids = {r.id: r.external_id for r in rows}
|
||||||
return {
|
return {
|
||||||
"conversations": [
|
"conversations": [
|
||||||
{
|
{
|
||||||
**conversations.public(r),
|
**conversations.public(r),
|
||||||
"title": r.title or implied_title(firsts.get(cast("int", r.id))),
|
"title": r.title or implied_title(firsts.get(cast("int", r.id))),
|
||||||
"last_item": _queue_item(latest.get(cast("int", r.id))),
|
"last_item": _queue_item(latest.get(cast("int", r.id))),
|
||||||
|
"context_tokens": contexts.get(r.external_id, 0),
|
||||||
|
"parent": public_ids.get(r.parent_id) if r.parent_id else None,
|
||||||
}
|
}
|
||||||
for r in rows
|
for r in rows
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@app.get("/search")
|
||||||
|
async def search(request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
|
query = (request.query_params.get("q") or "").strip()
|
||||||
|
if len(query) < SEARCH_MIN:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_400_BAD_REQUEST, f"q must be at least {SEARCH_MIN} chars"
|
||||||
|
)
|
||||||
|
limit = min(query_int(request, "limit", SEARCH_LIMIT), SEARCH_LIMIT)
|
||||||
|
rows = await conversations.search(query, limit=limit)
|
||||||
|
firsts = await conversations.first_user_texts(
|
||||||
|
cast("int", r.id) for r in rows if not r.title
|
||||||
|
)
|
||||||
|
contexts = await conversations.context_tokens_by_id(r.external_id for r in rows)
|
||||||
|
files = _grep_memory(memory_root, query, limit=limit) if memory_root else []
|
||||||
|
return {
|
||||||
|
"query": query,
|
||||||
|
"conversations": [
|
||||||
|
{
|
||||||
|
**conversations.public(r),
|
||||||
|
"title": r.title or implied_title(firsts.get(cast("int", r.id))),
|
||||||
|
"context_tokens": contexts.get(r.external_id, 0),
|
||||||
|
}
|
||||||
|
for r in rows
|
||||||
|
],
|
||||||
|
"files": files,
|
||||||
|
}
|
||||||
|
|
||||||
@app.post("/conversations", status_code=status.HTTP_201_CREATED)
|
@app.post("/conversations", status_code=status.HTTP_201_CREATED)
|
||||||
async def create_conversation(request: Request) -> dict[str, Any]:
|
async def create_conversation(request: Request) -> dict[str, Any]:
|
||||||
token = await require_token(request, runtime, scope=SCOPE)
|
token = await require_token(request, runtime, scope=SCOPE)
|
||||||
@@ -312,6 +360,45 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
conv = await conv_of(public_id)
|
conv = await conv_of(public_id)
|
||||||
return {"id": conv.external_id, "messages": await conversations.history(conv)}
|
return {"id": conv.external_id, "messages": await conversations.history(conv)}
|
||||||
|
|
||||||
|
@app.get("/conversations/{public_id}/context")
|
||||||
|
async def get_context(public_id: str, request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
|
conv = await conv_of(public_id)
|
||||||
|
agent = runtime.agents.get(conv.agent_name)
|
||||||
|
if not isinstance(agent, ClaudeAgent):
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_404_NOT_FOUND, f"no claude agent behind {public_id}"
|
||||||
|
)
|
||||||
|
entries = list(await conversations.entries(conv))
|
||||||
|
for sub in await conversations.subpaths(conv):
|
||||||
|
entries.extend(await conversations.entries(conv, subpath=sub))
|
||||||
|
out = compose(
|
||||||
|
agent,
|
||||||
|
as_kind(conv.kind),
|
||||||
|
entries,
|
||||||
|
context_tokens=await conversations.context_tokens(conv),
|
||||||
|
)
|
||||||
|
out["id"] = conv.external_id
|
||||||
|
out["turns"] = prompt_count(entries)
|
||||||
|
if request.query_params.get("prompt"):
|
||||||
|
sources = agent.prompt_for(as_kind(conv.kind))
|
||||||
|
out["prompt_text"] = (
|
||||||
|
assemble(sources) if sources is not None else agent.system_prompt
|
||||||
|
)
|
||||||
|
return out
|
||||||
|
|
||||||
|
@app.get("/vault/graph")
|
||||||
|
async def vault_graph(request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
|
if links is None:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_404_NOT_FOUND,
|
||||||
|
"no vault root: set ApiFrontend(vault_root=...)",
|
||||||
|
)
|
||||||
|
paths = [p for p in request.query_params.getlist("path") if p]
|
||||||
|
limit = min(query_int(request, "limit", GRAPH_LIMIT), GRAPH_LIMIT)
|
||||||
|
return {"root": str(links.root), **links.neighbours(paths, limit=limit)}
|
||||||
|
|
||||||
@app.get("/conversations/{public_id}/entries")
|
@app.get("/conversations/{public_id}/entries")
|
||||||
async def get_entries(public_id: str, request: Request) -> dict[str, Any]:
|
async def get_entries(public_id: str, request: Request) -> dict[str, Any]:
|
||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
@@ -497,6 +584,42 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
"text": result.text,
|
"text": result.text,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@app.get("/conversations/{public_id}/turns")
|
||||||
|
async def get_turns(public_id: str, request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
|
conv = await conv_of(public_id)
|
||||||
|
limit = min(query_int(request, "limit", TURNS_LIMIT), TURNS_MAX)
|
||||||
|
return {
|
||||||
|
"id": conv.external_id,
|
||||||
|
"turns": await conversations.turns(conv, limit=limit),
|
||||||
|
}
|
||||||
|
|
||||||
|
@app.post("/conversations/{public_id}/rewind")
|
||||||
|
async def post_rewind(public_id: str, request: Request) -> dict[str, Any]:
|
||||||
|
token = await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
conv = await conv_of(public_id)
|
||||||
|
turn = text_of(await body_of(request), "turn")
|
||||||
|
try:
|
||||||
|
result = await conversations.rewind(conv, turn)
|
||||||
|
except LookupError as exc:
|
||||||
|
raise HTTPException(status.HTTP_404_NOT_FOUND, str(exc)) from exc
|
||||||
|
except RuntimeError as exc:
|
||||||
|
raise HTTPException(status.HTTP_409_CONFLICT, str(exc)) from exc
|
||||||
|
out = {
|
||||||
|
"turns": result.turns,
|
||||||
|
"session_id": result.conversation.session_id,
|
||||||
|
"previous_session": result.previous_session,
|
||||||
|
}
|
||||||
|
await audit.log(
|
||||||
|
runtime,
|
||||||
|
actor=f"token:{token}",
|
||||||
|
kind="api_rewind",
|
||||||
|
agent_name=conv.agent_name,
|
||||||
|
conversation=conv.external_id,
|
||||||
|
**out,
|
||||||
|
)
|
||||||
|
return {"id": conv.external_id, **out}
|
||||||
|
|
||||||
@app.post("/conversations/{public_id}/bind")
|
@app.post("/conversations/{public_id}/bind")
|
||||||
async def post_bind(public_id: str, request: Request) -> dict[str, Any]:
|
async def post_bind(public_id: str, request: Request) -> dict[str, Any]:
|
||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
@@ -553,15 +676,74 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
return _sse(runtime, conversation_id=None)
|
return _sse(runtime, conversation_id=None)
|
||||||
|
|
||||||
|
def pool_of() -> SessionPool:
|
||||||
|
if runtime.pool is None:
|
||||||
|
raise HTTPException(status.HTTP_404_NOT_FOUND, "no session pool")
|
||||||
|
return runtime.pool
|
||||||
|
|
||||||
|
def session_of(key: str) -> Session:
|
||||||
|
session = pool_of().get(key)
|
||||||
|
if session is None:
|
||||||
|
raise HTTPException(status.HTTP_404_NOT_FOUND, f"no live session {key}")
|
||||||
|
return session
|
||||||
|
|
||||||
@app.get("/sessions")
|
@app.get("/sessions")
|
||||||
async def sessions(request: Request) -> dict[str, Any]:
|
async def sessions(request: Request) -> dict[str, Any]:
|
||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
pool = runtime.pool
|
if runtime.pool is None:
|
||||||
return {
|
return {
|
||||||
"rss": pool.rss() if pool is not None else None,
|
"rss": None,
|
||||||
"rss_limit": pool.rss_limit if pool is not None else None,
|
"rss_limit": None,
|
||||||
"sessions": pool.snapshot() if pool is not None else [],
|
"cpu": None,
|
||||||
}
|
"sessions": [],
|
||||||
|
"others": [],
|
||||||
|
}
|
||||||
|
return runtime.pool.census()
|
||||||
|
|
||||||
|
@app.get("/sessions/{key}/stderr")
|
||||||
|
async def session_stderr(key: str, request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
return {"key": key, "lines": list(session_of(key).stderr)}
|
||||||
|
|
||||||
|
@app.post("/sessions/{key}/interrupt")
|
||||||
|
async def session_interrupt(key: str, request: Request) -> dict[str, Any]:
|
||||||
|
token = await require_token(request, runtime, scope=SCOPE)
|
||||||
|
session = session_of(key)
|
||||||
|
interrupted = await pool_of().interrupt(key)
|
||||||
|
await audit.log(
|
||||||
|
runtime,
|
||||||
|
actor=f"token:{token}",
|
||||||
|
kind="session_interrupt",
|
||||||
|
agent_name=session.agent,
|
||||||
|
session=key,
|
||||||
|
interrupted=interrupted,
|
||||||
|
)
|
||||||
|
return {"key": key, "interrupted": interrupted}
|
||||||
|
|
||||||
|
@app.post("/sessions/{key}/close")
|
||||||
|
async def session_close(key: str, request: Request) -> dict[str, Any]:
|
||||||
|
token = await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
session = session_of(key)
|
||||||
|
await pool_of().close(key)
|
||||||
|
await audit.log(
|
||||||
|
runtime,
|
||||||
|
actor=f"token:{token}",
|
||||||
|
kind="session_close",
|
||||||
|
agent_name=session.agent,
|
||||||
|
session=key,
|
||||||
|
busy=session.busy,
|
||||||
|
)
|
||||||
|
return {"key": key, "closed": True}
|
||||||
|
|
||||||
|
@app.post("/processes/{pid}/kill")
|
||||||
|
async def process_kill(pid: int, request: Request) -> dict[str, Any]:
|
||||||
|
token = await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
if not await pool_of().kill(pid):
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_404_NOT_FOUND, f"no process {pid} under the gateway"
|
||||||
|
)
|
||||||
|
await audit.log(runtime, actor=f"token:{token}", kind="process_kill", pid=pid)
|
||||||
|
return {"pid": pid, "killed": True}
|
||||||
|
|
||||||
@app.get("/schedules")
|
@app.get("/schedules")
|
||||||
async def schedules(request: Request) -> dict[str, Any]:
|
async def schedules(request: Request) -> dict[str, Any]:
|
||||||
@@ -580,15 +762,27 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
return await scheduler_of().snapshot()
|
return await scheduler_of().snapshot()
|
||||||
|
|
||||||
|
def job_of(scheduler: Scheduler, name: str) -> Job:
|
||||||
|
job = scheduler.job(name)
|
||||||
|
if job is None:
|
||||||
|
raise HTTPException(status.HTTP_404_NOT_FOUND, f"no job {name!r}")
|
||||||
|
return job
|
||||||
|
|
||||||
@app.post("/jobs/{name}/run", status_code=status.HTTP_202_ACCEPTED)
|
@app.post("/jobs/{name}/run", status_code=status.HTTP_202_ACCEPTED)
|
||||||
async def run_job(name: str, request: Request) -> dict[str, Any]:
|
async def run_job(name: str, request: Request) -> dict[str, Any]:
|
||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
scheduler = scheduler_of()
|
scheduler = scheduler_of()
|
||||||
job = scheduler.job(name)
|
job = job_of(scheduler, name)
|
||||||
if job is None:
|
|
||||||
raise HTTPException(status.HTTP_404_NOT_FOUND, f"no job {name!r}")
|
|
||||||
return {"job": await scheduler.trigger(job, await body_of(request))}
|
return {"job": await scheduler.trigger(job, await body_of(request))}
|
||||||
|
|
||||||
|
@app.get("/jobs/{name}/runs")
|
||||||
|
async def job_runs(name: str, request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
|
scheduler = scheduler_of()
|
||||||
|
job = job_of(scheduler, name)
|
||||||
|
limit = min(int(request.query_params.get("limit", 50)), 500)
|
||||||
|
return {"job": job.name, "runs": await scheduler.runs(job.name, limit)}
|
||||||
|
|
||||||
@app.delete("/jobs/queue/{job_id}")
|
@app.delete("/jobs/queue/{job_id}")
|
||||||
async def cancel_job(job_id: int, request: Request) -> dict[str, Any]:
|
async def cancel_job(job_id: int, request: Request) -> dict[str, Any]:
|
||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
@@ -646,22 +840,51 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
await require_token(request, runtime, scope=SCOPE)
|
await require_token(request, runtime, scope=SCOPE)
|
||||||
rows = await conversations.rate_limits(limit=200)
|
rows = await conversations.rate_limits(limit=200)
|
||||||
latest: dict[str, RateLimit] = {}
|
latest: dict[str, RateLimit] = {}
|
||||||
|
unified: dict[str, tuple[RateLimit, dict[str, Any]]] = {}
|
||||||
for row in rows:
|
for row in rows:
|
||||||
latest.setdefault(row.window, row)
|
latest.setdefault(row.window, row)
|
||||||
|
for name, info in _unified_windows(row).items():
|
||||||
|
unified.setdefault(name, (row, info))
|
||||||
now = datetime.now(UTC)
|
now = datetime.now(UTC)
|
||||||
windows = []
|
windows = []
|
||||||
for window, row in latest.items():
|
for window in dict.fromkeys([*latest, *unified]):
|
||||||
length = WINDOWS.get(window)
|
row = latest.get(window)
|
||||||
since = (
|
seen = unified.get(window)
|
||||||
_aware(row.resets_at) - length
|
public = _limit_public(row) if row else _limit_blank(window, seen)
|
||||||
if length is not None and row.resets_at is not None
|
resets = row.resets_at if row else None
|
||||||
else _aware(row.ts)
|
if seen and _newer_figure(seen, row):
|
||||||
)
|
carrier, info = seen
|
||||||
|
public["utilization"] = info.get("utilization")
|
||||||
|
public["reported_at"] = _iso(carrier.ts)
|
||||||
|
resets = _unix_dt(info.get("resetsAt")) or resets
|
||||||
|
public["resets_at"] = _iso(resets)
|
||||||
|
length = _window_length(window)
|
||||||
|
anchor = row if row is not None else (seen[0] if seen else None)
|
||||||
|
if length is not None and resets is not None:
|
||||||
|
since = _aware(resets) - length
|
||||||
|
else:
|
||||||
|
since = _aware(anchor.ts) if anchor else now
|
||||||
gateway = _sum_usage(await _usage_rows(runtime, since, now))
|
gateway = _sum_usage(await _usage_rows(runtime, since, now))
|
||||||
gateway["since"] = since.isoformat(timespec="seconds")
|
gateway["since"] = since.isoformat(timespec="seconds")
|
||||||
windows.append({**_limit_public(row), "gateway": gateway})
|
known = next(
|
||||||
|
(
|
||||||
|
r
|
||||||
|
for r in rows
|
||||||
|
if r.window == window
|
||||||
|
and r.utilization is not None
|
||||||
|
and _aware(r.ts) >= since
|
||||||
|
),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
windows.append(
|
||||||
|
{
|
||||||
|
**public,
|
||||||
|
"gateway": gateway,
|
||||||
|
"last_known": _limit_public(known) if known else None,
|
||||||
|
}
|
||||||
|
)
|
||||||
windows.sort(
|
windows.sort(
|
||||||
key=lambda w: (WINDOWS.get(w["window"], timedelta.max), w["window"])
|
key=lambda w: (_window_length(w["window"]) or timedelta.max, w["window"])
|
||||||
)
|
)
|
||||||
return {"windows": windows, "history": [_limit_public(r) for r in rows[:50]]}
|
return {"windows": windows, "history": [_limit_public(r) for r in rows[:50]]}
|
||||||
|
|
||||||
@@ -694,6 +917,19 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
"content": content,
|
"content": content,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@app.get("/scopes")
|
||||||
|
async def scopes(request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
by_frontend = {
|
||||||
|
fe.scope: fe.name or fe.path or "" for fe in runtime.frontends if fe.scope
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
"scopes": [
|
||||||
|
{"scope": s, "frontend": by_frontend.get(s)}
|
||||||
|
for s in sorted(runtime.scopes)
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
@app.get("/tokens")
|
@app.get("/tokens")
|
||||||
async def tokens(request: Request) -> dict[str, Any]:
|
async def tokens(request: Request) -> dict[str, Any]:
|
||||||
await require_token(request, runtime, scope=ADMIN_SCOPE)
|
await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
@@ -708,7 +944,7 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
data = await body_of(request)
|
data = await body_of(request)
|
||||||
name = text_of(data, "name").strip()
|
name = text_of(data, "name").strip()
|
||||||
scope = str(data.get("scope") or "*")
|
scope = str(data.get("scope") or "*")
|
||||||
if scope not in VALID_SCOPES:
|
if scope not in runtime.scopes:
|
||||||
raise HTTPException(status.HTTP_400_BAD_REQUEST, f"invalid scope {scope!r}")
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, f"invalid scope {scope!r}")
|
||||||
plaintext = secrets.token_urlsafe(32)
|
plaintext = secrets.token_urlsafe(32)
|
||||||
async with runtime.db.session() as session:
|
async with runtime.db.session() as session:
|
||||||
@@ -746,6 +982,58 @@ def build_app(runtime: GatewayRuntime, *, memory_root: Path | None = None) -> Fa
|
|||||||
)
|
)
|
||||||
return {"id": token_id, "revoked": True}
|
return {"id": token_id, "revoked": True}
|
||||||
|
|
||||||
|
@app.get("/secrets")
|
||||||
|
async def secrets_list(request: Request) -> dict[str, Any]:
|
||||||
|
await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
store = runtime.secrets
|
||||||
|
rows = await store.rows() if store is not None else []
|
||||||
|
return {
|
||||||
|
"enabled": store is not None and store.enabled,
|
||||||
|
"secrets": [_secret_public(r) for r in rows],
|
||||||
|
}
|
||||||
|
|
||||||
|
@app.put("/secrets/{name}")
|
||||||
|
async def secret_put(name: str, request: Request) -> dict[str, Any]:
|
||||||
|
actor = await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
store = runtime.secrets
|
||||||
|
if store is None or not store.enabled:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_503_SERVICE_UNAVAILABLE, "SECRETS_KEY is not set"
|
||||||
|
)
|
||||||
|
data = await body_of(request)
|
||||||
|
description = data.get("description")
|
||||||
|
payload, filename = _secret_payload(data)
|
||||||
|
try:
|
||||||
|
row = await store.put(
|
||||||
|
name,
|
||||||
|
description=str(description) if description is not None else None,
|
||||||
|
data=payload,
|
||||||
|
filename=filename,
|
||||||
|
)
|
||||||
|
except SecretsDisabledError as exc:
|
||||||
|
raise HTTPException(status.HTTP_503_SERVICE_UNAVAILABLE, str(exc)) from exc
|
||||||
|
except ValueError as exc:
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(exc)) from exc
|
||||||
|
await audit.log(
|
||||||
|
runtime,
|
||||||
|
actor=f"token:{actor}",
|
||||||
|
kind="secret_set",
|
||||||
|
name=name,
|
||||||
|
replaced=payload is not None,
|
||||||
|
)
|
||||||
|
return {"secret": _secret_public(row)}
|
||||||
|
|
||||||
|
@app.delete("/secrets/{name}")
|
||||||
|
async def secret_delete(name: str, request: Request) -> dict[str, Any]:
|
||||||
|
actor = await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
|
store = runtime.secrets
|
||||||
|
if store is None or not await store.delete(name):
|
||||||
|
raise HTTPException(status.HTTP_404_NOT_FOUND, f"no secret {name}")
|
||||||
|
await audit.log(
|
||||||
|
runtime, actor=f"token:{actor}", kind="secret_delete", name=name
|
||||||
|
)
|
||||||
|
return {"name": name, "deleted": True}
|
||||||
|
|
||||||
@app.get("/audit")
|
@app.get("/audit")
|
||||||
async def audit_list(request: Request) -> dict[str, Any]:
|
async def audit_list(request: Request) -> dict[str, Any]:
|
||||||
await require_token(request, runtime, scope=ADMIN_SCOPE)
|
await require_token(request, runtime, scope=ADMIN_SCOPE)
|
||||||
@@ -943,6 +1231,59 @@ def _queue_item(item: InjectQueueItem | None) -> dict[str, Any] | None:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _window_length(name: str) -> timedelta | None:
|
||||||
|
if name.startswith("five_hour"):
|
||||||
|
return timedelta(hours=5)
|
||||||
|
if "seven_day" in name:
|
||||||
|
return timedelta(days=7)
|
||||||
|
return WINDOWS.get(name)
|
||||||
|
|
||||||
|
|
||||||
|
def _unix_dt(value: Any) -> datetime | None:
|
||||||
|
if isinstance(value, int | float) and value > 0:
|
||||||
|
return datetime.fromtimestamp(value, tz=UTC)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _unified_windows(row: RateLimit) -> dict[str, dict[str, Any]]:
|
||||||
|
"""Every window the event carried a figure for, not just the one it was about."""
|
||||||
|
raw = row.raw if isinstance(row.raw, dict) else {}
|
||||||
|
found = raw.get("unifiedWindows")
|
||||||
|
if not isinstance(found, dict):
|
||||||
|
return {}
|
||||||
|
return {
|
||||||
|
str(name): info
|
||||||
|
for name, info in found.items()
|
||||||
|
if isinstance(info, dict) and info.get("utilization") is not None
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _newer_figure(
|
||||||
|
seen: tuple[RateLimit, dict[str, Any]], row: RateLimit | None
|
||||||
|
) -> bool:
|
||||||
|
if row is None or row.utilization is None:
|
||||||
|
return True
|
||||||
|
return seen[0].id != row.id and _aware(seen[0].ts) > _aware(row.ts)
|
||||||
|
|
||||||
|
|
||||||
|
def _limit_blank(
|
||||||
|
window: str, seen: tuple[RateLimit, dict[str, Any]] | None
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
carrier = seen[0] if seen else None
|
||||||
|
return {
|
||||||
|
"id": None,
|
||||||
|
"ts": _iso(carrier.ts) if carrier else None,
|
||||||
|
"window": window,
|
||||||
|
"status": "allowed",
|
||||||
|
"utilization": None,
|
||||||
|
"resets_at": None,
|
||||||
|
"overage_status": None,
|
||||||
|
"overage_resets_at": None,
|
||||||
|
"agent": carrier.agent_name if carrier else None,
|
||||||
|
"session_id": carrier.session_id if carrier else None,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def _limit_public(row: RateLimit) -> dict[str, Any]:
|
def _limit_public(row: RateLimit) -> dict[str, Any]:
|
||||||
return {
|
return {
|
||||||
"id": row.id,
|
"id": row.id,
|
||||||
@@ -969,6 +1310,46 @@ def _token_public(row: Token) -> dict[str, Any]:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _secret_public(row: Secret) -> dict[str, Any]:
|
||||||
|
return {
|
||||||
|
"name": row.name,
|
||||||
|
"description": row.description,
|
||||||
|
"kind": "file" if row.filename else "value",
|
||||||
|
"filename": row.filename,
|
||||||
|
"size": row.size,
|
||||||
|
"created_at": _iso(row.created_at),
|
||||||
|
"updated_at": _iso(row.updated_at),
|
||||||
|
"used_at": _iso(row.used_at),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _secret_payload(data: dict[str, Any]) -> tuple[bytes | None, str | None]:
|
||||||
|
value, upload = data.get("value"), data.get("file")
|
||||||
|
if value is not None and upload is not None:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_400_BAD_REQUEST, "send `value` or `file`, not both"
|
||||||
|
)
|
||||||
|
if isinstance(value, str):
|
||||||
|
if not value:
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, "`value` is empty")
|
||||||
|
return value.encode(), None
|
||||||
|
if isinstance(upload, dict):
|
||||||
|
filename, content = upload.get("name"), upload.get("content")
|
||||||
|
if not isinstance(filename, str) or not isinstance(content, str):
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_400_BAD_REQUEST, "`file` needs `name` and base64 `content`"
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
return base64.b64decode(content, validate=True), filename
|
||||||
|
except binascii.Error as exc:
|
||||||
|
raise HTTPException(
|
||||||
|
status.HTTP_400_BAD_REQUEST, "`file.content` is not base64"
|
||||||
|
) from exc
|
||||||
|
if value is not None or upload is not None:
|
||||||
|
raise HTTPException(status.HTTP_400_BAD_REQUEST, "`value` must be a string")
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
|
||||||
def _detail(raw: str) -> Any:
|
def _detail(raw: str) -> Any:
|
||||||
try:
|
try:
|
||||||
return json.loads(raw) if raw else {}
|
return json.loads(raw) if raw else {}
|
||||||
@@ -992,6 +1373,35 @@ def _memory_path(root: Path, raw: str) -> Path:
|
|||||||
return target
|
return target
|
||||||
|
|
||||||
|
|
||||||
|
def _grep_memory(root: Path, query: str, *, limit: int) -> list[dict[str, Any]]:
|
||||||
|
needle = query.casefold()
|
||||||
|
hits: list[dict[str, Any]] = []
|
||||||
|
for path in sorted(root.rglob("*")):
|
||||||
|
if len(hits) >= limit:
|
||||||
|
break
|
||||||
|
if not path.is_file() or any(part.startswith(".") for part in path.parts):
|
||||||
|
continue
|
||||||
|
if path.suffix not in (".md", ".txt", ".json", ".yaml", ".yml"):
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
if path.stat().st_size > MEMORY_MAX_FILE:
|
||||||
|
continue
|
||||||
|
text = path.read_text(encoding="utf-8")
|
||||||
|
except (OSError, UnicodeDecodeError):
|
||||||
|
continue
|
||||||
|
rel = str(path.relative_to(root))
|
||||||
|
if needle in rel.casefold():
|
||||||
|
hits.append({"path": rel, "line": 0, "snippet": ""})
|
||||||
|
continue
|
||||||
|
for number, line in enumerate(text.splitlines(), 1):
|
||||||
|
if needle in line.casefold():
|
||||||
|
hits.append(
|
||||||
|
{"path": rel, "line": number, "snippet": line.strip()[:SNIPPET_MAX]}
|
||||||
|
)
|
||||||
|
break
|
||||||
|
return hits
|
||||||
|
|
||||||
|
|
||||||
def _tree(root: Path, directory: Path, *, depth: int) -> list[dict[str, Any]]:
|
def _tree(root: Path, directory: Path, *, depth: int) -> list[dict[str, Any]]:
|
||||||
if depth > MEMORY_MAX_DEPTH:
|
if depth > MEMORY_MAX_DEPTH:
|
||||||
return []
|
return []
|
||||||
|
|||||||
@@ -1,13 +1,7 @@
|
|||||||
"""Frontend ABC + the runtime context handed to ``configure``.
|
"""Frontend ABC and the runtime context handed to ``configure``.
|
||||||
|
|
||||||
A frontend is anything that routes inbound traffic into the gateway: an
|
A frontend routes inbound traffic into the gateway (an HTTP mount, a
|
||||||
HTTP surface mounted under its ``path`` on the single gateway port
|
poller, or both); ``GatewayRuntime`` carries the built state each needs.
|
||||||
(``frontends/root.py``), a poller (Telegram), or both. ``GatewayRuntime``
|
|
||||||
carries everything a frontend may need that isn't user-config: built
|
|
||||||
registries, per-agent backends, and the in-memory token store. The
|
|
||||||
user's ``/config/config.py`` defines a ``Gateway`` (lists); ``cli.main``
|
|
||||||
turns that into a ``GatewayRuntime`` and hands it to each frontend's
|
|
||||||
``configure``.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -16,16 +10,21 @@ from abc import ABC, abstractmethod
|
|||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
from beaver_gateway.frontends.files import SendFileError
|
||||||
|
from beaver_gateway.security.auth import BUILTIN_SCOPES
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
from starlette.types import ASGIApp
|
from starlette.types import ASGIApp
|
||||||
|
|
||||||
|
from beaver_gateway.app import AgentRegistry, McpRegistry
|
||||||
from beaver_gateway.backends.base import Backend
|
from beaver_gateway.backends.base import Backend
|
||||||
from beaver_gateway.core.auth import TokenStore
|
from beaver_gateway.conversations.kinds import Kind
|
||||||
from beaver_gateway.core.kinds import Kind
|
from beaver_gateway.frontends.turn_record import TurnRecord
|
||||||
from beaver_gateway.core.registry import AgentRegistry, McpRegistry
|
from beaver_gateway.keystore import SecretStore
|
||||||
from beaver_gateway.core.turn_record import TurnRecord
|
from beaver_gateway.security.auth import TokenStore
|
||||||
from beaver_gateway.storage import Database
|
from beaver_gateway.storage import Database
|
||||||
from beaver_gateway.storage.models import Conversation, ConversationBinding
|
from beaver_gateway.storage.models import Conversation, ConversationBinding
|
||||||
|
|
||||||
@@ -36,20 +35,8 @@ if TYPE_CHECKING:
|
|||||||
class GatewayRuntime:
|
class GatewayRuntime:
|
||||||
"""Post-build state of the gateway, shared with every frontend.
|
"""Post-build state of the gateway, shared with every frontend.
|
||||||
|
|
||||||
Backends are keyed by **agent name**, not type — one ``RaycastBackend``
|
Backends are keyed by agent name, not type: one backend instance may
|
||||||
instance can serve many ``RaycastAgent`` instances, but the lookup
|
serve several agents, so the indirection lives at lookup time.
|
||||||
site (an inbound request with ``model=<agent.name>``) already has
|
|
||||||
the name in hand, so the indirection lives one step earlier.
|
|
||||||
|
|
||||||
``mcp_internal_urls`` is filled in Phase 2.1: one loopback URL per
|
|
||||||
declared ``McpServer`` so ``ClaudeSdkBackend``
|
|
||||||
can pass them to ``BackendOptions.mcp_servers`` without re-running
|
|
||||||
discovery.
|
|
||||||
|
|
||||||
``db`` (Phase 4.1) is the shared :class:`Database` handle. Phase 4.2
|
|
||||||
will switch ``TokenStore`` to read from it; Phase 4.3 admin/audit
|
|
||||||
write through it. Phase 4.1 only attaches it — existing frontends
|
|
||||||
ignore it.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
agents: AgentRegistry
|
agents: AgentRegistry
|
||||||
@@ -58,68 +45,55 @@ class GatewayRuntime:
|
|||||||
token_store: TokenStore
|
token_store: TokenStore
|
||||||
db: Database
|
db: Database
|
||||||
mcp_internal_urls: Mapping[str, str] = field(default_factory=dict)
|
mcp_internal_urls: Mapping[str, str] = field(default_factory=dict)
|
||||||
# Phase 4.3 — AdminFrontend reads creds + cookie-signing key from
|
|
||||||
# the runtime so the user's ``config.py`` doesn't have to know
|
|
||||||
# anything about env wiring. Defaulted to empty so existing tests /
|
|
||||||
# call sites that don't touch the admin path keep building; the
|
|
||||||
# admin frontend ``configure()`` itself rejects empty values.
|
|
||||||
admin_user: str = ""
|
admin_user: str = ""
|
||||||
|
"""Operator login for the admin console, checked by ``AdminFrontend``."""
|
||||||
admin_pass: str = ""
|
admin_pass: str = ""
|
||||||
session_secret: str = ""
|
session_secret: str = ""
|
||||||
# The full sibling-frontends list, in declaration order. AdminFrontend
|
|
||||||
# uses it to advertise concrete bearer-endpoint URLs (host/port) on
|
|
||||||
# the dashboard so the operator can copy ready-to-use links / curl
|
|
||||||
# snippets. Other frontends ignore it.
|
|
||||||
frontends: Sequence[Frontend] = field(default_factory=tuple)
|
frontends: Sequence[Frontend] = field(default_factory=tuple)
|
||||||
# Frontends that finish a turn (Anthropic Messages, Markdown) iterate
|
"""Every frontend in declaration order, for advertising their URLs."""
|
||||||
# this list and ``await`` each handler with a ``TurnRecord``. Handlers
|
|
||||||
# are appended during ``configure()`` by frontends that want a
|
|
||||||
# cross-frontend chat archive — currently the markdown frontend's
|
|
||||||
# ``log_all_chats`` mode. Handler exceptions are caught at the call
|
|
||||||
# site; they never block the user-visible response.
|
|
||||||
#
|
|
||||||
# The field is typed as ``list[Any]`` rather than the precise
|
|
||||||
# ``list[TurnLogHandler]`` because the alias lives under TYPE_CHECKING
|
|
||||||
# to keep ``anthropic.types`` out of the runtime import graph for
|
|
||||||
# this base module.
|
|
||||||
turn_log_handlers: list[TurnLogHandler] = field(default_factory=list)
|
turn_log_handlers: list[TurnLogHandler] = field(default_factory=list)
|
||||||
# M1b: conversations service, event bus and the shared session pool.
|
"""Called with a ``TurnRecord`` after each turn; failures never reach the user."""
|
||||||
# ``Any`` for the same import-graph reason as above; ``None`` only in
|
|
||||||
# tests that build a runtime without them.
|
|
||||||
conversations: Any = None
|
conversations: Any = None
|
||||||
bus: Any = None
|
bus: Any = None
|
||||||
pool: Any = None
|
pool: Any = None
|
||||||
scheduler: Any = None
|
scheduler: Any = None
|
||||||
# External origin the reverse proxy puts in front of the gateway
|
|
||||||
# (``Gateway.public_url``); ``None`` means "derive from the request".
|
|
||||||
public_url: str | None = None
|
public_url: str | None = None
|
||||||
|
"""``Gateway.public_url``; ``None`` derives the origin from the request."""
|
||||||
|
scopes: frozenset[str] = BUILTIN_SCOPES
|
||||||
|
"""Every token scope this gateway knows: the builtins plus each frontend's."""
|
||||||
|
send_file_roots: tuple[Path, ...] = ()
|
||||||
|
"""Directories ``send_file`` may read, from ``SEND_FILE_ROOTS``."""
|
||||||
|
secrets: SecretStore | None = None
|
||||||
|
"""Sealed secrets for ``secret run``; ``None`` in setups without them."""
|
||||||
|
|
||||||
|
|
||||||
class Frontend(ABC):
|
class Frontend(ABC):
|
||||||
"""Routes inbound traffic into the gateway.
|
"""Routes inbound traffic into the gateway.
|
||||||
|
|
||||||
HTTP frontends set ``path`` and return their ASGI app from ``app()``;
|
HTTP frontends set ``path`` and return their ASGI app from ``app()``;
|
||||||
``cli`` mounts every such app under that path on the one gateway
|
these are mounted under that path on the one gateway port. ``serve()``
|
||||||
port, so ``/anthropic/v1/messages`` reaches the Anthropic frontend's
|
is for non-HTTP work (polling, vault mirrors) and defaults to nothing.
|
||||||
``/v1/messages``. ``serve()`` is for work outside HTTP - polling,
|
``landing`` marks the app that ``/`` redirects to.
|
||||||
vault mirrors - and defaults to nothing. ``landing`` marks the app
|
|
||||||
that ``/`` redirects to (the admin console).
|
``scope`` is the token scope this frontend's routes are gated by; a
|
||||||
|
frontend that names its own scope gets keys nobody else's token opens,
|
||||||
|
and ``POST /api/tokens`` will mint them (see ``security/auth.py``).
|
||||||
|
|
||||||
A frontend that shows conversations declares ``name`` (the binding
|
A frontend that shows conversations declares ``name`` (the binding
|
||||||
key) and ``kinds`` (which conversation kinds it shows);
|
key) and ``kinds`` (which conversation kinds it shows). The first
|
||||||
``core/conversations`` refuses to bind a conversation to a frontend
|
frontend whose ``materialize`` returns a binding is the *home* of
|
||||||
outside its declaration. The first frontend in declaration order
|
that kind, used by ``spawn`` for new conversations; ``agent_for``
|
||||||
whose ``materialize`` returns a binding is the *home* of that kind:
|
names the default agent for a kind. Stateless frontends (MCP, admin)
|
||||||
``spawn`` calls it so a new conversation gets a window (a vault file,
|
leave these at their defaults.
|
||||||
a Telegram topic). ``agent_for`` names the default agent for a kind
|
|
||||||
so callers may omit ``agent``. Stateless frontends (MCP, admin) keep
|
|
||||||
the defaults and stay outside the routing.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
name: str = ""
|
name: str = ""
|
||||||
kinds: tuple[Kind, ...] = ()
|
kinds: tuple[Kind, ...] = ()
|
||||||
path: str | None = None
|
path: str | None = None
|
||||||
landing: bool = False
|
landing: bool = False
|
||||||
|
scope: str | None = None
|
||||||
|
sends_files: bool = False
|
||||||
|
"""Whether ``send_file`` reaches a human through this frontend."""
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def configure(self, runtime: GatewayRuntime) -> None: ...
|
def configure(self, runtime: GatewayRuntime) -> None: ...
|
||||||
@@ -139,3 +113,19 @@ class Frontend(ABC):
|
|||||||
async def mark_closed(self, conv: Conversation) -> bool: # noqa: ARG002
|
async def mark_closed(self, conv: Conversation) -> bool: # noqa: ARG002
|
||||||
"""Show in the window that the conversation is over (a renamed topic)."""
|
"""Show in the window that the conversation is over (a renamed topic)."""
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
async def send_file(
|
||||||
|
self,
|
||||||
|
conv: Conversation, # noqa: ARG002
|
||||||
|
path: str, # noqa: ARG002
|
||||||
|
*,
|
||||||
|
caption: str = "", # noqa: ARG002
|
||||||
|
method: str = "auto", # noqa: ARG002
|
||||||
|
) -> str:
|
||||||
|
"""Put a file in this conversation's window; the answer goes to the model."""
|
||||||
|
msg = f"frontend {self.name!r} cannot send files"
|
||||||
|
raise SendFileError(msg)
|
||||||
|
|
||||||
|
def file_note(self) -> str:
|
||||||
|
"""Roots and ceiling, for the tool description the model reads."""
|
||||||
|
return ""
|
||||||
@@ -1,8 +1,7 @@
|
|||||||
"""Shared bearer-token verification for HTTP frontends.
|
"""Shared bearer-token verification for HTTP frontends.
|
||||||
|
|
||||||
Extracted from ``AnthropicMessagesFrontend`` so the markdown frontend
|
Reused by every bearer-protected frontend instead of duplicating the
|
||||||
(and any future bearer-protected frontend) can reuse one canonical
|
header-parsing logic.
|
||||||
verifier instead of copy-pasting the header-parsing dance.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -23,14 +22,11 @@ __all__ = ["require_token"]
|
|||||||
async def require_token(
|
async def require_token(
|
||||||
request: Request, runtime: GatewayRuntime, *, scope: str
|
request: Request, runtime: GatewayRuntime, *, scope: str
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Verify the request's bearer + scope, return the token's audit name.
|
"""Verify the request's bearer token and scope; return its audit name.
|
||||||
|
|
||||||
Accepts ``X-Api-Key: <token>`` (Anthropic SDK / LibreChat),
|
Checks ``X-Api-Key``, then ``Authorization: Bearer``, then ``?token=``.
|
||||||
``Authorization: Bearer <token>`` (curl, Cursor) and, when neither
|
401 on a missing/unknown token, 403 if the token's scope doesn't cover
|
||||||
header is present, ``?token=<token>`` (Komodo alerters). 401 on missing /
|
``scope``. Bootstrap tokens carry ``"*"`` and pass every scope check.
|
||||||
unknown token; 403 on a known token whose scope doesn't cover
|
|
||||||
``scope``. Bootstrap tokens implicitly carry ``"*"`` and pass every
|
|
||||||
scope check.
|
|
||||||
"""
|
"""
|
||||||
api_key = request.headers.get("x-api-key")
|
api_key = request.headers.get("x-api-key")
|
||||||
authorization = request.headers.get("authorization")
|
authorization = request.headers.get("authorization")
|
||||||
@@ -39,8 +35,6 @@ async def require_token(
|
|||||||
elif authorization:
|
elif authorization:
|
||||||
identity = await runtime.token_store.verify_bearer(authorization)
|
identity = await runtime.token_store.verify_bearer(authorization)
|
||||||
else:
|
else:
|
||||||
# Webhook senders that cannot set headers (Komodo alerters) put the
|
|
||||||
# token in the query string; the URL is not logged with it.
|
|
||||||
qs_token = request.query_params.get("token")
|
qs_token = request.query_params.get("token")
|
||||||
identity = await runtime.token_store.verify(qs_token) if qs_token else None
|
identity = await runtime.token_store.verify(qs_token) if qs_token else None
|
||||||
if identity is None:
|
if identity is None:
|
||||||
@@ -0,0 +1,95 @@
|
|||||||
|
"""What an agent may hand to a human: an allowed path, a ceiling, a method.
|
||||||
|
|
||||||
|
The ceiling and the method names belong to the frontend's API; the check
|
||||||
|
lives here so every frontend refuses the same way, in words the model reads.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import mimetypes
|
||||||
|
from dataclasses import dataclass, replace
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
__all__ = ["AUTO", "DEFAULTS", "METHODS", "SendFileError", "SendFiles", "method_for"]
|
||||||
|
|
||||||
|
AUTO = "auto"
|
||||||
|
METHODS = ("document", "photo", "audio", "voice", "video", "video_note", "animation")
|
||||||
|
|
||||||
|
_BY_SUFFIX = {
|
||||||
|
".jpg": "photo",
|
||||||
|
".jpeg": "photo",
|
||||||
|
".png": "photo",
|
||||||
|
".webp": "photo",
|
||||||
|
".gif": "animation",
|
||||||
|
".mp3": "audio",
|
||||||
|
".flac": "audio",
|
||||||
|
".m4a": "audio",
|
||||||
|
".wav": "audio",
|
||||||
|
".aac": "audio",
|
||||||
|
".ogg": "voice",
|
||||||
|
".oga": "voice",
|
||||||
|
".opus": "voice",
|
||||||
|
".mp4": "video",
|
||||||
|
".m4v": "video",
|
||||||
|
".mov": "video",
|
||||||
|
}
|
||||||
|
_BY_TYPE = {"image": "photo", "audio": "audio", "video": "video"}
|
||||||
|
|
||||||
|
|
||||||
|
class SendFileError(Exception):
|
||||||
|
"""Why a file was not sent, worded for the model that asked."""
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class SendFiles:
|
||||||
|
"""Which files a frontend accepts from an agent and how big they may be."""
|
||||||
|
|
||||||
|
roots: tuple[Path, ...] = ()
|
||||||
|
"""Directories a path must resolve inside; nothing outside is sendable."""
|
||||||
|
max_bytes: int = 50 * 1024 * 1024
|
||||||
|
"""Upload ceiling of the API behind the frontend; the Bot API gives 50 MB."""
|
||||||
|
|
||||||
|
def with_roots(self, *roots: Path) -> SendFiles:
|
||||||
|
return replace(self, roots=tuple(dict.fromkeys((*self.roots, *roots))))
|
||||||
|
|
||||||
|
@property
|
||||||
|
def note(self) -> str:
|
||||||
|
"""Roots and ceiling as one line for the tool description."""
|
||||||
|
where = ", ".join(str(r) for r in self.roots) or "none, so nothing is sendable"
|
||||||
|
return f"paths under {where}, up to {mb(self.max_bytes)} MB"
|
||||||
|
|
||||||
|
def resolve(self, raw: str) -> Path:
|
||||||
|
path = Path(raw.strip()).expanduser().resolve()
|
||||||
|
if not any(path.is_relative_to(r.expanduser().resolve()) for r in self.roots):
|
||||||
|
msg = f"{path} is outside the allowed roots ({self.note}); copy it there"
|
||||||
|
raise SendFileError(msg)
|
||||||
|
if not path.is_file():
|
||||||
|
msg = f"{path} is not a file"
|
||||||
|
raise SendFileError(msg)
|
||||||
|
size = path.stat().st_size
|
||||||
|
if size > self.max_bytes:
|
||||||
|
msg = (
|
||||||
|
f"{path.name} is {mb(size)} MB, over the {mb(self.max_bytes)} MB limit"
|
||||||
|
)
|
||||||
|
raise SendFileError(msg)
|
||||||
|
return path
|
||||||
|
|
||||||
|
|
||||||
|
DEFAULTS = SendFiles()
|
||||||
|
|
||||||
|
|
||||||
|
def method_for(path: Path, method: str = AUTO) -> str:
|
||||||
|
if method and method != AUTO:
|
||||||
|
if method not in METHODS:
|
||||||
|
msg = f"unknown method {method!r}; one of: {', '.join((AUTO, *METHODS))}"
|
||||||
|
raise SendFileError(msg)
|
||||||
|
return method
|
||||||
|
guessed = _BY_SUFFIX.get(path.suffix.lower())
|
||||||
|
if guessed:
|
||||||
|
return guessed
|
||||||
|
mime, _ = mimetypes.guess_type(path.name)
|
||||||
|
return _BY_TYPE.get(mime.split("/")[0], "document") if mime else "document"
|
||||||
|
|
||||||
|
|
||||||
|
def mb(size: int) -> str:
|
||||||
|
return f"{round(size / 1024 / 1024, 2):g}"
|
||||||
@@ -1,13 +1,7 @@
|
|||||||
"""Markdown frontend — turn-by-turn chat archive backed by ``.md`` files.
|
"""Markdown frontend — chat archive backed by ``.md`` files in an Obsidian vault.
|
||||||
|
|
||||||
The user maintains chats as plain markdown files in an Obsidian vault.
|
An Obsidian plugin POSTs ``{filename, content?}`` to ``/chat``; the
|
||||||
A plugin in Obsidian POSTs ``{filename, content?}`` to ``/chat`` and the
|
frontend runs the agent on a trailing user turn and appends the reply.
|
||||||
frontend parses the file, finds the last turn, and runs the agent if
|
|
||||||
the last turn is ``user``. The full response is appended back to the
|
|
||||||
file as an ``### Assistant:`` turn. With ``log_all_chats=True`` the
|
|
||||||
frontend also subscribes to every other frontend's turns and writes
|
|
||||||
them into ``{vault_path}/{logged_subdir}/`` so the vault is the single
|
|
||||||
chronological archive of all conversations.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from beaver_gateway.frontends.markdown.frontend import MarkdownFrontend
|
from beaver_gateway.frontends.markdown.frontend import MarkdownFrontend
|
||||||
|
|||||||
@@ -1,17 +1,7 @@
|
|||||||
"""Cross-frontend chat logger.
|
"""Cross-frontend chat logger.
|
||||||
|
|
||||||
When ``MarkdownFrontend(log_all_chats=True)`` is configured, every turn
|
Mirrors turns completed by other frontends into the vault as ``.md``
|
||||||
completed by any other frontend (currently the Anthropic Messages
|
files, matching a conversation's continuation by content-hash fingerprint.
|
||||||
frontend) is mirrored into the vault as a ``.md`` file. Subsequent
|
|
||||||
turns of the same conversation append to the same file — matched by a
|
|
||||||
content-hash fingerprint stored in YAML frontmatter.
|
|
||||||
|
|
||||||
The fingerprint hashes the message history *before* the new assistant
|
|
||||||
reply. So the next request's input history (which now includes the
|
|
||||||
prior assistant reply) hashes to the value we just persisted —
|
|
||||||
``hash(prev_input + [assistant_reply])`` — and the lookup hits the
|
|
||||||
same file. New conversations (no prior fingerprint match) get a fresh
|
|
||||||
file under ``{vault_path}/{logged_subdir}/{agent_name}/``.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -32,7 +22,7 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
from anthropic.types import MessageParam
|
from anthropic.types import MessageParam
|
||||||
|
|
||||||
from beaver_gateway.core.turn_record import TurnRecord
|
from beaver_gateway.frontends.turn_record import TurnRecord
|
||||||
|
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.frontends.markdown.crossfront")
|
_log = logging.getLogger("beaver_gateway.frontends.markdown.crossfront")
|
||||||
@@ -53,12 +43,8 @@ vault. ``None`` keeps ``{vault}/{logged_subdir}/{agent}/{YYYY-MM-DD}_{slug}.md``
|
|||||||
def fingerprint_messages(messages: Iterable[MessageParam]) -> str:
|
def fingerprint_messages(messages: Iterable[MessageParam]) -> str:
|
||||||
"""Stable, short hex hash of a conversation prefix.
|
"""Stable, short hex hash of a conversation prefix.
|
||||||
|
|
||||||
Built from ``(role, normalized_content)`` pairs only — so the
|
Hashes ``(role, text-only content)`` pairs so differently-shaped message
|
||||||
Markdown frontend's parser-shaped messages (text-only) and the
|
histories that carry the same text still fingerprint identically.
|
||||||
Anthropic frontend's raw ``messages`` payload (which may also be
|
|
||||||
string-only at v1) hash compatibly when they represent the same
|
|
||||||
conversation. Tool blocks / images would diverge, but those aren't
|
|
||||||
in the v1 ingest path.
|
|
||||||
"""
|
"""
|
||||||
h = hashlib.sha1(usedforsecurity=False)
|
h = hashlib.sha1(usedforsecurity=False)
|
||||||
for msg in messages:
|
for msg in messages:
|
||||||
@@ -85,12 +71,8 @@ def fingerprint_messages(messages: Iterable[MessageParam]) -> str:
|
|||||||
class CrossFrontendLogger:
|
class CrossFrontendLogger:
|
||||||
"""Maintains the fingerprint→file map and writes turns to disk.
|
"""Maintains the fingerprint→file map and writes turns to disk.
|
||||||
|
|
||||||
The map is in-process; on startup ``warm_index`` rebuilds it from
|
The map is in-process, rebuilt by ``warm_index`` on startup; all disk
|
||||||
YAML frontmatter of every file under ``logged_subdir``. A miss
|
writes funnel through one lock to sidestep races.
|
||||||
creates a new file, a hit appends to the existing one. All disk
|
|
||||||
work funnels through one ``asyncio.Lock`` because the writes are
|
|
||||||
cheap and serializing them sidesteps a class of races we don't need
|
|
||||||
to think about.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -105,20 +87,13 @@ class CrossFrontendLogger:
|
|||||||
self._index: dict[str, Path] = {}
|
self._index: dict[str, Path] = {}
|
||||||
self._lock = asyncio.Lock()
|
self._lock = asyncio.Lock()
|
||||||
self._chat_path_fn = chat_path
|
self._chat_path_fn = chat_path
|
||||||
# When the user supplies a custom path function, files can land
|
|
||||||
# anywhere in the vault — so we have to scan the whole vault on
|
|
||||||
# startup to rebuild the fingerprint→path map. With the default
|
|
||||||
# layout we can bound the scan to ``_logs/``.
|
|
||||||
self._scan_root = vault_path if chat_path is not None else self._root
|
self._scan_root = vault_path if chat_path is not None else self._root
|
||||||
|
|
||||||
def warm_index(self) -> None:
|
def warm_index(self) -> None:
|
||||||
"""Scan logged files synchronously, populating the fingerprint map.
|
"""Scan logged files synchronously, populating the fingerprint map.
|
||||||
|
|
||||||
Called from ``MarkdownFrontend.configure`` so the map is ready
|
Called before any cross-frontend turn arrives; a custom ``chat_path``
|
||||||
before any cross-frontend turn arrives. ``frontmatter.load``
|
forces a full-vault walk instead of scanning just ``logged_subdir``.
|
||||||
reads only enough of the file to parse the YAML head, so the
|
|
||||||
scan is cheap even on large vaults — but a custom ``log_path``
|
|
||||||
forces a full-vault walk; mention that in the constructor doc.
|
|
||||||
"""
|
"""
|
||||||
if not self._scan_root.exists():
|
if not self._scan_root.exists():
|
||||||
return
|
return
|
||||||
@@ -140,29 +115,19 @@ class CrossFrontendLogger:
|
|||||||
async def handle(self, record: TurnRecord) -> None:
|
async def handle(self, record: TurnRecord) -> None:
|
||||||
"""Append or create a logged file for ``record``.
|
"""Append or create a logged file for ``record``.
|
||||||
|
|
||||||
Records that the markdown frontend itself produced
|
Skips ``source == "markdown"`` records (already on disk); matches
|
||||||
(``source=="markdown"``) are skipped — those already live in the
|
the target file by fingerprinting ``input_messages`` sans the new turn.
|
||||||
user's hand-written file and shouldn't be duplicated into the
|
|
||||||
``_logs`` shadow tree.
|
|
||||||
"""
|
"""
|
||||||
if record.source == "markdown":
|
if record.source == "markdown":
|
||||||
return
|
return
|
||||||
|
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
# ``input_messages`` is the *full* history sent to the backend
|
|
||||||
# (last entry is the new user turn). Match against the prefix
|
|
||||||
# that excludes the new user turn — that's what the previous
|
|
||||||
# write stored as its fingerprint. Empty prefix is the
|
|
||||||
# well-known "brand new chat" sentinel.
|
|
||||||
prefix = record.input_messages[:-1]
|
prefix = record.input_messages[:-1]
|
||||||
prev_fp = fingerprint_messages(prefix) if prefix else None
|
prev_fp = fingerprint_messages(prefix) if prefix else None
|
||||||
target = self._index.get(prev_fp) if prev_fp else None
|
target = self._index.get(prev_fp) if prev_fp else None
|
||||||
if target is None:
|
if target is None:
|
||||||
target = self._new_file_path(record)
|
target = self._new_file_path(record)
|
||||||
|
|
||||||
# Build the full history including the assistant reply; the
|
|
||||||
# new fingerprint matches *that* prefix, so the next user
|
|
||||||
# turn (history grows by one user msg) will hit this file.
|
|
||||||
assistant_msg: MessageParam = {
|
assistant_msg: MessageParam = {
|
||||||
"role": "assistant",
|
"role": "assistant",
|
||||||
"content": _flatten_text(record.output_message),
|
"content": _flatten_text(record.output_message),
|
||||||
@@ -174,21 +139,19 @@ class CrossFrontendLogger:
|
|||||||
existing = target.read_text(encoding="utf-8")
|
existing = target.read_text(encoding="utf-8")
|
||||||
parsed = frontmatter.loads(existing)
|
parsed = frontmatter.loads(existing)
|
||||||
body = strip_trailing_user_scaffold(parsed.content)
|
body = strip_trailing_user_scaffold(parsed.content)
|
||||||
# We append only the *new* user turn (the last one in
|
|
||||||
# input_messages, since prior turns are already on disk)
|
|
||||||
# plus the assistant reply.
|
|
||||||
new_user = record.input_messages[-1]
|
new_user = record.input_messages[-1]
|
||||||
new_block = renderer.render_user_param(new_user)
|
new_block = renderer.render_user_param(new_user)
|
||||||
new_block = renderer.append_to_body(
|
new_block = renderer.append_to_body(
|
||||||
new_block, renderer.render_assistant_message(record.output_message)
|
new_block,
|
||||||
|
renderer.render_assistant_message(
|
||||||
|
record.output_message,
|
||||||
|
thinking_is_progress=record.thinking_is_progress,
|
||||||
|
),
|
||||||
)
|
)
|
||||||
new_body = renderer.append_to_body(body, new_block)
|
new_body = renderer.append_to_body(body, new_block)
|
||||||
metadata = dict(parsed.metadata)
|
metadata = dict(parsed.metadata)
|
||||||
else:
|
else:
|
||||||
# Materialize the whole conversation from scratch.
|
new_body = _render_full_history(record)
|
||||||
new_body = _render_full_history(
|
|
||||||
record.input_messages, record.output_message
|
|
||||||
)
|
|
||||||
metadata = {}
|
metadata = {}
|
||||||
|
|
||||||
new_body = renderer.append_to_body(new_body, renderer.USER_SCAFFOLD)
|
new_body = renderer.append_to_body(new_body, renderer.USER_SCAFFOLD)
|
||||||
@@ -197,21 +160,15 @@ class CrossFrontendLogger:
|
|||||||
metadata["fingerprint"] = new_fp
|
metadata["fingerprint"] = new_fp
|
||||||
metadata["source"] = record.source
|
metadata["source"] = record.source
|
||||||
self._write(target, metadata, new_body)
|
self._write(target, metadata, new_body)
|
||||||
# Maintain the index: drop the old fp (it's stale once we
|
|
||||||
# write the new turn), add the new one.
|
|
||||||
if prev_fp:
|
if prev_fp:
|
||||||
self._index.pop(prev_fp, None)
|
self._index.pop(prev_fp, None)
|
||||||
self._index[new_fp] = target
|
self._index[new_fp] = target
|
||||||
|
|
||||||
# ---- internals -----------------------------------------------------
|
|
||||||
|
|
||||||
def _new_file_path(self, record: TurnRecord) -> Path:
|
def _new_file_path(self, record: TurnRecord) -> Path:
|
||||||
"""Pick a fresh filename for a brand-new conversation.
|
"""Pick a fresh filename for a brand-new conversation.
|
||||||
|
|
||||||
With a user-supplied ``chat_path`` we delegate to it (joining a
|
Delegates to ``chat_path`` if set; otherwise
|
||||||
relative result with the vault root). Without one, we fall back
|
``{logged_subdir}/{agent}/{date}_{hex8}.md``.
|
||||||
to ``{logged_subdir}/{agent}/{date}_{hex8}.md`` and ensure the
|
|
||||||
``.md`` suffix in case the user picks a non-md extension by hand.
|
|
||||||
"""
|
"""
|
||||||
if self._chat_path_fn is not None:
|
if self._chat_path_fn is not None:
|
||||||
result = self._chat_path_fn(
|
result = self._chat_path_fn(
|
||||||
@@ -224,8 +181,6 @@ class CrossFrontendLogger:
|
|||||||
result.parent.mkdir(parents=True, exist_ok=True)
|
result.parent.mkdir(parents=True, exist_ok=True)
|
||||||
return result
|
return result
|
||||||
day = datetime.now(UTC).strftime("%Y-%m-%d")
|
day = datetime.now(UTC).strftime("%Y-%m-%d")
|
||||||
# Short hex from the input hash so two same-day chats sort
|
|
||||||
# stably and don't collide.
|
|
||||||
salt = fingerprint_messages(record.input_messages)[:8]
|
salt = fingerprint_messages(record.input_messages)[:8]
|
||||||
agent_dir = self._root / record.agent_name
|
agent_dir = self._root / record.agent_name
|
||||||
agent_dir.mkdir(parents=True, exist_ok=True)
|
agent_dir.mkdir(parents=True, exist_ok=True)
|
||||||
@@ -238,9 +193,6 @@ class CrossFrontendLogger:
|
|||||||
text = frontmatter.dumps(post) + "\n"
|
text = frontmatter.dumps(post) + "\n"
|
||||||
else:
|
else:
|
||||||
text = body if body.endswith("\n") else body + "\n"
|
text = body if body.endswith("\n") else body + "\n"
|
||||||
# Sync write inside the lock — keeps the implementation tiny;
|
|
||||||
# individual logged turns are small enough that the blocking
|
|
||||||
# write doesn't matter at human conversation rates.
|
|
||||||
path.write_text(text, encoding="utf-8")
|
path.write_text(text, encoding="utf-8")
|
||||||
|
|
||||||
|
|
||||||
@@ -254,10 +206,10 @@ def _flatten_text(message: Any) -> str:
|
|||||||
return "\n\n".join(c for c in chunks if c)
|
return "\n\n".join(c for c in chunks if c)
|
||||||
|
|
||||||
|
|
||||||
def _render_full_history(messages: list[MessageParam], assistant: Any) -> str:
|
def _render_full_history(record: TurnRecord) -> str:
|
||||||
"""Render an entire conversation (used when materializing a new logged file)."""
|
"""Render an entire conversation (used when materializing a new logged file)."""
|
||||||
blocks: list[str] = []
|
blocks: list[str] = []
|
||||||
for msg in messages:
|
for msg in record.input_messages:
|
||||||
role = msg.get("role")
|
role = msg.get("role")
|
||||||
if role == "user":
|
if role == "user":
|
||||||
blocks.append(renderer.render_user_param(msg))
|
blocks.append(renderer.render_user_param(msg))
|
||||||
@@ -265,7 +217,11 @@ def _render_full_history(messages: list[MessageParam], assistant: Any) -> str:
|
|||||||
content = msg.get("content", "")
|
content = msg.get("content", "")
|
||||||
text = content if isinstance(content, str) else _content_to_text(content)
|
text = content if isinstance(content, str) else _content_to_text(content)
|
||||||
blocks.append(f"### Assistant:\n\n{text.strip()}\n" if text.strip() else "")
|
blocks.append(f"### Assistant:\n\n{text.strip()}\n" if text.strip() else "")
|
||||||
blocks.append(renderer.render_assistant_message(assistant))
|
blocks.append(
|
||||||
|
renderer.render_assistant_message(
|
||||||
|
record.output_message, thinking_is_progress=record.thinking_is_progress
|
||||||
|
)
|
||||||
|
)
|
||||||
body = ""
|
body = ""
|
||||||
for block in blocks:
|
for block in blocks:
|
||||||
if not block:
|
if not block:
|
||||||
@@ -277,18 +233,13 @@ def _render_full_history(messages: list[MessageParam], assistant: Any) -> str:
|
|||||||
def strip_trailing_user_scaffold(body: str) -> str:
|
def strip_trailing_user_scaffold(body: str) -> str:
|
||||||
"""Drop a trailing empty ``### User:`` block if present.
|
"""Drop a trailing empty ``### User:`` block if present.
|
||||||
|
|
||||||
Cross-frontend turns aren't typed into the file by the human — they
|
Avoids leaving two ``### User:`` headers in a row when appending a
|
||||||
arrive whole from another frontend. If we leave the previous run's
|
turn that wasn't typed into the file by hand.
|
||||||
scaffold in place, we'd write the new user turn right after an
|
|
||||||
empty marker (visual noise, two ``### User:`` headers in a row).
|
|
||||||
Trim it and let the append flow add a fresh scaffold at the end.
|
|
||||||
"""
|
"""
|
||||||
stripped = body.rstrip()
|
stripped = body.rstrip()
|
||||||
marker = "### User:"
|
marker = "### User:"
|
||||||
if not stripped.endswith(marker):
|
if not stripped.endswith(marker):
|
||||||
return body
|
return body
|
||||||
# Walk back: the scaffold is the marker preceded by either start-of-file
|
|
||||||
# or an HR/blank line. Find the last newline before the marker, cut.
|
|
||||||
head = stripped[: -len(marker)].rstrip()
|
head = stripped[: -len(marker)].rstrip()
|
||||||
if head.endswith("---"):
|
if head.endswith("---"):
|
||||||
head = head[: -len("---")].rstrip()
|
head = head[: -len("---")].rstrip()
|
||||||
|
|||||||
@@ -1,34 +1,7 @@
|
|||||||
"""``MarkdownFrontend`` — chat-via-markdown-files frontend.
|
"""``MarkdownFrontend`` — chat-via-markdown-files frontend.
|
||||||
|
|
||||||
Wires:
|
``POST /chat`` (and SSE ``/chat/stream``) parses the vault file, runs
|
||||||
|
the agent on the trailing user turn, and appends the reply.
|
||||||
* ``POST /chat {filename, content?, agent?}`` — bearer-authenticated
|
|
||||||
trigger. The plugin in Obsidian fires this after the user edits a
|
|
||||||
``.md`` and the file gets synced (or with ``content`` to short-circuit
|
|
||||||
the sync delay). We parse the file, check the last turn — assistant
|
|
||||||
→ no-op, user → run the agent and append.
|
|
||||||
* ``GET /healthz`` — liveness.
|
|
||||||
|
|
||||||
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 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, <vault-relative path>)``; frontmatter carries only ``agent``
|
|
||||||
and ``conversation_id`` (§3.10), tool calls are never rendered. The
|
|
||||||
frontend is the home of ``deep``: ``materialize`` gives a conversation
|
|
||||||
spawned elsewhere its file, and :class:`.mirror.ChatMirror` keeps that
|
|
||||||
file in step with replies produced outside ``/chat``.
|
|
||||||
|
|
||||||
Cross-frontend logging: when ``log_all_chats=True``, ``configure()``
|
|
||||||
registers a handler on ``runtime.turn_log_handlers`` so every other
|
|
||||||
frontend's completed turns also land in the vault. The handler logic
|
|
||||||
lives in :mod:`.crossfront` so this module stays focused on the HTTP
|
|
||||||
shape.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -45,23 +18,10 @@ from fastapi import FastAPI, HTTPException, Request, status
|
|||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
from fastapi.responses import JSONResponse, StreamingResponse
|
from fastapi.responses import JSONResponse, StreamingResponse
|
||||||
|
|
||||||
from beaver_gateway.core import audit
|
from beaver_gateway.backends.capture import TurnCapture
|
||||||
from beaver_gateway.core.conversation_store import (
|
from beaver_gateway.frontends.accumulate import StreamAccumulator
|
||||||
diff_and_fork,
|
|
||||||
load_messages,
|
|
||||||
rewrite_messages,
|
|
||||||
)
|
|
||||||
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.base import Frontend
|
||||||
|
from beaver_gateway.frontends.bearer import require_token
|
||||||
from beaver_gateway.frontends.markdown import parser, renderer
|
from beaver_gateway.frontends.markdown import parser, renderer
|
||||||
from beaver_gateway.frontends.markdown.crossfront import CrossFrontendLogger
|
from beaver_gateway.frontends.markdown.crossfront import CrossFrontendLogger
|
||||||
from beaver_gateway.frontends.markdown.files import (
|
from beaver_gateway.frontends.markdown.files import (
|
||||||
@@ -69,12 +29,25 @@ from beaver_gateway.frontends.markdown.files import (
|
|||||||
reattach_frontmatter,
|
reattach_frontmatter,
|
||||||
write_atomic,
|
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.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:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Callable
|
from collections.abc import AsyncIterator, Callable
|
||||||
|
|
||||||
from beaver_gateway.core.kinds import Kind
|
from beaver_gateway.conversations.kinds import Kind
|
||||||
from beaver_gateway.frontends.base import GatewayRuntime
|
from beaver_gateway.frontends.base import GatewayRuntime
|
||||||
from beaver_gateway.storage.models import Conversation, ConversationBinding
|
from beaver_gateway.storage.models import Conversation, ConversationBinding
|
||||||
|
|
||||||
@@ -85,23 +58,17 @@ _log = logging.getLogger("beaver_gateway.frontends.markdown")
|
|||||||
__all__ = ["MarkdownFrontend"]
|
__all__ = ["MarkdownFrontend"]
|
||||||
|
|
||||||
|
|
||||||
# How often we re-render the assistant turn into the .md file while the
|
|
||||||
# backend stream is still open. Trades responsiveness (faster updates to
|
|
||||||
# Obsidian sync / Raycast tailers) against write amplification. Each
|
|
||||||
# ``RawContentBlockStopEvent`` also forces a flush regardless of the
|
|
||||||
# timer, so block boundaries always land in the file.
|
|
||||||
_STREAM_FLUSH_DEBOUNCE = 0.4
|
_STREAM_FLUSH_DEBOUNCE = 0.4
|
||||||
|
|
||||||
# Debounce for the SSE ``/chat/stream`` path. Network IO is cheaper than
|
|
||||||
# atomic file rewrites, so we send updates more frequently — the client
|
|
||||||
# wants the lowest possible latency and we control the renderer on the
|
|
||||||
# other end (the Obsidian plugin splices deltas into the editor, no
|
|
||||||
# disk round-trip).
|
|
||||||
_SSE_FLUSH_DEBOUNCE = 0.1
|
_SSE_FLUSH_DEBOUNCE = 0.1
|
||||||
|
|
||||||
|
|
||||||
class MarkdownFrontend(Frontend):
|
class MarkdownFrontend(Frontend):
|
||||||
"""FastAPI app behind ``POST /chat`` driven by Obsidian-vault files."""
|
"""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
|
name = FRONTEND
|
||||||
kinds = ("deep",)
|
kinds = ("deep",)
|
||||||
@@ -116,6 +83,11 @@ class MarkdownFrontend(Frontend):
|
|||||||
logged_subdir: str = "_logs",
|
logged_subdir: str = "_logs",
|
||||||
chat_path: Callable[[str, str, Path], Path] | None = None,
|
chat_path: Callable[[str, str, Path], Path] | None = 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.vault_path = Path(vault_path).expanduser().resolve()
|
||||||
self.default_agent = default_agent
|
self.default_agent = default_agent
|
||||||
self.log_all_chats = log_all_chats
|
self.log_all_chats = log_all_chats
|
||||||
@@ -123,10 +95,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
self.chat_path = chat_path
|
self.chat_path = chat_path
|
||||||
self._runtime: GatewayRuntime | None = None
|
self._runtime: GatewayRuntime | None = None
|
||||||
self._app: FastAPI | None = None
|
self._app: FastAPI | None = None
|
||||||
# Files currently being processed by an in-flight ``POST /chat``.
|
|
||||||
# Checked-and-added atomically in the request handler (no
|
|
||||||
# ``await`` between the check and the insert) so a concurrent
|
|
||||||
# request reliably loses the race to 409.
|
|
||||||
self._busy: set[Path] = set()
|
self._busy: set[Path] = set()
|
||||||
self._crossfront: CrossFrontendLogger | None = None
|
self._crossfront: CrossFrontendLogger | None = None
|
||||||
self._mirror: ChatMirror | None = None
|
self._mirror: ChatMirror | None = None
|
||||||
@@ -149,9 +117,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
logged_subdir=self.logged_subdir,
|
logged_subdir=self.logged_subdir,
|
||||||
chat_path=self.chat_path,
|
chat_path=self.chat_path,
|
||||||
)
|
)
|
||||||
# Scan the existing logged files synchronously here so the
|
|
||||||
# fingerprint→path map is populated before the first
|
|
||||||
# cross-frontend turn arrives. Cheap: frontmatter-only read.
|
|
||||||
self._crossfront.warm_index()
|
self._crossfront.warm_index()
|
||||||
runtime.turn_log_handlers.append(self._crossfront.handle)
|
runtime.turn_log_handlers.append(self._crossfront.handle)
|
||||||
self._app = self._build_app(runtime)
|
self._app = self._build_app(runtime)
|
||||||
@@ -175,16 +140,10 @@ class MarkdownFrontend(Frontend):
|
|||||||
async def serve(self) -> None:
|
async def serve(self) -> None:
|
||||||
await self.mirror.run()
|
await self.mirror.run()
|
||||||
|
|
||||||
# ---- app builder ---------------------------------------------------
|
|
||||||
|
|
||||||
def _build_app(self, runtime: GatewayRuntime) -> FastAPI:
|
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 = FastAPI(title="beaver-gateway / Markdown")
|
||||||
|
|
||||||
# ``/chat/stream`` is consumed via ``fetch`` from the Obsidian
|
|
||||||
# plugin (``requestUrl`` can't read a body incrementally), and
|
|
||||||
# ``fetch`` is subject to CORS. Auth is bearer-token so we don't
|
|
||||||
# need credentialed mode; allow any origin and the standard
|
|
||||||
# methods/headers. The other endpoints are happy to ride along.
|
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
allow_origins=["*"],
|
allow_origins=["*"],
|
||||||
@@ -226,8 +185,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
|
|
||||||
file_path = self._resolve_path(filename)
|
file_path = self._resolve_path(filename)
|
||||||
|
|
||||||
# Atomic check-and-claim: both ops run between awaits, so a
|
|
||||||
# second request can't slip into the same file slot.
|
|
||||||
if file_path in self._busy:
|
if file_path in self._busy:
|
||||||
return JSONResponse(
|
return JSONResponse(
|
||||||
status_code=status.HTTP_409_CONFLICT,
|
status_code=status.HTTP_409_CONFLICT,
|
||||||
@@ -248,12 +205,10 @@ class MarkdownFrontend(Frontend):
|
|||||||
|
|
||||||
@app.post("/chat/stream")
|
@app.post("/chat/stream")
|
||||||
async def chat_stream(request: Request) -> Any:
|
async def chat_stream(request: Request) -> Any:
|
||||||
# Same contract as ``/chat`` (bearer auth, identical body),
|
"""SSE variant of ``/chat``: ``delta`` events, one disk write at the end.
|
||||||
# but the response is ``text/event-stream`` and intermediate
|
|
||||||
# rendered states are pushed as ``delta`` events. The
|
The 409-conflict response is still plain JSON — the stream hasn't started.
|
||||||
# gateway-side disk write only happens once, at end of turn,
|
"""
|
||||||
# so streaming consumers (Obsidian plugin) and Obsidian Sync
|
|
||||||
# don't fight over the same file mid-stream.
|
|
||||||
token_name = await require_token(request, runtime, scope="messages")
|
token_name = await require_token(request, runtime, scope="messages")
|
||||||
try:
|
try:
|
||||||
body = await request.json()
|
body = await request.json()
|
||||||
@@ -276,8 +231,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
|
|
||||||
file_path = self._resolve_path(filename)
|
file_path = self._resolve_path(filename)
|
||||||
|
|
||||||
# 409 path stays JSON — the stream hasn't started yet, so
|
|
||||||
# the caller can read it the same way as on ``/chat``.
|
|
||||||
if file_path in self._busy:
|
if file_path in self._busy:
|
||||||
return JSONResponse(
|
return JSONResponse(
|
||||||
status_code=status.HTTP_409_CONFLICT,
|
status_code=status.HTTP_409_CONFLICT,
|
||||||
@@ -305,8 +258,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
|
|
||||||
return app
|
return app
|
||||||
|
|
||||||
# ---- dispatch ------------------------------------------------------
|
|
||||||
|
|
||||||
async def _handle_chat(
|
async def _handle_chat(
|
||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
@@ -317,6 +268,7 @@ class MarkdownFrontend(Frontend):
|
|||||||
content_override: Any,
|
content_override: Any,
|
||||||
agent_override: str | None,
|
agent_override: str | None,
|
||||||
) -> Any:
|
) -> Any:
|
||||||
|
"""Non-streaming ``/chat`` handler: parse, align, run, persist."""
|
||||||
write_disk = content_override is None
|
write_disk = content_override is None
|
||||||
if isinstance(content_override, str):
|
if isinstance(content_override, str):
|
||||||
file_text = content_override
|
file_text = content_override
|
||||||
@@ -340,8 +292,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
"or configure `default_agent`",
|
"or configure `default_agent`",
|
||||||
)
|
)
|
||||||
|
|
||||||
# When the parser produced no messages (file empty / only
|
|
||||||
# frontmatter), there's nothing to dispatch.
|
|
||||||
if not parsed.messages:
|
if not parsed.messages:
|
||||||
return {
|
return {
|
||||||
"status": "nothing_to_do",
|
"status": "nothing_to_do",
|
||||||
@@ -384,9 +334,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
msgs=len(parsed.messages),
|
msgs=len(parsed.messages),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Resolve / mint the conversation row, align incoming against
|
|
||||||
# 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(
|
conv, conv_external_id, stored_msgs = await self._resolve_conversation(
|
||||||
runtime=runtime,
|
runtime=runtime,
|
||||||
metadata=parsed.metadata,
|
metadata=parsed.metadata,
|
||||||
@@ -394,7 +341,11 @@ class MarkdownFrontend(Frontend):
|
|||||||
file_path=file_path,
|
file_path=file_path,
|
||||||
first_user_text=parsed.turns[0].text,
|
first_user_text=parsed.turns[0].text,
|
||||||
)
|
)
|
||||||
outcome = diff_and_fork(stored=stored_msgs, incoming=parsed.turns)
|
outcome = diff_and_fork(
|
||||||
|
stored=stored_msgs,
|
||||||
|
incoming=parsed.turns,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
|
)
|
||||||
capture = TurnCapture()
|
capture = TurnCapture()
|
||||||
events = runtime.conversations.turn(
|
events = runtime.conversations.turn(
|
||||||
conv,
|
conv,
|
||||||
@@ -409,6 +360,7 @@ class MarkdownFrontend(Frontend):
|
|||||||
file_path=file_path,
|
file_path=file_path,
|
||||||
parsed=parsed,
|
parsed=parsed,
|
||||||
model=agent.model or agent.name,
|
model=agent.model or agent.name,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
filename=filename,
|
filename=filename,
|
||||||
write_disk=write_disk,
|
write_disk=write_disk,
|
||||||
)
|
)
|
||||||
@@ -424,6 +376,7 @@ class MarkdownFrontend(Frontend):
|
|||||||
parsed=parsed,
|
parsed=parsed,
|
||||||
message=message,
|
message=message,
|
||||||
agent_name=agent.name,
|
agent_name=agent.name,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
conv_external_id=conv_external_id,
|
conv_external_id=conv_external_id,
|
||||||
write_disk=write_disk,
|
write_disk=write_disk,
|
||||||
)
|
)
|
||||||
@@ -437,13 +390,11 @@ class MarkdownFrontend(Frontend):
|
|||||||
message=message,
|
message=message,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Broadcast our own turn so other handlers (none today, but the
|
|
||||||
# symmetry is worth keeping) see what happened. ``source`` marks
|
|
||||||
# the origin so ``CrossFrontendLogger`` can skip its own files.
|
|
||||||
record = TurnRecord(
|
record = TurnRecord(
|
||||||
agent_name=agent.name,
|
agent_name=agent.name,
|
||||||
input_messages=list(parsed.messages),
|
input_messages=list(parsed.messages),
|
||||||
output_message=message,
|
output_message=message,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
system=None,
|
system=None,
|
||||||
source="markdown",
|
source="markdown",
|
||||||
)
|
)
|
||||||
@@ -460,9 +411,7 @@ class MarkdownFrontend(Frontend):
|
|||||||
"new_content": new_content,
|
"new_content": new_content,
|
||||||
}
|
}
|
||||||
|
|
||||||
# ---- streaming dispatch (SSE) --------------------------------------
|
async def _handle_chat_streaming( # noqa: PLR0915
|
||||||
|
|
||||||
async def _handle_chat_streaming( # noqa: PLR0915 — mirrors _handle_chat, splitting only doubles read cost
|
|
||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
runtime: GatewayRuntime,
|
runtime: GatewayRuntime,
|
||||||
@@ -474,21 +423,9 @@ class MarkdownFrontend(Frontend):
|
|||||||
) -> AsyncIterator[bytes]:
|
) -> AsyncIterator[bytes]:
|
||||||
"""SSE counterpart of :meth:`_handle_chat`.
|
"""SSE counterpart of :meth:`_handle_chat`.
|
||||||
|
|
||||||
Mirrors the same pipeline (resolve file → parse → resolve agent →
|
Errors surface as ``error`` frames (HTTP is already 200 by then); disk
|
||||||
run backend → persist), but emits ``event: delta`` frames as the
|
is only written once at the end, avoiding a write race with Obsidian Sync.
|
||||||
rendered turn grows and a single terminal ``event: done`` /
|
|
||||||
``event: error``. Errors that ``_handle_chat`` would surface as
|
|
||||||
``HTTPException`` go out as ``error`` frames here (the HTTP
|
|
||||||
envelope is already 200 by the time the stream starts).
|
|
||||||
|
|
||||||
Intermediate disk writes are deliberately skipped — only the
|
|
||||||
post-stream :meth:`_write_assistant_reply` lands on disk, so the
|
|
||||||
gateway-side vault and the plugin-side editor are the only
|
|
||||||
writers in their respective halves of Obsidian Sync. Final
|
|
||||||
content is identical on both sides, so Sync no-ops.
|
|
||||||
"""
|
"""
|
||||||
# With ``content`` the plugin is the only writer of the file
|
|
||||||
# (§3.10): the gateway never touches disk in that case.
|
|
||||||
write_disk = content_override is None
|
write_disk = content_override is None
|
||||||
if isinstance(content_override, str):
|
if isinstance(content_override, str):
|
||||||
file_text = content_override
|
file_text = content_override
|
||||||
@@ -604,7 +541,11 @@ class MarkdownFrontend(Frontend):
|
|||||||
len(stored_msgs),
|
len(stored_msgs),
|
||||||
len(parsed.turns),
|
len(parsed.turns),
|
||||||
)
|
)
|
||||||
outcome = diff_and_fork(stored=stored_msgs, incoming=parsed.turns)
|
outcome = diff_and_fork(
|
||||||
|
stored=stored_msgs,
|
||||||
|
incoming=parsed.turns,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
|
)
|
||||||
_log.info(
|
_log.info(
|
||||||
"chat/stream: file=%s diff_and_fork divergence=%s "
|
"chat/stream: file=%s diff_and_fork divergence=%s "
|
||||||
"backend_msgs=%d persist_msgs=%d",
|
"backend_msgs=%d persist_msgs=%d",
|
||||||
@@ -637,7 +578,9 @@ class MarkdownFrontend(Frontend):
|
|||||||
partial = acc.finalize(model=model)
|
partial = acc.finalize(model=model)
|
||||||
if not partial.content:
|
if not partial.content:
|
||||||
return None
|
return None
|
||||||
rendered = renderer.render_assistant_message(partial)
|
rendered = renderer.render_assistant_message(
|
||||||
|
partial, thinking_is_progress=agent.thinking_is_progress
|
||||||
|
)
|
||||||
new_body = renderer.append_to_body(parsed.body, rendered)
|
new_body = renderer.append_to_body(parsed.body, rendered)
|
||||||
return reattach_frontmatter(parsed.metadata, new_body)
|
return reattach_frontmatter(parsed.metadata, new_body)
|
||||||
|
|
||||||
@@ -653,23 +596,20 @@ class MarkdownFrontend(Frontend):
|
|||||||
or (now - last_flush) >= _SSE_FLUSH_DEBOUNCE
|
or (now - last_flush) >= _SSE_FLUSH_DEBOUNCE
|
||||||
):
|
):
|
||||||
payload = snapshot()
|
payload = snapshot()
|
||||||
# Skip duplicate snapshots — e.g. tool_use blocks
|
|
||||||
# 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:
|
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_payload = payload
|
||||||
last_flush = now
|
last_flush = now
|
||||||
except Exception as exc: # noqa: BLE001 — wire any backend failure as an SSE error frame
|
except Exception as exc: # noqa: BLE001
|
||||||
_log.exception("backend failed for %s", filename)
|
_log.exception("backend failed for %s", filename)
|
||||||
# Mirror the legacy path: write the last partial + an error
|
|
||||||
# callout to disk so other consumers (logs, file watchers)
|
|
||||||
# see what arrived. The client gets a clean SSE ``error``.
|
|
||||||
partial = acc.finalize(model=model)
|
partial = acc.finalize(model=model)
|
||||||
new_body = parsed.body
|
new_body = parsed.body
|
||||||
if partial.content:
|
if partial.content:
|
||||||
new_body = renderer.append_to_body(
|
new_body = renderer.append_to_body(
|
||||||
new_body, renderer.render_assistant_message(partial)
|
new_body,
|
||||||
|
renderer.render_assistant_message(
|
||||||
|
partial, thinking_is_progress=agent.thinking_is_progress
|
||||||
|
),
|
||||||
)
|
)
|
||||||
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
|
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
|
||||||
if write_disk:
|
if write_disk:
|
||||||
@@ -692,6 +632,7 @@ class MarkdownFrontend(Frontend):
|
|||||||
parsed=parsed,
|
parsed=parsed,
|
||||||
message=message,
|
message=message,
|
||||||
agent_name=agent.name,
|
agent_name=agent.name,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
conv_external_id=conv_external_id,
|
conv_external_id=conv_external_id,
|
||||||
write_disk=write_disk,
|
write_disk=write_disk,
|
||||||
)
|
)
|
||||||
@@ -709,6 +650,7 @@ class MarkdownFrontend(Frontend):
|
|||||||
agent_name=agent.name,
|
agent_name=agent.name,
|
||||||
input_messages=list(parsed.messages),
|
input_messages=list(parsed.messages),
|
||||||
output_message=message,
|
output_message=message,
|
||||||
|
thinking_is_progress=agent.thinking_is_progress,
|
||||||
system=None,
|
system=None,
|
||||||
source="markdown",
|
source="markdown",
|
||||||
)
|
)
|
||||||
@@ -728,8 +670,6 @@ class MarkdownFrontend(Frontend):
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
# ---- helpers -------------------------------------------------------
|
|
||||||
|
|
||||||
async def _stream_to_file(
|
async def _stream_to_file(
|
||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
@@ -737,22 +677,14 @@ class MarkdownFrontend(Frontend):
|
|||||||
file_path: Path,
|
file_path: Path,
|
||||||
parsed: parser.ParsedFile,
|
parsed: parser.ParsedFile,
|
||||||
model: str,
|
model: str,
|
||||||
|
thinking_is_progress: bool,
|
||||||
filename: str,
|
filename: str,
|
||||||
write_disk: bool = True,
|
write_disk: bool = True,
|
||||||
) -> Any:
|
) -> Any:
|
||||||
"""Drain ``events`` into a ``Message``, flushing partials to disk.
|
"""Drain ``events`` into a ``Message``, flushing partials to disk.
|
||||||
|
|
||||||
Flushes happen on each ``RawContentBlockStopEvent`` (natural
|
Flushes on each block boundary and on a debounce timer; on backend
|
||||||
block boundary, content is markdown-consistent) and on the
|
failure it still flushes a partial + error callout, then re-raises.
|
||||||
``_STREAM_FLUSH_DEBOUNCE`` timer between events. The partial
|
|
||||||
write keeps the as-parsed frontmatter; the post-stream final
|
|
||||||
write in ``_write_assistant_reply`` is what stamps the refreshed
|
|
||||||
fingerprint / agent / conversation_id.
|
|
||||||
|
|
||||||
On backend exception we still flush the last partial and append
|
|
||||||
an error callout, so the human sees both what arrived and why it
|
|
||||||
stopped. The exception propagates so ``_handle_chat`` can map it
|
|
||||||
to a 500.
|
|
||||||
"""
|
"""
|
||||||
acc = StreamAccumulator()
|
acc = StreamAccumulator()
|
||||||
|
|
||||||
@@ -762,7 +694,9 @@ class MarkdownFrontend(Frontend):
|
|||||||
partial = acc.finalize(model=model)
|
partial = acc.finalize(model=model)
|
||||||
if not partial.content:
|
if not partial.content:
|
||||||
return
|
return
|
||||||
rendered = renderer.render_assistant_message(partial)
|
rendered = renderer.render_assistant_message(
|
||||||
|
partial, thinking_is_progress=thinking_is_progress
|
||||||
|
)
|
||||||
new_body = renderer.append_to_body(parsed.body, rendered)
|
new_body = renderer.append_to_body(parsed.body, rendered)
|
||||||
await write_atomic(
|
await write_atomic(
|
||||||
file_path, reattach_frontmatter(parsed.metadata, new_body)
|
file_path, reattach_frontmatter(parsed.metadata, new_body)
|
||||||
@@ -785,7 +719,10 @@ class MarkdownFrontend(Frontend):
|
|||||||
new_body = parsed.body
|
new_body = parsed.body
|
||||||
if partial.content:
|
if partial.content:
|
||||||
new_body = renderer.append_to_body(
|
new_body = renderer.append_to_body(
|
||||||
new_body, renderer.render_assistant_message(partial)
|
new_body,
|
||||||
|
renderer.render_assistant_message(
|
||||||
|
partial, thinking_is_progress=thinking_is_progress
|
||||||
|
),
|
||||||
)
|
)
|
||||||
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
|
new_body = renderer.append_to_body(new_body, _render_error_block(exc))
|
||||||
if write_disk:
|
if write_disk:
|
||||||
@@ -802,11 +739,14 @@ class MarkdownFrontend(Frontend):
|
|||||||
parsed: parser.ParsedFile,
|
parsed: parser.ParsedFile,
|
||||||
message: Any,
|
message: Any,
|
||||||
agent_name: str,
|
agent_name: str,
|
||||||
|
thinking_is_progress: bool,
|
||||||
conv_external_id: str,
|
conv_external_id: str,
|
||||||
write_disk: bool = True,
|
write_disk: bool = True,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Render the assistant turn, refresh frontmatter, write if we own the file."""
|
"""Render the assistant turn, refresh frontmatter, write if we own the file."""
|
||||||
rendered = renderer.render_assistant_message(message)
|
rendered = renderer.render_assistant_message(
|
||||||
|
message, thinking_is_progress=thinking_is_progress
|
||||||
|
)
|
||||||
new_body = renderer.append_to_body(parsed.body, rendered)
|
new_body = renderer.append_to_body(parsed.body, rendered)
|
||||||
new_body = renderer.append_to_body(new_body, renderer.USER_SCAFFOLD)
|
new_body = renderer.append_to_body(new_body, renderer.USER_SCAFFOLD)
|
||||||
updated_metadata = dict(parsed.metadata)
|
updated_metadata = dict(parsed.metadata)
|
||||||
@@ -829,11 +769,8 @@ class MarkdownFrontend(Frontend):
|
|||||||
) -> tuple[Conversation, str, list[dict[str, Any]]]:
|
) -> tuple[Conversation, str, list[dict[str, Any]]]:
|
||||||
"""Resolve the ``deep`` conversation for this file + its stored messages.
|
"""Resolve the ``deep`` conversation for this file + its stored messages.
|
||||||
|
|
||||||
Frontmatter ``conversation_id`` wins; a file that lost it is found
|
Precedence: frontmatter ``conversation_id`` > existing path binding >
|
||||||
by its visible ``(markdown, path)`` binding, then by adopting the
|
an adopted unbound conversation with the same first prompt > a new one.
|
||||||
one unbound pre-SDK conversation that starts with the same prompt;
|
|
||||||
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
|
conversations = runtime.conversations
|
||||||
rel = file_path.relative_to(self.vault_path).as_posix()
|
rel = file_path.relative_to(self.vault_path).as_posix()
|
||||||
@@ -882,9 +819,8 @@ class MarkdownFrontend(Frontend):
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""Stamp the DB with the post-turn canonical Anthropic-shape history.
|
"""Stamp the DB with the post-turn canonical Anthropic-shape history.
|
||||||
|
|
||||||
Combines the matched/spliced prior state, the new user prompt,
|
Combines prior state + new user prompt + the backend's synthesized
|
||||||
and the synthesized assistant/tool cycle from the backend (or a
|
cycle, falling back to text-only if ``capture`` is empty.
|
||||||
text-only fallback for backends that left ``capture`` empty).
|
|
||||||
"""
|
"""
|
||||||
new_user_msg = {"role": "user", "content": new_user_text}
|
new_user_msg = {"role": "user", "content": new_user_text}
|
||||||
synthesized = capture.synthesized_messages or _fallback_synthesized(message)
|
synthesized = capture.synthesized_messages or _fallback_synthesized(message)
|
||||||
@@ -906,12 +842,11 @@ class MarkdownFrontend(Frontend):
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _resolve_path(self, filename: str) -> Path:
|
def _resolve_path(self, filename: str) -> Path:
|
||||||
"""Resolve ``filename`` under the vault; reject escapes."""
|
"""Resolve ``filename`` under the vault; reject escapes.
|
||||||
# ``filename`` may be relative or absolute; we always anchor
|
|
||||||
# under ``vault_path`` so absolute paths from outside the vault
|
Leading slashes are stripped first — ``Path.__truediv__`` would
|
||||||
# don't sneak through. ``Path("/foo/bar")`` combined with a
|
otherwise discard ``vault_path`` for an absolute ``filename``.
|
||||||
# vault path keeps the absolute side; we strip leading slashes
|
"""
|
||||||
# to coerce the rooted form into a relative path before joining.
|
|
||||||
rel = filename.lstrip("/")
|
rel = filename.lstrip("/")
|
||||||
if not rel.endswith(".md"):
|
if not rel.endswith(".md"):
|
||||||
rel = rel + ".md"
|
rel = rel + ".md"
|
||||||
@@ -928,12 +863,8 @@ class MarkdownFrontend(Frontend):
|
|||||||
def _fallback_synthesized(message: Any) -> list[dict[str, Any]]:
|
def _fallback_synthesized(message: Any) -> list[dict[str, Any]]:
|
||||||
"""Build a single-assistant ``synthesized_messages`` list from a raw ``Message``.
|
"""Build a single-assistant ``synthesized_messages`` list from a raw ``Message``.
|
||||||
|
|
||||||
For backends that don't populate a :class:`TurnCapture` (anthropic
|
Fallback for backends that don't populate :class:`TurnCapture` (anthropic
|
||||||
HTTP, raycast, …) we don't have access to per-tool-cycle
|
HTTP, raycast, …); the reply lands as one canonical-block message.
|
||||||
granularity, so the assistant reply lands in the DB as one
|
|
||||||
canonical-block message. Tool memory across cache misses would
|
|
||||||
degrade in that case, but those backends don't have the cache-miss
|
|
||||||
re-seed problem to begin with — they manage history client-side.
|
|
||||||
"""
|
"""
|
||||||
content: list[dict[str, Any]] = []
|
content: list[dict[str, Any]] = []
|
||||||
for block in getattr(message, "content", ()):
|
for block in getattr(message, "content", ()):
|
||||||
|
|||||||
+84
-207
@@ -1,41 +1,7 @@
|
|||||||
"""Stateful conversation history for the markdown frontend.
|
"""Stored conversation history for the markdown frontend.
|
||||||
|
|
||||||
The gateway used to be stateless about identity: claude-code-api's
|
Aligns the file the user is editing against the DB-stored Anthropic-shape
|
||||||
in-memory session pool was keyed by a fingerprint of the messages the
|
message history, splicing text edits or forking on structural changes.
|
||||||
gateway forwarded, and on a fingerprint miss the same fingerprint was
|
|
||||||
used to seed a fresh PTY's JSONL transcript. That worked as long as
|
|
||||||
the frontend could round-trip the *exact* content blocks the live
|
|
||||||
session had observed. The markdown frontend can't — the parser strips
|
|
||||||
``[!tool]-`` callouts because the human is allowed to edit the prose,
|
|
||||||
and the rendered tool callouts don't carry the canonical ``tool_use``
|
|
||||||
block fields anyway. So a continuation hit was *only* reliable for
|
|
||||||
turns that never used a tool; once tools entered the picture, every
|
|
||||||
subsequent turn missed the cache and reseeded from a tool-less
|
|
||||||
transcript, leading to "assistant doesn't remember the tool calls it
|
|
||||||
just made."
|
|
||||||
|
|
||||||
This module makes the gateway stateful for the markdown frontend (and
|
|
||||||
any other frontend that wants in). The DB stores the full
|
|
||||||
Anthropic-shape message list — text blocks, ``tool_use`` blocks,
|
|
||||||
``tool_result`` blocks, thinking signatures — exactly as
|
|
||||||
claude-code-api would have seen on the wire. Before each turn we
|
|
||||||
align the file the user is editing against the stored history:
|
|
||||||
|
|
||||||
* If the user just appended a new user turn at the bottom, we feed
|
|
||||||
the backend our stored-plus-new history and the fingerprint hits.
|
|
||||||
* If the user edited the *text* inside an assistant turn but left the
|
|
||||||
tool callouts alone, we splice the new text into the stored
|
|
||||||
``tool_use`` blocks and feed *that* — the fingerprint misses (text
|
|
||||||
differs), claude-code-api reseeds with a full transcript (tools and
|
|
||||||
all), the new live session has memory of the prior tool calls.
|
|
||||||
* If the user changed the *structure* (added/removed/reordered a tool
|
|
||||||
callout, edited an old user turn, etc.) we fork: take stored history
|
|
||||||
up to the divergence, take incoming text-only past the divergence.
|
|
||||||
The fingerprint misses; claude-code-api reseeds with a clean
|
|
||||||
truncated history; downstream turns continue from there.
|
|
||||||
|
|
||||||
"Divergence point" is found by walking the file's turns and the
|
|
||||||
stored display turns in lockstep. See :func:`diff_and_fork`.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -48,6 +14,7 @@ from typing import TYPE_CHECKING, Any, cast
|
|||||||
from sqlalchemy import delete
|
from sqlalchemy import delete
|
||||||
from sqlmodel import select
|
from sqlmodel import select
|
||||||
|
|
||||||
|
from beaver_gateway.backends.transcript import spoken_text
|
||||||
from beaver_gateway.storage.models import Conversation, ConversationMessage
|
from beaver_gateway.storage.models import Conversation, ConversationMessage
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -67,21 +34,12 @@ __all__ = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
# ---- types --------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
class ForkOutcome:
|
class ForkOutcome:
|
||||||
"""Result of aligning the incoming file against stored history.
|
"""Result of aligning the incoming file against stored history.
|
||||||
|
|
||||||
``messages`` is what the gateway feeds to the backend (already
|
``messages`` is the backend input; ``persist_messages`` is history to
|
||||||
includes the new user prompt at the tail). ``persist_messages``
|
store before the new reply; ``divergence_index`` is where they diverged.
|
||||||
is the canonical conversation state the gateway should hold in
|
|
||||||
the DB *up to but not including* the new assistant reply — the
|
|
||||||
caller appends the synthesized turn from the backend onto this
|
|
||||||
and writes the result back. ``divergence_index`` is the
|
|
||||||
display-turn index at which incoming first disagreed with stored
|
|
||||||
(``None`` if everything matched; the new tail is appended cleanly).
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
messages: list[MessageParam]
|
messages: list[MessageParam]
|
||||||
@@ -95,9 +53,6 @@ class ForkOutcome:
|
|||||||
return self.divergence_index is None and not self.edited
|
return self.divergence_index is None and not self.edited
|
||||||
|
|
||||||
|
|
||||||
# ---- public store API ---------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
async def load_conversation(
|
async def load_conversation(
|
||||||
session: AsyncSession, *, frontend: str, external_id: str
|
session: AsyncSession, *, frontend: str, external_id: str
|
||||||
) -> Conversation | None:
|
) -> Conversation | None:
|
||||||
@@ -115,9 +70,7 @@ async def mint_conversation(
|
|||||||
) -> Conversation:
|
) -> Conversation:
|
||||||
"""Create a fresh conversation row with a new uuid for external_id.
|
"""Create a fresh conversation row with a new uuid for external_id.
|
||||||
|
|
||||||
Caller is responsible for persisting the returned ``external_id`` on
|
Caller must persist the returned ``external_id`` so future requests can find it.
|
||||||
the frontend side (frontmatter, response header, …) so future
|
|
||||||
requests can find this conversation again.
|
|
||||||
"""
|
"""
|
||||||
row = Conversation(
|
row = Conversation(
|
||||||
frontend=frontend, external_id=str(uuid.uuid4()), agent_name=agent_name
|
frontend=frontend, external_id=str(uuid.uuid4()), agent_name=agent_name
|
||||||
@@ -144,9 +97,8 @@ async def load_messages(
|
|||||||
) -> list[dict[str, Any]]:
|
) -> list[dict[str, Any]]:
|
||||||
"""Return stored messages ordered by ``seq`` ascending.
|
"""Return stored messages ordered by ``seq`` ascending.
|
||||||
|
|
||||||
Each entry is a canonical Anthropic ``MessageParam`` dict — ``role``
|
Each entry is a canonical Anthropic ``MessageParam`` dict, the same
|
||||||
plus ``content`` (string or list of block dicts). The same shape
|
shape fed to the backend on continuation.
|
||||||
we feed to the backend on continuation.
|
|
||||||
"""
|
"""
|
||||||
stmt = (
|
stmt = (
|
||||||
select(ConversationMessage)
|
select(ConversationMessage)
|
||||||
@@ -164,10 +116,8 @@ async def load_messages(
|
|||||||
def _sanitize_content(content: Any) -> Any:
|
def _sanitize_content(content: Any) -> Any:
|
||||||
"""Strip wire-illegal fields from stored Anthropic content blocks.
|
"""Strip wire-illegal fields from stored Anthropic content blocks.
|
||||||
|
|
||||||
Older capture code emitted ``"is_error": null`` on ``tool_result``
|
Drops ``tool_result.is_error: null`` — the API rejects null there
|
||||||
blocks; the Anthropic API rejects null there (the field is optional
|
though the field is optional.
|
||||||
but, when present, must be boolean). We omit the key on read so
|
|
||||||
historical rows don't break continuation.
|
|
||||||
"""
|
"""
|
||||||
if not isinstance(content, list):
|
if not isinstance(content, list):
|
||||||
return content
|
return content
|
||||||
@@ -188,27 +138,17 @@ def _sanitize_content(content: Any) -> Any:
|
|||||||
async def rewrite_messages(
|
async def rewrite_messages(
|
||||||
session: AsyncSession, *, conversation_id: int, messages: list[dict[str, Any]]
|
session: AsyncSession, *, conversation_id: int, messages: list[dict[str, Any]]
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Replace the conversation's stored messages with ``messages``.
|
"""Replace the conversation's stored messages (full overwrite, no branch history).
|
||||||
|
|
||||||
The user said no branch history — we overwrite on fork. Cheap at
|
Deletes and flushes before inserting — SQLAlchemy's default INSERT-before-DELETE
|
||||||
our volume; if it ever matters we can switch to soft-delete +
|
flush order would otherwise collide with ``uq_msg_conv_seq``.
|
||||||
branch pointers.
|
|
||||||
"""
|
"""
|
||||||
# Bulk-delete and flush before inserting the new sequence: SQLAlchemy's
|
|
||||||
# unit-of-work flushes INSERTs before DELETEs by default, which would
|
|
||||||
# trip ``uq_msg_conv_seq`` when the new rows reuse the same seq numbers
|
|
||||||
# as the soon-to-be-deleted ones.
|
|
||||||
# SQLModel descriptors resolve to ColumnElement at runtime but to bare
|
|
||||||
# ``int`` in ty's stubs; the select-path at line 135 lives behind sqlmodel's
|
|
||||||
# own ``select`` overloads that hide it, but ``sqlalchemy.delete().where``
|
|
||||||
# uses the raw stubs.
|
|
||||||
await session.execute( # ty: ignore[deprecated]
|
await session.execute( # ty: ignore[deprecated]
|
||||||
delete(ConversationMessage).where(
|
delete(ConversationMessage).where(
|
||||||
ConversationMessage.conversation_id == conversation_id # ty: ignore[invalid-argument-type]
|
ConversationMessage.conversation_id == conversation_id # ty: ignore[invalid-argument-type]
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
await session.flush()
|
await session.flush()
|
||||||
# Insert the new sequence.
|
|
||||||
for seq, m in enumerate(messages):
|
for seq, m in enumerate(messages):
|
||||||
session.add(
|
session.add(
|
||||||
ConversationMessage(
|
ConversationMessage(
|
||||||
@@ -218,7 +158,6 @@ async def rewrite_messages(
|
|||||||
content_json=json.dumps(m["content"], separators=(",", ":")),
|
content_json=json.dumps(m["content"], separators=(",", ":")),
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
# Bump conversation.updated_at.
|
|
||||||
conv = await session.get(Conversation, conversation_id)
|
conv = await session.get(Conversation, conversation_id)
|
||||||
if conv is not None:
|
if conv is not None:
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
@@ -228,21 +167,11 @@ async def rewrite_messages(
|
|||||||
await session.commit()
|
await session.commit()
|
||||||
|
|
||||||
|
|
||||||
# ---- alignment ----------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
class _StoredDisplayTurn:
|
class _StoredDisplayTurn:
|
||||||
"""A "display turn" reconstructed from stored raw messages.
|
"""A "display turn" reconstructed from stored raw messages.
|
||||||
|
|
||||||
``role`` is ``"user"`` (single user-prompt message) or
|
Parser-equivalent view used to diff the file against the DB.
|
||||||
``"assistant"`` (one or more assistant messages, optionally
|
|
||||||
interleaved with user-only-tool_result messages). ``messages`` is
|
|
||||||
the slice of stored raw messages this display turn covers, in
|
|
||||||
order. ``spoken_text`` and ``skeleton`` are the
|
|
||||||
parser-equivalents for diff purposes; ``text_segment_count`` lets
|
|
||||||
us refuse a splice when the user edited across a tool boundary in
|
|
||||||
a way we can't safely undo.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
role: str
|
role: str
|
||||||
@@ -252,13 +181,13 @@ class _StoredDisplayTurn:
|
|||||||
text_segment_count: int
|
text_segment_count: int
|
||||||
|
|
||||||
|
|
||||||
def _group_display_turns(stored: list[dict[str, Any]]) -> list[_StoredDisplayTurn]:
|
def _group_display_turns(
|
||||||
|
stored: list[dict[str, Any]], *, thinking_is_progress: bool
|
||||||
|
) -> list[_StoredDisplayTurn]:
|
||||||
"""Walk raw stored messages, group them into Obsidian-visible turns.
|
"""Walk raw stored messages, group them into Obsidian-visible turns.
|
||||||
|
|
||||||
A user-prompt message (``role=user`` with string content, or list
|
A user-prompt message (no ``tool_result`` blocks) opens a new turn;
|
||||||
content with no ``tool_result`` blocks) opens a user display turn.
|
tool-result-only messages roll into the current assistant turn.
|
||||||
Otherwise it's a tool-result follow-up and rolls into the current
|
|
||||||
assistant display turn.
|
|
||||||
"""
|
"""
|
||||||
out: list[_StoredDisplayTurn] = []
|
out: list[_StoredDisplayTurn] = []
|
||||||
i = 0
|
i = 0
|
||||||
@@ -277,7 +206,6 @@ def _group_display_turns(stored: list[dict[str, Any]]) -> list[_StoredDisplayTur
|
|||||||
)
|
)
|
||||||
i += 1
|
i += 1
|
||||||
continue
|
continue
|
||||||
# Assistant display turn: collect consecutive non-prompt messages.
|
|
||||||
group: list[dict[str, Any]] = []
|
group: list[dict[str, Any]] = []
|
||||||
while i < len(stored):
|
while i < len(stored):
|
||||||
m = stored[i]
|
m = stored[i]
|
||||||
@@ -285,7 +213,9 @@ def _group_display_turns(stored: list[dict[str, Any]]) -> list[_StoredDisplayTur
|
|||||||
break
|
break
|
||||||
group.append(m)
|
group.append(m)
|
||||||
i += 1
|
i += 1
|
||||||
spoken, skeleton, text_count = _summarize_assistant_group(group)
|
spoken, skeleton, text_count = _summarize_assistant_group(
|
||||||
|
group, thinking_is_progress=thinking_is_progress
|
||||||
|
)
|
||||||
out.append(
|
out.append(
|
||||||
_StoredDisplayTurn(
|
_StoredDisplayTurn(
|
||||||
role="assistant",
|
role="assistant",
|
||||||
@@ -299,14 +229,16 @@ def _group_display_turns(stored: list[dict[str, Any]]) -> list[_StoredDisplayTur
|
|||||||
|
|
||||||
|
|
||||||
def _is_user_prompt(content: Any) -> bool:
|
def _is_user_prompt(content: Any) -> bool:
|
||||||
"""A user message is a *prompt* unless its content carries tool_result blocks."""
|
"""A user message is a *prompt* unless its content carries tool_result blocks.
|
||||||
|
|
||||||
|
Unknown content shapes are conservatively treated as a prompt.
|
||||||
|
"""
|
||||||
if isinstance(content, str):
|
if isinstance(content, str):
|
||||||
return True
|
return True
|
||||||
if isinstance(content, list):
|
if isinstance(content, list):
|
||||||
return not any(
|
return not any(
|
||||||
isinstance(b, dict) and b.get("type") == "tool_result" for b in content
|
isinstance(b, dict) and b.get("type") == "tool_result" for b in content
|
||||||
)
|
)
|
||||||
# Unknown shape — be conservative, treat as prompt.
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
@@ -324,17 +256,14 @@ def _user_prompt_text(content: Any) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def _summarize_assistant_group(
|
def _summarize_assistant_group(
|
||||||
group: list[dict[str, Any]],
|
group: list[dict[str, Any]], *, thinking_is_progress: bool
|
||||||
) -> tuple[str, list[str], int]:
|
) -> tuple[str, list[str], int]:
|
||||||
"""Compute (spoken_text, tool_skeleton, text_segment_count) for a display group.
|
"""Compute (spoken_text, tool_skeleton, text_segment_count) for a display group.
|
||||||
|
|
||||||
Mirrors what ``parser.parse_assistant_structure`` would produce when
|
Must mirror what the renderer writes and the parser reads back: text and
|
||||||
re-parsing the rendered version of this group: consecutive text
|
progress thinking collapse into one segment, tool_use becomes a skeleton
|
||||||
blocks across assistant messages collapse into one text segment;
|
entry, tool_result, reasoning and empty thinking are invisible.
|
||||||
tool_use blocks become skeleton entries; tool_result messages and
|
|
||||||
thinking blocks are invisible.
|
|
||||||
"""
|
"""
|
||||||
# See ``diff_and_fork`` for why the parser-type imports are deferred.
|
|
||||||
from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment
|
from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment
|
||||||
|
|
||||||
segments: list[TextSegment | ToolSegment] = []
|
segments: list[TextSegment | ToolSegment] = []
|
||||||
@@ -351,7 +280,6 @@ def _summarize_assistant_group(
|
|||||||
|
|
||||||
for msg in group:
|
for msg in group:
|
||||||
if msg["role"] == "user":
|
if msg["role"] == "user":
|
||||||
# tool_result message — boundary for text but emits no segment.
|
|
||||||
_flush()
|
_flush()
|
||||||
continue
|
continue
|
||||||
content = msg.get("content")
|
content = msg.get("content")
|
||||||
@@ -360,15 +288,12 @@ def _summarize_assistant_group(
|
|||||||
for blk in content:
|
for blk in content:
|
||||||
if not isinstance(blk, dict):
|
if not isinstance(blk, dict):
|
||||||
continue
|
continue
|
||||||
btype = blk.get("type")
|
text = spoken_text(blk, thinking_is_progress=thinking_is_progress)
|
||||||
if btype == "text":
|
if text:
|
||||||
text = str(blk.get("text", "")).strip()
|
pending.append(text)
|
||||||
if text:
|
elif blk.get("type") == "tool_use":
|
||||||
pending.append(text)
|
|
||||||
elif btype == "tool_use":
|
|
||||||
_flush()
|
_flush()
|
||||||
segments.append(ToolSegment(name=str(blk.get("name", ""))))
|
segments.append(ToolSegment(name=str(blk.get("name", ""))))
|
||||||
# thinking: skip silently
|
|
||||||
_flush()
|
_flush()
|
||||||
spoken_chunks = [s.text for s in segments if isinstance(s, TextSegment)]
|
spoken_chunks = [s.text for s in segments if isinstance(s, TextSegment)]
|
||||||
spoken = "\n\n".join(c for c in spoken_chunks if c).strip()
|
spoken = "\n\n".join(c for c in spoken_chunks if c).strip()
|
||||||
@@ -377,30 +302,17 @@ def _summarize_assistant_group(
|
|||||||
return spoken, skeleton, text_count
|
return spoken, skeleton, text_count
|
||||||
|
|
||||||
|
|
||||||
# ---- the core algorithm -------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def diff_and_fork(
|
def diff_and_fork(
|
||||||
*, stored: list[dict[str, Any]], incoming: list[ParsedTurn]
|
*,
|
||||||
|
stored: list[dict[str, Any]],
|
||||||
|
incoming: list[ParsedTurn],
|
||||||
|
thinking_is_progress: bool,
|
||||||
) -> ForkOutcome:
|
) -> ForkOutcome:
|
||||||
"""Align the incoming parsed file against stored history.
|
"""Align the incoming parsed file against stored history.
|
||||||
|
|
||||||
``stored`` is the raw Anthropic-shape message list from the DB
|
``incoming`` must end with a user turn (the new prompt); raises otherwise.
|
||||||
(one entry per ``ConversationMessage`` row). ``incoming`` is the
|
Segment-class imports below are deferred to avoid a cycle with ``parser``.
|
||||||
user-visible turn list from the markdown parser. The last
|
|
||||||
``incoming`` entry must be a user turn — that's the new prompt
|
|
||||||
triggering this request.
|
|
||||||
|
|
||||||
Returns a :class:`ForkOutcome` whose ``messages`` is what the
|
|
||||||
backend should run on and whose ``persist_messages`` is the
|
|
||||||
canonical history to store in the DB once the backend's
|
|
||||||
synthesized cycle is appended.
|
|
||||||
"""
|
"""
|
||||||
# ``parser`` lives under ``frontends/markdown/`` whose ``__init__``
|
|
||||||
# eagerly loads ``frontend.py``, which in turn imports this module
|
|
||||||
# — pulling the parser at module-import time creates a cycle. The
|
|
||||||
# helpers below import the segment classes lazily inside their own
|
|
||||||
# function bodies to break it.
|
|
||||||
if not incoming or incoming[-1].role != "user":
|
if not incoming or incoming[-1].role != "user":
|
||||||
msg = (
|
msg = (
|
||||||
"diff_and_fork expects incoming to end with a user turn "
|
"diff_and_fork expects incoming to end with a user turn "
|
||||||
@@ -409,24 +321,22 @@ def diff_and_fork(
|
|||||||
)
|
)
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
|
|
||||||
stored_groups = _group_display_turns(stored)
|
stored_groups = _group_display_turns(
|
||||||
|
stored, thinking_is_progress=thinking_is_progress
|
||||||
|
)
|
||||||
new_user_turn = incoming[-1]
|
new_user_turn = incoming[-1]
|
||||||
prior_incoming = incoming[:-1]
|
prior_incoming = incoming[:-1]
|
||||||
|
|
||||||
spliced_groups, divergence, edited = _walk_prefix(prior_incoming, stored_groups)
|
spliced_groups, divergence, edited = _walk_prefix(
|
||||||
|
prior_incoming, stored_groups, thinking_is_progress=thinking_is_progress
|
||||||
|
)
|
||||||
|
|
||||||
if divergence is None and len(prior_incoming) < len(stored_groups):
|
if divergence is None and len(prior_incoming) < len(stored_groups):
|
||||||
if _file_lags_store(stored_groups, len(prior_incoming), new_user_turn):
|
if _file_lags_store(stored_groups, len(prior_incoming), new_user_turn):
|
||||||
# Not a deletion — the file simply never received turns we
|
|
||||||
# already ran. Adopt the stored tail verbatim so history stays
|
|
||||||
# structured and its fingerprint still matches the live
|
|
||||||
# session's. See ``_file_lags_store`` for why this matters.
|
|
||||||
spliced_groups.extend(
|
spliced_groups.extend(
|
||||||
list(g.messages) for g in stored_groups[len(prior_incoming) :]
|
list(g.messages) for g in stored_groups[len(prior_incoming) :]
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# Incoming truncated stored (user deleted some past turns).
|
|
||||||
# Truncate stored to match.
|
|
||||||
divergence = len(prior_incoming)
|
divergence = len(prior_incoming)
|
||||||
|
|
||||||
backend_msgs, persist_msgs = _assemble_tail(
|
backend_msgs, persist_msgs = _assemble_tail(
|
||||||
@@ -448,26 +358,9 @@ def _file_lags_store(
|
|||||||
) -> bool:
|
) -> bool:
|
||||||
"""Is the shorter incoming file a stale view rather than a deletion?
|
"""Is the shorter incoming file a stale view rather than a deletion?
|
||||||
|
|
||||||
A file with fewer display turns than the DB has two possible causes,
|
Tell: if the DB already holds the submitted prompt at the position the
|
||||||
and they need opposite handling:
|
file stops at, it's a lagging render, not a deletion — misjudging this
|
||||||
|
forks history, breaks the fingerprint, and respawns the backend session.
|
||||||
* the user deleted trailing turns — we should truncate to match;
|
|
||||||
* the turn ran, was persisted, but its reply never made it back into
|
|
||||||
the ``.md`` (the render lost a race with the user's next prompt, or
|
|
||||||
the reply rendered to nothing visible). The file is simply behind.
|
|
||||||
|
|
||||||
The tell is the prompt the user is submitting right now: if the DB
|
|
||||||
already holds it at exactly the position the file stops at, this is a
|
|
||||||
re-submission of a turn we've already run, not a deletion. Nobody
|
|
||||||
deletes a turn and immediately retypes it verbatim.
|
|
||||||
|
|
||||||
Getting this wrong is expensive and self-sustaining. Forking here
|
|
||||||
flattens every post-divergence turn into plain text (losing tool_use /
|
|
||||||
tool_result structure), persists that flattened history, and changes
|
|
||||||
the conversation fingerprint — so the backend's session pool misses,
|
|
||||||
spawns a fresh ``claude``, reseeds it from a multi-MB JSONL, and
|
|
||||||
strands the previous process. The file still lags afterwards, so the
|
|
||||||
next turn does it again.
|
|
||||||
"""
|
"""
|
||||||
if prior_len >= len(stored_groups):
|
if prior_len >= len(stored_groups):
|
||||||
return False
|
return False
|
||||||
@@ -476,16 +369,15 @@ def _file_lags_store(
|
|||||||
|
|
||||||
|
|
||||||
def _walk_prefix(
|
def _walk_prefix(
|
||||||
prior_incoming: list[ParsedTurn], stored_groups: list[_StoredDisplayTurn]
|
prior_incoming: list[ParsedTurn],
|
||||||
|
stored_groups: list[_StoredDisplayTurn],
|
||||||
|
*,
|
||||||
|
thinking_is_progress: bool,
|
||||||
) -> tuple[list[list[dict[str, Any]]], int | None, bool]:
|
) -> tuple[list[list[dict[str, Any]]], int | None, bool]:
|
||||||
"""Walk incoming vs stored side-by-side until first divergence.
|
"""Walk incoming vs stored side-by-side until first divergence.
|
||||||
|
|
||||||
Returns the spliced/matched group list (one entry per matched
|
An empty incoming skeleton (no tool callouts rendered) means prose
|
||||||
display turn, each carrying the raw messages we'll feed to the
|
alone decides the match; a mismatch forces a fresh backend session.
|
||||||
backend for that turn), the divergence index (``None`` if all
|
|
||||||
of ``prior_incoming`` matched) and whether any assistant prose
|
|
||||||
was spliced in from the file - a rewritten reply keeps the
|
|
||||||
structure but must not resume the session that said otherwise.
|
|
||||||
"""
|
"""
|
||||||
from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment
|
from beaver_gateway.frontends.markdown.parser import TextSegment, ToolSegment
|
||||||
|
|
||||||
@@ -506,8 +398,6 @@ def _walk_prefix(
|
|||||||
s.name for s in inc.structure if isinstance(s, ToolSegment)
|
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))
|
inc_text_count = sum(1 for s in inc.structure if isinstance(s, TextSegment))
|
||||||
# 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:
|
if inc_skeleton and inc_skeleton != st.skeleton:
|
||||||
return spliced_groups, i, edited
|
return spliced_groups, i, edited
|
||||||
if inc.text == st.spoken_text:
|
if inc.text == st.spoken_text:
|
||||||
@@ -515,7 +405,9 @@ def _walk_prefix(
|
|||||||
continue
|
continue
|
||||||
if inc_skeleton and inc_text_count != st.text_segment_count:
|
if inc_skeleton and inc_text_count != st.text_segment_count:
|
||||||
return spliced_groups, i, edited
|
return spliced_groups, i, edited
|
||||||
spliced = _splice_assistant_group(stored_group=st, incoming=inc)
|
spliced = _splice_assistant_group(
|
||||||
|
stored_group=st, incoming=inc, thinking_is_progress=thinking_is_progress
|
||||||
|
)
|
||||||
if spliced is None:
|
if spliced is None:
|
||||||
return spliced_groups, i, edited
|
return spliced_groups, i, edited
|
||||||
spliced_groups.append(spliced)
|
spliced_groups.append(spliced)
|
||||||
@@ -550,7 +442,10 @@ def _assemble_tail(
|
|||||||
|
|
||||||
|
|
||||||
def _splice_assistant_group(
|
def _splice_assistant_group(
|
||||||
*, stored_group: _StoredDisplayTurn, incoming: ParsedTurn
|
*,
|
||||||
|
stored_group: _StoredDisplayTurn,
|
||||||
|
incoming: ParsedTurn,
|
||||||
|
thinking_is_progress: bool,
|
||||||
) -> list[dict[str, Any]] | None:
|
) -> list[dict[str, Any]] | None:
|
||||||
"""Apply the file's prose to a stored assistant turn.
|
"""Apply the file's prose to a stored assistant turn.
|
||||||
|
|
||||||
@@ -558,39 +453,29 @@ def _splice_assistant_group(
|
|||||||
markdown can't express. Falls back to rebuilding from the file's
|
markdown can't express. Falls back to rebuilding from the file's
|
||||||
structure when the text blocks don't line up one-to-one.
|
structure when the text blocks don't line up one-to-one.
|
||||||
"""
|
"""
|
||||||
preserved = _splice_in_place(stored_group=stored_group, incoming=incoming)
|
preserved = _splice_in_place(
|
||||||
|
stored_group=stored_group,
|
||||||
|
incoming=incoming,
|
||||||
|
thinking_is_progress=thinking_is_progress,
|
||||||
|
)
|
||||||
if preserved is not None:
|
if preserved is not None:
|
||||||
return preserved
|
return preserved
|
||||||
return _splice_by_rebuild(stored_group=stored_group, incoming=incoming)
|
return _splice_by_rebuild(stored_group=stored_group, incoming=incoming)
|
||||||
|
|
||||||
|
|
||||||
def _is_text_block(blk: object) -> bool:
|
|
||||||
if not isinstance(blk, dict):
|
|
||||||
return False
|
|
||||||
block = cast("dict[str, Any]", blk)
|
|
||||||
return block.get("type") == "text" and bool(str(block.get("text", "")).strip())
|
|
||||||
|
|
||||||
|
|
||||||
def _splice_in_place(
|
def _splice_in_place(
|
||||||
*, stored_group: _StoredDisplayTurn, incoming: ParsedTurn
|
*,
|
||||||
|
stored_group: _StoredDisplayTurn,
|
||||||
|
incoming: ParsedTurn,
|
||||||
|
thinking_is_progress: bool,
|
||||||
) -> list[dict[str, Any]] | None:
|
) -> list[dict[str, Any]] | None:
|
||||||
"""Copy the stored messages, swapping only their text block contents.
|
"""Copy the stored messages, swapping only their text block contents.
|
||||||
|
|
||||||
Rebuilding a turn from the file loses everything the markdown never
|
Preserves message/thinking-block boundaries that a rebuild would lose
|
||||||
carried — thinking blocks, and the message boundaries claude chose.
|
(and that a cache-hit continuation depends on); returns ``None`` if
|
||||||
Both matter: an ``assistant[thinking] + assistant[text]`` pair (what
|
text blocks aren't 1:1 with incoming segments or the prose runs through
|
||||||
claude emits for a reasoning turn) collapses into a single message,
|
signed thinking blocks, so the caller rebuilds instead.
|
||||||
so the history is one message shorter than the one the backend
|
|
||||||
pooled its live session under, and the next turn misses the cache
|
|
||||||
and respawns. Substituting in place keeps the message count and the
|
|
||||||
invisible blocks exactly as stored.
|
|
||||||
|
|
||||||
Returns ``None`` when stored text blocks and incoming text segments
|
|
||||||
aren't one-to-one — consecutive text blocks merge into a single
|
|
||||||
rendered segment, so there'd be no way to know how to split the
|
|
||||||
edited prose back apart. The caller then rebuilds instead.
|
|
||||||
"""
|
"""
|
||||||
# See ``diff_and_fork`` for why the parser-type import is deferred.
|
|
||||||
from beaver_gateway.frontends.markdown.parser import TextSegment
|
from beaver_gateway.frontends.markdown.parser import TextSegment
|
||||||
|
|
||||||
new_texts = [
|
new_texts = [
|
||||||
@@ -601,9 +486,12 @@ def _splice_in_place(
|
|||||||
content = msg.get("content")
|
content = msg.get("content")
|
||||||
if msg["role"] != "assistant" or not isinstance(content, list):
|
if msg["role"] != "assistant" or not isinstance(content, list):
|
||||||
continue
|
continue
|
||||||
positions.extend(
|
for bi, blk in enumerate(cast("list[dict[str, Any]]", content)):
|
||||||
(mi, bi) for bi, blk in enumerate(content) if _is_text_block(blk)
|
if not spoken_text(blk, thinking_is_progress=thinking_is_progress):
|
||||||
)
|
continue
|
||||||
|
if blk["type"] == "thinking":
|
||||||
|
return None
|
||||||
|
positions.append((mi, bi))
|
||||||
if len(positions) != len(new_texts):
|
if len(positions) != len(new_texts):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -626,18 +514,9 @@ def _splice_by_rebuild(
|
|||||||
) -> list[dict[str, Any]] | None:
|
) -> list[dict[str, Any]] | None:
|
||||||
"""Rebuild an assistant display turn with new text + stored tool_use blocks.
|
"""Rebuild an assistant display turn with new text + stored tool_use blocks.
|
||||||
|
|
||||||
Walks the incoming structure; for each ``TextSegment`` emits a
|
Matches tool_use blocks to incoming ``ToolSegment``s by position; returns
|
||||||
text block into the current assistant message; for each
|
``None`` (caller forks) if a matching stored ``tool_result`` is missing.
|
||||||
``ToolSegment`` consumes the next stored ``tool_use`` block (by
|
|
||||||
position), closes the current assistant message, emits the
|
|
||||||
matching ``tool_result`` user message, and opens a new assistant
|
|
||||||
message. Final ``TextSegment`` closes the last assistant message.
|
|
||||||
|
|
||||||
Returns ``None`` if we can't find a matching tool_result for some
|
|
||||||
tool_use (stored history is malformed) — caller falls back to
|
|
||||||
fork.
|
|
||||||
"""
|
"""
|
||||||
# See ``diff_and_fork`` for why this import is deferred.
|
|
||||||
from beaver_gateway.frontends.markdown.parser import TextSegment
|
from beaver_gateway.frontends.markdown.parser import TextSegment
|
||||||
|
|
||||||
tool_uses, tool_results_by_id = _harvest_tool_blocks(stored_group)
|
tool_uses, tool_results_by_id = _harvest_tool_blocks(stored_group)
|
||||||
@@ -664,8 +543,6 @@ def _splice_by_rebuild(
|
|||||||
if current_asst:
|
if current_asst:
|
||||||
spliced.append({"role": "assistant", "content": current_asst})
|
spliced.append({"role": "assistant", "content": current_asst})
|
||||||
elif not spliced:
|
elif not spliced:
|
||||||
# Defensive: assistant turn with no text and no tools makes no
|
|
||||||
# sense; caller will treat as fork.
|
|
||||||
return None
|
return None
|
||||||
return spliced
|
return spliced
|
||||||
|
|
||||||
@@ -1,13 +1,8 @@
|
|||||||
"""Vault files for ``deep`` conversations that were not typed into a file (§3.10).
|
"""Vault files for ``deep`` conversations that were not typed into a file.
|
||||||
|
|
||||||
``materialize`` is the markdown frontend's answer to ``spawn(kind=deep)``:
|
``materialize`` creates the file for a conversation spawned elsewhere;
|
||||||
a new file in the vault with ``agent`` + ``conversation_id`` frontmatter
|
``run`` appends replies produced by other origins and keeps the canonical
|
||||||
and the ``(markdown, path)`` binding. ``run`` tails the gateway bus and
|
history in sync so a continuation typed in Obsidian resumes the same session.
|
||||||
appends every ``reply`` of a markdown-bound conversation to its file -
|
|
||||||
the seed turn of a spawn, a message posted through ``/api``, a turn
|
|
||||||
that came in over ``/v1/messages`` - and stamps the same exchange into
|
|
||||||
the canonical history, so a continuation typed in Obsidian aligns
|
|
||||||
against the store and resumes the same SDK session instead of reseeding.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -19,8 +14,6 @@ from typing import TYPE_CHECKING, Any
|
|||||||
|
|
||||||
import frontmatter
|
import frontmatter
|
||||||
|
|
||||||
from beaver_gateway.core.conversation_store import load_messages, rewrite_messages
|
|
||||||
from beaver_gateway.core.turn_record import slugify
|
|
||||||
from beaver_gateway.frontends.markdown import renderer
|
from beaver_gateway.frontends.markdown import renderer
|
||||||
from beaver_gateway.frontends.markdown.crossfront import strip_trailing_user_scaffold
|
from beaver_gateway.frontends.markdown.crossfront import strip_trailing_user_scaffold
|
||||||
from beaver_gateway.frontends.markdown.files import (
|
from beaver_gateway.frontends.markdown.files import (
|
||||||
@@ -28,12 +21,14 @@ from beaver_gateway.frontends.markdown.files import (
|
|||||||
reattach_frontmatter,
|
reattach_frontmatter,
|
||||||
write_atomic,
|
write_atomic,
|
||||||
)
|
)
|
||||||
|
from beaver_gateway.frontends.markdown.history import load_messages, rewrite_messages
|
||||||
|
from beaver_gateway.frontends.turn_record import slugify
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Callable, Sequence
|
from collections.abc import Callable, Sequence
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from beaver_gateway.core.bus import Event
|
from beaver_gateway.events.bus import Event
|
||||||
from beaver_gateway.frontends.base import GatewayRuntime
|
from beaver_gateway.frontends.base import GatewayRuntime
|
||||||
from beaver_gateway.storage.models import Conversation, ConversationBinding
|
from beaver_gateway.storage.models import Conversation, ConversationBinding
|
||||||
|
|
||||||
|
|||||||
@@ -1,18 +1,7 @@
|
|||||||
"""Parse a markdown chat file into Anthropic ``MessageParam`` history.
|
"""Parse a markdown chat file into Anthropic ``MessageParam`` history.
|
||||||
|
|
||||||
The file format is documented in ``frontends/markdown/__init__.py``:
|
Turn markers are ``### User:`` / ``### Assistant:`` H3 headers; backend
|
||||||
``### User:`` / ``### Assistant:`` H3 headers split turns, optional
|
input drops thinking/tool callouts, keeping assistant turns text-only.
|
||||||
``---`` HRs between turns are visual-only, ``> [!thinking]-`` and
|
|
||||||
``> [!tool]- <name>`` callouts mark structured assistant content.
|
|
||||||
|
|
||||||
For backend consumption we strip thinking and tool_use callouts —
|
|
||||||
assistant turns become text-only. Rationale: history replay through
|
|
||||||
claude-code's JSONL injection only needs the *narrated* answer (the
|
|
||||||
thinking signatures expire and the original tool_results aren't
|
|
||||||
captured in the renderer's output, so a faithful tool_use round-trip
|
|
||||||
isn't possible today). The renderer keeps callouts in the file because
|
|
||||||
they're informational for the human reader; the parser drops them when
|
|
||||||
shaping the backend's input.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -51,11 +40,7 @@ class TextSegment:
|
|||||||
class ToolSegment:
|
class ToolSegment:
|
||||||
"""A ``> [!tool]- <name>`` callout placeholder.
|
"""A ``> [!tool]- <name>`` callout placeholder.
|
||||||
|
|
||||||
Only the tool ``name`` is captured — the " · summary" suffix on the
|
Only the tool name is captured; the summary suffix is decorative.
|
||||||
callout title and the JSON body inside the quote block are
|
|
||||||
decorative for the human reader; the canonical tool_use block lives
|
|
||||||
in the DB and is keyed by *position+name* against the structure
|
|
||||||
parsed here.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
name: str
|
name: str
|
||||||
@@ -64,26 +49,11 @@ class ToolSegment:
|
|||||||
AssistantSegment = TextSegment | ToolSegment
|
AssistantSegment = TextSegment | ToolSegment
|
||||||
|
|
||||||
|
|
||||||
# Turn marker — must be exactly ``### User:`` or ``### Assistant:`` on
|
|
||||||
# its own line. Trailing whitespace tolerated; nothing after the colon
|
|
||||||
# on the same line (any inline content would mean the user typed
|
|
||||||
# something that just happens to look like a header, and we'd rather
|
|
||||||
# misparse than silently fold inline content into a turn).
|
|
||||||
_TURN_RE = re.compile(r"^###\s+(User|Assistant):\s*$", re.MULTILINE)
|
_TURN_RE = re.compile(r"^###\s+(User|Assistant):\s*$", re.MULTILINE)
|
||||||
|
|
||||||
# Callout-start lines we strip from assistant turns when extracting
|
|
||||||
# text. We don't try to parse the contents — for backend input we just
|
|
||||||
# need to drop the whole quoted block.
|
|
||||||
_CALLOUT_START_RE = re.compile(r"^>\s+\[!(thinking|tool)\]")
|
_CALLOUT_START_RE = re.compile(r"^>\s+\[!(thinking|tool)\]")
|
||||||
|
|
||||||
# Tool-callout title line: ``> [!tool]- <name>`` or ``> [!tool]- <name> · <summary>``.
|
|
||||||
# We only need the ``<name>`` part for skeleton matching; the summary is
|
|
||||||
# decorative (built by ``renderer.summarize_tool_input`` from inputs the
|
|
||||||
# user can edit visually without semantic consequence).
|
|
||||||
_TOOL_TITLE_RE = re.compile(r"^>\s+\[!tool\]-\s*(.*?)\s*$")
|
_TOOL_TITLE_RE = re.compile(r"^>\s+\[!tool\]-\s*(.*?)\s*$")
|
||||||
# Renderer joins name + summary with " · " (U+00B7) — see
|
|
||||||
# ``renderer.summarize_tool_input``. We split on it to recover the
|
|
||||||
# bare tool name.
|
|
||||||
_TOOL_TITLE_SEP = " · "
|
_TOOL_TITLE_SEP = " · "
|
||||||
|
|
||||||
|
|
||||||
@@ -91,14 +61,8 @@ _TOOL_TITLE_SEP = " · "
|
|||||||
class ParsedTurn:
|
class ParsedTurn:
|
||||||
"""One turn extracted from the chat file.
|
"""One turn extracted from the chat file.
|
||||||
|
|
||||||
``role`` is ``"user"`` or ``"assistant"``. ``text`` is the spoken
|
``text`` is spoken content only; ``structure`` (assistant turns only)
|
||||||
content with callouts stripped and HRs dropped — used both as the
|
carries the ordered text/tool segments used to align with stored history.
|
||||||
backend's ``MessageParam.content`` (back-compat with the existing
|
|
||||||
parser shape) and as the diff key against stored turns.
|
|
||||||
``structure`` is non-empty only for assistant turns: an ordered
|
|
||||||
list of ``TextSegment`` / ``ToolSegment`` reflecting the visible
|
|
||||||
layout of the assistant block, used by the conversation store to
|
|
||||||
align with the canonical tool_use blocks held in DB.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
role: str
|
role: str
|
||||||
@@ -110,16 +74,8 @@ class ParsedTurn:
|
|||||||
class ParsedFile:
|
class ParsedFile:
|
||||||
"""Result of parsing a single chat ``.md``.
|
"""Result of parsing a single chat ``.md``.
|
||||||
|
|
||||||
``metadata`` is the YAML frontmatter as a plain dict (empty if the
|
``messages`` is text-only history for the backend; ``turns`` carries
|
||||||
file has none). ``messages`` is the conversation history shaped for
|
per-turn structure; ``body`` is the raw content after frontmatter.
|
||||||
``Backend.complete`` — assistant turns are text-only. ``turns`` is
|
|
||||||
1:1 with ``messages`` and carries the per-turn structure (for
|
|
||||||
assistant turns) that the conversation store needs to detect
|
|
||||||
text-only edits vs. structural forks. ``body`` is the raw markdown
|
|
||||||
content *after* the frontmatter is stripped; the renderer needs it
|
|
||||||
when it appends a new assistant turn so it can preserve whatever
|
|
||||||
the human typed verbatim (including any callouts or HRs they
|
|
||||||
added).
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
metadata: dict[str, Any]
|
metadata: dict[str, Any]
|
||||||
@@ -131,15 +87,9 @@ class ParsedFile:
|
|||||||
def parse(text: str) -> ParsedFile:
|
def parse(text: str) -> ParsedFile:
|
||||||
"""Parse a chat ``.md`` into ``(metadata, body, messages, turns)``.
|
"""Parse a chat ``.md`` into ``(metadata, body, messages, turns)``.
|
||||||
|
|
||||||
A file with no turn markers but non-empty body is treated as a
|
A bare file (no turn markers) is treated as a single user turn; a
|
||||||
single user turn — the friendly path for "user types into a new
|
tool-only assistant turn gets placeholder ``" "`` content since the
|
||||||
file and hits send" before any turn markers exist.
|
backend rejects an empty string.
|
||||||
|
|
||||||
Assistant turns that have *only* tool callouts (no spoken text) are
|
|
||||||
preserved here even though their ``MessageParam.content`` is empty
|
|
||||||
— the structure carries tool-segment information the conversation
|
|
||||||
store needs for skeleton matching. The renderer in practice always
|
|
||||||
emits at least a trailing text block, so this branch is defensive.
|
|
||||||
"""
|
"""
|
||||||
parsed = frontmatter.loads(text)
|
parsed = frontmatter.loads(text)
|
||||||
metadata = dict(parsed.metadata)
|
metadata = dict(parsed.metadata)
|
||||||
@@ -175,14 +125,6 @@ def parse(text: str) -> ParsedFile:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
elif has_tools:
|
elif has_tools:
|
||||||
# Tool-only assistant turn: nothing to feed the backend
|
|
||||||
# as ``content`` (it'd reject an empty string), but the
|
|
||||||
# structure must survive so the store can align it
|
|
||||||
# against stored tool_use blocks. We synthesize a
|
|
||||||
# single-space text content for backend round-trip; the
|
|
||||||
# conversation store will replace this payload with the
|
|
||||||
# canonical stored blocks before the backend ever sees
|
|
||||||
# it on a continuation.
|
|
||||||
messages.append({"role": "assistant", "content": " "})
|
messages.append({"role": "assistant", "content": " "})
|
||||||
parsed_turns.append(
|
parsed_turns.append(
|
||||||
ParsedTurn(role="assistant", text="", structure=tuple(structure))
|
ParsedTurn(role="assistant", text="", structure=tuple(structure))
|
||||||
@@ -196,18 +138,8 @@ def parse(text: str) -> ParsedFile:
|
|||||||
def parse_assistant_structure(raw: str) -> list[TextSegment | ToolSegment]:
|
def parse_assistant_structure(raw: str) -> list[TextSegment | ToolSegment]:
|
||||||
"""Walk an assistant turn body, return its ordered text/tool segments.
|
"""Walk an assistant turn body, return its ordered text/tool segments.
|
||||||
|
|
||||||
Tool callouts become :class:`ToolSegment` with just the tool name —
|
Tool callouts become :class:`ToolSegment` (name only); thinking callouts
|
||||||
the title's optional ``" · summary"`` suffix and the JSON body
|
and HR lines are stripped; boundary-empty text segments are dropped.
|
||||||
inside the quote block are decorative; the canonical tool_use
|
|
||||||
block is held in the conversation store. Thinking callouts are
|
|
||||||
stripped entirely (they were never round-trippable through the
|
|
||||||
file — signatures expire). HR separator lines drop out.
|
|
||||||
|
|
||||||
Empty / whitespace-only text segments at the boundaries (start,
|
|
||||||
end, between adjacent tool callouts) are dropped so the skeleton
|
|
||||||
is robust against renderer whitespace choices; a non-empty text
|
|
||||||
segment with surrounding whitespace is trimmed on both ends but
|
|
||||||
preserved.
|
|
||||||
"""
|
"""
|
||||||
segments: list[TextSegment | ToolSegment] = []
|
segments: list[TextSegment | ToolSegment] = []
|
||||||
pending_text: list[str] = []
|
pending_text: list[str] = []
|
||||||
@@ -216,9 +148,6 @@ def parse_assistant_structure(raw: str) -> list[TextSegment | ToolSegment]:
|
|||||||
if not pending_text:
|
if not pending_text:
|
||||||
return
|
return
|
||||||
joined = "\n".join(pending_text)
|
joined = "\n".join(pending_text)
|
||||||
# Collapse runs of >2 blank lines (created when we stripped a
|
|
||||||
# mid-block callout) into one so the diff against a re-render
|
|
||||||
# is stable.
|
|
||||||
cleaned = re.sub(r"\n{3,}", "\n\n", joined).strip()
|
cleaned = re.sub(r"\n{3,}", "\n\n", joined).strip()
|
||||||
pending_text.clear()
|
pending_text.clear()
|
||||||
if cleaned:
|
if cleaned:
|
||||||
@@ -231,7 +160,6 @@ def parse_assistant_structure(raw: str) -> list[TextSegment | ToolSegment]:
|
|||||||
callout_match = _CALLOUT_START_RE.match(line)
|
callout_match = _CALLOUT_START_RE.match(line)
|
||||||
if callout_match:
|
if callout_match:
|
||||||
kind = callout_match.group(1)
|
kind = callout_match.group(1)
|
||||||
# Capture tool name *before* advancing past the block.
|
|
||||||
if kind == "tool":
|
if kind == "tool":
|
||||||
title_match = _TOOL_TITLE_RE.match(line)
|
title_match = _TOOL_TITLE_RE.match(line)
|
||||||
title = title_match.group(1) if title_match else ""
|
title = title_match.group(1) if title_match else ""
|
||||||
@@ -239,9 +167,7 @@ def parse_assistant_structure(raw: str) -> list[TextSegment | ToolSegment]:
|
|||||||
_flush_text()
|
_flush_text()
|
||||||
segments.append(ToolSegment(name=name))
|
segments.append(ToolSegment(name=name))
|
||||||
else:
|
else:
|
||||||
# Thinking callout — drop the whole block, emit nothing.
|
|
||||||
_flush_text()
|
_flush_text()
|
||||||
# Skip the rest of the quote block.
|
|
||||||
while i < len(lines) and lines[i].lstrip().startswith(">"):
|
while i < len(lines) and lines[i].lstrip().startswith(">"):
|
||||||
i += 1
|
i += 1
|
||||||
continue
|
continue
|
||||||
@@ -257,11 +183,7 @@ def parse_assistant_structure(raw: str) -> list[TextSegment | ToolSegment]:
|
|||||||
def _segments_to_spoken_text(segments: list[TextSegment | ToolSegment]) -> str:
|
def _segments_to_spoken_text(segments: list[TextSegment | ToolSegment]) -> str:
|
||||||
r"""Reduce a structure list to the spoken-text view the backend sees.
|
r"""Reduce a structure list to the spoken-text view the backend sees.
|
||||||
|
|
||||||
Concatenates :class:`TextSegment` contents with ``\n\n`` between
|
Concatenates text segments with ``\n\n``, dropping tool segments.
|
||||||
them, dropping :class:`ToolSegment` entries. Equivalent to what
|
|
||||||
the pre-Conversation-store parser did — we keep that behavior so
|
|
||||||
existing fingerprints (frontmatter ``fingerprint`` field) stay
|
|
||||||
valid.
|
|
||||||
"""
|
"""
|
||||||
chunks = [s.text for s in segments if isinstance(s, TextSegment)]
|
chunks = [s.text for s in segments if isinstance(s, TextSegment)]
|
||||||
return "\n\n".join(c for c in chunks if c).strip()
|
return "\n\n".join(c for c in chunks if c).strip()
|
||||||
@@ -290,15 +212,11 @@ def resolve_agent(
|
|||||||
return default
|
return default
|
||||||
|
|
||||||
|
|
||||||
# ---- internals ---------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def _split_turns(body: str) -> list[tuple[str, str]]:
|
def _split_turns(body: str) -> list[tuple[str, str]]:
|
||||||
"""Walk turn markers, return ``[(role_lc, raw_body), ...]``.
|
"""Walk turn markers, return ``[(role_lc, raw_body), ...]``.
|
||||||
|
|
||||||
Body for each turn is everything between this marker and the next
|
A marker is an exact ``### User:`` / ``### Assistant:`` line; trailing
|
||||||
(or EOF). Leading marker line itself is dropped. We don't trim
|
content after the colon means it's not a marker, not a turn boundary.
|
||||||
whitespace here — that's per-role.
|
|
||||||
"""
|
"""
|
||||||
matches = list(_TURN_RE.finditer(body))
|
matches = list(_TURN_RE.finditer(body))
|
||||||
if not matches:
|
if not matches:
|
||||||
@@ -315,9 +233,8 @@ def _split_turns(body: str) -> list[tuple[str, str]]:
|
|||||||
def _strip_hrs(raw: str) -> str:
|
def _strip_hrs(raw: str) -> str:
|
||||||
"""Drop decorative ``---`` separator lines (whole-line HRs only).
|
"""Drop decorative ``---`` separator lines (whole-line HRs only).
|
||||||
|
|
||||||
A ``---`` mid-paragraph (rare, but possible) stays. Only lines that
|
Only lines that are exactly ``---`` (with optional surrounding
|
||||||
are *exactly* the HR after optional surrounding whitespace are
|
whitespace) are removed; a ``---`` mid-paragraph stays.
|
||||||
removed — those are the ones the renderer emits between turns.
|
|
||||||
"""
|
"""
|
||||||
lines = raw.splitlines()
|
lines = raw.splitlines()
|
||||||
kept = [ln for ln in lines if ln.strip() != "---"]
|
kept = [ln for ln in lines if ln.strip() != "---"]
|
||||||
|
|||||||
@@ -1,10 +1,7 @@
|
|||||||
"""Render Anthropic ``Message`` (and individual user turns) into markdown.
|
"""Render Anthropic ``Message`` (and user turns) into markdown.
|
||||||
|
|
||||||
The renderer is one-way: it produces the human-facing artifact in the
|
One-way: produces the human-facing artifact only. The parser strips
|
||||||
vault. The parser strips tool/thinking callouts when reshaping the file
|
tool/thinking callouts separately when reshaping history for backend replay.
|
||||||
for backend replay — so what we write here is purely for the human
|
|
||||||
reader (and for the cross-frontend logger, which materializes turns
|
|
||||||
from other frontends).
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -15,8 +12,6 @@ from typing import TYPE_CHECKING, Any, cast
|
|||||||
from anthropic.types import Message, TextBlock, ThinkingBlock
|
from anthropic.types import Message, TextBlock, ThinkingBlock
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Iterable
|
|
||||||
|
|
||||||
from anthropic.types import MessageParam
|
from anthropic.types import MessageParam
|
||||||
|
|
||||||
|
|
||||||
@@ -30,23 +25,11 @@ __all__ = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
# Empty ``### User:`` block appended after each assistant reply so the
|
|
||||||
# human has an obvious place to type the next turn. Parser drops empty
|
|
||||||
# user blocks, so this doesn't re-trigger dispatch on its own.
|
|
||||||
# Blank line after the header, like every rendered turn - the file stays
|
|
||||||
# symmetric whether the human or the gateway wrote the marker.
|
|
||||||
USER_SCAFFOLD = "### User:\n\n"
|
USER_SCAFFOLD = "### User:\n\n"
|
||||||
|
|
||||||
|
|
||||||
# Default 4-backtick fence so tool results that contain literal ```` ``` ````
|
|
||||||
# don't collide. JSON inputs use 3 backticks because they almost never
|
|
||||||
# contain ``` and we get language syntax highlighting in Obsidian for free.
|
|
||||||
FENCE = "````"
|
FENCE = "````"
|
||||||
|
|
||||||
# Input keys we dangle after the tool name in the callout title, best
|
|
||||||
# first — purely cosmetic. ``description`` wins because when a tool
|
|
||||||
# offers one it's a human-written summary of the call, which beats a
|
|
||||||
# truncated shell command or path.
|
|
||||||
_TITLE_KEYS = ("description", "path", "file", "filename", "url", "command", "query")
|
_TITLE_KEYS = ("description", "path", "file", "filename", "url", "command", "query")
|
||||||
|
|
||||||
|
|
||||||
@@ -60,47 +43,37 @@ def render_assistant_text(text: str) -> str:
|
|||||||
return f"### Assistant:\n\n{text.strip()}\n"
|
return f"### Assistant:\n\n{text.strip()}\n"
|
||||||
|
|
||||||
|
|
||||||
def render_assistant_message(message: Message) -> str:
|
def render_assistant_message(message: Message, *, thinking_is_progress: bool) -> str:
|
||||||
"""Render an assistant ``Message`` (with content blocks) into a turn block.
|
"""Render an assistant ``Message`` into a turn block.
|
||||||
|
|
||||||
Blocks render in their original order:
|
Tool-use blocks render to nothing. Progress thinking renders as text in
|
||||||
|
place and vanishes when empty; reasoning renders as a collapsed callout,
|
||||||
* ``ThinkingBlock`` → ``> [!thinking]-`` collapsed callout
|
of which only the first between two texts reaches the file.
|
||||||
* ``TextBlock`` → plain text (the spoken answer)
|
|
||||||
* ``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.
|
|
||||||
"""
|
"""
|
||||||
parts: list[str] = ["### Assistant:", ""]
|
parts: list[str] = ["### Assistant:", ""]
|
||||||
|
after_reasoning = False
|
||||||
for block in message.content:
|
for block in message.content:
|
||||||
lines = list(_render_block(block))
|
if isinstance(block, ThinkingBlock) and not thinking_is_progress:
|
||||||
if not lines:
|
if not after_reasoning:
|
||||||
# Tool calls render to nothing - no separator for them either,
|
parts.extend((*_render_reasoning(block.thinking or ""), ""))
|
||||||
# or every tool leaves a blank line behind.
|
after_reasoning = True
|
||||||
continue
|
continue
|
||||||
parts.extend(lines)
|
text = _spoken(block)
|
||||||
parts.append("")
|
if text:
|
||||||
|
parts.extend((text, ""))
|
||||||
|
after_reasoning = False
|
||||||
return "\n".join(parts).rstrip() + "\n"
|
return "\n".join(parts).rstrip() + "\n"
|
||||||
|
|
||||||
|
|
||||||
def render_user_param(param: MessageParam) -> str:
|
def render_user_param(param: MessageParam) -> str:
|
||||||
"""Render a ``MessageParam`` user message into a ``### User:`` block.
|
"""Render a ``MessageParam`` user message into a ``### User:`` block.
|
||||||
|
|
||||||
Used by the cross-frontend logger when materializing turns from
|
Tool_result blocks in the content list are dropped silently.
|
||||||
other frontends. Tool_result blocks in the content list are dropped
|
|
||||||
silently — the markdown view doesn't track them (see ``parser.py``).
|
|
||||||
"""
|
"""
|
||||||
content = param.get("content", "")
|
content = param.get("content", "")
|
||||||
if isinstance(content, str):
|
if isinstance(content, str):
|
||||||
text = content
|
text = content
|
||||||
else:
|
else:
|
||||||
# The Anthropic SDK types ``content`` as a union of typed-dict
|
|
||||||
# *Param classes plus pydantic block models — both shapes appear
|
|
||||||
# in practice (raw incoming JSON yields dicts, SDK-built params
|
|
||||||
# yield BaseModels). Treat each entry as a dict-like and pull
|
|
||||||
# ``text`` opportunistically.
|
|
||||||
chunks = [
|
chunks = [
|
||||||
str(blk.get("text", ""))
|
str(blk.get("text", ""))
|
||||||
for blk in content
|
for blk in content
|
||||||
@@ -111,11 +84,9 @@ def render_user_param(param: MessageParam) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def append_to_body(existing: str, new_block: str) -> str:
|
def append_to_body(existing: str, new_block: str) -> str:
|
||||||
"""Append ``new_block`` to ``existing`` with a decorative HR separator.
|
"""Append ``new_block`` to ``existing`` with a decorative ``---`` separator.
|
||||||
|
|
||||||
Preserves the original body verbatim (whitespace, callouts, any
|
The separator is visual only; the parser ignores it.
|
||||||
formatting the human added). The HR is purely visual: parser ignores
|
|
||||||
it.
|
|
||||||
"""
|
"""
|
||||||
head = existing.rstrip()
|
head = existing.rstrip()
|
||||||
if head:
|
if head:
|
||||||
@@ -124,18 +95,9 @@ def append_to_body(existing: str, new_block: str) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def summarize_tool_input(name: str, tool_input: object) -> str:
|
def summarize_tool_input(name: str, tool_input: object) -> str:
|
||||||
"""Build the ``[!tool]- <summary>`` title string.
|
"""Build the ``[!tool]- <summary>`` title, picking one salient input field."""
|
||||||
|
|
||||||
Tries to pick a single salient field (``description``, ``path``,
|
|
||||||
``command``, etc.) from the input dict so the collapsed callout shows
|
|
||||||
something meaningful in Obsidian. Falls back to just the tool name.
|
|
||||||
"""
|
|
||||||
if not isinstance(tool_input, dict):
|
if not isinstance(tool_input, dict):
|
||||||
return name
|
return name
|
||||||
# Anthropic ``ToolUseBlock.input`` is typed as ``object`` — the
|
|
||||||
# SDK's runtime value is always a JSON dict (str→Any), so a local
|
|
||||||
# cast keeps the rest of the function readable without sprinkling
|
|
||||||
# per-line type narrowing on every ``.get`` call.
|
|
||||||
d = cast("dict[str, Any]", tool_input)
|
d = cast("dict[str, Any]", tool_input)
|
||||||
for key in _TITLE_KEYS:
|
for key in _TITLE_KEYS:
|
||||||
value = d.get(key)
|
value = d.get(key)
|
||||||
@@ -145,34 +107,21 @@ def summarize_tool_input(name: str, tool_input: object) -> str:
|
|||||||
return name
|
return name
|
||||||
|
|
||||||
|
|
||||||
# ---- internals ---------------------------------------------------------
|
def _render_reasoning(text: str) -> tuple[str, ...]:
|
||||||
|
lines = text.strip().splitlines() or [""]
|
||||||
|
return ("> [!thinking]-", *(f"> {line}" if line else ">" for line in lines))
|
||||||
|
|
||||||
|
|
||||||
def _render_block(block: object) -> Iterable[str]:
|
def _spoken(block: object) -> str:
|
||||||
if isinstance(block, TextBlock):
|
if isinstance(block, TextBlock):
|
||||||
text = (block.text or "").strip()
|
return (block.text or "").strip()
|
||||||
if text:
|
|
||||||
yield text
|
|
||||||
return
|
|
||||||
if isinstance(block, ThinkingBlock):
|
if isinstance(block, ThinkingBlock):
|
||||||
yield from _render_thinking(block.thinking or "")
|
return (block.thinking or "").strip()
|
||||||
return
|
return ""
|
||||||
# Tool-use blocks and unknown block types never reach the file.
|
|
||||||
|
|
||||||
|
|
||||||
def _render_thinking(text: str) -> Iterable[str]:
|
|
||||||
yield "> [!thinking]-"
|
|
||||||
for line in text.strip().splitlines() or [""]:
|
|
||||||
yield f"> {line}" if line else ">"
|
|
||||||
|
|
||||||
|
|
||||||
def adaptive_fence(content: str) -> str:
|
def adaptive_fence(content: str) -> str:
|
||||||
"""Return a backtick fence at least one longer than the longest run in ``content``.
|
"""Return a backtick fence longer than any backtick run in ``content``."""
|
||||||
|
|
||||||
Currently unused (tool *results* aren't persisted yet) — kept here
|
|
||||||
so when result capture lands the rendering side already has the
|
|
||||||
primitive.
|
|
||||||
"""
|
|
||||||
longest = 0
|
longest = 0
|
||||||
for match in re.finditer(r"`+", content):
|
for match in re.finditer(r"`+", content):
|
||||||
longest = max(longest, len(match.group(0)))
|
longest = max(longest, len(match.group(0)))
|
||||||
|
|||||||
@@ -1,36 +1,7 @@
|
|||||||
"""External MCP frontend (Phase 3.1).
|
"""External MCP frontend.
|
||||||
|
|
||||||
A streamable-HTTP gateway in front of the internal MCP aggregator
|
Streamable-HTTP gateway that reverse-proxies the internal MCP aggregator,
|
||||||
(``beaver_gateway.mcp.internal_app``). The aggregator hosts every
|
adding bearer auth, an audit log line per request, and a discovery page.
|
||||||
declared ``McpServer`` (``python_tool``, stdio proxy, HTTP proxy)
|
|
||||||
under ``/mcp/<name>`` plus a flat ``/mcp/all`` bundle on
|
|
||||||
``127.0.0.1:INTERNAL_MCP_PORT`` — that's the *internal* shape.
|
|
||||||
|
|
||||||
This frontend re-exposes those namespaces on its own port directly at
|
|
||||||
``/<name>/`` (no ``/mcp/`` prefix in the external routes — the port
|
|
||||||
itself already disambiguates). Caddy / nginx / Cloudflare in front
|
|
||||||
typically strips a prefix back on: ``domain.com/mcp/* → :8001/*``,
|
|
||||||
controlled by the operator's reverse-proxy config and surfaced to the
|
|
||||||
admin dashboard via ``public_base_url``. Three additions on top of the
|
|
||||||
raw aggregator:
|
|
||||||
|
|
||||||
* **Bearer auth** — ``Authorization: Bearer <token>``, ``X-Api-Key``,
|
|
||||||
or ``?token=<…>`` query string. All three forms verify against the
|
|
||||||
same :class:`TokenStore` as ``AnthropicMessagesFrontend``.
|
|
||||||
* **Audit log** — one line per request (token name, namespace,
|
|
||||||
request method/path, response status). The DB-backed audit log lives
|
|
||||||
in Phase 4; for now we just emit a structured log line.
|
|
||||||
* **Discovery page** at ``GET /`` (auth-gated) — HTML rendered with a
|
|
||||||
tiny inline Jinja2 template listing every namespace plus copy-pastable
|
|
||||||
config snippets for Cursor / claude.ai / Claude Desktop.
|
|
||||||
|
|
||||||
Why a reverse proxy and not a second mount? FastMCP's session managers
|
|
||||||
are tied to the lifespan they were created in; running the same
|
|
||||||
aggregator under two uvicorn servers double-initializes state. Building
|
|
||||||
two parallel aggregators would double upstream connections (two
|
|
||||||
subprocesses for every stdio MCP, two HTTP clients for every remote).
|
|
||||||
A loopback proxy keeps one source of truth — the internal aggregator —
|
|
||||||
and lets us layer policy on the outside.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -45,10 +16,10 @@ from starlette.applications import Starlette
|
|||||||
from starlette.responses import HTMLResponse, JSONResponse, StreamingResponse
|
from starlette.responses import HTMLResponse, JSONResponse, StreamingResponse
|
||||||
from starlette.routing import Route
|
from starlette.routing import Route
|
||||||
|
|
||||||
from beaver_gateway.core import audit
|
|
||||||
from beaver_gateway.frontends._urls import external_base
|
|
||||||
from beaver_gateway.frontends.base import Frontend
|
from beaver_gateway.frontends.base import Frontend
|
||||||
|
from beaver_gateway.frontends.urls import external_base
|
||||||
from beaver_gateway.mcp.internal_app import ALL_NAMESPACE
|
from beaver_gateway.mcp.internal_app import ALL_NAMESPACE
|
||||||
|
from beaver_gateway.security import audit
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Mapping
|
from collections.abc import AsyncIterator, Mapping
|
||||||
@@ -61,10 +32,6 @@ if TYPE_CHECKING:
|
|||||||
_log = logging.getLogger("beaver_gateway.frontends.mcp_server")
|
_log = logging.getLogger("beaver_gateway.frontends.mcp_server")
|
||||||
|
|
||||||
|
|
||||||
# Hop-by-hop headers that must NOT be forwarded across an HTTP proxy
|
|
||||||
# (RFC 7230 §6.1). Bypassing this filter would break chunked transfer
|
|
||||||
# encoding when ``Content-Length`` arrives, or upstream-aware proxies
|
|
||||||
# would refuse the second hop's connection-pool reuse.
|
|
||||||
_HOP_BY_HOP_HEADERS = frozenset(
|
_HOP_BY_HOP_HEADERS = frozenset(
|
||||||
{
|
{
|
||||||
"connection",
|
"connection",
|
||||||
@@ -80,11 +47,6 @@ _HOP_BY_HOP_HEADERS = frozenset(
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
# Standard auth-bearing headers we *do not* forward to the internal app —
|
|
||||||
# the internal app is on loopback with no auth of its own, and forwarding
|
|
||||||
# the inbound bearer would only confuse it. Each method-specific MCP
|
|
||||||
# request from the upstream Cursor/etc. carries a fresh ``mcp-session-id``
|
|
||||||
# that we *must* forward.
|
|
||||||
_AUTH_HEADERS = frozenset({"authorization", "x-api-key"})
|
_AUTH_HEADERS = frozenset({"authorization", "x-api-key"})
|
||||||
|
|
||||||
|
|
||||||
@@ -119,17 +81,16 @@ class McpServerFrontend(Frontend):
|
|||||||
self._http = None
|
self._http = None
|
||||||
|
|
||||||
def _build_app(self, runtime: GatewayRuntime) -> Starlette: # noqa: ARG002
|
def _build_app(self, runtime: GatewayRuntime) -> Starlette: # noqa: ARG002
|
||||||
|
"""Build the Starlette app.
|
||||||
|
|
||||||
|
Literal routes (``/``, ``/healthz``) are listed before the
|
||||||
|
namespace wildcard so they win the match instead of being
|
||||||
|
swallowed by it; two routes per namespace cover both ``/x`` and
|
||||||
|
``/x/y`` since Starlette won't fold them into one.
|
||||||
|
"""
|
||||||
routes = [
|
routes = [
|
||||||
Route("/", self._discovery, methods=["GET"]),
|
Route("/", self._discovery, methods=["GET"]),
|
||||||
Route("/healthz", self._healthz, methods=["GET"]),
|
Route("/healthz", self._healthz, methods=["GET"]),
|
||||||
# Namespaces mount at the root of this port — the port
|
|
||||||
# itself already disambiguates this from any other gateway
|
|
||||||
# surface. Two routes per namespace so both the
|
|
||||||
# trailing-slash and sub-path forms work (``/time`` AND
|
|
||||||
# ``/time/foo``); Starlette doesn't fold them into one
|
|
||||||
# route automatically. The literal routes above (``/``,
|
|
||||||
# ``/healthz``) are listed first and win the match, so
|
|
||||||
# they're not eaten by ``/{namespace}``.
|
|
||||||
Route(
|
Route(
|
||||||
"/{namespace}",
|
"/{namespace}",
|
||||||
self._proxy_endpoint,
|
self._proxy_endpoint,
|
||||||
@@ -151,7 +112,7 @@ class McpServerFrontend(Frontend):
|
|||||||
token_name, err = await _verify_request(request, runtime)
|
token_name, err = await _verify_request(request, runtime)
|
||||||
if err is not None:
|
if err is not None:
|
||||||
return err
|
return err
|
||||||
assert token_name is not None # noqa: S101 — narrow for ty
|
assert token_name is not None # noqa: S101
|
||||||
base = external_base(request, runtime)
|
base = external_base(request, runtime)
|
||||||
html = _render_discovery_page(
|
html = _render_discovery_page(
|
||||||
base_url=base, namespaces=list(runtime.mcps), actor=token_name
|
base_url=base, namespaces=list(runtime.mcps), actor=token_name
|
||||||
@@ -165,7 +126,7 @@ class McpServerFrontend(Frontend):
|
|||||||
token_name, err = await _verify_request(request, runtime)
|
token_name, err = await _verify_request(request, runtime)
|
||||||
if err is not None:
|
if err is not None:
|
||||||
return err
|
return err
|
||||||
assert token_name is not None # noqa: S101 — narrow for ty
|
assert token_name is not None # noqa: S101
|
||||||
|
|
||||||
namespace = request.path_params["namespace"]
|
namespace = request.path_params["namespace"]
|
||||||
subpath = request.path_params.get("path", "")
|
subpath = request.path_params.get("path", "")
|
||||||
@@ -201,11 +162,12 @@ class McpServerFrontend(Frontend):
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _upstream_url(self, namespace: str, subpath: str) -> str | None:
|
def _upstream_url(self, namespace: str, subpath: str) -> str | None:
|
||||||
|
"""Resolve ``namespace`` to its internal loopback URL.
|
||||||
|
|
||||||
|
``ALL_NAMESPACE`` isn't in the URL map, so its URL is synthesized
|
||||||
|
from any per-domain entry's authority.
|
||||||
|
"""
|
||||||
runtime = self._require_runtime()
|
runtime = self._require_runtime()
|
||||||
# ``all`` is built by the aggregator unconditionally when at least
|
|
||||||
# one MCP is configured; the URL map only contains per-domain
|
|
||||||
# entries (see ``build_internal_app``), so we synthesize ``all``'s
|
|
||||||
# loopback URL from any per-domain URL's authority.
|
|
||||||
if namespace == ALL_NAMESPACE:
|
if namespace == ALL_NAMESPACE:
|
||||||
sample = next(iter(runtime.mcp_internal_urls.values()), None)
|
sample = next(iter(runtime.mcp_internal_urls.values()), None)
|
||||||
if sample is None:
|
if sample is None:
|
||||||
@@ -232,13 +194,8 @@ async def _verify_request(
|
|||||||
) -> tuple[str | None, JSONResponse | None]:
|
) -> tuple[str | None, JSONResponse | None]:
|
||||||
"""Accept ``Authorization: Bearer``, ``X-Api-Key``, or ``?token=``.
|
"""Accept ``Authorization: Bearer``, ``X-Api-Key``, or ``?token=``.
|
||||||
|
|
||||||
The third form is the escape hatch for clients that can only put
|
Returns ``(actor_name, None)`` on success, else ``(None, error_response)``
|
||||||
secrets in the URL (claude.ai's MCP config historically did this).
|
— 401 for a missing/unknown token, 403 for one whose scope doesn't cover this call.
|
||||||
All three roads end at the same :class:`TokenStore`. Returns
|
|
||||||
``(actor_name, None)`` on success, ``(None, 401|403)`` otherwise
|
|
||||||
— the caller forwards the response as-is. Splitting auth vs scope
|
|
||||||
failures matters: 401 says "send me a token", 403 says "this token
|
|
||||||
is real but not for this endpoint".
|
|
||||||
"""
|
"""
|
||||||
api_key = request.headers.get("x-api-key")
|
api_key = request.headers.get("x-api-key")
|
||||||
if api_key:
|
if api_key:
|
||||||
@@ -276,9 +233,7 @@ def _forbidden(scope: str, required: str) -> JSONResponse:
|
|||||||
def _join_subpath(base_url: str, subpath: str) -> str:
|
def _join_subpath(base_url: str, subpath: str) -> str:
|
||||||
"""Concatenate the loopback URL with the proxied sub-path.
|
"""Concatenate the loopback URL with the proxied sub-path.
|
||||||
|
|
||||||
``base_url`` always ends in ``/`` (the aggregator publishes URLs
|
``base_url`` always ends in ``/`` to avoid Starlette's redirect dance.
|
||||||
that way to avoid Starlette's 307 redirect dance); the sub-path is
|
|
||||||
appended verbatim, with the query string handled by the caller.
|
|
||||||
"""
|
"""
|
||||||
if subpath:
|
if subpath:
|
||||||
return base_url + subpath.lstrip("/")
|
return base_url + subpath.lstrip("/")
|
||||||
@@ -294,18 +249,14 @@ async def _reverse_proxy(
|
|||||||
actor: str,
|
actor: str,
|
||||||
runtime: GatewayRuntime,
|
runtime: GatewayRuntime,
|
||||||
) -> StreamingResponse | JSONResponse:
|
) -> StreamingResponse | JSONResponse:
|
||||||
"""Bidirectionally stream an MCP request between client ↔ internal aggregator.
|
"""Bidirectionally stream an MCP request between client and internal aggregator.
|
||||||
|
|
||||||
Streamable-HTTP MCP responses can be a long-running SSE stream
|
The audit row is written right after the upstream response headers
|
||||||
(tools that emit partial progress) or a one-shot JSON body; we
|
arrive, before relaying its body, so a client-truncated stream is
|
||||||
don't peek — just relay chunks as they arrive in either direction
|
still audited.
|
||||||
until both sides close.
|
|
||||||
"""
|
"""
|
||||||
qs = request.url.query
|
qs = request.url.query
|
||||||
if qs:
|
if qs:
|
||||||
# Drop ``?token=`` from the forwarded URL — internal app doesn't
|
|
||||||
# need it, and propagating creds further than necessary widens
|
|
||||||
# the leak surface (logs, metrics, traces all see query strings).
|
|
||||||
scrubbed = _scrub_query(qs, drop={"token"})
|
scrubbed = _scrub_query(qs, drop={"token"})
|
||||||
if scrubbed:
|
if scrubbed:
|
||||||
upstream_url = f"{upstream_url}?{scrubbed}"
|
upstream_url = f"{upstream_url}?{scrubbed}"
|
||||||
@@ -351,9 +302,6 @@ async def _reverse_proxy(
|
|||||||
request.url.path,
|
request.url.path,
|
||||||
upstream_resp.status,
|
upstream_resp.status,
|
||||||
)
|
)
|
||||||
# Audit at upstream-response time: status reflects the MCP call's
|
|
||||||
# outcome (200 / tool-error / 4xx). Streaming relay below may be
|
|
||||||
# cut short by the client, but the row is already in by then.
|
|
||||||
await audit.log(
|
await audit.log(
|
||||||
runtime,
|
runtime,
|
||||||
actor=f"token:{actor}",
|
actor=f"token:{actor}",
|
||||||
@@ -370,7 +318,6 @@ async def _reverse_proxy(
|
|||||||
async for chunk in upstream_resp.content.iter_any():
|
async for chunk in upstream_resp.content.iter_any():
|
||||||
yield chunk
|
yield chunk
|
||||||
except (aiohttp.ClientError, asyncio.CancelledError):
|
except (aiohttp.ClientError, asyncio.CancelledError):
|
||||||
# Caller hung up or upstream dropped — just stop relaying.
|
|
||||||
return
|
return
|
||||||
finally:
|
finally:
|
||||||
upstream_resp.release()
|
upstream_resp.release()
|
||||||
@@ -390,6 +337,11 @@ async def _request_body_iter(request: Request) -> AsyncIterator[bytes]:
|
|||||||
|
|
||||||
|
|
||||||
def _forward_headers(request: Request) -> dict[str, str]:
|
def _forward_headers(request: Request) -> dict[str, str]:
|
||||||
|
"""Drop hop-by-hop headers (RFC 7230 §6.1) and inbound auth headers.
|
||||||
|
|
||||||
|
The internal aggregator is loopback-only with no auth of its own,
|
||||||
|
so forwarding the caller's bearer would only confuse it.
|
||||||
|
"""
|
||||||
out: dict[str, str] = {}
|
out: dict[str, str] = {}
|
||||||
for key, value in request.headers.items():
|
for key, value in request.headers.items():
|
||||||
lowered = key.lower()
|
lowered = key.lower()
|
||||||
@@ -422,11 +374,7 @@ def _scrub_query(query: str, *, drop: frozenset[str] | set[str]) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def _render_discovery_page(*, base_url: str, namespaces: list[Any], actor: str) -> str:
|
def _render_discovery_page(*, base_url: str, namespaces: list[Any], actor: str) -> str:
|
||||||
"""Render the auth-gated namespace + config-snippet page.
|
"""Render the auth-gated namespace + config-snippet page."""
|
||||||
|
|
||||||
Inline HTML (no Jinja file) — keeps Phase 3 free of template-dir
|
|
||||||
plumbing that Phase 4's AdminFrontend will own.
|
|
||||||
"""
|
|
||||||
name_list = [getattr(ns, "name", str(ns)) for ns in namespaces]
|
name_list = [getattr(ns, "name", str(ns)) for ns in namespaces]
|
||||||
rows = (
|
rows = (
|
||||||
"\n".join(
|
"\n".join(
|
||||||
@@ -460,7 +408,7 @@ def _escape(value: str) -> str:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
_DISCOVERY_TEMPLATE = "\n".join( # noqa: FLY002 — readability beats one-string-blob
|
_DISCOVERY_TEMPLATE = "\n".join( # noqa: FLY002
|
||||||
[
|
[
|
||||||
"<!doctype html>",
|
"<!doctype html>",
|
||||||
'<html lang="en">',
|
'<html lang="en">',
|
||||||
|
|||||||
@@ -46,7 +46,7 @@ def build_root_app(
|
|||||||
]
|
]
|
||||||
for fe in mounted:
|
for fe in mounted:
|
||||||
app = fe.app()
|
app = fe.app()
|
||||||
assert app is not None # noqa: S101 - filtered above; narrows for ty
|
assert app is not None # noqa: S101
|
||||||
assert fe.path is not None # noqa: S101
|
assert fe.path is not None # noqa: S101
|
||||||
routes.append(Mount(fe.path, app=app, name=fe.name or fe.path.strip("/")))
|
routes.append(Mount(fe.path, app=app, name=fe.name or fe.path.strip("/")))
|
||||||
for path, app in (extra or {}).items():
|
for path, app in (extra or {}).items():
|
||||||
|
|||||||
@@ -29,10 +29,8 @@ async def events_with_heartbeat(
|
|||||||
) -> AsyncIterator[Any]:
|
) -> AsyncIterator[Any]:
|
||||||
"""Pass ``events`` through, yielding ``None`` after ``interval`` seconds of silence.
|
"""Pass ``events`` through, yielding ``None`` after ``interval`` seconds of silence.
|
||||||
|
|
||||||
One in-flight ``__anext__`` task is reused across timeouts: a second
|
Only one consumer may iterate the result at a time; cancelling the
|
||||||
consumer on the same async generator raises ``RuntimeError``.
|
outer scope cancels the in-flight upstream fetch instead of leaking it.
|
||||||
Cancellation of the outer scope cancels that task instead of leaving
|
|
||||||
it dangling.
|
|
||||||
"""
|
"""
|
||||||
src = events.__aiter__()
|
src = events.__aiter__()
|
||||||
next_task: asyncio.Task[Any] | None = None
|
next_task: asyncio.Task[Any] | None = None
|
||||||
@@ -1,5 +1,6 @@
|
|||||||
"""Telegram frontend (§3.8): General = master, topic = branch, drafts, inbox/outbox."""
|
"""Telegram frontend: General = master, topic = branch, drafts, inbox/outbox."""
|
||||||
|
|
||||||
from beaver_gateway.frontends.telegram.frontend import Attachments, TelegramFrontend
|
from beaver_gateway.frontends.telegram.frontend import Attachments, TelegramFrontend
|
||||||
|
from beaver_gateway.frontends.telegram.texts import TelegramTexts
|
||||||
|
|
||||||
__all__ = ["Attachments", "TelegramFrontend"]
|
__all__ = ["Attachments", "TelegramFrontend", "TelegramTexts"]
|
||||||
@@ -1,19 +1,15 @@
|
|||||||
"""One ``sendMessageDraft`` stream per running turn (§3.8).
|
"""One ``sendMessageDraft`` stream per running turn.
|
||||||
|
|
||||||
The client folds a draft into the message that follows only when their
|
Telegram folds a draft into the following message only when their texts are
|
||||||
texts are identical - so the last push is the final text itself, rendered
|
identical, so the last push is the final text verbatim. Drafts live 30s and
|
||||||
exactly as the outbox will send it, with no status line.
|
throttle to about one edit/s, so this refreshes on a timer, not every delta.
|
||||||
|
|
||||||
A draft is ephemeral and lives 30 s, Telegram throttles edits to about one
|
|
||||||
per second per chat, and thinking or a tool call would otherwise look like a
|
|
||||||
hang - so the draft opens with a status line straight away, is refreshed on
|
|
||||||
a timer rather than on every delta, and is kept alive while nothing changes.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import contextlib
|
import contextlib
|
||||||
|
import html
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
import zlib
|
import zlib
|
||||||
@@ -21,7 +17,7 @@ from typing import TYPE_CHECKING
|
|||||||
|
|
||||||
from aiogram.exceptions import TelegramAPIError, TelegramBadRequest
|
from aiogram.exceptions import TelegramAPIError, TelegramBadRequest
|
||||||
|
|
||||||
from beaver_gateway.frontends.telegram.render import to_html
|
from beaver_gateway.frontends.telegram.render import to_html, to_html_tail
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from aiogram import Bot
|
from aiogram import Bot
|
||||||
@@ -31,6 +27,7 @@ __all__ = ["Draft"]
|
|||||||
_log = logging.getLogger("beaver_gateway.frontends.telegram.drafts")
|
_log = logging.getLogger("beaver_gateway.frontends.telegram.drafts")
|
||||||
|
|
||||||
_TAIL = 3500
|
_TAIL = 3500
|
||||||
|
_REASONING_TAIL = 400
|
||||||
_KEEPALIVE = 20.0
|
_KEEPALIVE = 20.0
|
||||||
|
|
||||||
|
|
||||||
@@ -43,7 +40,7 @@ class Draft:
|
|||||||
thread_id: int | None,
|
thread_id: int | None,
|
||||||
turn_id: str,
|
turn_id: str,
|
||||||
interval: float = 0.7,
|
interval: float = 0.7,
|
||||||
status: str = "⏳ думаю",
|
status: str = "⏳ thinking",
|
||||||
) -> None:
|
) -> None:
|
||||||
self._bot = bot
|
self._bot = bot
|
||||||
self._chat_id = chat_id
|
self._chat_id = chat_id
|
||||||
@@ -52,12 +49,18 @@ class Draft:
|
|||||||
self._interval = interval
|
self._interval = interval
|
||||||
self.status = status
|
self.status = status
|
||||||
self.text = ""
|
self.text = ""
|
||||||
|
self.reasoning = ""
|
||||||
|
self._paused = False
|
||||||
self._dirty = True
|
self._dirty = True
|
||||||
self._broken = False
|
self._broken = False
|
||||||
self._last_sent = 0.0
|
self._last_sent = 0.0
|
||||||
self._task: asyncio.Task[None] | None = None
|
self._task: asyncio.Task[None] | None = None
|
||||||
self._inflight: asyncio.Future[None] | None = None
|
self._inflight: asyncio.Future[None] | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def draft_id(self) -> int:
|
||||||
|
return self._draft_id
|
||||||
|
|
||||||
def start(self) -> None:
|
def start(self) -> None:
|
||||||
if self._task is None:
|
if self._task is None:
|
||||||
self._task = asyncio.create_task(self._run())
|
self._task = asyncio.create_task(self._run())
|
||||||
@@ -66,10 +69,32 @@ class Draft:
|
|||||||
if status != self.status:
|
if status != self.status:
|
||||||
self.status = status
|
self.status = status
|
||||||
self._dirty = True
|
self._dirty = True
|
||||||
|
self._paused = bool(self.text.strip())
|
||||||
|
|
||||||
def append(self, text: str) -> None:
|
def append(self, text: str) -> None:
|
||||||
if text:
|
"""Add streamed text; a tool call in between starts a new paragraph.
|
||||||
|
|
||||||
|
The model ends a sentence, calls a tool, then goes on without a
|
||||||
|
break - the same text the final message shows as paragraphs.
|
||||||
|
"""
|
||||||
|
if not text:
|
||||||
|
return
|
||||||
|
if self._paused and not self.text.endswith("\n"):
|
||||||
|
self.text = self.text.rstrip() + "\n\n" + text.lstrip()
|
||||||
|
else:
|
||||||
self.text += text
|
self.text += text
|
||||||
|
self.reasoning = ""
|
||||||
|
self._paused = False
|
||||||
|
self._dirty = True
|
||||||
|
|
||||||
|
def think(self, text: str) -> None:
|
||||||
|
if text:
|
||||||
|
self.reasoning += text
|
||||||
|
self._dirty = True
|
||||||
|
|
||||||
|
def rethink(self) -> None:
|
||||||
|
if self.reasoning:
|
||||||
|
self.reasoning = ""
|
||||||
self._dirty = True
|
self._dirty = True
|
||||||
|
|
||||||
async def stop(self) -> None:
|
async def stop(self) -> None:
|
||||||
@@ -112,12 +137,12 @@ class Draft:
|
|||||||
async def _push(self) -> None:
|
async def _push(self) -> None:
|
||||||
self._dirty = False
|
self._dirty = False
|
||||||
self._last_sent = time.monotonic()
|
self._last_sent = time.monotonic()
|
||||||
text = self._render()
|
rendered, plain = self._render()
|
||||||
try:
|
try:
|
||||||
try:
|
try:
|
||||||
await self._send(to_html(text), "HTML")
|
await self._send(rendered, "HTML")
|
||||||
except TelegramBadRequest:
|
except TelegramBadRequest:
|
||||||
await self._send(text, None)
|
await self._send(plain, None)
|
||||||
except TelegramAPIError as exc:
|
except TelegramAPIError as exc:
|
||||||
self._broken = True
|
self._broken = True
|
||||||
_log.warning(
|
_log.warning(
|
||||||
@@ -131,8 +156,25 @@ class Draft:
|
|||||||
message_thread_id=self._thread_id,
|
message_thread_id=self._thread_id,
|
||||||
text=text,
|
text=text,
|
||||||
parse_mode=parse_mode,
|
parse_mode=parse_mode,
|
||||||
|
can_stop=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _render(self) -> str:
|
def _render(self) -> tuple[str, str]:
|
||||||
|
rendered, plain = [to_html(self.status)], [self.status]
|
||||||
|
thought = _tail(self.reasoning.strip(), _REASONING_TAIL)
|
||||||
|
if thought:
|
||||||
|
rendered.append(f"<blockquote>{html.escape(thought)}</blockquote>")
|
||||||
|
plain.append(thought)
|
||||||
tail = self.text[-_TAIL:]
|
tail = self.text[-_TAIL:]
|
||||||
return f"{self.status}\n\n{tail}" if tail.strip() else self.status
|
if tail.strip():
|
||||||
|
rendered.append(to_html_tail(self.text, _TAIL))
|
||||||
|
plain.append(tail)
|
||||||
|
return "\n\n".join(rendered), "\n\n".join(plain)
|
||||||
|
|
||||||
|
|
||||||
|
def _tail(text: str, limit: int) -> str:
|
||||||
|
if len(text) <= limit:
|
||||||
|
return text
|
||||||
|
cut = text[-limit:]
|
||||||
|
_, _, rest = cut.partition(" ")
|
||||||
|
return f"… {rest or cut}"
|
||||||
File diff suppressed because it is too large.
Load diff
@@ -1,8 +1,7 @@
|
|||||||
"""Long-polling inbox (§3.8).
|
"""Long-polling inbox.
|
||||||
|
|
||||||
Every update lands in ``telegram_updates`` before the offset moves past it;
|
Every update lands in ``telegram_updates`` before the offset moves past it;
|
||||||
a worker handles rows from the table, oldest first, and finishes whatever a
|
a worker handles rows from the table, oldest first, resuming after a restart.
|
||||||
previous process left unprocessed at startup.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -32,7 +31,7 @@ __all__ = ["Inbox"]
|
|||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.frontends.telegram.inbox")
|
_log = logging.getLogger("beaver_gateway.frontends.telegram.inbox")
|
||||||
|
|
||||||
ALLOWED_UPDATES = ("message", "callback_query")
|
ALLOWED_UPDATES = ("message", "callback_query", "stopped_message_generation")
|
||||||
|
|
||||||
|
|
||||||
def _unset(_value: object) -> None:
|
def _unset(_value: object) -> None:
|
||||||
|
|||||||
@@ -1,8 +1,7 @@
|
|||||||
"""Outbox (§3.8): a reply is a ``deliveries`` row first, a message second.
|
"""Outbox: a reply is a ``deliveries`` row first, a message second.
|
||||||
|
|
||||||
Rows are sent oldest first, retried with backoff on network errors and
|
Rows are sent oldest first, retried with backoff on network errors, resent
|
||||||
flood limits, resent as plain text when Telegram rejects our HTML, and
|
as plain text when Telegram rejects our HTML, and given up once it's gone.
|
||||||
given up only when Telegram says the window is gone.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -32,7 +31,7 @@ from beaver_gateway.storage.models import Delivery
|
|||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from aiogram import Bot
|
from aiogram import Bot
|
||||||
|
|
||||||
from beaver_gateway.core.bus import EventBus
|
from beaver_gateway.events.bus import EventBus
|
||||||
from beaver_gateway.storage.db import Database
|
from beaver_gateway.storage.db import Database
|
||||||
|
|
||||||
__all__ = ["Outbox"]
|
__all__ = ["Outbox"]
|
||||||
@@ -143,7 +142,7 @@ class Outbox:
|
|||||||
if "parse" in text and not row.plain:
|
if "parse" in text and not row.plain:
|
||||||
await self._retry(row, str(exc), delay=0.0, plain=True)
|
await self._retry(row, str(exc), delay=0.0, plain=True)
|
||||||
elif any(marker in text for marker in _GONE):
|
elif any(marker in text for marker in _GONE):
|
||||||
await self._fail(row, str(exc))
|
await self._fail(row, str(exc), gone=True)
|
||||||
else:
|
else:
|
||||||
await self._fail(row, str(exc))
|
await self._fail(row, str(exc))
|
||||||
except (TelegramNotFound, TelegramForbiddenError) as exc:
|
except (TelegramNotFound, TelegramForbiddenError) as exc:
|
||||||
@@ -180,7 +179,7 @@ class Outbox:
|
|||||||
if delay < 5.0:
|
if delay < 5.0:
|
||||||
self._wake.set()
|
self._wake.set()
|
||||||
|
|
||||||
async def _fail(self, row: Delivery, error: str) -> None:
|
async def _fail(self, row: Delivery, error: str, *, gone: bool = False) -> None:
|
||||||
_log.error(
|
_log.error(
|
||||||
"delivery #%s to %s/%s given up: %s",
|
"delivery #%s to %s/%s given up: %s",
|
||||||
row.id,
|
row.id,
|
||||||
@@ -196,6 +195,7 @@ class Outbox:
|
|||||||
chat_id=row.chat_id,
|
chat_id=row.chat_id,
|
||||||
thread_id=row.thread_id,
|
thread_id=row.thread_id,
|
||||||
error=error,
|
error=error,
|
||||||
|
gone=gone,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _mark(
|
async def _mark(
|
||||||
|
|||||||
@@ -4,96 +4,89 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import html
|
import html
|
||||||
import re
|
import re
|
||||||
from typing import Any
|
from dataclasses import dataclass
|
||||||
|
from typing import TYPE_CHECKING, Any, Literal
|
||||||
|
|
||||||
__all__ = ["LIMIT", "chunks", "status_label", "to_html"]
|
from beaver_gateway.frontends.telegram.texts import TOOL_LABELS
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Callable, Mapping
|
||||||
|
|
||||||
|
__all__ = ["LIMIT", "chunks", "status_label", "to_html", "to_html_tail"]
|
||||||
|
|
||||||
LIMIT = 4000
|
LIMIT = 4000
|
||||||
_FENCE = re.compile(r"```[^\n]*\n(.*?)(?:```|$)", re.DOTALL)
|
_SNAP = 200
|
||||||
_INLINE_CODE = re.compile(r"(`[^`\n]+`)")
|
|
||||||
_HEADING = re.compile(r"^#{1,6}\s+(.+)$", re.MULTILINE)
|
_FENCE_OPEN = re.compile(r"^\s*```(.*)$")
|
||||||
_BOLD = re.compile(r"\*\*(.+?)\*\*|__(.+?)__", re.DOTALL)
|
_FENCE_CLOSE = re.compile(r"^\s*```\s*$")
|
||||||
|
_LANG = re.compile(r"[\w+#.-]+")
|
||||||
|
_QUOTE = re.compile(r"^\s*(?:>\s?)+")
|
||||||
|
_TABLE_SEP = re.compile(r"^\s*\|?\s*:?-+:?\s*(?:\|\s*:?-+:?\s*)*\|?\s*$")
|
||||||
|
_CELL_SPLIT = re.compile(r"(?<!\\)\|")
|
||||||
|
_HEADING = re.compile(r"^\s{0,3}#{1,6}\s+(.+?)\s*#*\s*$")
|
||||||
|
_RULE = re.compile(r"^\s*([-*_])(?:\s*\1){2,}\s*$")
|
||||||
|
_TASK = re.compile(r"^(\s*)[-*+]\s+\[([ xX])\]\s*(.*)$")
|
||||||
|
_BULLET = re.compile(r"^(\s*)[-*+]\s+(.*)$")
|
||||||
|
_NUMBERED = re.compile(r"^(\s*)(\d+)[.)]\s+(.*)$")
|
||||||
|
_TOKEN = re.compile(
|
||||||
|
r"`[^`\n]+`|\[\[[^\]\n]+\]\]|!?\[[^\]]*\]\((?:https?|tg)://[^)\s]+\)"
|
||||||
|
)
|
||||||
|
_WIKILINK = re.compile(r"\[\[([^\]|#]+)(?:#([^\]|]+))?(?:\|([^\]]+))?\]\]")
|
||||||
|
_LINK = re.compile(r"(!?)\[([^\]]*)\]\(([^)\s]+)\)")
|
||||||
|
_BOLD = re.compile(r"\*\*(.+?)\*\*|__(.+?)__")
|
||||||
_ITALIC = re.compile(r"(?<![\w*])\*(?!\s)(.+?)(?<!\s)\*(?![\w*])")
|
_ITALIC = re.compile(r"(?<![\w*])\*(?!\s)(.+?)(?<!\s)\*(?![\w*])")
|
||||||
_ITALIC_U = re.compile(r"(?<![\w_])_(?!\s)(.+?)(?<!\s)_(?![\w_])")
|
_ITALIC_U = re.compile(r"(?<![\w_])_(?!\s)(.+?)(?<!\s)_(?![\w_])")
|
||||||
_LINK = re.compile(r"\[([^\]]+)\]\((https?://[^)\s]+)\)")
|
|
||||||
_BULLET = re.compile(r"^(\s*)[-*]\s+", re.MULTILINE)
|
|
||||||
_STRIKE = re.compile(r"~~(.+?)~~")
|
_STRIKE = re.compile(r"~~(.+?)~~")
|
||||||
|
_SPOILER = re.compile(r"\|\|(.+?)\|\|")
|
||||||
|
|
||||||
_LABELS: dict[str, str] = {
|
|
||||||
"Read": "читаю vault",
|
_Kind = Literal["fence", "quote", "table", "text"]
|
||||||
"Glob": "ищу файлы",
|
|
||||||
"Grep": "ищу в vault",
|
|
||||||
"Edit": "правлю файл",
|
@dataclass(slots=True)
|
||||||
"Write": "пишу файл",
|
class _Block:
|
||||||
"MultiEdit": "правлю файлы",
|
kind: _Kind
|
||||||
"Bash": "выполняю команду",
|
lines: list[str]
|
||||||
"WebSearch": "ищу в сети",
|
|
||||||
"WebFetch": "читаю страницу",
|
|
||||||
"Task": "запустил сабагента",
|
|
||||||
"Agent": "запустил сабагента",
|
|
||||||
"AskUserQuestion": "спрашиваю",
|
|
||||||
"TodoWrite": "планирую",
|
|
||||||
"Skill": "открываю скилл",
|
|
||||||
"mcp__gateway__spawn": "открываю разговор",
|
|
||||||
"mcp__gateway__read_conversation": "читаю разговор",
|
|
||||||
"mcp__gateway__say": "говорю",
|
|
||||||
"mcp__gateway__schedule": "ставлю напоминание",
|
|
||||||
"mcp__gateway__inject": "передаю в другой разговор",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def to_html(text: str) -> str:
|
def to_html(text: str) -> str:
|
||||||
out: list[str] = []
|
return "\n".join(_render(block) for block in _blocks(text)).strip()
|
||||||
pos = 0
|
|
||||||
for match in _FENCE.finditer(text):
|
|
||||||
out.append(_inline(text[pos : match.start()]))
|
|
||||||
out.append(f"<pre>{html.escape(match.group(1).rstrip())}</pre>")
|
|
||||||
pos = match.end()
|
|
||||||
out.append(_inline(text[pos:]))
|
|
||||||
return "".join(out).strip()
|
|
||||||
|
|
||||||
|
|
||||||
def _inline(text: str) -> str:
|
def to_html_tail(text: str, tail_chars: int) -> str:
|
||||||
parts = _INLINE_CODE.split(text)
|
start = max(len(text) - tail_chars, 0)
|
||||||
for i, part in enumerate(parts):
|
if start:
|
||||||
if i % 2:
|
newline = text.find("\n", start, start + _SNAP)
|
||||||
parts[i] = f"<code>{html.escape(part[1:-1])}</code>"
|
if newline >= 0:
|
||||||
continue
|
start = newline + 1
|
||||||
s = html.escape(part, quote=False)
|
opening = _open_fence(text[:start])
|
||||||
s = _HEADING.sub(r"<b>\1</b>", s)
|
tail = text[start:]
|
||||||
s = _BOLD.sub(lambda m: f"<b>{m.group(1) or m.group(2)}</b>", s)
|
return to_html(f"{opening}\n{tail}" if opening else tail)
|
||||||
s = _ITALIC.sub(r"<i>\1</i>", s)
|
|
||||||
s = _ITALIC_U.sub(r"<i>\1</i>", s)
|
|
||||||
s = _STRIKE.sub(r"<s>\1</s>", s)
|
|
||||||
s = _LINK.sub(r'<a href="\2">\1</a>', s)
|
|
||||||
parts[i] = _BULLET.sub(r"\1• ", s)
|
|
||||||
return "".join(parts)
|
|
||||||
|
|
||||||
|
|
||||||
def chunks(text: str, limit: int = LIMIT) -> list[str]:
|
def chunks(text: str, limit: int = LIMIT) -> list[str]:
|
||||||
text = text.strip()
|
text = text.strip()
|
||||||
if len(text) <= limit:
|
if not text:
|
||||||
return [text] if text else []
|
return []
|
||||||
|
if _fits(text, limit):
|
||||||
|
return [text]
|
||||||
out: list[str] = []
|
out: list[str] = []
|
||||||
while len(text) > limit:
|
cur: list[str] = []
|
||||||
cut = _cut_point(text, limit)
|
for block in _blocks(text):
|
||||||
out.append(text[:cut].rstrip())
|
for piece in _pieces(block, limit):
|
||||||
text = text[cut:].lstrip()
|
if cur and not _fits("\n".join([*cur, piece]), limit):
|
||||||
if text:
|
_flush(out, cur)
|
||||||
out.append(text)
|
cur.append(piece)
|
||||||
|
_flush(out, cur)
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
|
||||||
def _cut_point(text: str, limit: int) -> int:
|
def status_label(
|
||||||
for sep in ("\n\n", "\n", ". ", " "):
|
name: str,
|
||||||
cut = text.rfind(sep, limit // 2, limit)
|
tool_input: dict[str, Any] | None = None,
|
||||||
if cut > 0:
|
labels: Mapping[str, str] | None = None,
|
||||||
return cut + len(sep)
|
) -> str:
|
||||||
return limit
|
label = (labels if labels is not None else TOOL_LABELS).get(name)
|
||||||
|
|
||||||
|
|
||||||
def status_label(name: str, tool_input: dict[str, Any] | None = None) -> str:
|
|
||||||
label = _LABELS.get(name)
|
|
||||||
if label is not None:
|
if label is not None:
|
||||||
return label
|
return label
|
||||||
if name.startswith("mcp__"):
|
if name.startswith("mcp__"):
|
||||||
@@ -109,3 +102,245 @@ def status_label(name: str, tool_input: dict[str, Any] | None = None) -> str:
|
|||||||
hint = value.strip().splitlines()[0][:60]
|
hint = value.strip().splitlines()[0][:60]
|
||||||
break
|
break
|
||||||
return f"{name} {hint}" if hint else name
|
return f"{name} {hint}" if hint else name
|
||||||
|
|
||||||
|
|
||||||
|
def _blocks(text: str) -> list[_Block]:
|
||||||
|
lines = text.split("\n")
|
||||||
|
blocks: list[_Block] = []
|
||||||
|
i = 0
|
||||||
|
while i < len(lines):
|
||||||
|
line = lines[i]
|
||||||
|
if _FENCE_OPEN.match(line):
|
||||||
|
j = i + 1
|
||||||
|
while j < len(lines) and not _FENCE_CLOSE.match(lines[j]):
|
||||||
|
j += 1
|
||||||
|
end = min(j + 1, len(lines))
|
||||||
|
blocks.append(_Block("fence", lines[i:end]))
|
||||||
|
i = end
|
||||||
|
elif _QUOTE.match(line):
|
||||||
|
j = i + 1
|
||||||
|
while j < len(lines) and _QUOTE.match(lines[j]):
|
||||||
|
j += 1
|
||||||
|
blocks.append(_Block("quote", lines[i:j]))
|
||||||
|
i = j
|
||||||
|
elif _is_table_start(lines, i):
|
||||||
|
j = i + 2
|
||||||
|
while j < len(lines) and "|" in lines[j] and lines[j].strip():
|
||||||
|
j += 1
|
||||||
|
blocks.append(_Block("table", lines[i:j]))
|
||||||
|
i = j
|
||||||
|
elif not line.strip():
|
||||||
|
blocks.append(_Block("text", [line]))
|
||||||
|
i += 1
|
||||||
|
else:
|
||||||
|
j = i + 1
|
||||||
|
while j < len(lines) and _is_plain(lines, j):
|
||||||
|
j += 1
|
||||||
|
blocks.append(_Block("text", lines[i:j]))
|
||||||
|
i = j
|
||||||
|
return blocks
|
||||||
|
|
||||||
|
|
||||||
|
def _is_table_start(lines: list[str], i: int) -> bool:
|
||||||
|
if "|" not in lines[i] or i + 1 >= len(lines):
|
||||||
|
return False
|
||||||
|
sep = lines[i + 1]
|
||||||
|
return "|" in sep and _TABLE_SEP.match(sep) is not None
|
||||||
|
|
||||||
|
|
||||||
|
def _is_plain(lines: list[str], i: int) -> bool:
|
||||||
|
line = lines[i]
|
||||||
|
return bool(
|
||||||
|
line.strip()
|
||||||
|
and not _FENCE_OPEN.match(line)
|
||||||
|
and not _QUOTE.match(line)
|
||||||
|
and not _is_table_start(lines, i)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _open_fence(text: str) -> str | None:
|
||||||
|
opening: str | None = None
|
||||||
|
for line in text.split("\n"):
|
||||||
|
if opening is None:
|
||||||
|
if _FENCE_OPEN.match(line):
|
||||||
|
opening = line
|
||||||
|
elif _FENCE_CLOSE.match(line):
|
||||||
|
opening = None
|
||||||
|
return opening
|
||||||
|
|
||||||
|
|
||||||
|
def _render(block: _Block) -> str:
|
||||||
|
if block.kind == "fence":
|
||||||
|
return _render_fence(block.lines)
|
||||||
|
if block.kind == "quote":
|
||||||
|
body = "\n".join(_line(_QUOTE.sub("", line, 1)) for line in block.lines)
|
||||||
|
return f"<blockquote>{body}</blockquote>" if body.strip() else ""
|
||||||
|
if block.kind == "table":
|
||||||
|
return _render_table(block.lines)
|
||||||
|
return "\n".join(_line(line) for line in block.lines)
|
||||||
|
|
||||||
|
|
||||||
|
def _fence_parts(lines: list[str]) -> tuple[str, list[str]]:
|
||||||
|
body = lines[1:]
|
||||||
|
if body and _FENCE_CLOSE.match(body[-1]):
|
||||||
|
body = body[:-1]
|
||||||
|
return lines[0], body
|
||||||
|
|
||||||
|
|
||||||
|
def _render_fence(lines: list[str]) -> str:
|
||||||
|
head, body = _fence_parts(lines)
|
||||||
|
code = html.escape("\n".join(body).rstrip())
|
||||||
|
if not code:
|
||||||
|
return ""
|
||||||
|
lang = _language(head)
|
||||||
|
if lang:
|
||||||
|
return f'<pre><code class="language-{lang}">{code}</code></pre>'
|
||||||
|
return f"<pre>{code}</pre>"
|
||||||
|
|
||||||
|
|
||||||
|
def _language(head: str) -> str:
|
||||||
|
match = _FENCE_OPEN.match(head)
|
||||||
|
info = match.group(1).strip() if match else ""
|
||||||
|
token = info.split()[0] if info else ""
|
||||||
|
return token if _LANG.fullmatch(token) else ""
|
||||||
|
|
||||||
|
|
||||||
|
def _render_table(lines: list[str]) -> str:
|
||||||
|
rows = [_cells(line) for i, line in enumerate(lines) if i != 1]
|
||||||
|
width = max(len(row) for row in rows)
|
||||||
|
rows = [[*row, *[""] * (width - len(row))] for row in rows]
|
||||||
|
widths = [max(len(row[c]) for row in rows) for c in range(width)]
|
||||||
|
text = "\n".join(_table_row(row, widths) for row in rows)
|
||||||
|
return f"<pre>{html.escape(text)}</pre>"
|
||||||
|
|
||||||
|
|
||||||
|
def _table_row(row: list[str], widths: list[int]) -> str:
|
||||||
|
cells = (cell.ljust(w) for cell, w in zip(row, widths, strict=True))
|
||||||
|
return " | ".join(cells).rstrip(" |")
|
||||||
|
|
||||||
|
|
||||||
|
def _cells(line: str) -> list[str]:
|
||||||
|
inner = line.strip()
|
||||||
|
inner = inner.removeprefix("|").removesuffix("|")
|
||||||
|
return [cell.strip().replace("\\|", "|") for cell in _CELL_SPLIT.split(inner)]
|
||||||
|
|
||||||
|
|
||||||
|
def _line(line: str) -> str:
|
||||||
|
if _RULE.match(line):
|
||||||
|
return "———"
|
||||||
|
if match := _HEADING.match(line):
|
||||||
|
return f"<b>{_inline(match.group(1))}</b>"
|
||||||
|
if match := _TASK.match(line):
|
||||||
|
box = "☐" if match.group(2) == " " else "☑"
|
||||||
|
return f"{match.group(1)}{box} {_inline(match.group(3))}"
|
||||||
|
if match := _BULLET.match(line):
|
||||||
|
return f"{match.group(1)}• {_inline(match.group(2))}"
|
||||||
|
if match := _NUMBERED.match(line):
|
||||||
|
return f"{match.group(1)}{match.group(2)}. {_inline(match.group(3))}"
|
||||||
|
return _inline(line)
|
||||||
|
|
||||||
|
|
||||||
|
def _inline(text: str) -> str:
|
||||||
|
out: list[str] = []
|
||||||
|
pos = 0
|
||||||
|
for match in _TOKEN.finditer(text):
|
||||||
|
out.append(_format(text[pos : match.start()]))
|
||||||
|
out.append(_token(match.group()))
|
||||||
|
pos = match.end()
|
||||||
|
out.append(_format(text[pos:]))
|
||||||
|
return "".join(out)
|
||||||
|
|
||||||
|
|
||||||
|
def _token(token: str) -> str:
|
||||||
|
if token.startswith("`"):
|
||||||
|
return f"<code>{html.escape(token[1:-1])}</code>"
|
||||||
|
if token.startswith("[["):
|
||||||
|
return _wikilink(token)
|
||||||
|
match = _LINK.fullmatch(token)
|
||||||
|
if match is None:
|
||||||
|
return html.escape(token, quote=False)
|
||||||
|
href = html.escape(match.group(3), quote=True)
|
||||||
|
label = _format(match.group(2)) or href
|
||||||
|
return f'<a href="{href}">{label}</a>'
|
||||||
|
|
||||||
|
|
||||||
|
def _wikilink(token: str) -> str:
|
||||||
|
match = _WIKILINK.fullmatch(token)
|
||||||
|
if match is None:
|
||||||
|
return html.escape(token, quote=False)
|
||||||
|
note, heading, alias = (g.strip() if g else "" for g in match.groups())
|
||||||
|
label = alias or (f"{note} › {heading}" if heading else note) # noqa: RUF001
|
||||||
|
return f"<u>{html.escape(label, quote=False)}</u>"
|
||||||
|
|
||||||
|
|
||||||
|
def _format(text: str) -> str:
|
||||||
|
s = html.escape(text, quote=False)
|
||||||
|
s = _BOLD.sub(lambda m: f"<b>{m.group(1) or m.group(2)}</b>", s)
|
||||||
|
s = _ITALIC.sub(r"<i>\1</i>", s)
|
||||||
|
s = _ITALIC_U.sub(r"<i>\1</i>", s)
|
||||||
|
s = _STRIKE.sub(r"<s>\1</s>", s)
|
||||||
|
return _SPOILER.sub(r"<tg-spoiler>\1</tg-spoiler>", s)
|
||||||
|
|
||||||
|
|
||||||
|
def _fits(text: str, limit: int) -> bool:
|
||||||
|
return len(to_html(text)) <= limit
|
||||||
|
|
||||||
|
|
||||||
|
def _flush(out: list[str], cur: list[str]) -> None:
|
||||||
|
chunk = "\n".join(cur).strip()
|
||||||
|
if chunk:
|
||||||
|
out.append(chunk)
|
||||||
|
cur.clear()
|
||||||
|
|
||||||
|
|
||||||
|
def _pieces(block: _Block, limit: int) -> list[str]:
|
||||||
|
whole = "\n".join(block.lines)
|
||||||
|
if _fits(whole, limit):
|
||||||
|
return [whole]
|
||||||
|
if block.kind == "fence":
|
||||||
|
head, body = _fence_parts(block.lines)
|
||||||
|
return _pack(body, lambda ls: "\n".join([head, *ls, "```"]), limit)
|
||||||
|
if block.kind == "quote":
|
||||||
|
inner = [_QUOTE.sub("", line, 1) for line in block.lines]
|
||||||
|
return _pack(inner, lambda ls: "\n".join(f"> {s}" for s in ls), limit)
|
||||||
|
if block.kind == "table":
|
||||||
|
header = block.lines[:2]
|
||||||
|
return _pack(block.lines[2:], lambda ls: "\n".join(header + ls), limit)
|
||||||
|
return _pack(block.lines, "\n".join, limit)
|
||||||
|
|
||||||
|
|
||||||
|
def _pack(lines: list[str], wrap: Callable[[list[str]], str], limit: int) -> list[str]:
|
||||||
|
pieces: list[str] = []
|
||||||
|
cur: list[str] = []
|
||||||
|
for line in lines:
|
||||||
|
if cur and not _fits(wrap([*cur, line]), limit):
|
||||||
|
pieces.append(wrap(cur))
|
||||||
|
cur = []
|
||||||
|
if _fits(wrap([line]), limit):
|
||||||
|
cur.append(line)
|
||||||
|
else:
|
||||||
|
pieces.extend(wrap([part]) for part in _shatter(line, wrap, limit))
|
||||||
|
if cur:
|
||||||
|
pieces.append(wrap(cur))
|
||||||
|
return pieces
|
||||||
|
|
||||||
|
|
||||||
|
def _shatter(line: str, wrap: Callable[[list[str]], str], limit: int) -> list[str]:
|
||||||
|
parts: list[str] = []
|
||||||
|
while line:
|
||||||
|
cut = len(line)
|
||||||
|
while cut > 1 and not _fits(wrap([line[:cut].rstrip()]), limit):
|
||||||
|
rendered = len(to_html(wrap([line[:cut]])))
|
||||||
|
cut = min(cut - 1, _cut_point(line, max(1, cut * limit // rendered)))
|
||||||
|
if part := line[:cut].rstrip():
|
||||||
|
parts.append(part)
|
||||||
|
line = line[cut:].lstrip()
|
||||||
|
return parts
|
||||||
|
|
||||||
|
|
||||||
|
def _cut_point(text: str, limit: int) -> int:
|
||||||
|
for sep in ("\n\n", "\n", ". ", " "):
|
||||||
|
cut = text.rfind(sep, limit // 2, limit)
|
||||||
|
if cut > 0:
|
||||||
|
return cut + len(sep)
|
||||||
|
return limit
|
||||||
@@ -0,0 +1,107 @@
|
|||||||
|
"""Every string the Telegram frontend shows the user; English defaults, overridable."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
|
__all__ = ["TOOL_LABELS", "TelegramTexts"]
|
||||||
|
|
||||||
|
TOOL_LABELS: dict[str, str] = {
|
||||||
|
"Read": "reading",
|
||||||
|
"Glob": "looking for files",
|
||||||
|
"Grep": "searching",
|
||||||
|
"Edit": "editing a file",
|
||||||
|
"Write": "writing a file",
|
||||||
|
"MultiEdit": "editing files",
|
||||||
|
"Bash": "running a command",
|
||||||
|
"WebSearch": "searching the web",
|
||||||
|
"WebFetch": "reading a page",
|
||||||
|
"Task": "started a subagent",
|
||||||
|
"Agent": "started a subagent",
|
||||||
|
"AskUserQuestion": "asking",
|
||||||
|
"TodoWrite": "planning",
|
||||||
|
"Skill": "opening a skill",
|
||||||
|
"mcp__gateway__spawn": "opening a conversation",
|
||||||
|
"mcp__gateway__read_conversation": "reading a conversation",
|
||||||
|
"mcp__gateway__say": "speaking",
|
||||||
|
"mcp__gateway__schedule": "setting a reminder",
|
||||||
|
"mcp__gateway__inject": "passing to another conversation",
|
||||||
|
"mcp__gateway__send_file": "sending a file",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class TelegramTexts:
|
||||||
|
commands: tuple[tuple[str, str], ...] = (
|
||||||
|
("merge", "merge this branch into the master"),
|
||||||
|
("new", "new branch: in General a new topic, in a topic a fresh one here"),
|
||||||
|
("chat", "open a deep chat in the vault: /chat topic"),
|
||||||
|
("status", "this conversation"),
|
||||||
|
("help", "commands"),
|
||||||
|
)
|
||||||
|
help_head: str = "General is the master, any other topic is a branch."
|
||||||
|
branch: str = "branch"
|
||||||
|
closed_prefix: str = "✅ "
|
||||||
|
idle_prefix: str = "💤 "
|
||||||
|
waiting: str = "⏳ thinking"
|
||||||
|
thinking: str = "🤔 thinking"
|
||||||
|
writing: str = "✍️ writing"
|
||||||
|
tool: str = "⏳ {label}"
|
||||||
|
tool_labels: dict[str, str] = field(default_factory=lambda: dict(TOOL_LABELS))
|
||||||
|
attachment_sticker: str = "[attachment: sticker - not supported]"
|
||||||
|
attachment_animation: str = "[attachment: animation - not supported]"
|
||||||
|
attachment_kinds: dict[str, str] = field(
|
||||||
|
default_factory=lambda: {
|
||||||
|
"photo": "photo",
|
||||||
|
"document": "file",
|
||||||
|
"voice": "voice message",
|
||||||
|
"audio": "audio",
|
||||||
|
"video": "video",
|
||||||
|
"video_note": "video note",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
attachment_too_big: str = (
|
||||||
|
"[attachment: {kind} {name}, {mb} MB - over 20 MB, Telegram does not "
|
||||||
|
"serve it to bots]"
|
||||||
|
)
|
||||||
|
attachment_failed: str = "[attachment: {kind} {name} - download failed: {error}]"
|
||||||
|
attachment_size: str = ", {kb} KB"
|
||||||
|
attachment: str = "[attachment: {kind} {path}{size}]"
|
||||||
|
attachment_kept: str = (
|
||||||
|
"[attachment: {kind} {path}{size}; kept {days} days, move it into the "
|
||||||
|
"vault if it matters]"
|
||||||
|
)
|
||||||
|
merge_nothing: str = "nothing to merge: this is not an open branch"
|
||||||
|
merging: str = "🔀 merging into the master"
|
||||||
|
merged: str = "✅ merged into the master"
|
||||||
|
merge_failed: str = "⚠️ merge failed: {error}"
|
||||||
|
new_branch_topic: str = "🌿 branch «{title}» - in a new topic"
|
||||||
|
new_branch_here: str = "🌿 new branch on this topic"
|
||||||
|
chat_usage: str = "/chat <topic>"
|
||||||
|
chat_failed: str = "could not open: {error}"
|
||||||
|
chat_opened: str = "💬 deep chat: {where}"
|
||||||
|
status_unbound: str = "this topic is bound to nothing - write, and a branch opens"
|
||||||
|
status_line: str = "session: {live} · turn: {turn} · queued: {queued}"
|
||||||
|
status_live: str = "live"
|
||||||
|
status_dead: str = "none"
|
||||||
|
status_running: str = "running"
|
||||||
|
status_idle: str = "idle"
|
||||||
|
status_question: str = "❓ waiting for an answer"
|
||||||
|
status_pool: str = "pool: {sessions} sessions, rss {mb} MB"
|
||||||
|
status_outbox: str = "outbox: {pending} queued"
|
||||||
|
turn_failed: str = "⚠️ the turn failed, see the gateway logs"
|
||||||
|
background_failed: str = (
|
||||||
|
"⚠️ a background turn ({origin}) failed, see the gateway logs"
|
||||||
|
)
|
||||||
|
question_closed: str = "the question is already closed"
|
||||||
|
question_timeout: str = "⌛ time is up - answer as text"
|
||||||
|
question_answered: str = "✅ {answer}"
|
||||||
|
question_typed: str = "✍️ {text}"
|
||||||
|
done_button: str = "✅ done"
|
||||||
|
mirrored: str = "📝 from the panel:\n{text}"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def help(self) -> str:
|
||||||
|
return (
|
||||||
|
self.help_head + "\n" + "\n".join(f"/{c} - {d}" for c, d in self.commands)
|
||||||
|
)
|
||||||
@@ -0,0 +1,96 @@
|
|||||||
|
"""Cross-frontend turn record.
|
||||||
|
|
||||||
|
Frontends emit a :class:`TurnRecord` to every handler in
|
||||||
|
``GatewayRuntime.turn_log_handlers`` after a turn finishes.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from anthropic.types import Message, MessageParam
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["TurnRecord", "slugify"]
|
||||||
|
|
||||||
|
|
||||||
|
_SLUG_BAD_RE = re.compile(r"[^\w\s\-]+", flags=re.UNICODE)
|
||||||
|
_SLUG_SEP_RE = re.compile(r"[\s\-]+", flags=re.UNICODE)
|
||||||
|
|
||||||
|
|
||||||
|
def slugify(text: str, *, maxlen: int = 40) -> str:
|
||||||
|
"""Sanitize ``text`` for use as a filename fragment.
|
||||||
|
|
||||||
|
Truncates to ``maxlen`` and returns ``"untitled"`` for empty input.
|
||||||
|
"""
|
||||||
|
cleaned = _SLUG_BAD_RE.sub(" ", text).strip()
|
||||||
|
cleaned = _SLUG_SEP_RE.sub("-", cleaned).strip("-")
|
||||||
|
if not cleaned:
|
||||||
|
return "untitled"
|
||||||
|
if len(cleaned) > maxlen:
|
||||||
|
cleaned = cleaned[:maxlen].rstrip("-")
|
||||||
|
return cleaned or "untitled"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class TurnRecord:
|
||||||
|
"""One completed turn, as seen by a frontend.
|
||||||
|
|
||||||
|
``input_messages`` excludes the assistant reply; ``source`` names the
|
||||||
|
producing frontend so a cross-frontend logger can skip its own turns.
|
||||||
|
"""
|
||||||
|
|
||||||
|
agent_name: str
|
||||||
|
input_messages: list[MessageParam]
|
||||||
|
output_message: Message
|
||||||
|
thinking_is_progress: bool
|
||||||
|
system: str | None = None
|
||||||
|
source: str = "anthropic"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def first_user_text(self) -> str:
|
||||||
|
"""Empty string if the input history has no user turn."""
|
||||||
|
for msg in self.input_messages:
|
||||||
|
if msg.get("role") == "user":
|
||||||
|
return _text_of(msg.get("content", ""))
|
||||||
|
return ""
|
||||||
|
|
||||||
|
@property
|
||||||
|
def last_user_text(self) -> str:
|
||||||
|
for msg in reversed(self.input_messages):
|
||||||
|
if msg.get("role") == "user":
|
||||||
|
return _text_of(msg.get("content", ""))
|
||||||
|
return ""
|
||||||
|
|
||||||
|
@property
|
||||||
|
def assistant_text(self) -> str:
|
||||||
|
"""Plain text of the assistant reply, with thinking/tool blocks dropped."""
|
||||||
|
chunks = [
|
||||||
|
getattr(block, "text", "") or ""
|
||||||
|
for block in getattr(self.output_message, "content", ())
|
||||||
|
if getattr(block, "type", None) == "text"
|
||||||
|
]
|
||||||
|
return "\n\n".join(c for c in chunks if c)
|
||||||
|
|
||||||
|
|
||||||
|
def _text_of(content: object) -> str:
|
||||||
|
"""Flatten an Anthropic ``MessageParam.content`` field to a plain string.
|
||||||
|
|
||||||
|
Handles both raw-string and block-list shapes; unrecognized content is
|
||||||
|
silently skipped since this exists for naming, not faithful reconstruction.
|
||||||
|
"""
|
||||||
|
if isinstance(content, str):
|
||||||
|
return content
|
||||||
|
if isinstance(content, list):
|
||||||
|
parts: list[str] = []
|
||||||
|
for blk in content:
|
||||||
|
if not isinstance(blk, dict):
|
||||||
|
continue
|
||||||
|
d: Any = blk
|
||||||
|
if d.get("type") == "text":
|
||||||
|
parts.append(str(d.get("text", "")))
|
||||||
|
return "\n\n".join(p for p in parts if p)
|
||||||
|
return ""
|
||||||
File renamed without changes.
@@ -0,0 +1 @@
|
|||||||
|
"""Cron, webhook and event jobs, deferred injects, the subscription budget."""
|
||||||
@@ -0,0 +1,158 @@
|
|||||||
|
"""What a job is: its triggers, the budget it respects, and what a run may do."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Awaitable, Callable, Coroutine
|
||||||
|
|
||||||
|
from beaver_gateway.conversations.distill import LineCap
|
||||||
|
from beaver_gateway.conversations.injects import Priority
|
||||||
|
from beaver_gateway.conversations.service import Conversations, DistillResult
|
||||||
|
from beaver_gateway.jobs.scheduler import Scheduler
|
||||||
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
|
__all__ = ["Budget", "Job", "JobRun"]
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class Job:
|
||||||
|
name: str
|
||||||
|
run: Callable[[JobRun], Awaitable[None]]
|
||||||
|
cron: str | None = None
|
||||||
|
webhook: bool = False
|
||||||
|
events: tuple[str, ...] = ()
|
||||||
|
critical: bool = True
|
||||||
|
dedupe: bool = True
|
||||||
|
|
||||||
|
@property
|
||||||
|
def entrypoint(self) -> str:
|
||||||
|
return f"job:{self.name}"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class Budget:
|
||||||
|
threshold: float = 0.7
|
||||||
|
tokens: int | None = None
|
||||||
|
window: timedelta = timedelta(hours=5)
|
||||||
|
limit_window: str = "five_hour"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(slots=True)
|
||||||
|
class JobRun:
|
||||||
|
job: Job
|
||||||
|
trigger: str
|
||||||
|
payload: dict[str, Any]
|
||||||
|
scheduler: Scheduler
|
||||||
|
|
||||||
|
@property
|
||||||
|
def conversations(self) -> Conversations:
|
||||||
|
return self.scheduler.conversations
|
||||||
|
|
||||||
|
async def master(self) -> Conversation | None:
|
||||||
|
masters = await self.conversations.find(kind="master", status="open", limit=1)
|
||||||
|
return masters[0] if masters else None
|
||||||
|
|
||||||
|
async def inject_master(
|
||||||
|
self, text: str, *, urgency: Priority = "normal", origin: str | None = None
|
||||||
|
) -> bool:
|
||||||
|
master = await self.master()
|
||||||
|
if master is None:
|
||||||
|
_log.error(
|
||||||
|
"job %s: no open master, inject lost: %s", self.job.name, text[:200]
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
await self.conversations.inject(
|
||||||
|
master, text, urgency=urgency, origin=origin or self.job.name
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def inject_into(
|
||||||
|
self,
|
||||||
|
public_id: str | None,
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
urgency: Priority = "normal",
|
||||||
|
origin: str | None = None,
|
||||||
|
) -> bool:
|
||||||
|
"""Inject into that conversation, or the master once it is gone or over."""
|
||||||
|
conv = await self.conversations.get(public_id) if public_id else None
|
||||||
|
if conv is None or conv.status != "open":
|
||||||
|
return await self.inject_master(text, urgency=urgency, origin=origin)
|
||||||
|
await self.conversations.inject(
|
||||||
|
conv, text, urgency=urgency, origin=origin or self.job.name
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def spawn_job(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
agent: str,
|
||||||
|
text: str,
|
||||||
|
title: str | None = None,
|
||||||
|
line_cap: LineCap | None = None,
|
||||||
|
) -> Conversation:
|
||||||
|
"""A headless job turn; ``line_cap`` bounces a rewrite past the cap."""
|
||||||
|
return await self.conversations.spawn(
|
||||||
|
kind="job",
|
||||||
|
agent=agent,
|
||||||
|
seed="brief",
|
||||||
|
text=text,
|
||||||
|
title=title,
|
||||||
|
origin="job",
|
||||||
|
flags={"line_cap": line_cap.as_flags()} if line_cap else None,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def close_idle(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
kind: str = "deep",
|
||||||
|
days: int = 2,
|
||||||
|
limit: int = 3,
|
||||||
|
since: datetime | None = None,
|
||||||
|
) -> list[DistillResult]:
|
||||||
|
"""Close chats quiet for ``days``, at most ``limit`` per run."""
|
||||||
|
out: list[DistillResult] = []
|
||||||
|
for conv in await self.conversations.idle(
|
||||||
|
kind=kind, days=days, since=since, limit=limit
|
||||||
|
):
|
||||||
|
try:
|
||||||
|
out.append(
|
||||||
|
await self.conversations.distill(conv, reason=f"idle {days}d")
|
||||||
|
)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("closing idle %s failed", conv.external_id)
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def close_quiet(
|
||||||
|
self, *, kind: str = "branch", days: int = 3, limit: int | None = None
|
||||||
|
) -> list[Conversation]:
|
||||||
|
"""Close chats quiet for ``days`` without a digest, their windows marked."""
|
||||||
|
out: list[Conversation] = []
|
||||||
|
for conv in await self.conversations.idle(kind=kind, days=days, limit=limit):
|
||||||
|
if await self.conversations.busy(
|
||||||
|
conv
|
||||||
|
) or await self.conversations.queue.pending(cast("int", conv.id)):
|
||||||
|
continue
|
||||||
|
row = await self.conversations.close(conv)
|
||||||
|
await self.conversations.mark_closed(row)
|
||||||
|
out.append(row)
|
||||||
|
return out
|
||||||
|
|
||||||
|
async def retry_in(self, delay: timedelta) -> None:
|
||||||
|
await self.scheduler.trigger(
|
||||||
|
self.job, self.payload, delay=delay, trigger=self.trigger
|
||||||
|
)
|
||||||
|
|
||||||
|
async def rotate(self, *, force: str | None = None) -> list[Conversation]:
|
||||||
|
rotation = self.scheduler.rotation
|
||||||
|
return await rotation.tick(force=force) if rotation is not None else []
|
||||||
|
|
||||||
|
def background(self, coro: Coroutine[Any, Any, Any]) -> None:
|
||||||
|
self.scheduler.background(coro)
|
||||||
@@ -1,32 +1,35 @@
|
|||||||
"""Jobs and deferred injects on pgqueuer (§3.6, §4.5).
|
"""Jobs and deferred injects on pgqueuer.
|
||||||
|
|
||||||
A job is a name, a handler and its triggers: a cron expression, the
|
A job is a name, a handler and its triggers (cron, webhook, bus events); a
|
||||||
webhook ``/hooks/<name>``, gateway bus events. The executor is pgqueuer on
|
handler only queues work, and the turn itself runs in the conversation's worker.
|
||||||
the gateway's own Postgres (a dedicated autocommit connection for
|
|
||||||
LISTEN/NOTIFY), so cron ticks, webhook deliveries and the ``schedule``
|
|
||||||
tool's one-off injects all live in one table and survive a restart. A
|
|
||||||
handler only queues work for a conversation and returns; the turn itself
|
|
||||||
runs in the conversation's worker. Non-critical jobs step aside while the
|
|
||||||
subscription window is past its threshold.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import contextlib
|
import contextlib
|
||||||
|
import functools
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
import traceback
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from datetime import UTC, datetime, timedelta
|
from datetime import UTC, datetime, timedelta, tzinfo
|
||||||
from typing import TYPE_CHECKING, Any, cast
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
from zoneinfo import ZoneInfo
|
||||||
|
|
||||||
|
from croniter import croniter
|
||||||
from pgqueuer import PgQueuer, Queries
|
from pgqueuer import PgQueuer, Queries
|
||||||
|
from pgqueuer.core.executors import ScheduleExecutor
|
||||||
from pgqueuer.models import JobId
|
from pgqueuer.models import JobId
|
||||||
|
from sqlalchemy import func
|
||||||
|
from sqlmodel import col, select
|
||||||
from starlette.applications import Starlette
|
from starlette.applications import Starlette
|
||||||
from starlette.responses import JSONResponse
|
from starlette.responses import JSONResponse
|
||||||
from starlette.routing import Route
|
from starlette.routing import Route
|
||||||
|
|
||||||
from beaver_gateway.core.conversations import parse_at
|
from beaver_gateway.conversations.service import parse_at
|
||||||
|
from beaver_gateway.jobs.job import Budget, Job, JobRun
|
||||||
|
from beaver_gateway.storage.models import JobRunRecord
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Awaitable, Callable, Coroutine, Sequence
|
from collections.abc import Awaitable, Callable, Coroutine, Sequence
|
||||||
@@ -36,123 +39,32 @@ if TYPE_CHECKING:
|
|||||||
from pgqueuer.ports.driver import Driver
|
from pgqueuer.ports.driver import Driver
|
||||||
from starlette.requests import Request
|
from starlette.requests import Request
|
||||||
|
|
||||||
from beaver_gateway.core.conversations import Conversations, DistillResult
|
from beaver_gateway.conversations.injects import Priority
|
||||||
from beaver_gateway.core.distill import LineCap
|
from beaver_gateway.conversations.rotation import Rotation
|
||||||
from beaver_gateway.core.injects import Priority
|
from beaver_gateway.conversations.service import Conversations
|
||||||
from beaver_gateway.core.rotation import Rotation
|
|
||||||
from beaver_gateway.storage.models import Conversation
|
from beaver_gateway.storage.models import Conversation
|
||||||
|
|
||||||
__all__ = ["INJECT", "Budget", "Job", "JobRun", "Scheduler"]
|
__all__ = ["INJECT", "Budget", "Job", "JobRun", "LocalCron", "Scheduler", "next_run"]
|
||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.core.scheduler")
|
_log = logging.getLogger("beaver_gateway.jobs.scheduler")
|
||||||
|
|
||||||
INJECT = "inject"
|
INJECT = "inject"
|
||||||
RETRY = timedelta(minutes=15)
|
RETRY = timedelta(minutes=15)
|
||||||
|
ERROR_MAX = 2000
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
def next_run(expression: str, tz: tzinfo, now: datetime | None = None) -> datetime:
|
||||||
class Job:
|
"""The cron's next fire time, read in ``tz``, returned in UTC."""
|
||||||
name: str
|
start = (now or datetime.now(UTC)).astimezone(tz)
|
||||||
run: Callable[[JobRun], Awaitable[None]]
|
return croniter(expression, start_time=start).get_next(datetime).astimezone(UTC)
|
||||||
cron: str | None = None
|
|
||||||
webhook: bool = False
|
|
||||||
events: tuple[str, ...] = ()
|
|
||||||
critical: bool = True
|
|
||||||
dedupe: bool = True
|
|
||||||
|
|
||||||
@property
|
|
||||||
def entrypoint(self) -> str:
|
|
||||||
return f"job:{self.name}"
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass
|
||||||
class Budget:
|
class LocalCron(ScheduleExecutor):
|
||||||
threshold: float = 0.7
|
tz: tzinfo = UTC
|
||||||
tokens: int | None = None
|
|
||||||
window: timedelta = timedelta(hours=5)
|
|
||||||
limit_window: str = "five_hour"
|
|
||||||
|
|
||||||
|
def get_next(self) -> datetime:
|
||||||
@dataclass(slots=True)
|
return next_run(self.parameters.expression, self.tz)
|
||||||
class JobRun:
|
|
||||||
job: Job
|
|
||||||
trigger: str
|
|
||||||
payload: dict[str, Any]
|
|
||||||
scheduler: Scheduler
|
|
||||||
|
|
||||||
@property
|
|
||||||
def conversations(self) -> Conversations:
|
|
||||||
return self.scheduler.conversations
|
|
||||||
|
|
||||||
async def master(self) -> Conversation | None:
|
|
||||||
masters = await self.conversations.find(kind="master", status="open", limit=1)
|
|
||||||
return masters[0] if masters else None
|
|
||||||
|
|
||||||
async def inject_master(
|
|
||||||
self, text: str, *, urgency: Priority = "normal", origin: str | None = None
|
|
||||||
) -> bool:
|
|
||||||
master = await self.master()
|
|
||||||
if master is None:
|
|
||||||
_log.error(
|
|
||||||
"job %s: no open master, inject lost: %s", self.job.name, text[:200]
|
|
||||||
)
|
|
||||||
return False
|
|
||||||
await self.conversations.inject(
|
|
||||||
master, text, urgency=urgency, origin=origin or self.job.name
|
|
||||||
)
|
|
||||||
return True
|
|
||||||
|
|
||||||
async def spawn_job(
|
|
||||||
self,
|
|
||||||
*,
|
|
||||||
agent: str,
|
|
||||||
text: str,
|
|
||||||
title: str | None = None,
|
|
||||||
line_cap: LineCap | None = None,
|
|
||||||
) -> Conversation:
|
|
||||||
"""A headless job turn; ``line_cap`` bounces a rewrite past the cap."""
|
|
||||||
return await self.conversations.spawn(
|
|
||||||
kind="job",
|
|
||||||
agent=agent,
|
|
||||||
seed="brief",
|
|
||||||
text=text,
|
|
||||||
title=title,
|
|
||||||
origin="job",
|
|
||||||
flags={"line_cap": line_cap.as_flags()} if line_cap else None,
|
|
||||||
)
|
|
||||||
|
|
||||||
async def close_idle(
|
|
||||||
self,
|
|
||||||
*,
|
|
||||||
kind: str = "deep",
|
|
||||||
days: int = 2,
|
|
||||||
limit: int = 3,
|
|
||||||
since: datetime | None = None,
|
|
||||||
) -> list[DistillResult]:
|
|
||||||
"""§4.5: close chats quiet for ``days``, at most ``limit`` per run."""
|
|
||||||
out: list[DistillResult] = []
|
|
||||||
for conv in await self.conversations.idle(
|
|
||||||
kind=kind, days=days, since=since, limit=limit
|
|
||||||
):
|
|
||||||
try:
|
|
||||||
out.append(
|
|
||||||
await self.conversations.distill(conv, reason=f"idle {days}d")
|
|
||||||
)
|
|
||||||
except Exception: # noqa: BLE001
|
|
||||||
_log.exception("closing idle %s failed", conv.external_id)
|
|
||||||
return out
|
|
||||||
|
|
||||||
async def retry_in(self, delay: timedelta) -> None:
|
|
||||||
await self.scheduler.trigger(
|
|
||||||
self.job, self.payload, delay=delay, trigger=self.trigger
|
|
||||||
)
|
|
||||||
|
|
||||||
async def rotate(self) -> list[Conversation]:
|
|
||||||
rotation = self.scheduler.rotation
|
|
||||||
return await rotation.tick() if rotation is not None else []
|
|
||||||
|
|
||||||
def background(self, coro: Coroutine[Any, Any, Any]) -> None:
|
|
||||||
self.scheduler.background(coro)
|
|
||||||
|
|
||||||
|
|
||||||
class Scheduler:
|
class Scheduler:
|
||||||
@@ -165,9 +77,11 @@ class Scheduler:
|
|||||||
budget: Budget | None = None,
|
budget: Budget | None = None,
|
||||||
rotation: Rotation | None = None,
|
rotation: Rotation | None = None,
|
||||||
heartbeat: timedelta = timedelta(seconds=30),
|
heartbeat: timedelta = timedelta(seconds=30),
|
||||||
|
tz: str = "UTC",
|
||||||
) -> None:
|
) -> None:
|
||||||
self.conversations = conversations
|
self.conversations = conversations
|
||||||
self.rotation = rotation
|
self.rotation = rotation
|
||||||
|
self.tz = ZoneInfo(tz)
|
||||||
self.budget = budget or Budget()
|
self.budget = budget or Budget()
|
||||||
self._jobs = {job.name: job for job in jobs}
|
self._jobs = {job.name: job for job in jobs}
|
||||||
self._driver = driver
|
self._driver = driver
|
||||||
@@ -187,8 +101,6 @@ class Scheduler:
|
|||||||
def job(self, name: str) -> Job | None:
|
def job(self, name: str) -> Job | None:
|
||||||
return self._jobs.get(name)
|
return self._jobs.get(name)
|
||||||
|
|
||||||
# ---- lifecycle -----------------------------------------------------
|
|
||||||
|
|
||||||
async def start(self) -> None:
|
async def start(self) -> None:
|
||||||
if self._driver is not None:
|
if self._driver is not None:
|
||||||
queries = Queries(self._driver)
|
queries = Queries(self._driver)
|
||||||
@@ -229,7 +141,12 @@ class Scheduler:
|
|||||||
|
|
||||||
if job.cron:
|
if job.cron:
|
||||||
|
|
||||||
@pgq.schedule(job.name, job.cron, clean_old=True)
|
@pgq.schedule(
|
||||||
|
job.name,
|
||||||
|
job.cron,
|
||||||
|
clean_old=True,
|
||||||
|
executor_factory=functools.partial(LocalCron, tz=self.tz),
|
||||||
|
)
|
||||||
async def cron(_schedule: PgSchedule) -> None:
|
async def cron(_schedule: PgSchedule) -> None:
|
||||||
await self._dispatch(job, trigger="cron", payload={})
|
await self._dispatch(job, trigger="cron", payload={})
|
||||||
|
|
||||||
@@ -244,8 +161,6 @@ class Scheduler:
|
|||||||
self._dispatch(job, trigger="event", payload=dict(event))
|
self._dispatch(job, trigger="event", payload=dict(event))
|
||||||
)
|
)
|
||||||
|
|
||||||
# ---- dispatch ------------------------------------------------------
|
|
||||||
|
|
||||||
async def _dispatch(
|
async def _dispatch(
|
||||||
self, job: Job, *, trigger: str, payload: dict[str, Any]
|
self, job: Job, *, trigger: str, payload: dict[str, Any]
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -257,19 +172,33 @@ class Scheduler:
|
|||||||
return
|
return
|
||||||
started = datetime.now(UTC)
|
started = datetime.now(UTC)
|
||||||
bus.publish("job.start", job=job.name, trigger=trigger)
|
bus.publish("job.start", job=job.name, trigger=trigger)
|
||||||
status = "done"
|
status, error = "done", None
|
||||||
try:
|
try:
|
||||||
await job.run(JobRun(job, trigger, payload, self))
|
await job.run(JobRun(job, trigger, payload, self))
|
||||||
except Exception: # noqa: BLE001
|
except Exception as exc: # noqa: BLE001
|
||||||
status = "failed"
|
status, error = "failed", _error_text(exc)
|
||||||
_log.exception("job %s (%s) failed", job.name, trigger)
|
_log.exception("job %s (%s) failed", job.name, trigger)
|
||||||
self._runs[job.name] = {
|
record = JobRunRecord(
|
||||||
"trigger": trigger,
|
job=job.name,
|
||||||
"started_at": started.isoformat(timespec="seconds"),
|
trigger=trigger,
|
||||||
"status": status,
|
started_at=started,
|
||||||
}
|
finished_at=datetime.now(UTC),
|
||||||
|
status=status,
|
||||||
|
error=error,
|
||||||
|
payload=payload,
|
||||||
|
)
|
||||||
|
await self._record(record)
|
||||||
|
self._runs[job.name] = self._run_public(record)
|
||||||
bus.publish("job.end", job=job.name, trigger=trigger, status=status)
|
bus.publish("job.end", job=job.name, trigger=trigger, status=status)
|
||||||
|
|
||||||
|
async def _record(self, record: JobRunRecord) -> None:
|
||||||
|
try:
|
||||||
|
async with self.conversations.db.session() as session:
|
||||||
|
session.add(record)
|
||||||
|
await session.commit()
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("job %s: run not recorded", record.job)
|
||||||
|
|
||||||
async def trigger(
|
async def trigger(
|
||||||
self,
|
self,
|
||||||
job: Job,
|
job: Job,
|
||||||
@@ -306,8 +235,6 @@ class Scheduler:
|
|||||||
)
|
)
|
||||||
return int(ids[0]) if ids and ids[0] is not None else None
|
return int(ids[0]) if ids and ids[0] is not None else None
|
||||||
|
|
||||||
# ---- deferred injects ----------------------------------------------
|
|
||||||
|
|
||||||
async def schedule(
|
async def schedule(
|
||||||
self,
|
self,
|
||||||
conv: Conversation,
|
conv: Conversation,
|
||||||
@@ -325,14 +252,14 @@ class Scheduler:
|
|||||||
if self._queries is None:
|
if self._queries is None:
|
||||||
msg = "scheduler needs postgres; `schedule` is unavailable"
|
msg = "scheduler needs postgres; `schedule` is unavailable"
|
||||||
raise RuntimeError(msg)
|
raise RuntimeError(msg)
|
||||||
when = parse_at(at)
|
when = parse_at(at, self.tz).astimezone(self.tz)
|
||||||
delay = max(when - datetime.now(UTC), timedelta(0))
|
delay = max(when - datetime.now(UTC), timedelta(0))
|
||||||
payload = _encode(
|
payload = _encode(
|
||||||
{
|
{
|
||||||
"conversation": conv.external_id,
|
"conversation": conv.external_id,
|
||||||
"text": text,
|
"text": text,
|
||||||
"urgency": urgency,
|
"urgency": urgency,
|
||||||
"at": when.isoformat(timespec="seconds"),
|
"at": when.isoformat(timespec="minutes"),
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
ids = await self._queries.enqueue(
|
ids = await self._queries.enqueue(
|
||||||
@@ -347,7 +274,7 @@ class Scheduler:
|
|||||||
"schedule.created",
|
"schedule.created",
|
||||||
conversation_id=conv.external_id,
|
conversation_id=conv.external_id,
|
||||||
job=job_id,
|
job=job_id,
|
||||||
execute_at=when.isoformat(timespec="seconds"),
|
execute_at=when.isoformat(timespec="minutes"),
|
||||||
)
|
)
|
||||||
return job_id, when
|
return job_id, when
|
||||||
|
|
||||||
@@ -380,8 +307,6 @@ class Scheduler:
|
|||||||
origin="schedule",
|
origin="schedule",
|
||||||
)
|
)
|
||||||
|
|
||||||
# ---- budget --------------------------------------------------------
|
|
||||||
|
|
||||||
async def utilization(self) -> float | None:
|
async def utilization(self) -> float | None:
|
||||||
now = datetime.now(UTC)
|
now = datetime.now(UTC)
|
||||||
values: list[float] = []
|
values: list[float] = []
|
||||||
@@ -402,8 +327,6 @@ class Scheduler:
|
|||||||
utilization = await self.utilization()
|
utilization = await self.utilization()
|
||||||
return utilization is not None and utilization > self.budget.threshold
|
return utilization is not None and utilization > self.budget.threshold
|
||||||
|
|
||||||
# ---- introspection -------------------------------------------------
|
|
||||||
|
|
||||||
async def snapshot(self) -> dict[str, Any]:
|
async def snapshot(self) -> dict[str, Any]:
|
||||||
crons: dict[str, PgSchedule] = {}
|
crons: dict[str, PgSchedule] = {}
|
||||||
queue: list[dict[str, Any]] = []
|
queue: list[dict[str, Any]] = []
|
||||||
@@ -411,16 +334,22 @@ class Scheduler:
|
|||||||
for row in await self._queries.peek_schedule():
|
for row in await self._queries.peek_schedule():
|
||||||
crons[str(row.entrypoint)] = row
|
crons[str(row.entrypoint)] = row
|
||||||
queue = [
|
queue = [
|
||||||
_job_public(row) for row in await self._queries.browse_queue(limit=200)
|
_job_public(row, self.tz)
|
||||||
|
for row in await self._queries.browse_queue(limit=200)
|
||||||
]
|
]
|
||||||
utilization = await self.utilization()
|
utilization = await self.utilization()
|
||||||
|
last = await self._last_runs()
|
||||||
return {
|
return {
|
||||||
"enabled": self.enabled,
|
"enabled": self.enabled,
|
||||||
|
"tz": str(self.tz),
|
||||||
"utilization": utilization,
|
"utilization": utilization,
|
||||||
"throttled": utilization is not None
|
"throttled": utilization is not None
|
||||||
and utilization > self.budget.threshold,
|
and utilization > self.budget.threshold,
|
||||||
"threshold": self.budget.threshold,
|
"threshold": self.budget.threshold,
|
||||||
"jobs": [self._job_public(job, crons.get(job.name)) for job in self.jobs],
|
"jobs": [
|
||||||
|
self._job_public(job, crons.get(job.name), last.get(job.name))
|
||||||
|
for job in self.jobs
|
||||||
|
],
|
||||||
"queue": queue,
|
"queue": queue,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -428,27 +357,61 @@ class Scheduler:
|
|||||||
if self._queries is None:
|
if self._queries is None:
|
||||||
return []
|
return []
|
||||||
rows = await self._queries.browse_queue(limit=500, entrypoints=[INJECT])
|
rows = await self._queries.browse_queue(limit=500, entrypoints=[INJECT])
|
||||||
out = [_job_public(row) for row in rows]
|
out = [_job_public(row, self.tz) for row in rows]
|
||||||
if conv is not None:
|
if conv is not None:
|
||||||
out = [
|
out = [
|
||||||
j for j in out if j["payload"].get("conversation") == conv.external_id
|
j for j in out if j["payload"].get("conversation") == conv.external_id
|
||||||
]
|
]
|
||||||
return out
|
return out
|
||||||
|
|
||||||
def _job_public(self, job: Job, cron: PgSchedule | None) -> dict[str, Any]:
|
async def runs(self, name: str, limit: int = 50) -> list[dict[str, Any]]:
|
||||||
|
stmt = (
|
||||||
|
select(JobRunRecord)
|
||||||
|
.where(JobRunRecord.job == name)
|
||||||
|
.order_by(col(JobRunRecord.id).desc())
|
||||||
|
.limit(limit)
|
||||||
|
)
|
||||||
|
async with self.conversations.db.session() as session:
|
||||||
|
rows = (await session.exec(stmt)).all()
|
||||||
|
return [self._run_public(row) for row in rows]
|
||||||
|
|
||||||
|
async def _last_runs(self) -> dict[str, dict[str, Any]]:
|
||||||
|
newest = select(func.max(JobRunRecord.id)).group_by(JobRunRecord.job)
|
||||||
|
stmt = select(JobRunRecord).where(col(JobRunRecord.id).in_(newest))
|
||||||
|
try:
|
||||||
|
async with self.conversations.db.session() as session:
|
||||||
|
rows = (await session.exec(stmt)).all()
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("job runs unreadable, using memory")
|
||||||
|
return dict(self._runs)
|
||||||
|
return {row.job: self._run_public(row) for row in rows}
|
||||||
|
|
||||||
|
def _job_public(
|
||||||
|
self, job: Job, cron: PgSchedule | None, run: dict[str, Any] | None
|
||||||
|
) -> dict[str, Any]:
|
||||||
return {
|
return {
|
||||||
"name": job.name,
|
"name": job.name,
|
||||||
"cron": job.cron,
|
"cron": job.cron,
|
||||||
"webhook": job.webhook,
|
"webhook": job.webhook,
|
||||||
"events": list(job.events),
|
"events": list(job.events),
|
||||||
"critical": job.critical,
|
"critical": job.critical,
|
||||||
"next_run": _iso(cron.next_run) if cron is not None else None,
|
"next_run": _iso(cron.next_run, self.tz) if cron is not None else None,
|
||||||
"last_run": _iso(cron.last_run) if cron is not None else None,
|
"last_run": _iso(cron.last_run, self.tz) if cron is not None else None,
|
||||||
"status": str(cron.status) if cron is not None else None,
|
"status": str(cron.status) if cron is not None else None,
|
||||||
"run": self._runs.get(job.name),
|
"run": run or self._runs.get(job.name),
|
||||||
}
|
}
|
||||||
|
|
||||||
# ---- http ----------------------------------------------------------
|
def _run_public(self, row: JobRunRecord) -> dict[str, Any]:
|
||||||
|
return {
|
||||||
|
"id": row.id,
|
||||||
|
"job": row.job,
|
||||||
|
"trigger": row.trigger,
|
||||||
|
"started_at": _iso(row.started_at, self.tz),
|
||||||
|
"finished_at": _iso(row.finished_at, self.tz),
|
||||||
|
"status": row.status,
|
||||||
|
"error": row.error,
|
||||||
|
"payload": row.payload,
|
||||||
|
}
|
||||||
|
|
||||||
def app(self, authorize: Callable[[Request], Awaitable[Any]]) -> Starlette:
|
def app(self, authorize: Callable[[Request], Awaitable[Any]]) -> Starlette:
|
||||||
async def hook(request: Request) -> JSONResponse:
|
async def hook(request: Request) -> JSONResponse:
|
||||||
@@ -463,8 +426,6 @@ class Scheduler:
|
|||||||
|
|
||||||
return Starlette(routes=[Route("/{name}", hook, methods=["POST"])])
|
return Starlette(routes=[Route("/{name}", hook, methods=["POST"])])
|
||||||
|
|
||||||
# ---- internals -----------------------------------------------------
|
|
||||||
|
|
||||||
def background(self, coro: Coroutine[Any, Any, Any]) -> None:
|
def background(self, coro: Coroutine[Any, Any, Any]) -> None:
|
||||||
self._spawn(coro)
|
self._spawn(coro)
|
||||||
|
|
||||||
@@ -506,21 +467,28 @@ def _decode(payload: bytes | None) -> dict[str, Any]:
|
|||||||
return data if isinstance(data, dict) else {"raw": data}
|
return data if isinstance(data, dict) else {"raw": data}
|
||||||
|
|
||||||
|
|
||||||
def _job_public(row: PgJob) -> dict[str, Any]:
|
def _job_public(row: PgJob, tz: tzinfo) -> dict[str, Any]:
|
||||||
return {
|
return {
|
||||||
"id": int(row.id),
|
"id": int(row.id),
|
||||||
"entrypoint": row.entrypoint,
|
"entrypoint": row.entrypoint,
|
||||||
"status": str(row.status),
|
"status": str(row.status),
|
||||||
"execute_after": _iso(row.execute_after),
|
"execute_after": _iso(row.execute_after, tz),
|
||||||
"created": _iso(row.created),
|
"created": _iso(row.created, tz),
|
||||||
"attempts": row.attempts,
|
"attempts": row.attempts,
|
||||||
"payload": _decode(row.payload),
|
"payload": _decode(row.payload),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _error_text(exc: BaseException) -> str:
|
||||||
|
head = f"{type(exc).__name__}: {exc}"
|
||||||
|
return f"{head}\n{''.join(traceback.format_exception(exc))}"[:ERROR_MAX]
|
||||||
|
|
||||||
|
|
||||||
def _aware(value: datetime) -> datetime:
|
def _aware(value: datetime) -> datetime:
|
||||||
return value if value.tzinfo is not None else value.replace(tzinfo=UTC)
|
return value if value.tzinfo is not None else value.replace(tzinfo=UTC)
|
||||||
|
|
||||||
|
|
||||||
def _iso(value: datetime | None) -> str | None:
|
def _iso(value: datetime | None, tz: tzinfo = UTC) -> str | None:
|
||||||
return _aware(value).isoformat(timespec="seconds") if value is not None else None
|
if value is None:
|
||||||
|
return None
|
||||||
|
return _aware(value).astimezone(tz).isoformat(timespec="seconds")
|
||||||
@@ -0,0 +1,19 @@
|
|||||||
|
"""Secrets the model runs commands with: sealed in the database, handed out to root."""
|
||||||
|
|
||||||
|
from beaver_gateway.keystore.server import serve
|
||||||
|
from beaver_gateway.keystore.store import (
|
||||||
|
NAME,
|
||||||
|
SecretBox,
|
||||||
|
SecretsDisabledError,
|
||||||
|
SecretStore,
|
||||||
|
UnknownSecretError,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"NAME",
|
||||||
|
"SecretBox",
|
||||||
|
"SecretStore",
|
||||||
|
"SecretsDisabledError",
|
||||||
|
"UnknownSecretError",
|
||||||
|
"serve",
|
||||||
|
]
|
||||||
@@ -0,0 +1,289 @@
|
|||||||
|
"""``secret ls`` and ``secret run NAME... -- command``: stdlib only, run by path.
|
||||||
|
|
||||||
|
The model's uid calls this file; it asks the root half (the same file
|
||||||
|
through ``sudo``) for the values, which reads them from the gateway's
|
||||||
|
root-only socket. The command then runs as the model's uid with each
|
||||||
|
value in its own variable, files in a private tmpfs directory, and the
|
||||||
|
values masked in what it prints.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import contextlib
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import shutil
|
||||||
|
import signal
|
||||||
|
import socket
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import threading
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import IO, TYPE_CHECKING, Any
|
||||||
|
from urllib.parse import quote
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from io import BufferedReader
|
||||||
|
|
||||||
|
SOCKET = Path("/run/beaver-secrets/gateway.sock")
|
||||||
|
ROOT_HALF = ("sudo", "-n", "/usr/local/libexec/beaver-secret")
|
||||||
|
SHM = Path("/dev/shm") # noqa: S108
|
||||||
|
PREFIX = "beaver-secret-"
|
||||||
|
NAME = re.compile(r"[A-Z][A-Z0-9_]{0,63}")
|
||||||
|
MIN_MASKED = 6
|
||||||
|
HOLD = 1 << 16
|
||||||
|
CHUNK = 1 << 16
|
||||||
|
DRAIN = 2.0
|
||||||
|
|
||||||
|
USAGE = """usage:
|
||||||
|
secret ls names and descriptions
|
||||||
|
secret run NAME [NAME...] -- CMD run CMD with each secret in $NAME
|
||||||
|
|
||||||
|
A value secret arrives as its text; a file secret as the path to a
|
||||||
|
private copy that is removed when CMD exits. Values are masked in
|
||||||
|
CMD's output as <secret:NAME>.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class SecretError(Exception):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def request(payload: dict[str, Any], path: Path = SOCKET) -> dict[str, Any]:
|
||||||
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn:
|
||||||
|
try:
|
||||||
|
conn.connect(str(path))
|
||||||
|
except OSError as exc:
|
||||||
|
msg = f"the gateway's secrets socket is not reachable ({exc.strerror})"
|
||||||
|
raise SecretError(msg) from exc
|
||||||
|
conn.sendall(json.dumps(payload).encode() + b"\n")
|
||||||
|
chunks = []
|
||||||
|
while chunk := conn.recv(CHUNK):
|
||||||
|
chunks.append(chunk)
|
||||||
|
reply = json.loads(b"".join(chunks) or b"{}")
|
||||||
|
if not reply.get("ok"):
|
||||||
|
raise SecretError(reply.get("error") or "the gateway gave no answer")
|
||||||
|
return reply
|
||||||
|
|
||||||
|
|
||||||
|
def root_half(argv: list[str]) -> int:
|
||||||
|
if argv[:1] == ["ls"] and len(argv) == 1:
|
||||||
|
payload: dict[str, Any] = {"op": "ls"}
|
||||||
|
elif (
|
||||||
|
argv[:1] == ["get"]
|
||||||
|
and len(argv) > 1
|
||||||
|
and all(NAME.fullmatch(a) for a in argv[1:])
|
||||||
|
):
|
||||||
|
payload = {"op": "get", "names": argv[1:]}
|
||||||
|
else:
|
||||||
|
print("beaver-secret: ls | get NAME...", file=sys.stderr)
|
||||||
|
return 2
|
||||||
|
try:
|
||||||
|
reply = request(payload)
|
||||||
|
except SecretError as exc:
|
||||||
|
reply = {"ok": False, "error": str(exc)}
|
||||||
|
sys.stdout.write(json.dumps(reply))
|
||||||
|
return 0 if reply.get("ok") else 1
|
||||||
|
|
||||||
|
|
||||||
|
def ask(argv: list[str]) -> dict[str, Any]:
|
||||||
|
if os.geteuid() == 0:
|
||||||
|
return request(
|
||||||
|
{"op": "ls"} if argv == ["ls"] else {"op": "get", "names": argv[1:]}
|
||||||
|
)
|
||||||
|
proc = subprocess.run( # noqa: S603
|
||||||
|
[*ROOT_HALF, *argv], capture_output=True, check=False
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
reply = json.loads(proc.stdout)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
err = proc.stderr.decode(errors="replace").strip()
|
||||||
|
msg = f"the root helper failed: {err or f'exit {proc.returncode}'}"
|
||||||
|
raise SecretError(msg) from None
|
||||||
|
if not reply.get("ok"):
|
||||||
|
raise SecretError(reply.get("error") or "the root helper gave no answer")
|
||||||
|
return reply
|
||||||
|
|
||||||
|
|
||||||
|
def list_secrets() -> int:
|
||||||
|
rows = ask(["ls"])["secrets"]
|
||||||
|
if not rows:
|
||||||
|
print("no secrets yet; a human adds them in the admin, System -> Secrets")
|
||||||
|
return 0
|
||||||
|
width = max(len(r["name"]) for r in rows)
|
||||||
|
for r in rows:
|
||||||
|
kind = f"file {r['filename']}" if r.get("filename") else "value"
|
||||||
|
line = f"{r['name']:<{width}} {kind}"
|
||||||
|
if r.get("description"):
|
||||||
|
line += f" - {r['description']}"
|
||||||
|
print(line)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
class Masker:
|
||||||
|
"""Streams bytes through, replacing every needle; never splits one."""
|
||||||
|
|
||||||
|
def __init__(self, needles: dict[bytes, bytes]) -> None:
|
||||||
|
self._needles = needles
|
||||||
|
self._longest = max((len(n) for n in needles), default=0)
|
||||||
|
ordered = sorted(needles, key=lambda n: -len(n))
|
||||||
|
self._pattern = (
|
||||||
|
re.compile(b"|".join(re.escape(n) for n in ordered)) if ordered else None
|
||||||
|
)
|
||||||
|
self._pending = b""
|
||||||
|
|
||||||
|
def feed(self, data: bytes) -> bytes:
|
||||||
|
if self._pattern is None:
|
||||||
|
return data
|
||||||
|
self._pending += data
|
||||||
|
cut = self._pending.rfind(b"\n") + 1
|
||||||
|
if len(self._pending) - cut > HOLD:
|
||||||
|
cut = max(cut, len(self._pending) - self._longest + 1)
|
||||||
|
for match in self._pattern.finditer(self._pending):
|
||||||
|
if match.start() < cut < match.end():
|
||||||
|
cut = match.end()
|
||||||
|
head, self._pending = self._pending[:cut], self._pending[cut:]
|
||||||
|
return self._mask(head)
|
||||||
|
|
||||||
|
def flush(self) -> bytes:
|
||||||
|
tail, self._pending = self._pending, b""
|
||||||
|
return self._mask(tail)
|
||||||
|
|
||||||
|
def _mask(self, data: bytes) -> bytes:
|
||||||
|
if self._pattern is None or not data:
|
||||||
|
return data
|
||||||
|
return self._pattern.sub(lambda m: self._needles[m[0]], data)
|
||||||
|
|
||||||
|
|
||||||
|
def needles_for(values: dict[str, bytes]) -> dict[bytes, bytes]:
|
||||||
|
out: dict[bytes, bytes] = {}
|
||||||
|
for name, data in values.items():
|
||||||
|
mask = f"<secret:{name}>".encode()
|
||||||
|
try:
|
||||||
|
text = data.decode()
|
||||||
|
except UnicodeDecodeError:
|
||||||
|
continue
|
||||||
|
for part in text.splitlines() or [text]:
|
||||||
|
for form in {part.strip(), quote(part.strip(), safe="")}:
|
||||||
|
if len(form) >= MIN_MASKED:
|
||||||
|
out[form.encode()] = mask
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def sweep(root: Path) -> None:
|
||||||
|
for stale in root.glob(f"{PREFIX}*"):
|
||||||
|
pid = stale.name.removeprefix(PREFIX).split("-", 1)[0]
|
||||||
|
if not pid.isdigit() or _alive(int(pid)):
|
||||||
|
continue
|
||||||
|
shutil.rmtree(stale, ignore_errors=True)
|
||||||
|
|
||||||
|
|
||||||
|
def _alive(pid: int) -> bool:
|
||||||
|
try:
|
||||||
|
os.kill(pid, 0)
|
||||||
|
except ProcessLookupError:
|
||||||
|
return False
|
||||||
|
except PermissionError:
|
||||||
|
return True
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def _pump(src: BufferedReader, dst: IO[bytes], masker: Masker) -> None:
|
||||||
|
while chunk := src.read1(CHUNK):
|
||||||
|
dst.write(masker.feed(chunk))
|
||||||
|
dst.flush()
|
||||||
|
dst.write(masker.flush())
|
||||||
|
dst.flush()
|
||||||
|
|
||||||
|
|
||||||
|
def run(names: list[str], command: list[str], *, root: Path | None = None) -> int:
|
||||||
|
secrets = ask(["get", *names])["secrets"]
|
||||||
|
base = root or (SHM if SHM.is_dir() else Path(tempfile.gettempdir()))
|
||||||
|
sweep(base)
|
||||||
|
workdir = Path(tempfile.mkdtemp(prefix=f"{PREFIX}{os.getpid()}-", dir=base))
|
||||||
|
env = dict(os.environ)
|
||||||
|
values: dict[str, bytes] = {}
|
||||||
|
try:
|
||||||
|
for name in names:
|
||||||
|
item = secrets[name]
|
||||||
|
data = base64.b64decode(item["data"])
|
||||||
|
if item.get("filename"):
|
||||||
|
path = workdir / name / item["filename"]
|
||||||
|
path.parent.mkdir(mode=0o700)
|
||||||
|
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
||||||
|
with os.fdopen(fd, "wb") as fh:
|
||||||
|
fh.write(data)
|
||||||
|
env[name] = str(path)
|
||||||
|
else:
|
||||||
|
env[name] = data.decode(errors="surrogateescape")
|
||||||
|
values[name] = data
|
||||||
|
return _spawn(command, env, needles_for(values))
|
||||||
|
finally:
|
||||||
|
shutil.rmtree(workdir, ignore_errors=True)
|
||||||
|
|
||||||
|
|
||||||
|
def _spawn(command: list[str], env: dict[str, str], needles: dict[bytes, bytes]) -> int:
|
||||||
|
try:
|
||||||
|
child = subprocess.Popen( # noqa: S603
|
||||||
|
command, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE
|
||||||
|
)
|
||||||
|
except (OSError, ValueError) as exc:
|
||||||
|
reason = exc.strerror if isinstance(exc, OSError) else str(exc)
|
||||||
|
print(f"secret: cannot run {command[0]}: {reason}", file=sys.stderr)
|
||||||
|
return 127
|
||||||
|
forwarded = (
|
||||||
|
(signal.SIGINT, signal.SIGTERM, signal.SIGHUP)
|
||||||
|
if threading.current_thread() is threading.main_thread()
|
||||||
|
else ()
|
||||||
|
)
|
||||||
|
previous = {
|
||||||
|
sig: signal.signal(sig, lambda s, _f: child.send_signal(s)) for sig in forwarded
|
||||||
|
}
|
||||||
|
pumps = [
|
||||||
|
threading.Thread(target=_pump, args=(src, dst, Masker(needles)), daemon=True)
|
||||||
|
for src, dst in (
|
||||||
|
(child.stdout, sys.stdout.buffer),
|
||||||
|
(child.stderr, sys.stderr.buffer),
|
||||||
|
)
|
||||||
|
]
|
||||||
|
for pump in pumps:
|
||||||
|
pump.start()
|
||||||
|
try:
|
||||||
|
code = child.wait()
|
||||||
|
for pump in pumps:
|
||||||
|
pump.join(DRAIN)
|
||||||
|
finally:
|
||||||
|
for sig, handler in previous.items():
|
||||||
|
signal.signal(sig, handler)
|
||||||
|
return 128 - code if code < 0 else code
|
||||||
|
|
||||||
|
|
||||||
|
def main(argv: list[str] | None = None) -> int:
|
||||||
|
args = sys.argv[1:] if argv is None else argv
|
||||||
|
if args[:1] == ["--root"]:
|
||||||
|
return root_half(args[1:])
|
||||||
|
try:
|
||||||
|
if args == ["ls"]:
|
||||||
|
return list_secrets()
|
||||||
|
if args[:1] == ["run"] and "--" in args:
|
||||||
|
split = args.index("--")
|
||||||
|
names, command = args[1:split], args[split + 1 :]
|
||||||
|
bad = [n for n in names if not NAME.fullmatch(n)]
|
||||||
|
if names and command and not bad:
|
||||||
|
return run(names, command)
|
||||||
|
if bad:
|
||||||
|
print(f"secret: invalid name {bad[0]!r}", file=sys.stderr)
|
||||||
|
return 2
|
||||||
|
except SecretError as exc:
|
||||||
|
print(f"secret: {exc}", file=sys.stderr)
|
||||||
|
return 1
|
||||||
|
print(USAGE, end="", file=sys.stderr)
|
||||||
|
return 2
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
with contextlib.suppress(KeyboardInterrupt):
|
||||||
|
sys.exit(main())
|
||||||
@@ -0,0 +1,108 @@
|
|||||||
|
"""Root-only unix socket that hands sealed secrets to the ``secret`` helper."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import contextlib
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
from beaver_gateway.keystore.store import (
|
||||||
|
SecretsDisabledError,
|
||||||
|
UnknownSecretError,
|
||||||
|
check_name,
|
||||||
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Awaitable, Callable, Sequence
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from beaver_gateway.keystore.store import SecretStore
|
||||||
|
|
||||||
|
_log = logging.getLogger("beaver_gateway.keystore")
|
||||||
|
|
||||||
|
REQUEST_LIMIT = 1 << 16
|
||||||
|
|
||||||
|
|
||||||
|
async def serve(
|
||||||
|
store: SecretStore,
|
||||||
|
path: Path,
|
||||||
|
*,
|
||||||
|
on_use: Callable[[Sequence[str]], Awaitable[None]] | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Listen on ``path`` (mode 0600 in a 0700 directory) until cancelled."""
|
||||||
|
_clear(path)
|
||||||
|
|
||||||
|
async def handle(
|
||||||
|
reader: asyncio.StreamReader, writer: asyncio.StreamWriter
|
||||||
|
) -> None:
|
||||||
|
try:
|
||||||
|
line = await reader.readline()
|
||||||
|
reply = await answer(store, line, on_use=on_use)
|
||||||
|
writer.write(json.dumps(reply).encode() + b"\n")
|
||||||
|
await writer.drain()
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception("secret request failed")
|
||||||
|
finally:
|
||||||
|
writer.close()
|
||||||
|
with contextlib.suppress(Exception):
|
||||||
|
await writer.wait_closed()
|
||||||
|
|
||||||
|
server = await asyncio.start_unix_server(
|
||||||
|
handle, path=str(path), limit=REQUEST_LIMIT
|
||||||
|
)
|
||||||
|
_lock(path)
|
||||||
|
_log.info("secrets socket on %s", path)
|
||||||
|
async with server:
|
||||||
|
await server.serve_forever()
|
||||||
|
|
||||||
|
|
||||||
|
async def answer(
|
||||||
|
store: SecretStore,
|
||||||
|
line: bytes,
|
||||||
|
*,
|
||||||
|
on_use: Callable[[Sequence[str]], Awaitable[None]] | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
try:
|
||||||
|
request = json.loads(line)
|
||||||
|
op = request.get("op") if isinstance(request, dict) else None
|
||||||
|
if op == "ls":
|
||||||
|
return {"ok": True, "secrets": [_listed(r) for r in await store.rows()]}
|
||||||
|
if op == "get":
|
||||||
|
names = [check_name(str(n)) for n in request.get("names") or ()]
|
||||||
|
if not names:
|
||||||
|
return {"ok": False, "error": "no secret names given"}
|
||||||
|
revealed = await store.reveal(names)
|
||||||
|
if on_use is not None:
|
||||||
|
await on_use(names)
|
||||||
|
return {
|
||||||
|
"ok": True,
|
||||||
|
"secrets": {
|
||||||
|
r.name: {
|
||||||
|
"data": base64.b64encode(r.data).decode(),
|
||||||
|
"filename": r.filename,
|
||||||
|
}
|
||||||
|
for r in revealed
|
||||||
|
},
|
||||||
|
}
|
||||||
|
except UnknownSecretError as exc:
|
||||||
|
return {"ok": False, "error": f"unknown secret: {exc}"}
|
||||||
|
except (SecretsDisabledError, ValueError) as exc:
|
||||||
|
return {"ok": False, "error": str(exc)}
|
||||||
|
return {"ok": False, "error": "op must be `ls` or `get`"}
|
||||||
|
|
||||||
|
|
||||||
|
def _clear(path: Path) -> None:
|
||||||
|
path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||||
|
path.parent.chmod(0o700)
|
||||||
|
path.unlink(missing_ok=True)
|
||||||
|
|
||||||
|
|
||||||
|
def _lock(path: Path) -> None:
|
||||||
|
path.chmod(0o600)
|
||||||
|
|
||||||
|
|
||||||
|
def _listed(row: Any) -> dict[str, Any]:
|
||||||
|
return {"name": row.name, "description": row.description, "filename": row.filename}
|
||||||
@@ -0,0 +1,193 @@
|
|||||||
|
"""Seal, list, replace and reveal secrets; the key never leaves this module."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from cryptography.exceptions import InvalidTag
|
||||||
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||||
|
from sqlmodel import col, select
|
||||||
|
|
||||||
|
from beaver_gateway.storage.models import Secret
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Iterable, Sequence
|
||||||
|
|
||||||
|
from beaver_gateway.storage import Database
|
||||||
|
|
||||||
|
NAME = re.compile(r"[A-Z][A-Z0-9_]{0,63}")
|
||||||
|
"""A secret's name is the environment variable a command receives it under."""
|
||||||
|
|
||||||
|
KEY_BYTES = 32
|
||||||
|
NONCE_BYTES = 12
|
||||||
|
MAX_SIZE = 1 << 20
|
||||||
|
"""Largest value or file the store accepts, in bytes."""
|
||||||
|
|
||||||
|
|
||||||
|
class SecretsDisabledError(RuntimeError):
|
||||||
|
"""``SECRETS_KEY`` is unset or unusable, so nothing can be sealed or opened."""
|
||||||
|
|
||||||
|
|
||||||
|
class UnknownSecretError(LookupError):
|
||||||
|
"""A name that is not in the store."""
|
||||||
|
|
||||||
|
|
||||||
|
class SecretBox:
|
||||||
|
"""AES-256-GCM with the secret's name as associated data."""
|
||||||
|
|
||||||
|
__slots__ = ("_aead",)
|
||||||
|
|
||||||
|
def __init__(self, key: bytes) -> None:
|
||||||
|
if len(key) != KEY_BYTES:
|
||||||
|
msg = f"SECRETS_KEY must be {KEY_BYTES} bytes, got {len(key)}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
self._aead = AESGCM(key)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_hex(cls, text: str) -> SecretBox:
|
||||||
|
try:
|
||||||
|
key = bytes.fromhex(text.strip())
|
||||||
|
except ValueError as exc:
|
||||||
|
msg = "SECRETS_KEY must be hex, e.g. `openssl rand -hex 32`"
|
||||||
|
raise ValueError(msg) from exc
|
||||||
|
return cls(key)
|
||||||
|
|
||||||
|
def seal(self, name: str, data: bytes) -> bytes:
|
||||||
|
nonce = os.urandom(NONCE_BYTES)
|
||||||
|
return nonce + self._aead.encrypt(nonce, data, name.encode())
|
||||||
|
|
||||||
|
def open(self, name: str, sealed: bytes) -> bytes:
|
||||||
|
nonce, body = sealed[:NONCE_BYTES], sealed[NONCE_BYTES:]
|
||||||
|
return self._aead.decrypt(nonce, body, name.encode())
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class Revealed:
|
||||||
|
name: str
|
||||||
|
data: bytes
|
||||||
|
filename: str | None
|
||||||
|
|
||||||
|
|
||||||
|
def _now() -> datetime:
|
||||||
|
return datetime.now(UTC)
|
||||||
|
|
||||||
|
|
||||||
|
def check_name(name: str) -> str:
|
||||||
|
if not NAME.fullmatch(name):
|
||||||
|
msg = (
|
||||||
|
f"invalid secret name {name!r}: use A-Z, 0-9 and _, starting with a letter"
|
||||||
|
)
|
||||||
|
raise ValueError(msg)
|
||||||
|
return name
|
||||||
|
|
||||||
|
|
||||||
|
def clean_filename(filename: str) -> str:
|
||||||
|
base = os.path.basename(filename.replace("\\", "/")).strip() # noqa: PTH119
|
||||||
|
if base in ("", ".", ".."):
|
||||||
|
msg = f"invalid file name {filename!r}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
return base
|
||||||
|
|
||||||
|
|
||||||
|
class SecretStore:
|
||||||
|
"""Rows of :class:`Secret`; ``box`` is ``None`` when no key is configured."""
|
||||||
|
|
||||||
|
__slots__ = ("_box", "_db")
|
||||||
|
|
||||||
|
def __init__(self, db: Database, box: SecretBox | None) -> None:
|
||||||
|
self._db = db
|
||||||
|
self._box = box
|
||||||
|
|
||||||
|
@property
|
||||||
|
def enabled(self) -> bool:
|
||||||
|
return self._box is not None
|
||||||
|
|
||||||
|
def _require_box(self) -> SecretBox:
|
||||||
|
if self._box is None:
|
||||||
|
msg = "SECRETS_KEY is not set"
|
||||||
|
raise SecretsDisabledError(msg)
|
||||||
|
return self._box
|
||||||
|
|
||||||
|
async def rows(self) -> Sequence[Secret]:
|
||||||
|
async with self._db.session() as session:
|
||||||
|
result = await session.exec(select(Secret).order_by(col(Secret.name)))
|
||||||
|
return result.all()
|
||||||
|
|
||||||
|
async def put(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
*,
|
||||||
|
description: str | None = None,
|
||||||
|
data: bytes | None = None,
|
||||||
|
filename: str | None = None,
|
||||||
|
) -> Secret:
|
||||||
|
"""Create or replace; ``data=None`` keeps the value, sets the description."""
|
||||||
|
check_name(name)
|
||||||
|
if data is not None:
|
||||||
|
if len(data) > MAX_SIZE:
|
||||||
|
msg = f"secret {name} is {len(data)} bytes, the limit is {MAX_SIZE}"
|
||||||
|
raise ValueError(msg)
|
||||||
|
sealed = self._require_box().seal(name, data)
|
||||||
|
async with self._db.session() as session:
|
||||||
|
row = await session.get(Secret, name)
|
||||||
|
if row is None:
|
||||||
|
if data is None:
|
||||||
|
msg = f"secret {name} needs a value or a file"
|
||||||
|
raise ValueError(msg)
|
||||||
|
row = Secret(name=name, sealed=sealed)
|
||||||
|
if data is not None:
|
||||||
|
row.sealed = sealed
|
||||||
|
row.size = len(data)
|
||||||
|
row.filename = clean_filename(filename) if filename else None
|
||||||
|
if description is not None:
|
||||||
|
row.description = description.strip()
|
||||||
|
row.updated_at = _now()
|
||||||
|
session.add(row)
|
||||||
|
await session.commit()
|
||||||
|
await session.refresh(row)
|
||||||
|
return row
|
||||||
|
|
||||||
|
async def delete(self, name: str) -> bool:
|
||||||
|
async with self._db.session() as session:
|
||||||
|
row = await session.get(Secret, name)
|
||||||
|
if row is None:
|
||||||
|
return False
|
||||||
|
await session.delete(row)
|
||||||
|
await session.commit()
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def reveal(self, names: Iterable[str]) -> list[Revealed]:
|
||||||
|
box = self._require_box()
|
||||||
|
wanted = list(dict.fromkeys(names))
|
||||||
|
async with self._db.session() as session:
|
||||||
|
rows = {
|
||||||
|
r.name: r
|
||||||
|
for r in (
|
||||||
|
await session.exec(
|
||||||
|
select(Secret).where(col(Secret.name).in_(wanted))
|
||||||
|
)
|
||||||
|
).all()
|
||||||
|
}
|
||||||
|
missing = [n for n in wanted if n not in rows]
|
||||||
|
if missing:
|
||||||
|
raise UnknownSecretError(", ".join(missing))
|
||||||
|
out = [_open(box, rows[n]) for n in wanted]
|
||||||
|
now = _now()
|
||||||
|
for row in rows.values():
|
||||||
|
row.used_at = now
|
||||||
|
session.add(row)
|
||||||
|
await session.commit()
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def _open(box: SecretBox, row: Secret) -> Revealed:
|
||||||
|
try:
|
||||||
|
data = box.open(row.name, row.sealed)
|
||||||
|
except InvalidTag as exc:
|
||||||
|
msg = f"secret {row.name} does not open: SECRETS_KEY changed since it was set"
|
||||||
|
raise SecretsDisabledError(msg) from exc
|
||||||
|
return Revealed(name=row.name, data=data, filename=row.filename)
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
"""MCP server definitions and (later) the internal aggregator app."""
|
"""MCP server definitions and the internal aggregator app."""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
|||||||
@@ -1,11 +1,7 @@
|
|||||||
"""Proxy ``FastMCP`` servers for user-declared external MCPs (``stdio``/``http``).
|
"""Proxy ``FastMCP`` servers for user-declared external MCPs (``stdio``/``http``).
|
||||||
|
|
||||||
Both flavours end up as a ``FastMCPProxy`` instance, built via
|
Built via ``fastmcp.server.create_proxy``; the proxy lazily opens its
|
||||||
``fastmcp.server.create_proxy``. The proxy lazily opens the underlying
|
transport on first request, so unused MCPs cost nothing until routed to.
|
||||||
client transport when the first MCP request arrives, so we don't pay
|
|
||||||
for connections that nothing routes to. From the aggregator app's
|
|
||||||
point of view a proxy is indistinguishable from a regular ``FastMCP``
|
|
||||||
namespace — same ``http_app`` surface, same mount semantics.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|||||||
@@ -1,26 +1,7 @@
|
|||||||
"""Internal MCP aggregator — one ASGI app, N FastMCP namespaces + ``all``.
|
"""Internal MCP aggregator — one ASGI app, N FastMCP namespaces + ``all``.
|
||||||
|
|
||||||
Each ``McpServer`` declared in the user's config becomes its own
|
Each ``McpServer`` in the user's config is mounted under ``/mcp/<name>``;
|
||||||
``FastMCP`` instance (regular for ``python_tool``, ``FastMCPProxy`` for
|
``/mcp/all/`` unions every namespace's tools for single-MCP clients.
|
||||||
``stdio``/``http``) and is mounted under ``/mcp/<name>`` on a single
|
|
||||||
Starlette app. This app runs on ``127.0.0.1:INTERNAL_MCP_PORT`` (not
|
|
||||||
EXPOSE'd in Docker) so the ClaudeCode subprocess can reach each
|
|
||||||
namespace via loopback as a distinct MCP server URL — preserving
|
|
||||||
per-domain framing while costing only one process worth of RAM
|
|
||||||
(PRD §6).
|
|
||||||
|
|
||||||
Phase 3 adds ``/mcp/all/``: a single FastMCP whose tools are the union
|
|
||||||
of every namespace's tools, prefixed by FastMCP's ``namespace_<tool>``
|
|
||||||
convention (e.g. ``time_current_time``). It's the escape-hatch for
|
|
||||||
clients that can only configure one MCP server — discouraged for tool-
|
|
||||||
heavy setups (PRD §6 cites the ~95%→~71% tool-selection drop on flat
|
|
||||||
namespaces) but real and reachable.
|
|
||||||
|
|
||||||
The aggregator returns both the app and a ``{name: url}`` map; Phase
|
|
||||||
``ClaudeSdkBackend`` plugs the map directly into
|
|
||||||
``BackendOptions.mcp_servers``. ``/mcp/all/`` is NOT included in that
|
|
||||||
map — claude-code-agents always get per-domain framing; only the
|
|
||||||
external MCP frontend (Phase 3.1) reverse-proxies the flat endpoint.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -52,11 +33,13 @@ def build_internal_app(
|
|||||||
) -> tuple[Starlette, dict[str, str], dict[str, FastMCP]]:
|
) -> tuple[Starlette, dict[str, str], dict[str, FastMCP]]:
|
||||||
"""Build the aggregator ``Starlette`` app, per-namespace URL map, and server map.
|
"""Build the aggregator ``Starlette`` app, per-namespace URL map, and server map.
|
||||||
|
|
||||||
``host``/``port`` only flavour the URL strings handed back — actually
|
``host``/``port`` only flavour the URL strings handed back; the
|
||||||
listening on them is the caller's job (``cli.main`` runs a uvicorn
|
caller (``cli.main``) is what actually listens on them. Published
|
||||||
server in a TaskGroup). We accept the address here so callers don't
|
URLs carry a trailing slash so ``Mount`` doesn't 307-redirect the
|
||||||
have to format the URLs themselves and risk drifting from the
|
first request. Each
|
||||||
``/mcp/<name>`` convention.
|
namespace gets its own :class:`RedactingMiddleware` so a direct
|
||||||
|
``call_tool`` on a child (e.g. the Raycast backend does that) is
|
||||||
|
filtered too, not just requests through the mounted app.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
* Starlette app to serve via uvicorn.
|
* Starlette app to serve via uvicorn.
|
||||||
@@ -70,10 +53,6 @@ def build_internal_app(
|
|||||||
MCP tools into the Raycast wire).
|
MCP tools into the Raycast wire).
|
||||||
"""
|
"""
|
||||||
servers: dict[str, FastMCP] = {spec.name: _build_server(spec) for spec in mcps}
|
servers: dict[str, FastMCP] = {spec.name: _build_server(spec) for spec in mcps}
|
||||||
# The one place every tool result passes through on its way to the
|
|
||||||
# model. Attached per namespace rather than once at the top so that
|
|
||||||
# a direct ``call_tool`` on a child (the Raycast backend does that)
|
|
||||||
# is filtered too.
|
|
||||||
for server in servers.values():
|
for server in servers.values():
|
||||||
server.add_middleware(RedactingMiddleware())
|
server.add_middleware(RedactingMiddleware())
|
||||||
|
|
||||||
@@ -82,8 +61,6 @@ def build_internal_app(
|
|||||||
}
|
}
|
||||||
routes = [Mount(f"/mcp/{name}", app=app) for name, app in child_apps.items()]
|
routes = [Mount(f"/mcp/{name}", app=app) for name, app in child_apps.items()]
|
||||||
|
|
||||||
# /mcp/all — flat-namespace bundle. Skip when there's nothing to
|
|
||||||
# bundle so we don't pay for an empty session manager lifecycle.
|
|
||||||
all_app = None
|
all_app = None
|
||||||
if servers:
|
if servers:
|
||||||
all_server = _build_all_server(servers)
|
all_server = _build_all_server(servers)
|
||||||
@@ -93,11 +70,11 @@ def build_internal_app(
|
|||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(_parent: Starlette) -> AsyncIterator[None]:
|
async def lifespan(_parent: Starlette) -> AsyncIterator[None]:
|
||||||
# Each FastMCP http_app stores its session manager init in its
|
"""Enter every child app's lifespan.
|
||||||
# own lifespan. Without entering them the streamable-HTTP layer
|
|
||||||
# 500s on every request. AsyncExitStack composes them so all
|
Without it the streamable-HTTP layer 500s on every request;
|
||||||
# children come up together and unwind in reverse order on
|
children unwind in reverse order.
|
||||||
# shutdown.
|
"""
|
||||||
async with AsyncExitStack() as stack:
|
async with AsyncExitStack() as stack:
|
||||||
for child in child_apps.values():
|
for child in child_apps.values():
|
||||||
await stack.enter_async_context(child.router.lifespan_context(child))
|
await stack.enter_async_context(child.router.lifespan_context(child))
|
||||||
@@ -108,23 +85,22 @@ def build_internal_app(
|
|||||||
yield
|
yield
|
||||||
|
|
||||||
app = Starlette(routes=routes, lifespan=lifespan)
|
app = Starlette(routes=routes, lifespan=lifespan)
|
||||||
# Trailing slash on the published URL skips Starlette's
|
|
||||||
# 307 redirect from ``/mcp/<name>`` to ``/mcp/<name>/`` that
|
|
||||||
# ``Mount`` produces when a child route lives at ``/``.
|
|
||||||
urls = {name: f"http://{host}:{port}/mcp/{name}/" for name in servers}
|
urls = {name: f"http://{host}:{port}/mcp/{name}/" for name in servers}
|
||||||
return app, urls, servers
|
return app, urls, servers
|
||||||
|
|
||||||
|
|
||||||
def _build_server(spec: McpServerT) -> FastMCP:
|
def _build_server(spec: McpServerT) -> FastMCP:
|
||||||
"""Dispatch on the discriminated union to the matching builder."""
|
"""Dispatch on the discriminated union to the matching builder.
|
||||||
|
|
||||||
|
The final branch is unreachable while ``McpServerT`` stays closed;
|
||||||
|
it exists to keep type-narrowing honest if a variant is ever added.
|
||||||
|
"""
|
||||||
if isinstance(spec, PythonToolMcp):
|
if isinstance(spec, PythonToolMcp):
|
||||||
return build_python_tool_server(spec)
|
return build_python_tool_server(spec)
|
||||||
if isinstance(spec, StdioMcp):
|
if isinstance(spec, StdioMcp):
|
||||||
return build_stdio_proxy(spec)
|
return build_stdio_proxy(spec)
|
||||||
if isinstance(spec, HttpMcp):
|
if isinstance(spec, HttpMcp):
|
||||||
return build_http_proxy(spec)
|
return build_http_proxy(spec)
|
||||||
# `McpServerT` is a closed union; this is unreachable but keeps
|
|
||||||
# type-narrowing honest if a new variant lands without updates here.
|
|
||||||
msg = f"unsupported McpServer variant: {type(spec).__name__}"
|
msg = f"unsupported McpServer variant: {type(spec).__name__}"
|
||||||
raise TypeError(msg)
|
raise TypeError(msg)
|
||||||
|
|
||||||
|
|||||||
@@ -1,22 +1,8 @@
|
|||||||
"""Tolerant transports for upstream MCPs that don't strictly speak JSON-RPC.
|
"""Tolerant transports for upstream MCPs that don't strictly speak JSON-RPC.
|
||||||
|
|
||||||
Some real-world MCP servers print non-JSON chatter to stdout before / between
|
Some MCP servers print non-JSON chatter to stdout (banners, dependency-load
|
||||||
their actual JSON-RPC frames (``Processing...``, banners, dependency-load
|
messages) that the reference client forwards as exceptions, which bleed into
|
||||||
messages, etc.). The reference ``mcp.client.stdio.stdio_client`` parses every
|
client UIs as warnings. This transport drops those lines instead.
|
||||||
stdout line as JSON-RPC and ships any parse failure downstream as an
|
|
||||||
exception, which the MCP ``ClientSession`` then logs as a warning that bleeds
|
|
||||||
into client UIs (Cursor, Cline) when they connect through us.
|
|
||||||
|
|
||||||
``LenientStdioTransport`` re-implements the stdio-client wiring with one
|
|
||||||
behavioural change: lines that don't parse as JSON-RPC are *silently
|
|
||||||
dropped* (one ``DEBUG`` log entry, no exception forwarded). Downstream
|
|
||||||
consumers see only valid messages. We keep the rest of the contract identical
|
|
||||||
to the reference client, including the spec-mandated graceful shutdown
|
|
||||||
sequence (close stdin → wait → SIGTERM → SIGKILL).
|
|
||||||
|
|
||||||
The transport plugs into ``fastmcp.server.create_proxy`` just like
|
|
||||||
``StdioTransport`` does, so the rest of the aggregator doesn't need to
|
|
||||||
know which flavour it got.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -70,13 +56,15 @@ class LenientStdioTransport(ClientTransport):
|
|||||||
cwd: str | None = None,
|
cwd: str | None = None,
|
||||||
log_file: TextIO | None = None,
|
log_file: TextIO | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
|
"""``log_file`` takes an already-open ``TextIO``.
|
||||||
|
|
||||||
|
Unlike the upstream transport, this one does not open a ``Path``
|
||||||
|
for you.
|
||||||
|
"""
|
||||||
self.command = command
|
self.command = command
|
||||||
self.args = args
|
self.args = args
|
||||||
self.env = env
|
self.env = env
|
||||||
self.cwd = cwd
|
self.cwd = cwd
|
||||||
# TextIO only — pre-open Path callers themselves. The upstream
|
|
||||||
# ``StdioTransport`` opens Path for you, but we keep this thin so
|
|
||||||
# the type contract stays narrow and easy to validate.
|
|
||||||
self.log_file = log_file
|
self.log_file = log_file
|
||||||
|
|
||||||
@contextlib.asynccontextmanager
|
@contextlib.asynccontextmanager
|
||||||
@@ -100,7 +88,7 @@ class LenientStdioTransport(ClientTransport):
|
|||||||
|
|
||||||
|
|
||||||
@contextlib.asynccontextmanager
|
@contextlib.asynccontextmanager
|
||||||
async def _lenient_stdio_client( # noqa: PLR0915 — mirrors mcp.client.stdio.stdio_client
|
async def _lenient_stdio_client( # noqa: PLR0915
|
||||||
server: StdioServerParameters, errlog: TextIO = sys.stderr
|
server: StdioServerParameters, errlog: TextIO = sys.stderr
|
||||||
) -> AsyncIterator[
|
) -> AsyncIterator[
|
||||||
tuple[
|
tuple[
|
||||||
@@ -112,8 +100,8 @@ async def _lenient_stdio_client( # noqa: PLR0915 — mirrors mcp.client.stdio.s
|
|||||||
|
|
||||||
All differences from upstream live in ``stdout_reader``: lines that fail
|
All differences from upstream live in ``stdout_reader``: lines that fail
|
||||||
``JSONRPCMessage.model_validate_json`` are logged at DEBUG and skipped,
|
``JSONRPCMessage.model_validate_json`` are logged at DEBUG and skipped,
|
||||||
never forwarded as exceptions. This is what makes warning-noisy MCPs
|
never forwarded as exceptions. Shutdown still follows the MCP spec
|
||||||
quiet from the consumer's point of view.
|
sequence: close stdin, wait, SIGTERM, then SIGKILL.
|
||||||
"""
|
"""
|
||||||
read_stream_writer, read_stream = anyio.create_memory_object_stream[
|
read_stream_writer, read_stream = anyio.create_memory_object_stream[
|
||||||
SessionMessage | Exception
|
SessionMessage | Exception
|
||||||
@@ -159,7 +147,7 @@ async def _lenient_stdio_client( # noqa: PLR0915 — mirrors mcp.client.stdio.s
|
|||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
message = types.JSONRPCMessage.model_validate_json(stripped)
|
message = types.JSONRPCMessage.model_validate_json(stripped)
|
||||||
except Exception: # noqa: BLE001 — by design, see module doc
|
except Exception: # noqa: BLE001
|
||||||
_log.debug(
|
_log.debug(
|
||||||
"lenient stdio: dropped non-JSON line: %r",
|
"lenient stdio: dropped non-JSON line: %r",
|
||||||
stripped[:200],
|
stripped[:200],
|
||||||
@@ -192,7 +180,6 @@ async def _lenient_stdio_client( # noqa: PLR0915 — mirrors mcp.client.stdio.s
|
|||||||
try:
|
try:
|
||||||
yield read_stream, write_stream
|
yield read_stream, write_stream
|
||||||
finally:
|
finally:
|
||||||
# MCP spec stdio shutdown: close stdin → wait → SIGTERM → SIGKILL.
|
|
||||||
if process.stdin:
|
if process.stdin:
|
||||||
with contextlib.suppress(Exception):
|
with contextlib.suppress(Exception):
|
||||||
await process.stdin.aclose()
|
await process.stdin.aclose()
|
||||||
|
|||||||
@@ -1,26 +1,7 @@
|
|||||||
"""Redact credentials on the way out of an MCP tool.
|
"""Redact credentials in MCP tool output before it reaches the model.
|
||||||
|
|
||||||
A tool result is a wider channel than the log. It goes straight into the
|
Covers every MCP route; in-process tools like ``Bash``/``Read`` are
|
||||||
model's context, from there into the turn record in Postgres, and from
|
guarded separately by ``agents.policy``.
|
||||||
there into the markdown transcript in ``💬 чаты`` — which Obsidian Sync
|
|
||||||
carries off the machine. One ``komodo`` deploy returns the resolved
|
|
||||||
compose file, ``environment:`` block and all, so a single call can put
|
|
||||||
every secret of a stack into all four places at once.
|
|
||||||
|
|
||||||
One filter, not a list of exceptions: every MCP the model can reach is
|
|
||||||
built into a ``FastMCP`` by :mod:`beaver_gateway.mcp.internal_app` —
|
|
||||||
``python_tool`` bundles like komodo, stdio subprocesses, remote HTTP
|
|
||||||
servers — and every route into one of them runs its middleware chain.
|
|
||||||
That covers the per-namespace ``/mcp/<name>/`` mounts claude-code talks
|
|
||||||
to, the ``/mcp/all`` bundle, the external MCP frontend reverse-proxying
|
|
||||||
into both, and the Raycast backend's direct ``call_tool``. A new MCP in
|
|
||||||
``config.py`` is covered the day it is added, without anyone
|
|
||||||
remembering to list it here.
|
|
||||||
|
|
||||||
What this does not reach: tools that never touch a FastMCP server — the
|
|
||||||
gateway's own ``gateway`` tools, and everything claude-code runs inside
|
|
||||||
its own process (``Bash``, ``Read``). Those are guarded by
|
|
||||||
:mod:`beaver_gateway.core.policy` and the vault mounts instead.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -31,7 +12,7 @@ import mcp.types as mt
|
|||||||
from fastmcp.server.middleware import Middleware
|
from fastmcp.server.middleware import Middleware
|
||||||
from fastmcp.tools.base import ToolResult
|
from fastmcp.tools.base import ToolResult
|
||||||
|
|
||||||
from beaver_gateway.core.redact import redact, redact_data
|
from beaver_gateway.security.redact import redact, redact_data
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from fastmcp.server.middleware import CallNext, MiddlewareContext
|
from fastmcp.server.middleware import CallNext, MiddlewareContext
|
||||||
@@ -69,6 +50,12 @@ class RedactingMiddleware(Middleware):
|
|||||||
context: MiddlewareContext[mt.CallToolRequestParams],
|
context: MiddlewareContext[mt.CallToolRequestParams],
|
||||||
call_next: CallNext[mt.CallToolRequestParams, ToolResult],
|
call_next: CallNext[mt.CallToolRequestParams, ToolResult],
|
||||||
) -> ToolResult:
|
) -> ToolResult:
|
||||||
|
"""Mask the result.
|
||||||
|
|
||||||
|
A non-``None`` ``meta`` on the masked copy takes the
|
||||||
|
``CallToolResult`` path that skips output-schema validation,
|
||||||
|
since masked structured content may no longer match it.
|
||||||
|
"""
|
||||||
result = await call_next(context)
|
result = await call_next(context)
|
||||||
content = [_redact_block(block) for block in result.content]
|
content = [_redact_block(block) for block in result.content]
|
||||||
structured = redact_data(result.structured_content)
|
structured = redact_data(result.structured_content)
|
||||||
@@ -77,9 +64,5 @@ class RedactingMiddleware(Middleware):
|
|||||||
return ToolResult(
|
return ToolResult(
|
||||||
content=content,
|
content=content,
|
||||||
structured_content=structured,
|
structured_content=structured,
|
||||||
# Masked structured content no longer has to satisfy the
|
|
||||||
# tool's output schema; a non-None meta takes the
|
|
||||||
# ``CallToolResult`` path that skips that validation, the
|
|
||||||
# same trick ``ResponseLimitingMiddleware`` uses.
|
|
||||||
meta=result.meta if result.meta is not None else {},
|
meta=result.meta if result.meta is not None else {},
|
||||||
)
|
)
|
||||||
@@ -6,7 +6,7 @@ discriminated-union members so downstream code can ``match`` on ``kind``.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from collections.abc import Callable # noqa: TC003 — runtime use by pydantic
|
from collections.abc import Callable # noqa: TC003
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Annotated, Literal
|
from typing import TYPE_CHECKING, Annotated, Literal
|
||||||
|
|
||||||
@@ -20,39 +20,35 @@ class _BaseMcp(BaseModel):
|
|||||||
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
|
model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
|
||||||
|
|
||||||
name: str
|
name: str
|
||||||
|
"""The namespace this MCP server is mounted under."""
|
||||||
|
|
||||||
|
|
||||||
class StdioMcp(_BaseMcp):
|
class StdioMcp(_BaseMcp):
|
||||||
"""Subprocess MCP server we spawn and connect to over stdio.
|
"""Subprocess MCP server spawned and connected to over stdio."""
|
||||||
|
|
||||||
``lenient`` switches the upstream stdio reader to a tolerant variant
|
|
||||||
that silently drops non-JSON-RPC lines from the subprocess's stdout
|
|
||||||
(``Processing...``-style chatter, banners, dependency-load messages).
|
|
||||||
The reference ``mcp.client.stdio.stdio_client`` forwards those parse
|
|
||||||
failures as exceptions, which bleed into Cursor/Cline UIs as warnings
|
|
||||||
when they connect through us. Default ``False`` keeps the strict
|
|
||||||
contract — flip it on per-namespace for known-noisy upstreams.
|
|
||||||
"""
|
|
||||||
|
|
||||||
kind: Literal["stdio"] = "stdio"
|
kind: Literal["stdio"] = "stdio"
|
||||||
command: tuple[str, ...]
|
command: tuple[str, ...]
|
||||||
|
"""Argv used to spawn the subprocess."""
|
||||||
env: dict[str, str] | None = None
|
env: dict[str, str] | None = None
|
||||||
|
"""Extra environment variables for the subprocess."""
|
||||||
cwd: Path | None = None
|
cwd: Path | None = None
|
||||||
|
"""Working directory to spawn the subprocess in."""
|
||||||
lenient: bool = False
|
lenient: bool = False
|
||||||
|
"""Tolerate non-JSON-RPC stdout lines instead of raising (the strict
|
||||||
|
reader's parse errors surface as warnings in Cursor/Cline). Off by default."""
|
||||||
|
|
||||||
|
|
||||||
class HttpMcp(_BaseMcp):
|
class HttpMcp(_BaseMcp):
|
||||||
"""Remote MCP server reached over streamable HTTP.
|
"""Remote MCP server reached over streamable HTTP."""
|
||||||
|
|
||||||
``headers`` are forwarded on every request to the upstream MCP — handy
|
|
||||||
for upstreams that authenticate via custom header rather than a Bearer
|
|
||||||
token (``auth``).
|
|
||||||
"""
|
|
||||||
|
|
||||||
kind: Literal["http"] = "http"
|
kind: Literal["http"] = "http"
|
||||||
url: str
|
url: str
|
||||||
|
"""Streamable-HTTP endpoint URL."""
|
||||||
auth: str | None = None
|
auth: str | None = None
|
||||||
|
"""Bearer token sent as the upstream's ``Authorization`` header."""
|
||||||
headers: dict[str, str] | None = None
|
headers: dict[str, str] | None = None
|
||||||
|
"""Extra headers forwarded on every request; for upstreams that
|
||||||
|
authenticate via a custom header instead of ``auth``."""
|
||||||
|
|
||||||
|
|
||||||
class PythonToolMcp(_BaseMcp):
|
class PythonToolMcp(_BaseMcp):
|
||||||
@@ -60,13 +56,14 @@ class PythonToolMcp(_BaseMcp):
|
|||||||
|
|
||||||
kind: Literal["python_tool"] = "python_tool"
|
kind: Literal["python_tool"] = "python_tool"
|
||||||
tools: tuple[Callable[..., object], ...]
|
tools: tuple[Callable[..., object], ...]
|
||||||
|
"""The callables to expose as MCP tools."""
|
||||||
|
|
||||||
|
|
||||||
McpServerT = Annotated[StdioMcp | HttpMcp | PythonToolMcp, Field(discriminator="kind")]
|
McpServerT = Annotated[StdioMcp | HttpMcp | PythonToolMcp, Field(discriminator="kind")]
|
||||||
|
|
||||||
|
|
||||||
class McpServer:
|
class McpServer:
|
||||||
"""Factory facade matching the PRD-documented config surface."""
|
"""Factory facade for declaring MCP servers in config."""
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def stdio(
|
def stdio(
|
||||||
@@ -78,6 +75,7 @@ class McpServer:
|
|||||||
cwd: Path | str | None = None,
|
cwd: Path | str | None = None,
|
||||||
lenient: bool = False,
|
lenient: bool = False,
|
||||||
) -> StdioMcp:
|
) -> StdioMcp:
|
||||||
|
"""Declare a subprocess MCP server spawned over stdio."""
|
||||||
return StdioMcp(
|
return StdioMcp(
|
||||||
name=name,
|
name=name,
|
||||||
command=tuple(command),
|
command=tuple(command),
|
||||||
@@ -95,10 +93,12 @@ class McpServer:
|
|||||||
auth: str | None = None,
|
auth: str | None = None,
|
||||||
headers: dict[str, str] | None = None,
|
headers: dict[str, str] | None = None,
|
||||||
) -> HttpMcp:
|
) -> HttpMcp:
|
||||||
|
"""Declare a remote MCP server reached over streamable HTTP."""
|
||||||
return HttpMcp(name=name, url=url, auth=auth, headers=headers)
|
return HttpMcp(name=name, url=url, auth=auth, headers=headers)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def python_tool(
|
def python_tool(
|
||||||
cls, *, name: str, tools: Iterable[Callable[..., object]]
|
cls, *, name: str, tools: Iterable[Callable[..., object]]
|
||||||
) -> PythonToolMcp:
|
) -> PythonToolMcp:
|
||||||
|
"""Declare a namespace of Python callables exposed as MCP tools."""
|
||||||
return PythonToolMcp(name=name, tools=tuple(tools))
|
return PythonToolMcp(name=name, tools=tuple(tools))
|
||||||
@@ -1,8 +1,7 @@
|
|||||||
"""Wrap a ``PythonToolMcp`` spec into a mountable ``FastMCP`` instance.
|
"""Wrap a ``PythonToolMcp`` spec into a mountable ``FastMCP`` instance.
|
||||||
|
|
||||||
Each ``python_tool`` McpServer in the user's config becomes a separate
|
Each ``python_tool`` McpServer becomes its own ``FastMCP`` namespace so
|
||||||
``FastMCP`` namespace — one domain, one server URL — so models keep the
|
models keep the per-domain tool framing they were trained on.
|
||||||
per-domain framing they were trained on (see PRD §6).
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|||||||
@@ -0,0 +1 @@
|
|||||||
|
"""Bearer tokens, the audit log and credential redaction."""
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
"""Single entry point for writing :class:`AuditLog` rows.
|
||||||
|
|
||||||
|
``log()`` is fire-and-forget: it awaits the DB write so callers get
|
||||||
|
ordering, but never raises — failures are logged and swallowed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
|
from beaver_gateway.storage import append_audit
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from beaver_gateway.frontends.base import GatewayRuntime
|
||||||
|
|
||||||
|
|
||||||
|
_log = logging.getLogger("beaver_gateway.audit")
|
||||||
|
|
||||||
|
|
||||||
|
KNOWN_KINDS: frozenset[str] = frozenset(
|
||||||
|
{
|
||||||
|
"messages",
|
||||||
|
"mcp_call",
|
||||||
|
"tool_call",
|
||||||
|
"login_ok",
|
||||||
|
"login_failed",
|
||||||
|
"logout",
|
||||||
|
"token_create",
|
||||||
|
"token_revoke",
|
||||||
|
"secret_set",
|
||||||
|
"secret_delete",
|
||||||
|
"secret_use",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
"""Kinds the gateway currently emits; ``AuditLog.kind`` stays free-form."""
|
||||||
|
|
||||||
|
|
||||||
|
async def log(
|
||||||
|
runtime: GatewayRuntime,
|
||||||
|
*,
|
||||||
|
actor: str,
|
||||||
|
kind: str,
|
||||||
|
agent_name: str | None = None,
|
||||||
|
**detail: Any,
|
||||||
|
) -> None:
|
||||||
|
"""Best-effort audit insert. Never raises; DB failures are logged.
|
||||||
|
|
||||||
|
``actor`` is free-form (``"token:<name>"``, ``"admin:<user>"``,
|
||||||
|
``"anon"``); ``kind`` is a short tag (see :data:`KNOWN_KINDS`); and
|
||||||
|
``**detail`` is JSON-serialised into ``AuditLog.detail_json`` — keep
|
||||||
|
it small, not full request bodies.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
async with runtime.db.session() as session:
|
||||||
|
await append_audit(
|
||||||
|
session,
|
||||||
|
actor=actor,
|
||||||
|
kind=kind,
|
||||||
|
agent_name=agent_name,
|
||||||
|
detail=detail or None,
|
||||||
|
)
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
_log.exception(
|
||||||
|
"audit write failed: actor=%s kind=%s agent=%s", actor, kind, agent_name
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["KNOWN_KINDS", "log"]
|
||||||
@@ -1,32 +1,7 @@
|
|||||||
"""Bearer-token verification (Phase 4.2 — DB-backed with in-memory cache).
|
"""Bearer-token verification: a DB-backed store with an in-memory cache.
|
||||||
|
|
||||||
The store is fed by two sources:
|
Optional ``BOOTSTRAP_TOKENS`` env entries are checked first. Verify takes
|
||||||
|
a raw token or ``Authorization`` header and returns a :class:`TokenIdentity`.
|
||||||
1. **DB** (``Token`` table from Phase 4.1) — the primary source. Rows
|
|
||||||
carry Argon2id hashes; the admin UI (Phase 4.3) will be the only
|
|
||||||
writer at steady state.
|
|
||||||
2. **`BOOTSTRAP_TOKENS`** env — a name→plaintext map kept around for
|
|
||||||
first-run, disaster-recovery, and ``examples/`` smoke tests. These
|
|
||||||
entries live alongside DB rows in the cache and are never persisted.
|
|
||||||
|
|
||||||
Hot path is in-memory: at :meth:`start` we pull every non-revoked DB
|
|
||||||
row and stash it in a list; subsequent :meth:`verify` calls re-load
|
|
||||||
when the cache is older than ``ttl_seconds``. ``last_used_at`` updates
|
|
||||||
are coalesced into a small dict and flushed by a background task every
|
|
||||||
``flush_interval`` seconds — one transaction per flush rather than one
|
|
||||||
per request.
|
|
||||||
|
|
||||||
We can't index DB rows by a derived plaintext key because Argon2 salts
|
|
||||||
are random — so verify does a linear scan over candidates, calling
|
|
||||||
``argon2.PasswordHasher.verify`` on each. N is small by design (single
|
|
||||||
operator, ~10 tokens at most); the cost is irrelevant. The scan runs
|
|
||||||
through ``asyncio.to_thread`` to keep the event loop free of the ~50ms
|
|
||||||
KDF block.
|
|
||||||
|
|
||||||
The module knows nothing about HTTP frameworks. It takes a raw token
|
|
||||||
(or a verbatim ``Authorization`` header value) and returns a
|
|
||||||
:class:`TokenIdentity` (name + scope + db-id), or ``None`` for a miss.
|
|
||||||
Frontends own the 401 response shape.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -46,7 +21,7 @@ from argon2.exceptions import InvalidHashError, VerifyMismatchError
|
|||||||
from beaver_gateway.storage import list_active_tokens, touch_token
|
from beaver_gateway.storage import list_active_tokens, touch_token
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Mapping
|
from collections.abc import Iterable, Mapping
|
||||||
|
|
||||||
from beaver_gateway.storage import Database
|
from beaver_gateway.storage import Database
|
||||||
|
|
||||||
@@ -60,8 +35,8 @@ class TokenStoreError(ValueError):
|
|||||||
"""Malformed ``BOOTSTRAP_TOKENS`` value or duplicate token."""
|
"""Malformed ``BOOTSTRAP_TOKENS`` value or duplicate token."""
|
||||||
|
|
||||||
|
|
||||||
VALID_SCOPES: frozenset[str] = frozenset({"*", "messages", "mcp", "admin", "api"})
|
BUILTIN_SCOPES: frozenset[str] = frozenset({"*", "messages", "mcp", "admin", "api"})
|
||||||
"""The scopes a ``Token.scope`` may hold (Phase 4.3 admin UI enforces).
|
"""The scopes the gateway's own frontends define.
|
||||||
|
|
||||||
* ``*`` — wildcard, may use any frontend
|
* ``*`` — wildcard, may use any frontend
|
||||||
* ``messages`` — Anthropic Messages frontend only
|
* ``messages`` — Anthropic Messages frontend only
|
||||||
@@ -69,17 +44,25 @@ VALID_SCOPES: frozenset[str] = frozenset({"*", "messages", "mcp", "admin", "api"
|
|||||||
* ``api`` — conversations API (``/api``) and its SSE
|
* ``api`` — conversations API (``/api``) and its SSE
|
||||||
* ``admin`` — reserved for programmatic admin access; the AdminFrontend
|
* ``admin`` — reserved for programmatic admin access; the AdminFrontend
|
||||||
itself authenticates via session cookies, not bearer tokens, so this
|
itself authenticates via session cookies, not bearer tokens, so this
|
||||||
scope is unused today and kept for forward compatibility.
|
scope is unused today.
|
||||||
|
|
||||||
|
A setup's own frontend adds its scope on top; see :func:`scopes_with` and
|
||||||
|
``GatewayRuntime.scopes``, which is what ``POST /api/tokens`` validates.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def scopes_with(extra: Iterable[str | None]) -> frozenset[str]:
|
||||||
|
"""The builtin scopes plus every frontend's own; ``None`` entries drop out."""
|
||||||
|
return BUILTIN_SCOPES | {s for s in extra if s}
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
class TokenIdentity:
|
class TokenIdentity:
|
||||||
"""What :meth:`TokenStore.verify` resolves to on success.
|
"""What :meth:`TokenStore.verify` resolves to on success.
|
||||||
|
|
||||||
``token_id`` is the DB row id for persisted tokens, or ``None`` for
|
``token_id`` is the DB row id for persisted tokens, or ``None`` for
|
||||||
an env-bootstrap match (those have no DB row to touch). ``scope``
|
an env-bootstrap match (those have no DB row to touch). ``scope``
|
||||||
gates which frontend the token may hit (see :data:`VALID_SCOPES`);
|
gates which frontend the token may hit (see :data:`BUILTIN_SCOPES`);
|
||||||
bootstrap tokens implicitly get ``"*"``.
|
bootstrap tokens implicitly get ``"*"``.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -92,8 +75,8 @@ class TokenIdentity:
|
|||||||
|
|
||||||
``"*"`` is the wildcard; an exact match satisfies a single scope.
|
``"*"`` is the wildcard; an exact match satisfies a single scope.
|
||||||
Unknown ``required`` values intentionally fall through to a
|
Unknown ``required`` values intentionally fall through to a
|
||||||
strict equality check — callers should pass one of
|
strict equality check — callers should pass a scope the runtime
|
||||||
:data:`VALID_SCOPES`.
|
knows (:data:`BUILTIN_SCOPES` plus every frontend's own).
|
||||||
"""
|
"""
|
||||||
return self.scope in ("*", required)
|
return self.scope in ("*", required)
|
||||||
|
|
||||||
@@ -108,33 +91,23 @@ class _CachedToken:
|
|||||||
hashed_value: str
|
hashed_value: str
|
||||||
|
|
||||||
|
|
||||||
# Default Argon2id parameters from ``argon2-cffi`` are fine for our scope.
|
|
||||||
# They target ~50ms on a modern CPU — enough to make a stolen-hash brute
|
|
||||||
# force expensive, cheap enough to verify a handful per request.
|
|
||||||
_HASHER = PasswordHasher()
|
_HASHER = PasswordHasher()
|
||||||
|
|
||||||
|
|
||||||
def hash_token(plaintext: str) -> str:
|
def hash_token(plaintext: str) -> str:
|
||||||
"""Return an Argon2id hash for ``plaintext`` (admin / seed-only path).
|
"""Return an Argon2id hash for ``plaintext``, for admin/seed use."""
|
||||||
|
|
||||||
Phase 4.3 will call this when the admin creates a token; Phase 4.2
|
|
||||||
exposes it so smoke scripts can seed the DB without re-implementing
|
|
||||||
the same line.
|
|
||||||
"""
|
|
||||||
return _HASHER.hash(plaintext)
|
return _HASHER.hash(plaintext)
|
||||||
|
|
||||||
|
|
||||||
class TokenStore:
|
class TokenStore:
|
||||||
"""DB-backed verifier with in-memory cache + TTL + batched touches.
|
"""DB-backed verifier with an in-memory cache, TTL, and batched touches.
|
||||||
|
|
||||||
Construct in ``cli.main`` after :class:`Database` is up, then
|
Call ``await store.start()`` after construction to prime the cache
|
||||||
``await store.start()`` to prime the cache and spin up the flusher
|
and start the flusher; ``await store.stop()`` drains the touch queue.
|
||||||
task. ``await store.stop()`` on shutdown drains the touch queue.
|
|
||||||
|
|
||||||
Bootstrap entries (from ``BOOTSTRAP_TOKENS``) sit alongside DB rows
|
Bootstrap entries (``BOOTSTRAP_TOKENS``) are checked first and stay
|
||||||
in the same lookup path; we check them first, in constant time, so
|
usable even if the DB is unreachable; they never get flushed to
|
||||||
they remain usable even if the DB is unreachable. They never appear
|
``last_used_at`` since they have no DB row.
|
||||||
in ``last_used_at`` flushes because they have no DB row.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
__slots__ = (
|
__slots__ = (
|
||||||
@@ -159,8 +132,11 @@ class TokenStore:
|
|||||||
ttl_seconds: float = 30.0,
|
ttl_seconds: float = 30.0,
|
||||||
flush_interval: float = 5.0,
|
flush_interval: float = 5.0,
|
||||||
) -> None:
|
) -> None:
|
||||||
# Bootstrap is keyed by value internally so verify is O(1) over
|
"""Index bootstrap tokens for O(1) verify lookup.
|
||||||
# plaintext. Each value also keeps its name for audit lines.
|
|
||||||
|
Each one's scope defaults to ``"*"`` unless narrowed via
|
||||||
|
``name:value:scope``.
|
||||||
|
"""
|
||||||
by_value: dict[str, str] = {}
|
by_value: dict[str, str] = {}
|
||||||
for name, value in (bootstrap or {}).items():
|
for name, value in (bootstrap or {}).items():
|
||||||
if not name or not value:
|
if not name or not value:
|
||||||
@@ -174,9 +150,6 @@ class TokenStore:
|
|||||||
raise TokenStoreError(msg)
|
raise TokenStoreError(msg)
|
||||||
by_value[value] = name
|
by_value[value] = name
|
||||||
self._bootstrap_by_value: dict[str, str] = by_value
|
self._bootstrap_by_value: dict[str, str] = by_value
|
||||||
# Bootstrap entries are ``"*"`` unless the env narrows them
|
|
||||||
# (``name:value:scope``) - a webhook token in a URL should not be
|
|
||||||
# an admin token.
|
|
||||||
self._bootstrap_scopes: dict[str, str] = {
|
self._bootstrap_scopes: dict[str, str] = {
|
||||||
name: (bootstrap_scopes or {}).get(name, _BOOTSTRAP_SCOPE)
|
name: (bootstrap_scopes or {}).get(name, _BOOTSTRAP_SCOPE)
|
||||||
for name in by_value.values()
|
for name in by_value.values()
|
||||||
@@ -192,8 +165,6 @@ class TokenStore:
|
|||||||
self._touch_queue: dict[int, datetime] = {}
|
self._touch_queue: dict[int, datetime] = {}
|
||||||
self._flusher_task: asyncio.Task[None] | None = None
|
self._flusher_task: asyncio.Task[None] | None = None
|
||||||
|
|
||||||
# ---- bootstrap parsing (kept for `cli` / tests) ---------------------
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def parse_bootstrap(raw: str) -> dict[str, str]:
|
def parse_bootstrap(raw: str) -> dict[str, str]:
|
||||||
"""Parse ``name1:value1,name2:value2[:scope]`` (``BOOTSTRAP_TOKENS``)."""
|
"""Parse ``name1:value1,name2:value2[:scope]`` (``BOOTSTRAP_TOKENS``)."""
|
||||||
@@ -210,15 +181,9 @@ class TokenStore:
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_env(cls, raw: str, db: Database | None = None) -> TokenStore:
|
def from_env(cls, raw: str, db: Database | None = None) -> TokenStore:
|
||||||
"""Legacy entrypoint: bootstrap-only (or bootstrap + db).
|
"""Bootstrap-only (or bootstrap + db) convenience constructor."""
|
||||||
|
|
||||||
Phase 1.3 call sites still expect a one-liner; we keep the
|
|
||||||
classmethod so they don't have to learn the new constructor.
|
|
||||||
"""
|
|
||||||
return cls(db, bootstrap=cls.parse_bootstrap(raw))
|
return cls(db, bootstrap=cls.parse_bootstrap(raw))
|
||||||
|
|
||||||
# ---- lifecycle ------------------------------------------------------
|
|
||||||
|
|
||||||
async def start(self) -> None:
|
async def start(self) -> None:
|
||||||
"""Prime the cache and (if a DB is attached) start the flusher loop."""
|
"""Prime the cache and (if a DB is attached) start the flusher loop."""
|
||||||
await self._refresh()
|
await self._refresh()
|
||||||
@@ -242,19 +207,19 @@ class TokenStore:
|
|||||||
self._bootstrap_scopes[name] = scope
|
self._bootstrap_scopes[name] = scope
|
||||||
|
|
||||||
async def invalidate(self) -> None:
|
async def invalidate(self) -> None:
|
||||||
"""Force the next verify to re-read from DB (Phase 4.3 admin hook)."""
|
"""Force the next verify to re-read from DB."""
|
||||||
self._loaded_at = 0.0
|
self._loaded_at = 0.0
|
||||||
|
|
||||||
# ---- verify path ----------------------------------------------------
|
|
||||||
|
|
||||||
async def verify(self, token: str | None) -> TokenIdentity | None:
|
async def verify(self, token: str | None) -> TokenIdentity | None:
|
||||||
"""Return the matching identity, or ``None`` for unknown/empty tokens."""
|
"""Return the matching identity, or ``None`` for unknown/empty tokens.
|
||||||
|
|
||||||
|
Bootstrap entries are checked first via constant-time compare;
|
||||||
|
the DB cache list is swapped rather than mutated on refresh, so
|
||||||
|
a concurrent reload can't affect an in-progress scan.
|
||||||
|
"""
|
||||||
if not token:
|
if not token:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# Bootstrap first: constant-time compare per entry, never hits DB.
|
|
||||||
# `compare_digest` is overkill for a name→value lookup but cheap
|
|
||||||
# and removes one timing variable for free.
|
|
||||||
for value, name in self._bootstrap_by_value.items():
|
for value, name in self._bootstrap_by_value.items():
|
||||||
if hmac.compare_digest(token, value):
|
if hmac.compare_digest(token, value):
|
||||||
return TokenIdentity(
|
return TokenIdentity(
|
||||||
@@ -268,9 +233,6 @@ class TokenStore:
|
|||||||
|
|
||||||
await self._ensure_fresh()
|
await self._ensure_fresh()
|
||||||
|
|
||||||
# Snapshot the cache reference so a refresh mid-scan doesn't
|
|
||||||
# surprise us. List itself is immutable per refresh (we swap,
|
|
||||||
# not mutate).
|
|
||||||
cache = self._cache
|
cache = self._cache
|
||||||
for entry in cache:
|
for entry in cache:
|
||||||
try:
|
try:
|
||||||
@@ -306,17 +268,18 @@ class TokenStore:
|
|||||||
def __bool__(self) -> bool:
|
def __bool__(self) -> bool:
|
||||||
return bool(self._cache) or bool(self._bootstrap_by_value)
|
return bool(self._cache) or bool(self._bootstrap_by_value)
|
||||||
|
|
||||||
# ---- internals ------------------------------------------------------
|
|
||||||
|
|
||||||
async def _ensure_fresh(self) -> None:
|
async def _ensure_fresh(self) -> None:
|
||||||
|
"""Reload the cache when stale.
|
||||||
|
|
||||||
|
Double-checked under the lock so concurrent callers don't all
|
||||||
|
trigger a reload.
|
||||||
|
"""
|
||||||
if self._db is None:
|
if self._db is None:
|
||||||
return
|
return
|
||||||
now = time.monotonic()
|
now = time.monotonic()
|
||||||
if now - self._loaded_at <= self._ttl:
|
if now - self._loaded_at <= self._ttl:
|
||||||
return
|
return
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
# Re-check under the lock — first arrival reloaded, others
|
|
||||||
# should fall through.
|
|
||||||
now = time.monotonic()
|
now = time.monotonic()
|
||||||
if now - self._loaded_at <= self._ttl:
|
if now - self._loaded_at <= self._ttl:
|
||||||
return
|
return
|
||||||
@@ -331,8 +294,6 @@ class TokenStore:
|
|||||||
next_cache: list[_CachedToken] = []
|
next_cache: list[_CachedToken] = []
|
||||||
for row in rows:
|
for row in rows:
|
||||||
if row.id is None:
|
if row.id is None:
|
||||||
# Defensive: SQLModel will assign an id on insert; a
|
|
||||||
# None here would mean someone handed us an unsaved row.
|
|
||||||
continue
|
continue
|
||||||
next_cache.append(
|
next_cache.append(
|
||||||
_CachedToken(
|
_CachedToken(
|
||||||
@@ -353,30 +314,31 @@ class TokenStore:
|
|||||||
await self._flush_now()
|
await self._flush_now()
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
raise
|
raise
|
||||||
except Exception: # noqa: BLE001 — never let the flusher die silently
|
except Exception: # noqa: BLE001
|
||||||
_log.exception("token flusher crashed; touches will stop")
|
_log.exception("token flusher crashed; touches will stop")
|
||||||
|
|
||||||
async def _flush_now(self) -> None:
|
async def _flush_now(self) -> None:
|
||||||
|
"""Flush queued touches in one transaction.
|
||||||
|
|
||||||
|
Detaches the queue first so concurrent ``verify()`` writes land
|
||||||
|
in the next flush.
|
||||||
|
"""
|
||||||
if self._db is None or not self._touch_queue:
|
if self._db is None or not self._touch_queue:
|
||||||
return
|
return
|
||||||
# Detach the queue so concurrent verify() writes don't bleed
|
|
||||||
# into the in-flight transaction.
|
|
||||||
pending, self._touch_queue = self._touch_queue, {}
|
pending, self._touch_queue = self._touch_queue, {}
|
||||||
async with self._db.session() as session:
|
async with self._db.session() as session:
|
||||||
for token_id in pending:
|
for token_id in pending:
|
||||||
# We don't pass the timestamp through — `touch_token`
|
|
||||||
# stamps `now` itself, and we'd rather have one source
|
|
||||||
# of truth than reconcile clocks.
|
|
||||||
await touch_token(session, token_id=token_id)
|
await touch_token(session, token_id=token_id)
|
||||||
_log.debug("flushed %d token touch(es)", len(pending))
|
_log.debug("flushed %d token touch(es)", len(pending))
|
||||||
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"VALID_SCOPES",
|
"BUILTIN_SCOPES",
|
||||||
"TokenIdentity",
|
"TokenIdentity",
|
||||||
"TokenStore",
|
"TokenStore",
|
||||||
"TokenStoreError",
|
"TokenStoreError",
|
||||||
"hash_token",
|
"hash_token",
|
||||||
|
"scopes_with",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
Loaded 100 of 212 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user