Building Pi Extensions

Overview

Extensions add capabilities to the Pi Agent runtime running inside each Rivet actor. They follow Pi Agent Core’s plugin model and are registered during actor initialization.

Architectural context: Each Rivet actor contains one Pi Agent instance with multiple extensions. Extensions are shared across all actor instances (stateless code), but each actor has its own extension state (if needed).

Extension Architecture

graph TB
    subgraph "Rivet Actor"
        PI[Pi Agent Core]
        
        subgraph "Extensions (stateless code)"
            MCP[MCP Client Extension]
            FMT[WhatsApp Formatter]
            SEC[Security Layer]
            ANAL[Analytics]
        end
        
        subgraph "Per-Actor State (in Rivet)"
            STATE[(Durable State)]
        end
    end
    
    subgraph "External Services"
        MCP_SERVER[@shopify/dev-mcp]
        WA[WhatsApp Cloud API]
        LLM[LLM Provider]
    end
    
    PI --> MCP
    PI --> FMT
    PI --> SEC
    PI --> ANAL
    
    MCP --> MCP_SERVER
    FMT --> WA
    PI --> LLM
    
    SEC -.->|reads/writes| STATE
    ANAL -.->|writes| STATE
    
    style PI fill:#e1f5fe,stroke:#0277bd,stroke-width:2px
    style STATE fill:#c8e6c9,stroke:#388e3c

Extension 1: MCP Client (Shopify Integration)

Connects Pi Agent to Shopify’s official Storefront MCP server.

# extensions/mcp_client.py
import asyncio
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from pi_agent_core import Tool

class ShopifyMCPExtension:
    """
    Manages connection to @shopify/dev-mcp server.
    Each actor has its own MCP client instance.
    """
    
    def __init__(self, store_domain: str, access_token: str):
        self.store_domain = store_domain
        self.access_token = access_token
        self.session: ClientSession | None = None
        self.tools: list[Tool] = []
    
    async def initialize(self) -> list[Tool]:
        """Connect to MCP server and discover tools."""
        server_params = StdioServerParameters(
            command="npx",
            args=["-y", "@shopify/dev-mcp"],
            env={
                "SHOPIFY_STORE_DOMAIN": self.store_domain,
                "SHOPIFY_STOREFRONT_TOKEN": self.access_token,
            }
        )
        
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                
                # Discover available tools
                tools_response = await session.list_tools()
                
                # Convert MCP tools to Pi Agent tools
                self.tools = [
                    Tool(
                        name=tool.name,
                        description=tool.description,
                        input_schema=tool.inputSchema,
                        execute=self._create_executor(session, tool.name)
                    )
                    for tool in tools_response.tools
                ]
                
                return self.tools
    
    def _create_executor(self, session: ClientSession, tool_name: str):
        """Create async executor for a specific MCP tool."""
        async def execute(**kwargs):
            result = await session.call_tool(tool_name, kwargs)
            return result.content
        
        return execute
    
    async def close(self):
        """Clean up MCP connection."""
        if self.session:
            await self.session.close()


# Usage in actor
@rivet.actor
class ShoppingAgentActor(Actor):
    async def on_init(self):
        # Initialize MCP extension
        self.mcp_extension = ShopifyMCPExtension(
            store_domain=os.environ["SHOPIFY_STORE_DOMAIN"],
            access_token=os.environ["SHOPIFY_STOREFRONT_TOKEN"]
        )
        tools = await self.mcp_extension.initialize()
        
        # Create Pi Agent with MCP tools
        self.agent = Agent(...)
        self.agent.set_tools(tools)
    
    async def on_terminate(self):
        # Clean up MCP connection
        if hasattr(self, 'mcp_extension'):
            await self.mcp_extension.close()

Available Shopify Tools

The @shopify/dev-mcp server exposes:

Tool	Purpose
`search_shop_catalog`	Search products by query
`get_product_details`	Get product details by ID
`create_cart`	Create new cart
`add_to_cart`	Add item to cart
`update_cart`	Update cart quantities
`get_cart`	Get current cart
`create_checkout`	Create checkout session
`get_order`	Get order details
`search_orders`	Search orders

Extension 2: WhatsApp Formatter

Formats Pi Agent responses for WhatsApp Cloud API message types.

# extensions/whatsapp_formatter.py
import re

