oino_core/session/
model.rs

1//! Model and thinking-level control (`setModel`, `cycleModel`, `setThinkingLevel`,
2//! `cycleThinkingLevel`, `routedModel`, `_recordSelection`, virtual routing).
3
4use oino_ai::{JsonValue, Message, Model, ThinkingLevel};
5use oino_models::resolver::{
6    ThinkingSwitch, next_available_model, next_scoped_model, thinking_level_for_model_switch,
7};
8use oino_models::virtual_models::{
9    BranchEntryRef, ModelRoute, RouteReason, RouteResult, VIRTUAL_MODEL_STATE_ENTRY,
10    VirtualModelStateData, branch_selection, find_latest_response, is_virtual_api,
11    should_store_route_state, virtual_model_state,
12};
13use oino_session::{EntryKind, SessionEntry};
14use serde::{Deserialize, Serialize};
15use serde_json::json;
16use tokio_util::sync::CancellationToken;
17
18use super::{AgentSession, SessionOpError, lock};
19use crate::events::{AgentSessionEvent, ModelChangeSource};
20use crate::extensions::api::CallKind;
21
22pub use oino_models::resolver::CycleDirection;
23
24/// The model cycle result record.
25#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
26#[serde(rename_all = "camelCase")]
27pub struct CycleModelResult {
28    /// New model.
29    pub model: Model,
30    /// New thinking level.
31    pub thinking_level: ThinkingLevel,
32    /// Cycled through the `--models` scope.
33    pub is_scoped: bool,
34}
35
36/// Branch entries as the virtual-model helpers see them.
37pub(crate) fn branch_refs(branch: &[SessionEntry]) -> Vec<BranchEntryRef<'_>> {
38    branch
39        .iter()
40        .map(|entry| match &entry.kind {
41            EntryKind::ModelChange { provider, model_id } => {
42                BranchEntryRef::ModelChange { provider, model_id }
43            }
44            EntryKind::Message {
45                message: oino_ai::AgentMessage::Assistant(assistant),
46            } => BranchEntryRef::Assistant {
47                provider: &assistant.provider,
48                model: &assistant.model,
49                api: &assistant.api,
50            },
51            EntryKind::Custom { custom_type, data } => BranchEntryRef::Custom {
52                custom_type,
53                data: data.as_ref(),
54            },
55            _ => BranchEntryRef::Other,
56        })
57        .collect()
58}
59
60impl AgentSession {
61    /// Under a virtual selection: the physical model of the latest successful response.
62    #[must_use]
63    pub fn routed_model(&self) -> Option<Model> {
64        let model = self.model()?;
65        if !is_virtual_api(&model.api) {
66            return None;
67        }
68        let llm = oino_session::convert_to_llm(&self.messages());
69        let latest = find_latest_response(&llm)?;
70        self.inner
71            .models
72            .find_physical(&latest.provider, &latest.model)
73    }
74
75    /// Thinking levels the current model supports (all levels without a model).
76    #[must_use]
77    pub fn available_thinking_levels(&self) -> Vec<ThinkingLevel> {
78        self.model().map_or_else(
79            || ThinkingLevel::ALL.to_vec(),
80            |model| model.available_thinking_levels(),
81        )
82    }
83
84    fn thinking_for_switch(
85        &self,
86        target: Option<&Model>,
87        explicit: Option<ThinkingLevel>,
88    ) -> ThinkingLevel {
89        let settings = lock(&self.inner.settings).snapshot();
90        let levels = settings.model_thinking_levels();
91        thinking_level_for_model_switch(
92            target,
93            &ThinkingSwitch {
94                explicit,
95                model_thinking_levels: Some(&levels),
96                default_thinking_level: settings.default_thinking_level(),
97                current: Some(self.thinking_level()),
98            },
99        )
100    }
101
102    async fn emit_model_select(&self, next: &Model, previous: Option<&Model>, source: &str) {
103        if previous.is_some_and(|p| p.provider == next.provider && p.id == next.id) {
104            return;
105        }
106        let runner = self.runner();
107        if runner.has_handlers("model_select") {
108            let mut payload = json!({"type": "model_select", "model": next, "source": source});
109            if let Some(previous) = previous {
110                payload["previousModel"] = json!(previous);
111            }
112            runner.emit(payload).await;
113        }
114    }
115
116    fn apply_model(&self, model: &Model, persist: bool) -> Result<(), SessionOpError> {
117        self.inner.agent.set_model(model.clone());
118        lock(&self.inner.manager)
119            .append_model_change(&model.provider, &model.id)
120            .map_err(|e| SessionOpError::Failed(e.to_string()))?;
121        if persist {
122            let result = self
123                .with_settings(|s| s.set_default_model_and_provider(&model.provider, &model.id));
124            if let Err(error) = result {
125                tracing::warn!("saving default model failed: {error}");
126            }
127            self.add_persisted_default_to_scope(model);
128        }
129        Ok(())
130    }
131
132    fn add_persisted_default_to_scope(&self, model: &Model) {
133        let mut scope = lock(&self.inner.scoped_models);
134        if scope.is_empty()
135            || scope
136                .iter()
137                .any(|s| s.model.provider == model.provider && s.model.id == model.id)
138        {
139            return;
140        }
141        scope.push(oino_models::resolver::ScopedModel {
142            model: model.clone(),
143            thinking_level: None,
144        });
145        drop(scope);
146        let enabled = lock(&self.inner.settings).settings().enabled_models();
147        let Some(mut enabled) = enabled.filter(|e| !e.is_empty()) else {
148            return;
149        };
150        let reference = format!("{}/{}", model.provider, model.id);
151        if enabled.iter().any(|p| p.eq_ignore_ascii_case(&reference)) {
152            return;
153        }
154        enabled.push(reference);
155        if let Err(error) = self.with_settings(|s| s.set_enabled_models(Some(&enabled))) {
156            tracing::warn!("saving enabledModels failed: {error}");
157        }
158    }
159
160    /// Validates auth, records `model_change`, applies the thinking level.
161    ///
162    /// # Errors
163    ///
164    /// `No API key for {provider}/{id}` when the provider has no credentials.
165    pub async fn set_model(&self, model: Model, persist: bool) -> Result<Model, SessionOpError> {
166        if self.inner.models.check_auth(&model.provider).is_none() {
167            return Err(SessionOpError::Failed(format!(
168                "No API key for {}/{}",
169                model.provider, model.id
170            )));
171        }
172        let previous = self.model();
173        let thinking = self.thinking_for_switch(Some(&model), None);
174        self.apply_model(&model, persist)?;
175        self.set_thinking_level(thinking, false).await;
176        self.emit_model_select(&model, previous.as_ref(), "set")
177            .await;
178        self.emit_model_changed(ModelChangeSource::Set);
179        Ok(model)
180    }
181
182    /// C y c l e   M o d e l.
183    ///
184    /// # Errors
185    ///
186    /// Persisting the selection failed.
187    pub async fn cycle_model(
188        &self,
189        direction: CycleDirection,
190        persist: bool,
191    ) -> Result<Option<CycleModelResult>, SessionOpError> {
192        let available = self.inner.models.available();
193        let current = self.model();
194        let scope = self.scoped_models();
195        let (next, explicit, is_scoped) = if scope.is_empty() {
196            match next_available_model(&available, current.as_ref(), direction) {
197                Some(model) => (model, None, false),
198                None => return Ok(None),
199            }
200        } else {
201            match next_scoped_model(&scope, &available, current.as_ref(), direction) {
202                Some(scoped) => (scoped.model, scoped.thinking_level, true),
203                None => return Ok(None),
204            }
205        };
206        let thinking = self.thinking_for_switch(Some(&next), explicit);
207        self.apply_model(&next, persist)?;
208        self.set_thinking_level(thinking, false).await;
209        self.emit_model_select(&next, current.as_ref(), "cycle")
210            .await;
211        self.emit_model_changed(ModelChangeSource::Cycle);
212        Ok(Some(CycleModelResult {
213            model: next,
214            thinking_level: self.thinking_level(),
215            is_scoped,
216        }))
217    }
218
219    /// Clamps, records only changes; returns the effective level.
220    pub async fn set_thinking_level(&self, level: ThinkingLevel, persist: bool) -> ThinkingLevel {
221        let available = self.available_thinking_levels();
222        let effective = if available.contains(&level) {
223            level
224        } else if self.model().is_some() {
225            level.clamp(&available)
226        } else {
227            ThinkingLevel::Off
228        };
229        let previous = self.thinking_level();
230        self.inner.agent.set_thinking_level(effective);
231        if persist && let Err(error) = self.with_settings(|s| s.set_default_thinking_level(level)) {
232            tracing::warn!("saving defaultThinkingLevel failed: {error}");
233        }
234        if effective != previous {
235            if let Err(error) =
236                lock(&self.inner.manager).append_thinking_level_change(&effective.to_string())
237            {
238                tracing::error!("persisting thinking level failed: {error}");
239            }
240            self.emit(AgentSessionEvent::ThinkingLevelChanged { level: effective });
241            let runner = self.runner();
242            if runner.has_handlers("thinking_level_select") {
243                runner
244                    .emit(json!({"type": "thinking_level_select", "level": effective, "previousLevel": previous}))
245                    .await;
246            }
247        }
248        effective
249    }
250
251    /// Cycles the thinking level; `None` when the model has no reasoning.
252    pub async fn cycle_thinking_level(&self, persist: bool) -> Option<ThinkingLevel> {
253        if !self.model().is_some_and(|m| m.reasoning) {
254            return None;
255        }
256        let levels = self.available_thinking_levels();
257        let index = levels.iter().position(|l| *l == self.thinking_level());
258        let next = levels[index.map_or(0, |i| (i + 1) % levels.len())];
259        self.set_thinking_level(next, persist).await;
260        Some(next)
261    }
262
263    /// Appends `model_change` when the branch implies another model.
264    pub(crate) fn record_selection(&self) {
265        let Some(model) = self.model() else {
266            return;
267        };
268        let models = &self.inner.models;
269        let recorded = {
270            let manager = lock(&self.inner.manager);
271            let branch: Vec<SessionEntry> = manager.branch(None).into_iter().cloned().collect();
272            branch_selection(&branch_refs(&branch), |provider, id| {
273                models.virtual_model(provider, id).is_some()
274            })
275        };
276        let Some(recorded) = recorded else {
277            return;
278        };
279        if recorded.provider == model.provider && recorded.model_id == model.id {
280            return;
281        }
282        let recorded_virtual = models
283            .find(&recorded.provider, &recorded.model_id)
284            .is_some_and(|m| is_virtual_api(&m.api));
285        if !is_virtual_api(&model.api) && !recorded_virtual {
286            return;
287        }
288        if let Err(error) =
289            lock(&self.inner.manager).append_model_change(&model.provider, &model.id)
290        {
291            tracing::error!("recording model selection failed: {error}");
292        }
293    }
294
295    /// Routes a virtual model through its extension, storing
296    /// router state on the branch.
297    pub(crate) async fn route_virtual(
298        &self,
299        model: &Model,
300        messages: Vec<Message>,
301        reason: RouteReason,
302        thinking: ThinkingLevel,
303        failed: Option<oino_ai::AssistantMessage>,
304        _cancel: &CancellationToken,
305    ) -> Result<ModelRoute, String> {
306        let state: Option<JsonValue> = {
307            let manager = lock(&self.inner.manager);
308            let branch: Vec<SessionEntry> = manager.branch(None).into_iter().cloned().collect();
309            virtual_model_state(&branch_refs(&branch), &model.provider, &model.id).cloned()
310        };
311        let request = self
312            .inner
313            .models
314            .route_request(model, messages, reason, thinking, failed, state.clone())
315            .map_err(|e| e.to_string())?;
316        let registered = self
317            .inner
318            .models
319            .virtual_model(&model.provider, &model.id)
320            .ok_or_else(|| {
321                format!(
322                    "Virtual model {}/{} is not registered.",
323                    model.provider, model.id
324                )
325            })?;
326        let runner = self.runner();
327        let extension = runner
328            .set()
329            .get(&registered.extension_id)
330            .cloned()
331            .ok_or_else(|| format!("Extension {} is not loaded", registered.extension_id))?;
332        let cx = runner.call_context().with_kind(CallKind::Service);
333        let response = extension
334            .instance
335            .route_model(json!(request), cx)
336            .await
337            .map_err(|e| e.message)?;
338        let result: RouteResult = serde_json::from_value(response).map_err(|e| e.to_string())?;
339        let route = self
340            .inner
341            .models
342            .validate_route(model, result)
343            .map_err(|e| e.to_string())?;
344        if should_store_route_state(reason, state.as_ref(), route.state.as_ref()) {
345            let data = VirtualModelStateData {
346                provider: model.provider.clone(),
347                model_id: model.id.clone(),
348                state: route.state.clone().unwrap_or(JsonValue::Null),
349            };
350            let appended = lock(&self.inner.manager)
351                .append_custom_entry(VIRTUAL_MODEL_STATE_ENTRY, serde_json::to_value(data).ok());
352            if let Ok(id) = appended {
353                self.publish_entry(&id);
354            }
355        }
356        Ok(route)
357    }
358}
oino_core/session/model.rs

oino_core/session/
model.rs
