Authenticate using 3-legged OAuth with auth manager

from google.adk.agents.llm_agent import LlmAgent
from google.adk.auth.credential_manager import CredentialManager
from google.adk.integrations.agent_identity import GcpAuthProvider, GcpAuthProviderScheme
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.auth.auth_tool import AuthConfig

# Register the Google Cloud Auth Provider so the CredentialManager can use it.
CredentialManager.register_auth_provider(GcpAuthProvider())

# The URI to redirect the user to after consent is granted and the
# callback is received by the auth provider.
CONTINUE_URI = "https://YOUR_FRONTEND_URL/validateUserId"

# Create the Auth Provider scheme using the auth provider's full resource name.
auth_scheme = GcpAuthProviderScheme(
    name="projects/PROJECT_ID/locations/LOCATION/connectors/AUTH_PROVIDER_NAME",
    continue_uri=CONTINUE_URI
)

# Configure an MCP tool with the authentication scheme.
toolset = McpToolset(
    connection_params=StreamableHTTPConnectionParams(url="https://YOUR_MCP_SERVER_URL"),
    auth_scheme=auth_scheme,
)

# Initialize the agent with the authenticated tools.
agent = LlmAgent(
    name="AGENT_NAME",
    model="gemini-2.5-flash",
    instruction="AGENT_INSTRUCTIONS",
    tools=[toolset],
)

import os
from google.adk.agents import Agent
from google.adk.apps import App
from google.adk.auth.credential_manager import CredentialManager
from google.adk.integrations.agent_identity import GcpAuthProvider, GcpAuthProviderScheme
from google.adk.models import Gemini
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
import google.auth
from google.genai import types

_, project_id = google.auth.default()
os.environ["GOOGLE_CLOUD_PROJECT"] = "PROJECT_ID"
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"

bigquery_mcp_auth_provider_id = "AUTH_PROVIDER_NAME"
bigquery_mcp_endpoint = os.environ.get(
    "BIGQUERY_MCP_ENDPOINT", "https://bigquery.googleapis.com/mcp"
)

# Register Cloud de Confiance auth provider for Agent Identity Credentials service
CredentialManager.register_auth_provider(GcpAuthProvider())

# The URI to redirect the user to after consent is granted and the callback is received.
CONTINUE_URI = "http://127.0.0.1:8501/validateUserId"

bigquery_mcp_auth_scheme = GcpAuthProviderScheme(
    name=f"projects/{project_id}/locations/LOCATION/connectors/{bigquery_mcp_auth_provider_id}",
    scopes=["https://www.googleapis.com/auth/bigquery"],
    continue_uri=CONTINUE_URI,
)

bigquery_mcp_tools = McpToolset(
    connection_params=StreamableHTTPConnectionParams(url=bigquery_mcp_endpoint),
    auth_scheme=bigquery_mcp_auth_scheme,
    errlog=None,
)

root_agent = Agent(
    name="root_agent",
    model=Gemini(
        model="gemini-2.5-flash",
        retry_options=types.HttpRetryOptions(attempts=3),
    ),
    instruction=(
        "You are a helpful AI assistant designed to provide accurate and useful"
        " information. You can also use your BigQuery MCP tools to look up"
        " BigQuery data."
    ),
    tools=[bigquery_mcp_tools],
)

app = App(
    root_agent=root_agent,
    name="AGENT_NAME",
)

import httpx
from google.adk.agents.llm_agent import LlmAgent
from google.adk.auth.credential_manager import CredentialManager
from google.adk.integrations.agent_identity import GcpAuthProvider
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
from google.adk.apps import App
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_tool import AuthConfig
from google.adk.tools.authenticated_function_tool import AuthenticatedFunctionTool
from vertexai import agent_engines

# First, register Cloud de Confiance auth provider
CredentialManager.register_auth_provider(GcpAuthProvider())

# The URI to redirect the user to after consent is completed.
CONTINUE_URI = "WEB_APP_VALIDATE_USER_URI"

# Create Auth Config
spotify_auth_config = AuthConfig(
    auth_scheme=GcpAuthProviderScheme(
        name="projects/PROJECT_ID/locations/LOCATION/connectors/AUTH_PROVIDER_NAME",
        continue_uri=CONTINUE_URI
    )
)

# Use the Auth Config in Authenticated Function Tool
spotify_search_track_tool = AuthenticatedFunctionTool(
    func=spotify_search_track, auth_config=spotify_auth_config
)

# Sample function tool
async def spotify_search_track(credential: AuthCredential, query: str) -> str | list:
    token = None
    if credential.http and credential.http.credentials:
        token = credential.http.credentials.token

    if not token:
        return "Error: No authentication token available."

    async with httpx.AsyncClient() as client:
        response = await client.get(
            "https://api.spotify.com/v1/search",
            headers={"Authorization": f"Bearer {token}"},
            params={"q": query, "type": "track", "limit": 1},
        )
        # Add your own logic here

