feat: add hourly rate limits for agent triggers
All checks were successful
Lint and Syntax Check / build (pull_request) Successful in 6s
All checks were successful
Lint and Syntax Check / build (pull_request) Successful in 6s
- Add AGENT_HOURLY env var (default: 60 triggers per hour)
- Track hourly usage in SHSF database with key 'agent_{AGENT_USERNAME}_cooldown'
- Check and increment counter before sending notifications
- Add 'clear_limit' route to reset counter hourly via external trigger
Resolves #4
This commit is contained in:
BIN
__pycache__/main.cpython-313.pyc
Normal file
BIN
__pycache__/main.cpython-313.pyc
Normal file
Binary file not shown.
67
main.py
67
main.py
@@ -11,6 +11,7 @@ OPENCLAW_THINKING = os.getenv("OPENCLAW_THINKING", "low")
|
|||||||
DATABASE_STORAGE = os.getenv("DATABASE_STORAGE_NAME")
|
DATABASE_STORAGE = os.getenv("DATABASE_STORAGE_NAME")
|
||||||
AGENT_USERNAME = os.getenv("AGENT_USERNAME")
|
AGENT_USERNAME = os.getenv("AGENT_USERNAME")
|
||||||
AGENT_PROMPT_FILE = os.getenv("AGENT_PROMPT_FILE")
|
AGENT_PROMPT_FILE = os.getenv("AGENT_PROMPT_FILE")
|
||||||
|
AGENT_HOURLY = int(os.getenv("AGENT_HOURLY", "60"))
|
||||||
|
|
||||||
eventsToHandle = ["pull_request", "issues", "issue_comment"]
|
eventsToHandle = ["pull_request", "issues", "issue_comment"]
|
||||||
actionsToHandle = ["assigned", "created"]
|
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}")
|
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):
|
def fill_template(template, event_object, action_str, type_str):
|
||||||
fields = {
|
fields = {
|
||||||
"action_str": action_str,
|
"action_str": action_str,
|
||||||
@@ -129,7 +167,7 @@ def build_message(event_object, action_str, type_str):
|
|||||||
return message
|
return message
|
||||||
|
|
||||||
|
|
||||||
def sendToAgent(event_object):
|
def sendToAgent(event_object, db):
|
||||||
headers = {"x-openclaw-token": OPENCLAW_TOKEN, "Content-Type": "application/json"}
|
headers = {"x-openclaw-token": OPENCLAW_TOKEN, "Content-Type": "application/json"}
|
||||||
print(f"Preparing to send notification to Agent for {json.dumps(event_object)}")
|
print(f"Preparing to send notification to Agent for {json.dumps(event_object)}")
|
||||||
|
|
||||||
@@ -155,6 +193,11 @@ def sendToAgent(event_object):
|
|||||||
)
|
)
|
||||||
return
|
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)
|
message = build_message(event_object, action_str, type_str)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -355,7 +398,7 @@ def main(args):
|
|||||||
|
|
||||||
# Send to OpenClaw
|
# Send to OpenClaw
|
||||||
if action == "assigned":
|
if action == "assigned":
|
||||||
sendToAgent(event_object)
|
sendToAgent(event_object, db)
|
||||||
else:
|
else:
|
||||||
print(f"Action {action} is not configured to send to agent")
|
print(f"Action {action} is not configured to send to agent")
|
||||||
|
|
||||||
@@ -381,7 +424,7 @@ def main(args):
|
|||||||
"url": comment_data.get("html_url"),
|
"url": comment_data.get("html_url"),
|
||||||
}
|
}
|
||||||
|
|
||||||
sendToAgent(event_object)
|
sendToAgent(event_object, db)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"_shsf": "v2",
|
"_shsf": "v2",
|
||||||
@@ -389,6 +432,24 @@ def main(args):
|
|||||||
"_res": {"status": "received"},
|
"_res": {"status": "received"},
|
||||||
"_headers": {"Content-Type": "application/json"},
|
"_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:
|
else:
|
||||||
return {
|
return {
|
||||||
"_shsf": "v2",
|
"_shsf": "v2",
|
||||||
|
|||||||
Reference in New Issue
Block a user