diff --git a/src/beaver_gateway/agents/claude.py b/src/beaver_gateway/agents/claude.py index 4b04c45..38f969f 100644 --- a/src/beaver_gateway/agents/claude.py +++ b/src/beaver_gateway/agents/claude.py @@ -1,8 +1,9 @@ """Claude agent definition, backed by the Claude Agent SDK. The system prompt is either ``system_prompt`` verbatim or, when -``prompt_sources`` is set, the concatenation of those files assembled at -every session spawn (see ``core/prompt.py``). ``skill_sets`` are +``prompt_sources`` is set, the concatenation of those files (or +``(tag, file)`` pairs) assembled at every session spawn (see +``core/prompt.py``). ``skill_sets`` are directories of ``/SKILL.md`` folders; each becomes a local SDK plugin. Nothing from disk is loaded otherwise: the adapter runs with ``setting_sources=[]``. @@ -16,6 +17,7 @@ from pathlib import Path # noqa: TC003 - pydantic runtime from pydantic import BaseModel, ConfigDict, Field from beaver_gateway.agents.base import BaseAgent +from beaver_gateway.core.prompt import PromptSource # noqa: TC001 - pydantic runtime __all__ = ["ClaudeAgent", "ClaudeOptions"] @@ -49,6 +51,6 @@ class ClaudeOptions(BaseModel): class ClaudeAgent(BaseAgent): cwd: Path system_prompt: str = "" - prompt_sources: tuple[Path, ...] = () + prompt_sources: tuple[PromptSource, ...] = () skill_sets: tuple[Path, ...] = () options: ClaudeOptions = Field(default_factory=ClaudeOptions) diff --git a/src/beaver_gateway/backends/claude_sdk.py b/src/beaver_gateway/backends/claude_sdk.py index 64b2fa8..4589dcf 100644 --- a/src/beaver_gateway/backends/claude_sdk.py +++ b/src/beaver_gateway/backends/claude_sdk.py @@ -13,8 +13,9 @@ Events on the wire are the Anthropic ``MessageStreamEvent`` family: one 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, under that uid. +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 @@ -204,7 +205,7 @@ class ClaudeSdkBackend: self._mcp_disallowed = _mcp_disallowed(agent, mcp_tool_names or {}) self._sessions: dict[str, _Live] = {} self._reaper: asyncio.Task[None] | None = None - self._uid = _resolve_uid(self._runner.user) + self._uid, self._gid = _resolve_ids(self._runner.user) self._wrapper: Path | None = None @property @@ -411,6 +412,7 @@ class ClaudeSdkBackend: env = dict(opt.env) if self._runner.home is not None: env["HOME"] = str(self._runner.home) + env.setdefault("CLAUDE_CONFIG_DIR", str(self._runner.home / ".claude")) plugins = self._plugins() system_prompt = ( prompt_assembly.assemble(agent.prompt_sources) @@ -430,7 +432,6 @@ class ClaudeSdkBackend: cwd=str(agent.cwd), add_dirs=list(opt.add_dirs), env=env, - user=self._runner.user, cli_path=str(self._exec_wrapper(extra_keep=tuple(env))), include_partial_messages=opt.include_partial_messages, session_store=self._store, @@ -472,6 +473,8 @@ class ClaudeSdkBackend: target=json.dumps(target), keep=json.dumps(keep), prefixes=json.dumps(list(ENV_KEEP_PREFIXES)), + uid=json.dumps(self._uid), + gid=json.dumps(self._gid), ) digest = hashlib.sha256(script.encode("utf-8")).hexdigest()[:12] path = self._work_dir / f"claude-exec-{digest}.py" @@ -525,9 +528,15 @@ import sys TARGET = {target} KEEP = set({keep}) PREFIXES = tuple({prefixes}) +UID = {uid} +GID = {gid} env = {{ k: v for k, v in os.environ.items() if k in KEEP or k.startswith(PREFIXES) }} +if UID is not None and os.getuid() != UID: + os.setgroups([]) + os.setgid(GID) + os.setuid(UID) os.execve(TARGET, [TARGET, *sys.argv[1:]], env) """ @@ -550,12 +559,11 @@ def _claude_binary() -> str: return found -def _resolve_uid(user: str | None) -> int | None: +def _resolve_ids(user: str | None) -> tuple[int | None, int | None]: if user is None: - return None - if user.isdigit(): - return int(user) - return pwd.getpwnam(user).pw_uid + return None, None + record = pwd.getpwuid(int(user)) if user.isdigit() else pwd.getpwnam(user) + return record.pw_uid, record.pw_gid def _chown_tree(root: Path, uid: int) -> None: diff --git a/src/beaver_gateway/core/prompt.py b/src/beaver_gateway/core/prompt.py index 3b9a82c..a81193b 100644 --- a/src/beaver_gateway/core/prompt.py +++ b/src/beaver_gateway/core/prompt.py @@ -3,7 +3,9 @@ The gateway holds no prompt text: an agent names its granules (paths from ``config.py``) and :func:`assemble` concatenates them in that order, so the result is byte-for-byte identical for every session of the same agent as -long as the files are. Each granule's hash is logged at assembly so a +long as the files are. A source is a path, or a ``(tag, path)`` pair whose +content is wrapped in ``...`` - 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. """ @@ -17,18 +19,24 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Iterable -__all__ = ["assemble"] +__all__ = ["PromptSource", "assemble"] _log = logging.getLogger("beaver_gateway.core.prompt") +PromptSource = str | Path | tuple[str, str | Path] -def assemble(sources: Iterable[str | Path]) -> str: + +def assemble(sources: Iterable[PromptSource]) -> str: parts: list[str] = [] for source in sources: - path = Path(source) + tag, raw = source if isinstance(source, tuple) else (None, source) + path = Path(raw) text = path.read_text(encoding="utf-8").strip() digest = hashlib.sha256(text.encode("utf-8")).hexdigest()[:12] - _log.info("prompt granule %s sha=%s bytes=%d", path, digest, len(text)) - if text: - parts.append(text) + _log.info( + "prompt granule %s tag=%s sha=%s bytes=%d", path, tag, digest, len(text) + ) + if not text: + continue + parts.append(f"<{tag}>\n{text}\n" if tag else text) return "\n\n".join(parts) + "\n" diff --git a/src/beaver_gateway/storage/db.py b/src/beaver_gateway/storage/db.py index 1214274..331f22e 100644 --- a/src/beaver_gateway/storage/db.py +++ b/src/beaver_gateway/storage/db.py @@ -18,6 +18,7 @@ import json from datetime import UTC, datetime from typing import TYPE_CHECKING, Any +from sqlalchemy import inspect, text from sqlalchemy.ext.asyncio import create_async_engine from sqlmodel import SQLModel, select from sqlmodel.ext.asyncio.session import AsyncSession @@ -27,6 +28,7 @@ from beaver_gateway.storage.models import AuditLog, Token, Usage if TYPE_CHECKING: from collections.abc import Sequence + from sqlalchemy import Connection from sqlalchemy.ext.asyncio import AsyncEngine @@ -72,9 +74,14 @@ class Database: self._engine: AsyncEngine = create_async_engine(_to_async_url(url), echo=False) async def create_all(self) -> None: - """Idempotent ``CREATE TABLE IF NOT EXISTS`` for every model.""" + """Create missing tables, then add columns models gained since. + + Nullable columns only - that is the whole migration story until + Alembic is worth it. + """ async with self._engine.begin() as conn: await conn.run_sync(SQLModel.metadata.create_all) + await conn.run_sync(_add_missing_columns) def session(self) -> AsyncSession: """Open a fresh :class:`AsyncSession` (use as ``async with``). @@ -90,6 +97,19 @@ class Database: await self._engine.dispose() +def _add_missing_columns(conn: Connection) -> None: + inspector = inspect(conn) + for table in SQLModel.metadata.sorted_tables: + existing = {c["name"] for c in inspector.get_columns(table.name)} + for column in table.columns: + if column.name in existing: + continue + kind = column.type.compile(conn.dialect) + conn.execute( + text(f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}") + ) + + # ---- Token CRUD --------------------------------------------------------- diff --git a/tests/test_claude_sdk_backend.py b/tests/test_claude_sdk_backend.py index cda7100..1744d39 100644 --- a/tests/test_claude_sdk_backend.py +++ b/tests/test_claude_sdk_backend.py @@ -25,7 +25,12 @@ from claude_agent_sdk import ( from beaver_gateway.agents.base import ExposedMcp from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions -from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend, UsageEvent, fingerprint +from beaver_gateway.backends.claude_sdk import ( + ClaudeSdkBackend, + RunnerConfig, + UsageEvent, + fingerprint, +) from beaver_gateway.core.transcript import messages_from_entries from beaver_gateway.core.turn_capture import TurnCapture @@ -359,7 +364,9 @@ async def test_prompt_sources_are_assembled(cwd: Path) -> None: (cwd / "a.md").write_text("alpha\n") (cwd / "b.md").write_text("\nbeta\n\n") backend = _backend( - cwd, InMemorySessionStore(), prompt_sources=(cwd / "a.md", cwd / "b.md") + cwd, + InMemorySessionStore(), + prompt_sources=(("role", cwd / "a.md"), cwd / "b.md"), ) await _drain( backend.complete( @@ -368,7 +375,33 @@ async def test_prompt_sources_are_assembled(cwd: Path) -> None: conversation_id="c", ) ) - assert FakeClient.instances[0].options.system_prompt == "alpha\n\nbeta\n" + assert ( + FakeClient.instances[0].options.system_prompt + == "\nalpha\n\n\nbeta\n" + ) + + +async def test_runner_user_lands_in_wrapper(cwd: Path) -> None: + import os + import pwd + + me = pwd.getpwuid(os.getuid()) + backend = _backend(cwd, InMemorySessionStore()) + backend._runner = RunnerConfig(user=me.pw_name, home=cwd) + backend._uid, backend._gid = me.pw_uid, me.pw_gid + await _drain( + backend.complete( + agent=backend.agent, + messages=[{"role": "user", "content": "x"}], + conversation_id="c", + ) + ) + opts = FakeClient.instances[0].options + assert opts.env["HOME"] == str(cwd) + assert opts.env["CLAUDE_CONFIG_DIR"] == str(cwd / ".claude") + wrapper = Path(opts.cli_path).read_text() + assert f"UID = {me.pw_uid}" in wrapper + assert "os.setuid(UID)" in wrapper async def test_close_disconnects(cwd: Path) -> None: diff --git a/tests/test_db_migrate.py b/tests/test_db_migrate.py new file mode 100644 index 0000000..99580ab --- /dev/null +++ b/tests/test_db_migrate.py @@ -0,0 +1,33 @@ +import sqlite3 +import tempfile +from pathlib import Path + +from sqlmodel import select + +from beaver_gateway.storage import Database +from beaver_gateway.storage.models import Conversation + + +async def test_create_all_adds_missing_columns() -> None: + path = Path(tempfile.mkdtemp(prefix="beaver-migrate-")) / "old.db" + raw = sqlite3.connect(path) + raw.execute( + "CREATE TABLE conversations (id INTEGER PRIMARY KEY, frontend VARCHAR NOT NULL, " + "external_id VARCHAR NOT NULL, agent_name VARCHAR NOT NULL, " + "created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)" + ) + raw.execute( + "INSERT INTO conversations VALUES (1, 'markdown', 'x', 'a', '2026-01-01', '2026-01-01')" + ) + raw.commit() + raw.close() + + db = Database(f"sqlite:///{path}") + await db.create_all() + async with db.session() as session: + conv = (await session.exec(select(Conversation))).one() + assert conv.session_id is None + conv.session_id = "sid" + session.add(conv) + await session.commit() + await db.dispose()