diff --git a/contributing/samples/mcp_service_account_agent/README.md b/contributing/samples/mcp_service_account_agent/README.md new file mode 100644 index 00000000..519537c6 --- /dev/null +++ b/contributing/samples/mcp_service_account_agent/README.md @@ -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 +``` \ No newline at end of file diff --git a/contributing/samples/mcp_service_account_agent/__init__.py b/contributing/samples/mcp_service_account_agent/__init__.py new file mode 100644 index 00000000..c48963cd --- /dev/null +++ b/contributing/samples/mcp_service_account_agent/__init__.py @@ -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 diff --git a/contributing/samples/mcp_service_account_agent/agent.py b/contributing/samples/mcp_service_account_agent/agent.py new file mode 100644 index 00000000..dc3ebf7b --- /dev/null +++ b/contributing/samples/mcp_service_account_agent/agent.py @@ -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(), + ), + ), + ) + ], +) diff --git a/src/google/adk/auth/credential_manager.py b/src/google/adk/auth/credential_manager.py index 3503e1e4..c022ab69 100644 --- a/src/google/adk/auth/credential_manager.py +++ b/src/google/adk/auth/credential_manager.py @@ -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( diff --git a/tests/unittests/auth/test_credential_manager.py b/tests/unittests/auth/test_credential_manager.py index da3ea542..ab021d1e 100644 --- a/tests/unittests/auth/test_credential_manager.py +++ b/tests/unittests/auth/test_credential_manager.py @@ -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):