plugins/agent/engine/src/lib.rs

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

mod auto_context;
mod cleanup;
mod compact;
mod context_input;
mod continuation;
mod executor;
pub use executor::ExecutorInput;
mod handoff;
pub use handoff::{HandoffGate, HandoffReservation, HeldHandoff, PendingSeal};
mod history;
mod interactions;
mod model_attempt;
mod model_source;
pub use model_source::{ModelSource, PreparedModel};
pub mod pricing;
mod request_composition;
pub use history::project as project_model_history;
mod prune;
mod runner;
mod running;
mod steps;
pub use running::{RunningCompletion, RunningInvocation};

#[derive(Clone, Debug, thiserror::Error)]
#[error("invocation cleanup is unconfirmed")]
pub struct CleanupUnconfirmed;

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

use maka_agent_facts::event::CommitError;
use maka_agent_facts::event::Invocation;
use maka_agent_facts::execution::InvocationConfiguration;
use maka_agent_facts::input::MessageInput;
use maka_agent_facts::tools::ToolError;
use maka_event_log::{EventLog, StoreError};
use maka_model::{ModelError, ModelExecutor, ProviderConfig};
use maka_tools::CodeExecutor;
use maka_tools::ToolCatalog;
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;

#[derive(Debug, thiserror::Error)]
pub enum RunError {
    #[error("session already has an active invocation")]
    Busy,
    /// Preflight became stale; no invocation was opened or source claimed.
    #[error("Session tool loading changed during continuation preparation; prepare again")]
    ContinuationChanged,
    #[error("session requires reconciliation before another invocation: {0}")]
    ReconciliationRequired(String),
    #[error("invalid execution input: {0}")]
    InvalidInput(String),
    #[error("execution cancelled")]
    Cancelled,
    #[error("model output was incomplete; local tools from this response were not executed")]
    ModelIncomplete,
    #[error(transparent)]
    Commit(#[from] CommitError),
    #[error(transparent)]
    Store(#[from] StoreError),
    #[error(transparent)]
    Model(#[from] ModelError),
    #[error(transparent)]
    Tool(ToolError),
    /// Physical ownership is uncertain; an enclosing native owner must retain
    /// its Root and capacity rather than infer cleanup from business completion.
    #[error("execution cleanup is unconfirmed")]
    CleanupUnconfirmed,
    #[error("execution worker failed: {0}")]
    Internal(String),
}

impl RunError {
    /// Explicit unconfirmed cleanup in this error. Absence is not proof of
    /// cleanup; only the original RunningCompletion carries that observation.
    pub fn is_cleanup_unconfirmed(&self) -> bool {
        matches!(
            self,
            Self::CleanupUnconfirmed
                | Self::Tool(
                    ToolError::CleanupUnconfirmed(_)
                        | ToolError::OperationCleanupUnconfirmed { .. }
                )
        )
    }
}

impl From<ToolError> for RunError {
    fn from(error: ToolError) -> Self {
        if matches!(
            error,
            ToolError::CleanupUnconfirmed(_) | ToolError::OperationCleanupUnconfirmed { .. }
        ) {
            cleanup::unconfirmed();
        }
        Self::Tool(error)
    }
}

#[derive(Clone)]
pub struct RunInput<W = RunWork> {
    /// Optional for standalone adapters; Host executions supply a public provider source.
    pub model_source: Option<Arc<dyn ModelSource>>,
    pub model_revision: Option<maka_agent_facts::composition::SourceRevision>,
    /// Public access-path identity; not the transport protocol or adapter name.
    pub provider_id: String,
    pub invocation: Invocation,
    pub context: Option<maka_agent_facts::context::ModelRequestContext>,
    pub work: W,
    /// Stable domain admission identity, when this invocation was client-issued.
    pub request_fingerprint: Option<String>,
    pub provider: ProviderConfig,
    pub provider_options: Value,
    /// Resolved SDK reply budget; summaries additionally cap their text at 8K.
    pub main_output_limit: Option<u64>,
    /// Current selected-model capability; image bytes are projected per request.
    pub supports_vision: bool,
    pub configuration: InvocationConfiguration,
}

// Admission owns the submitted body. Running execution owns only settings and
// identity; canonical opening facts supply the body to bounded history reads.
type RunContext = RunInput<()>;

impl RunInput {
    fn into_context(self) -> (RunContext, RunWork) {
        let Self {
            work,
            model_source,
            model_revision,
            provider_id,
            invocation,
            context,
            request_fingerprint,
            provider,
            provider_options,
            main_output_limit,
            supports_vision,
            configuration,
        } = self;
        (
            RunInput {
                work: (),
                model_source,
                model_revision,
                provider_id,
                invocation,
                context,
                request_fingerprint,
                provider,
                provider_options,
                main_output_limit,
                supports_vision,
                configuration,
            },
            work,
        )
    }
}

#[derive(Clone)]
pub enum RunWork {
    Handoff {
        source: maka_agent_facts::continuation::RunBoundary,
        pause: Box<maka_agent_facts::handoff::HandoffPause>,
        tools: ToolCatalog,
    },
    Continuation {
        source: maka_agent_facts::continuation::RunBoundary,
        tools: ToolCatalog,
    },
    Message {
        message: MessageInput,
        source_messages: Vec<maka_agent_facts::message::RootSourceMessage>,
        tools: ToolCatalog,
        /// Only a newly submitted client Turn may observe a sealed, unknown prior tool effect.
        allow_prior_unknown: bool,
    },
    ContextCompact,
}

struct Inner {
    pricing: Option<Arc<dyn pricing::Pricing>>,
    log: Arc<EventLog>,
    model: ModelExecutor,
    cells: CodeExecutor,
    sessions: Mutex<std::collections::HashMap<String, SessionState>>,
    active: Mutex<HashSet<String>>,
    workers: TaskTracker,
}

#[derive(Clone, Default)]
struct SessionState {
    code_store: maka_code_runtime::CellStore,
    loaded: maka_tools::availability::LoadedTools,
}

impl Inner {
    fn session_state(&self, session: &str) -> SessionState {
        self.sessions
            .lock()
            .unwrap()
            .entry(session.into())
            .or_default()
            .clone()
    }

    fn loaded_snapshot(&self, session: &str) -> std::collections::BTreeSet<String> {
        self.sessions
            .lock()
            .unwrap()
            .get(session)
            .map(|state| state.loaded.snapshot())
            .unwrap_or_default()
    }

    fn clear_loaded(&self, session: &str) {
        if let Some(state) = self.sessions.lock().unwrap().get(session) {
            state.loaded.clear();
        }
    }
}

/// Owns session admission and the entire invocation lifetime, including effects
/// still draining after the caller stops waiting.
#[derive(Clone)]
pub struct Engine(Arc<Inner>);

/// Engine-derived replay proof and its immutable input. Preparation uses registered
/// tools; model steps capture live bindings only after durable admission revalidates
/// the discovery basis. A changed basis requires fresh preparation.
pub struct PreparedContinuation {
    engine: std::sync::Weak<Inner>,
    input: RunInput,
    claim: continuation::PreparedClaim,
}
impl PreparedContinuation {
    pub fn invocation(&self) -> &Invocation {
        &self.input.invocation
    }
}

impl Engine {
    pub fn new(log: Arc<EventLog>, model: ModelExecutor, cells: CodeExecutor) -> Self {
        Self::create(log, model, cells, None)
    }

    pub fn with_pricing(
        log: Arc<EventLog>,
        model: ModelExecutor,
        cells: CodeExecutor,
        pricing: Arc<dyn pricing::Pricing>,
    ) -> Self {
        Self::create(log, model, cells, Some(pricing))
    }

    fn create(
        log: Arc<EventLog>,
        model: ModelExecutor,
        cells: CodeExecutor,
        pricing: Option<Arc<dyn pricing::Pricing>>,
    ) -> Self {
        Self(Arc::new(Inner {
            pricing,
            log,
            model,
            cells,
            sessions: Mutex::default(),
            active: Mutex::new(HashSet::new()),
            workers: TaskTracker::new(),
        }))
    }

    pub async fn run(
        &self,
        input: RunInput,
        cancellation: CancellationToken,
    ) -> Result<Invocation, RunError> {
        self.start(input, cancellation).await?.wait().await
    }

    /// Host calls after session retirement has fenced and drained its runs.
    pub fn release_session_state(&self, session: &str) {
        self.0.sessions.lock().unwrap().remove(session);
    }

    /// Observe replay safety without reserving a source or creating a claim.
    /// Admission validates the claim against current canonical facts.
    pub async fn check_continuation(
        &self,
        input: &RunInput,
        cancellation: &CancellationToken,
    ) -> Result<(), RunError> {
        let (RunWork::Continuation { source, tools, .. } | RunWork::Handoff { source, tools, .. }) =
            &input.work
        else {
            return Err(RunError::InvalidInput("not a continuation request".into()));
        };
        self.0
            .log
            .prepare_prune_candidates(&input.invocation.session_id, None, 0, 0, None)
            .await?;
        continuation::inspect(
            &self.0,
            input,
            source,
            tools,
            cancellation,
            &self.0.loaded_snapshot(&input.invocation.session_id),
        )
        .await
        .map(|_| ())
    }

    pub async fn prepare_continuation(
        &self,
        input: RunInput,
        cancellation: &CancellationToken,
    ) -> Result<PreparedContinuation, RunError> {
        let (RunWork::Continuation { source, tools, .. } | RunWork::Handoff { source, tools, .. }) =
            &input.work
        else {
            return Err(RunError::InvalidInput("not a continuation request".into()));
        };
        self.0
            .log
            .prepare_prune_candidates(&input.invocation.session_id, None, 0, 0, None)
            .await?;
        let claim = continuation::prepare(&self.0, &input, source, tools, cancellation).await?;
        Ok(PreparedContinuation {
            engine: Arc::downgrade(&self.0),
            input,
            claim,
        })
    }

    pub async fn start_continuation(
        &self,
        prepared: PreparedContinuation,
        cancellation: CancellationToken,
    ) -> Result<RunningInvocation, RunError> {
        if !std::sync::Weak::ptr_eq(&prepared.engine, &Arc::downgrade(&self.0)) {
            return Err(RunError::InvalidInput(
                "continuation belongs to another Engine".into(),
            ));
        }
        self.start_model(prepared.input, Some(prepared.claim), cancellation)
            .await
    }

    /// The owner first stops admissions and requests cancellation. This barrier
    /// also covers workers whose admission waiter was dropped before receiving
    /// a RunningInvocation handle.
    /// Caller must fence admission and request cancellation first. Tracking task
    /// termination alone does not confirm successful resource cleanup; retain
    /// and observe each original RunningCompletion separately.
    pub async fn drain(&self) {
        self.0.workers.close();
        self.0.workers.wait().await;
    }

    /// Returns only after durable admission. The returned owner covers execution
    /// and draining; dropping it requests cancellation without abandoning effects.
    pub async fn start(
        &self,
        input: RunInput,
        cancellation: CancellationToken,
    ) -> Result<RunningInvocation, RunError> {
        self.start_model(input, None, cancellation).await
    }

    async fn start_model(
        &self,
        input: RunInput,
        claim: Option<continuation::PreparedClaim>,
        cancellation: CancellationToken,
    ) -> Result<RunningInvocation, RunError> {
        if input.main_output_limit.is_some_and(|limit| limit == 0) {
            return Err(RunError::InvalidInput("invalid Main output limit".into()));
        }
        if let Some(context) = &input.context {
            context
                .validate()
                .map_err(|reason| RunError::InvalidInput(reason.into()))?;
        }
        match &input.work {
            RunWork::Handoff { source, pause, .. } => {
                pause
                    .validate(&source.invocation)
                    .map_err(|reason| RunError::InvalidInput(reason.into()))?;
                if input.invocation != pause.intent.successor(&source.invocation)
                    || input.request_fingerprint.is_some()
                    || input.context != pause.execution.context
                    || input.provider_options != pause.execution.provider_options
                    || input.main_output_limit != pause.execution.main_output_limit
                    || input.supports_vision != pause.execution.supports_vision
                    || model_attempt::route_identity(&input)?
                        != pause.execution.replay.route_identity
                {
                    return Err(RunError::InvalidInput(
                        "handoff execution settings changed".into(),
                    ));
                }
            }
            RunWork::Continuation { source, .. } => {
                if source.invocation.session_id != input.invocation.session_id
                    || input
                        .request_fingerprint
                        .as_ref()
                        .is_none_or(|s| !maka_agent_facts::archive::valid_projection_digest(s))
                {
                    return Err(RunError::InvalidInput(
                        "invalid continuation request".into(),
                    ));
                }
            }
            RunWork::ContextCompact
                if input
                    .request_fingerprint
                    .as_ref()
                    .is_none_or(|value| value.is_empty()) =>
            {
                return Err(RunError::InvalidInput(
                    "compaction requires request fingerprint".into(),
                ));
            }
            RunWork::Message {
                message,
                source_messages,
                ..
            } => {
                maka_agent_facts::message::validate_sources(message, source_messages)
                    .map_err(|reason| RunError::InvalidInput(reason.into()))?;
            }
            _ => {}
        }
        if !std::path::Path::new(&input.configuration.cwd).is_absolute()
            || input.configuration.cwd.contains('\0')
            || input
                .configuration
                .model
                .as_ref()
                .is_some_and(|binding| binding.model != input.provider.model)
        {
            return Err(RunError::InvalidInput(
                "invalid invocation configuration".into(),
            ));
        }
        let handoff = (!matches!(input.work, RunWork::ContextCompact)).then(|| {
            let root = match &input.work {
                RunWork::Handoff { pause, .. } => &pause.intent.root_run_id,
                _ => &input.invocation.run_id,
            };
            HandoffGate::new(input.invocation.clone(), root.clone())
        });
        let inner = self.0.clone();
        let invocation = input.invocation.clone();
        let tool_names = Arc::new(match &input.work {
            RunWork::Message { tools, .. }
            | RunWork::Continuation { tools, .. }
            | RunWork::Handoff { tools, .. } => tools.names().into_iter().collect(),
            RunWork::ContextCompact => Default::default(),
        });
        let worker_handoff = handoff.clone();
        self.start_owned(
            invocation,
            tool_names,
            handoff,
            cancellation,
            move |cancellation, admitted| {
                runner::run(inner, input, cancellation, admitted, worker_handoff, claim)
            },
        )
        .await
    }

    async fn start_owned<F, Fut>(
        &self,
        invocation: Invocation,
        tool_names: Arc<HashSet<String>>,
        handoff: Option<HandoffGate>,
        cancellation: CancellationToken,
        run: F,
    ) -> Result<RunningInvocation, RunError>
    where
        F: FnOnce(CancellationToken, tokio::sync::oneshot::Sender<()>) -> Fut + Send + 'static,
        Fut: std::future::Future<Output = Result<Invocation, RunError>> + Send + 'static,
    {
        let session_id = invocation.session_id.clone();
        if !self
            .0
            .active
            .lock()
            .map_err(|_| RunError::Internal("admission poisoned".into()))?
            .insert(session_id.clone())
        {
            return Err(RunError::Busy);
        }
        let admission = Admission {
            inner: self.0.clone(),
            session_id,
            handoff: handoff.clone(),
        };
        let cancellation = cancellation.child_token();
        let cancel_on_drop = cancellation.clone().drop_guard();
        let worker_cancellation = cancellation.clone();
        let (admitted, ready) = tokio::sync::oneshot::channel();
        let cleanup = cleanup::Evidence::default();
        let worker = self.0.workers.spawn(async move {
            let _admission = admission;
            let result = cleanup.scope(run(worker_cancellation, admitted)).await;
            let physical = if cleanup.confirmed()
                && !result.as_ref().is_err_and(RunError::is_cleanup_unconfirmed)
            {
                Ok(())
            } else {
                Err(CleanupUnconfirmed)
            };
            RunningCompletion {
                result,
                cleanup: physical,
            }
        });
        if ready.await.is_err() {
            return match worker.await {
                Ok(completion) if completion.cleanup.is_err() => Err(RunError::CleanupUnconfirmed),
                Ok(completion) => match completion.result {
                    Err(error) => Err(error),
                    Ok(_) => Err(RunError::Internal(
                        "worker omitted durable admission".into(),
                    )),
                },
                Err(_) => Err(RunError::CleanupUnconfirmed),
            };
        }
        cancel_on_drop.disarm();
        Ok(RunningInvocation::new(
            invocation,
            tool_names,
            cancellation,
            handoff,
            worker,
        ))
    }
}

struct Admission {
    inner: Arc<Inner>,
    session_id: String,
    handoff: Option<HandoffGate>,
}

impl Drop for Admission {
    fn drop(&mut self) {
        if let Some(handoff) = &self.handoff {
            handoff.close();
        }
        self.inner.active.lock().unwrap().remove(&self.session_id);
    }
}