diff --git a/.gitignore b/.gitignore index 412c02b..84e3190 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ .shsf.json -prompt.txt \ No newline at end of file +prompt.txt +__pycache__/ diff --git a/README.md b/README.md index 8da71e7..70187e1 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ The following environment variables are required for the service to function: | `DATABASE_STORAGE_NAME` | The name of the database storage to use. | | `AGENT_USERNAME` | The Gitea username of the AI agent (e.g., `whateveryouragentisnamed`). | | `AGENT_PROMPT_FILE` | (Optional) Path to a custom Markdown template for notifications. If not provided, a default template will be used. | +| `AGENT_HOURLY` | (Optional) Maximum number of triggers allowed per hour. Defaults to `60`. Used for rate limiting. | ## How it works diff --git a/main.py b/main.py index 6a587d2..29287ed 100644 --- a/main.py +++ b/main.py @@ -11,6 +11,7 @@ OPENCLAW_THINKING = os.getenv("OPENCLAW_THINKING", "low") DATABASE_STORAGE = os.getenv("DATABASE_STORAGE_NAME") AGENT_USERNAME = os.getenv("AGENT_USERNAME") AGENT_PROMPT_FILE = os.getenv("AGENT_PROMPT_FILE") +AGENT_HOURLY = int(os.getenv("AGENT_HOURLY", "60")) eventsToHandle = ["pull_request", "issues", "issue_comment"] actionsToHandle = ["assigned", "created"] @@ -76,6 +77,43 @@ def delete_stored_assignees(db, event_type, repository_id, number): print(f"Failed to delete assignees from DB: {e}") +def get_rate_limit_key(): + return f"agent_{AGENT_USERNAME}_cooldown" + + +def check_and_increment_rate_limit(db): + """Check if rate limit is reached, increment counter if not. Returns True if allowed, False if limit reached.""" + key = get_rate_limit_key() + try: + current = db.get(DATABASE_STORAGE, key) + if current is None: + current = 0 + else: + current = int(current) + + if current >= AGENT_HOURLY: + print(f"Rate limit reached: {current}/{AGENT_HOURLY} hourly triggers used.") + return False + + db.set(DATABASE_STORAGE, key, str(current + 1)) + print(f"Rate limit check passed: {current + 1}/{AGENT_HOURLY} hourly triggers used.") + return True + except Exception as e: + print(f"Failed to check/increment rate limit: {e}") + # On error, allow the request to proceed + return True + + +def reset_rate_limit(db): + """Reset the hourly rate limit counter.""" + key = get_rate_limit_key() + try: + db.set(DATABASE_STORAGE, key, "0") + print(f"Rate limit reset for agent {AGENT_USERNAME}.") + except Exception as e: + print(f"Failed to reset rate limit: {e}") + + def fill_template(template, event_object, action_str, type_str): fields = { "action_str": action_str, @@ -129,7 +167,7 @@ def build_message(event_object, action_str, type_str): return message -def sendToAgent(event_object): +def sendToAgent(event_object, db): headers = {"x-openclaw-token": OPENCLAW_TOKEN, "Content-Type": "application/json"} print(f"Preparing to send notification to Agent for {json.dumps(event_object)}") @@ -155,6 +193,11 @@ def sendToAgent(event_object): ) return + # Check rate limit before sending + if not check_and_increment_rate_limit(db): + print("Rate limit reached. Skipping notification to agent.") + return + message = build_message(event_object, action_str, type_str) try: @@ -355,7 +398,7 @@ def main(args): # Send to OpenClaw if action == "assigned": - sendToAgent(event_object) + sendToAgent(event_object, db) else: print(f"Action {action} is not configured to send to agent") @@ -381,7 +424,7 @@ def main(args): "url": comment_data.get("html_url"), } - sendToAgent(event_object) + sendToAgent(event_object, db) return { "_shsf": "v2", @@ -389,6 +432,24 @@ def main(args): "_res": {"status": "received"}, "_headers": {"Content-Type": "application/json"}, } + elif route == "clear_limit": + # Reset the hourly rate limit counter + try: + reset_rate_limit(db) + return { + "_shsf": "v2", + "_code": 200, + "_res": {"status": "Rate limit reset"}, + "_headers": {"Content-Type": "application/json"}, + } + except Exception as e: + print(f"Failed to reset rate limit: {e}") + return { + "_shsf": "v2", + "_code": 500, + "_res": {"error": "Failed to reset rate limit"}, + "_headers": {"Content-Type": "application/json"}, + } else: return { "_shsf": "v2",