Merge pull request #78 from Contrast-Security-OSS/add_more_contrast_llm_logging

AIML-171 Add missing Contrast LLM logging

Author: JacobMagesHaskinsContrast

src/main.py

@@ -295,8 +295,17 @@ def main():  # noqa: C901
         )
         if initial_credit_info:
             log(initial_credit_info.to_log_message())
+            # Log any initial warnings
+            if initial_credit_info.should_log_warning():
+                warning_msg = initial_credit_info.get_credit_warning_message()
+                if initial_credit_info.is_exhausted:
+                    log(warning_msg, is_error=True)
+                    error_exit(FailureCategory.GENERAL_FAILURE.value)
+                else:
+                    log(warning_msg, is_warning=True)
         else:
-            debug_log("Could not retrieve initial credit tracking information")
+            log("Could not retrieve initial credit tracking information", is_error=True)
+            error_exit(FailureCategory.GENERAL_FAILURE.value)
 
     while True:
         telemetry_handler.reset_vuln_specific_telemetry()
@@ -327,6 +336,20 @@ def main():  # noqa: C901
             log(f"\n--- Reached max PR limit ({max_open_prs_setting}). Current open PRs: {current_open_pr_count}. Stopping processing. ---")
             break
 
+        # Check credit exhaustion for Contrast LLM usage
+        if config.USE_CONTRAST_LLM:
+            current_credit_info = contrast_api.get_credit_tracking(
+                contrast_host=config.CONTRAST_HOST,
+                contrast_org_id=config.CONTRAST_ORG_ID,
+                contrast_app_id=config.CONTRAST_APP_ID,
+                contrast_auth_key=config.CONTRAST_AUTHORIZATION_KEY,
+                contrast_api_key=config.CONTRAST_API_KEY
+            )
+            if current_credit_info and current_credit_info.is_exhausted:
+                log("\n--- Credits exhausted. Stopping processing. ---")
+                log("Credits have been exhausted. Contact your CSM to request additional credits.", is_error=True)
+                break
+
         # --- Fetch Next Vulnerability Data from API ---
         if config.CODING_AGENT == CodingAgents.SMARTFIX.name:
             # For SMARTFIX, get vulnerability with prompts
@@ -550,6 +573,16 @@ def main():  # noqa: C901
                 projected_credit_info = current_credit_info.with_incremented_usage()
                 updated_pr_body += projected_credit_info.to_pr_body_section()
 
+                # Show countdown message and warnings
+                credits_after = projected_credit_info.credits_remaining
+                log(f"Credit consumed. {credits_after} credits remaining")
+                if projected_credit_info.should_log_warning():
+                    warning_msg = projected_credit_info.get_credit_warning_message()
+                    if projected_credit_info.is_exhausted:
+                        log(warning_msg, is_error=True)
+                    else:
+                        log(warning_msg, is_warning=True)
+
         # Create a brief summary for the telemetry aiSummaryReport (limited to 255 chars in DB)
         # Generate an optimized summary using the dedicated function in telemetry_handler
         brief_summary = telemetry_handler.create_ai_summary_report(updated_pr_body)

src/smartfix/domains/workflow/credit_tracking.py

@@ -38,6 +38,16 @@ def credits_remaining(self) -> int:
         """Calculate remaining credits."""
         return self.max_credits - self.credits_used
 
+    @property
+    def is_exhausted(self) -> bool:
+        """Check if credits are exhausted."""
+        return self.credits_remaining <= 0
+
+    @property
+    def is_low(self) -> bool:
+        """Check if credits are running low (5 or fewer remaining)."""
+        return self.credits_remaining <= 5 and self.credits_remaining > 0
+
     def _format_timestamp(self, iso_timestamp: str) -> str:
         """Format ISO timestamp to human-readable format."""
         if not iso_timestamp:
@@ -60,6 +70,19 @@ def to_log_message(self) -> str:
         return (f"Credits: {self.credits_used}/{self.max_credits} used "
                 f"({self.credits_remaining} remaining). Trial expires {self.end_date}")
 
