Merge pull request #36 from Contrast-Security-OSS/AIML-51_update_merge_handler_for_copilot

AIML-51 Update merge handler for copilot

Author: JacobMagesHaskinsContrast

src/closed_handler.py

@@ -25,6 +25,7 @@
 from src import contrast_api
 from src.config import get_config  # Using get_config function instead of direct import
 from src.utils import debug_log, extract_remediation_id_from_branch, extract_remediation_id_from_labels, log
+from src.git_handler import extract_issue_number_from_branch
 import src.telemetry_handler as telemetry_handler
 
 def handle_closed_pr():
@@ -65,18 +66,27 @@ def handle_closed_pr():
 
     debug_log(f"Branch name: {branch_name}")
 
+    labels = pull_request.get("labels", [])
+
     # Extract remediation ID from branch name or PR labels
     remediation_id = None
 
     # Check if this is a branch created by external agent (e.g., GitHub Copilot)
     if branch_name.startswith("copilot/fix"):
         debug_log("Branch appears to be created by external agent. Extracting remediation ID from PR labels.")
-        # Get labels from the PR
-        labels = pull_request.get("labels", [])
         remediation_id = extract_remediation_id_from_labels(labels)
+        # Extract GitHub issue number from branch name
+        issue_number = extract_issue_number_from_branch(branch_name)
+        if issue_number:
+            telemetry_handler.update_telemetry("additionalAttributes.externalIssueNumber", issue_number)
+            debug_log(f"Extracted external issue number from branch name: {issue_number}")
+        else:
+            debug_log(f"Could not extract issue number from branch name: {branch_name}")
+        telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
     else:
         # Use original method for branches created by SmartFix
         remediation_id = extract_remediation_id_from_branch(branch_name)
+        telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "INTERNAL-SMARTFIX")
 
     if not remediation_id:
         if branch_name.startswith("copilot/fix"):
@@ -90,7 +100,6 @@ def handle_closed_pr():
     telemetry_handler.update_telemetry("additionalAttributes.remediationId", remediation_id)
 
     # Try to extract vulnerability UUID from PR labels
-    labels = pull_request.get("labels", [])
     vuln_uuid = "unknown"
 
     for label in labels:

src/external_coding_agent.py

@@ -53,7 +53,8 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str) -
             return False
 
         log(f"\n::group::--- Using External Coding Agent ({self.config.CODING_AGENT}) ---")
-        
+        telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
+
         # Hard-coded vulnerability label for now, will be passed as argument later
         vulnerability_label = f"contrast-vuln-id:VULN-{vuln_uuid}"
         remediation_label = f"smartfix-id:{remediation_id}"
@@ -75,11 +76,10 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str) -
                 log(f"Failed to create issue with labels {vulnerability_label}, {remediation_label}", is_error=True)
                 error_exit(remediation_id, FailureCategory.AGENT_FAILURE.value)
 
-        telemetry_handler.update_telemetry("additionalAttributes.githubIssueNumber", issue_number)
+        telemetry_handler.update_telemetry("additionalAttributes.externalIssueNumber", issue_number)
 
         # Poll for PR creation by the external agent
         log(f"Waiting for external agent to create a PR for issue #{issue_number}")
-        telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "EXTERNAL")
 
         # Poll for a PR to be created by the external agent (100 attempts, 5 seconds apart = ~8.3 minutes max)
         pr_info = self._poll_for_pr(issue_number, remediation_id, vulnerability_label, remediation_label, max_attempts=100, sleep_seconds=5)

src/git_handler.py

@@ -20,7 +20,8 @@
 import os
 import json
 import subprocess
-from typing import List
+import re
+from typing import List, Optional
 from src.utils import run_command, debug_log, log, error_exit
 from src.contrast_api import FailureCategory
 from src.config import get_config
@@ -648,6 +649,37 @@ def find_open_pr_for_issue(issue_number: int) -> dict:
         log(f"Error searching for PRs related to issue #{issue_number}: {e}", is_error=True)
         return None
 
