tracevault-core 0.16.1

Core library for TraceVault - AI code tracing and attribution
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

use async_trait::async_trait;
use serde::Serialize;
use std::sync::Arc;

use crate::permissions::Permission;

/// Configuration passed to enterprise extension registration.
pub struct EnterpriseConfig {
    pub encryption_key: Option<String>,
}

/// Describes which features are available in this edition.
#[derive(Debug, Clone, Serialize)]
pub struct FeatureFlags {
    pub edition: &'static str,
    pub compliance: bool,
    pub audit_trail: bool,
    pub sso: bool,
    pub story_generation: bool,
    pub advanced_analytics: bool,
    pub multi_org: bool,
    pub encryption_at_rest: bool,
    pub full_policy_engine: bool,
    pub advanced_redaction: bool,
    pub chat_search: bool,
}

/// The central registry of pluggable features.
pub struct ExtensionRegistry {
    pub features: FeatureFlags,
    pub encryption: Arc<dyn EncryptionProvider>,
    pub story: Arc<dyn StoryProvider>,
    pub pricing: Arc<dyn PricingProvider>,
    pub compliance: Arc<dyn ComplianceProvider>,
    pub permissions: Arc<dyn PermissionsProvider>,
    pub sso: Arc<dyn SsoProvider>,
}

impl Clone for ExtensionRegistry {
    fn clone(&self) -> Self {
        Self {
            features: self.features.clone(),
            encryption: Arc::clone(&self.encryption),
            story: Arc::clone(&self.story),
            pricing: Arc::clone(&self.pricing),
            compliance: Arc::clone(&self.compliance),
            permissions: Arc::clone(&self.permissions),
            sso: Arc::clone(&self.sso),
        }
    }
}

// -- Encryption --

pub trait EncryptionProvider: Send + Sync {
    fn encrypt(&self, plaintext: &str) -> Result<(String, String), String>;
    fn decrypt(&self, ciphertext_b64: &str, nonce_b64: &str) -> Result<String, String>;
    fn is_enabled(&self) -> bool;
}

// -- Story (LLM) --

#[async_trait]
pub trait StoryProvider: Send + Sync {
    async fn generate_story(&self, prompt: &str, max_tokens: u32) -> Result<String, String>;
    fn is_available(&self) -> bool;
    fn provider_name(&self) -> &str;
    fn model_name(&self) -> &str;
}

// -- Pricing --

pub trait PricingProvider: Send + Sync {
    fn estimate_cost(
        &self,
        model: &str,
        input_tokens: i64,
        output_tokens: i64,
        cache_read_tokens: i64,
        cache_write_tokens: i64,
    ) -> f64;

    fn estimate_cache_savings(&self, model: &str, cache_read_tokens: i64) -> f64;

    fn cost_from_model_usage(
        &self,
        model_usage: Option<&serde_json::Value>,
        fallback_model: Option<&str>,
        fallback_input: i64,
        fallback_output: i64,
        fallback_cache_read: i64,
        fallback_cache_write: i64,
    ) -> f64;
}

// -- Compliance --

pub trait ComplianceProvider: Send + Sync {
    fn is_available(&self) -> bool;
    fn available_modes(&self) -> Vec<String>;
}

// -- Permissions --

pub trait PermissionsProvider: Send + Sync {
    fn valid_roles(&self) -> &[&str];
    fn role_permissions(&self, role: &str) -> std::collections::HashSet<Permission>;
    fn has_permission(&self, role: &str, perm: Permission) -> bool;
    fn is_valid_role(&self, role: &str) -> bool;
}

// -- SSO --

#[derive(Debug, Clone)]
pub struct SsoUserInfo {
    pub subject: String,
    pub email: String,
    pub name: Option<String>,
}

#[async_trait]
pub trait SsoProvider: Send + Sync {
    fn is_enabled(&self) -> bool;

    /// Build the OIDC authorization URL that the user's browser should be redirected to.
    async fn authorization_url(
        &self,
        issuer_url: &str,
        client_id: &str,
        client_secret: &str,
        redirect_uri: &str,
        state: &str,
    ) -> Result<String, String>;

    /// Exchange an authorization code for user identity claims.
    async fn exchange_code(
        &self,
        issuer_url: &str,
        client_id: &str,
        client_secret: &str,
        redirect_uri: &str,
        code: &str,
    ) -> Result<SsoUserInfo, String>;
}