oxicode_ai/
role_routing.rs

1//! Role-routing provider — plugs role switching into the main agent loop.
2//!
3//! [`RoleRoutingProvider`] wraps a primary [`Provider`]. On each `stream()`
4//! call it builds [`RoleSignals`] from the request context, picks a role via
5//! [`decide_role`], resolves it to a concrete [`Model`] via [`RoleRegistry`],
6//! and delegates to that model's provider. When no roles are configured (the
7//! default), it is a transparent pass-through — **zero behavior change**, so
8//! the existing single-model agent path is untouched unless the user opts in
9//! via `[model_roles]` in settings.
10//!
11//! This is the "plug into the main loop" integration: because `Provider::
12//! stream()` receives the full [`Context`] on every call, the wrapper decides
13//! the model per request without any change to `oxicode-agent`'s agent loop (which
14//! otherwise resolves a single fixed `config.model_id`).
15
16use crate::providers::{Provider, StreamResult};
17use crate::role_switcher::{
18    DEFAULT_LONG_CONTEXT_THRESHOLD, RoleSignals, decide_role, resolve_role_to_model,
19};
20use crate::roles::RoleRegistry;
21use crate::{Context, Model, StreamOptions, ThinkingLevel};
22use parking_lot::RwLock;
23use std::pin::Pin;
24use std::sync::Arc;
25
26/// A [`Provider`] wrapper that routes each request to the model selected by
27/// the role-switching layer.
28///
29/// Construct it at the composition root (where the agent's primary provider is
30/// built) via [`RoleRoutingProvider::new`], then hand it to `Agent::new` in
31/// place of the raw provider.
32pub struct RoleRoutingProvider {
33    default: Arc<dyn Provider>,
34    registry: Arc<RwLock<RoleRegistry>>,
35}
36
37impl RoleRoutingProvider {
38    /// Wrap `default`. Role-selected models delegate to their own provider
39    /// (resolved via the global provider registry); anything else — including
40    /// an unset or unresolvable role — falls back to `default` with the
41    /// originally-requested model.
42    #[must_use]
43    pub fn new(default: Arc<dyn Provider>, registry: Arc<RwLock<RoleRegistry>>) -> Self {
44        Self { default, registry }
45    }
46
47    /// Mutate the live role registry (e.g. from the settings UI) so changes
48    /// apply to subsequent `stream()` calls without rewrapping the provider.
49    pub fn update_registry(&self, f: impl FnOnce(&mut RoleRegistry)) {
50        f(&mut self.registry.write());
51    }
52
53    /// Build the role signals for a request from its context + options.
54    ///
55    /// `explicit_override` and `current_tool` are not derivable from a
56    /// `stream()` call (there is no user pin or tool-execution context at the
57    /// LLM-call boundary), so they are left as `None`; the switching rests on
58    /// thinking, token count, and triviality.
59    #[must_use]
60    pub fn signals_from_request(
61        context: &Context,
62        options: &Option<StreamOptions>,
63    ) -> RoleSignals<'static> {
64        let thinking_enabled = options
65            .as_ref()
66            .and_then(|o| o.thinking_level)
67            .is_some_and(|level| level != ThinkingLevel::Off);
68        RoleSignals {
69            explicit_override: None,
70            current_tool: None,
71            thinking_enabled,
72            estimated_tokens: estimate_tokens(context),
73            long_context_threshold: DEFAULT_LONG_CONTEXT_THRESHOLD,
74            is_trivial: is_trivial(context),
75        }
76    }
77}
78
79/// Rough prompt token estimate: ~4 chars per token across all message text.
80fn estimate_tokens(context: &Context) -> usize {
81    context
82        .messages
83        .iter()
84        .map(|m| m.text_content().unwrap_or_default().len() / 4)
85        .sum()
86}
87
88/// A turn is "trivial" when the last message is short and has no code fence —
89/// a conservative heuristic for routing to the fast (`smol`) role.
90fn is_trivial(context: &Context) -> bool {
91    let last = context
92        .messages
93        .last()
94        .map(|m| m.text_content().unwrap_or_default())
95        .unwrap_or_default();
96    last.len() < 40 && !last.contains("```")
97}
98
99impl Provider for RoleRoutingProvider {
100    fn stream<'a>(
101        &'a self,
102        model: &'a Model,
103        context: &'a Context,
104        options: Option<StreamOptions>,
105    ) -> Pin<Box<dyn Future<Output = StreamResult> + Send + 'a>> {
106        let default = Arc::clone(&self.default);
107        let registry = Arc::clone(&self.registry);
108        Box::pin(async move {
109            // Decide the role model while holding the read lock, then drop the
110            // guard before any `.await` so the future stays `Send`.
111            let role_model = {
112                let reg = registry.read();
113                if reg.is_empty() {
114                    None
115                } else {
116                    let signals = Self::signals_from_request(context, &options);
117                    let role = decide_role(&signals);
118                    resolve_role_to_model(role, &reg)
119                        .filter(|m| m.provider != model.provider || m.id != model.id)
120                }
121            };
122            let Some(role_model) = role_model else {
123                return default.stream(model, context, options).await;
124            };
125            // Delegate to the role model's provider, gracefully degrading to the
126            // default model on any failure (e.g. a cross-provider role whose key
127            // isn't set) instead of failing the whole turn.
128            match crate::get_provider_arc(&role_model.provider) {
129                Some(provider) => {
130                    match provider.stream(&role_model, context, options.clone()).await {
131                        Ok(stream) => Ok(stream),
132                        Err(err) => {
133                            tracing::warn!(
134                                target: "role-router",
135                                error = %err,
136                                "role-model provider failed; falling back to default model"
137                            );
138                            default.stream(model, context, options).await
139                        }
140                    }
141                }
142                None => default.stream(model, context, options).await,
143            }
144        })
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use crate::{Context, Message, StreamOptions, ThinkingLevel};
152
153    fn ctx_with_last(text: &str) -> Context {
154        Context {
155            messages: vec![Message::user(text)],
156            ..Context::default()
157        }
158    }
159
160    #[test]
161    fn signals_thinking_from_options() {
162        let ctx = ctx_with_last("explain");
163        let opts = StreamOptions::default().thinking_level(ThinkingLevel::High);
164        let s = RoleRoutingProvider::signals_from_request(&ctx, &Some(opts));
165        assert!(s.thinking_enabled);
166    }
167
168    #[test]
169    fn signals_thinking_off_is_disabled() {
170        let ctx = ctx_with_last("explain");
171        let opts = StreamOptions::default().thinking_level(ThinkingLevel::Off);
172        let s = RoleRoutingProvider::signals_from_request(&ctx, &Some(opts));
173        assert!(!s.thinking_enabled);
174    }
175
176    #[test]
177    fn signals_no_options_is_not_thinking() {
178        let ctx = ctx_with_last("explain");
179        let s = RoleRoutingProvider::signals_from_request(&ctx, &None);
180        assert!(!s.thinking_enabled);
181    }
182
183    #[test]
184    fn signals_long_context_exceeds_threshold() {
185        // ~80k chars of text → ~20k tokens estimate crosses the 60k threshold
186        // only if we add enough; use a big message to push past 60_000 tokens.
187        let big = "x".repeat(60_000 * 4 + 100);
188        let ctx = ctx_with_last(&big);
189        let s = RoleRoutingProvider::signals_from_request(&ctx, &None);
190        assert!(
191            s.estimated_tokens > DEFAULT_LONG_CONTEXT_THRESHOLD,
192            "estimated {} should exceed {}",
193            s.estimated_tokens,
194            DEFAULT_LONG_CONTEXT_THRESHOLD
195        );
196    }
197
198    #[test]
199    fn signals_short_message_is_trivial() {
200        let ctx = ctx_with_last("hi");
201        let s = RoleRoutingProvider::signals_from_request(&ctx, &None);
202        assert!(s.is_trivial);
203    }
204
205    #[test]
206    fn signals_code_fence_is_not_trivial() {
207        let ctx = ctx_with_last("```\ncode\n```");
208        let s = RoleRoutingProvider::signals_from_request(&ctx, &None);
209        assert!(!s.is_trivial);
210    }
211
212    #[test]
213    fn estimate_tokens_scales_with_text() {
214        let small = ctx_with_last("hi");
215        let large = ctx_with_last(&"x".repeat(4_000));
216        assert!(estimate_tokens(&large) > estimate_tokens(&small));
217    }
218
219    #[test]
220    fn empty_registry_passes_through() {
221        // Structural check: an empty registry short-circuits to the default
222        // provider with the request model unchanged (verified by the early
223        // return in `stream`). The provider-selection logic (decide_role +
224        // resolve_role_to_model) is unit-tested in role_switcher.rs.
225        let r = RoleRegistry::new();
226        assert!(r.is_empty());
227    }
228}
oxicode_ai/role_routing.rs

oxicode_ai/
role_routing.rs
