Skip to main content

Rubrik Agent Cloud policy enforcement plugin for AI agent frameworks

Project description

rubrik-agent-cloud-policy-plugin

Lightweight Python package for enforcing Rubrik Agent Cloud tool-blocking policies in AI agent frameworks. Framework-agnostic core with a Google ADK adapter.

Install

pip install rubrik-agent-cloud-policy-plugin[adk]

Only dependencies: httpx + google-adk (which you already have).

Quick Start

Option A: Explicit plugin (one line)

from google.adk.agents import Agent
from google.adk.runners import Runner
from rubrik_agent_cloud_policy.adk import RubrikPolicyPlugin

runner = Runner(
    agent=my_agent,
    app_name="my_app",
    session_service=session_service,
    plugins=[RubrikPolicyPlugin()],
)

Option B: Auto-instrumentation (zero code changes)

from rubrik_agent_cloud_policy.auto_instrument import auto_instrument
auto_instrument()  # patches all future Runner instances

For fully zero-code instrumentation, use sitecustomize.py — it runs before any user code in all Python invocations. Add this to your site-packages/sitecustomize.py:

import os
if os.getenv("RUBRIK_AUTO_INSTRUMENT", "").lower() in ("true", "1"):
    from rubrik_agent_cloud_policy.auto_instrument import auto_instrument
    auto_instrument()

Then set the env var:

export RUBRIK_AUTO_INSTRUMENT=true

Note: PYTHONSTARTUP does not work — it only runs for interactive Python sessions, not for python3 script.py or Agent Engine containers. Use sitecustomize.py instead.

Option C: ADK_EXTRA_PLUGINS (zero code changes)

export ADK_EXTRA_PLUGINS='["rubrik_agent_cloud_policy.adk.RubrikPolicyPlugin"]'

All options read configuration from environment variables:

export RUBRIK_WEBHOOK_URL="https://webhooks.example.com/{tenant_short_code}/{webhook_uuid}/"
export RUBRIK_API_KEY="your-api-key"
export RUBRIK_POLICY_FAIL_OPEN="true"                  # default: true (fail-open)

Webhook URL Format

The webhook URL must include the tenant short code and webhook UUID in the path:

https://webhooks.{domain}/{tenant_short_code}/{webhook_uuid}/

This is the per-tenant webhook URL returned when a webhook is created via the Gateway API. The Istio AuthorizationPolicy only allows POST to specific endpoint paths — GET requests (including health checks) return 403.

How It Works

  1. LLM responds with tool calls (e.g. get_weather, delete_file)
  2. after_model_callback intercepts the response
  3. Plugin translates Gemini function calls to OpenAI format
  4. POSTs to the Rubrik webhook (/v1/after_completion/openai/v1)
  5. Webhook evaluates against configured policies, returns allowed tools
  6. Blocked tools are stripped from the response; explanation text is injected
  7. Agent receives only the allowed tool calls

If the webhook is unreachable, the plugin fails open by default (all tools allowed). Set RUBRIK_POLICY_FAIL_OPEN=false to fail closed.

before_tool_callback acts as a safety net for tools that bypass after_model_callback (e.g. when another plugin short-circuits via before_model_callback).

Deploy to Vertex AI Agent Engine

1. Build the wheel

cd policy-plugin
pip install build
python -m build
# produces dist/rubrik_agent_cloud_policy_plugin-0.1.1-py3-none-any.whl

2. Upload to GCS (Agent Engine needs HTTP access during build)

gsutil cp dist/rubrik_agent_cloud_policy_plugin-0.1.1-py3-none-any.whl gs://{bucket}/packages/
gsutil acl ch -u AllUsers:R gs://{bucket}/packages/rubrik_agent_cloud_policy_plugin-0.1.1-py3-none-any.whl

3. Deploy

import vertexai
from vertexai import agent_engines
from rubrik_agent_cloud_policy.adk import RubrikPolicyPlugin

vertexai.init(project="my-project", location="us-central1",
              staging_bucket="gs://my-bucket")

adk_app = agent_engines.AdkApp(
    agent=my_agent,
    plugins=[RubrikPolicyPlugin()],
)

