wactorz_agents/

dynamic_agent.rs

//! Dynamic runtime-script agent.
//!
//! [`DynamicAgent`] executes **Rhai** scripts generated by the LLM at runtime.
//! Rhai is a safe, sandboxed scripting language for Rust — it replaces the
//! Python `exec()` approach with proper isolation.
//!
//! The Rhai engine exposes a limited API to scripts:
//! - `agent_log(msg)` — log a message
//! - `agent_alert(msg)` — broadcast an alert
//! - `agent_state_get(key)` — read from persistent state
//! - `agent_state_set(key, value)` — write to persistent state
//!
//! Scripts must define a `fn main()` function that is called on each message.

use anyhow::Result;
use async_trait::async_trait;
use rhai::{AST, Engine};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc;

use crate::llm_agent::{LlmAgent, LlmConfig};
use wactorz_core::{Actor, ActorConfig, ActorMetrics, ActorState, EventPublisher, Message};

/// A dynamically scripted agent.
pub struct DynamicAgent {
    config: ActorConfig,
    /// The compiled Rhai AST for the current script.
    script_ast: Option<AST>,
    /// Raw source of the current script (for re-compilation after reload).
    script_source: String,
    /// Human-readable description, used as LLM system-prompt context.
    description: String,
    /// Optional LLM backend — called when the Rhai script produces no output.
    llm: Option<LlmAgent>,
    /// Persistent key-value state accessible from scripts.
    agent_state: std::sync::Arc<std::sync::Mutex<HashMap<String, serde_json::Value>>>,
    log_queue: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
    engine: Engine,
    state: ActorState,
    metrics: Arc<ActorMetrics>,
    mailbox_tx: mpsc::Sender<Message>,
    mailbox_rx: Option<mpsc::Receiver<Message>>,
    publisher: Option<EventPublisher>,
}

impl DynamicAgent {
    /// Create a new DynamicAgent with the given Rhai script source.
    pub fn new(config: ActorConfig, script_source: impl Into<String>) -> Self {
        let (tx, rx) = mpsc::channel(config.mailbox_capacity);
        let agent_state = std::sync::Arc::new(std::sync::Mutex::new(HashMap::<
            String,
            serde_json::Value,
        >::new()));
        let log_queue = std::sync::Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
        let mut agent = Self {
            config,
            script_ast: None,
            script_source: script_source.into(),
            description: String::new(),
            llm: None,
            agent_state,
            log_queue,
            engine: Engine::new(),
            state: ActorState::Initializing,
            metrics: Arc::new(ActorMetrics::new()),
            mailbox_tx: tx,
            mailbox_rx: Some(rx),
            publisher: None,
        };
        agent.register_api();
        agent
    }

    /// Attach an EventPublisher for MQTT output.
    pub fn with_publisher(mut self, p: EventPublisher) -> Self {
        self.publisher = Some(p);
        self
    }

    /// Attach an LLM backend used as fallback when the Rhai script produces no output.
    pub fn with_llm(mut self, llm_config: LlmConfig, description: String) -> Self {
        let llm_cfg = ActorConfig::new(format!("{}-llm", self.config.name));
        self.llm = Some(LlmAgent::new(llm_cfg, llm_config));
        self.description = description;
        self
    }

