trader-mcp/src/trader_mcp/strategy_sizing.py

from __future__ import annotations

from typing import Any


def _call_context_suggest_order_amount(context: Any, kwargs: dict[str, Any]) -> float | None:
    if not hasattr(context, "suggest_order_amount"):
        return None

    variants: list[dict[str, Any]] = []
    drop_sets = (
        (),
        ("quote_notional",),
        ("dust_collect",),
        ("order_size",),
        ("quote_notional", "dust_collect"),
        ("quote_notional", "order_size"),
        ("dust_collect", "order_size"),
        ("quote_notional", "dust_collect", "order_size"),
    )
    for drop_keys in drop_sets:
        variant = {key: value for key, value in kwargs.items() if key not in drop_keys}
        if variant not in variants:
            variants.append(variant)

    last_error: TypeError | None = None
    for variant in variants:
        try:
            return float(context.suggest_order_amount(**variant) or 0.0)
        except TypeError as exc:
            last_error = exc

    if last_error is not None:
        raise last_error
    return None


def suggest_quote_sized_amount(
    context: Any,
    *,
    side: str,
    price: float,
    levels: int,
    min_notional: float,
    fee_rate: float,
    order_notional_quote: float = 0.0,
    max_order_notional_quote: float = 0.0,
    dust_collect: bool = False,
    order_size: float = 0.0,
) -> float:
    side = str(side or "").strip().lower()
    price = float(price or 0.0)
    levels = int(levels or 0)
    min_notional = max(float(min_notional or 0.0), 0.0)
    fee_rate = max(float(fee_rate or 0.0), 0.0)
    order_notional_quote = max(float(order_notional_quote or 0.0), 0.0)
    max_order_notional_quote = max(float(max_order_notional_quote or 0.0), 0.0)
    order_size = max(float(order_size or 0.0), 0.0)

    if levels <= 0 or price <= 0 or side not in {"buy", "sell"}:
        return 0.0

    kwargs = {
        "side": side,
        "price": price,
        "levels": levels,
        "min_notional": min_notional,
        "fee_rate": fee_rate,
        "quote_notional": order_notional_quote,
        "max_notional_per_order": max_order_notional_quote,
        "dust_collect": bool(dust_collect),
        "order_size": order_size,
    }
    amount = _call_context_suggest_order_amount(context, kwargs)
    if amount is not None:
        if amount <= 0:
            return 0.0
        if side == "buy":
            if min_notional > 0 and (amount * price * (1 + fee_rate)) < min_notional:
                return 0.0
        else:
            if min_notional > 0 and (amount * price) < min_notional:
                return 0.0
        return amount

    if order_notional_quote <= 0:
        return 0.0

    effective_quote = order_notional_quote
    if max_order_notional_quote > 0:
        effective_quote = min(effective_quote, max_order_notional_quote)

    if side == "buy":
        min_quote_needed = min_notional * (1 + fee_rate)
        if min_notional > 0 and effective_quote < min_quote_needed:
            return 0.0
        amount = effective_quote / (price * (1 + fee_rate))
    else:
        amount = effective_quote / price
    min_amount = (min_notional / price) if side == "sell" and min_notional > 0 else 0.0
    if min_amount > 0 and amount < min_amount:
        return 0.0
    return max(amount, 0.0)


