futu_mcp/

state.rs

//! 共享状态：网关连接 + 订阅状态 + 授权

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::Ordering;

use anyhow::{Context, Result, anyhow};
use futu_auth::{KeyRecord, KeyStore, RuntimeCounters};
use futu_core::qot_symbol;
use futu_net::client::{ClientConfig, FutuClient, ReconnectingClient};
use futu_net::reconnect::ReconnectPolicy;
use futu_qot::types::Security;
use rmcp::{RoleServer, service::Peer};
use tokio::sync::{Mutex, Semaphore};

mod push_filter;
mod push_subscribers;
pub(crate) use push_subscribers::{LegacyPushServiceLease, PushDeliveryTarget};
#[cfg(test)]
pub(crate) use push_subscribers::{MAX_MODERN_PUSH_HANDLES, MODERN_PUSH_QUEUE_CAPACITY};
pub(crate) use push_subscribers::{parse_push_resource_uri, push_resource_uri};
#[cfg(test)]
mod tests;

use push_filter::{TradePushDecode, classify_trade_push, trd_market_int_to_str};
#[cfg(test)]
use push_filter::{
    extract_acc_id_and_market_from_push, is_trade_push_proto_id, subscriber_should_receive,
    subscriber_should_receive_with_market, subscriber_visible_to_caller,
};
use push_subscribers::{PushSubscriber, SubscriberDelivery};

use crate::qot_sdk_adapter;

const MCP_CONNECT_TOTAL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3);
const MCP_CONNECT_RETRY_DELAY: std::time::Duration = std::time::Duration::from_millis(200);
pub(crate) const LEGACY_PUSH_INFO_LEVEL_RANK: u8 = 1;

/// v1.4.38 Phase 5 helper: bytes → base64 (用于 push body 安全包进 JSON)
fn base64_encode_bytes(bytes: &[u8]) -> String {
    use base64::Engine as _;
    base64::engine::general_purpose::STANDARD.encode(bytes)
}

struct PendingPushDelivery {
    target: SubscriberDelivery,
    data: serde_json::Value,
    session_id: String,
    owner_key_id: Option<String>,
    proto_id: u32,
}

/// Single authorization admission shared by legacy logging and modern resource
/// delivery. A caller must pass this function before either target is created,
/// so decode-failed restricted trade bodies and current scope revocations can
/// never enter a modern queue.
fn push_delivery_is_authorized(
    subscriber: &PushSubscriber,
    key_store: &KeyStore,
    filter_registry: &futu_auth_pipeline::FilterRegistry,
    decode_result: &TradePushDecode,
    event_type: &'static str,
    push_acc_id: Option<u64>,
    push_trd_market_str: Option<&'static str>,
    proto_id: u32,
) -> bool {
    let (allowed_acc_ids, allowed_markets) = match &subscriber.delivery {
        SubscriberDelivery::LegacyPeer(_, _) => (
            subscriber.allowed_acc_ids_snapshot.clone(),
            subscriber.allowed_markets_snapshot.clone(),
        ),
        SubscriberDelivery::ModernResource(_) => {
            let Some(owner) = subscriber.owner_key_id.as_deref() else {
                return false;
            };
            let Some(current) = key_store
                .get_by_id_for_current_machine(owner)
                .filter(|record| {
                    !record.is_expired(chrono::Utc::now())
                        && record.scopes.contains(&futu_auth::Scope::AccRead)
                })
            else {
                return false;
            };
            (
                current.allowed_acc_ids.clone(),
                current.allowed_markets.clone(),
            )
        }
    };
    if matches!(decode_result, TradePushDecode::DecodeFailed) {
        let restricted = allowed_acc_ids
            .as_ref()
            .is_some_and(|allowed| !allowed.is_empty());
        if restricted {
            let key_id = subscriber.owner_key_id.as_deref().unwrap_or("<none>");
            futu_auth::metrics::bump_ws_filtered("trade_decode_failed", key_id);
            tracing::warn!(
                proto_id,
                key_id,
                "MCP trade push body decode failed; dropped before legacy/modern delivery for restricted key"
            );
            return false;
        }
    }
    let sub_state = (!subscriber.acc_ids.is_empty()).then_some(&subscriber.acc_ids);
    let ctx = futu_auth_pipeline::PushEventCtx {
        event_type,
        event_acc: push_acc_id,
        allowed_acc_ids: allowed_acc_ids.as_ref(),
        sub_state,
        event_trd_market: push_trd_market_str,
        allowed_markets: allowed_markets.as_ref(),
    };
    if filter_registry.should_drop_event(&ctx) {
        let key_id = subscriber.owner_key_id.as_deref().unwrap_or("<none>");
        futu_auth::metrics::bump_ws_filtered("trade_market", key_id);
        return false;
    }
    true
}

