Merge pull request modelcontextprotocol#11 from jkoelker/jk/fastmcp

refactor(tools): use fastmcp context

Author: jkoelker

src/schwab_mcp/context.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any, cast
 
 from schwab.client import AsyncClient
+from mcp.server.fastmcp import Context as MCPContext
 
 if TYPE_CHECKING:
     from schwab_mcp.tools._protocols import (
@@ -44,4 +45,7 @@ def __post_init__(self) -> None:
         self.transactions = cast(TransactionsClient, self.client)
 
 
-__all__ = ["SchwabServerContext"]
+SchwabContext = MCPContext[Any, SchwabServerContext, Any]
+
+
+__all__ = ["SchwabServerContext", "SchwabContext"]

src/schwab_mcp/tools/__init__.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import inspect
+
 from mcp.server.fastmcp import FastMCP
 from schwab.client import AsyncClient
 
-from schwab_mcp.tools.registry import iter_registered_tools, register
-from schwab_mcp.tools import utils as tool_utils
+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
@@ -15,22 +16,34 @@
 from schwab_mcp.tools import quotes as _quotes  # noqa: F401
 from schwab_mcp.tools import transactions as _txns  # noqa: F401
 
+_TOOL_MODULES = (
+    _tools,
+    _account,
+    _history,
+    _options,
+    _orders,
+    _quotes,
+    _txns,
+)
+
 
 def register_tools(server: FastMCP, client: AsyncClient, *, allow_write: bool) -> None:
     """Register all Schwab tools with the provided FastMCP server."""
-    tool_utils.set_write_enabled(allow_write)
-
-    for func in iter_registered_tools():
-        if getattr(func, "_write", False) and not allow_write:
-            continue
-        annotations = getattr(func, "_tool_annotations", None)
-        tool_kwargs = {
-            "name": func.__name__,
-            "description": func.__doc__,
-        }
-        if annotations is not None:
-            tool_kwargs["annotations"] = annotations
-        server.tool(**tool_kwargs)(func)
+    _ = 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"]

src/schwab_mcp/tools/account.py

@@ -2,42 +2,41 @@
 
 from typing import Annotated
 
-from mcp.server.fastmcp import Context
-
+from schwab_mcp.context import SchwabContext, SchwabServerContext
 from schwab_mcp.tools.registry import register
-from schwab_mcp.tools.utils import call, get_context
+from schwab_mcp.tools.utils import call
 
 
 @register
 async def get_account_numbers(
-    ctx: Context,
+    ctx: SchwabContext,
 ) -> str:
     """
     Returns mapping of account IDs to account hashes. Hashes required for account-specific calls. Use first.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     return await call(context.accounts.get_account_numbers)
 
 
 @register
 async def get_accounts(
-    ctx: Context,
+    ctx: SchwabContext,
 ) -> str:
     """
     Returns balances/info for all linked accounts (funds, cash, margin). Does not return hashes; use get_account_numbers first.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     return await call(context.accounts.get_accounts)
 
 
 @register
 async def get_accounts_with_positions(
-    ctx: Context,
+    ctx: SchwabContext,
 ) -> str:
     """
     Returns balances, info, and positions (holdings, cost, gain/loss) for all linked accounts. Does not return hashes; use get_account_numbers first.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     return await call(
         context.accounts.get_accounts,
         fields=[context.accounts.Account.Fields.POSITIONS],
@@ -46,25 +45,25 @@ async def get_accounts_with_positions(
 
 @register
 async def get_account(
-    ctx: Context,
+    ctx: SchwabContext,
     account_hash: Annotated[str, "Account hash for the Schwab account"],
 ) -> str:
     """
     Returns balance/info for a specific account via account_hash (from get_account_numbers). Includes funds, cash, margin info.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     return await call(context.accounts.get_account, account_hash)
 
 
 @register
 async def get_account_with_positions(
-    ctx: Context,
+    ctx: SchwabContext,
     account_hash: Annotated[str, "Account hash for the Schwab account"],
 ) -> str:
     """
     Returns balance, info, and positions for a specific account via account_hash. Includes holdings, quantity, cost basis, unrealized gain/loss.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     return await call(
         context.accounts.get_account,
         account_hash,
@@ -74,10 +73,10 @@ async def get_account_with_positions(
 
 @register
 async def get_user_preferences(
-    ctx: Context,
+    ctx: SchwabContext,
 ) -> str:
     """
     Returns user preferences (nicknames, display settings, notifications) for all linked accounts.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     return await call(context.accounts.get_user_preferences)

src/schwab_mcp/tools/history.py

@@ -3,10 +3,10 @@
 from typing import Annotated
 
 import datetime
-from mcp.server.fastmcp import Context
 
+from schwab_mcp.context import SchwabContext, SchwabServerContext
 from schwab_mcp.tools.registry import register
-from schwab_mcp.tools.utils import call, get_context
+from schwab_mcp.tools.utils import call
 
 
 def _parse_iso_datetime(value: str | None) -> datetime.datetime | None:
@@ -15,7 +15,7 @@ def _parse_iso_datetime(value: str | None) -> datetime.datetime | None:
 
 @register
 async def get_advanced_price_history(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
     period_type: Annotated[
         str | None, "Period type: DAY, MONTH, YEAR, YEAR_TO_DATE"
@@ -60,7 +60,7 @@ async def get_advanced_price_history(
       YEAR_TO_DATE: DAILY, WEEKLY (default)
     Dates must be in ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
@@ -97,7 +97,7 @@ async def get_advanced_price_history(
 
 @register
 async def get_price_history_every_minute(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
     start_datetime: Annotated[
         str | None, "Start date for history (ISO format, e.g., '2023-01-01T09:30:00')"
@@ -111,7 +111,7 @@ async def get_price_history_every_minute(
     """
     Get OHLCV price history per minute. For detailed intraday analysis. Max 48 days history. Dates ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
@@ -129,7 +129,7 @@ async def get_price_history_every_minute(
 
 @register
 async def get_price_history_every_five_minutes(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
     start_datetime: Annotated[
         str | None, "Start date for history (ISO format, e.g., '2023-01-01T09:30:00')"
@@ -143,7 +143,7 @@ async def get_price_history_every_five_minutes(
     """
     Get OHLCV price history per 5 minutes. Balance between detail and noise. Approx. 9 months history. Dates ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
@@ -161,7 +161,7 @@ async def get_price_history_every_five_minutes(
 
 @register
 async def get_price_history_every_ten_minutes(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
     start_datetime: Annotated[
         str | None, "Start date for history (ISO format, e.g., '2023-01-01T09:30:00')"
@@ -175,7 +175,7 @@ async def get_price_history_every_ten_minutes(
     """
     Get OHLCV price history per 10 minutes. Good for intraday trends/levels. Approx. 9 months history. Dates ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
@@ -193,7 +193,7 @@ async def get_price_history_every_ten_minutes(
 
 @register
 async def get_price_history_every_fifteen_minutes(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
     start_datetime: Annotated[
         str | None, "Start date for history (ISO format, e.g., '2023-01-01T09:30:00')"
@@ -207,7 +207,7 @@ async def get_price_history_every_fifteen_minutes(
     """
     Get OHLCV price history per 15 minutes. Shows significant intraday moves, filters noise. Approx. 9 months history. Dates ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
@@ -225,7 +225,7 @@ async def get_price_history_every_fifteen_minutes(
 
 @register
 async def get_price_history_every_thirty_minutes(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
     start_datetime: Annotated[
         str | None, "Start date for history (ISO format, e.g., '2023-01-01T09:30:00')"
@@ -239,7 +239,7 @@ async def get_price_history_every_thirty_minutes(
     """
     Get OHLCV price history per 30 minutes. For broader intraday trends, filters noise. Approx. 9 months history. Dates ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
@@ -257,7 +257,7 @@ async def get_price_history_every_thirty_minutes(
 
 @register
 async def get_price_history_every_day(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security to fetch price history for"],
     start_datetime: Annotated[
         str | None,
@@ -277,7 +277,7 @@ async def get_price_history_every_day(
     """
     Get daily OHLCV price history. For medium/long-term analysis. Extensive history (back to 1985 possible). Dates ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
@@ -295,7 +295,7 @@ async def get_price_history_every_day(
 
 @register
 async def get_price_history_every_week(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the security"],
     start_datetime: Annotated[
         str | None, "Start date for history (ISO format, e.g., '2023-01-01T00:00:00')"
@@ -309,7 +309,7 @@ async def get_price_history_every_week(
     """
     Get weekly OHLCV price history. For long-term analysis, major cycles. Extensive history (back to 1985 possible). Dates ISO format.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)

src/schwab_mcp/tools/options.py

@@ -3,10 +3,10 @@
 from typing import Annotated
 
 import datetime
-from mcp.server.fastmcp import Context
 
+from schwab_mcp.context import SchwabContext, SchwabServerContext
 from schwab_mcp.tools.registry import register
-from schwab_mcp.tools.utils import call, get_context
+from schwab_mcp.tools.utils import call
 
 
 def _parse_date(value: str | datetime.date | None) -> datetime.date | None:
@@ -21,7 +21,7 @@ def _parse_date(value: str | datetime.date | None) -> datetime.date | None:
 
 @register
 async def get_option_chain(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the underlying security (e.g., 'AAPL', 'SPY')"],
     contract_type: Annotated[
         str | None, "Type of option contracts: CALL, PUT, or ALL (default)"
@@ -47,7 +47,7 @@ async def get_option_chain(
     Params: symbol, contract_type (CALL/PUT/ALL), strike_count (default 25), include_quotes (bool), from_date (YYYY-MM-DD), to_date (YYYY-MM-DD).
     Limit data returned using strike_count and date parameters.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.options
 
     from_date_obj = _parse_date(from_date)
@@ -68,7 +68,7 @@ async def get_option_chain(
 
 @register
 async def get_advanced_option_chain(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the underlying security"],
     contract_type: Annotated[
         str | None, "Type of contracts: CALL, PUT, or ALL (default)"
@@ -123,7 +123,7 @@ async def get_advanced_option_chain(
     Params: symbol, contract_type, strike_count, include_quotes, strategy (SINGLE/ANALYTICAL/etc.), interval, strike, strike_range (ITM/NTM/etc.), from/to_date, volatility/underlying_price/interest_rate/days_to_expiration (for ANALYTICAL), exp_month, option_type (STANDARD/NON_STANDARD/ALL).
     Limit data returned using strike_count and date parameters.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.options
 
     from_date_obj = _parse_date(from_date)
@@ -158,12 +158,12 @@ async def get_advanced_option_chain(
 
 @register
 async def get_option_expiration_chain(
-    ctx: Context,
+    ctx: SchwabContext,
     symbol: Annotated[str, "Symbol of the underlying security"],
 ) -> str:
     """
     Returns available option expiration dates for a symbol, without contract details. Lightweight call to find available cycles. Param: symbol.
     """
-    context = get_context(ctx)
+    context: SchwabServerContext = ctx.request_context.lifespan_context
     client = context.options
     return await call(client.get_option_expiration_chain, symbol)