mentra 0.13.0

mod compact;
mod config;
mod events;
mod lifecycle;
mod pending;
mod pending_block;
mod round_strategy;
mod runner;
mod snapshot;
mod steering;
mod subagent;
mod task_state;
mod team;
mod terminal_output;
#[cfg(test)]
mod tests;
mod wait;

use std::{
    collections::HashSet,
    sync::{
        Arc, Mutex,
        atomic::{AtomicU64, Ordering},
    },
};

use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, watch};

use crate::{
    ContentBlock, Message,
    background::BackgroundNotification,
    error::RuntimeError,
    memory::journal::{AgentMemory, AgentMemoryState as MemoryState},
    provider::{Provider, ProviderId, ToolChoice},
    runtime::{
        LoadedAgentState, RuntimeIntrinsicTool, TaskItem,
        handle::{AgentExecutionConfig, AgentObserver, RuntimeHandle},
    },
    team::TeamMessage,
    transcript::{DelegationArtifact, DelegationEdge, TranscriptItem},
};

pub(crate) use team::parse_task_input;

pub use config::{
    AgentConfig, CompactionConfig, ContextCompactionConfig, MemoryConfig, TaskConfig,
    TeamAutonomyConfig, TeamConfig, ToolProfile, ToolResultPagingConfig, WorkspaceConfig,
};
pub use events::{
    AgentEvent, AgentSnapshot, AgentStatus, CompactionDetails, CompactionTrigger,
    ContextCompactionDetails, ContextCompactionTrigger, PendingToolUseSummary, SpawnedAgentStatus,
    SpawnedAgentSummary,
};
pub use pending::PendingAssistantTurn;
pub use round_strategy::{
    ReasoningChange, RoundAdjustment, RoundBoundary, RoundContext, RoundDecision, RoundStrategy,
    RoundToolResult,
};
use runner::TurnRunner;
pub use steering::{QueueMode, SteeringHandle};
pub(crate) use subagent::DisposableSubagentTemplate;
pub use terminal_output::{FinalOutput, TerminalOutputSpec};
pub use wait::{AgentWaitFuture, AgentWaitHandle};

static NEXT_AGENT_ID: AtomicU64 = AtomicU64::new(1);

/// Running or persisted agent managed by a [`crate::Runtime`].
pub struct Agent {
    id: String,
    runtime: RuntimeHandle,
    model: String,
    provider_id: ProviderId,
    name: String,
    config: AgentConfig,
    memory: AgentMemory,
    tasks: Vec<TaskItem>,
    rounds_since_task: usize,
    event_bus: AgentEventBus,
    snapshot: Arc<Mutex<AgentSnapshot>>,
    snapshot_tx: watch::Sender<AgentSnapshot>,
    provider: Arc<dyn Provider>,
    hidden_tools: HashSet<String>,
    terminal_tool_gate: Arc<Mutex<Option<String>>>,
    max_rounds: Option<usize>,
    inflight_background_notifications: Vec<BackgroundNotification>,
    inflight_team_messages: Vec<TeamMessage>,
    steering: SteeringHandle,
    inflight_steer: Vec<Vec<ContentBlock>>,
    inflight_follow_up: Vec<Vec<ContentBlock>>,
    teammate_identity: Option<TeammateIdentity>,
    idle_requested: bool,
    current_run_id: Option<String>,
    /// Full texts of results this agent received paged, keyed by
    /// `tool_use_id` — the backing store for `read_tool_result`. Empty and
    /// unused unless `config.tool_result_paging` is set.
    paged_tool_results: crate::tool::paging::PagedToolResults,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct TeammateIdentity {
    pub(crate) role: String,
    pub(crate) lead: String,
}

#[derive(Default)]
pub(crate) struct AgentSpawnOptions {
    pub(crate) hidden_tools: HashSet<String>,
    pub(crate) max_rounds: Option<usize>,
    pub(crate) teammate_identity: Option<TeammateIdentity>,
}

type AgentEventTap = Arc<dyn Fn(&AgentEvent) + Send + Sync>;

#[derive(Default)]
struct AgentEventTapRegistry {
    next_id: u64,
    taps: Vec<(u64, AgentEventTap)>,
}

pub(crate) struct AgentEventTapGuard {
    registry: Arc<Mutex<AgentEventTapRegistry>>,
    id: u64,
}

#[derive(Clone)]
pub(crate) struct AgentEventBus {
    tx: broadcast::Sender<AgentEvent>,
    taps: Arc<Mutex<AgentEventTapRegistry>>,
}

impl AgentEventBus {
    fn new(capacity: usize) -> Self {
        let (tx, _) = broadcast::channel(capacity);
        Self {
            tx,
            taps: Arc::new(Mutex::new(AgentEventTapRegistry::default())),
        }
    }

