1use 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#[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 #[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 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)]
101pub 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)]
134pub 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)]
194pub 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
302pub(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}