demo.py

"""
edgepilot demo
two agents, one stackql mcp server, two clouds

Query before mutate: the agents discover cluster ids, ruleset ids, and rule ids
via SELECT. Nothing about Confluent or Cloudflare ids is baked into env vars.
"""

import asyncio
import os
from datetime import datetime, timedelta, timezone
from anthropic import Anthropic
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

client = Anthropic()

RECON = """You are the recon agent for the edgepilot demo.

Use StackQL to gather three things:

1. The active Kafka cluster in Confluent environment {environment_id}:
     SELECT
       id                                                                              AS kafka_cluster_id,
       JSON_EXTRACT(spec, '$.region')                                                  AS region,
       LOWER(JSON_EXTRACT(spec, '$.cloud'))                                            AS cloud_provider,
       SPLIT_PART(SPLIT_PART(JSON_EXTRACT(spec, '$.http_endpoint'), '//', 2), '.', 1)  AS kafka_endpoint_id
     FROM confluent.managed_kafka_clusters.clusters
     WHERE environment = '{environment_id}'

2. Aggregated HTTP request analytics for Cloudflare zone {zone_id} (last 30 min):
     SELECT datetime, client_country_name, edge_response_status,
            client_request_http_method_name, requests, bytes
     FROM cloudflare.zones.http_requests_adaptive_groups
     WHERE zone_tag = '{zone_id}'
     AND since = '{since}'
     AND until = '{until}'
   Summarise: total requests, distinct countries, any non-2xx.

3. The current http_ratelimit phase ruleset for the zone:
     SELECT id,
            JSON_EXTRACT(rules, '$[0].id')                            AS rule_id,
            JSON_EXTRACT(rules, '$[0].description')                   AS description,
            JSON_EXTRACT(rules, '$[0].ratelimit.requests_per_period') AS threshold,
            JSON_EXTRACT(rules, '$[0].ratelimit.period')              AS period
     FROM cloudflare.rulesets.phases
     WHERE zone_id = '{zone_id}'
     AND ruleset_phase = 'http_ratelimit'

Return a brief plain-text summary. Include these values verbatim so the action
agent can reference them:
  - kafka_cluster_id, kafka_endpoint_id, region, cloud_provider
  - ruleset_id (the phase id), rule_id

Be terse."""

ACTION = """You are the action agent.

You will receive a recon report containing the values you need. Do exactly two
things using StackQL:

1. Tighten the rate limit rule by lowering its requests_per_period. The
   REPLACE replaces ALL rules in the http_ratelimit phase entrypoint with a
   single tightened rule:
     REPLACE cloudflare.rulesets.phases
     SET rules = '[
       {{
         "action": "block",
         "ratelimit": {{
           "characteristics": ["ip.src", "cf.colo.id"],
           "period": 10,
           "requests_per_period": 30,
           "mitigation_timeout": 10
         }},
         "expression": "(starts_with(http.request.uri.path, \\"/\\"))",
         "description": "edgepilot demo rule (managed by stackql-deploy)",
         "enabled": true
       }}
     ]'
     WHERE zone_id       = '{zone_id}'
     AND   ruleset_phase = 'http_ratelimit'

2. Log the decision to Kafka topic 'edge-autopilot-decisions'. Use the kafka
   cluster info from the recon report:
     INSERT INTO kafka.kafka.records(
       cluster_id, kafka_endpoint_id, region, cloud_provider,
       topic_name, key, value
     )
     SELECT
       '<kafka_cluster_id from recon>',
       '<kafka_endpoint_id from recon>',
       '<region from recon>',
       '<cloud_provider from recon>',
       'edge-autopilot-decisions',
       '{{"type":"STRING","data":"edgepilot.action"}}',
       '{{"type":"JSON","data":{{"agent":"action","action":"tighten_rate_limit","threshold":30,"reasoning":"<your reasoning>","timestamp":"<utc iso8601>"}}}}'

Be decisive. One REPLACE, one INSERT. No SELECTs."""


def ts():
    return datetime.now().strftime("%H:%M:%S.%f")[:11]


def log(agent, msg):
    print(f"[{ts()}]  {agent:<8}  {msg}")


async def run(name, session, tools, system, message):
    messages = [{"role": "user", "content": message}]

    while True:
        response = client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=1024,
            system=system,
            tools=tools,
            messages=messages,
        )

        tool_blocks = [b for b in response.content if b.type == "tool_use"]

        if not tool_blocks:
            for b in response.content:
                if hasattr(b, "text"):
                    log(name, b.text[:400])
                    return b.text
            return ""

        tool_results = []
        for block in tool_blocks:
            query_preview = str(block.input).replace("\n", " ")[:160]
            log(name, f"  {block.name}  {query_preview}")

            result = await session.call_tool(block.name, block.input)
            content = ""
            if result.content:
                first = result.content[0]
                content = first.text if hasattr(first, "text") else str(first)
            log(name, f"  -> {content[:200]}")

            tool_results.append({
                "type": "tool_result",
                "tool_use_id": block.id,
                "content": content,
            })

        messages.append({"role": "assistant", "content": response.content})
        messages.append({"role": "user", "content": tool_results})


async def main():
    zone_id        = os.environ["CLOUDFLARE_ZONE_ID"]
    environment_id = os.environ["CONFLUENT_ENVIRONMENT_ID"]

    until = datetime.now(timezone.utc).replace(microsecond=0)
    since = until - timedelta(minutes=30)
    since_s = since.isoformat().replace("+00:00", "Z")
    until_s = until.isoformat().replace("+00:00", "Z")

    server = StdioServerParameters(
        command="stackql",
        args=["mcp"],
        env={
            **os.environ,
            "CLOUDFLARE_API_TOKEN":       os.environ["CLOUDFLARE_API_TOKEN"],
            "CONFLUENT_CLOUD_API_KEY":    os.environ["CONFLUENT_CLOUD_API_KEY"],
            "CONFLUENT_CLOUD_API_SECRET": os.environ["CONFLUENT_CLOUD_API_SECRET"],
            "KAFKA_API_KEY":              os.environ["KAFKA_API_KEY"],
            "KAFKA_API_SECRET":           os.environ["KAFKA_API_SECRET"],
        },
    )

    print("\n=== edgepilot ===\n")

    async with stdio_client(server) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()

            tools_response = await session.list_tools()
            tools = [
                {
                    "name": t.name,
                    "description": t.description,
                    "input_schema": dict(t.inputSchema),
                }
                for t in tools_response.tools
            ]
            log("system", f"{len(tools)} stackql tools loaded\n")

            recon_output = await run(
                "recon",
                session,
                tools,
                RECON.format(
                    zone_id=zone_id,
                    environment_id=environment_id,
                    since=since_s,
                    until=until_s,
                ),
                f"Run recon on Cloudflare zone {zone_id} and Confluent environment {environment_id}.",
            )

            print()

            await run(
                "action",
                session,
                tools,
                ACTION.format(zone_id=zone_id),
                f"Recon report:\n\n{recon_output}",
            )

    print("\n=== done ===\n")


if __name__ == "__main__":
    asyncio.run(main())