futu_mcp/tools/

subscription.rs

//! MCP subscription tools (QOT subscribe/unsubscribe and trade push subscriber state).

use crate::handlers;
use crate::tool_args::{
    NoArgs, QuerySubscriptionReq, SubAccPushReq, SubscribeReq, UnsubAccPushReq, UnsubscribeReq,
};
use crate::tool_auth::CallerSnapshot;
use rmcp::{
    RoleServer, handler::server::wrapper::Parameters, service::RequestContext, tool, tool_router,
};

use super::{FutuServer, system::json_tool_output};

#[tool_router(router = subscription_tool_router, vis = "pub(crate)")]
impl FutuServer {
    #[tool(
        description = "Query current subscription state (subscribed types, quota used/remaining). Python SDK: OpenQuoteContext.query_subscription."
    )]
    async fn futu_query_subscription(
        &self,
        Parameters(req): Parameters<QuerySubscriptionReq>,
        req_ctx: RequestContext<RoleServer>,
    ) -> std::result::Result<String, String> {
        tracing::info!(tool = "futu_query_subscription");
        let client = self
            .read_client_or_err("futu_query_subscription", &req_ctx, None, None)
            .await?;
        Self::wrap_result(handlers::core::query_subscription(&client, req.is_req_all_conn).await)
    }

    #[tool(
        description = "Get current daemon used quota counters: subscribed quote quota and historical K-line quota. Python SDK: OpenQuoteContext.get_used_quota."
    )]
    async fn futu_get_used_quota(
        &self,
        Parameters(_req): Parameters<NoArgs>,
        req_ctx: RequestContext<RoleServer>,
    ) -> std::result::Result<String, String> {
        tracing::info!(tool = "futu_get_used_quota");
        let client = self
            .read_client_or_err("futu_get_used_quota", &req_ctx, None, None)
            .await?;
        Self::wrap_result(handlers::core::get_used_quota(&client).await)
    }

    /// v1.4.74 A1 BUG-011 fix: 加 `futu_subscribe` 补齐 MCP 对称性。
    ///
    /// MCP 之前有 `futu_unsubscribe` / `futu_sub_acc_push` / `futu_unsub_acc_push`
    /// / `futu_query_subscription` / `futu_push_subscriber_info` 但缺
    /// `futu_subscribe` 本身 → API 对称性被打破（"能取消但不能订阅"）。
    ///
    /// 架构：本 tool 触发 gateway 层订阅（CMD 3001 QOT_SUB），不返 push stream
    /// 本身。Push 数据通过 SSE notification 走（v1.4.58 MCP SSE basics），
    /// 客户端用 `futu_push_subscriber_info` 查询订阅状态。
    ///
    /// 对齐 Python SDK `OpenQuoteContext.subscribe`。
    #[tool(
        description = "Subscribe market data for given symbols + sub_types. Push data arrives via SSE notifications (HTTP mode). Python SDK: OpenQuoteContext.subscribe."
    )]
    async fn futu_subscribe(
        &self,
        Parameters(req): Parameters<SubscribeReq>,
        req_ctx: RequestContext<RoleServer>,
    ) -> std::result::Result<String, String> {
        req.validate()?;
        tracing::info!(
            tool = "futu_subscribe",
            symbols = ?req.symbols,
            sub_types = ?req.sub_types,
            is_first_push = req.is_first_push,
            is_reg_push = req.is_reg_push,
            extended_time = ?req.extended_time,
            session = ?req.session,
            is_sub_order_book_detail = ?req.is_sub_order_book_detail,
        );
        let client = self
            .read_client_or_err("futu_subscribe", &req_ctx, None, None)
            .await?;
        Self::wrap_result(
            handlers::core::subscribe(
                &client,
                &req.symbols,
                &req.sub_types,
                req.is_first_push,
                req.is_reg_push,
                req.extended_time,
                req.session,
                req.is_sub_order_book_detail,
            )
            .await,
        )
    }

    #[tool(
        description = "Unsubscribe market data (by symbol+type, or unsub_all to clear this connection). Python SDK: OpenQuoteContext.unsubscribe / unsubscribe_all."
    )]
    async fn futu_unsubscribe(
        &self,
        Parameters(req): Parameters<UnsubscribeReq>,
        req_ctx: RequestContext<RoleServer>,
    ) -> std::result::Result<String, String> {
        tracing::info!(
            tool = "futu_unsubscribe",
            count = req.symbols.len(),
            unsub_all = req.unsub_all
        );
        let client = self
            .read_client_or_err("futu_unsubscribe", &req_ctx, None, None)
            .await?;
        Self::wrap_result(
            handlers::core::unsubscribe(&client, &req.symbols, &req.sub_types, req.unsub_all).await,
        )
    }

    #[tool(
        description = "Subscribe account order / deal push for given trading accounts. HTTP-mode MCP clients receive pushes as LoggingMessage notifications with {kind, proto_id, body_base64}. Payload body is raw Futu protobuf; decode client-side. Python SDK: OpenTradeContext.sub_acc_push."
    )]
    async fn futu_sub_acc_push(
        &self,
        Parameters(req): Parameters<SubAccPushReq>,
        req_ctx: RequestContext<RoleServer>,
    ) -> std::result::Result<String, String> {
        // v1.4.102 codex 47 F2 (P1): 空 acc_ids 必拒. 之前空 list 进
        // register_push_subscriber, subscriber_should_receive 把空 set 视为
        // "subscribe-all" 全开 push, agent 调 {"acc_ids": []} 等于全订阅
        // (silent privilege escalation 反模式 D / pitfall #45).
        if req.acc_ids.is_empty() {
            return Err(
                "futu_sub_acc_push: acc_ids 必填非空. 之前空 list silent 全订阅 \
                 (v1.4.102 codex 47 F2 P1 fix). 调 futu_list_accounts 看可用 \
                 acc_id, 显式列出要订阅的账户."
                    .to_string(),
            );
        }

        // v1.4.103 (codex 50 F4 / 52 F4 / 54 F3 / 58 F2 — B6): 在调 daemon
        // 之前用 caller's KeyRecord 比对每个 req.acc_ids ⊆ allowed_acc_ids.
        // 之前 register-after-daemon-success 只防 daemon 失败的 race window,
        // 不防 narrow-scope key 让 daemon 全局订阅 acc B (subscriber 本身被
        // delivery filter 过滤掉, 但 daemon 已对 acc B 发起订阅副作用).
        //
        // v1.4.104 codex F3 (P2): 用 pipeline 返的 snapshot 做 ownership +
        // visibility — pipeline 授权 + push 注册同一 KeyRecord 实例, 避免
        // SIGHUP 之间 drift / fail-open. 取第 1 个 acc_id 的 snapshot 作 owner;
        // 后续 acc_ids 通过 pipeline 复 check (不重 capture, 同一 caller).
        let mut snap_opt: Option<CallerSnapshot> = None;
        for (idx, acc_id) in req.acc_ids.iter().enumerate() {
            let snap = self.require_acc_read_with_acc_id(
                "futu_sub_acc_push",
                &req_ctx,
                req.api_key.as_deref(),
                Some(*acc_id),
            )?;
            if idx == 0 {
                snap_opt = Some(snap);
            }
        }
        let Some(snap) = snap_opt else {
            return Err(serde_json::json!({
                "error": "futu_sub_acc_push: caller snapshot resolution failed after acc_ids validation",
                "status": "error",
                "hint": "retry after refreshing API key state; if it persists, check auth pipeline logs",
            })
            .to_string());
        };

        tracing::info!(tool = "futu_sub_acc_push", count = req.acc_ids.len());

        let client = self.client_or_err().await?;

        // v1.4.102 codex 47 F3 / 48 F3 (P2): register-after-daemon-success.
        // 之前先 register peer 后调 daemon, daemon 失败时 local subscriber 留下
        // 来 — 后续如果其他 caller 让 daemon sub 了同 acc 流, 这个本来失败的
        // session 仍能收 push (race window leak). 现在: 先调 daemon, 成功才
        // register local subscriber; daemon 失败 → 不留 local state.
        let daemon_resp = handlers::trade::sub_acc_push(&client, &req.acc_ids).await;
        match daemon_resp {
            Ok(_) => {
                // v1.4.104 codex F3 (P2): 用 pipeline snapshot 不重新 resolve.
                // 与 v1.4.103 codex F5.7 (P2) 防 SIGHUP race 行为对齐, 但现
                // owner_key_id + bearer_token 都从 snapshot 直取, 与 pipeline
                // 授权决策同一身份. 不再 await 后重 verify.
                let owner_key_id = snap.key_id.clone();
                let bearer_token = snap.bearer_token.clone();
                let acc_ids_set: std::collections::HashSet<u64> =
                    req.acc_ids.iter().copied().collect();
                let session_id = self
                    .state
                    .register_push_subscriber_with_owner(
                        req_ctx.peer.clone(),
                        acc_ids_set,
                        bearer_token,
                        owner_key_id,
                    )
                    .await;
                tracing::info!(
                    tool = "futu_sub_acc_push",
                    session_id = %session_id,
                    count = req.acc_ids.len(),
                    "v1.4.102 audit 47 F3: push subscriber registered after daemon success"
                );
                json_tool_output(
                    "futu_sub_acc_push",
                    &serde_json::json!({
                        "ok": true,
                        "subscribed_acc_ids": req.acc_ids,
                        "session_id": session_id,
                        "unsub_hint": format!(
                            "call `futu_unsub_acc_push` with session_id=\"{session_id}\" to stop receiving pushes; otherwise auto-purged after 4h"
                        ),
                    }),
                )
            }
            Err(e) => Self::tool_err(format!("futu_sub_acc_push: {e}")),
        }
    }

    #[tool(
        description = "Unsubscribe from account push notifications. Pass session_id from previous futu_sub_acc_push response (session_id field or unsub_hint). Returns {removed_count}. If session_id is not found, removed_count=0 (likely auto-purged or never registered)."
    )]
    async fn futu_unsub_acc_push(
        &self,
        Parameters(req): Parameters<UnsubAccPushReq>,
        req_ctx: RequestContext<RoleServer>,
    ) -> std::result::Result<String, String> {
        // v1.4.103 codex F4 (P1) fail-closed: scope check 用 caller-specific
        // (HTTP Bearer / api_key 优先) 而不是 process-wide startup key.
        // 之前 require_tool_scope 只看 startup, 受限 Bearer fall through.
        let snap = self.require_acc_read_with_acc_id(
            "futu_unsub_acc_push",
            &req_ctx,
            req.api_key.as_deref(),
            None,
        )?;
        let Some(ref session_id) = req.session_id else {
            return Self::tool_err(
                "session_id required. Get it from a previous futu_sub_acc_push response.",
            );
        };
        // v1.4.103 (codex 50 F6 / 53 F4 — B8): unsub session ownership check.
        // 之前任何 caller 拿到可见 session_id 即可 remove, 跨 caller 容易踩
        // (尤其 push_subscriber_info 列其他 session 后被恶意 caller unsub).
        // 现在: 解析 caller key (Bearer / startup), 比对 subscriber.owner_key_id.
        // v1.4.103 codex F4 (P1) fail-closed: invalid Bearer → 不 fall back
        // startup key. 因为前面 require_acc_read_with_acc_id 已经把 invalid
        // api_key / Bearer reject 掉了, 走到这里直接复用同一 auth snapshot 的
        // key_id 做 ownership check，避免重新解析时与 sub/register 身份漂移。
        let caller_key_id = snap.key_id;

        let removed = self
            .state
            .unregister_push_subscriber_with_owner_check(
                session_id.as_str(),
                caller_key_id.as_deref(),
            )
            .await;
        tracing::info!(
            tool = "futu_unsub_acc_push",
            removed = removed.is_ok(),
            session_id = %session_id,
            caller_key_id = ?caller_key_id,
            "v1.4.103 B8: push subscriber unregister with ownership check"
        );
        match removed {
            Ok(true) => json_tool_output("futu_unsub_acc_push", &serde_json::json!({
                "ok": true,
                "removed_count": 1,
                "session_id": session_id,
            })),
            Ok(false) => json_tool_output("futu_unsub_acc_push", &serde_json::json!({
                "ok": true,
                "removed_count": 0,
                "session_id": session_id,
                "hint": "session_id not found (likely 4h auto-purged or never registered)",
            })),
            Err(reason) => Err(serde_json::json!({
                "error": format!("ownership check failed: {reason}"),
                "status": "error",
                "hint": "only the original caller (matched by API key id) or an admin-scope key can unsub a session.",
            })
            .to_string()),
        }
    }

    #[tool(
        description = "Diagnostic — list active push subscriptions on this MCP server. Returns {total_count, subscriptions: [{session_id, acc_ids, age_secs}]}. Useful to verify whether futu_sub_acc_push registered, check auto-purge timing, or debug missing pushes."
    )]
    async fn futu_push_subscriber_info(
        &self,
        Parameters(_req): Parameters<NoArgs>,
        req_ctx: RequestContext<RoleServer>,
    ) -> std::result::Result<String, String> {
        // v1.4.103 codex F5 (P2): 加 RequestContext 让 caller-specific Bearer
        // 解析能生效. 之前 require_tool_scope 只看 startup key, narrow Bearer
        // caller 看到的 caller_allowed 是 startup key 的, 跨租户泄漏其他
        // agent 的订阅 acc_ids.
        //
        // v1.4.104 codex F3 (P2) fix: 用 pipeline 返的 snapshot 做 visibility
        // filter, 不再 re-resolve from Bearer/startup (避免 SIGHUP race —
        // pipeline 授权与 visibility 用同一 KeyRecord 实例).
        let snap =
            self.require_acc_read_with_acc_id("futu_push_subscriber_info", &req_ctx, None, None)?;
        let subs = self
            .state
            .push_subscribers_summary(snap.allowed_acc_ids.as_ref())
            .await;
        json_tool_output(
            "futu_push_subscriber_info",
            &serde_json::json!({
                "ok": true,
                "total_count": subs.len(),
                "subscriptions": subs.iter().map(|(sid, acc_ids, age)| serde_json::json!({
                    "session_id": sid,
                    "acc_ids": acc_ids.iter().collect::<Vec<_>>(),
                    "age_secs": age,
                })).collect::<Vec<_>>(),
            }),
        )
    }
}