#[allow(deprecated)]
async fn notify_legacy_push(
    peer: &Peer<RoleServer>,
    minimum_level: &std::sync::atomic::AtomicU8,
    data: serde_json::Value,
) -> Result<(), rmcp::service::ServiceError> {
    if minimum_level.load(Ordering::Acquire) > LEGACY_PUSH_INFO_LEVEL_RANK {
        return Ok(());
    }
    let params =
        rmcp::model::LoggingMessageNotificationParam::new(rmcp::model::LoggingLevel::Info, data)
            .with_logger("futu_push");
    peer.notify_logging_message(params).await
}

/// MCP server 运行时状态
#[derive(Clone)]
pub struct ServerState {
    /// [`Inner`] 共享可变状态（gateway 地址 + 懒加载的 [`FutuClient`]）
    inner: Arc<Mutex<Inner>>,
    /// 是否启用交易写工具（place/modify/cancel）。默认 false。旧开关，仅当
    /// `key_store.is_configured() == false` 时生效。
    enable_trading: bool,
    /// 是否允许对 real 环境下单。默认 false。旧开关，同上。
    allow_real_trading: bool,
    /// keys.json 加载的 KeyStore。`is_configured()` 为 true 时走 scope 授权模式。
    key_store: Arc<KeyStore>,
    /// 调用方传入的 API Key 对应的记录；None 表示未提供 key。
    authed_key: Option<Arc<KeyRecord>>,
    /// 交易密码所属登录账号。用于 `futu_unlock_trade` 从账号级 keychain
    /// `trade-password.<login-account>` 读取密码；None 时走 legacy/global/env 兼容路径。
    trade_pwd_account: Option<String>,
    /// 限额运行时（日累计计数器）
    counters: Arc<RuntimeCounters>,
    /// v1.4.38 Phase 5: MCP push 订阅者注册表（session_uuid → subscriber）。
    /// `futu_sub_acc_push` 工具在 HTTP 模式下调用时注册当前 session，daemon
    /// push 到 MCP 后按 acc_id filter 向注册的 peer 发
    /// `notify_logging_message`（server-initiated notification）。
    push_subscribers: Arc<Mutex<HashMap<String, PushSubscriber>>>,
    /// Modern resources retain bounded queues for up to 4h. A semaphore makes
    /// handle cardinality admission atomic and lets the tool fail before any
    /// daemon side effect. The fixed local limit is a memory-safety boundary,
    /// not backend/config data; raise it only after measuring worst-case push
    /// payload memory and concurrent client demand.
    modern_push_slots: Arc<Semaphore>,
}

/// ServerState 内部可变部分，加锁存放 gateway 地址 + 懒加载的 [`FutuClient`]。
struct Inner {
    /// 网关 TCP 地址（如 `127.0.0.1:11111`）
    gateway: String,
    /// 懒加载的底层连接；首次调用 [`ServerState::client`] 时建立，后续复用
    client: Option<Arc<FutuClient>>,
}

