//! Tonic gRPC server implementation for SessionAdmin service. use session_manager::SessionManager; use slog::Logger; use std::net::SocketAddr; use tokio::sync::broadcast; use tonic::{transport::Server, Request, Response, Status}; use super::generated::admin::{ GetUserInfoRequest, GetUserInfoResponse, GetUserStatusRequest, GetUserStatusResponse, GetWorkspaceOnlineUsersRequest, GetWorkspaceOnlineUsersResponse, IsUserOnlineRequest, IsUserOnlineResponse, KickUserFromWorkspaceRequest, KickUserFromWorkspaceResponse, KickUserRequest, KickUserResponse, ListUserSessionsRequest, ListUserSessionsResponse, ListWorkspaceSessionsRequest, ListWorkspaceSessionsResponse, }; use super::generated::admin_session_admin::session_admin_server::{ SessionAdmin, SessionAdminServer, }; use super::types::{parse_uuid, to_proto_info, to_proto_session, to_proto_status}; #[derive(Clone)] pub struct SessionAdminService { session_manager: SessionManager, } impl SessionAdminService { pub fn new(session_manager: SessionManager) -> Self { Self { session_manager } } } #[tonic::async_trait] impl SessionAdmin for SessionAdminService { async fn list_workspace_sessions( &self, req: Request, ) -> Result, Status> { let workspace_id = parse_uuid(&req.get_ref().workspace_id) .ok_or_else(|| Status::invalid_argument("invalid workspace_id"))?; let sessions = self .session_manager .get_workspace_sessions(&workspace_id) .await .map_err(|e| Status::internal(e.to_string()))?; let sessions = sessions.iter().map(to_proto_session).collect(); Ok(Response::new(ListWorkspaceSessionsResponse { sessions })) } async fn list_user_sessions( &self, req: Request, ) -> Result, Status> { let user_id = parse_uuid(&req.get_ref().user_id) .ok_or_else(|| Status::invalid_argument("invalid user_id"))?; let sessions = self .session_manager .get_user_sessions(&user_id) .await .map_err(|e| Status::internal(e.to_string()))?; let sessions = sessions.iter().map(to_proto_session).collect(); Ok(Response::new(ListUserSessionsResponse { sessions })) } async fn kick_user_from_workspace( &self, req: Request, ) -> Result, Status> { let r = req.get_ref(); let user_id = parse_uuid(&r.user_id).ok_or_else(|| Status::invalid_argument("invalid user_id"))?; let workspace_id = parse_uuid(&r.workspace_id) .ok_or_else(|| Status::invalid_argument("invalid workspace_id"))?; let kicked_count = self .session_manager .kick_user_from_workspace(&user_id, &workspace_id) .await .map_err(|e| Status::internal(e.to_string()))?; Ok(Response::new(KickUserFromWorkspaceResponse { kicked_count: kicked_count as u32, })) } async fn kick_user( &self, req: Request, ) -> Result, Status> { let user_id = parse_uuid(&req.get_ref().user_id) .ok_or_else(|| Status::invalid_argument("invalid user_id"))?; let kicked_count = self .session_manager .kick_user(&user_id) .await .map_err(|e| Status::internal(e.to_string()))?; Ok(Response::new(KickUserResponse { kicked_count: kicked_count as u32, })) } async fn get_user_status( &self, req: Request, ) -> Result, Status> { let user_id = parse_uuid(&req.get_ref().user_id) .ok_or_else(|| Status::invalid_argument("invalid user_id"))?; let status = self .session_manager .get_user_status(&user_id) .await .map_err(|e| Status::internal(e.to_string()))?; Ok(Response::new(GetUserStatusResponse { status: to_proto_status(status) as i32, })) } async fn get_user_info( &self, req: Request, ) -> Result, Status> { let user_id = parse_uuid(&req.get_ref().user_id) .ok_or_else(|| Status::invalid_argument("invalid user_id"))?; let info = self .session_manager .get_user_info(&user_id) .await .map_err(|e| Status::internal(e.to_string()))?; Ok(Response::new(GetUserInfoResponse { info: info.as_ref().map(to_proto_info), })) } async fn get_workspace_online_users( &self, req: Request, ) -> Result, Status> { let workspace_id = parse_uuid(&req.get_ref().workspace_id) .ok_or_else(|| Status::invalid_argument("invalid workspace_id"))?; let user_ids = self .session_manager .get_workspace_online_users(&workspace_id) .await .map_err(|e| Status::internal(e.to_string()))?; let user_ids = user_ids.iter().map(|u| u.to_string()).collect(); Ok(Response::new(GetWorkspaceOnlineUsersResponse { user_ids })) } async fn is_user_online( &self, req: Request, ) -> Result, Status> { let user_id = parse_uuid(&req.get_ref().user_id) .ok_or_else(|| Status::invalid_argument("invalid user_id"))?; let online = self .session_manager .is_user_online(&user_id) .await .map_err(|e| Status::internal(e.to_string()))?; Ok(Response::new(IsUserOnlineResponse { online })) } } /// Default gRPC admin port. pub const DEFAULT_GRPC_PORT: u16 = 9090; /// Start the Tonic gRPC server on the given address. pub async fn serve( addr: SocketAddr, session_manager: SessionManager, log: Logger, ) -> anyhow::Result<()> { let service = SessionAdminService::new(session_manager); let incoming = tonic::transport::server::TcpIncoming::bind(addr) .map_err(|e| anyhow::anyhow!("failed to bind TcpIncoming: {}", e))?; slog::info!(log, "Admin gRPC server listening on {}", addr); Server::builder() .add_service(SessionAdminServer::new(service)) .serve_with_incoming(incoming) .await?; Ok(()) } /// Spawn the gRPC server as a background task. pub fn spawn( addr: SocketAddr, session_manager: SessionManager, log: Logger, mut shutdown_rx: broadcast::Receiver<()>, ) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let result = serve(addr, session_manager, log).await; if let Err(e) = result { eprintln!("Admin gRPC server error: {}", e); } let _ = shutdown_rx.recv().await; }) }