feat: use Vertex AI MCP server (#4475)

* feat: use Vertex AI MCP server

* Address review comments

notebooks/official/mcp/vertex_ai_mcp_server.ipynb

@@ -0,0 +1,492 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "47T2pQIjCZ0l"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Copyright 2026 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6rxPLFytCZ0m"
+      },
+      "source": [
+        "# Using the Vertex AI MCP Server\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/mcp/vertex_ai_mcp_server.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fmcp%2Fvertex_ai_mcp_server.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/mcp/vertex_ai_mcp_server.ipynb\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/mcp/vertex_ai_mcp_server.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>\n",
+        "\n",
+        "<div style=\"clear: both;\"></div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ofM8TVK9CZ0m"
+      },
+      "source": [
+        "| Authors |\n",
+        "| --- |\n",
+        "| [Eric Dong](www.linkedin.com/in/edong186) |"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "intro"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to connect to the **Remote Model Context Protocol (MCP) server for Vertex AI** (https://us-central1-aiplatform.googleapis.com/mcp).\n",
+        "\n",
+        "Unlike local MCP servers that run as subprocesses, this server is hosted by Google Cloud. It allows LLMs to interact directly with your Google Cloud infrastructure, specifically Endpoint Management, Prompt Management, and Model Garden, using natural language."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zuTAj16yCZ0m"
+      },
+      "source": [
+        "## Getting Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DMEzzv4DCZ0m"
+      },
+      "source": [
+        "### Install libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "install"
+      },
+      "outputs": [],
+      "source": [
+        "%pip install --upgrade --quiet google-auth google-genai mcp jq"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kYhseimIOaXY"
+      },
+      "source": [
+        "> ⚠️ Note: You can ignore the pip's dependency errors."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HPi-nLRjCZ0n"
+      },
+      "source": [
+        "### Import libraries\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "dbuB1Mm6CZ0n"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "import google.auth\n",
+        "from google import genai\n",
+        "from google.auth.transport.requests import Request\n",
+        "from google.genai import types\n",
+        "from mcp import ClientSession\n",
+        "from mcp.client.streamable_http import streamablehttp_client"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rtbBEoaWCZ0n"
+      },
+      "source": [
+        "### Authenticate your notebook environment\n",
+        "\n",
+        "If you are running this notebook on Google Colab, run the cell below to authenticate your environment."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "id": "eJYzVGrICZ0n"
+      },
+      "outputs": [],
+      "source": [
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fA_trRo9CZ0n"
+      },
+      "source": [
+        "### Autenticate your Google Cloud Project for Vertex AI\n",
+        "\n",
+        "You can use a Google Cloud Project or an API Key for authentication. This tutorial uses a Google Cloud Project.\n",
+        "\n",
+        "- [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "id": "tGmr38fGCZ0n"
+      },
+      "outputs": [],
+      "source": [
+        "# fmt: off\n",
+        "# TODO: replace with template\n",
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
+        "# fmt: on\n",
+        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
+        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
+        "\n",
+        "LOCATION = \"us-central1\"\n",
+        "\n",
+        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6YACrFBgPdeS"
+      },
+      "source": [
+        "## Use the Vertex AI MCP endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "id": "LXg9rp_65AUL"
+      },
+      "outputs": [],
+      "source": [
+        "MCP_URL = \"https://us-central1-aiplatform.googleapis.com/mcp\"  # @param {type: \"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BCX2mwd1W4Sd"
+      },
+      "source": [
+        "## Enable MCP for your Project\n",
+        "\n",
+        "⚠️ You will encounter a 403 Forbidden or Permission Denied error if the MCP feature is not enabled for your specific service.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "id": "QwKUUSAuW7Ty"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\u001b[1;33mWARNING:\u001b[0m The service aiplatform.googleapis.com is already enabled for MCP.\n"
+          ]
+        }
+      ],
+      "source": [
+        "os.environ[\"CLOUDSDK_CORE_PROJECT\"] = PROJECT_ID\n",
+        "\n",
+        "!gcloud beta services mcp enable aiplatform.googleapis.com"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3bzK7GZgPBEA"
+      },
+      "source": [
+        "## Get the access token to authenticate against the MCP endpoint\n",
+        "\n",
+        "- Define the explicit scope required for Google Cloud APIs\n",
+        "- Pass the scopes into the `auth.default()` method\n",
+        "- The token will expire in 60 minutes by default."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 8,
+      "metadata": {
+        "id": "run-mcp"
+      },
+      "outputs": [],
+      "source": [
+        "def get_auth_token():\n",
+        "    scopes = [\"https://www.googleapis.com/auth/cloud-platform\"]\n",
+        "    credentials, _ = google.auth.default(scopes=scopes)\n",
+        "    credentials.refresh(Request())\n",
+        "    return credentials.token\n",
+        "\n",
+        "\n",
+        "token = get_auth_token()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "PseJfp-mRq7X"
+      },
+      "source": [
+        "## Using the MCP Server"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iEqCdFAQPuXU"
+      },
+      "source": [
+        "### **Example 1**: Use cURL\n",
+        "\n",
+        "This method is for low-level discovery, particularly useful for debugging.\n",
+        "\n",
+        "#### Push the variables we generated earlier into the system environment"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 9,
+      "metadata": {
+        "id": "ruAZVQFi2sj-"
+      },
+      "outputs": [],
+      "source": [
+        "os.environ[\"TOKEN\"] = token\n",
+        "os.environ[\"PROJECT_ID\"] = PROJECT_ID\n",
+        "os.environ[\"MCP_URL\"] = MCP_URL"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "48pVNH78QB9O"
+      },
+      "source": [
+        "#### Send an exact JSON-RPC payload to initialize the MCP connection\n",
+        "\n",
+        "- `X-Goog-User-Project` header is mandatory for billing and quota routing."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "id": "_OAWwxsJIRG8"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "create_endpoint\n",
+            "get_endpoint\n",
+            "list_endpoints\n",
+            "update_endpoint\n",
+            "delete_endpoint\n",
+            "create_prompt\n",
+            "get_prompt\n",
+            "list_prompts\n",
+            "update_prompt\n",
+            "delete_prompt\n",
+            "get_publisher_model\n",
+            "list_publisher_models\n",
+            "search_publisher_models\n"
+          ]
+        }
+      ],
+      "source": [
+        "%%bash\n",
+        "\n",
+        "curl -N -X POST \"$MCP_URL\" \\\n",
+        "  -H \"Authorization: Bearer $TOKEN\" \\\n",
+        "  -H \"X-Goog-User-Project: $PROJECT_ID\" \\\n",
+        "  -H \"Content-Type: application/json\" \\\n",
+        "  -H 'accept: application/json, text/event-stream' \\\n",
+        "  -d '{\n",
+        "    \"method\": \"tools/list\",\n",
+        "    \"jsonrpc\": \"2.0\",\n",
+        "    \"id\": 1\n",
+        "  }'  2>/dev/null >response.json\n",
+        "\n",
+        "jq -r \".result.tools[].name\" response.json"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bJXfUTj-R9qr"
+      },
+      "source": [
+        "### **Example 2**: Use SDK\n",
+        "\n",
+        "This method is for agentic approach. The Gen AI SDK, Agent Developement Kit(ADK) and other agent frameworks handle Automatic Function Calling (AFC) when you pass the session as a tool. The model will decide which tool to call based on your prompt.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "id": "BD-dIVcedXA8"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Available MCP Tools:\n",
+            "- create_endpoint: Creates an Endpoint.\n",
+            "- get_endpoint: Gets an Endpoint.\n",
+            "- list_endpoints: Lists Endpoints in a Location.\n",
+            "- update_endpoint: Updates an Endpoint.\n",
+            "- delete_endpoint: Deletes an Endpoint.\n",
+            "- create_prompt: Creates a new managed prompt. Saves the display_name, template text, system instructions, and inference settings (e.g. temperature). Use this to persist a prompt for future versioning.\n",
+            "- get_prompt: Retrieves full details of a prompt including metadata, labels, and configuration. Use this to inspect a prompt's settings or retrieve its full resource name before an update.\n",
+            "- list_prompts: Lists all managed prompts in a location. Use this to browse existing prompts or map a user-friendly display_name to a technical resource name.\n",
+            "- update_prompt: Updates an existing prompt's fields (display_name, instructions, settings, or labels). This performs a patch; fields not specified remain unchanged.\n",
+            "- delete_prompt: Deletes a Prompt.\n",
+            "- get_publisher_model: Retrieves the full metadata for a specific Model Garden publisher model, including supported machine learning tasks (e.g., classification, generation), input/output modalities (e.g., text, image), and available actions such as deployment or fine-tuning configurations. Use this tool to inspect the detailed capabilities, technical requirements, and documentation of a foundation or third-party model before using or deploying it.\n",
+            "- list_publisher_models: Lists available publisher models in Model Garden within a specified publisher's scope. This provides a broad overview of a publisher's portfolio, such as models offered by Google or third-party providers. Use this tool to browse the catalog and discover technical resource names for further investigation or deployment.\n",
+            "- search_publisher_models: Searches for Model Garden publisher models based on a specified query and publisher. Use this tool to find models that match specific keywords or categories within a publisher's scope, allowing for more targeted discovery of foundation or third-party models in Vertex AI Model Garden.\n",
+            "🤖 User Prompt: List my Vertex AI endpoints in project and location us-central1 and check if any managed prompts exist.\n",
+            "\n",
+            "--- Model Response ---\n",
+            "No Vertex AI endpoints or managed prompts were found in project and location `us-central1`.\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Initialize a Gen AI client\n",
+        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=\"global\")\n",
+        "\n",
+        "\n",
+        "async def run_mcp_agent(user_prompt: str):\n",
+        "    headers = {\n",
+        "        \"Authorization\": f\"Bearer {get_auth_token()}\",\n",
+        "        \"X-Goog-User-Project\": PROJECT_ID,\n",
+        "    }\n",
+        "\n",
+        "    # streamablehttp_client manages the persistent connection to the remote MCP endpoint\n",
+        "    async with streamablehttp_client(url=MCP_URL, headers=headers) as streams:\n",
+        "        async with ClientSession(streams[0], streams[1]) as session:\n",
+        "            await session.initialize()\n",
+        "\n",
+        "            tools = await session.list_tools()\n",
+        "            print(\"Available MCP Tools:\")\n",
+        "            for tool in tools.tools:\n",
+        "                print(f\"- {tool.name}: {tool.description}\")\n",
+        "\n",
+        "            print(f\"🤖 User Prompt: {user_prompt}\\n\")\n",
+        "\n",
+        "            # The 'tools=[session]' allows Gemini to 'see' and execute MCP tools\n",
+        "            response = await client.aio.models.generate_content(\n",
+        "                model=\"gemini-3.1-flash-lite-preview\",\n",
+        "                contents=user_prompt,\n",
+        "                config=types.GenerateContentConfig(\n",
+        "                    tools=[session],\n",
+        "                ),\n",
+        "            )\n",
+        "\n",
+        "            print(\"--- Model Response ---\")\n",
+        "            print(response.text)\n",
+        "\n",
+        "\n",
+        "# Execute the agent\n",
+        "prompt = f\"List my Vertex AI endpoints in project {PROJECT_ID} and location {LOCATION} and check if any managed prompts exist.\"\n",
+        "\n",
+        "await run_mcp_agent(prompt)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dX10XXKEgj4N"
+      },
+      "source": [
+        "## Learn More\n",
+        "You have successfully connected a Gemini model to the Vertex AI Remote MCP server to manage cloud resources using natural language. Explore the following resources:\n",
+        "\n",
+        "- [Vertex AI MCP Reference](https://docs.cloud.google.com/vertex-ai/docs/reference/mcp)\n",
+        "- [Managing MCP Servers](https://docs.cloud.google.com/mcp/enable-disable-mcp-servers)"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "vertex_ai_mcp_server.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}