safety layer: hard daily budget, user quotas, spend report, forgetme + privacy
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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 <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."
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"])
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user