From 27825daf26d01002d83f7f47f08f12ceb01af38e Mon Sep 17 00:00:00 2001
From: Renzo Pretto <renzopretto@microsoft.com>
Date: Thu, 16 Jul 2026 13:55:39 -0700
Subject: [PATCH 1/4] Add managed-identity auth lever to chaos-mcp server

The chaos-studio MCP server acquired ARM/Log Analytics tokens exclusively via
the operator's local `az login` session (the user principal). This adds a
lever to source tokens from an Azure Managed Identity instead, so the tools can
run unattended (CI, containers, AKS, VMs) with no interactive az login.

- CHAOS_MCP_AUTH_MODE=managed-identity (aliases msi/mi) switches the token
  source; default stays `cli`.
- CHAOS_MCP_MSI_CLIENT_ID optionally pins a user-assigned identity.
- MI tokens come from the App Service/Container Apps identity endpoint
  (IDENTITY_ENDPOINT + IDENTITY_HEADER) when present, else IMDS.
- _get_token() signature is unchanged, so existing callers/tests are unaffected.
- Adds test_auth_mode.py and documents the lever in the MCP README.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 027a1f0d-6e9a-476a-9f25-04dd0a61916f
---
 copilot-cli-plugin/mcp/README.md              |  42 ++++-
 copilot-cli-plugin/mcp/chaos_mcp/azure.py     | 105 +++++++++++-
 .../mcp/tests/test_auth_mode.py               | 156 ++++++++++++++++++
 3 files changed, 297 insertions(+), 6 deletions(-)
 create mode 100644 copilot-cli-plugin/mcp/tests/test_auth_mode.py

diff --git a/copilot-cli-plugin/mcp/README.md b/copilot-cli-plugin/mcp/README.md
index 5c02f48..a87a00d 100644
--- a/copilot-cli-plugin/mcp/README.md
+++ b/copilot-cli-plugin/mcp/README.md
@@ -3,10 +3,47 @@
 MCP server exposing Azure Chaos Studio v2 operations as agent-callable tools.
 
 The server relies on the user's local `az` CLI session for authentication
-rather than managing tokens itself, keeping the server stateless. See the
+rather than managing tokens itself, keeping the server stateless. It can also
+be pointed at an Azure Managed Identity for unattended runs — see
+[Authentication](#authentication). See the
 [copilot-cli-plugin README](../README.md) for the full tool table and how
 this package fits into the Chaos Studio Copilot CLI plugin.
 
+## Authentication
+
+By default every tool borrows the operator's local `az login` session — the
+**user principal** — via `az account get-access-token`.
+
+For unattended hosts (CI, containers, AKS, VMs) you can flip a lever to use an
+Azure **Managed Identity** instead, so the tools run with no interactive
+`az login`:
+
+| Env var | Values | Effect |
+|---|---|---|
+| `CHAOS_MCP_AUTH_MODE` | `cli` (default), `managed-identity` (aliases: `msi`, `mi`) | Selects the token source. |
+| `CHAOS_MCP_MSI_CLIENT_ID` | client id (optional) | Pins a **user-assigned** identity; omit to use the **system-assigned** identity. |
+
+In `managed-identity` mode the server acquires tokens from the Azure identity
+endpoint — the App Service / Container Apps / Functions endpoint
+(`IDENTITY_ENDPOINT` + `IDENTITY_HEADER`) when present, otherwise IMDS
+(`169.254.169.254`) for VMs / VMSS / AKS. Grant the identity the same ARM roles
+you would grant the user (managing `Microsoft.Chaos` workspaces and querying
+Azure Monitor).
+
+```json
+{
+  "mcpServers": {
+    "chaos-studio": {
+      "command": "chaos-mcp",
+      "env": {
+        "CHAOS_MCP_AUTH_MODE": "managed-identity",
+        "CHAOS_MCP_MSI_CLIENT_ID": "<user-assigned-identity-client-id>"
+      }
+    }
+  }
+}
+```
+
 ## Install
 
 ```bash
@@ -27,7 +64,8 @@ Register the server in your MCP client config (see
 - **Claude Code**: `claude mcp add chaos-studio -- chaos-mcp`
 - **Cursor**: add the block above to `.cursor/mcp.json`.
 
-The server requires an active `az login` session; tools return a structured
+The server requires an active `az login` session (or a managed identity — see
+[Authentication](#authentication)); tools return a structured
 `{"ok": false, "errorType": ...}` envelope (rather than raising) when auth or
 permissions are missing, so agents can remediate and retry.
 
diff --git a/copilot-cli-plugin/mcp/chaos_mcp/azure.py b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
index d78e887..1a507ab 100644
--- a/copilot-cli-plugin/mcp/chaos_mcp/azure.py
+++ b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
@@ -1,12 +1,16 @@
 """Thin wrappers around `az` CLI + ARM REST calls.
 
-The MCP server intentionally relies on the user's local `az` session for auth
+By default the MCP server relies on the operator's local `az` session for auth
 rather than managing tokens itself. This matches the skill's auth model and
-keeps the server stateless.
+keeps the server stateless. Setting ``CHAOS_MCP_AUTH_MODE=managed-identity``
+flips a lever so the server instead acquires tokens from an Azure Managed
+Identity — letting the skills run unattended (CI, containers, AKS, VMs) with
+no interactive `az login`.
 """
 from __future__ import annotations
 
 import json
+import os
 import shutil
 import subprocess
 import time
@@ -25,6 +29,35 @@
 # Production callers leave it None; pytest sets it to an httpx.MockTransport.
 _TEST_TRANSPORT: Any = None
 
+# -----------------------------------------------------------------------------
+# Authentication mode (the "user principal vs managed identity" lever)
+# -----------------------------------------------------------------------------
+# By default the server borrows the operator's local `az login` session (the
+# user principal). Set CHAOS_MCP_AUTH_MODE=managed-identity to instead acquire
+# tokens from an Azure Managed Identity, so the skills can run unattended (CI,
+# containers, AKS, VMs) with no `az` session. CHAOS_MCP_MSI_CLIENT_ID optionally
+# pins a specific user-assigned identity by client id (omit it to use the
+# system-assigned identity).
+AUTH_MODE_ENV = "CHAOS_MCP_AUTH_MODE"
+MSI_CLIENT_ID_ENV = "CHAOS_MCP_MSI_CLIENT_ID"
+
+_MANAGED_IDENTITY_ALIASES = frozenset(
+    {"managed-identity", "managed_identity", "managedidentity", "msi", "mi", "identity"}
+)
+
+# IMDS token endpoint (VMs, VMSS, AKS).
+IMDS_TOKEN_ENDPOINT = "http://169.254.169.254/metadata/identity/oauth2/token"
+IMDS_API_VERSION = "2018-02-01"
+# App Service / Container Apps / Functions inject IDENTITY_ENDPOINT +
+# IDENTITY_HEADER instead of exposing IMDS.
+APP_SERVICE_IDENTITY_API_VERSION = "2019-08-01"
+
+
+def _auth_mode() -> str:
+    """Return the configured auth mode: 'managed-identity' or 'cli' (default)."""
+    raw = (os.environ.get(AUTH_MODE_ENV) or "cli").strip().lower()
+    return "managed-identity" if raw in _MANAGED_IDENTITY_ALIASES else "cli"
+
 
 class AzureError(RuntimeError):
     """Raised when an ARM call or `az` invocation fails."""
@@ -65,12 +98,23 @@ def az_show_account() -> AzContext:
 
 
 def _get_token(resource: str = ARM_ENDPOINT) -> str:
-    """Acquire an access token for the given audience via the local `az` session.
+    """Acquire an access token for the given audience.
 
     `resource` is the token audience URL (e.g., `https://management.azure.com`
     for ARM or `https://api.loganalytics.io` for Log Analytics queries).
