oino_models/
virtual_models.rs

1//! Catalog entries that route each request to a physical model.
2//!
3//! Extensions register definitions (`register-virtual-model`) and implement the guest export
4//! `route-model`. This module holds the definitions, the route request/result wire types,
5//! validation of route results and the branch helpers for selection and router state. The
6//! routing call itself is made by `oino-core` through the extension.
7
8use indexmap::IndexMap;
9use oino_ai::{AssistantMessage, InputModality, JsonValue, Message, Model, ModelCost};
10use oino_ai::{StopReason, ThinkingLevel, ThinkingLevelMap};
11use serde::{Deserialize, Deserializer, Serialize};
12
13/// API id of virtual catalog entries. Requests for it fail unless routed.
14pub const VIRTUAL_MODEL_API: &str = "oino-virtual";
15
16/// `customType` of the session entry storing router state on a branch.
17pub const VIRTUAL_MODEL_STATE_ENTRY: &str = "oino.virtual-model-state";
18
19/// A virtual model definition without `route` (implemented by the guest export).
20#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
21#[serde(rename_all = "camelCase")]
22pub struct VirtualModelDefinition {
23    /// Provider the model is listed under (any id, also one with physical models).
24    pub provider: String,
25    /// Model id; must not be a physical model id of `provider`.
26    pub id: String,
27    /// Display name.
28    pub name: String,
29    /// Levels offered for selection (default `["off"]`).
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub thinking_levels: Option<Vec<ThinkingLevel>>,
32    /// Limit shown before the first response (unset = unknown, 0).
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub context_window: Option<u64>,
35    /// Output limit shown before the first response.
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub max_tokens: Option<u64>,
38    /// Input types offered (default text and images).
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub input: Option<Vec<InputModality>>,
41}
42
43/// Whether a model (or message `api`) names a virtual model.
44#[must_use]
45pub fn is_virtual_api(api: &str) -> bool {
46    api == VIRTUAL_MODEL_API
47}
48
49/// The catalog entry of a definition.
50#[must_use]
51pub fn create_virtual_model(definition: &VirtualModelDefinition) -> Model {
52    let levels = definition
53        .thinking_levels
54        .clone()
55        .unwrap_or_else(|| vec![ThinkingLevel::Off]);
56    let slot =
57        |level: ThinkingLevel| Some(levels.contains(&level).then(|| level.as_str().to_owned()));
58    Model {
59        id: definition.id.clone(),
60        name: definition.name.clone(),
61        api: VIRTUAL_MODEL_API.to_owned(),
62        provider: definition.provider.clone(),
63        base_url: String::new(),
64        input: definition
65            .input
66            .clone()
67            .unwrap_or_else(|| vec![InputModality::Text, InputModality::Image]),
68        input_limits: None,
69        cost: ModelCost::default(),
70        headers: None,
71        kind: None,
72        reasoning: levels.iter().any(|level| *level != ThinkingLevel::Off),
73        thinking_level_map: Some(ThinkingLevelMap {
74            off: slot(ThinkingLevel::Off),
75            minimal: slot(ThinkingLevel::Minimal),
76            low: slot(ThinkingLevel::Low),
77            medium: slot(ThinkingLevel::Medium),
78            high: slot(ThinkingLevel::High),
79            xhigh: slot(ThinkingLevel::Xhigh),
80            max: slot(ThinkingLevel::Max),
81        }),
82        prompt_cache: None,
83        context_window: definition.context_window.unwrap_or(0),
84        max_tokens: definition.max_tokens.unwrap_or(0),
85        sampling_params: None,
86        sampling_params_by_thinking_level: None,
87        compat: None,
88    }
89}
90
91/// A registered virtual model.
92#[derive(Clone, Debug, PartialEq)]
93pub struct RegisteredVirtualModel {
94    /// The definition as registered.
95    pub definition: VirtualModelDefinition,
96    /// Extension implementing `route-model`.
97    pub extension_id: String,
98    /// Catalog entry.
99    pub model: Model,
100}
101
102/// Errors registering a virtual model.
103#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
104pub enum VirtualModelError {
105    /// Empty provider or id.
106    #[error("Virtual model provider and id must not be empty.")]
107    EmptyId,
108    /// The id is a physical model of the provider.
109    #[error("Virtual model {provider}/{id} conflicts with a physical model.")]
110    ConflictsWithPhysical {
111        /// Provider id.
112        provider: String,
113        /// Model id.
114        id: String,
115    },
116}
117
118/// Virtual model definitions by provider, then id.
119#[derive(Clone, Debug, Default, PartialEq)]
120pub struct VirtualModelRegistry {
121    by_provider: IndexMap<String, IndexMap<String, RegisteredVirtualModel>>,
122}
123
124impl VirtualModelRegistry {
125    /// Registers or replaces a definition. `is_physical` reports whether `provider/id` is a
126    /// physical catalog model.
127    ///
128    /// # Errors
129    ///
130    /// [`VirtualModelError`] for empty ids or a physical conflict.
131    pub fn register(
132        &mut self,
133        extension_id: &str,
134        definition: VirtualModelDefinition,
135        is_physical: bool,
136    ) -> Result<(), VirtualModelError> {
137        if definition.provider.trim().is_empty() || definition.id.trim().is_empty() {
138            return Err(VirtualModelError::EmptyId);
139        }
140        if is_physical {
141            return Err(VirtualModelError::ConflictsWithPhysical {
142                provider: definition.provider,
143                id: definition.id,
144            });
145        }
146        let model = create_virtual_model(&definition);
147        self.by_provider
148            .entry(definition.provider.clone())
149            .or_default()
150            .insert(
151                definition.id.clone(),
152                RegisteredVirtualModel {
153                    definition,
154                    extension_id: extension_id.to_owned(),
155                    model,
156                },
157            );
158        Ok(())
159    }
160
161    /// Removes a definition; `false` when it was not registered.
162    pub fn unregister(&mut self, provider: &str, id: &str) -> bool {
163        let Some(models) = self.by_provider.get_mut(provider) else {
164            return false;
165        };
166        let removed = models.shift_remove(id).is_some();
167        if models.is_empty() {
168            self.by_provider.shift_remove(provider);
169        }
170        removed
171    }
172
173    /// Removes every definition registered by `extension_id`.
174    pub fn remove_extension(&mut self, extension_id: &str) {
175        for models in self.by_provider.values_mut() {
176            models.retain(|_, entry| entry.extension_id != extension_id);
177        }
178        self.by_provider.retain(|_, models| !models.is_empty());
179    }
180
181    /// A registered definition.
182    #[must_use]
183    pub fn get(&self, provider: &str, id: &str) -> Option<&RegisteredVirtualModel> {
184        self.by_provider.get(provider)?.get(id)
185    }
186
187    /// Catalog entries listed under `provider`.
188    pub fn models_for(&self, provider: &str) -> impl Iterator<Item = &Model> {
189        self.by_provider
190            .get(provider)
191            .into_iter()
192            .flat_map(|models| models.values().map(|entry| &entry.model))
193    }
194
195    /// Providers with at least one virtual model, in registration order.
196    pub fn provider_ids(&self) -> impl Iterator<Item = &str> {
197        self.by_provider.keys().map(String::as_str)
198    }
199
200    /// Every registered virtual model.
201    pub fn iter(&self) -> impl Iterator<Item = &RegisteredVirtualModel> {
202        self.by_provider.values().flat_map(IndexMap::values)
203    }
204}
205
206/// Why a request is routed.
207#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
208#[serde(rename_all = "snake_case")]
209pub enum RouteReason {
210    /// First request after a user-written message (prompt, steering, follow-up).
211    User,
212    /// Any other request of the agent loop.
213    Continuation,
214    /// Automatic retry after a failed request.
215    Retry,
216    /// Request outside the agent loop (compaction summary, extension call).
217    Direct,
218}
219
220/// Physical model and level of the latest successful response.
221#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
222#[serde(rename_all = "camelCase")]
223pub struct RoutePrevious {
224    /// Physical model.
225    pub model: Model,
226    /// Level it ran with.
227    #[serde(default, skip_serializing_if = "Option::is_none")]
228    pub thinking_level: Option<ThinkingLevel>,
229}
230
231/// The failed request of a retry.
232#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
233#[serde(rename_all = "camelCase")]
234pub struct RouteFailed {
235    /// Physical model.
236    pub model: Model,
237    /// Level it ran with.
238    #[serde(default, skip_serializing_if = "Option::is_none")]
239    pub thinking_level: Option<ThinkingLevel>,
240    /// The failed assistant message (`stopReason`, `errorMessage`).
241    pub message: AssistantMessage,
242}
243
244/// A route request (the abort signal is the host call's cancellation).
245#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
246#[serde(rename_all = "camelCase")]
247pub struct RouteRequest {
248    /// The selected virtual model.
249    pub model: Model,
250    /// The selected level (meaning is up to the router).
251    pub thinking_level: ThinkingLevel,
252    /// Why the request is made.
253    pub reason: RouteReason,
254    /// Latest successful physical response in `messages`.
255    #[serde(default, skip_serializing_if = "Option::is_none")]
256    pub previous: Option<RoutePrevious>,
257    /// For retries: the failed request.
258    #[serde(default, skip_serializing_if = "Option::is_none")]
259    pub failed: Option<RouteFailed>,
260    /// Router state last stored on the branch (absent for `direct`).
261    #[serde(default, skip_serializing_if = "Option::is_none")]
262    pub state: Option<JsonValue>,
263    /// Conversation for this request, including system messages.
264    pub messages: Vec<Message>,
265}
266
267/// `provider/id` named by a route result: either a `"provider/id"` string or a model-like
268/// object with `provider` and `id`.
269#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
270pub struct RouteTarget {
271    /// Provider id.
272    pub provider: String,
273    /// Model id.
274    pub id: String,
275}
276
277impl<'de> Deserialize<'de> for RouteTarget {
278    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
279        #[derive(Deserialize)]
280        #[serde(untagged)]
281        enum Raw {
282            Key(String),
283            Object { provider: String, id: String },
284        }
285        match Raw::deserialize(deserializer)? {
286            Raw::Object { provider, id } => Ok(Self { provider, id }),
287            Raw::Key(key) => key
288                .split_once('/')
289                .map(|(provider, id)| Self {
290                    provider: provider.to_owned(),
291                    id: id.to_owned(),
292                })
293                .ok_or_else(|| {
294                    serde::de::Error::custom(format!("route model \"{key}\" is not provider/id"))
295                }),
296        }
297    }
298}
299
300fn present<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<JsonValue>, D::Error> {
301    JsonValue::deserialize(deserializer).map(Some)
302}
303
304/// A route result as returned by the guest.
305#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
306#[serde(rename_all = "camelCase")]
307pub struct RouteResult {
308    /// Physical target.
309    pub model: RouteTarget,
310    /// Requested level (clamped to the target).
311    pub thinking_level: ThinkingLevel,
312    /// New router state; absent keeps the current state, `null` is stored as `null`.
313    #[serde(
314        default,
315        deserialize_with = "present",
316        skip_serializing_if = "Option::is_none"
317    )]
318    pub state: Option<JsonValue>,
319}
320
321/// A validated route: physical model with credentials and clamped level.
322#[derive(Clone, Debug, PartialEq)]
323pub struct ModelRoute {
324    /// Physical model.
325    pub model: Model,
326    /// Level clamped to `model`.
327    pub thinking_level: ThinkingLevel,
328    /// State returned by the router.
329    pub state: Option<JsonValue>,
330}
331
332/// Routing failures; the request ends with an error response.
333#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
334pub enum RouteError {
335    /// The selected virtual model is not registered.
336    #[error("Virtual model {provider}/{id} is not registered.")]
337    NotRegistered {
338        /// Provider id.
339        provider: String,
340        /// Model id.
341        id: String,
342    },
343    /// The target is unknown or virtual.
344    #[error("Virtual model {provider}/{id} routed to {target}, which is not a physical model.")]
345    NotPhysical {
346        /// Virtual provider.
347        provider: String,
348        /// Virtual id.
349        id: String,
350        /// `provider/id` of the target.
351        target: String,
352    },
353    /// The target's provider has no credentials.
354    #[error("Virtual model {provider}/{id} routed to {target}, which has no credentials.")]
355    NoCredentials {
356        /// Virtual provider.
357        provider: String,
358        /// Virtual id.
359        id: String,
360        /// `provider/id` of the target.
361        target: String,
362    },
363    /// The router itself failed (guest error or invalid result).
364    #[error("{message}")]
365    Router {
366        /// Error text from the guest.
367        message: String,
368    },
369}
370
371/// The latest assistant message that did not fail or abort.
372#[must_use]
373pub fn find_latest_response(messages: &[Message]) -> Option<&AssistantMessage> {
374    messages.iter().rev().find_map(|message| match message {
375        Message::Assistant(assistant)
376            if !matches!(
377                assistant.stop_reason,
378                StopReason::Error | StopReason::Aborted
379            ) =>
380        {
381            Some(assistant)
382        }
383        _ => None,
384    })
385}
386
387/// Data of a [`VIRTUAL_MODEL_STATE_ENTRY`] custom entry.
388#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
389#[serde(rename_all = "camelCase")]
390pub struct VirtualModelStateData {
391    /// Virtual model provider.
392    pub provider: String,
393    /// Virtual model id.
394    pub model_id: String,
395    /// Router state.
396    pub state: JsonValue,
397}
398
399/// Whether core stores the router state returned for a request before sending it.
400///
401/// Any returned object except `undefined` and `request.state` itself is stored. Object
402/// identity does not survive the WASM boundary, so Oino skips a state equal to the current one.
403/// `direct` requests never store state.
404#[must_use]
405pub fn should_store_route_state(
406    reason: RouteReason,
407    current: Option<&JsonValue>,
408    returned: Option<&JsonValue>,
409) -> bool {
410    match (reason, returned) {
411        (RouteReason::Direct, _) | (_, None) => false,
412        (_, Some(returned)) => current != Some(returned),
413    }
414}
415
416/// A session entry as seen by the branch helpers (core maps `oino-session` entries to this).
417#[derive(Clone, Copy, Debug, PartialEq)]
418pub enum BranchEntryRef<'a> {
419    /// `model_change`.
420    ModelChange {
421        /// Provider id.
422        provider: &'a str,
423        /// Model id.
424        model_id: &'a str,
425    },
426    /// `message` with an assistant message.
427    Assistant {
428        /// Provider of the response.
429        provider: &'a str,
430        /// Model of the response.
431        model: &'a str,
432        /// API of the response (virtual after a failed routing).
433        api: &'a str,
434    },
435    /// `custom` entry.
436    Custom {
437        /// `customType`.
438        custom_type: &'a str,
439        /// `data`.
440        data: Option<&'a JsonValue>,
441    },
442    /// Anything else.
443    Other,
444}
445
446/// A recorded model selection (`provider/modelId`).
447#[derive(Clone, Debug, PartialEq, Eq)]
448pub struct BranchSelection {
449    /// Provider id.
450    pub provider: String,
451    /// Model id.
452    pub model_id: String,
453}
454
455/// A virtual `model_change` holds until the next `model_change`;
456/// otherwise the latest physical response wins. `is_registered_virtual(provider, id)` reports a
457/// currently registered virtual model (unregistered ones fall back to the last response).
458#[must_use]
459pub fn branch_selection(
460    branch: &[BranchEntryRef<'_>],
461    is_registered_virtual: impl Fn(&str, &str) -> bool,
462) -> Option<BranchSelection> {
463    for (index, entry) in branch.iter().enumerate().rev() {
464        match *entry {
465            BranchEntryRef::ModelChange { provider, model_id } => {
466                return Some(BranchSelection {
467                    provider: provider.to_owned(),
468                    model_id: model_id.to_owned(),
469                });
470            }
471            BranchEntryRef::Assistant {
472                provider,
473                model,
474                api,
475            } if !is_virtual_api(api) => {
476                let change = branch[..index].iter().rev().find_map(|e| match *e {
477                    BranchEntryRef::ModelChange { provider, model_id } => {
478                        Some((provider, model_id))
479                    }
480                    _ => None,
481                });
482                return Some(match change {
483                    Some((p, m)) if is_registered_virtual(p, m) => BranchSelection {
484                        provider: p.to_owned(),
485                        model_id: m.to_owned(),
486                    },
487                    _ => BranchSelection {
488                        provider: provider.to_owned(),
489                        model_id: model.to_owned(),
490                    },
491                });
492            }
493            _ => {}
494        }
495    }
496    None
497}
498
499/// The latest router state stored on the branch for a model.
500#[must_use]
501pub fn virtual_model_state<'a>(
502    branch: &[BranchEntryRef<'a>],
503    provider: &str,
504    model_id: &str,
505) -> Option<&'a JsonValue> {
506    branch.iter().rev().find_map(|entry| match *entry {
507        BranchEntryRef::Custom {
508            custom_type,
509            data: Some(data),
510        } if custom_type == VIRTUAL_MODEL_STATE_ENTRY
511            && data.get("provider").and_then(JsonValue::as_str) == Some(provider)
512            && data.get("modelId").and_then(JsonValue::as_str) == Some(model_id) =>
513        {
514            data.get("state")
515        }
516        _ => None,
517    })
518}
519
520#[cfg(test)]
521mod tests {
522    use super::*;
523    use serde_json::json;
524
525    fn definition() -> VirtualModelDefinition {
526        serde_json::from_value(json!({"provider": "router", "id": "auto", "name": "Auto", "thinkingLevels": ["low", "high"]}))
527            .unwrap()
528    }
529
530    #[test]
531    fn catalog_entry_matches_pi() {
532        let model = create_virtual_model(&definition());
533        assert_eq!(
534            serde_json::to_value(&model).unwrap(),
535            json!({
536                "id": "auto", "name": "Auto", "api": "oino-virtual", "provider": "router", "baseUrl": "",
537                "input": ["text", "image"], "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
538                "reasoning": true,
539                "thinkingLevelMap": {"off": null, "minimal": null, "low": "low", "medium": null, "high": "high", "xhigh": null, "max": null},
540                "contextWindow": 0, "maxTokens": 0
541            })
542        );
543        assert_eq!(
544            model.available_thinking_levels(),
545            vec![ThinkingLevel::Low, ThinkingLevel::High]
546        );
547    }
548
549    #[test]
550    fn registry_conflicts_and_removal() {
551        let mut registry = VirtualModelRegistry::default();
552        assert_eq!(
553            registry
554                .register("ext", definition(), true)
555                .unwrap_err()
556                .to_string(),
557            "Virtual model router/auto conflicts with a physical model."
558        );
559        registry.register("ext", definition(), false).unwrap();
560        assert_eq!(registry.models_for("router").count(), 1);
561        registry.remove_extension("ext");
562        assert!(registry.get("router", "auto").is_none());
563    }
564
565    #[test]
566    fn route_result_accepts_key_or_object() {
567        let a: RouteResult =
568            serde_json::from_value(json!({"model": "anthropic/claude", "thinkingLevel": "high"}))
569                .unwrap();
570        let b: RouteResult = serde_json::from_value(json!({"model": {"provider": "anthropic", "id": "claude"}, "thinkingLevel": "high", "state": null})).unwrap();
571        assert_eq!(a.model, b.model);
572        assert_eq!(a.state, None);
573        assert_eq!(b.state, Some(JsonValue::Null));
574    }
575
576    #[test]
577    fn branch_selection_rules() {
578        let branch = [
579            BranchEntryRef::ModelChange {
580                provider: "router",
581                model_id: "auto",
582            },
583            BranchEntryRef::Assistant {
584                provider: "anthropic",
585                model: "claude",
586                api: "anthropic-messages",
587            },
588        ];
589        let selection = branch_selection(&branch, |p, m| p == "router" && m == "auto").unwrap();
590        assert_eq!(selection.model_id, "auto");
591        let selection = branch_selection(&branch, |_, _| false).unwrap();
592        assert_eq!(selection.model_id, "claude");
593        let state = json!({"provider": "router", "modelId": "auto", "state": {"phase": "plan"}});
594        let branch = [BranchEntryRef::Custom {
595            custom_type: VIRTUAL_MODEL_STATE_ENTRY,
596            data: Some(&state),
597        }];
598        assert_eq!(
599            virtual_model_state(&branch, "router", "auto"),
600            Some(&json!({"phase": "plan"}))
601        );
602    }
603
604    #[test]
605    fn state_storage_rules() {
606        let current = json!({"a": 1});
607        assert!(!should_store_route_state(
608            RouteReason::User,
609            Some(&current),
610            None
611        ));
612        assert!(!should_store_route_state(
613            RouteReason::User,
614            Some(&current),
615            Some(&json!({"a": 1}))
616        ));
617        assert!(should_store_route_state(
618            RouteReason::User,
619            Some(&current),
620            Some(&json!({"a": 2}))
621        ));
622        assert!(!should_store_route_state(
623            RouteReason::Direct,
624            None,
625            Some(&json!(1))
626        ));
627    }
628}
oino_models/virtual_models.rs

oino_models/
virtual_models.rs
