Skip to content

Commit d6b042d

Browse files
dougj-contrastdougj-contrast
authored
Merge pull request #57 from Contrast-Security-OSS/AIML-126-update-closed-and-merged-handlers-to-support-claude
AIML-126: update closed and merged handlers to support claude branch naming convention
2 parents 3e5c91b + db4b6f8 commit d6b042d

File tree

6 files changed

+413
-43
lines changed

6 files changed

+413
-43
lines changed

‎src/closed_handler.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,8 @@ def _extract_remediation_info(pull_request: dict) -> tuple:
7272
# Extract remediation ID from branch name or PR labels
7373
remediation_id = None
7474

75-
# Check if this is a branch created by external agent (e.g., GitHub Copilot)
76-
if branch_name.startswith("copilot/fix"):
75+
# Check if this is a branch created by external agent (e.g., GitHub Copilot or Claude Code)
76+
if branch_name.startswith("copilot/fix") or branch_name.startswith("claude/issue-"):
7777
debug_log("Branch appears to be created by external agent. Extracting remediation ID from PR labels.")
7878
remediation_id = extract_remediation_id_from_labels(labels)
7979
# Extract GitHub issue number from branch name
@@ -83,14 +83,18 @@ def _extract_remediation_info(pull_request: dict) -> tuple:
8383
debug_log(f"Extracted external issue number from branch name: {issue_number}")
8484
else:
8585
debug_log(f"Could not extract issue number from branch name: {branch_name}")
86-
telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
86+
87+
# Set the external coding agent in telemetry based on branch prefix
88+
coding_agent = "EXTERNAL-CLAUDE-CODE" if branch_name.startswith("claude/") else "EXTERNAL-COPILOT"
89+
debug_log(f"Determined external coding agent to be: {coding_agent}")
90+
telemetry_handler.update_telemetry("additionalAttributes.codingAgent", coding_agent)
8791
else:
8892
# Use original method for branches created by SmartFix
8993
remediation_id = extract_remediation_id_from_branch(branch_name)
9094
telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "INTERNAL-SMARTFIX")
9195

9296
if not remediation_id:
93-
if branch_name.startswith("copilot/fix"):
97+
if branch_name.startswith("copilot/fix") or branch_name.startswith("claude/issue-"):
9498
log(f"Error: Could not extract remediation ID from PR labels for external agent branch: {branch_name}", is_error=True)
9599
else:
96100
log(f"Error: Could not extract remediation ID from branch name: {branch_name}", is_error=True)

‎src/git_handler.py‎