-    Tokens are not cached — each call invokes `az account get-access-token`.
+
+    Uses the local `az` session (user principal) by default; when
+    ``CHAOS_MCP_AUTH_MODE`` selects a managed identity, tokens come from the
+    Azure identity endpoint instead. Tokens are not cached — each call acquires
+    a fresh token.
     """
+    if _auth_mode() == "managed-identity":
+        return _get_token_via_managed_identity(resource)
+    return _get_token_via_cli(resource)
+
+
+def _get_token_via_cli(resource: str) -> str:
+    """Acquire a token from the local `az` session (the user principal)."""
     proc = subprocess.run(
         [_az_path(), "account", "get-access-token", "--resource", resource, "-o", "json"],
         capture_output=True,
@@ -83,6 +127,59 @@ def _get_token(resource: str = ARM_ENDPOINT) -> str:
     return json.loads(proc.stdout)["accessToken"]
 
 
+def _get_token_via_managed_identity(resource: str) -> str:
+    """Acquire a token from an Azure Managed Identity (no `az` session needed).
+
+    Honors the App Service / Container Apps / Functions identity endpoint
+    (``IDENTITY_ENDPOINT`` + ``IDENTITY_HEADER``) when present, otherwise falls
+    back to the IMDS endpoint used by VMs, VMSS and AKS. Set
+    ``CHAOS_MCP_MSI_CLIENT_ID`` to target a specific user-assigned identity.
+    """
+    client_id = (os.environ.get(MSI_CLIENT_ID_ENV) or "").strip() or None
+    identity_endpoint = os.environ.get("IDENTITY_ENDPOINT")
+    identity_header = os.environ.get("IDENTITY_HEADER")
+
+    if identity_endpoint and identity_header:
+        url = identity_endpoint
+        params: dict[str, str] = {
+            "resource": resource,
+            "api-version": APP_SERVICE_IDENTITY_API_VERSION,
+        }
+        headers = {"X-IDENTITY-HEADER": identity_header}
+    else:
+        url = IMDS_TOKEN_ENDPOINT
+        params = {"resource": resource, "api-version": IMDS_API_VERSION}
+        headers = {"Metadata": "true"}
+    if client_id:
+        params["client_id"] = client_id
+
+    try:
+        resp = httpx.get(url, params=params, headers=headers, timeout=10.0)
+    except httpx.HTTPError as e:
+        raise AzureError(
+            f"Failed to reach the managed-identity token endpoint for {resource}: {e}. "
+            "Is a managed identity available on this host?"
+        ) from e
+
+    if resp.status_code != 200:
+        raise AzureError(
+            f"Managed-identity token request for {resource} failed with HTTP "
+            f"{resp.status_code}: {resp.text.strip()}"
+        )
+
+    try:
+        token = resp.json().get("access_token")
+    except Exception as e:  # noqa: BLE001
+        raise AzureError(
+            f"Malformed managed-identity token response for {resource}: {resp.text.strip()}"
+        ) from e
+    if not token:
+        raise AzureError(
+            f"Managed-identity token response for {resource} contained no access_token."
+        )
+    return token
+
+
 def az_get_arm_token() -> str:
     """Back-compat shim — acquire a token scoped to ARM."""
     return _get_token(ARM_ENDPOINT)
diff --git a/copilot-cli-plugin/mcp/tests/test_auth_mode.py b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
new file mode 100644
index 0000000..cf20310
--- /dev/null
+++ b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
@@ -0,0 +1,156 @@
+"""Unit tests for the auth-mode lever (`az` user principal vs managed identity).
+
+No network, no `az` shell-outs: `httpx.get` is monkeypatched and the managed
+identity env vars are set per test. Validates that:
+- the default / `cli` mode routes through the `az` CLI helper;
+- `managed-identity` mode acquires tokens from IMDS;
+- the App Service / Container Apps identity endpoint is preferred when present;
+- `CHAOS_MCP_MSI_CLIENT_ID` pins a user-assigned identity;
+- mode aliases (`msi`, `mi`) are honored;
+- error cases surface as `AzureError`.
+"""
+from __future__ import annotations
+
+import httpx
+import pytest
+
+from chaos_mcp import azure as az
+
+
+@pytest.fixture(autouse=True)
+def _clear_auth_env(monkeypatch):
+    """Ensure a clean auth environment for every test."""
+    for var in (
+        az.AUTH_MODE_ENV,
+        az.MSI_CLIENT_ID_ENV,
+        "IDENTITY_ENDPOINT",
+        "IDENTITY_HEADER",
+    ):
+        monkeypatch.delenv(var, raising=False)
+
+
+def _fake_get(recorder: dict):
+    def fake_get(url, params=None, headers=None, timeout=None):
+        recorder["url"] = url
+        recorder["params"] = dict(params or {})
+        recorder["headers"] = dict(headers or {})
+        return httpx.Response(200, json={"access_token": "mi-token", "expires_in": "3600"})
+
+    return fake_get
+
+
+# ---------------------------------------------------------------------------
+# Mode selection
+# ---------------------------------------------------------------------------
+
+
+def test_default_mode_is_cli(monkeypatch):
+    assert az._auth_mode() == "cli"
+    called: dict = {}
+
+    def fake_cli(resource):
+        called["resource"] = resource
+        return "cli-token"
+
+    monkeypatch.setattr(az, "_get_token_via_cli", fake_cli)
+    monkeypatch.setattr(
+        az,
+        "_get_token_via_managed_identity",
+        lambda *_a, **_k: pytest.fail("MI path must not run in cli mode"),
+    )
+    assert az._get_token(az.ARM_ENDPOINT) == "cli-token"
+    assert called["resource"] == az.ARM_ENDPOINT
+
+
+@pytest.mark.parametrize("value", ["managed-identity", "msi", "mi", "MANAGED-IDENTITY", " mi "])
+def test_managed_identity_aliases(monkeypatch, value):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, value)
+    assert az._auth_mode() == "managed-identity"
+
+
+# ---------------------------------------------------------------------------
+# IMDS path (VM / VMSS / AKS)
+# ---------------------------------------------------------------------------
+
+
+def test_managed_identity_uses_imds(monkeypatch):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+    rec: dict = {}
+    monkeypatch.setattr(az.httpx, "get", _fake_get(rec))
+
+    token = az._get_token(az.ARM_ENDPOINT)
+
+    assert token == "mi-token"
+    assert rec["url"] == az.IMDS_TOKEN_ENDPOINT
+    assert rec["headers"]["Metadata"] == "true"
+    assert rec["params"]["resource"] == az.ARM_ENDPOINT
+    assert rec["params"]["api-version"] == az.IMDS_API_VERSION
+    assert "client_id" not in rec["params"]
+
+
+def test_managed_identity_pins_user_assigned_client_id(monkeypatch):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+    monkeypatch.setenv(az.MSI_CLIENT_ID_ENV, "abc-123")
+    rec: dict = {}
+    monkeypatch.setattr(az.httpx, "get", _fake_get(rec))
+
+    az._get_token(az.LOG_ANALYTICS_ENDPOINT)
+
+    assert rec["params"]["client_id"] == "abc-123"
+    assert rec["params"]["resource"] == az.LOG_ANALYTICS_ENDPOINT
+
+
+# ---------------------------------------------------------------------------
+# App Service / Container Apps identity endpoint takes precedence
+# ---------------------------------------------------------------------------
+
+
+def test_managed_identity_prefers_app_service_endpoint(monkeypatch):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+    monkeypatch.setenv("IDENTITY_ENDPOINT", "http://localhost:42/token")
+    monkeypatch.setenv("IDENTITY_HEADER", "secret-header")
+    rec: dict = {}
+    monkeypatch.setattr(az.httpx, "get", _fake_get(rec))
+
+    az._get_token(az.ARM_ENDPOINT)
+
+    assert rec["url"] == "http://localhost:42/token"
+    assert rec["headers"]["X-IDENTITY-HEADER"] == "secret-header"
+    assert "Metadata" not in rec["headers"]
+    assert rec["params"]["api-version"] == az.APP_SERVICE_IDENTITY_API_VERSION
+
+
+# ---------------------------------------------------------------------------
+# Error handling
+# ---------------------------------------------------------------------------
+
+
+def test_managed_identity_non_200_raises(monkeypatch):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+    monkeypatch.setattr(
+        az.httpx,
+        "get",
+        lambda *a, **k: httpx.Response(400, text="no identity assigned"),
+    )
+    with pytest.raises(az.AzureError, match="HTTP 400"):
+        az._get_token(az.ARM_ENDPOINT)
+
+
+def test_managed_identity_transport_error_raises(monkeypatch):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+
+    def boom(*_a, **_k):
+        raise httpx.ConnectError("no route to IMDS")
+
+    monkeypatch.setattr(az.httpx, "get", boom)
+    with pytest.raises(az.AzureError, match="managed-identity token endpoint"):
+        az._get_token(az.ARM_ENDPOINT)
+
+
+def test_managed_identity_missing_token_raises(monkeypatch):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+    monkeypatch.setattr(
+        az.httpx, "get", lambda *a, **k: httpx.Response(200, json={"expires_in": "3600"})
+    )
+    with pytest.raises(az.AzureError, match="no access_token"):
+        az._get_token(az.ARM_ENDPOINT)

From 1fb9d957edd425556c89ae3e0b97b18b945496fa Mon Sep 17 00:00:00 2001
From: Renzo Pretto <renzopretto@microsoft.com>
Date: Thu, 16 Jul 2026 14:05:47 -0700
Subject: [PATCH 2/4] Make auth mode runtime-selectable via chaos_set_auth_mode
 tool

Customers should be able to choose user principal vs managed identity during
their Copilot session, not by editing config/env and restarting the server.

- Add chaos_set_auth_mode(mode, msi_client_id) and chaos_get_auth_mode() MCP
  tools. The choice is an in-memory, session-scoped override applied to every
  subsequent tool call.
- Precedence: runtime override > CHAOS_MCP_AUTH_MODE env > 'cli' default.
- azure.py gains set_auth_mode/reset_auth_mode/get_auth_config plus mode/
  client-id resolution helpers; _get_token* now read the effective values.
- Tests cover override precedence, reset-to-env, invalid mode, and the tool
  wrappers (33 passed). Tool count 13 -> 15.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 027a1f0d-6e9a-476a-9f25-04dd0a61916f
---
 copilot-cli-plugin/mcp/README.md              |  34 +++---
 copilot-cli-plugin/mcp/chaos_mcp/azure.py     | 109 ++++++++++++++++--
 copilot-cli-plugin/mcp/chaos_mcp/server.py    |  40 +++++++
 .../mcp/tests/test_auth_mode.py               | 103 ++++++++++++++++-
 .../mcp/tests/test_monitor_tools.py           |   8 +-
 5 files changed, 262 insertions(+), 32 deletions(-)

diff --git a/copilot-cli-plugin/mcp/README.md b/copilot-cli-plugin/mcp/README.md
index a87a00d..e4ae546 100644
--- a/copilot-cli-plugin/mcp/README.md
+++ b/copilot-cli-plugin/mcp/README.md
@@ -14,9 +14,23 @@ this package fits into the Chaos Studio Copilot CLI plugin.
 By default every tool borrows the operator's local `az login` session — the
 **user principal** — via `az account get-access-token`.
 
-For unattended hosts (CI, containers, AKS, VMs) you can flip a lever to use an
-Azure **Managed Identity** instead, so the tools run with no interactive
-`az login`:
+### Choose the mode during a session (no config change)
+
+The customer can switch between the user principal and a **Managed Identity**
+mid-conversation by asking the agent to call the `chaos_set_auth_mode` tool —
+no config edit or server restart needed. The choice applies to every subsequent
+tool call for the life of the session.
+
+| Tool | Purpose |
+|---|---|
+| `chaos_set_auth_mode(mode, msi_client_id?)` | `mode` = `cli` (user principal) or `managed-identity` (aliases `msi`/`mi`); `msi_client_id` optionally pins a user-assigned identity. |
+| `chaos_get_auth_mode()` | Report the effective `{mode, msiClientId, source}`. |
+
+### Startup default (optional env vars)
+
+For unattended hosts (CI, containers, AKS, VMs) you can also set the initial
+mode via env vars so no interactive `az login` is required. A runtime
+`chaos_set_auth_mode` call always takes precedence over these.
 
 | Env var | Values | Effect |
 |---|---|---|
@@ -30,20 +44,6 @@ endpoint — the App Service / Container Apps / Functions endpoint
 you would grant the user (managing `Microsoft.Chaos` workspaces and querying
 Azure Monitor).
 
-```json
-{
-  "mcpServers": {
-    "chaos-studio": {
-      "command": "chaos-mcp",
-      "env": {
-        "CHAOS_MCP_AUTH_MODE": "managed-identity",
-        "CHAOS_MCP_MSI_CLIENT_ID": "<user-assigned-identity-client-id>"
-      }
-    }
-  }
-}
-```
-
 ## Install
 
 ```bash
diff --git a/copilot-cli-plugin/mcp/chaos_mcp/azure.py b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
index 1a507ab..e4358cd 100644
--- a/copilot-cli-plugin/mcp/chaos_mcp/azure.py
+++ b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
@@ -32,18 +32,28 @@
 # -----------------------------------------------------------------------------
 # Authentication mode (the "user principal vs managed identity" lever)
 # -----------------------------------------------------------------------------
-# By default the server borrows the operator's local `az login` session (the
-# user principal). Set CHAOS_MCP_AUTH_MODE=managed-identity to instead acquire
-# tokens from an Azure Managed Identity, so the skills can run unattended (CI,
-# containers, AKS, VMs) with no `az` session. CHAOS_MCP_MSI_CLIENT_ID optionally
-# pins a specific user-assigned identity by client id (omit it to use the
-# system-assigned identity).
+# The server can authenticate either as the operator's local `az login` session
+# (the user principal, default) or as an Azure Managed Identity — so the tools
+# can run unattended (CI, containers, AKS, VMs) with no `az` session.
+#
+# The mode is chosen at three levels, highest precedence first:
+#   1. A runtime override set *during the session* via `set_auth_mode(...)`
+#      (surfaced to agents as the `chaos_set_auth_mode` MCP tool). This lets the
+#      customer flip between user principal and MI mid-conversation, no config
+#      edit or server restart required.
+#   2. The CHAOS_MCP_AUTH_MODE / CHAOS_MCP_MSI_CLIENT_ID env vars (a startup
+#      default, e.g. for headless deployments).
+#   3. The built-in default: `cli` (user principal).
 AUTH_MODE_ENV = "CHAOS_MCP_AUTH_MODE"
 MSI_CLIENT_ID_ENV = "CHAOS_MCP_MSI_CLIENT_ID"
 
+AUTH_MODE_CLI = "cli"
+AUTH_MODE_MANAGED_IDENTITY = "managed-identity"
+
 _MANAGED_IDENTITY_ALIASES = frozenset(
     {"managed-identity", "managed_identity", "managedidentity", "msi", "mi", "identity"}
 )
+_CLI_ALIASES = frozenset({"cli", "az", "user", "user-principal", "userprincipal", "default"})
 
 # IMDS token endpoint (VMs, VMSS, AKS).
 IMDS_TOKEN_ENDPOINT = "http://169.254.169.254/metadata/identity/oauth2/token"
@@ -52,11 +62,87 @@
 # IDENTITY_HEADER instead of exposing IMDS.
 APP_SERVICE_IDENTITY_API_VERSION = "2019-08-01"
 
+# Session-scoped runtime override (set via the chaos_set_auth_mode tool). None
+# means "not overridden — fall back to env / default". This is process-global
+# in-memory state: it lives for the life of the MCP server process (i.e. the
+# customer's session) and is never persisted to disk.
+_auth_mode_override: str | None = None
+_msi_client_id_override: str | None = None
+
+
+def _normalize_mode(value: str) -> str:
+    """Map a user-supplied mode string to a canonical mode, or raise.
+
+    Accepts the managed-identity aliases (managed-identity/msi/mi/...) and the
+    cli aliases (cli/az/user-principal/...).
+    """
+    raw = (value or "").strip().lower()
+    if raw in _MANAGED_IDENTITY_ALIASES:
+        return AUTH_MODE_MANAGED_IDENTITY
+    if raw in _CLI_ALIASES:
+        return AUTH_MODE_CLI
+    raise AzureError(
+        f"Unknown auth mode '{value}'. Use 'cli' (user principal) or "
+        "'managed-identity'."
+    )
+
+
+def set_auth_mode(mode: str, msi_client_id: str | None = None) -> dict:
+    """Set the session-scoped auth mode override and return the effective config.
+
+    `mode` is 'cli' (user principal) or 'managed-identity' (aliases accepted).
+    `msi_client_id` optionally pins a user-assigned identity when mode is
+    managed-identity; pass None/empty to use the system-assigned identity.
+    """
+    global _auth_mode_override, _msi_client_id_override
+    normalized = _normalize_mode(mode)
+    _auth_mode_override = normalized
+    if normalized == AUTH_MODE_MANAGED_IDENTITY:
+        _msi_client_id_override = (msi_client_id or "").strip() or None
+    else:
+        _msi_client_id_override = None
+    return get_auth_config()
+
+
+def reset_auth_mode() -> dict:
+    """Clear the runtime override so mode falls back to env vars / default."""
+    global _auth_mode_override, _msi_client_id_override
+    _auth_mode_override = None
+    _msi_client_id_override = None
+    return get_auth_config()
+
+
+def get_auth_config() -> dict:
+    """Return the effective auth configuration and where it came from."""
+    source = "override" if _auth_mode_override is not None else (
+        "env" if os.environ.get(AUTH_MODE_ENV) else "default"
+    )
+    return {
+        "mode": _auth_mode(),
+        "msiClientId": _msi_client_id(),
+        "source": source,
+    }
+
 
 def _auth_mode() -> str:
-    """Return the configured auth mode: 'managed-identity' or 'cli' (default)."""
+    """Return the effective auth mode: 'managed-identity' or 'cli' (default).
+
+    Precedence: runtime override > CHAOS_MCP_AUTH_MODE env var > 'cli'.
+    """
+    if _auth_mode_override is not None:
+        return _auth_mode_override
     raw = (os.environ.get(AUTH_MODE_ENV) or "cli").strip().lower()
-    return "managed-identity" if raw in _MANAGED_IDENTITY_ALIASES else "cli"
+    return AUTH_MODE_MANAGED_IDENTITY if raw in _MANAGED_IDENTITY_ALIASES else AUTH_MODE_CLI
+
+
+def _msi_client_id() -> str | None:
+    """Return the effective user-assigned MI client id, or None.
+
+    Precedence: runtime override > CHAOS_MCP_MSI_CLIENT_ID env var > None.
+    """
+    if _auth_mode_override is not None:
+        return _msi_client_id_override
+    return (os.environ.get(MSI_CLIENT_ID_ENV) or "").strip() or None
 
 
 class AzureError(RuntimeError):
@@ -132,10 +218,11 @@ def _get_token_via_managed_identity(resource: str) -> str:
 
     Honors the App Service / Container Apps / Functions identity endpoint
     (``IDENTITY_ENDPOINT`` + ``IDENTITY_HEADER``) when present, otherwise falls
-    back to the IMDS endpoint used by VMs, VMSS and AKS. Set
-    ``CHAOS_MCP_MSI_CLIENT_ID`` to target a specific user-assigned identity.
+    back to the IMDS endpoint used by VMs, VMSS and AKS. The user-assigned
+    identity (if any) comes from the runtime override or
+    ``CHAOS_MCP_MSI_CLIENT_ID``.
     """
