{AKS} AI Manager Grant the caller built-in roles on create and namespace add

src/aimanager/HISTORY.rst

@@ -3,6 +3,13 @@
 Release History
 ===============
 
+1.5.0
+++++++
+* ``az aimanager create`` and ``az aimanager namespace add``: On success, grant the caller the
+  built-in ``Azure AIManager Contributor`` and ``Azure AIManager and namespace RBAC Reader``
+  roles on the new resource (best-effort; requires Owner or User Access Administrator). Skipped
+  with ``--no-wait``.
+
 1.4.1
 ++++++
 * ``az aimanager modelsource`` and ``az aimanager namespace modeldeployment``: Accept

src/aimanager/azext_aimanager/_help.py

@@ -15,6 +15,10 @@
 helps['aimanager create'] = """
     type: command
     short-summary: Create an AI Manager resource.
+    long-summary: >
+        Once creation succeeds the caller is granted the built-in 'Azure AIManager Contributor'
+        and 'Azure AIManager and namespace RBAC Reader' roles on the new AI Manager (best-effort;
+        requires Owner or User Access Administrator). Skipped with --no-wait.
     examples:
         - name: Create an AI Manager
           text: az aimanager create --name my-ai-manager -g myrg -l eastus2
@@ -145,6 +149,10 @@
 helps['aimanager namespace add'] = """
     type: command
     short-summary: Add a namespace to an AI Manager.
+    long-summary: >
+        Once creation succeeds the caller is granted the built-in 'Azure AIManager Contributor'
+        and 'Azure AIManager and namespace RBAC Reader' roles on the new namespace (best-effort;
+        requires Owner or User Access Administrator). Skipped with --no-wait.
     examples:
         - name: Add a namespace
           text: az aimanager namespace add -m my-ai-manager -g myrg --name team-alpha

src/aimanager/azext_aimanager/_roleassignments.py

@@ -0,0 +1,114 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import base64
+import json
+import uuid
+
+from azure.cli.core.commands.client_factory import get_mgmt_service_client, get_subscription_id
+from azure.cli.core.profiles import ResourceType, get_sdk
+from azure.core.exceptions import HttpResponseError, ResourceExistsError
+from knack.log import get_logger
+
+from azext_aimanager.constants import AIMANAGER_ROLE_NAMES
+
+logger = get_logger(__name__)
+
+
+def _get_caller_identity(cli_ctx):
+    """Return the caller's Entra object ID and a best-guess principal type from the token.
+
+    The caller already exists in the directory, so there is no propagation delay to wait on.
+    Returns (None, None) if the object ID cannot be read.
+    """
+    from azure.cli.core._profile import Profile
+    try:
+        cred, _, _ = Profile(cli_ctx=cli_ctx).get_raw_token()
+        payload = cred[1].split('.')[1]
+        payload += '=' * (-len(payload) % 4)
+        claims = json.loads(base64.urlsafe_b64decode(payload))
+    except Exception:  # pylint: disable=broad-except
+        return None, None
+    # 'idtyp' == 'app' identifies an app-only (service principal) token; otherwise it is a user.
+    principal_type = 'ServicePrincipal' if claims.get('idtyp') == 'app' else 'User'
+    return claims.get('oid'), principal_type
+
+
+def assign_caller_roles(cmd, scope, role_definition_ids):
+    """Best-effort: grant the caller the given built-in roles on the scope.
+
+    Single attempt per role (the caller already exists, so the AAD-propagation retry the shared
+    AKS helper performs is unnecessary and would hang the create for ~2 minutes on the common
+    permission-denied path). Never raises: creating an AI Manager or namespace must not fail just
+    because the caller lacks permission to assign roles (that requires Owner or User Access
+    Administrator).
+    """
+    object_id, principal_type = _get_caller_identity(cmd.cli_ctx)
+    if not object_id:
+        logger.warning(
+            "Could not determine the caller's object ID; skipping role assignment on %s.", scope)
+        return
+
+    try:
+        subscription_id = get_subscription_id(cmd.cli_ctx)
+        role_assignment_create_parameters = get_sdk(
+            cmd.cli_ctx, ResourceType.MGMT_AUTHORIZATION,
+            'RoleAssignmentCreateParameters', mod='models', operation_group='role_assignments')
+        assignments_client = get_mgmt_service_client(
+            cmd.cli_ctx, ResourceType.MGMT_AUTHORIZATION).role_assignments
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.warning("Could not set up role assignment on %s: %s", scope, ex)
+        return
+    # The caller is a user or a service principal. Start with the token's hint and fall back to
+    # the other type if ARM reports a principal-type mismatch, rather than guessing wrong.
+    principal_types = [principal_type] + [t for t in ('User', 'ServicePrincipal') if t != principal_type]
+
+    for role_id in role_definition_ids:
+        matched_type = _assign_role(assignments_client, role_assignment_create_parameters,
+                                    subscription_id, scope, role_id, object_id, principal_types)
+        # Once we learn the caller's real type, try it first for the remaining roles.
+        if matched_type and principal_types[0] != matched_type:
+            principal_types = [matched_type] + [t for t in principal_types if t != matched_type]
+
+
+def _assign_role(assignments_client, params_model, subscription_id, scope, role_id,
+                 object_id, principal_types):
+    """Create one role assignment, trying each principal type. Returns the type that worked, or
+    None if the assignment could not be created."""
+    role_definition_id = (
+        f"/subscriptions/{subscription_id}"
+        f"/providers/Microsoft.Authorization/roleDefinitions/{role_id}")
+    last = len(principal_types) - 1
+    for index, principal_type in enumerate(principal_types):
+        try:
+            assignments_client.create(scope, str(uuid.uuid4()), params_model(
+                role_definition_id=role_definition_id,
+                principal_id=object_id,
+                principal_type=principal_type))
+            return principal_type
+        except ResourceExistsError:
+            return principal_type  # already assigned; idempotent
+        except HttpResponseError as ex:
+            message = ex.message or ""
+            code = getattr(getattr(ex, 'error', None), 'code', None)
+            if code == 'RoleAssignmentExists' or 'already exists' in message.lower():
+                return principal_type
+            if 'UnmatchedPrincipalType' in message and index < last:
+                continue  # wrong guess; try the next principal type
+            _warn_assignment_failed(scope, role_id, object_id)
+            return None
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.warning("Skipping role assignment '%s' on %s: %s", role_id, scope, ex)
+            return None
+    return None
+
+
+def _warn_assignment_failed(scope, role_id, object_id):
+    role_name = AIMANAGER_ROLE_NAMES.get(role_id, role_id)
+    logger.warning(
+        "Could not assign '%s' to the caller on %s. This is expected if you are not an Owner or "
+        "User Access Administrator. An administrator can grant it with:\n"
+        "  az role assignment create --assignee-object-id %s --role \"%s\" --scope %s",
+        role_name, scope, object_id, role_name, scope)

