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    pub chat_search: bool,
26}
27
28/// The central registry of pluggable features.
29pub struct ExtensionRegistry {
30    pub features: FeatureFlags,
31    pub encryption: Arc<dyn EncryptionProvider>,
32    pub story: Arc<dyn StoryProvider>,
33    pub pricing: Arc<dyn PricingProvider>,
34    pub compliance: Arc<dyn ComplianceProvider>,
35    pub permissions: Arc<dyn PermissionsProvider>,
36}
37
38impl Clone for ExtensionRegistry {
39    fn clone(&self) -> Self {
40        Self {
41            features: self.features.clone(),
42            encryption: Arc::clone(&self.encryption),
43            story: Arc::clone(&self.story),
44            pricing: Arc::clone(&self.pricing),
45            compliance: Arc::clone(&self.compliance),
46            permissions: Arc::clone(&self.permissions),
47        }
48    }
49}
50
51// -- Encryption --
52
53pub trait EncryptionProvider: Send + Sync {
54    fn encrypt(&self, plaintext: &str) -> Result<(String, String), String>;
55    fn decrypt(&self, ciphertext_b64: &str, nonce_b64: &str) -> Result<String, String>;
56    fn is_enabled(&self) -> bool;
57}
58
59// -- Story (LLM) --
60
61#[async_trait]
62pub trait StoryProvider: Send + Sync {
63    async fn generate_story(&self, prompt: &str, max_tokens: u32) -> Result<String, String>;
64    fn is_available(&self) -> bool;
65    fn provider_name(&self) -> &str;
66    fn model_name(&self) -> &str;
67}
68
69// -- Pricing --
70
71pub trait PricingProvider: Send + Sync {
72    fn estimate_cost(
73        &self,
74        model: &str,
75        input_tokens: i64,
76        output_tokens: i64,
77        cache_read_tokens: i64,
78        cache_write_tokens: i64,
79    ) -> f64;
80
81    fn estimate_cache_savings(&self, model: &str, cache_read_tokens: i64) -> f64;
82
83    fn cost_from_model_usage(
84        &self,
85        model_usage: Option<&serde_json::Value>,
86        fallback_model: Option<&str>,
87        fallback_input: i64,
88        fallback_output: i64,
89        fallback_cache_read: i64,
90        fallback_cache_write: i64,
91    ) -> f64;
92}
93
94// -- Compliance --
95
96pub trait ComplianceProvider: Send + Sync {
97    fn is_available(&self) -> bool;
98    fn available_modes(&self) -> Vec<String>;
99}
100
101// -- Permissions --
102
103pub trait PermissionsProvider: Send + Sync {
104    fn valid_roles(&self) -> &[&str];
105    fn role_permissions(&self, role: &str) -> std::collections::HashSet<Permission>;
106    fn has_permission(&self, role: &str, perm: Permission) -> bool;
107    fn is_valid_role(&self, role: &str) -> bool;
108}