company-ai-platform/app/modules/feishu/long_connection.py

import json
import logging
from socket import gethostname
from threading import Event, Thread
from time import monotonic
from typing import Any
from urllib.parse import urlsplit

from app.core.background.task_queue import enqueue_feishu_inbound_event
from app.core.config import get_settings
from app.core.constants import ActorValue
from app.core.database import SessionLocal
from app.application.feishu import FeishuEventService
from app.modules.feishu.constants import (
    FEISHU_DEFAULT_OPEN_API_DOMAIN,
    FeishuEventSource,
    FeishuEventTransport,
    FeishuResponseKey,
)
from app.modules.observability.constants import HeartbeatComponent, HeartbeatStatus
from app.modules.observability.service import ObservabilityService

logger = logging.getLogger(__name__)


def _sdk_domain(base_url: str) -> str:
    parsed = urlsplit(base_url)
    if not parsed.scheme or not parsed.netloc:
        return FEISHU_DEFAULT_OPEN_API_DOMAIN
    return f"{parsed.scheme}://{parsed.netloc}"


def _sdk_connection_is_open(client: Any) -> bool:
    """Return whether the SDK exposes a currently open WebSocket connection."""

    try:
        connection = getattr(client, "_conn", None)
    except Exception:
        return False
    if connection is None:
        return False

    try:
        state = getattr(connection, "state", None)
    except Exception:
        return False
    if state is not None:
        state_name = getattr(state, "name", None)
        state_text = str(state_name or state).strip().lower()
        if state_text == "open" or state_text.endswith(".open"):
            return True
        if any(
            marker in state_text
            for marker in ("connecting", "closing", "closed")
        ):
            return False

    try:
        closed = getattr(connection, "closed", None)
    except Exception:
        return False
    if isinstance(closed, bool):
        return not closed

    try:
        opened = getattr(connection, "open", None)
    except Exception:
        return False
    if isinstance(opened, bool):
        return opened

    try:
        if getattr(connection, "close_code", None) is not None:
            return False
        transport = getattr(connection, "transport", None)
        is_closing = getattr(transport, "is_closing", None)
        if callable(is_closing) and is_closing():
            return False
    except Exception:
        return False

    # A non-null SDK-private connection object is not enough evidence that the
    # WebSocket handshake completed or that the transport remains usable.
    return False


def _connection_heartbeat_status(client: Any) -> str:
    if _sdk_connection_is_open(client):
        return HeartbeatStatus.OK
    return HeartbeatStatus.DEGRADED


def _record_runtime_heartbeat(instance_id: str, status_value: str) -> None:
    db = SessionLocal()
    try:
        ObservabilityService(db).record_heartbeat(
            component=HeartbeatComponent.FEISHU_EVENTS,
            instance_id=instance_id,
            status_value=status_value,
            actor=ActorValue.FEISHU,
        )
    except Exception:
        db.rollback()
        logger.exception("Failed to record Feishu event process heartbeat")
    finally:
        db.close()


def _heartbeat_loop(
    stop_event: Event,
    instance_id: str,
    interval_seconds: int,
    client: Any,
) -> None:
    last_status: str | None = None
    next_heartbeat_at = 0.0
    while not stop_event.is_set():
        status_value = _connection_heartbeat_status(client)
        current_time = monotonic()
        if status_value != last_status or current_time >= next_heartbeat_at:
            _record_runtime_heartbeat(instance_id, status_value)
            last_status = status_value
            next_heartbeat_at = current_time + max(1, interval_seconds)
        stop_event.wait(1)


def _start_heartbeat_loop(client: Any) -> tuple[Event, Thread]:
    settings = get_settings()
    stop_event = Event()
    thread = Thread(
        target=_heartbeat_loop,
        args=(
            stop_event,
            gethostname(),
            settings.heartbeat_interval_seconds,
            client,
        ),
        name="feishu-events-heartbeat",
        daemon=True,
    )
    thread.start()
    return stop_event, thread


def _sdk_event_to_payload(event: Any) -> dict[str, Any]:
    try:
        from lark_oapi.core.json import JSON
    except ImportError as exc:
        raise RuntimeError("lark-oapi is required for Feishu long connection") from exc

    data = JSON.marshal(event)
    if not data:
        return {}
    return json.loads(data)


def _handle_message_event(event: Any) -> None:
    _handle_verified_sdk_event(event)


def _handle_app_ticket_event(event: Any) -> None:
    _handle_verified_sdk_event(event)


def _handle_verified_sdk_event(event: Any) -> None:
    payload = _sdk_event_to_payload(event)
    db = SessionLocal()
    try:
        acceptance = FeishuEventService(db).accept_verified_event(
            payload,
            source=FeishuEventSource.LONG_CONNECTION,
            auto_reply=True,
        )
    finally:
        db.close()
    if (
        acceptance.event_key is not None
        and acceptance.should_dispatch
        and get_settings().task_queue_enabled
    ):
        enqueue_feishu_inbound_event(
            acceptance.event_key,
            actor=ActorValue.FEISHU,
        )
    logger.info(
        "Accepted Feishu long connection event key=%s status=%s duplicate=%s",
        acceptance.event_key,
        acceptance.response.get(FeishuResponseKey.STATUS),
        bool(acceptance.response.get(FeishuResponseKey.DUPLICATE)),
    )


def run_long_connection() -> None:
    """Start the Feishu long connection client and block forever."""

    settings = get_settings()
    if settings.feishu_event_transport != FeishuEventTransport.LONG_CONNECTION:
        raise RuntimeError(
            "FEISHU_EVENT_TRANSPORT must be long_connection for this process"
        )
    if not settings.feishu_app_id or not settings.feishu_app_secret:
        raise RuntimeError("FEISHU_APP_ID and FEISHU_APP_SECRET are required")

    try:
        import lark_oapi as lark
    except ImportError as exc:
        raise RuntimeError("Install lark-oapi before starting Feishu long connection") from exc

    event_handler = (
        lark.EventDispatcherHandler.builder(
            settings.feishu_encrypt_key or "",
            settings.feishu_verification_token or "",
        )
        .register_p2_im_message_receive_v1(_handle_message_event)
        .register_p1_customized_event("app_ticket", _handle_app_ticket_event)
        .build()
    )
    client = lark.ws.Client(
        app_id=settings.feishu_app_id,
        app_secret=settings.feishu_app_secret,
        event_handler=event_handler,
        log_level=lark.LogLevel.WARNING,
        domain=_sdk_domain(settings.feishu_base_url),
    )
    logger.info("Starting Feishu long connection client")
    stop_event, heartbeat_thread = _start_heartbeat_loop(client)
    try:
        client.start()
    finally:
        stop_event.set()
        heartbeat_thread.join(timeout=1)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    run_long_connection()