Merge pull request #55 from Contrast-Security-OSS/AIML-124-enhance-reset-issue-to-support-claude

AIML-124: external_agent.generate_fixes() to support claude code on existing issue

Author: dougj-contrast

src/coding_agents.py

@@ -0,0 +1,26 @@
+# -
+# #%L
+# Contrast AI SmartFix
+# %%
+# Copyright (C) 2025 Contrast Security, Inc.
+# %%
+# Contact: [email protected]
+# License: Commercial
+# NOTICE: This Software and the patented inventions embodied within may only be
+# used as part of Contrast Security's commercial offerings. Even though it is
+# made available through public repositories, use of this Software is subject to
+# the applicable End User Licensing Agreement found at
+# https://www.contrastsecurity.com/enduser-terms-0317a or as otherwise agreed
+# between Contrast Security and the End User. The Software may not be reverse
+# engineered, modified, repackaged, sold, redistributed or otherwise used in a
+# way not consistent with the End User License Agreement.
+# #L%
+#
+
+from enum import Enum
+
+
+class CodingAgents(Enum):
+    SMARTFIX = "SMARTFIX"
+    GITHUB_COPILOT = "GITHUB_COPILOT"
+    CLAUDE_CODE = "CLAUDE_CODE"

src/config.py

@@ -22,6 +22,7 @@
 import json
 from pathlib import Path
 from typing import Optional, Any, Dict, List
+from src.coding_agents import CodingAgents
 
 
 def _log_config_message(message: str, is_error: bool = False, is_warning: bool = False):
@@ -65,7 +66,7 @@ def __init__(self, env: Dict[str, str] = os.environ, testing: bool = False):
 
         # --- AI Agent Configuration ---
         self.CODING_AGENT = self._get_coding_agent()
-        is_smartfix_coding_agent = self.CODING_AGENT == "SMARTFIX"
+        is_smartfix_coding_agent = self.CODING_AGENT == CodingAgents.SMARTFIX.name
 
         default_agent_model = ""
         if is_smartfix_coding_agent:
@@ -165,11 +166,13 @@ def _check_contrast_config_values_exist(self):
 
     def _get_coding_agent(self) -> str:
         coding_agent = self._get_env_var("CODING_AGENT", required=False, default="SMARTFIX")
-        valid_agents = ["SMARTFIX", "GITHUB_COPILOT", "CLAUDE_CODE"]
-        if coding_agent.upper() not in valid_agents:
-            _log_config_message(f"Warning: Invalid CODING_AGENT '{coding_agent}'. Must be one of {valid_agents}. Defaulting to 'SMARTFIX'.", is_warning=True)
-            return "SMARTFIX"
-        return coding_agent.upper()
+        try:
+            # Try to convert string to Enum
+            CodingAgents[coding_agent.upper()]
+            return coding_agent.upper()
+        except (KeyError, ValueError):
+            _log_config_message(f"Warning: Invalid CODING_AGENT '{coding_agent}'. Must be one of {[agent.name for agent in CodingAgents]}. Defaulting to '{CodingAgents.SMARTFIX.name}'.", is_warning=True)
+            return CodingAgents.SMARTFIX.name
 
     def _parse_and_validate_severities(self, json_str: Optional[str]) -> List[str]:
         default_severities = ["CRITICAL", "HIGH"]

src/external_coding_agent.py

@@ -24,8 +24,10 @@
 from src.config import Config
 from src import git_handler
 from src import telemetry_handler
+from src.coding_agents import CodingAgents
 
 
+@PendingDeprecationWarning
 class ExternalCodingAgent:
     """
     A class that interfaces with an external coding agent through an API or command line.
@@ -143,7 +145,7 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str, i
         Returns:
             bool: False if the CODING_AGENT is SMARTFIX, True otherwise
         """
-        if hasattr(self.config, 'CODING_AGENT') and self.config.CODING_AGENT == "SMARTFIX":
+        if hasattr(self.config, 'CODING_AGENT') and self.config.CODING_AGENT == CodingAgents.SMARTFIX.name:
             debug_log("SMARTFIX agent detected, ExternalCodingAgent.generate_fixes returning False")
             return False
 
@@ -155,7 +157,7 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str, i
         remediation_label = f"smartfix-id:{remediation_id}"
         issue_title = vuln_title
 
-        if self.config.CODING_AGENT == "CLAUDE_CODE":
+        if self.config.CODING_AGENT == CodingAgents.CLAUDE_CODE.name:
             debug_log("CLAUDE_CODE agent detected, tagging @claude in issue title for processing")
             issue_title = f"@claude fix: {issue_title}"
 
