Variet/gravity_control
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(bridge): resolve websocket zombie connection and bounding memory leaks

Browse Source
This commit is contained in:
Variet Worker
2026-03-23 21:11:52 +09:00
parent e21f71baf8
commit ecebec3906
10 changed files with 110 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
bot.py
View File

@@ -202,7 +202,7 @@ class GravityBot(commands.Bot):
self.conv_to_project: dict[str, str] = {} # conv_id → project
self.channel_to_project: dict[int, str] = {} # channel.id → project
self.session_status_messages: dict[str, int] = {} # conv_id → msg_id
self._sent_approval_ids: set[str] = set()
self._sent_approval_ids: dict[str, bool] = {} # request_id → bool
self._deferred_ids: dict[str, int] = {} # request_id → defer count
self._sent_commands: dict[str, str] = {} # request_id → command text (for MERGE edit detection)
self._ready_event = asyncio.Event()
@@ -255,6 +255,13 @@ class GravityBot(commands.Bot):
cmd_data["id"] = cmd_data.get("id", str(int(_time.time() * 1000)))
self.gateway.push_command(project, cmd_data)
def _cap_dict(self, d: dict, max_size: int = 5000):
"""Prevent memory leaks by capping dictionary sizes using insertion order (oldest first)."""
if len(d) >= max_size:
to_remove = len(d) - max_size + max_size // 10 # remove 10%
for k in list(d.keys())[:to_remove]:
d.pop(k, None)
@staticmethod
def _make_channel_name(project_name: str) -> str:
"""ag-gravity_control, ag-deriva, etc."""
@@ -650,11 +657,13 @@ class GravityBot(commands.Bot):
reject_commands = {"deny", "reject", "cancel", "decline", "dismiss", "stop"}
if req.command.strip().lower() in reject_commands:
logger.warning(f"Auto-approve BLOCKED: command='{req.command}' is reject-word — skipping")
self._sent_approval_ids.add(req.request_id)
self._cap_dict(self._sent_approval_ids)
self._sent_approval_ids[req.request_id] = True
phase1_processed += 1
continue
self._sent_approval_ids.add(req.request_id)
self._cap_dict(self._sent_approval_ids)
self._sent_approval_ids[req.request_id] = True
# Smart button_index: read buttons array from pending file
# file_permission buttons = [Allow Once(0), Allow This Conv(1), Deny(2)]
@@ -724,7 +733,9 @@ class GravityBot(commands.Bot):
channel = await self._get_channel(project)
if channel:
self._sent_approval_ids.add(req.request_id)
self._cap_dict(self._sent_approval_ids)
self._sent_approval_ids[req.request_id] = True
self._cap_dict(self._sent_commands)
self._sent_commands[req.request_id] = req.command
await self._send_approval_request(channel, req)
phase1_processed += 1
@@ -767,7 +778,7 @@ class GravityBot(commands.Bot):
self._deferred_ids.pop(rid, None)
self._sent_commands.pop(rid, None)
self._approval_messages.pop(rid, None)
self._sent_approval_ids.discard(rid)
self._sent_approval_ids.pop(rid, None)
phase2_processed += 1
elif status == "expired":
@@ -790,7 +801,7 @@ class GravityBot(commands.Bot):
f.unlink()
self._deferred_ids.pop(rid, None)
self._sent_commands.pop(rid, None)
self._sent_approval_ids.discard(rid)
self._sent_approval_ids.pop(rid, None)
phase2_processed += 1
elif status == "pending":
@@ -885,6 +896,7 @@ class GravityBot(commands.Bot):
pass
logger.info(f"Sent approval request: {request.request_id[:12]}")
self._cap_dict(self._approval_messages)
self._approval_messages[request.request_id] = msg.id # FIX #4: Track msg_id for auto_resolved lookup
# ─── Discord → IDE Text Relay + Multi-PC UX ───────────────────────────
@@ -1073,7 +1085,10 @@ class GravityBot(commands.Bot):
view = ApprovalView(self.bridge, request, buttons=buttons, hub=self.hub)
msg = await channel.send(embed=embed, view=view)
self._sent_approval_ids.add(request_id)
self._cap_dict(self._sent_approval_ids)
self._sent_approval_ids[request_id] = True
self._cap_dict(self._approval_messages)
self._approval_messages[request_id] = msg.id
logger.info(f"[HUB-PENDING] Sent approval: {request_id[:12]} project={project}")
@@ -1082,7 +1097,8 @@ class GravityBot(commands.Bot):
async def _auto_approve_via_hub(self, request: ApprovalRequest):
"""Auto-approve a pending request via Hub."""
self._sent_approval_ids.add(request.request_id)
self._cap_dict(self._sent_approval_ids)
self._sent_approval_ids[request.request_id] = True
delivered = False
if self.hub: