Skip to content

Commit c0fc0bc

Browse files
ayman3000ayman3000
authored
Merge branch 'main' into feature/ollama-llm
2 parents c07da48 + 2770012 commit c0fc0bc

File tree

5 files changed

+426
-0
lines changed

5 files changed

+426
-0
lines changed

‎contributing/samples/mcp_toolset_auth/README.md‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# MCP Toolset OAuth Authentication Sample
2+
3+
This sample demonstrates the toolset authentication feature where OAuth credentials are required for both tool listing and tool calling.
4+
5+
## Overview
6+
7+
The toolset authentication flow works in two phases:
8+
9+
1. **Phase 1**: When the agent tries to get tools from the MCP server without credentials, the toolset signals "authentication required" and returns an auth request event.
10+
11+
2. **Phase 2**: After the user provides OAuth credentials, the agent can successfully list and call tools.
12+
13+
## Files
14+
15+
- `oauth_mcp_server.py` - MCP server that requires Bearer token authentication
16+
- `agent.py` - Agent configuration with OAuth-protected MCP toolset
17+
- `main.py` - Test script demonstrating the two-phase auth flow
18+
19+
## Running the Sample
20+
21+
1. Start the MCP server in one terminal:
22+
23+
```bash
24+
PYTHONPATH=src python contributing/samples/mcp_toolset_auth/oauth_mcp_server.py
25+
```
26+
27+
2. Run the test script in another terminal:
28+
29+
```bash
30+
PYTHONPATH=src python contributing/samples/mcp_toolset_auth/main.py
31+
```
32+
33+
## Expected Behavior
34+
35+
1. First invocation yields an `adk_request_credential` function call
36+
2. The credential ID is `_adk_toolset_auth_McpToolset` to indicate toolset auth
37+
3. After providing the access token, the agent can list and call tools
38+
39+
## Testing with ADK Web UI
40+
41+
You can also test with the ADK web UI:
42+
43+
```bash
44+
adk web contributing/samples/mcp_toolset_auth
45+
```
46+
47+
Note: The web UI will display the auth request and you'll need to manually provide credentials.

‎contributing/samples/mcp_toolset_auth/__init__.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from . import agent

