feagi_services/impls/

agent_service_impl.rs

// Copyright 2025 Neuraville Inc.
// Licensed under the Apache License, Version 2.0

//! Agent service implementation
//!
//! This implementation delegates all agent management to the PNS AgentRegistry,
//! ensuring centralized coordination consistent with the Python implementation.

use async_trait::async_trait;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{debug, error, info, warn};

use crate::traits::agent_service::*;
use crate::traits::registration_handler::RegistrationHandlerTrait;
use crate::traits::RuntimeService as RuntimeServiceTrait;
use crate::types::agent_registry::AgentRegistry;
use crate::types::registration::RegistrationRequest;
use feagi_brain_development::ConnectomeManager;
use feagi_structures::genomic::cortical_area::CorticalID;

/// Implementation of the Agent service
pub struct AgentServiceImpl {
    connectome_manager: Arc<RwLock<ConnectomeManager>>,
    agent_registry: Arc<RwLock<AgentRegistry>>,
    registration_handler: Option<Arc<dyn RegistrationHandlerTrait>>,
    runtime_service: Arc<RwLock<Option<Arc<dyn RuntimeServiceTrait + Send + Sync>>>>,
}

impl AgentServiceImpl {
    pub fn new(
        connectome_manager: Arc<RwLock<ConnectomeManager>>,
        agent_registry: Arc<RwLock<AgentRegistry>>,
    ) -> Self {
        Self {
            connectome_manager,
            agent_registry,
            registration_handler: None,
            runtime_service: Arc::new(RwLock::new(None)),
        }
    }

    /// Create AgentServiceImpl with runtime service
    pub fn new_with_runtime(
        connectome_manager: Arc<RwLock<ConnectomeManager>>,
        agent_registry: Arc<RwLock<AgentRegistry>>,
        runtime_service: Arc<dyn RuntimeServiceTrait + Send + Sync>,
    ) -> Self {
        Self {
            connectome_manager,
            agent_registry,
            registration_handler: None,
            runtime_service: Arc::new(RwLock::new(Some(runtime_service))),
        }
    }

    /// Set the PNS registration handler for full transport negotiation
    /// Set registration handler (accepts trait object to break circular dependency)
    pub fn set_registration_handler(&mut self, handler: Arc<dyn RegistrationHandlerTrait>) {
        self.registration_handler = Some(handler);
        info!("🦀 [AGENT-SERVICE] Registration handler connected");
    }

    /// Set the runtime service for sensory injection (thread-safe, can be called after Arc wrapping)
    pub fn set_runtime_service(&self, runtime_service: Arc<dyn RuntimeServiceTrait + Send + Sync>) {
        let mut guard = self.runtime_service.write();
        if let Some(existing) = guard.as_ref() {
            if Arc::ptr_eq(existing, &runtime_service) {
                debug!(
                    "🦀 [AGENT-SERVICE] Runtime service already connected; ignoring duplicate bind"
                );
                return;
            }
            warn!("🦀 [AGENT-SERVICE] Runtime service replaced with a new instance");
        } else {
            info!("🦀 [AGENT-SERVICE] Runtime service connected");
        }
        *guard = Some(runtime_service);
    }
}