Lines changed: 31 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -764,7 +764,8 @@ def reset_issue(issue_number: int, remediation_label: str) -> bool:
764764
def find_open_pr_for_issue(issue_number: int) -> dict:
765765
"""
766766
Finds an open pull request associated with the given issue number.
767-
Specifically looks for PRs with branch names matching the pattern 'copilot/fix-{issue_number}'.
767+
Specifically looks for PRs with branch names matching the pattern 'copilot/fix-{issue_number}'
768+
or 'claude/issue-{issue_number}-'.
768769
769770
Args:
770771
issue_number: The issue number to find a PR for
@@ -775,8 +776,8 @@ def find_open_pr_for_issue(issue_number: int) -> dict:
775776
debug_log(f"Searching for open PR related to issue #{issue_number}")
776777
gh_env = get_gh_env()
777778

778-
# Use a search pattern that matches PRs with branch names following the 'copilot/fix-{issue_number}' pattern
779-
# IDEA: Whenever we start supporting other coding agents the proper branch prefix will need to be passed in
779+
# Use search patterns that match PRs with branch names for both Copilot and Claude Code
780+
# First try to find PRs with Copilot branch pattern
780781
search_pattern = f"head:copilot/fix-{issue_number}"
781782

782783
pr_list_command = [
@@ -792,8 +793,22 @@ def find_open_pr_for_issue(issue_number: int) -> dict:
792793
pr_list_output = run_command(pr_list_command, env=gh_env, check=False)
793794

794795
if not pr_list_output or pr_list_output.strip() == "[]":
795-
debug_log(f"No open PRs found for issue #{issue_number}")
796-
return None
796+
# Try again with claude branch pattern
797+
claude_search_pattern = f"head:claude/issue-{issue_number}-"
798+
claude_pr_list_command = [
799+
"gh", "pr", "list",
800+
"--repo", config.GITHUB_REPOSITORY,
801+
"--state", "open",
802+
"--search", claude_search_pattern,
803+
"--limit", "1",
804+
"--json", "number,url,title,headRefName,baseRefName,state"
805+
]
806+
807+
pr_list_output = run_command(claude_pr_list_command, env=gh_env, check=False)
808+
809+
if not pr_list_output or pr_list_output.strip() == "[]":
810+
debug_log(f"No open PRs found for issue #{issue_number} with either Copilot or Claude branch pattern")
811+
return None
797812

798813
prs_data = json.loads(pr_list_output)
799814

@@ -822,7 +837,8 @@ def find_open_pr_for_issue(issue_number: int) -> dict:
822837

823838
def extract_issue_number_from_branch(branch_name: str) -> Optional[int]:
824839
"""
825-
Extracts the GitHub issue number from a branch name with format 'copilot/fix-<issue_number>'.
840+
Extracts the GitHub issue number from a branch name with format 'copilot/fix-<issue_number>'
841+
or 'claude/issue-<issue_number>-YYYYMMDD-HHMM'.
826842
827843
Args:
828844
branch_name: The branch name to extract the issue number from
@@ -833,10 +849,14 @@ def extract_issue_number_from_branch(branch_name: str) -> Optional[int]:
833849
if not branch_name:
834850
return None
835851

836-
# Use regex to match the exact pattern: copilot/fix-<number>
837-
# This ensures we only match the expected format and extract just the number
838-
pattern = r'^copilot/fix-(\d+)$'
839-
match = re.match(pattern, branch_name)
852+
# Check for copilot branch format: copilot/fix-<number>
853+
copilot_pattern = r'^copilot/fix-(\d+)$'
854+
match = re.match(copilot_pattern, branch_name)
855+
856+
if not match:
857+
# Check for claude branch format: claude/issue-<number>-YYYYMMDD-HHMM
858+
claude_pattern = r'^claude/issue-(\d+)-\d{8}-\d{4}$'
859+
match = re.match(claude_pattern, branch_name)
840860

841861
if match:
842862
try:
@@ -845,8 +865,7 @@ def extract_issue_number_from_branch(branch_name: str) -> Optional[int]:
845865
if issue_number > 0:
846866
return issue_number
847867
except ValueError:
848-
# This shouldn't happen since \d+ only matches digits, but being safe
849-
debug_log(f"Failed to convert extracted issue number '{match.group(1)}' to int")
868+
debug_log(f"Failed to convert extracted issue number '{match.group(1)}' from copilot or claude branch to int")
850869
pass
851870

852871
return None

‎src/merge_handler.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,8 @@ def _extract_remediation_info(pull_request: dict) -> tuple:
7272
# Extract remediation ID from branch name or PR labels
7373
remediation_id = None
7474

75-
# Check if this is a branch created by external agent (e.g., GitHub Copilot)
76-
if branch_name.startswith("copilot/fix"):
75+
# Check if this is a branch created by external agent (e.g., GitHub Copilot or Claude Code)
76+
if branch_name.startswith("copilot/fix") or branch_name.startswith("claude/issue-"):
7777
debug_log("Branch appears to be created by external agent. Extracting remediation ID from PR labels.")
7878
remediation_id = extract_remediation_id_from_labels(labels)
7979
# Extract GitHub issue number from branch name
@@ -83,14 +83,18 @@ def _extract_remediation_info(pull_request: dict) -> tuple:
8383
debug_log(f"Extracted external issue number from branch name: {issue_number}")
8484
else:
8585
debug_log(f"Could not extract issue number from branch name: {branch_name}")
86-
telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
86+
87+
# Set the external coding agent in telemetry based on branch prefix
88+
coding_agent = "EXTERNAL-CLAUDE-CODE" if branch_name.startswith("claude/") else "EXTERNAL-COPILOT"
89+
debug_log(f"Determined external coding agent to be: {coding_agent}")
90+
telemetry_handler.update_telemetry("additionalAttributes.codingAgent", coding_agent)
8791
else:
8892
# Use original method for branches created by SmartFix
8993
remediation_id = extract_remediation_id_from_branch(branch_name)
9094
telemetry_handler.update_telemetry("additionalAttributes.codingAgent", "INTERNAL-SMARTFIX")
9195

