AIML-644: Migrate credit tracking to org-level endpoint

Add get_credit_tracking_org() using the org-scoped URL
(/organizations/{orgId}/credit-tracking) and migrate all four
get_credit_tracking() call sites in main.py to use it, removing
the CONTRAST_APP_ID dependency. The org-level route already exists
in the backend (RemediationCreditTrackingController). Add tests
covering the new function's success, URL shape, and error handling.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

Author: JacobMagesHaskinsContrast

src/contrast_api.py

@@ -1139,4 +1139,51 @@ def get_credit_tracking(contrast_host: str, contrast_org_id: str, contrast_app_i
         debug_log(f"Unexpected error calling credit-tracking API: {e}")
         return None
 
+
+def get_credit_tracking_org(contrast_host: str, contrast_org_id: str, contrast_auth_key: str, contrast_api_key: str) -> Optional[CreditTrackingResponse]:
+    """Get credit tracking information from the org-level Contrast API endpoint.
+
+    Args:
+        contrast_host: The Contrast Security host URL.
+        contrast_org_id: The organization ID.
+        contrast_auth_key: The Contrast authorization key.
+        contrast_api_key: The Contrast API key.
+
+    Returns:
+        CreditTrackingResponse object if successful, None if failed.
+    """
+    api_url = f"https://{normalize_host(contrast_host)}/api/v4/aiml-remediation/organizations/{contrast_org_id}/credit-tracking"
+
+    headers = {
+        "Authorization": contrast_auth_key,
+        "API-Key": contrast_api_key,
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "User-Agent": config.USER_AGENT
+    }
+
+    try:
+        debug_log(f"Fetching org-level credit tracking from: {api_url}")
+        response = requests.get(api_url, headers=headers, timeout=30)
+        response.raise_for_status()
+
+        debug_log(f"Org-level credit tracking API response status code: {response.status_code}")
+        debug_log(f"Raw org-level credit tracking response: {response.text}")
+
+        data = response.json()
+        return CreditTrackingResponse.from_api_response(data)
+
+    except requests.exceptions.HTTPError as e:
+        debug_log(f"HTTP error fetching org-level credit tracking: {e.response.status_code} - {e.response.text}")
+        return None
+    except requests.exceptions.RequestException as e:
+        debug_log(f"Request error fetching org-level credit tracking: {e}")
+        return None
+    except json.JSONDecodeError:
+        debug_log("Error decoding JSON response from org-level credit-tracking API.")
+        return None
+    except Exception as e:
+        debug_log(f"Unexpected error calling org-level credit-tracking API: {e}")
+        return None
+
 # %%

src/main.py

@@ -158,10 +158,9 @@ def _main_impl(vuln_count):  # noqa: C901
 
     # Log initial credit tracking status if using Contrast LLM (only for SMARTFIX agent)
     if config.CODING_AGENT == CodingAgents.SMARTFIX.name and config.USE_CONTRAST_LLM:
-        initial_credit_info = contrast_api.get_credit_tracking(
+        initial_credit_info = contrast_api.get_credit_tracking_org(
             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
         )
@@ -209,10 +208,9 @@ def _main_impl(vuln_count):  # noqa: C901
 
         # Check credit exhaustion for Contrast LLM usage
         if config.USE_CONTRAST_LLM:
-            current_credit_info = contrast_api.get_credit_tracking(
+            current_credit_info = contrast_api.get_credit_tracking_org(
                 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
             )
@@ -513,10 +511,9 @@ def _main_impl(vuln_count):  # noqa: C901
 
                 # Append credit tracking information to PR body if using Contrast LLM
                 if config.CODING_AGENT == CodingAgents.SMARTFIX.name and config.USE_CONTRAST_LLM:
-                    current_credit_info = contrast_api.get_credit_tracking(
+                    current_credit_info = contrast_api.get_credit_tracking_org(
                         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
                     )
@@ -599,10 +596,9 @@ def _main_impl(vuln_count):  # noqa: C901
 
                             # Log updated credit tracking status after PR notification (only for SMARTFIX agent)
                             if config.CODING_AGENT == CodingAgents.SMARTFIX.name and config.USE_CONTRAST_LLM:
-                                updated_credit_info = contrast_api.get_credit_tracking(
+                                updated_credit_info = contrast_api.get_credit_tracking_org(
                                     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
                                 )

test/contrast_api/test_contrast_api_credit_tracking.py

@@ -260,5 +260,78 @@ def test_get_credit_tracking_uses_bearer_authorization(self, mock_get):
         self.assertEqual(headers['Authorization'], 'test-auth-key')
 
 
+class TestContrastApiCreditTrackingOrg(unittest.TestCase):
+    """Test cases for org-level credit tracking (no app_id) in contrast_api module."""
+
+    def setUp(self):
+        self.sample_api_response = {
+            "organizationId": "12345678-1234-1234-1234-123456789abc",
+            "enabled": True,
+            "maxCredits": 50,
+            "creditsUsed": 7,
+            "startDate": "2024-10-01T14:30:00Z",
+            "endDate": "2024-11-12T14:30:00Z"
+        }
+
+    @patch('src.contrast_api.requests.get')
+    def test_get_credit_tracking_org_returns_valid_response(self, mock_get):
+        """Test that org-level endpoint returns properly structured CreditTrackingResponse."""
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = json.dumps(self.sample_api_response)
+        mock_response.json.return_value = self.sample_api_response
+        mock_get.return_value = mock_response
+
+        result = contrast_api.get_credit_tracking_org(
+            contrast_host="test.contrastsecurity.com",
+            contrast_org_id="test-org-id",
+            contrast_auth_key="test-auth-key",
+            contrast_api_key="test-api-key"
+        )
+
+        self.assertIsInstance(result, CreditTrackingResponse)
+        self.assertTrue(result.enabled)
+
+    @patch('src.contrast_api.requests.get')
+    def test_get_credit_tracking_org_url_has_no_app_id(self, mock_get):
+        """Test that the org-level endpoint URL does not include an application ID."""
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = json.dumps(self.sample_api_response)
+        mock_response.json.return_value = self.sample_api_response
+        mock_get.return_value = mock_response
+
+        contrast_api.get_credit_tracking_org(
+            contrast_host="test.contrastsecurity.com",
+            contrast_org_id="test-org-id",
+            contrast_auth_key="test-auth-key",
+            contrast_api_key="test-api-key"
+        )
+
+        called_url = mock_get.call_args[0][0]
+        self.assertIn("/organizations/test-org-id/credit-tracking", called_url)
+        self.assertNotIn("/applications/", called_url)
+
+    @patch('src.contrast_api.requests.get')
+    def test_get_credit_tracking_org_returns_none_on_http_error(self, mock_get):
+        """Test that org-level endpoint returns None on HTTP errors."""
+        mock_response = MagicMock()
+        mock_response.status_code = 404
+        mock_response.text = "Not Found"
+        mock_get.return_value = mock_response
+        mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
+            response=mock_response
+        )
+
+        result = contrast_api.get_credit_tracking_org(
+            contrast_host="test.contrastsecurity.com",
+            contrast_org_id="test-org-id",
+            contrast_auth_key="test-auth-key",
+            contrast_api_key="test-api-key"
+        )
+
+        self.assertIsNone(result)
+
+
 if __name__ == '__main__':
     unittest.main()