#[async_trait]
impl AgentService for AgentServiceImpl {
    async fn register_agent(
        &self,
        registration: AgentRegistration,
    ) -> AgentResult<AgentRegistrationResponse> {
        info!(
            "🦀 [AGENT-SERVICE] Registering agent: {} (type: {})",
            registration.agent_id, registration.agent_type
        );

        // If we have a registration handler, use it (gets full transport info)
        if let Some(handler) = &self.registration_handler {
            info!(
                "📝 [AGENT-SERVICE] Using PNS registration handler for full transport negotiation"
            );

            // Build PNS registration request
            let pns_request = RegistrationRequest {
                agent_id: registration.agent_id.clone(),
                agent_type: registration.agent_type.clone(),
                capabilities: serde_json::to_value(&registration.capabilities)
                    .unwrap_or(serde_json::Value::Object(serde_json::Map::new())),
                chosen_transport: registration.chosen_transport.clone(), // Pass through the agent's transport choice
            };

            // Call registration handler - use spawn_blocking to avoid blocking the async runtime
            // The process_registration method is sync and may block, so we offload it to a blocking thread pool
            let handler_clone = handler.clone();
            let pns_response = tokio::task::spawn_blocking(move || {
                handler_clone.process_registration(pns_request)
            })
            .await
            .map_err(|e| {
                AgentError::RegistrationFailed(format!("Registration task panicked: {:?}", e))
            })?
            .map_err(AgentError::RegistrationFailed)?;

            // Convert PNS transport configs to service transport configs
            let transports = pns_response.transports.map(|ts| {
                ts.into_iter()
                    .map(|t| TransportConfig {
                        transport_type: t.transport_type,
                        enabled: t.enabled,
                        ports: t.ports,
                        host: t.host,
                    })
                    .collect()
            });

            // Serialize cortical areas with proper error handling (don't panic!)
            let cortical_areas_json =
                serde_json::to_value(&pns_response.cortical_areas).map_err(|e| {
                    error!(
                        "❌ [AGENT-SERVICE] Failed to serialize cortical areas: {}",
                        e
                    );
                    AgentError::Internal(format!("Failed to serialize cortical areas: {}", e))
                })?;

            return Ok(AgentRegistrationResponse {
                status: pns_response.status,
                message: pns_response
                    .message
                    .unwrap_or_else(|| "Success".to_string()),
                success: true,
                transport: None,
                rates: None,
                transports,
                recommended_transport: pns_response.recommended_transport,
                shm_paths: pns_response.shm_paths,
                cortical_areas: cortical_areas_json,
            });
        }

        return Err(AgentError::RegistrationFailed(
            "Registration handler not available - required for FEAGI 2.0".to_string(),
        ));
    }

    async fn heartbeat(&self, request: HeartbeatRequest) -> AgentResult<()> {
        // Check if agent exists before attempting heartbeat
        let agent_exists = {
            let registry = self.agent_registry.read();
            registry.get(&request.agent_id).is_some()
        };

        if !agent_exists {
            // Log diagnostic information when agent not found
            let all_agents: Vec<String> = {
                let registry = self.agent_registry.read();
                registry
                    .get_all()
                    .iter()
                    .map(|a| a.agent_id.clone())
                    .collect()
            };
            warn!(
                "⚠️ [HEARTBEAT] Agent '{}' not found in registry. Registered agents ({}): {:?}",
                request.agent_id,
                all_agents.len(),
                all_agents
            );
            return Err(AgentError::NotFound(format!(
                "Agent {} not found in registry (total registered: {})",
                request.agent_id,
                all_agents.len()
            )));
        }

        // Agent exists - update heartbeat
        self.agent_registry
            .write()
            .heartbeat(&request.agent_id)
            .map_err(|e| {
                error!(
                    "❌ [HEARTBEAT] Failed to update heartbeat for '{}': {}",
                    request.agent_id, e
                );
                AgentError::NotFound(e)
            })?;
        Ok(())
    }

    async fn list_agents(&self) -> AgentResult<Vec<String>> {
        tracing::trace!(target: "feagi-services", "list_agents() called - acquiring registry read lock...");
        let registry = self.agent_registry.read();
        let agents = registry.get_all();
        let agent_ids: Vec<String> = agents.iter().map(|a| a.agent_id.clone()).collect();

        tracing::trace!(
            target: "feagi-services",
            "list_agents() found {} agents: {:?}",
            agent_ids.len(),
            agent_ids
        );
        tracing::trace!(
            target: "feagi-services",
            "AgentRegistry pointer: {:p}",
            &*self.agent_registry as *const _
        );

        Ok(agent_ids)
    }

