horus 0.6.9

//! Agent handles and the single linear command dispatch loop.

use std::collections::BTreeMap;
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;

use serde_json::Value;
use tokio::sync::mpsc;
use uuid::Uuid;

use crate::Error;
use crate::Result;
use crate::backend::checkpoint::CHECKPOINT_VERSION;
use crate::backend::checkpoint::Checkpoint;
use crate::backend::checkpoint::CheckpointStore;
use crate::backend::checkpoint::ExecutionOutcome;
use crate::backend::checkpoint::ExecutionRecord;
use crate::backend::model::ModelChoice;
use crate::backend::model::ModelInfo;
use crate::backend::model::ModelRouter;
use crate::backend::sandbox::Sandbox;
use crate::middleware::FrontendExtensions;
use crate::middleware::MiddlewareCommandContext;
use crate::middleware::MiddlewareStack;
use crate::middleware::tools::Catalog;
use crate::protocol::Event;
use crate::protocol::EventMsg;
use crate::protocol::ModelChangedEvent;
use crate::protocol::Op;
use crate::protocol::SessionConfiguredEvent;
use crate::protocol::SessionContext;
use crate::protocol::SessionResumeRequestedEvent;
use crate::protocol::Submission;
use crate::protocol::WarningEvent;

mod input;
mod startup;
mod tool_step;
mod turn;

pub use self::startup::create_agent;

const COMMAND_QUEUE_CAPACITY: usize = 64;
const EVENT_QUEUE_CAPACITY: usize = 256;
const MAX_DEFERRED_SUBMISSIONS: usize = 64;
const MAX_IDENTIFIER_BYTES: usize = 4 * 1024;
const MAX_OPERATION_BYTES: usize = 256;
const MAX_ATTACHMENT_REFERENCES: usize = 16;
const DEFAULT_INITIAL_REPLAY_BATCHES: usize = 100;

/// Default maximum number of primary model steps in one turn.
pub const DEFAULT_MAX_MODEL_STEPS: usize = 256;

/// Dependencies and policy for one agent session.
#[derive(Clone)]
pub struct AgentConfig {
    model: Arc<ModelRouter>,
    provider: String,
    sandbox: Arc<Sandbox>,
    checkpoints: Arc<dyn CheckpointStore>,
    middleware: MiddlewareStack,
    system_prompt: String,
    session_id: String,
    context_window: i64,
    default_context_window: i64,
    session_context: SessionContext,
    metadata: BTreeMap<String, Value>,
    metadata_configured: bool,
    model_route_configured: bool,
    initial_replay_batches: usize,
    max_model_steps: usize,
}

impl AgentConfig {
    /// Creates a complete agent configuration.
    pub fn new(
        model: Arc<ModelRouter>,
        sandbox: Arc<Sandbox>,
        checkpoints: Arc<dyn CheckpointStore>,
        middleware: MiddlewareStack,
        system_prompt: impl Into<String>,
    ) -> Self {
        let provider = model.default_provider().to_string();
        let session_id = Uuid::new_v4().to_string();
        Self {
            model,
            provider,
            sandbox,
            checkpoints,
            middleware,
            system_prompt: system_prompt.into(),
            session_id,
            context_window: 272_000,
            default_context_window: 272_000,
            session_context: SessionContext::default(),
            metadata: BTreeMap::new(),
            metadata_configured: false,
            model_route_configured: false,
            initial_replay_batches: DEFAULT_INITIAL_REPLAY_BATCHES,
            max_model_steps: DEFAULT_MAX_MODEL_STEPS,
        }
    }

    /// Sets a stable ID used to resume a checkpointed session.
    #[must_use]
    pub fn session_id(mut self, session_id: impl Into<String>) -> Self {
        self.session_id = session_id.into();
        self
    }

    /// Attaches trusted, frontend-visible labels to this session.
    #[must_use]
    pub fn session_context(mut self, context: SessionContext) -> Self {
        self.session_context = context;
        self
    }

    /// Sets the provider's context window for usage display and policy.
    #[must_use]
    pub fn context_window(mut self, context_window: i64) -> Self {
        self.context_window = context_window;
        self.default_context_window = context_window;
        self
    }

    /// Sets the maximum durable transcript batches rendered when a session opens.
    #[must_use]
    pub fn initial_replay_batches(mut self, max_batches: usize) -> Self {
        self.initial_replay_batches = max_batches;
        self
    }

    /// Sets the maximum number of primary model steps in one turn.
    #[must_use]
    pub fn max_model_steps(mut self, max_steps: usize) -> Self {
        self.max_model_steps = max_steps;
        self
    }

