safety layer: hard daily budget, user quotas, spend report, forgetme + privacy

This commit is contained in:
Oleksandr Kozachuk
2026-07-13 13:21:45 +02:00
parent f6c3e7d8e5
commit 02c989946b
11 changed files with 465 additions and 6 deletions
+10
View File
@@ -43,3 +43,13 @@ Decisions inside the set architecture. D-NNN, never renumbered.
per message instead of appending: histories are capped at per message instead of appending: histories are capped at
`history-limit` (~200-350 rows) and trims must be reflected; `history-limit` (~200-350 rows) and trims must be reflected;
correctness over micro-optimization. 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.
+9
View File
@@ -27,3 +27,12 @@ enable-game-info = true
# staff-alert-max-per-hour = 10 # staff-alert-max-per-hour = 10
# Staff commands (staff channel only): !bot pause | resume | images on|off # Staff commands (staff channel only): !bot pause | resume | images on|off
# | tasks on|off | quiet <minutes> | status # | tasks on|off | quiet <minutes> | 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."
+60
View File
@@ -19,6 +19,11 @@ from watchdog.observers import Observer
from .ai_responder import AIMessage from .ai_responder import AIMessage
from .openai_responder import OpenAIResponder 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): class ConfigFileHandler(FileSystemEventHandler):
def __init__(self, on_modified): 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"): if self.is_staff_channel(message.channel) and str(message.content).startswith("!bot"):
await self.handle_staff_command(message) await self.handle_staff_command(message)
return 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(): if not self.replies_allowed():
return return
if str(message.content).startswith("!wichtel"): if str(message.content).startswith("!wichtel"):
@@ -132,6 +145,24 @@ class FjerkroaBot(commands.Bot):
return return
await self.handle_message_through_responder(message) 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: def is_staff_channel(self, channel) -> bool:
staff = getattr(self, "staff_channel", None) staff = getattr(self, "staff_channel", None)
return staff is not None and getattr(channel, "id", None) == getattr(staff, "id", 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"]: elif args[:1] == ["status"]:
quiet_left = max(0, int(self.quiet_until - time.monotonic())) 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" 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}") logging.info(f"staff command {args}: {reply}")
await message.channel.send(reply, suppress_embeds=True) await message.channel.send(reply, suppress_embeds=True)
@@ -179,6 +215,12 @@ class FjerkroaBot(commands.Bot):
allowed += list(self.config.get("additional-responders", [])) allowed += list(self.config.get("additional-responders", []))
return channel_name in [name for name in allowed if name] 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: async def send_staff_alert(self, text: str) -> None:
"""Rate-limited, never silently dropped (OPS-07/08).""" """Rate-limited, never silently dropped (OPS-07/08)."""
if self.staff_channel is None: 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: if response.picture is not None and not self.images_enabled:
logging.info("image generation disabled by operator - sending text only") logging.info("image generation disabled by operator - sending text only")
response.picture = None 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( async def respond(
self, self,
@@ -385,6 +434,17 @@ class FjerkroaBot(commands.Bot):
# In case the message shouldn't be ignored, log the handling action # In case the message shouldn't be ignored, log the handling action
self.log_message_action("handle", message, channel_name) 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 # Get the AI responder based on the channel name
airesponder = self.get_ai_responder(channel_name) airesponder = self.get_ai_responder(channel_name)
+20
View File
@@ -10,6 +10,7 @@ import openai
from .ai_responder import AIResponder, exponential_backoff, pp, sanitize_external_text from .ai_responder import AIResponder, exponential_backoff, pp, sanitize_external_text
from .igdblib import IGDBQuery from .igdblib import IGDBQuery
from .leonardo_draw import LeonardoAIDrawMixIn from .leonardo_draw import LeonardoAIDrawMixIn
from .quota import QuotaLedger
# The response envelope, enforced server-side via structured outputs # The response envelope, enforced server-side via structured outputs
# (ENV-19). All fields required, closed object, nullable where the # (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", ""))) 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) # After a rate limit the next attempt runs on retry-model (ENV-15 / D2)
self._use_retry_model = False 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 # Initialize IGDB if enabled
self.igdb = None self.igdb = None
@@ -69,21 +72,36 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn):
logging.warning("❌ IGDB integration DISABLED - missing configuration or disabled in config") logging.warning("❌ IGDB integration DISABLED - missing configuration or disabled in config")
async def draw_openai(self, description: str) -> BytesIO: 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): for _ in range(3):
try: try:
response = await openai_image(self.client, prompt=description, n=1, size="1024x1024", model="dall-e-3") 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)}") logging.info(f"Drawed a picture with DALL-E on this description: {repr(description)}")
return response return response
except Exception as err: except Exception as err:
logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}") logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}")
raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries") 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]: async def chat(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]:
# Safety check for mock objects in tests # Safety check for mock objects in tests
if not isinstance(messages, list) or len(messages) == 0: if not isinstance(messages, list) or len(messages) == 0:
logging.warning("Invalid messages format in chat method") logging.warning("Invalid messages format in chat method")
return None, limit 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: try:
# Clean up any orphaned tool messages from previous conversations # Clean up any orphaned tool messages from previous conversations
clean_messages = [] clean_messages = []
@@ -134,6 +152,7 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn):
) )
result = await openai_chat(self.client, **chat_kwargs) result = await openai_chat(self.client, **chat_kwargs)
self._record_usage(result)
# Handle function calls if present # Handle function calls if present
message = result.choices[0].message 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}") logging.debug(f"🔧 Last few messages: {messages[-3:] if len(messages) > 3 else messages}")
final_result = await openai_chat(self.client, **final_chat_kwargs) final_result = await openai_chat(self.client, **final_chat_kwargs)
self._record_usage(final_result)
answer_obj = final_result.choices[0].message answer_obj = final_result.choices[0].message
logging.debug( logging.debug(
+27 -2
View File
@@ -13,7 +13,7 @@ from contextlib import closing
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
SCHEMA_VERSION = 1 SCHEMA_VERSION = 2
class PersistentStore: class PersistentStore:
@@ -27,14 +27,21 @@ class PersistentStore:
return conn return conn
def _init_db(self) -> None: def _init_db(self) -> None:
# Forward-only migrations keyed on user_version (PER-06)
self.db_path.parent.mkdir(parents=True, exist_ok=True) self.db_path.parent.mkdir(parents=True, exist_ok=True)
with closing(self._connect()) as conn, conn: 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( conn.execute(
"CREATE TABLE IF NOT EXISTS history (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL)" "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 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)") 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}") conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
os.chmod(self.db_path, 0o600) # conversation data (PER-04) os.chmod(self.db_path, 0o600) # conversation data (PER-04)
@@ -65,6 +72,24 @@ class PersistentStore:
(channel, content), (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: def migrate_pickles(self, channel: str, history_file: Path, memory_file: Path) -> None:
"""Import legacy pickles once; rename them *.migrated (PER-03).""" """Import legacy pickles once; rename them *.migrated (PER-03)."""
if self.load_history(channel) or self.load_memory(channel) is not None: if self.load_history(channel) or self.load_memory(channel) is not None:
+77
View File
@@ -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"])
+44
View File
@@ -31,3 +31,47 @@ and game descriptions are attacker-influenced input.
The `hack` envelope field remains as an advisory signal (logged, The `hack` envelope field remains as an advisory signal (logged,
staff-notified) but is no longer the defense. 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.
+6
View File
@@ -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 `!bot tasks off` disables bot-initiated posting (today: the boreness
loop; later: the FDB-011 scheduler); `!bot tasks on` restores. loop; later: the FDB-011 scheduler); `!bot tasks on` restores.
Bot-initiated posts also respect pause/quiet. 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.
+8 -1
View File
@@ -28,9 +28,16 @@ accident).
### PER-04 — Database hygiene (coverage: test) ### PER-04 — Database hygiene (coverage: test)
The database runs in WAL journal mode, carries `PRAGMA user_version` 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). 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) ### PER-05 — Writes run off the event loop (coverage: test)
History and memory persistence happen in a worker thread History and memory persistence happen in a worker thread
+3 -3
View File
@@ -10,7 +10,7 @@ from pathlib import Path
from unittest.mock import MagicMock from unittest.mock import MagicMock
from fjerkroa_bot.ai_responder import AIMessage, AIResponder 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_bdd_envelope import FakeModelResponder, envelope
from .test_main import TestBotBase from .test_main import TestBotBase
@@ -78,12 +78,12 @@ class TestPickleMigration(StoreBase):
class TestDatabaseHygiene(StoreBase): class TestDatabaseHygiene(StoreBase):
def test_wal_version_and_permissions(self): 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) PersistentStore(self.db_path)
conn = sqlite3.connect(self.db_path) conn = sqlite3.connect(self.db_path)
try: try:
self.assertEqual(conn.execute("PRAGMA journal_mode").fetchone()[0], "wal") 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: finally:
conn.close() conn.close()
mode = stat.S_IMODE(self.db_path.stat().st_mode) mode = stat.S_IMODE(self.db_path.stat().st_mode)
+201
View File
@@ -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()