hehe_server/
server.rs

1use axum::{
2    routing::{get, post},
3    Router,
4};
5use hehe_agent::Agent;
6use tower_http::cors::{Any, CorsLayer};
7use tower_http::trace::TraceLayer;
8use tracing::info;
9
10use crate::config::ServerConfig;
11use crate::error::Result;
12use crate::routes;
13use crate::state::AppState;
14
15pub struct Server {
16    config: ServerConfig,
17    state: AppState,
18}
19
20impl Server {
21    pub fn new(config: ServerConfig, agent: Agent) -> Self {
22        Self {
23            config,
24            state: AppState::new(agent),
25        }
26    }
27
28    pub fn router(&self) -> Router {
29        let cors = CorsLayer::new()
30            .allow_origin(Any)
31            .allow_methods(Any)
32            .allow_headers(Any);
33
34        Router::new()
35            .route("/health", get(routes::health))
36            .route("/ready", get(routes::ready))
37            .route("/api/v1/chat", post(routes::chat))
38            .route("/api/v1/chat/stream", post(routes::chat_stream))
39            .layer(cors)
40            .layer(TraceLayer::new_for_http())
41            .with_state(self.state.clone())
42    }
43
44    pub async fn run(self) -> Result<()> {
45        let addr = self.config.socket_addr();
46        let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
47            crate::error::ServerError::internal(format!("Failed to bind to {}: {}", addr, e))
48        })?;
49
50        info!("Server listening on {}", addr);
51
52        axum::serve(listener, self.router())
53            .await
54            .map_err(|e| crate::error::ServerError::internal(e.to_string()))?;
55
56        Ok(())
57    }
58
59    pub async fn run_with_shutdown<F>(self, shutdown: F) -> Result<()>
60    where
61        F: std::future::Future<Output = ()> + Send + 'static,
62    {
63        let addr = self.config.socket_addr();
64        let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
65            crate::error::ServerError::internal(format!("Failed to bind to {}: {}", addr, e))
66        })?;
67
68        info!("Server listening on {}", addr);
69
70        axum::serve(listener, self.router())
71            .with_graceful_shutdown(shutdown)
72            .await
73            .map_err(|e| crate::error::ServerError::internal(e.to_string()))?;
74
75        Ok(())
76    }
77
78    pub fn config(&self) -> &ServerConfig {
79        &self.config
80    }
81
82    pub fn state(&self) -> &AppState {
83        &self.state
84    }
85}
86
87pub async fn shutdown_signal() {
88    let ctrl_c = async {
89        tokio::signal::ctrl_c()
90            .await
91            .expect("Failed to install Ctrl+C handler");
92    };
93
94    #[cfg(unix)]
95    let terminate = async {
96        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
97            .expect("Failed to install signal handler")
98            .recv()
99            .await;
100    };
101
102    #[cfg(not(unix))]
103    let terminate = std::future::pending::<()>();
104
105    tokio::select! {
106        _ = ctrl_c => {},
107        _ = terminate => {},
108    }
109
110    info!("Shutdown signal received");
111}
hehe_server/server.rs

hehe_server/
server.rs