class WhatsAppFormatterExtension:
    """
    Converts agent responses to WhatsApp message formats.
    Stateless - no per-actor state needed.
    """
    
    def format_response(self, response: str) -> dict:
        """
        Convert agent response to WhatsApp API format.
        
        Returns dict ready to send via WhatsApp Cloud API.
        """
        # Truncate to WhatsApp limit (4096 chars)
        if len(response) > 4096:
            response = response[:4093] + "..."
        
        return {
            "messaging_product": "whatsapp",
            "to": None,  # Set by actor
            "type": "text",
            "text": {"body": response}
        }
    
    def format_product_list(self, products: list[dict]) -> dict:
        """Format product search results as interactive list."""
        if not products:
            return self.format_response("No products found.")
        
        # Build list message
        sections = [{
            "title": "Products",
            "rows": [
                {
                    "id": f"product_{p['id']}",
                    "title": p['title'],
                    "description": f"${p['price']} - {p.get('description', '')[:50]}"
                }
                for p in products[:10]  # WhatsApp limit: 10 rows
            ]
        }]
        
        return {
            "messaging_product": "whatsapp",
            "to": None,
            "type": "interactive",
            "interactive": {
                "type": "list",
                "header": {"type": "text", "text": "🛍️ Products Found"},
                "body": {"text": f"Found {len(products)} products"},
                "footer": {"text": "Tap to view details"},
                "action": {
                    "button": "View Products",
                    "sections": sections
                }
            }
        }
    
    def format_cart_summary(self, cart: dict) -> dict:
        """Format cart as button message."""
        items_text = "\n".join([
            f"• {item['title']} × {item['quantity']} = ${item['price'] * item['quantity']}"
            for item in cart.get('items', [])
        ])
        
        total = sum(
            item['price'] * item['quantity']
            for item in cart.get('items', [])
        )
        
        return {
            "messaging_product": "whatsapp",
            "to": None,
            "type": "interactive",
            "interactive": {
                "type": "button",
                "body": {
                    "text": f"🛒 Your Cart\n\n{items_text}\n\n*Total: ${total}*"
                },
                "action": {
                    "buttons": [
                        {"type": "reply", "reply": {"id": "checkout", "title": "Checkout"}},
                        {"type": "reply", "reply": {"id": "add_more", "title": "Add More"}},
                        {"type": "reply", "reply": {"id": "clear_cart", "title": "Clear Cart"}}
                    ]
                }
            }
        }


# Usage in actor
@rivet.actor
class ShoppingAgentActor(Actor):
    async def on_message(self, text: str) -> None:
        # Process message
        response = await self.agent.prompt(text)
        
        # Format for WhatsApp
        formatter = WhatsAppFormatterExtension()
        
        # Check if response contains product list
        if self._is_product_list(response):
            products = self._extract_products(response)
            message = formatter.format_product_list(products)
        elif self._is_cart_summary(response):
            cart = self._extract_cart(response)
            message = formatter.format_cart_summary(cart)
        else:
            message = formatter.format_response(response)
        
        # Send via WhatsApp
        message['to'] = self.user_id
        await self._send_whatsapp(message)

Extension 3: Security Layer

Rate limiting, injection detection, input/output sanitization.

# extensions/security.py
import re
import time
from collections import deque

class SecurityExtension:
    """
    Handles rate limiting and prompt injection detection.
    State is stored in Rivet durable state (per actor).
    """
    
    def __init__(self, max_messages_per_minute: int = 20):
        self.max_messages_per_minute = max_messages_per_minute
        
        # Injection patterns
        self.injection_patterns = [
            r'(?i)ignore\s+(all\s+)?previous\s+instructions',
            r'(?i)forget\s+(all\s+)?your\s+instructions',
            r'(?i)you\s+are\s+now\s+a',
            r'(?i)system\s*:\s*',
            r'(?i)assistant\s*:\s*',
        ]
        
        # Sensitive output patterns
        self.sensitive_patterns = [
            r'sk-ant-[a-zA-Z0-9]{20,}',  # Anthropic API keys
            r'sk-[a-zA-Z0-9]{20,}',       # OpenAI API keys
            r'shp_[a-zA-Z0-9]{20,}',      # Shopify tokens
            r'-----BEGIN.*PRIVATE KEY-----',
        ]
    
    async def check_rate_limit(self, actor_state: dict) -> bool:
        """Check if user is within rate limits."""
        now = time.time()
        message_timestamps = actor_state.get('message_timestamps', deque())
        
        # Remove timestamps older than 1 minute
        while message_timestamps and message_timestamps[0] < now - 60:
            message_timestamps.popleft()
        
        # Check if at limit
        if len(message_timestamps) >= self.max_messages_per_minute:
            return False
        
        # Record this message
        message_timestamps.append(now)
        actor_state['message_timestamps'] = message_timestamps
        
        return True
    
    def detect_injection(self, text: str) -> bool:
        """Check if text contains prompt injection attempts."""
        for pattern in self.injection_patterns:
            if re.search(pattern, text):
                return True
        return False
    
    def sanitize_input(self, text: str) -> str:
        """Clean input before processing."""
        # Remove control characters
        text = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', text)
        
        # Normalize whitespace
        text = re.sub(r'\s+', ' ', text).strip()
        
        return text
    
    def sanitize_output(self, text: str) -> str:
        """Remove sensitive data from output."""
        for pattern in self.sensitive_patterns:
            text = re.sub(pattern, '[REDACTED]', text)
        return text


