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 pub fn resolve(&self, id: &str) -> Result<ResolvedModel> {
641 let engine = self.get();
642 let model = engine.resolve_model(id)?;
643 let provider = self.cached_provider(&model.provider)?;
644 Ok(ResolvedModel {
645 model,
646 provider,
647 model_id: id.to_string(),
648 })
649 }
650
651 fn cached_provider(&self, name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
653 if let Some(p) = self.provider_cache.read().get(name) {
654 return Ok(Arc::clone(p));
655 }
656 let provider = self.get().create_provider(name)?;
657 self.provider_cache
658 .write()
659 .insert(name.to_string(), Arc::clone(&provider));
660 Ok(provider)
661 }
662}
663
664impl ModelResolver for EngineHandle {
665 fn resolve_default(&self) -> Result<ResolvedModel> {
666 EngineHandle::resolve_default(self)
667 }
668
669 fn resolve(&self, id: &str) -> Result<ResolvedModel> {
670 EngineHandle::resolve(self, id)
671 }
672}
673
674impl std::fmt::Debug for EngineHandle {
675 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
676 let engine = self.inner.read();
677 f.debug_struct("EngineHandle")
678 .field("current_model", &engine.default_model_id())
679 .finish()
680 }
681}
682
683#[cfg(test)]
688mod tests {
689 use super::*;
690 #[test]
691 fn resolve_default_reflects_hot_swap() {
692 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
696 let handle = EngineHandle::new(Arc::new(engine));
697 let r1 = handle.resolve_default().expect("initial resolve");
698 assert_eq!(r1.model_id, "anthropic/claude-sonnet-4-20250514");
699 assert_eq!(r1.model.provider, "anthropic");
700
701 handle.swap(OxiosEngine::new("openai/gpt-4o"));
702 let r2 = handle.resolve_default().expect("post-swap resolve");
703 assert_eq!(r2.model_id, "openai/gpt-4o");
704 assert_eq!(r2.model.provider, "openai");
705 }
706
707 #[test]
708 fn resolve_default_fails_for_unknown_model() {
709 let engine = OxiosEngine::new("zai-coding-plan/glm-5-turbo");
713 let handle = EngineHandle::new(Arc::new(engine));
714 assert!(handle.resolve_default().is_err());
715 }
716
717 #[test]
718 fn model_resolver_impl_delegates_to_resolve_default() {
719 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
722 let handle = EngineHandle::new(Arc::new(engine));
723 let via_trait: &dyn ModelResolver = &handle;
724 let r = via_trait.resolve_default().expect("trait resolve");
725 assert_eq!(r.model_id, "anthropic/claude-sonnet-4-20250514");
726 }
727
728 #[test]
729 fn test_resolve_model_with_provider_prefix() {
730 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
731 let model = engine.resolve_model("openai/gpt-4o").unwrap();
732 assert_eq!(model.provider, "openai");
733 assert_eq!(model.id, "gpt-4o");
734 }
735
736 #[test]
737 fn test_resolve_model_without_provider_prefix() {
738 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
739 let model = engine.resolve_model("claude-sonnet-4-20250514").unwrap();
740 assert_eq!(model.provider, "anthropic");
741 }
742
743 #[test]
744 fn test_default_model_id() {
745 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
746 assert_eq!(
747 engine.default_model_id(),
748 "anthropic/claude-sonnet-4-20250514"
749 );
750 }
751
752 #[test]
753 fn test_resolve_model_not_found() {
754 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
755 let result = engine.resolve_model("nonexistent/model-xyz");
756 assert!(result.is_err());
757 }
758
759 #[test]
760 fn test_create_provider_anthropic() {
761 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
762 let provider = engine.create_provider("anthropic");
763 assert!(provider.is_ok());
764 }
765
766 #[test]
767 fn test_create_provider_not_found() {
768 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
769 let result = engine.create_provider("nonexistent_provider");
770 assert!(result.is_err());
771 }
772
773 #[test]
774 fn test_builder_with_credential() {
775 let engine = OxiosEngine::builder()
776 .default_model("openai/gpt-4o")
777 .credential("openai", "sk-test", None)
778 .build();
779 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
780 }
781
782 #[test]
783 fn test_engine_provider_trait_on_engine() {
784 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
785 let provider: &dyn EngineProvider = &engine;
786 assert!(provider.create_provider("anthropic").is_ok());
787 assert!(provider.resolve_model("openai/gpt-4o").is_ok());
788 }
789
790 #[test]
793 fn test_engine_handle_get_returns_current() {
794 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
795 let handle = EngineHandle::new(Arc::new(engine));
796 let e = handle.get();
797 assert_eq!(e.default_model_id(), "anthropic/claude-sonnet-4-20250514");
798 }
799
800 #[test]
801 fn test_engine_handle_swap_updates() {
802 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
803 let handle = EngineHandle::new(Arc::new(engine));
804
805 let new_engine = OxiosEngine::new("openai/gpt-4o");
806 handle.swap(new_engine);
807
808 let e = handle.get();
809 assert_eq!(e.default_model_id(), "openai/gpt-4o");
810 }
811
812 #[test]
813 fn test_engine_handle_swap_preserves_old_arc() {
814 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
816 let handle = EngineHandle::new(Arc::new(engine));
817
818 let old = handle.get();
819 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
820
821 handle.swap(OxiosEngine::new("openai/gpt-4o"));
822
823 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
825
826 let current = handle.get();
828 assert_eq!(current.default_model_id(), "openai/gpt-4o");
829 }
830
831 #[test]
834 fn test_rfc014_phase_d_default_fields_are_none() {
835 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
839 assert!(engine.authorizer().is_none());
840 assert!(engine.tracer().is_none());
841 assert!(engine.cost_tracker().is_none());
842
843 let engine = OxiosEngine::from_config("anthropic/claude-sonnet-4-20250514", None);
844 assert!(engine.authorizer().is_none());
845 assert!(engine.tracer().is_none());
846 assert!(engine.cost_tracker().is_none());
847
848 let engine = OxiosEngine::builder()
849 .default_model("openai/gpt-4o")
850 .build();
851 assert!(engine.authorizer().is_none());
852 assert!(engine.tracer().is_none());
853 assert!(engine.cost_tracker().is_none());
854
855 let (engine, _rc) = OxiosEngine::builder()
856 .default_model("openai/gpt-4o")
857 .build_with_routing();
858 assert!(engine.authorizer().is_none());
859 assert!(engine.tracer().is_none());
860 assert!(engine.cost_tracker().is_none());
861 }
862
863 #[test]
864 fn test_rfc014_phase_d_with_tracer() {
865 let tracer = Arc::new(oxi_sdk::Tracer::new());
867 let engine = OxiosEngine::builder()
868 .default_model("openai/gpt-4o")
869 .with_tracer(tracer.clone())
870 .build();
871 assert!(engine.tracer().is_some());
872 assert!(engine.authorizer().is_none());
873 assert!(engine.cost_tracker().is_none());
874 }
875
876 #[test]
877 fn test_rfc014_phase_d_with_cost_tracker() {
878 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
883 let model_registry = oxi_for_registry.models_arc();
884 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
885 model_registry,
886 oxi_sdk::CostTrackerConfig::default(),
887 ));
888 let engine = OxiosEngine::builder()
889 .default_model("openai/gpt-4o")
890 .with_cost_tracker(cost_tracker)
891 .build();
892 assert!(engine.cost_tracker().is_some());
893 assert!(engine.authorizer().is_none());
894 assert!(engine.tracer().is_none());
895 }
896
897 #[test]
898 fn test_rfc014_phase_d_with_authorizer() {
899 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
901 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
902 let engine = OxiosEngine::builder()
903 .default_model("openai/gpt-4o")
904 .with_authorizer(authorizer)
905 .build();
906 assert!(engine.authorizer().is_some());
907 assert!(engine.tracer().is_none());
908 assert!(engine.cost_tracker().is_none());
909 }
910
911 #[test]
912 fn test_rfc014_phase_d_all_three_handles() {
913 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
917 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
918 let tracer = Arc::new(oxi_sdk::Tracer::new());
919 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
920 let model_registry = oxi_for_registry.models_arc();
921 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
922 model_registry,
923 oxi_sdk::CostTrackerConfig::default(),
924 ));
925
926 let engine = OxiosEngine::builder()
927 .default_model("openai/gpt-4o")
928 .api_key("openai", "sk-test")
929 .with_authorizer(authorizer)
930 .with_tracer(tracer)
931 .with_cost_tracker(cost_tracker)
932 .build();
933
934 assert!(engine.authorizer().is_some());
935 assert!(engine.tracer().is_some());
936 assert!(engine.cost_tracker().is_some());
937 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
938 }
939
940 #[tokio::test]
946 #[ignore]
947 async fn catalog_integration_init_and_resolve() {
948 let catalog = OxiosEngine::init_file_catalog()
950 .await
951 .expect("catalog init should succeed (SNAP is always embedded)");
952
953 assert!(
956 catalog.model_count_sync() > 0,
957 "catalog should expose models from the embedded snapshot"
958 );
959 assert!(!catalog.list_providers_sync().is_empty());
960
961 let engine = OxiosEngine::builder()
963 .default_model("anthropic/claude-sonnet-4-20250514")
964 .with_catalog(catalog)
965 .build();
966 let model = engine
967 .resolve_model("openai/gpt-4o")
968 .expect("catalog-backed resolve_model should succeed");
969 assert_eq!(model.provider, "openai");
970 assert_eq!(model.id, "gpt-4o");
971 }
972}