9296
if not remediation_id:
93-
if branch_name.startswith("copilot/fix"):
97+
if branch_name.startswith("copilot/fix") or branch_name.startswith("claude/issue-"):
9498
log(f"Error: Could not extract remediation ID from PR labels for external agent branch: {branch_name}", is_error=True)
9599
else:
96100
log(f"Error: Could not extract remediation ID from branch name: {branch_name}", is_error=True)

‎test/test_closed_handler.py‎

Lines changed: 103 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,26 @@
2727
# Add project root to path for imports
2828
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
2929

30+
# Define test environment variables used throughout the test file
31+
TEST_ENV_VARS = {
32+
'GITHUB_REPOSITORY': 'mock/repo',
33+
'GITHUB_TOKEN': 'mock-token',
34+
'BASE_BRANCH': 'main',
35+
'CONTRAST_HOST': 'test.contrastsecurity.com',
36+
'CONTRAST_ORG_ID': 'test-org-id',
37+
'CONTRAST_APP_ID': 'test-app-id',
38+
'CONTRAST_AUTHORIZATION_KEY': 'test-auth-key',
39+
'CONTRAST_API_KEY': 'test-api-key',
40+
'GITHUB_WORKSPACE': '/tmp',
41+
'RUN_TASK': 'closed',
42+
'BUILD_COMMAND': 'echo "Test build command"',
43+
'GITHUB_EVENT_PATH': '/tmp/github_event.json',
44+
'REPO_ROOT': '/tmp/test_repo',
45+
}
46+
47+
# Set environment variables before importing modules to prevent initialization errors
48+
os.environ.update(TEST_ENV_VARS)
49+
3050
# Now import project modules (after path modification)
3151
from src.config import reset_config, get_config # noqa: E402
3252
from src import closed_handler # noqa: E402
@@ -43,16 +63,8 @@ def setUp(self):
4363

4464
reset_config()
4565

46-
# Mock environment variables
47-
self.env_patcher = patch.dict(os.environ, {
48-
'CONTRAST_HOST': 'test.contrastsecurity.com',
49-
'CONTRAST_ORG_ID': 'test-org-id',
50-
'CONTRAST_APP_ID': 'test-app-id',
51-
'CONTRAST_AUTHORIZATION_KEY': 'test-auth-key',
52-
'CONTRAST_API_KEY': 'test-api-key',
53-
'GITHUB_EVENT_PATH': '/tmp/github_event.json',
54-
'REPO_ROOT': '/tmp/test_repo',
55-
})
66+
# Mock environment variables with complete required vars
67+
self.env_patcher = patch.dict(os.environ, TEST_ENV_VARS)
5668
self.env_patcher.start()
5769