src/aimanager/azext_aimanager/constants.py

@@ -8,6 +8,16 @@
 DELETE_POLICY_DELETE = "Delete"
 DELETE_POLICIES = [DELETE_POLICY_KEEP, DELETE_POLICY_DELETE]
 
+# Built-in role definition GUIDs granted to the caller when they create an AI Manager or a
+# namespace, so the creator can immediately manage (ARM) and read (K8S) the resource.
+AIMANAGER_CONTRIBUTOR_ROLE_ID = "413f2675-4911-4010-be3b-c720b43a3c59"  # Azure AIManager Contributor
+AIMANAGER_RBAC_READER_ROLE_ID = "9c77f8a7-b0b9-4462-844c-de6e66add8ba"  # Azure AIManager and namespace RBAC Reader
+AIMANAGER_CALLER_ROLE_IDS = [AIMANAGER_CONTRIBUTOR_ROLE_ID, AIMANAGER_RBAC_READER_ROLE_ID]
+AIMANAGER_ROLE_NAMES = {
+    AIMANAGER_CONTRIBUTOR_ROLE_ID: "Azure AIManager Contributor",
+    AIMANAGER_RBAC_READER_ROLE_ID: "Azure AIManager and namespace RBAC Reader",
+}
+
 MODEL_DEPLOYMENT_PERFORMANCE_MODES = ["Balanced", "Latency", "Throughput"]
 
 # Table output projections for the 'az aimanager model' commands.

src/aimanager/azext_aimanager/custom.py

@@ -6,6 +6,7 @@
 import os
 
 from azure.cli.core.azclierror import ClientRequestError, InvalidArgumentValueError
+from azure.cli.core.commands import LongRunningOperation
 from azure.cli.core.util import sdk_no_wait
 from azure.core import MatchConditions
 from azure.core.exceptions import ResourceNotFoundError
@@ -18,6 +19,8 @@
     parse_key_value_list,
     print_or_merge_credentials,
 )
+from azext_aimanager._roleassignments import assign_caller_roles
+from azext_aimanager.constants import AIMANAGER_CALLER_ROLE_IDS
 
 logger = get_logger(__name__)
 
