Support V1/V2 per-audience token acquisition for MCP servers

libraries/microsoft-agents-a365-tooling-extensions-agentframework/microsoft_agents_a365/tooling/extensions/agentframework/services/mcp_tool_registration_service.py

@@ -1,16 +1,19 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+from __future__ import annotations
+
 import logging
 import uuid
 from datetime import datetime, timezone
-from typing import Any, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, List, Optional, Sequence
 
-from agent_framework import RawAgent, Message, BaseHistoryProvider, MCPStreamableHTTPTool
-from agent_framework.azure import AzureOpenAIChatClient
-from agent_framework.openai import OpenAIChatClient
+from agent_framework import RawAgent, Message, HistoryProvider, MCPStreamableHTTPTool
 import httpx
 
+if TYPE_CHECKING:
+    from agent_framework.openai import OpenAIChatClient
+
 from microsoft_agents.hosting.core import Authorization, TurnContext
 
 from microsoft_agents_a365.runtime import OperationResult
@@ -22,6 +25,7 @@
 from microsoft_agents_a365.tooling.utils.constants import Constants
 from microsoft_agents_a365.tooling.utils.utility import (
     get_mcp_platform_authentication_scope,
+    is_development_environment,
 )
 
 
@@ -55,9 +59,9 @@ def __init__(self, logger: Optional[logging.Logger] = None):
 
     async def add_tool_servers_to_agent(
         self,
-        chat_client: Union[OpenAIChatClient, AzureOpenAIChatClient],
+        chat_client: OpenAIChatClient,
         agent_instructions: str,
-        initial_tools: List[Any],
+        initial_tools: List[object],
         auth: Authorization,
         auth_handler_name: str,
         turn_context: TurnContext,
@@ -67,7 +71,7 @@ async def add_tool_servers_to_agent(
         Add MCP tool servers to a RawAgent (mirrors .NET implementation).
 
         Args:
-            chat_client: The chat client instance (Union[OpenAIChatClient, AzureOpenAIChatClient])
+            chat_client: The chat client instance (OpenAIChatClient supports both OpenAI and Azure OpenAI)
             agent_instructions: Instructions for the agent behavior
             initial_tools: List of initial tools to add to the agent
             auth: Authorization context for token exchange
@@ -82,23 +86,31 @@ async def add_tool_servers_to_agent(
             Exception: If agent creation fails.
         """
         try:
-            # Exchange token if not provided
-            if not auth_token:
+            is_dev = is_development_environment()
+            if not auth_token and not is_dev:
+                # Only exchange a token in production; dev mode uses BEARER_TOKEN* env vars instead.
                 scopes = get_mcp_platform_authentication_scope()
                 authToken = await auth.exchange_token(turn_context, scopes, auth_handler_name)
                 auth_token = authToken.token
 
-            agentic_app_id = Utility.resolve_agent_identity(turn_context, auth_token)
+            # In dev mode, agentic_app_id is not needed for manifest-based discovery.
+            agentic_app_id = (
+                "" if is_dev else Utility.resolve_agent_identity(turn_context, auth_token)
+            )
 
             self._logger.info(f"Listing MCP tool servers for agent {agentic_app_id}")
 
             options = ToolOptions(orchestrator_name=self._orchestrator_name)
 
-            # Get MCP server configurations
+            # Get MCP server configurations — pass auth context so each server receives
+            # its own per-audience Authorization token (V1 = shared ATG, V2 = per-GUID).
             server_configs = await self._mcp_server_configuration_service.list_tool_servers(
                 agentic_app_id=agentic_app_id,
                 auth_token=auth_token,
                 options=options,
+                authorization=auth,
+                auth_handler_name=auth_handler_name,
+                turn_context=turn_context,
             )
 
             self._logger.info(f"Loaded {len(server_configs)} MCP server configurations")
@@ -112,16 +124,26 @@ async def add_tool_servers_to_agent(
                 server_name = config.mcp_server_name or config.mcp_server_unique_name
 
                 try:
-                    # Prepare auth headers
-                    headers = {}
-                    if auth_token:
-                        headers[Constants.Headers.AUTHORIZATION] = (
-                            f"{Constants.Headers.BEARER_PREFIX} {auth_token}"
+                    # Merge base (non-auth) headers with per-server headers from list_tool_servers.
+                    # server.headers already contains the correct per-audience Authorization token.
+                    base_headers = {
+                        Constants.Headers.USER_AGENT: Utility.get_user_agent_header(
+                            self._orchestrator_name
                         )
-
-                    headers[Constants.Headers.USER_AGENT] = Utility.get_user_agent_header(
-                        self._orchestrator_name
-                    )
+                    }
+                    server_headers = dict(config.headers) if config.headers else {}
+                    # Fall back to the shared discovery token when no per-server
+                    # Authorization header was attached (e.g. dev mode without
+                    # BEARER_TOKEN* env vars, or legacy V1 callers).
+                    if Constants.Headers.AUTHORIZATION not in server_headers and auth_token:
+                        server_headers[Constants.Headers.AUTHORIZATION] = (
+                            auth_token
+                            if auth_token.lower().startswith(
+                                f"{Constants.Headers.BEARER_PREFIX.lower()} "
+                            )
+                            else f"{Constants.Headers.BEARER_PREFIX} {auth_token}"
+                        )
+                    headers = {**base_headers, **server_headers}  # server auth takes precedence
 
                     # Create httpx client with auth headers configured
                     http_client = httpx.AsyncClient(
@@ -307,18 +329,18 @@ async def send_chat_history_messages(
 
     async def send_chat_history_from_store(
         self,
-        chat_message_store: BaseHistoryProvider,
+        chat_message_store: HistoryProvider,
         turn_context: TurnContext,
         tool_options: Optional[ToolOptions] = None,
     ) -> OperationResult:
         """
-        Send chat history from a BaseHistoryProvider to the MCP platform.
+        Send chat history from a HistoryProvider to the MCP platform.
 
         This is a convenience method that extracts messages from the store
         and delegates to send_chat_history_messages().
 
         Args:
-            chat_message_store: BaseHistoryProvider containing the conversation history.
+            chat_message_store: HistoryProvider containing the conversation history.
             turn_context: TurnContext from the Agents SDK containing conversation info.
             tool_options: Optional configuration for the request.
 

libraries/microsoft-agents-a365-tooling-extensions-agentframework/pyproject.toml

@@ -26,7 +26,7 @@ license = {text = "MIT"}
 dependencies = [
     "microsoft-agents-a365-tooling",
     "microsoft-agents-hosting-core",
-    "agent-framework-azure-ai",
+    "agent-framework",
     "azure-identity",
     "typing-extensions",
     "httpx",

libraries/microsoft-agents-a365-tooling-extensions-azureaifoundry/microsoft_agents_a365/tooling/extensions/azureaifoundry/services/mcp_tool_registration_service.py

@@ -27,7 +27,10 @@
     McpToolServerConfigurationService,
 )
 from microsoft_agents_a365.tooling.utils.constants import Constants
-from microsoft_agents_a365.tooling.utils.utility import get_mcp_platform_authentication_scope
+from microsoft_agents_a365.tooling.utils.utility import (
+    get_mcp_platform_authentication_scope,
+    is_development_environment,
+)
 
 
 class McpToolRegistrationService:
@@ -99,16 +102,18 @@ async def add_tool_servers_to_agent(
         if project_client is None:
             raise ValueError("project_client cannot be None")
 
-        if not auth_token:
+        is_dev = is_development_environment()
+        if not auth_token and not is_dev:
             scopes = get_mcp_platform_authentication_scope()
             authToken = await auth.exchange_token(context, scopes, auth_handler_name)
             auth_token = authToken.token
 
         try:
-            agentic_app_id = Utility.resolve_agent_identity(context, auth_token)
-            # Get the tool definitions and resources using the async implementation
+            agentic_app_id = "" if is_dev else Utility.resolve_agent_identity(context, auth_token)
+            # Get the tool definitions and resources — pass auth context so each server receives
+            # its own per-audience Authorization token (V1 = shared ATG, V2 = per-GUID).
             tool_definitions, tool_resources = await self._get_mcp_tool_definitions_and_resources(
-                agentic_app_id, auth_token or ""
+                agentic_app_id, auth_token or "", auth, auth_handler_name, context
             )
 
             # Update the agent with the tools
@@ -127,7 +132,12 @@ async def add_tool_servers_to_agent(
             raise
 
     async def _get_mcp_tool_definitions_and_resources(
-        self, agentic_app_id: str, auth_token: str
+        self,
+        agentic_app_id: str,
+        auth_token: str,
+        authorization: Authorization,
+        auth_handler_name: str,
+        turn_context: TurnContext,
     ) -> Tuple[List[McpTool], Optional[ToolResources]]:
         """
         Internal method to get MCP tool definitions and resources.
@@ -136,7 +146,10 @@ async def _get_mcp_tool_definitions_and_resources(
 
         Args:
             agentic_app_id: Agentic App ID for the agent.
-            auth_token: Authentication token to access the MCP servers.
+            auth_token: Authentication token used for gateway discovery.
+            authorization: Authorization context for per-audience token exchange.
+            auth_handler_name: Auth handler name for per-audience token exchange.
+            turn_context: TurnContext for per-audience token exchange.
 
         Returns:
             Tuple containing tool definitions and resources.
@@ -145,11 +158,17 @@ async def _get_mcp_tool_definitions_and_resources(
             self._logger.error("MCP server configuration service is not available")
             return ([], None)
 
-        # Get MCP server configurations
+        # Get MCP server configurations — pass auth context so each server receives
+        # its own per-audience Authorization token (V1 = shared ATG, V2 = per-GUID).
         options = ToolOptions(orchestrator_name=self._orchestrator_name)
         try:
             servers = await self._mcp_server_configuration_service.list_tool_servers(
-                agentic_app_id, auth_token, options
+                agentic_app_id,
+                auth_token,
+                options,
+                authorization=authorization,
+                auth_handler_name=auth_handler_name,
+                turn_context=turn_context,
             )
         except Exception as ex:
             self._logger.error(
@@ -191,14 +210,19 @@ async def _get_mcp_tool_definitions_and_resources(
             # Configure the tool
             mcp_tool.set_approval_mode("never")
 
-            # Set up authorization header
-            if auth_token:
-                header_value = (
+            # Set per-server headers: use per-audience token from list_tool_servers
+            # (V1 servers get the shared ATG token, V2 servers get their own audience token).
+            # Fall back to the shared discovery auth_token only if no per-server header is set.
+            server_headers = dict(server.headers) if server.headers else {}
+            auth_header = server_headers.get(Constants.Headers.AUTHORIZATION)
+            if not auth_header and auth_token:
+                auth_header = (
                     auth_token
                     if auth_token.lower().startswith(f"{Constants.Headers.BEARER_PREFIX.lower()} ")
                     else f"{Constants.Headers.BEARER_PREFIX} {auth_token}"
                 )
-                mcp_tool.update_headers(Constants.Headers.AUTHORIZATION, header_value)
+            if auth_header:
+                mcp_tool.update_headers(Constants.Headers.AUTHORIZATION, auth_header)
 
             mcp_tool.update_headers(
                 Constants.Headers.USER_AGENT, Utility.get_user_agent_header(self._orchestrator_name)

libraries/microsoft-agents-a365-tooling-extensions-googleadk/microsoft_agents_a365/tooling/extensions/googleadk/services/mcp_tool_registration_service.py

@@ -26,6 +26,7 @@
 from microsoft_agents_a365.tooling.utils.constants import Constants
 from microsoft_agents_a365.tooling.utils.utility import (
     get_mcp_platform_authentication_scope,
+    is_development_environment,
 )
 
 
@@ -76,19 +77,25 @@ async def add_tool_servers_to_agent(
         Returns:
             None
         """
-        if not auth_token:
+        is_dev = is_development_environment()
+        if not auth_token and not is_dev:
+            # Only exchange a token in production; dev mode uses BEARER_TOKEN* env vars instead.
             scopes = get_mcp_platform_authentication_scope()
             auth_token_obj = await auth.exchange_token(context, scopes, auth_handler_name)
             auth_token = auth_token_obj.token
 
-        agentic_app_id = Utility.resolve_agent_identity(context, auth_token)
+        # In dev mode, agentic_app_id is not needed for manifest-based discovery.
+        agentic_app_id = "" if is_dev else Utility.resolve_agent_identity(context, auth_token)
         self._logger.info(f"Listing MCP tool servers for agent {agentic_app_id}")
 
         options = ToolOptions(orchestrator_name=self._orchestrator_name)
         mcp_server_configs = await self._mcp_server_configuration_service.list_tool_servers(
             agentic_app_id=agentic_app_id,
             auth_token=auth_token,
             options=options,
+            authorization=auth,
+            auth_handler_name=auth_handler_name,
+            turn_context=context,
         )
 
         self._logger.info(f"Loaded {len(mcp_server_configs)} MCP server configurations")
@@ -104,10 +111,6 @@ async def add_tool_servers_to_agent(
 
         # Convert MCP server configs to McpToolset objects (only new ones)
         mcp_servers_info = []
-        mcp_server_headers = {
-            Constants.Headers.AUTHORIZATION: f"{Constants.Headers.BEARER_PREFIX} {auth_token}",
-            Constants.Headers.USER_AGENT: Utility.get_user_agent_header(self._orchestrator_name),
-        }
 
         for server_config in mcp_server_configs:
             # Skip if server URL already exists
@@ -119,10 +122,29 @@ async def add_tool_servers_to_agent(
                 continue
 
             try:
+                base_headers = {
+                    Constants.Headers.USER_AGENT: Utility.get_user_agent_header(
+                        self._orchestrator_name
+                    )
+                }
+                server_headers = dict(server_config.headers) if server_config.headers else {}
+                # Fall back to the shared discovery token when no per-server
+                # Authorization header was attached (e.g. dev mode without
+                # BEARER_TOKEN* env vars, or legacy V1 callers).
+                if Constants.Headers.AUTHORIZATION not in server_headers and auth_token:
+                    server_headers[Constants.Headers.AUTHORIZATION] = (
+                        auth_token
+                        if auth_token.lower().startswith(
+                            f"{Constants.Headers.BEARER_PREFIX.lower()} "
+                        )
+                        else f"{Constants.Headers.BEARER_PREFIX} {auth_token}"
+                    )
+                headers = {**base_headers, **server_headers}  # per-audience token takes precedence
+
                 server_info = McpToolset(
                     connection_params=StreamableHTTPConnectionParams(
                         url=server_config.url,
-                        headers=mcp_server_headers,
+                        headers=headers,
                     )
                 )