    fn now_ms() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64
    }

    /// Register the agent API functions into the Rhai engine.
    fn register_api(&mut self) {
        let state_r = std::sync::Arc::clone(&self.agent_state);
        let state_w = std::sync::Arc::clone(&self.agent_state);
        let log_q = std::sync::Arc::clone(&self.log_queue);
        let log_q2 = std::sync::Arc::clone(&self.log_queue);

        self.engine
            .register_fn("agent_state_get", move |key: &str| -> rhai::Dynamic {
                let map = state_r.lock().unwrap();
                match map.get(key) {
                    Some(v) => rhai::serde::to_dynamic(v.clone()).unwrap_or(rhai::Dynamic::UNIT),
                    None => rhai::Dynamic::UNIT,
                }
            });

        self.engine
            .register_fn("agent_state_set", move |key: &str, value: rhai::Dynamic| {
                if let Ok(json) = rhai::serde::from_dynamic::<serde_json::Value>(&value) {
                    state_w.lock().unwrap().insert(key.to_string(), json);
                }
            });

        self.engine.register_fn("agent_log", move |msg: &str| {
            log_q.lock().unwrap().push(msg.to_string());
        });

        self.engine.register_fn("agent_alert", move |msg: &str| {
            log_q2.lock().unwrap().push(format!("ALERT: {msg}"));
            tracing::warn!("DynamicAgent alert: {msg}");
        });

        // Type helpers: allow `"text" + any_value` and `str(x)` in scripts.
        // LLM-generated scripts often write `"= " + result` without calling
        // `.to_string()` first; these registrations make that safe.
        self.engine
            .register_fn("+", |a: String, b: rhai::Dynamic| -> String {
                format!("{a}{b}")
            });
        self.engine
            .register_fn("str", |v: rhai::Dynamic| -> String { v.to_string() });
    }

    /// Compile (or recompile) the stored script source.
    pub fn compile_script(&mut self) -> Result<()> {
        let ast = self
            .engine
            .compile(&self.script_source)
            .map_err(|e| anyhow::anyhow!("Rhai compile error: {e}"))?;
        self.script_ast = Some(ast);
        Ok(())
    }

    /// Hot-reload a new script, recompiling it immediately.
    pub fn reload_script(&mut self, new_source: impl Into<String>) -> Result<()> {
        self.script_source = new_source.into();
        self.compile_script()
    }

    /// LLM fallback when Rhai produces no output.
    async fn llm_respond(&mut self, user_input: &str) -> String {
        if let Some(llm) = &mut self.llm {
            let prompt = format!(
                "You are {}. {}\n\nUser: {}",
                self.config.name,
                if self.description.is_empty() {
                    "A helpful AI agent."
                } else {
                    &self.description
                },
                user_input,
            );
            llm.complete(&prompt)
                .await
                .unwrap_or_else(|e| format!("LLM error: {e}"))
        } else {
            format!(
                "I received '{}' but my script produced no response. \
                 Ask main-actor to give me a better script.",
                &user_input[..user_input.len().min(60)]
            )
        }
    }

    /// LLM fallback when Rhai errors out; includes the error for context.
    async fn llm_respond_with_error(&mut self, user_input: &str, script_err: &str) -> String {
        if let Some(llm) = &mut self.llm {
            let prompt = format!(
                "You are {}. {}\n\
                 Your Rhai script had an error: {}\n\
                 Respond helpfully to the user anyway.\n\nUser: {}",
                self.config.name,
                if self.description.is_empty() {
                    "A helpful AI agent."
                } else {
                    &self.description
                },
                script_err,
                user_input,
            );
            llm.complete(&prompt)
                .await
                .unwrap_or_else(|e| format!("LLM error: {e}"))
        } else {
            format!("script error: {script_err}")
        }
    }

    /// Invoke the script's `main(message_json)` function.
    ///
    /// Fallback: if the script does not define `fn main`, evaluates the whole
    /// script top-to-bottom with `msg` pre-populated in the scope.  This lets
    /// the LLM omit the function wrapper and still get a response.
    fn run_script(&self, message_json: &str) -> Result<String> {
        let ast = self
            .script_ast
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("script not compiled"))?;

        let mut scope = rhai::Scope::new();
        let call_result = self.engine.call_fn::<rhai::Dynamic>(
            &mut scope,
            ast,
            "main",
            (message_json.to_string(),),
        );

        match call_result {
            Ok(v) => Ok(v.to_string()),
            Err(e) if e.to_string().contains("not found") => {
                // fn main not defined — evaluate script directly with msg in scope
                let mut scope2 = rhai::Scope::new();
                scope2.push("msg", message_json.to_string());
                let v = self
                    .engine
                    .eval_ast_with_scope::<rhai::Dynamic>(&mut scope2, ast)
                    .unwrap_or(rhai::Dynamic::from("(no output)"));
                Ok(v.to_string())
            }
            Err(e) => Err(anyhow::anyhow!("Rhai error: {e}")),
        }
    }
}

#[async_trait]
impl Actor for DynamicAgent {
    fn id(&self) -> String {
        self.config.id.clone()
    }
    fn name(&self) -> &str {
        &self.config.name
    }
    fn state(&self) -> ActorState {
        self.state.clone()
    }
    fn metrics(&self) -> Arc<ActorMetrics> {
        Arc::clone(&self.metrics)
    }
    fn mailbox(&self) -> mpsc::Sender<Message> {
        self.mailbox_tx.clone()
    }
    fn is_protected(&self) -> bool {
        self.config.protected
    }

    async fn on_start(&mut self) -> Result<()> {
        self.compile_script()?;
        self.state = ActorState::Running;
        // Announce to the dashboard so it appears immediately
        if let Some(pub_) = &self.publisher {
            pub_.publish(
                wactorz_mqtt::topics::spawn(&self.config.id),
                &serde_json::json!({
                    "agentId":   self.config.id,
                    "agentName": self.config.name,
                    "agentType": "dynamic",
                    "timestampMs": Self::now_ms(),
                }),
            );
        }
        Ok(())
    }

