ForcePilot/backend/package/yuxi/channels/adapters/urbit/session.py

from __future__ import annotations

import time
from collections import defaultdict

from yuxi.channels.models import ChannelMessage
from yuxi.utils.logging_config import logger


def resolve_thread_key(msg: ChannelMessage, default_agent_id: str = "ChatbotAgent") -> str:
    chat_type = msg.metadata.get("chat_type", "group")
    chat_id = msg.identity.channel_chat_id

    if chat_type == "direct":
        return f"agent:{default_agent_id}:urbit:direct:{chat_id}"

    group_name = msg.metadata.get("group_name", chat_id)
    ch_type = msg.metadata.get("urbit_resource_type", "chat")
    return f"agent:{default_agent_id}:urbit:group:{group_name}:{ch_type}"


def resolve_session_route(msg: ChannelMessage, default_agent_id: str = "ChatbotAgent") -> str:
    return resolve_thread_key(msg, default_agent_id)


_unsafe_sessions: defaultdict[str, set[str]] = defaultdict(set)
_session_last_access: dict[str, float] = {}
_SESSION_TTL_S = 3600


def detect_unsafe_session(msg: ChannelMessage) -> list[str]:
    chat_type = msg.metadata.get("chat_type", "group")
    if chat_type != "direct":
        return []

    chat_id = msg.identity.channel_chat_id
    participant_key = f"dm:{chat_id}"
    sender_ship = msg.metadata.get("urbit_ship", "")

    if not sender_ship:
        return []

    _unsafe_sessions[participant_key].add(sender_ship)
    _session_last_access[participant_key] = time.monotonic()
    participants = _unsafe_sessions[participant_key]

    if len(participants) > 2:
        logger.warning(
            f"[Urbit] Unsafe DM session detected: "
            f"{participant_key} has {len(participants)} "
            f"participants: {participants}"
        )
        return list(participants)

    return []


def get_unsafe_sessions() -> dict[str, set[str]]:
    return dict(_unsafe_sessions)


def clear_unsafe_sessions() -> None:
    _unsafe_sessions.clear()
    _session_last_access.clear()


def cleanup_expired_sessions() -> int:
    now = time.monotonic()
    expired = [
        key for key, ts in _session_last_access.items()
        if now - ts > _SESSION_TTL_S
    ]
    for key in expired:
        _unsafe_sessions.pop(key, None)
        _session_last_access.pop(key, None)
    if expired:
        logger.debug(f"[Urbit] Cleaned up {len(expired)} expired unsafe session records")
    return len(expired)