-    client_id = (os.environ.get(MSI_CLIENT_ID_ENV) or "").strip() or None
+    client_id = _msi_client_id()
     identity_endpoint = os.environ.get("IDENTITY_ENDPOINT")
     identity_header = os.environ.get("IDENTITY_HEADER")
 
diff --git a/copilot-cli-plugin/mcp/chaos_mcp/server.py b/copilot-cli-plugin/mcp/chaos_mcp/server.py
index 005032f..1ceec7d 100644
--- a/copilot-cli-plugin/mcp/chaos_mcp/server.py
+++ b/copilot-cli-plugin/mcp/chaos_mcp/server.py
@@ -81,6 +81,46 @@ def _grant_reader(scope: str, principal_id: str) -> dict[str, Any]:
 # ---------------------------------------------------------------------------
 
 
+@mcp.tool()
+def chaos_set_auth_mode(
+    mode: str,
+    msi_client_id: str | None = None,
+) -> dict[str, Any]:
+    """Choose how the Chaos Studio tools authenticate to Azure, for the rest of
+    this session.
+
+    Call this when the customer wants the tools to act as an Azure **Managed
+    Identity** instead of their signed-in **user principal** (the default), or
+    to switch back. The choice is applied immediately to every subsequent tool
+    call and persists for the life of this MCP session — no config edit or
+    server restart needed.
+
+    Args:
+        mode: 'cli' to use the local `az login` session (the user principal), or
+            'managed-identity' to use an Azure Managed Identity (aliases 'msi',
+            'mi' are accepted).
+        msi_client_id: Optional client id of a user-assigned managed identity to
+            use when mode is 'managed-identity'. Omit to use the system-assigned
+            identity. Ignored in 'cli' mode.
+
+    Returns the effective auth configuration ({mode, msiClientId, source}).
+    """
+    try:
+        return _ok(az.set_auth_mode(mode, msi_client_id))
+    except az.AzureError as e:
+        return _err(e)
+
+
+@mcp.tool()
+def chaos_get_auth_mode() -> dict[str, Any]:
+    """Report how the Chaos Studio tools are currently authenticating.
+
+    Returns {mode, msiClientId, source} where `source` is 'override' (set this
+    session via chaos_set_auth_mode), 'env' (from CHAOS_MCP_AUTH_MODE), or
+    'default' (the built-in `cli` / user-principal default)."""
+    return _ok(az.get_auth_config())
+
+
 @mcp.tool()
 def chaos_create_workspace(
     subscription_id: str,
diff --git a/copilot-cli-plugin/mcp/tests/test_auth_mode.py b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
index cf20310..a45dccf 100644
--- a/copilot-cli-plugin/mcp/tests/test_auth_mode.py
+++ b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
@@ -19,7 +19,11 @@
 
 @pytest.fixture(autouse=True)
 def _clear_auth_env(monkeypatch):
-    """Ensure a clean auth environment for every test."""
+    """Ensure a clean auth environment (env vars AND runtime override) per test.
+
+    The runtime override is process-global in-memory state, so it MUST be reset
+    around every test or one test's `set_auth_mode` would poison the next.
+    """
     for var in (
         az.AUTH_MODE_ENV,
         az.MSI_CLIENT_ID_ENV,
@@ -27,6 +31,9 @@ def _clear_auth_env(monkeypatch):
         "IDENTITY_HEADER",
     ):
         monkeypatch.delenv(var, raising=False)
+    az.reset_auth_mode()
+    yield
+    az.reset_auth_mode()
 
 
 def _fake_get(recorder: dict):
@@ -154,3 +161,97 @@ def test_managed_identity_missing_token_raises(monkeypatch):
     )
     with pytest.raises(az.AzureError, match="no access_token"):
         az._get_token(az.ARM_ENDPOINT)
