Modules/_core.py

"""Shared helpers: call logging and a global rate limiter.

The limiter is deliberately global (not per IP). The Space exists to serve a
chat app a few calls per user turn; it must never become a high-volume proxy.
"""

from __future__ import annotations

import json
import os
import sys
import threading
import time
from collections import deque
from typing import Any


def _env_int(name: str, default: int) -> int:
    try:
        value = int(os.getenv(name, "") or default)
        return value if value > 0 else default
    except ValueError:
        return default


class RateLimiter:
    """Sliding-window limiter. Waits briefly for a slot, then gives up.

    Giving up (instead of sleeping for up to a minute like the old limiter)
    keeps worker threads free and hands the model a clear error string.
    """

    def __init__(self, per_minute: int, max_wait_seconds: float = 8.0) -> None:
        self.per_minute = per_minute
        self.max_wait = max_wait_seconds
        self._calls: deque[float] = deque()
        self._lock = threading.Lock()

    def acquire(self) -> bool:
        deadline = time.monotonic() + self.max_wait
        while True:
            with self._lock:
                now = time.monotonic()
                while self._calls and now - self._calls[0] >= 60.0:
                    self._calls.popleft()
                if len(self._calls) < self.per_minute:
                    self._calls.append(now)
                    return True
                wait = 60.0 - (now - self._calls[0])
            if time.monotonic() + wait > deadline:
                return False
            time.sleep(min(wait, 1.0))


SEARCH_LIMITER = RateLimiter(_env_int("TOOLS_SEARCH_PER_MINUTE", 30))
FETCH_LIMITER = RateLimiter(_env_int("TOOLS_FETCH_PER_MINUTE", 60))

RATE_LIMIT_MESSAGE = (
    "Error: this tool is receiving too many requests right now (global limit {limit}/minute). "
    "Wait about a minute before calling it again."
)


def _compact(val: Any) -> Any:
    if isinstance(val, (str, int, float, bool)) or val is None:
        return val if not isinstance(val, str) else val[:300]
    return repr(val)[:120]


def log_call_start(func_name: str, **kwargs: Any) -> None:
    try:
        payload = json.dumps({k: _compact(v) for k, v in kwargs.items()}, ensure_ascii=False)
        print(f"[TOOL CALL] {func_name} inputs: {payload[:800]}", flush=True, file=sys.__stdout__)
    except Exception:
        pass


def log_call_end(func_name: str, output_desc: str) -> None:
    try:
        print(f"[TOOL RESULT] {func_name} output: {output_desc[:500]}", flush=True, file=sys.__stdout__)
    except Exception:
        pass


def to_int(value: Any, default: int) -> int:
    """Coerce MCP numbers (which may arrive as floats or strings) to int."""
    try:
        return int(float(value))
    except (TypeError, ValueError):
        return default