feat(agent): implement AgentService with CRUD and gRPC stream handler

- Add AgentService trait with full CRUD operations

- Implement AgentServiceImpl backed by AgentRepo

- Refactor gRPC AgentServerService with real stream handler

- Add agent types and repository module
This commit is contained in:
GW_MC
2026-07-05 05:29:13 +00:00
parent 6d1df8dcca
commit e44a67f5a8
3 changed files with 312 additions and 12 deletions

View File

@@ -1,26 +1,120 @@
use nxmesh_proto::{AgentMessage, MasterMessage, agent_service_server::AgentService};
use chrono::Utc;
use nxmesh_proto::{
AgentMessage, MasterMessage, agent_service_server::AgentService as GrpcAgentService,
};
use sea_orm::DatabaseConnection;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tracing::{error, info};
use uuid::Uuid;
pub mod repo;
mod repo;
pub mod types;
pub use types::{AgentRecord, CreateAgentRecord, State, UpdateAgentRecord};
use crate::service::error::RepoError;
use repo::{AgentRepo, AgentRepoImpl};
#[cfg_attr(test, mockall::automock)]
#[async_trait::async_trait]
pub trait AgentService: Send + Sync + 'static {
async fn list(&self) -> Result<Vec<AgentRecord>, RepoError>;
async fn get(&self, id: Uuid) -> Result<Option<AgentRecord>, RepoError>;
async fn create(&self, rec: &CreateAgentRecord) -> Result<AgentRecord, RepoError>;
async fn update(
&self,
id: Uuid,
rec: &UpdateAgentRecord,
) -> Result<Option<AgentRecord>, RepoError>;
async fn delete(&self, id: Uuid) -> Result<bool, RepoError>;
}
pub struct AgentServiceImpl {
repo: Box<dyn AgentRepo>,
}
impl AgentServiceImpl {
pub fn new(db: DatabaseConnection) -> Self {
Self {
repo: Box::new(AgentRepoImpl::new(db)),
}
}
}
#[async_trait::async_trait]
impl AgentService for AgentServiceImpl {
async fn list(&self) -> Result<Vec<AgentRecord>, RepoError> {
self.repo.list().await
}
async fn get(&self, id: Uuid) -> Result<Option<AgentRecord>, RepoError> {
self.repo.get(id).await
}
async fn create(&self, rec: &CreateAgentRecord) -> Result<AgentRecord, RepoError> {
self.repo.create(rec).await
}
async fn update(
&self,
id: Uuid,
rec: &UpdateAgentRecord,
) -> Result<Option<AgentRecord>, RepoError> {
self.repo.update(id, rec).await
}
async fn delete(&self, id: Uuid) -> Result<bool, RepoError> {
self.repo.delete(id).await
}
}
#[derive(Debug, Default)]
pub struct AgentServerService {}
#[async_trait::async_trait]
impl AgentService for AgentServerService {
#[doc = " Server streaming response type for the Stream method."]
type StreamStream = tonic::codec::Streaming<MasterMessage>;
impl GrpcAgentService for AgentServerService {
type StreamStream =
tokio_stream::wrappers::ReceiverStream<std::result::Result<MasterMessage, tonic::Status>>;
#[doc = " Stream establishes a persistent connection for real-time communication"]
#[allow(
mismatched_lifetime_syntaxes,
clippy::type_complexity,
clippy::type_repetition_in_bounds
)]
async fn stream(
&self,
request: tonic::Request<tonic::Streaming<AgentMessage>>,
) -> Result<tonic::Response<Self::StreamStream>, tonic::Status> {
todo!()
let mut inbound = request.into_inner();
let (tx, rx) = mpsc::channel::<std::result::Result<MasterMessage, tonic::Status>>(32);
let outbound = ReceiverStream::new(rx);
tokio::spawn(async move {
loop {
match inbound.message().await {
Ok(Some(msg)) => {
info!("Received AgentMessage: {:?}", msg);
let ack = MasterMessage {
timestamp: Utc::now().timestamp_millis(),
message_id: Uuid::new_v4().to_string(),
payload: None,
};
if let Err(e) = tx.send(Ok(ack)).await {
error!("Failed to send MasterMessage ack: {:?}", e);
return;
}
}
Ok(None) => {
info!("Agent closed the outbound stream");
return;
}
Err(e) => {
error!("Error receiving AgentMessage: {:?}", e);
return;
}
}
}
});
Ok(tonic::Response::new(outbound))
}
async fn connection_test(

View File

@@ -0,0 +1,146 @@
use sea_orm::{ActiveModelTrait, ActiveValue::Set, DatabaseConnection, EntityTrait};
use uuid::Uuid;
use super::types::{AgentRecord, CreateAgentRecord, State, UpdateAgentRecord};
use crate::{
db::entities::agents::{ActiveModel as AgentActiveModel, Entity as Agent},
service::error::RepoError,
};
#[async_trait::async_trait]
pub(crate) trait AgentRepo: Send + Sync + 'static {
async fn list(&self) -> Result<Vec<AgentRecord>, RepoError>;
async fn get(&self, id: Uuid) -> Result<Option<AgentRecord>, RepoError>;
async fn create(&self, rec: &CreateAgentRecord) -> Result<AgentRecord, RepoError>;
async fn update(
&self,
id: Uuid,
rec: &UpdateAgentRecord,
) -> Result<Option<AgentRecord>, RepoError>;
async fn delete(&self, id: Uuid) -> Result<bool, RepoError>;
}
pub(crate) struct AgentRepoImpl {
db: DatabaseConnection,
}
impl AgentRepoImpl {
pub fn new(db: DatabaseConnection) -> Self {
Self { db }
}
}
#[async_trait::async_trait]
impl AgentRepo for AgentRepoImpl {
async fn list(&self) -> Result<Vec<AgentRecord>, RepoError> {
let agents = Agent::find().all(&self.db).await?;
Ok(agents
.into_iter()
.map(|m| AgentRecord {
id: m.id,
name: m.name,
ip_address: m.ip_address,
state: m.state.into(),
deployment_mode: m.deployment_mode,
last_seen_at: m.last_seen_at.map(|dt| dt.to_string()),
labels: m.labels,
created_at: m.created_at.to_string(),
updated_at: m.updated_at.to_string(),
})
.collect())
}
async fn get(&self, id: Uuid) -> Result<Option<AgentRecord>, RepoError> {
let agent = Agent::find_by_id(id).one(&self.db).await?;
Ok(agent.map(|m| AgentRecord {
id: m.id,
name: m.name,
ip_address: m.ip_address,
state: m.state.into(),
deployment_mode: m.deployment_mode,
last_seen_at: m.last_seen_at.map(|dt| dt.to_string()),
labels: m.labels,
created_at: m.created_at.to_string(),
updated_at: m.updated_at.to_string(),
}))
}
async fn create(&self, rec: &CreateAgentRecord) -> Result<AgentRecord, RepoError> {
let new_agent = AgentActiveModel {
id: Set(Uuid::new_v4()),
name: Set(rec.name.clone()),
ip_address: Set(rec.ip_address.clone()),
state: Set(State::Active.into()),
deployment_mode: Set(None),
last_seen_at: Set(None),
labels: Set(None),
created_at: Set(chrono::Utc::now().into()),
updated_at: Set(chrono::Utc::now().into()),
..Default::default()
};
let agent = new_agent.insert(&self.db).await?;
Ok(AgentRecord {
id: agent.id,
name: agent.name,
ip_address: agent.ip_address,
state: agent.state.into(),
deployment_mode: agent.deployment_mode,
last_seen_at: agent.last_seen_at.map(|dt| dt.to_string()),
labels: agent.labels,
created_at: agent.created_at.to_string(),
updated_at: agent.updated_at.to_string(),
})
}
async fn update(
&self,
id: Uuid,
rec: &UpdateAgentRecord,
) -> Result<Option<AgentRecord>, RepoError> {
let existing = match Agent::find_by_id(id).one(&self.db).await? {
Some(agent) => agent,
None => return Ok(None),
};
let mut agent: AgentActiveModel = AgentActiveModel::from(existing);
if let Some(name) = &rec.name {
agent.name = Set(name.clone());
}
if let Some(ip_address) = &rec.ip_address {
agent.ip_address = Set(Some(ip_address.clone()));
}
if let Some(state) = &rec.state {
agent.state = Set(String::from(*state));
}
if let Some(deployment_mode) = &rec.deployment_mode {
agent.deployment_mode = Set(Some(deployment_mode.clone()));
}
if let Some(labels) = &rec.labels {
agent.labels = Set(Some(labels.clone()));
}
agent.updated_at = Set(chrono::Utc::now().into());
let updated = agent.update(&self.db).await?;
Ok(Some(AgentRecord {
id: updated.id,
name: updated.name,
ip_address: updated.ip_address,
state: updated.state.into(),
deployment_mode: updated.deployment_mode,
last_seen_at: updated.last_seen_at.map(|dt| dt.to_string()),
labels: updated.labels,
created_at: updated.created_at.to_string(),
updated_at: updated.updated_at.to_string(),
}))
}
async fn delete(&self, id: Uuid) -> Result<bool, RepoError> {
let result = Agent::delete_by_id(id).exec(&self.db).await?;
Ok(result.rows_affected > 0)
}
}