impl ServerState {
    /// 创建默认 state：`enable_trading=false` / `allow_real_trading=false` /
    /// 空 [`KeyStore`] / 无 authed_key。使用 `with_*` 链式方法注入额外能力。
    pub fn new(gateway: String) -> Self {
        Self {
            inner: Arc::new(Mutex::new(Inner {
                gateway,
                client: None,
            })),
            enable_trading: false,
            allow_real_trading: false,
            key_store: Arc::new(KeyStore::empty()),
            authed_key: None,
            trade_pwd_account: None,
            counters: Arc::new(RuntimeCounters::new()),
            push_subscribers: Arc::new(Mutex::new(HashMap::new())),
            modern_push_slots: Arc::new(Semaphore::new(push_subscribers::MAX_MODERN_PUSH_HANDLES)),
        }
    }

    /// 启用交易写工具（构造器式链式设置）
    pub fn with_trading(mut self, enable_trading: bool, allow_real_trading: bool) -> Self {
        self.enable_trading = enable_trading;
        self.allow_real_trading = allow_real_trading;
        self
    }

    /// 设置 KeyStore（新授权模式）
    pub fn with_key_store(mut self, store: Arc<KeyStore>) -> Self {
        self.key_store = store;
        self
    }

    /// 设置已通过验证的 API Key 记录
    pub fn with_authed_key(mut self, key: Option<Arc<KeyRecord>>) -> Self {
        self.authed_key = key;
        self
    }

    /// 设置交易密码所属登录账号（MCP 只连 gateway，本身无法可靠推断 daemon
    /// 的 login account；由 CLI/env/config 显式注入）。
    pub fn with_trade_pwd_account(mut self, account: Option<String>) -> Self {
        self.trade_pwd_account = account;
        self
    }

    /// 是否启用了 scope 授权模式
    pub fn is_scope_mode(&self) -> bool {
        self.key_store.is_configured()
    }

    /// 交易写工具开关（legacy mode）。
    pub fn enable_trading(&self) -> bool {
        self.enable_trading
    }

    /// real 环境交易写工具开关（legacy mode）。
    pub fn allow_real_trading(&self) -> bool {
        self.allow_real_trading
    }

    /// 当前 MCP API key store。返回共享引用，避免调用方替换 runtime storage。
    pub fn key_store(&self) -> &Arc<KeyStore> {
        &self.key_store
    }

    /// startup 阶段验证过的 key 快照；调用方需要 fresh record 时仍应按 id 回查 key store。
    pub fn authed_key(&self) -> Option<Arc<KeyRecord>> {
        self.authed_key.clone()
    }

    /// 交易密码所属登录账号。
    pub fn trade_pwd_account(&self) -> Option<&str> {
        self.trade_pwd_account.as_deref()
    }

    /// 限额运行时计数器。返回共享引用，避免调用方替换 runtime storage。
    pub fn counters(&self) -> &Arc<RuntimeCounters> {
        &self.counters
    }

    /// 当前配置的 gateway 地址。
    pub async fn gateway(&self) -> String {
        self.inner.lock().await.gateway.clone()
    }