    pub(crate) fn send(&self, event: AgentEvent) {
        let taps = {
            let registry = self.taps.lock().expect("agent event tap registry poisoned");
            registry
                .taps
                .iter()
                .map(|(_, tap)| Arc::clone(tap))
                .collect::<Vec<_>>()
        };
        for tap in taps {
            tap(&event);
        }
        let _ = self.tx.send(event);
    }

    pub(crate) fn subscribe(&self) -> broadcast::Receiver<AgentEvent> {
        self.tx.subscribe()
    }

    pub(crate) fn register_tap(
        &self,
        tap: impl Fn(&AgentEvent) + Send + Sync + 'static,
    ) -> AgentEventTapGuard {
        let mut registry = self.taps.lock().expect("agent event tap registry poisoned");
        let id = registry.next_id;
        registry.next_id += 1;
        registry.taps.push((id, Arc::new(tap)));
        AgentEventTapGuard {
            registry: Arc::clone(&self.taps),
            id,
        }
    }
}

impl Drop for AgentEventTapGuard {
    fn drop(&mut self) {
        let mut registry = self
            .registry
            .lock()
            .expect("agent event tap registry poisoned");
        registry.taps.retain(|(tap_id, _)| *tap_id != self.id);
    }
}

impl Agent {
    pub(crate) fn new(
        runtime: RuntimeHandle,
        model: String,
        name: String,
        config: AgentConfig,
        provider: Arc<dyn Provider>,
        options: AgentSpawnOptions,
    ) -> Result<Self, RuntimeError> {
        let AgentSpawnOptions {
            hidden_tools,
            max_rounds,
            teammate_identity,
        } = options;
        let store = runtime.store();
        let agent_id = format!(
            "agent-{:x}-{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos(),
            NEXT_AGENT_ID.fetch_add(1, Ordering::Relaxed)
        );
        let memory = AgentMemory::new(agent_id.clone(), store.clone(), MemoryState::default());
        let event_bus = AgentEventBus::new(256);
        let memory_view = memory.snapshot_view();
        let snapshot = AgentSnapshot {
            history_len: memory_view.history_len,
            current_text: memory_view.current_text,
            pending_tool_uses: memory_view.pending_tool_uses,
            ..Default::default()
        };
        let snapshot = Arc::new(Mutex::new(snapshot));
        let (snapshot_tx, _) =
            watch::channel(snapshot.lock().expect("agent snapshot poisoned").clone());
        let mut agent = Self {
            id: agent_id,
            runtime,
            model,
            provider_id: provider.descriptor().id,
            name,
            config,
            memory,
            tasks: Vec::new(),
            rounds_since_task: 0,
            event_bus,
            snapshot,
            snapshot_tx,
            provider,
            hidden_tools,
            terminal_tool_gate: Arc::new(Mutex::new(None)),
            max_rounds,
            inflight_background_notifications: Vec::new(),
            inflight_team_messages: Vec::new(),
            steering: SteeringHandle::new(),
            inflight_steer: Vec::new(),
            inflight_follow_up: Vec::new(),
            teammate_identity,
            idle_requested: false,
            current_run_id: None,
            paged_tool_results: Default::default(),
        };
        agent
            .runtime
            .store()
            .create_agent(&agent.persisted_record(), agent.memory.state())?;
        let execution_config = AgentExecutionConfig {
            name: agent.name.clone(),
            team_dir: agent.config.team.team_dir.clone(),
            tasks_dir: agent.config.task.tasks_dir.clone(),
            base_dir: agent.config.workspace.base_dir.clone(),
            memory_tool_search_limit: agent.config.memory.tool_search_limit,
            auto_route_shell: agent.config.workspace.auto_route_shell,
            is_teammate: agent.teammate_identity.is_some(),
        };
        let observer = AgentObserver {
            events: agent.event_bus.clone(),
            snapshot_tx: agent.snapshot_tx.clone(),
            snapshot: Arc::clone(&agent.snapshot),
        };
        agent
            .runtime
            .register_agent(&agent.id, &agent.name, execution_config, &observer)?;
        agent.register_tool_result_pager();
        agent.refresh_tasks_from_disk()?;
        Ok(agent)
    }

