1use 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
15static ROUTER: OnceLock<Router> = OnceLock::new();
17
18pub fn set_global_router(router: Router) {
20 let _ = ROUTER.set(router);
21}
22
23pub 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
54pub 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
63pub 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}