def cap_amount_to_balance_target(
    *,
    suggested_amount: float,
    side: str,
    price: float,
    fee_rate: float,
    balance_target: float,
    base_available: float,
    counter_available: float,
    min_notional: float = 0.0,
) -> float:
    amount = max(float(suggested_amount or 0.0), 0.0)
    side = str(side or "").strip().lower()
    price = float(price or 0.0)
    fee_rate = max(float(fee_rate or 0.0), 0.0)
    target = min(max(float(balance_target if balance_target is not None else 1.0), 0.0), 1.0)
    base_available = max(float(base_available or 0.0), 0.0)
    counter_available = max(float(counter_available or 0.0), 0.0)
    min_notional = max(float(min_notional or 0.0), 0.0)

    if amount <= 0 or price <= 0 or side not in {"buy", "sell"}:
        return 0.0

    def meets_min_notional(candidate: float) -> bool:
        if min_notional <= 0:
            return True
        if side == "buy":
            return (candidate * price * (1 + fee_rate)) >= min_notional
        return (candidate * price) >= min_notional

    if side == "buy":
        max_amount = counter_available / (price * (1 + fee_rate)) if counter_available > 0 else 0.0
    else:
        max_amount = base_available

    amount = min(amount, max(max_amount, 0.0))
    if amount <= 0:
        return 0.0
    if not meets_min_notional(amount):
        return 0.0
    if target >= 1.0:
        return amount

    base_value = base_available * price
    total_value = base_value + counter_available
    if total_value <= 0:
        return 0.0

    capped_amount = amount
    if side == "buy":
        remaining_quote = (target * total_value) - base_value
        if remaining_quote <= 0:
            return 0.0
        target_amount = remaining_quote / (price * (1 + target * fee_rate))
        capped_amount = min(amount, target_amount)
    else:
        target_base_ratio = 1.0 - target
        remaining_quote = base_value - (target_base_ratio * total_value)
        if remaining_quote <= 0:
            return 0.0
        denominator = price * max(1.0 - (target_base_ratio * fee_rate), 1e-12)
        target_amount = remaining_quote / denominator
        capped_amount = min(amount, target_amount)

    if capped_amount <= 0:
        return 0.0
    if not meets_min_notional(capped_amount):
        return 0.0
    return max(capped_amount, 0.0)


def suggest_rebalance_amount(
    *,
    side: str,
    price: float,
    fee_rate: float,
    base_available: float,
    counter_available: float,
    target_ratio: float,
    step_ratio: float,
    balance_tolerance: float,
    min_order_notional_quote: float = 0.0,
    max_order_notional_quote: float = 0.0,
) -> float:
    side = str(side or "").strip().lower()
    price = float(price or 0.0)
    fee_rate = max(float(fee_rate or 0.0), 0.0)
    base_available = max(float(base_available or 0.0), 0.0)
    counter_available = max(float(counter_available or 0.0), 0.0)
    target_ratio = min(max(float(target_ratio or 0.0), 0.0), 1.0)
    step_ratio = max(float(step_ratio or 0.0), 0.0)
    balance_tolerance = max(float(balance_tolerance or 0.0), 0.0)
    min_order_notional_quote = max(float(min_order_notional_quote or 0.0), 0.0)
    max_order_notional_quote = max(float(max_order_notional_quote or 0.0), 0.0)

    if side not in {"buy", "sell"} or price <= 0:
        return 0.0

    base_value = base_available * price
    total_value = base_value + counter_available
    if total_value <= 0:
        return 0.0

    current_ratio = base_value / total_value
    drift = abs(current_ratio - target_ratio)
    if drift <= balance_tolerance:
        return 0.0

    target_quote = total_value * min(drift, step_ratio)
    if min_order_notional_quote > 0:
        target_quote = max(target_quote, min_order_notional_quote)
    if max_order_notional_quote > 0:
        target_quote = min(target_quote, max_order_notional_quote)
    if target_quote <= 0:
        return 0.0
    if min_order_notional_quote > 0 and target_quote < min_order_notional_quote:
        return 0.0

    if side == "buy":
        max_affordable_quote = counter_available
        if max_affordable_quote < min_order_notional_quote * (1 + fee_rate):
            return 0.0
        capped_quote = min(target_quote, max_affordable_quote)
        amount = capped_quote / (price * (1 + fee_rate))
    else:
        max_affordable_quote = base_available * price
        if max_affordable_quote < min_order_notional_quote:
            return 0.0
        capped_quote = min(target_quote, max_affordable_quote)
        amount = capped_quote / (price * (1 + fee_rate))

    if amount <= 0:
        return 0.0
    return max(amount, 0.0)