Skip to main content

weft_core/api/
rpc.rs

1//! 统一 RPC 端点(/rpc) + 内部 dispatch(供 FFI 直接调用不走网络)。
2
3use axum::body::Body;
4use axum::extract::State;
5use axum::http::{Method, Request};
6use axum::response::IntoResponse;
7use axum::routing::Router;
8use axum::Json;
9use serde::{Deserialize, Serialize};
10use std::sync::OnceLock;
11use tower::ServiceExt;
12
13use crate::api::openai_compat::AppState;
14
15/// 全局 Router 引用(build_router 后设一次,/rpc 和 FFI rpc_call 共用)。
16static ROUTER: OnceLock<Router> = OnceLock::new();
17
18/// 在 core 启动时调用(build_router 后),注册全局 router 供内部 dispatch 用。
19pub fn set_global_router(router: Router) {
20    let _ = ROUTER.set(router);
21}
22
23/// 进程内就绪检查:ROUTER 已注册即表示 FFI dispatch 可用(与 HTTP listener 无关)。
24/// FFI start_core 用它判断就绪,取代依赖 HTTP health 的脆弱探测。
25pub fn router_ready() -> bool {
26    ROUTER.get().is_some()
27}
28
29#[derive(Debug, Deserialize)]
30pub struct RequestEnvelope {
31    pub id: String,
32    #[serde(default = "default_method")]
33    pub method: String,
34    pub path: String,
35    #[serde(default)]
36    pub headers: std::collections::HashMap<String, String>,
37    #[serde(default)]
38    pub body: serde_json::Value,
39}
40
41fn default_method() -> String {
42    "CALL".into()
43}
44
45#[derive(Debug, Serialize)]
46pub struct ResponseEnvelope {
47    pub id: String,
48    pub status: u16,
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub headers: Option<std::collections::HashMap<String, String>>,
51    pub body: serde_json::Value,
52}
53
54/// POST /rpc — 统一 RPC 入口(内部 dispatch,不走网络回环)。
55pub async fn rpc_endpoint(
56    State(state): State<AppState>,
57    Json(envelope): Json<RequestEnvelope>,
58) -> impl IntoResponse {
59    let token = state.runtime_token.clone().unwrap_or_default();
60    Json(dispatch_internal(envelope, &token).await)
61}
62
63/// 内部 dispatch:直接调 router.oneshot(不走网络)。
64/// 供 /rpc endpoint 和 FFI rpc_call 共用。
65pub async fn dispatch_internal(envelope: RequestEnvelope, token: &str) -> ResponseEnvelope {
66    let router = match ROUTER.get() {
67        Some(r) => r.clone(),
68        None => {
69            return ResponseEnvelope {
70                id: envelope.id,
71                status: 503,
72                headers: None,
73                body: serde_json::json!({ "error": "router not initialized" }),
74            };
75        }
76    };
77
78    let method = match envelope.method.to_uppercase().as_str() {
79        "QUERY" | "GET" => Method::GET,
80        "DELETE" => Method::DELETE,
81        "PUT" => Method::PUT,
82        _ => Method::POST,
83    };
84
85    let path = if envelope.path.starts_with('/') {
86        envelope.path.clone()
87    } else {
88        format!("/{}", envelope.path)
89    };
90
91    let body_bytes = serde_json::to_vec(&envelope.body).unwrap_or_default();
92
93    let mut req_builder = Request::builder()
94        .method(method)
95        .uri(&path)
96        .header("content-type", "application/json")
97        .header("authorization", format!("Bearer {}", token));
98
99    for (k, v) in &envelope.headers {
100        if k.to_lowercase() != "authorization" {
101            req_builder = req_builder.header(k.as_str(), v.as_str());
102        }
103    }
104
105    let internal_req = match req_builder.body(Body::from(body_bytes)) {
106        Ok(r) => r,
107        Err(e) => {
108            return ResponseEnvelope {
109                id: envelope.id,
110                status: 400,
111                headers: None,
112                body: serde_json::json!({ "error": format!("bad request: {e}") }),
113            };
114        }
115    };
116
117    match router.oneshot(internal_req).await {
118        Ok(response) => {
119            let status = response.status().as_u16();
120            let body_bytes = axum::body::to_bytes(response.into_body(), 10 * 1024 * 1024)
121                .await
122                .unwrap_or_default();
123            let body: serde_json::Value =
124                serde_json::from_slice(&body_bytes).unwrap_or(serde_json::Value::Null);
125            ResponseEnvelope {
126                id: envelope.id,
127                status,
128                headers: None,
129                body,
130            }
131        }
132        Err(e) => ResponseEnvelope {
133            id: envelope.id,
134            status: 500,
135            headers: None,
136            body: serde_json::json!({ "error": format!("dispatch error: {e}") }),
137        },
138    }
139}