@@ -186,6 +188,14 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str, i
 
         telemetry_handler.update_telemetry("additionalAttributes.externalIssueNumber", issue_number)
 
+        if self.config.CODING_AGENT == CodingAgents.CLAUDE_CODE.name:
+            # temporary short-circuit for Claude until we implement the PR processing logic
+            log("Claude agent processing support is not implemented as of yet so stop processing and log agent failure", is_error=True)
+            telemetry_handler.update_telemetry("resultInfo.prCreated", False)
+            telemetry_handler.update_telemetry("resultInfo.failureReason", "Claude processing not implemented")
+            telemetry_handler.update_telemetry("resultInfo.failureCategory", FailureCategory.AGENT_FAILURE.name)
+            error_exit(remediation_id, FailureCategory.AGENT_FAILURE.value)
+
         # Poll for PR creation by the external agent
         log(f"Waiting for external agent to create a PR for issue #{issue_number}")
 

src/git_handler.py

@@ -25,6 +25,7 @@
 from src.utils import run_command, debug_log, log, error_exit
 from src.contrast_api import FailureCategory
 from src.config import get_config
+from src.coding_agents import CodingAgents
 config = get_config()
 
 
@@ -540,7 +541,7 @@ def create_issue(title: str, body: str, vuln_label: str, remediation_label: str)
             issue_number = int(os.path.basename(issue_url.strip()))
             log(f"Issue number extracted: {issue_number}")
 
-            if config.CODING_AGENT == "CLAUDE_CODE":
+            if config.CODING_AGENT == CodingAgents.CLAUDE_CODE.name:
                 debug_log("CLAUDE_CODE agent detected no need to edit issue for assignment")
                 return issue_number
 
@@ -629,7 +630,8 @@ def reset_issue(issue_number: int, remediation_label: str) -> bool:
     Resets a GitHub issue by:
     1. Removing all existing labels that start with "smartfix-id:"
     2. Adding the specified remediation label
-    3. Unassigning the @Copilot user and reassigning the issue to @Copilot
+    3. If coding agent is CoPilot then unassigning the @Copilot user and reassigning the issue to @Copilot
+    4. If coding agent is Claude Code then adding a comment to notify @claude to reprocess the issue
 
     The reset will not occur if there's an open PR for the issue.
 
@@ -709,6 +711,25 @@ def reset_issue(issue_number: int, remediation_label: str) -> bool:
         run_command(add_label_command, env=gh_env, check=True)
         log(f"Added new remediation label to issue #{issue_number}")
 
