use crate::{AgentCapabilities, AgentDescriptor, AuthToken, FeagiApiVersion};
use feagi_io::traits_and_enums::shared::{
TransportProtocolEndpoint, TransportProtocolImplementation,
};
use feagi_io::AgentID;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentRegistrationMessage {
ClientRequestRegistration(RegistrationRequest),
ServerRespondsRegistration(RegistrationResponse),
ClientRequestDeregistration(DeregistrationRequest),
ServerRespondsDeregistration(DeregistrationResponse),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RegistrationRequest {
agent_descriptor: AgentDescriptor,
auth_token: AuthToken,
requested_capabilities: Vec<AgentCapabilities>,
connection_protocol: TransportProtocolImplementation,
api_version: FeagiApiVersion,
}
impl RegistrationRequest {
pub const MAX_REQUEST_SIZE: usize = 1024;
pub fn new(
agent_descriptor: AgentDescriptor,
auth_token: AuthToken,
requested_capabilities: Vec<AgentCapabilities>,
connection_protocol: TransportProtocolImplementation,
) -> Self {
Self {
agent_descriptor,
auth_token,
requested_capabilities,
connection_protocol,
api_version: FeagiApiVersion::get_current_api_version(),
}
}
pub fn api_version(&self) -> &FeagiApiVersion {
&self.api_version
}
pub fn agent_descriptor(&self) -> &AgentDescriptor {
&self.agent_descriptor
}
pub fn auth_token(&self) -> &AuthToken {
&self.auth_token
}
pub fn requested_capabilities(&self) -> &[AgentCapabilities] {
&self.requested_capabilities
}
pub fn connection_protocol(&self) -> &TransportProtocolImplementation {
&self.connection_protocol
}
}
#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RegistrationResponse {
FailedInvalidRequest, FailedInvalidAuth, AlreadyRegistered,
Success(
AgentID,
HashMap<AgentCapabilities, TransportProtocolEndpoint>,
),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct DeregistrationRequest {
reason: Option<String>,
}
impl DeregistrationRequest {
pub fn new(reason: Option<String>) -> Self {
Self { reason }
}
pub fn reason(&self) -> Option<&str> {
self.reason.as_deref()
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DeregistrationResponse {
Success,
NotRegistered,
}
#[cfg(test)]
mod tests {
use super::{AgentRegistrationMessage, DeregistrationRequest, DeregistrationResponse};
#[test]
fn deregistration_request_round_trip_serialization_preserves_reason() {
let request = AgentRegistrationMessage::ClientRequestDeregistration(
DeregistrationRequest::new(Some("shutdown".to_string())),
);
let encoded =
serde_json::to_string(&request).expect("deregistration request should serialize");
let decoded: AgentRegistrationMessage =
serde_json::from_str(&encoded).expect("deregistration request should deserialize");
assert_eq!(request, decoded);
}
#[test]
fn deregistration_response_round_trip_serialization_preserves_variant() {
let response = AgentRegistrationMessage::ServerRespondsDeregistration(
DeregistrationResponse::NotRegistered,
);
let encoded =
serde_json::to_string(&response).expect("deregistration response should serialize");
let decoded: AgentRegistrationMessage =
serde_json::from_str(&encoded).expect("deregistration response should deserialize");
assert_eq!(response, decoded);
}
}