Agent Capability Toggles: Enabling and Disabling Tools Per Customer or Plan

Why Capability Toggles for Agents

SaaS products gate features by plan — free users get basic functionality, paid users unlock advanced capabilities. The same principle applies to AI agents. A free-tier agent might answer questions from a knowledge base. A pro-tier agent might also search the web, execute code, and analyze uploaded files. Capability toggles let you build one agent codebase that dynamically adjusts its powers based on who is using it.

This is different from feature flags. Feature flags control rollout of new features to everyone. Capability toggles control which features are available to specific customers based on their subscription plan.

Defining the Capability Model

Start by defining what capabilities exist and which plans include them.

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional


class Plan(Enum):
    FREE = "free"
    PRO = "pro"
    ENTERPRISE = "enterprise"


@dataclass
class Capability:
    name: str
    display_name: str
    description: str
    min_plan: Plan
    daily_limit: Optional[int] = None
    requires_setup: bool = False


CAPABILITIES = {
    "knowledge_search": Capability(
        name="knowledge_search",
        display_name="Knowledge Base Search",
        description="Search through uploaded documents and FAQ.",
        min_plan=Plan.FREE,
    ),
    "web_search": Capability(
        name="web_search",
        display_name="Web Search",
        description="Search the internet for current information.",
        min_plan=Plan.PRO,
        daily_limit=100,
    ),
    "code_execution": Capability(
        name="code_execution",
        display_name="Code Execution",
        description="Run Python code to analyze data or perform calculations.",
        min_plan=Plan.PRO,
        daily_limit=50,
    ),
    "file_analysis": Capability(
        name="file_analysis",
        display_name="File Analysis",
        description="Upload and analyze CSV, Excel, and PDF files.",
        min_plan=Plan.PRO,
    ),
    "custom_tools": Capability(
        name="custom_tools",
        display_name="Custom API Integrations",
        description="Connect your own APIs as agent tools.",
        min_plan=Plan.ENTERPRISE,
    ),
    "multi_agent": Capability(
        name="multi_agent",
        display_name="Multi-Agent Workflows",
        description="Chain multiple specialized agents for complex tasks.",
        min_plan=Plan.ENTERPRISE,
    ),
}

PLAN_HIERARCHY = {Plan.FREE: 0, Plan.PRO: 1, Plan.ENTERPRISE: 2}

Capability Resolver

The resolver checks a customer's plan, their usage limits, and any custom overrides to determine which capabilities are available right now.

from datetime import date


@dataclass
class CustomerContext:
    customer_id: str
    plan: Plan
    custom_overrides: dict[str, bool] = field(default_factory=dict)
    usage_today: dict[str, int] = field(default_factory=dict)


class CapabilityResolver:
    def __init__(self, capabilities: dict[str, Capability]):
        self._capabilities = capabilities

    def resolve(self, ctx: CustomerContext) -> dict[str, bool]:
        result = {}
        for name, cap in self._capabilities.items():
            # Check custom overrides first
            if name in ctx.custom_overrides:
                result[name] = ctx.custom_overrides[name]
                continue

            # Check plan level
            if PLAN_HIERARCHY[ctx.plan] < PLAN_HIERARCHY[cap.min_plan]:
                result[name] = False
                continue

            # Check daily limits
            if cap.daily_limit is not None:
                used = ctx.usage_today.get(name, 0)
                if used >= cap.daily_limit:
                    result[name] = False
                    continue

            result[name] = True

        return result

    def get_upgrade_suggestions(self, ctx: CustomerContext) -> list[dict]:
        suggestions = []
        for name, cap in self._capabilities.items():
            if PLAN_HIERARCHY[ctx.plan] < PLAN_HIERARCHY[cap.min_plan]:
                suggestions.append({
                    "capability": cap.display_name,
                    "description": cap.description,
                    "required_plan": cap.min_plan.value,
                    "current_plan": ctx.plan.value,
                })
        return suggestions

Dynamic Tool List Building

When the agent starts a conversation, build its tool list based on the resolved capabilities. Tools that the customer cannot access are simply not registered.

from typing import Callable


# Simulated tool registry
TOOL_REGISTRY: dict[str, Callable] = {}


def register_tool(capability_name: str):
    def decorator(func: Callable) -> Callable:
        TOOL_REGISTRY[capability_name] = func
        return func
    return decorator


@register_tool("knowledge_search")
def search_knowledge_base(query: str) -> str:
    return f"Knowledge results for: {query}"


