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    pub sso: Arc<dyn SsoProvider>,
37}
38
39impl Clone for ExtensionRegistry {
40    fn clone(&self) -> Self {
41        Self {
42            features: self.features.clone(),
43            encryption: Arc::clone(&self.encryption),
44            story: Arc::clone(&self.story),
45            pricing: Arc::clone(&self.pricing),
46            compliance: Arc::clone(&self.compliance),
47            permissions: Arc::clone(&self.permissions),
48            sso: Arc::clone(&self.sso),
49        }
50    }
51}
52
53// -- Encryption --
54
55pub trait EncryptionProvider: Send + Sync {
56    fn encrypt(&self, plaintext: &str) -> Result<(String, String), String>;
57    fn decrypt(&self, ciphertext_b64: &str, nonce_b64: &str) -> Result<String, String>;
58    fn is_enabled(&self) -> bool;
59}
60
61// -- Story (LLM) --
62
63#[async_trait]
64pub trait StoryProvider: Send + Sync {
65    async fn generate_story(&self, prompt: &str, max_tokens: u32) -> Result<String, String>;
66    fn is_available(&self) -> bool;
67    fn provider_name(&self) -> &str;
68    fn model_name(&self) -> &str;
69}
70
71// -- Pricing --
72
73pub trait PricingProvider: Send + Sync {
74    fn estimate_cost(
75        &self,
76        model: &str,
77        input_tokens: i64,
78        output_tokens: i64,
79        cache_read_tokens: i64,
80        cache_write_tokens: i64,
81    ) -> f64;
82
83    fn estimate_cache_savings(&self, model: &str, cache_read_tokens: i64) -> f64;
84
85    fn cost_from_model_usage(
86        &self,
87        model_usage: Option<&serde_json::Value>,
88        fallback_model: Option<&str>,
89        fallback_input: i64,
90        fallback_output: i64,
91        fallback_cache_read: i64,
92        fallback_cache_write: i64,
93    ) -> f64;
94}
95
96// -- Compliance --
97
98pub trait ComplianceProvider: Send + Sync {
99    fn is_available(&self) -> bool;
100    fn available_modes(&self) -> Vec<String>;
101}
102
103// -- Permissions --
104
105pub trait PermissionsProvider: Send + Sync {
106    fn valid_roles(&self) -> &[&str];
107    fn role_permissions(&self, role: &str) -> std::collections::HashSet<Permission>;
108    fn has_permission(&self, role: &str, perm: Permission) -> bool;
109    fn is_valid_role(&self, role: &str) -> bool;
110}
111
112// -- SSO --
113
114#[derive(Debug, Clone)]
115pub struct SsoUserInfo {
116    pub subject: String,
117    pub email: String,
118    pub name: Option<String>,
119}
120
121#[async_trait]
122pub trait SsoProvider: Send + Sync {
123    fn is_enabled(&self) -> bool;
124
125    /// Build the OIDC authorization URL that the user's browser should be redirected to.
126    async fn authorization_url(
127        &self,
128        issuer_url: &str,
129        client_id: &str,
130        client_secret: &str,
131        redirect_uri: &str,
132        state: &str,
133    ) -> Result<String, String>;
134
135    /// Exchange an authorization code for user identity claims.
136    async fn exchange_code(
137        &self,
138        issuer_url: &str,
139        client_id: &str,
140        client_secret: &str,
141        redirect_uri: &str,
142        code: &str,
143    ) -> Result<SsoUserInfo, String>;
144}