Skip to content

chaos-mcp: let customers choose user-principal vs managed-identity auth at runtime - #26

Merged
Nikhil Kaul (nikhilkaul1234) merged 4 commits into
mainfrom
user/renzopretto/mcp-managed-identity-auth
Jul 17, 2026
Merged

chaos-mcp: let customers choose user-principal vs managed-identity auth at runtime#26
Nikhil Kaul (nikhilkaul1234) merged 4 commits into
mainfrom
user/renzopretto/mcp-managed-identity-auth

Conversation

@RenzoPrettoMS

@RenzoPrettoMS RenzoPrettoMS commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

What

Lets the customer choose, during their Copilot session, whether the chaos-studio MCP 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:

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() Reports effective {mode, msiClientId, source}.

Precedence: runtime override > CHAOS_MCP_AUTH_MODE env (startup default) > cli. The single choke point _get_token() keeps its signature, so all existing callers/tests are unaffected.

In managed-identity mode 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 existing httpx).

Tests

tests/test_auth_mode.py covers 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 the az user principal).

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
@RenzoPrettoMS RenzoPrettoMS changed the title chaos-mcp: add managed-identity auth lever (MI vs user principal) chaos-mcp: let customers choose user-principal vs managed-identity auth at runtime Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_TRANSPORT seam for httpx but the MI path uses httpx.get directly 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.

Comment thread copilot-cli-plugin/mcp/chaos_mcp/azure.py Outdated
Comment thread copilot-cli-plugin/mcp/chaos_mcp/azure.py
Comment thread copilot-cli-plugin/mcp/README.md
- 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@nikhilkaul1234
Nikhil Kaul (nikhilkaul1234) merged commit 4fc7e0c into main Jul 17, 2026
8 checks passed
@nikhilkaul1234
Nikhil Kaul (nikhilkaul1234) deleted the user/renzopretto/mcp-managed-identity-auth branch July 17, 2026 22:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants