molo-agent 0.4.0

//! Application-level event set for [`ReActAgent`](crate::agent::ReActAgent).
//!
//! These are ReAct's own event types, not framework-level events: other
//! Agent implementations (planning / retrieval / sub-agent orchestration)
//! define their own event types and simply implement
//! [`AgentEvent`](crate::agent::AgentEvent) — the event set is up to each
//! Agent; the framework does not prescribe a unified set.
//! ReAct's event set is **closed**, expressed as a single enum
//! [`ReActEvent`] — consumers do one downcast followed by an exhaustive
//! match, rather than downcasting per type. Once an
//! [`EventChannel`](crate::event_channel::EventChannel) is attached,
//! ReActAgent publishes these events on both the streaming and non-streaming
//! run paths.
//!
//! # Examples
//!
//! Attach a broadcast channel and subscribe to the sequence of event names
//! for one run:
//!
//! ```
//! # extern crate molo_agent as molo;
//! # #[tokio::main]
//! # async fn main() -> Result<(), molo::AgentError> {
//! use molo::agent::AgentEvent;
//! use molo::event_channel::{BroadcastEventChannel, EventChannel, EventReceiver};
//! use molo::provider::{FakeProvider, FakeReply};
//! use molo::{react_agent, Agent};
//!
//! let channel = BroadcastEventChannel::new(64);
//! let mut rx = channel.subscribe();
//! let mut agent = react_agent!(FakeProvider::new([FakeReply::Text("Hello".into())]))
//!     .with_event_channel(channel);
//!
//! agent.run("hi").await?;
//! drop(agent); // channel closed: once the subscriber drains the remaining events, recv returns None
//! while let Some(event) = rx.recv().await {
//!     println!("{}", event.name());
//! }
//! # Ok(())
//! # }
//! ```

use super::{AgentEvent, RunSummary};
use crate::observability::{AgentEventRecord, EventSeverity, RedactionRecord};
use crate::tool::{RegistryError, ToolResult};
use crate::{AgentError, UserInput};

/// A single event from ReActAgent (application level; the event name is
/// provided per variant via [`AgentEvent::name`]).
#[derive(Debug, Clone, PartialEq)]
pub enum ReActEvent {
    /// The run has started (the user input is already recorded). Marks the
    /// start of a run segment in a long-lived channel.
    RunStarted {
        /// Observability identifier for this run: every trace span carries a
        /// `run.id` attribute with the same value, which correlates the event
        /// stream with observability data; subsequent events in the same
        /// segment implicitly belong to this run.
        run_id: String,
        /// The user input for this run.
        input: UserInput,
    },
    /// An increment of the model's reply text (same content as
    /// [`MessageChunk::Delta`](super::MessageChunk::Delta) — pushed to
    /// pipeline-only subscribers so they need not pull the stream).
    Delta {
        /// Text increment.
        text: String,
    },
    /// An increment of the model's reasoning. Note that
    /// [`MessageChunk`](crate::agent::MessageChunk) has no Reasoning variant
    /// (see its docs); pipeline events are each Agent's own choice, and ReAct
    /// chooses to push reasoning.
    Reasoning {
        /// Reasoning increment.
        text: String,
    },
    /// A tool has started executing (the loop received the model's tool call
    /// request and is about to execute it).
    ToolStarted {
        /// The original id of this call, paired with
        /// [`ToolCompleted`](ReActEvent::ToolCompleted).
        id: String,
        /// The tool name.
        name: String,
        /// The arguments JSON generated by the model.
        arguments: String,
    },
    /// A tool has finished dispatching; `result` is the return value of
    /// [`ToolRegistry::call`](crate::tool::ToolRegistry::call) — Ok(Output)
    /// is immediate output, Ok(Effect) is a side-effect request, Err is a
    /// registry/tool failure (Display is the error text fed back to the
    /// model).
    ToolCompleted {
        /// The original id of this call, paired with
        /// [`ToolStarted`](ReActEvent::ToolStarted).
        id: String,
        /// The tool name.
        name: String,
        /// The dispatch result (Ok/Err status is classified by the registry,
        /// see [`RegistryError`](crate::tool::RegistryError)).
        result: Result<ToolResult, RegistryError>,
    },
    /// The run has ended (published on success, cancellation, and error) —
    /// the end of a run segment in a long-lived channel, carrying the
    /// execution summary for this run (also accumulated on the non-streaming
    /// path).
    RunEnded {
        /// Execution summary for this run (same semantics as the streaming
        /// `Done` chunk).
        summary: RunSummary,
        /// None = normal end; Some = cancelled / errored (cancellation is
        /// [`AgentError::Cancelled`](crate::AgentError::Cancelled)).
        error: Option<AgentError>,
    },
}