    /// Sets durable framework-internal metadata used by installed capabilities.
    ///
    /// On resume, calling this replaces the saved metadata. Omitting it preserves
    /// the saved value.
    #[must_use]
    pub fn metadata(mut self, metadata: BTreeMap<String, Value>) -> Self {
        self.metadata = metadata;
        self.metadata_configured = true;
        self
    }

    /// Selects a registered model route and optional reasoning effort.
    pub fn model_route(mut self, route: &str, reasoning_effort: Option<&str>) -> Result<Self> {
        self.select_model_with_reasoning(route, reasoning_effort)?;
        self.model_route_configured = true;
        Ok(self)
    }

    /// Makes the configured router default replace a saved route on resume.
    #[must_use]
    pub fn override_saved_model_route(mut self) -> Self {
        self.model_route_configured = true;
        self
    }

    fn select_model(&mut self, route: &str) -> Result<ModelChoice> {
        self.select_model_with_reasoning(route, None)
    }

    fn select_model_with_reasoning(
        &mut self,
        route: &str,
        reasoning_effort: Option<&str>,
    ) -> Result<ModelChoice> {
        let choice = self.model.resolve_choice(route, reasoning_effort)?.clone();
        self.provider.clone_from(&choice.route);
        self.context_window = choice.context_window.unwrap_or(self.default_context_window);
        Ok(choice)
    }
}

/// Cloneable command side of a running agent.
#[derive(Clone)]
pub struct AgentSender {
    commands: mpsc::Sender<Submission>,
}

impl AgentSender {
    /// Sends a submission with a caller-controlled correlation ID.
    pub fn send(&self, submission: Submission) -> Result<()> {
        validate_submission(&submission)?;
        self.commands
            .try_send(submission)
            .map_err(|error| match error {
                mpsc::error::TrySendError::Full(_) => {
                    Error::Busy("agent command queue is full".into())
                }
                mpsc::error::TrySendError::Closed(_) => {
                    Error::Stopped("agent command channel closed".into())
                }
            })
    }

    /// Submits a command and returns its correlation ID.
    pub fn submit(&self, op: Op) -> Result<String> {
        let id = Uuid::new_v4().to_string();
        self.send(Submission { id: id.clone(), op })?;
        Ok(id)
    }
}

/// Validates one submission before callers perform more expensive boundary work.
pub fn validate_submission(submission: &Submission) -> Result<()> {
    validate_identifier("submission ID", &submission.id, MAX_IDENTIFIER_BYTES)?;
    match &submission.op {
        Op::UserInput { text, attachments } => validate_user_input(text, attachments),
        Op::ActiveInput {
            operation,
            turn_id,
            text,
        } => {
            validate_identifier("active operation", operation, MAX_OPERATION_BYTES)?;
            validate_identifier("turn ID", turn_id, MAX_IDENTIFIER_BYTES)?;
            validate_user_input(text, &[])
        }
        Op::Interrupt { turn_id } => validate_identifier("turn ID", turn_id, MAX_IDENTIFIER_BYTES),
        Op::ExecApproval { id, .. } => validate_identifier("approval ID", id, MAX_IDENTIFIER_BYTES),
        Op::CapabilityCommand {
            capability,
            command,
            arguments,
            input,
            target,
        } => {
            validate_identifier("capability ID", capability, MAX_OPERATION_BYTES)?;
            validate_identifier("command", command, MAX_OPERATION_BYTES)?;
            if arguments.len() > crate::protocol::MAX_CAPABILITY_INPUT_BYTES {
                return Err(Error::Config(
                    "middleware command arguments exceed size limit".into(),
                ));
            }
            if input
                .as_ref()
                .is_some_and(|input| input.len() > crate::protocol::MAX_CAPABILITY_INPUT_BYTES)
            {
                return Err(Error::Config(
                    "middleware command input exceeds size limit".into(),
                ));
            }
            if target.is_some_and(|target| target.batch_item_count == 0) {
                return Err(Error::Config(
                    "message target item count must be positive".into(),
                ));
            }
            Ok(())
        }
        Op::SetModel { route } => validate_identifier("model route", route, MAX_IDENTIFIER_BYTES),
        Op::ResumeSession { session_id } => {
            validate_identifier("session ID", session_id, MAX_IDENTIFIER_BYTES)
        }
    }
}