5870
self.config = get_config()
@@ -272,6 +284,87 @@ def test_handle_closed_pr_integration(self, mock_init_telemetry, mock_load_event
272284
mock_extract_vuln.assert_called_once_with([])
273285
mock_notify.assert_called_once_with("REM-123", 123)
274286
mock_send_telemetry.assert_called_once()
287+
288+
@patch('src.telemetry_handler.update_telemetry')
289+
def test_extract_remediation_info_copilot_branch(self, mock_update_telemetry):
290+
"""Test _extract_remediation_info with Copilot branch"""
291+
with patch('src.closed_handler.extract_issue_number_from_branch') as mock_extract_issue:
292+
with patch('src.closed_handler.extract_remediation_id_from_labels') as mock_extract_remediation_id:
293+
# Setup
294+
mock_extract_issue.return_value = 42
295+
mock_extract_remediation_id.return_value = "REM-456"
296+
297+
pull_request = {
298+
"head": {"ref": "copilot/fix-42"},
299+
"labels": [{"name": "smartfix-id:REM-456"}]
300+
}
301+
302+
# Execute
303+
result = closed_handler._extract_remediation_info(pull_request)
304+
305+
# Assert
306+
self.assertEqual(result, ("REM-456", [{"name": "smartfix-id:REM-456"}]))
307+
mock_extract_issue.assert_called_once_with("copilot/fix-42")
308+
mock_extract_remediation_id.assert_called_once_with([{"name": "smartfix-id:REM-456"}])
309+
310+
# Verify telemetry updates
311+
mock_update_telemetry.assert_any_call("additionalAttributes.externalIssueNumber", 42)
312+
mock_update_telemetry.assert_any_call("additionalAttributes.codingAgent", "EXTERNAL-COPILOT")
313+
314+
@patch('src.telemetry_handler.update_telemetry')
315+
def test_extract_remediation_info_claude_branch(self, mock_update_telemetry):
316+
"""Test _extract_remediation_info with Claude Code branch"""
317+
with patch('src.closed_handler.extract_issue_number_from_branch') as mock_extract_issue:
318+
with patch('src.closed_handler.extract_remediation_id_from_labels') as mock_extract_remediation_id:
319+
# Setup
320+
mock_extract_issue.return_value = 75
321+
mock_extract_remediation_id.return_value = "REM-789"
322+
323+
pull_request = {
324+
"head": {"ref": "claude/issue-75-20250908-1723"},
325+
"labels": [{"name": "smartfix-id:REM-789"}]
326+
}
327+
328+
# Execute
329+
result = closed_handler._extract_remediation_info(pull_request)
330+
331+
# Assert
332+
self.assertEqual(result, ("REM-789", [{"name": "smartfix-id:REM-789"}]))
333+
mock_extract_issue.assert_called_once_with("claude/issue-75-20250908-1723")
334+
mock_extract_remediation_id.assert_called_once_with([{"name": "smartfix-id:REM-789"}])
335+
336+
# Verify telemetry updates - key assertions for Claude Code
337+
mock_update_telemetry.assert_any_call("additionalAttributes.externalIssueNumber", 75)
338+
mock_update_telemetry.assert_any_call("additionalAttributes.codingAgent", "EXTERNAL-CLAUDE-CODE")
339+
340+
@patch('src.telemetry_handler.update_telemetry')
341+
def test_extract_remediation_info_claude_branch_no_issue_number(self, mock_update_telemetry):
342+
"""Test _extract_remediation_info with Claude Code branch without extractable issue number"""
343+
with patch('src.closed_handler.extract_issue_number_from_branch') as mock_extract_issue:
344+
with patch('src.closed_handler.extract_remediation_id_from_labels') as mock_extract_remediation_id:
345+
# Setup - simulate issue number not found
346+
mock_extract_issue.return_value = None
347+
mock_extract_remediation_id.return_value = "REM-789"
348+
349+
pull_request = {
350+
"head": {"ref": "claude/issue-75-20250908-1723"},
351+
"labels": [{"name": "smartfix-id:REM-789"}]
352+
}
353+
354+
# Execute
355+
result = closed_handler._extract_remediation_info(pull_request)
356+
357+
# Assert
358+
self.assertEqual(result, ("REM-789", [{"name": "smartfix-id:REM-789"}]))
359+
mock_extract_issue.assert_called_once_with("claude/issue-75-20250908-1723")
360+
mock_extract_remediation_id.assert_called_once_with([{"name": "smartfix-id:REM-789"}])
361+
362+
# Should NOT call update_telemetry for externalIssueNumber, but SHOULD call it for codingAgent
363+
# Verify it is not called with externalIssueNumber
364+
for call in mock_update_telemetry.call_args_list:
365+
self.assertNotEqual(call[0][0], "additionalAttributes.externalIssueNumber")
366+
# But should still identify as Claude Code agent
367+
mock_update_telemetry.assert_any_call("additionalAttributes.codingAgent", "EXTERNAL-CLAUDE-CODE")
275368

276369

277370
if __name__ == '__main__':

0 commit comments

Comments
 (0)