Aap 62693 job launch

awx/main/tasks/jobs.py

@@ -93,6 +93,13 @@
 # Django flags
 from flags.state import flag_enabled
 
+# Workload Identity imports
+from ansible_base.resource_registry.workload_identity_client import (
+    WorkloadIdentityClient,
+    TokenRequestError,
+    get_workload_identity_client,
+)
+
 # Workload Identity
 from ansible_base.lib.workload_identity.controller import AutomationControllerJobScope
 
@@ -514,6 +521,10 @@ def pre_run_hook(self, instance, private_data_dir):
         # Before task is started, ensure that job_event partitions exist
         create_partition(instance.event_class._meta.db_table, start=instance.created)
 
+        # Request workload identity JWT token if feature is enabled and job needs it
+        if flag_enabled("FEATURE_OIDC_WORKLOAD_IDENTITY_ENABLED") and self._job_needs_workload_jwt(instance):
+            self._request_workload_identity_token(instance)
+
     def post_run_hook(self, instance, status):
         """
         Hook for any steps to run before job/task is marked as complete.
@@ -532,6 +543,90 @@ def final_run_hook(self, instance, status, private_data_dir):
         if instance.spawned_by_workflow:
             ScheduleWorkflowManager().schedule()
 
+    def _job_needs_workload_jwt(self, instance):
+        """
+        Determine if this job needs a workload identity JWT token.
+
+        Currently checks if the job has any vault credentials, as Vault is the
+        primary use case for workload identity JWTs (ANSTRAT-1558).
+
+        Args:
+            instance: The UnifiedJob instance being executed
+
+        Returns:
+            bool: True if job needs a JWT token
+        """
+        # Check if job has credentials attribute (not all job types do)
+        if not hasattr(instance, 'credentials'):
+            return False
+
+        # Check if any vault credentials are present
+        # TODO: Extend this to check for other credential types that support workload identity
+        # when AAP-62694 (JWT-enabled credential plugin) is implemented
+        vault_creds = instance.credentials.filter(credential_type__kind='vault').exists()
+
+        if vault_creds:
+            logger.debug(f"Job {instance.id} has vault credentials, will request workload JWT")
+            return True
+
+        # TODO: Check for other credential types that require workload identity
+        # e.g., AWS, Azure, GCP when those integrations are implemented
+
+        return False
+
+    def _request_workload_identity_token(self, instance):
+        """
+        Request a workload identity JWT token from Gateway for this job.
+
+        This method is called during pre_run_hook if the feature flag is enabled
+        and the job needs a JWT (determined by _job_needs_workload_jwt).
+        The JWT token is logged but not currently used for anything (see AAP-62694).
+
+        Args:
+            instance: The UnifiedJob instance being executed
+        """
+        try:
+            # Populate JWT claims from the job
+            claims = populate_claims_for_workload(instance)
+
+            # TODO: Make scope and audience configurable
+            # For now, hardcode the scope as defined in the OIDC spec
+            scope = "aap_controller_automation_job"
+            # TODO: Audience should be configured per integration (Vault, AWS, etc.)
+            audience = "https://vault.example.com"
+
+            # Get the workload identity client (uses service token auth)
+            client = get_workload_identity_client()
+
+            # Request the JWT token
+            logger.info(f"Requesting workload identity token for job {instance.id} ({instance.name})")
+            response = client.request_workload_jwt(
+                claims=claims,
+                scope=scope,
+                audience=audience
+            )
+
+            logger.info(
+                f"Successfully obtained workload identity token for job {instance.id}. "
+                f"JWT length: {len(response.jwt)} characters"
+            )
+            logger.debug(f"JWT token: {response.jwt[:50]}...")
+
+            # TODO (AAP-62694): Store JWT for use by credential plugins
+            # For now, we just log that we got it
+
+        except TokenRequestError as e:
+            logger.error(
+                f"Failed to obtain workload identity token for job {instance.id}: {e}"
+            )
+            # Don't fail the job if JWT request fails - this is optional functionality
+
+        except Exception as e:
+            logger.exception(
+                f"Unexpected error requesting workload identity token for job {instance.id}: {e}"
+            )
+            # Don't fail the job on unexpected errors
+
     def should_use_fact_cache(self):
         return False
 

awx/main/tests/unit/tasks/test_jobs.py

@@ -427,3 +427,115 @@ def test_populate_claims_for_adhoc_command(workload_attrs, expected_claims):
 
     claims = jobs.populate_claims_for_workload(adhoc_command)
     assert claims == expected_claims
+
+
+@pytest.mark.django_db
+@mock.patch('awx.main.tasks.jobs.get_workload_identity_client')
+def test_request_workload_identity_token_success(mock_get_client, job_template_factory, credential):
+    """Test successful workload identity token request."""
+    # Mock the client and response
+    from ansible_base.resource_registry.workload_identity_client import WorkloadIdentityTokenResponse
+    mock_client = mock.Mock()
+    mock_client.request_workload_jwt.return_value = WorkloadIdentityTokenResponse(
+        jwt='eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.signature'
+    )
+    mock_get_client.return_value = mock_client
+
+    # Create a test job with vault credential
+    jt = job_template_factory('test-jt')
+    job = jt.create_unified_job()
+    job.id = 123
+    job.name = 'Test Job'
+    job.credentials.add(credential)
+
+    # Create a BaseTask instance and call the method
+    task = jobs.BaseTask()
+    task._request_workload_identity_token(job)
+
+    # Verify client was called with claims
+    mock_client.request_workload_jwt.assert_called_once()
+    call_args = mock_client.request_workload_jwt.call_args
+    assert 'claims' in call_args.kwargs
+    assert 'scope' in call_args.kwargs
+    assert 'audience' in call_args.kwargs
+    assert call_args.kwargs['scope'] == 'aap_controller_automation_job'
+
+
+@pytest.mark.django_db
+@mock.patch('awx.main.tasks.jobs.get_workload_identity_client')
+@mock.patch('awx.main.tasks.jobs.flag_enabled')
+@mock.patch('awx.main.tasks.jobs.logger')
+def test_request_workload_identity_token_failure(mock_logger, mock_flag_enabled, mock_get_client):
+    """Test workload identity token request failure handling."""
+    # Enable the feature flag
+    mock_flag_enabled.return_value = True
+
+    # Mock the client to raise an error
+    from ansible_base.resource_registry.workload_identity_client import TokenRequestError
+    mock_client = mock.Mock()
+    mock_client.request_workload_jwt.side_effect = TokenRequestError("Gateway unavailable")
+    mock_get_client.return_value = mock_client
+
+    # Create a test job
+    job = Job(id=456, name='Failing Job')
+
+    # Create a BaseTask instance and call the method
+    task = jobs.BaseTask()
+    task._request_workload_identity_token(job)
+
+    # Verify error was logged
+    mock_logger.error.assert_called()
+    error_message = mock_logger.error.call_args[0][0]
+    assert 'Failed to obtain workload identity token' in error_message
+    assert '456' in error_message
+
+
+@pytest.mark.django_db
+@mock.patch('awx.main.tasks.jobs.get_workload_identity_client')
+@mock.patch('awx.main.tasks.jobs.flag_enabled')
+def test_request_workload_identity_token_feature_flag_disabled(mock_flag_enabled, mock_get_client):
+    """Test that JWT request is skipped when feature flag is disabled."""
+    # Disable the feature flag
+    mock_flag_enabled.return_value = False
+
+    # Create a test job
+    job = Job(id=789, name='Test Job')
+
+    # Create a BaseTask instance and call pre_run_hook
+    task = jobs.BaseTask()
+    task.pre_run_hook(job, '/tmp/private_data')
+
+    # Verify client was never called
+    mock_get_client.assert_not_called()
+
+
+@pytest.mark.django_db
+def test_job_needs_workload_jwt_with_vault_credential(job_template_factory, vault_credential):
+    """Test that jobs with vault credentials need JWT tokens."""
+    jt = job_template_factory('test-jt')
+    job = jt.create_unified_job()
+    job.credentials.add(vault_credential)
+
+    task = jobs.BaseTask()
+    assert task._job_needs_workload_jwt(job) is True
+
+
+@pytest.mark.django_db
+def test_job_needs_workload_jwt_without_vault_credential(job_template_factory, machine_credential):
+    """Test that jobs without vault credentials don't need JWT tokens."""
+    jt = job_template_factory('test-jt')
+    job = jt.create_unified_job()
+    job.credentials.add(machine_credential)
+
+    task = jobs.BaseTask()
+    assert task._job_needs_workload_jwt(job) is False
+
+
+@pytest.mark.django_db
+def test_job_needs_workload_jwt_no_credentials(job_template_factory):
+    """Test that jobs without any credentials don't need JWT tokens."""
+    jt = job_template_factory('test-jt')
+    job = jt.create_unified_job()
+
+    task = jobs.BaseTask()
+    assert task._job_needs_workload_jwt(job) is False

awx/settings/development_defaults.py

@@ -69,3 +69,6 @@
 # ======================!!!!!!! FOR DEVELOPMENT ONLY !!!!!!!=================================
 
 FEATURE_INDIRECT_NODE_COUNTING_ENABLED = True
+
+# Feature flag for workload identity JWT integration (AAP-62693)
+FEATURE_OIDC_WORKLOAD_IDENTITY_ENABLED = True