Merge pull request 'feat: hourly rate limits for triggers' (#8) from feat/hourly-ratelimits into main
All checks were successful
Lint and Syntax Check / build (push) Successful in 5s
All checks were successful
Lint and Syntax Check / build (push) Successful in 5s
Reviewed-on: #8
This commit was merged in pull request #8.
This commit is contained in:
1
.gitignore
vendored
1
.gitignore
vendored
@@ -1,2 +1,3 @@
|
||||
.shsf.json
|
||||
prompt.txt
|
||||
__pycache__/
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
67
main.py
67
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",
|
||||
|
||||
Reference in New Issue
Block a user