diff --git a/FEATURE_SUGGESTIONS/FEATURE_MCP_CLAUDE_INTEGRATION.md b/FEATURE_SUGGESTIONS/FEATURE_MCP_CLAUDE_INTEGRATION.md index e3a4ec419..9817635c2 100644 --- a/FEATURE_SUGGESTIONS/FEATURE_MCP_CLAUDE_INTEGRATION.md +++ b/FEATURE_SUGGESTIONS/FEATURE_MCP_CLAUDE_INTEGRATION.md @@ -1,9 +1,16 @@ # חיבור Claude.ai ל‑CodeKeeper דרך MCP — מסמך תכנון -> **סטטוס:** תכנון (Draft) — לפני כתיבת קוד. ממתין לאישור כיוון. +> **סטטוס:** פאזה 0 (MVP קריאה‑בלבד + PAT) **מומשה** — ראו `mcp_server/`. פאזות 1–3 עדיין בתכנון. > **ענף פיתוח:** `claude/mcp-codekeeper-webapp-ldnzsg` > **מתי להשתמש:** לפני מימוש חיבור MCP; מסמך זה הוא מקור האמת לתכנון. -> **ראו גם:** [CodeBot – Project Docs](https://amirbiron.github.io/CodeBot/), `CLAUDE.md` (מדיניות מחייבת). +> **ראו גם:** `mcp_server/README.md` (שימוש), [CodeBot – Project Docs](https://amirbiron.github.io/CodeBot/), `CLAUDE.md` (מדיניות מחייבת). + +> **מצב מימוש (פאזה 0):** נכתבה חבילת `mcp_server/` (שרת FastMCP קריאה‑בלבד עם 7 כלים +> ל‑`code_snippets` + `collections`), אימות PAT מעל קולקשן `mcp_tokens`, סקריפט +> `scripts/mcp_issue_token.py`, וטסטים `tests/test_mcp_*.py`. **הערת סטייה מהמסמך:** +> מאגר הטוקנים נמצא ב‑`mcp_server/token_store.py` (ולא `database/mcp_tokens.py`), כדי +> שהמודולים יהיו נטולי תלויות כבדות וניתנים לבדיקה בבידוד. פקודת הבוט `/connect_claude` +> כבר זמינה (הנפקת טוקן מתוך טלגרם); OAuth (Claude.ai) וכלי כתיבה יגיעו בפאזות הבאות. --- @@ -88,7 +95,7 @@ db.delete_file(user_id, file_name) # מחיקה רכה (recycle ## 4. ארכיטקטורה מוצעת -``` +```text ┌────────────┐ OAuth 2.1 (PKCE) ┌──────────────────────────┐ │ Claude.ai │ ───────────────────► │ שרת MCP חדש (שירות נפרד) │ │ (Connector)│ Streamable HTTP │ Python + ASGI (uvicorn) │ @@ -117,6 +124,7 @@ db.delete_file(user_id, file_name) # מחיקה רכה (recycle ## 5. הכלים והמשאבים (MCP surface) ### 5.1 Tools + | כלי | קלט | פלט | נשען על | |-----|-----|-----|---------| | `list_files` | `limit`, `page`, `language?`, `tag?` | רשימת מטא‑דאטה (בלי `code`) | `get_user_files` | diff --git a/main.py b/main.py index b8fde56f2..6e7395a22 100644 --- a/main.py +++ b/main.py @@ -777,6 +777,72 @@ async def _send_direct_admins(context: ContextTypes.DEFAULT_TYPE, text: str) -> return False +async def connect_claude_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + """מנפיק טוקן אישי (PAT) לחיבור הקבצים של המשתמש ל‑Claude דרך MCP (קריאה בלבד).""" + try: + message = update.message or update.effective_message + if message is None: + return + + # אבטחה: הטוקן סודי — מנפיקים רק בצ'אט פרטי כדי שלא ידלוף בקבוצה. + chat = update.effective_chat + if getattr(chat, "type", "private") != "private": + await message.reply_text("🔒 הפקודה זמינה בצ'אט פרטי בלבד (הטוקן סודי).") + return + + user = update.effective_user + user_id = getattr(user, "id", None) + if not user_id: + await message.reply_text("לא זוהה משתמש.") + return + + raw = None + try: + from src.infrastructure.composition.webapp_container import get_files_facade + + raw = get_files_facade().issue_mcp_token(int(user_id), label="Claude") + except Exception: + logger.error("connect_claude: token issue failed", exc_info=True) + raw = None + + if not raw: + await message.reply_text("אירעה שגיאה ביצירת הטוקן. נסו שוב מאוחר יותר.") + return + + base = (os.getenv("MCP_SERVER_URL") or "https://YOUR-MCP-HOST").rstrip("/") + add_cmd = ( + f'claude mcp add --transport http codekeeper {base}/mcp ' + f'--header "Authorization: Bearer {raw}"' + ) + text = ( + "🔌 חיבור הקבצים שלך ל‑Claude (MCP)\n\n" + "הטוקן האישי שלך (יוצג פעם אחת בלבד — שמור אותו):\n" + f"{raw}\n\n" + "לחיבור מ‑Claude Code (העתק‑הדבק):\n" + f"{add_cmd}\n\n" + "⚠️ אל תשתפו את הטוקן. כרגע החיבור עובד מול Claude Code / Desktop " + "(קריאה בלבד); תמיכה ב‑Claude.ai בוובאפ תגיע בהמשך." + ) + try: + await message.reply_text(text, parse_mode=ParseMode.HTML) + except Exception: + # נפילת פרסום HTML לא תגרום לאובדן הטוקן — שולחים גרסת טקסט. + plain = ( + text.replace("", "") + .replace("", "") + .replace("", "") + .replace("", "") + ) + await message.reply_text(plain) + except Exception: + logger.error("connect_claude_command failed", exc_info=True) + try: + if update and update.message: + await update.message.reply_text("אירעה שגיאה. נסו שוב מאוחר יותר.") + except Exception: + pass + + async def admin_report_command(update: Update, context: ContextTypes.DEFAULT_TYPE): """קבלת דיווח משתמש ושליחה לאדמינים.""" try: @@ -2109,6 +2175,12 @@ class HelpSection(TypedDict): HELP_SECTIONS: list[HelpSection] = [ + { + "title": "🔌 חיבור ל‑Claude (MCP)", + "entries": [ + {"commands": ("connect_claude",), "description": "חיבור הקבצים שלך ל‑Claude (טוקן MCP)"}, + ], + }, { "title": "🔔 תזכורות", "entries": [ @@ -3823,6 +3895,7 @@ async def _global_callback_guard(update: Update, context: ContextTypes.DEFAULT_T self.application.add_handler(CommandHandler("stats", self.stats_command)) self.application.add_handler(CommandHandler("check", self.check_commands)) self.application.add_handler(CommandHandler("admin", admin_report_command)) + self.application.add_handler(CommandHandler("connect_claude", connect_claude_command)) # ChatOps: /jobs (Background Jobs Monitor) async def jobs_command(update: Update, context: ContextTypes.DEFAULT_TYPE): diff --git a/mcp_server/README.md b/mcp_server/README.md new file mode 100644 index 000000000..24c076dc2 --- /dev/null +++ b/mcp_server/README.md @@ -0,0 +1,107 @@ +# CodeKeeper MCP Server (פאזה 0 — קריאה בלבד) + +שרת [MCP](https://modelcontextprotocol.io) קטן שחושף את **הקבצים והאוספים** +השמורים של המשתמש ל‑Claude (Claude Code / Claude Desktop כרגע; Claude.ai דרך +OAuth בפאזה הבאה). קריאה בלבד — אין כלי כתיבה/מחיקה. + +> תכנון מלא: `FEATURE_SUGGESTIONS/FEATURE_MCP_CLAUDE_INTEGRATION.md` + +--- + +## מה זה עושה + +ניגש ישירות לשכבת ה‑DB הקיימת (`database.db` + `CollectionsManager`), מסונן +תמיד לפי ה‑`user_id` שנגזר מהטוקן. מכבד את חוק ה‑Smart Projection: רשימות/חיפוש +לא מחזירים את שדה ה‑`code` הכבד — תוכן מלא רק ב‑`get_file`. + +### הכלים (Tools) + +| כלי | תיאור | +|-----|-------| +| `list_files` | רשימת קבצים (מטא‑דאטה בלבד), עם עימוד | +| `search_code` | חיפוש טקסט בקוד → מטא‑דאטה של קבצים תואמים | +| `get_file` | תוכן מלא של קובץ לפי `file_name` או `file_id` (אופציונלי: גרסה) | +| `list_versions` | היסטוריית גרסאות של קובץ (מטא‑דאטה) | +| `list_collections` | האוספים של המשתמש | +| `get_collection` | אוסף בודד לפי id | +| `get_collection_items` | הקבצים בתוך אוסף (עם עימוד/סינון תיקייה) | + +--- + +## אימות (Personal Access Token) + +הזדהות היא `Authorization: Bearer `. הטוקן נשמר כ‑**hash בלבד** בקולקשן +`mcp_tokens`, קשור ל‑`user_id`, וניתן לביטול. ה‑`user_id` נגזר תמיד מהטוקן — +לעולם לא מקלט הלקוח. + +### הנפקת טוקן + +**הדרך הפשוטה — מתוך הבוט:** שלחו `/connect_claude` בצ'אט פרטי עם הבוט. תקבלו טוקן +מוכן + פקודת חיבור מוכנה להעתקה. (הפקודה זמינה בצ'אט פרטי בלבד כדי שהטוקן לא ידלוף.) + +**לאופס/בדיקות (CLI):** +```bash +MONGODB_URL="..." python scripts/mcp_issue_token.py --user-id --label "Claude Desktop" +``` +הטוקן מוצג **פעם אחת בלבד** — שמרו אותו. + +--- + +## הרצה מקומית +```bash +pip install -r requirements/development.txt +MONGODB_URL="..." uvicorn mcp_server.app:app --host 0.0.0.0 --port 8000 +``` +- נקודת הקצה של MCP: `POST/GET /mcp` (Streamable HTTP). +- בריאות (ללא אימות): `GET /healthz`. + +--- + +## חיבור מ‑Claude Code +```bash +claude mcp add --transport http codekeeper http://localhost:8000/mcp \ + --header "Authorization: Bearer " +``` + +## חיבור מ‑Claude Desktop (`claude_desktop_config.json`) +```json +{ + "mcpServers": { + "codekeeper": { + "type": "http", + "url": "https:///mcp", + "headers": { "Authorization": "Bearer " } + } + } +} +``` + +> **Claude.ai (וובאפ):** דורש OAuth 2.1 — זו פאזה 1, עדיין לא כאן. PAT עובד היום +> מול Claude Code / Claude Desktop. + +--- + +## פריסה (Render) +שירות web נפרד (ASGI), משתף את אותו `MONGODB_URL` ואת אותם סודות: +```text +Start command: uvicorn mcp_server.app:app --host 0.0.0.0 --port $PORT +Health check: /healthz +``` +ENV: `MONGODB_URL` (משותף), ו‑`MCP_SERVER_URL` (ה‑URL הציבורי של השירות — משמש את +`/connect_claude` בבוט כדי לבנות את פקודת החיבור). מומלץ שירות נפרד (ולא בתוך הוובאפ) +כי ה‑worker של הוובאפ יחיד ורגיש לזיכרון. + +--- + +## מבנה הקוד + +| קובץ | תפקיד | +|------|-------| +| `token_store.py` | ניהול PAT (הנפקה/אימות/ביטול) מעל `mcp_tokens` | +| `backend.py` | גישה לנתונים + סריאליזציה (Smart Projection, בדיקת בעלות) | +| `handlers.py` | לוגיקת הכלים הטהורה (ולידציה/clamping) — יעד הטסטים | +| `auth.py` | middleware של Bearer + `current_user_id(ctx)` | +| `server.py` | חיווט FastMCP: כלים + אפליקציית ASGI מאומתת | +| `app.py` | נקודת כניסה: `app = create_app()` | + +טסטים: `tests/test_mcp_*.py` (fakes ידניים, בלי MongoDB אמיתי). diff --git a/mcp_server/__init__.py b/mcp_server/__init__.py new file mode 100644 index 000000000..b8c45a3ff --- /dev/null +++ b/mcp_server/__init__.py @@ -0,0 +1,16 @@ +"""CodeKeeper MCP server package. + +A small, self-contained Model Context Protocol (MCP) server that exposes a +user's saved code files and collections to MCP clients (Claude Code / Claude +Desktop today; Claude.ai via OAuth in a later phase). + +Design goals for this package: +- Read-only for now (Phase 0). No write/delete tools. +- Every module here imports only light dependencies at module top-level, so the + pieces can be unit-tested in isolation without pulling in the whole app. + The heavy ``database`` package is imported lazily inside ``backend`` / + ``app`` (only when actually serving), never at import time. +- Authentication is a Personal Access Token (Bearer) verified against the + ``mcp_tokens`` collection. The authenticated ``user_id`` is derived from the + token only, never from client input. +""" diff --git a/mcp_server/app.py b/mcp_server/app.py new file mode 100644 index 000000000..641e6f59c --- /dev/null +++ b/mcp_server/app.py @@ -0,0 +1,38 @@ +"""ASGI entrypoint for the CodeKeeper MCP server. + +Run with:: + + uvicorn mcp_server.app:app --host 0.0.0.0 --port $PORT + +Importing this module connects to MongoDB (via the shared ``database`` layer), +so it is intentionally NOT imported by the unit tests — those exercise the +lighter ``handlers`` / ``backend`` / ``token_store`` modules directly. +""" + +from __future__ import annotations + +import os + +from .backend import ProductionBackend +from .server import build_app +from .token_store import MCPTokenStore +from .wiring import resolve_mongo + + +def create_app(): + from database import db as db_manager # lazy heavy import + + mongo = resolve_mongo(db_manager) + if mongo is None: + raise RuntimeError( + "MongoDB is not available (database.db is None). " + "Set MONGODB_URL for the MCP service." + ) + + token_store = MCPTokenStore(mongo) + backend = ProductionBackend(db_manager=db_manager, mongo_db=mongo) + name = os.getenv("MCP_SERVER_NAME", "CodeKeeper") + return build_app(backend, token_store, name=name) + + +app = create_app() diff --git a/mcp_server/auth.py b/mcp_server/auth.py new file mode 100644 index 000000000..9cedb01c5 --- /dev/null +++ b/mcp_server/auth.py @@ -0,0 +1,85 @@ +"""PAT (Bearer) authentication for the MCP HTTP app. + +A Starlette middleware verifies ``Authorization: Bearer `` against the +token store and injects the authenticated ``user_id`` onto ``request.state``. +Tools then read it via :func:`current_user_id` (same request object — no +contextvar propagation to worry about). + +This is deliberately simple (Phase 0). Phase 1 replaces it with OAuth 2.1. +""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable +from typing import Any + +import anyio +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse + +logger = logging.getLogger(__name__) + +# Paths that must remain reachable without a token (health checks, etc.). +EXEMPT_PATHS = {"/healthz", "/health", "/"} + +_WWW_AUTH = 'Bearer realm="CodeKeeper MCP"' + + +def _unauthorized(code: str) -> JSONResponse: + return JSONResponse( + {"error": code}, + status_code=401, + headers={"WWW-Authenticate": _WWW_AUTH}, + ) + + +class PATAuthMiddleware(BaseHTTPMiddleware): + """Authenticate every non-exempt request with a Personal Access Token.""" + + def __init__( + self, app: Any, token_store: Any, *, exempt_paths: Iterable[str] = EXEMPT_PATHS + ) -> None: + super().__init__(app) + self._store = token_store + self._exempt = set(exempt_paths) + + async def dispatch(self, request: Request, call_next): + if request.url.path in self._exempt: + return await call_next(request) + + auth = request.headers.get("authorization", "") + if not auth.lower().startswith("bearer "): + return _unauthorized("missing_bearer_token") + + token = auth.split(" ", 1)[1].strip() + try: + # verify() is a sync DB call — keep it off the event loop. + principal = await anyio.to_thread.run_sync(self._store.verify, token) + except Exception: + # Log the failure (never the token) so DB/connectivity issues are + # diagnosable in production, then fail closed. + logger.warning("MCP token verification raised an error", exc_info=True) + principal = None + + if not principal: + return _unauthorized("invalid_token") + + request.state.user_id = int(principal["user_id"]) + request.state.scopes = list(principal.get("scopes") or []) + return await call_next(request) + + +def current_user_id(ctx: Any) -> int: + """Return the authenticated user's id from an MCP tool ``Context``. + + Raises ``PermissionError`` if the request was not authenticated (which + should be impossible once :class:`PATAuthMiddleware` is installed, but we + fail closed rather than leak data). + """ + request = getattr(getattr(ctx, "request_context", None), "request", None) + user_id = getattr(getattr(request, "state", None), "user_id", None) + if user_id is None: + raise PermissionError("unauthenticated") + return int(user_id) diff --git a/mcp_server/backend.py b/mcp_server/backend.py new file mode 100644 index 000000000..06dc52ff0 --- /dev/null +++ b/mcp_server/backend.py @@ -0,0 +1,173 @@ +"""Data-access layer for the MCP tools. + +Tools/handlers depend on a duck-typed "backend" (any object exposing the read +methods below), so they can be unit-tested with a fake. ``ProductionBackend`` +wraps the real in-process database layer (``database.db`` + +``CollectionsManager``) and is imported lazily so this module stays light. + +All read paths are ``user_id``-scoped. The one method that touches a +non-user-scoped DB call (``get_file_by_id``) re-checks ownership here. + +Per the project's "Smart Projection" rule, list/search results never carry the +heavy ``code``/``content`` fields — full content is returned only by +``get_file`` for an explicit single-file fetch. +""" + +from __future__ import annotations + +import datetime as _dt +from typing import Any + +_HEAVY_FIELDS = ("code", "content", "raw_data", "raw_content") + + +def _json_safe(value: Any) -> Any: + """Recursively convert Mongo/BSON types to JSON-friendly values.""" + if isinstance(value, (_dt.datetime, _dt.date)): + try: + return value.isoformat() + except Exception: + return str(value) + if type(value).__name__ == "ObjectId": # avoid importing bson + return str(value) + if isinstance(value, dict): + return {str(k): _json_safe(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(v) for v in value] + return value + + +def _clean(doc: dict[str, Any], *, include_code: bool = False) -> dict[str, Any]: + """Serialize a file document. Drops heavy fields unless ``include_code``.""" + out: dict[str, Any] = {} + for key, val in (doc or {}).items(): + if key == "_id": + out["id"] = str(val) + continue + if not include_code and key in _HEAVY_FIELDS: + continue + out[key] = _json_safe(val) + # Friendlier alias without dropping the original field. + if "programming_language" in out: + out.setdefault("language", out["programming_language"]) + return out + + +def _full(doc: dict[str, Any]) -> dict[str, Any]: + """Serialize a single file WITH content (regular ``code`` or large ``content``).""" + out = _clean(doc, include_code=True) + if not out.get("code") and out.get("content"): + out["code"] = out["content"] + return out + + +def _strip_heavy(value: Any) -> Any: + """Recursively drop heavy content fields from an already-serialized value.""" + if isinstance(value, dict): + return {k: _strip_heavy(v) for k, v in value.items() if k not in _HEAVY_FIELDS} + if isinstance(value, list): + return [_strip_heavy(v) for v in value] + return value + + +class ProductionBackend: + """Backend backed by the real in-process ``database`` layer. + + Heavy imports (``database``) happen lazily on first use so importing this + module never drags in the whole application. + """ + + def __init__( + self, db_manager: Any = None, mongo_db: Any = None, collections_manager: Any = None + ) -> None: + self._dbm = db_manager + self._mongo = mongo_db + self._cm = collections_manager + + # -- lazy wiring ------------------------------------------------------- + def _require_dbm(self) -> Any: + if self._dbm is None: + from database import db as _db # lazy heavy import + + self._dbm = _db + return self._dbm + + def _collections(self) -> Any: + if self._cm is None: + from database.collections_manager import CollectionsManager # lazy + + mongo = ( + self._mongo if self._mongo is not None else getattr(self._require_dbm(), "db", None) + ) + if mongo is None: + raise RuntimeError("MongoDB handle unavailable for collections") + self._cm = CollectionsManager(mongo) + return self._cm + + # -- files ------------------------------------------------------------- + def list_files(self, user_id: int, *, page: int = 1, per_page: int = 50) -> dict[str, Any]: + files, total = self._require_dbm().get_regular_files_paginated(user_id, page, per_page) + return { + "files": [_clean(f) for f in (files or [])], + "total": int(total or 0), + "page": page, + "per_page": per_page, + } + + def search_code( + self, user_id: int, *, query: str, language: str | None = None, limit: int = 20 + ) -> list[dict[str, Any]]: + rows = self._require_dbm().search_code( + user_id, query, programming_language=language, limit=limit + ) + return [_clean(r) for r in (rows or [])] + + def get_file( + self, + user_id: int, + *, + file_name: str | None = None, + file_id: str | None = None, + version: int | None = None, + ) -> dict[str, Any] | None: + dbm = self._require_dbm() + if file_id: + doc = dbm.get_file_by_id(file_id) + # get_file_by_id is NOT user-scoped -> enforce ownership explicitly. + if not doc or int(doc.get("user_id", -1)) != int(user_id): + return None + elif file_name and version is not None: + doc = dbm.get_version(user_id, file_name, int(version)) + elif file_name: + doc = dbm.get_latest_version(user_id, file_name) + else: + return None + return _full(doc) if doc else None + + def list_versions(self, user_id: int, *, file_name: str) -> list[dict[str, Any]]: + return [_clean(v) for v in (self._require_dbm().get_all_versions(user_id, file_name) or [])] + + # -- collections ------------------------------------------------------- + def list_collections(self, user_id: int, *, limit: int = 100) -> dict[str, Any]: + return self._collections().list_collections(user_id, limit=limit) + + def get_collection(self, user_id: int, *, collection_id: str) -> dict[str, Any]: + return self._collections().get_collection(user_id, collection_id) + + def get_collection_items( + self, + user_id: int, + *, + collection_id: str, + page: int = 1, + per_page: int = 50, + folder: str | None = None, + ) -> dict[str, Any]: + result = self._collections().get_collection_items( + user_id, collection_id, page=page, per_page=per_page, folder_filter=folder + ) + # Defense-in-depth: collection items are file *references* (no code today), + # but never let a heavy content field slip through if the manager changes. + if isinstance(result, dict) and isinstance(result.get("items"), list): + result["items"] = [_strip_heavy(item) for item in result["items"]] + return result diff --git a/mcp_server/handlers.py b/mcp_server/handlers.py new file mode 100644 index 000000000..b7ad9c34e --- /dev/null +++ b/mcp_server/handlers.py @@ -0,0 +1,96 @@ +"""Pure tool handlers. + +These are plain functions (no MCP / no Starlette imports) that validate and +clamp inputs, then delegate to a ``Backend``. Keeping them separate from the +FastMCP wiring makes the business logic trivially unit-testable. + +Every handler takes an authoritative, server-derived ``user_id`` — callers must +never pass a client-supplied user id here. +""" + +from __future__ import annotations + +from typing import Any + +MAX_PER_PAGE = 200 +MAX_SEARCH_LIMIT = 100 +MAX_COLLECTIONS_LIMIT = 500 + + +def _clamp(value: Any, lo: int, hi: int, default: int) -> int: + try: + ivalue = int(value) + except (TypeError, ValueError): + return default + return max(lo, min(hi, ivalue)) + + +def list_files(backend: Any, user_id: int, *, page: int = 1, per_page: int = 50) -> dict[str, Any]: + return backend.list_files( + user_id, + page=_clamp(page, 1, 10**9, 1), + per_page=_clamp(per_page, 1, MAX_PER_PAGE, 50), + ) + + +def search_code( + backend: Any, user_id: int, *, query: str, language: str | None = None, limit: int = 20 +) -> list[dict[str, Any]]: + query = (query or "").strip() + if not query: + return [] + return backend.search_code( + user_id, + query=query, + language=(language or None), + limit=_clamp(limit, 1, MAX_SEARCH_LIMIT, 20), + ) + + +def get_file( + backend: Any, + user_id: int, + *, + file_name: str | None = None, + file_id: str | None = None, + version: int | None = None, +) -> dict[str, Any] | None: + if not file_name and not file_id: + return None + return backend.get_file(user_id, file_name=file_name, file_id=file_id, version=version) + + +def list_versions(backend: Any, user_id: int, *, file_name: str) -> list[dict[str, Any]]: + if not file_name: + return [] + return backend.list_versions(user_id, file_name=file_name) + + +def list_collections(backend: Any, user_id: int, *, limit: int = 100) -> dict[str, Any]: + return backend.list_collections(user_id, limit=_clamp(limit, 1, MAX_COLLECTIONS_LIMIT, 100)) + + +def get_collection(backend: Any, user_id: int, *, collection_id: str) -> dict[str, Any]: + if not collection_id: + return {"ok": False, "error": "missing_collection_id"} + return backend.get_collection(user_id, collection_id=collection_id) + + +def get_collection_items( + backend: Any, + user_id: int, + *, + collection_id: str, + page: int = 1, + per_page: int = 50, + folder: str | None = None, +) -> dict[str, Any]: + if not collection_id: + return {"ok": False, "error": "missing_collection_id"} + return backend.get_collection_items( + user_id, + collection_id=collection_id, + page=_clamp(page, 1, 10**9, 1), + per_page=_clamp(per_page, 1, MAX_PER_PAGE, 50), + folder=folder, + ) diff --git a/mcp_server/server.py b/mcp_server/server.py new file mode 100644 index 000000000..17e5cd8cc --- /dev/null +++ b/mcp_server/server.py @@ -0,0 +1,104 @@ +"""FastMCP server wiring: tools + resources + the authenticated ASGI app. + +``build_mcp`` registers the read-only tools against an injected ``Backend``. +``build_app`` returns a Starlette ASGI app (Streamable HTTP) wrapped with PAT +auth plus an unauthenticated ``/healthz`` endpoint for platform health checks. + +Tools are defined as **sync** functions on purpose: FastMCP runs sync tools in a +worker thread, so the blocking (pymongo) backend calls never stall the event +loop, and the tool can still read ``ctx.request_context.request.state``. +""" + +from __future__ import annotations + +from typing import Any + +from mcp.server.fastmcp import Context, FastMCP +from starlette.responses import JSONResponse +from starlette.routing import Route + +from . import handlers +from .auth import PATAuthMiddleware, current_user_id + +_INSTRUCTIONS = ( + "Access the current user's private code files and collections stored in " + "CodeKeeper. Use search_code / list_files to find files (metadata only), " + "and get_file to read full contents. All data is scoped to the " + "authenticated user; this server is read-only." +) + + +def build_mcp(backend: Any, *, name: str = "CodeKeeper") -> FastMCP: + mcp: FastMCP = FastMCP(name, instructions=_INSTRUCTIONS, stateless_http=True) + + @mcp.tool(description="List the user's saved code files (metadata only, no code).") + def list_files(ctx: Context, page: int = 1, per_page: int = 50) -> dict: + return handlers.list_files(backend, current_user_id(ctx), page=page, per_page=per_page) + + @mcp.tool(description="Search the user's code by text; returns file metadata (no content).") + def search_code(ctx: Context, query: str, language: str | None = None, limit: int = 20) -> dict: + results = handlers.search_code( + backend, current_user_id(ctx), query=query, language=language, limit=limit + ) + return {"query": query, "count": len(results), "results": results} + + @mcp.tool(description="Get a file's full content by name or id (optional version number).") + def get_file( + ctx: Context, + file_name: str | None = None, + file_id: str | None = None, + version: int | None = None, + ) -> dict: + doc = handlers.get_file( + backend, current_user_id(ctx), file_name=file_name, file_id=file_id, version=version + ) + if doc is None: + return {"found": False} + return {"found": True, "file": doc} + + @mcp.tool(description="List all saved versions of a file by file_name (metadata only).") + def list_versions(ctx: Context, file_name: str) -> dict: + versions = handlers.list_versions(backend, current_user_id(ctx), file_name=file_name) + return {"file_name": file_name, "count": len(versions), "versions": versions} + + @mcp.tool(description="List the user's collections (named folders of files).") + def list_collections(ctx: Context, limit: int = 100) -> dict: + return handlers.list_collections(backend, current_user_id(ctx), limit=limit) + + @mcp.tool(description="Get a single collection by its id.") + def get_collection(ctx: Context, collection_id: str) -> dict: + return handlers.get_collection(backend, current_user_id(ctx), collection_id=collection_id) + + @mcp.tool(description="List files in a collection (paginated); optional folder filter.") + def get_collection_items( + ctx: Context, + collection_id: str, + page: int = 1, + per_page: int = 50, + folder: str | None = None, + ) -> dict: + return handlers.get_collection_items( + backend, + current_user_id(ctx), + collection_id=collection_id, + page=page, + per_page=per_page, + folder=folder, + ) + + return mcp + + +async def _healthz(_request): + return JSONResponse({"status": "ok", "service": "codekeeper-mcp"}) + + +def build_app(backend: Any, token_store: Any, *, name: str = "CodeKeeper"): + """Build the authenticated Streamable-HTTP ASGI app.""" + mcp = build_mcp(backend, name=name) + app = mcp.streamable_http_app() # Starlette app exposing POST/GET /mcp + # Unauthenticated health endpoint for the hosting platform. + app.router.routes.append(Route("/healthz", _healthz, methods=["GET"])) + # Auth guards everything except the exempt paths (see PATAuthMiddleware). + app.add_middleware(PATAuthMiddleware, token_store=token_store) + return app diff --git a/mcp_server/token_store.py b/mcp_server/token_store.py new file mode 100644 index 000000000..ad055902e --- /dev/null +++ b/mcp_server/token_store.py @@ -0,0 +1,153 @@ +"""Personal Access Token (PAT) store for the MCP server. + +A PAT lets a non-browser client (an MCP client such as Claude Code) authenticate +as a specific Telegram user without a session cookie. This mirrors the existing +one-time ``webapp_tokens`` bridge (see ``conversation_handlers.py``), but the +tokens here are long-lived, reusable, scoped, and revocable. + +Security choices: +- We store only a SHA-256 **hash** of each token, never the raw value. The raw + token is shown to the user exactly once (at issue time). +- Tokens carry 256 bits of entropy (``secrets.token_urlsafe(32)``), so a single + SHA-256 is sufficient here — unlike low-entropy passwords, no slow KDF needed. +- ``user_id`` is always authoritative and server-derived from the stored token. + +The store depends only on a duck-typed pymongo collection handle, so it can be +unit-tested with a tiny fake collection and adds no import-time dependencies. +""" + +from __future__ import annotations + +import hashlib +import secrets +from collections.abc import Iterable +from datetime import UTC, datetime, timedelta +from typing import Any + +TOKEN_PREFIX = "ckmcp_" # human-recognizable, greppable prefix +_RAW_ENTROPY_BYTES = 32 # 256-bit +DEFAULT_SCOPES = ("read",) +COLLECTION_NAME = "mcp_tokens" + + +def _now() -> datetime: + return datetime.now(UTC) + + +def _as_aware(dt: Any) -> datetime | None: + """Coerce a possibly tz-naive datetime (as Mongo may return) to aware UTC.""" + if not isinstance(dt, datetime): + return None + if dt.tzinfo is None: + return dt.replace(tzinfo=UTC) + return dt + + +def _iso(dt: Any) -> str | None: + aware = _as_aware(dt) + return aware.isoformat() if aware else None + + +def hash_token(raw_token: str) -> str: + """Return the stable lookup hash for a raw token.""" + return hashlib.sha256((raw_token or "").encode("utf-8")).hexdigest() + + +class MCPTokenStore: + """CRUD-lite store over the ``mcp_tokens`` collection.""" + + def __init__(self, db: Any, *, collection_name: str = COLLECTION_NAME) -> None: + # ``db`` is a pymongo Database handle (or a compatible fake). + self._coll = db[collection_name] + self._ensure_indexes() + + def _ensure_indexes(self) -> None: + # Best-effort; token operations must never fail because indexing did. + try: + self._coll.create_index("token_hash", unique=True) + self._coll.create_index("user_id") + except Exception: + pass + + # -- write ------------------------------------------------------------- + def issue( + self, + user_id: int, + *, + label: str | None = None, + scopes: Iterable[str] = DEFAULT_SCOPES, + ttl_days: int | None = None, + ) -> str: + """Create a new token and return the RAW value (shown once).""" + raw = TOKEN_PREFIX + secrets.token_urlsafe(_RAW_ENTROPY_BYTES) + expires_at = _now() + timedelta(days=ttl_days) if ttl_days else None + doc = { + "token_hash": hash_token(raw), + # first few chars only, for display in a "connections" list + "token_prefix": raw[: len(TOKEN_PREFIX) + 6], + "user_id": int(user_id), + "scopes": list(scopes) or list(DEFAULT_SCOPES), + "label": (label or "").strip() or "Claude", + "created_at": _now(), + "last_used_at": None, + "expires_at": expires_at, + "revoked": False, + } + self._coll.insert_one(doc) + return raw + + def revoke(self, user_id: int, token_id: Any) -> bool: + """Revoke a token owned by ``user_id`` (by its string id).""" + query_id: Any = token_id + try: # ids are ObjectId in real Mongo, plain values in fakes + from bson import ObjectId + + query_id = ObjectId(str(token_id)) + except Exception: + query_id = token_id + res = self._coll.update_one( + {"_id": query_id, "user_id": int(user_id)}, + {"$set": {"revoked": True, "revoked_at": _now()}}, + ) + return bool(getattr(res, "modified_count", 0)) + + # -- read -------------------------------------------------------------- + def verify(self, raw_token: str) -> dict[str, Any] | None: + """Return ``{"user_id", "scopes"}`` for a valid token, else ``None``. + + A token is valid iff it exists, is not revoked, and is not expired. + """ + if not raw_token: + return None + doc = self._coll.find_one({"token_hash": hash_token(raw_token), "revoked": {"$ne": True}}) + if not doc: + return None + expires_at = _as_aware(doc.get("expires_at")) + if expires_at is not None and expires_at < _now(): + return None + try: # best-effort "last used" touch; never block auth on it + self._coll.update_one({"_id": doc["_id"]}, {"$set": {"last_used_at": _now()}}) + except Exception: + pass + return { + "user_id": int(doc["user_id"]), + "scopes": list(doc.get("scopes") or DEFAULT_SCOPES), + } + + def list_for_user(self, user_id: int) -> list[dict[str, Any]]: + """List a user's tokens (metadata only — never the hash or raw value).""" + out: list[dict[str, Any]] = [] + for d in self._coll.find({"user_id": int(user_id)}): + out.append( + { + "id": str(d.get("_id")), + "token_prefix": d.get("token_prefix"), + "label": d.get("label"), + "scopes": list(d.get("scopes") or []), + "created_at": _iso(d.get("created_at")), + "last_used_at": _iso(d.get("last_used_at")), + "expires_at": _iso(d.get("expires_at")), + "revoked": bool(d.get("revoked", False)), + } + ) + return out diff --git a/mcp_server/wiring.py b/mcp_server/wiring.py new file mode 100644 index 000000000..89c7fc7a7 --- /dev/null +++ b/mcp_server/wiring.py @@ -0,0 +1,28 @@ +"""Small wiring helpers shared between the ASGI entrypoint and CLI scripts. + +Kept free of heavy/side-effecting imports so both ``mcp_server.app`` and +``scripts/mcp_issue_token.py`` can reuse it without triggering a Mongo connect. +""" + +from __future__ import annotations + +from typing import Any + + +def resolve_mongo(db_manager: Any) -> Any: + """Return the pymongo Database handle, forcing a lazy connection if needed.""" + mongo = getattr(db_manager, "db", None) + if mongo is not None: + return mongo + # Some managers connect lazily; nudge them, then re-read. + for attr in ("connect", "_get_repo", "ensure_connection"): + fn = getattr(db_manager, attr, None) + if callable(fn): + try: + fn() + except Exception: + pass + mongo = getattr(db_manager, "db", None) + if mongo is not None: + return mongo + return None diff --git a/requirements/base.txt b/requirements/base.txt index 0d0f743d2..d189bab8d 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -77,6 +77,10 @@ uvicorn==0.38.0 whitenoise==6.11.0 Flask-Compress==1.20 +# MCP server (Model Context Protocol) — exposes user files to Claude via mcp_server/ +# Compatible with the existing pins above (uvicorn>=0.31, pydantic>=2.11, httpx 0.28.x, starlette). +mcp==1.28.1 + # Rate limiting (Flask + generic) Flask-Limiter==4.0.0 limits==5.6.0 diff --git a/scripts/mcp_issue_token.py b/scripts/mcp_issue_token.py new file mode 100644 index 000000000..0bd5c3114 --- /dev/null +++ b/scripts/mcp_issue_token.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Issue a Personal Access Token (PAT) for the CodeKeeper MCP server. + +This is ops/testing tooling — it lets you mint a token for a given Telegram +user id without going through the bot. The raw token is printed ONCE; store it +in your MCP client's ``Authorization: Bearer `` header. + +Usage: + MONGODB_URL=... python scripts/mcp_issue_token.py --user-id 12345 \ + [--label "Claude Desktop"] [--ttl-days 90] +""" + +from __future__ import annotations + +import argparse +import os +import sys + + +def main() -> int: + parser = argparse.ArgumentParser(description="Issue a CodeKeeper MCP access token") + parser.add_argument( + "--user-id", type=int, required=True, help="Telegram user id (owner of the files)" + ) + parser.add_argument("--label", default="Claude", help="Human label for this token") + parser.add_argument( + "--ttl-days", type=int, default=None, help="Optional expiry in days (default: no expiry)" + ) + args = parser.parse_args() + + if not os.getenv("MONGODB_URL"): + print("ERROR: MONGODB_URL is not set.", file=sys.stderr) + return 2 + + # Make the repo importable when run directly. + sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + from database import db as db_manager # noqa: E402 + from mcp_server.token_store import MCPTokenStore # noqa: E402 + from mcp_server.wiring import resolve_mongo # noqa: E402 + + mongo = resolve_mongo(db_manager) + if mongo is None: + print("ERROR: could not connect to MongoDB.", file=sys.stderr) + return 3 + + store = MCPTokenStore(mongo) + raw = store.issue(args.user_id, label=args.label, ttl_days=args.ttl_days) + print(raw) + print( + f"\nIssued for user_id={args.user_id} (label={args.label!r}). " + "Store it now — it will not be shown again.", + file=sys.stderr, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/infrastructure/composition/files_facade.py b/src/infrastructure/composition/files_facade.py index 556937dd2..f10eadb64 100644 --- a/src/infrastructure/composition/files_facade.py +++ b/src/infrastructure/composition/files_facade.py @@ -639,6 +639,28 @@ def insert_webapp_login_token(self, token_doc: Dict[str, Any]) -> bool: logger.error("insert_webapp_login_token failed", exc_info=True) return False + def issue_mcp_token( + self, user_id: int, *, label: str = "Claude", ttl_days: Optional[int] = None + ) -> Optional[str]: + """ + Issue a long-lived MCP Personal Access Token bound to ``user_id`` and + return the RAW token (shown once). Returns None on failure. + + Keeps the MCP token-store wiring in one place so handlers don't reach + for raw PyMongo. See ``mcp_server/token_store.py``. + """ + try: + mongo_db = self.get_mongo_db() + if mongo_db is None: + return None + from mcp_server.token_store import MCPTokenStore # lazy import + + store = MCPTokenStore(mongo_db) + return store.issue(int(user_id), label=label, ttl_days=ttl_days) + except Exception: + logger.error("issue_mcp_token failed", exc_info=True) + return None + def list_active_user_ids(self) -> Optional[List[int]]: """ Return user ids eligible for admin broadcast (non-blocked users). diff --git a/tests/test_mcp_auth_middleware.py b/tests/test_mcp_auth_middleware.py new file mode 100644 index 000000000..ff4ed95f0 --- /dev/null +++ b/tests/test_mcp_auth_middleware.py @@ -0,0 +1,68 @@ +"""Unit tests for the PAT auth middleware (Starlette-level, no server/port).""" + +import pytest + +pytest.importorskip("starlette") + +from starlette.requests import Request # noqa: E402 +from starlette.responses import PlainTextResponse # noqa: E402 + +from mcp_server.auth import PATAuthMiddleware # noqa: E402 + + +class _FakeStore: + def __init__(self, mapping): + self.mapping = mapping + + def verify(self, token): + return self.mapping.get(token) + + +def _request(path="/mcp", headers=None): + raw_headers = [(k.lower().encode(), v.encode()) for k, v in (headers or [])] + scope = { + "type": "http", + "method": "POST", + "path": path, + "headers": raw_headers, + "query_string": b"", + "server": ("test", 80), + "scheme": "http", + } + return Request(scope) + + +async def _call_next(request): + return PlainTextResponse(f"ok:{getattr(request.state, 'user_id', None)}") + + +def _mw(store): + return PATAuthMiddleware(app=None, token_store=store) + + +async def test_missing_bearer_returns_401(): + resp = await _mw(_FakeStore({})).dispatch(_request(headers=[]), _call_next) + assert resp.status_code == 401 + assert resp.headers.get("WWW-Authenticate", "").startswith("Bearer") + + +async def test_invalid_token_returns_401(): + store = _FakeStore({"good": {"user_id": 5, "scopes": ["read"]}}) + resp = await _mw(store).dispatch( + _request(headers=[("authorization", "Bearer bad")]), _call_next + ) + assert resp.status_code == 401 + + +async def test_valid_token_injects_user_id(): + store = _FakeStore({"good": {"user_id": 5, "scopes": ["read"]}}) + resp = await _mw(store).dispatch( + _request(headers=[("authorization", "Bearer good")]), _call_next + ) + assert resp.status_code == 200 + assert resp.body == b"ok:5" + + +async def test_healthz_is_exempt(): + resp = await _mw(_FakeStore({})).dispatch(_request(path="/healthz", headers=[]), _call_next) + assert resp.status_code == 200 diff --git a/tests/test_mcp_backend.py b/tests/test_mcp_backend.py new file mode 100644 index 000000000..a2f87350e --- /dev/null +++ b/tests/test_mcp_backend.py @@ -0,0 +1,136 @@ +"""Unit tests for ProductionBackend: serialization + the critical ownership check.""" + +from mcp_server.backend import ProductionBackend, _full + + +class _FakeDbManager: + def __init__(self, files=None, by_id=None, versions=None, search=None): + self._files = files or [] + self._by_id = by_id or {} + self._versions = versions or [] + self._search = search or [] + + def get_regular_files_paginated(self, user_id, page, per_page): + return list(self._files), len(self._files) + + def search_code(self, user_id, query, programming_language=None, limit=20): + return list(self._search) + + def get_file_by_id(self, file_id): + return self._by_id.get(file_id) + + def get_latest_version(self, user_id, file_name): + return next((f for f in self._files if f.get("file_name") == file_name), None) + + def get_version(self, user_id, file_name, version): + return next( + ( + v + for v in self._versions + if v.get("file_name") == file_name and v.get("version") == version + ), + None, + ) + + def get_all_versions(self, user_id, file_name): + return [v for v in self._versions if v.get("file_name") == file_name] + + +def test_list_files_excludes_heavy_code_field(): + dbm = _FakeDbManager( + files=[ + { + "_id": "a", + "file_name": "x.py", + "code": "secret", + "programming_language": "python", + "file_size": 6, + } + ] + ) + out = ProductionBackend(db_manager=dbm).list_files(7, page=1, per_page=50) + assert out["total"] == 1 + f = out["files"][0] + assert "code" not in f + assert f["id"] == "a" + assert f["file_name"] == "x.py" + assert f["language"] == "python" + + +def test_search_excludes_code(): + dbm = _FakeDbManager( + search=[{"_id": "s1", "file_name": "y.py", "programming_language": "python"}] + ) + rows = ProductionBackend(db_manager=dbm).search_code(7, query="y") + assert rows and "code" not in rows[0] + assert rows[0]["id"] == "s1" + + +def test_get_file_by_id_enforces_ownership(): + dbm = _FakeDbManager( + by_id={"f1": {"_id": "f1", "user_id": 999, "file_name": "o.py", "code": "x"}} + ) + be = ProductionBackend(db_manager=dbm) + # Requesting user is not the owner -> denied. + assert be.get_file(7, file_id="f1") is None + # Owner gets full content. + got = be.get_file(999, file_id="f1") + assert got is not None and got["code"] == "x" + + +def test_get_file_by_name_returns_full_content(): + dbm = _FakeDbManager( + files=[{"_id": "a", "file_name": "x.py", "code": "print(1)", "user_id": 7}] + ) + doc = ProductionBackend(db_manager=dbm).get_file(7, file_name="x.py") + assert doc["code"] == "print(1)" + + +def test_get_specific_version(): + dbm = _FakeDbManager( + versions=[{"_id": "v2", "file_name": "x.py", "version": 2, "code": "v2code"}] + ) + doc = ProductionBackend(db_manager=dbm).get_file(7, file_name="x.py", version=2) + assert doc["version"] == 2 and doc["code"] == "v2code" + + +def test_large_file_content_mapped_to_code(): + assert _full({"_id": "a", "content": "blob"})["code"] == "blob" + + +def test_collections_delegate_to_manager(): + class _FakeCM: + def list_collections(self, user_id, limit=100): + return {"seen": (user_id, limit)} + + def get_collection(self, user_id, collection_id): + return {"seen": (user_id, collection_id)} + + def get_collection_items( + self, user_id, collection_id, page=1, per_page=20, folder_filter=None + ): + return {"seen": (user_id, collection_id, page, per_page, folder_filter)} + + be = ProductionBackend(collections_manager=_FakeCM()) + assert be.list_collections(3, limit=50)["seen"] == (3, 50) + assert be.get_collection(3, collection_id="c1")["seen"] == (3, "c1") + assert be.get_collection_items(3, collection_id="c1", page=2, per_page=10, folder="f")[ + "seen" + ] == (3, "c1", 2, 10, "f") + + +def test_collection_items_strip_heavy_fields(): + class _LeakyCM: + def get_collection_items( + self, user_id, collection_id, page=1, per_page=20, folder_filter=None + ): + return { + "ok": True, + "items": [{"id": "1", "file_name": "a.py", "code": "LEAK", "content": "LEAK2"}], + } + + be = ProductionBackend(collections_manager=_LeakyCM()) + out = be.get_collection_items(3, collection_id="c1") + item = out["items"][0] + assert "code" not in item and "content" not in item + assert item["file_name"] == "a.py" diff --git a/tests/test_mcp_handlers.py b/tests/test_mcp_handlers.py new file mode 100644 index 000000000..1119b0689 --- /dev/null +++ b/tests/test_mcp_handlers.py @@ -0,0 +1,79 @@ +"""Unit tests for the pure tool handlers (input validation + clamping).""" + +from mcp_server import handlers + + +class _RecordingBackend: + def __init__(self): + self.calls = [] + + def list_files(self, user_id, *, page, per_page): + self.calls.append(("list_files", user_id, page, per_page)) + return {"files": [], "total": 0, "page": page, "per_page": per_page} + + def search_code(self, user_id, *, query, language, limit): + self.calls.append(("search", user_id, query, language, limit)) + return [] + + def get_file(self, user_id, *, file_name, file_id, version): + self.calls.append(("get_file", user_id, file_name, file_id, version)) + return None + + def list_versions(self, user_id, *, file_name): + self.calls.append(("versions", user_id, file_name)) + return [] + + def list_collections(self, user_id, *, limit): + self.calls.append(("list_coll", user_id, limit)) + return {} + + def get_collection(self, user_id, *, collection_id): + self.calls.append(("get_coll", user_id, collection_id)) + return {} + + def get_collection_items(self, user_id, *, collection_id, page, per_page, folder): + self.calls.append(("items", user_id, collection_id, page, per_page, folder)) + return {} + + +def test_list_files_clamps_page_and_per_page(): + be = _RecordingBackend() + handlers.list_files(be, 1, page=0, per_page=99999) + assert be.calls[0] == ("list_files", 1, 1, 200) # page floored, per_page capped + + +def test_search_empty_query_short_circuits(): + be = _RecordingBackend() + assert handlers.search_code(be, 1, query=" ") == [] + assert be.calls == [] # backend not touched + + +def test_search_limit_capped(): + be = _RecordingBackend() + handlers.search_code(be, 1, query="x", limit=10_000) + assert be.calls[0] == ("search", 1, "x", None, 100) + + +def test_get_file_requires_an_identifier(): + be = _RecordingBackend() + assert handlers.get_file(be, 1) is None + assert be.calls == [] + + +def test_list_versions_requires_name(): + be = _RecordingBackend() + assert handlers.list_versions(be, 1, file_name="") == [] + assert be.calls == [] + + +def test_get_collection_items_missing_id_errors_without_call(): + be = _RecordingBackend() + out = handlers.get_collection_items(be, 1, collection_id="") + assert out["ok"] is False + assert be.calls == [] + + +def test_collections_limit_capped(): + be = _RecordingBackend() + handlers.list_collections(be, 1, limit=10_000) + assert be.calls[0] == ("list_coll", 1, 500) diff --git a/tests/test_mcp_server_build.py b/tests/test_mcp_server_build.py new file mode 100644 index 000000000..ef7c56bd3 --- /dev/null +++ b/tests/test_mcp_server_build.py @@ -0,0 +1,59 @@ +"""Smoke tests for the FastMCP wiring (tools registered, health route present).""" + +import pytest + +pytest.importorskip("mcp") +pytest.importorskip("starlette") + +from mcp_server.server import build_app, build_mcp # noqa: E402 + +_EXPECTED_TOOLS = { + "list_files", + "search_code", + "get_file", + "list_versions", + "list_collections", + "get_collection", + "get_collection_items", +} + + +class _FakeBackend: + def list_files(self, *a, **k): + return {} + + def search_code(self, *a, **k): + return [] + + def get_file(self, *a, **k): + return None + + def list_versions(self, *a, **k): + return [] + + def list_collections(self, *a, **k): + return {} + + def get_collection(self, *a, **k): + return {} + + def get_collection_items(self, *a, **k): + return {} + + +class _FakeStore: + def verify(self, token): + return None + + +async def test_all_tools_are_registered(): + mcp = build_mcp(_FakeBackend()) + tools = await mcp.list_tools() + names = {t.name for t in tools} + assert _EXPECTED_TOOLS <= names + + +def test_build_app_exposes_healthz_route(): + app = build_app(_FakeBackend(), _FakeStore()) + paths = {getattr(r, "path", None) for r in app.routes} + assert "/healthz" in paths diff --git a/tests/test_mcp_token_store.py b/tests/test_mcp_token_store.py new file mode 100644 index 000000000..30f84f7b6 --- /dev/null +++ b/tests/test_mcp_token_store.py @@ -0,0 +1,126 @@ +"""Unit tests for the MCP Personal Access Token store. + +Uses a tiny hand-rolled fake Mongo collection (the repo convention — no +mongomock dependency). +""" + +from datetime import UTC, datetime, timedelta + +from mcp_server.token_store import TOKEN_PREFIX, MCPTokenStore, hash_token + + +class _Result: + def __init__(self, modified=0): + self.modified_count = modified + + +class _FakeCollection: + def __init__(self): + self.docs = [] + self._seq = 0 + + def create_index(self, *a, **k): + return "idx" + + def insert_one(self, doc): + self._seq += 1 + doc = dict(doc) + doc.setdefault("_id", self._seq) + self.docs.append(doc) + return _Result() + + @staticmethod + def _match(doc, query): + for key, cond in query.items(): + if isinstance(cond, dict) and "$ne" in cond: + if doc.get(key) == cond["$ne"]: + return False + elif doc.get(key) != cond: + return False + return True + + def find_one(self, query): + for d in self.docs: + if self._match(d, query): + return d + return None + + def find(self, query): + return [d for d in self.docs if self._match(d, query)] + + def update_one(self, query, update): + for d in self.docs: + if self._match(d, query): + d.update(update.get("$set", {})) + return _Result(1) + return _Result(0) + + +class _FakeDB: + def __init__(self): + self._collections = {} + + def __getitem__(self, name): + return self._collections.setdefault(name, _FakeCollection()) + + +def _store(): + return MCPTokenStore(_FakeDB()) + + +def test_issue_returns_prefixed_raw_and_stores_hash_only(): + store = _store() + raw = store.issue(111, label="Claude") + assert raw.startswith(TOKEN_PREFIX) + doc = store._coll.docs[0] + assert doc["token_hash"] == hash_token(raw) + assert "token" not in doc # raw value is never persisted + assert doc["user_id"] == 111 + assert doc["scopes"] == ["read"] + + +def test_verify_roundtrip(): + store = _store() + raw = store.issue(222) + assert store.verify(raw) == {"user_id": 222, "scopes": ["read"]} + + +def test_verify_wrong_or_empty_token_returns_none(): + store = _store() + store.issue(1) + assert store.verify("ckmcp_nope") is None + assert store.verify("") is None + + +def test_revoked_token_is_denied(): + store = _store() + raw = store.issue(5) + token_id = store._coll.docs[0]["_id"] + assert store.revoke(5, token_id) is True + assert store.verify(raw) is None + + +def test_expired_token_is_denied(): + store = _store() + raw = store.issue(7, ttl_days=1) + store._coll.docs[0]["expires_at"] = datetime.now(UTC) - timedelta(days=2) + assert store.verify(raw) is None + + +def test_verify_updates_last_used(): + store = _store() + raw = store.issue(8) + assert store._coll.docs[0]["last_used_at"] is None + store.verify(raw) + assert store._coll.docs[0]["last_used_at"] is not None + + +def test_list_for_user_hides_secrets(): + store = _store() + store.issue(9, label="A") + rows = store.list_for_user(9) + assert len(rows) == 1 + assert "token_hash" not in rows[0] and "token" not in rows[0] + assert rows[0]["label"] == "A" + assert rows[0]["token_prefix"].startswith(TOKEN_PREFIX) + assert store.list_for_user(999) == []