---
title: Building Pi Extensions
---
# 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

```mermaid
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.

```python
# 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.

```python
# 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.

```python
# 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.

```python
# 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:

```python
@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

1. **Stateless code, stateful data**: Extensions should be stateless classes. Store state in Rivet durable state.

2. **Initialize in `on_init`**: Set up connections (MCP, databases) during actor initialization.

3. **Clean up in `on_terminate`**: Close connections when actor terminates.

4. **Use per-actor instances**: Each actor gets its own extension instances (no sharing across actors).

5. **Handle failures gracefully**: Extensions should catch exceptions and not crash the actor.

6. **Log important events**: Use structured logging for debugging and monitoring.

## Extension Testing

Test extensions in isolation before integrating into actors:

```python
# 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'
```

## See Also

- [Pi Agent Setup](/pi-agent-setup) — Pi Agent Core integration
- [Pi Agent Core API Reference](/pi-agent-core-api-reference) — Tool and agent API
- [Rivet Actor Model](/rivet-actor-model) — Actor implementation
- [Concurrency & Security](/concurrency-security) — Security details