+
+
+# ---------------------------------------------------------------------------
+# Runtime override — the customer chooses mid-session (chaos_set_auth_mode)
+# ---------------------------------------------------------------------------
+
+
+def test_runtime_override_switches_to_managed_identity():
+    assert az._auth_mode() == "cli"
+    cfg = az.set_auth_mode("managed-identity", "uami-client-id")
+    assert cfg == {
+        "mode": "managed-identity",
+        "msiClientId": "uami-client-id",
+        "source": "override",
+    }
+    assert az._auth_mode() == "managed-identity"
+    assert az._msi_client_id() == "uami-client-id"
+
+
+def test_runtime_override_beats_env(monkeypatch):
+    # Env selects MI, but the customer overrides back to the user principal.
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+    monkeypatch.setenv(az.MSI_CLIENT_ID_ENV, "from-env")
+    assert az._auth_mode() == "managed-identity"
+
+    az.set_auth_mode("cli")
+    assert az._auth_mode() == "cli"
+    # cli mode drops any user-assigned client id.
+    assert az._msi_client_id() is None
+    assert az.get_auth_config()["source"] == "override"
+
+
+def test_runtime_reset_falls_back_to_env(monkeypatch):
+    monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
+    az.set_auth_mode("cli")
+    assert az._auth_mode() == "cli"
+    az.reset_auth_mode()
+    assert az._auth_mode() == "managed-identity"
+    assert az.get_auth_config()["source"] == "env"
+
+
+def test_set_auth_mode_rejects_unknown():
+    with pytest.raises(az.AzureError, match="Unknown auth mode"):
+        az.set_auth_mode("kerberos")
+
+
+def test_get_auth_config_default_source():
+    cfg = az.get_auth_config()
+    assert cfg == {"mode": "cli", "msiClientId": None, "source": "default"}
+
+
+def test_override_drives_actual_token_acquisition(monkeypatch):
+    """Flipping mode at runtime must route _get_token through the MI path."""
+    rec: dict = {}
+    monkeypatch.setattr(az.httpx, "get", _fake_get(rec))
+    monkeypatch.setattr(
+        az, "_get_token_via_cli", lambda *_a, **_k: pytest.fail("should use MI after override")
+    )
+    az.set_auth_mode("mi", "runtime-uami")
+    assert az._get_token(az.ARM_ENDPOINT) == "mi-token"
+    assert rec["params"]["client_id"] == "runtime-uami"
+
+
+# ---------------------------------------------------------------------------
+# MCP tool wrappers (server.chaos_set_auth_mode / chaos_get_auth_mode)
+# ---------------------------------------------------------------------------
+
+
+def test_tool_set_and_get_auth_mode():
+    from chaos_mcp import server as srv
+
+    set_result = srv.chaos_set_auth_mode("managed-identity", "tool-uami")
+    assert set_result["ok"] is True
+    assert set_result["result"]["mode"] == "managed-identity"
+    assert set_result["result"]["msiClientId"] == "tool-uami"
+
+    get_result = srv.chaos_get_auth_mode()
+    assert get_result == {
+        "ok": True,
+        "result": {
+            "mode": "managed-identity",
+            "msiClientId": "tool-uami",
+            "source": "override",
+        },
+    }
+
+
+def test_tool_set_auth_mode_invalid_returns_error_envelope():
+    from chaos_mcp import server as srv
+
+    result = srv.chaos_set_auth_mode("nope")
+    assert result["ok"] is False
+    assert result["errorType"] == "AzureError"
+    assert "Unknown auth mode" in result["error"]
diff --git a/copilot-cli-plugin/mcp/tests/test_monitor_tools.py b/copilot-cli-plugin/mcp/tests/test_monitor_tools.py
index d3eafa3..6f2a0e2 100644
--- a/copilot-cli-plugin/mcp/tests/test_monitor_tools.py
+++ b/copilot-cli-plugin/mcp/tests/test_monitor_tools.py
@@ -327,8 +327,8 @@ def handler(req: httpx.Request) -> httpx.Response:
 # ---------------------------------------------------------------------------
 
 