+    def get_credit_warning_message(self) -> str:
+        """Get warning message for credit status, with color formatting."""
+        if self.is_exhausted:
+            return "Credits have been exhausted. Contact your CSM to request additional credits."
+        elif self.is_low:
+            # Yellow text formatting for low credits warning
+            return f"\033[0;33m{self.credits_remaining} credits remaining \033[0m"
+        return ""
+
+    def should_log_warning(self) -> bool:
+        """Check if a warning should be logged."""
+        return self.is_exhausted or self.is_low
+
     def to_pr_body_section(self) -> str:
         """Format credit information for PR body append."""
         if not self.enabled:

test/test_credit_tracking.py

@@ -69,6 +69,44 @@ def test_credits_remaining_property(self):
         response.credits_used = 50
         self.assertEqual(response.credits_remaining, 0)  # 50 - 50
 
+    def test_is_exhausted_property(self):
+        """Test credit exhaustion detection."""
+        response = CreditTrackingResponse.from_api_response(self.sample_api_data)
+
+        # Not exhausted with 43 remaining
+        self.assertFalse(response.is_exhausted)
+
+        # Exhausted at exactly 0 remaining
+        response.credits_used = 50
+        self.assertTrue(response.is_exhausted)
+
+        # Exhausted when over limit
+        response.credits_used = 55
+        self.assertTrue(response.is_exhausted)
+
+    def test_is_low_property(self):
+        """Test low credit detection."""
+        response = CreditTrackingResponse.from_api_response(self.sample_api_data)
+
+        # Not low with 43 remaining
+        self.assertFalse(response.is_low)
+
+        # Low with exactly 5 remaining
+        response.credits_used = 45
+        self.assertTrue(response.is_low)
+
+        # Low with 1 remaining
+        response.credits_used = 49
+        self.assertTrue(response.is_low)
+
+        # Not low when exhausted (0 remaining)
+        response.credits_used = 50
+        self.assertFalse(response.is_low)
+
+        # Not low with 6 remaining
+        response.credits_used = 44
+        self.assertFalse(response.is_low)
+
     def test_to_log_message_enabled(self):
         """Test log message formatting when enabled."""
         response = CreditTrackingResponse.from_api_response(self.sample_api_data)
@@ -77,6 +115,64 @@ def test_to_log_message_enabled(self):
         expected = "Credits: 7/50 used (43 remaining). Trial expires 2024-11-12T14:30:00Z"
         self.assertEqual(message, expected)
 
+    def test_get_credit_warning_message(self):
+        """Test credit warning message generation."""
+        response = CreditTrackingResponse.from_api_response(self.sample_api_data)
+
+        # Test exhausted credits
+        response.credits_used = 50
+        warning_msg = response.get_credit_warning_message()
+
+        self.assertEqual(warning_msg, "Credits have been exhausted. Contact your CSM to request additional credits.")
+
+        # Test low credits (with color formatting)
+        response.credits_used = 45  # 5 remaining
+        warning_msg = response.get_credit_warning_message()
+
+        self.assertIn("5 credits remaining", warning_msg)
+        self.assertIn("\033[0;33m", warning_msg)  # Yellow color code
+        self.assertIn("\033[0m", warning_msg)   # Reset color code
+
+        # Test normal credits (no warning)
+        response.credits_used = 44  # 6 remaining
+        warning_msg = response.get_credit_warning_message()
+
+        self.assertEqual(warning_msg, "")
+
+    def test_should_log_warning(self):
+        """Test warning condition detection."""
+        response = CreditTrackingResponse.from_api_response(self.sample_api_data)
+
+        # Normal state - no warning
+        self.assertFalse(response.should_log_warning())
+
+        # Low credits - should warn
+        response.credits_used = 45
+        self.assertTrue(response.should_log_warning())
+
+        # Exhausted - should warn
+        response.credits_used = 50
+        self.assertTrue(response.should_log_warning())
+
+    def test_basic_functionality_only(self):
+        """Test that we only have basic client-side functionality."""
+        response = CreditTrackingResponse.from_api_response(self.sample_api_data)
+
+        # Verify the basic properties work
+        self.assertEqual(response.credits_remaining, 43)
+        self.assertFalse(response.is_exhausted)
+        self.assertFalse(response.is_low)
+
+        # Verify exhaustion detection
+        response.credits_used = 50
+        self.assertTrue(response.is_exhausted)
+        self.assertFalse(response.is_low)
+
+        # Verify low credit detection
+        response.credits_used = 46  # 4 remaining
+        self.assertFalse(response.is_exhausted)
+        self.assertTrue(response.is_low)
+
     def test_to_log_message_disabled(self):
         """Test log message formatting when disabled."""
         response = CreditTrackingResponse.from_api_response(self.disabled_api_data)
@@ -90,11 +186,12 @@ def test_to_pr_body_section_enabled(self):
         response = CreditTrackingResponse.from_api_response(self.sample_api_data)
         pr_section = response.to_pr_body_section()
 
+        # Test new format matches documentation spec
         self.assertIn("### Contrast LLM Credits", pr_section)
-        self.assertIn("**Used:** 7/50", pr_section)
-        self.assertIn("**Remaining:** 43", pr_section)
-        self.assertIn("**Trial Period:** Oct 01, 2024 to Nov 12, 2024", pr_section)
-        self.assertTrue(pr_section.startswith("\n---\n"))
+        self.assertIn("- **Used:** 7/50", pr_section)
+        self.assertIn("- **Remaining:** 43", pr_section)
+        # Should include trial period dates
+        self.assertIn("- **Trial Period:** Oct 01, 2024 to Nov 12, 2024", pr_section)
 
     def test_to_pr_body_section_disabled(self):
         """Test PR body section formatting when disabled."""
@@ -155,15 +252,19 @@ def test_edge_cases(self):
 
         response = CreditTrackingResponse.from_api_response(data)
         pr_section = response.to_pr_body_section()
-        self.assertIn("**Trial Period:** Jan 15, 2025 to Mar 31, 2025", pr_section)
+        self.assertIn("- **Used:** 0/50", pr_section)
+        self.assertIn("- **Remaining:** 50", pr_section)
+        self.assertIn("- **Trial Period:** Jan 15, 2025 to Mar 31, 2025", pr_section)
 
         # Test empty dates edge case
         data["startDate"] = ""
         data["endDate"] = ""
 
         response = CreditTrackingResponse.from_api_response(data)
         pr_section = response.to_pr_body_section()
-        self.assertIn("**Trial Period:** Unknown to Unknown", pr_section)
+        self.assertIn("- **Used:** 0/50", pr_section)
+        self.assertIn("- **Remaining:** 50", pr_section)
+        self.assertIn("- **Trial Period:** Unknown to Unknown", pr_section)
 
 
 if __name__ == '__main__':

test/test_main.py

@@ -161,24 +161,18 @@ def test_duplicate_vuln_with_open_pr_skips_cleanly(self):
             with patch('src.git_handler.generate_label_details') as mock_label:
                 mock_label.return_value = ('contrast-vuln-id:TEST-VULN-UUID-123', 'color', 'desc')
 
-                # Mock error_exit to track if it's called (it shouldn't be)
-                with patch('src.main.error_exit') as mock_error_exit:
-                    with patch.dict('os.environ', self.env_vars, clear=True):
-                        # Run main and capture output
-                        with io.StringIO() as buf, contextlib.redirect_stdout(buf):
-                            main()
-                            output = buf.getvalue()
-
-                        # Verify the vulnerability was skipped both times
-                        self.assertIn("Skipping vulnerability TEST-VULN-UUID-123", output)
-                        self.assertIn("Already skipped TEST-VULN-UUID-123 before, breaking loop", output)
-
-                        # CRITICAL: Verify error_exit was NOT called
-                        # The bug would have triggered error_exit on the second iteration
-                        mock_error_exit.assert_not_called()
-
-                        # Verify the loop broke cleanly
-                        self.assertIn("No vulnerabilities were processed in this run", output)
+                with patch.dict('os.environ', self.env_vars, clear=True):
+                    # Run main and capture output
+                    with io.StringIO() as buf, contextlib.redirect_stdout(buf):
+                        main()
+                        output = buf.getvalue()
+
+                    # Verify the vulnerability was skipped both times
+                    self.assertIn("Skipping vulnerability TEST-VULN-UUID-123", output)
+                    self.assertIn("Already skipped TEST-VULN-UUID-123 before, breaking loop", output)
+
+                    # Verify the loop broke cleanly
+                    self.assertIn("No vulnerabilities were processed in this run", output)
 
 
 if __name__ == '__main__':