1use anyhow::Result;
19use oxi_sdk::{
20 CatalogConfig, FileModelCatalog, ModelCatalog, Oxi, OxiBuilder, ProviderPool, RateLimitPolicy,
21};
22use std::sync::Arc;
23
24use oxios_ouroboros::{ModelResolver, ResolvedModel};
25
26use crate::credential::{CredentialStore, discover_auth_store_providers};
27
28pub struct OxiosEngine {
42 oxi: Oxi,
43 default_model_id: String,
44 routing_control: Option<oxi_sdk::RoutingControl>,
46 pools: parking_lot::RwLock<std::collections::HashMap<String, Arc<dyn oxi_sdk::Provider>>>,
49 authorizer: Option<Arc<oxi_sdk::Authorizer>>,
54 tracer: Option<Arc<oxi_sdk::Tracer>>,
55 cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
56}
57
58impl OxiosEngine {
59 pub fn new(default_model_id: impl Into<String>) -> Self {
64 let model_id = default_model_id.into();
65 let oxi = OxiBuilder::new().with_builtins().build();
66 Self {
67 oxi,
68 default_model_id: model_id,
69 routing_control: None,
70 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
71 authorizer: None,
73 tracer: None,
74 cost_tracker: None,
75 }
76 }
77
78 pub fn from_config(default_model_id: impl Into<String>, config_api_key: Option<&str>) -> Self {
90 Self::from_config_with_catalog_opt(default_model_id, config_api_key, None)
91 }
92
93 pub fn from_config_with_catalog(
102 default_model_id: impl Into<String>,
103 config_api_key: Option<&str>,
104 catalog: Arc<dyn ModelCatalog>,
105 ) -> Self {
106 Self::from_config_with_catalog_opt(default_model_id, config_api_key, Some(catalog))
107 }
108
109 fn from_config_with_catalog_opt(
110 default_model_id: impl Into<String>,
111 config_api_key: Option<&str>,
112 catalog: Option<Arc<dyn ModelCatalog>>,
113 ) -> Self {
114 let model_id = default_model_id.into();
115
116 let primary_provider = model_id
118 .split_once('/')
119 .map(|(p, _)| p)
120 .unwrap_or("anthropic");
121
122 let mut builder = OxiBuilder::new().with_builtins();
123
124 let mut providers_to_try: Vec<String> = vec![
129 "anthropic".into(),
130 "openai".into(),
131 "google".into(),
132 "deepseek".into(),
133 "xai".into(),
134 "groq".into(),
135 "openrouter".into(),
136 "mistral".into(),
137 "cerebras".into(),
138 "fireworks".into(),
139 "github-copilot".into(),
140 "huggingface".into(),
141 "together".into(),
142 "minimax".into(),
143 "moonshotai".into(),
144 "kimi-coding".into(),
145 "zai".into(),
146 "opencode".into(),
147 ];
148
149 if let Ok(extra) = discover_auth_store_providers() {
152 for p in extra {
153 if !providers_to_try.contains(&p) {
154 providers_to_try.push(p);
155 }
156 }
157 }
158
159 let primary_owned = primary_provider.to_string();
161 if !providers_to_try.contains(&primary_owned) {
162 providers_to_try.push(primary_owned);
163 }
164
165 for provider in &providers_to_try {
166 let config_key = if provider == primary_provider {
169 config_api_key
170 } else {
171 None
172 };
173
174 if let Some((key, source)) = CredentialStore::resolve(provider, config_key) {
175 tracing::debug!(
176 provider,
177 source = ?source,
178 "Injected credential into engine"
179 );
180 builder = builder.api_key(provider, key);
181 }
182 }
183
184 let builder = match catalog {
185 Some(cat) => builder.with_catalog(cat),
186 None => builder,
187 };
188 let oxi = builder.build();
189 Self {
190 oxi,
191 default_model_id: model_id,
192 routing_control: None,
193 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
194 authorizer: None,
196 tracer: None,
197 cost_tracker: None,
198 }
199 }
200
201 pub fn builder() -> OxiosEngineBuilder {
230 OxiosEngineBuilder {
231 inner: OxiBuilder::new().with_builtins(),
232 default_model_id: "anthropic/claude-sonnet-4-20250514".to_string(),
233 authorizer: None,
235 tracer: None,
236 cost_tracker: None,
237 }
238 }
239
240 pub fn catalog_config() -> CatalogConfig {
248 let home = dirs::home_dir()
249 .unwrap_or_else(|| std::path::PathBuf::from("."))
250 .join(".oxios");
251 CatalogConfig {
252 cache_path: home.join("cache/models-dev.json"),
253 etag_path: home.join("cache/models-dev.json.etag"),
254 override_path: home.join("catalog/overrides.toml"),
255 snapshot_path: home.join("cache/models-dev.json"),
256 local_discovery_urls: Vec::new(),
258 ..CatalogConfig::default()
259 }
260 }
261
262 pub async fn init_file_catalog() -> Result<Arc<dyn ModelCatalog>> {
272 let catalog: Arc<dyn ModelCatalog> =
273 FileModelCatalog::init(Self::catalog_config())
274 .await
275 .map_err(|e| anyhow::anyhow!("Failed to initialize model catalog: {e}"))?;
276 Ok(catalog)
277 }
278
279 pub fn oxi(&self) -> &Oxi {
284 &self.oxi
285 }
286
287 pub fn authorizer(&self) -> Option<&Arc<oxi_sdk::Authorizer>> {
292 self.authorizer.as_ref()
293 }
294
295 pub fn tracer(&self) -> Option<&Arc<oxi_sdk::Tracer>> {
300 self.tracer.as_ref()
301 }
302
303 pub fn cost_tracker(&self) -> Option<&Arc<oxi_sdk::CostTracker>> {
308 self.cost_tracker.as_ref()
309 }
310
311 pub fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
313 self.oxi.resolve_model(model_id)
314 }
315
316 pub fn create_provider(&self, name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
318 self.oxi.create_provider(name)
319 }
320
321 pub fn default_model_id(&self) -> &str {
323 &self.default_model_id
324 }
325
326 pub fn routing_control(&self) -> Option<&oxi_sdk::RoutingControl> {
328 self.routing_control.as_ref()
329 }
330
331 pub fn pooled_provider(&self, name: &str, rpm: u32) -> Result<Arc<dyn oxi_sdk::Provider>> {
339 {
341 let pools = self.pools.read();
342 if let Some(pooled) = pools.get(name) {
343 return Ok(pooled.clone());
344 }
345 }
346
347 let base = self.create_provider(name)?;
349 let policy = RateLimitPolicy::rpm(rpm);
350 let pool = ProviderPool::new(base, policy, name);
351 let pooled: Arc<dyn oxi_sdk::Provider> = Arc::new(pool);
352
353 {
355 let mut pools = self.pools.write();
356 pools.insert(name.to_string(), pooled.clone());
357 }
358
359 tracing::info!(provider = name, rpm, "Created provider pool");
360 Ok(pooled)
361 }
362}
363
364pub struct OxiosEngineBuilder {
370 inner: OxiBuilder,
371 default_model_id: String,
372 authorizer: Option<Arc<oxi_sdk::Authorizer>>,
375 tracer: Option<Arc<oxi_sdk::Tracer>>,
376 cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
377}
378
379impl OxiosEngineBuilder {
380 pub fn default_model(mut self, model_id: impl Into<String>) -> Self {
382 self.default_model_id = model_id.into();
383 self
384 }
385
386 pub fn api_key(self, provider: &str, key: impl Into<String>) -> Self {
388 Self {
389 inner: self.inner.api_key(provider, key),
390 default_model_id: self.default_model_id,
391 authorizer: self.authorizer,
392 tracer: self.tracer,
393 cost_tracker: self.cost_tracker,
394 }
395 }
396
397 pub fn credential(
399 self,
400 provider: &str,
401 api_key: impl Into<String>,
402 base_url: Option<&str>,
403 ) -> Self {
404 Self {
405 inner: self.inner.credential(provider, api_key, base_url),
406 default_model_id: self.default_model_id,
407 authorizer: self.authorizer,
408 tracer: self.tracer,
409 cost_tracker: self.cost_tracker,
410 }
411 }
412
413 pub fn provider(self, name: &str, p: impl oxi_sdk::Provider + 'static) -> Self {
415 Self {
416 inner: self.inner.provider(name, p),
417 default_model_id: self.default_model_id,
418 authorizer: self.authorizer,
419 tracer: self.tracer,
420 cost_tracker: self.cost_tracker,
421 }
422 }
423
424 pub fn build(self) -> OxiosEngine {
426 OxiosEngine {
427 oxi: self.inner.build(),
428 default_model_id: self.default_model_id,
429 routing_control: None,
430 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
431 authorizer: self.authorizer,
433 tracer: self.tracer,
434 cost_tracker: self.cost_tracker,
435 }
436 }
437
438 pub fn build_with_routing(self) -> (OxiosEngine, oxi_sdk::RoutingControl) {
442 use oxi_sdk::RoutingControl;
443
444 let routing_config = oxi_sdk::routing::RoutingConfig::default();
445 let routing_control = RoutingControl::new(routing_config);
446 let engine = OxiosEngine {
447 oxi: self.inner.build(),
448 default_model_id: self.default_model_id,
449 routing_control: Some(routing_control.clone()),
450 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
451 authorizer: self.authorizer,
453 tracer: self.tracer,
454 cost_tracker: self.cost_tracker,
455 };
456 (engine, routing_control)
457 }
458
459 pub fn with_authorizer(mut self, authorizer: Arc<oxi_sdk::Authorizer>) -> Self {
471 self.authorizer = Some(authorizer);
472 self
473 }
474
475 pub fn with_tracer(mut self, tracer: Arc<oxi_sdk::Tracer>) -> Self {
478 self.tracer = Some(tracer);
479 self
480 }
481
482 pub fn with_cost_tracker(mut self, cost_tracker: Arc<oxi_sdk::CostTracker>) -> Self {
485 self.cost_tracker = Some(cost_tracker);
486 self
487 }
488
489 pub fn with_catalog(mut self, catalog: Arc<dyn oxi_sdk::ModelCatalog>) -> Self {
500 self.inner = self.inner.with_catalog(catalog);
501 self
502 }
503}
504
505pub trait EngineProvider: Send + Sync {
513 fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>>;
515
516 fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model>;
518
519 fn default_model_id(&self) -> &str;
521}
522
523impl EngineProvider for OxiosEngine {
524 fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
525 self.create_provider(provider_name)
526 }
527
528 fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
529 self.resolve_model(model_id)
530 }
531
532 fn default_model_id(&self) -> &str {
533 &self.default_model_id
534 }
535}
536
537impl std::fmt::Debug for OxiosEngine {
538 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
539 f.debug_struct("OxiosEngine")
540 .field("default_model_id", &self.default_model_id)
541 .field("routing_enabled", &self.routing_control.is_some())
542 .finish()
543 }
544}
545
546pub struct EngineHandle {
568 inner: parking_lot::RwLock<Arc<OxiosEngine>>,
569 provider_cache:
573 parking_lot::RwLock<std::collections::HashMap<String, Arc<dyn oxi_sdk::Provider>>>,
574}
575impl EngineHandle {
576 pub fn new(engine: Arc<OxiosEngine>) -> Self {
578 Self {
579 inner: parking_lot::RwLock::new(engine),
580 provider_cache: parking_lot::RwLock::new(std::collections::HashMap::new()),
581 }
582 }
583
584 pub fn get(&self) -> Arc<OxiosEngine> {
589 Arc::clone(&self.inner.read())
590 }
591
592 pub fn swap(&self, new_engine: OxiosEngine) {
597 {
598 let mut guard = self.inner.write();
599 let old_id = guard.default_model_id().to_string();
600 *guard = Arc::new(new_engine);
601 tracing::info!(
602 old_model = %old_id,
603 new_model = %guard.default_model_id(),
604 "Engine hot-swapped"
605 );
606 }
607 self.provider_cache.write().clear();
610 tracing::debug!("Provider cache cleared on engine swap");
611 }
612
613 pub fn resolve_default(&self) -> Result<ResolvedModel> {
624 let engine = self.get();
625 let model_id = engine.default_model_id().to_string();
626 let model = engine.resolve_model(&model_id)?;
627 let provider = self.cached_provider(&model.provider)?;
628 Ok(ResolvedModel {
629 model,
630 provider,
631 model_id,
632 })
633 }
634
635 fn cached_provider(&self, name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
637 if let Some(p) = self.provider_cache.read().get(name) {
638 return Ok(Arc::clone(p));
639 }
640 let provider = self.get().create_provider(name)?;
641 self.provider_cache
642 .write()
643 .insert(name.to_string(), Arc::clone(&provider));
644 Ok(provider)
645 }
646}
647
648impl ModelResolver for EngineHandle {
649 fn resolve_default(&self) -> Result<ResolvedModel> {
650 EngineHandle::resolve_default(self)
651 }
652}
653
654impl std::fmt::Debug for EngineHandle {
655 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
656 let engine = self.inner.read();
657 f.debug_struct("EngineHandle")
658 .field("current_model", &engine.default_model_id())
659 .finish()
660 }
661}
662
663#[cfg(test)]
668mod tests {
669 use super::*;
670 #[test]
671 fn resolve_default_reflects_hot_swap() {
672 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
676 let handle = EngineHandle::new(Arc::new(engine));
677 let r1 = handle.resolve_default().expect("initial resolve");
678 assert_eq!(r1.model_id, "anthropic/claude-sonnet-4-20250514");
679 assert_eq!(r1.model.provider, "anthropic");
680
681 handle.swap(OxiosEngine::new("openai/gpt-4o"));
682 let r2 = handle.resolve_default().expect("post-swap resolve");
683 assert_eq!(r2.model_id, "openai/gpt-4o");
684 assert_eq!(r2.model.provider, "openai");
685 }
686
687 #[test]
688 fn resolve_default_fails_for_unknown_model() {
689 let engine = OxiosEngine::new("zai-coding-plan/glm-5-turbo");
693 let handle = EngineHandle::new(Arc::new(engine));
694 assert!(handle.resolve_default().is_err());
695 }
696
697 #[test]
698 fn model_resolver_impl_delegates_to_resolve_default() {
699 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
702 let handle = EngineHandle::new(Arc::new(engine));
703 let via_trait: &dyn ModelResolver = &handle;
704 let r = via_trait.resolve_default().expect("trait resolve");
705 assert_eq!(r.model_id, "anthropic/claude-sonnet-4-20250514");
706 }
707
708 #[test]
709 fn test_resolve_model_with_provider_prefix() {
710 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
711 let model = engine.resolve_model("openai/gpt-4o").unwrap();
712 assert_eq!(model.provider, "openai");
713 assert_eq!(model.id, "gpt-4o");
714 }
715
716 #[test]
717 fn test_resolve_model_without_provider_prefix() {
718 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
719 let model = engine.resolve_model("claude-sonnet-4-20250514").unwrap();
720 assert_eq!(model.provider, "anthropic");
721 }
722
723 #[test]
724 fn test_default_model_id() {
725 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
726 assert_eq!(
727 engine.default_model_id(),
728 "anthropic/claude-sonnet-4-20250514"
729 );
730 }
731
732 #[test]
733 fn test_resolve_model_not_found() {
734 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
735 let result = engine.resolve_model("nonexistent/model-xyz");
736 assert!(result.is_err());
737 }
738
739 #[test]
740 fn test_create_provider_anthropic() {
741 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
742 let provider = engine.create_provider("anthropic");
743 assert!(provider.is_ok());
744 }
745
746 #[test]
747 fn test_create_provider_not_found() {
748 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
749 let result = engine.create_provider("nonexistent_provider");
750 assert!(result.is_err());
751 }
752
753 #[test]
754 fn test_builder_with_credential() {
755 let engine = OxiosEngine::builder()
756 .default_model("openai/gpt-4o")
757 .credential("openai", "sk-test", None)
758 .build();
759 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
760 }
761
762 #[test]
763 fn test_engine_provider_trait_on_engine() {
764 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
765 let provider: &dyn EngineProvider = &engine;
766 assert!(provider.create_provider("anthropic").is_ok());
767 assert!(provider.resolve_model("openai/gpt-4o").is_ok());
768 }
769
770 #[test]
773 fn test_engine_handle_get_returns_current() {
774 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
775 let handle = EngineHandle::new(Arc::new(engine));
776 let e = handle.get();
777 assert_eq!(e.default_model_id(), "anthropic/claude-sonnet-4-20250514");
778 }
779
780 #[test]
781 fn test_engine_handle_swap_updates() {
782 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
783 let handle = EngineHandle::new(Arc::new(engine));
784
785 let new_engine = OxiosEngine::new("openai/gpt-4o");
786 handle.swap(new_engine);
787
788 let e = handle.get();
789 assert_eq!(e.default_model_id(), "openai/gpt-4o");
790 }
791
792 #[test]
793 fn test_engine_handle_swap_preserves_old_arc() {
794 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
796 let handle = EngineHandle::new(Arc::new(engine));
797
798 let old = handle.get();
799 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
800
801 handle.swap(OxiosEngine::new("openai/gpt-4o"));
802
803 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
805
806 let current = handle.get();
808 assert_eq!(current.default_model_id(), "openai/gpt-4o");
809 }
810
811 #[test]
814 fn test_rfc014_phase_d_default_fields_are_none() {
815 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
819 assert!(engine.authorizer().is_none());
820 assert!(engine.tracer().is_none());
821 assert!(engine.cost_tracker().is_none());
822
823 let engine = OxiosEngine::from_config("anthropic/claude-sonnet-4-20250514", None);
824 assert!(engine.authorizer().is_none());
825 assert!(engine.tracer().is_none());
826 assert!(engine.cost_tracker().is_none());
827
828 let engine = OxiosEngine::builder()
829 .default_model("openai/gpt-4o")
830 .build();
831 assert!(engine.authorizer().is_none());
832 assert!(engine.tracer().is_none());
833 assert!(engine.cost_tracker().is_none());
834
835 let (engine, _rc) = OxiosEngine::builder()
836 .default_model("openai/gpt-4o")
837 .build_with_routing();
838 assert!(engine.authorizer().is_none());
839 assert!(engine.tracer().is_none());
840 assert!(engine.cost_tracker().is_none());
841 }
842
843 #[test]
844 fn test_rfc014_phase_d_with_tracer() {
845 let tracer = Arc::new(oxi_sdk::Tracer::new());
847 let engine = OxiosEngine::builder()
848 .default_model("openai/gpt-4o")
849 .with_tracer(tracer.clone())
850 .build();
851 assert!(engine.tracer().is_some());
852 assert!(engine.authorizer().is_none());
853 assert!(engine.cost_tracker().is_none());
854 }
855
856 #[test]
857 fn test_rfc014_phase_d_with_cost_tracker() {
858 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
863 let model_registry = oxi_for_registry.models_arc();
864 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
865 model_registry,
866 oxi_sdk::CostTrackerConfig::default(),
867 ));
868 let engine = OxiosEngine::builder()
869 .default_model("openai/gpt-4o")
870 .with_cost_tracker(cost_tracker)
871 .build();
872 assert!(engine.cost_tracker().is_some());
873 assert!(engine.authorizer().is_none());
874 assert!(engine.tracer().is_none());
875 }
876
877 #[test]
878 fn test_rfc014_phase_d_with_authorizer() {
879 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
881 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
882 let engine = OxiosEngine::builder()
883 .default_model("openai/gpt-4o")
884 .with_authorizer(authorizer)
885 .build();
886 assert!(engine.authorizer().is_some());
887 assert!(engine.tracer().is_none());
888 assert!(engine.cost_tracker().is_none());
889 }
890
891 #[test]
892 fn test_rfc014_phase_d_all_three_handles() {
893 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
897 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
898 let tracer = Arc::new(oxi_sdk::Tracer::new());
899 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
900 let model_registry = oxi_for_registry.models_arc();
901 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
902 model_registry,
903 oxi_sdk::CostTrackerConfig::default(),
904 ));
905
906 let engine = OxiosEngine::builder()
907 .default_model("openai/gpt-4o")
908 .api_key("openai", "sk-test")
909 .with_authorizer(authorizer)
910 .with_tracer(tracer)
911 .with_cost_tracker(cost_tracker)
912 .build();
913
914 assert!(engine.authorizer().is_some());
915 assert!(engine.tracer().is_some());
916 assert!(engine.cost_tracker().is_some());
917 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
918 }
919
920 #[tokio::test]
926 #[ignore]
927 async fn catalog_integration_init_and_resolve() {
928 let catalog = OxiosEngine::init_file_catalog()
930 .await
931 .expect("catalog init should succeed (SNAP is always embedded)");
932
933 assert!(
936 catalog.model_count_sync() > 0,
937 "catalog should expose models from the embedded snapshot"
938 );
939 assert!(!catalog.list_providers_sync().is_empty());
940
941 let engine = OxiosEngine::builder()
943 .default_model("anthropic/claude-sonnet-4-20250514")
944 .with_catalog(catalog)
945 .build();
946 let model = engine
947 .resolve_model("openai/gpt-4o")
948 .expect("catalog-backed resolve_model should succeed");
949 assert_eq!(model.provider, "openai");
950 assert_eq!(model.id, "gpt-4o");
951 }
952}