Complete STDIO Server

"""
stdio_transport.py - STDIO transport for MCP servers.

This module builds on mcp_core.py from Lesson 4.
"""

import sys
import json
import logging

# Import from Lesson 4
from mcp_core import handle_request


# =============================================================================
# Section 1: Logging Setup (stderr only!)
# =============================================================================

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[logging.StreamHandler(sys.stderr)]  # CRITICAL: stderr only!
)
logger = logging.getLogger("mcp-stdio")

# =============================================================================
# Section 2: STDIO Transport
# =============================================================================

def write_response(response: dict) -> None:
    """Write a JSON-RPC response to stdout."""
    # Use compact JSON - no pretty printing for STDIO!
    json_str = json.dumps(response, separators=(',', ':'))
    sys.stdout.write(json_str + "\n")
    sys.stdout.flush()  # CRITICAL: Always flush!
    logger.debug(f"Sent response for id={response.get('id')}")


def run_stdio_server() -> None:
    """
    Run MCP server using STDIO transport.

    Reads JSON-RPC messages from stdin (one per line).
    Writes responses to stdout.
    Logs to stderr.
    """
    logger.info("STDIO MCP Server starting...")

    while True:
        try:
            # Read one line from stdin
            line = sys.stdin.readline()

            # EOF - client disconnected
            if not line:
                logger.info("EOF received, shutting down")
                break

            # Skip empty lines
            line = line.strip()
            if not line:
                continue

            # Parse JSON
            try:
                request = json.loads(line)
            except json.JSONDecodeError as e:
                logger.error(f"JSON parse error: {e}")
                continue

            # Use shared handler from mcp_core
            response = handle_request(request)

            # Send response (None for notifications)
            if response is not None:
                write_response(response)

        except KeyboardInterrupt:
            logger.info("Interrupted, shutting down")
            break
        except Exception as e:
            logger.error(f"Unexpected error: {e}", exc_info=True)

    logger.info("Server shutdown complete")

# =============================================================================
# Run if executed directly
# =============================================================================

if __name__ == "__main__":
    run_stdio_server()

"""
stdio_example.py - STDIO server using tools.py

Run: python stdio_example.py
Test: echo '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}' | python stdio_example.py
"""

from datetime import datetime

# Import from our modules
from mcp_core import router
from tools import ToolRegistry, register_tool_handlers, text_result
from stdio_transport import run_stdio_server

# =============================================================================
# Register tools
# =============================================================================

tool_registry = ToolRegistry()


@tool_registry.register(
    name="echo",
    description="Echo back a message",
    input_schema={
        "type": "object",
        "properties": {"message": {"type": "string", "description": "Message to echo"}},
        "required": ["message"]
    }
)
def echo(args: dict) -> dict:
    return text_result(args.get("message", ""))


@tool_registry.register(
    name="get_time",
    description="Get current server time",
    input_schema={"type": "object", "properties": {}, "required": []}
)
def get_time(args: dict) -> dict:
    return text_result(f"Current time: {datetime.now().isoformat()}")


@tool_registry.register(
    name="add_numbers",
    description="Add two numbers",
    input_schema={
        "type": "object",
        "properties": {
            "a": {"type": "number", "description": "First number"},
            "b": {"type": "number", "description": "Second number"}
        },
        "required": ["a", "b"]
    }
)
def add_numbers(args: dict) -> dict:
    a = args.get("a", 0)
    b = args.get("b", 0)
    return text_result(f"{a} + {b} = {a + b}")


# Register with the MCP router
register_tool_handlers(router, tool_registry)


# =============================================================================
# Run the server
# =============================================================================

if __name__ == "__main__":
    run_stdio_server()