-def test_server_lists_thirteen_tools():
-    """Importing server.py should register all 13 tools on the FastMCP instance."""
+def test_server_lists_all_tools():
+    """Importing server.py should register all 15 tools on the FastMCP instance."""
     import asyncio
     from chaos_mcp import server as srv
 
@@ -337,4 +337,6 @@ def test_server_lists_thirteen_tools():
     assert "monitor_query_metrics" in names
     assert "monitor_query_logs" in names
     assert "monitor_search_activity_log" in names
-    assert len(names) == 13, f"expected 13 tools, got {len(names)}: {sorted(names)}"
+    assert "chaos_set_auth_mode" in names
+    assert "chaos_get_auth_mode" in names
+    assert len(names) == 15, f"expected 15 tools, got {len(names)}: {sorted(names)}"

From 6ebcf5a19455a7f88042dffe7418209e7b6c813f Mon Sep 17 00:00:00 2001
From: Renzo Pretto <renzopretto@microsoft.com>
Date: Thu, 16 Jul 2026 15:47:33 -0700
Subject: [PATCH 3/4] Address PR review: proxy-safe IMDS, client-id env
 fallback, audit note

- azure.py: pass trust_env=False on the managed-identity token request so IMDS
  (169.254.169.254) is never routed through HTTP(S)_PROXY and the App Service
  X-IDENTITY-HEADER secret can't leak through a proxy.
