diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4542ea53..79ffd76f 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -70,6 +70,7 @@ jobs: tests/db/test_pg_sections.py \ tests/db/test_pg_auth_queries.py \ tests/db/test_pg_integrations.py \ + tests/db/test_memory_context_schema.py \ tests/api/ \ -v diff --git a/Dockerfile b/Dockerfile index 68f5e33c..a1e3cdc1 100644 --- a/Dockerfile +++ b/Dockerfile @@ -20,7 +20,7 @@ FROM python:3.11-slim # cron is needed for anchor trigger scheduling in the bot container. # curl is needed for the NodeSource setup script. RUN apt-get update && apt-get install -y --no-install-recommends \ - gcc libffi-dev cron git gosu curl \ + gcc libffi-dev cron git gosu curl redis-server \ && rm -rf /var/lib/apt/lists/* # Install Node.js 20 via NodeSource so npm has a properly self-contained @@ -34,6 +34,19 @@ RUN curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \ ARG CLAUDE_CODE_VERSION=2.1.116 RUN npm install -g @anthropic-ai/claude-code@${CLAUDE_CODE_VERSION} +# Home directory template for the agent pool manager. +# Running `claude --version` once causes the CLI to write its initial config +# (~/.claude.json, ~/.claude/) into the temp home. We snapshot that state so +# each warm subprocess gets a pre-seeded home dir, eliminating first-run setup +# time and ~/.claude.json lock contention on concurrent spawns. +# If the CLI produces no files (e.g. newer version skips auto-init), the +# template dir is left empty — the isolation benefit still applies. +RUN HOME=/tmp/claude-template \ + claude --version 2>/dev/null || true \ + && mkdir -p /etc/claude-home-template \ + && cp -r /tmp/claude-template/. /etc/claude-home-template/ 2>/dev/null || true \ + && chmod -R a+rX /etc/claude-home-template + # Non-root user for running services. UID 1000 matches the default Pi user # so bind-mounted /data files are accessible without permission issues. RUN useradd -m -u 1000 -s /bin/bash tether @@ -45,6 +58,17 @@ RUN mkdir -p /run/tether/creds \ && chown tether:tether /run/tether/creds \ && chmod 0700 /run/tether/creds +# Agent pool manager home directory base. +# /var/lib/ is root-owned — the tether user cannot mkdir here at runtime, +# so the pool base dir must be pre-created and handed off in the image. +# Also hand ownership of the Claude home template to tether so that +# initialize() can write .claude.json into it at first boot (auth seeding). +RUN mkdir -p /var/lib/tether/claude-homes \ + && chown tether:tether /var/lib/tether/claude-homes \ + && chmod 755 /var/lib/tether/claude-homes \ + && chown tether:tether /etc/claude-home-template \ + && chmod 755 /etc/claude-home-template + WORKDIR /app # External dependencies (cached layer — only reruns when requirements.txt changes) diff --git a/agent_pool_manager/config.py b/agent_pool_manager/config.py index f23a86dc..1294c9a8 100644 --- a/agent_pool_manager/config.py +++ b/agent_pool_manager/config.py @@ -39,6 +39,23 @@ class AgentPoolConfig: control_response_timeout_seconds: float = 60.0 """Seconds the pool waits for a control_response before denying the tool call.""" + connect_timeout_seconds: float = 15.0 + """Max seconds to wait for ClaudeSDKClient.connect() before treating the spawn + as failed. Must be shorter than prime_timeout_seconds so a hung transport + is detected and killed before the priming phase would time out anyway.""" + + initialize_timeout_ms: int = 120000 + """Value injected as CLAUDE_CODE_STREAM_CLOSE_TIMEOUT (ms) in subprocess env. + The SDK uses this as its internal initialize timeout (floor: 60 000 ms). + Default: 120 000 ms (2 min).""" + + home_dir_base: str = "/var/lib/tether/claude-homes" + """Directory where isolated per-subprocess home dirs are created.""" + + home_dir_template: str = "/etc/claude-home-template" + """Directory copied into each home dir at initialization. + If it does not exist, home dirs start empty (lock-contention fix still applies).""" + _FIELD_NAMES = {f.name for f in fields(AgentPoolConfig)} diff --git a/agent_pool_manager/homes.py b/agent_pool_manager/homes.py new file mode 100644 index 00000000..f4704c49 --- /dev/null +++ b/agent_pool_manager/homes.py @@ -0,0 +1,172 @@ +"""Isolated per-subprocess home directory pool. + +Each subprocess gets its own ``$HOME`` to eliminate ``~/.claude.json`` lock +contention when two subprocesses spawn simultaneously. Dirs are pre-created +at pool startup and reset to a template state before each assignment so stale +CLI state from a previous subprocess never bleeds into the next one. + +Design notes: +- Dirs are never deleted during normal operation — only reset. +- Reset (``shutil.rmtree`` + ``shutil.copytree``) runs in a thread-pool + executor so it does not block the asyncio event loop during spawn. +- TTL sweep evicts stale assignments (e.g. a subprocess that was lost + without going through ``_terminate``). +""" +from __future__ import annotations + +import asyncio +import logging +import shutil +import time +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from agent_pool_manager.config import AgentPoolConfig + +log = logging.getLogger(__name__) + +# Path to the running container user's .claude.json. Isolated to a constant +# so tests can monkeypatch without touching the filesystem at /home/tether. +_AUTH_SOURCE_PATH: Path = Path("/home/tether/.claude.json") + + +def _reset_home(home: Path, template: Path | None) -> None: + """Synchronous reset called from a thread executor. + + Removes all content in *home* then either copies *template* in (if it + exists) or leaves *home* as an empty directory. + """ + # Wipe and recreate so the dir is empty/fresh regardless of template. + shutil.rmtree(home, ignore_errors=True) + home.mkdir(parents=True, exist_ok=True) + if template is not None and template.is_dir(): + # copytree requires the destination not to exist; use dirs_exist_ok + # (Python 3.8+) so we can copy *into* the already-created home dir. + shutil.copytree(str(template), str(home), dirs_exist_ok=True) + + +class HomeDirPool: + """Pool of isolated home directories for agent subprocesses. + + Call ``await initialize()`` once before first use. + Call ``await acquire()`` to obtain a home dir for a new subprocess. + Call ``release(path)`` when the subprocess terminates. + Call ``await sweep()`` periodically (from RefillLoop.run_once) to evict + stale assignments that were never released. + """ + + def __init__(self, config: AgentPoolConfig) -> None: + self._config = config + self._base = Path(config.home_dir_base) + template_path = Path(config.home_dir_template) + self._template: Path | None = template_path if template_path.is_dir() else None + self._dirs: list[Path] = [] + self._available: asyncio.Queue[Path] = asyncio.Queue() + # str(path) → expiry timestamp (time.monotonic) + self._checked_out: dict[str, float] = {} + + async def initialize(self) -> None: + """Create home dirs and seed from template. Idempotent on re-call.""" + count = self._config.capacity_total + self._base.mkdir(parents=True, exist_ok=True) + + # Re-discover template now (it may have been created since __init__) + template_path = Path(self._config.home_dir_template) + self._template = template_path if template_path.is_dir() else None + + # Seed template with auth credentials at runtime if not already present. + # /home/tether/.claude.json is written by `claude setup-token` after + # the container first boots — it is never baked into the image (Fly + # secret). Copying it into the template here means every warm + # subprocess starts with valid credentials without an extra setup step. + # Skips gracefully when: template dir absent, source absent, or + # template already has its own .claude.json. + if self._template is not None: + template_claude_json = self._template / ".claude.json" + if not template_claude_json.exists() and _AUTH_SOURCE_PATH.exists(): + try: + shutil.copy2(_AUTH_SOURCE_PATH, template_claude_json) + log.info( + "home_pool.seeded_template source=%s dest=%s", + _AUTH_SOURCE_PATH, + template_claude_json, + ) + except OSError: + log.warning( + "home_pool.seed_failed source=%s dest=%s", + _AUTH_SOURCE_PATH, + template_claude_json, + exc_info=True, + ) + + for i in range(count): + home = self._base / f"home-{i}" + home.mkdir(parents=True, exist_ok=True) + # Skip reset for dirs currently checked out by a live subprocess — + # wiping them would corrupt the running CLI process's HOME. + if str(home) not in self._checked_out: + # Seed from template in executor to avoid blocking the loop. + await asyncio.to_thread(_reset_home, home, self._template) + if home not in self._dirs: + self._dirs.append(home) + # Populate the available queue; clear first to avoid duplication on re-init. + + # Rebuild _available from dirs not currently checked out. + while not self._available.empty(): + try: + self._available.get_nowait() + except asyncio.QueueEmpty: + break + checked_out_paths = set(self._checked_out.keys()) + for home in self._dirs: + if str(home) not in checked_out_paths: + self._available.put_nowait(home) + + log.info( + "home_pool.initialized count=%d template=%s available=%d", + count, + self._template, + self._available.qsize(), + ) + + async def acquire(self) -> Path: + """Return a home dir, reset to template state. + + Blocks (up to the caller's timeout) if the pool is exhausted. + Raises ``asyncio.QueueEmpty`` if called with nowait and no dirs are + available — callers should use ``asyncio.wait_for`` for timeouts. + """ + home = await self._available.get() + await asyncio.to_thread(_reset_home, home, self._template) + # Use 2× subprocess max_age as home TTL so sweep never evicts a home + # whose subprocess might still be running (drain-on-touch retirement + # means a subprocess can live past its nominal max_age). + expiry = time.monotonic() + self._config.max_age_seconds * 2 + self._checked_out[str(home)] = expiry + log.debug("home_pool.acquired path=%s", home) + return home + + def release(self, home: Path) -> None: + """Return *home* to the available queue. Idempotent.""" + key = str(home) + if key not in self._checked_out: + # Already released or never acquired — silently ignore. + return + self._checked_out.pop(key) + self._available.put_nowait(home) + log.debug("home_pool.released path=%s", home) + + async def sweep(self) -> None: + """Evict stale assignments that were never released (TTL-based safety net).""" + now = time.monotonic() + evicted = [k for k, expiry in list(self._checked_out.items()) if now > expiry] + for key in evicted: + self._checked_out.pop(key, None) + path = Path(key) + self._available.put_nowait(path) + log.warning("home_pool.sweep_evict path=%s", path) + + def available_count(self) -> int: + """Number of home dirs currently available (not checked out).""" + return self._available.qsize() diff --git a/agent_pool_manager/metrics.py b/agent_pool_manager/metrics.py new file mode 100644 index 00000000..59375e39 --- /dev/null +++ b/agent_pool_manager/metrics.py @@ -0,0 +1,192 @@ +"""Lightweight Prometheus-format metrics registry for the agent pool manager. + +No external dependency — renders the standard Prometheus text exposition format +(https://prometheus.io/docs/instrumenting/exposition_formats/). + +Intended usage: + metrics = PoolMetrics() + metrics.acquire_total.inc() + metrics.acquire_latency_seconds.observe(0.123) + text = metrics.render_text() # feed to GET /metrics +""" +from __future__ import annotations + +import math +import threading +import time +from dataclasses import dataclass, field +from typing import Callable + + +# --------------------------------------------------------------------------- +# Primitive metric types +# --------------------------------------------------------------------------- + +class Counter: + """Monotonically increasing counter.""" + + def __init__(self, name: str, help_text: str) -> None: + self.name = name + self.help_text = help_text + self._value: float = 0.0 + self._lock = threading.Lock() + + def inc(self, amount: float = 1.0) -> None: + with self._lock: + self._value += amount + + @property + def value(self) -> float: + return self._value + + def render_text(self) -> str: + return ( + f"# HELP {self.name} {self.help_text}\n" + f"# TYPE {self.name} counter\n" + f"{self.name} {self._value}\n" + ) + + +_DEFAULT_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0) + + +class Histogram: + """Prometheus-style histogram with configurable bucket boundaries.""" + + def __init__( + self, + name: str, + help_text: str, + buckets: tuple[float, ...] = _DEFAULT_BUCKETS, + ) -> None: + self.name = name + self.help_text = help_text + self._buckets = sorted(buckets) + self._counts: list[float] = [0.0] * len(self._buckets) + self._inf_count: float = 0.0 + self._sum: float = 0.0 + self._total_count: float = 0.0 + self._lock = threading.Lock() + + def observe(self, value: float) -> None: + with self._lock: + self._sum += value + self._total_count += 1.0 + self._inf_count += 1.0 + # _counts[i] stores the *cumulative* count of observations <= bound[i]. + # Every observation increments all buckets whose bound it fits within. + # render_text() emits _counts[i] directly — no re-accumulation needed. + for i, bound in enumerate(self._buckets): + if value <= bound: + self._counts[i] += 1.0 + + def render_text(self) -> str: + lines = [ + f"# HELP {self.name} {self.help_text}", + f"# TYPE {self.name} histogram", + ] + # _counts[i] is already the cumulative count for le=bound[i] — + # emit directly without re-accumulating (which would double-count). + for i, bound in enumerate(self._buckets): + lines.append(f'{self.name}_bucket{{le="{bound}"}} {self._counts[i]}') + lines.append(f'{self.name}_bucket{{le="+Inf"}} {self._inf_count}') + lines.append(f"{self.name}_sum {self._sum}") + lines.append(f"{self.name}_count {self._total_count}") + return "\n".join(lines) + "\n" + + +class Gauge: + """Point-in-time gauge, driven by a callable for dynamic values.""" + + def __init__(self, name: str, help_text: str, fn: Callable[[], float]) -> None: + self.name = name + self.help_text = help_text + self._fn = fn + + def render_text(self) -> str: + value = self._fn() + return ( + f"# HELP {self.name} {self.help_text}\n" + f"# TYPE {self.name} gauge\n" + f"{self.name} {value}\n" + ) + + +# --------------------------------------------------------------------------- +# PoolMetrics — the registry +# --------------------------------------------------------------------------- + +class PoolMetrics: + """All Prometheus metrics for the agent pool manager. + + Attach to a Pool instance via ``pool._metrics = metrics`` so the pool can + record events. The server renders via ``metrics.render_text()``. + """ + + def __init__(self, pool: "Pool | None" = None) -> None: # type: ignore[name-defined] + self._pool = pool + + self.acquire_total = Counter( + "pool_acquire_total", + "Total successful subprocess acquisitions", + ) + self.acquire_timeout_total = Counter( + "pool_acquire_timeout_total", + "Total acquire attempts that timed out (PoolExhausted)", + ) + self.release_total = Counter( + "pool_release_total", + "Total subprocess releases", + ) + self.expire_total = Counter( + "pool_expire_total", + "Total subprocesses drained due to TTL expiry", + ) + self.refill_total = Counter( + "pool_refill_total", + "Total subprocess spawns triggered by refill loop", + ) + self.spawn_guard_rejection_total = Counter( + "pool_spawn_guard_rejection_total", + "Total spawn attempts rejected by the auth guard (missing CLAUDE_CODE_OAUTH_TOKEN)", + ) + self.acquire_latency_seconds = Histogram( + "pool_acquire_latency_seconds", + "Time from acquire() call to handle returned (seconds)", + ) + + def attach_pool(self, pool: "Pool") -> None: # type: ignore[name-defined] + """Attach this metrics instance to a pool for gauge rendering.""" + self._pool = pool + + def render_text(self) -> str: + """Render all metrics in Prometheus text exposition format.""" + parts = [ + self.acquire_total.render_text(), + self.acquire_timeout_total.render_text(), + self.release_total.render_text(), + self.expire_total.render_text(), + self.refill_total.render_text(), + self.spawn_guard_rejection_total.render_text(), + self.acquire_latency_seconds.render_text(), + ] + + # Dynamic gauges from pool state (if attached) + if self._pool is not None: + status = self._pool.status() + parts += [ + _gauge_text("pool_size_warm", "Current warm subprocess count", status.get("total_warm", 0)), + _gauge_text("pool_size_active", "Current active (handed-out) subprocess count", status.get("total_active", 0)), + _gauge_text("pool_size_warming", "Current subprocesses being spawned/primed", status.get("total_warming", 0)), + _gauge_text("pool_capacity_total", "Configured total subprocess capacity", status.get("capacity_total", 0)), + ] + + return "".join(parts) + + +def _gauge_text(name: str, help_text: str, value: float) -> str: + return ( + f"# HELP {name} {help_text}\n" + f"# TYPE {name} gauge\n" + f"{name} {value}\n" + ) diff --git a/agent_pool_manager/pool.py b/agent_pool_manager/pool.py index ea3455a1..f63208a3 100644 --- a/agent_pool_manager/pool.py +++ b/agent_pool_manager/pool.py @@ -2,12 +2,14 @@ from __future__ import annotations import asyncio +import dataclasses import datetime import logging +import os import time import uuid from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any from claude_agent_sdk import ClaudeSDKClient from claude_agent_sdk.types import ( @@ -19,10 +21,129 @@ from .config import AgentPoolConfig from .control import ControlBridge, ControlTimeout +from .homes import HomeDirPool + +if TYPE_CHECKING: + from pathlib import Path + from .metrics import PoolMetrics log = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Diagnostic helpers — sensitive value redaction. +# +# The warm spawn path runs in production with the user's OAuth token in env. +# We need visibility into what's actually being passed to the subprocess, +# but must never log the raw token. ``_redact_env`` shows the key names +# and a short prefix of each value (8 chars) so we can confirm shape without +# leaking secrets. +# --------------------------------------------------------------------------- + +_SENSITIVE_ENV_KEYS = frozenset({ + "CLAUDE_CODE_OAUTH_TOKEN", + "ANTHROPIC_API_KEY", + "OPENAI_API_KEY", + "TETHER_JWT_SECRET", + "VAULT_KEY", +}) + + +def _redact_env(env: dict | None) -> dict: + """Return a redacted copy of an env dict — sensitive values become ````. + + Non-sensitive values are passed through unchanged. Used in diagnostic + logs so the env composition is visible without leaking secrets. + """ + if not env: + return {} + out: dict[str, str] = {} + for k, v in env.items(): + sval = str(v) if v is not None else "" + if k in _SENSITIVE_ENV_KEYS: + prefix = sval[:8] if len(sval) >= 8 else sval + out[k] = f"" + else: + out[k] = sval + return out + + +def _mcp_servers_form(mcp_servers: Any) -> str: + """Return a short string describing the form of an mcp_servers value. + + The Claude SDK accepts ``dict[str, McpServerConfig] | str | Path`` but + we historically pass a ``list[str]`` from ``_V2_0_OPTIONS``. The list + form falls through to ``str(value)`` in the SDK, which is suspected + to be a contributor to the 15 s warm-spawn hang. + """ + if mcp_servers is None: + return "None" + if isinstance(mcp_servers, dict): + return f"dict(keys={list(mcp_servers.keys())})" + if isinstance(mcp_servers, list): + return f"list({mcp_servers!r})" + if isinstance(mcp_servers, (str, bytes)): + return f"str(len={len(mcp_servers)})" + return f"other(type={type(mcp_servers).__name__})" + + +def _options_summary(options: dict[str, Any]) -> dict: + """Return a compact, redacted summary of the options dict for logging. + + Keys whose values are large or sensitive are summarised rather than + dumped in full — this keeps log lines readable in fly.io's log stream. + """ + return { + "model": options.get("model"), + "allowed_tools_count": len(options.get("allowed_tools", []) or []), + "max_turns": options.get("max_turns"), + "permission_mode": options.get("permission_mode"), + "mcp_servers_form": _mcp_servers_form(options.get("mcp_servers")), + "env_keys": sorted((options.get("env") or {}).keys()), + "env_redacted": _redact_env(options.get("env")), + "extra_keys": sorted( + k for k in options.keys() + if k not in {"model", "allowed_tools", "max_turns", "permission_mode", + "mcp_servers", "env"} + ), + } + + +# MCP server URL for the tether MCP service (supervisord, port 5001). +_MCP_TETHER_URL = "http://localhost:5001/sse" + + +def _expand_mcp_placeholders(options: dict[str, Any], mcp_key: str) -> dict[str, Any]: + """Expand the ``['tether']`` MCP placeholder into a real SSE config dict. + + The static ``_V2_0_OPTIONS`` in ``bot.agent_dispatch`` carries + ``mcp_servers=['tether']`` as a stable hash-stable placeholder. The SDK + expects ``dict[str, McpServerConfig]``; passing a list causes it to fall + through to ``str(value)`` which produces ``--mcp-config "['tether']"`` on + the CLI — an unparseable value that causes a 15 s connect hang. + + This function is called at ``_spawn_and_prime`` time (not at options-dict + creation time) so the hash computed from the placeholder stays stable + across the warm endpoint and dispatch_v2_0 callers. + + Returns a shallow-copied options dict with ``mcp_servers`` replaced. + Does NOT mutate the input dict. + """ + mcp_servers = options.get("mcp_servers") + if not (isinstance(mcp_servers, list) and "tether" in mcp_servers): + return options # already correct form or absent — no copy needed + + result = dict(options) + result["mcp_servers"] = { + "tether": { + "type": "sse", + "url": _MCP_TETHER_URL, + "headers": {"Authorization": f"Bearer {mcp_key}"}, + } + } + return result + + class PoolExhausted(Exception): """Raised when no warm subprocess becomes available within the timeout.""" @@ -60,6 +181,13 @@ class Subprocess: last_used_at: float = field(default_factory=time.monotonic) in_use: bool = False callback_ctx: _CallbackContext = field(default_factory=_CallbackContext) + # Ephemeral MCP api_key created at spawn time; revoked on _terminate. + # None when pool has no DB access or no user_id was available at spawn. + mcp_key_id: str | None = None + mcp_user_id: str | None = None + # Isolated home directory assigned at spawn time; released on _terminate. + # None when home pool is not configured. + home_path: "Path | None" = None def is_expired(self, max_age_seconds: int) -> bool: return (time.monotonic() - self.spawned_at) > max_age_seconds @@ -72,7 +200,7 @@ class Pool: (drain-on-touch), not via a background timer. """ - def __init__(self, config: AgentPoolConfig) -> None: + def __init__(self, config: AgentPoolConfig, *, pg_pool: Any = None) -> None: self.config = config # warm queue per options_hash self._warm: dict[str, asyncio.Queue[Subprocess]] = {} @@ -87,6 +215,31 @@ def __init__(self, config: AgentPoolConfig) -> None: self.control_bridge = ControlBridge( timeout_seconds=config.control_response_timeout_seconds ) + # optional metrics instance — attach via pool._metrics = metrics + self._metrics: "PoolMetrics | None" = None + # optional asyncpg pool for ephemeral MCP key creation/revocation + self._pg_pool: Any = pg_pool + # optional home directory pool — set via initialize_home_pool() + self._home_pool: HomeDirPool | None = None + + # The Python SDK reads CLAUDE_CODE_STREAM_CLOSE_TIMEOUT from the + # manager process os.environ (not from the subprocess env dict). + # Set it here so the SDK's _send_control_request timeout is honoured. + # Use setdefault so an operator-supplied env var takes precedence. + os.environ.setdefault( + "CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", + str(config.initialize_timeout_ms), + ) + + async def initialize_home_pool(self) -> None: + """Create and seed the home directory pool. + + Must be called once before _spawn_and_prime if home isolation is + desired. Safe to call multiple times (idempotent). + """ + if self._home_pool is None: + self._home_pool = HomeDirPool(self.config) + await self._home_pool.initialize() # ------------------------------------------------------------------ # Public API @@ -97,6 +250,7 @@ async def acquire( options_hash: str, options: dict[str, Any], timeout: float | None = None, + user_id: str | None = None, ) -> tuple[str, dict[str, Any]]: """Hand out a warm subprocess for the given options_hash. @@ -111,46 +265,94 @@ async def acquire( if timeout is None: timeout = self.config.acquire_default_timeout + t_start = time.monotonic() self._options_cache[options_hash] = options - deadline = time.monotonic() + timeout + deadline = t_start + timeout queue = self._get_or_create_queue(options_hash) exhausted = PoolExhausted( f"No warm subprocess for hash {options_hash!r} within {timeout}s" ) - while True: - remaining = deadline - time.monotonic() - if remaining <= 0: - raise exhausted + try: + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise exhausted - try: - sub = queue.get_nowait() - except asyncio.QueueEmpty: - # Wait for a warm item to appear, then re-check the deadline. try: - sub = await asyncio.wait_for(queue.get(), timeout=min(remaining, 0.1)) - except asyncio.TimeoutError: + sub = queue.get_nowait() + except asyncio.QueueEmpty: + # Wait for a warm item to appear, then re-check the deadline. + try: + sub = await asyncio.wait_for(queue.get(), timeout=min(remaining, 0.1)) + except asyncio.TimeoutError: + continue + + # Drain expired entries + if sub.is_expired(self.config.max_age_seconds): + latency_ms = (time.monotonic() - t_start) * 1000 + log.info( + "pool.expire options_hash=%s user_id=%s age_s=%.1f latency_ms=%.1f", + options_hash, user_id, + time.monotonic() - sub.spawned_at, latency_ms, + ) + if self._metrics: + self._metrics.expire_total.inc() + asyncio.create_task(self._terminate(sub)) continue - # Drain expired entries - if sub.is_expired(self.config.max_age_seconds): - log.debug("Draining expired subprocess for hash %s", options_hash) - asyncio.create_task(self._terminate(sub)) - continue - - # Hand it out - handle_id = str(uuid.uuid4()) - sub.in_use = True - sub.last_used_at = time.monotonic() - sub.callback_ctx.handle_id = handle_id - async with self._lock: - self._active[handle_id] = sub + # User isolation: the warm queue is keyed by options_hash, which is + # intentionally computed from the static placeholder so the hash + # stays stable across users. But each subprocess now carries an + # ephemeral MCP Bearer token minted for the user that triggered its + # warm spawn. A subprocess spawned for user A must not serve user B. + # If there is a user mismatch, terminate and try the next entry. + if ( + sub.mcp_user_id is not None + and user_id is not None + and sub.mcp_user_id != user_id + ): + log.info( + "pool.user_mismatch options_hash=%s sub_user=%s req_user=%s" + " — discarding subprocess to prevent cross-user MCP key leak", + options_hash, sub.mcp_user_id, user_id, + ) + asyncio.create_task(self._terminate(sub)) + continue - meta = { - "subprocess_pid": _extract_pid(sub.proc), - "ready_at": datetime.datetime.utcnow().isoformat() + "Z", - } - return handle_id, meta + # Hand it out + handle_id = str(uuid.uuid4()) + sub.in_use = True + sub.last_used_at = time.monotonic() + sub.callback_ctx.handle_id = handle_id + async with self._lock: + self._active[handle_id] = sub + + latency_s = time.monotonic() - t_start + latency_ms = latency_s * 1000 + log.info( + "pool.acquire handle_id=%s user_id=%s options_hash=%s latency_ms=%.1f", + handle_id, user_id, options_hash, latency_ms, + ) + if self._metrics: + self._metrics.acquire_total.inc() + self._metrics.acquire_latency_seconds.observe(latency_s) + + meta = { + "subprocess_pid": _extract_pid(sub.proc), + "ready_at": datetime.datetime.utcnow().isoformat() + "Z", + } + return handle_id, meta + + except PoolExhausted: + latency_ms = (time.monotonic() - t_start) * 1000 + log.info( + "pool.acquire_timeout options_hash=%s user_id=%s timeout_s=%.2f latency_ms=%.1f", + options_hash, user_id, timeout, latency_ms, + ) + if self._metrics: + self._metrics.acquire_timeout_total.inc() + raise async def release(self, handle_id: str, reusable: bool = False) -> None: """Release a handle. @@ -166,6 +368,13 @@ async def release(self, handle_id: str, reusable: bool = False) -> None: sub.in_use = False sub.last_used_at = time.monotonic() + log.info( + "pool.release handle_id=%s options_hash=%s reusable=%s", + handle_id, sub.options_hash, reusable, + ) + if self._metrics: + self._metrics.release_total.inc() + if reusable and not sub.is_expired(self.config.max_age_seconds): queue = self._get_or_create_queue(sub.options_hash) await queue.put(sub) @@ -225,27 +434,73 @@ def status(self) -> dict[str, Any]: # Internal helpers — also used by RefillLoop and tests # ------------------------------------------------------------------ - async def _inject_warm(self, options_hash: str, options: dict[str, Any]) -> None: + async def _inject_warm( + self, + options_hash: str, + options: dict[str, Any], + *, + user_id: str | None = None, + ) -> None: """Spawn, prime, and push one subprocess to the warm queue. Used directly by RefillLoop and by tests via FakeClient patch. Silently no-ops at capacity. """ - if not await self._try_inject_warm(options_hash, options): + if not await self._try_inject_warm(options_hash, options, user_id=user_id): log.debug("Capacity full — skipping inject for hash %s", options_hash) - async def _try_inject_warm(self, options_hash: str, options: dict[str, Any]) -> bool: + async def _try_inject_warm( + self, + options_hash: str, + options: dict[str, Any], + *, + user_id: str | None = None, + ) -> bool: """Attempt to spawn-and-prime one subprocess. - Returns True if spawned, False if at capacity. + Returns True if spawned, False if at capacity or if the spawn guard fires. + + Spawn guard: if options['env'] does not contain CLAUDE_CODE_OAUTH_TOKEN, + the spawn is rejected with a WARNING. Subprocesses launched without OAuth + credentials time out after connect_timeout_seconds (15 s) with no useful + output, waste pool capacity, and leave asyncio "Task exception was never + retrieved" errors in the logs. The guard fires before the warming counter + is incremented, so rejected attempts do not count against capacity. + + All known spawn paths (RefillLoop.hint, RefillLoop.run_once, _inject_warm) + converge here, so this single check covers the entire spawn surface. """ + log.info( + "pool.inject_warm_entry options_hash=%s warm=%d active=%d warming=%d capacity=%d", + options_hash, + self.warm_count(options_hash), + len(self._active), + self.warming_count(options_hash), + self.config.capacity_total, + ) + + token = (options.get("env") or {}).get("CLAUDE_CODE_OAUTH_TOKEN") + if not token: + log.warning( + "pool.spawn_guard: options_hash=%s missing CLAUDE_CODE_OAUTH_TOKEN in env" + " — skipping spawn to prevent auth-timeout waste", + options_hash, + ) + if self._metrics: + self._metrics.spawn_guard_rejection_total.inc() + return False + async with self._lock: if self.total_count() >= self.config.capacity_total: + log.info( + "pool.inject_warm_capacity_full options_hash=%s total=%d capacity=%d", + options_hash, self.total_count(), self.config.capacity_total, + ) return False self._warming[options_hash] = self._warming.get(options_hash, 0) + 1 try: - sub = await self._spawn_and_prime(options_hash, options) + sub = await self._spawn_and_prime(options_hash, options, user_id=user_id) finally: async with self._lock: self._warming[options_hash] = max( @@ -255,29 +510,229 @@ async def _try_inject_warm(self, options_hash: str, options: dict[str, Any]) -> queue = self._get_or_create_queue(options_hash) await queue.put(sub) self._options_cache[options_hash] = options - log.debug("Primed subprocess ready for hash %s", options_hash) + log.info("pool.refill options_hash=%s warm_depth=%d", options_hash, queue.qsize()) + if self._metrics: + self._metrics.refill_total.inc() return True async def _spawn_and_prime( - self, options_hash: str, options: dict[str, Any] + self, + options_hash: str, + options: dict[str, Any], + *, + user_id: str | None = None, ) -> Subprocess: - """Spawn a ClaudeSDKClient, connect, and send the priming prompt.""" + """Spawn a ClaudeSDKClient, connect, and send the priming prompt. + + Diagnostic logging: this method is the prime suspect for the 15 s + warm-spawn hang in prod. We log: + + * The options summary (redacted) before constructing sdk_options + * Timing checkpoints around ``client.connect()`` and priming + * The full exception details on connect failure (type, message, + time spent, subprocess PID if any) + * Subprocess stderr lines piped via the SDK's ``stderr`` callback, + so the CLI's own error output is visible in fly.io logs + + Without these, ``connect()`` can hang for 15 s with no observable + cause from the application side. + """ + log.info( + "pool.spawn_start options_hash=%s summary=%r", + options_hash, + _options_summary(options), + ) + ctx = _CallbackContext() + + # Ephemeral MCP key injection — expand ['tether'] placeholder into a + # real SSE config dict with a per-spawn Bearer token. The key is + # created here (not at options-dict creation time) so the options hash + # computed from the placeholder stays stable across warm-endpoint and + # dispatch_v2_0 callers. + mcp_key_id: str | None = None + if self._pg_pool is not None and user_id: # truthy: excludes None and "" + try: + import db.postgres as pg + from db.pg_queries.api_keys import create_key as _create_key + async with pg.get_conn(self._pg_pool, user_id=user_id) as _conn: + _raw_key, _key_rec = await _create_key( + _conn, user_id=user_id, name=f"pool_mcp_{options_hash[:8]}" + ) + mcp_key_id = _key_rec["id"] + options = _expand_mcp_placeholders(options, _raw_key) + log.info( + "pool.mcp_key_created options_hash=%s key_id=%s", + options_hash, mcp_key_id, + ) + except Exception: + log.warning( + "pool.mcp_key_create_failed options_hash=%s" + " — spawning without MCP auth injection", + options_hash, + exc_info=True, + ) + # Strip the ['tether'] placeholder so the subprocess doesn't hang + # waiting for a tether MCP server it cannot authenticate with. + # An empty dict tells the SDK "no MCP servers" — clean start. + options = dict(options) + options["mcp_servers"] = {} + + # Inject initialize timeout — tells the SDK how long to wait before + # treating a connect() call as failed. Use setdefault so a caller- + # supplied value is not overridden. + env = dict(options.get("env") or {}) + env.setdefault( + "CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", + str(self.config.initialize_timeout_ms), + ) + options = dict(options) + options["env"] = env + + # Assign an isolated home directory if the home pool is available. + # This eliminates ~/.claude.json file-lock contention when multiple + # subprocesses spawn simultaneously. + home_path: "Path | None" = None + if self._home_pool is not None: + try: + home_path = await asyncio.wait_for( + self._home_pool.acquire(), + timeout=self.config.connect_timeout_seconds, + ) + except asyncio.TimeoutError: + log.warning( + "pool.home_acquire_timeout options_hash=%s timeout_s=%.1f" + " — spawning without isolated home dir", + options_hash, + self.config.connect_timeout_seconds, + ) + if home_path is not None: + env = dict(options["env"]) + env["HOME"] = str(home_path) + options["env"] = env + + # Wire subprocess stderr to our logger so CLI errors surface in + # fly.io's log stream. Without this, stderr is dropped on the floor. + def _stderr_cb(line: str) -> None: + log.info( + "pool.subprocess_stderr options_hash=%s line=%s", + options_hash, line.rstrip(), + ) + + options_with_stderr = dict(options) + # Don't overwrite a caller-supplied stderr callback. + options_with_stderr.setdefault("stderr", _stderr_cb) + + # Turn on the CLI's debug-to-stderr mode so the SDK protocol traces + # also flow through our stderr callback. Without this we only see + # whatever the CLI writes to stderr on its own (which is typically + # nothing on a successful run and very little even on errors). + # This is the single most useful flag for diagnosing a connect() hang. + existing_extra = dict(options_with_stderr.get("extra_args") or {}) + existing_extra.setdefault("debug-to-stderr", None) + options_with_stderr["extra_args"] = existing_extra + sdk_options = self._build_sdk_options( - options, + options_with_stderr, can_use_tool=self._make_forwarding_callback(ctx), ) client = ClaudeSDKClient(options=sdk_options) - await client.connect() - # Prime: send a cheap prompt so the subprocess pays its init cost now + # If we created an ephemeral MCP key and the spawn fails (connect + # timeout, OAuth failure, prime error, etc.) the Subprocess dataclass + # is never constructed so _terminate is never called. Guard here to + # ensure the key is revoked on any failure after creation. try: - await asyncio.wait_for( - self._do_prime(client), - timeout=self.config.prime_timeout_seconds, + t_connect_start = time.monotonic() + try: + await asyncio.wait_for( + client.connect(), + timeout=self.config.connect_timeout_seconds, + ) + except Exception as exc: + elapsed = time.monotonic() - t_connect_start + pid = _extract_pid(client) + log.warning( + "pool.connect_failed options_hash=%s exc_type=%s msg=%s" + " elapsed_s=%.2f pid=%s timeout_s=%.1f", + options_hash, + type(exc).__name__, + str(exc) or "", + elapsed, + pid, + self.config.connect_timeout_seconds, + ) + # Kill the underlying subprocess so a failed spawn doesn't leave a + # zombie Claude CLI process holding memory until the OS cleans it up. + # This covers both auth failures (70 s timeout) and other errors. + try: + transport = getattr(client, "_transport", None) + proc = getattr(transport, "_process", None) if transport is not None else None + if proc is not None: + proc.kill() + # asyncio.subprocess.Process.wait() is a coroutine; a + # synchronous subprocess.Popen.wait() is not. Guard against + # the sync case so a future SDK transport change doesn't + # silently suppress the TypeError and leave zombies. + if asyncio.iscoroutinefunction(getattr(proc, "wait", None)): + await proc.wait() + else: + log.warning( + "pool.cleanup: proc.wait() is not a coroutine for hash=%s" + " — skipping await; subprocess may remain as zombie", + options_hash, + ) + except Exception: + pass + raise + + connect_elapsed = time.monotonic() - t_connect_start + log.info( + "pool.connect_done options_hash=%s elapsed_s=%.2f pid=%s", + options_hash, connect_elapsed, _extract_pid(client), ) - except asyncio.TimeoutError: - log.warning("Priming timed out for hash %s — keeping unprimed client", options_hash) + + # Prime: send a cheap prompt so the subprocess pays its init cost now + t_prime_start = time.monotonic() + try: + await asyncio.wait_for( + self._do_prime(client), + timeout=self.config.prime_timeout_seconds, + ) + log.info( + "pool.prime_done options_hash=%s elapsed_s=%.2f", + options_hash, time.monotonic() - t_prime_start, + ) + except asyncio.TimeoutError: + log.warning( + "pool.prime_timeout options_hash=%s elapsed_s=%.2f timeout_s=%d" + " — keeping unprimed client", + options_hash, + time.monotonic() - t_prime_start, + self.config.prime_timeout_seconds, + ) + + except Exception: + # Spawn failed after key was created — revoke key to prevent orphan rows. + if mcp_key_id is not None and self._pg_pool is not None and user_id: + try: + import db.postgres as pg + from db.pg_queries.api_keys import revoke_key as _revoke_key + async with pg.get_conn(self._pg_pool, user_id=user_id) as _conn: + await _revoke_key(_conn, mcp_key_id, user_id) + log.info( + "pool.mcp_key_revoked_on_spawn_failure options_hash=%s key_id=%s", + options_hash, mcp_key_id, + ) + except Exception: + log.warning( + "pool.mcp_key_revoke_on_spawn_failure_failed key_id=%s", + mcp_key_id, + exc_info=True, + ) + if home_path is not None and self._home_pool is not None: + self._home_pool.release(home_path) + raise now = time.monotonic() return Subprocess( @@ -287,6 +742,9 @@ async def _spawn_and_prime( spawned_at=now, primed_at=now, callback_ctx=ctx, + mcp_key_id=mcp_key_id, + mcp_user_id=user_id, + home_path=home_path, ) @staticmethod @@ -333,11 +791,32 @@ def _build_sdk_options( options: dict[str, Any], can_use_tool: Any = None, ) -> ClaudeAgentOptions: - """Convert the options dict to ClaudeAgentOptions, ignoring unknown keys.""" - known = {f for f in dir(ClaudeAgentOptions) if not f.startswith("_")} + """Convert the options dict to ClaudeAgentOptions, ignoring unknown keys. + + Diagnostic logging: emit which keys were passed through to the SDK + and which were filtered out. This helps explain unexpected SDK + behaviour when callers pass dict-shaped options that don't map 1:1 + to ``ClaudeAgentOptions`` fields. + + Field enumeration uses ``dataclasses.fields()``, NOT ``dir()`` — + ``dir()`` on a dataclass class omits fields declared with + ``default_factory`` (env, mcp_servers, allowed_tools, extra_args, + plugins, add_dirs, betas, disallowed_tools). Using ``dir()`` here + silently dropped the user's OAuth token from the subprocess env, + which was the root cause of the 15 s warm-spawn timeout in prod. + """ + known = {f.name for f in dataclasses.fields(ClaudeAgentOptions)} filtered = {k: v for k, v in options.items() if k in known} + dropped = sorted(k for k in options.keys() if k not in known) if can_use_tool is not None: filtered["can_use_tool"] = can_use_tool + + log.info( + "pool.sdk_options known_keys=%s dropped_keys=%s mcp_servers_form=%s", + sorted(filtered.keys()), + dropped, + _mcp_servers_form(filtered.get("mcp_servers")), + ) return ClaudeAgentOptions(**filtered) def _get_or_create_queue(self, options_hash: str) -> asyncio.Queue[Subprocess]: @@ -345,9 +824,35 @@ def _get_or_create_queue(self, options_hash: str) -> asyncio.Queue[Subprocess]: self._warm[options_hash] = asyncio.Queue() return self._warm[options_hash] - @staticmethod - async def _terminate(sub: Subprocess) -> None: + async def _terminate(self, sub: Subprocess) -> None: + # Disconnect first — this guarantees no further MCP calls will be made + # by the subprocess before we invalidate its credential. Revoking + # before disconnect would leave a window where in-flight MCP requests + # receive 401 responses. try: await sub.proc.disconnect() except Exception: pass + + # Revoke ephemeral MCP key after the subprocess is disconnected. + # Failure must not propagate — log and continue. + if sub.mcp_key_id is not None and self._pg_pool is not None and sub.mcp_user_id: + try: + import db.postgres as pg + from db.pg_queries.api_keys import revoke_key as _revoke_key + async with pg.get_conn(self._pg_pool, user_id=sub.mcp_user_id) as _conn: + await _revoke_key(_conn, sub.mcp_key_id, sub.mcp_user_id) + log.info( + "pool.mcp_key_revoked options_hash=%s key_id=%s", + sub.options_hash, sub.mcp_key_id, + ) + except Exception: + log.warning( + "pool.mcp_key_revoke_failed key_id=%s", + sub.mcp_key_id, + exc_info=True, + ) + + # Return the isolated home directory to the pool. + if sub.home_path is not None and self._home_pool is not None: + self._home_pool.release(sub.home_path) diff --git a/agent_pool_manager/refill.py b/agent_pool_manager/refill.py index 2a389743..a94651e8 100644 --- a/agent_pool_manager/refill.py +++ b/agent_pool_manager/refill.py @@ -21,38 +21,87 @@ def __init__(self, pool: Pool) -> None: self._pool = pool # options_hash → options dict self._registry: dict[str, dict[str, Any]] = {} + # options_hash → user_id (stored alongside options for key injection) + self._user_ids: dict[str, str] = {} self._hint_event = asyncio.Event() self._running = False self._task: asyncio.Task | None = None - def register(self, options_hash: str, options: dict[str, Any]) -> None: + def register( + self, + options_hash: str, + options: dict[str, Any], + *, + user_id: str | None = None, + ) -> None: """Register a hash for periodic refill. Idempotent.""" + already_known = options_hash in self._registry self._registry[options_hash] = options - - async def hint(self, options_hash: str, options: dict[str, Any]) -> None: + if user_id: # truthy: excludes None and "" to prevent UUID parse errors + self._user_ids[options_hash] = user_id + log.info( + "refill.register options_hash=%s already_known=%s registry_size=%d", + options_hash, already_known, len(self._registry), + ) + + async def hint( + self, + options_hash: str, + options: dict[str, Any], + *, + user_id: str | None = None, + ) -> None: """Signal that a user will likely need a subprocess soon. Registers the hash and triggers an immediate refill cycle. """ - self.register(options_hash, options) + self.register(options_hash, options, user_id=user_id) self._hint_event.set() # Run one fill cycle immediately for this hash deficit = self._deficit(options_hash) + log.info( + "refill.hint options_hash=%s deficit=%d target_depth=%d" + " warm=%d warming=%d", + options_hash, + deficit, + self._pool.config.target_depth_per_hash, + self._pool.warm_count(options_hash), + self._pool.warming_count(options_hash), + ) for _ in range(deficit): asyncio.create_task( - self._pool._try_inject_warm(options_hash, options) + self._pool._try_inject_warm(options_hash, options, user_id=user_id) ) async def run_once(self) -> None: """Run one full refill cycle across all registered hashes.""" for options_hash, options in list(self._registry.items()): + # Defensive: normalise "" → None so create_key never receives an + # empty UUID regardless of how the value entered _user_ids. + user_id = self._user_ids.get(options_hash) or None deficit = self._deficit(options_hash) + if deficit > 0: + log.info( + "refill.run_once options_hash=%s deficit=%d warm=%d warming=%d", + options_hash, deficit, + self._pool.warm_count(options_hash), + self._pool.warming_count(options_hash), + ) for _ in range(deficit): - accepted = await self._pool._try_inject_warm(options_hash, options) + accepted = await self._pool._try_inject_warm( + options_hash, options, user_id=user_id + ) if not accepted: - log.debug("Capacity full during refill for hash %s", options_hash) + log.info( + "refill.capacity_full options_hash=%s — breaking refill cycle", + options_hash, + ) break + # TTL sweep — evict any stale home dir assignments each cycle. + if self._pool._home_pool is not None: + await self._pool._home_pool.sweep() + async def run(self) -> None: """Continuous refill loop — run as an asyncio task.""" self._running = True diff --git a/agent_pool_manager/server.py b/agent_pool_manager/server.py index d1a084ac..08865197 100644 --- a/agent_pool_manager/server.py +++ b/agent_pool_manager/server.py @@ -5,26 +5,152 @@ import dataclasses import json import logging +import os +import re +import time +import urllib.parse +import uuid from contextlib import asynccontextmanager from typing import Any, AsyncIterator from fastapi import FastAPI, HTTPException, Request -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse from pydantic import BaseModel, Field from .config import AgentPoolConfig, load_pool_config +from .metrics import PoolMetrics from .pool import Pool, PoolExhausted from .refill import RefillLoop log = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Setup-token pexpect helpers +# --------------------------------------------------------------------------- + +_ANTHROPIC_URL_RE = re.compile(r"https://\S+") +_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]") +_OAUTH_TOKEN_RE = re.compile(r"sk-ant-[A-Za-z0-9_-]+") +_VALID_ANTHROPIC_NETLOCS = {"console.anthropic.com", "claude.com"} +_SETUP_TTL = 600 # seconds + +# Pending setup-token sessions: session_id → {"child": ..., "created_at": float} +_setup_token_sessions: dict[str, dict] = {} + + +def _extract_anthropic_url(text: str) -> str | None: + joined = text.replace("\r", "").replace("\n", "") + for match in _ANTHROPIC_URL_RE.finditer(joined): + candidate = match.group(0).rstrip(".,;)") + parsed = urllib.parse.urlparse(candidate) + if parsed.scheme == "https" and parsed.netloc in _VALID_ANTHROPIC_NETLOCS: + return candidate + return None + + +def _start_pexpect_sync(env: dict) -> tuple: + """Spawn ``claude setup-token`` in a PTY, wait for the auth URL. + + Returns ``(child, url)`` on success or ``(None, None)`` on failure. + The child stays alive waiting for the OAuth code via :func:`_complete_pexpect_sync`. + """ + import pexpect # lazy import — keeps module importable when pexpect absent + + log.info("setup_token/start: spawning claude setup-token") + try: + child = pexpect.spawn("claude", args=["setup-token"], env=env, dimensions=(24, 500)) + except Exception: + log.exception("setup_token/start: pexpect.spawn failed") + return None, None + + buf = "" + deadline = time.time() + 60 + while time.time() < deadline: + try: + chunk = child.read_nonblocking(4096, timeout=1) + buf += chunk.decode(errors="replace") + clean = _ANSI_ESCAPE_RE.sub("", buf) + url = _extract_anthropic_url(clean) + if url: + log.info("setup_token/start: auth URL extracted: %s", url) + return child, url + except pexpect.TIMEOUT: + continue + except pexpect.EOF: + log.warning("setup_token/start: EOF before URL found; output: %r", buf[-500:]) + break + except Exception: + log.exception("setup_token/start: unexpected error reading child output") + break + + log.error("setup_token/start: timed out waiting for auth URL") + try: + child.close(force=True) + except Exception: + pass + return None, None + + +def _complete_pexpect_sync(child, code: str) -> tuple[str, str]: + """Send the OAuth code to the waiting child and capture the OAuth token. + + Returns ``(result, token)`` where result is ``"ok"``, ``"failed"``, + ``"timeout"``, or ``"error"``. + """ + import pexpect # lazy import + + log.info("setup_token/complete: sending code to pexpect child") + try: + child.send(code + "\r") + except Exception: + log.exception("setup_token/complete: pexpect send failed") + return "error", "" + + buf = b"" + token: str | None = None + sent_confirm = False + deadline = time.time() + 120 + + while time.time() < deadline: + try: + chunk = child.read_nonblocking(4096, timeout=1) + buf += chunk + clean = _ANSI_ESCAPE_RE.sub("", chunk.decode(errors="replace")) + match = _OAUTH_TOKEN_RE.search(clean) + if match: + token = match.group(0) + break + if not sent_confirm and b"\r\r\n\r\r\n" in buf: + child.send("\r") + sent_confirm = True + except pexpect.TIMEOUT: + continue + except pexpect.EOF: + full_clean = _ANSI_ESCAPE_RE.sub("", buf.decode(errors="replace")) + match = _OAUTH_TOKEN_RE.search(full_clean) + if match: + token = match.group(0) + break + except Exception: + log.exception("setup_token/complete: unexpected error") + return "error", "" + else: + return "timeout", "" + + try: + child.close() + except Exception: + pass + + return ("ok", token) if token else ("failed", "") + # --------------------------------------------------------------------------- # Request / response models # --------------------------------------------------------------------------- class AcquireRequest(BaseModel): - user_id: str + user_id: str = Field(..., min_length=1) options_hash: str options: dict[str, Any] = Field(default_factory=dict) timeout_seconds: float | None = None @@ -32,6 +158,10 @@ class AcquireRequest(BaseModel): class ReleaseRequest(BaseModel): reusable: bool = False + # Optional token tracking — passed by the layer after a turn completes + user_id: str | None = None + input_tokens: int | None = None + output_tokens: int | None = None class QueryRequest(BaseModel): @@ -40,7 +170,7 @@ class QueryRequest(BaseModel): class HintRequest(BaseModel): - user_id: str + user_id: str = Field(..., min_length=1) options_hash: str options: dict[str, Any] = Field(default_factory=dict) @@ -52,6 +182,49 @@ class ControlResponseRequest(BaseModel): denial_message: str | None = None +class SetupTokenCompleteRequest(BaseModel): + session_id: str = Field(..., min_length=1) + code: str = Field(..., min_length=1) + + +# --------------------------------------------------------------------------- +# Token usage async write — no-op if DB pool unavailable +# --------------------------------------------------------------------------- + +# Strong references to in-flight token-write tasks prevent GC from silently +# dropping them before they complete (asyncio creates a weak reference only). +_token_write_tasks: set[asyncio.Task] = set() + + +def write_token_usage_async(user_id: str, input_tokens: int, output_tokens: int) -> None: + """Fire-and-forget token usage DB write. + + Creates an asyncio task to write token counts to the ``token_usage`` table. + The task is intentionally not awaited — callers on the hot path (release + endpoint) must not block on the DB write. + + A strong reference is kept in ``_token_write_tasks`` until the task + completes, preventing the GC from silently dropping in-flight writes. + + If the event loop is not running or the DB layer is unavailable, the + failure is logged at WARNING level (not silently swallowed). + """ + try: + from db.pg_queries.token_usage import record_token_usage + task = asyncio.create_task( + record_token_usage(user_id, input_tokens, output_tokens), + name=f"token-usage-{user_id}", + ) + _token_write_tasks.add(task) + task.add_done_callback(_token_write_tasks.discard) + except RuntimeError as exc: + # RuntimeError: no running event loop — should not happen in an async + # handler but log at warning so it's visible, not silently dropped. + log.warning("write_token_usage_async: no event loop — token write dropped: %s", exc) + except Exception: + log.warning("write_token_usage_async: failed to schedule token write", exc_info=True) + + # --------------------------------------------------------------------------- # App factory # --------------------------------------------------------------------------- @@ -59,25 +232,62 @@ class ControlResponseRequest(BaseModel): def build_app( pool: Pool | None = None, refill: RefillLoop | None = None, + metrics: PoolMetrics | None = None, ) -> FastAPI: """Build the FastAPI app. - ``pool`` and ``refill`` may be injected for testing; if omitted they are - constructed from the TetherConfig at startup. + ``pool``, ``refill``, and ``metrics`` may be injected for testing; if + omitted they are constructed from the TetherConfig at startup. """ @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncIterator[None]: if pool is not None: app.state.pool = pool app.state.refill = refill or RefillLoop(pool) + app.state.metrics = metrics or PoolMetrics() else: try: from config.loader import TetherConfig cfg = load_pool_config(TetherConfig()) except Exception: cfg = AgentPoolConfig() - app.state.pool = Pool(cfg) + + # Attempt to initialise a Postgres pool for ephemeral MCP key + # creation/revocation. Gracefully degrades to no-key mode if + # DATABASE_URL is absent (e.g. Pi bot running SQLite-only). + pg_pool = None + try: + from db.postgres import create_pool as _create_pg_pool + pg_pool = await _create_pg_pool() + log.info("agent_pool_manager: Postgres pool initialised for MCP key injection") + except Exception: + log.warning( + "agent_pool_manager: Postgres pool unavailable" + " — MCP key injection disabled; subprocesses will receive" + " list-form mcp_servers (may cause connect hang on non-Pi deploys)", + exc_info=True, + ) + + app.state.pool = Pool(cfg, pg_pool=pg_pool) app.state.refill = RefillLoop(app.state.pool) + app.state.metrics = PoolMetrics() + + # Wire metrics into pool so acquire/release/refill events are recorded + app.state.pool._metrics = app.state.metrics + app.state.metrics.attach_pool(app.state.pool) + + # Initialise the home directory pool so isolated HOME dirs are available + # before the first warm spawn. Gracefully degrades if the base dir + # does not exist on this deployment (Pi vs Fly). + try: + await app.state.pool.initialize_home_pool() + log.info("agent_pool_manager: home directory pool initialised") + except Exception: + log.warning( + "agent_pool_manager: home directory pool unavailable" + " — subprocesses will share default HOME (lock contention possible)", + exc_info=True, + ) app.state.refill.start() log.info("Agent pool manager started") @@ -98,6 +308,7 @@ async def acquire(request: Request, req: AcquireRequest) -> JSONResponse: req.options_hash, req.options, timeout=req.timeout_seconds, + user_id=req.user_id, ) except PoolExhausted: return JSONResponse( @@ -196,6 +407,10 @@ async def release(handle_id: str, req: ReleaseRequest, request: Request) -> None raise HTTPException(status_code=404, detail="handle not found") await the_pool.release(handle_id, reusable=req.reusable) + # Fire-and-forget token write when the caller reports usage + if req.user_id and req.input_tokens is not None and req.output_tokens is not None: + write_token_usage_async(req.user_id, req.input_tokens, req.output_tokens) + # ----------------------------------------------------------------------- # GET /status # ----------------------------------------------------------------------- @@ -204,13 +419,36 @@ async def status(request: Request) -> dict: the_pool: Pool = request.app.state.pool return the_pool.status() + # ----------------------------------------------------------------------- + # GET /metrics — Prometheus text exposition format + # ----------------------------------------------------------------------- + @app.get("/metrics") + async def metrics_endpoint(request: Request) -> PlainTextResponse: + """Expose pool metrics in Prometheus text format. + + Counters, histograms, and pool size gauges are included. + Scrape with Prometheus or read directly for debugging. + """ + the_metrics: PoolMetrics = request.app.state.metrics + text = the_metrics.render_text() + return PlainTextResponse(content=text, media_type="text/plain; version=0.0.4") + # ----------------------------------------------------------------------- # POST /hint # ----------------------------------------------------------------------- @app.post("/hint", status_code=202) async def hint(req: HintRequest, request: Request) -> dict: the_refill: RefillLoop = request.app.state.refill - asyncio.create_task(the_refill.hint(req.options_hash, req.options)) + # Diagnostic: log hint receipt with redacted summary so we can confirm + # what reached the pool service and correlate with refill-side logs. + from .pool import _options_summary # local import — keep module deps flat + log.info( + "pool_server.hint_recv user_id=%s options_hash=%s summary=%r", + req.user_id, req.options_hash, _options_summary(req.options), + ) + asyncio.create_task( + the_refill.hint(req.options_hash, req.options, user_id=req.user_id) + ) return {"queued": True} # ----------------------------------------------------------------------- @@ -236,9 +474,93 @@ async def control_response( detail=f"request_id {req.request_id!r} not found or already resolved", ) + # ----------------------------------------------------------------------- + # POST /setup-token + # + # Spawns ``claude setup-token`` in a PTY, waits for the Anthropic auth + # URL, and returns it along with a session_id the caller uses to submit + # the OAuth code via POST /setup-token/complete. + # + # Security note: this endpoint is unauthenticated at the pool-manager + # layer (consistent with /acquire and /hint). Access control is enforced + # by the API service that proxies here — callers must not expose the pool + # manager port (5002) to untrusted networks. + # ----------------------------------------------------------------------- + @app.post("/setup-token") + async def setup_token_start(request: Request) -> JSONResponse: + # Sweep expired sessions before starting a new one + now = time.time() + expired = [sid for sid, s in _setup_token_sessions.items() + if now - s["created_at"] > _SETUP_TTL] + for sid in expired: + entry = _setup_token_sessions.pop(sid, None) + if entry: + loop = asyncio.get_running_loop() + loop.run_in_executor(None, lambda c=entry["child"]: _close_child(c)) + + env = {**os.environ} + loop = asyncio.get_running_loop() + child, url = await loop.run_in_executor(None, _start_pexpect_sync, env) + + if child is None or url is None: + return JSONResponse( + status_code=503, + content={"error": "setup_token_failed", "detail": "claude setup-token did not produce an auth URL"}, + ) + + session_id = str(uuid.uuid4()) + _setup_token_sessions[session_id] = {"child": child, "created_at": now} + log.info("setup_token/start: session_id=%s url=%s", session_id, url) + return JSONResponse({"session_id": session_id, "url": url}) + + # ----------------------------------------------------------------------- + # POST /setup-token/complete + # + # Sends the OAuth code to the waiting pexpect child identified by + # session_id and returns the resulting OAuth token. + # ----------------------------------------------------------------------- + @app.post("/setup-token/complete") + async def setup_token_complete(req: SetupTokenCompleteRequest) -> JSONResponse: + entry = _setup_token_sessions.pop(req.session_id, None) + if entry is None: + raise HTTPException(status_code=404, detail="session not found or expired") + + child = entry["child"] + loop = asyncio.get_running_loop() + result, token = await loop.run_in_executor( + None, _complete_pexpect_sync, child, req.code + ) + log.info("setup_token/complete: session_id=%s result=%s", req.session_id, result) + return JSONResponse({"result": result, "token": token if token else None}) + + # ----------------------------------------------------------------------- + # DELETE /setup-token/{session_id} + # + # Cancels a pending setup-token session, killing the pexpect child. + # Called by the API proxy when a user restarts /start to prevent + # orphaned subprocess accumulation. + # ----------------------------------------------------------------------- + @app.delete("/setup-token/{session_id}", status_code=204) + async def setup_token_cancel(session_id: str, request: Request) -> None: + entry = _setup_token_sessions.pop(session_id, None) + if entry is None: + raise HTTPException(status_code=404, detail="session not found or expired") + child = entry["child"] + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, _close_child, child) + log.info("setup_token/cancel: session_id=%s child reaped", session_id) + return app +def _close_child(child) -> None: + """Kill a pexpect child — used for TTL cleanup.""" + try: + child.close(force=True) + except Exception: + pass + + # --------------------------------------------------------------------------- # Serialisation helper # --------------------------------------------------------------------------- diff --git a/api/main.py b/api/main.py index dfc50692..1f7bf279 100644 --- a/api/main.py +++ b/api/main.py @@ -41,8 +41,10 @@ from api.routes import ical as ical_routes from api.routes import conversations as conversations_routes from api.routes import internal as internal_routes +from api.routes import pool as pool_routes from api.ws import manager from api.auth import auth_dependency, decode_jwt +from api.redis_pubsub import subscribe_and_forward, get_redis_url from api.limiter import limiter from db.pool_middleware import lifespan as _pool_lifespan from db.pg_queries.errors import StaleReadError @@ -221,6 +223,11 @@ async def integrity_handler(request, exc): prefix="/api/internal", include_in_schema=False, ) + app.include_router( + pool_routes.router, + prefix="/api/internal/pool", + include_in_schema=False, + ) # --- Premium plugin hook --- try: @@ -289,12 +296,16 @@ async def websocket_endpoint(websocket: WebSocket): await manager.connect(websocket, user_id) if is_admin: manager.register_only(websocket, "__bot__") + redis_task = asyncio.create_task( + subscribe_and_forward(websocket, user_id, redis_url=get_redis_url()) + ) try: while True: await websocket.receive_text() except (WebSocketDisconnect, RuntimeError): pass finally: + redis_task.cancel() manager.disconnect(websocket, user_id) if is_admin: manager.disconnect(websocket, "__bot__") @@ -336,12 +347,16 @@ async def websocket_endpoint(websocket: WebSocket): manager.register_only(websocket, user_id) if is_admin: manager.register_only(websocket, "__bot__") + redis_task = asyncio.create_task( + subscribe_and_forward(websocket, user_id, redis_url=get_redis_url()) + ) try: while True: await websocket.receive_text() except (WebSocketDisconnect, RuntimeError): pass finally: + redis_task.cancel() manager.disconnect(websocket, user_id) if is_admin: manager.disconnect(websocket, "__bot__") diff --git a/api/redis_pubsub.py b/api/redis_pubsub.py new file mode 100644 index 00000000..6f84d08a --- /dev/null +++ b/api/redis_pubsub.py @@ -0,0 +1,146 @@ +"""Redis pub/sub helpers for cross-process event delivery. + +Architecture +------------ +The interactive-agent-layer (port 5003) and the API (port 8000) run as +separate supervisord programs in the same container. The layer's +WSPublisher dual-writes events to both in-process asyncio queues and a +Redis channel. This module provides the subscriber side: per connected +WebSocket, a background task subscribes to the user's channel and +forwards events as WS frames to the browser. + +Channel key: ``user:{user_id}:events`` + +Graceful degradation +-------------------- +If ``REDIS_URL`` is not set, the subscription task is never started and +the app runs without cross-process event delivery. Background events +(trial_usage_update, Beacon notifications) won't reach connected +browsers until Redis is configured — but the app starts and serves +requests normally. + +Future migration to managed Redis +---------------------------------- +Change ``REDIS_URL`` from ``redis://localhost:6379`` (supervisord Redis) +to a managed Upstash/ElastiCache URL with TLS +(``rediss://user:pass@host:port``). No code changes needed — the Redis +client handles auth and TLS from the URL scheme. + +Fly provisioning (when ready) +------------------------------ + fly redis create --name tether-redis --org personal + fly secrets set REDIS_URL="" --app tether-prod + fly secrets set REDIS_URL="" --app tether-dev + +Until then the supervisord ``[program:redis]`` stanza in supervisord.conf +provides a local Redis instance with ``REDIS_URL=redis://localhost:6379``. +""" +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +import os +from typing import Any + +from shared.redis_channels import channel_for # noqa: F401 — re-exported for API consumers + +logger = logging.getLogger(__name__) + + +def get_redis_url() -> str | None: + """Return REDIS_URL from environment, or None if unset.""" + return os.environ.get("REDIS_URL") or None + + +async def subscribe_and_forward( + websocket: Any, + user_id: str, + *, + redis_url: str | None = None, + server: Any = None, # fakeredis FakeServer for testing +) -> None: + """Subscribe to user:{user_id}:events and forward events to websocket. + + Runs until cancelled (typically when the WS disconnects). Each event + published to the channel is deserialized and sent as a JSON WS frame. + WS send failures are swallowed so a disconnected client doesn't crash + the subscriber task. + + Parameters + ---------- + websocket: + An open WebSocket with a ``send_json(dict)`` async method. + user_id: + The authenticated user whose channel to subscribe to. + redis_url: + Redis connection URL. Defaults to ``get_redis_url()``. + server: + Optional fakeredis FakeServer — used in tests to share state + between publisher and subscriber without a real Redis process. + """ + import redis.asyncio as aioredis + + url = redis_url or get_redis_url() + if url is None and server is None: + logger.warning( + "subscribe_and_forward: REDIS_URL not set — " + "background events will not reach user_id=%s", + user_id, + ) + # Block forever (keeps caller's task alive until cancelled) + await asyncio.Future() + return + + # Build client: real Redis from URL, or fakeredis via server param + if server is not None: + import fakeredis.aioredis as faioredis + client: Any = faioredis.FakeRedis(server=server) + else: + client = aioredis.from_url(url) + + pubsub = client.pubsub() + channel = channel_for(user_id) + await pubsub.subscribe(channel) + logger.debug( + "subscribe_and_forward: subscribed to %s for user_id=%s", channel, user_id + ) + # Use get_message(timeout=0) + asyncio.sleep rather than pubsub.listen(). + # pubsub.listen() blocks on an internal queue.get() that does not cleanly + # propagate CancelledError in Python 3.11, causing the task to hang on + # cancellation. asyncio.sleep() is the sole cancellation point here and + # always propagates CancelledError regardless of redis-py internals. + # + # Python 3.11 edge case: subscribe() can absorb the initial CancelledError + # from task.cancel() (it catches it internally for cleanup but still + # completes the subscription). When this happens _must_cancel is not set, + # so CancelledError is not automatically re-thrown at future await points. + # We detect this via task.cancelling() (added in 3.11) and raise manually. + _task = asyncio.current_task() + try: + while True: + message = await pubsub.get_message( + ignore_subscribe_messages=False, timeout=0 + ) + if message is not None and message["type"] == "message": + try: + event = json.loads(message["data"]) + await websocket.send_json(event) + except Exception as exc: + logger.debug( + "subscribe_and_forward: WS send failed user_id=%s: %s", + user_id, exc, + ) + else: + await asyncio.sleep(0.01) + # If subscribe() absorbed our cancellation signal, honour it now. + if _task is not None and getattr(_task, "cancelling", lambda: 0)(): + raise asyncio.CancelledError() + finally: + with contextlib.suppress(Exception): + await asyncio.wait_for(pubsub.unsubscribe(channel), timeout=2.0) + with contextlib.suppress(Exception): + await asyncio.wait_for(pubsub.aclose(), timeout=2.0) + with contextlib.suppress(Exception): + await asyncio.wait_for(client.aclose(), timeout=2.0) diff --git a/api/routes/bot.py b/api/routes/bot.py index 29594b0a..361a182b 100644 --- a/api/routes/bot.py +++ b/api/routes/bot.py @@ -254,6 +254,36 @@ async def status_fn(msg: str) -> None: ) raise # Let the session task propagate disconnect to the outer handler + # Async event callback for streamed layer events (text deltas, + # permission requests, etc.). agent_text_delta events are sent + # as agent_text_delta frames so the browser can render them + # incrementally. Other event types (including turn_complete) are + # forwarded as-is. + # turn_complete_sent tracks whether event_fn already forwarded a + # turn_complete from the session — if so, the response_parts path + # below skips its own turn_complete to prevent a double-send that + # would poison the frontend's incoming queue for the next message. + turn_complete_sent: list[bool] = [False] + + async def event_fn(event: dict) -> None: + try: + etype = event.get("type") + if etype == "agent_text_delta": + delta = event.get("delta", "") + if delta: + await websocket.send_json({"type": "agent_text_delta", "delta": delta}) + else: + if etype == "turn_complete": + turn_complete_sent[0] = True + await websocket.send_json(event) + except Exception as e: + logger.debug( + "bot_chat: event_fn send failed (client likely disconnected)," + " user_id=%s: %s", + user_id, e, + ) + raise + # Capture responses delivered via send_fn. handle_message always # calls send_fn(final) and returns None — the return value is not # used for response delivery. This list is local to each message @@ -274,6 +304,8 @@ def capture_send_fn(msg: str) -> None: user_id=user_id, vault=getattr(websocket.app.state, "vault", None), status_fn=status_fn, + event_fn=event_fn, + is_admin=websocket.state.is_admin, ) ) @@ -299,7 +331,7 @@ def capture_send_fn(msg: str) -> None: # further status frames from the session arrive after the ack. await _cancel_and_wait(session_task) await websocket.send_json({"type": "status", "content": "Stopped."}) - await websocket.send_json({"type": "done"}) + await websocket.send_json({"type": "turn_complete", "final_text": "", "session_id": ""}) session_task = None continue # Back to outer loop — ready for next message @@ -355,7 +387,7 @@ def capture_send_fn(msg: str) -> None: "send another message to continue." ), }) - await websocket.send_json({"type": "done"}) + await websocket.send_json({"type": "turn_complete", "final_text": "", "session_id": ""}) session_task = None continue except Exception as e: @@ -367,15 +399,19 @@ def capture_send_fn(msg: str) -> None: "type": "error", "message": "Something went wrong. Please try again.", }) - await websocket.send_json({"type": "done"}) + await websocket.send_json({"type": "turn_complete", "final_text": "", "session_id": ""}) session_task = None continue session_task = None - response = "\n\n".join(response_parts) if response_parts else None - if response: - await websocket.send_json({"type": "chunk", "content": response}) - await websocket.send_json({"type": "done"}) + if not turn_complete_sent[0]: + response = "\n\n".join(response_parts) if response_parts else "" + logger.info( + "bot_chat: sending turn_complete final_text_chars=%d user_id=%s", + len(response), + user_id, + ) + await websocket.send_json({"type": "turn_complete", "final_text": response, "session_id": ""}) except WebSocketDisconnect: await _cancel_and_wait(session_task) if session_task else None diff --git a/api/routes/conversations.py b/api/routes/conversations.py index 0b2ea91d..8e1e2d9a 100644 --- a/api/routes/conversations.py +++ b/api/routes/conversations.py @@ -2,6 +2,7 @@ Endpoints: GET /conversations list conversations (RLS-scoped) + GET /conversations/index lightweight index (id, title, parent_context_node_id, updated_at, message_count) POST /conversations create conversation GET /conversations/{id} get single conversation PATCH /conversations/{id} patch conversation @@ -14,10 +15,15 @@ integer primary key of conversation_history. This gives stable pages under concurrent inserts, unlike offset-based pagination. The tradeoff is that clients must treat the cursor as an opaque integer string. See PR description. + +Note on route ordering: /conversations/index MUST be registered before +/conversations/{conversation_id} so FastAPI does not match the literal +string "index" as a conversation_id path parameter. """ from __future__ import annotations import json +from typing import Literal from fastapi import APIRouter, Depends, HTTPException, Query, Request from pydantic import BaseModel @@ -28,6 +34,7 @@ get_conversation, list_conversations, list_conversation_messages, + list_conversations_index, update_conversation, ) from db.pg_queries.nodes import get_node @@ -44,6 +51,10 @@ _NOTIF_MODES = {"all", "focus", "quiet", "off"} _CHANNELS = {"telegram", "email", "push"} +# Valid conversation states — open/closed existed from Phase B; pending/rejected +# added in Phase 3 for Beacon-initiated conversations (spec §7.1). +ConversationState = Literal["open", "closed", "pending", "rejected"] + class ConversationCreate(BaseModel): name: str @@ -55,7 +66,7 @@ class ConversationCreate(BaseModel): class ConversationPatch(BaseModel): name: str | None = None priority: str | None = None - state: str | None = None + state: ConversationState | None = None context_node_id: str | None = None @@ -135,6 +146,26 @@ async def create_conv( return await _attach_folder_name(conn, conv) +# --------------------------------------------------------------------------- +# GET /conversations/index +# --------------------------------------------------------------------------- +# NOTE: registered before /{conversation_id} so "index" is not matched as an id. + + +@router.get("/conversations/index") +async def get_conversations_index( + request: Request, + _auth=Depends(auth_dependency), + conn=Depends(get_db_conn), +): + """Lightweight index: id, title, parent_context_node_id, updated_at, message_count. + + No message bodies. One DB query. Used by the frontend to build tree views + without fetching full conversation detail for each item. + """ + return await list_conversations_index(conn, user_id=request.state.user_id) + + # --------------------------------------------------------------------------- # GET /conversations/{conversation_id} # --------------------------------------------------------------------------- diff --git a/api/routes/integrations.py b/api/routes/integrations.py index 8936ec88..9369f1cd 100644 --- a/api/routes/integrations.py +++ b/api/routes/integrations.py @@ -18,10 +18,7 @@ import asyncio import logging -import os -import re import time -import urllib.parse import asyncpg import httpx @@ -64,36 +61,9 @@ _start_locks: dict[str, asyncio.Lock] = {} _SETUP_TTL = 600 # seconds -_ANTHROPIC_URL_RE = re.compile(r"https://\S+") -# Strip ANSI CSI escape sequences emitted by TUI programs on a PTY. -# Note: does not strip OSC8 hyperlinks (\x1b]8;;URL\x07...\x1b]8;;\x07) — if -# claude setup-token ever uses them, extend this regex or post-process the URL. -_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]") - -# `claude setup-token` prints a long-lived OAuth token to stdout (sk-ant-oat-…) -# rather than writing a credentials file. The trailing run is URL-safe base64 -# plus dashes/underscores; we capture the longest such run after the prefix. -_ANTHROPIC_URL_SCHEME = "https" -# claude setup-token may emit URLs on either domain depending on version. -_VALID_ANTHROPIC_NETLOCS = {"console.anthropic.com", "claude.com"} -_OAUTH_TOKEN_RE = re.compile(r"sk-ant-[A-Za-z0-9_-]+") - - -def _extract_anthropic_url(text: str) -> str | None: - """Extract and strictly validate the Anthropic auth URL from subprocess output. - - Strips CR/LF before matching so PTY line-wrapping doesn't split the URL. - Uses urlparse to check scheme and netloc exactly — prevents substring-match - bypasses where a valid domain appears at an arbitrary position in a crafted URL. - """ - # PTY output uses CRLF; strip both so the URL regex can match across wrap points. - joined = text.replace("\r", "").replace("\n", "") - for match in _ANTHROPIC_URL_RE.finditer(joined): - candidate = match.group(0).rstrip(".,;)") # strip trailing punctuation - parsed = urllib.parse.urlparse(candidate) - if parsed.scheme == _ANTHROPIC_URL_SCHEME and parsed.netloc in _VALID_ANTHROPIC_NETLOCS: - return candidate - return None +# Pool manager base URL — pexpect subprocess work is delegated to this service. +_POOL_MANAGER_BASE_URL = "http://127.0.0.1:5002" + # --------------------------------------------------------------------------- @@ -108,173 +78,14 @@ class AnthropicCompleteBody(BaseModel): code: str -# --------------------------------------------------------------------------- -# Anthropic OAuth helpers -# --------------------------------------------------------------------------- - -def _close_child_sync(child) -> None: - """Kill a pexpect child and close its PTY file descriptor.""" - try: - child.close(force=True) - except Exception: - pass - - -async def _reap_child(child) -> None: - """Async wrapper: close pexpect child in an executor to avoid blocking the event loop.""" - loop = asyncio.get_running_loop() - try: - await loop.run_in_executor(None, _close_child_sync, child) - except Exception: - pass - - -def _start_pexpect_sync(env: dict) -> tuple: - """Spawn ``claude setup-token`` in a PTY, wait for the auth URL, return ``(child, url)``. - - The child process stays alive after this returns, waiting for the user to - paste the OAuth code via :func:`_complete_pexpect_sync`. - - A wide PTY (220 cols) prevents the URL from line-wrapping, which would - break the URL regex. Returns ``(None, None)`` on any failure. - """ - import pexpect # lazy — keeps module importable when pexpect is not installed - - logger.info("anthropic/start: spawning claude setup-token") - try: - child = pexpect.spawn( - "claude", - args=["setup-token"], - env=env, - dimensions=(24, 500), - ) - logger.debug("anthropic/start: pexpect.spawn ok, pid=%s", child.pid) - except Exception: - logger.exception("pexpect.spawn failed for claude setup-token") - return None, None - - buf = "" - deadline = time.time() + 30 - while time.time() < deadline: - try: - chunk = child.read_nonblocking(4096, timeout=1) - # Accumulate raw bytes first, then strip ANSI from the full buffer - # so CSI sequences that straddle read boundaries are handled correctly. - buf += chunk.decode(errors="replace") - clean = _ANSI_ESCAPE_RE.sub("", buf) - logger.debug("anthropic/start: pexpect buffer so far (clean): %r", clean[-500:]) - url = _extract_anthropic_url(clean) - if url: - logger.info("anthropic/start: auth URL extracted: %s", url) - return child, url - except pexpect.TIMEOUT: - continue - except pexpect.EOF: - logger.warning( - "anthropic/start: pexpect EOF before URL found; full output: %r", - _ANSI_ESCAPE_RE.sub("", buf), - ) - break - except Exception: - logger.exception("Unexpected error reading pexpect child output") - break - - logger.error( - "anthropic/start: gave up waiting for auth URL after 30s; " - "last buffer (clean, last 500 chars): %r", - _ANSI_ESCAPE_RE.sub("", buf)[-500:], - ) - _close_child_sync(child) - return None, None - - -def _complete_pexpect_sync(child, code: str) -> tuple[str, str]: - """Send the OAuth code to the waiting pexpect child and capture the OAuth token. - - Streams child output live (one INFO log per chunk) so we can see exactly - what the CLI emits after the code is sent. Scans each chunk for the token - pattern so we catch it as soon as it appears rather than waiting for EOF. - - Returns a (result, token) tuple where result is one of: - ``"ok"`` — token found; token string is non-empty - ``"failed"`` — child exited without printing a token - ``"timeout"`` — 120 seconds elapsed without token or EOF - ``"error"`` — unexpected error - """ - import pexpect # lazy — keeps module importable when pexpect is not installed - - logger.info("anthropic/complete: sending code (length=%d) to pexpect child", len(code)) - try: - child.send(code + '\r') # \r alone = one Enter in PTY canonical mode - except Exception: - logger.exception("pexpect send failed — child likely already exited") - return "error", "" - - buf = b"" - token: str | None = None - sent_confirm = False - deadline = time.time() + 120 - - while time.time() < deadline: - try: - chunk = child.read_nonblocking(4096, timeout=1) - buf += chunk - clean = _ANSI_ESCAPE_RE.sub("", chunk.decode(errors="replace")) - safe = _OAUTH_TOKEN_RE.sub("sk-ant-***REDACTED***", clean) - logger.info("anthropic/complete: child output chunk: %r", safe) - match = _OAUTH_TOKEN_RE.search(clean) - if match: - token = match.group(0) - logger.info("anthropic/complete: token found in output stream") - break - # After the code echo the CLI may sit at a secondary "press Enter" - # prompt. Detect the two-blank-line pattern that follows the - # asterisk echo and send a confirming \r once. - if not sent_confirm and b"\r\r\n\r\r\n" in buf: - logger.info("anthropic/complete: sending confirm \\r after code echo") - child.send('\r') - sent_confirm = True - except pexpect.TIMEOUT: - continue - except pexpect.EOF: - logger.info("anthropic/complete: child EOF") - # Token may have arrived in the same read as EOF — scan full buffer. - full_clean = _ANSI_ESCAPE_RE.sub("", buf.decode(errors="replace")) - match = _OAUTH_TOKEN_RE.search(full_clean) - if match: - token = match.group(0) - logger.info("anthropic/complete: token found at EOF") - break - except Exception: - logger.exception("pexpect read_nonblocking failed") - return "error", "" - else: - full_clean = _ANSI_ESCAPE_RE.sub("", buf.decode(errors="replace")) - logger.warning( - "anthropic/complete: timed out waiting for token; output so far: %r", - full_clean[-500:], - ) - return "timeout", "" - - try: - child.close() - except Exception: - pass - - if token: - return "ok", token - logger.warning( - "anthropic/complete: child exited without printing token; full output: %r", - _ANSI_ESCAPE_RE.sub("", buf.decode(errors="replace"))[-500:], - ) - return "failed", "" async def _sweep_expired_setups() -> None: - """Kill and remove any pending setups older than _SETUP_TTL seconds. + """Remove pending setup entries older than _SETUP_TTL seconds. - Each killed child is reaped via an asyncio task so PTY FDs are released - without blocking the current request. + Session cleanup on the pool manager side is handled by the pool manager + itself; this sweep only removes stale session_id references from our + local ``_pending_setups`` dict. """ now = time.time() expired_users = [ @@ -282,10 +93,8 @@ async def _sweep_expired_setups() -> None: if now - entry["started_at"] > _SETUP_TTL ] for uid in expired_users: - entry = _pending_setups.pop(uid) - child = entry.get("child") - if child is not None: - asyncio.create_task(_reap_child(child)) + _pending_setups.pop(uid, None) + logger.debug("anthropic/sweep: removed expired setup for user_id=%s", uid) # --------------------------------------------------------------------------- @@ -673,37 +482,58 @@ async def anthropic_start( async def _anthropic_start_locked(request: Request, user_id: str) -> dict: """Inner /start logic — called while holding the per-user _start_locks entry. - Spawns ``claude setup-token`` in a PTY via :func:`_start_pexpect_sync` (run - in a thread-pool executor so the event loop is not blocked). The pexpect child - stays alive waiting for the user to paste their OAuth code. + Proxies to the agent pool manager's POST /setup-token endpoint, which + spawns ``claude setup-token`` in a PTY and returns the auth URL and a + session_id. The session_id is stored in ``_pending_setups`` so the + /complete handler can forward the OAuth code to the correct subprocess. """ logger.info("anthropic/start: request for user_id=%s", user_id) if user_id in _pending_setups: - logger.info("anthropic/start: killing existing pending setup for user_id=%s", user_id) - old = _pending_setups.pop(user_id) - asyncio.create_task(_reap_child(old["child"])) - - env_override = dict(os.environ) + old_entry = _pending_setups.pop(user_id) + old_session_id = old_entry.get("session_id") + logger.info( + "anthropic/start: canceling stale pending setup for user_id=%s session_id=%s", + user_id, old_session_id, + ) + if old_session_id: + try: + async with httpx.AsyncClient(timeout=5.0) as client: + await client.delete(f"{_POOL_MANAGER_BASE_URL}/setup-token/{old_session_id}") + except Exception: + logger.warning( + "anthropic/start: failed to cancel stale session_id=%s (continuing)", + old_session_id, + exc_info=True, + ) - loop = asyncio.get_running_loop() try: - child, url = await loop.run_in_executor( - None, _start_pexpect_sync, env_override - ) + async with httpx.AsyncClient(timeout=90.0) as client: + resp = await client.post(f"{_POOL_MANAGER_BASE_URL}/setup-token", json={}) except Exception as exc: - logger.exception("Unexpected error from _start_pexpect_sync: %s", exc) - raise HTTPException(status_code=502, detail="Failed to spawn setup process") + logger.exception("anthropic/start: pool manager request failed: %s", exc) + raise HTTPException(status_code=502, detail="Failed to reach setup-token service") - if url is None: - logger.error("anthropic/start: failed to extract auth URL for user_id=%s", user_id) + if resp.status_code != 200: + logger.error( + "anthropic/start: pool manager returned %d for user_id=%s body=%r", + resp.status_code, user_id, resp.text, + ) raise HTTPException(status_code=502, detail="Auth URL not found in claude output") + data = resp.json() + session_id = data.get("session_id") + url = data.get("url") + + if not session_id or not url: + logger.error("anthropic/start: pool manager response missing fields: %r", data) + raise HTTPException(status_code=502, detail="Invalid response from setup-token service") + _pending_setups[user_id] = { - "child": child, + "session_id": session_id, "started_at": time.time(), } - logger.info("anthropic/start: success, auth URL ready for user_id=%s", user_id) + logger.info("anthropic/start: success, auth URL ready for user_id=%s session_id=%s", user_id, session_id) return {"url": url, "expires_in": _SETUP_TTL} @@ -718,40 +548,53 @@ async def anthropic_complete( body: AnthropicCompleteBody, _auth: dict = Depends(auth_dependency), ): - """Submit the OAuth code to the waiting pexpect child and persist credentials.""" + """Forward the OAuth code to the pool manager and persist the returned token.""" user_id = request.state.user_id logger.info("anthropic/complete: request for user_id=%s, code_length=%d", user_id, len(body.code)) - # Pop atomically: prevents two concurrent /complete calls from racing on the - # same child process or temp dir. entry = _pending_setups.pop(user_id, None) if entry is None: logger.warning("anthropic/complete: no pending setup found for user_id=%s", user_id) raise HTTPException(status_code=404, detail="No pending Anthropic setup for this user") - child = entry["child"] + session_id = entry["session_id"] age = time.time() - entry["started_at"] - logger.debug("anthropic/complete: pending setup age=%.1fs", age) + logger.debug("anthropic/complete: pending setup age=%.1fs session_id=%s", age, session_id) - loop = asyncio.get_running_loop() - result, token = await loop.run_in_executor(None, _complete_pexpect_sync, child, body.code) - logger.info("anthropic/complete: pexpect result=%r for user_id=%s", result, user_id) + try: + async with httpx.AsyncClient(timeout=150.0) as client: + resp = await client.post( + f"{_POOL_MANAGER_BASE_URL}/setup-token/complete", + json={"session_id": session_id, "code": body.code}, + ) + except Exception as exc: + logger.exception("anthropic/complete: pool manager request failed: %s", exc) + raise HTTPException(status_code=502, detail="Failed to reach setup-token service") + + if resp.status_code != 200: + logger.error( + "anthropic/complete: pool manager returned %d for user_id=%s body=%r", + resp.status_code, user_id, resp.text, + ) + raise HTTPException(status_code=502, detail="Setup token completion failed") + + data = resp.json() + result = data.get("result") + token = data.get("token") or "" + logger.info("anthropic/complete: pool manager result=%r for user_id=%s", result, user_id) if result == "error": - asyncio.create_task(_reap_child(child)) raise HTTPException(status_code=502, detail="Setup process closed unexpectedly") if result == "timeout": - asyncio.create_task(_reap_child(child)) raise HTTPException(status_code=504, detail="Setup process timed out") - # child already reaped by _complete_pexpect_sync. if result == "failed": logger.warning("anthropic/complete: setup process reported failure for user_id=%s", user_id) return {"ok": False, "error": "setup failed"} - # result == "ok": token was extracted from the output stream by _complete_pexpect_sync. - logger.debug("anthropic/complete: OAuth token extracted (len=%d)", len(token)) + # result == "ok": token extracted by the pool manager's pexpect handler. + logger.debug("anthropic/complete: OAuth token received (len=%d)", len(token)) vault = request.app.state.vault if vault is None: diff --git a/api/routes/nodes.py b/api/routes/nodes.py index adfdf1ca..201db75e 100644 --- a/api/routes/nodes.py +++ b/api/routes/nodes.py @@ -12,6 +12,7 @@ get_auto_archivable_nodes, archive_node, get_user_setting, ) +from db.pg_queries.nodes import list_nodes_index from db.pool_middleware import get_db_conn from api.ws import manager from api.auth import auth_dependency @@ -142,6 +143,19 @@ async def search_sections_route(_auth=Depends(auth_dependency), return await search_sections(conn, q.strip(), node_id=node_id) +@router.get("/nodes/index") +async def get_nodes_index(request: Request, _auth=Depends(auth_dependency), + conn: asyncpg.Connection = Depends(get_db_conn)): + """Lightweight index: id, title, parent_id, path, child_count. + + No section data. One recursive-CTE DB query. Used by the frontend to + build the node tree without fetching full node detail for each item. + + NOTE: registered before /{node_id} so "index" is not matched as a node id. + """ + return await list_nodes_index(conn, user_id=request.state.user_id) + + @router.get("/nodes/{node_id}") async def get_node_route(node_id: str, _auth=Depends(auth_dependency), conn: asyncpg.Connection = Depends(get_db_conn)): diff --git a/api/routes/pool.py b/api/routes/pool.py new file mode 100644 index 00000000..3d772243 --- /dev/null +++ b/api/routes/pool.py @@ -0,0 +1,183 @@ +"""Pool warm endpoint — user-facing hint to pre-warm a subprocess for the caller. + +POST /api/internal/pool/warm + Auth: cookie JWT (user-session auth, NOT X-Internal-Token cron auth) + Body: { "agent_version": str } + Response: 202 { "hinted": bool, "options_hash": str } + +Always returns 202 — warming is best-effort and must never block the frontend. +Pool failure is logged and reflected in hinted=false for observability. +""" +from __future__ import annotations + +import hashlib +import json +import logging +from typing import Any + +from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import BaseModel + +from api.auth import auth_dependency + +log = logging.getLogger(__name__) + +router = APIRouter() + +# --------------------------------------------------------------------------- +# Canonical options per agent version. +# +# These MUST match what bot.agent_dispatch passes to the layer on session/start. +# If _V2_0_OPTIONS changes in agent_dispatch, update this mapping too — and +# the test test_warm_options_hash_matches_layer_algorithm will catch any drift. +# --------------------------------------------------------------------------- + +def _get_agent_options() -> dict[str, dict[str, Any]]: + """Lazy import so agent_dispatch doesn't load at module import time.""" + try: + from bot.agent_dispatch import _V2_0_OPTIONS + v2_options: dict[str, Any] = _V2_0_OPTIONS + except ImportError: + v2_options = {} + + return { + "tether-agent-1.0": {}, + "tether-agent-2.0": v2_options, + # 2.5 goes through the premium handler directly (not the layer), + # so pool warming via layer is not applicable. We still accept 2.5 + # and warm with a minimal options set so the FE can fire unconditionally. + "tether-agent-2.5": v2_options, + } + + +def _compute_options_hash(options: dict[str, Any]) -> str: + """SHA-256 canonical-JSON hash, truncated to 16 hex chars. + + Algorithm is identical to interactive_agent_layer.session._stable_options_hash. + Both must stay in sync — the test test_warm_options_hash_matches_layer_algorithm + enforces this. + """ + canonical = json.dumps(options, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode()).hexdigest()[:16] + + +def _get_pool_client(request: Request): + """Return the pool client, lazily constructing and caching on app.state. + + Tests inject a mock by setting ``app.state.pool_client`` before requests. + In production the client is constructed once from config and reused, so + httpx connection pooling applies across requests. + """ + client = getattr(request.app.state, "pool_client", None) + if client is not None: + return client + + # Construct from config on first use and cache for the app's lifetime. + try: + from config.loader import config + base_url: str = config.get("agent_pool.base_url", "http://127.0.0.1:5002") + except Exception: + base_url = "http://127.0.0.1:5002" + + from agent_pool_manager.client import PoolClient + pool_client = PoolClient(base_url=base_url) + request.app.state.pool_client = pool_client + return pool_client + + +# --------------------------------------------------------------------------- +# Request / response models +# --------------------------------------------------------------------------- + +class WarmRequest(BaseModel): + agent_version: str + + +# --------------------------------------------------------------------------- +# Endpoint +# --------------------------------------------------------------------------- + +@router.post("/warm", status_code=202) +async def pool_warm( + body: WarmRequest, + request: Request, + _auth: dict = Depends(auth_dependency), +) -> dict: + """Pre-warm a pool subprocess for the authenticated user. + + Returns 202 always — warming is best-effort. ``hinted=false`` signals + that the pool call failed (pool unreachable), but the FE should not retry + aggressively; the next real acquire will fall through to cold-start. + """ + user_id: str = request.state.user_id + + log.info( + "pool_warm.entry user_id=%s agent_version=%s", + user_id, body.agent_version, + ) + + agent_options = _get_agent_options() + options = agent_options.get(body.agent_version) + if options is None: + raise HTTPException( + status_code=400, + detail=f"Unknown agent_version: {body.agent_version!r}. " + f"Valid versions: {sorted(agent_options)}", + ) + + vault = request.app.state.vault + if vault is None: + log.error( + "pool_warm: vault is not configured — set VAULT_KEY in the environment. " + "Pool warming requires vault to inject OAuth credentials into subprocesses." + ) + return {"hinted": False, "options_hash": _compute_options_hash(options)} + + try: + async with vault.materialize(user_id) as env_dict: + options_for_hint = {**options, "env": env_dict} + except ValueError: + log.warning( + "pool_warm: no vault credentials for user_id=%s agent_version=%s" + " — user must connect their Anthropic account first", + user_id, + body.agent_version, + ) + return {"hinted": False, "options_hash": _compute_options_hash(options)} + + # Hash computed after env injection so the warm partition matches the + # acquire-side hash in interactive_agent_layer. Partitions are per-user — + # a subprocess authenticated as user A must not serve user B. + options_hash = _compute_options_hash(options_for_hint) + + # Diagnostic: surface what we're about to hint so we can correlate with + # pool-side logs. Env values are redacted but key names are visible — + # this confirms what the subprocess will receive. + env_keys = sorted((options_for_hint.get("env") or {}).keys()) + mcp_servers = options_for_hint.get("mcp_servers") + log.info( + "pool_warm.hint_send user_id=%s agent_version=%s options_hash=%s" + " env_keys=%s mcp_servers_type=%s mcp_servers_value=%r", + user_id, body.agent_version, options_hash, + env_keys, type(mcp_servers).__name__, mcp_servers, + ) + + pool_client = _get_pool_client(request) + + hinted = False + try: + await pool_client.hint(user_id, options_hash, options_for_hint) + hinted = True + log.info( + "pool_warm.hint_ok user_id=%s options_hash=%s", + user_id, options_hash, + ) + except Exception as exc: + log.warning( + "pool_warm: pool hint failed user_id=%s agent_version=%s: %s", + user_id, + body.agent_version, + exc, + ) + + return {"hinted": hinted, "options_hash": options_hash} diff --git a/api/routes/preferences.py b/api/routes/preferences.py index f505b158..cb6de536 100644 --- a/api/routes/preferences.py +++ b/api/routes/preferences.py @@ -1,32 +1,94 @@ +from __future__ import annotations +from typing import Literal from pydantic import BaseModel from fastapi import APIRouter, Depends, HTTPException, Request from db.pool_middleware import get_db_conn from db.pg_queries.preferences import upsert_user_preference, get_user_preferences +from db.pg_queries.notifications import ( + get_notification_routing_with_defaults, + set_notification_routing, +) from api.auth import auth_dependency router = APIRouter(prefix="/user/preferences", tags=["preferences"]) # NOTE: This router is included with prefix="/api" in main.py → final paths are /api/user/preferences +# --------------------------------------------------------------------------- +# Notification routing validation types +# --------------------------------------------------------------------------- + +_VALID_NOTIFICATION_TYPES = Literal[ + "anchor_ping", "task_followup", "beacon", "meeting_event", "scheduling_update" +] +_RoutingMode = Literal["thread_by_key", "fixed", "bot_decides", "new_each"] +_RoutingPriority = Literal["normal", "important", "urgent"] +_RoutingChannel = Literal["telegram", "web", "discord", "slack"] + + +class NotificationRoutingEntry(BaseModel): + mode: _RoutingMode + priority: _RoutingPriority + external: list[_RoutingChannel] + key_template: str | None = None # only relevant for thread_by_key mode + conversation_id: str | None = None # only relevant for fixed mode + + +# Pydantic enforces that only these 5 keys appear in the dict. +NotificationRouting = dict[_VALID_NOTIFICATION_TYPES, NotificationRoutingEntry] + + +# --------------------------------------------------------------------------- +# Request / response models +# --------------------------------------------------------------------------- class PreferencesBody(BaseModel): theme: str | None = None mode: str | None = None + notification_routing: NotificationRouting | None = None + +# --------------------------------------------------------------------------- +# Routes +# --------------------------------------------------------------------------- @router.get("") -async def get_preferences(request: Request, _auth=Depends(auth_dependency), conn=Depends(get_db_conn)): +async def get_preferences( + request: Request, + _auth=Depends(auth_dependency), + conn=Depends(get_db_conn), +): user_id = request.state.user_id prefs = await get_user_preferences(conn, user_id) - return {"theme": prefs.get("theme"), "mode": prefs.get("mode")} + routing = await get_notification_routing_with_defaults(conn, user_id) + return { + "theme": prefs.get("theme"), + "mode": prefs.get("mode"), + "notification_routing": routing, + } @router.patch("") -async def patch_preferences(body: PreferencesBody, request: Request, _auth=Depends(auth_dependency), conn=Depends(get_db_conn)): - if body.theme is None and body.mode is None: - raise HTTPException(status_code=400, detail="At least one field (theme, mode) must be provided") +async def patch_preferences( + body: PreferencesBody, + request: Request, + _auth=Depends(auth_dependency), + conn=Depends(get_db_conn), +): + if body.theme is None and body.mode is None and body.notification_routing is None: + raise HTTPException( + status_code=400, + detail="At least one field (theme, mode, notification_routing) must be provided", + ) user_id = request.state.user_id if body.theme is not None: await upsert_user_preference(conn, user_id, "theme", body.theme) if body.mode is not None: await upsert_user_preference(conn, user_id, "mode", body.mode) + if body.notification_routing is not None: + # Serialize validated Pydantic models back to plain dicts for storage. + routing_dict = { + k: v.model_dump(exclude_none=True) + for k, v in body.notification_routing.items() + } + await set_notification_routing(conn, user_id, routing_dict) return {"ok": True} diff --git a/api/routes/settings.py b/api/routes/settings.py index a551c2b4..a984cf3a 100644 --- a/api/routes/settings.py +++ b/api/routes/settings.py @@ -22,5 +22,5 @@ async def list_settings(request: Request, _auth=Depends(auth_dependency), async def put_setting(key: str, body: SetSettingBody, request: Request, _auth=Depends(auth_dependency), conn: asyncpg.Connection = Depends(get_db_conn)): - await set_user_setting(conn, request.state.user_id, key, body.value) + await set_user_setting(conn, key, body.value) return {"ok": True} diff --git a/bot/agent_dispatch.py b/bot/agent_dispatch.py index 196ca7a8..81724d6a 100644 --- a/bot/agent_dispatch.py +++ b/bot/agent_dispatch.py @@ -4,31 +4,91 @@ requested agent_version: tether-agent-1.0 → existing JSON-mutation pipeline (handle_message) - tether-agent-2.0 → stub notice + 1.0 fallback (not yet wired) - tether-agent-2.5 → stub notice + 1.0 fallback (not yet wired) + tether-agent-2.0 → real LayerClient pipeline; 1.0 fallback on error/disabled + tether-agent-2.5 → premium session for paid/admin users; 1.0 fallback for free unknown / None → treated as tether-agent-2.0 (picker default) -The stub-then-fallback pattern ensures users selecting 2.0 or 2.5 still get -a response while the real pipelines are built out. The stub is delivered via -send_fn so it joins the 1.0 response in a single chunk frame on the client -(the WS handler accumulates all send_fn calls and joins them with "\\n\\n"). +The 2.0 pipeline calls the interactive-agent-layer service, which handles +streaming, tool-use translation, and permission gating. SSE events from the +layer are forwarded to the WS client via status_fn; the final response is +delivered via send_fn when turn_complete arrives. + +The 2.5 pipeline routes paid users and admin users to the premium handler +(Session + Beacon + RAG). Free users receive an upgrade notice and fall back +to tether-agent-1.0. Admin users bypass the subscription check entirely. + +Fallback behaviour: if the 2.0 layer is disabled (config) or unreachable +(HTTP error), dispatch silently falls back to handle_message so users always +get a response. Note: The Telegram polling path (bot/message_handler.py) calls handle_message directly — it bypasses this dispatcher (Telegram has no picker UI). """ from __future__ import annotations +import asyncio +import contextlib import logging from collections.abc import Callable from typing import Any +import httpx + from bot.message_handler import handle_message +from config.loader import config +from interactive_agent_layer.client import LayerClient logger = logging.getLogger(__name__) _DEFAULT_VERSION = "tether-agent-2.0" -_STUB_VERSIONS: frozenset[str] = frozenset({"tether-agent-2.0", "tether-agent-2.5"}) -_KNOWN_VERSIONS: frozenset[str] = _STUB_VERSIONS | {"tether-agent-1.0"} + + +class _SendTracker: + """Wraps a send_fn callback to record whether it was ever called. + + Used by _dispatch_v25 to detect whether the premium handler streamed + any output before raising an exception. If it did, we skip the 1.0 + fallback (handle_message) entirely to avoid double-send: the user has + already received premium output and calling handle_message would splice + a second, unrelated response on top of it. + """ + + def __init__(self, fn: Callable[[str], None]) -> None: + self._fn = fn + self.called: bool = False + + def __call__(self, text: str) -> None: + self.called = True + self._fn(text) + + +_KNOWN_VERSIONS: frozenset[str] = frozenset( + {"tether-agent-1.0", "tether-agent-2.0", "tether-agent-2.5"} +) + +# MCP tools available to tether-agent-2.0 (basic tether MCP only, no premium). +_V2_0_OPTIONS: dict[str, Any] = { + "model": "claude-haiku-4-5-20251001", + "allowed_tools": [ + "upsert_tasks", + "upsert_context", + "delete_tasks", + "delete_context", + "read_context", + "read_tasks", + "get_plan", + "get_anchors", + "search", + ], + "max_turns": 2, + "permission_mode": "auto", + "mcp_servers": ["tether"], # basic tether MCP only; no premium tools +} + + +def _layer_enabled() -> bool: + """Return True if the interactive-agent-layer is enabled in config.""" + return config.get_bool("agent_layer.enabled", True) def _stub_message(version: str) -> str: @@ -38,25 +98,341 @@ def _stub_message(version: str) -> str: ) +async def _dispatch_v2_0( + text: str, + send_fn: Callable[[str], None], + pool: Any, + user_id: str, + vault: Any = None, + status_fn: Any = None, + event_fn: Any = None, +) -> None: + """Run the tether-agent-2.0 pipeline via the interactive-agent-layer. + + Starts a layer session, runs one turn, and routes events: + - agent_text_delta / unknown event types → event_fn (async, for direct WS + forwarding; skips send_fn at turn_complete if any delta was sent) + - status / agent_action → status_fn + - turn_complete → send_fn(final_text) unless deltas were already streamed + + Falls back to handle_message (1.0 pipeline) when: + - agent_layer.enabled is false in config + - the layer service is unreachable or returns an HTTP error + + Session cleanup (end_session) is always attempted in the finally block so + the layer doesn't hold dangling sessions on error. On asyncio cancellation, + interrupt() is signalled before end_session so the pool can reclaim the + subprocess quickly. + """ + if not _layer_enabled(): + logger.info( + "dispatch_v2_0: agent_layer disabled, falling back to 1.0 user_id=%s", + user_id, + ) + await handle_message(text, send_fn, pool, user_id, vault=vault, status_fn=status_fn) + return + + base_url: str = config.get("agent_layer.base_url", "http://127.0.0.1:5003") + layer = LayerClient(base_url) + session_id: str | None = None + delta_sent = False + + # Inject the user's OAuth token so the pool subprocess can authenticate. + # Never mutate _V2_0_OPTIONS; always copy. + if vault is None: + logger.error( + "dispatch_v2_0: vault is not configured — VAULT_KEY must be set. " + "Falling back to 1.0 pipeline for user_id=%s", + user_id, + ) + await handle_message(text, send_fn, pool, user_id, vault=vault, status_fn=status_fn) + return + + try: + async with vault.materialize(user_id) as env_dict: + options: dict[str, Any] = {**_V2_0_OPTIONS, "env": dict(env_dict)} + except ValueError: + logger.warning( + "dispatch_v2_0: no vault credentials for user_id=%s" + " — user must connect Anthropic account. Falling back to 1.0 pipeline.", + user_id, + ) + await handle_message(text, send_fn, pool, user_id, vault=vault, status_fn=status_fn) + return + + try: + session_id = await layer.start_session( + user_id=user_id, + user_ws_id=user_id, # proxy: use user_id until per-connection IDs land + agent_version="tether-agent-2.0", + options=options, + user_message=text, + ) + + async for event in layer.turn(session_id, text): + etype = event.get("type") + + if etype == "turn_complete": + if not delta_sent: + send_fn(event.get("final_text", "")) + break + + if etype == "turn_error": + # Layer emitted an in-band error (e.g. pool exhausted) — log the + # real cause and fall back to 1.0 so the user still gets a response. + logger.warning( + "dispatch_v2_0: layer turn error, falling back to 1.0" + " user_id=%s: %s", + user_id, + event.get("message", "unknown"), + ) + await handle_message( + text, send_fn, pool, user_id, vault=vault, status_fn=status_fn + ) + return + + if etype in ("status", "agent_action"): + if status_fn is not None: + if etype == "status": + msg = event.get("message", "") + if msg: + await status_fn(msg) + else: + action = event.get("action", "") + if action: + await status_fn(action) + continue + + # agent_text_delta, permission_request, and any future event types + # are forwarded via event_fn for direct WS delivery. + if event_fn is not None: + await event_fn(event) + if etype == "agent_text_delta" and event.get("delta"): + delta_sent = True + + except asyncio.CancelledError: + if session_id is not None: + with contextlib.suppress(Exception): + await layer.interrupt(session_id) + raise + + except httpx.HTTPError as exc: + # Raw transport failure (layer unreachable, connection reset mid-stream, + # etc.) — distinct from turn_error which is an in-band error from a + # running layer. Both fall back to 1.0 so the user still gets a response. + logger.warning( + "dispatch_v2_0: layer turn transport error, falling back to 1.0" + " user_id=%s: %s", + user_id, + exc, + ) + await handle_message(text, send_fn, pool, user_id, vault=vault, status_fn=status_fn) + + finally: + if session_id is not None: + with contextlib.suppress(Exception): + await layer.end_session(session_id) + + +async def _dispatch_v25( + text: str, + send_fn: Callable[[str], None], + pool: Any, + user_id: str, + *, + vault: Any = None, + status_fn: Any = None, + is_admin: bool = False, + pool_client: Any = None, +) -> None: + """Handle tether-agent-2.5: premium session for paid/admin users, 1.0 fallback for free. + + Paid users and admin users are routed to the premium handler (Session + Beacon + RAG). + Admin users bypass the subscription check entirely — no subscription row required. + Free users (non-admin, non-paid) receive an upgrade notice and fall back to tether-agent-1.0. + If tether-premium is not installed, paid/admin users also fall back to 1.0. + + pool_client is the agent-pool-manager PoolClient. When provided it is threaded through + to the premium handler so PipelineBackend can acquire warm subprocesses instead of spawning + inline. When absent, it is created from config automatically so callers that don't have + it on hand still get pool routing. + + This function is a clean boundary that maps to a future HTTP endpoint in the + self-hosted premium access plan (phase P1+). + """ + import db.postgres as pg + from db.pg_queries.subscriptions import get_user_is_paid + + is_paid = False + if not is_admin: + # Admin users skip the subscription DB check entirely. + try: + async with pg.get_conn(pool, user_id) as conn: + is_paid = await get_user_is_paid(conn) + except Exception: + # Intentional fail-closed policy: on transient DB errors, treat the user + # as free-tier rather than granting premium access. Fail-open (granting + # access on error) was explicitly rejected as a security/billing risk. + # Sending a neutral "try again" message was also rejected because it + # gives users no actionable path. The 1.0 fallback ensures a response. + logger.warning( + "dispatch_v25: subscription check failed for user_id=%s — defaulting to free", + user_id, + ) + + if is_admin or is_paid: + # Resolve pool_client: use caller-supplied one, or create from config. + # Mirror the LayerClient pattern in _dispatch_v2_0: construct lazily from + # config rather than requiring callers to pass it explicitly. + effective_pool_client = pool_client + if effective_pool_client is None: + try: + from agent_pool_manager.client import from_config + from config.loader import config as _cfg + effective_pool_client = from_config(_cfg) + except Exception: + logger.warning( + "dispatch_v25: could not create pool_client from config for user_id=%s" + " — premium handler will run without pool routing", + user_id, + exc_info=True, + ) + + # Wrap send_fn so we can detect whether the premium handler streamed + # any output before raising. If it did, we must NOT run handle_message + # (1.0 fallback) — doing so would splice a second response on top of + # partial premium output the user has already received. + tracked_send = _SendTracker(send_fn) + try: + from tether_premium.register import get_premium_handler + from db.pg_queries import get_anchors + from bot.handler_utils import get_current_anchor + from bot.llm import _llm_env_extras, _llm_user_id + + async with pg.get_conn(pool, user_id) as conn: + anchors = await get_anchors(conn) + current_anchor = get_current_anchor(anchors) + + # Inner coroutine to call the handler — defined once and used in + # both the vault-materialized and no-vault branches below. + async def _invoke_handler() -> Any: + return await get_premium_handler()( + text, pool, user_id, anchors, current_anchor, + send_fn=tracked_send, status_fn=status_fn, + pool_client=effective_pool_client, + ) + + # Set _llm_user_id contextvar so PipelineBackend._complete_via_pool() uses + # the current request's user_id at call time, not the frozen one baked into + # the singleton LLMRouter/PipelineBackend instance. This is the fix for the + # singleton user_id capture bug: without this, all requests after the first + # would route pool handles to the first caller's user_id. + uid_token = _llm_user_id.set(user_id) + try: + # Materialize vault credentials and set _llm_env_extras so every LLM + # call inside the premium handler (PipelineBackend, AgentSDKBackend) + # inherits the user's OAuth token via the subprocess env. + # Pattern mirrors handle_message() in bot/message_handler.py. + if vault is not None: + async with vault.with_lock(user_id): + # Sentinel flag: True once we have entered the vault.materialize + # context manager. Used in except ValueError below to distinguish + # "no credentials stored" (flag=False, safe to call handler without + # vault) from "ValueError raised inside the handler itself" (flag=True, + # must re-raise — otherwise the handler is invoked a second time). + _vault_materialized = False + try: + async with vault.materialize(user_id) as env_extras: + _vault_materialized = True + token = _llm_env_extras.set(dict(env_extras)) + try: + response = await _invoke_handler() + finally: + _llm_env_extras.reset(token) + except ValueError: + if _vault_materialized: + # ValueError came from inside the handler, not from + # vault.materialize(). Re-raise so the outer except + # Exception handler deals with it (and runs 1.0 fallback). + raise + # No credentials stored for this user — run without env injection. + logger.warning( + "dispatch_v25: no vault credentials for user_id=%s" + " — running premium handler without OAuth env", + user_id, + ) + response = await _invoke_handler() + else: + response = await _invoke_handler() + finally: + _llm_user_id.reset(uid_token) + + tracked_send(response or "") + logger.info( + "dispatch_v25: response delivered via tracked_send, chars=%d user_id=%s", + len(response or ""), + user_id, + ) + return + except (ImportError, NotImplementedError) as exc: + # Log exc_info so the actual missing symbol or not-implemented site + # appears in the log — without this, the error is silent and the + # "premium not available" message is misleading (implies the package + # is absent when it may be an import symbol mismatch inside the handler). + logger.warning( + "dispatch_v25: premium handler raised %s for user_id=%s" + " — falling back to 1.0: %s", + type(exc).__name__, + user_id, + exc, + exc_info=True, + ) + except Exception: + logger.exception( + "dispatch_v25: premium handler raised for user_id=%s — falling back to 1.0", + user_id, + ) + + if tracked_send.called: + # Premium already streamed output to the user. Skip handle_message to + # avoid double-send. Running handle_message here would also risk DB + # side effects (task mutations etc.) on a message that premium already + # partially handled. + logger.warning( + "dispatch_v25: premium handler streamed then raised for user_id=%s" + " — suppressing 1.0 fallback to avoid double-send", + user_id, + ) + return + else: + send_fn( + "tether-agent-2.5 is available on the Pro plan — you're currently on " + "the free plan. Routing to tether-agent-1.0 for this message." + ) + + await handle_message(text, send_fn, pool, user_id, vault=vault, status_fn=status_fn) + + async def dispatch_message( agent_version: str | None, text: str, send_fn: Callable[[str], None], pool: Any, user_id: str, + *, vault: Any = None, status_fn: Any = None, + event_fn: Any = None, + is_admin: bool = False, ) -> None: """Dispatch a user message to the correct pipeline based on agent_version. For tether-agent-1.0, delegates directly to handle_message with no stub. - For tether-agent-2.0/2.5 (not yet wired), sends a stub notice via send_fn - and then falls back to the 1.0 pipeline so the user still gets a response. + For tether-agent-2.0, calls the interactive-agent-layer real pipeline with + a silent fallback to 1.0 on error or when the layer is disabled. + For tether-agent-2.5, routes to _dispatch_v25 (paid/admin = premium; free = 1.0 fallback). Unknown or None versions default to tether-agent-2.0 and log a warning. - Both the stub and the 1.0 response are accumulated by the WS handler's - capture_send_fn and joined into a single chunk frame on the client. - Args: agent_version: The version string from the WS message, or None if absent. text: The user message text. @@ -65,6 +441,10 @@ async def dispatch_message( user_id: Authenticated user ID. vault: Optional credential vault for per-user LLM auth. status_fn: Optional async callback for real-time status pushes. + event_fn: Optional async callback for streamed layer events (text deltas, + permission requests, etc.) forwarded directly to the WS client. + is_admin: When True, bypass subscription check for 2.5 dispatch (admin users + have no subscription row but must reach the premium handler). """ version = agent_version if agent_version in _KNOWN_VERSIONS else _DEFAULT_VERSION if version != agent_version: @@ -75,12 +455,29 @@ async def dispatch_message( user_id, ) - if version in _STUB_VERSIONS: - send_fn(_stub_message(version)) - logger.warning( - "dispatch_message: %s not yet wired — stub sent, falling back to 1.0 user_id=%s", - version, - user_id, + if version == "tether-agent-1.0": + await handle_message(text, send_fn, pool, user_id, vault=vault, status_fn=status_fn) + return + + if version == "tether-agent-2.0": + await _dispatch_v2_0( + text, send_fn, pool, user_id, + vault=vault, status_fn=status_fn, event_fn=event_fn, ) + return + if version == "tether-agent-2.5": + await _dispatch_v25( + text, send_fn, pool, user_id, + vault=vault, status_fn=status_fn, is_admin=is_admin, + ) + return + + # Catch-all for any future known versions not yet wired + send_fn(_stub_message(version)) + logger.warning( + "dispatch_message: %s not yet wired — stub sent, falling back to 1.0 user_id=%s", + version, + user_id, + ) await handle_message(text, send_fn, pool, user_id, vault=vault, status_fn=status_fn) diff --git a/bot/llm.py b/bot/llm.py index 3dab8975..37b569bc 100644 --- a/bot/llm.py +++ b/bot/llm.py @@ -24,6 +24,16 @@ "_llm_env_extras", default=None ) +# Per-request user ID — set by _dispatch_v25 before invoking the premium handler so +# that PipelineBackend.complete() uses the *current* user_id at call time instead of +# the frozen one baked into the singleton LLMRouter/PipelineBackend instance. +# This is the fix for the singleton user_id capture bug: without this contextvar, +# all requests after the first would use the first caller's user_id when acquiring +# pool handles. Never set at module level; always use .set()/.reset(). +_llm_user_id: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "_llm_user_id", default=None +) + logger = logging.getLogger(__name__) @dataclass @@ -64,11 +74,122 @@ def is_available(self) -> bool: ... # --------------------------------------------------------------------------- class PipelineBackend(LLMBackend): - """Invokes the Claude agent SDK. Always available as fallback.""" + """Invokes the Claude agent SDK. + + When constructed with pool_client + user_id, routes through the agent-pool-manager + (acquire → query_stream → release) instead of spawning a subprocess inline. This + is the preferred path for the 2.5 premium pipeline where subprocesses are kept warm + and reused across requests. + + Without pool_client (default), falls back to the original inline SDK spawn. This + preserves backward compatibility for the 1.0 pipeline and any caller that constructs + PipelineBackend directly without pool wiring. + """ + + def __init__( + self, + pool_client: object | None = None, + user_id: str | None = None, + ) -> None: + self._pool_client = pool_client + self._user_id = user_id def is_available(self) -> bool: return True + def _build_prompt( + self, + messages: list[dict], + system: str | list[str], + ) -> str: + """Flatten system + messages into a single text prompt.""" + parts = [] + if isinstance(system, list): + parts.append("\n".join(system)) + else: + parts.append(system) + for msg in messages: + role = msg.get("role", "user") + content = msg.get("content", "") + parts.append(f"\n[{role}]\n{content}") + return "\n".join(parts) + + async def _complete_via_pool( + self, + prompt: str, + model: str, + env: dict[str, str], + ) -> str: + """Acquire a warm handle from the pool, stream result, release. + + Uses _llm_user_id contextvar at call time (not self._user_id) so that the + singleton PipelineBackend correctly routes each request to its own user's pool + handle rather than the first caller's user_id that was baked in at construction. + + Only accumulates text from the final result event (subtype=success). The pool + also emits intermediate assistant events with the same text — using both sources + would double-count the output. + + reusable is set to True only on clean completion; on any error the handle is + released with reusable=False so the pool knows not to reuse a potentially + corrupted subprocess. + """ + from interactive_agent_layer.session import _stable_options_hash + + # Prefer contextvar user_id (set by _dispatch_v25 per-request) over the + # instance attribute (frozen at singleton construction time). + effective_user_id = _llm_user_id.get() or self._user_id + + options: dict = { + "model": model, + "permission_mode": "bypassPermissions", + "env": env, + } + options_hash = _stable_options_hash(options) + + handle_id = await self._pool_client.acquire( # type: ignore[union-attr] + effective_user_id, options_hash, options + ) + ok = False + try: + output: str = "" + async for event in self._pool_client.query_stream(handle_id, prompt): # type: ignore[union-attr] + etype = event.get("type") + # Use only the final result event — it is the authoritative complete + # response from the pool subprocess. Intermediate assistant events + # carry the same text; accumulating both would double the output. + if etype == "result" and event.get("subtype") == "success": + output = event.get("result", "") + break + ok = True + return output.strip() + finally: + # reusable=True only on clean completion; poisoned handles must not re-enter pool + await self._pool_client.release(handle_id, reusable=ok) # type: ignore[union-attr] + + async def _complete_inline( + self, + prompt: str, + model: str, + env: dict[str, str], + ) -> str: + """Original inline SDK spawn path (fallback when no pool_client).""" + from claude_agent_sdk import query, ClaudeAgentOptions, AssistantMessage, TextBlock + + opts = ClaudeAgentOptions( + model=model, + permission_mode="bypassPermissions", + env=env, + ) + + parts: list[str] = [] + async for msg in query(prompt=prompt, options=opts): + if isinstance(msg, AssistantMessage): + for block in msg.content: + if isinstance(block, TextBlock): + parts.append(block.text) + return "".join(parts).strip() + async def complete( self, messages: list[dict], @@ -79,40 +200,28 @@ async def complete( thinking_budget: int = 8000, max_tokens: int = 8096, ) -> LLMResponse: - from claude_agent_sdk import query, ClaudeAgentOptions, AssistantMessage, TextBlock - - prompt_parts = [] - if isinstance(system, list): - prompt_parts.append("\n".join(system)) - else: - prompt_parts.append(system) - for msg in messages: - role = msg.get("role", "user") - content = msg.get("content", "") - prompt_parts.append(f"\n[{role}]\n{content}") - prompt = "\n".join(prompt_parts) + prompt = self._build_prompt(messages, system) env: dict[str, str] = {} extras = _llm_env_extras.get() if extras: env.update(extras) - opts = ClaudeAgentOptions( - model=model, - permission_mode="bypassPermissions", - env=env, - ) + # Determine effective user_id — prefer the per-request contextvar (set by + # _dispatch_v25) over the instance attribute (frozen at construction for singletons). + effective_user_id = _llm_user_id.get() or self._user_id - async def _collect() -> str: - parts: list[str] = [] - async for msg in query(prompt=prompt, options=opts): - if isinstance(msg, AssistantMessage): - for block in msg.content: - if isinstance(block, TextBlock): - parts.append(block.text) - return "".join(parts).strip() + if self._pool_client is not None and effective_user_id is not None: + output = await asyncio.wait_for( + self._complete_via_pool(prompt, model, env), + timeout=180, + ) + else: + output = await asyncio.wait_for( + self._complete_inline(prompt, model, env), + timeout=180, + ) - output = await asyncio.wait_for(_collect(), timeout=180) return LLMResponse( content=output, tool_calls=[], diff --git a/cloud-version.txt b/cloud-version.txt index 098db4a5..f084e93d 100644 --- a/cloud-version.txt +++ b/cloud-version.txt @@ -1 +1 @@ -0.3.0a1 +0.7.0a3 diff --git a/config/agent_translations.yaml b/config/agent_translations.yaml index a4c42ba3..bb604d84 100644 --- a/config/agent_translations.yaml +++ b/config/agent_translations.yaml @@ -13,6 +13,44 @@ search: type: background phrase: "Searching for '{query}'" +# Memory-context tools (2.5 agent) +read_context: + type: background + phrase: "Looking at {node_title}" + +read_memory: + type: background + phrase: "Checking memory" + +read_node_memory: + type: background + phrase: "Reading notes on {node_title}" + +search_context: + type: background + phrase: "Searching context" + +search_memory: + type: background + phrase: "Searching memory" + +grep_context: + type: background + phrase: "Scanning context" + +# MCP-namespaced aliases used by 2.0 agent via tether MCP server +mcp__tether__upsert_tasks: + type: background + phrase: "Updating tasks" + +mcp__tether__get_anchors: + type: background + phrase: "Reading your schedule" + +mcp__tether__get_plan: + type: background + phrase: "Reading your plan" + # Premium internal tools — kept opaque consult_advisor: type: background_hidden @@ -34,29 +72,49 @@ send_status_update: type: passthrough # Channel 2 — user-action tools (permission-eligible) +# kind: "user_section_edit" | "destructive" | "read_out_of_scope" + upsert_tasks: type: user_action phrase_short: "Updating tasks" permission_summary: "Update {count} tasks" permission_detail_field: "tasks" + kind: "user_section_edit" delete_tasks: type: user_action phrase_short: "Removing tasks" permission_summary: "Delete {count} items" permission_detail_field: "operations" + kind: "destructive" upsert_context: type: user_action phrase_short: "Updating context" permission_summary: "Update context: {subject}" permission_detail_field: "nodes" + kind: "user_section_edit" delete_context: type: user_action phrase_short: "Removing context" permission_summary: "Delete context: {subject}" permission_detail_field: "operations" + kind: "destructive" + +write_node_memory: + type: user_action + phrase_short: "Noting: {title}" + permission_summary: "Save note: {title}" + permission_detail_field: "value" + kind: "user_section_edit" + +propose_user_memory_write: + type: user_action + phrase_short: "Proposing memory update" + permission_summary: "Propose memory: {key}" + permission_detail_field: "value" + kind: "user_section_edit" # Fallback _unknown: diff --git a/config/app_config.yaml b/config/app_config.yaml index fdf28756..2efcd35b 100644 --- a/config/app_config.yaml +++ b/config/app_config.yaml @@ -3,10 +3,10 @@ models: orchestrator: claude-sonnet-4-5 - meta_eval: claude-haiku-4-5 - quick_classifier: claude-haiku-4-5 + meta_eval: claude-haiku-4-5-20251001 + quick_classifier: claude-haiku-4-5-20251001 response_builder: claude-sonnet-4-5 - satisfaction_eval: claude-haiku-4-5 + satisfaction_eval: claude-haiku-4-5-20251001 pipeline: history_exchanges: 5 @@ -30,8 +30,12 @@ agent_pool: base_url: "http://127.0.0.1:5002" # override for remote pool deployments enabled: true # false = callers fall back to direct spawn (local dev) control_response_timeout_seconds: 60 # seconds pool waits for control_response before denying + initialize_timeout_ms: 240000 # CLAUDE_CODE_STREAM_CLOSE_TIMEOUT injected into subprocess env (ms) + home_dir_base: /var/lib/tether/claude-homes # isolated home dirs for subprocesses + home_dir_template: /etc/claude-home-template # seed state copied into each home dir at init llm: + layer_base_url: "http://127.0.0.1:5003" use_v3: false v2_fallback: true mcp_server_url: "http://localhost:5001/sse" diff --git a/db/migrations/versions/f6232ca1d274_merge_stream_a_and_stream_b_migration_.py b/db/migrations/versions/f6232ca1d274_merge_stream_a_and_stream_b_migration_.py new file mode 100644 index 00000000..48fd7f35 --- /dev/null +++ b/db/migrations/versions/f6232ca1d274_merge_stream_a_and_stream_b_migration_.py @@ -0,0 +1,28 @@ +"""merge stream A and stream B migration heads + +Revision ID: f6232ca1d274 +Revises: l1m2n3o4p5q6, k1l2m3n4o5p6 +Create Date: 2026-06-04 12:35:25.837234 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'f6232ca1d274' +down_revision: Union[str, Sequence[str], None] = ('l1m2n3o4p5q6', 'k1l2m3n4o5p6') +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + pass + + +def downgrade() -> None: + """Downgrade schema.""" + pass diff --git a/db/migrations/versions/h2b3c4d5e6f7_token_usage.py b/db/migrations/versions/h2b3c4d5e6f7_token_usage.py new file mode 100644 index 00000000..1dbba258 --- /dev/null +++ b/db/migrations/versions/h2b3c4d5e6f7_token_usage.py @@ -0,0 +1,38 @@ +"""Add token_usage table to record input/output tokens per user per turn. + +Feeds the trial counter display and eventual billing enforcement. +Token counts come from the agent SDK ResultMessage.usage dict, captured +on handle release by the pool manager. + +Revision ID: h2b3c4d5e6f7 +Revises: g1a2b3c4d5e6 +Create Date: 2026-05-21 +""" +from typing import Sequence, Union + +from alembic import op + +revision: str = "h2b3c4d5e6f7" +down_revision: Union[str, Sequence[str], None] = "g1a2b3c4d5e6" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute(""" + CREATE TABLE token_usage ( + id BIGSERIAL PRIMARY KEY, + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + recorded_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """) + # Index for per-user rollup queries (billing, trial enforcement) + op.execute(""" + CREATE INDEX idx_token_usage_user_id ON token_usage (user_id, recorded_at DESC) + """) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS token_usage") diff --git a/db/migrations/versions/i1j2k3l4m5n6_beacon_phase3_prep.py b/db/migrations/versions/i1j2k3l4m5n6_beacon_phase3_prep.py new file mode 100644 index 00000000..f61458e8 --- /dev/null +++ b/db/migrations/versions/i1j2k3l4m5n6_beacon_phase3_prep.py @@ -0,0 +1,296 @@ +"""Beacon Phase 3 prep — memory tables, conversation lifecycle extension + +Creates the 6 Beacon memory tables per spec §5: + - beacon_dispatches (L1 ephemeral — per-run dispatch records) + - beacon_decisions (L1 — triage decision audit log) + - beacon_suppressions (L1.5 — silent-exit suppression registry) + - beacon_memory (L2 — Beacon's freeform working memory) + - beacon_durable_memory (L3 — compacted long-term patterns) + - beacon_compaction_log (compaction audit) + +All tables have: + - user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE + - ENABLE ROW LEVEL SECURITY + FORCE ROW LEVEL SECURITY + - Policy: USING (user_id = current_setting('app.current_user_id', true)::uuid) + +Also extends conversations table: + - handle TEXT nullable, partial unique index (user_id, handle) WHERE handle IS NOT NULL + - expires_at TIMESTAMPTZ nullable (used by Beacon to auto-archive pending convs) + +Note: the state column is TEXT; 'pending' and 'rejected' are now valid values alongside +'open' and 'closed'. No DB-level enum constraint — state validation is enforced at the +API layer (ConversationPatch Pydantic model). + +Also backfills conversation_history.source: + Any existing rows with source='notification' are updated to source='chat'. + As of this migration the 'notification' value is deprecated — only 'chat', 'assistant', + and 'system' are valid going forward. The backfill is idempotent. + +Revision ID: i1j2k3l4m5n6 +Revises: h2b3c4d5e6f7 +Create Date: 2026-05-22 +""" +from typing import Sequence, Union + +from alembic import op + +revision: str = "i1j2k3l4m5n6" +down_revision: Union[str, Sequence[str], None] = "h2b3c4d5e6f7" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ------------------------------------------------------------------ + # L1 ephemeral — beacon_dispatches + # ------------------------------------------------------------------ + op.execute( + """ + CREATE TABLE beacon_dispatches ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + checkpoint_type TEXT NOT NULL, + mode TEXT NOT NULL, + conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + dispatched_at TIMESTAMPTZ DEFAULT now(), + prompt_summary TEXT, + priority TEXT NOT NULL DEFAULT 'normal', + dispatched_agent_role TEXT NOT NULL DEFAULT 'default', + state TEXT NOT NULL DEFAULT 'active', + last_message_at TIMESTAMPTZ, + concluded_at TIMESTAMPTZ, + evaluated_at TIMESTAMPTZ, + rejection_memory_key TEXT, + notes TEXT + ) + """ + ) + op.execute( + """ + CREATE INDEX beacon_dispatches_active + ON beacon_dispatches (user_id, state) + WHERE state = 'active' + """ + ) + op.execute( + """ + CREATE INDEX beacon_dispatches_pending_eval + ON beacon_dispatches (concluded_at) + WHERE state = 'concluded' AND evaluated_at IS NULL + """ + ) + op.execute("ALTER TABLE beacon_dispatches ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE beacon_dispatches FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY beacon_dispatches_isolation ON beacon_dispatches + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ + # L1 ephemeral — beacon_decisions + # ------------------------------------------------------------------ + op.execute( + """ + CREATE TABLE beacon_decisions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + checkpoint_type TEXT NOT NULL, + mode TEXT NOT NULL, + decided_at TIMESTAMPTZ DEFAULT now(), + action TEXT NOT NULL, + reason TEXT, + beacon_run_id UUID + ) + """ + ) + op.execute( + """ + CREATE INDEX beacon_decisions_user_recent + ON beacon_decisions (user_id, decided_at DESC) + """ + ) + op.execute("ALTER TABLE beacon_decisions ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE beacon_decisions FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY beacon_decisions_isolation ON beacon_decisions + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ + # L1.5 — beacon_suppressions + # ------------------------------------------------------------------ + op.execute( + """ + CREATE TABLE beacon_suppressions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + scope_key TEXT NOT NULL, + reason TEXT, + source TEXT NOT NULL, + created_at TIMESTAMPTZ DEFAULT now(), + expires_at TIMESTAMPTZ + ) + """ + ) + # Note: the spec suggested a partial index predicate of + # "WHERE expires_at IS NULL OR expires_at > now()" but Postgres requires + # index predicates to use only IMMUTABLE functions. now() is STABLE, not + # IMMUTABLE, so it cannot appear in a WHERE clause of an index definition. + # The index covers all rows; runtime queries filter by expires_at as needed. + op.execute( + """ + CREATE INDEX beacon_suppressions_lookup + ON beacon_suppressions (user_id, scope_key) + """ + ) + op.execute("ALTER TABLE beacon_suppressions ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE beacon_suppressions FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY beacon_suppressions_isolation ON beacon_suppressions + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ + # L2 — beacon_memory (Beacon's freeform working memory) + # ------------------------------------------------------------------ + op.execute( + """ + CREATE TABLE beacon_memory ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + key TEXT NOT NULL, + value TEXT NOT NULL, + updated_at TIMESTAMPTZ DEFAULT now(), + last_read_at TIMESTAMPTZ, + UNIQUE (user_id, key) + ) + """ + ) + op.execute( + """ + CREATE INDEX beacon_memory_user_key + ON beacon_memory (user_id, key text_pattern_ops) + """ + ) + op.execute("ALTER TABLE beacon_memory ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE beacon_memory FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY beacon_memory_isolation ON beacon_memory + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ + # L3 — beacon_durable_memory (compacted long-term patterns) + # ------------------------------------------------------------------ + op.execute( + """ + CREATE TABLE beacon_durable_memory ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + key TEXT NOT NULL, + value TEXT NOT NULL, + source TEXT NOT NULL, + evidence JSONB, + confidence TEXT DEFAULT 'medium', + created_at TIMESTAMPTZ DEFAULT now(), + updated_at TIMESTAMPTZ DEFAULT now(), + UNIQUE (user_id, key) + ) + """ + ) + op.execute( + """ + CREATE INDEX beacon_durable_memory_user_key + ON beacon_durable_memory (user_id, key text_pattern_ops) + """ + ) + op.execute("ALTER TABLE beacon_durable_memory ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE beacon_durable_memory FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY beacon_durable_memory_isolation ON beacon_durable_memory + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ + # Compaction audit log + # ------------------------------------------------------------------ + op.execute( + """ + CREATE TABLE beacon_compaction_log ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + trigger_type TEXT NOT NULL, + started_at TIMESTAMPTZ DEFAULT now(), + completed_at TIMESTAMPTZ, + surfaces_touched JSONB, + tokens_before INT, + tokens_after INT, + notes TEXT + ) + """ + ) + op.execute("ALTER TABLE beacon_compaction_log ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE beacon_compaction_log FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY beacon_compaction_log_isolation ON beacon_compaction_log + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ + # conversations table extension (spec §7.4) + # handle: slug for @-routing, unique per user when set + # expires_at: Beacon uses this to auto-archive pending convs + # ------------------------------------------------------------------ + op.execute( + "ALTER TABLE conversations ADD COLUMN handle TEXT" + ) + op.execute( + "ALTER TABLE conversations ADD COLUMN expires_at TIMESTAMPTZ" + ) + op.execute( + """ + CREATE UNIQUE INDEX conversations_handle_user + ON conversations (user_id, handle) + WHERE handle IS NOT NULL + """ + ) + + # ------------------------------------------------------------------ + # Backfill conversation_history.source (spec §8.3) + # 'notification' is deprecated — rows that carry it are updated to 'chat'. + # The backfill is idempotent: re-running on a clean DB is a no-op. + # ------------------------------------------------------------------ + op.execute( + """ + UPDATE conversation_history + SET source = 'chat' + WHERE source = 'notification' + """ + ) + + +def downgrade() -> None: + # conversations extensions + op.execute("DROP INDEX IF EXISTS conversations_handle_user") + op.execute("ALTER TABLE conversations DROP COLUMN IF EXISTS expires_at") + op.execute("ALTER TABLE conversations DROP COLUMN IF EXISTS handle") + + # Beacon tables (reverse creation order — foreign-key safe) + op.execute("DROP TABLE IF EXISTS beacon_compaction_log") + op.execute("DROP TABLE IF EXISTS beacon_durable_memory") + op.execute("DROP TABLE IF EXISTS beacon_memory") + op.execute("DROP TABLE IF EXISTS beacon_suppressions") + op.execute("DROP TABLE IF EXISTS beacon_decisions") + op.execute("DROP TABLE IF EXISTS beacon_dispatches") diff --git a/db/migrations/versions/j1k2l3m4n5o6_session_notes.py b/db/migrations/versions/j1k2l3m4n5o6_session_notes.py new file mode 100644 index 00000000..3ad69c5f --- /dev/null +++ b/db/migrations/versions/j1k2l3m4n5o6_session_notes.py @@ -0,0 +1,53 @@ +"""Add session_notes table — singleton per-user storage for bot session summaries. + +Replaces the previous pattern of writing to ~/.tether-config/.session-notes.md, +which breaks in Fly containers due to filesystem permissions. The table holds one +row per user; content is accumulated across sessions (append) or rewritten (LLM +summarization pass) by tether-premium's memory pipeline. + +Design: + - user_id UUID PRIMARY KEY — singleton row, user is the natural key + - content TEXT NOT NULL DEFAULT '' — empty string is a valid reset state + - updated_at TIMESTAMPTZ — refreshed on every upsert for auditing/expiry + +Security: + - ENABLE ROW LEVEL SECURITY + FORCE ROW LEVEL SECURITY + - Policy: USING (user_id = current_setting('app.current_user_id', true)::uuid) + - ON DELETE CASCADE from users table — notes are cleaned up when user is deleted + +Revision ID: j1k2l3m4n5o6 +Revises: i1j2k3l4m5n6 +Create Date: 2026-05-23 +""" +from typing import Sequence, Union + +from alembic import op + +revision: str = "j1k2l3m4n5o6" +down_revision: Union[str, Sequence[str], None] = "i1j2k3l4m5n6" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE session_notes ( + user_id UUID PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE, + content TEXT NOT NULL DEFAULT '', + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + op.execute("ALTER TABLE session_notes ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE session_notes FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY session_notes_isolation ON session_notes + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS session_notes") diff --git a/db/migrations/versions/k1l2m3n4o5p6_permission_grants.py b/db/migrations/versions/k1l2m3n4o5p6_permission_grants.py new file mode 100644 index 00000000..3ac4c880 --- /dev/null +++ b/db/migrations/versions/k1l2m3n4o5p6_permission_grants.py @@ -0,0 +1,80 @@ +"""Add permission_grants table — per-conversation tool permission grants. + +When a user approves a permission_request for a specific (target, kind), a grant +is stored here so the PermissionGate can skip the interactive flow on subsequent +calls within the same conversation. + +Design: + - Grants live as long as the conversation does; no explicit TTL. + - expires_at is NULL by default; may be used for future time-bounded grants. + - The (user_id, conversation_id, target, kind) tuple is the natural lookup key. + - ON DELETE CASCADE from users ensures no orphan grants. + +Security: + - ENABLE ROW LEVEL SECURITY + FORCE ROW LEVEL SECURITY + - Policy: USING (user_id = current_setting('app.current_user_id', true)::uuid) + +Revision ID: k1l2m3n4o5p6 +Revises: j1k2l3m4n5o6 +Create Date: 2026-06-02 +""" +from __future__ import annotations + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "k1l2m3n4o5p6" +down_revision = "j1k2l3m4n5o6" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "permission_grants", + sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True, + server_default=sa.text("gen_random_uuid()")), + sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("conversation_id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("target", sa.Text, nullable=False), + sa.Column("kind", sa.Text, nullable=False), + sa.Column("granted_at", sa.TIMESTAMP(timezone=True), nullable=False, + server_default=sa.text("now()")), + sa.Column("expires_at", sa.TIMESTAMP(timezone=True), nullable=True), + ) + + # Index for the primary lookup: does a grant exist for this (user, conv, target, kind)? + op.create_index( + "ix_permission_grants_lookup", + "permission_grants", + ["user_id", "conversation_id", "target", "kind"], + ) + + # Foreign key back to users (cascade on delete) + op.create_foreign_key( + "fk_permission_grants_user_id", + "permission_grants", + "users", + ["user_id"], + ["id"], + ondelete="CASCADE", + ) + + # RLS + op.execute("ALTER TABLE permission_grants ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE permission_grants FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY permission_grants_user_isolation ON permission_grants + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + +def downgrade() -> None: + op.execute("DROP POLICY IF EXISTS permission_grants_user_isolation ON permission_grants") + op.drop_index("ix_permission_grants_lookup", table_name="permission_grants") + op.drop_constraint("fk_permission_grants_user_id", "permission_grants", type_="foreignkey") + op.drop_table("permission_grants") diff --git a/db/migrations/versions/l1m2n3o4p5q6_memory_context_v2.py b/db/migrations/versions/l1m2n3o4p5q6_memory_context_v2.py new file mode 100644 index 00000000..e5c7d622 --- /dev/null +++ b/db/migrations/versions/l1m2n3o4p5q6_memory_context_v2.py @@ -0,0 +1,237 @@ +"""Memory-context v2 schema additions. + +Adds the infrastructure required for agent memory + context access system: + +1. node_sections gains: + - origin TEXT NOT NULL DEFAULT 'user' + Tracks who wrote the section: 'user' | 'conversation_agent' | 'system' + - visible_to_user BOOL NOT NULL DEFAULT true + Controls whether the section surfaces in user-facing context reads. + Bot-internal notes (origin='conversation_agent', visible_to_user=false) + are hidden from the UI but readable by agents with the right scope. + +2. node_data_summary — M-level summarization cache per node. + Beacon populates this table (Stream D); agents read from it (Stream A tools). + Each row is one M-level for one node. value carries the full key tree at + that level: {keys: {key_name: {value, expands_to_M(k+1)}}}. + Degrades gracefully when no rows exist (tools fall back to node_sections). + +3. user_memory — L2 working memory for user facts/patterns/preferences. + Written by Beacon; read by interactive 2.5 agents at session start. + Mirrored from beacon_memory (Beacon spec §5.1). + +4. user_durable_memory — L3 compacted long-term user patterns. + Written by Beacon compaction only (monthly+). + Mirrored from beacon_durable_memory (Beacon spec §5.1). + +5. pending_memory_writes — staging table for propose_user_memory_write. + MCP tool propose_user_memory_write() inserts here; Beacon evaluator + reviews and commits or discards after the conversation concludes. + +6. node_read_log — read-credit tracking per conversation. + Every read_context / read_node_memory call inserts a row. + write_node_memory consults this table for advisory read-before-write check + (v1: warns on violation; v2: hard block once Stream C is stable). + +All new tables have: + - ENABLE ROW LEVEL SECURITY + - FORCE ROW LEVEL SECURITY + - USING (user_id = current_setting('app.current_user_id', true)::uuid) + +Revision ID: l1m2n3o4p5q6 +Revises: j1k2l3m4n5o6 +Create Date: 2026-06-03 +Note: k1l2m3n4o5p6 is taken by Stream B (permission_grants); renamed to l1m2n3o4p5q6. +""" +from typing import Sequence, Union + +from alembic import op + +revision: str = "l1m2n3o4p5q6" +down_revision: Union[str, Sequence[str], None] = "j1k2l3m4n5o6" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ------------------------------------------------------------------ # + # 1. node_sections — add origin + visible_to_user # + # ------------------------------------------------------------------ # + op.execute( + """ + ALTER TABLE node_sections + ADD COLUMN origin TEXT NOT NULL DEFAULT 'user', + ADD COLUMN visible_to_user BOOL NOT NULL DEFAULT true + """ + ) + + # ------------------------------------------------------------------ # + # 2. node_data_summary # + # ------------------------------------------------------------------ # + op.execute( + """ + CREATE TABLE node_data_summary ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + node_id UUID NOT NULL REFERENCES context_nodes(id) ON DELETE CASCADE, + level_ordinal INT NOT NULL, + value JSONB NOT NULL, + abstract TEXT, + source_checksum TEXT NOT NULL, + generated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (node_id, level_ordinal) + ) + """ + ) + op.execute( + "CREATE INDEX idx_node_data_summary_node ON node_data_summary (node_id)" + ) + op.execute( + "CREATE INDEX idx_node_data_summary_user ON node_data_summary (user_id)" + ) + op.execute("ALTER TABLE node_data_summary ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE node_data_summary FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY node_data_summary_isolation ON node_data_summary + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ # + # 3. user_memory (L2 — working memory for user facts/patterns) # + # ------------------------------------------------------------------ # + op.execute( + """ + CREATE TABLE user_memory ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + key TEXT NOT NULL, + value TEXT NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + last_read_at TIMESTAMPTZ, + UNIQUE (user_id, key) + ) + """ + ) + op.execute( + "CREATE INDEX user_memory_user_key ON user_memory (user_id, key text_pattern_ops)" + ) + op.execute("ALTER TABLE user_memory ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE user_memory FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY user_memory_isolation ON user_memory + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ # + # 4. user_durable_memory (L3 — compacted long-term user patterns) # + # ------------------------------------------------------------------ # + op.execute( + """ + CREATE TABLE user_durable_memory ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + key TEXT NOT NULL, + value TEXT NOT NULL, + source TEXT NOT NULL, + evidence JSONB, + confidence TEXT NOT NULL DEFAULT 'medium', + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (user_id, key) + ) + """ + ) + op.execute( + "CREATE INDEX user_durable_memory_user_key ON user_durable_memory (user_id, key text_pattern_ops)" + ) + op.execute("ALTER TABLE user_durable_memory ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE user_durable_memory FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY user_durable_memory_isolation ON user_durable_memory + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ # + # 5. pending_memory_writes — staging for propose_user_memory_write # + # ------------------------------------------------------------------ # + op.execute( + """ + CREATE TABLE pending_memory_writes ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + conversation_id UUID REFERENCES conversations(id) ON DELETE SET NULL, + key TEXT NOT NULL, + value TEXT NOT NULL, + reason TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + reviewed_at TIMESTAMPTZ + ) + """ + ) + op.execute( + "CREATE INDEX pending_memory_writes_user_status ON pending_memory_writes (user_id, status)" + ) + op.execute("ALTER TABLE pending_memory_writes ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE pending_memory_writes FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY pending_memory_writes_isolation ON pending_memory_writes + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + # ------------------------------------------------------------------ # + # 6. node_read_log — per-conversation read-credit tracking # + # ------------------------------------------------------------------ # + op.execute( + """ + CREATE TABLE node_read_log ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + conversation_id UUID REFERENCES conversations(id) ON DELETE CASCADE, + node_id UUID NOT NULL REFERENCES context_nodes(id) ON DELETE CASCADE, + level_ordinal INT NOT NULL, + title TEXT, + read_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + op.execute( + """ + CREATE INDEX node_read_log_conv_node ON node_read_log + (conversation_id, node_id, level_ordinal) + """ + ) + op.execute( + "CREATE INDEX node_read_log_user ON node_read_log (user_id)" + ) + op.execute("ALTER TABLE node_read_log ENABLE ROW LEVEL SECURITY") + op.execute("ALTER TABLE node_read_log FORCE ROW LEVEL SECURITY") + op.execute( + """ + CREATE POLICY node_read_log_isolation ON node_read_log + USING (user_id = current_setting('app.current_user_id', true)::uuid) + """ + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS node_read_log") + op.execute("DROP TABLE IF EXISTS pending_memory_writes") + op.execute("DROP TABLE IF EXISTS user_durable_memory") + op.execute("DROP TABLE IF EXISTS user_memory") + op.execute("DROP TABLE IF EXISTS node_data_summary") + op.execute( + """ + ALTER TABLE node_sections + DROP COLUMN IF EXISTS visible_to_user, + DROP COLUMN IF EXISTS origin + """ + ) diff --git a/db/pg_auth_queries.py b/db/pg_auth_queries.py index 79729d30..58913f44 100644 --- a/db/pg_auth_queries.py +++ b/db/pg_auth_queries.py @@ -371,10 +371,10 @@ async def clear_telegram_bot_token( async def get_most_recent_telegram_user(conn: asyncpg.Connection) -> dict | None: - """Return the user with the most recent telegram link or message turn. + """Return the user with the most recent telegram link. - 'Most recent' = latest created_at in telegram_connections (for users who have a telegram_chat_id set). - Falls back to users ordered by conversation_history recency if available. + Orders by u.created_at (users table) — telegram_connections has no + created_at column (only user_id and telegram_chat_id). Returns dict with 'id' (str UUID) and 'telegram_chat_id' (str), or None. """ @@ -385,7 +385,7 @@ async def get_most_recent_telegram_user(conn: asyncpg.Connection) -> dict | None JOIN telegram_connections tc ON tc.user_id = u.id WHERE tc.telegram_chat_id IS NOT NULL AND tc.telegram_chat_id != '' - ORDER BY tc.created_at DESC NULLS LAST + ORDER BY u.created_at DESC NULLS LAST LIMIT 1 """ ) diff --git a/db/pg_queries/__init__.py b/db/pg_queries/__init__.py index cfee53d9..17eb1d86 100644 --- a/db/pg_queries/__init__.py +++ b/db/pg_queries/__init__.py @@ -33,7 +33,7 @@ from db.pg_queries.sections import ( get_sections, get_section, upsert_section, append_section, delete_section, list_section_files, create_section_file, rename_section_file, - reorder_section_files, search_sections, + reorder_section_files, search_sections, grep_sections, ) from db.pg_queries.milestones import ( create_milestone, get_milestones, patch_milestone, delete_milestone, @@ -75,3 +75,14 @@ get_user_is_paid, get_user_is_paid_by_id, ) +from db.pg_queries.memory import ( + list_user_memory, get_user_memory_entry, upsert_user_memory, delete_user_memory, + list_user_durable_memory, get_user_durable_memory_entry, upsert_user_durable_memory, + insert_pending_memory_write, get_pending_memory_write, + list_pending_memory_writes, review_pending_memory_write, +) +from db.pg_queries.node_memory import ( + get_node_summary, list_node_summary_levels, upsert_node_summary, + log_node_read, has_read_node_in_conversation, get_conversation_reads, + get_context_node_id_for_conversation, get_node_tree_distance, +) diff --git a/db/pg_queries/conversations.py b/db/pg_queries/conversations.py index 0efbc562..97a929a2 100644 --- a/db/pg_queries/conversations.py +++ b/db/pg_queries/conversations.py @@ -272,7 +272,7 @@ async def list_conversation_messages( """Return messages for a conversation, newest first, with cursor pagination. Uses `before_id` cursor: returns rows with id < before_id so callers can - paginate backwards through history. Returns {"items": [...], "has_more": bool}. + paginate backwards through history. Returns {"messages": [...], "has_more": bool}. Note: `before_id` is an integer-typed primary key on conversation_history — this cursor approach is stable under concurrent inserts, unlike offset-based @@ -292,7 +292,7 @@ async def list_conversation_messages( SELECT ch.id::text AS id, ch.role, - ch.body AS content, + ch.body, ch.conversation_id::text AS conversation_id, ch.ts AS created_at, ch.source, @@ -307,5 +307,41 @@ async def list_conversation_messages( ) has_more = len(rows) > limit - items = [dict(r) for r in rows[:limit]] - return {"items": items, "has_more": has_more} + messages = [dict(r) for r in rows[:limit]] + return {"messages": messages, "has_more": has_more} + + +async def list_conversations_index( + conn: asyncpg.Connection, + *, + user_id: str, +) -> list[dict]: + """Return a lightweight index of all conversations for the given user. + + Returns [{id, title, parent_context_node_id, state, priority, updated_at, message_count}]. + One query — no per-row N+1. Message bodies and full detail are excluded. + Used by the frontend to populate conversation trees quickly. + + state and priority are included so the frontend can render pending badges + and priority dots on first paint without a follow-up upgrade call. + """ + rows = await conn.fetch( + """ + SELECT + c.id::text AS id, + c.name AS title, + c.context_node_id::text AS parent_context_node_id, + c.state, + c.priority, + COALESCE(c.last_message_at, c.created_at) AS updated_at, + COUNT(ch.id)::int AS message_count + FROM conversations c + LEFT JOIN conversation_history ch ON ch.conversation_id = c.id + WHERE c.user_id = $1::uuid + GROUP BY c.id, c.name, c.context_node_id, c.state, c.priority, + c.last_message_at, c.created_at + ORDER BY COALESCE(c.last_message_at, c.created_at) DESC + """, + user_id, + ) + return [dict(r) for r in rows] diff --git a/db/pg_queries/memory.py b/db/pg_queries/memory.py new file mode 100644 index 00000000..55183c52 --- /dev/null +++ b/db/pg_queries/memory.py @@ -0,0 +1,341 @@ +"""Async Postgres queries — user_memory, user_durable_memory, pending_memory_writes. + +user_memory: L2 working memory (user facts, patterns, preferences). + Written by Beacon; read by interactive 2.5 agents at session start (full M=4). + Keyed by hierarchical convention: 'preferences/morning_routine', 'facts/work/role', etc. + +user_durable_memory: L3 compacted long-term user patterns. + Written by Beacon compaction only (monthly+). More stable than L2. + +pending_memory_writes: proposals from MCP tool propose_user_memory_write(). + Beacon post-conversation evaluator reviews and commits or discards. + +All three tables carry RLS — callers must have set app.current_user_id before +invoking these functions. +""" +from __future__ import annotations + +import asyncpg + + +# ───────────────────────────────────────────────────────────────────────────── +# user_memory (L2) +# ───────────────────────────────────────────────────────────────────────────── + + +async def list_user_memory( + conn: asyncpg.Connection, + *, + prefix: str | None = None, + M: int = 2, +) -> list[dict]: + """Return user_memory entries for the current RLS user. + + M controls how much data is returned per entry: + M=1 key only + M=2 key + ~50-char preview of value + M=3 key + ~200-char truncated value + M=4 key + full value + + Optional prefix filters by key prefix (e.g., 'preferences/'). + """ + if M == 1: + select = "key" + elif M == 2: + select = "key, left(value, 50) AS value" + elif M == 3: + select = "key, left(value, 200) AS value" + else: + select = "key, value, updated_at, last_read_at" + + if prefix: + rows = await conn.fetch( + f"SELECT {select} FROM user_memory WHERE key LIKE $1 || '%' ORDER BY key", + prefix, + ) + else: + rows = await conn.fetch( + f"SELECT {select} FROM user_memory ORDER BY key", + ) + return [dict(r) for r in rows] + + +async def get_user_memory_entry( + conn: asyncpg.Connection, + key: str, +) -> dict | None: + """Fetch a single user_memory entry by exact key. + + Updates last_read_at as a side effect (tracks when Beacon last accessed it). + """ + row = await conn.fetchrow( + """ + UPDATE user_memory + SET last_read_at = now() + WHERE user_id = current_setting('app.current_user_id', true)::uuid + AND key = $1 + RETURNING key, value, updated_at, last_read_at + """, + key, + ) + return dict(row) if row else None + + +async def upsert_user_memory( + conn: asyncpg.Connection, + key: str, + value: str, +) -> dict: + """Write a user_memory entry (insert or update). Returns the final row.""" + row = await conn.fetchrow( + """ + INSERT INTO user_memory (user_id, key, value, updated_at) + VALUES (current_setting('app.current_user_id', true)::uuid, $1, $2, now()) + ON CONFLICT (user_id, key) DO UPDATE + SET value = EXCLUDED.value, updated_at = now() + RETURNING key, value, updated_at + """, + key, value, + ) + return dict(row) + + +async def delete_user_memory( + conn: asyncpg.Connection, + key: str, +) -> bool: + """Delete a user_memory entry. Returns True if a row was deleted.""" + result = await conn.execute( + """ + DELETE FROM user_memory + WHERE user_id = current_setting('app.current_user_id', true)::uuid + AND key = $1 + """, + key, + ) + return result.split()[-1] != "0" + + +# ───────────────────────────────────────────────────────────────────────────── +# user_durable_memory (L3) +# ───────────────────────────────────────────────────────────────────────────── + + +async def list_user_durable_memory( + conn: asyncpg.Connection, + *, + prefix: str | None = None, + M: int = 2, +) -> list[dict]: + """Return user_durable_memory entries for the current RLS user. + + Same M-level semantics as list_user_memory. + """ + if M == 1: + select = "key" + elif M == 2: + select = "key, left(value, 50) AS value" + elif M == 3: + select = "key, left(value, 200) AS value" + else: + select = "key, value, source, confidence, created_at, updated_at" + + if prefix: + rows = await conn.fetch( + f"SELECT {select} FROM user_durable_memory WHERE key LIKE $1 || '%' ORDER BY key", + prefix, + ) + else: + rows = await conn.fetch( + f"SELECT {select} FROM user_durable_memory ORDER BY key", + ) + return [dict(r) for r in rows] + + +async def get_user_durable_memory_entry( + conn: asyncpg.Connection, + key: str, +) -> dict | None: + """Fetch a single user_durable_memory entry by exact key (full data).""" + row = await conn.fetchrow( + """ + SELECT key, value, source, evidence, confidence, created_at, updated_at + FROM user_durable_memory + WHERE user_id = current_setting('app.current_user_id', true)::uuid + AND key = $1 + """, + key, + ) + return dict(row) if row else None + + +async def upsert_user_durable_memory( + conn: asyncpg.Connection, + key: str, + value: str, + source: str, + *, + evidence: dict | None = None, + confidence: str = "medium", +) -> dict: + """Write a user_durable_memory entry (insert or update). Returns the final row.""" + row = await conn.fetchrow( + """ + INSERT INTO user_durable_memory (user_id, key, value, source, evidence, confidence, updated_at) + VALUES ( + current_setting('app.current_user_id', true)::uuid, + $1, $2, $3, $4, $5, now() + ) + ON CONFLICT (user_id, key) DO UPDATE + SET value = EXCLUDED.value, + source = EXCLUDED.source, + evidence = EXCLUDED.evidence, + confidence = EXCLUDED.confidence, + updated_at = now() + RETURNING key, value, source, confidence, updated_at + """, + key, value, source, evidence, confidence, + ) + return dict(row) + + +# ───────────────────────────────────────────────────────────────────────────── +# pending_memory_writes +# ───────────────────────────────────────────────────────────────────────────── + + +async def insert_pending_memory_write( + conn: asyncpg.Connection, + key: str, + value: str, + reason: str, + *, + conversation_id: str | None = None, +) -> str: + """Insert a pending_memory_write proposal. Returns the new UUID.""" + import uuid as _uuid + conv_uuid = _uuid.UUID(conversation_id) if conversation_id else None + row = await conn.fetchrow( + """ + INSERT INTO pending_memory_writes (user_id, conversation_id, key, value, reason) + VALUES ( + current_setting('app.current_user_id', true)::uuid, + $1, $2, $3, $4 + ) + RETURNING id::text + """, + conv_uuid, key, value, reason, + ) + return row["id"] + + +async def get_pending_memory_write( + conn: asyncpg.Connection, + proposal_id: str, +) -> dict | None: + """Fetch a single pending_memory_write by ID.""" + import uuid as _uuid + row = await conn.fetchrow( + """ + SELECT id::text, key, value, reason, status, conversation_id::text, created_at, reviewed_at + FROM pending_memory_writes + WHERE id = $1::uuid + AND user_id = current_setting('app.current_user_id', true)::uuid + """, + _uuid.UUID(proposal_id), + ) + return dict(row) if row else None + + +async def list_pending_memory_writes( + conn: asyncpg.Connection, + *, + status: str = "pending", +) -> list[dict]: + """Return pending_memory_write proposals with the given status.""" + rows = await conn.fetch( + """ + SELECT id::text, key, value, reason, status, conversation_id::text, created_at + FROM pending_memory_writes + WHERE user_id = current_setting('app.current_user_id', true)::uuid + AND status = $1 + ORDER BY created_at DESC + """, + status, + ) + return [dict(r) for r in rows] + + +async def review_pending_memory_write( + conn: asyncpg.Connection, + proposal_id: str, + new_status: str, +) -> bool: + """Mark a proposal as 'accepted' or 'rejected'. Returns True if found.""" + import uuid as _uuid + result = await conn.execute( + """ + UPDATE pending_memory_writes + SET status = $1, reviewed_at = now() + WHERE id = $2::uuid + AND user_id = current_setting('app.current_user_id', true)::uuid + """, + new_status, _uuid.UUID(proposal_id), + ) + return result.split()[-1] != "0" + + +# ───────────────────────────────────────────────────────────────────────────── +# Search +# ───────────────────────────────────────────────────────────────────────────── + + +async def search_user_memory( + conn: asyncpg.Connection, + query: str, + scope: str = "user", + limit: int = 20, +) -> list[dict]: + """Search user memory entries by key or value using ILIKE. + + scope: 'user' → user_memory (L2), 'user_durable' → user_durable_memory (L3). + query: Search string — matched against key (exact, prefix, substring) and value. + + Returns entries ordered by relevance: + score 2 — exact key match + score 1 — key prefix or key substring match + score 0 — value-only match + + Each entry: {key, value, score}. + """ + _VALID_SCOPES = {"user": "user_memory", "user_durable": "user_durable_memory"} + if scope not in _VALID_SCOPES: + raise ValueError(f"invalid scope: {scope!r}; must be 'user' or 'user_durable'") + table = _VALID_SCOPES[scope] + pattern = f"%{query}%" + + rows = await conn.fetch( + f""" + SELECT + key, + value, + CASE + WHEN key = $1 THEN 2 + WHEN key ILIKE $2 THEN 1 + ELSE 0 + END AS score + FROM {table} + WHERE key ILIKE $2 OR value ILIKE $2 + ORDER BY score DESC, key + LIMIT $3 + """, + query, pattern, limit, + ) + return [ + { + "key": r["key"], + "value": r["value"], + "score": r["score"], + } + for r in rows + ] diff --git a/db/pg_queries/node_memory.py b/db/pg_queries/node_memory.py new file mode 100644 index 00000000..161af882 --- /dev/null +++ b/db/pg_queries/node_memory.py @@ -0,0 +1,225 @@ +"""Async Postgres queries — node_data_summary, node_read_log, and scope helpers. + +node_data_summary: M-level summarization cache per context_node. + Populated by Beacon (Stream D summarization — NOT this module's concern). + Read by MCP tools to serve M-level detail without loading all sections. + Degrades gracefully when no rows exist (callers fall back to node_sections). + +node_read_log: per-conversation read-credit tracking. + Inserted by read_context / read_node_memory on every access. + Consulted by write_node_memory for advisory read-before-write check. + +All queries use RLS — callers must have set app.current_user_id. +""" +from __future__ import annotations + +import uuid as _uuid + +import asyncpg + + +# ───────────────────────────────────────────────────────────────────────────── +# node_data_summary +# ───────────────────────────────────────────────────────────────────────────── + + +async def get_node_summary( + conn: asyncpg.Connection, + node_id: str, + level_ordinal: int, +) -> dict | None: + """Return the summary row for (node_id, level_ordinal), or None if absent. + + Callers should degrade gracefully when None is returned: + - For M < last: no summary cached yet (Beacon hasn't summarized) + - For M = last: fall back to node_sections directly + """ + row = await conn.fetchrow( + """ + SELECT node_id::text, level_ordinal, value, abstract, source_checksum, generated_at + FROM node_data_summary + WHERE node_id = $1::uuid AND level_ordinal = $2 + """, + _uuid.UUID(node_id), level_ordinal, + ) + return dict(row) if row else None + + +async def list_node_summary_levels( + conn: asyncpg.Connection, + node_id: str, +) -> list[int]: + """Return all available level_ordinal values for a node (sorted ascending). + + Useful for read_context to know which M-levels are cached vs. must fall back. + """ + rows = await conn.fetch( + """ + SELECT level_ordinal + FROM node_data_summary + WHERE node_id = $1::uuid + ORDER BY level_ordinal + """, + _uuid.UUID(node_id), + ) + return [r["level_ordinal"] for r in rows] + + +async def upsert_node_summary( + conn: asyncpg.Connection, + node_id: str, + level_ordinal: int, + value: dict, + source_checksum: str, + *, + abstract: str | None = None, +) -> dict: + """Upsert a summary row. Called by Beacon summarization (Stream D). + + Not called by read tools — this is write-only from the summarizer side. + """ + row = await conn.fetchrow( + """ + INSERT INTO node_data_summary + (user_id, node_id, level_ordinal, value, abstract, source_checksum, generated_at) + VALUES ( + current_setting('app.current_user_id', true)::uuid, + $1::uuid, $2, $3::jsonb, $4, $5, now() + ) + ON CONFLICT (node_id, level_ordinal) DO UPDATE + SET value = EXCLUDED.value, + abstract = EXCLUDED.abstract, + source_checksum = EXCLUDED.source_checksum, + generated_at = now() + RETURNING node_id::text, level_ordinal, abstract, source_checksum, generated_at + """, + _uuid.UUID(node_id), level_ordinal, value, abstract, source_checksum, + ) + return dict(row) + + +# ───────────────────────────────────────────────────────────────────────────── +# node_read_log +# ───────────────────────────────────────────────────────────────────────────── + + +async def log_node_read( + conn: asyncpg.Connection, + node_id: str, + level_ordinal: int, + *, + conversation_id: str | None = None, + title: str | None = None, +) -> None: + """Record that the current user read node_id at level_ordinal in this conversation. + + Called by read_context and read_node_memory on every access. + conversation_id may be None for admin/debug calls without conversation context. + """ + conv_uuid = _uuid.UUID(conversation_id) if conversation_id else None + await conn.execute( + """ + INSERT INTO node_read_log (user_id, conversation_id, node_id, level_ordinal, title) + VALUES ( + current_setting('app.current_user_id', true)::uuid, + $1, $2::uuid, $3, $4 + ) + """, + conv_uuid, _uuid.UUID(node_id), level_ordinal, title, + ) + + +async def has_read_node_in_conversation( + conn: asyncpg.Connection, + node_id: str, + conversation_id: str, +) -> bool: + """Return True if any read of node_id was logged for this conversation. + + Used by write_node_memory advisory read-before-write check: + v1: log WARNING if False but allow the write + v2: hard block if False once Stream C callers always provide conversation_id + """ + count = await conn.fetchval( + """ + SELECT COUNT(*) + FROM node_read_log + WHERE user_id = current_setting('app.current_user_id', true)::uuid + AND conversation_id = $1::uuid + AND node_id = $2::uuid + """, + _uuid.UUID(conversation_id), _uuid.UUID(node_id), + ) + return (count or 0) > 0 + + +async def get_context_node_id_for_conversation( + conn: asyncpg.Connection, + conversation_id: str, +) -> str | None: + """Return the context_node_id linked to a conversation, or None if unlinked.""" + row = await conn.fetchrow( + "SELECT context_node_id::text FROM conversations WHERE id = $1::uuid", + _uuid.UUID(conversation_id), + ) + if not row: + return None + return row["context_node_id"] + + +async def get_node_tree_distance( + conn: asyncpg.Connection, + from_id: str, + to_id: str, + max_N: int, +) -> int | None: + """Return the tree distance (edges) between two context nodes, or None if > max_N. + + Uses a recursive CTE that walks both parent and child edges, bounded by max_N. + Returns 0 if from_id == to_id. + """ + if from_id == to_id: + return 0 + + row = await conn.fetchrow( + """ + WITH RECURSIVE reachable(node_id, dist) AS ( + SELECT $1::uuid, 0 + UNION ALL + SELECT + CASE WHEN cn.parent_id = r.node_id THEN cn.id + ELSE cn.parent_id + END, + r.dist + 1 + FROM reachable r + JOIN context_nodes cn + ON (cn.id = r.node_id AND cn.parent_id IS NOT NULL) + OR (cn.parent_id = r.node_id) + WHERE r.dist < $3 + ) + SELECT dist FROM reachable WHERE node_id = $2::uuid LIMIT 1 + """, + _uuid.UUID(from_id), _uuid.UUID(to_id), max_N, + ) + return row["dist"] if row else None + + +async def get_conversation_reads( + conn: asyncpg.Connection, + conversation_id: str, +) -> list[dict]: + """Return all node read credits for a conversation (newest first). + + Used for diagnostics and advisory enforcement reports. + """ + rows = await conn.fetch( + """ + SELECT node_id::text, level_ordinal, title, read_at + FROM node_read_log + WHERE user_id = current_setting('app.current_user_id', true)::uuid + AND conversation_id = $1::uuid + ORDER BY read_at DESC + """, + _uuid.UUID(conversation_id), + ) + return [dict(r) for r in rows] diff --git a/db/pg_queries/nodes.py b/db/pg_queries/nodes.py index 5aef282c..22e1593a 100644 --- a/db/pg_queries/nodes.py +++ b/db/pg_queries/nodes.py @@ -154,6 +154,45 @@ async def get_all_node_paths(conn: asyncpg.Connection) -> list[str]: return [r["path"] for r in rows] +async def list_nodes_index(conn: asyncpg.Connection, *, user_id: str) -> list[dict]: + """Return a lightweight index of all non-archived context nodes for a user. + + Returns [{id, title, parent_id, path, child_count}]. + One recursive CTE query — no per-row N+1. Section data is excluded. + Used by the frontend to populate the node tree quickly. + + Filters explicitly by user_id in addition to the RLS policy on the + connection, matching the pattern used by list_conversations_index. + """ + rows = await conn.fetch( + """ + WITH RECURSIVE tree(id, parent_id, name, path) AS ( + SELECT id, parent_id, name, name::text AS path + FROM context_nodes + WHERE parent_id IS NULL AND archived = FALSE AND user_id = $1::uuid + UNION ALL + SELECT cn.id, cn.parent_id, cn.name, tree.path || '/' || cn.name + FROM context_nodes cn + JOIN tree ON cn.parent_id = tree.id + WHERE cn.archived = FALSE AND cn.user_id = $1::uuid + ) + SELECT + t.id::text AS id, + t.name AS title, + t.parent_id::text AS parent_id, + t.path, + COUNT(c.id)::int AS child_count + FROM tree t + LEFT JOIN context_nodes c ON c.parent_id = t.id AND c.archived = FALSE + AND c.user_id = $1::uuid + GROUP BY t.id, t.name, t.parent_id, t.path + ORDER BY t.path + """, + user_id, + ) + return [dict(r) for r in rows] + + async def get_children( conn: asyncpg.Connection, parent_id: str | None = None, @@ -328,6 +367,61 @@ async def get_milestone_nodes( return [_node(r) for r in rows] +# ── Hop distance ───────────────────────────────────────────────────────────── + + +async def get_node_hop_distance( + conn: asyncpg.Connection, + from_node_id: str, + to_node_id: str, +) -> int | None: + """Return the undirected tree distance (hop count) between two context nodes. + + Uses an LCA (lowest common ancestor) approach: walk ancestor chains from + both nodes and find the shortest combined path via their common ancestor. + + Returns: + int — number of hops on the shortest tree path. + None — nodes are in separate trees (no common ancestor). + + Performance note: for large trees this CTE scans O(depth) rows per node. + A materialized path column or a dedicated ancestor table would reduce this + to O(1) lookups at the cost of write overhead. Consider adding + ``ltree`` indexing or a closure table if this becomes a hot path. + """ + result = await conn.fetchval( + """ + WITH RECURSIVE + from_ancestors(id, depth) AS ( + SELECT id, 0 AS depth + FROM context_nodes + WHERE id = $1 + UNION ALL + SELECT cn.parent_id, fa.depth + 1 + FROM context_nodes cn + JOIN from_ancestors fa ON cn.id = fa.id + WHERE cn.parent_id IS NOT NULL + ), + to_ancestors(id, depth) AS ( + SELECT id, 0 AS depth + FROM context_nodes + WHERE id = $2 + UNION ALL + SELECT cn.parent_id, ta.depth + 1 + FROM context_nodes cn + JOIN to_ancestors ta ON cn.id = ta.id + WHERE cn.parent_id IS NOT NULL + ) + SELECT MIN(fa.depth + ta.depth) + FROM from_ancestors fa + JOIN to_ancestors ta ON fa.id = ta.id + """, + _uuid.UUID(from_node_id), + _uuid.UUID(to_node_id), + ) + return int(result) if result is not None else None + + # ── Node-task linking ───────────────────────────────────────────────────────── async def link_task_to_node(conn: asyncpg.Connection, node_id: str, task_id: str) -> None: diff --git a/db/pg_queries/sections.py b/db/pg_queries/sections.py index 76972b8c..91d95bd1 100644 --- a/db/pg_queries/sections.py +++ b/db/pg_queries/sections.py @@ -8,7 +8,8 @@ async def get_sections(conn: asyncpg.Connection, node_id: str) -> list[dict]: rows = await conn.fetch( """ - SELECT section_type, name, body, position, version, updated_at + SELECT section_type, name, body, position, version, updated_at, + origin, visible_to_user FROM node_sections WHERE node_id = $1 ORDER BY section_type, position """, @@ -22,7 +23,8 @@ async def get_section( ) -> dict | None: row = await conn.fetchrow( """ - SELECT section_type, name, body, position, version, updated_at + SELECT section_type, name, body, position, version, updated_at, + origin, visible_to_user FROM node_sections WHERE node_id = $1 AND section_type = $2 AND name = $3 """, @@ -37,7 +39,15 @@ async def upsert_section( section_type: str, body: str, name: str = "main", + *, + origin: str = "user", + visible_to_user: bool = True, ) -> dict: + """Insert or replace a node section. + + origin: 'user' | 'conversation_agent' | 'system' — who authored this section. + visible_to_user: False hides the section from user-facing reads. + """ nid = _uuid.UUID(node_id) user_uuid = await conn.fetchval( "SELECT current_setting('app.current_user_id', true)::uuid" @@ -48,13 +58,19 @@ async def upsert_section( ) row = await conn.fetchrow( """ - INSERT INTO node_sections (user_id, node_id, section_type, name, body, position, updated_at) - VALUES ($1, $2, $3, $4, $5, $6, now()) + INSERT INTO node_sections + (user_id, node_id, section_type, name, body, position, + origin, visible_to_user, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now()) ON CONFLICT (node_id, section_type, name) DO UPDATE - SET body = EXCLUDED.body, updated_at = now(), version = node_sections.version + 1 - RETURNING section_type, name, body, version, updated_at + SET body = EXCLUDED.body, + origin = EXCLUDED.origin, + visible_to_user = EXCLUDED.visible_to_user, + updated_at = now(), + version = node_sections.version + 1 + RETURNING section_type, name, body, version, updated_at, origin, visible_to_user """, - user_uuid, nid, section_type, name, body, next_pos, + user_uuid, nid, section_type, name, body, next_pos, origin, visible_to_user, ) await conn.execute( "UPDATE context_nodes SET updated_at = now() WHERE id = $1", nid @@ -68,13 +84,20 @@ async def append_section( section_type: str, content: str, name: str = "main", + *, + origin: str = "user", + visible_to_user: bool = True, ) -> dict: + """Append content to an existing section (or create it). Preserves existing body.""" existing = await get_section(conn, node_id, section_type, name=name) if existing and existing["body"]: new_body = existing["body"] + "\n\n" + content else: new_body = content - return await upsert_section(conn, node_id, section_type, new_body, name=name) + return await upsert_section( + conn, node_id, section_type, new_body, name=name, + origin=origin, visible_to_user=visible_to_user, + ) async def delete_section( @@ -230,3 +253,147 @@ async def search_sections( "name": r["name"], "snippet": r["snippet"]} for r in rows ] + + +async def grep_sections( + conn: asyncpg.Connection, + pattern: str, + node_ids: list[str] | None = None, + limit: int = 20, +) -> list[dict]: + """ILIKE text search over node section bodies. + + pattern: SQL ILIKE pattern (caller wraps in % if needed). + node_ids: Optional list of node UUID strings to restrict results. + limit: Max rows returned (capped at 100 by callers). + + Returns list of dicts: {node_id, node_name, section_type, section_name, + snippet, origin, visible_to_user}. + """ + if node_ids: + rows = await conn.fetch( + """ + SELECT + ns.node_id::text, + cn.name AS node_name, + ns.section_type, + ns.name AS section_name, + left(ns.body, 300) AS snippet, + ns.origin, + ns.visible_to_user + FROM node_sections ns + JOIN context_nodes cn ON cn.id = ns.node_id + WHERE ns.node_id = ANY($1::uuid[]) + AND ns.body ILIKE $2 + ORDER BY cn.name, ns.section_type, ns.name + LIMIT $3 + """, + node_ids, pattern, limit, + ) + else: + rows = await conn.fetch( + """ + SELECT + ns.node_id::text, + cn.name AS node_name, + ns.section_type, + ns.name AS section_name, + left(ns.body, 300) AS snippet, + ns.origin, + ns.visible_to_user + FROM node_sections ns + JOIN context_nodes cn ON cn.id = ns.node_id + WHERE ns.body ILIKE $1 + ORDER BY cn.name, ns.section_type, ns.name + LIMIT $2 + """, + pattern, limit, + ) + return [ + { + "node_id": r["node_id"], + "node_name": r["node_name"], + "section_type": r["section_type"], + "section_name": r["section_name"], + "snippet": r["snippet"], + "origin": r["origin"], + "visible_to_user": r["visible_to_user"], + } + for r in rows + ] + + +async def search_sections_fts( + conn: asyncpg.Connection, + query: str, + node_ids: list[str] | None = None, + limit: int = 20, +) -> list[dict]: + """Full-text search over node section bodies using Postgres tsvector. + + Uses the trigger-maintained search_vector column (GIN-indexed). + Returns results ordered by ts_rank descending (most relevant first). + + query: Natural language search terms (passed to plainto_tsquery). + node_ids: Optional list of node UUID strings to restrict search. + limit: Max rows returned (capped at caller's discretion). + + Returns list of dicts: {node_id, node_name, section_type, section_name, + snippet, score}. + """ + if node_ids: + rows = await conn.fetch( + """ + SELECT + ns.node_id, + cn.name AS node_name, + ns.section_type, + ns.name AS section_name, + ts_headline( + 'english', ns.body, + plainto_tsquery('english', $1), + 'StartSel=, StopSel=, MaxWords=30, MinWords=10' + ) AS snippet, + ts_rank(ns.search_vector, plainto_tsquery('english', $1)) AS score + FROM node_sections ns + JOIN context_nodes cn ON cn.id = ns.node_id + WHERE ns.search_vector @@ plainto_tsquery('english', $1) + AND ns.node_id = ANY($2::uuid[]) + ORDER BY score DESC + LIMIT $3 + """, + query, node_ids, limit, + ) + else: + rows = await conn.fetch( + """ + SELECT + ns.node_id, + cn.name AS node_name, + ns.section_type, + ns.name AS section_name, + ts_headline( + 'english', ns.body, + plainto_tsquery('english', $1), + 'StartSel=, StopSel=, MaxWords=30, MinWords=10' + ) AS snippet, + ts_rank(ns.search_vector, plainto_tsquery('english', $1)) AS score + FROM node_sections ns + JOIN context_nodes cn ON cn.id = ns.node_id + WHERE ns.search_vector @@ plainto_tsquery('english', $1) + ORDER BY score DESC + LIMIT $2 + """, + query, limit, + ) + return [ + { + "node_id": str(r["node_id"]), + "node_name": r["node_name"], + "section_type": r["section_type"], + "section_name": r["section_name"], + "snippet": r["snippet"], + "score": float(r["score"]), + } + for r in rows + ] diff --git a/db/pg_queries/session_notes.py b/db/pg_queries/session_notes.py new file mode 100644 index 00000000..d0788a9c --- /dev/null +++ b/db/pg_queries/session_notes.py @@ -0,0 +1,49 @@ +"""Async Postgres queries — session_notes table. + +session_notes: user_id UUID PRIMARY KEY, content TEXT, updated_at TIMESTAMPTZ + +One row per user — a singleton accumulator for bot session summaries. +Callers own the append/rewrite logic; this layer is a plain get/upsert pair. +RLS ensures each connection can only see its own row. +""" +from __future__ import annotations + +import asyncpg + + +async def get_session_notes(conn: asyncpg.Connection) -> str | None: + """Return the session notes content for the current RLS user. + + Returns None when no row exists or when content is empty/whitespace-only. + The caller can treat None as "no notes yet" and fall back to a template. + """ + row = await conn.fetchrow( + """ + SELECT content + FROM session_notes + WHERE user_id = current_setting('app.current_user_id', true)::uuid + """ + ) + if not row: + return None + content = row["content"] + return content if content and content.strip() else None + + +async def upsert_session_notes(conn: asyncpg.Connection, content: str) -> None: + """Write session notes for the current RLS user (insert or replace). + + On first call: inserts a new row. + On subsequent calls: replaces content and refreshes updated_at. + Empty string is a valid reset value (get_session_notes returns None for it). + """ + await conn.execute( + """ + INSERT INTO session_notes (user_id, content, updated_at) + VALUES (current_setting('app.current_user_id', true)::uuid, $1, now()) + ON CONFLICT (user_id) DO UPDATE + SET content = EXCLUDED.content, + updated_at = now() + """, + content, + ) diff --git a/db/pg_queries/token_usage.py b/db/pg_queries/token_usage.py new file mode 100644 index 00000000..0fbb8f2c --- /dev/null +++ b/db/pg_queries/token_usage.py @@ -0,0 +1,112 @@ +"""Async Postgres queries — token_usage table. + +Token counts are captured from the agent SDK ResultMessage.usage dict on +handle release and written here asynchronously (off the hot path). + +These feed: + - The trial counter display (remaining turns for free-tier users) + - Future billing enforcement + - Analytics / cost attribution +""" +from __future__ import annotations + +import logging + +log = logging.getLogger(__name__) + + +async def record_token_usage( + user_id: str, + input_tokens: int, + output_tokens: int, + *, + pg_pool: object | None = None, +) -> None: + """Insert one token_usage row for the given user. + + Accepts an optional ``pg_pool`` (asyncpg pool) — if not provided, the + function attempts to acquire one from the app-level singleton. Silent + no-op if no pool is available (avoids crashing the release path). + + ``input_tokens`` and ``output_tokens`` map directly to the SDK's + ``ResultMessage.usage`` dict keys. + """ + try: + pool = pg_pool or _get_default_pool() + if pool is None: + log.debug("token_usage: no pg_pool available, skipping write") + return + async with pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO token_usage (user_id, input_tokens, output_tokens) + VALUES ($1, $2, $3) + """, + user_id, + input_tokens, + output_tokens, + ) + log.debug( + "token_usage recorded user_id=%s in=%d out=%d", + user_id, input_tokens, output_tokens, + ) + except Exception: + # Never raise from the async write — the release path must not fail + log.warning( + "token_usage write failed for user_id=%s — dropping record", + user_id, exc_info=True, + ) + + +def _get_default_pool() -> object | None: + """Best-effort fetch of the app-level asyncpg pool singleton. + + Returns None if the pool is not yet initialised (e.g. in unit tests). + """ + try: + from db.connection import get_pool # type: ignore[import] + return get_pool() + except Exception: + return None + + +async def get_token_usage_totals( + conn: object, + user_id: str, + since_iso: str | None = None, +) -> dict[str, int]: + """Return aggregated token totals for a user. + + ``since_iso``: ISO-8601 timestamp lower bound (e.g. start of current month). + Returns ``{"input_tokens": N, "output_tokens": N, "row_count": N}``. + """ + if since_iso is not None: + row = await conn.fetchrow( # type: ignore[attr-defined] + """ + SELECT + COALESCE(SUM(input_tokens), 0) AS input_tokens, + COALESCE(SUM(output_tokens), 0) AS output_tokens, + COUNT(*) AS row_count + FROM token_usage + WHERE user_id = $1 AND recorded_at >= $2::timestamptz + """, + user_id, + since_iso, + ) + else: + row = await conn.fetchrow( # type: ignore[attr-defined] + """ + SELECT + COALESCE(SUM(input_tokens), 0) AS input_tokens, + COALESCE(SUM(output_tokens), 0) AS output_tokens, + COUNT(*) AS row_count + FROM token_usage + WHERE user_id = $1 + """, + user_id, + ) + return { + "input_tokens": int(row["input_tokens"]), + "output_tokens": int(row["output_tokens"]), + "row_count": int(row["row_count"]), + } diff --git a/db/postgres.py b/db/postgres.py index ff844824..d36de966 100644 --- a/db/postgres.py +++ b/db/postgres.py @@ -52,8 +52,9 @@ async def create_pool() -> asyncpg.Pool: if _pool is None: _pool = await asyncpg.create_pool( dsn=os.environ["DATABASE_URL"], - min_size=2, + min_size=0, max_size=10, + max_inactive_connection_lifetime=60, command_timeout=30, init=register_jsonb_codec, ) diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh index e7be3487..4f42bb94 100644 --- a/docker-entrypoint.sh +++ b/docker-entrypoint.sh @@ -97,4 +97,5 @@ else echo "[entrypoint] CONFIG_SOURCE_URL not set — booting from baked-in defaults + env vars" fi +echo "0.0.0.0 statsig.anthropic.com" >> /etc/hosts exec "$@" diff --git a/docs/interactive-agent-layer.md b/docs/interactive-agent-layer.md new file mode 100644 index 00000000..0f78ccf6 --- /dev/null +++ b/docs/interactive-agent-layer.md @@ -0,0 +1,75 @@ +# Interactive Agent Layer — Architecture Notes + +This document describes the cross-process event delivery model for the interactive agent layer service (`tether/interactive_agent_layer/`). + +## Process boundary + +The interactive agent layer runs as a separate supervisord program from the API service: + +```ini +# supervisord.conf +[program:interactive-agent-layer] # port 5003 — its own OS process +[program:tether-api] # port 8000 — separate OS process +``` + +Each supervisord program is an independent OS process with its own event loop and memory space. Co-location does not mean in-process. The in-process `WSPublisher` (asyncio queues) cannot deliver events across this boundary. + +## Active-turn events: SSE → dispatch event_fn → WebSocket + +Events generated **during an active agent turn** (while `POST /session/{id}/turn` is streaming) flow cross-process via HTTP: + +``` +layer (process A) + ↓ yields on /session/{id}/turn SSE stream +dispatch / bot handler (process A or B) + ↓ LayerClient.turn() reads SSE incrementally (no buffering) + ↓ event_fn callback fires per event +API /ws handler (process B) + ↓ forwards to open WebSocket connection +frontend +``` + +Event types delivered this way: + +| Event | When | +|---|---| +| `agent_text_delta` | Streaming text fragment | +| `agent_action` | Tool-use translated to friendly phrase | +| `permission_request` | User approval prompt (user_action tools) | +| `status` | Pipeline status update | +| `turn_complete` | Turn finished | +| `session_ended` | Session terminated or interrupted | +| `trial_usage_update` | Live trial count (2.5, free-tier) | + +`permission_request` was wired onto the SSE stream in Wave 4 (PR #394). `PermissionGate` enqueues events into an `outbound_events: asyncio.Queue` that `run_turn` drains alongside pool SSE events, yielding them on the turn stream. + +## Background events: Redis pub/sub via WSPublisher dual-write + +Events fired **outside an active turn** (background lifecycle, async quota updates) cannot use the SSE stream. These use Redis pub/sub (PR #404): + +``` +WSPublisher.push() + ├─ in-process asyncio queues → in-process subscribers (tests, co-located callers) + └─ Redis PUBLISH tether:ws:{user_ws_id} → API /ws handler subscribes → WebSocket +``` + +- **In-process path** is preserved — tests and local dev work without Redis +- **Graceful degradation** when `REDIS_URL` is absent — Redis publish silently skipped +- **Channel key:** `tether:ws:{user_ws_id}` + +## Permission round-trip + +``` +pool control_request SSE + ↓ session.run_turn() sees event == "control_request" + ↓ asyncio.create_task(_handle_control_request) + ↓ PermissionGate.can_use_tool() → puts permission_request in outbound_events queue + ↓ _drain_until_done() yields permission_request on turn SSE stream + ↓ dispatch event_fn → user WebSocket → frontend shows Approve/Deny +user clicks + ↓ POST /permission/{request_id}/respond → layer HTTP API + ↓ resolves asyncio.Future in session.permission_pending + ↓ PermissionGate returns Allow/Deny + ↓ send_control_response to pool → pool unblocks can_use_tool callback + ↓ SDK proceeds or skips the tool call +``` diff --git a/fly.dev.toml b/fly.dev.toml index 9c7eed18..bf14842d 100644 --- a/fly.dev.toml +++ b/fly.dev.toml @@ -15,6 +15,13 @@ primary_region = "sjc" # dev machine can fully sleep — no persistent bot polling needed min_machines_running = 0 +# Redis provisioning (when ready for managed Redis): +# fly redis create --name tether-redis --org personal +# fly secrets set REDIS_URL="" --app tether-dev +# fly secrets set REDIS_URL="" --app tether-prod +# Until then, supervisord runs a local redis-server (see supervisord.conf [program:redis]) +# and REDIS_URL is set to redis://localhost:6379 in the api and interactive-agent-layer programs. + [env] PYTHON_ENV = "development" diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 598359e6..1c0f1d06 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -9,7 +9,7 @@ "version": "0.0.0", "dependencies": { "@vueuse/core": "^14.2.1", - "dompurify": "^3.4.0", + "dompurify": "^3.4.11", "marked": "^18.0.2", "pinia": "^3.0.4", "vue": "^3.5.12", @@ -1852,10 +1852,9 @@ "license": "MIT" }, "node_modules/dompurify": { - "version": "3.4.1", - "resolved": "https://registry.npmjs.org/dompurify/-/dompurify-3.4.1.tgz", - "integrity": "sha512-JahakDAIg1gyOm7dlgWSDjV4n7Ip2PKR55NIT6jrMfIgLFgWo81vdr1/QGqWtFNRqXP9UV71oVePtjqS2ebnPw==", - "license": "(MPL-2.0 OR Apache-2.0)", + "version": "3.4.11", + "resolved": "https://registry.npmjs.org/dompurify/-/dompurify-3.4.11.tgz", + "integrity": "sha512-zhlUV12GsaRzMsf9q5M254YhA4+VuF0fG+QFqu6aYpoGlKtz+w8//jBcGVYBgQkR5GHjUomejY84AV+/uPbWdw==", "optionalDependencies": { "@types/trusted-types": "^2.0.7" } diff --git a/frontend/package.json b/frontend/package.json index 202ae970..2fe75809 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -12,7 +12,7 @@ }, "dependencies": { "@vueuse/core": "^14.2.1", - "dompurify": "^3.4.0", + "dompurify": "^3.4.11", "marked": "^18.0.2", "pinia": "^3.0.4", "vue": "^3.5.12", diff --git a/frontend/src/App.vue b/frontend/src/App.vue index 42a6c499..d0f9c6f2 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -1,17 +1,33 @@ diff --git a/frontend/src/components/AgentActionPill.vue b/frontend/src/components/AgentActionPill.vue new file mode 100644 index 00000000..9a5cff76 --- /dev/null +++ b/frontend/src/components/AgentActionPill.vue @@ -0,0 +1,35 @@ + + + diff --git a/frontend/src/components/AgentBehaviorSection.vue b/frontend/src/components/AgentBehaviorSection.vue new file mode 100644 index 00000000..e4fee88f --- /dev/null +++ b/frontend/src/components/AgentBehaviorSection.vue @@ -0,0 +1,73 @@ + + + diff --git a/frontend/src/components/AgentPicker.vue b/frontend/src/components/AgentPicker.vue index 3fb8e519..1bfa590f 100644 --- a/frontend/src/components/AgentPicker.vue +++ b/frontend/src/components/AgentPicker.vue @@ -4,12 +4,6 @@ import { useAgentPickerStore } from '../stores/agentPicker' import { useAuthStore } from '../stores/auth' import type { AgentVersion } from '../stores/agentPicker' -withDefaults(defineProps<{ - trialMessagesLeft?: number -}>(), { - trialMessagesLeft: 10, -}) - const store = useAgentPickerStore() const authStore = useAuthStore() const open = ref(false) @@ -18,26 +12,44 @@ const rootEl = ref(null) // Premium users bypass trial counter and BYOK modal entirely. const isPremium = computed(() => authStore.user?.is_paid ?? false) +// Live trial count from WS events; null until first event arrives (show fallback). +const trialRemaining = computed(() => store.trialMessagesRemaining ?? 10) + +// True when current provider leaks 2.5 internals — 2.5 shown as unavailable. +const providerIsLeaky = computed(() => store.isLeakyProvider) + +// True when free user has exhausted their monthly 2.5 trial quota. +const trialExhausted = computed( + () => !isPremium.value && store.trialMessagesRemaining === 0, +) + const AGENTS: Array<{ id: AgentVersion; label: string; sublabel: string }> = [ { id: 'tether-agent-1.0', label: 'tether-agent-1.0', sublabel: 'Classic · free' }, { id: 'tether-agent-2.0', label: 'tether-agent-2.0', sublabel: 'Modern · free' }, { id: 'tether-agent-2.5', label: 'tether-agent-2.5', sublabel: 'Premium' }, ] +/** + * Whether the 2.5 option is locked (unselectable). + * Locks when: provider is leaky, OR trial is exhausted. + * Premium users are never locked. + */ +function is25Locked(id: AgentVersion): boolean { + if (id !== 'tether-agent-2.5') return false + if (isPremium.value) return false + return providerIsLeaky.value || trialExhausted.value +} + function toggleOpen() { open.value = !open.value } async function select(version: AgentVersion) { + if (is25Locked(version)) return // silently ignore clicks on locked option open.value = false await store.setAgent(version) } -// "Stay on 2.0" — cancel the pending 2.5 selection, keep current agent. -function stayOn20() { - store.cancelByokModal() -} - // Close dropdown when clicking outside the component's root. function onDocumentClick(e: MouseEvent) { if (rootEl.value && !rootEl.value.contains(e.target as Node)) { @@ -66,45 +78,72 @@ onBeforeUnmount(() => document.removeEventListener('mousedown', onDocumentClick)
diff --git a/frontend/src/components/BotChat.vue b/frontend/src/components/BotChat.vue index efda09b0..2e59cc4f 100644 --- a/frontend/src/components/BotChat.vue +++ b/frontend/src/components/BotChat.vue @@ -5,6 +5,8 @@ import { useAgentPickerStore } from '../stores/agentPicker' import MessageBubble from './MessageBubble.vue' import AgentPicker from './AgentPicker.vue' import PermissionModal from './PermissionModal.vue' +import AgentActionPill from './AgentActionPill.vue' +import StatusIndicator from './StatusIndicator.vue' const emit = defineEmits<{ close: [] }>() @@ -78,9 +80,13 @@ function onKeydown(e: KeyboardEvent) {

Send a message to start chatting.

+

{{ chatStore.statusMessage }}

+

+ Session timed out — send a message to resume +

@@ -88,11 +94,13 @@ function onKeydown(e: KeyboardEvent) { class="border-t border-[--border-1] p-3 flex gap-2 flex-shrink-0" @submit.prevent="onSubmit" > +