@register_tool("web_search")
def search_web(query: str) -> str:
    return f"Web results for: {query}"


@register_tool("code_execution")
def execute_code(code: str) -> str:
    return f"Executed code, output: ..."


def build_agent_tools(ctx: CustomerContext) -> list[Callable]:
    resolver = CapabilityResolver(CAPABILITIES)
    enabled = resolver.resolve(ctx)

    tools = []
    for cap_name, is_enabled in enabled.items():
        if is_enabled and cap_name in TOOL_REGISTRY:
            tools.append(TOOL_REGISTRY[cap_name])

    return tools

Upgrade Prompts in Agent Responses

When a user tries to use a capability they do not have access to, the agent should gracefully explain the limitation and suggest an upgrade rather than silently failing.

class CapabilityGatekeeper:
    def __init__(self, resolver: CapabilityResolver):
        self._resolver = resolver

    def check_or_suggest(self, ctx: CustomerContext, capability: str) -> dict:
        resolved = self._resolver.resolve(ctx)

        if resolved.get(capability, False):
            return {"allowed": True}

        cap = CAPABILITIES.get(capability)
        if not cap:
            return {"allowed": False, "reason": "Unknown capability"}

        # Check if it is a plan issue or a limit issue
        if PLAN_HIERARCHY[ctx.plan] < PLAN_HIERARCHY[cap.min_plan]:
            return {
                "allowed": False,
                "reason": "plan_upgrade_required",
                "message": (
                    f"{cap.display_name} is available on the "
                    f"{cap.min_plan.value.title()} plan and above. "
                    f"You are currently on the {ctx.plan.value.title()} plan."
                ),
                "upgrade_to": cap.min_plan.value,
            }

        if cap.daily_limit is not None:
            used = ctx.usage_today.get(capability, 0)
            if used >= cap.daily_limit:
                return {
                    "allowed": False,
                    "reason": "daily_limit_reached",
                    "message": (
                        f"You have used all {cap.daily_limit} "
                        f"{cap.display_name} requests for today. "
                        f"Limits reset at midnight UTC."
                    ),
                    "used": used,
                    "limit": cap.daily_limit,
                }

        return {"allowed": False, "reason": "unknown"}

Usage Tracking

Track capability usage per customer per day to enforce limits.

import redis
from datetime import datetime


class UsageTracker:
    def __init__(self, redis_client: redis.Redis):
        self._redis = redis_client

    def _key(self, customer_id: str, capability: str) -> str:
        today = datetime.utcnow().strftime("%Y-%m-%d")
        return f"usage:{customer_id}:{capability}:{today}"

    def increment(self, customer_id: str, capability: str) -> int:
        key = self._key(customer_id, capability)
        count = self._redis.incr(key)
        if count == 1:
            # Set expiry to end of day plus buffer
            self._redis.expire(key, 90000)  # 25 hours
        return count

    def get_usage(self, customer_id: str, capability: str) -> int:
        key = self._key(customer_id, capability)
        val = self._redis.get(key)
        return int(val) if val else 0

    def get_all_usage(self, customer_id: str) -> dict[str, int]:
        today = datetime.utcnow().strftime("%Y-%m-%d")
        result = {}
        for cap_name in CAPABILITIES:
            key = f"usage:{customer_id}:{cap_name}:{today}"
            val = self._redis.get(key)
            if val:
                result[cap_name] = int(val)
        return result

FAQ

How do I handle custom overrides for specific enterprise customers?

Store custom overrides in the database alongside the customer record. The CustomerContext.custom_overrides dict lets you enable capabilities that a plan would not normally include (for example, granting a free-tier customer temporary access to web search for a trial) or disable capabilities that the plan includes (for compliance reasons).

Should I remove tools from the agent entirely or just block them at runtime?

Remove them entirely. If you register a tool with the LLM but block it at runtime, the model might still try to call it and produce confusing error messages. By only registering tools the customer has access to, the model never even knows about unavailable capabilities. This produces cleaner conversations.

How do I handle mid-conversation plan upgrades?

Reload the customer context at the start of each conversation turn. If the customer upgrades mid-conversation, the new tools become available on their next message. You can also proactively notify them by including a system message like "New capabilities are now available" when the tool list changes between turns.

---

#CapabilityToggles #AIAgents #SaaSPlans #FeatureGating #Python #AgenticAI #LearnAI #AIEngineering