    pub(crate) fn from_loaded(
        runtime: RuntimeHandle,
        mut state: LoadedAgentState,
        provider: Arc<dyn Provider>,
    ) -> Result<Self, RuntimeError> {
        let mut memory = AgentMemory::new(state.record.id.clone(), runtime.store(), state.memory);
        let recovery = memory.recover()?;
        if recovery.interrupted {
            state.record.status = AgentStatus::Interrupted;
            runtime.store().update_run_state(
                recovery
                    .interrupted_run_id
                    .as_deref()
                    .expect("recovery should include run id"),
                "interrupted",
                Some("recovered after interruption"),
            )?;
            runtime.store().save_agent_record(&state.record)?;
        }
        let memory_view = memory.snapshot_view();
        let snapshot = AgentSnapshot {
            status: state.record.status.clone(),
            history_len: memory_view.history_len,
            current_text: memory_view.current_text,
            pending_tool_uses: memory_view.pending_tool_uses,
            pending_team_messages: 0,
            subagents: state.record.subagents.clone(),
            ..Default::default()
        };
        let snapshot = Arc::new(Mutex::new(snapshot));
        let (snapshot_tx, _) =
            watch::channel(snapshot.lock().expect("agent snapshot poisoned").clone());
        let event_bus = AgentEventBus::new(256);
        let mut agent = Self {
            id: state.record.id.clone(),
            runtime,
            model: state.record.model.clone(),
            provider_id: state.record.provider_id.clone(),
            name: state.record.name.clone(),
            config: state.record.config.clone(),
            memory,
            tasks: Vec::new(),
            rounds_since_task: state.record.rounds_since_task,
            event_bus,
            snapshot,
            snapshot_tx,
            provider,
            hidden_tools: state.record.hidden_tools,
            terminal_tool_gate: Arc::new(Mutex::new(None)),
            max_rounds: state.record.max_rounds,
            inflight_background_notifications: Vec::new(),
            inflight_team_messages: Vec::new(),
            steering: SteeringHandle::new(),
            inflight_steer: Vec::new(),
            inflight_follow_up: Vec::new(),
            teammate_identity: state.record.teammate_identity,
            idle_requested: state.record.idle_requested,
            current_run_id: None,
            paged_tool_results: Default::default(),
        };
        let execution_config = AgentExecutionConfig {
            name: agent.name.clone(),
            team_dir: agent.config.team.team_dir.clone(),
            tasks_dir: agent.config.task.tasks_dir.clone(),
            base_dir: agent.config.workspace.base_dir.clone(),
            memory_tool_search_limit: agent.config.memory.tool_search_limit,
            auto_route_shell: agent.config.workspace.auto_route_shell,
            is_teammate: agent.teammate_identity.is_some(),
        };
        let observer = AgentObserver {
            events: agent.event_bus.clone(),
            snapshot_tx: agent.snapshot_tx.clone(),
            snapshot: Arc::clone(&agent.snapshot),
        };
        agent
            .runtime
            .register_agent(&agent.id, &agent.name, execution_config, &observer)?;
        agent.register_tool_result_pager();
        agent.refresh_tasks_from_disk()?;
        Ok(agent)
    }

    /// Returns the agent's display name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns the stable persisted agent identifier.
    pub fn id(&self) -> &str {
        &self.id
    }

    /// Returns the model identifier used by the agent.
    pub fn model(&self) -> &str {
        &self.model
    }