    async fn handle_message(&mut self, message: Message) -> Result<()> {
        // Extract text content — pass plain text to Rhai so simple scripts work.
        let content = match &message.payload {
            wactorz_core::message::MessageType::Text { content } => content.clone(),
            wactorz_core::message::MessageType::Task { description, .. } => description.clone(),
            _ => return Ok(()),
        };
        tracing::info!(
            "[{}] recv: {:?}",
            self.config.name,
            &content[..content.len().min(120)]
        );

        let script_result = tokio::task::block_in_place(|| self.run_script(&content));

        // Drain script log queue
        let logs: Vec<String> = self.log_queue.lock().unwrap().drain(..).collect();
        for log in &logs {
            tracing::info!("[{}] script log: {log}", self.config.name);
            if let Some(pub_) = &self.publisher {
                pub_.publish(
                    wactorz_mqtt::topics::logs(&self.config.id),
                    &serde_json::json!({"message": log}),
                );
            }
        }

        // Determine the final response:
        //   1. Rhai returned a useful string → use it
        //   2. Rhai returned "()" / "" OR errored → LLM fallback (if configured)
        //   3. No LLM → publish error/empty notice so the user isn't left hanging
        let response: String = match script_result {
            Ok(r) if !r.is_empty() && r != "()" => {
                tracing::info!(
                    "[{}] script → {:?}",
                    self.config.name,
                    &r[..r.len().min(120)]
                );
                r
            }
            Ok(r) => {
                tracing::info!(
                    "[{}] script returned {:?} — trying LLM fallback",
                    self.config.name,
                    r
                );
                self.llm_respond(&content).await
            }
            Err(e) => {
                tracing::warn!("[{}] script error: {e}", self.config.name);
                self.llm_respond_with_error(&content, &e.to_string()).await
            }
        };

        if !response.is_empty()
            && let Some(pub_) = &self.publisher
        {
            pub_.publish(
                wactorz_mqtt::topics::chat(&self.config.id),
                &serde_json::json!({
                    "from":        self.config.name,
                    "to":          "user",
                    "content":     response,
                    "timestampMs": Self::now_ms(),
                }),
            );
        }
        Ok(())
    }

    async fn on_heartbeat(&mut self) -> Result<()> {
        if let Some(pub_) = &self.publisher {
            pub_.publish(
                wactorz_mqtt::topics::heartbeat(&self.config.id),
                &serde_json::json!({
                    "agentId":   self.config.id,
                    "agentName": self.config.name,
                    "state":     self.state,
                    "timestampMs": Self::now_ms(),
                }),
            );
        }
        Ok(())
    }

    async fn on_stop(&mut self) -> Result<()> {
        if let Some(pub_) = &self.publisher {
            pub_.publish(
                wactorz_mqtt::topics::status(&self.config.id),
                &serde_json::json!({
                    "agentId":   self.config.id,
                    "agentName": self.config.name,
                    "state":     "stopped",
                    "timestampMs": Self::now_ms(),
                }),
            );
        }
        Ok(())
    }

    async fn run(&mut self) -> Result<()> {
        self.on_start().await?;
        let mut rx = self
            .mailbox_rx
            .take()
            .ok_or_else(|| anyhow::anyhow!("DynamicAgent already running"))?;
        let mut hb = tokio::time::interval(std::time::Duration::from_secs(
            self.config.heartbeat_interval_secs,
        ));
        hb.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        let mut paused = false;
        loop {
            tokio::select! {
                biased;
                msg = rx.recv() => {
                    match msg {
                        None => break,
                        Some(m) => {
                            self.metrics.record_received();
                            use wactorz_core::message::{ActorCommand, MessageType};
                            match &m.payload {
                                MessageType::Command { command: ActorCommand::Stop } => break,

                                MessageType::Command { command: ActorCommand::Pause } if !paused => {
                                    paused = true;
                                    self.state = ActorState::Paused;
                                    tracing::info!("[{}] paused", self.config.name);
                                    if let Some(pub_) = &self.publisher {
                                        pub_.publish(
                                            wactorz_mqtt::topics::status(&self.config.id),
                                            &serde_json::json!({
                                                "agentId":   self.config.id,
                                                "agentName": self.config.name,
                                                "state":     "paused",
                                                "timestampMs": Self::now_ms(),
                                            }),
                                        );
                                    }
                                }

                                MessageType::Command { command: ActorCommand::Resume } if paused => {
                                    paused = false;
                                    self.state = ActorState::Running;
                                    tracing::info!("[{}] resumed", self.config.name);
                                    if let Some(pub_) = &self.publisher {
                                        pub_.publish(
                                            wactorz_mqtt::topics::status(&self.config.id),
                                            &serde_json::json!({
                                                "agentId":   self.config.id,
                                                "agentName": self.config.name,
                                                "state":     "running",
                                                "timestampMs": Self::now_ms(),
                                            }),
                                        );
                                    }
                                }

                                _ if !paused => {
                                    match self.handle_message(m).await {
                                        Ok(_) => self.metrics.record_processed(),
                                        Err(e) => {
                                            tracing::error!("[{}] {e}", self.config.name);
                                            self.metrics.record_failed();
                                        }
                                    }
                                }

                                _ => {} // paused — drop non-command messages
                            }
                        }
                    }
                }
                _ = hb.tick() => {
                    if !paused {
                        self.metrics.record_heartbeat();
                        if let Err(e) = self.on_heartbeat().await {
                            tracing::error!("[{}] heartbeat: {e}", self.config.name);
                        }
                    }
                }
            }
        }
        self.state = ActorState::Stopped;
        self.on_stop().await
    }
}