@@ -30,6 +33,40 @@ def _get_model(cmd, name, operation_group):
     )
 
 
+def _aimanager_scope(cli_ctx, resource_group_name, ai_manager_name):
+    from azure.cli.core.commands.client_factory import get_subscription_id
+    return (
+        f"/subscriptions/{get_subscription_id(cli_ctx)}/resourceGroups/{resource_group_name}"
+        f"/providers/Microsoft.ContainerService/aiManagers/{ai_manager_name}")
+
+
+def _namespace_scope(cli_ctx, resource_group_name, ai_manager_name, namespace_name):
+    return (
+        _aimanager_scope(cli_ctx, resource_group_name, ai_manager_name)
+        + f"/namespaces/{namespace_name}")
+
+
+def _grant_caller_roles_on_success(cmd, poller, no_wait, scope):
+    """Wait for the create to succeed, then grant the caller the built-in roles (best-effort).
+
+    Returns the value the command should return: the completed resource when we waited for it,
+    otherwise the poller. With --no-wait the grant is skipped and the poller is returned, since
+    the command returns before the operation completes and success cannot be confirmed.
+    """
+    if no_wait:
+        logger.warning(
+            "--no-wait was set, so the caller's role assignments on %s were skipped. Re-run "
+            "without --no-wait, or assign the roles manually.", scope)
+        return poller
+    result = LongRunningOperation(cmd.cli_ctx)(poller)  # blocks until Succeeded; raises on failure
+    try:
+        assign_caller_roles(cmd, scope, AIMANAGER_CALLER_ROLE_IDS)
+    except Exception as ex:  # pylint: disable=broad-except
+        # Role assignment is best-effort: never fail a successful create/add because of it.
+        logger.warning("Could not assign the caller's roles on %s: %s", scope, ex)
+    return result
+
+
 # region AI Manager
 
 def _construct_aimanager(cmd, location, tags, delete_policy, identity=None):