- azure.py: _msi_client_id() now falls back to CHAOS_MCP_MSI_CLIENT_ID when a
  managed-identity switch doesn't name a client id, instead of silently dropping
  the env-pinned identity; returns None outside MI mode. Docstring/tool text
  updated to match.
- README: note that MI-mode actions are attributed to the identity (audit) and
  that chaos_set_auth_mode should not be blanket-auto-approved.
- Tests: assert trust_env=False; cover env-pin fallback, explicit override, and
  cli-mode client-id suppression (36 passed).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 027a1f0d-6e9a-476a-9f25-04dd0a61916f
---
 copilot-cli-plugin/mcp/README.md              |  2 ++
 copilot-cli-plugin/mcp/chaos_mcp/azure.py     | 16 ++++++++---
 copilot-cli-plugin/mcp/chaos_mcp/server.py    |  5 ++--
 .../mcp/tests/test_auth_mode.py               | 27 ++++++++++++++++++-
 4 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/copilot-cli-plugin/mcp/README.md b/copilot-cli-plugin/mcp/README.md
index e4ae546..6d304bd 100644
--- a/copilot-cli-plugin/mcp/README.md
+++ b/copilot-cli-plugin/mcp/README.md
@@ -26,6 +26,8 @@ tool call for the life of the session.
 | `chaos_set_auth_mode(mode, msi_client_id?)` | `mode` = `cli` (user principal) or `managed-identity` (aliases `msi`/`mi`); `msi_client_id` optionally pins a user-assigned identity. |
 | `chaos_get_auth_mode()` | Report the effective `{mode, msiClientId, source}`. |
 
