From 3b5232c14f48e1d5b170f3698d91639b079722c8 Mon Sep 17 00:00:00 2001
From: "Xiang (Sean) Zhou" <seanzhougoogle@google.com>
Date: Sun, 18 May 2025 16:59:31 -0700
Subject: [PATCH] feat: add sample mcp agent that connects to mcp server via
 sse endpoint directly

PiperOrigin-RevId: 760388717
---
 contributing/samples/mcp_sse_agent/README.md  |  8 ++
 .../{mcp_agent => mcp_sse_agent}/__init__.py  |  0
 contributing/samples/mcp_sse_agent/agent.py   | 58 ++++++++++++++
 .../mcp_sse_agent/filesystem_server.py        | 80 +++++++++++++++++++
 .../mcp_stdio_server_agent/__init__.py        | 15 ++++
 .../agent.py                                  |  0
 6 files changed, 161 insertions(+)
 create mode 100644 contributing/samples/mcp_sse_agent/README.md
 rename contributing/samples/{mcp_agent => mcp_sse_agent}/__init__.py (100%)
 create mode 100755 contributing/samples/mcp_sse_agent/agent.py
 create mode 100644 contributing/samples/mcp_sse_agent/filesystem_server.py
 create mode 100755 contributing/samples/mcp_stdio_server_agent/__init__.py
 rename contributing/samples/{mcp_agent => mcp_stdio_server_agent}/agent.py (100%)

diff --git a/contributing/samples/mcp_sse_agent/README.md b/contributing/samples/mcp_sse_agent/README.md
new file mode 100644
index 0000000..1c211dd
--- /dev/null
+++ b/contributing/samples/mcp_sse_agent/README.md
@@ -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
+```
+
diff --git a/contributing/samples/mcp_agent/__init__.py b/contributing/samples/mcp_sse_agent/__init__.py
similarity index 100%
rename from contributing/samples/mcp_agent/__init__.py
rename to contributing/samples/mcp_sse_agent/__init__.py
diff --git a/contributing/samples/mcp_sse_agent/agent.py b/contributing/samples/mcp_sse_agent/agent.py
new file mode 100755
index 0000000..888a88b
--- /dev/null
+++ b/contributing/samples/mcp_sse_agent/agent.py
@@ -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',
+            ],
+        )
+    ],
+)
diff --git a/contributing/samples/mcp_sse_agent/filesystem_server.py b/contributing/samples/mcp_sse_agent/filesystem_server.py
new file mode 100644
index 0000000..9a04f97
--- /dev/null
+++ b/contributing/samples/mcp_sse_agent/filesystem_server.py
@@ -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!")
diff --git a/contributing/samples/mcp_stdio_server_agent/__init__.py b/contributing/samples/mcp_stdio_server_agent/__init__.py
new file mode 100755
index 0000000..c48963c
--- /dev/null
+++ b/contributing/samples/mcp_stdio_server_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_agent/agent.py b/contributing/samples/mcp_stdio_server_agent/agent.py
similarity index 100%
rename from contributing/samples/mcp_agent/agent.py
rename to contributing/samples/mcp_stdio_server_agent/agent.py