From 02c989946b17a3b95eddadc9e593ed0ca5d78a32 Mon Sep 17 00:00:00 2001 From: Oleksandr Kozachuk Date: Mon, 13 Jul 2026 13:21:45 +0200 Subject: [PATCH] safety layer: hard daily budget, user quotas, spend report, forgetme + privacy --- DECISIONS.md | 10 ++ config.toml | 9 ++ fjerkroa_bot/discord_bot.py | 60 +++++++++ fjerkroa_bot/openai_responder.py | 20 +++ fjerkroa_bot/persistence.py | 29 +++- fjerkroa_bot/quota.py | 77 +++++++++++ specs/SPEC-003-safety.md | 44 ++++++ specs/SPEC-006-operator-controls.md | 6 + specs/SPEC-009-persistence.md | 9 +- tests/test_spec_per.py | 6 +- tests/test_spec_quota.py | 201 ++++++++++++++++++++++++++++ 11 files changed, 465 insertions(+), 6 deletions(-) create mode 100644 fjerkroa_bot/quota.py create mode 100644 tests/test_spec_quota.py diff --git a/DECISIONS.md b/DECISIONS.md index b9b5a5e..1749038 100644 --- a/DECISIONS.md +++ b/DECISIONS.md @@ -43,3 +43,13 @@ Decisions inside the set architecture. D-NNN, never renumbered. per message instead of appending: histories are capped at `history-limit` (~200-350 rows) and trims must be reflected; correctness over micro-optimization. +- **D-012** — Budget spend is *estimated* from configured per-token/ + per-image prices, not fetched from the billing API: deterministic, + testable, no extra scopes. Dashboard hard limits (FDB-001) stay the + outer safety net; this ledger is the inner, immediate one. User + image quota counts at grant time (reservation), global image spend + at generation time. +- **D-013** — `!forgetme` v1 purges history rows only; the + single-string channel memory cannot be selectively cleaned. Full + fact-level erasure ships with FDB-007 structured memory — stated in + the user-facing confirmation, not hidden. diff --git a/config.toml b/config.toml index 028014b..b4605cf 100644 --- a/config.toml +++ b/config.toml @@ -27,3 +27,12 @@ enable-game-info = true # staff-alert-max-per-hour = 10 # Staff commands (staff channel only): !bot pause | resume | images on|off # | tasks on|off | quiet | status +# Cost governance (SPEC-003 SAF-04..07) — budget is a HARD cap, fail-closed: +# daily-budget-usd = 2.0 +# price-input-per-m = 1.0 # USD per 1M input tokens (gpt-5.6-luna) +# price-output-per-m = 6.0 # USD per 1M output tokens +# price-per-image = 0.05 +# user-daily-messages = 200 +# user-daily-images = 10 +# Privacy (SAF-08/09): users can always run !forgetme and !privacy +# privacy-notice = "I keep recent messages and a summary. !forgetme deletes yours." diff --git a/fjerkroa_bot/discord_bot.py b/fjerkroa_bot/discord_bot.py index 5d770a4..eccb0eb 100644 --- a/fjerkroa_bot/discord_bot.py +++ b/fjerkroa_bot/discord_bot.py @@ -19,6 +19,11 @@ from watchdog.observers import Observer from .ai_responder import AIMessage from .openai_responder import OpenAIResponder +DEFAULT_PRIVACY_NOTICE = ( + "I keep recent channel messages and a short conversation summary to answer better. " + "Type !forgetme to remove your messages from my history. Questions: ask the staff." +) + class ConfigFileHandler(FileSystemEventHandler): def __init__(self, on_modified): @@ -125,6 +130,14 @@ class FjerkroaBot(commands.Bot): if self.is_staff_channel(message.channel) and str(message.content).startswith("!bot"): await self.handle_staff_command(message) return + # user-rights commands work even while paused (SAF-08/09) + content = str(message.content).strip().lower() + if content.startswith("!forgetme"): + await self.forget_user(message) + return + if content.startswith("!privacy"): + await message.channel.send(self.config.get("privacy-notice", DEFAULT_PRIVACY_NOTICE), suppress_embeds=True) + return if not self.replies_allowed(): return if str(message.content).startswith("!wichtel"): @@ -132,6 +145,24 @@ class FjerkroaBot(commands.Bot): return await self.handle_message_through_responder(message) + async def forget_user(self, message: Message) -> None: + """Purge the requesting user's messages everywhere (SAF-08).""" + user = message.author.name + removed = 0 + for responder in [self.airesponder, *self.aichannels.values()]: + before = len(responder.history) + responder.history = [item for item in responder.history if f'"user": "{user}"' not in str(item.get("content", ""))] + removed += before - len(responder.history) + await responder._persist_history() + if self.airesponder.store is not None: + removed += self.airesponder.store.delete_history_of_user(user) + logging.info(f"forgetme: removed {removed} history entries for {user}") + await message.channel.send( + f"Removed your messages from my history ({removed} entries). " + "Note: channel memory summaries may still hold traces until the structured-memory upgrade.", + suppress_embeds=True, + ) + def is_staff_channel(self, channel) -> bool: staff = getattr(self, "staff_channel", None) return staff is not None and getattr(channel, "id", None) == getattr(staff, "id", None) @@ -166,6 +197,11 @@ class FjerkroaBot(commands.Bot): elif args[:1] == ["status"]: quiet_left = max(0, int(self.quiet_until - time.monotonic())) reply = f"replies={self.replies_enabled} images={self.images_enabled} tasks={self.tasks_enabled} quiet_left={quiet_left}s" + elif args[:1] == ["spend"]: + ledger = self.airesponder.ledger + tokens_in, tokens_out = ledger.tokens_today() + budget = self.config.get("daily-budget-usd", "none") + reply = f"spend today: ${ledger.spent_usd():.2f} (tokens {tokens_in}/{tokens_out}, images {ledger.images_today()}), budget: {budget}" logging.info(f"staff command {args}: {reply}") await message.channel.send(reply, suppress_embeds=True) @@ -179,6 +215,12 @@ class FjerkroaBot(commands.Bot): allowed += list(self.config.get("additional-responders", [])) return channel_name in [name for name in allowed if name] + async def _budget_alert_once(self) -> None: + today = time.strftime("%Y-%m-%d") + if getattr(self, "_budget_alert_day", None) != today: + self._budget_alert_day = today + await self.send_staff_alert("Daily budget exhausted - bot stays silent until midnight (SAF-04).") + async def send_staff_alert(self, text: str) -> None: """Rate-limited, never silently dropped (OPS-07/08).""" if self.staff_channel is None: @@ -366,6 +408,13 @@ class FjerkroaBot(commands.Bot): if response.picture is not None and not self.images_enabled: logging.info("image generation disabled by operator - sending text only") response.picture = None + # Per-user daily image quota (SAF-07) + if response.picture is not None and message.user != "system" and "user-daily-images" in self.config: + if self.airesponder.ledger.user_images(message.user) >= int(self.config["user-daily-images"]): + logging.warning(f"user {message.user} over daily image quota - stripping picture") + response.picture = None + else: + self.airesponder.ledger.count_user_image(message.user) async def respond( self, @@ -385,6 +434,17 @@ class FjerkroaBot(commands.Bot): # In case the message shouldn't be ignored, log the handling action self.log_message_action("handle", message, channel_name) + # Hard daily budget, fail-closed; staff hears once per day (SAF-04) + if not self.airesponder.ledger.budget_ok(): + await self._budget_alert_once() + return + + # Per-user daily message quota; system (bot-initiated) exempt (SAF-06) + if message.user != "system" and "user-daily-messages" in self.config: + if self.airesponder.ledger.count_user_message(message.user) > int(self.config["user-daily-messages"]): + logging.warning(f"user {message.user} over daily message quota - ignoring") + return + # Get the AI responder based on the channel name airesponder = self.get_ai_responder(channel_name) diff --git a/fjerkroa_bot/openai_responder.py b/fjerkroa_bot/openai_responder.py index 78571ac..7da9999 100644 --- a/fjerkroa_bot/openai_responder.py +++ b/fjerkroa_bot/openai_responder.py @@ -10,6 +10,7 @@ import openai from .ai_responder import AIResponder, exponential_backoff, pp, sanitize_external_text from .igdblib import IGDBQuery from .leonardo_draw import LeonardoAIDrawMixIn +from .quota import QuotaLedger # The response envelope, enforced server-side via structured outputs # (ENV-19). All fields required, closed object, nullable where the @@ -48,6 +49,8 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn): self.client = openai.AsyncOpenAI(api_key=self.config.get("openai-token", self.config.get("openai-key", ""))) # After a rate limit the next attempt runs on retry-model (ENV-15 / D2) self._use_retry_model = False + # Daily usage metering + hard budget, fail-closed (SAF-04/05) + self.ledger = QuotaLedger(self.store, lambda: self.config) # Initialize IGDB if enabled self.igdb = None @@ -69,21 +72,36 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn): logging.warning("❌ IGDB integration DISABLED - missing configuration or disabled in config") async def draw_openai(self, description: str) -> BytesIO: + if not self.ledger.budget_ok(): + raise RuntimeError("daily budget exhausted - refusing image call") for _ in range(3): try: response = await openai_image(self.client, prompt=description, n=1, size="1024x1024", model="dall-e-3") + self.ledger.add_images(1) logging.info(f"Drawed a picture with DALL-E on this description: {repr(description)}") return response except Exception as err: logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}") raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries") + def _record_usage(self, result: Any) -> None: + usage = getattr(result, "usage", None) + prompt_tokens = getattr(usage, "prompt_tokens", None) + completion_tokens = getattr(usage, "completion_tokens", None) + if isinstance(prompt_tokens, int) and isinstance(completion_tokens, int): + self.ledger.add_tokens(prompt_tokens, completion_tokens) + async def chat(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]: # Safety check for mock objects in tests if not isinstance(messages, list) or len(messages) == 0: logging.warning("Invalid messages format in chat method") return None, limit + # Hard daily budget, fail-closed (SAF-04) + if not self.ledger.budget_ok(): + logging.error("daily budget exhausted - refusing model call") + return None, limit + try: # Clean up any orphaned tool messages from previous conversations clean_messages = [] @@ -134,6 +152,7 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn): ) result = await openai_chat(self.client, **chat_kwargs) + self._record_usage(result) # Handle function calls if present message = result.choices[0].message @@ -204,6 +223,7 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn): logging.debug(f"🔧 Last few messages: {messages[-3:] if len(messages) > 3 else messages}") final_result = await openai_chat(self.client, **final_chat_kwargs) + self._record_usage(final_result) answer_obj = final_result.choices[0].message logging.debug( diff --git a/fjerkroa_bot/persistence.py b/fjerkroa_bot/persistence.py index dae1ecc..213676c 100644 --- a/fjerkroa_bot/persistence.py +++ b/fjerkroa_bot/persistence.py @@ -13,7 +13,7 @@ from contextlib import closing from pathlib import Path from typing import Any, Dict, List, Optional -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 class PersistentStore: @@ -27,14 +27,21 @@ class PersistentStore: return conn def _init_db(self) -> None: + # Forward-only migrations keyed on user_version (PER-06) self.db_path.parent.mkdir(parents=True, exist_ok=True) with closing(self._connect()) as conn, conn: - if conn.execute("PRAGMA user_version").fetchone()[0] < SCHEMA_VERSION: + version = conn.execute("PRAGMA user_version").fetchone()[0] + if version < 1: conn.execute( "CREATE TABLE IF NOT EXISTS history (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL)" ) conn.execute("CREATE INDEX IF NOT EXISTS history_channel ON history (channel)") conn.execute("CREATE TABLE IF NOT EXISTS memory (channel TEXT PRIMARY KEY, content TEXT NOT NULL)") + if version < 2: + conn.execute( + "CREATE TABLE IF NOT EXISTS usage (day TEXT NOT NULL, key TEXT NOT NULL, value REAL NOT NULL, PRIMARY KEY (day, key))" + ) + if version < SCHEMA_VERSION: conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") os.chmod(self.db_path, 0o600) # conversation data (PER-04) @@ -65,6 +72,24 @@ class PersistentStore: (channel, content), ) + def usage_add(self, day: str, key: str, amount: float) -> None: + with closing(self._connect()) as conn, conn: + conn.execute( + "INSERT INTO usage (day, key, value) VALUES (?, ?, ?) ON CONFLICT(day, key) DO UPDATE SET value = value + excluded.value", + (day, key, amount), + ) + + def usage_get(self, day: str, key: str) -> float: + with closing(self._connect()) as conn: + row = conn.execute("SELECT value FROM usage WHERE day = ? AND key = ?", (day, key)).fetchone() + return float(row[0]) if row else 0.0 + + def delete_history_of_user(self, user: str) -> int: + """Remove persisted rows carrying this user's messages (SAF-08).""" + with closing(self._connect()) as conn, conn: + cursor = conn.execute("DELETE FROM history WHERE content LIKE ?", (f'%"user": "{user}"%',)) + return cursor.rowcount + def migrate_pickles(self, channel: str, history_file: Path, memory_file: Path) -> None: """Import legacy pickles once; rename them *.migrated (PER-03).""" if self.load_history(channel) or self.load_memory(channel) is not None: diff --git a/fjerkroa_bot/quota.py b/fjerkroa_bot/quota.py new file mode 100644 index 0000000..4845a3b --- /dev/null +++ b/fjerkroa_bot/quota.py @@ -0,0 +1,77 @@ +"""Daily usage metering + hard budget (SPEC-003, SAF-04..07). + +The ledger estimates spend from configured prices and answers the one +question that matters fail-closed: may the bot still call the API +today? Counters live in the store's usage table when a store exists +(SAF-05), else in memory (degraded but safe). +""" + +import time +from typing import Callable, Dict, Optional, Tuple + +from .persistence import PersistentStore + +DEFAULT_PRICE_INPUT_PER_M = 1.0 +DEFAULT_PRICE_OUTPUT_PER_M = 6.0 +DEFAULT_PRICE_PER_IMAGE = 0.05 + + +class QuotaLedger: + def __init__(self, store: Optional[PersistentStore], config_getter: Callable[[], Dict]) -> None: + self.store = store + self._config = config_getter + self._memory: Dict[Tuple[str, str], float] = {} + + @staticmethod + def _day() -> str: + return time.strftime("%Y-%m-%d") + + def _add(self, key: str, amount: float) -> None: + if self.store is not None: + self.store.usage_add(self._day(), key, amount) + else: + slot = (self._day(), key) + self._memory[slot] = self._memory.get(slot, 0.0) + amount + + def _get(self, key: str) -> float: + if self.store is not None: + return self.store.usage_get(self._day(), key) + return self._memory.get((self._day(), key), 0.0) + + def add_tokens(self, prompt_tokens: int, completion_tokens: int) -> None: + self._add("tokens-in", prompt_tokens) + self._add("tokens-out", completion_tokens) + + def add_images(self, count: int = 1) -> None: + self._add("images", count) + + def tokens_today(self) -> Tuple[int, int]: + return int(self._get("tokens-in")), int(self._get("tokens-out")) + + def images_today(self) -> int: + return int(self._get("images")) + + def count_user_message(self, user: str) -> int: + self._add(f"msg-user:{user}", 1) + return int(self._get(f"msg-user:{user}")) + + def count_user_image(self, user: str) -> int: + self._add(f"img-user:{user}", 1) + return int(self._get(f"img-user:{user}")) + + def user_images(self, user: str) -> int: + return int(self._get(f"img-user:{user}")) + + def spent_usd(self) -> float: + config = self._config() + tokens_in, tokens_out = self.tokens_today() + price_in = float(config.get("price-input-per-m", DEFAULT_PRICE_INPUT_PER_M)) + price_out = float(config.get("price-output-per-m", DEFAULT_PRICE_OUTPUT_PER_M)) + price_image = float(config.get("price-per-image", DEFAULT_PRICE_PER_IMAGE)) + return (tokens_in * price_in + tokens_out * price_out) / 1_000_000.0 + self.images_today() * price_image + + def budget_ok(self) -> bool: + config = self._config() + if "daily-budget-usd" not in config: + return True + return self.spent_usd() < float(config["daily-budget-usd"]) diff --git a/specs/SPEC-003-safety.md b/specs/SPEC-003-safety.md index 64cb5c0..f3c8e93 100644 --- a/specs/SPEC-003-safety.md +++ b/specs/SPEC-003-safety.md @@ -31,3 +31,47 @@ and game descriptions are attacker-influenced input. The `hack` envelope field remains as an advisory signal (logged, staff-notified) but is no longer the defense. + +### SAF-04 — Hard daily budget, fail-closed (coverage: test) + +When `daily-budget-usd` is configured and today's estimated spend +reaches it, no further model or image calls happen: the responder +refuses before calling the API, and the bot answers nothing until +midnight. Staff get exactly one alert per day about the silence. +Cost estimation uses `price-input-per-m` (default 1.0), +`price-output-per-m` (default 6.0) and `price-per-image` (default +0.05). A budget of 0 means fully silent — fail-closed by +construction. + +### SAF-05 — Usage is metered and survives restarts (coverage: test) + +Token counts (prompt/completion) from every model call and every +generated image are recorded per calendar day; with a store +configured the counters persist across restarts (usage table, +schema v2). + +### SAF-06 — Per-user daily message quota (coverage: test) + +When `user-daily-messages` is configured, messages beyond the cap +from one user on one day are ignored (logged, no model call). The +`system` user (bot-initiated flows) is exempt. + +### SAF-07 — Per-user daily image quota (coverage: test) + +When `user-daily-images` is configured, picture requests beyond the +user's daily cap are stripped from the response (the text answer +still goes out). + +### SAF-08 — !forgetme purges a user's history (coverage: test) + +`!forgetme` removes the requesting user's messages from all live +responder histories and from the store, then confirms in-channel. +v1 limitation, stated in the confirmation: the single-string channel +memory may still contain summarized traces — full fact-level erasure +arrives with the structured memory (FDB-007). + +### SAF-09 — !privacy states the data practice (coverage: test) + +`!privacy` answers with the configured `privacy-notice` (a default +notice ships in code): what is stored, that `!forgetme` exists. +Works even while the bot is paused. diff --git a/specs/SPEC-006-operator-controls.md b/specs/SPEC-006-operator-controls.md index 7230844..8ba19a9 100644 --- a/specs/SPEC-006-operator-controls.md +++ b/specs/SPEC-006-operator-controls.md @@ -56,3 +56,9 @@ the alert text is written to the error log — never silently dropped. `!bot tasks off` disables bot-initiated posting (today: the boreness loop; later: the FDB-011 scheduler); `!bot tasks on` restores. Bot-initiated posts also respect pause/quiet. + +### OPS-10 — Spend report (coverage: test) + +`!bot spend` answers in the staff channel with today's estimated +spend in USD, token and image counts, and the configured budget. +Management sees the cost, not just the cap. diff --git a/specs/SPEC-009-persistence.md b/specs/SPEC-009-persistence.md index e2cc1ef..511835d 100644 --- a/specs/SPEC-009-persistence.md +++ b/specs/SPEC-009-persistence.md @@ -28,9 +28,16 @@ accident). ### PER-04 — Database hygiene (coverage: test) The database runs in WAL journal mode, carries `PRAGMA user_version` -= schema version (currently 1) for future migrations/rollback += the store's current schema version for the migration/rollback policy, and the file is chmod 0600 (it stores conversation data). +### PER-06 — Schema migrations run forward automatically (coverage: test) + +Opening a database with an older `user_version` applies the missing +migration steps in order (v1 → v2 adds the usage table) and preserves +existing rows. Deploy rollback policy: never roll binaries back past +a schema bump without restoring the pre-deploy backup. + ### PER-05 — Writes run off the event loop (coverage: test) History and memory persistence happen in a worker thread diff --git a/tests/test_spec_per.py b/tests/test_spec_per.py index 3a0e0db..7021ede 100644 --- a/tests/test_spec_per.py +++ b/tests/test_spec_per.py @@ -10,7 +10,7 @@ from pathlib import Path from unittest.mock import MagicMock from fjerkroa_bot.ai_responder import AIMessage, AIResponder -from fjerkroa_bot.persistence import PersistentStore +from fjerkroa_bot.persistence import SCHEMA_VERSION, PersistentStore from .test_bdd_envelope import FakeModelResponder, envelope from .test_main import TestBotBase @@ -78,12 +78,12 @@ class TestPickleMigration(StoreBase): class TestDatabaseHygiene(StoreBase): def test_wal_version_and_permissions(self): - """PER-04: WAL mode, user_version = 1, file mode 0600.""" + """PER-04: WAL mode, user_version = current schema version, file mode 0600.""" PersistentStore(self.db_path) conn = sqlite3.connect(self.db_path) try: self.assertEqual(conn.execute("PRAGMA journal_mode").fetchone()[0], "wal") - self.assertEqual(conn.execute("PRAGMA user_version").fetchone()[0], 1) + self.assertEqual(conn.execute("PRAGMA user_version").fetchone()[0], SCHEMA_VERSION) finally: conn.close() mode = stat.S_IMODE(self.db_path.stat().st_mode) diff --git a/tests/test_spec_quota.py b/tests/test_spec_quota.py new file mode 100644 index 0000000..6339c8d --- /dev/null +++ b/tests/test_spec_quota.py @@ -0,0 +1,201 @@ +"""Unit coverage for FDB-014: SAF-04..09, OPS-10, PER-06.""" + +import json +import sqlite3 +import tempfile +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +from discord import TextChannel + +from fjerkroa_bot.ai_responder import AIMessage, AIResponse +from fjerkroa_bot.openai_responder import OpenAIResponder +from fjerkroa_bot.persistence import SCHEMA_VERSION, PersistentStore +from fjerkroa_bot.quota import QuotaLedger + +from .test_bdd_envelope import envelope +from .test_spec_ops import OpsBase + + +def ledger_with_store(tmp, config): + store = PersistentStore(Path(tmp) / "bot.db") + return QuotaLedger(store, lambda: config) + + +class TestBudgetFailClosed(unittest.IsolatedAsyncioTestCase): + def test_spend_math_and_cutoff(self): + """SAF-04: spend estimate from configured prices; budget reached -> not ok.""" + config = {"daily-budget-usd": 0.01} + ledger = QuotaLedger(None, lambda: config) + self.assertTrue(ledger.budget_ok()) + ledger.add_tokens(20000, 0) # 20k in * $1/M = $0.02 >= $0.01 + self.assertFalse(ledger.budget_ok()) + + def test_zero_budget_is_silent(self): + """SAF-04: budget 0 -> fail-closed immediately.""" + ledger = QuotaLedger(None, lambda: {"daily-budget-usd": 0}) + self.assertFalse(ledger.budget_ok()) + + def test_no_budget_key_means_unlimited(self): + """SAF-04: without daily-budget-usd the gate stays open.""" + ledger = QuotaLedger(None, lambda: {}) + ledger.add_tokens(10_000_000, 10_000_000) + self.assertTrue(ledger.budget_ok()) + + async def test_chat_refuses_over_budget(self): + """SAF-04: exhausted budget -> chat() returns None without an API call.""" + config = {"openai-token": "t", "model": "m", "system": "s", "history-limit": 5, "daily-budget-usd": 0} + responder = OpenAIResponder(config, "chat") + with patch("fjerkroa_bot.openai_responder.openai_chat", new_callable=AsyncMock) as chat_mock: + answer, _ = await responder.chat([{"role": "user", "content": "hi"}], 10) + self.assertIsNone(answer) + chat_mock.assert_not_awaited() + + +class TestBudgetStaffAlert(OpsBase): + async def test_alert_once_and_silence(self): + """SAF-04: one staff alert per day, responder never called while exhausted.""" + self.bot.config["daily-budget-usd"] = 0 + self.bot.send_message_with_typing = AsyncMock() + origin = MagicMock(spec=TextChannel) + await self.bot.respond(AIMessage("alice", "hei", "chat"), origin) + await self.bot.respond(AIMessage("alice", "hei again", "chat"), origin) + self.bot.send_message_with_typing.assert_not_awaited() + self.assertEqual(self.bot.staff_channel.send.await_count, 1) + + +class TestUsageMetering(unittest.TestCase): + def test_usage_persists_across_instances(self): + """SAF-05: token/image counters survive a restart via the usage table.""" + with tempfile.TemporaryDirectory() as tmp: + config = {"daily-budget-usd": 100} + ledger = ledger_with_store(tmp, config) + ledger.add_tokens(1000, 500) + ledger.add_images(2) + reborn = ledger_with_store(tmp, config) + self.assertGreater(reborn.spent_usd(), 0) + self.assertEqual(reborn.images_today(), 2) + + +class TestChatRecordsUsage(unittest.IsolatedAsyncioTestCase): + async def test_tokens_recorded_from_api_usage(self): + """SAF-05: chat() feeds prompt/completion token counts into the ledger.""" + config = {"openai-token": "t", "model": "m", "system": "s", "history-limit": 5} + responder = OpenAIResponder(config, "chat") + message = Mock(content=envelope(answer="x", answer_needed=True), role="assistant", tool_calls=None, refusal=None) + usage = Mock(prompt_tokens=1234, completion_tokens=56) + with patch("fjerkroa_bot.openai_responder.openai_chat", new_callable=AsyncMock) as chat_mock: + chat_mock.return_value = Mock(choices=[Mock(message=message)], usage=usage) + await responder.chat([{"role": "user", "content": "hi"}], 10) + self.assertEqual(responder.ledger.tokens_today(), (1234, 56)) + + +class TestMessageQuota(OpsBase): + async def test_user_over_message_cap_is_ignored(self): + """SAF-06: messages over user-daily-messages are dropped before the model.""" + self.bot.config["user-daily-messages"] = 2 + self.bot.send_message_with_typing = AsyncMock(return_value=AIResponse(None, False, "chat", None, None, False, False)) + origin = MagicMock(spec=TextChannel) + for _ in range(3): + await self.bot.respond(AIMessage("alice", "hei", "chat"), origin) + self.assertEqual(self.bot.send_message_with_typing.await_count, 2) + + async def test_system_user_exempt(self): + """SAF-06: bot-initiated (system) messages bypass the user quota.""" + self.bot.config["user-daily-messages"] = 1 + self.bot.send_message_with_typing = AsyncMock(return_value=AIResponse(None, False, "chat", None, None, False, False)) + origin = MagicMock(spec=TextChannel) + for _ in range(3): + await self.bot.respond(AIMessage("system", "impulse", "chat", True, False), origin) + self.assertEqual(self.bot.send_message_with_typing.await_count, 3) + + +class TestImageQuota(OpsBase): + async def test_picture_stripped_over_cap(self): + """SAF-07: picture requests over user-daily-images are stripped, text kept.""" + self.bot.config["user-daily-images"] = 1 + self.bot.send_answer_with_typing = AsyncMock() + origin = MagicMock(spec=TextChannel) + for _ in range(2): + response = AIResponse("here", True, "chat", None, "a cat", False, False) + self.bot.send_message_with_typing = AsyncMock(return_value=response) + await self.bot.respond(AIMessage("alice", "draw", "chat"), origin) + first = self.bot.send_answer_with_typing.await_args_list[0].args[0] + second = self.bot.send_answer_with_typing.await_args_list[1].args[0] + self.assertEqual(first.picture, "a cat") + self.assertIsNone(second.picture) + + +class TestForgetMe(OpsBase): + async def test_forgetme_purges_live_history_and_confirms(self): + """SAF-08: !forgetme removes the user's rows from live history + confirms.""" + entry = {"role": "user", "content": json.dumps({"user": "alice", "message": "secret", "channel": "chat"})} + other = {"role": "user", "content": json.dumps({"user": "bob", "message": "stays", "channel": "chat"})} + self.bot.airesponder.history = [dict(entry), dict(other)] + message = self.public_msg("!forgetme") + message.author.name = "alice" + await self.bot.on_message(message) + contents = [item["content"] for item in self.bot.airesponder.history] + self.assertFalse(any('"alice"' in content for content in contents)) + self.assertTrue(any('"bob"' in content for content in contents)) + message.channel.send.assert_awaited_once() + + def test_store_purge_by_user(self): + """SAF-08: the store deletes persisted rows containing the user's messages.""" + with tempfile.TemporaryDirectory() as tmp: + store = PersistentStore(Path(tmp) / "bot.db") + store.save_history( + "chat", + [ + {"role": "user", "content": json.dumps({"user": "alice", "message": "x"})}, + {"role": "user", "content": json.dumps({"user": "bob", "message": "y"})}, + ], + ) + store.delete_history_of_user("alice") + remaining = store.load_history("chat") + self.assertEqual(len(remaining), 1) + self.assertIn('"bob"', remaining[0]["content"]) + + +class TestPrivacyNotice(OpsBase): + async def test_privacy_answers_even_when_paused(self): + """SAF-09: !privacy answers with the notice, also while paused.""" + self.bot.replies_enabled = False + self.bot.config["privacy-notice"] = "We store recent messages. Use !forgetme." + message = self.public_msg("!privacy") + await self.bot.on_message(message) + message.channel.send.assert_awaited_once() + self.assertIn("!forgetme", message.channel.send.await_args.args[0]) + + +class TestSpendCommand(OpsBase): + async def test_spend_report(self): + """OPS-10: !bot spend reports estimated USD + counters + budget.""" + await self.bot.on_message(self.staff_msg("!bot spend")) + self.bot.staff_channel.send.assert_awaited() + text = self.bot.staff_channel.send.await_args.args[0] + self.assertIn("$", text) + self.assertIn("tokens", text) + + +class TestSchemaMigration(unittest.TestCase): + def test_v1_database_upgrades_to_current(self): + """PER-06: v1 db gains the usage table, keeps rows, bumps user_version.""" + with tempfile.TemporaryDirectory() as tmp: + db_path = Path(tmp) / "bot.db" + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE history (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL)") + conn.execute("CREATE TABLE memory (channel TEXT PRIMARY KEY, content TEXT NOT NULL)") + conn.execute("INSERT INTO history (channel, role, content) VALUES ('chat', 'user', 'kept')") + conn.execute("PRAGMA user_version = 1") + conn.commit() + conn.close() + store = PersistentStore(db_path) + self.assertEqual(store.load_history("chat"), [{"role": "user", "content": "kept"}]) + store.usage_add("2026-07-13", "tokens-in", 5) + check = sqlite3.connect(db_path) + try: + self.assertEqual(check.execute("PRAGMA user_version").fetchone()[0], SCHEMA_VERSION) + finally: + check.close()