Interactive Brokers MCP Server

"""Shared helpers for IBKR service objects."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from functools import wraps

from ib_async import IB


def ensure_connection(func):
    """
    Decorator to ensure IB connection before executing service methods.

    This decorator checks if the IB client is connected before calling the decorated method.
    If not connected, it raises a RuntimeError.

    Supports both synchronous and asynchronous methods.
    Can be combined with @allow_local_mode to skip check when mode="local".

    Usage:
        class PortfolioManager(IBService):
            @ensure_connection
            async def get_account_summary(self):
                # Async method implementation
                pass

            @ensure_connection
            def get_positions(self):
                # Sync method implementation
                pass
    """
    if asyncio.iscoroutinefunction(func):
        @wraps(func)
        async def async_wrapper(self, *args, **kwargs):
            """Wrapper for async methods that checks IB connection before executing."""
            # Skip check if ib is None (local mode)
            if self.ib is None:
                return await func(self, *args, **kwargs)
            if not self.ib.isConnected():
                raise RuntimeError("Not connected to Interactive Brokers")
            return await func(self, *args, **kwargs)
        return async_wrapper
    else:
        @wraps(func)
        def sync_wrapper(self, *args, **kwargs):
            """Wrapper for sync methods that checks IB connection before executing."""
            # Skip check if ib is None (local mode)
            if self.ib is None:
                return func(self, *args, **kwargs)
            if not self.ib.isConnected():
                raise RuntimeError("Not connected to Interactive Brokers")
            return func(self, *args, **kwargs)
        return sync_wrapper


@dataclass(slots=True)
class IBService:
    """Base class for IBKR service helpers that need the shared IB client.

    This base class provides the IB client instance to all service classes.
    Methods that require an active connection should be decorated with @ensure_connection.
    """

    ib: IB