feat: add sample mcp agent that connects to mcp server via sse endpoint directly

PiperOrigin-RevId: 760388717
This commit is contained in:
Xiang (Sean) Zhou 2025-05-18 16:59:31 -07:00 committed by Copybara-Service
parent 9e767b3fe1
commit 3b5232c14f
6 changed files with 161 additions and 0 deletions

View File

@ -0,0 +1,8 @@
This agent connects to a local MCP server via sse.
To run this agent, start the local MCP server first by :
```bash
uv run filesystem_server.py
```

View File

@ -0,0 +1,58 @@
# 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.
import os
from google.adk.agents.llm_agent import LlmAgent
from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
from google.adk.tools.mcp_tool.mcp_toolset import SseServerParams
_allowed_path = os.path.dirname(os.path.abspath(__file__))
root_agent = LlmAgent(
model='gemini-2.0-flash',
name='enterprise_assistant',
instruction=f"""\
Help user accessing their file systems.
Allowed directory: {_allowed_path}
""",
tools=[
MCPToolset(
connection_params=SseServerParams(
url='http://localhost:3000/sse',
headers={'Accept': 'text/event-stream'},
),
# don't want agent to do write operation
# you can also do below
# tool_filter=lambda tool, ctx=None: tool.name
# not in [
# 'write_file',
# 'edit_file',
# 'create_directory',
# 'move_file',
# ],
tool_filter=[
'read_file',
'read_multiple_files',
'list_directory',
'directory_tree',
'search_files',
'get_file_info',
'list_allowed_directories',
],
)
],
)

View File

@ -0,0 +1,80 @@
# 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.
import asyncio
import os
from pathlib import Path
import sys
from mcp.server.fastmcp import FastMCP
# Create an MCP server with a name
mcp = FastMCP("Filesystem Server", host="localhost", port=3000)
# Add a tool to read file contents
@mcp.tool(description="Read contents of a file")
def read_file(filepath: str) -> str:
"""Read and return the contents of a file."""
with open(filepath, "r") as f:
return f.read()
# Add a tool to list directory contents
@mcp.tool(description="List contents of a directory")
def list_directory(dirpath: str) -> list:
"""List all files and directories in the given directory."""
return os.listdir(dirpath)
# Add a tool to get current working directory
@mcp.tool(description="Get current working directory")
def get_cwd() -> str:
"""Return the current working directory."""
return str(Path.cwd())
# Graceful shutdown handler
async def shutdown(signal, loop):
"""Cleanup tasks tied to the service's shutdown."""
print(f"\nReceived exit signal {signal.name}...")
# Get all running tasks
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
# Cancel all tasks
for task in tasks:
task.cancel()
print(f"Cancelling {len(tasks)} outstanding tasks")
await asyncio.gather(*tasks, return_exceptions=True)
# Stop the loop
loop.stop()
print("Shutdown complete!")
# Main entry point with graceful shutdown handling
if __name__ == "__main__":
try:
# The MCP run function ultimately uses asyncio.run() internally
mcp.run(transport="sse")
except KeyboardInterrupt:
print("\nServer shutting down gracefully...")
# The asyncio event loop has already been stopped by the KeyboardInterrupt
print("Server has been shut down.")
except Exception as e:
print(f"Unexpected error: {e}")
sys.exit(1)
finally:
print("Thank you for using the Filesystem MCP Server!")

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