fn validate_user_input(
    text: &str,
    attachments: &[crate::protocol::SessionFileReference],
) -> Result<()> {
    if text.trim().is_empty() && attachments.is_empty() {
        return Err(Error::Config("user input cannot be empty".into()));
    }
    if text.len() > crate::protocol::MAX_USER_INPUT_BYTES {
        return Err(Error::Config("user input exceeds size limit".into()));
    }
    if attachments.len() > MAX_ATTACHMENT_REFERENCES {
        return Err(Error::Config(format!(
            "user input cannot reference more than {MAX_ATTACHMENT_REFERENCES} attachments"
        )));
    }
    let mut attachment_ids = std::collections::BTreeSet::new();
    for attachment in attachments {
        if !attachment_ids.insert(&attachment.id) {
            return Err(Error::Config(
                "attachment IDs must be unique per message".into(),
            ));
        }
        if Uuid::parse_str(&attachment.id).is_err() {
            return Err(Error::Config("attachment ID must be a UUID".into()));
        }
        validate_identifier("attachment name", &attachment.name, 255)?;
        validate_identifier("attachment media type", &attachment.media_type, 127)?;
        if attachment.size == 0 {
            return Err(Error::Config("attachment size must be positive".into()));
        }
    }
    Ok(())
}

fn validate_identifier(name: &str, value: &str, limit: usize) -> Result<()> {
    if value.trim().is_empty() {
        return Err(Error::Config(format!("{name} cannot be empty")));
    }
    if value.len() > limit {
        return Err(Error::Config(format!("{name} exceeds size limit")));
    }
    Ok(())
}

/// Bidirectional handle consumed by a frontend.
pub struct Agent {
    sender: AgentSender,
    events: mpsc::Receiver<Event>,
    frontend: FrontendExtensions,
    session: SessionConfiguredEvent,
    model: ModelInfo,
    model_choices: Vec<ModelChoice>,
    tool_count: usize,
    next_before_sequence: Option<u64>,
}

impl Agent {
    /// Returns a cloneable command sender.
    #[must_use]
    pub fn sender(&self) -> AgentSender {
        self.sender.clone()
    }

    /// Receives the next agent event.
    ///
    /// Frontends must keep draining events while the agent is running. Durable
    /// lifecycle events apply backpressure when this receiver is not polled.
    pub async fn next_event(&mut self) -> Option<Event> {
        self.events.recv().await
    }

    /// Returns the commands and status data exported by installed middleware.
    #[must_use]
    pub fn frontend(&self) -> &FrontendExtensions {
        &self.frontend
    }

    /// Returns the immutable session descriptor emitted at startup.
    #[must_use]
    pub fn session(&self) -> &SessionConfiguredEvent {
        &self.session
    }

    /// Returns frontend-safe settings for the selected model route.
    #[must_use]
    pub fn model(&self) -> &ModelInfo {
        &self.model
    }

    /// Returns the selected model route at frontend startup.
    #[must_use]
    pub fn model_route(&self) -> &str {
        &self.session.model.route
    }

    /// Returns every model route exposed to frontend selectors.
    #[must_use]
    pub fn model_choices(&self) -> &[ModelChoice] {
        &self.model_choices
    }

    /// Returns the number of tools registered for this agent.
    #[must_use]
    pub const fn tool_count(&self) -> usize {
        self.tool_count
    }

    /// Returns the cursor immediately preceding the initial transcript replay.
    #[must_use]
    pub const fn next_before_sequence(&self) -> Option<u64> {
        self.next_before_sequence
    }

    /// Separates command and event halves for a frontend event loop.
    ///
    /// The returned event receiver must be drained while commands are active.
    #[must_use]
    pub fn into_parts(self) -> (AgentSender, mpsc::Receiver<Event>) {
        (self.sender, self.events)
    }
}

struct Runner {
    config: AgentConfig,
    system_prompt: Arc<str>,
    catalog: Catalog,
    state: Checkpoint,
    review_session_id: String,
    transcript_delta: Vec<Value>,
    deferred: VecDeque<Submission>,
    events: mpsc::Sender<Event>,
}

