Test Agent Engine Threat Detection

This page explains how to verify that Agent Engine Threat Detection is working by intentionally triggering detectors and checking for findings. Agent Engine Threat Detection is a built-in service of Security Command Center.

Before you begin

To detect potential threats to your Vertex AI Agent Engine agents, ensure that the Agent Engine Threat Detection service is enabled in Security Command Center.

Set up environment

To test Agent Engine Threat Detection, set up a demo agent that simulates malicious activity.

Create project and activate shell

Select or create a Google Cloud project to use for testing.

To test detectors, you can use the Google Cloud console and Cloud Shell.

Go to the Google Cloud console.

Go to the Google Cloud console
Select the project that you will use to test.
Click Activate Cloud Shell.

You can also run the testing instructions from a local shell.

Set up Vertex AI Agent Engine

If you haven't used Vertex AI Agent Engine in this project before, Set up the Vertex AI Agent Engine environment before starting. Record the name of the Cloud Storage staging bucket that you created. We also recommend that you follow the Vertex AI Agent Engine quickstart to learn how to develop and deploy an agent with the Vertex AI SDK.

Create test script

You will create several files that will be used to deploy a new agent for testing. Use a text editor for this, such as nano.

Create a new file named requirements.txt with the following contents.

google-cloud-aiplatform[agent_engines]
google-adk
google-genai
aiohttp
cloudpickle
pydantic

Create a new empty file called installation_scripts/install.sh. Some tests require adding content to this file.

Create a new file called main.py, with the following contents. Replace the PROJECT_ID, LOCATION, and STAGING_BUCKET variables. The staging bucket name must include the gs:// prefix.

import asyncio
import os
import subprocess
import socket
import vertexai
from vertexai import Client, agent_engines
from google.adk.agents import llm_agent
from google.adk.sessions.in_memory_session_service import InMemorySessionService

# Replace with your own project, location, and staging bucket.
LOCATION = "LOCATION"
PROJECT_ID = "PROJECT_ID"
# Staging bucket must have gs:// prefix
STAGING_BUCKET = "STAGING_BUCKET"

client = Client(project=PROJECT_ID, location=LOCATION)

def _run_command(args, **kwargs):
  output = f"Called {' '.join(args)}\n"
  try:
    res = subprocess.run(args, capture_output=True, text=True, **kwargs)
    if res.stdout:
      output += f"Result: {res.stdout.strip()}\n"
    if res.stderr:
      output += f"Error: {res.stderr.strip()}\n"
  except subprocess.TimeoutExpired:
    output += "Command timed out as expected."
  return output

# Tool to simulate threats. The function body will be replaced for individual
# detector tests.
def threat_detection_test():
  output = _run_command(["sleep", "60"])
  output += _run_command(["echo", "this is a fake threat"])
  output += _run_command(["sleep", "10"])
  return output

root_agent = llm_agent.Agent(
    model="gemini-2.5-flash",
    name="threat_detection_test_agent",
    description="Runs threat detection test.",
    instruction="""
      You are an agent that runs a threat detection test using a fake malicious
      command.
    """,
    tools=[threat_detection_test],
)

async def main():
  vertexai.init(
    project=PROJECT_ID,
    location=LOCATION,
    staging_bucket=STAGING_BUCKET,
  )
  app = agent_engines.AdkApp(
      agent=root_agent, session_service_builder=InMemorySessionService
  )
  remote_agent = client.agent_engines.create(
      agent=app,
      config={
          "display_name": "scc_threat_test_agent",
          "identity_type": vertexai.types.IdentityType.AGENT_IDENTITY,
          "requirements": [
              "google-cloud-aiplatform[agent_engines,adk]",
              "cloudpickle",
              "pydantic",
          ],
          "staging_bucket": STAGING_BUCKET,
          "extra_packages": [
              "installation_scripts/install.sh",
          ],
      },
  )
  print("Deployed agent: ", remote_agent.api_resource.name)

  try:
      async for event in remote_agent.async_stream_query(
          user_id="threat_detection_tester",
          message="Run the threat detection test",
      ):
        print(event)
  finally:
      client.agent_engines.delete(name=remote_agent.api_resource.name, force=True)

if __name__ == "__main__":
  asyncio.run(main())

Set up virtual environment

Create and activate a Python virtual environment.

  python3 -m venv env
  source env/bin/activate

Install necessary dependencies in the virtual environment.
```
pip install -r requirements.txt
```

Test the script

Run the script from inside the virtual environment with python3 main.py. This command will take several minutes to build, deploy, and execute the test agent.

The script outputs the resource name of the deployed agent, and a few JSON objects including an LLM response and other metadata. If you encounter permission errors or deployment errors at this stage, consult troubleshooting for next steps.

Test detectors

To test Agent Engine Threat Detection detectors, replace the code in the threat_detection_test function with code that simulates an attack. The script can take a long time to deploy and query the agent. To speed up testing, you can combine the contents of several of these functions together.