Skip to main content

tracevault_core/
extensions.rs

1use async_trait::async_trait;
2use serde::Serialize;
3use std::sync::Arc;
4
5use crate::permissions::Permission;
6
7/// Configuration passed to enterprise extension registration.
8pub struct EnterpriseConfig {
9    pub encryption_key: Option<String>,
10}
11
12/// Describes which features are available in this edition.
13#[derive(Debug, Clone, Serialize)]
14pub struct FeatureFlags {
15    pub edition: &'static str,
16    pub compliance: bool,
17    pub audit_trail: bool,
18    pub sso: bool,
19    pub story_generation: bool,
20    pub advanced_analytics: bool,
21    pub multi_org: bool,
22    pub encryption_at_rest: bool,
23    pub full_policy_engine: bool,
24    pub advanced_redaction: bool,
25}
26
27/// The central registry of pluggable features.
28pub struct ExtensionRegistry {
29    pub features: FeatureFlags,
30    pub encryption: Arc<dyn EncryptionProvider>,
31    pub story: Arc<dyn StoryProvider>,
32    pub pricing: Arc<dyn PricingProvider>,
33    pub compliance: Arc<dyn ComplianceProvider>,
34    pub permissions: Arc<dyn PermissionsProvider>,
35}
36
37impl Clone for ExtensionRegistry {
38    fn clone(&self) -> Self {
39        Self {
40            features: self.features.clone(),
41            encryption: Arc::clone(&self.encryption),
42            story: Arc::clone(&self.story),
43            pricing: Arc::clone(&self.pricing),
44            compliance: Arc::clone(&self.compliance),
45            permissions: Arc::clone(&self.permissions),
46        }
47    }
48}
49
50// -- Encryption --
51
52pub trait EncryptionProvider: Send + Sync {
53    fn encrypt(&self, plaintext: &str) -> Result<(String, String), String>;
54    fn decrypt(&self, ciphertext_b64: &str, nonce_b64: &str) -> Result<String, String>;
55    fn is_enabled(&self) -> bool;
56}
57
58// -- Story (LLM) --
59
60#[async_trait]
61pub trait StoryProvider: Send + Sync {
62    async fn generate_story(&self, prompt: &str, max_tokens: u32) -> Result<String, String>;
63    fn is_available(&self) -> bool;
64    fn provider_name(&self) -> &str;
65    fn model_name(&self) -> &str;
66}
67
68// -- Pricing --
69
70pub trait PricingProvider: Send + Sync {
71    fn estimate_cost(
72        &self,
73        model: &str,
74        input_tokens: i64,
75        output_tokens: i64,
76        cache_read_tokens: i64,
77        cache_write_tokens: i64,
78    ) -> f64;
79
80    fn estimate_cache_savings(&self, model: &str, cache_read_tokens: i64) -> f64;
81
82    fn cost_from_model_usage(
83        &self,
84        model_usage: Option<&serde_json::Value>,
85        fallback_model: Option<&str>,
86        fallback_input: i64,
87        fallback_output: i64,
88        fallback_cache_read: i64,
89        fallback_cache_write: i64,
90    ) -> f64;
91}
92
93// -- Compliance --
94
95pub trait ComplianceProvider: Send + Sync {
96    fn is_available(&self) -> bool;
97    fn available_modes(&self) -> Vec<String>;
98}
99
100// -- Permissions --
101
102pub trait PermissionsProvider: Send + Sync {
103    fn valid_roles(&self) -> &[&str];
104    fn role_permissions(&self, role: &str) -> std::collections::HashSet<Permission>;
105    fn has_permission(&self, role: &str, perm: Permission) -> bool;
106    fn is_valid_role(&self, role: &str) -> bool;
107}