MetaCrate/crates/metacrate-grid-agent/src/main.rs

use metacrate_grid_agent::{
    AgentService, ConfigLoader, GridEventKind, ObservableEvent, OperatingMode,
};
use std::error::Error;
use std::fmt;
use std::path::PathBuf;
#[cfg(feature = "live-grid")]
use std::sync::Arc;
#[cfg(feature = "live-grid")]
use std::time::{SystemTime, UNIX_EPOCH};

const MAX_ARGUMENTS: usize = 8;

#[derive(Debug)]
struct CliError(String);

impl fmt::Display for CliError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl Error for CliError {}

#[derive(Default)]
struct Options {
    config: Option<PathBuf>,
    check_config: bool,
    run_once: bool,
}

fn options() -> Result<Option<Options>, CliError> {
    let mut result = Options::default();
    let mut arguments = std::env::args_os().skip(1);
    let mut count = 0;
    while let Some(argument) = arguments.next() {
        count += 1;
        if count > MAX_ARGUMENTS {
            return Err(CliError(format!(
                "at most {MAX_ARGUMENTS} command-line arguments are accepted"
            )));
        }
        if argument == "--help" || argument == "-h" {
            println!(
                "metacrate-grid-agent [--config PATH] [--check-config | --run-once]\n\
                 Configuration precedence: defaults < JSON < secret files < environment."
            );
            return Ok(None);
        }
        if argument == "--check-config" {
            result.check_config = true;
        } else if argument == "--run-once" {
            result.run_once = true;
        } else if argument == "--config" {
            let path = arguments
                .next()
                .ok_or_else(|| CliError("--config requires a path".into()))?;
            count += 1;
            result.config = Some(PathBuf::from(path));
        } else {
            let argument = PathBuf::from(argument);
            return Err(CliError(format!(
                "unknown argument {}; use --help",
                argument.display()
            )));
        }
    }
    if result.check_config && result.run_once {
        return Err(CliError(
            "--check-config and --run-once are mutually exclusive".into(),
        ));
    }
    Ok(Some(result))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let Some(options) = options()? else {
        return Ok(());
    };
    let mut loader = ConfigLoader::new();
    if let Some(path) = options.config {
        loader = loader.with_file(path);
    }
    let config = loader.load()?;
    if options.check_config {
        println!("configuration is valid for {:?} mode", config.mode);
        return Ok(());
    }
    if config.mode != OperatingMode::OfflineFake {
        #[cfg(feature = "live-grid")]
        return run_live(config, options.run_once).await;
        #[cfg(not(feature = "live-grid"))]
        return Err(CliError(
            "live grid mode requires rebuilding with --features live-grid".into(),
        )
        .into());
    }

    let startup_timeout = config.timeouts.startup;
    let mut handle = AgentService::offline(config)?.start()?;
    if options.run_once {
        tokio::time::timeout(startup_timeout, async {
            loop {
                match handle.next_event().await {
                    Some(ObservableEvent::Grid(event))
                        if event.kind == GridEventKind::BackendReady =>
                    {
                        return Ok::<(), CliError>(());
                    }
                    Some(_) => {}
                    None => {
                        return Err(CliError("service stopped before backend readiness".into()));
                    }
                }
            }
        })
        .await
        .map_err(|_| CliError("timed out waiting for offline backend readiness".into()))??;
        handle.shutdown().await?;
        println!("grid agent completed one offline startup/shutdown cycle");
        return Ok(());
    }
    println!("grid agent started in offline/fake mode; press Ctrl-C to stop");
    loop {
        tokio::select! {
            signal = tokio::signal::ctrl_c() => {
                signal?;
                break;
            }
            event = handle.next_event() => {
                if event.is_none() {
                    break;
                }
            }
        }
    }
    handle.shutdown().await?;
    println!("grid agent stopped cleanly");
    Ok(())
}

