feat: remove not needed encryption of secrets (#1123)

* feat: remove not needed encryption of secrets

Instead use an uuid generator as we do for pii, and reuse
same session store mechanism

Closes: #929

* fix tests

* unify interface in sensitive data

* add missing tests

* changes from rebase

* fixes from review

* fixes in tests

* fix tests

Author: yrobla

src/codegate/cli.py

@@ -16,7 +16,7 @@
 from codegate.config import Config, ConfigurationError
 from codegate.db.connection import init_db_sync, init_session_if_not_exists
 from codegate.pipeline.factory import PipelineFactory
-from codegate.pipeline.secrets.manager import SecretsManager
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 from codegate.providers import crud as provendcrud
 from codegate.providers.copilot.provider import CopilotProvider
 from codegate.server import init_app
@@ -331,8 +331,8 @@ def serve(  # noqa: C901
             click.echo("Existing Certificates are already present.")
 
         # Initialize secrets manager and pipeline factory
-        secrets_manager = SecretsManager()
-        pipeline_factory = PipelineFactory(secrets_manager)
+        sensitive_data_manager = SensitiveDataManager()
+        pipeline_factory = PipelineFactory(sensitive_data_manager)
 
         app = init_app(pipeline_factory)
 

src/codegate/pipeline/base.py

@@ -12,34 +12,23 @@
 from codegate.clients.clients import ClientType
 from codegate.db.models import Alert, AlertSeverity, Output, Prompt
 from codegate.extract_snippets.message_extractor import CodeSnippet
-from codegate.pipeline.secrets.manager import SecretsManager
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 
 logger = structlog.get_logger("codegate")
 
 
 @dataclass
 class PipelineSensitiveData:
-    manager: SecretsManager
+    manager: SensitiveDataManager
     session_id: str
-    api_key: Optional[str] = None
     model: Optional[str] = None
-    provider: Optional[str] = None
-    api_base: Optional[str] = None
 
     def secure_cleanup(self):
         """Securely cleanup sensitive data for this session"""
         if self.manager is None or self.session_id == "":
             return
-
         self.manager.cleanup_session(self.session_id)
         self.session_id = ""
-
-        # Securely wipe the API key using the same method as secrets manager
-        if self.api_key is not None:
-            api_key_bytes = bytearray(self.api_key.encode())
-            self.manager.crypto.wipe_bytearray(api_key_bytes)
-            self.api_key = None
-
         self.model = None
 
 
@@ -274,19 +263,19 @@ class InputPipelineInstance:
     def __init__(
         self,
         pipeline_steps: List[PipelineStep],
-        secret_manager: SecretsManager,
+        sensitive_data_manager: SensitiveDataManager,
         is_fim: bool,
         client: ClientType = ClientType.GENERIC,
     ):
         self.pipeline_steps = pipeline_steps
-        self.secret_manager = secret_manager
+        self.sensitive_data_manager = sensitive_data_manager
         self.is_fim = is_fim
         self.context = PipelineContext(client=client)
 
         # we create the sesitive context here so that it is not shared between individual requests
         # TODO: could we get away with just generating the session ID for an instance?
         self.context.sensitive = PipelineSensitiveData(
-            manager=self.secret_manager,
+            manager=self.sensitive_data_manager,
             session_id=str(uuid.uuid4()),
         )
         self.context.metadata["is_fim"] = is_fim
@@ -343,20 +332,20 @@ class SequentialPipelineProcessor:
     def __init__(
         self,
         pipeline_steps: List[PipelineStep],
-        secret_manager: SecretsManager,
+        sensitive_data_manager: SensitiveDataManager,
         client_type: ClientType,
         is_fim: bool,
     ):
         self.pipeline_steps = pipeline_steps
-        self.secret_manager = secret_manager
+        self.sensitive_data_manager = sensitive_data_manager
         self.is_fim = is_fim
         self.instance = self._create_instance(client_type)
 
     def _create_instance(self, client_type: ClientType) -> InputPipelineInstance:
         """Create a new pipeline instance for processing a request"""
         return InputPipelineInstance(
             self.pipeline_steps,
-            self.secret_manager,
+            self.sensitive_data_manager,
             self.is_fim,
             client_type,
         )

src/codegate/pipeline/factory.py

@@ -12,18 +12,18 @@
     PiiRedactionNotifier,
     PiiUnRedactionStep,
 )
-from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.pipeline.secrets.secrets import (
     CodegateSecrets,
     SecretRedactionNotifier,
     SecretUnredactionStep,
 )
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 from codegate.pipeline.system_prompt.codegate import SystemPrompt
 
 
 class PipelineFactory:
-    def __init__(self, secrets_manager: SecretsManager):
-        self.secrets_manager = secrets_manager
+    def __init__(self, sensitive_data_manager: SensitiveDataManager):
+        self.sensitive_data_manager = sensitive_data_manager
 
     def create_input_pipeline(self, client_type: ClientType) -> SequentialPipelineProcessor:
         input_steps: List[PipelineStep] = [
@@ -32,7 +32,7 @@ def create_input_pipeline(self, client_type: ClientType) -> SequentialPipelinePr
             # and without obfuscating the secrets, we'd leak the secrets during those
             # later steps
             CodegateSecrets(),
-            CodegatePii(),
+            CodegatePii(self.sensitive_data_manager),
             CodegateCli(),
             CodegateContextRetriever(),
             SystemPrompt(
@@ -41,19 +41,19 @@ def create_input_pipeline(self, client_type: ClientType) -> SequentialPipelinePr
         ]
         return SequentialPipelineProcessor(
             input_steps,
-            self.secrets_manager,
+            self.sensitive_data_manager,
             client_type,
             is_fim=False,
         )
 
     def create_fim_pipeline(self, client_type: ClientType) -> SequentialPipelineProcessor:
         fim_steps: List[PipelineStep] = [
             CodegateSecrets(),
-            CodegatePii(),
+            CodegatePii(self.sensitive_data_manager),
         ]
         return SequentialPipelineProcessor(
             fim_steps,
-            self.secrets_manager,
+            self.sensitive_data_manager,
             client_type,
             is_fim=True,
         )

src/codegate/pipeline/pii/analyzer.py

@@ -1,47 +1,16 @@
-import uuid
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, List, Optional
 
 import structlog
 from presidio_analyzer import AnalyzerEngine
 from presidio_anonymizer import AnonymizerEngine
 
 from codegate.db.models import AlertSeverity
 from codegate.pipeline.base import PipelineContext
+from codegate.pipeline.sensitive_data.session_store import SessionStore
 
 logger = structlog.get_logger("codegate.pii.analyzer")
 
 
-class PiiSessionStore:
-    """
-    A class to manage PII (Personally Identifiable Information) session storage.
-
-    Attributes:
-        session_id (str): The unique identifier for the session. If not provided, a new UUID
-        is generated. mappings (Dict[str, str]): A dictionary to store mappings between UUID
-        placeholders and PII.
-
-    Methods:
-        add_mapping(pii: str) -> str:
-            Adds a PII string to the session store and returns a UUID placeholder for it.
-
-        get_pii(uuid_placeholder: str) -> str:
-            Retrieves the PII string associated with the given UUID placeholder. If the placeholder
-            is not found, returns the placeholder itself.
-    """
-
-    def __init__(self, session_id: str = None):
-        self.session_id = session_id or str(uuid.uuid4())
-        self.mappings: Dict[str, str] = {}
-
-    def add_mapping(self, pii: str) -> str:
-        uuid_placeholder = f"<{str(uuid.uuid4())}>"
-        self.mappings[uuid_placeholder] = pii
-        return uuid_placeholder
-
-    def get_pii(self, uuid_placeholder: str) -> str:
-        return self.mappings.get(uuid_placeholder, uuid_placeholder)
-
-
 class PiiAnalyzer:
     """
     PiiAnalyzer class for analyzing and anonymizing text containing PII.
@@ -52,12 +21,12 @@ class PiiAnalyzer:
             Get or create the singleton instance of PiiAnalyzer.
         analyze:
             text (str): The text to analyze for PII.
-            Tuple[str, List[Dict[str, Any]], PiiSessionStore]: The anonymized text, a list of
+            Tuple[str, List[Dict[str, Any]], SessionStore]: The anonymized text, a list of
             found PII details, and the session store.
             entities (List[str]): The PII entities to analyze for.
         restore_pii:
             anonymized_text (str): The text with anonymized PII.
-            session_store (PiiSessionStore): The PiiSessionStore used for anonymization.
+            session_store (SessionStore): The SessionStore used for anonymization.
             str: The text with original PII restored.
     """
 
@@ -95,13 +64,11 @@ def __init__(self):
         # Create analyzer with custom NLP engine
         self.analyzer = AnalyzerEngine(nlp_engine=nlp_engine)
         self.anonymizer = AnonymizerEngine()
-        self.session_store = PiiSessionStore()
+        self.session_store = SessionStore()
 
         PiiAnalyzer._instance = self
 
-    def analyze(
-        self, text: str, context: Optional[PipelineContext] = None
-    ) -> Tuple[str, List[Dict[str, Any]], PiiSessionStore]:
+    def analyze(self, text: str, context: Optional[PipelineContext] = None) -> List:
         # Prioritize credit card detection first
         entities = [
             "PHONE_NUMBER",
@@ -125,81 +92,30 @@ def analyze(
             language="en",
             score_threshold=0.3,  # Lower threshold to catch more potential matches
         )
+        return analyzer_results
 
-        # Track found PII
-        found_pii = []
-
-        # Only anonymize if PII was found
-        if analyzer_results:
-            # Log each found PII instance and anonymize
-            anonymized_text = text
-            for result in analyzer_results:
-                pii_value = text[result.start : result.end]
-                uuid_placeholder = self.session_store.add_mapping(pii_value)
-                pii_info = {
-                    "type": result.entity_type,
-                    "value": pii_value,
-                    "score": result.score,
-                    "start": result.start,
-                    "end": result.end,
-                    "uuid_placeholder": uuid_placeholder,
-                }
-                found_pii.append(pii_info)
-                anonymized_text = anonymized_text.replace(pii_value, uuid_placeholder)
-
-                # Log each PII detection with its UUID mapping
-                logger.info(
-                    "PII detected and mapped",
-                    pii_type=result.entity_type,
-                    score=f"{result.score:.2f}",
-                    uuid=uuid_placeholder,
-                    # Don't log the actual PII value for security
-                    value_length=len(pii_value),
-                    session_id=self.session_store.session_id,
-                )
-
-            # Log summary of all PII found in this analysis
-            if found_pii and context:
-                # Create notification string for alert
-                notify_string = (
-                    f"**PII Detected** 🔒\n"
-                    f"- Total PII Found: {len(found_pii)}\n"
-                    f"- Types Found: {', '.join(set(p['type'] for p in found_pii))}\n"
-                )
-                context.add_alert(
-                    self._name,
-                    trigger_string=notify_string,
-                    severity_category=AlertSeverity.CRITICAL,
-                )
-
-                logger.info(
-                    "PII analysis complete",
-                    total_pii_found=len(found_pii),
-                    pii_types=[p["type"] for p in found_pii],
-                    session_id=self.session_store.session_id,
-                )
-
-            # Return the anonymized text, PII details, and session store
-            return anonymized_text, found_pii, self.session_store
-
-        # If no PII found, return original text, empty list, and session store
-        return text, [], self.session_store
-
-    def restore_pii(self, anonymized_text: str, session_store: PiiSessionStore) -> str:
+    def restore_pii(self, session_id: str, anonymized_text: str) -> str:
         """
         Restore the original PII (Personally Identifiable Information) in the given anonymized text.
 
         This method replaces placeholders in the anonymized text with their corresponding original
-        PII values using the mappings stored in the provided PiiSessionStore.
+        PII values using the mappings stored in the provided SessionStore.
 
         Args:
             anonymized_text (str): The text containing placeholders for PII.
-            session_store (PiiSessionStore): The session store containing mappings of placeholders
+            session_id (str): The session id containing mappings of placeholders
             to original PII.
 
         Returns:
             str: The text with the original PII restored.
         """
-        for uuid_placeholder, original_pii in session_store.mappings.items():
+        session_data = self.session_store.get_by_session_id(session_id)
+        if not session_data:
+            logger.warning(
+                "No active PII session found for given session ID. Unable to restore PII."
+            )
+            return anonymized_text
+
+        for uuid_placeholder, original_pii in session_data.items():
             anonymized_text = anonymized_text.replace(uuid_placeholder, original_pii)
         return anonymized_text

src/codegate/pipeline/pii/manager.py

@@ -1,18 +1,21 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
+import uuid
 
 import regex as re
 import structlog
 from litellm import ChatCompletionRequest, ChatCompletionSystemMessage, ModelResponse
 from litellm.types.utils import Delta, StreamingChoices
 
 from codegate.config import Config
+from codegate.db.models import AlertSeverity
 from codegate.pipeline.base import (
     PipelineContext,
     PipelineResult,
     PipelineStep,
 )
 from codegate.pipeline.output import OutputPipelineContext, OutputPipelineStep
-from codegate.pipeline.pii.manager import PiiManager
+from codegate.pipeline.pii.analyzer import PiiAnalyzer
+from codegate.pipeline.sensitive_data.manager import SensitiveData, SensitiveDataManager
 from codegate.pipeline.systemmsg import add_or_update_system_message
 
 logger = structlog.get_logger("codegate")
@@ -25,7 +28,7 @@ class CodegatePii(PipelineStep):
 
     Methods:
         __init__:
-            Initializes the CodegatePii pipeline step and sets up the PiiManager.
+            Initializes the CodegatePii pipeline step and sets up the SensitiveDataManager.
 
         name:
             Returns the name of the pipeline step.
@@ -37,14 +40,15 @@ class CodegatePii(PipelineStep):
             Processes the chat completion request to detect and redact PII. Updates the request with
             anonymized text and stores PII details in the context metadata.
 
-        restore_pii(anonymized_text: str) -> str:
-            Restores the original PII from the anonymized text using the PiiManager.
+        restore_pii(session_id: str, anonymized_text: str) -> str:
+            Restores the original PII from the anonymized text using the SensitiveDataManager.
     """
 
-    def __init__(self):
+    def __init__(self, sensitive_data_manager: SensitiveDataManager):
         """Initialize the CodegatePii pipeline step."""
         super().__init__()
-        self.pii_manager = PiiManager()
+        self.sensitive_data_manager = sensitive_data_manager
+        self.analyzer = PiiAnalyzer.get_instance()
 
     @property
     def name(self) -> str:
@@ -65,6 +69,68 @@ def _get_redacted_snippet(self, message: str, pii_details: List[Dict[str, Any]])
 
         return message[start:end]
 
+    def process_results(
+        self, session_id: str, text: str, results: List, context: PipelineContext
+    ) -> Tuple[List, str]:
+        # Track found PII
+        found_pii = []
+
+        # Log each found PII instance and anonymize
+        anonymized_text = text
+        for result in results:
+            pii_value = text[result.start : result.end]
+
+            # add to session store
+            obj = SensitiveData(original=pii_value, service="pii", type=result.entity_type)
+            uuid_placeholder = self.sensitive_data_manager.store(session_id, obj)
+            anonymized_text = anonymized_text.replace(pii_value, uuid_placeholder)
+
+            # Add to found PII list
+            pii_info = {
+                "type": result.entity_type,
+                "value": pii_value,
+                "score": result.score,
+                "start": result.start,
+                "end": result.end,
+                "uuid_placeholder": uuid_placeholder,
+            }
+            found_pii.append(pii_info)
+
+            # Log each PII detection with its UUID mapping
+            logger.info(
+                "PII detected and mapped",
+                pii_type=result.entity_type,
+                score=f"{result.score:.2f}",
+                uuid=uuid_placeholder,
+                # Don't log the actual PII value for security
+                value_length=len(pii_value),
+                session_id=session_id,
+            )
+
+        # Log summary of all PII found in this analysis
+        if found_pii and context:
+            # Create notification string for alert
+            notify_string = (
+                f"**PII Detected** 🔒\n"
+                f"- Total PII Found: {len(found_pii)}\n"
+                f"- Types Found: {', '.join(set(p['type'] for p in found_pii))}\n"
+            )
+            context.add_alert(
+                self.name,
+                trigger_string=notify_string,
+                severity_category=AlertSeverity.CRITICAL,
+            )
+
+            logger.info(
+                "PII analysis complete",
+                total_pii_found=len(found_pii),
+                pii_types=[p["type"] for p in found_pii],
+                session_id=session_id,
+            )
+
+        # Return the anonymized text, PII details, and session store
+        return found_pii, anonymized_text
+
     async def process(
         self, request: ChatCompletionRequest, context: PipelineContext
     ) -> PipelineResult:
@@ -75,33 +141,39 @@ async def process(
         total_pii_found = 0
         all_pii_details: List[Dict[str, Any]] = []
         last_redacted_text = ""
+        session_id = context.sensitive.session_id
 
         for i, message in enumerate(new_request["messages"]):
             if "content" in message and message["content"]:
                 # This is where analyze and anonymize the text
                 original_text = str(message["content"])
-                anonymized_text, pii_details = self.pii_manager.analyze(original_text, context)
-
-                if pii_details:
-                    total_pii_found += len(pii_details)
-                    all_pii_details.extend(pii_details)
-                    new_request["messages"][i]["content"] = anonymized_text
-
-                    # If this is a user message, grab the redacted snippet!
-                    if message.get("role") == "user":
-                        last_redacted_text = self._get_redacted_snippet(
-                            anonymized_text, pii_details
-                        )
+                results = self.analyzer.analyze(original_text, context)
+                if results:
+                    pii_details, anonymized_text = self.process_results(
+                        session_id, original_text, results, context
+                    )
+
+                    if pii_details:
+                        total_pii_found += len(pii_details)
+                        all_pii_details.extend(pii_details)
+                        new_request["messages"][i]["content"] = anonymized_text
+
+                        # If this is a user message, grab the redacted snippet!
+                        if message.get("role") == "user":
+                            last_redacted_text = self._get_redacted_snippet(
+                                anonymized_text, pii_details
+                            )
 
         logger.info(f"Total PII instances redacted: {total_pii_found}")
 
         # Store the count, details, and redacted text in context metadata
         context.metadata["redacted_pii_count"] = total_pii_found
         context.metadata["redacted_pii_details"] = all_pii_details
         context.metadata["redacted_text"] = last_redacted_text
+        context.metadata["session_id"] = session_id
 
         if total_pii_found > 0:
-            context.metadata["pii_manager"] = self.pii_manager
+            context.metadata["sensitive_data_manager"] = self.sensitive_data_manager
 
             system_message = ChatCompletionSystemMessage(
                 content=Config.get_config().prompts.pii_redacted,
@@ -113,8 +185,31 @@ async def process(
 
         return PipelineResult(request=new_request, context=context)
 
-    def restore_pii(self, anonymized_text: str) -> str:
-        return self.pii_manager.restore_pii(anonymized_text)
+    def restore_pii(self, session_id: str, anonymized_text: str) -> str:
+        """
+        Restore the original PII (Personally Identifiable Information) in the given anonymized text.
+
+        This method replaces placeholders in the anonymized text with their corresponding original
+        PII values using the mappings stored in the provided SessionStore.
+
+        Args:
+            anonymized_text (str): The text containing placeholders for PII.
+            session_id (str): The session id containing mappings of placeholders
+            to original PII.
+
+        Returns:
+            str: The text with the original PII restored.
+        """
+        session_data = self.sensitive_data_manager.get_by_session_id(session_id)
+        if not session_data:
+            logger.warning(
+                "No active PII session found for given session ID. Unable to restore PII."
+            )
+            return anonymized_text
+
+        for uuid_placeholder, original_pii in session_data.items():
+            anonymized_text = anonymized_text.replace(uuid_placeholder, original_pii)
+        return anonymized_text
 
 
 class PiiUnRedactionStep(OutputPipelineStep):
@@ -136,12 +231,12 @@ class PiiUnRedactionStep(OutputPipelineStep):
     """
 
     def __init__(self):
-        self.redacted_pattern = re.compile(r"<([0-9a-f-]{0,36})>")
+        self.redacted_pattern = re.compile(r"#([0-9a-f-]{0,36})#")
         self.complete_uuid_pattern = re.compile(
             r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
         )  # noqa: E501
-        self.marker_start = "<"
-        self.marker_end = ">"
+        self.marker_start = "#"
+        self.marker_end = "#"
 
     @property
     def name(self) -> str:
@@ -151,7 +246,7 @@ def _is_complete_uuid(self, uuid_str: str) -> bool:
         """Check if the string is a complete UUID"""
         return bool(self.complete_uuid_pattern.match(uuid_str))
 
-    async def process_chunk(
+    async def process_chunk(  # noqa: C901
         self,
         chunk: ModelResponse,
         context: OutputPipelineContext,
@@ -162,6 +257,10 @@ async def process_chunk(
             return [chunk]
 
         content = chunk.choices[0].delta.content
+        session_id = input_context.sensitive.session_id
+        if not session_id:
+            logger.error("Could not get any session id, cannot process pii")
+            return [chunk]
 
         # Add current chunk to buffer
         if context.prefix_buffer:
@@ -172,13 +271,13 @@ async def process_chunk(
         current_pos = 0
         result = []
         while current_pos < len(content):
-            start_idx = content.find("<", current_pos)
+            start_idx = content.find(self.marker_start, current_pos)
             if start_idx == -1:
                 # No more markers!, add remaining content
                 result.append(content[current_pos:])
                 break
 
-            end_idx = content.find(">", start_idx)
+            end_idx = content.find(self.marker_end, start_idx + 1)
             if end_idx == -1:
                 # Incomplete marker, buffer the rest
                 context.prefix_buffer = content[current_pos:]
@@ -190,16 +289,18 @@ async def process_chunk(
 
             # Extract potential UUID if it's a valid format!
             uuid_marker = content[start_idx : end_idx + 1]
-            uuid_value = uuid_marker[1:-1]  # Remove < >
+            uuid_value = uuid_marker[1:-1]  # Remove # #
 
             if self._is_complete_uuid(uuid_value):
                 # Get the PII manager from context metadata
                 logger.debug(f"Valid UUID found: {uuid_value}")
-                pii_manager = input_context.metadata.get("pii_manager") if input_context else None
-                if pii_manager and pii_manager.session_store:
+                sensitive_data_manager = (
+                    input_context.metadata.get("sensitive_data_manager") if input_context else None
+                )
+                if sensitive_data_manager and sensitive_data_manager.session_store:
                     # Restore original value from PII manager
                     logger.debug("Attempting to restore PII from UUID marker")
-                    original = pii_manager.session_store.get_pii(uuid_marker)
+                    original = sensitive_data_manager.get_original_value(session_id, uuid_marker)
                     logger.debug(f"Restored PII: {original}")
                     result.append(original)
                 else:

src/codegate/pipeline/pii/pii.py

@@ -16,8 +16,8 @@
     PipelineStep,
 )
 from codegate.pipeline.output import OutputPipelineContext, OutputPipelineStep
-from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.pipeline.secrets.signatures import CodegateSignatures, Match
+from codegate.pipeline.sensitive_data.manager import SensitiveData, SensitiveDataManager
 from codegate.pipeline.systemmsg import add_or_update_system_message
 
 logger = structlog.get_logger("codegate")
@@ -171,25 +171,35 @@ def obfuscate(self, text: str, snippet: Optional[CodeSnippet]) -> tuple[str, Lis
 class SecretsEncryptor(SecretsModifier):
     def __init__(
         self,
-        secrets_manager: SecretsManager,
+        sensitive_data_manager: SensitiveDataManager,
         context: PipelineContext,
         session_id: str,
     ):
-        self._secrets_manager = secrets_manager
+        self._sensitive_data_manager = sensitive_data_manager
         self._session_id = session_id
         self._context = context
         self._name = "codegate-secrets"
+
         super().__init__()
 
     def _hide_secret(self, match: Match) -> str:
         # Encrypt and store the value
-        encrypted_value = self._secrets_manager.store_secret(
-            match.value,
-            match.service,
-            match.type,
-            self._session_id,
+        if not self._session_id:
+            raise ValueError("Session id must be provided")
+
+        if not match.value:
+            raise ValueError("Value must be provided")
+        if not match.service:
+            raise ValueError("Service must be provided")
+        if not match.type:
+            raise ValueError("Secret type must be provided")
+
+        obj = SensitiveData(original=match.value, service=match.service, type=match.type)
+        uuid_placeholder = self._sensitive_data_manager.store(self._session_id, obj)
+        logger.debug(
+            "Stored secret", service=match.service, type=match.type, placeholder=uuid_placeholder
         )
-        return f"REDACTED<${encrypted_value}>"
+        return f"REDACTED<{uuid_placeholder}>"
 
     def _notify_secret(
         self, match: Match, code_snippet: Optional[CodeSnippet], protected_text: List[str]
@@ -251,7 +261,7 @@ def _redact_text(
         self,
         text: str,
         snippet: Optional[CodeSnippet],
-        secrets_manager: SecretsManager,
+        sensitive_data_manager: SensitiveDataManager,
         session_id: str,
         context: PipelineContext,
     ) -> tuple[str, List[Match]]:
@@ -260,14 +270,14 @@ def _redact_text(
 
         Args:
             text: The text to protect
-            secrets_manager: ..
+            sensitive_data_manager: ..
             session_id: ..
             context: The pipeline context to be able to log alerts
         Returns:
             Tuple containing protected text with encrypted values and the count of redacted secrets
         """
         # Find secrets in the text
-        text_encryptor = SecretsEncryptor(secrets_manager, context, session_id)
+        text_encryptor = SecretsEncryptor(sensitive_data_manager, context, session_id)
         return text_encryptor.obfuscate(text, snippet)
 
     async def process(
@@ -287,8 +297,10 @@ async def process(
         if "messages" not in request:
             return PipelineResult(request=request, context=context)
 
-        secrets_manager = context.sensitive.manager
-        if not secrets_manager or not isinstance(secrets_manager, SecretsManager):
+        sensitive_data_manager = context.sensitive.manager
+        if not sensitive_data_manager or not isinstance(
+            sensitive_data_manager, SensitiveDataManager
+        ):
             raise ValueError("Secrets manager not found in context")
         session_id = context.sensitive.session_id
         if not session_id:
@@ -305,15 +317,15 @@ async def process(
         for i, message in enumerate(new_request["messages"]):
             if "content" in message and message["content"]:
                 redacted_content, secrets_matched = self._redact_message_content(
-                    message["content"], secrets_manager, session_id, context
+                    message["content"], sensitive_data_manager, session_id, context
                 )
                 new_request["messages"][i]["content"] = redacted_content
                 if i > last_assistant_idx:
                     total_matches += secrets_matched
         new_request = self._finalize_redaction(context, total_matches, new_request)
         return PipelineResult(request=new_request, context=context)
 
-    def _redact_message_content(self, message_content, secrets_manager, session_id, context):
+    def _redact_message_content(self, message_content, sensitive_data_manager, session_id, context):
         # Extract any code snippets
         extractor = MessageCodeExtractorFactory.create_snippet_extractor(context.client)
         snippets = extractor.extract_snippets(message_content)
@@ -322,7 +334,7 @@ def _redact_message_content(self, message_content, secrets_manager, session_id,
 
         for snippet in snippets:
             redacted_snippet, secrets_matched = self._redact_text(
-                snippet, snippet, secrets_manager, session_id, context
+                snippet, snippet, sensitive_data_manager, session_id, context
             )
             redacted_snippets[snippet.code] = redacted_snippet
             total_matches.extend(secrets_matched)
@@ -336,7 +348,7 @@ def _redact_message_content(self, message_content, secrets_manager, session_id,
             if start_index > last_end:
                 non_snippet_part = message_content[last_end:start_index]
                 redacted_part, secrets_matched = self._redact_text(
-                    non_snippet_part, "", secrets_manager, session_id, context
+                    non_snippet_part, "", sensitive_data_manager, session_id, context
                 )
                 non_snippet_parts.append(redacted_part)
                 total_matches.extend(secrets_matched)
@@ -347,7 +359,7 @@ def _redact_message_content(self, message_content, secrets_manager, session_id,
         if last_end < len(message_content):
             remaining_text = message_content[last_end:]
             redacted_remaining, secrets_matched = self._redact_text(
-                remaining_text, "", secrets_manager, session_id, context
+                remaining_text, "", sensitive_data_manager, session_id, context
             )
             non_snippet_parts.append(redacted_remaining)
             total_matches.extend(secrets_matched)
@@ -428,9 +440,14 @@ async def process_chunk(
             encrypted_value = match.group(1)
             if encrypted_value.startswith("$"):
                 encrypted_value = encrypted_value[1:]
+
+            session_id = input_context.sensitive.session_id
+            if not session_id:
+                raise ValueError("Session ID not found in context")
+
             original_value = input_context.sensitive.manager.get_original_value(
+                session_id,
                 encrypted_value,
-                input_context.sensitive.session_id,
             )
 
             if original_value is None:

src/codegate/pipeline/secrets/gatecrypto.py

@@ -0,0 +1,50 @@
+import json
+from typing import Dict, Optional
+import pydantic
+import structlog
+from codegate.pipeline.sensitive_data.session_store import SessionStore
+
+logger = structlog.get_logger("codegate")
+
+
+class SensitiveData(pydantic.BaseModel):
+    """Represents sensitive data with additional metadata."""
+
+    original: str
+    service: Optional[str] = None
+    type: Optional[str] = None
+
+
+class SensitiveDataManager:
+    """Manages encryption, storage, and retrieval of secrets"""
+
+    def __init__(self):
+        self.session_store = SessionStore()
+
+    def store(self, session_id: str, value: SensitiveData) -> Optional[str]:
+        if not session_id or not value.original:
+            return None
+        return self.session_store.add_mapping(session_id, value.model_dump_json())
+
+    def get_by_session_id(self, session_id: str) -> Optional[Dict]:
+        if not session_id:
+            return None
+        data = self.session_store.get_by_session_id(session_id)
+        return SensitiveData.model_validate_json(data) if data else None
+
+    def get_original_value(self, session_id: str, uuid_placeholder: str) -> Optional[str]:
+        if not session_id:
+            return None
+        secret_entry_json = self.session_store.get_mapping(session_id, uuid_placeholder)
+        return (
+            SensitiveData.model_validate_json(secret_entry_json).original
+            if secret_entry_json
+            else None
+        )
+
+    def cleanup_session(self, session_id: str):
+        if session_id:
+            self.session_store.cleanup_session(session_id)
+
+    def cleanup(self):
+        self.session_store.cleanup()

src/codegate/pipeline/secrets/manager.py

@@ -0,0 +1,33 @@
+from typing import Dict, Optional
+import uuid
+
+
+class SessionStore:
+    """
+    A generic session store for managing data protection.
+    """
+
+    def __init__(self):
+        self.sessions: Dict[str, Dict[str, str]] = {}
+
+    def add_mapping(self, session_id: str, data: str) -> str:
+        uuid_placeholder = f"#{str(uuid.uuid4())}#"
+        if session_id not in self.sessions:
+            self.sessions[session_id] = {}
+        self.sessions[session_id][uuid_placeholder] = data
+        return uuid_placeholder
+
+    def get_by_session_id(self, session_id: str) -> Optional[Dict]:
+        return self.sessions.get(session_id, None)
+
+    def get_mapping(self, session_id: str, uuid_placeholder: str) -> Optional[str]:
+        return self.sessions.get(session_id, {}).get(uuid_placeholder)
+
+    def cleanup_session(self, session_id: str):
+        """Clears all stored mappings for a specific session."""
+        if session_id in self.sessions:
+            del self.sessions[session_id]
+
+    def cleanup(self):
+        """Clears all stored mappings for all sessions."""
+        self.sessions.clear()

src/codegate/pipeline/secrets/secrets.py

@@ -17,7 +17,7 @@
 from codegate.pipeline.base import PipelineContext
 from codegate.pipeline.factory import PipelineFactory
 from codegate.pipeline.output import OutputPipelineInstance
-from codegate.pipeline.secrets.manager import SecretsManager
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 from codegate.providers.copilot.mapping import PIPELINE_ROUTES, VALIDATED_ROUTES, PipelineType
 from codegate.providers.copilot.pipeline import (
     CopilotChatPipeline,
@@ -200,7 +200,7 @@ def __init__(self, loop: asyncio.AbstractEventLoop):
         self.ca = CertificateAuthority.get_instance()
         self.cert_manager = TLSCertDomainManager(self.ca)
         self._closing = False
-        self.pipeline_factory = PipelineFactory(SecretsManager())
+        self.pipeline_factory = PipelineFactory(SensitiveDataManager())
         self.input_pipeline: Optional[CopilotPipeline] = None
         self.fim_pipeline: Optional[CopilotPipeline] = None
         # the context as provided by the pipeline

src/codegate/pipeline/sensitive_data/manager.py

@@ -3,44 +3,7 @@
 import pytest
 from presidio_analyzer import RecognizerResult
 
-from codegate.pipeline.pii.analyzer import PiiAnalyzer, PiiSessionStore
-
-
-class TestPiiSessionStore:
-    def test_init_with_session_id(self):
-        session_id = "test-session"
-        store = PiiSessionStore(session_id)
-        assert store.session_id == session_id
-        assert store.mappings == {}
-
-    def test_init_without_session_id(self):
-        store = PiiSessionStore()
-        assert isinstance(store.session_id, str)
-        assert len(store.session_id) > 0
-        assert store.mappings == {}
-
-    def test_add_mapping(self):
-        store = PiiSessionStore()
-        pii = "test@example.com"
-        placeholder = store.add_mapping(pii)
-
-        assert placeholder.startswith("<")
-        assert placeholder.endswith(">")
-        assert store.mappings[placeholder] == pii
-
-    def test_get_pii_existing(self):
-        store = PiiSessionStore()
-        pii = "test@example.com"
-        placeholder = store.add_mapping(pii)
-
-        result = store.get_pii(placeholder)
-        assert result == pii
-
-    def test_get_pii_nonexistent(self):
-        store = PiiSessionStore()
-        placeholder = "<nonexistent>"
-        result = store.get_pii(placeholder)
-        assert result == placeholder
+from codegate.pipeline.pii.analyzer import PiiAnalyzer
 
 
 class TestPiiAnalyzer:
@@ -104,68 +67,31 @@ def test_singleton_pattern(self):
         with pytest.raises(RuntimeError, match="Use PiiAnalyzer.get_instance()"):
             PiiAnalyzer()
 
-    def test_analyze_no_pii(self, analyzer, mock_analyzer_engine):
-        text = "Hello world"
-        mock_analyzer_engine.analyze.return_value = []
-
-        result_text, found_pii, session_store = analyzer.analyze(text)
-
-        assert result_text == text
-        assert found_pii == []
-        assert isinstance(session_store, PiiSessionStore)
-
-    def test_analyze_with_pii(self, analyzer, mock_analyzer_engine):
-        text = "My email is test@example.com"
-        email_pii = RecognizerResult(
-            entity_type="EMAIL_ADDRESS",
-            start=12,
-            end=28,
-            score=1.0,  # EmailRecognizer returns a score of 1.0
-        )
-        mock_analyzer_engine.analyze.return_value = [email_pii]
-
-        result_text, found_pii, session_store = analyzer.analyze(text)
-
-        assert len(found_pii) == 1
-        pii_info = found_pii[0]
-        assert pii_info["type"] == "EMAIL_ADDRESS"
-        assert pii_info["value"] == "test@example.com"
-        assert pii_info["score"] == 1.0
-        assert pii_info["start"] == 12
-        assert pii_info["end"] == 28
-        assert "uuid_placeholder" in pii_info
-        # Verify the placeholder was used to replace the PII
-        placeholder = pii_info["uuid_placeholder"]
-        assert result_text == f"My email is {placeholder}"
-        # Verify the mapping was stored
-        assert session_store.get_pii(placeholder) == "test@example.com"
-
     def test_restore_pii(self, analyzer):
-        session_store = PiiSessionStore()
         original_text = "test@example.com"
-        placeholder = session_store.add_mapping(original_text)
-        anonymized_text = f"My email is {placeholder}"
+        session_id = "session-id"
 
-        restored_text = analyzer.restore_pii(anonymized_text, session_store)
+        placeholder = analyzer.session_store.add_mapping(session_id, original_text)
+        anonymized_text = f"My email is {placeholder}"
+        restored_text = analyzer.restore_pii(session_id, anonymized_text)
 
         assert restored_text == f"My email is {original_text}"
 
     def test_restore_pii_multiple(self, analyzer):
-        session_store = PiiSessionStore()
         email = "test@example.com"
         phone = "123-456-7890"
-        email_placeholder = session_store.add_mapping(email)
-        phone_placeholder = session_store.add_mapping(phone)
+        session_id = "session-id"
+        email_placeholder = analyzer.session_store.add_mapping(session_id, email)
+        phone_placeholder = analyzer.session_store.add_mapping(session_id, phone)
         anonymized_text = f"Email: {email_placeholder}, Phone: {phone_placeholder}"
 
-        restored_text = analyzer.restore_pii(anonymized_text, session_store)
+        restored_text = analyzer.restore_pii(session_id, anonymized_text)
 
         assert restored_text == f"Email: {email}, Phone: {phone}"
 
     def test_restore_pii_no_placeholders(self, analyzer):
-        session_store = PiiSessionStore()
         text = "No PII here"
-
-        restored_text = analyzer.restore_pii(text, session_store)
+        session_id = "session-id"
+        restored_text = analyzer.restore_pii(session_id, text)
 
         assert restored_text == text

src/codegate/pipeline/sensitive_data/session_store.py

@@ -4,9 +4,10 @@
 from litellm import ChatCompletionRequest, ModelResponse
 from litellm.types.utils import Delta, StreamingChoices
 
-from codegate.pipeline.base import PipelineContext
+from codegate.pipeline.base import PipelineContext, PipelineSensitiveData
 from codegate.pipeline.output import OutputPipelineContext
 from codegate.pipeline.pii.pii import CodegatePii, PiiRedactionNotifier, PiiUnRedactionStep
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 
 
 class TestCodegatePii:
@@ -19,8 +20,9 @@ def mock_config(self):
             yield mock_config
 
     @pytest.fixture
-    def pii_step(self, mock_config):
-        return CodegatePii()
+    def pii_step(self):
+        mock_sensitive_data_manager = MagicMock()
+        return CodegatePii(mock_sensitive_data_manager)
 
     def test_name(self, pii_step):
         assert pii_step.name == "codegate-pii"
@@ -51,57 +53,6 @@ async def test_process_no_messages(self, pii_step):
         assert result.request == request
         assert result.context == context
 
-    @pytest.mark.asyncio
-    async def test_process_with_pii(self, pii_step):
-        original_text = "My email is test@example.com"
-        request = ChatCompletionRequest(
-            model="test-model", messages=[{"role": "user", "content": original_text}]
-        )
-        context = PipelineContext()
-
-        # Mock the PII manager's analyze method
-        placeholder = "<test-uuid>"
-        pii_details = [
-            {
-                "type": "EMAIL_ADDRESS",
-                "value": "test@example.com",
-                "score": 1.0,
-                "start": 12,
-                "end": 27,
-                "uuid_placeholder": placeholder,
-            }
-        ]
-        anonymized_text = f"My email is {placeholder}"
-        pii_step.pii_manager.analyze = MagicMock(return_value=(anonymized_text, pii_details))
-
-        result = await pii_step.process(request, context)
-
-        # Verify the user message was anonymized
-        user_messages = [m for m in result.request["messages"] if m["role"] == "user"]
-        assert len(user_messages) == 1
-        assert user_messages[0]["content"] == anonymized_text
-
-        # Verify metadata was updated
-        assert result.context.metadata["redacted_pii_count"] == 1
-        assert len(result.context.metadata["redacted_pii_details"]) == 1
-        # The redacted text should be just the placeholder since that's what _get_redacted_snippet returns  # noqa: E501
-        assert result.context.metadata["redacted_text"] == placeholder
-        assert "pii_manager" in result.context.metadata
-
-        # Verify system message was added
-        system_messages = [m for m in result.request["messages"] if m["role"] == "system"]
-        assert len(system_messages) == 1
-        assert system_messages[0]["content"] == "PII has been redacted"
-
-    def test_restore_pii(self, pii_step):
-        anonymized_text = "My email is <test-uuid>"
-        original_text = "My email is test@example.com"
-        pii_step.pii_manager.restore_pii = MagicMock(return_value=original_text)
-
-        restored = pii_step.restore_pii(anonymized_text)
-
-        assert restored == original_text
-
 
 class TestPiiUnRedactionStep:
     @pytest.fixture
@@ -148,7 +99,7 @@ async def test_process_chunk_with_uuid(self, unredaction_step):
                 StreamingChoices(
                     finish_reason=None,
                     index=0,
-                    delta=Delta(content=f"Text with <{uuid}>"),
+                    delta=Delta(content=f"Text with #{uuid}#"),
                     logprobs=None,
                 )
             ],
@@ -157,17 +108,16 @@ async def test_process_chunk_with_uuid(self, unredaction_step):
             object="chat.completion.chunk",
         )
         context = OutputPipelineContext()
-        input_context = PipelineContext()
+        manager = SensitiveDataManager()
+        sensitive = PipelineSensitiveData(manager=manager, session_id="session-id")
+        input_context = PipelineContext(sensitive=sensitive)
 
         # Mock PII manager in input context
-        mock_pii_manager = MagicMock()
-        mock_session = MagicMock()
-        mock_session.get_pii = MagicMock(return_value="test@example.com")
-        mock_pii_manager.session_store = mock_session
-        input_context.metadata["pii_manager"] = mock_pii_manager
+        mock_sensitive_data_manager = MagicMock()
+        mock_sensitive_data_manager.get_original_value = MagicMock(return_value="test@example.com")
+        input_context.metadata["sensitive_data_manager"] = mock_sensitive_data_manager
 
         result = await unredaction_step.process_chunk(chunk, context, input_context)
-
         assert result[0].choices[0].delta.content == "Text with test@example.com"
 
 

src/codegate/providers/copilot/provider.py

@@ -7,13 +7,13 @@
 
 from codegate.pipeline.base import PipelineContext, PipelineSensitiveData
 from codegate.pipeline.output import OutputPipelineContext
-from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.pipeline.secrets.secrets import (
     SecretsEncryptor,
     SecretsObfuscator,
     SecretUnredactionStep,
 )
 from codegate.pipeline.secrets.signatures import CodegateSignatures, Match
+from codegate.pipeline.sensitive_data.manager import SensitiveData, SensitiveDataManager
 
 
 class TestSecretsModifier:
@@ -69,9 +69,11 @@ class TestSecretsEncryptor:
     def setup(self, temp_yaml_file):
         CodegateSignatures.initialize(temp_yaml_file)
         self.context = PipelineContext()
-        self.secrets_manager = SecretsManager()
+        self.sensitive_data_manager = SensitiveDataManager()
         self.session_id = "test_session"
-        self.encryptor = SecretsEncryptor(self.secrets_manager, self.context, self.session_id)
+        self.encryptor = SecretsEncryptor(
+            self.sensitive_data_manager, self.context, self.session_id
+        )
 
     def test_hide_secret(self):
         # Create a test match
@@ -87,12 +89,12 @@ def test_hide_secret(self):
 
         # Test secret hiding
         hidden = self.encryptor._hide_secret(match)
-        assert hidden.startswith("REDACTED<$")
+        assert hidden.startswith("REDACTED<")
         assert hidden.endswith(">")
 
         # Verify the secret was stored
-        encrypted_value = hidden[len("REDACTED<$") : -1]
-        original = self.secrets_manager.get_original_value(encrypted_value, self.session_id)
+        encrypted_value = hidden[len("REDACTED<") : -1]
+        original = self.sensitive_data_manager.get_original_value(self.session_id, encrypted_value)
         assert original == "AKIAIOSFODNN7EXAMPLE"
 
     def test_obfuscate(self):
@@ -101,7 +103,7 @@ def test_obfuscate(self):
         protected, matched_secrets = self.encryptor.obfuscate(text, None)
 
         assert len(matched_secrets) == 1
-        assert "REDACTED<$" in protected
+        assert "REDACTED<" in protected
         assert "AKIAIOSFODNN7EXAMPLE" not in protected
         assert "Other text" in protected
 
@@ -171,25 +173,24 @@ def setup_method(self):
         """Setup fresh instances for each test"""
         self.step = SecretUnredactionStep()
         self.context = OutputPipelineContext()
-        self.secrets_manager = SecretsManager()
+        self.sensitive_data_manager = SensitiveDataManager()
         self.session_id = "test_session"
 
         # Setup input context with secrets manager
         self.input_context = PipelineContext()
         self.input_context.sensitive = PipelineSensitiveData(
-            manager=self.secrets_manager, session_id=self.session_id
+            manager=self.sensitive_data_manager, session_id=self.session_id
         )
 
     @pytest.mark.asyncio
     async def test_complete_marker_processing(self):
         """Test processing of a complete REDACTED marker"""
         # Store a secret
-        encrypted = self.secrets_manager.store_secret(
-            "secret_value", "test_service", "api_key", self.session_id
-        )
+        obj = SensitiveData(original="secret_value", service="test_service", type="api_key")
+        encrypted = self.sensitive_data_manager.store(self.session_id, obj)
 
         # Add content with REDACTED marker to buffer
-        self.context.buffer.append(f"Here is the REDACTED<${encrypted}> in text")
+        self.context.buffer.append(f"Here is the REDACTED<{encrypted}> in text")
 
         # Process a chunk
         result = await self.step.process_chunk(
@@ -204,7 +205,7 @@ async def test_complete_marker_processing(self):
     async def test_partial_marker_buffering(self):
         """Test handling of partial REDACTED markers"""
         # Add partial marker to buffer
-        self.context.buffer.append("Here is REDACTED<$")
+        self.context.buffer.append("Here is REDACTED<")
 
         # Process a chunk
         result = await self.step.process_chunk(
@@ -218,7 +219,7 @@ async def test_partial_marker_buffering(self):
     async def test_invalid_encrypted_value(self):
         """Test handling of invalid encrypted values"""
         # Add content with invalid encrypted value
-        self.context.buffer.append("Here is REDACTED<$invalid_value> in text")
+        self.context.buffer.append("Here is REDACTED<invalid_value> in text")
 
         # Process chunk
         result = await self.step.process_chunk(
@@ -227,7 +228,7 @@ async def test_invalid_encrypted_value(self):
 
         # Should keep the REDACTED marker for invalid values
         assert len(result) == 1
-        assert result[0].choices[0].delta.content == "Here is REDACTED<$invalid_value> in text"
+        assert result[0].choices[0].delta.content == "Here is REDACTED<invalid_value> in text"
 
     @pytest.mark.asyncio
     async def test_missing_context(self):
@@ -271,17 +272,16 @@ async def test_no_markers(self):
     async def test_wrong_session(self):
         """Test unredaction with wrong session ID"""
         # Store secret with one session
-        encrypted = self.secrets_manager.store_secret(
-            "secret_value", "test_service", "api_key", "different_session"
-        )
+        obj = SensitiveData(original="test_service", service="api_key", type="different_session")
+        encrypted = self.sensitive_data_manager.store("different_session", obj)
 
         # Try to unredact with different session
-        self.context.buffer.append(f"Here is the REDACTED<${encrypted}> in text")
+        self.context.buffer.append(f"Here is the REDACTED<{encrypted}> in text")
 
         result = await self.step.process_chunk(
             create_model_response("text"), self.context, self.input_context
         )
 
         # Should keep REDACTED marker when session doesn't match
         assert len(result) == 1
-        assert result[0].choices[0].delta.content == f"Here is the REDACTED<${encrypted}> in text"
+        assert result[0].choices[0].delta.content == f"Here is the REDACTED<{encrypted}> in text"

tests/pipeline/pii/test_analyzer.py

@@ -0,0 +1,48 @@
+import json
+from unittest.mock import MagicMock, patch
+import pytest
+from codegate.pipeline.sensitive_data.manager import SensitiveData, SensitiveDataManager
+from codegate.pipeline.sensitive_data.session_store import SessionStore
+
+
+class TestSensitiveDataManager:
+    @pytest.fixture
+    def mock_session_store(self):
+        """Mock the SessionStore instance used within SensitiveDataManager."""
+        return MagicMock(spec=SessionStore)
+
+    @pytest.fixture
+    def manager(self, mock_session_store):
+        """Patch SensitiveDataManager to use the mocked SessionStore."""
+        with patch.object(SensitiveDataManager, "__init__", lambda self: None):
+            manager = SensitiveDataManager()
+            manager.session_store = mock_session_store  # Manually inject the mock
+            return manager
+
+    def test_store_success(self, manager, mock_session_store):
+        """Test storing a SensitiveData object successfully."""
+        session_id = "session-123"
+        sensitive_data = SensitiveData(original="secret_value", service="AWS", type="API_KEY")
+
+        # Mock session store behavior
+        mock_session_store.add_mapping.return_value = "uuid-123"
+
+        result = manager.store(session_id, sensitive_data)
+
+        # Verify correct function calls
+        mock_session_store.add_mapping.assert_called_once_with(
+            session_id, sensitive_data.model_dump_json()
+        )
+        assert result == "uuid-123"
+
+    def test_store_invalid_session_id(self, manager):
+        """Test storing data with an invalid session ID (should return None)."""
+        sensitive_data = SensitiveData(original="secret_value", service="AWS", type="API_KEY")
+        result = manager.store("", sensitive_data)  # Empty session ID
+        assert result is None
+
+    def test_store_missing_original_value(self, manager):
+        """Test storing data without an original value (should return None)."""
+        sensitive_data = SensitiveData(original="", service="AWS", type="API_KEY")  # Empty original
+        result = manager.store("session-123", sensitive_data)
+        assert result is None

tests/pipeline/pii/test_pi.py

@@ -0,0 +1,114 @@
+import uuid
+import pytest
+from codegate.pipeline.sensitive_data.session_store import SessionStore
+
+
+class TestSessionStore:
+    @pytest.fixture
+    def session_store(self):
+        """Fixture to create a fresh SessionStore instance before each test."""
+        return SessionStore()
+
+    def test_add_mapping_creates_uuid(self, session_store):
+        """Test that add_mapping correctly stores data and returns a UUID."""
+        session_id = "session-123"
+        data = "test-data"
+
+        uuid_placeholder = session_store.add_mapping(session_id, data)
+
+        # Ensure the returned placeholder follows the expected format
+        assert uuid_placeholder.startswith("#") and uuid_placeholder.endswith("#")
+        assert len(uuid_placeholder) > 2  # Should have a UUID inside
+
+        # Verify data is correctly stored
+        stored_data = session_store.get_mapping(session_id, uuid_placeholder)
+        assert stored_data == data
+
+    def test_add_mapping_creates_unique_uuids(self, session_store):
+        """Ensure multiple calls to add_mapping generate unique UUIDs."""
+        session_id = "session-123"
+        data1 = "data1"
+        data2 = "data2"
+
+        uuid_placeholder1 = session_store.add_mapping(session_id, data1)
+        uuid_placeholder2 = session_store.add_mapping(session_id, data2)
+
+        assert uuid_placeholder1 != uuid_placeholder2  # UUIDs must be unique
+
+        # Ensure data is correctly stored
+        assert session_store.get_mapping(session_id, uuid_placeholder1) == data1
+        assert session_store.get_mapping(session_id, uuid_placeholder2) == data2
+
+    def test_get_by_session_id(self, session_store):
+        """Test retrieving all stored mappings for a session ID."""
+        session_id = "session-123"
+        data1 = "data1"
+        data2 = "data2"
+
+        uuid1 = session_store.add_mapping(session_id, data1)
+        uuid2 = session_store.add_mapping(session_id, data2)
+
+        stored_session_data = session_store.get_by_session_id(session_id)
+
+        assert uuid1 in stored_session_data
+        assert uuid2 in stored_session_data
+        assert stored_session_data[uuid1] == data1
+        assert stored_session_data[uuid2] == data2
+
+    def test_get_by_session_id_not_found(self, session_store):
+        """Test get_by_session_id when session does not exist (should return None)."""
+        session_id = "non-existent-session"
+        assert session_store.get_by_session_id(session_id) is None
+
+    def test_get_mapping_success(self, session_store):
+        """Test retrieving a specific mapping."""
+        session_id = "session-123"
+        data = "test-data"
+
+        uuid_placeholder = session_store.add_mapping(session_id, data)
+
+        assert session_store.get_mapping(session_id, uuid_placeholder) == data
+
+    def test_get_mapping_not_found(self, session_store):
+        """Test retrieving a mapping that does not exist (should return None)."""
+        session_id = "session-123"
+        uuid_placeholder = "#non-existent-uuid#"
+
+        assert session_store.get_mapping(session_id, uuid_placeholder) is None
+
+    def test_cleanup_session(self, session_store):
+        """Test that cleanup_session removes all data for a session ID."""
+        session_id = "session-123"
+        session_store.add_mapping(session_id, "test-data")
+
+        # Ensure session exists before cleanup
+        assert session_store.get_by_session_id(session_id) is not None
+
+        session_store.cleanup_session(session_id)
+
+        # Ensure session is removed after cleanup
+        assert session_store.get_by_session_id(session_id) is None
+
+    def test_cleanup_session_non_existent(self, session_store):
+        """Test cleanup_session on a non-existent session (should not raise errors)."""
+        session_id = "non-existent-session"
+        session_store.cleanup_session(session_id)  # Should not fail
+        assert session_store.get_by_session_id(session_id) is None
+
+    def test_cleanup(self, session_store):
+        """Test global cleanup removes all stored sessions."""
+        session_id1 = "session-1"
+        session_id2 = "session-2"
+
+        session_store.add_mapping(session_id1, "data1")
+        session_store.add_mapping(session_id2, "data2")
+
+        # Ensure sessions exist before cleanup
+        assert session_store.get_by_session_id(session_id1) is not None
+        assert session_store.get_by_session_id(session_id2) is not None
+
+        session_store.cleanup()
+
+        # Ensure all sessions are removed after cleanup
+        assert session_store.get_by_session_id(session_id1) is None
+        assert session_store.get_by_session_id(session_id2) is None

tests/pipeline/pii/test_pii_manager.py

@@ -14,19 +14,13 @@
 
 from codegate import __version__
 from codegate.pipeline.factory import PipelineFactory
-from codegate.pipeline.secrets.manager import SecretsManager
+from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
 from codegate.providers.registry import ProviderRegistry
 from codegate.server import init_app
 from src.codegate.cli import UvicornServer, cli
 from src.codegate.codegate_logging import LogFormat, LogLevel
 
 
-@pytest.fixture
-def mock_secrets_manager():
-    """Create a mock secrets manager."""
-    return MagicMock(spec=SecretsManager)
-
-
 @pytest.fixture
 def mock_provider_registry():
     """Create a mock provider registry."""
@@ -96,9 +90,9 @@ def test_version_endpoint(mock_fetch_latest_version, test_client: TestClient) ->
     assert response_data["is_latest"] is False
 
 
-@patch("codegate.pipeline.secrets.manager.SecretsManager")
+@patch("codegate.pipeline.sensitive_data.manager.SensitiveDataManager")
 @patch("codegate.server.get_provider_registry")
-def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_factory) -> None:
+def test_provider_registration(mock_registry, mock_pipeline_factory) -> None:
     """Test that all providers are registered correctly."""
     init_app(mock_pipeline_factory)