diff --git a/server.py b/server.py index 0f7d15a..b595b3d 100644 --- a/server.py +++ b/server.py @@ -14,6 +14,21 @@ import schwab_scraper.unified_api as api # Initialize FastMCP mcp = FastMCP("SchwabScraper") +def serialize(obj: Any) -> str: + """Safely serialize Pydantic models or datclasses to JSON string.""" + if hasattr(obj, "model_dump_json"): + return obj.model_dump_json() + elif hasattr(obj, "model_dump"): + return json.dumps(obj.model_dump(), default=str) + elif isinstance(obj, list): + # Handle lists of models + return json.dumps([ + o.model_dump() if hasattr(o, "model_dump") else o + for o in obj + ], default=str) + return json.dumps(obj, default=str) + + @mcp.tool() async def get_session_status(debug: bool = False) -> str: """Get the current session status of the Schwab scraper. @@ -22,7 +37,7 @@ async def get_session_status(debug: bool = False) -> str: debug: Enable debug logging """ result = await api.get_session_status(debug=debug) - return json.dumps(result) + return serialize(result) @mcp.tool() async def list_accounts(debug: bool = False) -> str: @@ -32,7 +47,7 @@ async def list_accounts(debug: bool = False) -> str: debug: Enable debug logging """ result = await api.list_accounts(debug=debug) - return json.dumps(result) + return serialize(result) @mcp.tool() async def get_account_overview(account: Optional[str] = None, debug: bool = False) -> str: @@ -43,7 +58,7 @@ async def get_account_overview(account: Optional[str] = None, debug: bool = Fals debug: Enable debug logging """ result = await api.get_account_overview(account=account, debug=debug) - return json.dumps(result) + return serialize(result) @mcp.tool() async def get_positions(account: Optional[str] = None, include_non_equity: bool = False, debug: bool = False) -> str: @@ -55,7 +70,7 @@ async def get_positions(account: Optional[str] = None, include_non_equity: bool debug: Enable debug logging """ result = await api.get_positions(account=account, include_non_equity=include_non_equity, debug=debug) - return json.dumps(result) + return serialize(result) @mcp.tool() async def get_transactions( @@ -81,7 +96,7 @@ async def get_transactions( time_period=time_period, debug=debug ) - return json.dumps(result) + return serialize(result) @mcp.tool() async def get_morningstar_data(ticker: str, debug: bool = False) -> str: @@ -92,7 +107,7 @@ async def get_morningstar_data(ticker: str, debug: bool = False) -> str: debug: Enable debug logging """ result = await api.get_morningstar_data(ticker, debug=debug) - return json.dumps(result) + return serialize(result) @mcp.tool() async def upload_cookies(cookies_json: str) -> str: