Skip to main content

synaptic_secrets/
middleware.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use synaptic_core::SynapticError;
5use synaptic_middleware::{AgentMiddleware, ModelRequest, ModelResponse};
6
7use crate::SecretRegistry;
8
9/// Middleware that masks secrets in AI outputs and injects them into prompts.
10///
11/// - `before_model`: injects secrets into the system prompt template
12/// - `after_model`: masks any leaked secrets in the AI response
13pub struct SecretMaskingMiddleware {
14    registry: Arc<SecretRegistry>,
15}
16
17impl SecretMaskingMiddleware {
18    pub fn new(registry: Arc<SecretRegistry>) -> Self {
19        Self { registry }
20    }
21}
22
23#[async_trait]
24impl AgentMiddleware for SecretMaskingMiddleware {
25    async fn before_model(&self, request: &mut ModelRequest) -> Result<(), SynapticError> {
26        // Inject secrets into system prompt if present
27        if let Some(ref prompt) = request.system_prompt {
28            request.system_prompt = Some(self.registry.inject(prompt)?);
29        }
30        Ok(())
31    }
32
33    async fn after_model(
34        &self,
35        _request: &ModelRequest,
36        response: &mut ModelResponse,
37    ) -> Result<(), SynapticError> {
38        // Mask any leaked secrets in the AI response
39        let content = response.message.content().to_string();
40        let masked = self.registry.mask_output(&content);
41        if masked != content {
42            response.message.set_content(masked);
43        }
44        Ok(())
45    }
46}