schwab-mcp-custom/server.py

import json
import os
from typing import Optional, Any

from fastmcp import FastMCP
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route, Mount
import uvicorn

# Import the unified API from the schwab_scraper dependency
import schwab_scraper.unified_api as api

# Initialize FastMCP
mcp = FastMCP("SchwabScraper")

@mcp.tool()
async def get_session_status(debug: bool = False) -> str:
    """Get the current session status of the Schwab scraper.
    
    Args:
        debug: Enable debug logging
    """
    result = await api.get_session_status(debug=debug)
    return json.dumps(result)

@mcp.tool()
async def list_accounts(debug: bool = False) -> str:
    """List all Schwab accounts.
    
    Args:
        debug: Enable debug logging
    """
    result = await api.list_accounts(debug=debug)
    return json.dumps(result)

@mcp.tool()
async def get_account_overview(account: Optional[str] = None, debug: bool = False) -> str:
    """Get the overview for a specific account.
    
    Args:
        account: Account summary or ID (optional)
        debug: Enable debug logging
    """
    result = await api.get_account_overview(account=account, debug=debug)
    return json.dumps(result)

@mcp.tool()
async def get_positions(account: Optional[str] = None, include_non_equity: bool = False, debug: bool = False) -> str:
    """Get positions for a specific account.
    
    Args:
        account: Account summary or ID (optional)
        include_non_equity: Whether to include non-equity positions
        debug: Enable debug logging
    """
    result = await api.get_positions(account=account, include_non_equity=include_non_equity, debug=debug)
    return json.dumps(result)

@mcp.tool()
async def get_transactions(
    account: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    time_period: Optional[str] = None,
    debug: bool = False
) -> str:
    """Get transaction history.
    
    Args:
        account: Account ID (optional)
        start_date: Start date for transactions (optional)
        end_date: End date for transactions (optional)
        time_period: Time period (e.g., '1D', '1M') (optional)
        debug: Enable debug logging
    """
    result = await api.get_transaction_history(
        account=account,
        start_date=start_date,
        end_date=end_date,
        time_period=time_period,
        debug=debug
    )
    return json.dumps(result)

@mcp.tool()
async def get_morningstar_data(ticker: str, debug: bool = False) -> str:
    """Get Morningstar data for a ticker.
    
    Args:
        ticker: Stock ticker symbol
        debug: Enable debug logging
    """
    result = await api.get_morningstar_data(ticker, debug=debug)
    return json.dumps(result)

@mcp.tool()
async def upload_cookies(cookies_json: str) -> str:
    """Upload session cookies to the server to assist with authentication.
    
    Args:
        cookies_json: JSON string of cookies exported from a browser (Playwright format)
    """
    try:
        # Validate JSON
        cookies = json.loads(cookies_json)
        # Write to cookies.json
        with open("cookies.json", "w") as f:
            json.dump(cookies, f)
        return json.dumps({"status": "success", "message": "cookies.json updated successfully"})
    except Exception as e:
        return json.dumps({"status": "error", "message": str(e)})

@mcp.tool()
async def api_call(endpoint: str, method: str = "GET", params: str = "{}") -> str:
    """Executes a raw API call to the Schwab service (Dummy implementation).
    
    Refer to the 'api-reference' resource for available endpoints and parameters.
    
    Args:
        endpoint: The API path
        method: HTTP method (GET, POST, etc.)
        params: JSON string of parameters/body
    """
    return json.dumps({"status": "not_implemented", "message": "API pass-through not supported for scraper"})

@mcp.resource("service://api-reference")
def get_api_docs() -> str:
    """Returns the API documentation for using the 'api_call' tool."""
    return "Schwab Scraper MCP Server - Unified API Documentation\n\nThis server provides tools to interact with Schwab accounts via scraping. The 'api_call' tool is a placeholder."

async def health(request):
    """Health check endpoint."""
    return JSONResponse({"status": "ok"})

# Create the Starlette application
mcp_app = mcp.http_app()
app = Starlette(
    routes=[
        Route("/health", health),
        Mount("/", app=mcp_app)
    ],
    lifespan=mcp_app.lifespan
)

if __name__ == "__main__":
    port = int(os.getenv("PORT", 8160))
    uvicorn.run(app, host="0.0.0.0", port=port)