+        # If using CLAUDE_CODE, skip reassignment and tag @claude in comment
+        if config.CODING_AGENT == CodingAgents.CLAUDE_CODE.name:
+            debug_log("CLAUDE_CODE agent detected need to add a comment and tag @claude for reprocessing")
+            # Add a comment to the existing issue to notify @claude to reprocess
+            comment:str = f"@claude reprocess this issue with the new remediation label: {remediation_label} and attempt a fix."
+            comment_command = [
+                "gh", "issue", "comment",
+                "--repo", config.GITHUB_REPOSITORY,
+                str(issue_number),
+                "--create-if-none",
+                "--edit-last",
+                "--body", comment
+            ]
+
+            # add a new comment and use the @claude handle to reprocess the issue
+            run_command(comment_command, env=gh_env, check=True)
+            log(f"Added new comment tagging @claude to issue #{issue_number}")
+            return True
+
         # Unassign from @Copilot (if assigned)
         unassign_command = [
             "gh", "issue", "edit",

src/github/external_coding_agent.py

@@ -24,6 +24,7 @@
 from src.config import Config
 from src import git_handler
 from src import telemetry_handler
+from src.coding_agents import CodingAgents
 
 
 class ExternalCodingAgent:
@@ -143,7 +144,7 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str, i
         Returns:
             bool: False if the CODING_AGENT is SMARTFIX, True otherwise
         """
-        if hasattr(self.config, 'CODING_AGENT') and self.config.CODING_AGENT == "SMARTFIX":
+        if hasattr(self.config, 'CODING_AGENT') and self.config.CODING_AGENT == CodingAgents.SMARTFIX.name:
             debug_log("SMARTFIX agent detected, ExternalCodingAgent.generate_fixes returning False")
             return False
 
@@ -155,6 +156,10 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str, i
         remediation_label = f"smartfix-id:{remediation_id}"
         issue_title = vuln_title
 
+        if self.config.CODING_AGENT == CodingAgents.CLAUDE_CODE.name:
+            debug_log("CLAUDE_CODE agent detected, tagging @claude in issue title for processing")
+            issue_title = f"@claude fix: {issue_title}"
+
         # Use the provided issue_body or fall back to default
         if issue_body is None:
             log(f"Failed to generate issue body for vulnerability id {vuln_uuid}", is_error=True)
@@ -182,6 +187,14 @@ def generate_fixes(self, vuln_uuid: str, remediation_id: str, vuln_title: str, i
 
         telemetry_handler.update_telemetry("additionalAttributes.externalIssueNumber", issue_number)
 
+        if self.config.CODING_AGENT == CodingAgents.CLAUDE_CODE.name:
+            # temporary short-circuit for Claude until we implement the PR processing logic
+            log("Claude agent processing support is not implemented as of yet so stop processing and log agent failure", is_error=True)
+            telemetry_handler.update_telemetry("resultInfo.prCreated", False)
+            telemetry_handler.update_telemetry("resultInfo.failureReason", "Claude processing not implemented")
+            telemetry_handler.update_telemetry("resultInfo.failureCategory", FailureCategory.AGENT_FAILURE.name)
+            error_exit(remediation_id, FailureCategory.AGENT_FAILURE.value)
+
         # Poll for PR creation by the external agent
         log(f"Waiting for external agent to create a PR for issue #{issue_number}")
 

src/main.py

@@ -28,6 +28,7 @@
 
 # Import configurations and utilities
 from src.config import get_config
+from src.coding_agents import CodingAgents
 from src.utils import debug_log, log, error_exit
 from src import telemetry_handler
 from src.qa_handler import run_build_command
@@ -310,7 +311,7 @@ def main():  # noqa: C901
             break
 
         # --- Fetch Next Vulnerability Data from API ---
-        if config.CODING_AGENT == "SMARTFIX":
+        if config.CODING_AGENT == CodingAgents.SMARTFIX.name:
             # For SMARTFIX, get vulnerability with prompts
             log("\n::group::--- Fetching next vulnerability and prompts from Contrast API ---")
             vulnerability_data = contrast_api.get_vulnerability_with_prompts(
@@ -382,7 +383,7 @@ def main():  # noqa: C901
         log(f"\n\033[0;33m Selected vuln to fix: {vuln_title} \033[0m")
 
         # --- Check if we need to use the external coding agent ---
-        if config.CODING_AGENT != "SMARTFIX":
+        if config.CODING_AGENT != CodingAgents.SMARTFIX.name:
             external_agent = ExternalCodingAgent(config)
             # Assemble the issue body from vulnerability details
             issue_body = external_agent.assemble_issue_body(vulnerability_data)

test/test_git_handler.py

@@ -50,6 +50,7 @@
 # Import with testing=True
 from src.config import get_config, reset_config  # noqa: E402
 from src import git_handler  # noqa: E402
+from src.coding_agents import CodingAgents  # noqa: E402
 
 
 class TestGitHandler(unittest.TestCase):
@@ -206,7 +207,8 @@ def test_create_issue_failure(self, mock_log, mock_ensure_label, mock_run_comman
     @patch('src.git_handler.ensure_label')
     @patch('src.git_handler.log')
     @patch('src.git_handler.debug_log')
-    def test_reset_issue_success(self, mock_debug_log, mock_log, mock_ensure_label, mock_find_open_pr, mock_run_command, mock_check_issues):
+    @patch('src.git_handler.config')
+    def test_reset_issue_success(self, mock_config, mock_debug_log, mock_log, mock_ensure_label, mock_find_open_pr, mock_run_command, mock_check_issues):
         """Test resetting a GitHub issue when successful"""
         # Setup
         issue_number = 42
@@ -215,6 +217,10 @@ def test_reset_issue_success(self, mock_debug_log, mock_log, mock_ensure_label,
         # Mock that no open PR exists
         mock_find_open_pr.return_value = None
         mock_check_issues.return_value = True
+        
+        # Explicitly configure for SMARTFIX agent
+        mock_config.CODING_AGENT = CodingAgents.SMARTFIX.name
+        mock_config.GITHUB_REPOSITORY = 'mock/repo'
 
         # Mock successful issue view with labels
         mock_run_command.side_effect = [
@@ -298,6 +304,111 @@ def test_reset_issue_with_open_pr(self, mock_log, mock_find_open_pr):
             "Cannot reset issue #42 because it has an open PR #123: https://github.com/mock/repo/pull/123",
             is_error=True
         )
+        
+    @patch('src.git_handler.check_issues_enabled')
+    @patch('src.git_handler.run_command')
+    @patch('src.git_handler.find_open_pr_for_issue')
+    @patch('src.git_handler.ensure_label')
+    @patch('src.git_handler.log')
+    @patch('src.git_handler.debug_log')
+    @patch('src.git_handler.config')
+    def test_reset_issue_claude_code(self, mock_config, mock_debug_log, mock_log, mock_ensure_label, mock_find_open_pr, mock_run_command, mock_check_issues):
+        """Test resetting a GitHub issue when using Claude Code agent"""
+        # Setup
+        issue_number = 42
+        remediation_label = "smartfix-id:5678"
+
+        # Mock that no open PR exists
+        mock_find_open_pr.return_value = None
+        mock_check_issues.return_value = True
+        
+        # Configure the mock to use CLAUDE_CODE
+        mock_config.CODING_AGENT = CodingAgents.CLAUDE_CODE.name
+        mock_config.GITHUB_REPOSITORY = 'mock/repo'
+
+        # Mock successful issue view with labels and other API calls
+        mock_run_command.side_effect = [
+            # First call - issue view response
+            json.dumps({"labels": [{"name": "contrast-vuln-id:VULN-1234"}, {"name": "smartfix-id:OLD-REM"}]}),
+            # Second call - remove label response
+            "",
+            # Third call - add label response
+            "",
+            # Fourth call - comment with @claude tag
+            ""
+        ]
+        mock_ensure_label.return_value = True
+
+        # Execute
+        result = git_handler.reset_issue(issue_number, remediation_label)
+
+        # Assert
+        mock_check_issues.assert_called_once()
+        self.assertEqual(mock_run_command.call_count, 4)  # Should call run_command 4 times (view, remove label, add label, add comment)
+        self.assertTrue(result)
+        
+        # Check that Claude-specific logic was executed
+        mock_debug_log.assert_any_call("CLAUDE_CODE agent detected need to add a comment and tag @claude for reprocessing")
+        mock_log.assert_any_call(f"Added new comment tagging @claude to issue #{issue_number}")
+        
+        # Verify the comment command
+        comment_command_call = mock_run_command.call_args_list[3]
+        comment_command = comment_command_call[0][0]
+        
+        # Verify command structure
+        self.assertEqual(comment_command[0], "gh")
+        self.assertEqual(comment_command[1], "issue")
+        self.assertEqual(comment_command[2], "comment")
+        self.assertEqual(comment_command[5], str(issue_number))
+        
+        # Verify comment body contains '@claude' and the remediation label
+        comment_body = comment_command[-1]
+        self.assertIn("@claude", comment_body)
+        self.assertIn(remediation_label, comment_body)
+        
+    @patch('src.git_handler.check_issues_enabled')
+    @patch('src.git_handler.run_command')
+    @patch('src.git_handler.find_open_pr_for_issue')
+    @patch('src.git_handler.ensure_label')
+    @patch('src.git_handler.log')
+    @patch('src.git_handler.config')
+    def test_reset_issue_claude_code_error(self, mock_config, mock_log, mock_ensure_label, mock_find_open_pr, mock_run_command, mock_check_issues):
+        """Test resetting a GitHub issue when using Claude Code agent but an error occurs"""
+        # Setup
+        issue_number = 42
+        remediation_label = "smartfix-id:5678"
+
+        # Mock that no open PR exists
+        mock_find_open_pr.return_value = None
+        mock_check_issues.return_value = True
+        
+        # Configure the mock to use CLAUDE_CODE
+        mock_config.CODING_AGENT = CodingAgents.CLAUDE_CODE.name
+        mock_config.GITHUB_REPOSITORY = 'mock/repo'
+
+        # Mock successful label operations but comment command fails
+        mock_run_command.side_effect = [
+            # First call - issue view response
+            json.dumps({"labels": [{"name": "contrast-vuln-id:VULN-1234"}, {"name": "smartfix-id:OLD-REM"}]}),
+            # Second call - remove label response
+            "",
+            # Third call - add label response
+            "",
+            # Fourth call - comment command fails
+            Exception("Failed to comment")
+        ]
+        mock_ensure_label.return_value = True
+
+        # Execute
+        result = git_handler.reset_issue(issue_number, remediation_label)
+
+        # Assert
+        mock_check_issues.assert_called_once()
+        self.assertEqual(mock_run_command.call_count, 4)  # Should still call run_command 4 times
+        self.assertFalse(result)  # Should return False due to the error
+        
+        # Verify error was logged
+        mock_log.assert_any_call(f"Failed to reset issue #{issue_number}: Failed to comment", is_error=True)
 
     @patch('src.git_handler.run_command')
     @patch('src.git_handler.debug_log')