Ran via tools/grade_run.py against LM Studio (no clipboard). Results: TTS: 80 Minor Flaws PASSES (real N-worker concurrency) <- Qwen 49, didn't parse Rust: 72 Minor Flaws COMPILES CLEAN (0 errs w/ deps) <- Qwen 50, 7 real errors Webhook: 55 Critical uses forbidden aiohttp (won't run) <- Qwen 75, passed Automation: 48 Critical SyntaxError (global-after-assign) <- first run for both DECISIVE head-to-head: Gemma generalizes where Qwen fails (TTS, Rust), but Qwen beats it on stdlib-discipline prompts (webhook). The two are COMPLEMENTARY local offloads, not redundant. Fixed: grade_run.py extractor (markdown/prose wrapping, multi-fence lang selection), TTFT-null handling in generator. TTFT capture from LM Studio API still needs the right stats key (left null + noted). Co-Authored-By: Claude <noreply@anthropic.com>
276 lines
9.0 KiB
Python
276 lines
9.0 KiB
Python
import asyncio
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import time
|
|
import random
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
from typing import Dict, Optional
|
|
|
|
import aiohttp
|
|
from aiohttp import web
|
|
|
|
# --- CONFIGURATION ---
|
|
SHARED_SECRET = "super_secret_twitch_key"
|
|
DISCORD_WEBHOOK_URL = "https://discord.com/api/webhooks/mock"
|
|
RATE_LIMIT_CAPACITY = 5
|
|
RATE_LIMIT_REFILL_PER_SEC = 5
|
|
IDEMPOTENCY_TTL_SECONDS = 300 # 5 minutes
|
|
|
|
# --- MODELS & STATE ---
|
|
|
|
@dataclass
|
|
class DiscordResponse:
|
|
status: int
|
|
retry_after: Optional[float] = None
|
|
|
|
class TokenBucket:
|
|
"""Implements the Token Bucket algorithm for rate limiting."""
|
|
def __init__(self, capacity: float, refill_rate: float):
|
|
self.capacity = capacity
|
|
self.refill_rate = refill_rate
|
|
self.tokens = capacity
|
|
self.last_update = time.monotonic()
|
|
self.lock = asyncio.Lock()
|
|
|
|
async def consume(self):
|
|
async with self.lock:
|
|
while self.tokens < 1:
|
|
# Refill tokens based on time passed
|
|
now = time.monotonic()
|
|
passed = now - self.last_update
|
|
self.tokens = min(self.capacity, self.tokens + (passed * self.refill_rate))
|
|
self.last_update = now
|
|
|
|
if self.tokens < 1:
|
|
# Wait for enough time to pass to get at least 1 token
|
|
wait_time = (1 - self.tokens) / self.refill_rate
|
|
await asyncio.sleep(wait_time)
|
|
|
|
self.tokens -= 1
|
|
|
|
class IdempotencyStore:
|
|
"""In-memory store to prevent duplicate processing."""
|
|
def __init__(self, ttl: int):
|
|
self.store: Dict[str, float] = {}
|
|
self.ttl = ttl
|
|
|
|
def is_duplicate(self, event_id: str) -> bool:
|
|
now = time.time()
|
|
if event_id in self.store:
|
|
if now - self.store[event_id] < self.ttl:
|
|
return True
|
|
self.store[event_id] = now
|
|
return False
|
|
|
|
async def cleanup_loop(self):
|
|
"""Background task to evict stale entries."""
|
|
while True:
|
|
await asyncio.sleep(60)
|
|
now = time.time()
|
|
expired = [k for k, v in self.store.items() if now - v > self.ttl]
|
|
for k in expired:
|
|
del self.store[k]
|
|
|
|
# --- MOCK DISCORD SERVICE ---
|
|
|
|
async def discord_send(payload: dict) -> DiscordResponse:
|
|
"""
|
|
Mocks a Discord webhook.
|
|
Succeeds ~95% of the time. Returns 429 occasionally.
|
|
"""
|
|
# Simulate network latency
|
|
await asyncio.sleep(0.05)
|
|
|
|
rand = random.random()
|
|
if rand < 0.05: # 5% chance of rate limit
|
|
return DiscordResponse(status=429, retry_after=1.0)
|
|
elif rand < 0.07: # 2% chance of other error
|
|
return DiscordResponse(status=500)
|
|
|
|
# Success
|
|
return DiscordResponse(status=204)
|
|
|
|
# --- CORE BRIDGE LOGIC ---
|
|
|
|
class WebhookBridge:
|
|
def __init__(self):
|
|
self.idempotency = IdempotencyStore(IDEMPOTENCY_TTL_SECONDS)
|
|
self.limiter = TokenBucket(RATE_LIMIT_CAPACITY, RATE_LIMIT_REFILL_PER_SEC)
|
|
self.queue = asyncio.Queue()
|
|
self.forward_count = 0 # For testing assertions
|
|
|
|
async def worker(self):
|
|
"""Background worker that processes the queue and respects rate limits."""
|
|
while True:
|
|
payload = await self.queue.get()
|
|
try:
|
|
await self._process_forward(payload)
|
|
except Exception as e:
|
|
print(f"[Worker Error] {e}")
|
|
finally:
|
|
self.queue.task_done()
|
|
|
|
async def _process_forward(self, payload: dict):
|
|
"""Handles the actual POST to Discord with retry logic."""
|
|
await self.limiter.consume()
|
|
|
|
resp = await discord_send(payload)
|
|
|
|
if resp.status == 429:
|
|
print(f"[Rate Limit] Discord returned 429. Retrying after {resp.retry_after}s")
|
|
await asyncio.sleep(resp.retry_after or 1)
|
|
# Retry once
|
|
resp = await discord_send(payload)
|
|
|
|
if resp.status in (200, 204):
|
|
self.forward_count += 1
|
|
else:
|
|
print(f"[Error] Discord failed with status {resp.status}")
|
|
|
|
async def handle_webhook(self, request: web.Request) -> web.Response:
|
|
"""HTTP Handler for POST /webhook."""
|
|
# 1. Verify Signature
|
|
signature_header = request.headers.get("X-Signature")
|
|
if not signature_header:
|
|
return web.Response(text="Missing X-Signature", status=401)
|
|
|
|
body = await request.read()
|
|
expected_sig = hmac.new(
|
|
SHARED_SECRET.encode(),
|
|
body,
|
|
hashlib.sha256
|
|
).hexdigest()
|
|
|
|
if not hmac.compare_digest(expected_sig, signature_header):
|
|
return web.Response(text="Invalid Signature", status=401)
|
|
|
|
# 2. Parse JSON
|
|
try:
|
|
data = json.loads(body)
|
|
event_id = data.get("event_id")
|
|
if not event_id:
|
|
return web.Response(text="Missing event_id", status=400)
|
|
except json.JSONDecodeError:
|
|
return web.Response(text="Invalid JSON", status=400)
|
|
|
|
# 3. Idempotency Check
|
|
if self.idempotency.is_duplicate(event_id):
|
|
return web.Response(text="Duplicate event", status=200)
|
|
|
|
# 4. Enqueue for forwarding
|
|
await self.queue.put(data)
|
|
return web.Response(text="Accepted", status=202)
|
|
|
|
# --- TEST SUITE ---
|
|
|
|
async def run_tests(bridge: WebhookBridge, server_url: str):
|
|
print("\n--- Starting Tests ---")
|
|
async with aiohttp.ClientSession() as session:
|
|
|
|
# Helper to sign requests
|
|
def sign(payload_dict):
|
|
body = json.dumps(payload_dict).encode()
|
|
sig = hmac.new(SHARED_SECRET.encode(), body, hashlib.sha256).hexdigest()
|
|
return body, sig
|
|
|
|
# a) Valid Request
|
|
print("Test A: Valid signature...")
|
|
payload_a = {"event_id": "evt_1", "type": "chat", "data": {"msg": "hi"}}
|
|
body, sig = sign(payload_a)
|
|
async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": sig}) as r:
|
|
assert r.status == 202
|
|
await asyncio.sleep(0.5) # Wait for worker
|
|
assert bridge.forward_count == 1
|
|
print("✅ Passed")
|
|
|
|
# b) Tampered Signature
|
|
print("Test B: Tampered signature...")
|
|
payload_b = {"event_id": "evt_2", "type": "chat"}
|
|
body, _ = sign(payload_b)
|
|
async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": "wrong_sig"}) as r:
|
|
assert r.status == 401
|
|
print("✅ Passed")
|
|
|
|
# c) Idempotency (Replay)
|
|
print("Test C: Replay event_id...")
|
|
payload_c = {"event_id": "evt_1", "type": "chat"} # Same ID as Test A
|
|
body, sig = sign(payload_c)
|
|
async with session.post(f"{server_url}/webhook", data=body, headers={"X-Signature": sig}) as r:
|
|
assert r.status == 200 # Returns 200 for duplicates per requirements
|
|
await asyncio.sleep(0.5)
|
|
assert bridge.forward_count == 1 # Count should NOT have increased
|
|
print("✅ Passed")
|
|
|
|
# d) Rate Limiting (Burst)
|
|
print("Test D: Burst > 5 events/sec...")
|
|
bridge.forward_count = 0 # Reset count for this test
|
|
start_time = time.monotonic()
|
|
|
|
# Send 10 events rapidly
|
|
tasks = []
|
|
for i in range(10):
|
|
p = {"event_id": f"burst_{i}", "type": "chat"}
|
|
b, s = sign(p)
|
|
tasks.append(session.post(f"{server_url}/webhook", data=b, headers={"X-Signature": s}))
|
|
|
|
await asyncio.gather(*tasks)
|
|
|
|
# Wait for all to be processed by worker
|
|
await bridge.queue.join()
|
|
duration = time.monotonic() - start_time
|
|
|
|
# 10 events at 5/sec should take at least ~1.8-2.0 seconds
|
|
# (First 5 instant, next 5 delayed by 1s each)
|
|
print(f" Burst of 10 took {duration:.2f}s")
|
|
assert duration >= 1.0
|
|
print("✅ Passed")
|
|
|
|
print("--- All Tests Passed Successfully ---\n")
|
|
|
|
# --- MAIN ENTRYPOINT ---
|
|
|
|
async def main():
|
|
bridge = WebhookBridge()
|
|
app = web.Application()
|
|
app.router.add_post('/webhook', bridge.handle_webhook)
|
|
|
|
runner = web.AppRunner(app)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, 'localhost', 8080)
|
|
await site.start()
|
|
|
|
print("Server started at http://localhost:8080")
|
|
print("Press Ctrl+C to stop.")
|
|
|
|
# Start background tasks
|
|
worker_task = asyncio.create_task(bridge.worker())
|
|
cleanup_task = asyncio.create_task(bridge.idempotency.cleanup_loop())
|
|
|
|
# Run tests automatically
|
|
try:
|
|
await run_tests(bridge, "http://localhost:8080")
|
|
except AssertionError as e:
|
|
print(f"❌ Test Failed!")
|
|
raise e
|
|
except Exception as e:
|
|
print(f"❌ Error during tests: {e}")
|
|
raise e
|
|
|
|
# Keep running until interrupted
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(3600)
|
|
except asyncio.CancelledError:
|
|
pass
|
|
finally:
|
|
worker_task.cancel()
|
|
cleanup_task.cancel()
|
|
await runner.cleanup()
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
pass |