‎contributing/samples/mcp_toolset_auth/agent.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Agent that uses MCP toolset requiring OAuth authentication.
16+
17+
This agent demonstrates the toolset authentication feature where OAuth
18+
credentials are required for both tool listing and tool calling.
19+
"""
20+
21+
from __future__ import annotations
22+
23+
from fastapi.openapi.models import OAuth2
24+
from fastapi.openapi.models import OAuthFlowAuthorizationCode
25+
from fastapi.openapi.models import OAuthFlows
26+
from google.adk.agents import LlmAgent
27+
from google.adk.auth.auth_credential import AuthCredential
28+
from google.adk.auth.auth_credential import AuthCredentialTypes
29+
from google.adk.auth.auth_credential import OAuth2Auth
30+
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
31+
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
32+
33+
# OAuth2 auth scheme with authorization code flow
34+
# This specifies the OAuth metadata needed for the full OAuth flow
35+
auth_scheme = OAuth2(
36+
flows=OAuthFlows(
37+
authorizationCode=OAuthFlowAuthorizationCode(
38+
authorizationUrl='https://example.com/oauth/authorize',
39+
tokenUrl='https://example.com/oauth/token',
40+
scopes={'read': 'Read access', 'write': 'Write access'},
41+
)
42+
)
43+
)
44+
45+
# OAuth credential with client credentials (used for token exchange)
46+
# In a real scenario, this would be used to obtain the access token
47+
auth_credential = AuthCredential(
48+
auth_type=AuthCredentialTypes.OAUTH2,
49+
oauth2=OAuth2Auth(
50+
client_id='test_client_id',
51+
client_secret='test_client_secret',
52+
),
53+
)
54+
55+
# Create the MCP toolset with OAuth authentication
56+
mcp_toolset = McpToolset(
57+
connection_params=StreamableHTTPConnectionParams(
58+
url='http://localhost:3001/mcp',
59+
),
60+
auth_scheme=auth_scheme,
61+
auth_credential=auth_credential,
62+
)
63+
64+
# Define the agent that uses the OAuth-protected MCP toolset
65+
root_agent = LlmAgent(
66+
model='gemini-2.0-flash',
67+
name='oauth_mcp_agent',
68+
instruction="""You are a helpful assistant that can access user information.
69+
70+
You have access to tools that require authentication:
71+
- get_user_profile: Get profile information for a specific user
72+
- list_users: List all available users
73+
74+
When the user asks about users, use these tools to help them.""",
75+
tools=[mcp_toolset],
76+
)

‎contributing/samples/mcp_toolset_auth/main.py‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Test script for MCP Toolset OAuth Authentication Flow.
16+
17+
This script demonstrates the two-phase tool discovery flow:
18+
1. First invocation: Agent tries to get tools, auth is required, returns auth
19+
request event (adk_request_credential)
20+
2. User provides OAuth credentials (simulated)
21+
3. Second invocation: Agent has credentials, can list and call tools
22+
23+
Usage:
24+
# Start the MCP server first (in another terminal):
25+
PYTHONPATH=src python contributing/samples/mcp_toolset_auth/oauth_mcp_server.py
26+
27+
# Run the demo:
28+
PYTHONPATH=src python contributing/samples/mcp_toolset_auth/main.py
29+
"""
30+
31+
from __future__ import annotations
32+
33+
import asyncio
34+
35+
from agent import auth_credential
36+
from agent import auth_scheme
37+
from agent import mcp_toolset
38+
from agent import root_agent
39+
from google.adk.auth.auth_credential import AuthCredential
40+
from google.adk.auth.auth_credential import AuthCredentialTypes
41+
from google.adk.auth.auth_credential import OAuth2Auth
42+
from google.adk.auth.auth_tool import AuthConfig
43+
from google.adk.runners import Runner
44+
from google.adk.sessions.in_memory_session_service import InMemorySessionService
45+
from google.genai import types
46+
47+
48+
async def run_demo():
49+
"""Run demo with real MCP server."""
50+
print('=' * 60)
51+
print('MCP Toolset OAuth Authentication Demo')
52+
print('=' * 60)
53+
print('\nNote: Make sure the MCP server is running:')
54+
print(' python oauth_mcp_server.py\n')
55+
56+
# Create session service and runner
57+
session_service = InMemorySessionService()
58+
runner = Runner(
59+
agent=root_agent,
60+
app_name='toolset_auth_demo',
61+
session_service=session_service,
62+
)
63+
64+
# Create a session
65+
session = await session_service.create_session(
66+
app_name='toolset_auth_demo',
67+
user_id='test_user',
68+
)
69+
70+
print(f'Session created: {session.id}')
71+
print('\n--- Phase 1: Initial request (no credentials) ---\n')
72+
73+
# First invocation - should trigger auth request
74+
user_message = 'List all users'
75+
print(f'User: {user_message}')
76+
77+
events = []
78+
auth_function_call_id = None
79+
max_events = 10
80+
81+
try:
82+
async for event in runner.run_async(
83+
session_id=session.id,
84+
user_id='test_user',
85+
new_message=types.Content(
86+
role='user',
87+
parts=[types.Part(text=user_message)],
88+
),
89+
):
90+
events.append(event)
91+
print(f'\nEvent from {event.author}:')
92+
if event.content and event.content.parts:
93+
for part in event.content.parts:
94+
if part.text:
95+
print(f' Text: {part.text}')
96+
if part.function_call:
97+
print(f' Function call: {part.function_call.name}')
98+
if part.function_call.name == 'adk_request_credential':
99+
auth_function_call_id = part.function_call.id
100+
101+
if len(events) >= max_events:
102+
print(f'\n** SAFETY LIMIT ({max_events} events) **')
103+
break
104+
105+
except Exception as e:
106+
print(f'\nError: {e}')
107+
print('Make sure the MCP server is running!')
108+
await mcp_toolset.close()
109+
return
110+
111+
if auth_function_call_id:
112+
print('\n** Auth request detected! **')
113+
print('\n--- Phase 2: Provide OAuth credentials ---\n')
114+
115+
# Simulate user providing OAuth credentials after completing OAuth flow
116+
auth_response = AuthConfig(
117+
auth_scheme=auth_scheme,
118+
raw_auth_credential=auth_credential,
119+
exchanged_auth_credential=AuthCredential(
120+
auth_type=AuthCredentialTypes.OAUTH2,
121+
oauth2=OAuth2Auth(
122+
access_token='test_access_token_12345',
123+
),
124+
),
125+
)
126+
127+
print('Providing access token: test_access_token_12345')
128+
129+
auth_response_message = types.Content(
130+
role='user',
131+
parts=[
132+
types.Part(
133+
function_response=types.FunctionResponse(
134+
name='adk_request_credential',
135+
id=auth_function_call_id,
136+
response=auth_response.model_dump(exclude_none=True),
137+
)
138+
)
139+
],
140+
)
141+
142+
async for event in runner.run_async(
143+
session_id=session.id,
144+
user_id='test_user',
145+
new_message=auth_response_message,
146+
):
147+
print(f'\nEvent from {event.author}:')
148+
if event.content and event.content.parts:
149+
for part in event.content.parts:
150+
if part.text:
151+
text = (
152+
part.text[:200] + '...' if len(part.text) > 200 else part.text
153+
)
154+
print(f' Text: {text}')
155+
if part.function_call:
156+
print(f' Function call: {part.function_call.name}')
157+
else:
158+
print('\n** No auth request - credentials may already be available **')
159+
160+
print('\n' + '=' * 60)
161+
print('Demo completed')
162+
print('=' * 60)
163+
164+
await mcp_toolset.close()
165+
166+
167+
if __name__ == '__main__':
168+
asyncio.run(run_demo())

0 commit comments

Comments
 (0)