diff --git a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/connector/client/user_token_client.py b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/connector/client/user_token_client.py index c0330524..289af112 100644 --- a/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/connector/client/user_token_client.py +++ b/libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/connector/client/user_token_client.py @@ -9,6 +9,7 @@ from microsoft_agents.hosting.core.connector import UserTokenClientBase from microsoft_agents.activity import ( + ChannelId, TokenOrSignInResourceResponse, TokenResponse, TokenStatus, @@ -112,6 +113,15 @@ class UserToken(UserTokenBase): def __init__(self, client: ClientSession): self.client = client + @staticmethod + def _base_channel_id(channel_id: Optional[str]) -> Optional[str]: + """Return the Bot Framework channel without an optional sub-channel.""" + if not channel_id or not channel_id.strip(): + return channel_id + + base_channel_id = ChannelId(channel_id).channel + return base_channel_id or channel_id + async def get_token( self, user_id: str, @@ -120,6 +130,8 @@ async def get_token( code: Optional[str] = None, ) -> TokenResponse: + channel_id = self._base_channel_id(channel_id) + with spans.GetUserToken( connection_name=connection_name, user_id=user_id, channel_id=channel_id ) as span: @@ -155,6 +167,8 @@ async def _get_token_or_sign_in_resource( ) -> TokenOrSignInResourceResponse: """Get token or sign-in resource for a user.""" + channel_id = self._base_channel_id(channel_id) + with spans.GetTokenOrSignInResource( connection_name=connection_name, user_id=user_id, channel_id=channel_id ) as span: @@ -192,6 +206,8 @@ async def get_aad_tokens( ) -> dict[str, TokenResponse]: """Get AAD tokens for a user.""" + channel_id = self._base_channel_id(channel_id) + with spans.GetAadTokens( connection_name=connection_name, user_id=user_id, channel_id=channel_id ) as span: @@ -221,6 +237,8 @@ async def sign_out( ) -> None: """Sign out user from a connection.""" + channel_id = self._base_channel_id(channel_id) + with spans.SignOut( user_id=user_id, connection_name=connection_name, channel_id=channel_id ) as span: @@ -249,6 +267,8 @@ async def get_token_status( ) -> list[TokenStatus]: """Get token status for a user.""" + channel_id = self._base_channel_id(channel_id) + with spans.GetTokenStatus(user_id=user_id, channel_id=channel_id) as span: params = {"userId": user_id} @@ -281,6 +301,8 @@ async def exchange_token( ) -> TokenResponse: """Exchange token for a user.""" + channel_id = self._base_channel_id(channel_id) + with spans.ExchangeToken( connection_name=connection_name, user_id=user_id, channel_id=channel_id ) as span: diff --git a/tests/hosting_core/connector/test_user_token_client.py b/tests/hosting_core/connector/test_user_token_client.py new file mode 100644 index 00000000..d030e03d --- /dev/null +++ b/tests/hosting_core/connector/test_user_token_client.py @@ -0,0 +1,108 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +"""Tests for UserToken Bot Framework operations.""" + +import pytest +from aiohttp import ClientSession, web +from aiohttp.test_utils import TestServer + +from microsoft_agents.hosting.core.connector.client.user_token_client import UserToken + + +class TestUserTokenBaseChannel: + """Bot Framework token operations use the base channel partition.""" + + @pytest.mark.asyncio + async def test_all_token_operations_normalize_composite_channel(self): + captured = [] + + async def handler(request): + captured.append( + (request.method, request.path, request.query.get("channelId")) + ) + path = request.path.lower() + if path.endswith("/gettokenorsigninresource"): + return web.json_response({"tokenResponse": {"token": "token"}}) + if path.endswith("/gettoken") or path.endswith("/exchange"): + return web.json_response({"token": "token"}) + if path.endswith("/getaadtokens"): + return web.json_response({"resource": {"token": "token"}}) + if path.endswith("/gettokenstatus"): + return web.json_response([]) + if path.endswith("/signout"): + return web.Response(status=204) + raise AssertionError(f"Unexpected token operation: {request.path}") + + app = web.Application() + app.router.add_route("*", "/{tail:.*}", handler) + server = TestServer(app) + await server.start_server() + try: + async with ClientSession(base_url=server.make_url("/")) as session: + user_token = UserToken(session) + args = { + "user_id": "user", + "connection_name": "connection", + "channel_id": "msteams:COPILOT", + } + await user_token.get_token(**args) + await user_token._get_token_or_sign_in_resource(**args, state="state") + await user_token.get_aad_tokens(**args) + await user_token.sign_out(**args) + await user_token.get_token_status( + user_id="user", channel_id="msteams:COPILOT" + ) + await user_token.exchange_token(**args) + finally: + await server.close() + + assert len(captured) == 6 + assert all(channel_id == "msteams" for _, _, channel_id in captured) + + @pytest.mark.asyncio + async def test_optional_channel_is_omitted_when_none(self): + captured = [] + + async def handler(request): + captured.append(request.query.get("channelId")) + path = request.path.lower() + if path.endswith("/gettoken"): + return web.json_response({"token": "token"}) + if path.endswith("/getaadtokens"): + return web.json_response({"resource": {"token": "token"}}) + if path.endswith("/gettokenstatus"): + return web.json_response([]) + if path.endswith("/signout"): + return web.Response(status=204) + raise AssertionError(f"Unexpected token operation: {request.path}") + + app = web.Application() + app.router.add_route("*", "/{tail:.*}", handler) + server = TestServer(app) + await server.start_server() + try: + async with ClientSession(base_url=server.make_url("/")) as session: + user_token = UserToken(session) + await user_token.get_token("user", "connection") + await user_token.get_aad_tokens("user", "connection") + await user_token.sign_out("user", "connection") + await user_token.get_token_status("user") + finally: + await server.close() + + assert captured == [None, None, None, None] + + @pytest.mark.parametrize( + ("channel_id", "expected"), + [ + ("msteams", "msteams"), + ("msteams:COPILOT", "msteams"), + ("msteams:", "msteams"), + (":COPILOT", ":COPILOT"), + (" ", " "), + (None, None), + ], + ) + def test_base_channel_id(self, channel_id, expected): + assert UserToken._base_channel_id(channel_id) == expected