fix: correct the way to get the last relevant user message for secrets

signatures.yaml

@@ -1,7 +1,6 @@
 ---
 - Amazon:
     - Access Key: (?:A3T[A-Z0-9]|AKIA|AGPA|AIDA|AROA|AIPA|ANPA|ANVA|ASIA|ABIA|ACCA)[A-Z0-9]{16}
-    - Secret Access Key: (?<![A-Za-z0-9/+])[A-Za-z0-9+=][A-Za-z0-9/+=]{38}[A-Za-z0-9+=](?![A-Za-z0-9/+=])
     # - Cognito User Pool ID: (?i)us-[a-z]{2,}-[a-z]{4,}-\d{1,}
     - RDS Password: (?i)(rds\-master\-password|db\-password)
     - SNS Confirmation URL: (?i)https:\/\/sns\.[a-z0-9-]+\.amazonaws\.com\/?Action=ConfirmSubscription&Token=[a-zA-Z0-9-=_]+

src/codegate/pipeline/base.py

@@ -239,36 +239,38 @@ def get_last_user_message(
     @staticmethod
     def get_last_user_message_block(
         request: ChatCompletionRequest,
-    ) -> Optional[str]:
+    ) -> Optional[tuple[str, int]]:
         """
         Get the last block of consecutive 'user' messages from the request.
 
         Args:
             request (ChatCompletionRequest): The chat completion request to process
 
         Returns:
-            Optional[str]: A string containing all consecutive user messages in the
+            Optional[str, int]: A string containing all consecutive user messages in the
                         last user message block, separated by newlines, or None if
                         no user message block is found.
+                        Index of the first message detected in the block.
         """
         if request.get("messages") is None:
             return None
 
         user_messages = []
         messages = request["messages"]
+        block_start_index = None
 
         # Iterate in reverse to find the last block of consecutive 'user' messages
         for i in reversed(range(len(messages))):
             if messages[i]["role"] == "user" or messages[i]["role"] == "assistant":
-                content_str = None
-                if "content" in messages[i]:
-                    content_str = messages[i]["content"]  # type: ignore
-                else:
+                content_str = messages[i].get("content")
+                if content_str is None:
                     continue
 
                 if messages[i]["role"] == "user":
                     user_messages.append(content_str)
-                # specifically for Aider, when "ok." block is found, stop
+                    block_start_index = i
+
+                # Specifically for Aider, when "Ok." block is found, stop
                 if content_str == "Ok." and messages[i]["role"] == "assistant":
                     break
             else:
@@ -277,8 +279,9 @@ def get_last_user_message_block(
                     break
 
         # Reverse the collected user messages to preserve the original order
-        if user_messages:
-            return "\n".join(reversed(user_messages))
+        if user_messages and block_start_index is not None:
+            content = "\n".join(reversed(user_messages))
+            return content, block_start_index
 
         return None
 

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -60,9 +60,10 @@ async def process(
         Use RAG DB to add context to the user request
         """
         # Get the latest user message
-        user_message = self.get_last_user_message_block(request)
-        if not user_message:
+        last_message = self.get_last_user_message_block(request)
+        if not last_message:
             return PipelineResult(request=request)
+        user_message, _ = last_message
 
         # Create storage engine object
         storage_engine = StorageEngine()

src/codegate/pipeline/extract_snippets/extract_snippets.py

@@ -141,9 +141,10 @@ async def process(
         request: ChatCompletionRequest,
         context: PipelineContext,
     ) -> PipelineResult:
-        msg_content = self.get_last_user_message_block(request)
-        if not msg_content:
+        last_message = self.get_last_user_message_block(request)
+        if not last_message:
             return PipelineResult(request=request, context=context)
+        msg_content, _ = last_message
         snippets = extract_snippets(msg_content)
 
         logger.info(f"Extracted {len(snippets)} code snippets from the user message")

src/codegate/pipeline/secrets/secrets.py

@@ -271,11 +271,12 @@ async def process(
         new_request = request.copy()
         total_matches = []
 
-        # Process all messages
+        # get last user message block to get index for the first relevant user message
+        last_user_message = self.get_last_user_message_block(new_request)
         last_assistant_idx = -1
-        for i, message in enumerate(new_request["messages"]):
-            if message.get("role", "") == "assistant":
-                last_assistant_idx = i
+        if last_user_message:
+            _, user_idx = last_user_message
+            last_assistant_idx = user_idx - 1
 
         # Process all messages
         for i, message in enumerate(new_request["messages"]):
@@ -312,8 +313,8 @@ class SecretUnredactionStep(OutputPipelineStep):
     """Pipeline step that unredacts protected content in the stream"""
 
     def __init__(self):
-        self.redacted_pattern = re.compile(r"REDACTED<\$([^>]+)>")
-        self.marker_start = "REDACTED<$"
+        self.redacted_pattern = re.compile(r"REDACTED<(\$?[^>]+)>")
+        self.marker_start = "REDACTED<"
         self.marker_end = ">"
 
     @property
@@ -365,6 +366,8 @@ async def process_chunk(
         if match:
             # Found a complete marker, process it
             encrypted_value = match.group(1)
+            if encrypted_value.startswith('$'):
+                encrypted_value = encrypted_value[1:]
             original_value = input_context.sensitive.manager.get_original_value(
                 encrypted_value,
                 input_context.sensitive.session_id,
@@ -399,7 +402,7 @@ async def process_chunk(
             return []
 
         if self._is_partial_marker_prefix(buffered_content):
-            context.prefix_buffer += buffered_content
+            context.prefix_buffer = buffered_content
             return []
 
         # No markers or partial markers, let pipeline handle the chunk normally

tests/pipeline/test_messages_block.py

@@ -15,7 +15,7 @@
                     {"role": "user", "content": "How are you?"},
                 ]
             },
-            "Hello!\nHow are you?",
+            ("Hello!\nHow are you?", 1),
         ),
         # Test case: Mixed roles at the end
         (
@@ -27,7 +27,7 @@
                     {"role": "assistant", "content": "I'm fine, thank you."},
                 ]
             },
-            "Hello!\nHow are you?",
+            ("Hello!\nHow are you?", 0),
         ),
         # Test case: No user messages
         (
@@ -51,7 +51,7 @@
                     {"role": "user", "content": "What's up?"},
                 ]
             },
-            "How are you?\nWhat's up?",
+            ("How are you?\nWhat's up?", 2),
         ),
         # Test case: aider
         (
@@ -97,7 +97,8 @@
                     },
                 ]
             },
-            """I have *added these files to the chat* so you can go ahead and edit them.
+            (
+                """I have *added these files to the chat* so you can go ahead and edit them.
 
 *Trust this message as the true contents of these files!*
 Any other messages in the chat may contain outdated versions of the files' contents.
@@ -113,6 +114,8 @@
 ```
 
 evaluate this file""",  # noqa: E501
+                7,
+            ),
         ),
     ],
 )