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