refactor(tools): use FastMCP registration hooks

Author: jkoelker

src/schwab_mcp/tools/__init__.py

@@ -1,20 +1,15 @@
 from __future__ import annotations
 
-import inspect
-
 from mcp.server.fastmcp import FastMCP
 from schwab.client import AsyncClient
 
-from schwab_mcp.tools.registry import register
-
-# Import tool modules for their registration side effects
-from schwab_mcp.tools import tools as _tools  # noqa: F401
-from schwab_mcp.tools import account as _account  # noqa: F401
-from schwab_mcp.tools import history as _history  # noqa: F401
-from schwab_mcp.tools import options as _options  # noqa: F401
-from schwab_mcp.tools import orders as _orders  # noqa: F401
-from schwab_mcp.tools import quotes as _quotes  # noqa: F401
-from schwab_mcp.tools import transactions as _txns  # noqa: F401
+from schwab_mcp.tools import account as _account
+from schwab_mcp.tools import history as _history
+from schwab_mcp.tools import options as _options
+from schwab_mcp.tools import orders as _orders
+from schwab_mcp.tools import quotes as _quotes
+from schwab_mcp.tools import tools as _tools
+from schwab_mcp.tools import transactions as _txns
 
 _TOOL_MODULES = (
     _tools,
@@ -32,18 +27,10 @@ def register_tools(server: FastMCP, client: AsyncClient, *, allow_write: bool) -
     _ = client
 
     for module in _TOOL_MODULES:
-        for _, func in inspect.getmembers(module, inspect.iscoroutinefunction):
-            if not getattr(func, "_registered_tool", False):
-                continue
-            if getattr(func, "_write", False) and not allow_write:
-                continue
-            annotations = getattr(func, "_tool_annotations", None)
-            server.add_tool(
-                func,
-                name=func.__name__,
-                description=func.__doc__,
-                annotations=annotations,
-            )
-
-
-__all__ = ["register_tools", "register"]
+        register_module = getattr(module, "register", None)
+        if register_module is None:
+            raise AttributeError(f"Tool module {module.__name__} missing register()")
+        register_module(server, allow_write=allow_write)
+
+
+__all__ = ["register_tools"]

src/schwab_mcp/tools/_registration.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+from mcp.types import ToolAnnotations
+
+ToolFn = Callable[..., Awaitable[Any]]
+
+
+def register_tool(
+    server: FastMCP,
+    func: ToolFn,
+    *,
+    write: bool = False,
+    annotations: ToolAnnotations | None = None,
+) -> None:
+    """Register a Schwab tool using FastMCP's decorator plumbing."""
+
+    tool_annotations = annotations
+    if tool_annotations is None:
+        if write:
+            tool_annotations = ToolAnnotations(
+                readOnlyHint=False,
+                destructiveHint=True,
+            )
+        else:
+            tool_annotations = ToolAnnotations(
+                readOnlyHint=True,
+            )
+    else:
+        update: dict[str, Any] = {}
+        if tool_annotations.readOnlyHint is None:
+            update["readOnlyHint"] = not write
+        if write and tool_annotations.destructiveHint is None:
+            update["destructiveHint"] = True
+        if update:
+            tool_annotations = tool_annotations.model_copy(update=update)
+
+    server.tool(
+        name=func.__name__,
+        description=func.__doc__,
+        annotations=tool_annotations,
+    )(func)
+
+
+__all__ = ["register_tool"]

src/schwab_mcp/tools/account.py

@@ -2,12 +2,13 @@
 
 from typing import Annotated
 
+from mcp.server.fastmcp import FastMCP
+
 from schwab_mcp.context import SchwabContext, SchwabServerContext
-from schwab_mcp.tools.registry import register
+from schwab_mcp.tools._registration import register_tool
 from schwab_mcp.tools.utils import JSONType, call
 
 
-@register
 async def get_account_numbers(
     ctx: SchwabContext,
 ) -> JSONType:
@@ -18,7 +19,6 @@ async def get_account_numbers(
     return await call(context.accounts.get_account_numbers)
 
 
-@register
 async def get_accounts(
     ctx: SchwabContext,
 ) -> JSONType:
@@ -29,7 +29,6 @@ async def get_accounts(
     return await call(context.accounts.get_accounts)
 
 
-@register
 async def get_accounts_with_positions(
     ctx: SchwabContext,
 ) -> JSONType:
@@ -43,7 +42,6 @@ async def get_accounts_with_positions(
     )
 
 
-@register
 async def get_account(
     ctx: SchwabContext,
     account_hash: Annotated[str, "Account hash for the Schwab account"],
@@ -55,7 +53,6 @@ async def get_account(
     return await call(context.accounts.get_account, account_hash)
 
 
-@register
 async def get_account_with_positions(
     ctx: SchwabContext,
     account_hash: Annotated[str, "Account hash for the Schwab account"],
@@ -71,7 +68,6 @@ async def get_account_with_positions(
     )
 
 
-@register
 async def get_user_preferences(
     ctx: SchwabContext,
 ) -> JSONType:
@@ -80,3 +76,19 @@ async def get_user_preferences(
     """
     context: SchwabServerContext = ctx.request_context.lifespan_context
     return await call(context.accounts.get_user_preferences)
+
+
+_READ_ONLY_TOOLS = (
+    get_account_numbers,
+    get_accounts,
+    get_accounts_with_positions,
+    get_account,
+    get_account_with_positions,
+    get_user_preferences,
+)
+
+
+def register(server: FastMCP, *, allow_write: bool) -> None:
+    _ = allow_write
+    for func in _READ_ONLY_TOOLS:
+        register_tool(server, func)

src/schwab_mcp/tools/history.py

@@ -3,17 +3,17 @@
 from typing import Annotated
 
 import datetime
+from mcp.server.fastmcp import FastMCP
 
 from schwab_mcp.context import SchwabContext, SchwabServerContext
-from schwab_mcp.tools.registry import register
+from schwab_mcp.tools._registration import register_tool
 from schwab_mcp.tools.utils import JSONType, call
 
 
 def _parse_iso_datetime(value: str | None) -> datetime.datetime | None:
     return datetime.datetime.fromisoformat(value) if value is not None else None
 
 
-@register
 async def get_advanced_price_history(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
@@ -95,7 +95,6 @@ async def get_advanced_price_history(
     )
 
 
-@register
 async def get_price_history_every_minute(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
@@ -127,7 +126,6 @@ async def get_price_history_every_minute(
     )
 
 
-@register
 async def get_price_history_every_five_minutes(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
@@ -159,7 +157,6 @@ async def get_price_history_every_five_minutes(
     )
 
 
-@register
 async def get_price_history_every_ten_minutes(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
@@ -191,7 +188,6 @@ async def get_price_history_every_ten_minutes(
     )
 
 
-@register
 async def get_price_history_every_fifteen_minutes(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
@@ -223,7 +219,6 @@ async def get_price_history_every_fifteen_minutes(
     )
 
 
-@register
 async def get_price_history_every_thirty_minutes(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
@@ -255,7 +250,6 @@ async def get_price_history_every_thirty_minutes(
     )
 
 
-@register
 async def get_price_history_every_day(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security to fetch price history for"],
@@ -293,7 +287,6 @@ async def get_price_history_every_day(
     )
 
 
-@register
 async def get_price_history_every_week(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
@@ -323,3 +316,21 @@ async def get_price_history_every_week(
         need_extended_hours_data=extended_hours,
         need_previous_close=previous_close,
     )
+
+
+_READ_ONLY_TOOLS = (
+    get_advanced_price_history,
+    get_price_history_every_minute,
+    get_price_history_every_five_minutes,
+    get_price_history_every_ten_minutes,
+    get_price_history_every_fifteen_minutes,
+    get_price_history_every_thirty_minutes,
+    get_price_history_every_day,
+    get_price_history_every_week,
+)
+
+
+def register(server: FastMCP, *, allow_write: bool) -> None:
+    _ = allow_write
+    for func in _READ_ONLY_TOOLS:
+        register_tool(server, func)

src/schwab_mcp/tools/options.py

@@ -3,9 +3,10 @@
 from typing import Annotated
 
 import datetime
+from mcp.server.fastmcp import FastMCP
 
 from schwab_mcp.context import SchwabContext, SchwabServerContext
-from schwab_mcp.tools.registry import register
+from schwab_mcp.tools._registration import register_tool
 from schwab_mcp.tools.utils import JSONType, call
 
 
@@ -19,7 +20,6 @@ def _parse_date(value: str | datetime.date | None) -> datetime.date | None:
     return datetime.datetime.strptime(value, "%Y-%m-%d").date()
 
 
-@register
 async def get_option_chain(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the underlying security (e.g., 'AAPL', 'SPY')"],
@@ -66,7 +66,6 @@ async def get_option_chain(
     )
 
 
-@register
 async def get_advanced_option_chain(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the underlying security"],
@@ -156,7 +155,6 @@ async def get_advanced_option_chain(
     )
 
 
-@register
 async def get_option_expiration_chain(
     ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the underlying security"],
@@ -167,3 +165,16 @@ async def get_option_expiration_chain(
     context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.options
     return await call(client.get_option_expiration_chain, symbol)
+
+
+_READ_ONLY_TOOLS = (
+    get_option_chain,
+    get_advanced_option_chain,
+    get_option_expiration_chain,
+)
+
+
+def register(server: FastMCP, *, allow_write: bool) -> None:
+    _ = allow_write
+    for func in _READ_ONLY_TOOLS:
+        register_tool(server, func)