Skip to content

fix(samples): refresh the access token in the BigQuery MCP sample #7218

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vishal-bulbule wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(samples): refresh the access token in the BigQuery MCP sample #7218

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions contributing/samples/integrations/bigquery_mcp/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,12 @@
# limitations under the License.

from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.readonly_context import ReadonlyContext
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.utils import _mtls_utils
import google.auth
import google.auth.transport.requests

BIGQUERY_AGENT_NAME = "adk_sample_bigquery_mcp_agent"
BIGQUERY_MCP_ENDPOINT = _mtls_utils.get_api_endpoint(
Expand All @@ -29,14 +31,25 @@
# Initialize the tools to use the application default credentials.
# https://cloud.google.com/docs/authentication/provide-credentials-adc
credentials, project_id = google.auth.default(scopes=[BIGQUERY_SCOPE])
credentials.refresh(google.auth.transport.requests.Request())
oauth_token = credentials.token


def _auth_headers(ctx: ReadonlyContext) -> dict[str, str]:
"""Returns the auth header, refreshing the access token when it expires.

Access tokens are valid for about an hour, so a token fetched once at import
would stop working in a long-running ``adk web`` or ``adk api_server``
process.
"""
if not credentials.valid:
credentials.refresh(google.auth.transport.requests.Request())
return {"Authorization": f"Bearer {credentials.token}"}


bigquery_mcp_toolset = McpToolset(
connection_params=StreamableHTTPConnectionParams(
url=BIGQUERY_MCP_ENDPOINT,
headers={"Authorization": f"Bearer {oauth_token}"},
)
),
header_provider=_auth_headers,
)

# The variable name `root_agent` determines what your root agent is for the
Expand Down
Loading