    /// 获取（或懒加载）网关客户端
    pub async fn client(&self) -> Result<Arc<FutuClient>> {
        let gateway = {
            let guard = self.inner.lock().await;
            if let Some(c) = &guard.client {
                return Ok(c.clone());
            }
            guard.gateway.clone()
        };

        let config = ClientConfig {
            addr: gateway.clone(),
            client_ver: env!("CARGO_PKG_VERSION").to_string(),
            client_id: "futu-mcp".to_string(),
            recv_notify: false,
            rsa_key: None,
        };
        let policy =
            ReconnectPolicy::new(MCP_CONNECT_RETRY_DELAY, MCP_CONNECT_RETRY_DELAY, Some(1));
        let mut reconnector = ReconnectingClient::new(config).with_policy(policy);
        let connect_result =
            tokio::time::timeout(MCP_CONNECT_TOTAL_TIMEOUT, reconnector.connect()).await;
        let (client, mut push_rx, _info) = match connect_result {
            Ok(result) => {
                result.with_context(|| format!("connect to futu gateway at {gateway}"))?
            }
            Err(_) => {
                return Err(anyhow!(
                    "connect to futu gateway at {gateway} timed out after {}s",
                    MCP_CONNECT_TOTAL_TIMEOUT.as_secs()
                ));
            }
        };

        let arc = Arc::new(client);
        {
            let mut guard = self.inner.lock().await;
            if let Some(c) = &guard.client {
                return Ok(c.clone());
            }
            guard.client = Some(arc.clone());
        }

        // v1.4.38 Phase 5 (100%): 按 acc_ids 过滤的 push broadcast
        //
        // 流程：
        // 1. push_rx 收 daemon 转发的 push
        // 2. 对 TRD_UPDATE_ORDER (2208) / TRD_UPDATE_ORDER_FILL (2218) 解包
        //    提取 acc_id
        // 3. 遍历订阅者，**只推给 acc_ids 匹配的**（或订阅者 acc_ids 空 = 不
        //    过滤，所有 acc 都收）
        // 4. 行情 push（QOT_UPDATE_*）无 acc_id 语义，广播给所有订阅者
        //
        // Per-session 独立 spawn notify，避免一个慢 session 阻塞其他
        let subs_for_push = Arc::downgrade(&self.push_subscribers);
        let key_store_for_push = Arc::downgrade(&self.key_store);
        // v1.4.105 F5 fix (codex review C4 USER_ACK B): MCP push filter 改用
        // FilterRegistry::should_drop_event 单一注册中心 (跟 4 surface 一致),
        // 替代之前 inline subscriber_should_receive_with_market. 防 sibling-route
        // bypass — 加新 push event filter 维度只在 install_defaults 注册一次,
        // MCP 自动覆盖.
        let filter_registry =
            std::sync::Arc::new(futu_auth_pipeline::FilterRegistry::with_defaults());
        tokio::spawn(async move {
            while let Some(push) = push_rx.recv().await {
                let Some(subs_for_push) = subs_for_push.upgrade() else {
                    break;
                };
                let Some(key_store_for_push) = key_store_for_push.upgrade() else {
                    break;
                };
                let subscribers = {
                    let subs = subs_for_push.lock().await;
                    if subs.is_empty() {
                        Vec::new()
                    } else {
                        subs.iter()
                            .map(|(session_id, sub)| (session_id.clone(), sub.clone()))
                            .collect::<Vec<_>>()
                    }
                };
                if subscribers.is_empty() {
                    continue; // fast path: no listeners, drop
                }
                // v1.4.105 T-C2 + v1.4.106 codex 0932 F6/F7: classify push by proto_id
                // (set membership), 不再靠 body decode 成功推断. trade body decode
                // 失败现在归 TradePushDecode::DecodeFailed (event_type="trade",
                // 无 acc/market gate 信息) — restricted key 应 drop, unrestricted
                // 透传带 decode_status="failed".
                let decode_result = classify_trade_push(push.proto_id, &push.body);
                let (push_acc_id, push_trd_market, decode_status, event_type) = match &decode_result
                {
                    TradePushDecode::NotTrade => (None, None, "ok", "quote"),
                    TradePushDecode::Decoded { acc_id, trd_market } => {
                        (Some(*acc_id), Some(*trd_market), "ok", "trade")
                    }
                    TradePushDecode::DecodeFailed => (None, None, "failed", "trade"),
                };
                let push_trd_market_str = push_trd_market.map(trd_market_int_to_str);
                // v1.4.106 codex 0932 F7 [P3]: payload 加 event_type / trd_market /
                // decode_status — 让客户端不需要按 proto_id 自己 derive (4 surface 一致).
                // body_base64 后向兼容保留.
                let payload = serde_json::json!({
                    "kind": "futu_push",
                    "proto_id": push.proto_id,
                    "acc_id": push_acc_id,
                    "event_type": event_type,
                    "trd_market": push_trd_market_str,
                    "decode_status": decode_status,
                    "body_base64": base64_encode_bytes(&push.body),
                });
                let deliveries = {
                    let mut deliveries = Vec::with_capacity(subscribers.len());
                    for (session_id, sub) in subscribers.iter() {
                        if !push_delivery_is_authorized(
                            sub,
                            &key_store_for_push,
                            &filter_registry,
                            &decode_result,
                            event_type,
                            push_acc_id,
                            push_trd_market_str,
                            push.proto_id,
                        ) {
                            continue;
                        }
                        deliveries.push(PendingPushDelivery {
                            target: sub.delivery.clone(),
                            data: payload.clone(),
                            session_id: session_id.clone(),
                            owner_key_id: sub.owner_key_id.clone(),
                            proto_id: push.proto_id,
                        });
                    }
                    deliveries
                };
                for delivery in deliveries {
                    match delivery.target {
                        SubscriberDelivery::LegacyPeer(peer, minimum_level) => {
                            tokio::spawn(async move {
                                let result =
                                    notify_legacy_push(&peer, &minimum_level, delivery.data)
                                        .await
                                        .map_err(|error| error.to_string());
                                if let Err(err) = result {
                                    tracing::warn!(
                                        proto_id = delivery.proto_id,
                                        session_id = delivery.session_id,
                                        owner_key_id = delivery.owner_key_id.as_deref().unwrap_or("<none>"),
                                        error = %err,
                                        "mcp push notification send failed"
                                    );
                                }
                            });
                        }
                        SubscriberDelivery::ModernResource(resource) => {
                            if let Some(work) = resource.enqueue(delivery.data) {
                                tokio::spawn(async move {
                                    resource.send_pending_notification(work).await;
                                });
                            }
                        }
                    }
                }
            }
        });

        // v1.4.39 Phase 5 stale cleanup: 5 分钟跑一次，移除 registered_at > 4h
        // 的订阅者。避免长跑 daemon 累积陈旧 subscriber（客户端断开 /  rmcp
        // session gone 但没显式 unregister 的情况）。
        let subs_for_purge = Arc::downgrade(&self.push_subscribers);
        tokio::spawn(async move {
            use std::time::Duration;
            const PURGE_INTERVAL: Duration = Duration::from_secs(5 * 60);
            let mut ticker = tokio::time::interval(PURGE_INTERVAL);
            ticker.tick().await; // skip the immediate first tick
            loop {
                ticker.tick().await;
                let Some(subs_for_purge) = subs_for_purge.upgrade() else {
                    break;
                };
                let now = std::time::Instant::now();
                let removed = {
                    let mut subs = subs_for_purge.lock().await;
                    let expired = subs
                        .iter()
                        .filter_map(|(handle, sub)| {
                            now.checked_duration_since(sub.registered_at)
                                .is_some_and(|age| age >= push_subscribers::PUSH_SUBSCRIBER_MAX_AGE)
                                .then_some(handle.clone())
                        })
                        .collect::<Vec<_>>();
                    expired
                        .into_iter()
                        .filter_map(|handle| subs.remove(&handle))
                        .collect::<Vec<_>>()
                };
                let purged = removed.len();
                for subscriber in removed {
                    subscriber.close_modern_resource();
                }
                if purged > 0 {
                    let remaining = subs_for_purge.lock().await.len();
                    tracing::info!(
                        purged,
                        remaining,
                        max_age_secs = push_subscribers::PUSH_SUBSCRIBER_MAX_AGE.as_secs(),
                        "v1.4.39 Phase 5: purged stale push subscribers (> 4h registered)"
                    );
                }
            }
        });

        Ok(arc)
    }