    /// Updates the model and provider used for future turns, then persists the
    /// new agent record so resumed sessions continue with the same setting.
    pub fn set_model(&mut self, model: crate::ModelInfo) -> Result<(), RuntimeError> {
        let provider = self
            .runtime
            .get_provider(Some(&model.provider))
            .ok_or_else(|| RuntimeError::ProviderNotFound(Some(model.provider.clone())))?;
        self.model = model.id;
        self.provider_id = provider.descriptor().id;
        self.provider = provider;
        self.persist_agent_record()
    }

    /// Updates the reasoning options requested on future turns, then persists the
    /// agent record so resumed sessions continue with the same setting.
    ///
    /// Mirrors [`set_model`](Self::set_model): a stateful override threaded into
    /// every subsequent model request (the runner reads
    /// `config.provider_request_options.reasoning` live). It composes with
    /// `set_model` for **per-phase tiering** — e.g. run the gather rounds at a low
    /// reasoning effort, then raise the effort (and switch to a stronger model) for
    /// a final synthesis turn on the same agent, without re-spawning and losing the
    /// gathered context. `None` clears any configured reasoning, restoring the
    /// provider's default effort.
    pub fn set_reasoning(
        &mut self,
        reasoning: Option<crate::provider::ReasoningOptions>,
    ) -> Result<(), RuntimeError> {
        self.config.provider_request_options.reasoning = reasoning;
        self.persist_agent_record()
    }

    /// Returns the effective agent configuration.
    pub fn config(&self) -> &AgentConfig {
        &self.config
    }

    /// Returns the committed transcript history.
    pub fn history(&self) -> &[Message] {
        self.memory.history()
    }

    /// Returns the canonical transcript items stored for this agent.
    pub fn transcript(&self) -> &crate::AgentTranscript {
        self.memory.transcript()
    }

    /// The transcript entry the next turn will continue from.
    pub fn leaf(&self) -> Option<&crate::transcript::EntryId> {
        self.transcript().leaf()
    }

    /// Returns to an earlier entry, so the next turn explores a new path from
    /// there.
    ///
    /// The abandoned entries stay in the transcript, reachable through
    /// [`children`](Self::children) — nothing is deleted, so the path just
    /// left can be returned to the same way. Returns how many entries left
    /// the active path.
    pub fn branch_from(
        &mut self,
        entry: &crate::transcript::EntryId,
    ) -> Result<usize, RuntimeError> {
        self.memory.branch_from(entry)
    }

    /// The entries recorded as continuing from `entry`. More than one means
    /// the conversation branched there.
    pub fn children(
        &self,
        entry: &crate::transcript::EntryId,
    ) -> Vec<&crate::transcript::TranscriptItem> {
        self.transcript().children(entry)
    }

    fn append_transcript_item(&mut self, item: TranscriptItem) -> Result<(), RuntimeError> {
        self.memory.append_transcript_item(item)
    }

    pub(crate) fn record_canonical_context(
        &mut self,
        content: impl Into<String>,
    ) -> Result<(), RuntimeError> {
        self.append_transcript_item(TranscriptItem::canonical_context(Message::user(
            ContentBlock::text(content.into()),
        )))
    }

    pub(crate) fn record_delegation_request(
        &mut self,
        content: impl Into<String>,
        delegation: DelegationArtifact,
        edge: Option<DelegationEdge>,
    ) -> Result<(), RuntimeError> {
        self.append_transcript_item(TranscriptItem::delegation_request(
            Message::user(ContentBlock::text(content.into())),
            delegation,
            edge,
        ))
    }

    pub(crate) fn record_delegation_result(
        &mut self,
        content: impl Into<String>,
        delegation: DelegationArtifact,
        edge: Option<DelegationEdge>,
    ) -> Result<(), RuntimeError> {
        self.append_transcript_item(TranscriptItem::delegation_result(
            Message::user(ContentBlock::text(content.into())),
            delegation,
            edge,
        ))
    }

    pub(crate) fn memory_revision(&self) -> u64 {
        self.memory.revision()
    }

    pub(crate) fn memory_engine(&self) -> Arc<crate::memory::MemoryEngine> {
        self.runtime.memory_engine()
    }

    /// Returns whether this agent is a persistent teammate rather than the lead agent.
    pub fn is_teammate(&self) -> bool {
        self.teammate_identity.is_some()
    }

