From 5cd970a8e174eece80bc70c43c81283de003809d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 8 Oct 2026 23:05:04 +0000 Subject: [PATCH] Group admins count as admins; merge /usage into /credits; ADMIN_ONLY locks DMs and Sharesight Admins are users in ADMIN_CHAT_IDS plus administrators of groups the bot has seen (cached five minutes). ADMIN_ONLY (default on) ignores DMs from non-admins and keeps Sharesight tools and data from them, including on the NEEDS_TOOLS rerun. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_014k4kywyN5dAo39izUPvxbw --- CLAUDE.md | 2 +- README.md | 16 ++++++++--- lib/app.py | 8 ++---- lib/ask.py | 14 +++++++-- lib/config.py | 5 +++- lib/context.py | 3 +- lib/telegram/access.py | 38 +++++++++++++++++++++++++ lib/telegram/commands.py | 34 +++++++++++----------- lib/telegram/handlers.py | 19 +++++++++++-- tests/conftest.py | 2 +- tests/test_access.py | 61 ++++++++++++++++++++++++++++++++++++++++ tests/test_app.py | 9 ++++-- tests/test_config.py | 5 ++++ tests/test_features.py | 43 +++++++++++++++++++++------- tests/test_handlers.py | 49 ++++++++++++++++++++++++++++++++ tests/test_post.py | 7 ++--- 16 files changed, 260 insertions(+), 55 deletions(-) create mode 100644 lib/telegram/access.py create mode 100644 tests/test_access.py diff --git a/CLAUDE.md b/CLAUDE.md index 32389d5..fac1272 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -24,7 +24,7 @@ MCP data tools and web search. User-facing docs are in `README.md`; this file is `openrouter.py` + `zai.py` (the only provider-specific code), `runner.py` (THE tool loop), `policy.py` (THE retry ladder), `gate.py` (which tools a message needs). - `prompts.py` all prompts. `ask.py` one entry point: route -> Request -> policy.ask -> usage -> NEEDS_TOOLS rerun. - `telegram/`: `handlers.py` (trigger -> request -> reply), `send.py`, `draft.py` (typing + streaming - drafts), `commands.py` (/credits, /usage). `features/`: `holding_news.py`, `movers.py`, `dm_buttons.py`, `post.py`. + drafts), `commands.py` (/credits), `access.py` (who is an admin). `features/`: `holding_news.py`, `movers.py`, `dm_buttons.py`, `post.py`. - `app.py` wiring/startup; `context.py` `Ctx` (settings, store, backend, registry, bot). - `lib/` holds only what the bot needs to run. Developer helpers live in `tools/`: `report.py` (`make prompt-report`), `capture.py` (real MCP output for fixtures); the bot never imports them. diff --git a/README.md b/README.md index 4b1b21d..d637524 100644 --- a/README.md +++ b/README.md @@ -160,8 +160,15 @@ Each is off by default. The holding-news DM and the movers reply are switched on |---|---|---| | `/start` | DM | Greeting (and the buttons if enabled) | | `/holdingnews` | DM, if holding news is on | Run the holding-news check now | -| `/credits` | a user ID in `ADMIN_CHAT_IDS` | Provider balance (OpenRouter; xAI and z.ai have none to show) | -| `/usage` | a user ID in `ADMIN_CHAT_IDS` | Requests, tokens (and % cached) and cost per model, last 24 h and 7 days | +| `/credits` | an admin (see below), in a group or a DM | Provider balance (OpenRouter; xAI and z.ai have none to show), then requests, tokens (and % cached) and cost per model for the last 24 h and 7 days | + +**Admins** are the user IDs in `ADMIN_CHAT_IDS` plus anyone who administers a group the bot is in (checked with +Telegram, cached for five minutes). The bot only knows a group once someone has written in it, or it was added, since +it started tracking groups. Admins who post anonymously (as the group) can't be recognised. With `ADMIN_ONLY` on (the +default), only admins can DM the bot, and Sharesight data is for admins only: anyone else asking about a portfolio +gets "Portfolio data is for admins only." Holding-news DMs, buttons and down-server messages the bot sends itself are +unaffected. With nobody in `ADMIN_CHAT_IDS` and no known group, nobody can DM the bot or reach Sharesight: set +`ADMIN_ONLY=off` to allow everyone. ## Caching and token saving **Provider prompt caching** (xAI and OpenRouter keep a conversation on the server holding its cached @@ -198,7 +205,7 @@ Each answer's log line shows the cached percentage. symbol, tags and Yahoo URLs dropped), lines cut to 240 characters and the bot's own to 160. - **Smaller prompts:** the system prompt is about half its previous size, the repeated instruction tail lives in the (cached) system prompt, images go at low detail. -- **Measure it:** every request is recorded in the `usage` table (`/usage`), and +- **Measure it:** every request is recorded in the `usage` table (`/credits`), and `make prompt-report` prints where a typical request's tokens go. ## Running several instances @@ -237,7 +244,8 @@ Set `TELEGRAM_BOT_TOKEN` and exactly one of the three LLM keys. ### Advanced and optional | Variable | Meaning | |---|---| -| `ADMIN_CHAT_IDS` | Telegram IDs of the bot's admins (comma/space separated). Your user ID gets a message when an MCP server goes down and may use `/credits` and `/usage`; a group's chat ID gets the down-server messages only. Empty (default) = no alerts and no admin commands. | +| `ADMIN_CHAT_IDS` | Telegram IDs of the bot's admins (comma/space separated). Your user ID gets a message when an MCP server goes down and may use `/credits`; a group's chat ID gets the down-server messages only. Admins of the groups the bot is in count as admins too. | +| `ADMIN_ONLY` | on|off. Only admins may DM the bot or get Sharesight data (default on). | | `DB_PATH` | SQLite file for chat history (default chat_log.db). Instances may share it. | | `HISTORY_LIMIT` | Messages of chat history in the prompt; the window is HISTORY_LIMIT to 1.5x (default 20). | | `MAX_TOKENS` | Reply cap in tokens, reasoning included (default 3000). | diff --git a/lib/app.py b/lib/app.py index 9cb6b3c..d8b6bd5 100644 --- a/lib/app.py +++ b/lib/app.py @@ -123,16 +123,14 @@ async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: return wrapper # Commands first: the first matching handler wins, and the catch-all would swallow them. - if st.admin_chats: - app.add_handler(CommandHandler("credits", command(commands.credits))) - app.add_handler(CommandHandler("usage", command(commands.usage))) + app.add_handler(CommandHandler("credits", command(commands.credits))) app.add_handler(MessageHandler( (filters.ChatType.GROUPS | filters.ChatType.PRIVATE) & (filters.UpdateType.MESSAGE | filters.UpdateType.EDITED_MESSAGE) & ~filters.StatusUpdate.ALL, handlers.on_message)) if st.holding_news: # the Unsubscribe / Undo buttons on holding-news messages app.add_handler(CallbackQueryHandler(command(holding_news.on_button), pattern=r"^hn:")) - if st.post_to_groups: # keeps the list of groups the bot can post in + if st.post_to_groups or st.admin_only: # keeps the list of groups the bot is in app.add_handler(ChatMemberHandler(command(post.on_membership), ChatMemberHandler.MY_CHAT_MEMBER)) app.add_error_handler(on_error) return app @@ -140,7 +138,7 @@ async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: def allowed_updates(st: config.Settings) -> list[str]: return (["message", "edited_message"] + (["callback_query"] if st.holding_news else []) - + (["my_chat_member"] if st.post_to_groups else [])) + + (["my_chat_member"] if st.post_to_groups or st.admin_only else [])) def git_hash() -> str: diff --git a/lib/ask.py b/lib/ask.py index e528b21..9c46368 100644 --- a/lib/ask.py +++ b/lib/ask.py @@ -37,17 +37,25 @@ def forward(text: str) -> None: return forward +def _public(servers: list) -> list: + """The MCP servers anyone may use: not the ones gated on portfolio questions (Sharesight).""" + return [s for s in servers if s.cfg.get("gate") != "portfolio"] + + async def ask(ctx: Ctx, parts: list[dict], route: Route, *, must_search: bool = False, cache_key: str | None = None, on_text=None, kind: str = "chat", - chat_id: int = 0, extra_tools=()) -> Answer: + chat_id: int = 0, extra_tools=(), admin: bool = True) -> Answer: """Ask the configured provider and record the usage. `route` says which tools to offer; must_search forces a search (and offers nothing else). A message the gate gave fewer tools - than exist that gets back NEEDS_TOOLS is asked again with everything on.""" + than exist that gets back NEEDS_TOOLS is asked again with everything on. A non-admin never + gets the portfolio tools, in either request.""" st = ctx.st if must_search: if not st.search: raise RuntimeError("search is disabled, but this request needs it") route = Route([], True) + elif not admin: + route = dataclasses.replace(route, servers=_public(route.servers)) saver = st.token_saver req = Request( system=system_prompt(ctx.first_name, route, ctx.backend.search_what, saver=saver), @@ -65,7 +73,7 @@ async def ask(ctx: Ctx, parts: list[dict], route: Route, *, must_search: bool = await _record(ctx, req.model, kind, chat_id, answer) if (route.simple or route.partial) and wants_tools(answer.text): log.info("Gate offered fewer tools, the model asked for more: asking again with everything on") - full = Route(ctx.registry.up(), st.search) + full = Route(ctx.registry.up() if admin else _public(ctx.registry.up()), st.search) req = dataclasses.replace( req, system=system_prompt(ctx.first_name, full, ctx.backend.search_what, saver=saver), model=st.model, tools=[t for s in full.servers for t in s.tools] + [t for t, _ in extra_tools], diff --git a/lib/config.py b/lib/config.py index 15f8f4e..3368092 100644 --- a/lib/config.py +++ b/lib/config.py @@ -23,7 +23,8 @@ "REASONING": "Reasoning effort: low, medium or high; empty = the model's default (low on z.ai, which cannot turn thinking off).", "SEARCH": "on|off. Web search (and X search on xAI; z.ai's search API on z.ai). Default on.", # Advanced and optional - "ADMIN_CHAT_IDS": "Telegram IDs of the bot's admins (comma/space separated). Your user ID gets MCP-down alerts and may use /credits and /usage; a group's chat ID gets the alerts only. Empty = neither.", + "ADMIN_CHAT_IDS": "Telegram IDs of the bot's admins (comma/space separated). Your user ID gets MCP-down alerts and may use /credits; a group's chat ID gets the alerts only. Admins of the groups the bot is in count as admins too.", + "ADMIN_ONLY": "on|off. Only admins may DM the bot or get Sharesight data (default on).", "DB_PATH": "SQLite file for chat history (default chat_log.db). Instances may share it.", "HISTORY_LIMIT": "Messages of chat history in the prompt; the window is HISTORY_LIMIT to 1.5x (default 20).", "MAX_TOKENS": "Reply cap in tokens, reasoning included (default 3000).", @@ -73,6 +74,7 @@ class Settings: tz: ZoneInfo mcp_config: str admin_chats: frozenset[int] # user IDs (alerts + admin commands) and group chat IDs (alerts only) + admin_only: bool movers_bots: frozenset[str] dm_buttons: bool post_to_groups: bool @@ -166,6 +168,7 @@ def load(env: Mapping[str, str] | None = None) -> Settings: tz=ZoneInfo(env.get("BOT_TZ", "UTC")), mcp_config=env.get("MCP_CONFIG", "mcp_servers.json"), admin_chats=_ids(env.get("ADMIN_CHAT_IDS", "")), + admin_only=_flag(env, "ADMIN_ONLY", default=True), movers_bots=frozenset( u.lower().lstrip("@") for u in env.get("MOVERS_BOTS", "").replace(",", " ").split() ), diff --git a/lib/context.py b/lib/context.py index c196bae..6363a18 100644 --- a/lib/context.py +++ b/lib/context.py @@ -2,7 +2,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from typing import Any @@ -19,6 +19,7 @@ class Ctx: backend: Backend registry: Registry bot: Any = None # telegram.Bot, set at startup + group_admins: dict = field(default_factory=dict) # chat id -> (admin user IDs, expiry); see telegram/access.py @property def first_name(self) -> str: diff --git a/lib/telegram/access.py b/lib/telegram/access.py new file mode 100644 index 0000000..30ab598 --- /dev/null +++ b/lib/telegram/access.py @@ -0,0 +1,38 @@ +"""Who counts as an admin: a user in ADMIN_CHAT_IDS, or an admin of a group the bot has seen.""" + +from __future__ import annotations + +import asyncio +import logging +import time + +from telegram.error import TelegramError + +from ..context import Ctx + +log = logging.getLogger("bot") +CACHE_SECONDS = 300 # a demoted admin loses access within this long + + +async def _group_admins(ctx: Ctx, chat_id: int) -> frozenset[int]: + cached = ctx.group_admins.get(chat_id) + if cached and cached[1] > time.monotonic(): + return cached[0] + try: + admins = frozenset(a.user.id for a in await ctx.bot.get_chat_administrators(chat_id)) + except TelegramError as e: + log.warning("Can't read the admins of chat %s: %s", chat_id, e) + admins = frozenset() + ctx.group_admins[chat_id] = (admins, time.monotonic() + CACHE_SECONDS) + return admins + + +async def is_admin(ctx: Ctx, user_id: int) -> bool: + """True for a user in ADMIN_CHAT_IDS or an administrator of any group the bot knows. Anonymous + admins post as the group itself, so they can't be recognised.""" + if user_id in ctx.st.admin_chats: + return True + for chat_id, _ in await asyncio.to_thread(ctx.store.chats): + if user_id in await _group_admins(ctx, chat_id): + return True + return False diff --git a/lib/telegram/commands.py b/lib/telegram/commands.py index fc17218..f2ac424 100644 --- a/lib/telegram/commands.py +++ b/lib/telegram/commands.py @@ -1,4 +1,4 @@ -"""Owner commands: /credits (provider balance, where there is one) and /usage (token ledger).""" +"""Admin command: /credits (provider balance, where there is one, and the token ledger).""" from __future__ import annotations @@ -10,16 +10,20 @@ from telegram.ext import ContextTypes from ..context import Ctx +from .access import is_admin log = logging.getLogger("bot") async def _gate(ctx: Ctx, update: Update) -> bool: """Log the command message (these handlers take it before on_message would) and say - whether the sender is an admin (a user ID in ADMIN_CHAT_IDS).""" + whether the sender is an admin (see access.is_admin).""" await asyncio.to_thread(ctx.store.save, update.effective_message) user = update.effective_user - return bool(user and user.id in ctx.st.admin_chats) + if user and await is_admin(ctx, user.id): + return True + log.info("Ignored %s from user %s: not an admin", update.effective_message.text, user and user.id) + return False async def _reply(ctx: Ctx, update: Update, text: str) -> None: @@ -27,18 +31,6 @@ async def _reply(ctx: Ctx, update: Update, text: str) -> None: await asyncio.to_thread(ctx.store.save, sent) -async def credits(ctx: Ctx, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - """/credits - admins only: how much provider credit is left (OpenRouter).""" - if not await _gate(ctx, update): - return - try: - text = await ctx.backend.credits() or "This provider has no credit balance to show." - except Exception as e: - log.exception("Credit check failed") - text = f"Couldn't fetch credits: {e}" - await _reply(ctx, update, text) - - def usage_text(rows: list[tuple], label: str) -> str: if not rows: return f"No usage recorded in the {label}." @@ -50,12 +42,18 @@ def usage_text(rows: list[tuple], label: str) -> str: return "\n".join(lines) -async def usage(ctx: Ctx, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - """/usage - admins only: tokens and cost per model for the last 24 hours and 7 days.""" +async def credits(ctx: Ctx, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """/credits - admins only: the provider balance (OpenRouter has one), then tokens and cost per + model for the last 24 hours and 7 days.""" if not await _gate(ctx, update): return + try: + balance = await ctx.backend.credits() or "This provider has no credit balance to show." + except Exception as e: + log.exception("Credit check failed") + balance = f"Couldn't fetch credits: {e}" now = int(time.time()) - parts = [] + parts = [balance] for label, secs in (("last 24 hours", 86_400), ("last 7 days", 7 * 86_400)): rows = await asyncio.to_thread(ctx.store.usage_summary, now - secs) parts.append(usage_text(rows, label)) diff --git a/lib/telegram/handlers.py b/lib/telegram/handlers.py index cf5dc47..d7d8009 100644 --- a/lib/telegram/handlers.py +++ b/lib/telegram/handlers.py @@ -23,10 +23,12 @@ from ..store import is_dm from ..textfmt import md_to_html, strip_disclaimer from ..tickers import link_tickers +from .access import is_admin from .draft import Draft, start_typing from .send import reply_chunks, reply_html log = logging.getLogger("bot") +PORTFOLIO_ADMINS_ONLY = "Portfolio data is for admins only." READ_IMAGE_WORDS = re.compile(r"\b(read|text|says?|chart|table|number|screenshot|ocr|zoom|detail|label)\b", re.I) @@ -77,7 +79,7 @@ def _record(self, msg: Message) -> None: self.ctx.store.remember_user(msg) self.ctx.store.save(msg) chat = getattr(msg, "chat", None) - if self.ctx.st.post_to_groups and chat and chat.type in post.GROUP_TYPES and chat.title: + if chat and chat.type in post.GROUP_TYPES and chat.title: self.ctx.store.remember_chat(chat.id, chat.title) async def on_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: @@ -90,10 +92,16 @@ async def on_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) - trig = self._trigger(msg) if trig is None: return + if trig.private and self.ctx.st.admin_only and not await self._is_admin(msg): + log.info("Ignored a DM from user %s: not an admin", msg.from_user and msg.from_user.id) + return if await self._command(msg, trig): return await self._answer(msg, trig) + async def _is_admin(self, msg: Message) -> bool: + return bool(msg.from_user and await is_admin(self.ctx, msg.from_user.id)) + def _trigger(self, msg: Message) -> Trigger | None: """Is this message for the bot? Other bots are logged but never answered (avoids bot loops), except movers lists, which are answered once.""" @@ -111,7 +119,7 @@ def _trigger(self, msg: Message) -> Trigger | None: return Trigger(text, private, is_movers, reply_target, replied_to_bot) async def _command(self, msg: Message, trig: Trigger) -> bool: - """Private-chat commands handled here (/credits and /usage have their own handlers).""" + """Private-chat commands handled here (/credits has its own handler).""" if not trig.private or not trig.text.startswith("/"): return False command = trig.text.split()[0].split("@")[0].lower() @@ -216,6 +224,10 @@ async def _generate(self, msg: Message, trig: Trigger, draft: Draft | None) -> s else: context = trig.text + " " + (describe(trig.reply_target) if trig.reply_target else "") chosen = route(context, st, ctx.registry) + admin = not st.admin_only or await self._is_admin(msg) + if not admin and is_portfolio_question(context, st) and any( + s.cfg.get("gate") == "portfolio" for s in ctx.registry.up()): + return PORTFOLIO_ADMINS_ONLY extra = [] if st.post_to_groups and trig.private and post.wants_post(trig.text): extra = [post.tool(ctx, msg.from_user.id, msg.from_user.full_name or "")] @@ -225,7 +237,8 @@ async def _generate(self, msg: Message, trig: Trigger, draft: Draft | None) -> s parts = [{"type": "text", "text": await self._chat_text(msg, trig, fresh=fresh)}] parts += await self._images(msg, trig) answer = await ask(ctx, parts, chosen, - cache_key=f"chat-{chat_id}", on_text=on_text, chat_id=chat_id, extra_tools=extra) + cache_key=f"chat-{chat_id}", on_text=on_text, chat_id=chat_id, extra_tools=extra, + admin=admin) return format_answer(answer.text, movers_bots=st.movers_bots, summary=trig.movers) diff --git a/tests/conftest.py b/tests/conftest.py index d0df5b2..573ea0a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,7 @@ from lib import config from lib.store import Store -BASE_ENV = {"TELEGRAM_BOT_TOKEN": "1:x", "OPENROUTER_API_KEY": "key"} +BASE_ENV = {"TELEGRAM_BOT_TOKEN": "1:x", "OPENROUTER_API_KEY": "key", "ADMIN_ONLY": "off"} # admin-only has its own tests @pytest.fixture diff --git a/tests/test_access.py b/tests/test_access.py new file mode 100644 index 0000000..99b11b2 --- /dev/null +++ b/tests/test_access.py @@ -0,0 +1,61 @@ +import time + +from telegram.error import TelegramError + +from lib import config +from lib.context import Ctx +from lib.telegram import access + +from .conftest import FakeBot +from .fakes import ScriptedBackend, registry + + +def make_ctx(env, store, **extra): + return Ctx(config.load({**env, **extra}), store, ScriptedBackend([]), registry(), FakeBot()) + + +async def test_listed_user_is_admin_without_any_lookup(env, store): + ctx = make_ctx(env, store, ADMIN_CHAT_IDS="5, -100") + assert await access.is_admin(ctx, 5) and not await access.is_admin(ctx, 6) + + +async def test_admin_of_a_known_group_is_admin(env, store): + ctx = make_ctx(env, store) + store.remember_chat(-1, "One") + store.remember_chat(-2, "Two") + ctx.bot.admins = {-1: [3], -2: [4]} + assert await access.is_admin(ctx, 4) and not await access.is_admin(ctx, 9) + + +async def test_unknown_groups_and_unreadable_groups_grant_nothing(env, store): + ctx = make_ctx(env, store) + ctx.bot.admins = {-1: [3]} # the bot has never seen group -1 + assert not await access.is_admin(ctx, 3) + store.remember_chat(-2, "Gone") # known, but the bot was removed (Telegram raises) + assert not await access.is_admin(ctx, 3) + + +async def test_admin_lists_are_cached_then_refreshed(env, store, monkeypatch): + ctx = make_ctx(env, store) + store.remember_chat(-1, "One") + ctx.bot.admins = {-1: [3]} + assert await access.is_admin(ctx, 3) + ctx.bot.admins = {-1: []} # demoted + assert await access.is_admin(ctx, 3) # still cached + now = time.monotonic() + monkeypatch.setattr(access.time, "monotonic", lambda: now + access.CACHE_SECONDS + 1) + assert not await access.is_admin(ctx, 3) + + +async def test_telegram_error_is_not_admin_and_is_cached(env, store): + ctx = make_ctx(env, store) + store.remember_chat(-1, "One") + calls = [] + + async def failing(chat_id): + calls.append(chat_id) + raise TelegramError("kicked") + + ctx.bot.get_chat_administrators = failing + assert not await access.is_admin(ctx, 3) and not await access.is_admin(ctx, 4) + assert calls == [-1] diff --git a/tests/test_app.py b/tests/test_app.py index df766fe..1152f29 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -30,13 +30,15 @@ def build(**extra): return app.build_app(Ctx(st, store, ScriptedBackend([]), registry())), st plain, st = build() - assert handler_kinds(plain) == ["MessageHandler"] and app.allowed_updates(st) == ["message", "edited_message"] + assert handler_kinds(plain) == ["CommandHandler", "MessageHandler"] and app.allowed_updates(st) == ["message", "edited_message"] full, st = build(ADMIN_CHAT_IDS="5", SHARESIGHT_HOLDING_NEWS_RECIPIENTS="Pf:alice") kinds = handler_kinds(full) - assert kinds == ["CommandHandler", "CommandHandler", "MessageHandler", "CallbackQueryHandler"] + assert kinds == ["CommandHandler", "MessageHandler", "CallbackQueryHandler"] assert "callback_query" in app.allowed_updates(st) - assert {c for h in full.handlers[0] if isinstance(h, CommandHandler) for c in h.commands} == {"credits", "usage"} + assert {c for h in full.handlers[0] if isinstance(h, CommandHandler) for c in h.commands} == {"credits"} assert any(isinstance(h, MessageHandler) for h in full.handlers[0]) + locked, st = build(ADMIN_ONLY="on") # remembers the groups it is in, to know who their admins are + assert "ChatMemberHandler" in handler_kinds(locked) and "my_chat_member" in app.allowed_updates(st) assert any(isinstance(h, CallbackQueryHandler) for h in full.handlers[0]) @@ -57,6 +59,7 @@ def test_main_accepts_and_ignores_a_label_and_unknown_args(monkeypatch, tmp_path monkeypatch.setenv("TELEGRAM_BOT_TOKEN", "1:x") monkeypatch.setenv("OPENROUTER_API_KEY", "b") monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setenv("ADMIN_ONLY", "off") monkeypatch.setenv("DB_PATH", str(tmp_path / "x.db")) monkeypatch.setenv("MCP_CONFIG", str(tmp_path / "none.json")) monkeypatch.setattr(app.Application, "run_polling", lambda self, **kw: started.update(kw)) diff --git a/tests/test_config.py b/tests/test_config.py index fd08a76..2f975d4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -65,6 +65,11 @@ def test_admin_chats_and_recipients(env): assert s.admin_chats == {-1, -2, 3} and s.holding_news_recipients == {"bob": "me"} +def test_admin_only_is_on_unless_turned_off(env): + assert not config.load(env).admin_only + assert config.load({k: v for k, v in env.items() if k != "ADMIN_ONLY"}).admin_only + + def test_no_names_in_code_defaults(env): s = config.load(env) assert s.holding_news_recipients == {} and s.movers_bots == frozenset() and s.portfolio_names == frozenset() diff --git a/tests/test_features.py b/tests/test_features.py index eb91a04..7d3f2e9 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -165,22 +165,45 @@ def test_movers_helpers(): # ---- commands -------------------------------------------------------------------------- -async def test_credits_and_usage_are_admin_only_and_logged(env, store): +async def test_credits_shows_balance_and_usage_to_admins_only(env, store): st = config.load({**env, "ADMIN_CHAT_IDS": "5, -100"}) ctx = Ctx(st, store, ScriptedBackend([]), registry(), FakeBot()) store.add_usage(-1, "m", "chat", 1, 1000, 500, 20, 0.5) - stranger = user_msg(ctx.bot, "/usage", user_id=9, mid=1) - await commands.usage(ctx, update_for(stranger), None) - assert stranger.replies == [] and store.history(-100, 20)[0].text == "/usage" - owner = user_msg(ctx.bot, "/usage", user_id=5, mid=2) - await commands.usage(ctx, update_for(owner), None) - assert "m: 1 requests, in 1,000 (50% cached)" in owner.replies[0]["text"] - owner2 = user_msg(ctx.bot, "/credits", user_id=5, mid=3) - await commands.credits(ctx, update_for(owner2), None) - assert "no credit balance" in owner2.replies[0]["text"] + stranger = user_msg(ctx.bot, "/credits", user_id=9, mid=1) + await commands.credits(ctx, update_for(stranger), None) + assert stranger.replies == [] and store.history(-100, 20)[0].text == "/credits" + owner = user_msg(ctx.bot, "/credits", user_id=5, mid=2) + await commands.credits(ctx, update_for(owner), None) + text = owner.replies[0]["text"] + assert "no credit balance" in text and "m: 1 requests, in 1,000 (50% cached)" in text assert commands.usage_text([], "last 24 hours") == "No usage recorded in the last 24 hours." +async def test_credits_still_shows_usage_when_the_balance_fails(env, store): + class Broken(ScriptedBackend): + async def credits(self): + raise RuntimeError("boom") + + ctx = Ctx(config.load({**env, "ADMIN_CHAT_IDS": "5"}), store, Broken([]), registry(), FakeBot()) + store.add_usage(-1, "m", "chat", 1, 10, 0, 2, 0.1) + msg = user_msg(ctx.bot, "/credits", user_id=5) + await commands.credits(ctx, update_for(msg), None) + assert "Couldn't fetch credits: boom" in msg.replies[0]["text"] and "m: 1 requests" in msg.replies[0]["text"] + + +async def test_group_admin_can_run_credits_in_the_group_and_in_a_dm(env, store): + ctx = Ctx(config.load(env), store, ScriptedBackend([]), registry(), FakeBot()) + store.remember_chat(-100, "Stocks") + ctx.bot.admins[-100] = [7] + for chat_id, mid in ((-100, 1), (7, 2)): + msg = user_msg(ctx.bot, "/credits", user_id=7, chat_id=chat_id, mid=mid) + await commands.credits(ctx, update_for(msg), None) + assert "Usage" in msg.replies[0]["text"] or "No usage" in msg.replies[0]["text"] + other = user_msg(ctx.bot, "/credits", user_id=8, mid=3) + await commands.credits(ctx, update_for(other), None) + assert other.replies == [] + + @pytest.mark.parametrize("cached,tin,pct", [(0, 0, 0), (50, 100, 50)]) def test_usage_text_percentages(cached, tin, pct): assert f"({pct}% cached)" in commands.usage_text([("m", 1, tin, cached, 1, 0.0)], "x") diff --git a/tests/test_handlers.py b/tests/test_handlers.py index ce0f71a..f3ae829 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -217,3 +217,52 @@ def test_dm_buttons_hash_changes_with_presets(monkeypatch): before = dm_buttons.keyboard_hash() monkeypatch.setitem(dm_buttons.PRESETS, "New", "x") assert dm_buttons.keyboard_hash() != before + + +# ---- admin-only (the default) --------------------------------------------------------- +async def test_dm_from_a_non_admin_is_logged_but_ignored(env, store): + ctx, backend, h = make_ctx(env, store, [step("hi")], ADMIN_ONLY="on", ADMIN_CHAT_IDS="5") + stranger = user_msg(ctx.bot, "hello", chat_id=9, user_id=9) + await run(h, stranger) + assert stranger.replies == [] and backend.seen == [] and store.history(9, 20) + admin = user_msg(ctx.bot, "hello", chat_id=5, user_id=5) + await run(h, admin) + assert admin.replies[0]["text"] == "hi" + + +async def test_group_admin_can_dm_but_admin_only_off_lets_everyone_in(env, store): + ctx, backend, h = make_ctx(env, store, [step("hi"), step("hi")], ADMIN_ONLY="on") + store.remember_chat(-100, "Stocks") + ctx.bot.admins[-100] = [7] + m = user_msg(ctx.bot, "hello", chat_id=7, user_id=7) + await run(h, m) + assert m.replies[0]["text"] == "hi" + ctx, backend, h = make_ctx(env, store, [step("hi")], ADMIN_ONLY="off") + open_dm = user_msg(ctx.bot, "hello", chat_id=9, user_id=9, mid=5) + await run(h, open_dm) + assert open_dm.replies[0]["text"] == "hi" + + +async def test_non_admin_gets_no_portfolio_data_in_a_group(env, store): + sharesight = FakeMcp("sharesight", gate="portfolio", tools=("list_portfolios",)) + ctx, backend, h = make_ctx(env, store, [step("ok")], servers=[sharesight], ADMIN_ONLY="on", + ADMIN_CHAT_IDS="5") + asked = user_msg(ctx.bot, "@stockbot how is my portfolio doing?", user_id=9) + await run(h, asked) + assert asked.replies[0]["text"] == "Portfolio data is for admins only." and backend.seen == [] + admin = user_msg(ctx.bot, "@stockbot how is my portfolio doing?", user_id=5, mid=2) + await run(h, admin) + assert backend.seen[0]["tools"] == ["sharesight__list_portfolios"] + + +async def test_non_admin_never_gets_portfolio_tools_even_on_the_rerun(env, store): + yahoo, sharesight = FakeMcp("yahoo"), FakeMcp("sharesight", gate="portfolio", tools=("list_portfolios",)) + ctx, backend, h = make_ctx(env, store, [step("ok")], servers=[yahoo, sharesight], ADMIN_ONLY="on") + m = user_msg(ctx.bot, "@stockbot how is NVDA looking?", user_id=9) + await run(h, m) + assert backend.seen[0]["tools"] == ["yahoo__get_quote"] + ctx, backend, h = make_ctx(env, store, [step("NEEDS_TOOLS"), step("done")], servers=[yahoo, sharesight], + ADMIN_ONLY="on") + m = user_msg(ctx.bot, "@stockbot how is NVDA looking?", user_id=9) + await run(h, m) + assert backend.seen[1]["tools"] == ["yahoo__get_quote"] and m.replies[0]["text"] == "done" diff --git a/tests/test_post.py b/tests/test_post.py index 66e85e0..27f13d2 100644 --- a/tests/test_post.py +++ b/tests/test_post.py @@ -113,14 +113,11 @@ async def test_flag_off_or_a_group_chat_or_a_normal_message_gets_no_post_tool(en assert all("send_to_group" not in s["tools"] for s in backend.seen) -async def test_groups_are_remembered_from_messages_only_when_the_flag_is_on(env, store): - ctx, backend, h = make_ctx(env, store, [], POST_TO_GROUPS_FROM_DM="on") +async def test_groups_are_remembered_from_their_messages(env, store): + ctx, backend, h = make_ctx(env, store, []) await run(h, user_msg(ctx.bot, "hello all", chat_id=GROUP, chat_title="Finance Alliance")) await run(h, user_msg(ctx.bot, "hi", chat_id=ME, user_id=ME, mid=2)) # a DM is not a group assert store.chats() == [(GROUP, "Finance Alliance")] - ctx2, _, h2 = make_ctx(env, store, []) - await run(h2, user_msg(ctx2.bot, "hello", chat_id=OTHER, chat_title="Chess", mid=3)) - assert store.chats() == [(GROUP, "Finance Alliance")] async def test_membership_updates_add_and_remove_groups(env, store):