{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 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": "JAPoU8Sm5E6e"
      },
      "source": [
        "# MCP on Vertex AI Agent Engine with custom installation scripts\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/tutorial_mcp_on_agent_engine.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%2Fgenerative-ai%2Fmain%2Fagents%2Fagent_engine%2Ftutorial_mcp_on_agent_engine.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/generative-ai/main/agents/agent_engine/tutorial_mcp_on_agent_engine.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/generative-ai/blob/main/agents/agent_engine/tutorial_mcp_on_agent_engine.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",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/tutorial_mcp_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/tutorial_mcp_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/tutorial_mcp_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/tutorial_mcp_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/agent_engine/tutorial_mcp_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| Author(s) |\n",
        "| --- |\n",
        "| [Shawn Yang](https://github.com/shawn-yang-google), [Ivan Nardini](https://github.com/inardini) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Vertex AI Agent Engine provides **support for custom installation scripts**. This means you can now run shell scripts you need during your agent's build process. You can install custom binaries, compile code from source, or set up any other specialized dependency your agent requires.\n",
        "\n",
        "Most importantly, this feature means unlocks a new pattern for deploying agents with complex, standardized tools. In fact, you can now **deploy agents using the Model Context Protocol (MCP) on Vertex AI Agent Engine**.\n",
        "\n",
        "This guide shows you how to build, test, and deploy a custom Reddit agent that uses an MCP server, demonstrating this new pattern for creating custom agents on Vertex AI Agent Engine.\n",
        "\n",
        "**What You'll Learn**\n",
        "\n",
        "  * How to set up your local environment for agent development.\n",
        "  * How to define and configure a toolset using the Model Context Protocol (MCP).\n",
        "  * How to build and test an agent locally using the Google Agent Development Kit (ADK).\n",
        "  * How to use a single, unified function for interactive testing in both local and remote environments.\n",
        "  * How to use the \"bring your own installation script\" feature to deploy the agent to Vertex AI Agent Engine.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK and other required packages\n",
        "\n",
        "Run the following command to install the Vertex AI SDK with the necessary extras for Agent Engine and the ADK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install \"google-cloud-aiplatform[agent_engines,adk]>=1.101.0\" \"aiofiles\" --force-reinstall --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wuDxcQBeWHwh"
      },
      "source": [
        "**Important:** After the installation completes, you must restart your Colab or notebook kernel for the new packages to be recognized.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Next, authenticate your account and initialize the Vertex AI SDK. This allows your environment to interact with your Google Cloud project.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
        "\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}\n",
        "\n",
        "# Set environment variables required for ADK\n",
        "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"TRUE\"\n",
        "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n",
        "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = LOCATION\n",
        "\n",
        "# Initialize Vertex AI\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1-2undfwqcHv"
      },
      "source": [
        "### (Optional) Get Reddit API Credentials\n",
        "\n",
        "Our agent needs credentials to access the Reddit API. You'll need to create a \"script\" app in your Reddit account preferences to get a **client ID** and **client secret**. You will also need a **refresh token**.\n",
        "\n",
        "**Important Security Note:** We will handle these credentials securely using environment variables. **Never hardcode secrets directly in your code.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rGhCn4qxqiZ8"
      },
      "outputs": [],
      "source": [
        "REDDIT_CLIENT_ID = \"[your-reddit-client-id]\"  # @param {type: \"string\", placeholder: \"[your-reddit-client-id]\", isTemplate: true}\n",
        "REDDIT_CLIENT_SECRET = \"[your-reddit-client-secret]\"  # @param {type: \"string\", placeholder: \"[your-reddit-client-secret]\", isTemplate: true}\n",
        "REDDIT_REFRESH_TOKEN = \"[your-reddit-refresh-token]\"  # @param {type: \"string\", placeholder: \"[your-reddit-refresh-token]\", isTemplate: true}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BaTpAWU_urge"
      },
      "source": [
        "## Import required libraries\n",
        "\n",
        "Import the necessary libraries from the ADK, Vertex AI SDK, and Python's standard library.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mwd3-fcsuvHl"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "\n",
        "logging.getLogger(\"google_adk\").setLevel(logging.CRITICAL)\n",
        "logging.getLogger(\"asyncio\").setLevel(logging.CRITICAL)\n",
        "import os\n",
        "import uuid\n",
        "from typing import Any, Iterator, Optional\n",
        "\n",
        "import aiofiles\n",
        "from google.adk.agents import LlmAgent\n",
        "from google.adk.runners import Runner\n",
        "from google.adk.sessions import InMemorySessionService\n",
        "from google.adk.tools.mcp_tool import StdioConnectionParams\n",
        "from google.adk.tools.mcp_tool.mcp_toolset import (MCPToolset,\n",
        "                                                   StdioServerParameters)\n",
        "from google.genai import types\n",
        "from vertexai import agent_engines\n",
        "from vertexai.agent_engines import AgentEngine\n",
        "from vertexai.preview.reasoning_engines import AdkApp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S5fnMisOCxBZ"
      },
      "source": [
        "### Helper function\n",
        "\n",
        "This tutorial uses a helper function, `chat_loop`, to provide a consistent interactive chat interface for testing the agent in both local and deployed environments. This function abstracts the differences between interacting with a local ADK Runner and a remote AgentEngine instance.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sgp2_wsP9tlw"
      },
      "outputs": [],
      "source": [
        "def chat_loop(\n",
        "    app, user_id: Optional[str] = None, session_id: Optional[str] = None\n",
        ") -> None:\n",
        "    \"\"\"Interactive chat loop for AI applications.\"\"\"\n",
        "\n",
        "    # Simple setup\n",
        "    user_id = user_id or f\"u_{uuid.uuid4().hex[:8]}\"\n",
        "\n",
        "    # Handle session based on app type\n",
        "    if isinstance(app, (AdkApp, AgentEngine)):\n",
        "        # Only create session if session_id is not provided\n",
        "        if not session_id:\n",
        "            session = app.create_session(user_id=user_id)\n",
        "            # Handle both dict and object session responses\n",
        "            if isinstance(session, dict):\n",
        "                session_id = session[\"id\"]\n",
        "            else:\n",
        "                session_id = session.id\n",
        "\n",
        "        def query_fn(msg: str):\n",
        "            return app.stream_query(user_id=user_id, session_id=session_id, message=msg)\n",
        "\n",
        "    elif isinstance(app, Runner):\n",
        "        session_id = session_id or f\"s_{uuid.uuid4().hex[:8]}\"\n",
        "\n",
        "        def query_fn(msg: str):\n",
        "            return app.run(\n",
        "                user_id=user_id,\n",
        "                session_id=session_id,\n",
        "                new_message=types.Content(role=\"user\", parts=[types.Part(text=msg)]),\n",
        "            )\n",
        "\n",
        "    else:\n",
        "        raise TypeError(\n",
        "            f\"Unsupported app type: {type(app)}. Expected AdkApp, AgentEngine, or Runner.\"\n",
        "        )\n",
        "\n",
        "    _print_startup_info(user_id, session_id)\n",
        "\n",
        "    # Main loop\n",
        "    while True:\n",
        "        try:\n",
        "            user_input = input(\"\\nYou: \").strip()\n",
        "            if not user_input or user_input.lower() in {\"quit\", \"exit\", \"bye\"}:\n",
        "                break\n",
        "\n",
        "            print(\"\\nAssistant: \", end=\"\", flush=True)\n",
        "\n",
        "            response = _get_response_text(query_fn(user_input))\n",
        "            print(response or \"(No response generated)\")\n",
        "\n",
        "        except (KeyboardInterrupt, EOFError):\n",
        "            print(\"\\n\\n🛑 Chat interrupted.\")\n",
        "            break\n",
        "        except Exception as e:\n",
        "            print(f\"\\n❌ Error: {e}\")\n",
        "            continue\n",
        "\n",
        "    print(\"\\n👋 Goodbye!\")\n",
        "\n",
        "\n",
        "def _print_startup_info(user_id: str, session_id: str) -> None:\n",
        "    \"\"\"Print startup information.\"\"\"\n",
        "    print(\"\\n🚀 Starting chat...\")\n",
        "    print(f\"👤 User ID: {user_id}\")\n",
        "    print(f\"📁 Session ID: {session_id}\")\n",
        "    print(\"💬 Type 'exit' or 'quit' to end.\")\n",
        "    print(\"-\" * 50)\n",
        "\n",
        "\n",
        "def _get_response_text(events: Iterator[Any]) -> str:\n",
        "    \"\"\"Extract response text from event stream.\"\"\"\n",
        "    responses = []\n",
        "\n",
        "    for event in events:\n",
        "        # Handle dict-like events (AgentEngine format)\n",
        "        if isinstance(event, dict):\n",
        "            text = _extract_from_dict_event(event)\n",
        "        # Handle object-like events\n",
        "        else:\n",
        "            text = _extract_from_object_event(event)\n",
        "\n",
        "        if text:\n",
        "            responses.append(text)\n",
        "            # Print streaming text in real-time\n",
        "            print(text, end=\"\", flush=True)\n",
        "\n",
        "    return \"\".join(responses)\n",
        "\n",
        "\n",
        "def _extract_from_dict_event(event: dict) -> Optional[str]:\n",
        "    \"\"\"Extract text from dictionary-style events (AgentEngine format).\"\"\"\n",
        "    # Handle AgentEngine response format\n",
        "    if \"parts\" in event and \"role\" in event:\n",
        "        # Only extract text from model responses, skip function calls/responses\n",
        "        if event.get(\"role\") == \"model\":\n",
        "            parts = event.get(\"parts\", [])\n",
        "            text_parts = []\n",
        "\n",
        "            for part in parts:\n",
        "                # Extract text content, skip function calls\n",
        "                if isinstance(part, dict) and \"text\" in part:\n",
        "                    text_parts.append(part[\"text\"])\n",
        "\n",
        "            return \"\".join(text_parts) if text_parts else None\n",
        "        return None\n",
        "\n",
        "    # Handle other dict formats\n",
        "    content = event.get(\"content\", {})\n",
        "    if isinstance(content, str):\n",
        "        return content\n",
        "\n",
        "    parts = content.get(\"parts\", [])\n",
        "    if not parts:\n",
        "        return None\n",
        "\n",
        "    text_parts = []\n",
        "    for part in parts:\n",
        "        if isinstance(part, dict) and \"text\" in part:\n",
        "            text_parts.append(part[\"text\"])\n",
        "\n",
        "    return \"\".join(text_parts) if text_parts else None\n",
        "\n",
        "\n",
        "def _extract_from_object_event(event: Any) -> Optional[str]:\n",
        "    \"\"\"Extract text from object-style events.\"\"\"\n",
        "    # Handle string content directly\n",
        "    content = getattr(event, \"content\", None)\n",
        "    if isinstance(content, str):\n",
        "        return content\n",
        "\n",
        "    # Handle content with parts\n",
        "    if content and hasattr(content, \"parts\"):\n",
        "        text_parts = []\n",
        "        for part in content.parts:\n",
        "            if hasattr(part, \"text\") and part.text:\n",
        "                text_parts.append(part.text)\n",
        "        return \"\".join(text_parts) if text_parts else None\n",
        "\n",
        "    # Handle direct text attribute\n",
        "    return getattr(event, \"text\", None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oCh91WIcXITc"
      },
      "source": [
        "## Building the Agent Locally\n",
        "\n",
        "Before deployment, the agent is built and tested in the local environment. This allows for rapid iteration and debugging.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_XIzMtR_rcUN"
      },
      "source": [
        "### Prepare the installation script\n",
        "\n",
        "The agent uses an external tool, `mcp-reddit`, which must be installed in the runtime environment. A shell script will be created to handle this installation. Vertex AI Agent Engine will execute this script during the deployment process.\n",
        "\n",
        "First, create a local directory to store the script.\n",
        "\n",
        "Then, the script performs the following actions:\n",
        "- Updates the package index and installs curl.\n",
        "- Installs uv, a fast Python package manager, to accelerate dependency installation.\n",
        "- Adds uv to the system's PATH.\n",
        "- Uses uv to install the mcp-reddit package directly from its Git repository.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-9KwhvQtr0-f"
      },
      "outputs": [],
      "source": [
        "!mkdir -p installation_scripts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q5f_iAu-k9e5"
      },
      "outputs": [],
      "source": [
        "install_local_mcp_file = \"\"\"\n",
        "#!/bin/bash\n",
        "# Exit immediately if a command exits with a non-zero status.\n",
        "set -e\n",
        "\n",
        "echo \"Installing MCP Reddit Server\"\n",
        "\n",
        "# Install uv (a fast Python package manager)\n",
        "apt-get update\n",
        "apt-get install -y curl\n",
        "echo \"Installing uv...\"\n",
        "curl -LsSf https://astral.sh/uv/install.sh | sh\n",
        "\n",
        "# Add uv to PATH for current session\n",
        "export PATH=\"$HOME/.local/bin:$PATH\"\n",
        "\n",
        "# Install the mcp-reddit tool using the command from its documentation\n",
        "echo \"Installing mcp-reddit using uv...\"\n",
        "uv pip install \"git+https://github.com/adhikasp/mcp-reddit.git\" --system\n",
        "echo \"MCP Reddit Server installation complete.\"\n",
        "\"\"\"\n",
        "\n",
        "with open(\"installation_scripts/install_local_mcp.sh\", \"w\") as f:\n",
        "    f.write(install_local_mcp_file)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OYmyvnTO2j9f"
      },
      "outputs": [],
      "source": [
        "!chmod +x installation_scripts/install_local_mcp.sh && ./installation_scripts/install_local_mcp.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jgTU1HTasbvO"
      },
      "source": [
        "### Defining the Agent\n",
        "\n",
        "Here, we define the `LlmAgent` using ADK. We provide a detailed prompt and connect it to our `MCPToolset`, pointing it to the runner script. We define our agent's using a factory function. The main test loop will then call this function to create an agent instance, providing it with the necessary async log file handle at runtime.\n",
        "\n",
        "> **Notice**: To solve the `fileno` error in Colab, we must redirect the tool's error stream to a file that supports asynchronous operations. We will use the aiofiles library for this. This requires our test to run inside an async function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Lha44dp9CCk"
      },
      "outputs": [],
      "source": [
        "def create_agent(errlog):\n",
        "    root_agent = LlmAgent(\n",
        "        model=\"gemini-2.5-flash\",\n",
        "        name=\"reddit_assistant_agent\",\n",
        "        instruction=\"Help the user fetch reddit info.\",\n",
        "        tools=[\n",
        "            MCPToolset(\n",
        "                connection_params=StdioConnectionParams(\n",
        "                    server_params=StdioServerParameters(\n",
        "                        command=\"mcp-reddit\",\n",
        "                    ),\n",
        "                ),\n",
        "                errlog=errlog,  # Required only in colab environment\n",
        "            )\n",
        "        ],\n",
        "    )\n",
        "    return root_agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1t6ceQ0IuVPg"
      },
      "source": [
        "### Test the agent locally\n",
        "\n",
        "Test the agent on the local machine to verify its functionality before deploying to the cloud.\n",
        "\n",
        "The local test setup involves:\n",
        "- Instantiating in-memory versions of the SessionService, MemoryService, and ArtifactService for lightweight local testing.\n",
        "- Creating an async file handle for the tool's error log.\n",
        "- Initializing the ADK Runner, which orchestrates the interaction between the user and the agent.\n",
        "- Calling the chat_loop helper to start an interactive conversation with the agent.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z6UiAGbT-jLp"
      },
      "outputs": [],
      "source": [
        "async def test_agent_locally():\n",
        "\n",
        "    # Create the session service\n",
        "    session_service = InMemorySessionService()\n",
        "\n",
        "    # Create the specific session where the conversation will happen\n",
        "    user_id = \"runner_user_01\"\n",
        "    session = await session_service.create_session(\n",
        "        app_name=\"MyRunnerApp\", user_id=user_id\n",
        "    )\n",
        "\n",
        "    # Initialize the Runner with the correct session service\n",
        "    errlog = await aiofiles.open(\n",
        "        \"error.log\", \"w+\"\n",
        "    )  # Required only in colab environment.\n",
        "    root_agent = create_agent(errlog)\n",
        "    runner = Runner(\n",
        "        agent=root_agent, app_name=\"MyRunnerApp\", session_service=session_service\n",
        "    )\n",
        "\n",
        "    try:\n",
        "        chat_loop(runner, user_id, session.id)\n",
        "    finally:\n",
        "        # Ensure the log file is always closed\n",
        "        await errlog.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d4c5b3b2f21"
      },
      "outputs": [],
      "source": [
        "await test_agent_locally()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ShSFCOuFgb_I"
      },
      "source": [
        "## Deploying to Vertex AI Agent Engine\n",
        "\n",
        "With our agent tested and working locally, it's time to deploy our MCP-enabled agent.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iEpQB0kb_Cln"
      },
      "source": [
        "### Create the agent module\n",
        "\n",
        "Since ADK 1.0.0, the `MCPToolset` contains non-serializable (non-pickleable) state, such as thread locks.\n",
        "\n",
        "To deploy an agent with such a toolset, it must be wrapped in a `ModuleAgent`. This approach involves defining the agent in a separate Python file (`root_agent.py`) and referencing it by the module and variable name during deployment.\n",
        "\n",
        "\n",
        "This agent module file defines the full application, including service builders for `VertexAiSessionService`, which are the cloud-based counterparts to the in-memory services used for local testing. The agent itself is wrapped in an AdkApp object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UwF1E-twlNdr"
      },
      "outputs": [],
      "source": [
        "root_agent_file = f\"\"\"\n",
        "import os\n",
        "from google.adk.agents import LlmAgent\n",
        "from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, StdioServerParameters\n",
        "from google.adk.tools.mcp_tool import StdioConnectionParams\n",
        "from vertexai.preview.reasoning_engines import AdkApp\n",
        "\n",
        "# Set variable\n",
        "PROJECT_ID=os.getenv('PROJECT_ID', '{PROJECT_ID}')\n",
        "LOCATION=os.getenv('LOCATION', '{LOCATION}')\n",
        "BUCKET_NAME=os.getenv('BUCKET_NAME', '{BUCKET_NAME}')\n",
        "REDDIT_CLIENT_ID=os.getenv('REDDIT_CLIENT_ID', '{REDDIT_CLIENT_ID}')\n",
        "REDDIT_CLIENT_SECRET=os.getenv('REDDIT_CLIENT_SECRET', '{REDDIT_CLIENT_SECRET}')\n",
        "REDDIT_REFRESH_TOKEN=os.getenv('REDDIT_REFRESH_TOKEN', '{REDDIT_REFRESH_TOKEN}')\n",
        "\n",
        "# Define session builder\n",
        "def session_service_builder():\n",
        "  # This is needed to ensure InitGoogle and AdkApp setup is called first.\n",
        "  from google.adk.sessions import VertexAiSessionService\n",
        "  return VertexAiSessionService(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "agent_app = AdkApp(\n",
        "    agent=LlmAgent(\n",
        "        model='gemini-2.5-flash',\n",
        "        name='reddit_assistant_agent',\n",
        "        instruction='Help the user fetch reddit info.',\n",
        "        tools=[\n",
        "            MCPToolset(\n",
        "                connection_params=StdioConnectionParams(\n",
        "                    server_params=StdioServerParameters(\n",
        "                      command=\"mcp-reddit\",\n",
        "                    ),\n",
        "                    timeout=60,\n",
        "                ),\n",
        "            )\n",
        "        ],\n",
        "    ),\n",
        "    session_service_builder=session_service_builder\n",
        ")\n",
        "\"\"\"\n",
        "\n",
        "with open(\"root_agent.py\", \"w\") as f:\n",
        "    f.write(root_agent_file)\n",
        "f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ebj7YOeSskCW"
      },
      "source": [
        "### Deploy the agent\n",
        "\n",
        "The agent_engines.create() function is used to deploy the agent. The configuration specifies several key parameters:\n",
        "\n",
        "- `agent_engine`: An `agent_engines.ModuleAgent` is used, pointing to the `root_agent.py` file (module_name) and the `agent_app` variable within it (agent_name).\n",
        "\n",
        "- `requirements`: A list of standard Python package dependencies.\n",
        "\n",
        "- `extra_packages`: A list of local files to be included in the build. This must include the agent module file (`root_agent.py`) and the custom installation script.\n",
        "\n",
        "- `env_vars`: A dictionary of environment variables to be set in the deployed container, used here to securely pass the Reddit API credentials.\n",
        "\n",
        "- `build_options`: This dictionary specifies custom build-time operations. The installation key is used to provide a list of scripts that will be executed before the application server starts. This is where the custom installation script is specified.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8E6NRMJu_cGt"
      },
      "outputs": [],
      "source": [
        "remote_app = agent_engines.create(\n",
        "    display_name=\"reddit_assistant_agent\",\n",
        "    description=\"A Reddit assistant agent with MCP\",\n",
        "    agent_engine=agent_engines.ModuleAgent(\n",
        "        module_name=\"root_agent\",\n",
        "        agent_name=\"agent_app\",\n",
        "        register_operations={\n",
        "            \"\": [\"get_session\", \"list_sessions\", \"create_session\", \"delete_session\"],\n",
        "            \"async\": [\n",
        "                \"async_get_session\",\n",
        "                \"async_list_sessions\",\n",
        "                \"async_create_session\",\n",
        "                \"async_delete_session\",\n",
        "            ],\n",
        "            \"stream\": [\"stream_query\", \"streaming_agent_run_with_events\"],\n",
        "            \"async_stream\": [\"async_stream_query\"],\n",
        "        },\n",
        "    ),\n",
        "    requirements=[\"google-cloud-aiplatform[agent_engines,adk]>=1.101.0\"],\n",
        "    extra_packages=[\n",
        "        \"root_agent.py\",\n",
        "        \"installation_scripts/install_local_mcp.sh\",\n",
        "    ],\n",
        "    env_vars={\n",
        "        \"PROJECT_ID\": PROJECT_ID,\n",
        "        \"LOCATION\": LOCATION,\n",
        "        \"REDDIT_CLIENT_ID\": REDDIT_CLIENT_ID,\n",
        "        \"REDDIT_CLIENT_SECRET\": REDDIT_CLIENT_SECRET,\n",
        "        \"REDDIT_REFRESH_TOKEN\": REDDIT_REFRESH_TOKEN,\n",
        "    },\n",
        "    build_options={\n",
        "        \"installation\": [\n",
        "            \"installation_scripts/install_local_mcp.sh\",\n",
        "        ],\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OWP0Vijjs_0T"
      },
      "source": [
        "After deployment, the chat_loop function can be used again to interact with the now remotely-hosted agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i7psp3vVBbG4"
      },
      "outputs": [],
      "source": [
        "chat_loop(remote_app)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid incurring ongoing charges to your Google Cloud account for the resources used in this tutorial, delete the resources you created. This section provides commands to delete the deployed Agent Engine and the associated Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a715oFmGKu4Y"
      },
      "outputs": [],
      "source": [
        "delete_agent_engine = True\n",
        "delete_bucket = True\n",
        "\n",
        "if delete_agent_engine:\n",
        "    agent_engines = agent_engines.list(filter=\"display_name=reddit_assistant_agent\")\n",
        "    for agent_engine in agent_engines:\n",
        "        agent_engine.delete(force=True)\n",
        "\n",
        "if delete_bucket:\n",
        "    !gsutil rm -r {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_mcp_on_agent_engine.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