    pub(crate) fn tasks(&self) -> &[TaskItem] {
        &self.tasks
    }

    /// Returns the most recent committed message, if any.
    pub fn last_message(&self) -> Option<&Message> {
        self.memory.last_message()
    }

    /// Subscribes to the agent's transient event stream.
    pub fn subscribe_events(&self) -> broadcast::Receiver<AgentEvent> {
        self.event_bus.subscribe()
    }

    /// Watches the current agent snapshot for state updates.
    pub fn watch_snapshot(&self) -> watch::Receiver<AgentSnapshot> {
        self.snapshot_tx.subscribe()
    }

    pub(crate) fn tools(&self) -> Arc<[crate::tool::ProviderToolSpec]> {
        let terminal_tool = self
            .terminal_tool_gate
            .lock()
            .expect("terminal tool gate poisoned")
            .clone();
        self.runtime
            .tools()
            .iter()
            .filter(|tool| {
                if let Some(name) = terminal_tool.as_ref() {
                    name == &tool.name
                        && self.runtime.tool_is_visible_to_agent(&tool.name, &self.id)
                } else {
                    self.can_use_tool(&tool.name)
                }
            })
            .cloned()
            .collect::<Vec<_>>()
            .into()
    }

    pub(crate) fn can_use_tool(&self, name: &str) -> bool {
        if !self.runtime.tool_is_visible_to_agent(name, &self.id) {
            return false;
        }

        if self
            .terminal_tool_gate
            .lock()
            .expect("terminal tool gate poisoned")
            .as_deref()
            == Some(name)
        {
            return true;
        }

        if self.hidden_tools.contains(name) {
            return false;
        }

        if !self.config.tool_profile.allows(name) {
            return false;
        }

        if name == RuntimeIntrinsicTool::Idle.to_string() {
            return self.teammate_identity.is_some();
        }

        // The pager's reader exists for the model only while there can be
        // paged results to read. Registration is runtime-wide (the registry
        // is keyed by tool name), so this per-agent gate — not registration —
        // is what keeps the tool out of an unpaged agent's roster, even when
        // a paging agent shares the same runtime.
        if name == crate::tool::paging::READ_TOOL_RESULT_TOOL {
            return self.config.tool_result_paging.is_some();
        }

        true
    }

    pub(crate) fn runtime_handle(&self) -> RuntimeHandle {
        self.runtime.clone()
    }

    /// Registers the pager's reader when this agent enables paging. The tool
    /// itself is stateless — it resolves both the retained results and the
    /// page size from the calling agent's context — so one registration
    /// serves every paging agent on the runtime, and re-registering is a
    /// no-op.
    fn register_tool_result_pager(&self) {
        if self.config.tool_result_paging.is_some() {
            self.runtime.register_tool(crate::tool::ReadToolResultTool);
        }
    }

    /// Retains the full text of a result that entered the transcript paged,
    /// so `read_tool_result` can serve its later windows. In memory only, for
    /// this agent's lifetime — a result the model never asks to continue
    /// simply goes away with the agent.
    pub(crate) fn record_paged_tool_result(&self, tool_use_id: &str, full: &str) {
        self.paged_tool_results.record(tool_use_id, full);
    }

    /// Returns a retained full result by `tool_use_id`. Only this agent's own
    /// paged results are reachable: the store is per-agent, so one agent can
    /// never read another's.
    pub(crate) fn paged_tool_result(&self, tool_use_id: &str) -> Option<Arc<str>> {
        self.paged_tool_results.get(tool_use_id)
    }

    pub(crate) fn max_rounds(&self) -> Option<usize> {
        self.max_rounds
    }

    pub(crate) fn tool_choice(&self) -> Option<ToolChoice> {
        if let Some(name) = self
            .terminal_tool_gate
            .lock()
            .expect("terminal tool gate poisoned")
            .clone()
        {
            return Some(ToolChoice::Tool { name });
        }

        match self.config.tool_choice.clone() {
            Some(ToolChoice::Tool { name }) if !self.can_use_tool(&name) => Some(ToolChoice::Auto),
            other => other,
        }
    }
}