#[cfg(feature = "live-grid")]
#[allow(clippy::too_many_lines)]
async fn run_live(
    config: metacrate_grid_agent::AgentConfig,
    run_once: bool,
) -> Result<(), Box<dyn Error>> {
    use metacrate_grid_agent::{
        GridSessionBackend, LibremetaverseClientOwner, SessionObservation, SessionState,
        SessionSupervisor,
    };

    let connection = config.grid.clone().ok_or_else(|| {
        CliError("validated live configuration did not contain a grid connection".into())
    })?;
    let owner = LibremetaverseClientOwner::new()?;
    let mut live = start_live_interactions(&config, &owner)?;
    let backend = match owner.session_backend_with_agent_services(
        connection,
        live.interaction.ingress(),
        live.perception.clone(),
        live.behavior.ingress(),
    ) {
        Ok(backend) => backend,
        Err(error) => {
            live.interaction.shutdown().await?;
            live.behavior.shutdown().await?;
            return Err(error.into());
        }
    };
    let erased: Arc<dyn GridSessionBackend> = Arc::new(backend);
    let mut handle = SessionSupervisor::new(
        erased,
        config.reconnect,
        config.limits.control_queue,
        config.limits.observable_queue,
    )?
    .start();

    if run_once {
        let readiness = tokio::time::timeout(config.timeouts.startup, async {
            loop {
                match handle.next_observation().await {
                    Some(SessionObservation::Transition { status, .. }) if status.agent_ready => {
                        return Ok::<(), CliError>(());
                    }
                    Some(SessionObservation::Transition {
                        status:
                            metacrate_grid_agent::SessionStatus {
                                state: SessionState::AuthenticationBlocked,
                                ..
                            },
                        ..
                    }) => {
                        return Err(CliError(
                            "grid authentication/configuration requires operator action".into(),
                        ));
                    }
                    Some(_) => {}
                    None => {
                        return Err(CliError(
                            "session supervisor stopped before readiness".into(),
                        ));
                    }
                }
            }
        })
        .await;
        let readiness = match readiness {
            Ok(result) => result,
            Err(_) => Err(CliError("timed out waiting for full grid readiness".into())),
        };
        if let Err(error) = readiness {
            let session_result = handle.shutdown().await;
            let interaction_result = live.interaction.shutdown().await;
            let behavior_result = live.behavior.shutdown().await;
            session_result?;
            interaction_result?;
            behavior_result?;
            return Err(error.into());
        }
        let session_result = handle.shutdown().await;
        let interaction_result = live.interaction.shutdown().await;
        let behavior_result = live.behavior.shutdown().await;
        session_result?;
        interaction_result?;
        behavior_result?;
        println!("grid agent completed one supervised login/logout cycle");
        return Ok(());
    }

    println!("grid agent session supervisor started; press Ctrl-C to stop");
    let mut signal_error = None;
    loop {
        tokio::select! {
            signal = tokio::signal::ctrl_c() => {
                if let Err(error) = signal {
                    signal_error = Some(error);
                }
                break;
            }
            event = handle.next_observation() => {
                let Some(event) = event else { break; };
                if let SessionObservation::Transition { status, reason, retry_in } = event {
                    println!(
                        "grid session state={} generation={} transport_connected={} agent_ready={} reason={reason:?} retry_in={retry_in:?}",
                        status.state.as_str(),
                        status.generation,
                        status.transport_connected,
                        status.agent_ready,
                    );
                }
            }
            event = live.interaction.next_observation() => {
                let Some(event) = event else { break; };
                println!("grid interaction event={event:?}");
            }
            event = live.perception_observations.recv() => {
                let Some(event) = event else { break; };
                println!("grid perception event={event:?}");
            }
            event = live.behavior.next_observation() => {
                let Some(event) = event else { break; };
                println!("grid behavior event={event:?}");
            }
        }
    }
    let session_result = handle.shutdown().await;
    let interaction_result = live.interaction.shutdown().await;
    let behavior_result = live.behavior.shutdown().await;
    session_result?;
    interaction_result?;
    behavior_result?;
    if let Some(error) = signal_error {
        return Err(error.into());
    }
    println!("grid agent stopped cleanly");
    Ok(())
}

#[cfg(feature = "live-grid")]
struct LiveInteractions {
    interaction: metacrate_grid_agent::InteractionHandle,
    perception: metacrate_grid_agent::PerceptionIngress,
    perception_observations:
        tokio::sync::mpsc::Receiver<metacrate_grid_agent::PerceptionObservation>,
    behavior: metacrate_grid_agent::BehaviorHandle,
}