impl AgentEvent for ReActEvent {
    fn name(&self) -> &'static str {
        match self {
            ReActEvent::RunStarted { .. } => "run.started",
            ReActEvent::Delta { .. } => "delta",
            ReActEvent::Reasoning { .. } => "reasoning",
            ReActEvent::ToolStarted { .. } => "tool.started",
            ReActEvent::ToolCompleted { .. } => "tool.completed",
            ReActEvent::RunEnded { .. } => "run.ended",
        }
    }

    fn to_record(&self) -> Option<AgentEventRecord> {
        let record = match self {
            ReActEvent::RunStarted { run_id, input } => AgentEventRecord::new(
                self.name(),
                EventSeverity::Info,
                serde_json::json!({
                    "input.kind": user_input_kind(input),
                    "input.bytes": user_input_bytes(input),
                }),
            )
            .with_run_id(run_id.clone())
            .with_redactions(vec![omitted("input")]),
            ReActEvent::Delta { text } => AgentEventRecord::new(
                self.name(),
                EventSeverity::Debug,
                serde_json::json!({
                    "delta.bytes": text.len(),
                }),
            )
            .with_redactions(vec![omitted("delta.text")]),
            ReActEvent::Reasoning { text } => AgentEventRecord::new(
                self.name(),
                EventSeverity::Debug,
                serde_json::json!({
                    "reasoning.bytes": text.len(),
                }),
            )
            .with_redactions(vec![omitted("reasoning.text")]),
            ReActEvent::ToolStarted {
                id,
                name,
                arguments,
            } => AgentEventRecord::new(
                self.name(),
                EventSeverity::Info,
                serde_json::json!({
                    "tool.id": id,
                    "tool.name": name,
                    "tool.arguments_bytes": arguments.len(),
                }),
            )
            .with_redactions(vec![omitted("tool.arguments")]),
            ReActEvent::ToolCompleted { id, name, result } => AgentEventRecord::new(
                self.name(),
                if result.is_ok() {
                    EventSeverity::Info
                } else {
                    EventSeverity::Warn
                },
                tool_completed_payload(id, name, result),
            ),
            ReActEvent::RunEnded { summary, error } => {
                let severity = if error.is_some() {
                    EventSeverity::Error
                } else {
                    EventSeverity::Info
                };
                let mut payload = run_summary_payload(summary);
                if let serde_json::Value::Object(object) = &mut payload {
                    object.insert(
                        "status".to_string(),
                        serde_json::json!(if error.is_some() { "error" } else { "ok" }),
                    );
                    if let Some(error) = error {
                        object.insert(
                            "error.kind".to_string(),
                            serde_json::json!(agent_error_kind(error)),
                        );
                    }
                }
                AgentEventRecord::new(self.name(), severity, payload)
            }
        };
        Some(record)
    }
}

fn omitted(field: &str) -> RedactionRecord {
    RedactionRecord {
        field: field.to_string(),
        reason: "raw content omitted".to_string(),
    }
}

fn user_input_kind(input: &UserInput) -> &'static str {
    match input {
        UserInput::Text(_) => "text",
        UserInput::Blocks(_) => "blocks",
        _ => "unknown",
    }
}

