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

refactor(server): share schwab client via lifespan context

Author: jkoelker

src/schwab_mcp/context.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, cast
+
+from schwab.client import AsyncClient
+
+if TYPE_CHECKING:
+    from schwab_mcp.tools._protocols import (
+        AccountClient,
+        OptionsClient,
+        OrdersClient,
+        PriceHistoryClient,
+        QuotesClient,
+        ToolsClient,
+        TransactionsClient,
+    )
+else:  # pragma: no cover - runtime only
+    AccountClient = OptionsClient = OrdersClient = PriceHistoryClient = QuotesClient = (
+        ToolsClient
+    ) = TransactionsClient = Any
+
+
+@dataclass(slots=True)
+class SchwabServerContext:
+    """Typed application context shared via FastMCP lifespan."""
+
+    client: AsyncClient
+    tools: ToolsClient = field(init=False)
+    accounts: AccountClient = field(init=False)
+    price_history: PriceHistoryClient = field(init=False)
+    options: OptionsClient = field(init=False)
+    orders: OrdersClient = field(init=False)
+    quotes: QuotesClient = field(init=False)
+    transactions: TransactionsClient = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.tools = cast(ToolsClient, self.client)
+        self.accounts = cast(AccountClient, self.client)
+        self.price_history = cast(PriceHistoryClient, self.client)
+        self.options = cast(OptionsClient, self.client)
+        self.orders = cast(OrdersClient, self.client)
+        self.quotes = cast(QuotesClient, self.client)
+        self.transactions = cast(TransactionsClient, self.client)
+
+
+__all__ = ["SchwabServerContext"]

src/schwab_mcp/server.py

@@ -1,13 +1,41 @@
 from __future__ import annotations
 
+import logging
 import sys
-from typing import Optional
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import AsyncContextManager, Callable, Optional
 
 import mcp.types as types
 from mcp.server.fastmcp import FastMCP
 from schwab.client import AsyncClient
 
 from schwab_mcp.tools import register_tools
+from schwab_mcp.context import SchwabServerContext
+
+
+logger = logging.getLogger(__name__)
+
+
+def _client_lifespan(
+    client: AsyncClient,
+) -> Callable[[FastMCP], AsyncContextManager[SchwabServerContext]]:
+    """Create a FastMCP lifespan context that exposes the Schwab async client."""
+
+    @asynccontextmanager
+    async def lifespan(_: FastMCP) -> AsyncGenerator[SchwabServerContext, None]:
+        context = SchwabServerContext(client=client)
+        try:
+            yield context
+        finally:
+            try:
+                await client.close_async_session()
+            except Exception:
+                logger.exception(
+                    "Failed to close Schwab async client session during shutdown."
+                )
+
+    return lifespan
 
 
 class SchwabMCPServer:
@@ -19,7 +47,7 @@ def __init__(
         client: AsyncClient,
         jesus_take_the_wheel: bool = False,
     ) -> None:
-        self._server = FastMCP(name=name)
+        self._server = FastMCP(name=name, lifespan=_client_lifespan(client))
         register_tools(self._server, client, allow_write=jesus_take_the_wheel)
 
     async def run(self) -> None:

src/schwab_mcp/tools/__init__.py

@@ -15,10 +15,10 @@
 from schwab_mcp.tools import quotes as _quotes  # noqa: F401
 from schwab_mcp.tools import transactions as _txns  # noqa: F401
 
+
 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)
-    setattr(server, "_schwab_client", client)
 
     for func in iter_registered_tools():
         if getattr(func, "_write", False) and not allow_write:

src/schwab_mcp/tools/_protocols.py

@@ -18,21 +18,37 @@ class PriceHistoryNamespace(Protocol):
 class PriceHistoryClient(Protocol):
     PriceHistory: PriceHistoryNamespace
 
