Add max-context-population input to override Azure PowerShell MaxCont…

README.md

@@ -20,6 +20,7 @@
     - [`allow-no-subscriptions`](#allow-no-subscriptions)
     - [`audience`](#audience)
     - [`auth-type`](#auth-type)
+    - [`max-context-population`](#max-context-population)
   - [Workflow Examples](#workflow-examples)
     - [Login With OpenID Connect (OIDC) \[Recommended\]](#login-with-openid-connect-oidc-recommended)
     - [Login With a Service Principal Secret](#login-with-a-service-principal-secret)
@@ -157,6 +158,7 @@ uses: azure/login@<full-length-commit-sha> # v3.0.2
 |allow-no-subscriptions|false|boolean|false|if login without subscription is allowed|
 |audience|false|string|api://AzureADTokenExchange|the audience to get the JWT ID token from GitHub OIDC provider|
 |auth-type|false|string|SERVICE_PRINCIPAL|the auth type|
+|max-context-population|false|integer||only used when `enable-AzPSSession` is `true`; overrides the Azure PowerShell `MaxContextPopulation`. Defaults to the Azure PowerShell default of 25 when unset.|
 |mask-client-id|false|boolean|true|if the `client-id` value is masked in workflow logs|
 
 ### `client-id`
@@ -174,6 +176,8 @@ Refer to [Login With OpenID Connect (OIDC)](#login-with-openid-connect-oidc-reco
 
 ### `mask-client-id`
 
+_Available in `azure/login@v3`._
+
 The input parameter `mask-client-id` controls whether the login client id is registered as a secret and masked in the workflow logs. It defaults to `true`.
 
 Set it to `false` when the client id is not treated as sensitive and masking gets in the way, for example when the same value appears in log output or command results that you need to read.
@@ -268,6 +272,28 @@ The input parameter `auth-type` specifies the type of authentication. The defaul
 
 Refer to [Login With System-assigned Managed Identity](#login-with-system-assigned-managed-identity) and [Login With User-assigned Managed Identity](#login-with-user-assigned-managed-identity) for its usage.
 
+### `max-context-population`
+
+_Available in `azure/login@v3`._
+
+The input parameter `max-context-population` is only used when [`enable-AzPSSession`](#enable-azpssession) is `true`. It overrides the Azure PowerShell `MaxContextPopulation` value that `Connect-AzAccount` uses, which controls how many subscription contexts are loaded into the session.
+
+Azure PowerShell loads a maximum of 25 subscription contexts by default. When the identity has access to more than 25 subscriptions, only a subset is loaded, so commands that enumerate or target subscriptions outside that subset may behave inconsistently. Set `max-context-population` to `-1` to load all subscriptions, or to a positive integer (1 to 2147483647) to load a specific number. When it is unset, the Azure PowerShell default of 25 applies and behavior is unchanged.
+
+```yaml
+  - name: Azure login
+    uses: azure/login@v3
+    with:
+      client-id: ${{ vars.AZURE_CLIENT_ID }}
+      tenant-id: ${{ vars.AZURE_TENANT_ID }}
+      subscription-id: ${{ vars.AZURE_SUBSCRIPTION_ID }}
+      enable-AzPSSession: true
+      max-context-population: -1
+```
+
+> [!NOTE]
+> Loading all subscription contexts with `-1` makes `Connect-AzAccount` slower when the identity can access a large number of subscriptions, because every subscription is enumerated during login. Set it only when your workflow needs contexts beyond the default 25.
+
 ## Workflow Examples
 
 ### Login With OpenID Connect (OIDC) [Recommended]

__tests__/LoginConfig.test.ts

@@ -405,4 +405,35 @@ describe("LoginConfig Test", () => {
         expect(loginConfig.subscriptionId).toBe("");
     });
 
+    async function initWithMaxContextPopulation(value: string): Promise<LoginConfig> {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'true');
+        setEnv('auth-type', 'SERVICE_PRINCIPAL');
+        setEnv('tenant-id', 'tenant-id');
+        setEnv('subscription-id', 'subscription-id');
+        setEnv('client-id', 'client-id');
+        setEnv('max-context-population', value);
+        const loginConfig = new LoginConfig();
+        await loginConfig.initialize();
+        return loginConfig;
+    }
+
+    test.each(['-1', '1', '25', '2147483647'])('validate accepts max-context-population=%s', async (value) => {
+        const loginConfig = await initWithMaxContextPopulation(value);
+        loginConfig.validate();
+        expect(loginConfig.maxContextPopulation).toBe(value);
+    });
+
+    test.each(['0', '-2', '1e3', '0x10', '5.0', '2147483648', 'abc'])('validate rejects invalid max-context-population=%s', async (value) => {
+        const loginConfig = await initWithMaxContextPopulation(value);
+        testValidateWithErrorMessage(loginConfig, "for 'max-context-population'. It must be -1");
+    });
+
+    test('whitespace-only max-context-population is treated as unset', async () => {
+        const loginConfig = await initWithMaxContextPopulation('   ');
+        loginConfig.validate();
+        expect(loginConfig.maxContextPopulation).toBe('');
+    });
+
 });

__tests__/PowerShell/AzPSScriptBuilder.test.ts

@@ -160,6 +160,49 @@ describe("Building the Az PS login invocation", () => {
         });
     });
 
+    test('max-context-population set: value passed as -MaxContextPopulation param', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'true');
+        setEnv('auth-type', 'SERVICE_PRINCIPAL');
+        setEnv('max-context-population', '-1');
+        const creds = {
+            'clientId': 'client-id',
+            'clientSecret': 'client-secret',
+            'tenantId': 'tenant-id',
+            'subscriptionId': 'subscription-id'
+        };
+        setEnv('creds', JSON.stringify(creds));
+
+        const loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSScriptBuilder.getAzPSLoginInvocation(loginConfig).then(({ args }) => {
+            expect(args).toEqual(expect.arrayContaining([
+                '-MaxContextPopulation', '-1',
+            ]));
+        });
+    });
+
+    test('max-context-population unset: -MaxContextPopulation param omitted', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'true');
+        setEnv('auth-type', 'SERVICE_PRINCIPAL');
+        const creds = {
+            'clientId': 'client-id',
+            'clientSecret': 'client-secret',
+            'tenantId': 'tenant-id',
+            'subscriptionId': 'subscription-id'
+        };
+        setEnv('creds', JSON.stringify(creds));
+
+        const loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSScriptBuilder.getAzPSLoginInvocation(loginConfig).then(({ args }) => {
+            expect(args).not.toContain('-MaxContextPopulation');
+        });
+    });
+
     test('SECURITY: adversarial ArmEndpoint travels as a discrete argv element', () => {
         setEnv('environment', 'azurestack');
         setEnv('enable-AzPSSession', 'true');

action.yml

@@ -34,6 +34,9 @@ inputs:
     description: 'The type of authentication. Supported values are SERVICE_PRINCIPAL, IDENTITY. Default value is SERVICE_PRINCIPAL'
     required: false
     default: 'SERVICE_PRINCIPAL'
+  max-context-population:
+    description: 'Only used when enable-AzPSSession is true. Overrides the Azure PowerShell MaxContextPopulation used by Connect-AzAccount (the number of subscription contexts loaded). Set to -1 to load all subscriptions, or a positive integer. When unset, the Azure PowerShell default of 25 applies.'
+    required: false
   mask-client-id:
     description: 'Set this value to false to stop registering the client-id as a secret, so it is not masked in workflow logs'
     required: false

src/PowerShell/AzPSLogin.ps1

@@ -14,7 +14,9 @@ param(
 
     [string]$ApplicationId,
 
-    [string]$ArmEndpoint
+    [string]$ArmEndpoint,
+
+    [int]$MaxContextPopulation
 )
 
 $ErrorActionPreference = 'Stop'
@@ -34,6 +36,7 @@ try {
     }
     if ($Tenant)       { $connectArgs['Tenant']       = $Tenant }
     if ($Subscription) { $connectArgs['Subscription'] = $Subscription }
+    if ($PSBoundParameters.ContainsKey('MaxContextPopulation')) { $connectArgs['MaxContextPopulation'] = $MaxContextPopulation }
 
     if ($AuthType -eq 'SERVICE_PRINCIPAL') {
         $connectArgs['ServicePrincipal'] = $true

src/PowerShell/AzPSScriptBuilder.ts

@@ -53,6 +53,9 @@ export default class AzPSScriptBuilder {
         if (loginConfig.environment.toLowerCase() === 'azurestack') {
             args.push('-ArmEndpoint', loginConfig.resourceManagerEndpointUrl);
         }
+        if (loginConfig.maxContextPopulation) {
+            args.push('-MaxContextPopulation', loginConfig.maxContextPopulation);
+        }
 
         if (loginConfig.authType === LoginConfig.AUTH_TYPE_SERVICE_PRINCIPAL) {
             args.push('-ApplicationId', loginConfig.servicePrincipalId);

src/common/LoginConfig.ts

@@ -25,6 +25,7 @@ export class LoginConfig {
     enableAzPSSession: boolean;
     audience: string;
     federatedToken: string;
+    maxContextPopulation: string;
     maskClientId: boolean;
 
     async initialize() {
@@ -42,6 +43,7 @@ export class LoginConfig {
 
         this.audience = core.getInput('audience', { required: false });
         this.federatedToken = null;
+        this.maxContextPopulation = core.getInput('max-context-population', { required: false }).trim();
 
         this.maskClientId = core.getInput('mask-client-id').toLowerCase() !== "false";
         if (this.maskClientId) {
@@ -110,6 +112,19 @@ export class LoginConfig {
         if (!this.subscriptionId && !this.allowNoSubscriptionsLogin) {
             throw new Error("Ensure 'subscription-id' is supplied or 'allow-no-subscriptions' is 'true'.");
         }
+        if (this.maxContextPopulation) {
+            // Validate the raw string (not Number(), which accepts 1e3/0x10/5.0
+            // and out-of-range values that PowerShell's [int] then rejects).
+            const INT32_MAX = 2147483647;
+            const isValid = this.maxContextPopulation === '-1'
+                || (/^[1-9][0-9]*$/.test(this.maxContextPopulation) && Number(this.maxContextPopulation) <= INT32_MAX);
+            if (!isValid) {
+                throw new Error(`Invalid value '${this.maxContextPopulation}' for 'max-context-population'. It must be -1 (load all subscription contexts) or a positive integer between 1 and ${INT32_MAX}.`);
+            }
+            if (!this.enableAzPSSession) {
+                core.warning("'max-context-population' is only applied when 'enable-AzPSSession' is 'true'. It has no effect on Azure CLI login and will be ignored.");
+            }
+        }
     }
 
     mask(parameterValue: string) {