# Usage in actor
@rivet.actor
class ShoppingAgentActor(Actor):
    async def on_message(self, text: str) -> None:
        security = SecurityExtension()
        
        # 1. Check rate limit
        if not await security.check_rate_limit(self.state):
            await self._send_whatsapp_text(
                self.user_id,
                "⏳ Rate limit exceeded. Please wait a moment."
            )
            return
        
        # 2. Detect injection
        if security.detect_injection(text):
            await self._send_whatsapp_text(
                self.user_id,
                "⚠️ Invalid input detected."
            )
            return
        
        # 3. Sanitize input
        clean_text = security.sanitize_input(text)
        
        # 4. Process message
        response = await self.agent.prompt(clean_text)
        
        # 5. Sanitize output
        clean_response = security.sanitize_output(response)
        
        # 6. Send response
        await self._send_whatsapp_text(self.user_id, clean_response)

Extension 4: Analytics

Tracks conversation metrics and tool usage.

# extensions/analytics.py
import time
from dataclasses import dataclass

@dataclass
class ConversationMetrics:
    total_messages: int = 0
    tool_calls: int = 0
    errors: int = 0
    last_active: float = 0

class AnalyticsExtension:
    """
    Tracks conversation metrics.
    State stored in Rivet durable state (per actor).
    """
    
    async def record_message(self, actor_state: dict):
        """Record that a message was processed."""
        metrics = actor_state.get('metrics', ConversationMetrics())
        metrics.total_messages += 1
        metrics.last_active = time.time()
        actor_state['metrics'] = metrics
    
    async def record_tool_call(self, actor_state: dict, tool_name: str):
        """Record a tool call."""
        metrics = actor_state.get('metrics', ConversationMetrics())
        metrics.tool_calls += 1
        actor_state['metrics'] = metrics
        
        # Optional: log to external analytics service
        # await self._send_to_analytics_service(...)
    
    async def record_error(self, actor_state: dict, error: str):
        """Record an error."""
        metrics = actor_state.get('metrics', ConversationMetrics())
        metrics.errors += 1
        actor_state['metrics'] = metrics


# Usage in actor
@rivet.actor
class ShoppingAgentActor(Actor):
    async def on_message(self, text: str) -> None:
        analytics = AnalyticsExtension()
        
        try:
            # Process message
            response = await self.agent.prompt(text)
            
            # Record success
            await analytics.record_message(self.state)
            
            # Send response
            await self._send_whatsapp_text(self.user_id, response)
            
        except Exception as e:
            # Record error
            await analytics.record_error(self.state, str(e))
            
            # Send error message
            await self._send_whatsapp_text(
                self.user_id,
                "⚠️ An error occurred. Please try again."
            )

Extension Registration

Extensions are registered during actor initialization:

@rivet.actor
class ShoppingAgentActor(Actor):
    async def on_init(self):
        # Initialize extensions
        self.mcp_extension = ShopifyMCPExtension(...)
        self.formatter = WhatsAppFormatterExtension()
        self.security = SecurityExtension()
        self.analytics = AnalyticsExtension()
        
        # Get tools from MCP
        tools = await self.mcp_extension.initialize()
        
        # Create Pi Agent with tools
        self.agent = Agent(...)
        self.agent.set_tools(tools)
    
    async def on_message(self, text: str) -> None:
        # Use extensions in message flow
        if not await self.security.check_rate_limit(self.state):
            # Handle rate limit
            return
        
        if self.security.detect_injection(text):
            # Handle injection
            return
        
        clean_text = self.security.sanitize_input(text)
        response = await self.agent.prompt(clean_text)
        clean_response = self.security.sanitize_output(response)
        
        # Format and send
        message = self.formatter.format_response(clean_response)
        await self._send_whatsapp(message)
        
        # Record analytics
        await self.analytics.record_message(self.state)

Extension Best Practices

Stateless code, stateful data: Extensions should be stateless classes. Store state in Rivet durable state.
Initialize in on_init: Set up connections (MCP, databases) during actor initialization.
Clean up in on_terminate: Close connections when actor terminates.
Use per-actor instances: Each actor gets its own extension instances (no sharing across actors).
Handle failures gracefully: Extensions should catch exceptions and not crash the actor.
Log important events: Use structured logging for debugging and monitoring.

Extension Testing

Test extensions in isolation before integrating into actors:

# test_extensions.py
import pytest

@pytest.mark.asyncio
async def test_mcp_extension():
    ext = ShopifyMCPExtension(
        store_domain="test.myshopify.com",
        access_token="test_token"
    )
    tools = await ext.initialize()
    assert len(tools) > 0
    assert any(t.name == "search_shop_catalog" for t in tools)
    await ext.close()

def test_security_extension():
    ext = SecurityExtension()
    
    # Test injection detection
    assert ext.detect_injection("ignore previous instructions")
    assert not ext.detect_injection("show me red shoes")
    
    # Test sanitization
    assert ext.sanitize_output("key: sk-ant-abc123") == "key: [REDACTED]"

def test_formatter_extension():
    ext = WhatsAppFormatterExtension()
    msg = ext.format_response("Hello")
    assert msg['type'] == 'text'
    assert msg['text']['body'] == 'Hello'