plugins/agent/native/src/owner/execution.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.
 */

use super::*;
use maka_agent_facts::event::Fact;
use maka_agent_facts::event::InvocationOutcome;
use maka_event_log::turns::{InvocationState, TurnBoundary};

pub(super) fn project((boundary, high_water): (TurnBoundary, u64)) -> Execution {
    Execution {
        session_id: boundary.invocation.session_id,
        turn_id: boundary.invocation.turn_id,
        run_id: boundary.invocation.run_id,
        invocation_id: boundary.invocation.invocation_id,
        high_water,
        state: match boundary.state {
            InvocationState::Admitted => ExecutionState::Admitted,
            InvocationState::Running => ExecutionState::Running,
            InvocationState::WaitingForUser => ExecutionState::Waiting,
            InvocationState::Ended { outcome, .. } => match outcome {
                InvocationOutcome::HandoffPaused { .. } => ExecutionState::Paused,
                InvocationOutcome::ContextCompactFinished {
                    outcome: maka_agent_facts::context::CompactOutcome::Failed { .. },
                } => ExecutionState::Failed,
                InvocationOutcome::Completed | InvocationOutcome::ContextCompactFinished { .. } => {
                    ExecutionState::Completed
                }
                InvocationOutcome::Failed { .. } => ExecutionState::Failed,
                InvocationOutcome::Cancelled { .. } => ExecutionState::Cancelled,
            },
        },
    }
}
fn parked(reason: ResumeParkReason) -> ResumeResult {
    ResumeResult::Parked { reason }
}

impl Owner {
    pub async fn work(&self, input: WorkRead) -> Result<Work, String> {
        use maka_event_log::message_resolution::MessageExecution as Stored;
        self.session(&input.session_id).await?;
        let offset = input.cursor.as_ref().map_or(0, |cursor| cursor.offset);
        let origin = match input.target {
            WorkTarget::Message { message_id } => {
                maka_event_log::observation::WorkOrigin::Message(message_id)
            }
            WorkTarget::Invocation { invocation_id } => {
                maka_event_log::observation::WorkOrigin::Invocation(invocation_id)
            }
        };
        let observed = self
            .store
            .work_observation(
                &input.session_id,
                origin,
                offset,
                input
                    .cursor
                    .as_ref()
                    .map(|cursor| cursor.invocation_id.as_str()),
            )
            .await
            .map_err(|_| "agent_work_unavailable")?;
        let (shared, boundary) = match observed.execution {
            Stored::Missing => return Ok(Work::Missing),
            Stored::Pending => return Ok(Work::Pending),
            Stored::Cancelled => return Ok(Work::Cancelled),
            Stored::Owned(boundary) => (false, boundary),
            Stored::Shared(boundary) => (true, boundary),
        };
        if input
            .cursor
            .as_ref()
            .is_some_and(|cursor| cursor.invocation_id != boundary.invocation.invocation_id)
        {
            return Ok(Work::Changed);
        }
        let answer = observed.answer.map(|answer| {
            let length = answer.text.len();
            AnswerPage {
                text: answer.text,
                next: (!answer.complete).then(|| AnswerCursor {
                    invocation_id: boundary.invocation.invocation_id.clone(),
                    offset: offset + length as u64,
                }),
            }
        });
        Ok(Work::Delivered {
            shared,
            execution: project((*boundary, observed.high_water)),
            answer,
            pending_interactions: observed.interactions.len() as u32,
        })
    }

    pub async fn execution(&self, input: ExecutionRead) -> Result<ExecutionResult, String> {
        self.session(&input.session_id).await?;
        let boundary = self
            .store
            .invocation_boundary(&input.session_id, &input.invocation_id)
            .await
            .map_err(|_| "agent_execution_unavailable")?;
        Ok(match boundary {
            Some(boundary) => ExecutionResult::Found {
                execution: project(boundary),
            },
            None => ExecutionResult::Missing,
        })
    }

