everruns-core 0.43.0

//! Store-backed turn-context resolution for runtime hosts.

use std::sync::Arc;

use crate::ResolvedExecutionSnapshot;
use crate::capabilities::{
    CapabilityRegistry, collect_message_filters_only, collect_message_filters_only_with_context,
    provider_managed_reduction_budget,
};
use crate::execution_loading::{AgentStore, HarnessStore, SessionStore};
use crate::message::{RuntimeMessage, RuntimeMessageRole};
use crate::message_filter::MessageQuery;
use crate::message_retriever::MessageRetriever;
use crate::provider_resolution::ProviderStore;
use crate::runtime_context::{
    AssembledTurnContext, ResolvedModelExecution, ResolvedTurnContextInput, TurnContextRequest,
    TurnContextResolver, assemble_resolved_turn_context, resolve_snapshot_capabilities,
};
use crate::session_files::SessionFileSystem;
use everruns_contracts::driver_registry::{ChatDriver, DriverRegistry};
use everruns_contracts::error::{AgentLoopError, Result};
use everruns_contracts::hosted_mcp::{HostedMcpDriver, HostedMcpResolver};
use everruns_contracts::model_spec::ModelSpec;
use everruns_contracts::provider::DriverId;
use everruns_contracts::tool_types::ToolDefinition;
use everruns_contracts::typed_id::{AgentId, HarnessId, ModelId, SessionId};

use crate::host::execution_snapshot::load_execution_snapshot;

/// Store-backed resolver used by direct atom callers and custom hosts.
#[derive(Clone)]
pub struct StoreTurnContextResolver {
    harness_store: Arc<dyn HarnessStore>,
    agent_store: Arc<dyn AgentStore>,
    session_store: Arc<dyn SessionStore>,
    message_retriever: Arc<dyn MessageRetriever>,
    provider_store: Arc<dyn ProviderStore>,
    capability_registry: CapabilityRegistry,
    driver_registry: DriverRegistry,
    file_store: Option<Arc<dyn SessionFileSystem>>,
    session_storage: Option<Arc<dyn crate::session_services::SessionStorageStore>>,
    hosted_mcp: Option<Arc<dyn HostedMcpResolver>>,
}

impl StoreTurnContextResolver {
    /// Construct a resolver from org-scoped runtime stores and registries.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        harness_store: Arc<dyn HarnessStore>,
        agent_store: Arc<dyn AgentStore>,
        session_store: Arc<dyn SessionStore>,
        message_retriever: Arc<dyn MessageRetriever>,
        provider_store: Arc<dyn ProviderStore>,
        capability_registry: CapabilityRegistry,
        driver_registry: DriverRegistry,
    ) -> Self {
        Self {
            harness_store,
            agent_store,
            session_store,
            message_retriever,
            provider_store,
            capability_registry,
            driver_registry,
            file_store: None,
            session_storage: None,
            hosted_mcp: None,
        }
    }

    /// Supply the session filesystem used by dynamic prompt capabilities.
    pub fn with_file_store(mut self, file_store: Arc<dyn SessionFileSystem>) -> Self {
        self.file_store = Some(file_store);
        self
    }

    /// Supply the session key/value store read by prompt capabilities whose
    /// contribution is persisted session state — `channel_context` and the
    /// `ThreadContext` it renders (EVE-977).
    pub fn with_session_storage(
        mut self,
        session_storage: Arc<dyn crate::session_services::SessionStorageStore>,
    ) -> Self {
        self.session_storage = Some(session_storage);
        self
    }

    /// Resolve registered MCP servers that OpenAI calls as a hosted tool
    /// (EVE-1115). Without it such servers fail the turn with a clear error.
    pub fn with_hosted_mcp_resolver(
        mut self,
        resolver: Option<Arc<dyn HostedMcpResolver>>,
    ) -> Self {
        self.hosted_mcp = resolver;
        self
    }
}

#[async_trait::async_trait]
impl TurnContextResolver for StoreTurnContextResolver {
    async fn resolve_turn_context(
        &self,
        request: TurnContextRequest,
    ) -> Result<AssembledTurnContext> {
        let snapshot = load_execution_snapshot(
            self.harness_store.as_ref(),
            self.agent_store.as_ref(),
            self.session_store.as_ref(),
            request.session_id,
        )
        .await?;
        validate_requested_topology(&snapshot, request.harness_id, request.agent_id)?;
        let mut assembled = assemble_from_snapshot(
            snapshot,
            self.message_retriever.as_ref(),
            self.provider_store.as_ref(),
            &self.capability_registry,
            &self.driver_registry,
            &request.mcp_tool_definitions,
            self.file_store.clone(),
            self.session_storage.clone(),
            request.allow_provider_managed_reduction,
            AssemblyMode::RequireMessages,
        )
        .await?;
        // Credentials join the call inside the driver, below the engine, so
        // they never reach events or the persisted config.
        if let Some(resolver) = &self.hosted_mcp {
            let driver = assembled.model.driver.clone();
            assembled.model.driver = Arc::new(HostedMcpDriver::new(driver, resolver.clone()));
        }
        Ok(assembled)
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum AssemblyMode {
    RequireMessages,
    AllowEmptyMessages,
}

#[allow(clippy::too_many_arguments)]
/// Load and assemble a reason context, requiring at least one model-visible
/// message.
pub async fn assemble_turn_context(
    harness_store: &dyn HarnessStore,
    agent_store: &dyn AgentStore,
    session_store: &dyn SessionStore,
    message_retriever: &dyn MessageRetriever,
    provider_store: &dyn ProviderStore,
    capability_registry: &CapabilityRegistry,
    driver_registry: &DriverRegistry,
    session_id: SessionId,
    harness_id: HarnessId,
    agent_id: Option<AgentId>,
    mcp_tool_definitions: &[ToolDefinition],
    file_store: Option<Arc<dyn SessionFileSystem>>,
    session_storage: Option<Arc<dyn crate::session_services::SessionStorageStore>>,
) -> Result<AssembledTurnContext> {
    let snapshot =
        load_execution_snapshot(harness_store, agent_store, session_store, session_id).await?;
    validate_requested_topology(&snapshot, harness_id, agent_id)?;
    assemble_from_snapshot(
        snapshot,
        message_retriever,
        provider_store,
        capability_registry,
        driver_registry,
        mcp_tool_definitions,
        file_store,
        session_storage,
        true,
        AssemblyMode::RequireMessages,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
/// Load and assemble a context for inspection, allowing empty history.
pub async fn inspect_turn_context(
    harness_store: &dyn HarnessStore,
    agent_store: &dyn AgentStore,
    session_store: &dyn SessionStore,
    message_retriever: &dyn MessageRetriever,
    provider_store: &dyn ProviderStore,
    capability_registry: &CapabilityRegistry,
    driver_registry: &DriverRegistry,
    session_id: SessionId,
    harness_id: HarnessId,
    agent_id: Option<AgentId>,
    mcp_tool_definitions: &[ToolDefinition],
    file_store: Option<Arc<dyn SessionFileSystem>>,
    session_storage: Option<Arc<dyn crate::session_services::SessionStorageStore>>,
) -> Result<AssembledTurnContext> {
    let snapshot =
        load_execution_snapshot(harness_store, agent_store, session_store, session_id).await?;
    validate_requested_topology(&snapshot, harness_id, agent_id)?;
    assemble_from_snapshot(
        snapshot,
        message_retriever,
        provider_store,
        capability_registry,
        driver_registry,
        mcp_tool_definitions,
        file_store,
        session_storage,
        true,
        AssemblyMode::AllowEmptyMessages,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
pub(crate) async fn inspect_turn_context_for_session(
    harness_store: &dyn HarnessStore,
    agent_store: &dyn AgentStore,
    session_store: &dyn SessionStore,
    message_retriever: &dyn MessageRetriever,
    provider_store: &dyn ProviderStore,
    capability_registry: &CapabilityRegistry,
    driver_registry: &DriverRegistry,
    session_id: SessionId,
    mcp_tool_definitions: &[ToolDefinition],
    file_store: Option<Arc<dyn SessionFileSystem>>,
    session_storage: Option<Arc<dyn crate::session_services::SessionStorageStore>>,
) -> Result<AssembledTurnContext> {
    let snapshot =
        load_execution_snapshot(harness_store, agent_store, session_store, session_id).await?;
    assemble_from_snapshot(
        snapshot,
        message_retriever,
        provider_store,
        capability_registry,
        driver_registry,
        mcp_tool_definitions,
        file_store,
        session_storage,
        true,
        AssemblyMode::AllowEmptyMessages,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
/// Assemble a reason context from a snapshot a host already loaded.
pub async fn assemble_turn_context_from_snapshot(
    snapshot: ResolvedExecutionSnapshot,
    message_retriever: &dyn MessageRetriever,
    provider_store: &dyn ProviderStore,
    capability_registry: &CapabilityRegistry,
    driver_registry: &DriverRegistry,
    mcp_tool_definitions: &[ToolDefinition],
    file_store: Option<Arc<dyn SessionFileSystem>>,
    session_storage: Option<Arc<dyn crate::session_services::SessionStorageStore>>,
) -> Result<AssembledTurnContext> {
    assemble_from_snapshot(
        snapshot,
        message_retriever,
        provider_store,
        capability_registry,
        driver_registry,
        mcp_tool_definitions,
        file_store,
        session_storage,
        true,
        AssemblyMode::RequireMessages,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
async fn assemble_from_snapshot(
    snapshot: ResolvedExecutionSnapshot,
    message_retriever: &dyn MessageRetriever,
    provider_store: &dyn ProviderStore,
    capability_registry: &CapabilityRegistry,
    driver_registry: &DriverRegistry,
    mcp_tool_definitions: &[ToolDefinition],
    file_store: Option<Arc<dyn SessionFileSystem>>,
    session_storage: Option<Arc<dyn crate::session_services::SessionStorageStore>>,
    allow_provider_managed_reduction: bool,
    mode: AssemblyMode,
) -> Result<AssembledTurnContext> {
    let resolved = resolve_snapshot_capabilities(&snapshot, capability_registry);
    let filters = collect_message_filters_only(
        &resolved.effective_overlay.capabilities,
        capability_registry,
    );
    let mut query = MessageQuery::new(snapshot.session_id);
    filters.apply_message_filters(&mut query);
    let initial_history = message_retriever.load_filtered_history(query).await?;
    let mut messages = initial_history.messages.clone();
    filters.apply_post_load_filters(&mut messages);
    if messages.is_empty() && mode == AssemblyMode::RequireMessages {
        return Err(AgentLoopError::NoMessages);
    }

    let controls_model_id = latest_model_override(&messages);
    let (model, resolved_model_id) =
        resolve_model(provider_store, controls_model_id, snapshot.default_model_id).await?;
    let mut model = resolve_model_execution_for_session(
        provider_store,
        driver_registry,
        model,
        Some(snapshot.session_id),
    )
    .await?;
    let provider_managed_reduction_option = allow_provider_managed_reduction
        .then(|| {
            provider_managed_reduction_budget(
                &resolved.resolved_capability_configs,
                capability_registry,
            )
        })
        .flatten()
        .and_then(|budget| {
            model.driver.provider_managed_reduction_option(
                &everruns_contracts::ProviderEndpoint::default(),
                &model.model,
                budget,
            )
        });
    model.provider_managed_reduction_option = provider_managed_reduction_option;

    let history = if model.provider_managed_reduction_option.is_some() {
        let filters = collect_message_filters_only_with_context(
            &resolved.resolved_capability_configs,
            capability_registry,
            true,
        );
        let mut query = MessageQuery::new(snapshot.session_id);
        filters.apply_message_filters(&mut query);
        let mut history = message_retriever.load_filtered_history(query).await?;
        filters.apply_post_load_filters(&mut history.messages);
        history
    } else {
        initial_history
    };
    let messages = if model.provider_managed_reduction_option.is_some() {
        history.messages.clone()
    } else {
        messages
    };

    assemble_resolved_turn_context(
        ResolvedTurnContextInput {
            snapshot,
            messages,
            message_source_sequence: history.source_sequence,
            model,
            resolved_model_id,
            mcp_tool_definitions: mcp_tool_definitions.to_vec(),
        },
        capability_registry,
        file_store,
        session_storage,
    )
    .await
}

fn validate_requested_topology(
    snapshot: &ResolvedExecutionSnapshot,
    harness_id: HarnessId,
    agent_id: Option<AgentId>,
) -> Result<()> {
    if snapshot.harness_id != harness_id || snapshot.agent_id != agent_id {
        return Err(AgentLoopError::config(format!(
            "resolved topology mismatch for session {}",
            snapshot.session_id
        )));
    }
    Ok(())
}

fn latest_model_override(messages: &[RuntimeMessage]) -> Option<ModelId> {
    messages
        .iter()
        .rev()
        .find(|message| message.role == RuntimeMessageRole::User)
        .and_then(|message| message.controls.as_ref())
        .and_then(|controls| controls.model_id)
}

async fn resolve_model(
    provider_store: &dyn ProviderStore,
    controls_model_id: Option<ModelId>,
    snapshot_model_id: Option<ModelId>,
) -> Result<(ModelSpec, Option<ModelId>)> {
    for model_id in [controls_model_id, snapshot_model_id].into_iter().flatten() {
        if let Some(model) = provider_store.get_model_spec(model_id).await? {
            return Ok((model, Some(model_id)));
        }
    }
    let model = provider_store
        .get_default_model_spec()
        .await?
        .ok_or_else(AgentLoopError::model_not_configured)?;
    Ok((model, None))
}

/// Resolve provider construction independently from credential-free model identity.
pub(crate) async fn resolve_model_execution(
    provider_store: &dyn ProviderStore,
    driver_registry: &DriverRegistry,
    spec: ModelSpec,
) -> Result<ResolvedModelExecution> {
    resolve_model_execution_for_session(provider_store, driver_registry, spec, None).await
}

async fn resolve_model_execution_for_session(
    provider_store: &dyn ProviderStore,
    driver_registry: &DriverRegistry,
    spec: ModelSpec,
    session: Option<SessionId>,
) -> Result<ResolvedModelExecution> {
    let config = match session {
        Some(session) => {
            provider_store
                .get_provider_config_for_session(&spec.provider, session)
                .await?
        }
        None => provider_store.get_provider_config(&spec.provider).await?,
    }
    .unwrap_or_else(|| {
        everruns_contracts::driver_registry::ProviderConfig::for_provider(
            spec.provider.clone(),
            DriverId::external(spec.provider.as_str()),
        )
    });
    let provider_type = config.provider_type.clone();
    let driver: Arc<dyn ChatDriver> = Arc::from(driver_registry.create_chat_driver(&config)?);
    Ok(ResolvedModelExecution {
        model: spec.model,
        provider: spec.provider,
        provider_type,
        driver,
        provider_managed_reduction_option: None,
    })
}

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

    #[tokio::test]
    async fn missing_model_resolution_returns_typed_terminal_error() {
        let store = InMemoryProviderStore::new();

        let error = resolve_model(&store, None, None).await.unwrap_err();

        assert!(matches!(error, AgentLoopError::ModelNotConfigured));
        assert!(error.is_non_retryable());
    }

    #[tokio::test]
    async fn explicit_model_resolution_wins_over_the_default() {
        let store = InMemoryProviderStore::with_default(ModelSpec::on("default", "fallback")).await;
        let model_id = ModelId::new();
        store
            .add_model(model_id, ModelSpec::on("explicit", "selected"))
            .await;

        let (model, resolved_id) = resolve_model(&store, Some(model_id), None).await.unwrap();

        assert_eq!(model.provider.as_str(), "explicit");
        assert_eq!(model.model, "selected");
        assert_eq!(resolved_id, Some(model_id));
    }

    #[tokio::test]
    async fn default_model_resolution_has_no_message_override() {
        let store = InMemoryProviderStore::with_default(ModelSpec::on("default", "fallback")).await;

        let (model, resolved_id) = resolve_model(&store, None, None).await.unwrap();

        assert_eq!(model.provider.as_str(), "default");
        assert_eq!(model.model, "fallback");
        assert_eq!(resolved_id, None);
    }
}