+> **Attribution & approval.** In managed-identity mode every ARM action — including the Reader role assignments made during workspace creation — is recorded in the Azure activity log as the **identity**, not the human operator. Review your audit/attribution requirements before switching. Because `chaos_set_auth_mode` is agent-callable, your MCP client's tool-approval prompt is the control that stops a session from being silently flipped onto a privileged host identity — do **not** blanket-auto-approve this tool.
+
 ### Startup default (optional env vars)
 
 For unattended hosts (CI, containers, AKS, VMs) you can also set the initial
diff --git a/copilot-cli-plugin/mcp/chaos_mcp/azure.py b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
index e4358cd..de21f0c 100644
--- a/copilot-cli-plugin/mcp/chaos_mcp/azure.py
+++ b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
@@ -138,9 +138,15 @@ def _auth_mode() -> str:
 def _msi_client_id() -> str | None:
     """Return the effective user-assigned MI client id, or None.
 
-    Precedence: runtime override > CHAOS_MCP_MSI_CLIENT_ID env var > None.
+    Only meaningful in managed-identity mode (returns None otherwise). Within MI
+    mode the precedence is: explicit runtime override > CHAOS_MCP_MSI_CLIENT_ID
+    env var > None (the system-assigned identity). A mode switch that doesn't
+    name a client id therefore keeps any env-pinned identity rather than
+    silently dropping it.
     """
-    if _auth_mode_override is not None:
+    if _auth_mode() != AUTH_MODE_MANAGED_IDENTITY:
+        return None
+    if _msi_client_id_override is not None:
         return _msi_client_id_override
     return (os.environ.get(MSI_CLIENT_ID_ENV) or "").strip() or None
 
@@ -241,7 +247,11 @@ def _get_token_via_managed_identity(resource: str) -> str:
         params["client_id"] = client_id
 
     try:
