encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: Fix credential manager so that it supports the ServiceAccountCredentialExchanger

Browse Source 
This fixes MCP authentication for gcloud service accounts. Previously it was failing to authenticate tool calls.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 826639044
This commit is contained in:
Kathy Wu
2025-10-31 14:55:06 -07:00
committed by Copybara-Service
parent a0df75b6fa
commit e8526f7e06
5 changed files with 178 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contributing/samples/mcp_service_account_agent/README.md
+55
View File
@@ -0,0 +1,55 @@
# MCP Service Account Agent Sample
This agent demonstrates how to connect to a remote MCP server using a gcloud service account for authentication. It uses Streamable HTTP for communication.
## Setup
Before running the agent, you need to configure the MCP server URL and your service account credentials in `agent.py`.
1. **Configure MCP Server URL:**
Update the `MCP_SERVER_URL` variable with the URL of your MCP server instance.
```python
# agent.py
# TODO: Update this to the production MCP server url and scopes.
MCP_SERVER_URL = "https://test.sandbox.googleapis.com/mcp"
```
2. **Set up Service Account Credentials:**
- Obtain the JSON key file for your gcloud service account.
- In `agent.py`, find the `ServiceAccountCredential` object and populate its parameters (e.g., `project_id`, `private_key`, `client_email`, etc.) with the corresponding values from your JSON key file.
```python
# agent.py
# TODO: Update this to the user's service account credentials.
auth_credential=AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
service_account=ServiceAccount(
service_account_credential=ServiceAccountCredential(
type_="service_account",
project_id="example",
private_key_id="123",
private_key="123",
client_email="test@example.iam.gserviceaccount.com",
client_id="123",
auth_uri="https://accounts.google.com/o/oauth2/auth",
token_uri="https://oauth2.googleapis.com/token",
auth_provider_x509_cert_url=(
"https://www.googleapis.com/oauth2/v1/certs"
),
client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/example.iam.gserviceaccount.com",
universe_domain="googleapis.com",
),
scopes=SCOPES.keys(),
),
),
```
## Running the Agent
Once configured, you can run the agent.
For example:
```bash
adk web
```
contributing/samples/mcp_service_account_agent/__init__.py
+15
View File
@@ -0,0 +1,15 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from . import agent
contributing/samples/mcp_service_account_agent/agent.py
+74
View File
@@ -0,0 +1,74 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowClientCredentials
from fastapi.openapi.models import OAuthFlows
from google.adk.agents.llm_agent import LlmAgent
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import ServiceAccount
from google.adk.auth.auth_credential import ServiceAccountCredential
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPServerParams
from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
# TODO: Update this to the production MCP server url and scopes.
MCP_SERVER_URL = "https://test.sandbox.googleapis.com/mcp"
SCOPES = {"https://www.googleapis.com/auth/cloud-platform": ""}
root_agent = LlmAgent(
model="gemini-2.0-flash",
name="enterprise_assistant",
instruction="""
Help the user with the tools available to you.
""",
tools=[
MCPToolset(
connection_params=StreamableHTTPServerParams(
url=MCP_SERVER_URL,
),
auth_scheme=OAuth2(
flows=OAuthFlows(
clientCredentials=OAuthFlowClientCredentials(
tokenUrl="https://oauth2.googleapis.com/token",
scopes=SCOPES,
)
)
),
# TODO: Update this to the user's service account credentials.
auth_credential=AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
service_account=ServiceAccount(
service_account_credential=ServiceAccountCredential(
type_="service_account",
project_id="example",
private_key_id="123",
private_key="123",
client_email="test@example.iam.gserviceaccount.com",
client_id="123",
auth_uri="https://accounts.google.com/o/oauth2/auth",
token_uri="https://oauth2.googleapis.com/token",
auth_provider_x509_cert_url=(
"https://www.googleapis.com/oauth2/v1/certs"
),
client_x509_cert_url="https://www.googleapis.com/robot/v1/metadata/x509/example.iam.gserviceaccount.com",
universe_domain="googleapis.com",
),
scopes=SCOPES.keys(),
),
),
)
],
)
src/google/adk/auth/credential_manager.py
+16 -4
View File
@@ -20,6 +20,7 @@ from typing import Optional
from fastapi.openapi.models import OAuth2
from ..agents.callback_context import CallbackContext
from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger
from ..utils.feature_decorator import experimental
from .auth_credential import AuthCredential
from .auth_credential import AuthCredentialTypes
@@ -84,7 +85,6 @@ class CredentialManager:
self._discovery_manager = OAuth2DiscoveryManager()
# Register default exchangers and refreshers
# TODO: support service account credential exchanger
from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger
from .refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
@@ -96,6 +96,12 @@ class CredentialManager:
AuthCredentialTypes.OPEN_ID_CONNECT, oauth2_exchanger
)
# TODO: Move ServiceAccountCredentialExchanger to the auth module
self._exchanger_registry.register(
AuthCredentialTypes.SERVICE_ACCOUNT,
ServiceAccountCredentialExchanger(),
)
oauth2_refresher = OAuth2CredentialRefresher()
self._refresher_registry.register(
AuthCredentialTypes.OAUTH2, oauth2_refresher
@@ -207,9 +213,15 @@ class CredentialManager:
if not exchanger:
return credential, False
exchanged_credential = await exchanger.exchange(
credential, self._auth_config.auth_scheme
)
if isinstance(exchanger, ServiceAccountCredentialExchanger):
exchanged_credential = exchanger.exchange_credential(
self._auth_config.auth_scheme, credential
)
else:
exchanged_credential = await exchanger.exchange(
credential, self._auth_config.auth_scheme
)
return exchanged_credential, True
async def _refresh_credential(
tests/unittests/auth/test_credential_manager.py
+18 -20
View File
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest.mock import ANY
from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch
@@ -30,6 +31,7 @@ from google.adk.auth.auth_schemes import AuthSchemeType
from google.adk.auth.auth_schemes import ExtendedOAuth2
from google.adk.auth.auth_tool import AuthConfig
from google.adk.auth.credential_manager import CredentialManager
from google.adk.auth.credential_manager import ServiceAccountCredentialExchanger
from google.adk.auth.oauth2_discovery import AuthorizationServerMetadata
import pytest
@@ -422,36 +424,32 @@ class TestCredentialManager:
await manager._validate_credential()
@pytest.mark.asyncio
async def test_exchange_credentials_service_account(self):
async def test_exchange_credentials_service_account(
self, service_account_credential, oauth2_auth_scheme
):
"""Test _exchange_credential with service account credential."""
mock_service_account = Mock(spec=ServiceAccount)
mock_credential = Mock(spec=AuthCredential)
mock_credential.auth_type = AuthCredentialTypes.SERVICE_ACCOUNT
auth_config = Mock(spec=AuthConfig)
auth_config.auth_scheme = Mock()
auth_config.auth_scheme = oauth2_auth_scheme
# Mock exchanger
mock_exchanger = Mock()
mock_exchanger.exchange = AsyncMock(return_value=mock_credential)
exchanged_credential = Mock(spec=AuthCredential)
manager = CredentialManager(auth_config)
# Mock the exchanger registry to return our mock exchanger
with patch.object(
manager._exchanger_registry,
"get_exchanger",
return_value=mock_exchanger,
):
ServiceAccountCredentialExchanger,
"exchange_credential",
return_value=exchanged_credential,
autospec=True,
) as mock_exchange_credential:
result, was_exchanged = await manager._exchange_credential(
mock_credential
service_account_credential
)
mock_exchanger.exchange.assert_called_once_with(
mock_credential, auth_config.auth_scheme
)
assert result == mock_credential
assert was_exchanged is True
mock_exchange_credential.assert_called_once_with(
ANY, oauth2_auth_scheme, service_account_credential
)
assert result == exchanged_credential
assert was_exchanged is True
@pytest.mark.asyncio
async def test_exchange_credential_no_exchanger(self):