    /// Open a transient C++-compatible internal-UI transport for the pre-login
    /// Verification tool. The ordinary cached MCP client must keep normal
    /// InitConnect semantics and wait for a real login identity.
    pub async fn verification_client(&self) -> Result<Arc<FutuClient>> {
        let gateway = self.gateway().await;
        let config = ClientConfig {
            addr: gateway.clone(),
            client_ver: env!("CARGO_PKG_VERSION").to_string(),
            client_id: futu_core::INTERNAL_UI_CLIENT_ID.to_string(),
            recv_notify: false,
            rsa_key: None,
        };
        let policy =
            ReconnectPolicy::new(MCP_CONNECT_RETRY_DELAY, MCP_CONNECT_RETRY_DELAY, Some(1));
        let mut reconnector = ReconnectingClient::new(config).with_policy(policy);
        let connect_result =
            tokio::time::timeout(MCP_CONNECT_TOTAL_TIMEOUT, reconnector.connect()).await;
        let (client, _push_rx, _info) = match connect_result {
            Ok(result) => result.with_context(|| {
                format!("connect Verification transport to futu gateway at {gateway}")
            })?,
            Err(_) => {
                return Err(anyhow!(
                    "connect Verification transport to futu gateway at {gateway} timed out after {}s",
                    MCP_CONNECT_TOTAL_TIMEOUT.as_secs()
                ));
            }
        };
        Ok(Arc::new(client))
    }
}