@@ -68,14 +105,18 @@ def create_aimanager(cmd,
     headers = get_aks_custom_headers(aks_custom_headers)
     ai_manager = _construct_aimanager(cmd, location, tags, delete_policy)
 
-    return sdk_no_wait(
+    poller = sdk_no_wait(
         no_wait,
         client.begin_create_or_update,
         resource_group_name,
         ai_manager_name,
         ai_manager,
         headers=headers,
     )
+    # Grant the caller the built-in roles once creation succeeds (best-effort).
+    return _grant_caller_roles_on_success(
+        cmd, poller, no_wait,
+        _aimanager_scope(cmd.cli_ctx, resource_group_name, ai_manager_name))
 
 
 # pylint: disable=unused-argument
@@ -213,7 +254,7 @@ def add_aimanager_namespace(cmd,
     namespace_config = _construct_namespace(
         cmd, parse_key_value_list(labels), parse_key_value_list(annotations))
 
-    return sdk_no_wait(
+    poller = sdk_no_wait(
         no_wait,
         client.begin_create_or_update,
         resource_group_name,
@@ -222,6 +263,10 @@ def add_aimanager_namespace(cmd,
         namespace_config,
         headers=headers,
     )
+    # Grant the caller the built-in roles on the namespace once creation succeeds (best-effort).
+    return _grant_caller_roles_on_success(
+        cmd, poller, no_wait,
+        _namespace_scope(cmd.cli_ctx, resource_group_name, ai_manager_name, namespace_name))
 
 
 # pylint: disable=unused-argument

src/aimanager/azext_aimanager/tests/latest/test_aimanager.py

@@ -0,0 +1,113 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
+
+from azext_aimanager import custom
+from azext_aimanager.constants import AIMANAGER_CALLER_ROLE_IDS
+
+SUB_PATCH = "azure.cli.core.commands.client_factory.get_subscription_id"
+
+AIMANAGER_SCOPE = ("/subscriptions/sub/resourceGroups/rg"
+                   "/providers/Microsoft.ContainerService/aiManagers/aim")
+NAMESPACE_SCOPE = AIMANAGER_SCOPE + "/namespaces/team-alpha"
+
+
+class TestCallerRoleWiring(unittest.TestCase):
+
+    def setUp(self):
+        self.cmd = SimpleNamespace(cli_ctx=object())
+        self.client = MagicMock()
+        self.client.get.side_effect = ResourceNotFoundError()  # resource does not already exist
+
+    @patch.object(custom, "LongRunningOperation")
+    @patch(SUB_PATCH, return_value="sub")
+    @patch.object(custom, "assign_caller_roles")
+    @patch.object(custom, "_construct_aimanager", return_value=object())
+    def test_create_assigns_roles_on_aimanager_scope(self, _construct, mock_assign, _sub, mock_lro):
+        mock_lro.return_value = lambda poller: poller  # waiting returns the resource
+
+        custom.create_aimanager(self.cmd, self.client, "rg", "aim", location="eastus2")
+
+        mock_lro.assert_called_once()  # waited for creation to succeed
+        mock_assign.assert_called_once()
+        _cmd, scope, roles = mock_assign.call_args.args
+        self.assertEqual(scope, AIMANAGER_SCOPE)
+        self.assertEqual(roles, AIMANAGER_CALLER_ROLE_IDS)
+
+    @patch.object(custom, "logger")
+    @patch.object(custom, "LongRunningOperation")
+    @patch(SUB_PATCH, return_value="sub")
+    @patch.object(custom, "assign_caller_roles")
+    @patch.object(custom, "_construct_aimanager", return_value=object())
+    def test_create_skips_roles_with_no_wait(self, _construct, mock_assign, _sub, mock_lro, mock_logger):
+        custom.create_aimanager(
+            self.cmd, self.client, "rg", "aim", location="eastus2", no_wait=True)
+
+        mock_assign.assert_not_called()
+        mock_lro.assert_not_called()
+        mock_logger.warning.assert_called_once()  # warns that the grant was skipped under --no-wait
+
+    @patch.object(custom, "LongRunningOperation")
+    @patch(SUB_PATCH, return_value="sub")
+    @patch.object(custom, "assign_caller_roles")
+    @patch.object(custom, "_construct_namespace", return_value=object())
+    def test_namespace_add_assigns_roles_on_namespace_scope(self, _construct, mock_assign, _sub, mock_lro):
+        mock_lro.return_value = lambda poller: poller
+
+        custom.add_aimanager_namespace(self.cmd, self.client, "rg", "aim", "team-alpha")
+
+        mock_lro.assert_called_once()
+        mock_assign.assert_called_once()
+        _cmd, scope, roles = mock_assign.call_args.args
+        self.assertEqual(scope, NAMESPACE_SCOPE)
+        self.assertEqual(roles, AIMANAGER_CALLER_ROLE_IDS)
+
+    @patch.object(custom, "logger")
+    @patch.object(custom, "LongRunningOperation")
+    @patch(SUB_PATCH, return_value="sub")
+    @patch.object(custom, "assign_caller_roles")
+    @patch.object(custom, "_construct_namespace", return_value=object())
+    def test_namespace_add_skips_roles_with_no_wait(self, _construct, mock_assign, _sub, mock_lro, mock_logger):
+        custom.add_aimanager_namespace(
+            self.cmd, self.client, "rg", "aim", "team-alpha", no_wait=True)
+
+        mock_assign.assert_not_called()
+        mock_lro.assert_not_called()
+        mock_logger.warning.assert_called_once()  # warns that the grant was skipped under --no-wait
+
+    @patch.object(custom, "LongRunningOperation")
+    @patch(SUB_PATCH, return_value="sub")
+    @patch.object(custom, "assign_caller_roles", side_effect=RuntimeError("role setup failed"))
+    @patch.object(custom, "_construct_aimanager", return_value=object())
+    def test_create_does_not_fail_when_role_assignment_errors(self, _construct, _assign, _sub, mock_lro):
+        mock_lro.return_value = lambda poller: "created-resource"
+
+        # A successful create must not fail because the (best-effort) role grant blew up.
+        result = custom.create_aimanager(self.cmd, self.client, "rg", "aim", location="eastus2")
+
+        self.assertEqual(result, "created-resource")
+
+    @patch.object(custom, "LongRunningOperation")
+    @patch(SUB_PATCH, return_value="sub")
+    @patch.object(custom, "assign_caller_roles")
+    @patch.object(custom, "_construct_aimanager", return_value=object())
+    def test_create_surfaces_lro_failure_and_skips_grant(self, _construct, mock_assign, _sub, mock_lro):
+        def _raise(_poller):
+            raise HttpResponseError(message="provisioning failed")
+        mock_lro.return_value = _raise
+
+        # A failed create must surface the error and must not grant roles.
+        with self.assertRaises(HttpResponseError):
+            custom.create_aimanager(self.cmd, self.client, "rg", "aim", location="eastus2")
+        mock_assign.assert_not_called()
+
+
+if __name__ == '__main__':
+    unittest.main()