WHEEL_URL = "https://storage.googleapis.com/{bucket}/packages/rubrik_agent_cloud_policy_plugin-0.1.1-py3-none-any.whl"

remote_app = agent_engines.create(
    agent_engine=adk_app,
    display_name="my-agent",
    requirements=[
        f"rubrik-agent-cloud-policy-plugin[adk] @ {WHEEL_URL}",
        "google-cloud-aiplatform>=1.142.0",
        "google-genai>=1.51.0",
    ],
    env_vars={
        "RUBRIK_WEBHOOK_URL": "https://webhooks.example.com/{tenant}/{webhook}/",
    },
)

Important: Do NOT set GOOGLE_API_KEY in env_vars — Agent Engine uses Vertex AI service account auth. Setting an API key conflicts with the session service (Project/location and API key are mutually exclusive).

cloudpickle note

Agent Engine serializes plugins locally with cloudpickle and deserializes on the remote container. RubrikPolicyClient implements __reduce__ so env vars like RUBRIK_WEBHOOK_URL are re-read on the remote side — not baked in from the local environment.

Direct Client Usage (Framework-Agnostic)

from rubrik_agent_cloud_policy import RubrikPolicyClient, ToolCall

client = RubrikPolicyClient(webhook_url="https://webhooks.example.com/tenant/webhook/")

result = await client.evaluate_tool_calls([
    ToolCall(name="get_weather", arguments='{"city": "London"}'),
    ToolCall(name="delete_file", arguments='{"path": "/etc/passwd"}'),
])

print(result.allowed)      # [ToolCall(name='get_weather', ...)]
print(result.blocked)      # [ToolCall(name='delete_file', ...)]
print(result.explanation)  # "Tool 'delete_file' blocked by Rubrik Agent Cloud"

Development

cd policy-plugin
uv venv --python 3.12 .venv && source .venv/bin/activate
uv pip install -e ".[dev]"
pytest tests/ -v   # 44 tests

Architecture

rubrik_agent_cloud_policy/
  client.py            # Framework-agnostic core (webhook client + format translation)
  adk.py               # Google ADK BasePlugin adapter
  auto_instrument.py   # Monkey-patch Runner for zero-code instrumentation
  # Future:
  # langchain.py       # LangChain callback handler
  # crewai.py          # CrewAI adapter

The core client handles all webhook communication and OpenAI format translation. Framework adapters are thin wrappers that convert framework-specific tool representations to ToolCall objects and apply PolicyResult back.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

rubrik_agent_cloud_policy_plugin-0.1.1.tar.gz (15.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

File details

Details for the file rubrik_agent_cloud_policy_plugin-0.1.1.tar.gz.

File metadata

File hashes

Hashes for rubrik_agent_cloud_policy_plugin-0.1.1.tar.gz
Algorithm Hash digest
SHA256 f509b0e0bd1da90b32ed3889f530b9fd4cb7660b26832e526c56a7e9d8cdfdcc
MD5 e587bc94680ca943ab8973000646c6e8
BLAKE2b-256 44911d38be36e5c7bcaf74a053785256c80366f2f279d86a9d02974a7e1d4f1a

See more details on using hashes here.

Provenance

The following attestation bundles were made for rubrik_agent_cloud_policy_plugin-0.1.1.tar.gz:

Publisher: policy-plugin-release.yml on rubrikinc/rubrik-agent-cloud

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file rubrik_agent_cloud_policy_plugin-0.1.1-py3-none-any.whl.

File metadata

File hashes

Hashes for rubrik_agent_cloud_policy_plugin-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 eb093f1a4af8283d081a87ed93766e54846d89368da8903d81b07531dd7f86e7
MD5 7f2e9ccb86bc8151a6948d7e0aaaf481
BLAKE2b-256 8a9d9de89dffd98c14bd204f9adbbc82e371a796af1a2a15d3b09aaf404d00ad

See more details on using hashes here.

Provenance

The following attestation bundles were made for rubrik_agent_cloud_policy_plugin-0.1.1-py3-none-any.whl:

Publisher: policy-plugin-release.yml on rubrikinc/rubrik-agent-cloud

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page