impl Runner {
    async fn run(&mut self, mut commands: mpsc::Receiver<Submission>) -> Result<()> {
        if let Some(pending) = self.state.pending_approval.clone() {
            let submission_id = pending.submission_id.clone();
            if let Err(error) = self.resume_pending(&mut commands, pending).await {
                self.fail_turn(&submission_id, error).await?;
            }
        }
        loop {
            let submission = match self.deferred.pop_front() {
                Some(submission) => submission,
                None => {
                    let Some(submission) = commands.recv().await else {
                        return Ok(());
                    };
                    submission
                }
            };
            match submission.op {
                Op::UserInput { text, attachments } => {
                    if let Err(error) = self
                        .start_turn(&mut commands, submission.id.clone(), text, attachments)
                        .await
                    {
                        self.fail_turn(&submission.id, error).await?;
                    }
                }
                Op::ActiveInput { .. } => {
                    self.emit(
                        submission.id,
                        EventMsg::Warning(WarningEvent {
                            message: "there is no active turn".into(),
                        }),
                    )
                    .await?;
                }
                Op::Interrupt { .. } => {
                    self.emit(
                        submission.id,
                        EventMsg::Warning(WarningEvent {
                            message: "no active turn to interrupt".into(),
                        }),
                    )
                    .await?;
                }
                Op::ExecApproval { .. } => {
                    self.emit(
                        submission.id,
                        EventMsg::Warning(WarningEvent {
                            message: "no approval request is active".into(),
                        }),
                    )
                    .await?;
                }
                Op::CapabilityCommand {
                    capability,
                    command,
                    arguments,
                    input,
                    target,
                } => {
                    self.capability_command(
                        submission.id,
                        capability,
                        command,
                        arguments,
                        input,
                        target,
                    )
                    .await?;
                }
                Op::SetModel { route } => {
                    self.set_model(submission.id, route).await?;
                }
                Op::ResumeSession { session_id } => {
                    self.request_resume(submission.id, session_id).await?;
                }
            }
        }
    }

    async fn set_model(&mut self, submission_id: String, route: String) -> Result<()> {
        let choice = match self.config.select_model(&route) {
            Ok(choice) => choice,
            Err(error) => {
                self.emit(
                    submission_id,
                    EventMsg::Warning(WarningEvent {
                        message: error.to_string(),
                    }),
                )
                .await?;
                return Ok(());
            }
        };
        self.state.model_route = Some(choice.route.clone());
        self.save().await?;
        self.emit(
            submission_id,
            EventMsg::ModelChanged(ModelChangedEvent {
                route: choice.route,
                model: choice.model,
                reasoning_effort: choice.reasoning_effort,
                model_context_window: Some(self.config.context_window),
            }),
        )
        .await?;
        Ok(())
    }

    async fn request_resume(&self, submission_id: String, session_id: String) -> Result<()> {
        let result = async {
            if session_id.trim().is_empty() {
                return Err(Error::Config("session ID cannot be empty".into()));
            }
            let checkpoint = self
                .config
                .checkpoints
                .load(&session_id)
                .await?
                .ok_or_else(|| Error::Unknown(format!("session `{session_id}`")))?;
            if checkpoint.version != CHECKPOINT_VERSION || checkpoint.session_id != session_id {
                return Err(Error::Checkpoint(
                    "checkpoint does not match the requested session".into(),
                ));
            }
            Ok(checkpoint.session_context)
        }
        .await;
        match result {
            Ok(context) => self.emit(
                submission_id,
                EventMsg::SessionResumeRequested(SessionResumeRequestedEvent {
                    session_id,
                    context,
                }),
            ),
            Err(error) => self.emit(
                submission_id,
                EventMsg::Warning(WarningEvent {
                    message: error.to_string(),
                }),
            ),
        }
        .await
    }

    async fn capability_command(
        &mut self,
        submission_id: String,
        capability: String,
        command: String,
        arguments: String,
        input: Option<String>,
        target: Option<crate::protocol::MessageTarget>,
    ) -> Result<()> {
        let output = self
            .config
            .middleware
            .command(
                &capability,
                MiddlewareCommandContext {
                    command: &command,
                    arguments: &arguments,
                    input: input.as_deref(),
                    target,
                    session_id: &self.config.session_id,
                    session_context: &self.config.session_context,
                    checkpoint: &self.state,
                    checkpoints: Arc::clone(&self.config.checkpoints),
                },
            )
            .await
            .map(|output| output.events);
        match output {
            Ok(events) => {
                for event in events {
                    self.emit(&submission_id, EventMsg::Frontend(event)).await?;
                }
            }
            Err(error) => {
                self.emit(
                    submission_id,
                    EventMsg::Warning(WarningEvent {
                        message: error.to_string(),
                    }),
                )
                .await?
            }
        }
        Ok(())
    }

    async fn save(&mut self) -> Result<u64> {
        self.persist(None).await
    }

    async fn save_execution(&mut self, execution: &ExecutionRecord) -> Result<u64> {
        self.persist(Some(execution)).await
    }