-    def get_advanced_price_history(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_advanced_price_history(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_price_history_every_minute(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_price_history_every_minute(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_price_history_every_five_minutes(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_price_history_every_five_minutes(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_price_history_every_ten_minutes(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_price_history_every_ten_minutes(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_price_history_every_fifteen_minutes(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_price_history_every_fifteen_minutes(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_price_history_every_thirty_minutes(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_price_history_every_thirty_minutes(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_price_history_every_day(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_price_history_every_day(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_price_history_every_week(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_price_history_every_week(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
 
 class OptionsNamespace(Protocol):
@@ -48,7 +64,9 @@ class OptionsClient(Protocol):
 
     def get_option_chain(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
 
-    def get_option_expiration_chain(self, symbol: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_option_expiration_chain(
+        self, symbol: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
 
 class QuoteFieldsNamespace(Protocol):
@@ -122,11 +140,17 @@ class OrderNamespace(Protocol):
 class OrdersClient(Protocol):
     Order: OrderNamespace
 
-    def get_orders_for_account(self, account_hash: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_orders_for_account(
+        self, account_hash: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def get_order(self, order_id: str, account_hash: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_order(
+        self, order_id: str, account_hash: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
-    def cancel_order(self, order_id: str, account_hash: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def cancel_order(
+        self, order_id: str, account_hash: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...
 
     def place_order(self, account_hash: str, **kwargs: Any) -> Awaitable[Any]: ...
 
@@ -144,4 +168,6 @@ class TransactionsClient(Protocol):
 
     def get_transactions(self, account_hash: str, **kwargs: Any) -> Awaitable[Any]: ...
 
-    def get_transaction(self, account_hash: str, transaction_id: str, **kwargs: Any) -> Awaitable[Any]: ...
+    def get_transaction(
+        self, account_hash: str, transaction_id: str, **kwargs: Any
+    ) -> Awaitable[Any]: ...

src/schwab_mcp/tools/account.py

@@ -5,7 +5,7 @@
 from mcp.server.fastmcp import Context
 
 from schwab_mcp.tools.registry import register
-from schwab_mcp.tools.utils import call, get_account_client
+from schwab_mcp.tools.utils import call, get_context
 
 
 @register
@@ -15,8 +15,8 @@ async def get_account_numbers(
     """
     Returns mapping of account IDs to account hashes. Hashes required for account-specific calls. Use first.
     """
-    client = get_account_client(ctx)
-    return await call(client.get_account_numbers)
+    context = get_context(ctx)
+    return await call(context.accounts.get_account_numbers)
 
 
 @register
@@ -26,8 +26,8 @@ async def get_accounts(
     """
     Returns balances/info for all linked accounts (funds, cash, margin). Does not return hashes; use get_account_numbers first.
     """
-    client = get_account_client(ctx)
-    return await call(client.get_accounts)
+    context = get_context(ctx)
+    return await call(context.accounts.get_accounts)
 
 
 @register
@@ -37,8 +37,11 @@ async def get_accounts_with_positions(
     """
     Returns balances, info, and positions (holdings, cost, gain/loss) for all linked accounts. Does not return hashes; use get_account_numbers first.
     """
-    client = get_account_client(ctx)
-    return await call(client.get_accounts, fields=[client.Account.Fields.POSITIONS])
+    context = get_context(ctx)
+    return await call(
+        context.accounts.get_accounts,
+        fields=[context.accounts.Account.Fields.POSITIONS],
+    )
 
 
 @register
@@ -49,8 +52,8 @@ async def get_account(
     """
     Returns balance/info for a specific account via account_hash (from get_account_numbers). Includes funds, cash, margin info.
     """
-    client = get_account_client(ctx)
-    return await call(client.get_account, account_hash)
+    context = get_context(ctx)
+    return await call(context.accounts.get_account, account_hash)
 
 
 @register
@@ -61,9 +64,11 @@ async def get_account_with_positions(
     """
     Returns balance, info, and positions for a specific account via account_hash. Includes holdings, quantity, cost basis, unrealized gain/loss.
     """
-    client = get_account_client(ctx)
+    context = get_context(ctx)
     return await call(
-        client.get_account, account_hash, fields=[client.Account.Fields.POSITIONS]
+        context.accounts.get_account,
+        account_hash,
+        fields=[context.accounts.Account.Fields.POSITIONS],
     )
 
 
@@ -74,5 +79,5 @@ async def get_user_preferences(
     """
     Returns user preferences (nicknames, display settings, notifications) for all linked accounts.
     """
-    client = get_account_client(ctx)
-    return await call(client.get_user_preferences)
+    context = get_context(ctx)
+    return await call(context.accounts.get_user_preferences)

src/schwab_mcp/tools/history.py

@@ -6,7 +6,7 @@
 from mcp.server.fastmcp import Context
 
 from schwab_mcp.tools.registry import register
-from schwab_mcp.tools.utils import call, get_price_history_client
+from schwab_mcp.tools.utils import call, get_context
 
 
 def _parse_iso_datetime(value: str | None) -> datetime.datetime | None:
@@ -17,7 +17,9 @@ def _parse_iso_datetime(value: str | None) -> datetime.datetime | None:
 async def get_advanced_price_history(
     ctx: Context,
     symbol: Annotated[str, "Symbol of the security"],
-    period_type: Annotated[str | None, "Period type: DAY, MONTH, YEAR, YEAR_TO_DATE"] = None,
+    period_type: Annotated[
+        str | None, "Period type: DAY, MONTH, YEAR, YEAR_TO_DATE"
+    ] = None,
     period: Annotated[
         str | None,
         (
@@ -26,7 +28,8 @@ async def get_advanced_price_history(
         ),
     ] = None,
     frequency_type: Annotated[
-        str | None, "Frequency type: MINUTE (for DAY), DAILY/WEEKLY (for MONTH/YTD), DAILY/WEEKLY/MONTHLY (for YEAR)"
+        str | None,
+        "Frequency type: MINUTE (for DAY), DAILY/WEEKLY (for MONTH/YTD), DAILY/WEEKLY/MONTHLY (for YEAR)",
     ] = None,
     frequency: Annotated[
         int | str | None,
@@ -57,16 +60,15 @@ async def get_advanced_price_history(
       YEAR_TO_DATE: DAILY, WEEKLY (default)
     Dates must be in ISO format.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)
 
     # Normalize enum-like strings
     period_type_enum = (
-        client.PriceHistory.PeriodType[period_type.upper()]
-        if period_type
-        else None
+        client.PriceHistory.PeriodType[period_type.upper()] if period_type else None
     )
     period_enum = client.PriceHistory.Period[period.upper()] if period else None
     frequency_type_enum = (
@@ -109,7 +111,8 @@ async def get_price_history_every_minute(
     """
     Get OHLCV price history per minute. For detailed intraday analysis. Max 48 days history. Dates ISO format.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)
@@ -140,7 +143,8 @@ 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.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)
@@ -171,7 +175,8 @@ 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.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)
@@ -202,7 +207,8 @@ 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.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)
@@ -233,7 +239,8 @@ 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.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)
@@ -270,7 +277,8 @@ 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.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)
@@ -301,7 +309,8 @@ 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.
     """
-    client = get_price_history_client(ctx)
+    context = get_context(ctx)
+    client = context.price_history
 
     start_dt = _parse_iso_datetime(start_datetime)
     end_dt = _parse_iso_datetime(end_datetime)