aap.py

"""
Agent Ad Protocol (AAP) — reference implementation.
Drop-in ad support for any MCP server.

Usage:
    from aap import AAPProvider
    provider = AAPProvider.from_json("example_ads.json")
    enriched = provider.inject_ad(tool_response, context_category="finance")
"""

import hashlib
import hmac
import json
import time
import uuid
from dataclasses import dataclass, field
from typing import Any

# Shared secret — in production, load from env / config.
_DEFAULT_SECRET = "aap-demo-secret-change-me"


@dataclass
class Ad:
    id: str
    advertiser: str
    headline: str
    body: str
    cta_url: str
    categories: list[str]
    cpm_cents: int = 100  # cost per 1000 impressions, in cents
    daily_budget: int = 1000  # max impressions/day
    active: bool = True


@dataclass
class Impression:
    token: str
    ad_id: str
    timestamp: float
    verified: bool = False


class AAPProvider:
    """Manages ad inventory, injection, and verification."""

    def __init__(self, ads: list[Ad] | None = None, secret: str = _DEFAULT_SECRET,
                 verify_url: str = "http://localhost:5100/verify"):
        self.ads: list[Ad] = ads or []
        self.secret = secret.encode()
        self.verify_url = verify_url
        # In-memory stores (swap for a DB in production)
        self.impressions: dict[str, Impression] = {}
        self.daily_counts: dict[str, int] = {}  # ad_id -> count today

    # --- Loading -----------------------------------------------------------

    @classmethod
    def from_json(cls, path: str, **kwargs) -> "AAPProvider":
        with open(path) as f:
            raw = json.load(f)
        ads = [Ad(**entry) for entry in raw]
        return cls(ads=ads, **kwargs)

    # --- Ad selection ------------------------------------------------------

    def _eligible(self, category: str) -> list[Ad]:
        return [
            ad for ad in self.ads
            if ad.active
            and category in ad.categories
            and self.daily_counts.get(ad.id, 0) < ad.daily_budget
        ]

    def select_ad(self, category: str) -> Ad | None:
        eligible = self._eligible(category)
        if not eligible:
            return None
        # Simple: pick highest CPM.  Production: auction / weighted random.
        return max(eligible, key=lambda a: a.cpm_cents)

    # --- Token generation --------------------------------------------------

    def _sign(self, payload: str) -> str:
        return hmac.new(self.secret, payload.encode(), hashlib.sha256).hexdigest()[:16]

    def _make_token(self, ad_id: str) -> str:
        uid = uuid.uuid4().hex[:12]
        ts = str(int(time.time()))
        payload = f"{uid}.{ad_id}.{ts}"
        sig = self._sign(payload)
        return f"{payload}.{sig}"

    # --- Injection ---------------------------------------------------------

    def inject_ad(self, response: dict[str, Any], context_category: str) -> dict[str, Any]:
        """Add an `_aap` field to a tool response dict. Returns the response unchanged if no ad matches."""
        ad = self.select_ad(context_category)
        if ad is None:
            return response

        token = self._make_token(ad.id)
        self.impressions[token] = Impression(token=token, ad_id=ad.id, timestamp=time.time())
        self.daily_counts[ad.id] = self.daily_counts.get(ad.id, 0) + 1

        response["_aap"] = {
            "v": 1,
            "ad": {
                "advertiser": ad.advertiser,
                "headline": ad.headline,
                "body": ad.body,
                "cta_url": ad.cta_url,
            },
            "impression_token": token,
            "verify_url": f"{self.verify_url}?token={token}",
            "expires": int(time.time()) + 3600,
        }
        return response

    # --- Verification ------------------------------------------------------

    def verify_impression(self, token: str) -> dict[str, Any]:
        """Validate a token and mark the impression as verified."""
        parts = token.split(".")
        if len(parts) != 4:
            return {"ok": False, "error": "malformed token"}

        uid, ad_id, ts_str, sig = parts

        # Check signature
        payload = f"{uid}.{ad_id}.{ts_str}"
        expected = self._sign(payload)
        if not hmac.compare_digest(sig, expected):
            return {"ok": False, "error": "invalid signature"}

        # Check expiry (1 hour)
        try:
            ts = int(ts_str)
        except ValueError:
            return {"ok": False, "error": "invalid timestamp"}
        if time.time() - ts > 3600:
            return {"ok": False, "error": "token expired"}

        # Mark verified
        imp = self.impressions.get(token)
        if imp:
            imp.verified = True

        return {"ok": True, "ad_id": ad_id, "verified": True}

    # --- Stats -------------------------------------------------------------

    def stats(self) -> dict[str, Any]:
        total = len(self.impressions)
        verified = sum(1 for i in self.impressions.values() if i.verified)
        return {
            "total_impressions": total,
            "verified_impressions": verified,
            "verification_rate": round(verified / total, 2) if total else 0,
            "daily_counts": dict(self.daily_counts),
        }