-        resp = httpx.get(url, params=params, headers=headers, timeout=10.0)
+        # trust_env=False: never route the IMDS / identity-endpoint request
+        # through HTTP(S)_PROXY. 169.254.169.254 must be reached directly, and
+        # proxying would also leak the App Service X-IDENTITY-HEADER secret. The
+        # Azure SDKs bypass proxies for IMDS for the same reasons.
+        resp = httpx.get(url, params=params, headers=headers, timeout=10.0, trust_env=False)
     except httpx.HTTPError as e:
         raise AzureError(
             f"Failed to reach the managed-identity token endpoint for {resource}: {e}. "
diff --git a/copilot-cli-plugin/mcp/chaos_mcp/server.py b/copilot-cli-plugin/mcp/chaos_mcp/server.py
index 1ceec7d..83f9e53 100644
--- a/copilot-cli-plugin/mcp/chaos_mcp/server.py
+++ b/copilot-cli-plugin/mcp/chaos_mcp/server.py
@@ -100,8 +100,9 @@ def chaos_set_auth_mode(
             'managed-identity' to use an Azure Managed Identity (aliases 'msi',
             'mi' are accepted).
         msi_client_id: Optional client id of a user-assigned managed identity to
-            use when mode is 'managed-identity'. Omit to use the system-assigned
-            identity. Ignored in 'cli' mode.
+            use when mode is 'managed-identity'. Omit to keep any identity pinned
+            by CHAOS_MCP_MSI_CLIENT_ID, or to fall back to the system-assigned
+            identity when none is pinned. Ignored in 'cli' mode.
 
     Returns the effective auth configuration ({mode, msiClientId, source}).
     """
diff --git a/copilot-cli-plugin/mcp/tests/test_auth_mode.py b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
index a45dccf..474b93c 100644
--- a/copilot-cli-plugin/mcp/tests/test_auth_mode.py
+++ b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
@@ -37,10 +37,11 @@ def _clear_auth_env(monkeypatch):
 
 
 def _fake_get(recorder: dict):
-    def fake_get(url, params=None, headers=None, timeout=None):
+    def fake_get(url, params=None, headers=None, timeout=None, **kwargs):
         recorder["url"] = url
         recorder["params"] = dict(params or {})
         recorder["headers"] = dict(headers or {})
+        recorder["kwargs"] = dict(kwargs)
         return httpx.Response(200, json={"access_token": "mi-token", "expires_in": "3600"})
 
     return fake_get
@@ -93,6 +94,8 @@ def test_managed_identity_uses_imds(monkeypatch):
     assert rec["params"]["resource"] == az.ARM_ENDPOINT
     assert rec["params"]["api-version"] == az.IMDS_API_VERSION
     assert "client_id" not in rec["params"]
+    # Must bypass any HTTP(S)_PROXY when talking to IMDS.
+    assert rec["kwargs"]["trust_env"] is False
 
 
 def test_managed_identity_pins_user_assigned_client_id(monkeypatch):
@@ -193,6 +196,28 @@ def test_runtime_override_beats_env(monkeypatch):
     assert az.get_auth_config()["source"] == "override"
 
 
+def test_runtime_switch_to_mi_without_client_id_keeps_env_pin(monkeypatch):
+    """Switching to MI without naming a client id must NOT drop an env pin."""
+    monkeypatch.setenv(az.MSI_CLIENT_ID_ENV, "env-pinned-uami")
+    az.set_auth_mode("managed-identity")  # no client id supplied
+    assert az._auth_mode() == "managed-identity"
+    assert az._msi_client_id() == "env-pinned-uami"
+    assert az.get_auth_config()["msiClientId"] == "env-pinned-uami"
+
+
+def test_runtime_switch_to_mi_explicit_client_id_overrides_env(monkeypatch):
+    monkeypatch.setenv(az.MSI_CLIENT_ID_ENV, "env-pinned-uami")
+    az.set_auth_mode("managed-identity", "explicit-uami")
+    assert az._msi_client_id() == "explicit-uami"
+
+
+def test_cli_mode_reports_no_client_id_even_with_env(monkeypatch):
+    monkeypatch.setenv(az.MSI_CLIENT_ID_ENV, "env-pinned-uami")
+    assert az._auth_mode() == "cli"
+    assert az._msi_client_id() is None
+    assert az.get_auth_config()["msiClientId"] is None
+
+
 def test_runtime_reset_falls_back_to_env(monkeypatch):
     monkeypatch.setenv(az.AUTH_MODE_ENV, "managed-identity")
     az.set_auth_mode("cli")

From 18fcc346f778c6896c25fc1ccd373b2ef16af9f6 Mon Sep 17 00:00:00 2001
From: Renzo Pretto <renzopretto@microsoft.com>
Date: Thu, 16 Jul 2026 16:02:54 -0700
Subject: [PATCH 4/4] Add Microsoft MIT license headers to changed source files

Prepend the standard Microsoft copyright/MIT license header to the Python
files touched by this PR (azure.py, server.py, and the two test modules).
Headers are leading comments so module docstrings remain intact; 36 tests pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 027a1f0d-6e9a-476a-9f25-04dd0a61916f
---
 copilot-cli-plugin/mcp/chaos_mcp/azure.py          | 3 +++
 copilot-cli-plugin/mcp/chaos_mcp/server.py         | 3 +++
 copilot-cli-plugin/mcp/tests/test_auth_mode.py     | 3 +++
 copilot-cli-plugin/mcp/tests/test_monitor_tools.py | 3 +++
 4 files changed, 12 insertions(+)

diff --git a/copilot-cli-plugin/mcp/chaos_mcp/azure.py b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
index de21f0c..069024b 100644
--- a/copilot-cli-plugin/mcp/chaos_mcp/azure.py
+++ b/copilot-cli-plugin/mcp/chaos_mcp/azure.py
@@ -1,3 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 """Thin wrappers around `az` CLI + ARM REST calls.
 
 By default the MCP server relies on the operator's local `az` session for auth
diff --git a/copilot-cli-plugin/mcp/chaos_mcp/server.py b/copilot-cli-plugin/mcp/chaos_mcp/server.py
index 83f9e53..01ddc7a 100644
--- a/copilot-cli-plugin/mcp/chaos_mcp/server.py
+++ b/copilot-cli-plugin/mcp/chaos_mcp/server.py
@@ -1,3 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 """MCP server exposing Azure Chaos Studio v2 operations as agent-callable tools.
 
 Mirrors the Copilot CLI plugin's PowerShell skills (create-workspace,
diff --git a/copilot-cli-plugin/mcp/tests/test_auth_mode.py b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
index 474b93c..a4c7f2e 100644
--- a/copilot-cli-plugin/mcp/tests/test_auth_mode.py
+++ b/copilot-cli-plugin/mcp/tests/test_auth_mode.py
@@ -1,3 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 """Unit tests for the auth-mode lever (`az` user principal vs managed identity).
 
 No network, no `az` shell-outs: `httpx.get` is monkeypatched and the managed
diff --git a/copilot-cli-plugin/mcp/tests/test_monitor_tools.py b/copilot-cli-plugin/mcp/tests/test_monitor_tools.py
index 6f2a0e2..c28d6c4 100644
--- a/copilot-cli-plugin/mcp/tests/test_monitor_tools.py
+++ b/copilot-cli-plugin/mcp/tests/test_monitor_tools.py
@@ -1,3 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 """Unit tests for the three Azure Monitor MCP tools.
 
 Uses httpx.MockTransport — no network, no `az` shell-outs (token acquisition