-
Notifications
You must be signed in to change notification settings - Fork 3.4k
[Core] Differentiate Copilot agent requests from manual requests by adding session ID into token claims #33309
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Draft
xuming-ms
wants to merge
8
commits into
Azure:dev
Choose a base branch
from
xuming-ms:ming/agentic-usage-differenciation
base: dev
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Draft
Changes from all commits
Commits
Show all changes
8 commits
Select commit
Hold shift + click to select a range
6a0e719
Differentiate Copilot agent requests from manual requests by adding s…
2f5464b
Fix docstring accuracy and add UserCredential agentic session integra…
0f47de4
Add telemetry for agentic session
c7de4f4
Add e2e test for agentic session differentiation
043b3e8
upgrade msal python
add0903
Use claim challenge for only the brocker path
369632d
e2e test login for a specific tenant
d0c712d
Remove debug e2e test
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,67 @@ | ||
| # -------------------------------------------------------------------------------------------- | ||
| # Copyright (c) Microsoft Corporation. All rights reserved. | ||
| # Licensed under the MIT License. See License.txt in the project root for license information. | ||
| # -------------------------------------------------------------------------------------------- | ||
|
|
||
| """ | ||
| Support for Entra Agentic Sessions. | ||
|
|
||
| When CLI runs inside an agent context (e.g., Copilot, Azure MCP), the orchestrator sets the | ||
| COPILOT_AGENT_SESSION_ID environment variable. CLI reads it and passes it to MSAL as both: | ||
| - A query parameter (`client_session`) so ESTS can identify the agentic session | ||
| - A claims challenge so ESTS embeds an agentic marker claim in the token (and MSAL bypasses | ||
| the access token cache to ensure a fresh, agent-tagged token is always fetched) | ||
|
|
||
| This enables downstream systems (RBAC, Defender, Purview) to enforce differentiated policies | ||
| for agent-driven vs. human-driven operations. | ||
| """ | ||
|
|
||
| import json | ||
| import os | ||
|
|
||
| from knack.log import get_logger | ||
|
|
||
| logger = get_logger(__name__) | ||
|
|
||
| COPILOT_AGENT_SESSION_ID = "COPILOT_AGENT_SESSION_ID" | ||
|
|
||
|
|
||
| def build_agentic_session_params(): | ||
| """Read COPILOT_AGENT_SESSION_ID and build the agentic claims challenge. | ||
|
|
||
| :returns: (session_id, claims_challenge) — both None when env var is not set. | ||
| """ | ||
| session_id = os.environ.get(COPILOT_AGENT_SESSION_ID) or None | ||
| if not session_id: | ||
| return None, None | ||
|
|
||
| logger.debug("Agentic session detected (COPILOT_AGENT_SESSION_ID is set)") | ||
|
|
||
| claims_challenge = json.dumps({ | ||
| "access_token": { | ||
| "xms_cli_sid": {"values": [session_id]} | ||
| } | ||
| }) | ||
| return session_id, claims_challenge | ||
|
|
||
|
|
||
| def merge_access_token_claims(existing_claims, new_claims): | ||
| """Merge new claims into an existing claims_challenge JSON string. | ||
|
|
||
| :param existing_claims: Existing claims_challenge JSON string (or None). | ||
| :param new_claims: New claims_challenge JSON string to merge in. Must not be None or empty, | ||
| and must contain a non-empty ``access_token`` object. | ||
| :returns: Merged claims_challenge JSON string. | ||
| :raises ValueError: If ``new_claims`` is None, empty, or does not contain a non-empty | ||
| ``access_token`` object. | ||
| """ | ||
| if not new_claims: | ||
| raise ValueError("new_claims must not be None or empty") | ||
| new_access_token = json.loads(new_claims).get("access_token") | ||
| if not new_access_token: | ||
| raise ValueError("new_claims must contain a non-empty access_token") | ||
|
|
||
| claims_dict = json.loads(existing_claims) if existing_claims else {} | ||
| claims_dict["access_token"] = claims_dict.get("access_token") or {} | ||
| claims_dict["access_token"].update(new_access_token) | ||
| return json.dumps(claims_dict) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
215 changes: 215 additions & 0 deletions
215
src/azure-cli-core/azure/cli/core/auth/tests/test_agentic_session.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,215 @@ | ||
| # -------------------------------------------------------------------------------------------- | ||
| # Copyright (c) Microsoft Corporation. All rights reserved. | ||
| # Licensed under the MIT License. See License.txt in the project root for license information. | ||
| # -------------------------------------------------------------------------------------------- | ||
|
|
||
| import json | ||
| import os | ||
| import unittest | ||
| from unittest.mock import patch | ||
|
|
||
| from azure.cli.core.auth.agentic_session import ( | ||
| COPILOT_AGENT_SESSION_ID, | ||
| build_agentic_session_params, | ||
| merge_access_token_claims, | ||
| ) | ||
|
|
||
|
|
||
| class TestBuildAgenticSessionParams(unittest.TestCase): | ||
|
|
||
| def test_returns_none_when_env_not_set(self): | ||
| with patch.dict(os.environ, {}, clear=True): | ||
| session_id, claims = build_agentic_session_params() | ||
| self.assertIsNone(session_id) | ||
| self.assertIsNone(claims) | ||
|
|
||
| def test_returns_none_when_env_is_empty_string(self): | ||
| with patch.dict(os.environ, {COPILOT_AGENT_SESSION_ID: ""}): | ||
| session_id, claims = build_agentic_session_params() | ||
| self.assertIsNone(session_id) | ||
| self.assertIsNone(claims) | ||
|
|
||
| def test_returns_session_id_and_claims(self): | ||
| with patch.dict(os.environ, {COPILOT_AGENT_SESSION_ID: "sess-456"}): | ||
| session_id, claims = build_agentic_session_params() | ||
| self.assertEqual(session_id, "sess-456") | ||
| parsed = json.loads(claims) | ||
| self.assertEqual(parsed["access_token"]["xms_cli_sid"]["values"], ["sess-456"]) | ||
|
|
||
| def _agentic_claims(session_id="s1"): | ||
| return json.dumps({"access_token": {"xms_cli_sid": {"values": [session_id]}}}) | ||
|
|
||
|
|
||
| class TestMergeAccessTokenClaims(unittest.TestCase): | ||
|
|
||
| # --- Validation --- | ||
|
|
||
| def test_raises_when_new_claims_is_none(self): | ||
| with self.assertRaises(ValueError): | ||
| merge_access_token_claims(None, None) | ||
|
|
||
| def test_raises_when_new_access_token_is_null(self): | ||
| new = json.dumps({"access_token": None}) | ||
| with self.assertRaises(ValueError): | ||
| merge_access_token_claims(None, new) | ||
|
|
||
| # --- Merging --- | ||
|
|
||
| def test_merges_into_none(self): | ||
| result = merge_access_token_claims(None, _agentic_claims("s1")) | ||
| claims = json.loads(result) | ||
| self.assertEqual(len(claims), 1) | ||
| self.assertEqual(len(claims["access_token"]), 1) | ||
| self.assertEqual(claims["access_token"]["xms_cli_sid"], {"values": ["s1"]}) | ||
|
|
||
| def test_merges_into_existing(self): | ||
| existing = json.dumps({"access_token": {"nbf": {"essential": True, "value": "999"}}}) | ||
| result = merge_access_token_claims(existing, _agentic_claims("s1")) | ||
| merged = json.loads(result) | ||
| self.assertEqual(len(merged), 1) | ||
| self.assertEqual(len(merged["access_token"]), 2) | ||
| self.assertEqual(merged["access_token"]["nbf"], {"essential": True, "value": "999"}) | ||
| self.assertEqual(merged["access_token"]["xms_cli_sid"], {"values": ["s1"]}) | ||
|
|
||
| def test_preserves_non_access_token_keys(self): | ||
| existing = json.dumps({ | ||
| "access_token": {"nbf": {"essential": True}}, | ||
| "id_token": {"auth_time": {"essential": True}} | ||
| }) | ||
| result = merge_access_token_claims(existing, _agentic_claims()) | ||
| merged = json.loads(result) | ||
| self.assertEqual(len(merged), 2) | ||
| self.assertEqual(len(merged["access_token"]), 2) | ||
| self.assertEqual(merged["id_token"], {"auth_time": {"essential": True}}) | ||
| self.assertEqual(merged["access_token"]["nbf"], {"essential": True}) | ||
| self.assertEqual(merged["access_token"]["xms_cli_sid"], {"values": ["s1"]}) | ||
|
|
||
| def test_new_claims_overwrites_existing_key(self): | ||
| existing = json.dumps({"access_token": {"xms_cli_sid": {"values": ["old"]}}}) | ||
| result = merge_access_token_claims(existing, _agentic_claims("new")) | ||
| merged = json.loads(result) | ||
| self.assertEqual(len(merged), 1) | ||
| self.assertEqual(len(merged["access_token"]), 1) | ||
| self.assertEqual(merged["access_token"]["xms_cli_sid"], {"values": ["new"]}) | ||
|
|
||
| def test_creates_access_token_when_missing_in_existing(self): | ||
| existing = json.dumps({"id_token": {"auth_time": {"essential": True}}}) | ||
| result = merge_access_token_claims(existing, _agentic_claims()) | ||
| merged = json.loads(result) | ||
| self.assertEqual(len(merged), 2) | ||
| self.assertEqual(len(merged["access_token"]), 1) | ||
| self.assertEqual(merged["id_token"], {"auth_time": {"essential": True}}) | ||
| self.assertEqual(merged["access_token"]["xms_cli_sid"], {"values": ["s1"]}) | ||
|
|
||
| def test_handles_null_access_token_in_existing(self): | ||
| existing = json.dumps({"access_token": None}) | ||
| result = merge_access_token_claims(existing, _agentic_claims()) | ||
| merged = json.loads(result) | ||
| self.assertEqual(len(merged), 1) | ||
| self.assertEqual(len(merged["access_token"]), 1) | ||
| self.assertEqual(merged["access_token"]["xms_cli_sid"], {"values": ["s1"]}) | ||
|
|
||
|
|
||
| class TestUserCredentialAgenticSession(unittest.TestCase): | ||
| """Verify that UserCredential.acquire_token merges agentic claims and passes | ||
| client_session param when COPILOT_AGENT_SESSION_ID is set.""" | ||
|
|
||
| def _build_user_credential(self, enable_broker=False): | ||
| """Build a UserCredential with mocked MSAL app.""" | ||
| from unittest.mock import MagicMock, PropertyMock | ||
| from azure.cli.core.auth.msal_credentials import UserCredential | ||
|
|
||
| cred = object.__new__(UserCredential) | ||
|
|
||
| cred._msal_app = MagicMock() | ||
| cred._msal_app.client_id = "test-client-id" | ||
| cred._msal_app._enable_broker = enable_broker | ||
| type(cred._msal_app).authority = PropertyMock(return_value=MagicMock( | ||
| instance="login.microsoftonline.com", | ||
| tenant="test-tenant", | ||
| is_adfs=False, | ||
| )) | ||
| cred._account = { | ||
| "home_account_id": "uid.utid", | ||
| "username": "user@test.com", | ||
| } | ||
| return cred | ||
|
|
||
| @patch.dict(os.environ, {COPILOT_AGENT_SESSION_ID: "agent-sess-1"}) | ||
| def test_non_broker_passes_data_only(self): | ||
| """Non-broker path: client_session in data for ext_cache_key, no claims_challenge.""" | ||
| cred = self._build_user_credential(enable_broker=False) | ||
| cred._msal_app.acquire_token_silent_with_error.return_value = { | ||
| "access_token": "agent-tagged-token", | ||
| "token_type": "Bearer", | ||
| "expires_in": 3600, | ||
| } | ||
|
|
||
| result = cred.acquire_token(["https://management.azure.com/.default"]) | ||
|
|
||
| self.assertEqual(result["access_token"], "agent-tagged-token") | ||
|
|
||
| call_kwargs = cred._msal_app.acquire_token_silent_with_error.call_args | ||
| self.assertIsNone(call_kwargs.kwargs.get("claims_challenge")) | ||
| self.assertEqual(call_kwargs.kwargs["data"], {"client_session": "agent-sess-1"}) | ||
| self.assertEqual(call_kwargs.kwargs["params"], {"client_session": "agent-sess-1"}) | ||
|
|
||
| @patch.dict(os.environ, {COPILOT_AGENT_SESSION_ID: "agent-sess-1"}) | ||
| def test_broker_passes_claims_and_data(self): | ||
| """Broker path: claims_challenge with xms_cli_sid AND client_session in data.""" | ||
| cred = self._build_user_credential(enable_broker=True) | ||
| cred._msal_app.acquire_token_silent_with_error.return_value = { | ||
| "access_token": "agent-tagged-token", | ||
| "token_type": "Bearer", | ||
| "expires_in": 3600, | ||
| } | ||
|
|
||
| result = cred.acquire_token(["https://management.azure.com/.default"]) | ||
|
|
||
| self.assertEqual(result["access_token"], "agent-tagged-token") | ||
|
|
||
| call_kwargs = cred._msal_app.acquire_token_silent_with_error.call_args | ||
| claims = json.loads(call_kwargs.kwargs["claims_challenge"]) | ||
| self.assertEqual(claims["access_token"]["xms_cli_sid"]["values"], ["agent-sess-1"]) | ||
| self.assertEqual(call_kwargs.kwargs["data"], {"client_session": "agent-sess-1"}) | ||
| self.assertEqual(call_kwargs.kwargs["params"], {"client_session": "agent-sess-1"}) | ||
|
|
||
| @patch.dict(os.environ, {}, clear=True) | ||
| def test_no_agentic_params_without_env(self): | ||
| """When COPILOT_AGENT_SESSION_ID is not set, no agentic params are added.""" | ||
| cred = self._build_user_credential(enable_broker=False) | ||
| cred._msal_app.acquire_token_silent_with_error.return_value = { | ||
| "access_token": "normal-token", | ||
| "token_type": "Bearer", | ||
| "expires_in": 3600, | ||
| } | ||
|
|
||
| result = cred.acquire_token(["https://management.azure.com/.default"]) | ||
|
|
||
| self.assertEqual(result["access_token"], "normal-token") | ||
|
|
||
| call_kwargs = cred._msal_app.acquire_token_silent_with_error.call_args | ||
| self.assertIsNone(call_kwargs.kwargs.get("claims_challenge")) | ||
| self.assertNotIn("params", call_kwargs.kwargs) | ||
|
|
||
| @patch.dict(os.environ, {COPILOT_AGENT_SESSION_ID: "agent-sess-2"}) | ||
| def test_broker_merges_with_existing_claims(self): | ||
| """Broker path: agentic claims are merged with existing claims_challenge.""" | ||
| cred = self._build_user_credential(enable_broker=True) | ||
| cred._msal_app.acquire_token_silent_with_error.return_value = { | ||
| "access_token": "token", | ||
| "token_type": "Bearer", | ||
| "expires_in": 3600, | ||
| } | ||
|
|
||
| existing_claims = json.dumps({"access_token": {"nbf": {"essential": True, "value": "999"}}}) | ||
| cred.acquire_token(["scope"], claims_challenge=existing_claims) | ||
|
|
||
| call_kwargs = cred._msal_app.acquire_token_silent_with_error.call_args | ||
| claims = json.loads(call_kwargs.kwargs["claims_challenge"]) | ||
| self.assertEqual(claims["access_token"]["nbf"], {"essential": True, "value": "999"}) | ||
| self.assertEqual(claims["access_token"]["xms_cli_sid"]["values"], ["agent-sess-2"]) | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| unittest.main() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.