Skip to main content

everruns_host/
runtime_context.rs

1//! Store-backed turn-context resolution for runtime hosts.
2
3use std::sync::Arc;
4
5use everruns_core::ResolvedExecutionSnapshot;
6use everruns_core::capabilities::{CapabilityRegistry, collect_message_filters_only};
7use everruns_core::execution_loading::{AgentStore, HarnessStore, SessionStore};
8use everruns_core::message::{Message, MessageRole};
9use everruns_core::message_filter::MessageQuery;
10use everruns_core::message_retriever::MessageRetriever;
11use everruns_core::provider_resolution::ProviderStore;
12use everruns_core::runtime_context::{
13    AssembledTurnContext, ResolvedModelExecution, ResolvedTurnContextInput, TurnContextRequest,
14    TurnContextResolver, assemble_resolved_turn_context, resolve_snapshot_capabilities,
15};
16use everruns_core::session_files::SessionFileSystem;
17use everruns_provider::driver_registry::{ChatDriver, DriverRegistry};
18use everruns_provider::error::{AgentLoopError, Result};
19use everruns_provider::model_spec::ModelSpec;
20use everruns_provider::provider::DriverId;
21use everruns_provider::tool_types::ToolDefinition;
22use everruns_provider::typed_id::{AgentId, HarnessId, ModelId, SessionId};
23
24use crate::execution_snapshot::load_execution_snapshot;
25
26/// Store-backed resolver used by direct atom callers and custom hosts.
27#[derive(Clone)]
28pub struct StoreTurnContextResolver {
29    harness_store: Arc<dyn HarnessStore>,
30    agent_store: Arc<dyn AgentStore>,
31    session_store: Arc<dyn SessionStore>,
32    message_retriever: Arc<dyn MessageRetriever>,
33    provider_store: Arc<dyn ProviderStore>,
34    capability_registry: CapabilityRegistry,
35    driver_registry: DriverRegistry,
36    file_store: Option<Arc<dyn SessionFileSystem>>,
37}
38
39impl StoreTurnContextResolver {
40    /// Construct a resolver from org-scoped runtime stores and registries.
41    #[allow(clippy::too_many_arguments)]
42    pub fn new(
43        harness_store: Arc<dyn HarnessStore>,
44        agent_store: Arc<dyn AgentStore>,
45        session_store: Arc<dyn SessionStore>,
46        message_retriever: Arc<dyn MessageRetriever>,
47        provider_store: Arc<dyn ProviderStore>,
48        capability_registry: CapabilityRegistry,
49        driver_registry: DriverRegistry,
50    ) -> Self {
51        Self {
52            harness_store,
53            agent_store,
54            session_store,
55            message_retriever,
56            provider_store,
57            capability_registry,
58            driver_registry,
59            file_store: None,
60        }
61    }
62
63    /// Supply the session filesystem used by dynamic prompt capabilities.
64    pub fn with_file_store(mut self, file_store: Arc<dyn SessionFileSystem>) -> Self {
65        self.file_store = Some(file_store);
66        self
67    }
68}
69
70#[async_trait::async_trait]
71impl TurnContextResolver for StoreTurnContextResolver {
72    async fn resolve_turn_context(
73        &self,
74        request: TurnContextRequest,
75    ) -> Result<AssembledTurnContext> {
76        assemble_turn_context(
77            self.harness_store.as_ref(),
78            self.agent_store.as_ref(),
79            self.session_store.as_ref(),
80            self.message_retriever.as_ref(),
81            self.provider_store.as_ref(),
82            &self.capability_registry,
83            &self.driver_registry,
84            request.session_id,
85            request.harness_id,
86            request.agent_id,
87            &request.mcp_tool_definitions,
88            self.file_store.clone(),
89        )
90        .await
91    }
92}
93
94#[derive(Clone, Copy, Debug, Eq, PartialEq)]
95enum AssemblyMode {
96    RequireMessages,
97    AllowEmptyMessages,
98}
99
100#[allow(clippy::too_many_arguments)]
101/// Load and assemble a reason context, requiring at least one model-visible
102/// message.
103pub async fn assemble_turn_context(
104    harness_store: &dyn HarnessStore,
105    agent_store: &dyn AgentStore,
106    session_store: &dyn SessionStore,
107    message_retriever: &dyn MessageRetriever,
108    provider_store: &dyn ProviderStore,
109    capability_registry: &CapabilityRegistry,
110    driver_registry: &DriverRegistry,
111    session_id: SessionId,
112    harness_id: HarnessId,
113    agent_id: Option<AgentId>,
114    mcp_tool_definitions: &[ToolDefinition],
115    file_store: Option<Arc<dyn SessionFileSystem>>,
116) -> Result<AssembledTurnContext> {
117    let snapshot =
118        load_execution_snapshot(harness_store, agent_store, session_store, session_id).await?;
119    validate_requested_topology(&snapshot, harness_id, agent_id)?;
120    assemble_from_snapshot(
121        snapshot,
122        message_retriever,
123        provider_store,
124        capability_registry,
125        driver_registry,
126        mcp_tool_definitions,
127        file_store,
128        AssemblyMode::RequireMessages,
129    )
130    .await
131}
132
133#[allow(clippy::too_many_arguments)]
134/// Load and assemble a context for inspection, allowing empty history.
135pub async fn inspect_turn_context(
136    harness_store: &dyn HarnessStore,
137    agent_store: &dyn AgentStore,
138    session_store: &dyn SessionStore,
139    message_retriever: &dyn MessageRetriever,
140    provider_store: &dyn ProviderStore,
141    capability_registry: &CapabilityRegistry,
142    driver_registry: &DriverRegistry,
143    session_id: SessionId,
144    harness_id: HarnessId,
145    agent_id: Option<AgentId>,
146    mcp_tool_definitions: &[ToolDefinition],
147    file_store: Option<Arc<dyn SessionFileSystem>>,
148) -> Result<AssembledTurnContext> {
149    let snapshot =
150        load_execution_snapshot(harness_store, agent_store, session_store, session_id).await?;
151    validate_requested_topology(&snapshot, harness_id, agent_id)?;
152    assemble_from_snapshot(
153        snapshot,
154        message_retriever,
155        provider_store,
156        capability_registry,
157        driver_registry,
158        mcp_tool_definitions,
159        file_store,
160        AssemblyMode::AllowEmptyMessages,
161    )
162    .await
163}
164
165#[allow(clippy::too_many_arguments)]
166pub(crate) async fn inspect_turn_context_for_session(
167    harness_store: &dyn HarnessStore,
168    agent_store: &dyn AgentStore,
169    session_store: &dyn SessionStore,
170    message_retriever: &dyn MessageRetriever,
171    provider_store: &dyn ProviderStore,
172    capability_registry: &CapabilityRegistry,
173    driver_registry: &DriverRegistry,
174    session_id: SessionId,
175    mcp_tool_definitions: &[ToolDefinition],
176    file_store: Option<Arc<dyn SessionFileSystem>>,
177) -> Result<AssembledTurnContext> {
178    let snapshot =
179        load_execution_snapshot(harness_store, agent_store, session_store, session_id).await?;
180    assemble_from_snapshot(
181        snapshot,
182        message_retriever,
183        provider_store,
184        capability_registry,
185        driver_registry,
186        mcp_tool_definitions,
187        file_store,
188        AssemblyMode::AllowEmptyMessages,
189    )
190    .await
191}
192
193#[allow(clippy::too_many_arguments)]
194/// Assemble a reason context from a snapshot a host already loaded.
195pub async fn assemble_turn_context_from_snapshot(
196    snapshot: ResolvedExecutionSnapshot,
197    message_retriever: &dyn MessageRetriever,
198    provider_store: &dyn ProviderStore,
199    capability_registry: &CapabilityRegistry,
200    driver_registry: &DriverRegistry,
201    mcp_tool_definitions: &[ToolDefinition],
202    file_store: Option<Arc<dyn SessionFileSystem>>,
203) -> Result<AssembledTurnContext> {
204    assemble_from_snapshot(
205        snapshot,
206        message_retriever,
207        provider_store,
208        capability_registry,
209        driver_registry,
210        mcp_tool_definitions,
211        file_store,
212        AssemblyMode::RequireMessages,
213    )
214    .await
215}
216
217#[allow(clippy::too_many_arguments)]
218async fn assemble_from_snapshot(
219    snapshot: ResolvedExecutionSnapshot,
220    message_retriever: &dyn MessageRetriever,
221    provider_store: &dyn ProviderStore,
222    capability_registry: &CapabilityRegistry,
223    driver_registry: &DriverRegistry,
224    mcp_tool_definitions: &[ToolDefinition],
225    file_store: Option<Arc<dyn SessionFileSystem>>,
226    mode: AssemblyMode,
227) -> Result<AssembledTurnContext> {
228    let resolved = resolve_snapshot_capabilities(&snapshot, capability_registry);
229    let filters = collect_message_filters_only(
230        &resolved.effective_overlay.capabilities,
231        capability_registry,
232    );
233    let mut query = MessageQuery::new(snapshot.session_id);
234    filters.apply_message_filters(&mut query);
235    let history = message_retriever.load_filtered_history(query).await?;
236    let mut messages = history.messages;
237    filters.apply_post_load_filters(&mut messages);
238    if messages.is_empty() && mode == AssemblyMode::RequireMessages {
239        return Err(AgentLoopError::NoMessages);
240    }
241
242    let controls_model_id = latest_model_override(&messages);
243    let (model, resolved_model_id) =
244        resolve_model(provider_store, controls_model_id, snapshot.default_model_id).await?;
245    let model = resolve_model_execution(provider_store, driver_registry, model).await?;
246
247    assemble_resolved_turn_context(
248        ResolvedTurnContextInput {
249            snapshot,
250            messages,
251            message_source_sequence: history.source_sequence,
252            model,
253            resolved_model_id,
254            mcp_tool_definitions: mcp_tool_definitions.to_vec(),
255        },
256        capability_registry,
257        file_store,
258    )
259    .await
260}
261
262fn validate_requested_topology(
263    snapshot: &ResolvedExecutionSnapshot,
264    harness_id: HarnessId,
265    agent_id: Option<AgentId>,
266) -> Result<()> {
267    if snapshot.harness_id != harness_id || snapshot.agent_id != agent_id {
268        return Err(AgentLoopError::config(format!(
269            "resolved topology mismatch for session {}",
270            snapshot.session_id
271        )));
272    }
273    Ok(())
274}
275
276fn latest_model_override(messages: &[Message]) -> Option<ModelId> {
277    messages
278        .iter()
279        .rev()
280        .find(|message| message.role == MessageRole::User)
281        .and_then(|message| message.controls.as_ref())
282        .and_then(|controls| controls.model_id)
283}
284
285async fn resolve_model(
286    provider_store: &dyn ProviderStore,
287    controls_model_id: Option<ModelId>,
288    snapshot_model_id: Option<ModelId>,
289) -> Result<(ModelSpec, Option<ModelId>)> {
290    for model_id in [controls_model_id, snapshot_model_id].into_iter().flatten() {
291        if let Some(model) = provider_store.get_model_spec(model_id).await? {
292            return Ok((model, Some(model_id)));
293        }
294    }
295    let model = provider_store
296        .get_default_model_spec()
297        .await?
298        .ok_or_else(AgentLoopError::model_not_configured)?;
299    Ok((model, None))
300}
301
302/// Resolve provider construction independently from credential-free model identity.
303pub(crate) async fn resolve_model_execution(
304    provider_store: &dyn ProviderStore,
305    driver_registry: &DriverRegistry,
306    spec: ModelSpec,
307) -> Result<ResolvedModelExecution> {
308    let config = provider_store
309        .get_provider_config(&spec.provider)
310        .await?
311        .unwrap_or_else(|| {
312            everruns_provider::driver_registry::ProviderConfig::for_provider(
313                spec.provider.clone(),
314                DriverId::external(spec.provider.as_str()),
315            )
316        });
317    let provider_type = config.provider_type.clone();
318    let driver: Arc<dyn ChatDriver> = Arc::from(driver_registry.create_chat_driver(&config)?);
319    Ok(ResolvedModelExecution {
320        model: spec.model,
321        provider: spec.provider,
322        provider_type,
323        driver,
324    })
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330    use crate::InMemoryProviderStore;
331
332    #[tokio::test]
333    async fn missing_model_resolution_returns_typed_terminal_error() {
334        let store = InMemoryProviderStore::new();
335
336        let error = resolve_model(&store, None, None).await.unwrap_err();
337
338        assert!(matches!(error, AgentLoopError::ModelNotConfigured));
339        assert!(error.is_non_retryable());
340    }
341
342    #[tokio::test]
343    async fn explicit_model_resolution_wins_over_the_default() {
344        let store = InMemoryProviderStore::with_default(ModelSpec::on("default", "fallback")).await;
345        let model_id = ModelId::new();
346        store
347            .add_model(model_id, ModelSpec::on("explicit", "selected"))
348            .await;
349
350        let (model, resolved_id) = resolve_model(&store, Some(model_id), None).await.unwrap();
351
352        assert_eq!(model.provider.as_str(), "explicit");
353        assert_eq!(model.model, "selected");
354        assert_eq!(resolved_id, Some(model_id));
355    }
356
357    #[tokio::test]
358    async fn default_model_resolution_has_no_message_override() {
359        let store = InMemoryProviderStore::with_default(ModelSpec::on("default", "fallback")).await;
360
361        let (model, resolved_id) = resolve_model(&store, None, None).await.unwrap();
362
363        assert_eq!(model.provider.as_str(), "default");
364        assert_eq!(model.model, "fallback");
365        assert_eq!(resolved_id, None);
366    }
367}