diff --git a/CHANGELOG.md b/CHANGELOG.md index 1398150..354b3d0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,35 +5,67 @@ All notable changes to this project will be documented in this file. See [Conven ## [Unreleased] ### Features -- Initial release of aiofmp -- Async-first Python client for Financial Modeling Prep API -- Built-in MCP (Model Context Protocol) server -- 160+ MCP tools across 22+ API categories -- CLI interface for easy MCP server management -- Comprehensive test coverage -- Type hints throughout the codebase -- Support for both STDIO and HTTP transports - -### API Categories -- Analyst: Financial estimates, ratings, price targets -- Calendar: Earnings, dividends, IPOs, economic events -- Chart: Historical price data and technical analysis -- Company: Company profiles, key metrics, corporate information -- Commodity: Commodity prices, quotes, historical data -- COT: Commitment of Traders reports -- Crypto: Cryptocurrency prices, quotes, market data -- DCF: Discounted Cash Flow valuations -- Directory: Symbol lists, exchanges, sectors, reference data -- Economics: Economic indicators, treasury rates, macro data -- ETF: ETF holdings, performance, analysis -- Forex: Foreign exchange rates and currency data -- Form 13F: Institutional holdings and filings -- Indexes: Stock market indices and performance -- Insider Trades: Insider trading activity and statistics -- Market Performance: Sector performance, market movers, P/E ratios -- News: Financial news, press releases, market updates -- Quote: Real-time quotes, price changes, market data -- Search: Symbol search, company search, stock screening -- Senate: Congressional trading disclosures -- Statements: Financial statements, ratios, metrics -- Technical Indicators: Moving averages, RSI, technical analysis tools +- **MCP server: selective tool registration.** New `--tools` / `--exclude-tools` CLI flags (with `AIOFMP_MCP_TOOLS` / `AIOFMP_MCP_EXCLUDE_TOOLS` env equivalents) let users restrict which of the 177 MCP tools get registered. The spec grammar mixes category-level and per-tool granularity: `chart` or `chart(*)` for a whole category, `chart(get_intraday_1hour,get_historical_price_full)` for specific tools, comma-separated. When both flags are set, the include set is the universe and exclude prunes from it. +- **MCP server: `--list-tools` flag** prints the full inventory of available categories and tool names, then exits. Does not require an API key. + +## [1.2.0] - 2026-05-24 + +### Features +- **Harvester.** New `aiofmp harvest` / `aiofmp harvest-status` CLI commands: a long-running daemon that proactively warms the local Parquet cache for 17 categories on per-category intervals (chart EOD/intraday, statements, news, analyst estimates/snapshots, insider trades, form 13F, commodities, forex, indexes, economics, DCF, technical indicators). +- **Plan-aware operation.** Reads an FMP plan tier (`basic`/`starter`/`premium`/`ultimate`) from `harvester.yaml`, paces requests via a sliding-window per-minute rate limiter, applies a US-only symbol filter on Starter, and auto-disables categories or endpoints that are entirely paywalled on the chosen plan (e.g. `form13f`, `period=quarter` for `key_metrics`, the `1min` intraday timeframe, the `press_releases` news variant). +- **Per-category paywall short-circuit with re-probe TTL.** After N consecutive HTTP 402s within a cycle, the harvester ends the cycle as `PARTIAL` and skips subsequent cycles for ~24h before re-probing — avoids burning requests against a fully-paywalled endpoint. +- **Bandwidth budget tracker.** SQLite-backed monthly ledger attributes bytes to the category that made the request; soft cap pauses categories until next month; hard cap raises `FMPBudgetError`. +- **SQLite state store.** Per-cycle bookkeeping, last-seen-date checkpoints, symbol-universe catalogs with TTL refresh, and the bandwidth ledger live in `/harvester.sqlite`. +- **Symbol catalog.** Lazily discovers six universes (financial symbols, actively trading, ETFs, commodities, forex pairs, indexes) and refreshes them on a configurable interval; supports payload-aware filters so the indexes universe can drop non-USD listings using FMP's `currency` field. +- **Caching pattern coverage.** New `PAGE_WALK` temporal pattern in `CachedClient` (harvester-write, user-read) for analyst estimates, insider trades, and form 13F. New `SnapshotStore` for the P4 single-row-per-entity pattern used by analyst ratings and DCF. +- **Multi-category cache sharing.** Chart endpoints (`historical-price-eod/*`, intraday) are now registered under `commodity`, `forex`, and `indexes` categories too, so the shared FMP wire endpoint hits the same Parquet store regardless of which SDK category was called. +- **Bandwidth callback + ContextVar.** `FmpClient.on_response_size` reports each response body size, with `current_harvest_category` ContextVar attributing the bytes to the category that initiated the call (or `"user"` outside a harvester cycle). +- **`aiofmp` CLI group.** Existing `aiofmp-mcp-server` entry point is preserved; new `aiofmp` umbrella command wraps `harvest` and `harvest-status`. + +### Fixes +- Atomic Parquet writes (write to `.tmp`, then rename) so an interrupted flush can't leave a 0-byte or partially-written file in the cache. +- Recover from 0-byte / corrupt Parquet files on read (delete and refetch) rather than erroring out. +- Sanitize Parquet records before write: stringify columns containing integers outside the float64-safe range (|x| > 2^53), stringify mixed string/numeric columns that appear during append, and null out empty struct values (`{"data": {}}`) that pyarrow can't infer a schema for. +- Per-minute rate-limit margin (~17%) under the documented plan caps to absorb residual server-side 429s. +- 429 and 5xx are transparently retried inside `_make_request` with bounded backoff; everything else (auth, paywall, parse, budget) propagates immediately so callers can apply per-exception policy. +- Per-category retry policy (`retry.on_429`, `retry.on_5xx`) wraps each `run_cycle` independently of the inner client retries. +- Cooperative cancellation: harvester cycles check `should_stop()` between symbols and pages so SIGINT propagates promptly. +- Statements safety-net seeds its checkpoint on first run and is bookkept under a separate `statements_safetynet` category row in the state store. +- News pagination uses a safety cap and short-circuits on empty pages. +- Symbol-catalog refresh treats permanent (4xx other than 429/402) errors as terminal and does not retry them. +- Use FMP's own economic-indicator names (`GDP`, `realGDP`, `CPI`, ...) — FRED codes like `UNRATE` / `FEDFUNDS` return `"Invalid name"`. +- Correct `financial_symbols` endpoint path. +- Quiet log lines for expected paywalls (info-level skip instead of warning per item) and at-startup category auto-disables. +- `aiofmp` CLI entry point is registered. + +### Tests +- End-to-end smoke test covering all 17 harvester categories. +- Live-run regression coverage for paywall-memory paths, the mixed-column sanitization, and the 5xx retry policy. + +### Docs +- Harvester README section plus `examples/harvester.example.yaml` tuned for Starter plan defaults. + +## [1.1.0] - 2026-03-05 + +### Features +- **CachedClient.** A drop-in caching wrapper around `FmpClient`. Intercepts cacheable time-series methods, stores data in local Parquet files, and only fetches missing date ranges from the API on subsequent calls. +- Supported caching patterns: + - **Date-range endpoints** (chart, economics, calendar, news, technical indicators, ...): gap detection fetches only the missing slices. + - **Period-based endpoints** (income statements, balance sheets, ratios, ...): fetches fresh data and merges with stored history by date. + - **Non-cacheable endpoints** (profiles, quotes, search, ...): pass through directly to the FMP API. +- Pluggable `StorageBackend` interface; ships with `ParquetStorage` (default location: `~/.aiofmp/cache`, override via `AIOFMP_CACHE_FILE_PATH`). +- `EndpointRegistry` describes which SDK methods are cacheable, the temporal pattern, the storage key, and the response date field. +- MCP server gains an `--cached` flag / `AIOFMP_CACHED` env var that switches the underlying client to `CachedClient` so MCP tool calls benefit from the same cache. + +## [1.0.0] - 2025-10-21 + +### Features +- Initial release of aiofmp. +- Async-first Python client for the Financial Modeling Prep API built on asyncio + aiohttp. +- Category-based organization with 22 API categories: analyst, calendar, chart, commodity, company, COT, crypto, DCF, directory, economics, ETF, forex, form 13F, indexes, insider trades, market performance, news, quote, search, senate, statements, technical indicators. +- Built-in MCP (Model Context Protocol) server (`aiofmp-mcp-server`) exposing 177 tools to AI assistants via FastMCP. +- Both STDIO and HTTP transports for the MCP server. +- Comprehensive error handling with typed exceptions: `FMPError`, `FMPAuthenticationError`, `FMPRateLimitError`, `FMPResponseError`, `FMPServerError`, `FMPPaywallError`. +- Connection pooling, configurable timeouts, exponential-backoff retries, and a per-process concurrent-request semaphore. +- Full type hints throughout the codebase. +- 500+ unit tests covering every category and the MCP tool layer. diff --git a/README.md b/README.md index eaeab8d..357af12 100644 --- a/README.md +++ b/README.md @@ -182,6 +182,39 @@ aiofmp-mcp-server --cached When `--cached` is enabled, the MCP server uses `CachedClient` under the hood. Time-series API calls are cached locally in Parquet files so repeated queries for the same historical data don't consume API quota. Cache directory defaults to `~/.aiofmp/cache` (override with `AIOFMP_CACHE_FILE_PATH`). +#### Selective Tool Registration + +By default the server exposes 177 tools across 22 categories. To reduce the surface area an AI agent has to reason about, restrict which tools are registered using `--tools` (allowlist) and/or `--exclude-tools` (denylist). The spec grammar mixes category-level and per-tool selection: + +| Spec | Meaning | +|---|---| +| `chart` | all tools in the chart category | +| `chart(*)` | same — explicit form | +| `chart(get_intraday_1hour)` | just that one chart tool | +| `chart(get_intraday_1hour,get_historical_price_full)` | exactly those chart tools | +| `*` | every tool in every category (default) | + +A spec is a comma-separated list of these entries. When both flags are set, `--tools` defines the universe and `--exclude-tools` prunes from it. + +```bash +# Only chart, quote, and search are visible to the AI +aiofmp-mcp-server --tools "chart(*),quote(*),search(*)" + +# A single tool from quote, all of search +aiofmp-mcp-server --tools "quote(get_stock_quote),search" + +# Everything except form13f, senate, and chart.get_intraday_1min +aiofmp-mcp-server --exclude-tools "form13f,senate,chart(get_intraday_1min)" + +# Allowlist with a carve-out +aiofmp-mcp-server --tools "chart(*),quote(*)" --exclude-tools "chart(get_intraday_1min)" + +# Discover available categories and tool names +aiofmp-mcp-server --list-tools +``` + +Environment variable equivalents (useful for Claude Desktop configs): `AIOFMP_MCP_TOOLS` and `AIOFMP_MCP_EXCLUDE_TOOLS`. Invalid category or tool names fail fast at startup with the valid list printed. + #### Claude Desktop Integration Add to your Claude Desktop configuration (`claude_desktop_config.json`): @@ -237,6 +270,13 @@ aiofmp-mcp-server --text-content # Enable local caching aiofmp-mcp-server --cached +# Restrict the tools that get registered +aiofmp-mcp-server --tools "chart(*),quote(get_stock_quote),search" +aiofmp-mcp-server --exclude-tools "form13f,senate" + +# List available categories and tools +aiofmp-mcp-server --list-tools + # All options aiofmp-mcp-server --transport http --host localhost --port 3000 --log-level INFO --api-key your_key --text-content --cached ``` @@ -249,6 +289,9 @@ aiofmp-mcp-server --transport http --host localhost --port 3000 --log-level INFO - `--api-key`: FMP API key (can also be set via `FMP_API_KEY` environment variable) - `--text-content`: Include text content alongside structured content in MCP tool responses (default: text content is empty when structured content is present) - `--cached`: Enable CachedClient to cache time-series data locally in Parquet files, minimizing API calls (default: off). Cache dir: `~/.aiofmp/cache` (override with `AIOFMP_CACHE_FILE_PATH`) +- `--tools`: Restrict which MCP tools are registered. Spec syntax: `category` or `category(*)` for a whole category, `category(tool1,tool2)` for specific tools; comma-separated. Env: `AIOFMP_MCP_TOOLS`. Default: all tools. +- `--exclude-tools`: Same syntax as `--tools`, but prunes from the include set. Env: `AIOFMP_MCP_EXCLUDE_TOOLS`. +- `--list-tools`: Print available categories and tool names, then exit. Does not require `FMP_API_KEY`. ### Available API Categories @@ -397,6 +440,8 @@ See `docs/superpowers/specs/2026-05-15-harvester-v2-design.md` for the full desi | `MCP_LOG_LEVEL` | Logging level | `INFO` | No | | `AIOFMP_CACHED` | Enable CachedClient (`true`/`false`) | `false` | No | | `AIOFMP_CACHE_FILE_PATH` | Cache directory for Parquet files | `~/.aiofmp/cache` | No | +| `AIOFMP_MCP_TOOLS` | Restrict MCP tools (allowlist; same grammar as `--tools`) | unset | No | +| `AIOFMP_MCP_EXCLUDE_TOOLS` | Prune MCP tools (denylist; same grammar as `--exclude-tools`) | unset | No | ### MCP Server Modes diff --git a/aiofmp/cli.py b/aiofmp/cli.py index 48e92b0..a9abcdf 100644 --- a/aiofmp/cli.py +++ b/aiofmp/cli.py @@ -58,6 +58,34 @@ def cli() -> None: default=False, help="Enable CachedClient: cache time-series data locally in Parquet files to minimize API calls. Cache dir defaults to ~/.aiofmp/cache (override with AIOFMP_CACHE_FILE_PATH).", ) +@click.option( + "--tools", + "tools_spec", + default=None, + help=( + "Restrict which MCP tools are registered. Spec syntax: " + "'category' or 'category(*)' for all tools in a category, " + "'category(tool1,tool2)' for specific tools, comma-separated. " + "Example: --tools 'chart(*),quote(get_stock_quote),search'. " + "Use --list-tools to discover names. Env: AIOFMP_MCP_TOOLS." + ), +) +@click.option( + "--exclude-tools", + "exclude_tools_spec", + default=None, + help=( + "Same syntax as --tools, but prunes from the include set. " + "Useful with or without --tools. Env: AIOFMP_MCP_EXCLUDE_TOOLS." + ), +) +@click.option( + "--list-tools", + "list_tools", + is_flag=True, + default=False, + help="Print available categories and tools, then exit.", +) def mcp_server( transport: str, host: str, @@ -66,6 +94,9 @@ def mcp_server( api_key: str | None, text_content: bool, cached: bool, + tools_spec: str | None, + exclude_tools_spec: str | None, + list_tools: bool, ): """ Start the aiofmp MCP server. @@ -95,6 +126,40 @@ def mcp_server( # Set logging level logging.getLogger().setLevel(getattr(logging, log_level.upper())) + # --list-tools is a discovery short-circuit; no server start, no API key needed. + if list_tools: + from .mcp_selection import format_inventory + + click.echo(format_inventory()) + return + + # Validate tool selection specs early so a bad CLI value errors out before + # we open any client sessions. We re-parse them inside register_tools(), + # but doing it here lets click report errors with proper exit codes. + if tools_spec is not None or exclude_tools_spec is not None: + from .mcp_selection import compute_selection, get_tool_inventory, parse_spec + + inventory = get_tool_inventory() + try: + include = parse_spec(tools_spec, inventory) if tools_spec else None + exclude = ( + parse_spec(exclude_tools_spec, inventory) + if exclude_tools_spec + else None + ) + effective = compute_selection(include, exclude, inventory) + except ValueError as e: + raise click.UsageError(str(e)) from e + if not effective: + raise click.UsageError( + "Tool selection is empty after combining --tools and --exclude-tools; " + "no tools would be registered." + ) + if tools_spec: + os.environ["AIOFMP_MCP_TOOLS"] = tools_spec + if exclude_tools_spec: + os.environ["AIOFMP_MCP_EXCLUDE_TOOLS"] = exclude_tools_spec + # Set API key if provided if api_key: os.environ["FMP_API_KEY"] = api_key diff --git a/aiofmp/mcp_selection.py b/aiofmp/mcp_selection.py new file mode 100644 index 0000000..18eae77 --- /dev/null +++ b/aiofmp/mcp_selection.py @@ -0,0 +1,210 @@ +"""Tool-selection grammar for the MCP server. + +Lets users restrict which of the 177 MCP tools get registered. Supports +category-level and per-tool granularity via a small spec grammar: + + chart -> all tools in chart + chart(*) -> all tools in chart + chart(get_intraday_1hour) -> just that one tool + chart(get_intraday_1hour,get_historical_price_full) -> those tools + * -> everything (default) + +A spec is a comma-separated list of such entries. Two specs can be combined: +the ``--tools`` allowlist defines the universe and ``--exclude-tools`` prunes +from it. + +The tool inventory is discovered by AST-scanning ``aiofmp/*_tools.py`` for +``@mcp.tool``-decorated ``(async )def`` functions. This avoids importing the +tool modules (which would side-effect the global ``mcp`` instance) just to +validate a CLI flag. +""" + +from __future__ import annotations + +import ast +import functools +import pathlib + +#: ``None`` in the allow-set means "all tools in this category". +#: A ``frozenset`` means exactly those tool names. +CategoryAllow = frozenset[str] | None +ParsedSpec = dict[str, CategoryAllow] + + +@functools.cache +def get_tool_inventory() -> dict[str, frozenset[str]]: + """Return ``{category: frozenset(tool_names)}`` for every ``*_tools.py``. + + Result is cached for the process lifetime. The scan reads files only; + it does not import the tool modules. + """ + inv: dict[str, frozenset[str]] = {} + here = pathlib.Path(__file__).parent + for path in sorted(here.glob("*_tools.py")): + if path.name == "mcp_tools.py": + continue + tree = ast.parse(path.read_text()) + tools: list[str] = [] + for node in tree.body: + if not isinstance(node, ast.AsyncFunctionDef | ast.FunctionDef): + continue + for dec in node.decorator_list: + target = dec.func if isinstance(dec, ast.Call) else dec + if ( + isinstance(target, ast.Attribute) + and target.attr == "tool" + and isinstance(target.value, ast.Name) + and target.value.id == "mcp" + ): + tools.append(node.name) + break + category = path.stem.removesuffix("_tools") + inv[category] = frozenset(tools) + return inv + + +def _split_top_level(spec: str) -> list[str]: + """Split ``spec`` on commas, respecting parentheses.""" + parts: list[str] = [] + depth = 0 + start = 0 + for i, ch in enumerate(spec): + if ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth < 0: + raise ValueError(f"unmatched ')' at position {i} in spec {spec!r}") + elif ch == "," and depth == 0: + parts.append(spec[start:i]) + start = i + 1 + if depth != 0: + raise ValueError(f"unmatched '(' in spec {spec!r}") + parts.append(spec[start:]) + return [p.strip() for p in parts if p.strip()] + + +def _parse_entry( + entry: str, inventory: dict[str, frozenset[str]] +) -> tuple[str, CategoryAllow]: + """Parse one entry like ``chart`` or ``chart(*)`` or ``chart(t1,t2)``.""" + if "(" not in entry: + name = entry.strip() + return (name, None) + + if not entry.endswith(")"): + raise ValueError( + f"malformed entry {entry!r}; expected 'category' or 'category(...)'" + ) + name, _, inside = entry[:-1].partition("(") + name = name.strip() + inside = inside.strip() + if not inside: + raise ValueError( + f"empty parens in {entry!r}; use '{name}(*)' or list tool names" + ) + if inside == "*": + return (name, None) + + tool_names = [t.strip() for t in inside.split(",")] + if any(not t for t in tool_names): + raise ValueError(f"empty tool name in {entry!r}") + return (name, frozenset(tool_names)) + + +def parse_spec( + spec: str, inventory: dict[str, frozenset[str]] | None = None +) -> ParsedSpec: + """Parse a comma-separated tool spec into ``{category: allow-set}``. + + Validates every category and every tool name against the inventory. + """ + if inventory is None: + inventory = get_tool_inventory() + entries = _split_top_level(spec) + if not entries: + raise ValueError("empty spec") + + result: ParsedSpec = {} + for entry in entries: + cat, allow = _parse_entry(entry, inventory) + if cat == "*": + if len(entries) != 1: + raise ValueError("'*' must be the only entry when used") + return dict.fromkeys(inventory) + if not cat: + raise ValueError(f"empty category name in {entry!r}") + if cat not in inventory: + raise ValueError(f"unknown category {cat!r}; valid: {sorted(inventory)}") + if cat in result: + raise ValueError(f"duplicate category {cat!r} in spec") + if allow is not None: + unknown = allow - inventory[cat] + if unknown: + raise ValueError( + f"unknown tool(s) in {cat!r}: {sorted(unknown)}; " + f"valid tools: {sorted(inventory[cat])}" + ) + result[cat] = allow + return result + + +def compute_selection( + include: ParsedSpec | None, + exclude: ParsedSpec | None, + inventory: dict[str, frozenset[str]] | None = None, +) -> dict[str, frozenset[str]]: + """Combine ``include`` and ``exclude`` into a per-category allow-set. + + Semantics: + - ``include=None`` → start with everything in the inventory. + - ``include`` set → start with just those (categories, tools). + - ``exclude`` then prunes from the working set; categories that end up + empty are dropped entirely. + + The returned dict maps each surviving category to the exact frozenset of + tools that should remain registered. Categories absent from the result + should not be imported at all. + """ + if inventory is None: + inventory = get_tool_inventory() + + if include is None: + working: dict[str, set[str]] = {c: set(inventory[c]) for c in inventory} + else: + working = {} + for cat, allow in include.items(): + working[cat] = set(inventory[cat]) if allow is None else set(allow) + + if exclude: + for cat, allow in exclude.items(): + if cat not in working: + continue + if allow is None: + del working[cat] + else: + working[cat] -= allow + if not working[cat]: + del working[cat] + + return {c: frozenset(t) for c, t in working.items() if t} + + +def format_inventory( + inventory: dict[str, frozenset[str]] | None = None, + restrict: list[str] | None = None, +) -> str: + """Human-readable inventory listing for ``--list-tools``.""" + if inventory is None: + inventory = get_tool_inventory() + lines: list[str] = [] + cats = sorted(inventory) if restrict is None else restrict + for cat in cats: + if cat not in inventory: + lines.append(f"{cat}: unknown category") + continue + tools = sorted(inventory[cat]) + lines.append(f"{cat} ({len(tools)}):") + for t in tools: + lines.append(f" - {t}") + return "\n".join(lines) diff --git a/aiofmp/mcp_server.py b/aiofmp/mcp_server.py index c237a0f..d0b53c2 100644 --- a/aiofmp/mcp_server.py +++ b/aiofmp/mcp_server.py @@ -6,12 +6,19 @@ """ import asyncio +import importlib import logging import os import sys from fastmcp import FastMCP +from .mcp_selection import ( + compute_selection, + get_tool_inventory, + parse_spec, +) + # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -20,36 +27,63 @@ mcp = FastMCP("FMP MCP Server") -def register_tools(): - """Register all MCP tools from the various modules.""" - try: - # Import and register tools from each category - from . import ( # noqa: F401 - analyst_tools, - calendar_tools, - chart_tools, - commodity_tools, - company_tools, - cot_tools, - crypto_tools, - dcf_tools, - directory_tools, - economics_tools, - etf_tools, - forex_tools, - form13f_tools, - indexes_tools, - insider_trades_tools, - market_performance_tools, - news_tools, - quote_tools, - search_tools, - senate_tools, - statements_tools, - technical_indicators_tools, +def _resolve_selection_from_env() -> dict[str, frozenset[str]]: + """Read ``AIOFMP_MCP_TOOLS`` / ``AIOFMP_MCP_EXCLUDE_TOOLS`` and resolve. + + Returns the per-category allow-set the registrar should honor. Empty env + vars mean "no restriction on this side". + """ + inventory = get_tool_inventory() + include_spec = os.getenv("AIOFMP_MCP_TOOLS", "").strip() + exclude_spec = os.getenv("AIOFMP_MCP_EXCLUDE_TOOLS", "").strip() + include = parse_spec(include_spec, inventory) if include_spec else None + exclude = parse_spec(exclude_spec, inventory) if exclude_spec else None + return compute_selection(include, exclude, inventory) + + +def register_tools(selection: dict[str, frozenset[str]] | None = None) -> None: + """Register MCP tools, optionally restricted to a per-category allow-set. + + Args: + selection: ``{category: frozenset(tool_names)}`` of tools to keep. When + ``None``, the selection is resolved from the + ``AIOFMP_MCP_TOOLS`` / ``AIOFMP_MCP_EXCLUDE_TOOLS`` env vars, falling + back to "register everything" if neither is set. + + Categories not present in ``selection`` are never imported. Within an + imported category, tools that aren't in the allow-set are removed via + ``mcp.remove_tool``. + """ + inventory = get_tool_inventory() + + if selection is None: + selection = _resolve_selection_from_env() + + if not selection: + logger.warning( + "MCP tool selection is empty; no tools will be registered. " + "Check AIOFMP_MCP_TOOLS / AIOFMP_MCP_EXCLUDE_TOOLS." ) + return + + try: + for cat in selection: + importlib.import_module(f"aiofmp.{cat}_tools") - logger.info("Successfully registered all MCP tools") + kept = 0 + removed = 0 + for cat, allowed in selection.items(): + for tool_name in inventory[cat] - allowed: + mcp.remove_tool(tool_name) + removed += 1 + kept += len(allowed) + + logger.info( + "Registered %d MCP tool(s) across %d categor(y/ies); pruned %d", + kept, + len(selection), + removed, + ) except ImportError as e: logger.error(f"Failed to import tool modules: {e}") raise diff --git a/tests/test_cli_mcp_server.py b/tests/test_cli_mcp_server.py new file mode 100644 index 0000000..b727163 --- /dev/null +++ b/tests/test_cli_mcp_server.py @@ -0,0 +1,88 @@ +"""Tests for the `aiofmp-mcp-server` CLI (tool-selection flags + listing).""" + +from __future__ import annotations + +import os +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from aiofmp.cli import cli + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +class TestListTools: + def test_list_tools_short_circuits_without_api_key(self, runner): + """--list-tools must work without FMP_API_KEY so users can discover names.""" + with patch.dict(os.environ, {}, clear=True): + result = runner.invoke(cli, ["mcp-server", "--list-tools"]) + assert result.exit_code == 0 + assert "chart" in result.output + assert "quote" in result.output + assert " - get_stock_quote" in result.output + + +class TestSpecValidation: + def test_unknown_category_errors(self, runner): + with patch.dict(os.environ, {"FMP_API_KEY": "k"}, clear=True): + result = runner.invoke(cli, ["mcp-server", "--tools", "bogus_cat"]) + assert result.exit_code != 0 + assert "unknown category" in result.output + + def test_unknown_tool_errors(self, runner): + with patch.dict(os.environ, {"FMP_API_KEY": "k"}, clear=True): + result = runner.invoke( + cli, ["mcp-server", "--tools", "quote(no_such_tool)"] + ) + assert result.exit_code != 0 + assert "unknown tool" in result.output + + def test_empty_effective_selection_errors(self, runner): + with patch.dict(os.environ, {"FMP_API_KEY": "k"}, clear=True): + result = runner.invoke( + cli, + [ + "mcp-server", + "--tools", + "quote(get_stock_quote)", + "--exclude-tools", + "quote(get_stock_quote)", + ], + ) + assert result.exit_code != 0 + assert "empty" in result.output.lower() + + def test_valid_spec_exports_env_and_runs(self, runner): + """Valid --tools should be exported as AIOFMP_MCP_TOOLS for register_tools.""" + with patch.dict(os.environ, {"FMP_API_KEY": "k"}, clear=True): + with patch("aiofmp.cli.asyncio.run") as mock_run: + mock_run.return_value = None + result = runner.invoke( + cli, + [ + "mcp-server", + "--tools", + "quote(get_stock_quote)", + ], + ) + # CLI exports the spec into env before calling asyncio.run + # so register_tools (inside run_server) can pick it up. + assert os.environ.get("AIOFMP_MCP_TOOLS") == "quote(get_stock_quote)" + assert result.exit_code == 0 + mock_run.assert_called_once() + + def test_exclude_only_exports_env(self, runner): + with patch.dict(os.environ, {"FMP_API_KEY": "k"}, clear=True): + with patch("aiofmp.cli.asyncio.run") as mock_run: + mock_run.return_value = None + result = runner.invoke( + cli, + ["mcp-server", "--exclude-tools", "form13f,senate"], + ) + assert os.environ.get("AIOFMP_MCP_EXCLUDE_TOOLS") == "form13f,senate" + assert result.exit_code == 0 diff --git a/tests/test_mcp_selection.py b/tests/test_mcp_selection.py new file mode 100644 index 0000000..e82a35a --- /dev/null +++ b/tests/test_mcp_selection.py @@ -0,0 +1,203 @@ +"""Tests for the MCP tool-selection grammar and inventory.""" + +from __future__ import annotations + +import pytest + +from aiofmp.mcp_selection import ( + compute_selection, + format_inventory, + get_tool_inventory, + parse_spec, +) + + +@pytest.fixture +def fake_inventory() -> dict[str, frozenset[str]]: + """Small inventory we control, so parser tests don't drift with real tool names.""" + return { + "chart": frozenset({"get_intraday_1min", "get_intraday_1hour", "get_full"}), + "quote": frozenset({"get_stock_quote", "get_aftermarket_quote"}), + "search": frozenset({"search_symbols", "search_companies"}), + } + + +class TestInventory: + def test_inventory_has_22_categories(self): + inv = get_tool_inventory() + assert len(inv) == 22 + + def test_inventory_total_tool_count(self): + inv = get_tool_inventory() + total = sum(len(v) for v in inv.values()) + assert total == 177 + + def test_inventory_contains_known_tools(self): + inv = get_tool_inventory() + assert "get_stock_quote" in inv["quote"] + assert "search_symbols" in inv["search"] + assert "get_intraday_1hour" in inv["chart"] + + def test_inventory_excludes_helper_module(self): + inv = get_tool_inventory() + assert "mcp" not in inv # mcp_tools.py has no @mcp.tool defs + + def test_inventory_is_cached(self): + assert get_tool_inventory() is get_tool_inventory() + + +class TestParseSpec: + def test_bare_category(self, fake_inventory): + result = parse_spec("chart", fake_inventory) + assert result == {"chart": None} + + def test_star_in_parens(self, fake_inventory): + result = parse_spec("chart(*)", fake_inventory) + assert result == {"chart": None} + + def test_single_tool(self, fake_inventory): + result = parse_spec("chart(get_intraday_1hour)", fake_inventory) + assert result == {"chart": frozenset({"get_intraday_1hour"})} + + def test_multiple_tools(self, fake_inventory): + result = parse_spec( + "chart(get_intraday_1hour,get_intraday_1min)", fake_inventory + ) + assert result == { + "chart": frozenset({"get_intraday_1hour", "get_intraday_1min"}) + } + + def test_multiple_categories(self, fake_inventory): + result = parse_spec("chart(*),quote(get_stock_quote),search", fake_inventory) + assert result == { + "chart": None, + "quote": frozenset({"get_stock_quote"}), + "search": None, + } + + def test_whitespace_tolerated(self, fake_inventory): + result = parse_spec( + " chart ( get_intraday_1hour , get_full ) , quote ", fake_inventory + ) + assert result == { + "chart": frozenset({"get_intraday_1hour", "get_full"}), + "quote": None, + } + + def test_universal_star(self, fake_inventory): + result = parse_spec("*", fake_inventory) + assert result == dict.fromkeys(fake_inventory) + + def test_star_with_other_entries_errors(self, fake_inventory): + with pytest.raises(ValueError, match="must be the only entry"): + parse_spec("*,chart", fake_inventory) + + def test_unknown_category_errors(self, fake_inventory): + with pytest.raises(ValueError, match="unknown category 'bogus'"): + parse_spec("bogus", fake_inventory) + + def test_unknown_tool_errors(self, fake_inventory): + with pytest.raises(ValueError, match="unknown tool"): + parse_spec("chart(no_such_tool)", fake_inventory) + + def test_empty_parens_errors(self, fake_inventory): + with pytest.raises(ValueError, match="empty parens"): + parse_spec("chart()", fake_inventory) + + def test_empty_tool_name_errors(self, fake_inventory): + with pytest.raises(ValueError, match="empty tool name"): + parse_spec("chart(get_full,)", fake_inventory) + + def test_duplicate_category_errors(self, fake_inventory): + with pytest.raises(ValueError, match="duplicate category"): + parse_spec("chart(*),chart(get_full)", fake_inventory) + + def test_unmatched_open_paren_errors(self, fake_inventory): + with pytest.raises(ValueError, match="unmatched '\\('"): + parse_spec("chart(get_full", fake_inventory) + + def test_unmatched_close_paren_errors(self, fake_inventory): + with pytest.raises(ValueError, match="unmatched '\\)'"): + parse_spec("chart)", fake_inventory) + + def test_empty_spec_errors(self, fake_inventory): + with pytest.raises(ValueError, match="empty spec"): + parse_spec("", fake_inventory) + + def test_only_commas_errors(self, fake_inventory): + with pytest.raises(ValueError, match="empty spec"): + parse_spec(",,", fake_inventory) + + +class TestComputeSelection: + def test_no_filters_returns_full_inventory(self, fake_inventory): + result = compute_selection(None, None, fake_inventory) + assert {c: set(t) for c, t in result.items()} == { + c: set(t) for c, t in fake_inventory.items() + } + + def test_include_only_category(self, fake_inventory): + include = parse_spec("chart", fake_inventory) + result = compute_selection(include, None, fake_inventory) + assert set(result) == {"chart"} + assert result["chart"] == fake_inventory["chart"] + + def test_include_specific_tools(self, fake_inventory): + include = parse_spec("chart(get_full)", fake_inventory) + result = compute_selection(include, None, fake_inventory) + assert result == {"chart": frozenset({"get_full"})} + + def test_exclude_only_whole_category(self, fake_inventory): + exclude = parse_spec("chart", fake_inventory) + result = compute_selection(None, exclude, fake_inventory) + assert "chart" not in result + assert "quote" in result + assert "search" in result + + def test_exclude_only_specific_tool(self, fake_inventory): + exclude = parse_spec("chart(get_intraday_1min)", fake_inventory) + result = compute_selection(None, exclude, fake_inventory) + assert "get_intraday_1min" not in result["chart"] + assert "get_intraday_1hour" in result["chart"] + + def test_include_and_exclude_combined(self, fake_inventory): + include = parse_spec("chart(*),quote(*)", fake_inventory) + exclude = parse_spec("chart(get_intraday_1min)", fake_inventory) + result = compute_selection(include, exclude, fake_inventory) + assert set(result) == {"chart", "quote"} + assert "get_intraday_1min" not in result["chart"] + assert "get_intraday_1hour" in result["chart"] + assert result["quote"] == fake_inventory["quote"] + + def test_exclude_prunes_to_empty_drops_category(self, fake_inventory): + include = parse_spec("chart(get_full)", fake_inventory) + exclude = parse_spec("chart(get_full)", fake_inventory) + result = compute_selection(include, exclude, fake_inventory) + assert "chart" not in result + + def test_exclude_unrelated_category_is_noop(self, fake_inventory): + include = parse_spec("chart", fake_inventory) + exclude = parse_spec("quote", fake_inventory) + result = compute_selection(include, exclude, fake_inventory) + assert set(result) == {"chart"} + + +class TestFormatInventory: + def test_format_includes_all_categories(self, fake_inventory): + out = format_inventory(fake_inventory) + assert "chart" in out + assert "quote" in out + assert "search" in out + + def test_format_lists_tools_indented(self, fake_inventory): + out = format_inventory(fake_inventory) + assert " - get_stock_quote" in out + + def test_format_restrict_to_one_category(self, fake_inventory): + out = format_inventory(fake_inventory, restrict=["chart"]) + assert "chart" in out + assert "quote" not in out + + def test_format_unknown_category_reported(self, fake_inventory): + out = format_inventory(fake_inventory, restrict=["bogus"]) + assert "unknown category" in out diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 7b48085..4ee94b3 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -12,7 +12,12 @@ from aiofmp.base import FMPAuthenticationError from aiofmp.fmp_client import get_fmp_client, reset_fmp_client -from aiofmp.mcp_server import main, run_server +from aiofmp.mcp_server import ( + _resolve_selection_from_env, + main, + register_tools, + run_server, +) class TestMCPServer: @@ -178,5 +183,84 @@ def test_port_conversion(self): assert isinstance(port, int) +class TestRegisterToolsSelection: + """Tests for tool-selection behavior in register_tools().""" + + def test_resolve_selection_no_env_returns_full_inventory(self): + with patch.dict(os.environ, {}, clear=True): + sel = _resolve_selection_from_env() + assert len(sel) == 22 + assert sum(len(v) for v in sel.values()) == 177 + + def test_resolve_selection_include_only(self): + with patch.dict(os.environ, {"AIOFMP_MCP_TOOLS": "quote,search"}, clear=True): + sel = _resolve_selection_from_env() + assert set(sel) == {"quote", "search"} + + def test_resolve_selection_exclude_only(self): + with patch.dict( + os.environ, {"AIOFMP_MCP_EXCLUDE_TOOLS": "form13f,senate"}, clear=True + ): + sel = _resolve_selection_from_env() + assert "form13f" not in sel + assert "senate" not in sel + assert "quote" in sel + + def test_resolve_selection_include_and_exclude(self): + with patch.dict( + os.environ, + { + "AIOFMP_MCP_TOOLS": "chart(*),quote(*)", + "AIOFMP_MCP_EXCLUDE_TOOLS": "chart(get_intraday_1min)", + }, + clear=True, + ): + sel = _resolve_selection_from_env() + assert set(sel) == {"chart", "quote"} + assert "get_intraday_1min" not in sel["chart"] + assert "get_intraday_1hour" in sel["chart"] + + def test_resolve_selection_per_tool(self): + with patch.dict( + os.environ, + {"AIOFMP_MCP_TOOLS": "quote(get_stock_quote)"}, + clear=True, + ): + sel = _resolve_selection_from_env() + assert sel == {"quote": frozenset({"get_stock_quote"})} + + def test_register_tools_imports_only_selected_categories(self): + """register_tools() must import exactly the modules in the selection + and call mcp.remove_tool for tools outside each category's allow-set. + + Mocks importlib + remove_tool so the global `mcp` instance shared with + other tests isn't mutated. + """ + from aiofmp.mcp_selection import get_tool_inventory + from aiofmp.mcp_server import mcp as mcp_instance + + selection = { + "quote": frozenset({"get_stock_quote"}), + "search": frozenset(), # whole-category drop after import + } + with patch("aiofmp.mcp_server.importlib.import_module") as mock_import: + with patch.object(mcp_instance, "remove_tool") as mock_remove: + register_tools(selection=selection) + + imported = {call.args[0] for call in mock_import.call_args_list} + assert imported == {"aiofmp.quote_tools", "aiofmp.search_tools"} + + removed = {call.args[0] for call in mock_remove.call_args_list} + inv = get_tool_inventory() + expected_removed = (inv["quote"] - {"get_stock_quote"}) | inv["search"] + assert removed == set(expected_removed) + + def test_register_tools_empty_selection_is_noop(self, caplog): + register_tools(selection={}) + assert any( + "no tools will be registered" in rec.message for rec in caplog.records + ) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/uv.lock b/uv.lock index 6a3aee5..bdae3f8 100644 --- a/uv.lock +++ b/uv.lock @@ -4,7 +4,7 @@ requires-python = ">=3.10" [[package]] name = "aiofmp" -version = "1.1.0" +version = "1.2.0" source = { editable = "." } dependencies = [ { name = "aiohttp" },