#[cfg(feature = "live-grid")]
fn start_live_interactions(
    config: &metacrate_grid_agent::AgentConfig,
    owner: &metacrate_grid_agent::LibremetaverseClientOwner,
) -> Result<LiveInteractions, Box<dyn Error>> {
    use metacrate_grid_agent::{
        AuthorizedBackendRouter, AuthorizedToolBackend, BehaviorBackend, BehaviorController,
        ConversationStore, InteractionCoordinator, LlmClient, LlmTransportLimits,
        MemoryPolicyAudit, PerceptionBackend, PolicyGateway, PolicyLimits, PolicyLlmResponder,
        ToolLoopLimits, behavior_policy_tools, perception_policy_tools,
    };

    let transport_limits = LlmTransportLimits {
        request_timeout: config.timeouts.request,
        total_timeout: config.interaction.model_timeout,
        max_prompt_bytes: config.limits.max_body_bytes,
        max_response_bytes: config.limits.max_body_bytes,
        max_concurrent_requests: config.interaction.max_concurrent_inference,
        ..LlmTransportLimits::default()
    };
    let conversation = Arc::new(ConversationStore::from_config(config)?);
    let perception = Arc::new(PerceptionBackend::new(
        Arc::new(owner.world_snapshot_source()),
        Arc::clone(&conversation),
        config.limits.observable_queue,
    )?);
    let perception_ingress = perception.ingress();
    let perception_observations = perception
        .take_observations()
        .ok_or_else(|| CliError("perception observation receiver already claimed".into()))?;
    let behavior = BehaviorController::new(
        config.behavior.clone(),
        Arc::new(owner.embodiment_sink()),
        config.limits.control_queue,
        config.limits.observable_queue,
        config.timeouts.shutdown,
    )?
    .start();
    let behavior_ingress = behavior.ingress();
    let client = Arc::new(LlmClient::new(config.llm.clone(), transport_limits)?);
    let audit = Arc::new(MemoryPolicyAudit::new(config.limits.observable_queue)?);
    let mut tools = perception_policy_tools()?;
    tools.extend(behavior_policy_tools(&config.behavior)?);
    let routes = tools
        .iter()
        .map(|tool| tool.definition.name.as_str().to_owned())
        .map(|name| {
            let backend: Arc<dyn AuthorizedToolBackend> = if name.starts_with("behavior_") {
                Arc::new(BehaviorBackend::new(behavior_ingress.clone()))
            } else {
                perception.clone()
            };
            (name, backend)
        })
        .collect::<Vec<_>>();
    let gateway = Arc::new(PolicyGateway::new(
        config.authorized_avatar_uuids.clone(),
        tools,
        PolicyLimits::default(),
        audit,
    )?);
    let loop_limits = ToolLoopLimits {
        max_tool_calls_per_turn: config.limits.max_tool_calls,
        max_tool_calls_per_session: config.limits.max_tool_calls,
        max_history_messages: config.conversation.limits.max_turns_per_session,
        max_history_bytes: config.conversation.limits.max_session_bytes,
        wall_clock_timeout: config.interaction.model_timeout,
        ..ToolLoopLimits::default()
    };
    let now = Arc::new(|| {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_or(0, |duration| duration.as_secs())
    });
    let routed_backend: Arc<dyn AuthorizedToolBackend> =
        Arc::new(AuthorizedBackendRouter::new(routes)?);
    let responder = Arc::new(PolicyLlmResponder::new(
        client,
        gateway,
        routed_backend,
        loop_limits,
        now,
    )?);
    let sink = Arc::new(owner.interaction_sink());
    let pacer: Arc<dyn metacrate_grid_agent::ResponsePacer> = Arc::new(behavior_ingress);
    let interaction = InteractionCoordinator::new(
        config.interaction.clone(),
        libremetaverse_types::UUID::zero(),
        config.authorized_avatar_uuids.clone(),
        conversation,
        responder,
        sink,
        config.limits.grid_event_queue,
        config.limits.observable_queue,
        config.timeouts.shutdown,
    )?
    .with_response_pacer(pacer)
    .start();
    Ok(LiveInteractions {
        interaction,
        perception: perception_ingress,
        perception_observations,
        behavior,
    })
}