gitdataai/libs/rpc/admin/client.rs

//! Tonic gRPC client wrapper for SessionAdmin service.

use session_manager::{OnlineStatus, SessionInfo, UserSession};
use tonic::transport::Channel;
use uuid::Uuid;

use super::generated::admin::{
    ListWorkspaceSessionsRequest, ListUserSessionsRequest,
    KickUserFromWorkspaceRequest, KickUserRequest,
    GetUserStatusRequest, GetUserInfoRequest,
    GetWorkspaceOnlineUsersRequest, IsUserOnlineRequest,
};
use super::generated::admin_session_admin::session_admin_client::SessionAdminClient;
use super::types::from_proto_status;

/// Auto-generated gRPC client type.
pub type GrpcClient = SessionAdminClient<Channel>;

/// Thin wrapper around the generated SessionAdminClient.
pub struct AdminGrpcClient {
    inner: GrpcClient,
}

impl AdminGrpcClient {
    /// Connect to the gRPC server at the given URI.
    pub async fn connect(uri: tonic::codegen::http::Uri) -> anyhow::Result<Self> {
        let inner = SessionAdminClient::connect(uri).await?;
        Ok(Self { inner })
    }

    /// Wrap an existing channel.
    pub fn new(inner: GrpcClient) -> Self {
        Self { inner }
    }

    pub async fn list_workspace_sessions(
        &mut self,
        workspace_id: Uuid,
    ) -> anyhow::Result<Vec<UserSession>> {
        let req = tonic::Request::new(ListWorkspaceSessionsRequest {
            workspace_id: workspace_id.to_string(),
        });
        let res = self.inner.list_workspace_sessions(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        Ok(res.into_inner().sessions
            .into_iter()
            .map(from_proto_session)
            .collect())
    }

    pub async fn list_user_sessions(
        &mut self,
        user_id: Uuid,
    ) -> anyhow::Result<Vec<UserSession>> {
        let req = tonic::Request::new(ListUserSessionsRequest {
            user_id: user_id.to_string(),
        });
        let res = self.inner.list_user_sessions(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        Ok(res.into_inner().sessions
            .into_iter()
            .map(from_proto_session)
            .collect())
    }

    pub async fn kick_user_from_workspace(
        &mut self,
        user_id: Uuid,
        workspace_id: Uuid,
    ) -> anyhow::Result<usize> {
        let req = tonic::Request::new(KickUserFromWorkspaceRequest {
            user_id: user_id.to_string(),
            workspace_id: workspace_id.to_string(),
        });
        let res = self.inner.kick_user_from_workspace(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        Ok(res.into_inner().kicked_count as usize)
    }

    pub async fn kick_user(&mut self, user_id: Uuid) -> anyhow::Result<usize> {
        let req = tonic::Request::new(KickUserRequest {
            user_id: user_id.to_string(),
        });
        let res = self.inner.kick_user(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        Ok(res.into_inner().kicked_count as usize)
    }

    pub async fn get_user_status(&mut self, user_id: Uuid) -> anyhow::Result<OnlineStatus> {
        let req = tonic::Request::new(GetUserStatusRequest {
            user_id: user_id.to_string(),
        });
        let res = self.inner.get_user_status(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        let status = super::generated::admin::OnlineStatus::try_from(res.get_ref().status)
            .unwrap_or(super::generated::admin::OnlineStatus::Offline);
        Ok(from_proto_status(status))
    }

    pub async fn get_user_info(
        &mut self,
        user_id: Uuid,
    ) -> anyhow::Result<Option<SessionInfo>> {
        let req = tonic::Request::new(GetUserInfoRequest {
            user_id: user_id.to_string(),
        });
        let res = self.inner.get_user_info(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        Ok(res.into_inner().info.map(|i| from_proto_info(&i)))
    }

    pub async fn get_workspace_online_users(
        &mut self,
        workspace_id: Uuid,
    ) -> anyhow::Result<Vec<Uuid>> {
        let req = tonic::Request::new(GetWorkspaceOnlineUsersRequest {
            workspace_id: workspace_id.to_string(),
        });
        let res = self.inner.get_workspace_online_users(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        res.into_inner().user_ids
            .into_iter()
            .map(|s| Uuid::parse_str(&s))
            .collect::<Result<Vec<_>, _>>()
            .map_err(|e| anyhow::anyhow!("invalid UUID in response: {}", e))
    }

    pub async fn is_user_online(&mut self, user_id: Uuid) -> anyhow::Result<bool> {
        let req = tonic::Request::new(IsUserOnlineRequest {
            user_id: user_id.to_string(),
        });
        let res = self.inner.is_user_online(req).await
            .map_err(|e| anyhow::anyhow!("gRPC error: {}", e))?;
        Ok(res.into_inner().online)
    }
}

// ---------------------------------------------------------------------------
// Proto → session_manager type conversions
// ---------------------------------------------------------------------------

fn from_proto_session(s: super::generated::admin::UserSession) -> UserSession {
    UserSession {
        session_id: Uuid::parse_str(&s.session_id).unwrap_or_default(),
        user_id: Uuid::parse_str(&s.user_id).unwrap_or_default(),
        workspace_id: Uuid::parse_str(&s.workspace_id).unwrap_or_default(),
        ip_address: s.ip_address,
        user_agent: s.user_agent,
        connected_at: s.connected_at
            .and_then(|ts| chrono::DateTime::from_timestamp(ts.seconds, ts.nanos as u32))
            .unwrap_or_else(chrono::Utc::now),
        last_heartbeat: s.last_heartbeat
            .and_then(|ts| chrono::DateTime::from_timestamp(ts.seconds, ts.nanos as u32))
            .unwrap_or_else(chrono::Utc::now),
    }
}

fn from_proto_info(info: &super::generated::admin::SessionInfo) -> SessionInfo {
    SessionInfo {
        user_id: Uuid::parse_str(&info.user_id).unwrap_or_default(),
        session_count: info.session_count as usize,
        workspaces: info.workspaces
            .iter()
            .filter_map(|w| Uuid::parse_str(w).ok())
            .collect(),
        latest_session: info.latest_session.as_ref().map(|s| from_proto_session(s.clone())),
    }
}