    pub async fn resume(
        self: &Arc<Self>,
        input: Resume,
        operation: &maka_plugin_core::operation::RecoveryRef,
    ) -> Result<ResumeResult, String> {
        use ResumeParkReason as Reason;
        let mut runs = self.runs.lock().await;
        let session = self.session(&input.session_id).await?;
        let identity = self.operation_id(operation);
        // Engine receipts and application operations use different digest encodings.
        let fingerprint = format!(
            "sha256:{}",
            hex::encode(Sha256::digest(operation.fingerprint.as_bytes()))
        );
        let invocation = Invocation {
            session_id: input.session_id.clone(),
            turn_id: format!("t_{identity}"),
            run_id: format!("r_{identity}"),
            invocation_id: format!("i_{identity}"),
        };
        // Acceptance wins over today's model availability, workspace and archive
        // state. The canonical opening is the receipt; there is no resume ledger.
        if let Some((accepted, source, recorded_fingerprint)) = self
            .store
            .continuation_receipt(&invocation.invocation_id)
            .await
            .map_err(|_| "agent_execution_unavailable")?
        {
            if accepted != invocation
                || recorded_fingerprint != fingerprint
                || source.invocation.session_id != input.session_id
                || source.invocation.invocation_id != input.source_invocation_id
                || source.high_water != input.source_high_water
            {
                return Err("agent_resume_conflict".into());
            }
            let boundary = self
                .store
                .invocation_boundary(&input.session_id, &invocation.invocation_id)
                .await
                .map_err(|_| "agent_execution_unavailable")?
                .ok_or("agent_execution_unavailable")?;
            return Ok(ResumeResult::Started {
                execution: project(boundary),
            });
        }
        if runs.closing || self.stop.is_cancelled() {
            return Ok(parked(Reason::Closed));
        }
        if session.archived {
            return Ok(parked(Reason::Archived));
        }
        if runs.controls.contains_key(&input.session_id)
            || !self
                .store
                .message_queue(&input.session_id)
                .await
                .map_err(|_| "agent_queue_unavailable")?
                .entries
                .is_empty()
        {
            return Ok(parked(Reason::Busy));
        }
        let Some((boundary, high_water)) = self
            .store
            .invocation_boundary(&input.session_id, &input.source_invocation_id)
            .await
            .map_err(|_| "agent_execution_unavailable")?
        else {
            return Ok(parked(Reason::SourceMissing));
        };
        if high_water != input.source_high_water {
            return Ok(parked(Reason::SourceChanged));
        }
        if !matches!(
            boundary.state,
            InvocationState::Ended {
                outcome: InvocationOutcome::Failed { .. } | InvocationOutcome::Cancelled { .. },
                ..
            }
        ) {
            return Ok(parked(Reason::SourceNotResumable));
        }
        let prefix = match self
            .store
            .run_source(&input.session_id, &boundary.invocation.run_id, None)
            .await
        {
            Ok(Some(prefix)) => prefix,
            Ok(None) => return Ok(parked(Reason::SourceMissing)),
            Err(_) => return Err("agent_execution_unavailable".into()),
        };
        if prefix.boundary.invocation != boundary.invocation
            || prefix.boundary.high_water != input.source_high_water
            || !matches!(&prefix.last.event.fact, Fact::InvocationEnded { .. })
        {
            return Ok(parked(Reason::SourceChanged));
        }
        let source_boundary = prefix.boundary;
        if self
            .store
            .continuation_for_source(&source_boundary)
            .await
            .map_err(|_| "agent_execution_unavailable")?
            .is_some()
        {
            return Ok(parked(Reason::AlreadyContinued));
        }
        let Some(mut configuration) = self
            .store
            .invocation_configuration(&boundary.invocation)
            .await
            .map_err(|_| "agent_execution_unavailable")?
        else {
            return Ok(parked(Reason::SourceNotResumable));
        };
        if configuration.model.is_none()
            || configuration.orchestration_mode.as_str() == plan::EXECUTION
        {
            return Ok(parked(Reason::SourceNotResumable));
        }
        let workspace = &session.configuration.workspace;
        if configuration.cwd != workspace.path
            || configuration.workspace_identity.as_ref() != Some(&workspace.identity)
            || workspace::validate(&self.workspace_roots, workspace, &self.managed_worktrees)
                .await
                .is_err()
        {
            return Ok(parked(Reason::WorkspaceChanged));
        }
        let (mut source, prepared) = crate::model::Source::select(
            self.configuration.clone(),
            self.catalog.clone(),
            self.credentials.clone(),
            input.model,
            input.thinking_level,
            input.session_id.clone(),
            self.revision.clone(),
        )
        .await
        .map_err(public_error)?;
        // A continuation retains the original tool protocol even when the
        // user selects a model whose default differs from the original Run.
        source.accepted.tool_mode = configuration.tool_mode;
        source.accepted.code_language = configuration.code_language;
        source.accepted.editing_tools = configuration.editing_tools;
        let graph = match self
            .resume_graph(&input.session_id, &configuration.orchestration_mode)
            .await
        {
            Ok(graph) => graph,
            Err(_) => return Ok(parked(Reason::SourceNotResumable)),
        };
        let basis = MessageBasis {
            language: configuration.language,
            compaction_prompt: configuration.compaction_prompt.clone(),
            graph,
            collaboration_mode: configuration.collaboration_mode,
            background_kind: BackgroundKind::Schedule,
            tool_discovery: configuration.tool_discovery,
            allowed_tools: session.configuration.tools.clone(),
            background: None,
            internal_tools: self.internal_tools.clone(),
            internal_bindings: self.internal_bindings.clone(),
            // Ordinary Agents retain their existing queue compatibility. A
            // coordinator continuation cannot silently adopt a newer behavior.
            behavior: self
                .behavior
                .as_ref()
                .and(configuration.system_prompt.clone()),
            model: Some(source.accepted.clone()),
            executor: None,
            workspace: workspace.clone(),
            fingerprint: fingerprint.clone(),
            tools: self.tool_basis.clone(),
        };
        configuration.model = Some(source.accepted.binding.clone());
        configuration.thinking_level = source.accepted.thinking_level;
        // Replay the original prepared prompt. Native service tools continue to
        // resolve the live Session permission boundary through their actual owner.
        let tools = self.selected_tools(
            &basis,
            configuration.tool_discovery,
            configuration.editing_tools,
        );
        let run = RunInput {
            invocation: invocation.clone(),
            model_source: Some(Arc::new(source)),
            model_revision: Some(prepared.revision),
            provider_id: prepared.provider_id,
            context: prepared.context,
            provider: prepared.provider,
            provider_options: prepared.options,
            main_output_limit: prepared.main_output_limit,
            supports_vision: prepared.supports_vision,
            request_fingerprint: Some(fingerprint),
            configuration,
            work: RunWork::Continuation {
                source: source_boundary,
                tools,
            },
        };
        let prepared = match self.engine.prepare_continuation(run, &self.stop).await {
            Ok(prepared) => prepared,
            Err(
                RunError::InvalidInput(_)
                | RunError::ReconciliationRequired(_)
                | RunError::Store(
                    maka_event_log::StoreError::InvalidTransition(_)
                    | maka_event_log::StoreError::HistoryUnsettled
                    | maka_event_log::StoreError::HistoryProviderEffectUnresolved,
                ),
            ) => {
                return Ok(parked(Reason::SafetyCheckFailed));
            }
            Err(error) => {
                if error.is_cleanup_unconfirmed() {
                    runs.quarantined.insert(input.session_id.clone());
                    runs.closing = true;
                    self.stop.cancel();
                }
                return Err(public_error(error));
            }
        };
        let running = match self
            .engine
            .start_continuation(prepared, self.stop.child_token())
            .await
        {
            Ok(running) => running,
            Err(RunError::Busy) => return Ok(parked(Reason::Busy)),
            Err(
                RunError::ContinuationChanged
                | RunError::InvalidInput(_)
                | RunError::ReconciliationRequired(_),
            ) => {
                return Ok(parked(Reason::SafetyCheckFailed));
            }
            Err(error) => {
                if error.is_cleanup_unconfirmed() {
                    runs.quarantined.insert(input.session_id.clone());
                }
                runs.closing = true;
                self.stop.cancel();
                return Err(public_error(error));
            }
        };
        runs.controls.insert(
            input.session_id,
            Control {
                compacting: false,
                invocation: invocation.clone(),
                basis,
                cancellation: running.cancellation(),
                failure: None,
            },
        );
        self.workers.spawn(self.clone().drive(running));
        let boundary = self
            .store
            .invocation_boundary(&invocation.session_id, &invocation.invocation_id)
            .await
            .map_err(|_| "agent_execution_unavailable")?
            .ok_or("agent_execution_unavailable")?;
        Ok(ResumeResult::Started {
            execution: project(boundary),
        })
    }
}