mattermost-dl/mm_search_mcp.py

#!/usr/bin/env python3
"""MCP server (stdio) exposing mm_search.py search capabilities.

Cursor / Claude Desktop config example (~/.cursor/mcp.json):

    {
      "mcpServers": {
        "mattermost-search": {
          "command": "/path/to/mattermost-dl/.venv/bin/python",
          "args": ["/path/to/mattermost-dl/mm_search_mcp.py"],
          "cwd": "/path/to/mattermost-dl",
          "env": {
            "MM_SEARCH_ROOT": "results",
            "MM_SEARCH_DB": "mattermost_index.sqlite"
          }
        }
      }
    }

Requires: pip install fastmcp
"""
from __future__ import annotations

import datetime as dt
import os
from typing import Any, Literal, Optional

from fastmcp import FastMCP

import mm_search as mm_search_lib

mcp = FastMCP("mattermost-search")

SearchMode = Literal["message", "file_name", "file_content"]


def _env_root() -> str:
    return os.environ.get("MM_SEARCH_ROOT", "results")


def _env_db() -> str:
    return os.environ.get("MM_SEARCH_DB", "mattermost_index.sqlite")


def _parse_date(value: str) -> dt.datetime:
    try:
        return mm_search_lib.parse_iso_date(value)
    except Exception as exc:
        raise ValueError(f"invalid date '{value}': {exc}") from exc


def _collect_matches(
    *,
    query: str,
    root: str,
    user: Optional[str],
    channel: Optional[str],
    team: Optional[str],
    after: Optional[dt.datetime],
    before: Optional[dt.datetime],
    regex: bool,
    search_mode: SearchMode,
    ignore_case: Optional[bool],
    limit: Optional[int],
    db: str,
    no_index: bool,
) -> list[dict[str, Any]]:
    if search_mode != "message" and regex:
        raise ValueError("search_mode file_name/file_content do not support regex=True")

    effective_ignore_case = ignore_case if ignore_case is not None else (not regex)
    matcher = mm_search_lib.make_matcher(query, regex, effective_ignore_case)
    use_index = (not no_index) and os.path.exists(db)

    if search_mode in ("file_name", "file_content"):
        if not use_index:
            raise ValueError(
                "Attachment search (file_name/file_content) requires sqlite index. "
                "Create it with mm_index.py and do not set no_index=True."
            )
        mode = "name" if search_mode == "file_name" else "content"
        matches_iter = mm_search_lib.search_attachments_index(
            db_path=db,
            query=query,
            user=user,
            channel=channel,
            team=team,
            after=after,
            before=before,
            limit=limit,
            mode=mode,
        )
        output_limit = None
    elif use_index and not regex:
        matches_iter = mm_search_lib.search_index(
            db_path=db,
            query=query,
            user=user,
            channel=channel,
            team=team,
            after=after,
            before=before,
            limit=limit,
            ignore_case=effective_ignore_case,
        )
        output_limit = None
    else:
        files = list(mm_search_lib.iter_result_files(root))
        matches_iter = (
            m
            for path in files
            for m in mm_search_lib.search_file(
                path=path,
                matcher=matcher,
                user=user,
                channel=channel,
                team=team,
                after=after,
                before=before,
            )
        )
        output_limit = limit

    out: list[dict[str, Any]] = []
    count = 0
    for match in matches_iter:
        if output_limit is not None and count >= output_limit:
            break
        out.append(match)
        count += 1
    return out


def _format_human(matches: list[dict[str, Any]], show_file: bool) -> str:
    lines: list[str] = []
    for m in matches:
        p = m["post"]
        created = p.get("created", "")
        username = p.get("username", "")
        msg = (p.get("message") or "").replace("\n", "\\n")
        attachment = m.get("attachment")
        attachment_suffix = ""
        if isinstance(attachment, dict) and attachment.get("filename"):
            attachment_suffix = f" [file: {attachment['filename']}]"
        if show_file:
            file_rel = os.path.relpath(m["file"])
            idx = p.get("idx", "?")
            prefix = f"{file_rel}:{idx}: "
        else:
            prefix = ""
        channel = m.get("channel") or "?"
        team = m.get("team") or ""
        team_prefix = f"[{team}]" if team else ""
        lines.append(f"{prefix}{team_prefix}[{channel}] {created} {username}: {msg}{attachment_suffix}")
    return "\n".join(lines)


@mcp.tool
def mm_search(
    query: str,
    root: Optional[str] = None,
    user: Optional[str] = None,
    channel: Optional[str] = None,
    team: Optional[str] = None,
    after: Optional[str] = None,
    before: Optional[str] = None,
    regex: bool = False,
    search_mode: SearchMode = "message",
    ignore_case: Optional[bool] = None,
    limit: Optional[int] = None,
    db: Optional[str] = None,
    no_index: bool = False,
    output: Literal["json", "human"] = "json",
    show_file: bool = False,
) -> dict[str, Any]:
    """Search Mattermost export JSON files produced by mattermost-dl.py.

    Mirrors mm_search.py CLI. Returns structured matches or human-readable lines.

    Args:
        query: Search string or regex (when regex=True).
        root: Root directory with export results (default: MM_SEARCH_ROOT or "results").
        user: Filter by username (exact match).
        channel: Filter by channel display_name or name (exact match).
        team: Filter by team name (exact match).
        after: Only messages on/after this date (YYYY-MM-DD or ISO).
        before: Only messages on/before this date (YYYY-MM-DD or ISO).
        regex: Treat query as regular expression (message scan only).
        search_mode: "message" (default), "file_name", or "file_content".
        ignore_case: Case-insensitive search. Default: True for substring, False for regex.
        limit: Maximum number of matches to return.
        db: SQLite index path (default: MM_SEARCH_DB or "mattermost_index.sqlite").
        no_index: Do not use sqlite index even if DB exists.
        output: "json" returns match objects; "human" returns CLI-style text lines.
        show_file: Include source filename and per-file index in human output.
    """
    resolved_root = root or _env_root()
    resolved_db = db or _env_db()
    parsed_after = _parse_date(after) if after else None
    parsed_before = _parse_date(before) if before else None

    matches = _collect_matches(
        query=query,
        root=resolved_root,
        user=user,
        channel=channel,
        team=team,
        after=parsed_after,
        before=parsed_before,
        regex=regex,
        search_mode=search_mode,
        ignore_case=ignore_case,
        limit=limit,
        db=resolved_db,
        no_index=no_index,
    )

    result: dict[str, Any] = {
        "count": len(matches),
        "used_index": (not no_index)
        and os.path.exists(resolved_db)
        and (search_mode in ("file_name", "file_content") or not regex),
        "search_mode": search_mode,
        "root": resolved_root,
        "db": resolved_db,
    }

    if output == "human":
        result["text"] = _format_human(matches, show_file=show_file)
    else:
        result["matches"] = matches

    return result


@mcp.tool
def mm_search_index_status(
    db: Optional[str] = None,
    root: Optional[str] = None,
) -> dict[str, Any]:
    """Report whether the sqlite index and export root are available."""
    resolved_db = db or _env_db()
    resolved_root = root or _env_root()
    db_exists = os.path.exists(resolved_db)
    root_exists = os.path.isdir(resolved_root)

    json_count = 0
    if root_exists:
        json_count = sum(1 for _ in mm_search_lib.iter_result_files(resolved_root))

    return {
        "db": resolved_db,
        "db_exists": db_exists,
        "root": resolved_root,
        "root_exists": root_exists,
        "json_files": json_count,
        "index_required_for": ["file_name", "file_content"],
    }


if __name__ == "__main__":
    mcp.run(transport="stdio")