    async fn persist(&mut self, execution: Option<&ExecutionRecord>) -> Result<u64> {
        let previous_sequence = self.state.sequence;
        self.state.sequence += 1;
        if let Err(error) = self
            .config
            .checkpoints
            .save(&self.state, &self.transcript_delta, execution)
            .await
        {
            self.state.sequence = previous_sequence;
            return Err(error);
        }
        self.transcript_delta.clear();
        Ok(self.state.sequence)
    }

    fn record_model_call(&mut self) -> Result<()> {
        let active =
            self.state.active_execution.as_mut().ok_or_else(|| {
                Error::Checkpoint("model called without an active execution".into())
            })?;
        active.model_calls = active
            .model_calls
            .checked_add(1)
            .ok_or_else(|| Error::Checkpoint("execution model-call count overflow".into()))?;
        Ok(())
    }

    fn record_usage(&mut self, usage: &crate::protocol::TokenUsage) -> Result<()> {
        let mut total_usage = self.state.total_usage.clone();
        total_usage.checked_add(usage).ok_or_else(|| {
            Error::Provider("provider token usage exceeds the supported range".into())
        })?;
        let active = self.state.active_execution.as_mut().ok_or_else(|| {
            Error::Checkpoint("usage recorded without an active execution".into())
        })?;
        let mut execution_usage = active.usage.clone();
        execution_usage.checked_add(usage).ok_or_else(|| {
            Error::Provider("provider token usage exceeds the supported range".into())
        })?;
        self.state.total_usage = total_usage;
        active.usage = execution_usage;
        Ok(())
    }

    fn record_tools(&mut self, tool_calls: u64, failed_tool_calls: u64) -> Result<()> {
        let active = self.state.active_execution.as_mut().ok_or_else(|| {
            Error::Checkpoint("tools recorded without an active execution".into())
        })?;
        active.tool_calls = active
            .tool_calls
            .checked_add(tool_calls)
            .ok_or_else(|| Error::Checkpoint("execution tool-call count overflow".into()))?;
        active.failed_tool_calls = active
            .failed_tool_calls
            .checked_add(failed_tool_calls)
            .ok_or_else(|| Error::Checkpoint("execution failed-tool count overflow".into()))?;
        Ok(())
    }

    fn finish_execution(&mut self, outcome: ExecutionOutcome) -> Result<ExecutionRecord> {
        self.state.finish_execution(outcome, unix_timestamp_ms()?)
    }

    async fn finish_and_save_execution(&mut self, outcome: ExecutionOutcome) -> Result<u64> {
        let active_execution = self.state.active_execution.clone();
        let execution_stats = self.state.execution_stats.clone();
        let execution = self.finish_execution(outcome)?;
        match self.save_execution(&execution).await {
            Ok(sequence) => Ok(sequence),
            Err(error) => {
                self.state.active_execution = active_execution;
                self.state.execution_stats = execution_stats;
                Err(error)
            }
        }
    }

    fn push_context(&mut self, item: Value) {
        self.state.context.push(item.clone());
        self.transcript_delta.push(item);
    }

    fn extend_context(&mut self, items: Vec<Value>) {
        self.state.context.extend(items.iter().cloned());
        self.transcript_delta.extend(items);
    }

    async fn emit(&self, submission_id: impl Into<String>, msg: EventMsg) -> Result<()> {
        send_event(
            &self.events,
            Event {
                submission_id: Some(submission_id.into()),
                msg,
            },
        )
        .await
    }
}

fn unix_timestamp_ms() -> Result<i64> {
    let elapsed = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_err(|_| Error::Checkpoint("system clock predates the Unix epoch".into()))?;
    i64::try_from(elapsed.as_millis())
        .map_err(|_| Error::Checkpoint("system clock exceeds the supported range".into()))
}

async fn send_event(events: &mpsc::Sender<Event>, event: Event) -> Result<()> {
    events
        .send(event)
        .await
        .map_err(|_| Error::Stopped("frontend event channel closed".into()))
}

fn try_send_event(events: &mpsc::Sender<Event>, event: Event) -> Result<()> {
    events.try_send(event).map_err(|error| match error {
        mpsc::error::TrySendError::Full(_) => Error::Stopped("frontend event queue is full".into()),
        mpsc::error::TrySendError::Closed(_) => {
            Error::Stopped("frontend event channel closed".into())
        }
    })
}

#[cfg(test)]
#[path = "runtime_tests.rs"]
mod tests;