    async fn get_agent_properties(&self, agent_id: &str) -> AgentResult<AgentProperties> {
        let registry = self.agent_registry.read();
        let agent = registry
            .get(agent_id)
            .ok_or_else(|| AgentError::NotFound(format!("Agent {} not found", agent_id)))?;

        // Extract properties from agent info
        let agent_ip = agent
            .metadata
            .get("agent_ip")
            .and_then(|v| v.as_str())
            .unwrap_or("127.0.0.1")
            .to_string();

        let agent_data_port = agent
            .metadata
            .get("agent_data_port")
            .and_then(|v| v.as_u64())
            .unwrap_or(0) as u16;

        let agent_version = agent
            .metadata
            .get("agent_version")
            .and_then(|v| v.as_str())
            .unwrap_or("unknown")
            .to_string();

        let controller_version = agent
            .metadata
            .get("controller_version")
            .and_then(|v| v.as_str())
            .unwrap_or("unknown")
            .to_string();

        let agent_router_address = format!("tcp://{}:{}", agent_ip, agent_data_port);

        // Build full capabilities map using feagi-sensorimotor format: {"input": [...], "output": [...]}
        let mut capabilities = HashMap::new();

        // Add vision capability if present
        if let Some(ref vision) = agent.capabilities.vision {
            capabilities.insert(
                "vision".to_string(),
                serde_json::to_value(vision).unwrap_or(serde_json::Value::Null),
            );
        }

        // Add output capability (feagi-sensorimotor format)
        if let Some(ref motor) = agent.capabilities.motor {
            // Convert motor capability to "output" format: array of cortical IDs
            let output_areas: Vec<String> = motor.source_cortical_areas.clone();
            capabilities.insert(
                "output".to_string(),
                serde_json::to_value(output_areas).unwrap_or(serde_json::Value::Null),
            );
        }

        // Add visualization capability if present
        if let Some(ref viz) = agent.capabilities.visualization {
            capabilities.insert(
                "visualization".to_string(),
                serde_json::to_value(viz).unwrap_or(serde_json::Value::Null),
            );
        }

        // Add input capability (feagi-sensorimotor format)
        // Note: cortical_mappings removed - device registrations are handled separately
        if let Some(ref _sensory) = agent.capabilities.sensory {
            // Return empty array since cortical_mappings no longer exist
            // Device registrations should be accessed via device_registrations endpoint
            capabilities.insert(
                "input".to_string(),
                serde_json::to_value(Vec::<String>::new()).unwrap_or(serde_json::Value::Null),
            );
        }

        // Add custom capabilities
        for (k, v) in agent.capabilities.custom.iter() {
            capabilities.insert(k.clone(), v.clone());
        }

        Ok(AgentProperties {
            agent_type: agent.agent_type.to_string(),
            agent_ip,
            agent_data_port,
            agent_router_address,
            agent_version,
            controller_version,
            capabilities,
            chosen_transport: agent.chosen_transport.clone(),
        })
    }

    async fn get_shared_memory_info(
        &self,
    ) -> AgentResult<HashMap<String, HashMap<String, serde_json::Value>>> {
        // TODO: Implement shared memory tracking in agent registry
        Ok(HashMap::new())
    }

    async fn deregister_agent(&self, agent_id: &str) -> AgentResult<()> {
        self.agent_registry
            .write()
            .deregister(agent_id)
            .map_err(AgentError::NotFound)?;

        info!(
            "✅ [AGENT-SERVICE] Agent '{}' deregistered successfully",
            agent_id
        );
        Ok(())
    }

