Fix SecretsManager

For some reason we coded up the SecretsManager so that it held only one
secret per session. Let's store a dict instead.

Author: jhrozek

src/codegate/pipeline/secrets/manager.py

@@ -21,7 +21,7 @@ class SecretsManager:
 
     def __init__(self):
         self.crypto = CodeGateCrypto()
-        self._session_store: dict[str, SecretEntry] = {}
+        self._session_store: dict[str, dict[str, SecretEntry]] = {}
         self._encrypted_to_session: dict[str, str] = {}  # Reverse lookup index
 
     def store_secret(self, value: str, service: str, secret_type: str, session_id: str) -> str:
@@ -41,12 +41,14 @@ def store_secret(self, value: str, service: str, secret_type: str, session_id: s
         encrypted_value = self.crypto.encrypt_token(value, session_id)
 
         # Store mappings
-        self._session_store[session_id] = SecretEntry(
+        session_secrets = self._session_store.get(session_id, {})
+        session_secrets[encrypted_value] = SecretEntry(
             original=value,
             encrypted=encrypted_value,
             service=service,
             secret_type=secret_type,
         )
+        self._session_store[session_id] = session_secrets
         self._encrypted_to_session[encrypted_value] = session_id
 
         logger.debug("Stored secret", service=service, type=secret_type, encrypted=encrypted_value)
@@ -58,7 +60,9 @@ def get_original_value(self, encrypted_value: str, session_id: str) -> Optional[
         try:
             stored_session_id = self._encrypted_to_session.get(encrypted_value)
             if stored_session_id == session_id:
-                return self._session_store[session_id].original
+                session_secrets = self._session_store[session_id].get(encrypted_value)
+                if session_secrets:
+                    return session_secrets.original
         except Exception as e:
             logger.error("Error retrieving secret", error=str(e))
         return None
@@ -71,9 +75,10 @@ def cleanup(self):
         """Securely wipe sensitive data"""
         try:
             # Convert and wipe original values
-            for entry in self._session_store.values():
-                original_bytes = bytearray(entry.original.encode())
-                self.crypto.wipe_bytearray(original_bytes)
+            for secrets in self._session_store.values():
+                for entry in secrets.values():
+                    original_bytes = bytearray(entry.original.encode())
+                    self.crypto.wipe_bytearray(original_bytes)
 
             # Clear the dictionaries
             self._session_store.clear()
@@ -92,9 +97,9 @@ def cleanup_session(self, session_id: str):
         """
         try:
             # Get the secret entry for the session
-            entry = self._session_store.get(session_id)
+            secrets = self._session_store.get(session_id, {})
 
-            if entry:
+            for entry in secrets.values():
                 # Securely wipe the original value
                 original_bytes = bytearray(entry.original.encode())
                 self.crypto.wipe_bytearray(original_bytes)

src/codegate/pipeline/secrets/secrets.py

@@ -362,6 +362,7 @@ async def process_chunk(
         if match:
             # Found a complete marker, process it
             encrypted_value = match.group(1)
+            print("----> encrypted_value: ", encrypted_value)
             original_value = input_context.sensitive.manager.get_original_value(
                 encrypted_value,
                 input_context.sensitive.session_id,
@@ -370,6 +371,8 @@ async def process_chunk(
             if original_value is None:
                 # If value not found, leave as is
                 original_value = match.group(0)  # Keep the REDACTED marker
+            else:
+                print("----> original_value: ", original_value)
 
             # Post an alert with the redacted content
             input_context.add_alert(self.name, trigger_string=encrypted_value)

src/codegate/providers/ollama/completion_handler.py

@@ -16,6 +16,7 @@ async def ollama_stream_generator(
     """OpenAI-style SSE format"""
     try:
         async for chunk in stream:
+            print(chunk)
             try:
                 yield f"{chunk.model_dump_json()}\n\n"
             except Exception as e:

src/codegate/providers/ollama/provider.py

@@ -38,6 +38,7 @@ def _setup_routes(self):
         """
         Sets up Ollama API routes.
         """
+
         @self.router.get(f"/{self.provider_route_name}/api/tags")
         async def get_tags(request: Request):
             """

tests/pipeline/secrets/test_manager.py

@@ -1,6 +1,6 @@
 import pytest
 
-from codegate.pipeline.secrets.manager import SecretEntry, SecretsManager
+from codegate.pipeline.secrets.manager import SecretsManager
 
 
 class TestSecretsManager:
@@ -21,11 +21,8 @@ def test_store_secret(self):
 
         # Verify the secret was stored
         stored = self.manager.get_by_session_id(self.test_session)
-        assert isinstance(stored, SecretEntry)
-        assert stored.original == self.test_value
-        assert stored.encrypted == encrypted
-        assert stored.service == self.test_service
-        assert stored.secret_type == self.test_type
+        assert isinstance(stored, dict)
+        assert stored[encrypted].original == self.test_value
 
         # Verify encrypted value can be retrieved
         retrieved = self.manager.get_original_value(encrypted, self.test_session)
@@ -86,10 +83,15 @@ def test_multiple_secrets_same_session(self):
         encrypted1 = self.manager.store_secret("secret1", "service1", "type1", self.test_session)
         encrypted2 = self.manager.store_secret("secret2", "service2", "type2", self.test_session)
 
-        # Latest secret should be retrievable
+        # Latest secret should be retrievable in the session
         stored = self.manager.get_by_session_id(self.test_session)
-        assert stored.original == "secret2"
-        assert stored.encrypted == encrypted2
+        assert isinstance(stored, dict)
+        assert stored[encrypted1].original == "secret1"
+        assert stored[encrypted2].original == "secret2"
+
+        # Both secrets should be retrievable directly
+        assert self.manager.get_original_value(encrypted1, self.test_session) == "secret1"
+        assert self.manager.get_original_value(encrypted2, self.test_session) == "secret2"
 
         # Both encrypted values should map to the session
         assert self.manager._encrypted_to_session[encrypted1] == self.test_session
@@ -119,15 +121,14 @@ def test_secure_cleanup(self):
 
         # Get reference to stored data before cleanup
         stored = self.manager.get_by_session_id(self.test_session)
-        original_value = stored.original
+        assert len(stored) == 1
 
         # Perform cleanup
         self.manager.cleanup()
 
         # Verify the original string was overwritten, not just removed
         # This test is a bit tricky since Python strings are immutable,
         # but we can at least verify the data is no longer accessible
-        assert original_value not in str(self.manager._session_store)
         assert self.test_value not in str(self.manager._session_store)
 
     def test_session_isolation(self):