// ========== symbol 解析 ==========

pub fn parse_symbol(s: &str) -> Result<Security> {
    let parsed = qot_symbol::parse_qot_symbol_parts(s).map_err(|err| anyhow!("{err}"))?;
    Ok(qot_sdk_adapter::security_from_parsed_symbol(parsed))
}

/// 格式化 Security 为 "MARKET.CODE"
pub fn format_symbol(sec: &Security) -> String {
    qot_symbol::format_qot_symbol(sec.market as i32, &sec.code)
}

/// v1.4.90 P2-C: audit log Option<T> 序列化助手。
///
/// **背景**：之前 audit log 把 `Option<f64>` 用 `?req.price`（tracing 的 Debug
/// shorthand）记录，渲染成 JSON 字符串 `"Some(400.0)"` / `"None"`，下游 jq /
/// DuckDB 数值聚合炸（aggregator 期望 `400.0` number 或 `null`）。
///
/// **修法**：用 NaN sentinel 把 `Option<f64>` flatten 成 `f64`，tracing-subscriber
/// 的 JSON formatter 内部走 `serde_json::Value::from(f64::NAN)` →
/// `Number::from_f64(NaN) = None` → `Value::Null`。
/// 整数 / 字符串同理（i32 → f64 NaN sentinel；&str → "" 哨兵）。
///
/// 验证依据：
/// - `tracing_subscriber::fmt::format::json` line 501 `record_f64` 直接调
///   `serde_json::Value::from(value)`
/// - `serde_json::Value::from(f64)` impl: `Number::from_f64(f).map_or(Value::Null, Value::Number)`
pub mod audit_fmt {
    /// `Option<f64>` → `f64`（None → NaN）。tracing JSON 渲染 NaN 为 `null`。
    #[inline]
    pub fn opt_f64(v: Option<f64>) -> f64 {
        v.unwrap_or(f64::NAN)
    }

    /// `Option<i32>` → `f64`（None → NaN，Some(n) → n as f64）。
    /// i32 ≤ 2^31 < 2^52 mantissa，无精度损失。
    #[inline]
    pub fn opt_i32(v: Option<i32>) -> f64 {
        v.map(f64::from).unwrap_or(f64::NAN)
    }

    /// `Option<&str>` → `&str`（None → ""）。"" 哨兵在 audit 上下文里足以区分
    /// 不传 vs 传空（因为 Symbol / owner 等业务字段不会是空字符串）。
    #[inline]
    pub fn opt_str(v: Option<&str>) -> &str {
        v.unwrap_or("")
    }
}