    async fn manual_stimulation(
        &self,
        stimulation_payload: HashMap<String, Vec<Vec<i32>>>,
        mode: ManualStimulationMode,
    ) -> AgentResult<HashMap<String, serde_json::Value>> {
        // Use RuntimeService for sensory injection (service layer, not direct NPU access)
        let runtime_service = self
            .runtime_service
            .read()
            .as_ref()
            .ok_or_else(|| {
                AgentError::Internal(
                    "Runtime service not available - cannot inject stimuli".to_string(),
                )
            })?
            .clone();

        let mut result = HashMap::new();
        let mut total_stimulated = 0usize;
        let mut requested_coordinates = 0usize;
        let mut successful_areas = 0;
        let mut failed_areas = Vec::new();
        let mut coordinates_not_found = 0;

        // Default potential for manual stimulation (high enough to trigger firing)
        const DEFAULT_POTENTIAL: f32 = 100.0;

        // First pass: validate all cortical areas and build injection data
        let mut injection_requests: Vec<(String, Vec<(u32, u32, u32, f32)>)> = Vec::new();

        {
            let manager = self.connectome_manager.read();

            for (cortical_id, coordinates) in stimulation_payload.iter() {
                requested_coordinates += coordinates.len();
                let cortical_id_typed = match CorticalID::try_from_base_64(cortical_id) {
                    Ok(id) => id,
                    Err(e) => {
                        failed_areas.push(cortical_id.clone());
                        warn!("Invalid cortical ID '{}': {}", cortical_id, e);
                        continue;
                    }
                };

                match manager.get_cortical_area(&cortical_id_typed) {
                    Some(_area) => {
                        // Build xyzp_data for this cortical area (coordinates with potential)
                        let mut xyzp_data = Vec::new();

                        for coord in coordinates {
                            if coord.len() != 3 {
                                warn!(
                                    "Invalid coordinate format: {:?} (expected [x, y, z])",
                                    coord
                                );
                                coordinates_not_found += 1;
                                continue;
                            }

                            let x = coord[0] as u32;
                            let y = coord[1] as u32;
                            let z = coord[2] as u32;

                            // Add coordinate with default potential
                            xyzp_data.push((x, y, z, DEFAULT_POTENTIAL));
                        }

                        if !xyzp_data.is_empty() {
                            injection_requests.push((cortical_id.clone(), xyzp_data));
                            successful_areas += 1;
                        }
                    }
                    None => {
                        failed_areas.push(cortical_id.clone());
                    }
                }
            }
        } // Drop manager lock here before await

        // Second pass: perform injections (no locks held)
        for (cortical_id, xyzp_data) in injection_requests {
            match runtime_service
                .inject_sensory_by_coordinates(
                    &cortical_id,
                    &xyzp_data,
                    match mode {
                        ManualStimulationMode::Candidate => {
                            crate::traits::runtime_service::ManualStimulationMode::Candidate
                        }
                        ManualStimulationMode::ForceFire => {
                            crate::traits::runtime_service::ManualStimulationMode::ForceFire
                        }
                    },
                )
                .await
            {
                Ok(injected_count) => {
                    total_stimulated += injected_count;
                    if injected_count < xyzp_data.len() {
                        coordinates_not_found += xyzp_data.len() - injected_count;
                    }
                }
                Err(e) => {
                    error!(
                        "❌ [MANUAL-STIMULATION] Failed to inject for area {}: {}",
                        cortical_id, e
                    );
                    coordinates_not_found += xyzp_data.len();
                }
            }
        }

        result.insert(
            "success".to_string(),
            serde_json::json!(failed_areas.is_empty() && coordinates_not_found == 0),
        );
        result.insert(
            "total_coordinates".to_string(),
            serde_json::json!(total_stimulated),
        );
        result.insert(
            "requested_coordinates".to_string(),
            serde_json::json!(requested_coordinates),
        );
        result.insert(
            "matched_coordinates".to_string(),
            serde_json::json!(total_stimulated),
        );
        result.insert(
            "unique_neuron_ids".to_string(),
            serde_json::json!(total_stimulated),
        );
        result.insert(
            "mode".to_string(),
            serde_json::json!(match mode {
                ManualStimulationMode::Candidate => "candidate",
                ManualStimulationMode::ForceFire => "force_fire",
            }),
        );
        result.insert(
            "successful_areas".to_string(),
            serde_json::json!(successful_areas),
        );
        result.insert("failed_areas".to_string(), serde_json::json!(failed_areas));

        if coordinates_not_found > 0 {
            result.insert(
                "coordinates_not_found".to_string(),
                serde_json::json!(coordinates_not_found),
            );
        }

        if !failed_areas.is_empty() {
            result.insert(
                "error".to_string(),
                serde_json::json!(format!("Some cortical areas not found: {:?}", failed_areas)),
            );
        }

        Ok(result)
    }

    fn try_set_runtime_service(&self, runtime_service: Arc<dyn RuntimeServiceTrait + Send + Sync>) {
        self.set_runtime_service(runtime_service);
    }
}