View File

@@ -0,0 +1,60 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum State {
Active,
Inactive,
Unreachable,
Unknown,
Disabled,
}
impl From<State> for String {
fn from(state: State) -> Self {
match state {
State::Active => "active".to_string(),
State::Inactive => "inactive".to_string(),
State::Unreachable => "unreachable".to_string(),
State::Unknown => "unknown".to_string(),
State::Disabled => "disabled".to_string(),
}
}
}
impl From<String> for State {
fn from(s: String) -> Self {
match s.as_str() {
"active" => State::Active,
"inactive" => State::Inactive,
"unreachable" => State::Unreachable,
"unknown" => State::Unknown,
"disabled" => State::Disabled,
_ => State::Inactive,
}
}
}
pub struct AgentRecord {
pub id: uuid::Uuid,
pub name: String,
pub ip_address: Option<String>,
pub state: State,
pub deployment_mode: Option<String>,
pub last_seen_at: Option<String>,
pub labels: Option<serde_json::Value>,
pub created_at: String,
pub updated_at: String,
}
pub struct CreateAgentRecord {
pub name: String,
pub ip_address: Option<String>,
}
pub struct UpdateAgentRecord {
pub name: Option<String>,
pub ip_address: Option<String>,
pub state: Option<State>,
pub deployment_mode: Option<String>,
pub labels: Option<serde_json::Value>,
}