agent = LlmAgent(
    name="AGENT_NAME",
    model="gemini-2.5-flash",
    instruction="AGENT_INSTRUCTIONS",
    tools=[spotify_search_track_tool],
)

app = App(
    name="APP_NAME",
    root_agent=agent,
)

vertex_app = agent_engines.AdkApp(app_name=app)

from google.adk.agents.llm_agent import LlmAgent
from google.adk.auth.credential_manager import CredentialManager
from google.adk.integrations.agent_identity import GcpAuthProvider
from google.adk.integrations.agent_identity import GcpAuthProviderScheme
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.auth.auth_tool import AuthConfig
from google.adk.integrations.agent_registry import AgentRegistry

# First, register Cloud de Confiance auth provider
CredentialManager.register_auth_provider(GcpAuthProvider())

# The URI to redirect the user to after consent is completed.
CONTINUE_URI="WEB_APP_VALIDATE_USER_URI"

# Create Cloud de Confiance auth provider by providing auth provider full resource name
auth_scheme = GcpAuthProviderScheme(
    name="projects/PROJECT_ID/locations/LOCATION/connectors/AUTH_PROVIDER_NAME",
    continue_uri=CONTINUE_URI
)

# Set Agent Registry
registry = AgentRegistry(project_id="PROJECT_ID", location="global")

toolset = registry.get_mcp_toolset(mcp_server_name="projects/PROJECT_ID/locations/global/mcpServers/agentregistry-00000000-0000-0000-0000-000000000000", auth_scheme=auth_scheme )

# Example MCP tool
toolset = McpToolset(
    connection_params=StreamableHTTPConnectionParams(url="MCP_URL"),
    auth_scheme=auth_scheme,
)

agent = LlmAgent(
    name="AGENT_NAME",
    model="MODEL_NAME",
    instruction="AGENT_INSTRUCTIONS",
    tools=[toolset],
)

if (fc := get_auth_request_function_call(event_data)):
    print("--> Authentication required by agent.")
    try:
        auth_config = get_auth_config(fc)
        auth_uri, consent_nonce = handle_adk_request_credential(
            auth_config, AUTH_PROVIDER_NAME, request.user_id
        )
        if auth_uri:
            event_data['popup_auth_uri'] = auth_uri
            fc_id = (
                fc.get('id') if isinstance(fc, dict)
                else getattr(fc, 'id', None)
            )
            event_data['auth_request_function_call_id'] = fc_id
            event_data['auth_config'] = auth_config.model_dump()

            # Store session state
            if session_id:
                consent_sessions[session_id] = {
                    "user_id": request.user_id,
                    "consent_nonce": consent_nonce
                }
    except Exception as e:
        print(f"Error handling adk_request_credential: {e}")
        # Optionally, add logic to inform the user about the error.

def handle_adk_request_credential(auth_config, auth_provider_name, user_id):
    ec = auth_config.exchanged_auth_credential
    if ec and ec.oauth2:
        oauth2 = ec.oauth2
        return oauth2.auth_uri, oauth2.nonce
    return None, None

@app.api_route("/validateUserId", methods=["GET"])
async def validate_user(request: Request):
    auth_provider_name = request.query_params.get("auth_provider_name")
    session_id = request.cookies.get("session_id")
    session = consent_sessions.get(session_id, {})

    payload = {
        "userId": session.get("user_id"),
        "userIdValidationState": request.query_params.get(
            "user_id_validation_state"
        ),
        "consentNonce": session.get("consent_nonce"),
    }

    base_url = "https://iamconnectorcredentials.googleapis.com/v1alpha"
    finalize_url = f"{base_url}/{auth_provider_name}/credentials:finalize"

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.post(finalize_url, json=payload)
            resp.raise_for_status()
    except httpx.HTTPError as e:
        err_text = e.response.text if hasattr(e, "response") else str(e)
        status = e.response.status_code if hasattr(e, "response") else 500
        return HTMLResponse(err_text, status_code=status)

    return HTMLResponse("""
        <script>
            window.close();
        </script>
        <p>Success. You can close this window.</p>
    """)

if (request.is_auth_resume and session.auth_request_function_call_id
    and session.auth_config):
    auth_content = types.Content(
        role='user',
        parts=[
            types.Part(
                function_response=types.FunctionResponse(
                    id=session.auth_request_function_call_id,
                    name='adk_request_credential',
                    response=session.auth_config
                )
            )
        ],
    )
    # Send message to agent
    async for event in agent.async_stream_query(
        user_id=request.user_id,
        message=auth_content,
        session_id=session_id,
    ):
        # ...