diff --git a/server.py b/server.py index 53d88f7..0322635 100644 --- a/server.py +++ b/server.py @@ -1,7 +1,9 @@ +import io import json import logging import os import time +from contextlib import contextmanager from typing import Optional, Any, Tuple from fastmcp import FastMCP @@ -12,6 +14,45 @@ import uvicorn import schwab_scraper.unified_api as api + +# --------------------------------------------------------------------------- +# Log capture helper — scraper logs go to stderr which is invisible in MCP +# stdio mode. This captures them so we can return them in the response. +# --------------------------------------------------------------------------- +@contextmanager +def capture_logs(logger_name: str = "schwab_scraper", level: int = logging.DEBUG): + """Context manager that captures log output to a string buffer. + + Yields the buffer so callers can read captured logs after the block. + """ + logger = logging.getLogger(logger_name) + # Ensure the logger will actually process messages at this level + old_level = logger.level + if old_level > level: + logger.setLevel(level) + + buf = io.StringIO() + handler = logging.StreamHandler(buf) + handler.setLevel(level) + # Match the format used by the scraper + handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")) + logger.addHandler(handler) + + try: + yield buf + finally: + logger.removeHandler(handler) + logger.setLevel(old_level) + + +def _enrich_with_logs(result: dict, log_buffer: io.StringIO, debug: bool) -> dict: + """Attach captured logs to a result dict when debug=True or on error.""" + logs = log_buffer.getvalue() + if logs and (debug or not result.get("success", False)): + result["logs"] = logs + return result + + # --------------------------------------------------------------------------- # Monkey-patch mcp.shared.session.RequestResponder to work around a # cancellation race in mcp==1.27.0 (github.com/modelcontextprotocol/ @@ -177,9 +218,11 @@ async def login( "data": None, }) - result = await api.login(username=username, password=password, debug=debug) - success = result.get("success", False) - login_manager.record_attempt(success) + with capture_logs(level=logging.DEBUG if debug else logging.INFO) as log_buf: + result = await api.login(username=username, password=password, debug=debug) + success = result.get("success", False) + login_manager.record_attempt(success) + result = _enrich_with_logs(result, log_buf, debug) return serialize(result) @@ -190,7 +233,9 @@ async def refresh_session(debug: bool = False) -> str: Args: debug: Enable debug logging """ - result = await api.refresh_session(debug=debug) + with capture_logs(level=logging.DEBUG if debug else logging.INFO) as log_buf: + result = await api.refresh_session(debug=debug) + result = _enrich_with_logs(result, log_buf, debug) return serialize(result)