chaos-mcp: let customers choose user-principal vs managed-identity auth at runtime - #26
Conversation
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
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
Nikhil Kaul (nikhilkaul1234)
left a comment
There was a problem hiding this comment.
Nice clean change overall - the single _get_token() choke point is the right seam, and the test coverage is thorough (aliases, both endpoints, precedence, error envelopes, tool count). Left three inline comments: one real bug I'd fix before merging (proxy handling on the IMDS call), one code/docstring mismatch on client-id precedence, and a suggested README note on attribution. Happy to approve once the proxy one is addressed.
Smaller things, fine to leave as-is:
- The module already has the
_TEST_TRANSPORTseam for httpx but the MI path useshttpx.getdirectly and the tests monkeypatch it instead - works, just inconsistent with the file's own convention. - No retry on IMDS transient failures (MS guidance suggests retrying 404/410/429/5xx) and no token caching, so a chatty session hits IMDS per tool call and IMDS does throttle. Fine for now given the existing no-caching design, but worth a hardening item.
- The override is process-global mutable state with no lock; harmless at tool-call cadence, just noting it's deliberate.
- 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
Nikhil Kaul (nikhilkaul1234)
left a comment
There was a problem hiding this comment.
All three addressed in 6ebcf5a - trust_env=False on the identity request (with a regression test asserting it), client-id fallback matching the documented precedence (nice extra coverage on the env-pin cases), and the attribution/auto-approve note in the README. CI green across the matrix. LGTM, thanks for the quick turnaround!
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
Nikhil Kaul (nikhilkaul1234)
left a comment
There was a problem hiding this comment.
Re-approving - the only change since my last review is the MIT license headers on the four touched files (verified the diff is headers only). Still LGTM.
What
Lets the customer choose, during their Copilot session, whether the
chaos-studioMCP tools authenticate as their signed-in user principal (default) or as an Azure Managed Identity — no code/config edit and no server restart.Why
Every tool acquired ARM / Log Analytics tokens via
az account get-access-token(the signed-in user). Customers need to run the tools as a Managed Identity (unattended hosts: CI, containers, AKS, VMs), and to flip between the two on demand.How
Two new MCP tools, backed by an in-memory session-scoped override in
chaos_mcp/azure.py:chaos_set_auth_mode(mode, msi_client_id?)mode=cli(user principal) ormanaged-identity(aliasesmsi/mi);msi_client_idoptionally pins a user-assigned identity.chaos_get_auth_mode(){mode, msiClientId, source}.Precedence: runtime override >
CHAOS_MCP_AUTH_MODEenv (startup default) >cli. The single choke point_get_token()keeps its signature, so all existing callers/tests are unaffected.In
managed-identitymode tokens come from the App Service / Container Apps / Functions identity endpoint (IDENTITY_ENDPOINT+IDENTITY_HEADER) when present, otherwise IMDS (169.254.169.254) for VMs / VMSS / AKS. No new dependencies (uses existinghttpx).Tests
tests/test_auth_mode.pycovers mode selection, IMDS + App Service endpoint, user-assigned client-id pinning, runtime override precedence over env, reset-to-env, invalid-mode rejection, and the tool wrappers. Tool count updated 13 -> 15. Full suite: 33 passed. Default behavior unchanged (still theazuser principal).