+def extract_issue_number_from_branch(branch_name: str) -> Optional[int]:
+    """
+    Extracts the GitHub issue number from a branch name with format 'copilot/fix-<issue_number>'.
+    
+    Args:
+        branch_name: The branch name to extract the issue number from
+        
+    Returns:
+        Optional[int]: The issue number if found and valid, None otherwise
+    """
+    if not branch_name:
+        return None
+    
+    # Use regex to match the exact pattern: copilot/fix-<number>
+    # This ensures we only match the expected format and extract just the number
+    pattern = r'^copilot/fix-(\d+)$'
+    match = re.match(pattern, branch_name)
+    
+    if match:
+        try:
+            issue_number = int(match.group(1))
+            # Validate that it's a positive number (GitHub issue numbers start from 1)
+            if issue_number > 0:
+                return issue_number
+        except ValueError:
+            # This shouldn't happen since \d+ only matches digits, but being safe
+            debug_log(f"Failed to convert extracted issue number '{match.group(1)}' to int")
+            pass
+    
+    return None
+
 def add_labels_to_pr(pr_number: int, labels: List[str]) -> bool:
     """
     Add labels to an existing pull request.

src/main.py

@@ -351,6 +351,8 @@ def main():
                 contrast_api.send_telemetry_data()
             continue  # Skip the built-in SmartFix code and PR creation
 
+        telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "INTERNAL-SMARTFIX")
+
         # --- Run AI Fix Agent (SmartFix) ---
         ai_fix_summary_full = agent_handler.run_ai_fix_agent(
             config.REPO_ROOT, fix_system_prompt, fix_user_prompt, remediation_id

src/merge_handler.py

@@ -24,7 +24,8 @@
 # Import from src package to ensure correct module resolution
 from src import contrast_api
 from src.config import get_config  # Using get_config function instead of direct import
-from src.utils import debug_log, extract_remediation_id_from_branch, log
+from src.utils import debug_log, extract_remediation_id_from_branch, extract_remediation_id_from_labels, log
+from src.git_handler import extract_issue_number_from_branch
 import src.telemetry_handler as telemetry_handler
 
 def handle_merged_pr():
@@ -65,18 +66,27 @@ def handle_merged_pr():
 
     debug_log(f"Branch name: {branch_name}")
 
+    labels = pull_request.get("labels", [])
+
     # Extract remediation ID from branch name or PR labels
     remediation_id = None
 
     # Check if this is a branch created by external agent (e.g., GitHub Copilot)
     if branch_name.startswith("copilot/fix"):
         debug_log("Branch appears to be created by external agent. Extracting remediation ID from PR labels.")
-        # Get labels from the PR
-        labels = pull_request.get("labels", [])
         remediation_id = extract_remediation_id_from_labels(labels)
+        # Extract GitHub issue number from branch name
+        issue_number = extract_issue_number_from_branch(branch_name)
+        if issue_number:
+            telemetry_handler.update_telemetry("additionalAttributes.externalIssueNumber", issue_number)
+            debug_log(f"Extracted external issue number from branch name: {issue_number}")
+        else:
+            debug_log(f"Could not extract issue number from branch name: {branch_name}")
+        telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
     else:
         # Use original method for branches created by SmartFix
         remediation_id = extract_remediation_id_from_branch(branch_name)
+        telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "INTERNAL-SMARTFIX")
 
     if not remediation_id:
         if branch_name.startswith("copilot/fix"):
@@ -90,7 +100,6 @@ def handle_merged_pr():
     telemetry_handler.update_telemetry("additionalAttributes.remediationId", remediation_id)
 
     # Try to extract vulnerability UUID from PR labels
-    labels = pull_request.get("labels", [])
     vuln_uuid = "unknown"
 
     for label in labels:

test/test_external_coding_agent.py

@@ -135,7 +135,7 @@ def test_generate_fixes_with_external_agent_pr_created(self, mock_log, mock_debu
         mock_log.assert_any_call(f"External agent created PR #123 at https://github.com/owner/repo/pull/123")
 
         # Verify telemetry updates
-        mock_update_telemetry.assert_any_call("additionalAttributes.codingAgent", "EXTERNAL")
+        mock_update_telemetry.assert_any_call("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
         mock_update_telemetry.assert_any_call("resultInfo.prCreated", True)
         mock_update_telemetry.assert_any_call("additionalAttributes.prStatus", "OPEN")
         mock_update_telemetry.assert_any_call("additionalAttributes.prNumber", 123)
@@ -187,7 +187,7 @@ def test_generate_fixes_with_external_agent_pr_timeout(self, mock_log, mock_debu
         mock_log.assert_any_call("External agent failed to create a PR within the timeout period", is_error=True)
 
         # Verify telemetry updates
-        mock_update_telemetry.assert_any_call("additionalAttributes.codingAgent", "EXTERNAL")
+        mock_update_telemetry.assert_any_call("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
         mock_update_telemetry.assert_any_call("resultInfo.prCreated", False)
         mock_update_telemetry.assert_any_call("resultInfo.failureReason", "PR creation timeout")
         mock_update_telemetry.assert_any_call("resultInfo.failureCategory", "AGENT_FAILURE")

test/test_git_handler.py

@@ -388,6 +388,52 @@ def test_add_labels_to_pr_success(self, mock_debug_log, mock_log, mock_run_comma
         mock_log.assert_any_call("Adding labels to PR #123: ['contrast-vuln-id:VULN-12345', 'smartfix-id:remediation-67890']")
         mock_log.assert_any_call("Successfully added labels to PR #123: ['contrast-vuln-id:VULN-12345', 'smartfix-id:remediation-67890']")
 
+    def test_extract_issue_number_from_branch_success(self):
+        """Test extracting issue number from valid copilot branch name"""
+        # Test cases with valid branch names
+        test_cases = [
+            ("copilot/fix-123", 123),
+            ("copilot/fix-1", 1),
+            ("copilot/fix-999999", 999999),
+            ("copilot/fix-42", 42),
+        ]
+        
+        for branch_name, expected_issue_number in test_cases:
+            with self.subTest(branch_name=branch_name):
+                result = git_handler.extract_issue_number_from_branch(branch_name)
+                self.assertEqual(result, expected_issue_number)
+
+    def test_extract_issue_number_from_branch_invalid(self):
+        """Test extracting issue number from invalid branch names"""
+        # Test cases with invalid branch names
+        invalid_branches = [
+            "main",                           # Wrong branch name
+            "feature/new-feature",           # Wrong branch name
+            "copilot/fix-",                  # Missing issue number
+            "copilot/fix-abc",               # Non-numeric issue number
+            "copilot/fix-123abc",            # Invalid format
+            "copilot/fix-123-extra",         # Extra parts
+            "smartfix/remediation-123",      # Different prefix
+            "",                              # Empty string
+        ]
+        
+        for branch_name in invalid_branches:
+            with self.subTest(branch_name=branch_name):
+                result = git_handler.extract_issue_number_from_branch(branch_name)
+                self.assertIsNone(result)
+
+    def test_extract_issue_number_from_branch_edge_cases(self):
+        """Test edge cases for extracting issue number from branch name"""
+        # Test edge cases
+        edge_cases = [
+            ("copilot/fix-2147483647", 2147483647),  # Large number (max 32-bit int)
+        ]
+        
+        for branch_name, expected_issue_number in edge_cases:
+            with self.subTest(branch_name=branch_name):
+                result = git_handler.extract_issue_number_from_branch(branch_name)
+                self.assertEqual(result, expected_issue_number)
+
     @patch('src.git_handler.ensure_label')
     @patch('src.git_handler.run_command')
     @patch('src.git_handler.log')