apollo/
agent_http.rs

1use std::net::SocketAddr;
2use std::sync::Arc;
3
4use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
5use axum::extract::State;
6use axum::http::{Method, StatusCode};
7use axum::routing::{get, post};
8use axum::{Json, Router};
9use serde::{Deserialize, Serialize};
10use tokio::sync::mpsc;
11use tower_http::cors::{Any, CorsLayer};
12
13use crate::agent::stream::AgentStreamEvent;
14use crate::agent::AgentRunner;
15use crate::channels::http_inject::HttpInjectChannel;
16use crate::channels::IncomingMessage;
17
18#[derive(Debug, Deserialize)]
19pub struct ChatRequestBody {
20    pub message: String,
21    #[serde(default = "default_chat_id")]
22    pub chat_id: String,
23}
24
25fn default_chat_id() -> String {
26    "embed".into()
27}
28
29#[derive(Debug, Serialize)]
30pub struct ChatResponseBody {
31    pub response: String,
32}
33
34pub async fn chat_once(
35    runner: &AgentRunner,
36    message: &str,
37    chat_id: &str,
38) -> anyhow::Result<String> {
39    let msg = IncomingMessage {
40        id: uuid::Uuid::new_v4().to_string(),
41        sender_id: "http".into(),
42        sender_name: Some("HTTP".into()),
43        chat_id: chat_id.to_string(),
44        text: message.to_string(),
45        is_group: false,
46        reply_to: None,
47        timestamp: chrono::Utc::now(),
48    };
49    let channel = HttpInjectChannel::new();
50    runner.handle_message(&msg, &channel).await
51}
52
53pub fn http_listen_addr() -> SocketAddr {
54    let port = std::env::var("APOLLO_HTTP_PORT")
55        .ok()
56        .and_then(|p| p.parse().ok())
57        .unwrap_or(31338);
58    SocketAddr::from(([127, 0, 0, 1], port))
59}
60
61fn cors_layer() -> CorsLayer {
62    CorsLayer::new()
63        .allow_origin(Any)
64        .allow_methods([Method::GET, Method::POST, Method::OPTIONS])
65        .allow_headers(Any)
66}
67
68pub fn spawn_http_server(runner: Arc<AgentRunner>) {
69    let addr = http_listen_addr();
70    tokio::spawn(async move {
71        let app = Router::new()
72            .route("/health", get(|| async { "ok" }))
73            .route("/v1/chat", post(chat_handler))
74            .route("/v1/chat/stream", get(ws_chat_upgrade))
75            .layer(cors_layer())
76            .with_state(runner);
77        let listener = match tokio::net::TcpListener::bind(addr).await {
78            Ok(l) => l,
79            Err(e) => {
80                tracing::error!("apollo http bind {}: {}", addr, e);
81                return;
82            }
83        };
84        tracing::info!(
85            "apollo agent HTTP http://{}/v1/chat · WS /v1/chat/stream",
86            addr
87        );
88        if let Err(e) = axum::serve(listener, app).await {
89            tracing::error!("apollo http server: {}", e);
90        }
91    });
92}
93
94async fn chat_handler(
95    State(runner): State<Arc<AgentRunner>>,
96    Json(body): Json<ChatRequestBody>,
97) -> Result<Json<ChatResponseBody>, StatusCode> {
98    if body.message.trim().is_empty() {
99        return Err(StatusCode::BAD_REQUEST);
100    }
101    match chat_once(&runner, body.message.trim(), &body.chat_id).await {
102        Ok(response) => Ok(Json(ChatResponseBody { response })),
103        Err(e) => {
104            tracing::error!("http chat: {}", e);
105            Err(StatusCode::INTERNAL_SERVER_ERROR)
106        }
107    }
108}
109
110async fn ws_chat_upgrade(
111    ws: WebSocketUpgrade,
112    State(runner): State<Arc<AgentRunner>>,
113) -> impl axum::response::IntoResponse {
114    ws.on_upgrade(move |socket| handle_ws_chat(socket, runner))
115}
116
117async fn handle_ws_chat(mut socket: WebSocket, runner: Arc<AgentRunner>) {
118    let Some(Ok(Message::Text(text))) = socket.recv().await else {
119        return;
120    };
121    let Ok(body) = serde_json::from_str::<ChatRequestBody>(&text) else {
122        let _ = socket
123            .send(Message::Text(
124                serde_json::json!({"type":"error","message":"invalid JSON"}).to_string(),
125            ))
126            .await;
127        return;
128    };
129    if body.message.trim().is_empty() {
130        let _ = socket
131            .send(Message::Text(
132                serde_json::json!({"type":"error","message":"empty message"}).to_string(),
133            ))
134            .await;
135        return;
136    }
137
138    let (stream_tx, mut stream_rx) = mpsc::unbounded_channel::<AgentStreamEvent>();
139    let runner_bg = Arc::clone(&runner);
140    let message = body.message.trim().to_string();
141    let chat_id = body.chat_id.clone();
142    let mut chat_task = tokio::spawn(async move {
143        runner_bg.set_stream_sink(Some(stream_tx));
144        let result = chat_once(&runner_bg, &message, &chat_id).await;
145        runner_bg.set_stream_sink(None);
146        result
147    });
148
149    loop {
150        tokio::select! {
151            Some(ev) = stream_rx.recv() => {
152                if let Ok(json) = serde_json::to_string(&ev) {
153                    if socket.send(Message::Text(json)).await.is_err() {
154                        return;
155                    }
156                }
157            }
158            result = &mut chat_task => {
159                match result {
160                    Ok(Ok(response)) => {
161                        let payload = serde_json::to_string(&AgentStreamEvent::Done {
162                            response: response.clone(),
163                        })
164                        .unwrap_or_else(|_| {
165                            serde_json::json!({"type":"done","response": response}).to_string()
166                        });
167                        let _ = socket.send(Message::Text(payload)).await;
168                    }
169                    Ok(Err(e)) => {
170                        let payload = serde_json::to_string(&AgentStreamEvent::Error {
171                            message: e.to_string(),
172                        })
173                        .unwrap_or_else(|_| {
174                            serde_json::json!({"type":"error","message": e.to_string()}).to_string()
175                        });
176                        let _ = socket.send(Message::Text(payload)).await;
177                    }
178                    Err(e) => {
179                        let payload = serde_json::json!({"type":"error","message": e.to_string()});
180                        let _ = socket.send(Message::Text(payload.to_string())).await;
181                    }
182                }
183                break;
184            }
185        }
186    }
187
188    while let Ok(ev) = stream_rx.try_recv() {
189        if let Ok(json) = serde_json::to_string(&ev) {
190            let _ = socket.send(Message::Text(json)).await;
191        }
192    }
193}
apollo/agent_http.rs

apollo/
agent_http.rs
