diff --git a/README.md b/README.md index 20d5da6..deeaf2d 100644 --- a/README.md +++ b/README.md @@ -152,9 +152,28 @@ Feature-complete across the planned phases: The scheduled job keeps its 30-minute overall lease and deadline, which can interrupt the four-call workflow before all per-call timeouts elapse. The selector defaults to `legacy` and - does not change deployed channels, schedules, or budgets. Existing budget checks compare - already-reported usage, so they are not hard spend reservations; [#117](https://github.com/R055LE/roger/issues/117) - tracks reservation enforcement before activating this profile in production. + does not change deployed channels, schedules, or budgets. Selecting the profile also requires + positive `DAILY_USD_CURATED`, `DAILY_ATTEMPTS_CURATED`, and `DAILY_WORKFLOWS_CURATED`; all three + default to zero. One workflow is a scheduled editorial evaluation or manual preview, with up to + four model attempts. These limits are independent of source observations and posted messages. + + Each selected-profile attempt reserves its request-specific input bound and full output ceiling + before provider I/O. The input bound includes UTF-8 messages, the complete schema, and framing, + up to 32,768 tokens; no evidence is trimmed to fit the budget. The dollar reservation includes + the higher cache-write input price. Admission compares used plus outstanding amounts against + both daily caps, including legacy Curated usage. Trustworthy completed usage settles once; + malformed, interrupted, or ambiguous calls keep their reservation. Reported overruns retain the + larger amount and hold all further selected-profile admissions: the previous bound was wrong. + Higher caps, midnight, and restarts do not release that hold; explicit reconciliation is required. + Reservations never expire or automatically replay. + + Budget days use the host-local day, as existing usage does; posting and source-observation limits + use `TZ`. Attempts count on admission day and actual settled charges on completion day. A workflow + continued after midnight counts on each day it makes a call. Unresolved reservations carry across + midnight and restarts until explicitly reconciled; there is no automatic release or reconciliation + command. `/status` shows used, reserved, and unknown amounts plus the last model or budget refusal. + The content-free accounting ledger follows existing usage retention and is not automatically + pruned. Legacy model profiles keep their existing budget behavior. Runs as a non-root, read-only-rootfs container. More than 300 tests cover the guard rules, the tool loop (including channel creation with access presets and the confirm-gated edit, post, and reorder diff --git a/compose.yaml b/compose.yaml index dc0d4de..6a0652e 100644 --- a/compose.yaml +++ b/compose.yaml @@ -48,6 +48,8 @@ services: DAILY_USD_AMBIENT: ${DAILY_USD_AMBIENT:-0} DAILY_USD_GIGABRAIN: ${DAILY_USD_GIGABRAIN:-0} DAILY_USD_CURATED: ${DAILY_USD_CURATED:-0} + DAILY_ATTEMPTS_CURATED: ${DAILY_ATTEMPTS_CURATED:-0} + DAILY_WORKFLOWS_CURATED: ${DAILY_WORKFLOWS_CURATED:-0} ADMIN_MAX_TOOL_CALLS: ${ADMIN_MAX_TOOL_CALLS:-10} ADMIN_MAX_TURNS: ${ADMIN_MAX_TURNS:-14} GIGABRAIN_MAX_TOOL_CALLS: ${GIGABRAIN_MAX_TOOL_CALLS:-10} diff --git a/roger.env.example b/roger.env.example index de4615f..b1da564 100644 --- a/roger.env.example +++ b/roger.env.example @@ -42,6 +42,10 @@ DAILY_USD_ADMIN=0 DAILY_USD_AMBIENT=0 DAILY_USD_GIGABRAIN=0 DAILY_USD_CURATED=0 +# Selected Curated profiles require positive USD, attempt, and workflow limits. +# One workflow is a scheduled evaluation or manual preview, with at most four calls. +DAILY_ATTEMPTS_CURATED=0 +DAILY_WORKFLOWS_CURATED=0 # --- admin tool loop bounds (§2.9) --- # hard cap on tool calls / model round-trips within a single admin request. Raise these if diff --git a/roger/bot.py b/roger/bot.py index 7b711bc..02c815d 100644 --- a/roger/bot.py +++ b/roger/bot.py @@ -363,6 +363,10 @@ def _format_status( curated_last_check: dict[str, Any] | None = None, curated_last_decision: dict[str, Any] | None = None, curated_last_input: dict[str, Any] | None = None, + curated_budget: dict[str, Any] | None = None, + curated_attempt_cap: int = 0, + curated_workflow_cap: int = 0, + curated_profile_last_call: dict[str, Any] | None = None, ) -> str: """Render the /status readout body (pure). The caller wraps it in a code block.""" usd_caps = usd_caps or {} @@ -390,6 +394,21 @@ def _format_status( if curated_configured and curated_check_interval_minutes: curated = f"checks every {curated_check_interval_minutes} min ({tz} daily limits)" lines.append(f"curated: {curated}") + if curated_budget is not None: + lines.append("curated model budget (host day; unresolved calls carry forward):") + for category in ("used", "reserved", "unknown"): + lines.append( + f" {category}: {curated_budget[f'{category}_tokens']:,} tokens; " + f"${curated_budget[f'{category}_usd']:.6f}" + ) + lines.append( + f"curated model attempts: {curated_budget['attempts_today']}/{curated_attempt_cap}; " + f"workflows: {curated_budget['workflows_today']}/{curated_workflow_cap}" + ) + if curated_budget["overrun_attempts"]: + lines.append("curated reservation hold: reservation_overrun; reconciliation required") + if curated_profile_last_call: + lines.append(f"curated last model result: {curated_profile_last_call['status']}") if curated_configured and curated_state is not None: lines.append( f"curated allowance: {curated_state['remaining_posts']}/{curated_max_posts_per_day} " @@ -462,6 +481,14 @@ async def gather_status(*, store: Store, settings: Settings, guild: Any) -> str: curated_last_check=json.loads(await store.get_meta("curated_last_check") or "{}"), curated_last_decision=json.loads(await store.get_meta("curated_last_decision") or "{}"), curated_last_input=json.loads(await store.get_meta("curated_last_input") or "{}"), + curated_budget=(await store.curated_budget_state() + if getattr(settings, "curated_request_profile", "legacy") != "legacy" + else None), + curated_attempt_cap=getattr(settings, "daily_attempts_curated", 0), + curated_workflow_cap=getattr(settings, "daily_workflows_curated", 0), + curated_profile_last_call=json.loads( + await store.get_meta("curated_profile_last_call") or "{}", + ), tz=settings.tz, ) diff --git a/roger/brains/curated.py b/roger/brains/curated.py index 92dbd52..683949f 100644 --- a/roger/brains/curated.py +++ b/roger/brains/curated.py @@ -545,13 +545,15 @@ async def finish(status: str, *, retry_at: float = 0) -> None: post = None else: post = await draft([entry for _, entry in ready], llm) - except BudgetExceeded: - local = clock().astimezone(timezone) + except BudgetExceeded as exc: + local = clock().astimezone() if exc.reason else clock().astimezone(timezone) tomorrow = datetime.datetime.combine( - local.date() + datetime.timedelta(days=1), datetime.time(), tzinfo=timezone, + local.date() + datetime.timedelta(days=1), datetime.time(), + tzinfo=None if exc.reason else timezone, ) await finish("budget", retry_at=tomorrow.timestamp()) - return await result("budget exceeded; skipped") + reason = f" ({exc.reason})" if exc.reason else "" + return await result(f"budget exceeded{reason}; skipped") except LLMConfigError: await finish("configuration") return await result("curated brain not configured") @@ -686,8 +688,9 @@ async def preview_curated_job( batch.entries, history, llm, now=now, max_age_hours=settings.scout_max_age_hours, ) - except BudgetExceeded: - return {"status": "budget exceeded; skipped"} + except BudgetExceeded as exc: + reason = f" ({exc.reason})" if exc.reason else "" + return {"status": f"budget exceeded{reason}; skipped"} except LLMConfigError: return {"status": "curated brain not configured"} except CuratedModelError as exc: diff --git a/roger/brains/stories.py b/roger/brains/stories.py index 1687c01..d56921c 100644 --- a/roger/brains/stories.py +++ b/roger/brains/stories.py @@ -14,7 +14,7 @@ from roger.brains.curated import DraftError, _json_object, _response_text, _safe_link from roger.identity import ROGER_IDENTITY -from roger.llm import LLM +from roger.llm import LLM, curated_workflow MAX_SOURCES = 8 SOURCE_EXCERPT_CAP = 1_600 @@ -959,6 +959,7 @@ async def _review( return parse_review(_story_response_text(response), decision, sources) +@curated_workflow() async def decide( entries: list[dict[str, Any]], history: dict[str, list[dict[str, Any]]], diff --git a/roger/config.py b/roger/config.py index 0ed45eb..8e48a3f 100644 --- a/roger/config.py +++ b/roger/config.py @@ -60,6 +60,9 @@ class Settings(BaseSettings): daily_usd_ambient: float = Field(default=0.0, ge=0, allow_inf_nan=False) daily_usd_gigabrain: float = Field(default=0.0, ge=0, allow_inf_nan=False) daily_usd_curated: float = Field(default=0.0, ge=0, allow_inf_nan=False) + # Selected profiles require explicit workflow and provider-attempt allowances. + daily_attempts_curated: int = Field(default=0, ge=0, le=2**31 - 1) + daily_workflows_curated: int = Field(default=0, ge=0, le=2**31 - 1) # --- admin tool loop bounds (§2.9) --- admin_max_tool_calls: int = Field(default=10, gt=0) @@ -145,6 +148,14 @@ def _timezone_must_be_loadable(cls, value: str) -> str: def _curated_profile_requires_stories(self) -> Settings: if self.curated_request_profile != "legacy" and not self.curated_developing_stories: raise ValueError("CURATED_REQUEST_PROFILE requires CURATED_DEVELOPING_STORIES") + if self.curated_request_profile != "legacy" and ( + self.daily_usd_curated <= 0 or self.daily_attempts_curated <= 0 + or self.daily_workflows_curated <= 0 + ): + raise ValueError( + "CURATED_REQUEST_PROFILE requires positive DAILY_USD_CURATED, " + "DAILY_ATTEMPTS_CURATED, and DAILY_WORKFLOWS_CURATED" + ) return self @property diff --git a/roger/llm.py b/roger/llm.py index 7da6b38..fa32ba4 100644 --- a/roger/llm.py +++ b/roger/llm.py @@ -12,6 +12,10 @@ import asyncio import json import logging +import uuid +from contextlib import asynccontextmanager +from contextvars import ContextVar +from decimal import Decimal from typing import Any from openai import ( @@ -40,6 +44,17 @@ REQUEST_TIMEOUT_S = 60.0 MAX_ATTEMPTS = 3 _BASE_BACKOFF_S = 0.5 +_curated_workflow_id: ContextVar[str | None] = ContextVar("curated_workflow_id", default=None) + + +@asynccontextmanager +async def curated_workflow(): + """One editorial evaluation, including its review and optional revision.""" + token = _curated_workflow_id.set(uuid.uuid4().hex) + try: + yield _curated_workflow_id.get() + finally: + _curated_workflow_id.reset(token) def _retry_after_seconds(exc: Exception) -> float | None: @@ -69,8 +84,13 @@ def __init__(self, status: str) -> None: class BudgetExceeded(RuntimeError): - def __init__(self, brain: str, used: float, cap: float, *, unit: str = "tokens") -> None: - if unit == "usd": + def __init__( + self, brain: str, used: float, cap: float, *, unit: str = "tokens", + reason: str | None = None, + ) -> None: + if reason is not None: + message = f"{brain} reservation denied: {reason}" + elif unit == "usd": message = f"{brain} daily $ budget exceeded (${used:.4f} >= ${cap:.4f})" else: message = f"{brain} daily token budget exceeded ({used} >= {cap})" @@ -79,6 +99,7 @@ def __init__(self, brain: str, used: float, cap: float, *, unit: str = "tokens") self.used = used self.cap = cap self.unit = unit + self.reason = reason class LLM: @@ -98,6 +119,8 @@ def __init__(self, settings: Settings, store: Store) -> None: "curated": settings.curated_models, } self._curated_profile = settings.curated_request_profile + self._curated_attempt_cap = settings.daily_attempts_curated + self._curated_workflow_cap = settings.daily_workflows_curated self._curated_review_models = settings.curated_review_models self._curated_public = settings.curated_channel_id is not None self._caps = { @@ -204,6 +227,9 @@ async def _complete_curated_profile( self, messages: list[dict[str, Any]], response_format: dict[str, Any], *, review: bool, ) -> Any: profile = curated_protocol.CURATED_REVIEWER if review else curated_protocol.CURATED_WRITER + workflow_id = _curated_workflow_id.get() + if workflow_id is None: + raise LLMConfigError("selected Curated profile requires an editorial workflow") try: kwargs = curated_protocol.request_kwargs( profile, messages, response_format, max_tokens=profile.max_tokens, @@ -214,22 +240,45 @@ async def _complete_curated_profile( input_bound = curated_protocol.input_token_upper_bound(messages, response_format) if input_bound > curated_protocol.MAX_INPUT_TOKENS: raise LLMConfigError("Curated request exceeds the input reservation maximum") - used = await self._store.usage_today("curated") - if used >= self._caps["curated"]: - raise BudgetExceeded("curated", used, self._caps["curated"]) - spent = await self._store.cost_today("curated") - if self._usd_caps["curated"] > 0 and spent >= self._usd_caps["curated"]: - raise BudgetExceeded("curated", spent, self._usd_caps["curated"], unit="usd") + reservation = { + "workflow_id": workflow_id, "role": "reviewer" if review else "writer", + "profile": profile.name, "input_tokens": input_bound, + "output_tokens": profile.max_tokens, + "cost_usd": ( + input_bound * max(profile.prompt_usd_per_million, + profile.prompt_cache_write_usd_per_million) + + profile.max_tokens * profile.completion_usd_per_million + ) / Decimal(1_000_000), + "token_cap": self._caps["curated"], "usd_cap": self._usd_caps["curated"], + "max_attempts": self._curated_attempt_cap, "max_workflows": self._curated_workflow_cap, + "timeout_seconds": profile.timeout, + } + attempt_id = await self._store.reserve_curated_attempt(**reservation) kwargs["extra_headers"] = curated_protocol.PROFILE_HEADERS diagnostic = { "profile": self._curated_profile, "role": "reviewer" if review else "writer", "model": profile.model, "status": "pending", "usage": None, } + if attempt_id is None: + state = await self._store.curated_budget_state(request=reservation) + reason = state["denial_reason"] or "budget_changed" + diagnostic["status"] = f"budget_{reason}" + await self._store.set_meta("curated_profile_last_call", json.dumps(diagnostic)) + metrics.LLM_BUDGET_EXCEEDED.labels("curated", reason).inc() + raise BudgetExceeded( + "curated", state["used_tokens"], self._caps["curated"], reason=reason, + ) await self._store.set_meta("curated_profile_last_call", json.dumps(diagnostic)) metrics.LLM_REQUESTS.labels("curated").inc() try: response = await self._client.chat.completions.create(**kwargs) + except asyncio.CancelledError: + await self._store.mark_curated_attempt_uncertain(attempt_id) + diagnostic["status"] = "request_interrupted_uncertain" + await self._store.set_meta("curated_profile_last_call", json.dumps(diagnostic)) + raise except Exception: + await self._store.mark_curated_attempt_uncertain(attempt_id) diagnostic["status"] = "request_error_uncertain" await self._store.set_meta("curated_profile_last_call", json.dumps(diagnostic)) metrics.LLM_ERRORS.labels("curated", "request_error_uncertain").inc() @@ -243,11 +292,22 @@ async def _complete_curated_profile( or (usage["reasoning_tokens"] is not None and usage["reasoning_tokens"] > usage["output_tokens"])): status = "invalid_usage_accounting" - else: - await self._store.add_usage( - "curated", usage["input_tokens"], usage["output_tokens"], - cost_usd=float(usage["cost_usd"]), - ) + # Content rejection can still have trustworthy accounting. Verify the route + # independently because the response classifier checks content failures first. + accounting_route = curated_protocol._response_status( + {**visible, "content": "accounting", "finish_reason": "stop", "refusal": None}, + profile, + ) + raw_usage = getattr(response, "usage", None) + settled = await self._store.finish_curated_attempt( + attempt_id, input_tokens=getattr(raw_usage, "prompt_tokens", None), + output_tokens=getattr(raw_usage, "completion_tokens", None), + cost_usd=getattr(raw_usage, "cost", None), + uncertain=(usage is None or accounting_route != "completed" + or status == "invalid_usage_accounting"), + ) + if settled == "overrun": + status = "reservation_overrun" cache = visible["usage_details"] if status == "completed" and any( cache[key] is not None and cache[key] > 0 diff --git a/roger/store.py b/roger/store.py index 9eed757..375fbdf 100644 --- a/roger/store.py +++ b/roger/store.py @@ -8,8 +8,12 @@ from __future__ import annotations import json +import math import os +import re import time +import uuid +from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal, InvalidOperation from enum import StrEnum from typing import Any @@ -106,6 +110,28 @@ class AuditStatus(StrEnum): PRIMARY KEY (date, brain) ); +CREATE TABLE IF NOT EXISTS curated_model_attempt ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + workflow_id TEXT NOT NULL, + admission_day TEXT NOT NULL, + role TEXT NOT NULL, + profile TEXT NOT NULL, + reserved_input INTEGER NOT NULL, + reserved_output INTEGER NOT NULL, + reserved_cost_units INTEGER NOT NULL, + state TEXT NOT NULL DEFAULT 'reserved', + actual_input INTEGER, + actual_output INTEGER, + actual_cost_usd TEXT, + actual_cost_units INTEGER, + settled_day TEXT, + created_at REAL NOT NULL, + deadline REAL NOT NULL, + settled_at REAL +); +CREATE INDEX IF NOT EXISTS curated_attempt_workflow ON curated_model_attempt(workflow_id); +CREATE INDEX IF NOT EXISTS curated_attempt_day ON curated_model_attempt(admission_day); + CREATE TABLE IF NOT EXISTS ambient_log ( id INTEGER PRIMARY KEY, ts REAL NOT NULL, @@ -141,6 +167,107 @@ class AuditStatus(StrEnum): CURATED_CHECK_LEASE_SECONDS = 30 * 60 +_USD_UNITS = 1_000_000_000 + + +def _cost_units(value: Any, rounding: str = ROUND_CEILING) -> int: + try: + amount = Decimal(str(value)) + if not amount.is_finite() or amount < 0: + raise ValueError("cost must be finite and nonnegative") + units = int((amount * _USD_UNITS).to_integral_value(rounding=rounding)) + if units > 2**63 - 1: + raise ValueError("cost exceeds storage limit") + return units + except (InvalidOperation, TypeError) as exc: + raise ValueError("invalid cost") from exc + + +def _curated_attempt_params( + *, workflow_id: str, role: str, profile: str, input_tokens: int, output_tokens: int, + cost_usd: Any, token_cap: int, usd_cap: Any, max_attempts: int, max_workflows: int, + timeout_seconds: float, now: float | None = None, +) -> dict[str, Any]: + for name, value in (("workflow_id", workflow_id), ("profile", profile)): + if not isinstance(value, str) or not re.fullmatch(r"[a-zA-Z0-9_-]{1,80}", value): + raise ValueError(f"invalid {name}") + if role not in {"writer", "reviewer"}: + raise ValueError("invalid curated role") + for value in (input_tokens, output_tokens, token_cap, max_attempts, max_workflows): + if type(value) is not int or not 0 <= value <= 2**31 - 1: + raise ValueError("token and attempt limits must be nonnegative integers") + if not 0 < input_tokens <= 32_768 or not 0 < output_tokens <= 16_384: + raise ValueError("invalid curated token reservation") + now = time.time() if now is None else now + if not math.isfinite(now) or not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError("invalid curated reservation deadline") + return { + "workflow_id": workflow_id, "role": role, "profile": profile, + "input_tokens": input_tokens, "output_tokens": output_tokens, + "cost_units": _cost_units(cost_usd), "token_cap": token_cap, + "usd_cap_units": _cost_units(usd_cap, ROUND_FLOOR), + "max_attempts": max_attempts, "max_workflows": max_workflows, + "day": _today(), "now": now, "deadline": now + timeout_seconds, + } + + +_CURATED_BUDGET_SQL = """ +WITH ledger AS ( + SELECT *, MAX(reserved_input, COALESCE(actual_input, 0)) + + MAX(reserved_output, COALESCE(actual_output, 0)) AS exposure_tokens, + MAX(reserved_cost_units, COALESCE(actual_cost_units, 0)) AS exposure_cost_units + FROM curated_model_attempt +), budget AS ( + SELECT + COALESCE((SELECT tokens_in + tokens_out FROM usage + WHERE brain = 'curated' AND date = :day), 0) + + COALESCE(SUM(CASE WHEN state = 'settled' AND settled_day = :day + THEN actual_input + actual_output ELSE 0 END), 0) AS used_tokens, + curated_cost_units(COALESCE((SELECT cost_usd FROM usage + WHERE brain = 'curated' AND date = :day), 0)) + + COALESCE(SUM(CASE WHEN state = 'settled' AND settled_day = :day + THEN actual_cost_units ELSE 0 END), 0) AS used_cost_units, + COALESCE((SELECT cost_usd FROM usage WHERE brain = 'curated' AND date = :day), 0) + + COALESCE(SUM(CASE WHEN state = 'settled' AND settled_day = :day + THEN CAST(actual_cost_usd AS REAL) ELSE 0 END), 0) AS used_usd, + COALESCE(SUM(CASE WHEN state = 'reserved' AND deadline > :now + THEN exposure_tokens ELSE 0 END), 0) AS reserved_tokens, + COALESCE(SUM(CASE WHEN state = 'reserved' AND deadline > :now + THEN exposure_cost_units ELSE 0 END), 0) AS reserved_cost_units, + COALESCE(SUM(CASE WHEN state IN ('uncertain', 'overrun') + OR (state = 'reserved' AND deadline <= :now) + THEN exposure_tokens ELSE 0 END), 0) AS unknown_tokens, + COALESCE(SUM(CASE WHEN state IN ('uncertain', 'overrun') + OR (state = 'reserved' AND deadline <= :now) + THEN exposure_cost_units ELSE 0 END), 0) AS unknown_cost_units, + COUNT(CASE WHEN state IN ('uncertain', 'overrun') + OR (state = 'reserved' AND deadline <= :now) THEN 1 END) AS unknown_attempts, + COUNT(CASE WHEN state = 'overrun' THEN 1 END) AS overrun_attempts, + COUNT(CASE WHEN admission_day = :day THEN 1 END) AS attempts_today, + COUNT(DISTINCT CASE WHEN admission_day = :day THEN workflow_id END) AS workflows_today, + COUNT(CASE WHEN workflow_id = :workflow_id THEN 1 END) AS workflow_attempts, + COUNT(CASE WHEN workflow_id = :workflow_id AND admission_day = :day THEN 1 END) + AS workflow_attempts_today + FROM ledger +) +""" + +_CURATED_DENIAL_SQL = """ +CASE + WHEN overrun_attempts > 0 THEN 'reservation_overrun' + WHEN :token_cap <= 0 OR :usd_cap_units <= 0 OR :max_attempts <= 0 OR :max_workflows <= 0 + THEN 'disabled' + WHEN workflow_attempts >= 4 THEN 'workflow_attempts' + WHEN attempts_today >= :max_attempts THEN 'daily_attempts' + WHEN workflow_attempts_today = 0 AND workflows_today >= :max_workflows THEN 'daily_workflows' + WHEN used_tokens + reserved_tokens + unknown_tokens + :input_tokens + :output_tokens > :token_cap + THEN 'tokens' + WHEN used_cost_units + reserved_cost_units + unknown_cost_units + :cost_units > :usd_cap_units + THEN 'usd' + ELSE NULL +END +""" + _CLAIM_CURATED_SQL = """ WITH selected AS ( SELECT id, action, event_key, story_ids_json @@ -249,21 +376,34 @@ class Store: def __init__(self, path: str) -> None: self._path = path self._db: aiosqlite.Connection | None = None + self._curated_db: aiosqlite.Connection | None = None async def open(self) -> Store: parent = os.path.dirname(self._path) if parent: os.makedirs(parent, exist_ok=True) - self._db = await aiosqlite.connect(self._path) + memory = self._path == ":memory:" + path = f"file:roger-{uuid.uuid4().hex}?mode=memory&cache=shared" if memory else self._path + self._db = await aiosqlite.connect(path, uri=memory) self._db.row_factory = aiosqlite.Row await self._db.execute("PRAGMA journal_mode=WAL") await self._db.execute("PRAGMA foreign_keys=ON") await self._db.executescript(_SCHEMA + _CURATED_DELIVERY_SCHEMA + ";") await self._migrate() await self._db.commit() + # Legacy methods can roll back the shared connection between awaits. The + # admission ledger must commit independently before any provider request starts. + self._curated_db = await aiosqlite.connect(path, uri=memory, isolation_level=None) + self._curated_db.row_factory = aiosqlite.Row + await self._curated_db.create_function( + "curated_cost_units", 1, _cost_units, deterministic=True, + ) return self async def close(self) -> None: + if self._curated_db is not None: + await self._curated_db.close() + self._curated_db = None if self._db is not None: await self._db.close() self._db = None @@ -274,6 +414,12 @@ def _conn(self) -> aiosqlite.Connection: raise RuntimeError("Store used before open()") return self._db + @property + def _curated_conn(self) -> aiosqlite.Connection: + if self._curated_db is None: + raise RuntimeError("Store used before open()") + return self._curated_db + async def _migrate(self) -> None: """Additive, idempotent migrations for columns ``CREATE TABLE IF NOT EXISTS`` can't add. @@ -378,6 +524,8 @@ async def journal_mode(self) -> str: async def usage_today(self, brain: str) -> int: """Total (in + out) tokens recorded for ``brain`` today. Drives the budget gate.""" + if brain == "curated": + return int((await self.curated_budget_state())["used_tokens"]) cursor = await self._conn.execute( "SELECT tokens_in + tokens_out FROM usage WHERE date = ? AND brain = ?", (_today(), brain), @@ -401,12 +549,98 @@ async def add_usage( async def cost_today(self, brain: str) -> float: """Actual USD charged for ``brain`` today (OpenRouter-reported cost, summed).""" + if brain == "curated": + return float((await self.curated_budget_state())["used_usd"]) cursor = await self._conn.execute( "SELECT cost_usd FROM usage WHERE date = ? AND brain = ?", (_today(), brain) ) row = await cursor.fetchone() return float(row[0]) if row and row[0] is not None else 0.0 + async def reserve_curated_attempt(self, **request: Any) -> int | None: + """Commit a bound before provider I/O. All unresolved days count against today's caps. + + ``_curated_attempt_params`` defines the admission fields. Admission and settlement + use the existing host-local budget day, independently of the posting timezone. + The separate connection commits the conditional INSERT before returning its ID. + """ + params = _curated_attempt_params(**request) + rows = await self._curated_conn.execute_fetchall( + _CURATED_BUDGET_SQL # noqa: S608 - only code-owned SQL fragments + + """ + INSERT INTO curated_model_attempt + (workflow_id, admission_day, role, profile, reserved_input, reserved_output, + reserved_cost_units, created_at, deadline) + SELECT :workflow_id, :day, :role, :profile, :input_tokens, :output_tokens, + :cost_units, :now, :deadline FROM budget + WHERE (""" + _CURATED_DENIAL_SQL + ") IS NULL RETURNING id", params, + ) + return int(rows[0][0]) if rows else None + + async def finish_curated_attempt( + self, attempt_id: int, *, input_tokens: Any = None, output_tokens: Any = None, + cost_usd: Any = None, uncertain: bool = False, now: float | None = None, + ) -> str | None: + """Settle once, or keep uncertain exposure. A known overrun retains the larger charge. + + Callers supply only trustworthy provider usage, even when the content was rejected. + Partial valid usage is retained without releasing any unaccounted reservation. + Like the legacy usage totals, this payload-free ledger is not automatically pruned. + """ + actual_input = (input_tokens if type(input_tokens) is int + and 0 <= input_tokens <= 2**63 - 1 else None) + actual_output = (output_tokens if type(output_tokens) is int + and 0 <= output_tokens <= 2**63 - 1 else None) + try: + actual_cost_units = _cost_units(cost_usd) + actual_cost = str(Decimal(str(cost_usd))) + except ValueError: + actual_cost_units, actual_cost = None, None + complete = all(value is not None for value in ( + actual_input, actual_output, actual_cost_units, + )) + rows = await self._curated_conn.execute_fetchall( + """UPDATE curated_model_attempt SET + state = CASE WHEN :input > reserved_input OR :output > reserved_output + OR :cost_units > reserved_cost_units THEN 'overrun' + WHEN NOT :complete OR :uncertain THEN 'uncertain' + ELSE 'settled' END, + actual_input = :input, actual_output = :output, + actual_cost_usd = :cost, actual_cost_units = :cost_units, + settled_day = :day, settled_at = :now + WHERE id = :id AND state = 'reserved' RETURNING state""", + {"id": attempt_id, "complete": complete, "uncertain": uncertain, "input": actual_input, + "output": actual_output, "cost": actual_cost, "cost_units": actual_cost_units, + "day": _today(), "now": time.time() if now is None else now}, + ) + return str(rows[0][0]) if rows else None + + async def mark_curated_attempt_uncertain(self, attempt_id: int) -> str | None: + return await self.finish_curated_attempt(attempt_id) + + async def curated_budget_state( + self, *, now: float | None = None, request: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Used, reserved, and unknown exposure; optionally diagnose a refused request. + + A diagnosis reflects this snapshot. Concurrent settlement may clear the gate after + a refusal. Deadline expiry only changes the displayed category, never the liability. + """ + params = {"day": _today(), "now": time.time() if now is None else now, "workflow_id": ""} + if request is not None: + params = _curated_attempt_params(**request) + if now is not None: + params["now"] = now + rows = await self._curated_conn.execute_fetchall( + _CURATED_BUDGET_SQL + "SELECT *, " # noqa: S608 - only code-owned SQL fragments + + (_CURATED_DENIAL_SQL if request is not None else "NULL") + + " AS denial_reason FROM budget", params, + ) + state = dict(rows[0]) + state["reserved_usd"] = state["reserved_cost_units"] / _USD_UNITS + state["unknown_usd"] = state["unknown_cost_units"] / _USD_UNITS + return state + # --- small key/value bot state (presence outfit, etc.) --- async def get_meta(self, key: str) -> str | None: diff --git a/tests/test_config.py b/tests/test_config.py index 91a56ac..a2d9660 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -154,6 +154,8 @@ def test_curated_defaults(monkeypatch): assert settings.curated_review_models == [] assert settings.daily_tokens_curated == 30_000 assert settings.daily_usd_curated == 0.0 + assert settings.daily_attempts_curated == 0 + assert settings.daily_workflows_curated == 0 assert settings.curated_check_interval_minutes == 0 assert settings.curated_max_posts_per_day == 1 assert settings.curated_min_spacing_minutes == 60 @@ -175,6 +177,9 @@ def test_curated_profile_reports_effective_models_without_changing_legacy_config monkeypatch.setenv("MODEL_CURATED_REVIEW", "legacy/reviewer") monkeypatch.setenv("CURATED_REQUEST_PROFILE", "sol-astra-v1") monkeypatch.setenv("CURATED_DEVELOPING_STORIES", "true") + monkeypatch.setenv("DAILY_USD_CURATED", "2") + monkeypatch.setenv("DAILY_ATTEMPTS_CURATED", "8") + monkeypatch.setenv("DAILY_WORKFLOWS_CURATED", "2") settings = Settings() assert settings.curated_models == ["openai/gpt-6.1-sol-20260929"] assert settings.curated_review_models == ["openai/gpt-6-astra-20260903"] @@ -189,6 +194,20 @@ def test_unknown_curated_profile_is_rejected(monkeypatch): Settings() +@pytest.mark.parametrize("missing", [ + "DAILY_USD_CURATED", "DAILY_ATTEMPTS_CURATED", "DAILY_WORKFLOWS_CURATED", +]) +def test_selected_profile_requires_explicit_positive_budget_limits(monkeypatch, missing): + _set_required(monkeypatch) + for key, value in { + "CURATED_REQUEST_PROFILE": "sol-astra-v1", "CURATED_DEVELOPING_STORIES": "true", + "DAILY_USD_CURATED": "2", "DAILY_ATTEMPTS_CURATED": "8", "DAILY_WORKFLOWS_CURATED": "2", + }.items(): + monkeypatch.setenv(key, "0" if key == missing else value) + with pytest.raises(ValidationError, match="requires positive DAILY_USD_CURATED"): + Settings() + + @pytest.mark.parametrize("setting", ["CURATED_CHECK_INTERVAL_MINUTES", "CURATED_MIN_SPACING_MINUTES"]) @pytest.mark.parametrize("value", ["-1", "1441"]) diff --git a/tests/test_llm.py b/tests/test_llm.py index 92c87ca..4465658 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -2,6 +2,7 @@ import asyncio import json +from decimal import Decimal from types import SimpleNamespace import httpx @@ -19,6 +20,7 @@ CuratedModelError, LLMConfigError, _retry_after_seconds, + curated_workflow, ) from roger.store import Store @@ -38,6 +40,9 @@ async def _no_sleep(*_args, **_kwargs): def _env(monkeypatch, **extra): + if extra.get("CURATED_REQUEST_PROFILE") == "sol-astra-v1": + extra = {"DAILY_USD_CURATED": "2", "DAILY_ATTEMPTS_CURATED": "8", + "DAILY_WORKFLOWS_CURATED": "2", **extra} for key, value in {**_REQUIRED, **extra}.items(): monkeypatch.setenv(key, value) @@ -571,6 +576,7 @@ async def test_curated_budget_exceeded_before_network(monkeypatch, tmp_path): @pytest.mark.parametrize("review", [False, True]) +@curated_workflow() async def test_selected_curated_profile_uses_evaluated_output_limits(monkeypatch, tmp_path, review): _env(monkeypatch, MODEL_CURATED="legacy/writer", MODEL_CURATED_REVIEW="legacy/reviewer", CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") @@ -627,6 +633,7 @@ def test_production_curated_profiles_match_evaluated_contract(): @pytest.mark.parametrize("review", [False, True]) +@curated_workflow() async def test_production_profile_serializes_exact_request(monkeypatch, tmp_path, review): _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") store = await Store(str(tmp_path / "wire.db")).open() @@ -684,6 +691,7 @@ async def handle(request): ("cache_write", "prompt_cache_policy_mismatch"), ("wrong_model", "model_mismatch"), ("blank", "blank_response"), ]) +@curated_workflow() async def test_profile_response_failures_preserve_safe_usage( monkeypatch, tmp_path, kind, status, ): @@ -719,12 +727,16 @@ async def create(**kwargs): assert json.loads(diagnostic)["status"] == status assert "private source text" not in diagnostic assert "unexpected/private-model" not in diagnostic - assert await store.cost_today("curated") == (0 if kind == "missing_usage" else 0.01) + uncertain = kind in {"missing_usage", "null_choices", "empty_choices", "wrong_model"} + assert await store.cost_today("curated") == (0 if uncertain else 0.01) + state = await store.curated_budget_state() + assert state["unknown_attempts"] == int(uncertain) finally: await llm._client.close() await store.close() +@curated_workflow() async def test_profile_transport_has_one_attempt_and_sanitized_uncertainty(monkeypatch, tmp_path): _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") store = await Store(str(tmp_path / "transport.db")).open() @@ -784,3 +796,211 @@ async def get(path, **kwargs): finally: await llm._client.close() await store.close() + + +async def _profile_call(llm, *, content="finding"): + return await llm.complete( + "curated", [{"role": "user", "content": content}], curated_story=True, + response_format=stories.DECISION_RESPONSE_FORMAT, + ) + + +@pytest.mark.parametrize("gate", ["tokens", "usd"]) +@curated_workflow() +async def test_profile_reserves_full_request_before_first_network_call(monkeypatch, tmp_path, gate): + profile = curated_protocol.CURATED_WRITER + bound = curated_protocol.input_token_upper_bound( + [{"role": "user", "content": "finding"}], stories.DECISION_RESPONSE_FORMAT, + ) + # This would cover ordinary input pricing, but not the declared cache-write price. + base_cost = (bound * profile.prompt_usd_per_million + + profile.max_tokens * profile.completion_usd_per_million) / Decimal(1_000_000) + _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true", + **({"DAILY_TOKENS_CURATED": str(bound + profile.max_tokens - 1)} if gate == "tokens" + else {"DAILY_USD_CURATED": str(base_cost)})) + store = await Store(str(tmp_path / "preflight-reservation.db")).open() + llm = LLM(Settings(), store) + + async def create(**kwargs): + raise AssertionError("denied reservation must not call provider") + + monkeypatch.setattr(llm._client.chat.completions, "create", create) + try: + with pytest.raises(BudgetExceeded, match=f"reservation denied: {gate}"): + await _profile_call(llm) + state = await store.curated_budget_state() + assert state["used_tokens"] == state["attempts_today"] == 0 + diagnostic = json.loads(await store.get_meta("curated_profile_last_call")) + assert diagnostic["status"] == f"budget_{gate}" + finally: + await llm._client.close() + await store.close() + + +async def test_profile_rejects_calls_without_a_workflow(monkeypatch, tmp_path): + _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") + store = await Store(str(tmp_path / "context.db")).open() + llm = LLM(Settings(), store) + try: + with pytest.raises(LLMConfigError, match="requires an editorial workflow"): + await _profile_call(llm) + assert (await store.curated_budget_state())["attempts_today"] == 0 + finally: + await llm._client.close() + await store.close() + + +@curated_workflow() +async def test_profile_counts_utf8_and_schema_before_reserving(monkeypatch, tmp_path): + _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") + store = await Store(str(tmp_path / "input.db")).open() + llm = LLM(Settings(), store) + try: + with pytest.raises(LLMConfigError, match="input reservation maximum"): + await _profile_call(llm, content="é" * 15_000) + assert (await store.curated_budget_state())["attempts_today"] == 0 + finally: + await llm._client.close() + await store.close() + + +@curated_workflow() +async def test_profile_fifth_call_in_one_workflow_is_denied(monkeypatch, tmp_path): + _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") + store = await Store(str(tmp_path / "four-calls.db")).open() + llm = LLM(Settings(), store) + calls = 0 + + async def create(**kwargs): + nonlocal calls + calls += 1 + return ChatCompletion.model_validate(_profile_response()) + + monkeypatch.setattr(llm._client.chat.completions, "create", create) + try: + for _ in range(4): + await _profile_call(llm) + with pytest.raises(BudgetExceeded, match="workflow_attempts"): + await _profile_call(llm) + assert calls == 4 + state = await store.curated_budget_state() + assert state["attempts_today"] == 4 + assert state["workflows_today"] == 1 + assert state["used_tokens"] == 600 + finally: + await llm._client.close() + await store.close() + + +async def test_profile_concurrent_workflows_have_independent_context(monkeypatch, tmp_path): + _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true", + DAILY_TOKENS_CURATED="100000") + store = await Store(str(tmp_path / "concurrent.db")).open() + llm = LLM(Settings(), store) + both_started = asyncio.Event() + calls = 0 + + async def create(**kwargs): + nonlocal calls + calls += 1 + if calls == 2: + both_started.set() + await both_started.wait() + return ChatCompletion.model_validate(_profile_response()) + + async def workflow(): + async with curated_workflow(): + return await _profile_call(llm) + + monkeypatch.setattr(llm._client.chat.completions, "create", create) + try: + await asyncio.wait_for(asyncio.gather(workflow(), workflow()), timeout=5) + assert calls == 2 + state = await store.curated_budget_state() + assert state["workflows_today"] == state["attempts_today"] == 2 + with pytest.raises(BudgetExceeded, match="daily_workflows"): + await workflow() + assert calls == 2 + with pytest.raises(LLMConfigError, match="editorial workflow"): + await _profile_call(llm) + finally: + await llm._client.close() + await store.close() + + +async def test_profile_cancellation_keeps_reservation_after_reopen(monkeypatch, tmp_path): + _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") + path = str(tmp_path / "cancel.db") + store = await Store(path).open() + llm = LLM(Settings(), store) + started = asyncio.Event() + + async def create(**kwargs): + started.set() + await asyncio.Future() + + async def workflow(): + async with curated_workflow(): + return await _profile_call(llm) + + monkeypatch.setattr(llm._client.chat.completions, "create", create) + try: + task = asyncio.create_task(workflow()) + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + diagnostic = json.loads(await store.get_meta("curated_profile_last_call")) + assert diagnostic["status"] == "request_interrupted_uncertain" + finally: + await llm._client.close() + await store.close() + reopened = await Store(path).open() + try: + state = await reopened.curated_budget_state() + assert state["unknown_attempts"] == 1 + assert state["unknown_tokens"] > 16_384 + assert state["unknown_usd"] > 0.16384 + assert state["attempts_today"] == state["workflows_today"] == 1 + finally: + await reopened.close() + + +@pytest.mark.parametrize("missing_tokens", [False, True]) +@curated_workflow() +async def test_profile_reported_overrun_keeps_actual_cost_even_with_partial_usage( + monkeypatch, tmp_path, missing_tokens, +): + _env(monkeypatch, CURATED_REQUEST_PROFILE="sol-astra-v1", CURATED_DEVELOPING_STORIES="true") + store = await Store(str(tmp_path / "overrun.db")).open() + llm = LLM(Settings(), store) + response = ChatCompletion.model_validate(_profile_response()) + response.usage.cost = 1.25 + if missing_tokens: + response.usage.prompt_tokens = None + calls = 0 + + async def create(**kwargs): + nonlocal calls + calls += 1 + return response + + monkeypatch.setattr(llm._client.chat.completions, "create", create) + try: + with pytest.raises(CuratedModelError, match="reservation_overrun"): + await _profile_call(llm) + state = await store.curated_budget_state() + assert state["unknown_usd"] == 1.25 + assert state["unknown_tokens"] > 16_384 + assert state["used_usd"] == 0 + rows = await store._conn.execute_fetchall( + "SELECT actual_cost_usd FROM curated_model_attempt", + ) + assert rows[0][0] == "1.25" + async with curated_workflow(): + with pytest.raises(BudgetExceeded, match="reservation_overrun"): + await _profile_call(llm) + assert calls == 1 + finally: + await llm._client.close() + await store.close() diff --git a/tests/test_status.py b/tests/test_status.py index dadc021..10aca02 100644 --- a/tests/test_status.py +++ b/tests/test_status.py @@ -205,3 +205,58 @@ async def test_gather_status_reports_channel_and_guild_state(tmp_path): assert "curated: unconfigured" in body finally: await store.close() + + +async def test_status_shows_curated_reserved_unknown_and_last_budget_reason(tmp_path): + store = await Store(str(tmp_path / "budget.db")).open() + request = dict( + role="writer", profile="sol", input_tokens=100, output_tokens=200, + cost_usd="0.01", token_cap=1000, usd_cap="0.1", max_attempts=8, max_workflows=2, + timeout_seconds=1200, + ) + try: + settled = await store.reserve_curated_attempt(workflow_id="a", **request) + await store.finish_curated_attempt( + settled, input_tokens=10, output_tokens=20, cost_usd="0.001", + ) + unknown = await store.reserve_curated_attempt(workflow_id="a", **request) + await store.mark_curated_attempt_uncertain(unknown) + await store.reserve_curated_attempt(workflow_id="b", **request) + await store.set_meta("curated_profile_last_call", json.dumps({"status": "budget_tokens"})) + body = await gather_status( + store=store, settings=_settings( + curated_request_profile="sol-astra-v1", daily_attempts_curated=8, + daily_workflows_curated=2, + ), guild=_fake_guild(channels={42: _FakeChannel()}), + ) + assert "used: 30 tokens; $0.001000" in body + assert "reserved: 300 tokens; $0.010000" in body + assert "unknown: 300 tokens; $0.010000" in body + assert "curated model attempts: 3/8; workflows: 2/2" in body + assert "curated last model result: budget_tokens" in body + assert "host day; unresolved calls carry forward" in body + finally: + await store.close() + + +async def test_status_shows_persistent_overrun_hold_without_a_new_call(tmp_path): + store = await Store(str(tmp_path / "overrun.db")).open() + try: + attempt = await store.reserve_curated_attempt( + workflow_id="a", role="writer", profile="sol", input_tokens=100, output_tokens=200, + cost_usd="0.01", token_cap=1000, usd_cap="0.1", max_attempts=8, max_workflows=2, + timeout_seconds=1200, + ) + await store.finish_curated_attempt( + attempt, input_tokens=10, output_tokens=20, cost_usd="0.02", + ) + body = await gather_status( + store=store, settings=_settings( + curated_request_profile="sol-astra-v1", daily_attempts_curated=8, + daily_workflows_curated=2, + ), guild=_fake_guild(channels={42: _FakeChannel()}), + ) + assert "unknown: 300 tokens; $0.020000" in body + assert "curated reservation hold: reservation_overrun; reconciliation required" in body + finally: + await store.close() diff --git a/tests/test_store.py b/tests/test_store.py index 251ae2b..89016e0 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -10,6 +10,340 @@ from roger.store import CURATED_CHECK_LEASE_SECONDS, RETENTION_DAYS, AuditStatus, Store +def _attempt(workflow_id="workflow-a", **overrides): + args = dict( + workflow_id=workflow_id, role="writer", profile="sol", input_tokens=100, + output_tokens=200, cost_usd="0.01", token_cap=1000, usd_cap="0.1", + max_attempts=10, max_workflows=10, timeout_seconds=1200, now=1000, + ) + args.update(overrides) + return args + + +@pytest.mark.parametrize("shared_connection", [False, True]) +@pytest.mark.parametrize("limit", [ + {"token_cap": 800}, {"usd_cap": "0.02"}, {"max_attempts": 1}, {"max_workflows": 1}, +]) +async def test_curated_reservations_cannot_spend_the_same_last_budget( + tmp_path, shared_connection, limit, +): + path = str(tmp_path / "roger.db") + first = await Store(path).open() + second = first if shared_connection else await Store(path).open() + try: + await first.add_usage("curated", 400, 100, 0.01) + claims = await asyncio.gather( + first.reserve_curated_attempt(**_attempt("a", **limit)), + second.reserve_curated_attempt(**_attempt("b", **limit)), + first.add_usage("admin", 10, 20, 0.1), + ) + assert sum(claim is not None for claim in claims[:2]) == 1 + state = await first.curated_budget_state(now=1001) + assert state["used_tokens"] == 500 + assert state["reserved_tokens"] == 300 + assert state["attempts_today"] == state["workflows_today"] == 1 + finally: + if second is not first: + await second.close() + await first.close() + + +@pytest.mark.parametrize("settlement", [False, True]) +async def test_curated_ledger_survives_other_store_methods_rolling_back( + tmp_path, monkeypatch, settlement, +): + store = await Store(str(tmp_path / "roger.db")).open() + try: + claim = await store.reserve_curated_attempt(**_attempt()) if settlement else None + await store._conn.execute("BEGIN") + original = store._conn.execute_fetchall + + async def interleave_rollback(sql, parameters=None): + rows = await original(sql, parameters) + if "RETURNING" in sql: + with pytest.raises(RuntimeError, match="claim missing"): + await store.mark_curated_sent(-1, 1) + return rows + + monkeypatch.setattr(store._conn, "execute_fetchall", interleave_rollback) + if settlement: + assert await store.finish_curated_attempt( + claim, input_tokens=10, output_tokens=20, cost_usd="0.001", + ) == "settled" + else: + claim = await store.reserve_curated_attempt(**_attempt()) + assert claim is not None + with pytest.raises(RuntimeError, match="claim missing"): + await store.mark_curated_sent(-1, 1) + rows = await original("SELECT state FROM curated_model_attempt WHERE id = ?", (claim,)) + assert [row[0] for row in rows] == ["settled" if settlement else "reserved"] + finally: + await store.close() + + +async def test_curated_memory_store_shares_usage_and_ledger(): + store = await Store(":memory:").open() + try: + await store.add_usage("curated", 100, 200, 0.01) + assert await store.reserve_curated_attempt(**_attempt(token_cap=599)) is None + claim = await store.reserve_curated_attempt(**_attempt(token_cap=600)) + assert claim is not None + assert (await store.curated_budget_state(now=1001))["reserved_tokens"] == 300 + assert await store.finish_curated_attempt( + claim, input_tokens=10, output_tokens=20, cost_usd="0.001", + ) == "settled" + assert await store.usage_today("curated") == 330 + finally: + await store.close() + + +async def test_curated_status_read_cannot_hold_a_snapshot_across_admission(tmp_path, monkeypatch): + path = str(tmp_path / "roger.db") + first, second = await Store(path).open(), await Store(path).open() + read_started, release_read = asyncio.Event(), asyncio.Event() + original_execute = first._curated_conn.execute + original_fetchall = first._curated_conn.execute_fetchall + + async def execute(sql, parameters=None): + cursor = await original_execute(sql, parameters) + if "AS denial_reason FROM budget" in sql: + read_started.set() + await release_read.wait() + return cursor + + async def execute_fetchall(sql, parameters=None): + rows = await original_fetchall(sql, parameters) + if "AS denial_reason FROM budget" in sql: + read_started.set() + await release_read.wait() + return rows + + monkeypatch.setattr(first._curated_conn, "execute", execute) + monkeypatch.setattr(first._curated_conn, "execute_fetchall", execute_fetchall) + status = asyncio.create_task(first.curated_budget_state(now=1001)) + try: + await asyncio.wait_for(read_started.wait(), timeout=5) + assert await second.reserve_curated_attempt(**_attempt("other-connection")) is not None + assert await first.reserve_curated_attempt(**_attempt("status-connection")) is not None + finally: + release_read.set() + await status + await first.close() + await second.close() + + +@pytest.mark.parametrize(("limit", "reason"), [ + ({"token_cap": 299}, "tokens"), ({"usd_cap": "0.009999999999"}, "usd"), + ({"max_attempts": 0}, "disabled"), ({"max_workflows": 0}, "disabled"), + ({"token_cap": 0}, "disabled"), ({"usd_cap": 0}, "disabled"), +]) +async def test_curated_admission_reports_the_current_gate(tmp_path, limit, reason): + store = await Store(str(tmp_path / "roger.db")).open() + try: + request = _attempt(**limit) + assert await store.reserve_curated_attempt(**request) is None + state = await store.curated_budget_state(request=request) + assert state["denial_reason"] == reason + assert state["attempts_today"] == 0 + finally: + await store.close() + + +async def test_curated_workflow_and_attempt_limits_are_independent(tmp_path): + store = await Store(str(tmp_path / "roger.db")).open() + try: + for _ in range(4): + claim = await store.reserve_curated_attempt(**_attempt(max_workflows=1)) + assert claim is not None + assert await store.finish_curated_attempt( + claim, input_tokens=10, output_tokens=20, cost_usd="0.001", + ) == "settled" + request = _attempt(max_workflows=1) + assert await store.reserve_curated_attempt(**request) is None + state = await store.curated_budget_state(request=request) + assert state["denial_reason"] == "workflow_attempts" + assert state["attempts_today"] == 4 + assert state["workflows_today"] == 1 + request = _attempt("new-workflow", max_workflows=1) + assert await store.reserve_curated_attempt(**request) is None + assert (await store.curated_budget_state(request=request))["denial_reason"] == ( + "daily_workflows" + ) + request = _attempt("new-workflow", max_attempts=4) + assert await store.reserve_curated_attempt(**request) is None + assert (await store.curated_budget_state(request=request))["denial_reason"] == ( + "daily_attempts" + ) + finally: + await store.close() + + +@pytest.mark.parametrize("shared_connection", [False, True]) +async def test_curated_settlement_is_exactly_once_and_adds_legacy_usage( + tmp_path, shared_connection, +): + path = str(tmp_path / "roger.db") + first = await Store(path).open() + second = first if shared_connection else await Store(path).open() + try: + await first.add_usage("curated", 30, 20, 0.005) + claim = await first.reserve_curated_attempt(**_attempt()) + results = await asyncio.gather(*[ + store.finish_curated_attempt( + claim, input_tokens=10, output_tokens=20, cost_usd="0.00123456789", + ) for store in (first, second) + ]) + assert sorted(str(result) for result in results) == ["None", "settled"] + assert await first.usage_today("curated") == 80 + assert await first.cost_today("curated") == pytest.approx(0.00623456789) + state = await first.curated_budget_state(now=1001) + assert state["reserved_tokens"] == state["unknown_tokens"] == 0 + assert state["used_cost_units"] == 6_234_568 + assert await first.mark_curated_attempt_uncertain(claim) is None + assert await first.usage_today("admin") == 0 + finally: + if second is not first: + await second.close() + await first.close() + + +@pytest.mark.parametrize("usage", [ + {}, {"input_tokens": -1, "output_tokens": 10, "cost_usd": "0.001"}, + {"input_tokens": True, "output_tokens": 10, "cost_usd": "0.001"}, + {"input_tokens": 10, "output_tokens": 1.5, "cost_usd": "0.001"}, + {"input_tokens": 10, "output_tokens": 20, "cost_usd": float("nan")}, + {"input_tokens": 10, "output_tokens": 20, "cost_usd": float("inf")}, + {"input_tokens": 10, "output_tokens": 20, "cost_usd": -0.001}, + {"input_tokens": 10, "output_tokens": 20, "cost_usd": "unusable"}, + {"input_tokens": 10, "output_tokens": 20, "cost_usd": "0.001", "uncertain": True}, +]) +async def test_curated_untrustworthy_outcomes_keep_full_reservation(tmp_path, usage): + store = await Store(str(tmp_path / "roger.db")).open() + try: + claim = await store.reserve_curated_attempt(**_attempt()) + assert await store.finish_curated_attempt(claim, **usage) == "uncertain" + assert await store.finish_curated_attempt( + claim, input_tokens=1, output_tokens=1, cost_usd=0, + ) is None + state = await store.curated_budget_state(now=1001) + assert state["used_tokens"] == state["reserved_tokens"] == 0 + assert state["unknown_tokens"] == 300 + assert state["unknown_usd"] == 0.01 + assert state["unknown_attempts"] == 1 + assert await store.reserve_curated_attempt(**_attempt("new", token_cap=599)) is None + finally: + await store.close() + + +@pytest.mark.parametrize("usage", [ + {"input_tokens": 150, "output_tokens": 250, "cost_usd": "0.02"}, + {"cost_usd": "0.02"}, +]) +async def test_curated_known_overrun_never_shrinks_to_reserved_charge(tmp_path, usage): + store = await Store(str(tmp_path / "roger.db")).open() + try: + claim = await store.reserve_curated_attempt(**_attempt()) + assert await store.finish_curated_attempt(claim, **usage) == "overrun" + state = await store.curated_budget_state(now=1001) + assert state["unknown_tokens"] == (400 if "input_tokens" in usage else 300) + assert state["unknown_usd"] == 0.02 + assert await store.reserve_curated_attempt(**_attempt("new", usd_cap="0.029")) is None + rows = await store._conn.execute_fetchall( + "SELECT actual_cost_usd FROM curated_model_attempt WHERE id = ?", (claim,), + ) + assert rows[0][0] == "0.02" + finally: + await store.close() + + +async def test_curated_overrun_blocks_new_workflows_after_midnight_and_cap_increase( + tmp_path, monkeypatch, +): + path = str(tmp_path / "roger.db") + monkeypatch.setattr("roger.store._today", lambda: "2026-10-01") + store = await Store(path).open() + claim = await store.reserve_curated_attempt(**_attempt()) + assert await store.finish_curated_attempt( + claim, input_tokens=10, output_tokens=20, cost_usd="0.02", + ) == "overrun" + await store.close() + monkeypatch.setattr("roger.store._today", lambda: "2026-10-02") + store = await Store(path).open() + try: + request = _attempt("new-workflow", token_cap=100_000, usd_cap="10") + assert await store.reserve_curated_attempt(**request) is None + state = await store.curated_budget_state(request=request) + assert state["denial_reason"] == "reservation_overrun" + assert state["overrun_attempts"] == 1 + assert state["attempts_today"] == 0 + assert state["unknown_usd"] == 0.02 + finally: + await store.close() + + +async def test_curated_unknowns_survive_restart_midnight_and_prune(tmp_path, monkeypatch): + path = str(tmp_path / "roger.db") + monkeypatch.setattr("roger.store._today", lambda: "2026-10-01") + store = await Store(path).open() + claim = await store.reserve_curated_attempt(**_attempt("carried")) + pending = await store.reserve_curated_attempt(**_attempt("unknown")) + await store.close() + monkeypatch.setattr("roger.store._today", lambda: "2026-10-02") + store = await Store(path).open() + try: + state = await store.curated_budget_state(now=1100) + assert state["attempts_today"] == state["workflows_today"] == 0 + assert state["reserved_tokens"] == 600 + assert state["unknown_tokens"] == 0 + state = await store.curated_budget_state(now=2200) + assert state["reserved_tokens"] == 0 + assert state["unknown_tokens"] == 600 + assert await store.reserve_curated_attempt(**_attempt("new", token_cap=899)) is None + assert await store.finish_curated_attempt( + claim, input_tokens=10, output_tokens=20, cost_usd="0.001", now=2201, + ) == "settled" + assert await store.usage_today("curated") == 30 + assert await store.cost_today("curated") == 0.001 + assert await store.reserve_curated_attempt(**_attempt("carried", now=2202)) is not None + state = await store.curated_budget_state(now=2202) + assert state["attempts_today"] == state["workflows_today"] == 1 + assert state["unknown_tokens"] == 300 + await store.prune(now=time.time() + 400 * 86400) + assert (await store.curated_budget_state(now=2202))["unknown_tokens"] == 300 + assert await store.mark_curated_attempt_uncertain(pending) == "uncertain" + finally: + await store.close() + + +async def test_curated_fractional_nanocharges_round_against_admission(tmp_path): + store = await Store(str(tmp_path / "roger.db")).open() + try: + await store.add_usage("curated", 0, 0, 0.0000000001) + request = _attempt(cost_usd="0.0000000001", usd_cap="0.0000000019") + assert await store.reserve_curated_attempt(**request) is None + assert (await store.curated_budget_state(request=request))["denial_reason"] == "usd" + request["usd_cap"] = "0.000000002" + assert await store.reserve_curated_attempt(**request) is not None + finally: + await store.close() + + +@pytest.mark.parametrize("invalid", [ + {"workflow_id": ""}, {"workflow_id": "secret/request"}, {"profile": "x" * 81}, + {"role": "admin"}, {"input_tokens": 32_769}, {"output_tokens": 16_385}, + {"input_tokens": True}, {"max_attempts": -1}, {"cost_usd": "NaN"}, + {"usd_cap": -1}, {"timeout_seconds": 0}, {"now": float("inf")}, +]) +async def test_curated_invalid_reservation_cannot_create_a_ledger_row(tmp_path, invalid): + store = await Store(str(tmp_path / "roger.db")).open() + try: + with pytest.raises(ValueError): + await store.reserve_curated_attempt(**_attempt(**invalid)) + assert (await store.curated_budget_state())["attempts_today"] == 0 + finally: + await store.close() + + async def test_record_audit_persists(tmp_path): store = await Store(str(tmp_path / "roger.db")).open() try: diff --git a/tests/test_stories.py b/tests/test_stories.py index 7e76ab6..b3b04c7 100644 --- a/tests/test_stories.py +++ b/tests/test_stories.py @@ -632,6 +632,7 @@ async def collect(*args, **kwargs): settings = Settings( discord_token="test", openrouter_api_key="test", owner_id=1, guild_id=2, # noqa: S106 curated_request_profile="sol-astra-v1", curated_developing_stories=True, + daily_usd_curated=2, daily_attempts_curated=8, daily_workflows_curated=1, scout_digest_dir=str(tmp_path), ) store = await Store(str(tmp_path / "profile-preview.db")).open() @@ -672,12 +673,33 @@ async def create(**kwargs): assert [call["model"] for call in calls] == [ curated_protocol.CURATED_WRITER.model, curated_protocol.CURATED_REVIEWER.model, ] - assert await store.usage_today("curated") == 300 - assert await store.cost_today("curated") == 0.02 + assert await store.usage_today("curated") == (150 if review_failure else 300) + assert await store.cost_today("curated") == (0.01 if review_failure else 0.02) + budget = await store.curated_budget_state() + assert budget["attempts_today"] == 2 + assert budget["workflows_today"] == 1 + assert budget["unknown_attempts"] == int(review_failure) for table in ("curated_story", "curated_story_decision", "curated_delivery", "seen", "curated_observation"): cursor = await store._conn.execute(f"SELECT COUNT(*) FROM {table}") # noqa: S608 assert (await cursor.fetchone())[0] == 0 + settings.curated_channel_id = 42 + + async def send(**kwargs): + raise AssertionError("the scheduled check shares the preview's workflow allowance") + + client = SimpleNamespace(get_channel=lambda _: SimpleNamespace(send=send)) + scheduled = await run_curated_job( + client=client, settings=settings, llm=llm, store=store, now=now, + ) + assert scheduled["status"] == "budget exceeded (daily_workflows); skipped" + assert len(calls) == 2 + assert (await store.curated_budget_state())["workflows_today"] == 1 + preview_denied = await preview_curated_job( + settings=settings, llm=llm, store=store, developing_stories=True, + ) + assert preview_denied["status"] == "budget exceeded (daily_workflows); skipped" + assert len(calls) == 2 finally: await llm._client.close() await store.close() @@ -705,6 +727,7 @@ async def decide(*args, **kwargs): discord_token="test", openrouter_api_key="test", owner_id=1, guild_id=2, # noqa: S106 curated_channel_id=42, curated_request_profile="sol-astra-v1", curated_developing_stories=True, curated_min_spacing_minutes=0, + daily_usd_curated=2, daily_attempts_curated=8, daily_workflows_curated=2, scout_digest_dir=str(tmp_path), ) store = await Store(str(tmp_path / "profile-job.db")).open()