fn user_input_bytes(input: &UserInput) -> usize {
    match input {
        UserInput::Text(text) => text.len(),
        UserInput::Blocks(blocks) => blocks
            .iter()
            .map(|block| serde_json::to_vec(block).map_or(0, |bytes| bytes.len()))
            .sum(),
        _ => 0,
    }
}

fn tool_completed_payload(
    id: &str,
    name: &str,
    result: &Result<ToolResult, RegistryError>,
) -> serde_json::Value {
    match result {
        Ok(ToolResult::Output(output)) => serde_json::json!({
            "tool.id": id,
            "tool.name": name,
            "status": "ok",
            "effect.requested": false,
            "output.bytes": output.content.len(),
            "artifacts": output.artifacts.len(),
        }),
        Ok(ToolResult::Effect(effect)) => serde_json::json!({
            "tool.id": id,
            "tool.name": name,
            "status": "ok",
            "effect.requested": true,
            "effect.id": effect.id,
            "effect.kind": format!("{:?}", effect.kind),
            "effect.risk": format!("{:?}", effect.risk),
        }),
        Ok(_) => serde_json::json!({
            "tool.id": id,
            "tool.name": name,
            "status": "ok",
            "effect.requested": false,
        }),
        Err(error) => serde_json::json!({
            "tool.id": id,
            "tool.name": name,
            "status": "error",
            "error.kind": registry_error_kind(error),
        }),
    }
}

fn run_summary_payload(summary: &RunSummary) -> serde_json::Value {
    serde_json::json!({
        "rounds": summary.rounds,
        "tool_calls": summary.tool_calls,
        "usage.prompt_tokens": summary.usage.prompt_tokens,
        "usage.completion_tokens": summary.usage.completion_tokens,
        "usage.total_tokens": summary.usage.total_tokens,
        "usage.omitted": summary.usage_omitted,
        "finish.reason": summary.finish_reason.as_ref().map(|reason| format!("{reason:?}")),
        "latency_ms": summary.latency.as_millis() as u64,
        "provider.model": summary.provider_model,
    })
}

fn registry_error_kind(error: &RegistryError) -> &'static str {
    match error {
        RegistryError::NotFound(_) => "not_found",
        RegistryError::InvalidArguments(_) => "invalid_arguments",
        RegistryError::Execution { .. } => "execution",
        RegistryError::NameCollision { .. } => "name_collision",
        RegistryError::SourceNameMismatch { .. } => "source_name_mismatch",
    }
}

fn agent_error_kind(error: &AgentError) -> &'static str {
    match error {
        AgentError::Memory(_) => "memory",
        AgentError::Provider(_) => "provider",
        AgentError::TooManyToolRounds(_) => "too_many_tool_rounds",
        AgentError::StructuredRetriesExhausted(_) => "structured_retries_exhausted",
        AgentError::StructuredParse(_) => "structured_parse",
        AgentError::Cancelled => "cancelled",
        AgentError::DeadlineExceeded => "deadline_exceeded",
        AgentError::EffectRequiresHarness(_) => "effect_requires_harness",
        AgentError::InvalidStep(_) => "invalid_step",
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn event_record_omits_raw_delta() {
        let event = ReActEvent::Delta {
            text: "secret-token".to_string(),
        };

        let record = event.to_record().expect("ReAct events expose records");
        let json = serde_json::to_string(&record).unwrap();

        assert!(json.contains("delta.bytes"));
        assert!(!json.contains("secret-token"));
        assert_eq!(record.redactions[0].field, "delta.text");
    }

    #[test]
    fn tool_started_record_omits_raw_arguments() {
        let event = ReActEvent::ToolStarted {
            id: "call-1".to_string(),
            name: "write_file".to_string(),
            arguments: r#"{"token":"secret-token"}"#.to_string(),
        };

        let record = event.to_record().expect("ReAct events expose records");
        let json = serde_json::to_string(&record).unwrap();

        assert!(json.contains("tool.arguments_bytes"));
        assert!(!json.contains("secret-token"));
        assert_eq!(record.redactions[0].field, "tool.arguments");
    }
}