refactor(extension): 모듈 분리 + Hub 통합 테스트 #task-395
- extension.ts 3,446→1,289줄 (-63%) - step-probe.ts (1,435줄): setupMonitor, processResponseFile, tryApprovalStrategies - observer-script.ts (687줄): DOM observer script - ws-client.ts (390줄): WSBridgeClient - step-utils.ts (114줄): step 파싱 유틸 - auth.py (115줄): JWT + registration code - hub.py (581줄): WSHub + per-client queue - Hub WS 연동 테스트 통과 (auth, chat, register) - VSIX v0.4.0 빌드
This commit is contained in:
Variet Worker
@@ -1,35 +1,59 @@
|
|||||||
# Architecture
|
# Architecture
|
||||||
|
|
||||||
> 이 프로젝트의 아키텍처를 설명하는 문서입니다.
|
|
||||||
> AI 에이전트는 구현 전 이 문서를 반드시 확인합니다.
|
> AI 에이전트는 구현 전 이 문서를 반드시 확인합니다.
|
||||||
|
|
||||||
## 프로젝트 개요
|
## 프로젝트 개요
|
||||||
|
|
||||||
<!-- 프로젝트의 목적과 핵심 기능을 간략히 서술 -->
|
Antigravity AI 코딩 에이전트의 Discord 연동 시스템.
|
||||||
|
- AG Extension ↔ WebSocket Hub ↔ Discord Bot (실시간)
|
||||||
(프로젝트 설명을 여기에 작성하세요)
|
- AG Extension ↔ 파일 bridge ↔ Collector ↔ Gateway ↔ Discord Bot (레거시)
|
||||||
|
|
||||||
## 디렉토리 구조
|
## 디렉토리 구조
|
||||||
|
|
||||||
```
|
```
|
||||||
project-root/
|
gravity_control/
|
||||||
├── src/ # 소스 코드
|
├── auth.py # JWT 토큰 관리
|
||||||
├── tests/ # 테스트
|
├── hub.py # WebSocket Hub (메시지 라우팅, 인스턴스 관리)
|
||||||
├── docs/ # 문서
|
├── bot.py # Discord 봇 (승인 UI, 채널 관리, Hub 핸들러)
|
||||||
├── .agents/ # AI 에이전트 설정
|
├── gateway.py # HTTP REST + /ws endpoint
|
||||||
└── ...
|
├── bridge.py # 파일 기반 IPC (레거시)
|
||||||
|
├── collector.py # 원격 파일→HTTP 릴레이 (Phase 2 삭제 예정)
|
||||||
|
├── watcher.py # Brain 디렉토리 변경 감시
|
||||||
|
├── config.py # 환경변수 설정
|
||||||
|
├── main.py # 진입점 (Bot + Hub + Watcher 시작)
|
||||||
|
├── parser.py # Markdown→Discord 파서
|
||||||
|
├── extension/src/
|
||||||
|
│ ├── extension.ts # 메인 Extension (3,300줄, 점진적 모듈화 진행)
|
||||||
|
│ ├── ws-client.ts # WSBridgeClient (Hub 연결, 재연결, 메시지 큐)
|
||||||
|
│ └── sdk/ # Antigravity SDK (로컬 임베드)
|
||||||
|
└── .agents/references/ # AI 에이전트 레퍼런스 문서
|
||||||
```
|
```
|
||||||
|
|
||||||
## 핵심 모듈
|
## 핵심 모듈
|
||||||
|
|
||||||
<!-- 각 모듈의 역할과 의존 관계를 설명 -->
|
|
||||||
|
|
||||||
| 모듈 | 역할 | 의존성 |
|
| 모듈 | 역할 | 의존성 |
|
||||||
|------|------|--------|
|
|------|------|--------|
|
||||||
| (모듈명) | (역할 설명) | (의존하는 모듈) |
|
| hub.py | WS 연결 관리, 메시지 라우팅, 인스턴스 번호 | auth.py |
|
||||||
|
| auth.py | JWT 토큰 생성/검증, registration code | - |
|
||||||
|
| bot.py | Discord UI, 승인 처리, 채팅 릴레이 | hub.py, bridge.py, parser.py |
|
||||||
|
| gateway.py | HTTP REST API + /ws endpoint | hub.py |
|
||||||
|
| ws-client.ts | Extension→Hub WS 클라이언트 | - |
|
||||||
|
| extension.ts | AG SDK 연동, step 감시, DOM observer | ws-client.ts, sdk/ |
|
||||||
|
|
||||||
## 데이터 흐름
|
## 데이터 흐름
|
||||||
|
|
||||||
<!-- 주요 데이터 흐름을 Mermaid 다이어그램이나 텍스트로 설명 -->
|
```
|
||||||
|
[Extension]
|
||||||
(데이터 흐름을 여기에 작성하세요)
|
│
|
||||||
|
┌────┴────┐
|
||||||
|
│ WS Hub │ ← 실시간 (preferred)
|
||||||
|
│ (ws-client.ts → hub.py)
|
||||||
|
└────┬────┘
|
||||||
|
│ ┌─────────────┐
|
||||||
|
├───────────────────→│ Discord Bot │→ Discord
|
||||||
|
│ └─────────────┘
|
||||||
|
┌────┴────┐
|
||||||
|
│파일 bridge│ ← 레거시 fallback
|
||||||
|
│(Collector → Gateway)
|
||||||
|
└─────────┘
|
||||||
|
```
|
||||||
|
|||||||
@@ -42,5 +42,8 @@
|
|||||||
| DISCORD_TOKEN | Discord 봇 토큰 | (필수) |
|
| DISCORD_TOKEN | Discord 봇 토큰 | (필수) |
|
||||||
| DISCORD_GUILD_ID | Discord 서버 ID | (필수) |
|
| DISCORD_GUILD_ID | Discord 서버 ID | (필수) |
|
||||||
| BRAIN_PATH | AG 브레인 경로 | `~/.gemini/antigravity/brain` |
|
| BRAIN_PATH | AG 브레인 경로 | `~/.gemini/antigravity/brain` |
|
||||||
| BOT_MODE | 봇 모드 (local/remote) | `local` |
|
| BOT_MODE | 봇 모드 (local/remote/gateway) | `local` |
|
||||||
| REMOTE_BRIDGE_URL | 원격 브릿지 URL | (remote 모드 전용) |
|
| REMOTE_BRIDGE_URL | 원격 브릿지 URL | (remote 모드 전용) |
|
||||||
|
| GATEWAY_API_KEY | Gateway REST API 인증 키 | (gateway 모드) |
|
||||||
|
| GRAVITY_HUB_SECRET | WS Hub JWT 서명 시크릿 | (자동생성 가능) |
|
||||||
|
| GRAVITY_REGISTRATION_CODE | Extension 등록 코드 | (미설정 시 인증 생략) |
|
||||||
|
|||||||
127
auth.py
Normal file
127
auth.py
Normal file
@@ -0,0 +1,127 @@
|
|||||||
|
"""Authentication module — JWT token management for WebSocket Hub.
|
||||||
|
|
||||||
|
Two-stage auth:
|
||||||
|
1. Extension connects with a registration code (built into .vsix at build time)
|
||||||
|
2. Hub validates the code and issues a short-lived JWT session token
|
||||||
|
3. Subsequent reconnections use the JWT token directly
|
||||||
|
|
||||||
|
The master secret is stored server-side only (GRAVITY_HUB_SECRET env var).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import hashlib
|
||||||
|
import hmac
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import time
|
||||||
|
from base64 import urlsafe_b64decode, urlsafe_b64encode
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Defaults
|
||||||
|
DEFAULT_TOKEN_TTL = 86400 # 24 hours
|
||||||
|
DEFAULT_REGISTRATION_CODE = "" # Set via GRAVITY_REGISTRATION_CODE env var
|
||||||
|
|
||||||
|
|
||||||
|
class TokenManager:
|
||||||
|
"""Manages JWT-like token creation and verification.
|
||||||
|
|
||||||
|
Uses HMAC-SHA256 for signing. Tokens contain:
|
||||||
|
- project: project name scope
|
||||||
|
- pc: PC identifier (hostname or custom name)
|
||||||
|
- iat: issued at (unix timestamp)
|
||||||
|
- exp: expiration (unix timestamp)
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, secret: str = "", registration_code: str = ""):
|
||||||
|
self.secret = secret or os.getenv("GRAVITY_HUB_SECRET", "")
|
||||||
|
self.registration_code = registration_code or os.getenv(
|
||||||
|
"GRAVITY_REGISTRATION_CODE", DEFAULT_REGISTRATION_CODE
|
||||||
|
)
|
||||||
|
if not self.secret:
|
||||||
|
# Auto-generate a secret if not set (ephemeral — tokens invalid after restart)
|
||||||
|
self.secret = hashlib.sha256(os.urandom(32)).hexdigest()
|
||||||
|
logger.warning(
|
||||||
|
"[AUTH] No GRAVITY_HUB_SECRET set — generated ephemeral secret. "
|
||||||
|
"Tokens will be invalid after server restart."
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_registration_code(self, code: str) -> bool:
|
||||||
|
"""Check if the provided registration code matches."""
|
||||||
|
if not self.registration_code:
|
||||||
|
# No registration code configured → allow all connections
|
||||||
|
logger.warning("[AUTH] No registration code configured — accepting all")
|
||||||
|
return True
|
||||||
|
return hmac.compare_digest(code, self.registration_code)
|
||||||
|
|
||||||
|
def create_token(
|
||||||
|
self, project: str, pc_name: str, ttl: int = DEFAULT_TOKEN_TTL
|
||||||
|
) -> str:
|
||||||
|
"""Create a signed token for a specific project and PC.
|
||||||
|
|
||||||
|
Returns a base64-encoded string: {header}.{payload}.{signature}
|
||||||
|
"""
|
||||||
|
now = int(time.time())
|
||||||
|
payload = {
|
||||||
|
"project": project,
|
||||||
|
"pc": pc_name,
|
||||||
|
"iat": now,
|
||||||
|
"exp": now + ttl,
|
||||||
|
}
|
||||||
|
payload_b64 = _b64_encode(json.dumps(payload))
|
||||||
|
signature = self._sign(payload_b64)
|
||||||
|
return f"{payload_b64}.{signature}"
|
||||||
|
|
||||||
|
def verify_token(self, token: str) -> dict | None:
|
||||||
|
"""Verify and decode a token.
|
||||||
|
|
||||||
|
Returns the payload dict if valid, None if invalid or expired.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
parts = token.split(".")
|
||||||
|
if len(parts) != 2:
|
||||||
|
return None
|
||||||
|
|
||||||
|
payload_b64, signature = parts
|
||||||
|
expected_sig = self._sign(payload_b64)
|
||||||
|
|
||||||
|
if not hmac.compare_digest(signature, expected_sig):
|
||||||
|
logger.warning("[AUTH] Invalid token signature")
|
||||||
|
return None
|
||||||
|
|
||||||
|
payload = json.loads(_b64_decode(payload_b64))
|
||||||
|
|
||||||
|
# Check expiration
|
||||||
|
if payload.get("exp", 0) < time.time():
|
||||||
|
logger.info(f"[AUTH] Token expired for {payload.get('pc', '?')}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
return payload
|
||||||
|
except (json.JSONDecodeError, ValueError, KeyError) as e:
|
||||||
|
logger.warning(f"[AUTH] Token decode error: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _sign(self, data: str) -> str:
|
||||||
|
"""HMAC-SHA256 sign and return base64."""
|
||||||
|
sig = hmac.new(
|
||||||
|
self.secret.encode("utf-8"),
|
||||||
|
data.encode("utf-8"),
|
||||||
|
hashlib.sha256,
|
||||||
|
).digest()
|
||||||
|
return _b64_encode_bytes(sig)
|
||||||
|
|
||||||
|
|
||||||
|
def _b64_encode(data: str) -> str:
|
||||||
|
"""URL-safe base64 encode a string, no padding."""
|
||||||
|
return urlsafe_b64encode(data.encode("utf-8")).rstrip(b"=").decode("ascii")
|
||||||
|
|
||||||
|
|
||||||
|
def _b64_encode_bytes(data: bytes) -> str:
|
||||||
|
"""URL-safe base64 encode bytes, no padding."""
|
||||||
|
return urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
||||||
|
|
||||||
|
|
||||||
|
def _b64_decode(data: str) -> str:
|
||||||
|
"""URL-safe base64 decode, handles missing padding."""
|
||||||
|
padded = data + "=" * (4 - len(data) % 4)
|
||||||
|
return urlsafe_b64decode(padded).decode("utf-8")
|
||||||
322
bot.py
322
bot.py
@@ -5,9 +5,15 @@ Multi-project channel architecture:
|
|||||||
- Each conversation maps to a project via conv_to_project dict
|
- Each conversation maps to a project via conv_to_project dict
|
||||||
- Extension registers projects via bridge/pending/ files
|
- Extension registers projects via bridge/pending/ files
|
||||||
- Commands include project_name for routing to correct IDE window
|
- Commands include project_name for routing to correct IDE window
|
||||||
|
|
||||||
|
Multi-PC UX:
|
||||||
|
- When multiple AG instances are active, messages get instance numbers (PC #1, #2)
|
||||||
|
- Users can target specific instances with !N <message> (e.g. !2 hello)
|
||||||
|
- When only one instance is active, natural conversation without numbers
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import re
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
@@ -184,17 +190,41 @@ class GravityBot(commands.Bot):
|
|||||||
self._processed_message_ids: deque[int] = deque(maxlen=200) # dedup for Gateway event replay
|
self._processed_message_ids: deque[int] = deque(maxlen=200) # dedup for Gateway event replay
|
||||||
self._approval_messages: dict[str, int] = {} # FIX #4: request_id → discord message_id (for auto_resolved lookup)
|
self._approval_messages: dict[str, int] = {} # FIX #4: request_id → discord message_id (for auto_resolved lookup)
|
||||||
self.gateway = None # Set by main.py in gateway mode
|
self.gateway = None # Set by main.py in gateway mode
|
||||||
|
self.hub = None # Set by main.py in gateway mode (WSHub instance)
|
||||||
|
|
||||||
def _write_command(self, project: str, text: str, **kwargs):
|
def _write_command(self, project: str, text: str, *,
|
||||||
"""Write command to bridge AND push to gateway (if gateway mode)."""
|
target_instance: int | None = None, **kwargs):
|
||||||
self.bridge.write_command(project, text, **kwargs)
|
"""Write command to bridge AND push to gateway/hub.
|
||||||
if self.gateway:
|
|
||||||
import time
|
Args:
|
||||||
self.gateway.push_command(project, {
|
target_instance: If set, send only to this instance number (via Hub).
|
||||||
"id": str(int(time.time() * 1000)),
|
If None, broadcast to all instances.
|
||||||
|
"""
|
||||||
|
cmd_data = {
|
||||||
"text": text,
|
"text": text,
|
||||||
"project_name": kwargs.get('project_name', project),
|
"project_name": kwargs.get('project_name', project),
|
||||||
})
|
}
|
||||||
|
|
||||||
|
# Hub route (preferred if available)
|
||||||
|
if self.hub:
|
||||||
|
import time as _time
|
||||||
|
cmd_data["id"] = str(int(_time.time() * 1000))
|
||||||
|
msg = {"type": "command", "data": cmd_data}
|
||||||
|
if target_instance is not None:
|
||||||
|
asyncio.create_task(
|
||||||
|
self.hub.send_to_instance(project, target_instance, msg)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
asyncio.create_task(
|
||||||
|
self.hub.broadcast_to_project(project, msg)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Legacy routes (file bridge + gateway HTTP)
|
||||||
|
self.bridge.write_command(project, text, **kwargs)
|
||||||
|
if self.gateway:
|
||||||
|
import time as _time
|
||||||
|
cmd_data["id"] = cmd_data.get("id", str(int(_time.time() * 1000)))
|
||||||
|
self.gateway.push_command(project, cmd_data)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _make_channel_name(project_name: str) -> str:
|
def _make_channel_name(project_name: str) -> str:
|
||||||
@@ -206,6 +236,8 @@ class GravityBot(commands.Bot):
|
|||||||
self.pending_approval_scanner.start()
|
self.pending_approval_scanner.start()
|
||||||
self.chat_snapshot_scanner.start()
|
self.chat_snapshot_scanner.start()
|
||||||
self._register_slash_commands()
|
self._register_slash_commands()
|
||||||
|
# Register Hub handlers (if Hub is available, set after setup_hook by main.py)
|
||||||
|
asyncio.get_event_loop().call_soon(self._register_hub_handlers)
|
||||||
logger.info("Bot setup complete")
|
logger.info("Bot setup complete")
|
||||||
|
|
||||||
def _register_slash_commands(self):
|
def _register_slash_commands(self):
|
||||||
@@ -826,7 +858,33 @@ class GravityBot(commands.Bot):
|
|||||||
logger.info(f"Sent approval request: {request.request_id[:12]}")
|
logger.info(f"Sent approval request: {request.request_id[:12]}")
|
||||||
self._approval_messages[request.request_id] = msg.id # FIX #4: Track msg_id for auto_resolved lookup
|
self._approval_messages[request.request_id] = msg.id # FIX #4: Track msg_id for auto_resolved lookup
|
||||||
|
|
||||||
# ─── Discord → IDE Text Relay ─────────────────────────────────────
|
# ─── Discord → IDE Text Relay + Multi-PC UX ───────────────────────────
|
||||||
|
|
||||||
|
def _get_instance_header(self, project: str, instance_number: int) -> str:
|
||||||
|
"""Format instance header based on active count.
|
||||||
|
|
||||||
|
Single instance: empty string (natural conversation)
|
||||||
|
Multiple instances: **[PC #N]** prefix
|
||||||
|
"""
|
||||||
|
if not self.hub:
|
||||||
|
return ""
|
||||||
|
active = self.hub.get_active_count(project)
|
||||||
|
if active <= 1:
|
||||||
|
return ""
|
||||||
|
return f"**[PC #{instance_number}]** "
|
||||||
|
|
||||||
|
def _parse_instance_target(self, text: str) -> tuple[int | None, str]:
|
||||||
|
"""Parse !N prefix from message text.
|
||||||
|
|
||||||
|
Returns (target_instance, remaining_text).
|
||||||
|
'!2 hello' -> (2, 'hello')
|
||||||
|
'hello' -> (None, 'hello')
|
||||||
|
'!stop' -> (None, '!stop') # special commands not treated as targeting
|
||||||
|
"""
|
||||||
|
match = re.match(r'^!(\d+)\s+(.+)', text, re.DOTALL)
|
||||||
|
if match:
|
||||||
|
return int(match.group(1)), match.group(2).strip()
|
||||||
|
return None, text
|
||||||
|
|
||||||
async def on_message(self, message: discord.Message):
|
async def on_message(self, message: discord.Message):
|
||||||
if message.author == self.user:
|
if message.author == self.user:
|
||||||
@@ -845,19 +903,24 @@ class GravityBot(commands.Bot):
|
|||||||
|
|
||||||
text = message.content.strip()
|
text = message.content.strip()
|
||||||
|
|
||||||
|
# Parse !N instance targeting (before special commands)
|
||||||
|
target_instance, actual_text = self._parse_instance_target(text)
|
||||||
|
|
||||||
# Special command: !stop — cancel AI work
|
# Special command: !stop — cancel AI work
|
||||||
if text == "!stop":
|
if actual_text == "!stop":
|
||||||
self._write_command(project, "!stop", project_name=project)
|
self._write_command(project, "!stop", target_instance=target_instance,
|
||||||
|
project_name=project)
|
||||||
|
target_label = f" (PC #{target_instance})" if target_instance else ""
|
||||||
embed = discord.Embed(
|
embed = discord.Embed(
|
||||||
title="⏹️ AI 작업 중지",
|
title="⏹️ AI 작업 중지",
|
||||||
description=f"프로젝트: **{project}**\n중지 요청을 Extension에 전달했습니다.",
|
description=f"프로젝트: **{project}**{target_label}\n중지 요청을 Extension에 전달했습니다.",
|
||||||
color=discord.Color.orange(),
|
color=discord.Color.orange(),
|
||||||
)
|
)
|
||||||
await message.channel.send(embed=embed)
|
await message.channel.send(embed=embed)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Special command: !auto — toggle auto-approve
|
# Special command: !auto — toggle auto-approve
|
||||||
if text == "!auto":
|
if actual_text == "!auto":
|
||||||
# Toggle per-project auto-approve
|
# Toggle per-project auto-approve
|
||||||
if project in self.auto_approve_projects:
|
if project in self.auto_approve_projects:
|
||||||
self.auto_approve_projects.discard(project)
|
self.auto_approve_projects.discard(project)
|
||||||
@@ -865,7 +928,8 @@ class GravityBot(commands.Bot):
|
|||||||
else:
|
else:
|
||||||
self.auto_approve_projects.add(project)
|
self.auto_approve_projects.add(project)
|
||||||
enabled = True
|
enabled = True
|
||||||
self._write_command(project, f"!auto {'on' if enabled else 'off'}", project_name=project)
|
self._write_command(project, f"!auto {'on' if enabled else 'off'}",
|
||||||
|
target_instance=target_instance, project_name=project)
|
||||||
emoji = "🟢" if enabled else "🔴"
|
emoji = "🟢" if enabled else "🔴"
|
||||||
mode = "자동 승인" if enabled else "수동 승인"
|
mode = "자동 승인" if enabled else "수동 승인"
|
||||||
embed = discord.Embed(
|
embed = discord.Embed(
|
||||||
@@ -877,18 +941,240 @@ class GravityBot(commands.Bot):
|
|||||||
await message.channel.send(embed=embed)
|
await message.channel.send(embed=embed)
|
||||||
return
|
return
|
||||||
|
|
||||||
# General text relay — routed by project
|
# General text relay — routed by project (+ optional instance targeting)
|
||||||
if text:
|
if actual_text:
|
||||||
self._write_command(project, text, project_name=project)
|
self._write_command(project, actual_text, target_instance=target_instance,
|
||||||
|
project_name=project)
|
||||||
await message.add_reaction("📨")
|
await message.add_reaction("📨")
|
||||||
|
target_label = f" PC #{target_instance}" if target_instance else ""
|
||||||
embed = discord.Embed(
|
embed = discord.Embed(
|
||||||
description=f"📨 → **{project}** IDE에 전달됨\n`{text[:100]}`",
|
description=f"📨 → **{project}**{target_label} IDE에 전달됨\n`{actual_text[:100]}`",
|
||||||
color=discord.Color.blurple(),
|
color=discord.Color.blurple(),
|
||||||
)
|
)
|
||||||
await message.channel.send(embed=embed, delete_after=10)
|
await message.channel.send(embed=embed, delete_after=10)
|
||||||
|
|
||||||
await self.process_commands(message)
|
await self.process_commands(message)
|
||||||
|
|
||||||
|
# ─── Hub Event Handlers ──────────────────────────────────────────
|
||||||
|
|
||||||
|
def _register_hub_handlers(self):
|
||||||
|
"""Register callbacks on the Hub for Extension->Bot messages."""
|
||||||
|
if not self.hub:
|
||||||
|
return
|
||||||
|
self.hub.set_bot_handlers(
|
||||||
|
on_pending=self._hub_on_pending,
|
||||||
|
on_chat=self._hub_on_chat,
|
||||||
|
on_register=self._hub_on_register,
|
||||||
|
on_auto_resolve=self._hub_on_auto_resolve,
|
||||||
|
on_brain_event=self._hub_on_brain_event,
|
||||||
|
)
|
||||||
|
logger.info("[BOT] Hub handlers registered")
|
||||||
|
|
||||||
|
async def _hub_on_pending(self, project: str, data: dict):
|
||||||
|
"""Handle pending approval from Hub (Extension->Hub->Bot)."""
|
||||||
|
try:
|
||||||
|
request_id = data.get("request_id", "")
|
||||||
|
if not request_id:
|
||||||
|
return
|
||||||
|
|
||||||
|
# Skip if already sent
|
||||||
|
if request_id in self._sent_approval_ids:
|
||||||
|
return
|
||||||
|
|
||||||
|
# Check auto_resolved status
|
||||||
|
status = data.get("status", "pending")
|
||||||
|
if status in ("auto_resolved", "expired"):
|
||||||
|
await self._handle_auto_resolved(request_id, status)
|
||||||
|
return
|
||||||
|
|
||||||
|
instance_number = data.get("_instance_number", 0)
|
||||||
|
pc_name = data.get("_pc_name", "")
|
||||||
|
header = self._get_instance_header(project, instance_number)
|
||||||
|
|
||||||
|
# Build approval request
|
||||||
|
request = ApprovalRequest(
|
||||||
|
request_id=request_id,
|
||||||
|
command=data.get("command", ""),
|
||||||
|
description=data.get("description", ""),
|
||||||
|
project_name=project,
|
||||||
|
step_type=data.get("step_type", ""),
|
||||||
|
status=status,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Auto-approve check
|
||||||
|
if project in self.auto_approve_projects:
|
||||||
|
await self._auto_approve_via_hub(request)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Send to Discord
|
||||||
|
channel = await self._get_channel(project)
|
||||||
|
if not channel:
|
||||||
|
logger.warning(f"[HUB-PENDING] No channel for project={project}")
|
||||||
|
return
|
||||||
|
|
||||||
|
buttons = data.get("buttons", [])
|
||||||
|
desc_parts = []
|
||||||
|
if header:
|
||||||
|
desc_parts.append(header)
|
||||||
|
desc_parts.append(f"**명령:** `{request.command[:200]}`")
|
||||||
|
if buttons:
|
||||||
|
btn_names = [b.get("text", "?") for b in buttons]
|
||||||
|
desc_parts.append(f"**선택지:** {' / '.join(btn_names)}")
|
||||||
|
if request.description:
|
||||||
|
desc_parts.append(request.description[:500])
|
||||||
|
|
||||||
|
embed = discord.Embed(
|
||||||
|
title="⚠️ 승인 요청",
|
||||||
|
description="\n".join(desc_parts),
|
||||||
|
color=discord.Color.orange(),
|
||||||
|
timestamp=datetime.now(timezone.utc),
|
||||||
|
)
|
||||||
|
embed.set_footer(text=f"ID: {request_id}")
|
||||||
|
|
||||||
|
view = ApprovalView(self.bridge, request, buttons=buttons)
|
||||||
|
msg = await channel.send(embed=embed, view=view)
|
||||||
|
|
||||||
|
self._sent_approval_ids.add(request_id)
|
||||||
|
self._approval_messages[request_id] = msg.id
|
||||||
|
logger.info(f"[HUB-PENDING] Sent approval: {request_id[:12]} project={project}")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[HUB-PENDING] Error: {e}")
|
||||||
|
|
||||||
|
async def _auto_approve_via_hub(self, request: ApprovalRequest):
|
||||||
|
"""Auto-approve a pending request via Hub."""
|
||||||
|
if self.hub:
|
||||||
|
await self.hub.send_response_to_pending_owner(request.request_id, {
|
||||||
|
"type": "response",
|
||||||
|
"data": {
|
||||||
|
"request_id": request.request_id,
|
||||||
|
"approved": True,
|
||||||
|
"button_index": 0,
|
||||||
|
"step_type": request.step_type,
|
||||||
|
"project_name": request.project_name,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
# Also write via legacy bridge
|
||||||
|
self.bridge.write_response(UserResponse(
|
||||||
|
request_id=request.request_id, approved=True,
|
||||||
|
step_type=request.step_type,
|
||||||
|
project_name=request.project_name,
|
||||||
|
))
|
||||||
|
logger.info(f"[HUB-AUTO] Auto-approved: {request.request_id[:12]}")
|
||||||
|
|
||||||
|
async def _hub_on_chat(self, project: str, data: dict):
|
||||||
|
"""Handle chat snapshot from Hub (Extension->Hub->Bot->Discord)."""
|
||||||
|
try:
|
||||||
|
content = data.get("content", "")
|
||||||
|
attached_files = data.get("attached_files", [])
|
||||||
|
if not content and not attached_files:
|
||||||
|
return
|
||||||
|
|
||||||
|
instance_number = data.get("_instance_number", 0)
|
||||||
|
header = self._get_instance_header(project, instance_number)
|
||||||
|
|
||||||
|
channel = await self._get_channel(project)
|
||||||
|
if not channel:
|
||||||
|
return
|
||||||
|
|
||||||
|
import io as _io
|
||||||
|
discord_files = []
|
||||||
|
for af in attached_files:
|
||||||
|
af_name = af.get("name", "document.md")
|
||||||
|
af_content = af.get("content", "")
|
||||||
|
if af_content:
|
||||||
|
discord_files.append(discord.File(
|
||||||
|
_io.BytesIO(af_content.encode("utf-8")),
|
||||||
|
filename=af_name,
|
||||||
|
))
|
||||||
|
|
||||||
|
display_content = f"{header}{content}" if header else content
|
||||||
|
|
||||||
|
FILE_ATTACH_THRESHOLD = 4000
|
||||||
|
if len(display_content) > FILE_ATTACH_THRESHOLD:
|
||||||
|
summary = display_content[:500].rsplit('\n', 1)[0]
|
||||||
|
embed = discord.Embed(
|
||||||
|
title="💬 AI 대화 내용",
|
||||||
|
description=f"{summary}\n\n📎 *전체 내용은 첨부 파일 참조* ({len(content):,}자)",
|
||||||
|
color=discord.Color.purple(),
|
||||||
|
timestamp=datetime.now(timezone.utc),
|
||||||
|
)
|
||||||
|
discord_files.append(discord.File(
|
||||||
|
_io.BytesIO(content.encode("utf-8")),
|
||||||
|
filename="chat_message.md",
|
||||||
|
))
|
||||||
|
await channel.send(embed=embed, files=discord_files)
|
||||||
|
else:
|
||||||
|
embed = discord.Embed(
|
||||||
|
title="💬 AI 대화 내용",
|
||||||
|
description=display_content,
|
||||||
|
color=discord.Color.purple(),
|
||||||
|
timestamp=datetime.now(timezone.utc),
|
||||||
|
)
|
||||||
|
await channel.send(
|
||||||
|
embed=embed,
|
||||||
|
files=discord_files if discord_files else discord.utils.MISSING,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(f"[HUB-CHAT] Sent to #{channel.name} ({len(content)} chars)")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[HUB-CHAT] Error: {e}")
|
||||||
|
|
||||||
|
async def _hub_on_register(self, data: dict):
|
||||||
|
"""Handle session registration from Hub."""
|
||||||
|
conv_id = data.get("conversation_id", "")
|
||||||
|
project = data.get("project_name", "")
|
||||||
|
if conv_id and project:
|
||||||
|
self.conv_to_project[conv_id] = project
|
||||||
|
logger.info(f"[HUB-REG] {conv_id[:8]} → {project}")
|
||||||
|
|
||||||
|
async def _hub_on_auto_resolve(self, project: str, data: dict):
|
||||||
|
"""Handle auto_resolve notification from Hub."""
|
||||||
|
request_id = data.get("request_id", "")
|
||||||
|
if request_id:
|
||||||
|
await self._handle_auto_resolved(request_id, "auto_resolved")
|
||||||
|
|
||||||
|
async def _hub_on_brain_event(self, project: str, data: dict):
|
||||||
|
"""Handle brain event from Hub (Extension->Hub->Bot->Discord)."""
|
||||||
|
try:
|
||||||
|
from watcher import BrainEvent, EventType
|
||||||
|
event = BrainEvent(
|
||||||
|
event_type=EventType(data.get("event_type", "file_changed")),
|
||||||
|
conversation_id=data.get("conversation_id", ""),
|
||||||
|
file_name=data.get("file_name", ""),
|
||||||
|
file_path=None,
|
||||||
|
content=data.get("content", ""),
|
||||||
|
timestamp=data.get("timestamp", time.time()),
|
||||||
|
)
|
||||||
|
await self.event_queue.put(event)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[HUB-EVENT] Error: {e}")
|
||||||
|
|
||||||
|
async def _handle_auto_resolved(self, request_id: str, status: str):
|
||||||
|
"""Edit Discord message to show auto-resolved/expired status."""
|
||||||
|
msg_id = self._approval_messages.get(request_id)
|
||||||
|
if not msg_id:
|
||||||
|
return
|
||||||
|
# Find the channel containing this message
|
||||||
|
for channel in self.project_channels.values():
|
||||||
|
try:
|
||||||
|
msg = await channel.fetch_message(msg_id)
|
||||||
|
embed = msg.embeds[0] if msg.embeds else None
|
||||||
|
if embed:
|
||||||
|
if status == "auto_resolved":
|
||||||
|
embed.color = discord.Color.green()
|
||||||
|
embed.set_footer(text="✅ 자동 해결됨")
|
||||||
|
else:
|
||||||
|
embed.color = discord.Color.greyple()
|
||||||
|
embed.set_footer(text="⏰ 만료됨")
|
||||||
|
await msg.edit(embed=embed, view=None)
|
||||||
|
self._approval_messages.pop(request_id, None)
|
||||||
|
break
|
||||||
|
except (discord.NotFound, discord.Forbidden):
|
||||||
|
continue
|
||||||
|
except Exception:
|
||||||
|
break
|
||||||
|
|
||||||
# ─── Chat Snapshot Scanner ─────────────────────────────────────────
|
# ─── Chat Snapshot Scanner ─────────────────────────────────────────
|
||||||
|
|
||||||
@tasks.loop(seconds=5)
|
@tasks.loop(seconds=5)
|
||||||
|
|||||||
@@ -48,6 +48,10 @@ class Config:
|
|||||||
REMOTE_BRIDGE_URL: str = os.getenv("REMOTE_BRIDGE_URL", "")
|
REMOTE_BRIDGE_URL: str = os.getenv("REMOTE_BRIDGE_URL", "")
|
||||||
GATEWAY_API_KEY: str = os.getenv("GATEWAY_API_KEY", "")
|
GATEWAY_API_KEY: str = os.getenv("GATEWAY_API_KEY", "")
|
||||||
|
|
||||||
|
# WebSocket Hub
|
||||||
|
GRAVITY_HUB_SECRET: str = os.getenv("GRAVITY_HUB_SECRET", "") # JWT signing secret
|
||||||
|
GRAVITY_REGISTRATION_CODE: str = os.getenv("GRAVITY_REGISTRATION_CODE", "") # Extension auth
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def validate(cls) -> list[str]:
|
def validate(cls) -> list[str]:
|
||||||
"""Return list of configuration errors."""
|
"""Return list of configuration errors."""
|
||||||
|
|||||||
5
docs/devlog/2026-03-17.md
Normal file
5
docs/devlog/2026-03-17.md
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
# Devlog — 2026-03-17
|
||||||
|
|
||||||
|
| # | 시간 | 작업 | 커밋 | 상태 |
|
||||||
|
|---|------|------|------|------|
|
||||||
|
| 009 | 00:00~06:38 | Extension 모듈 분리 + Hub 통합 테스트 + VSIX v0.4.0 빌드 | `TBD` | ✅ |
|
||||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
@@ -2,7 +2,7 @@
|
|||||||
"name": "gravity-bridge",
|
"name": "gravity-bridge",
|
||||||
"displayName": "Gravity Bridge",
|
"displayName": "Gravity Bridge",
|
||||||
"description": "Antigravity ↔ Discord 브리지 연동 확장",
|
"description": "Antigravity ↔ Discord 브리지 연동 확장",
|
||||||
"version": "0.3.16",
|
"version": "0.4.0",
|
||||||
"publisher": "variet",
|
"publisher": "variet",
|
||||||
"engines": {
|
"engines": {
|
||||||
"vscode": "^1.100.0"
|
"vscode": "^1.100.0"
|
||||||
@@ -68,6 +68,16 @@
|
|||||||
"type": "string",
|
"type": "string",
|
||||||
"default": "",
|
"default": "",
|
||||||
"description": "프로젝트 이름 (기본: git remote 레포명)"
|
"description": "프로젝트 이름 (기본: git remote 레포명)"
|
||||||
|
},
|
||||||
|
"gravityBridge.hubUrl": {
|
||||||
|
"type": "string",
|
||||||
|
"default": "",
|
||||||
|
"description": "WebSocket Hub URL (예: wss://your-server.com/ws)"
|
||||||
|
},
|
||||||
|
"gravityBridge.registrationCode": {
|
||||||
|
"type": "string",
|
||||||
|
"default": "",
|
||||||
|
"description": "Hub 등록 코드 (서버에서 발급)"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
698
extension/src/observer-script.ts
Normal file
698
extension/src/observer-script.ts
Normal file
@@ -0,0 +1,698 @@
|
|||||||
|
/**
|
||||||
|
* Approval Observer Script — injected into AG's renderer process.
|
||||||
|
*
|
||||||
|
* This is a self-contained JavaScript string template that runs in the
|
||||||
|
* browser context (no Node.js APIs). It scans the DOM for approval buttons,
|
||||||
|
* reports them to the HTTP bridge, and handles trigger clicks.
|
||||||
|
*
|
||||||
|
* Extracted from extension.ts for maintainability.
|
||||||
|
*/
|
||||||
|
|
||||||
|
export function generateApprovalObserverScript(_port: number): string {
|
||||||
|
// Port is hardcoded as fallback, but renderer also reads ag-bridge-ports.json for multi-bridge
|
||||||
|
return `
|
||||||
|
// ── Gravity Bridge v3: Approval Observer (deep DOM traversal — iframes, webviews, shadow DOMs) ──
|
||||||
|
(function(){
|
||||||
|
'use strict';
|
||||||
|
var BASE='',_obs=false,_sent={},_ready=false;
|
||||||
|
var _scanScheduled=false,_lastScanTs=0;
|
||||||
|
var THROTTLE_MS=100;
|
||||||
|
var CLEANUP_MS=300000;
|
||||||
|
var _domDumped=false;
|
||||||
|
|
||||||
|
function log(m){console.log('[GB Observer] '+m);}
|
||||||
|
log('v3 Script loaded — deep DOM traversal enabled');
|
||||||
|
|
||||||
|
// ── Deep DOM Traversal: find buttons across ALL boundaries ──
|
||||||
|
// Searches: main document → iframes (contentDocument) → webview elements → shadow DOMs
|
||||||
|
function deepFindButtons(patterns){
|
||||||
|
var results=[];
|
||||||
|
// 1. Main document buttons
|
||||||
|
collectButtons(document,results,patterns,'main');
|
||||||
|
// 2. Iframe traversal (try contentDocument — works if same-origin or webSecurity off)
|
||||||
|
var iframes=document.querySelectorAll('iframe');
|
||||||
|
for(var i=0;i<iframes.length;i++){
|
||||||
|
try{
|
||||||
|
var idoc=iframes[i].contentDocument||iframes[i].contentWindow.document;
|
||||||
|
if(idoc){collectButtons(idoc,results,patterns,'iframe#'+i+'('+iframes[i].className.substring(0,30)+')');}
|
||||||
|
}catch(e){
|
||||||
|
// Cross-origin — can't access. Log only on first dom dump
|
||||||
|
if(!_domDumped)log('iframe#'+i+' cross-origin: '+e.message.substring(0,60));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// 3. Webview elements (Electron <webview> tag — has executeJavaScript)
|
||||||
|
var webviews=document.querySelectorAll('webview');
|
||||||
|
for(var w=0;w<webviews.length;w++){
|
||||||
|
try{
|
||||||
|
var wvDoc=webviews[w].contentDocument;
|
||||||
|
if(wvDoc){collectButtons(wvDoc,results,patterns,'webview#'+w);}
|
||||||
|
}catch(e){
|
||||||
|
if(!_domDumped)log('webview#'+w+' access error: '+e.message.substring(0,60));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return results;
|
||||||
|
}
|
||||||
|
|
||||||
|
function collectButtons(doc,results,patterns,source){
|
||||||
|
if(!doc||!doc.querySelectorAll)return;
|
||||||
|
var btns=doc.querySelectorAll('button');
|
||||||
|
for(var i=0;i<btns.length;i++){
|
||||||
|
var b=btns[i];
|
||||||
|
if(b.disabled||b.hidden)continue;
|
||||||
|
try{if(!b.offsetParent&&b.style.display!=='fixed')continue;}catch(e){}
|
||||||
|
var txt=(b.textContent||'').trim();
|
||||||
|
if(!txt)continue;
|
||||||
|
for(var p=0;p<patterns.length;p++){
|
||||||
|
if(patterns[p].test(txt)){
|
||||||
|
results.push({btn:b,text:txt,source:source});
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// 4. Recurse into shadow DOMs
|
||||||
|
try{
|
||||||
|
var allEls=doc.querySelectorAll('*');
|
||||||
|
for(var j=0;j<allEls.length;j++){
|
||||||
|
var sr=allEls[j].shadowRoot;
|
||||||
|
if(sr)collectButtons(sr,results,patterns,source+'>shadow');
|
||||||
|
}
|
||||||
|
}catch(e){}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Deep DOM Inspector (recursive, POSTs results to bridge) ──
|
||||||
|
function runDeepInspect(){
|
||||||
|
var result={timestamp:new Date().toISOString(),windowURL:window.location.href,windowOrigin:window.location.origin,windowProtocol:window.location.protocol,framesCount:window.frames.length,nodes:[]};
|
||||||
|
log('DEEP-INSPECT: starting recursive DOM analysis...');
|
||||||
|
|
||||||
|
function inspectDoc(doc,depth,label){
|
||||||
|
var node={label:label,depth:depth,accessible:true,url:'',buttons:[],roleBtns:[],iframes:[],webviews:[],shadowDOMs:0,totalElements:0};
|
||||||
|
if(!doc){node.accessible=false;node.error='null document';result.nodes.push(node);return;}
|
||||||
|
try{node.url=(doc.URL||doc.documentURI||'unknown').substring(0,200);}catch(e){node.url='blocked';}
|
||||||
|
try{node.title=(doc.title||'').substring(0,100);}catch(e){}
|
||||||
|
try{node.readyState=doc.readyState;}catch(e){}
|
||||||
|
|
||||||
|
// CSP
|
||||||
|
try{
|
||||||
|
var csp=doc.querySelectorAll('meta[http-equiv="Content-Security-Policy"]');
|
||||||
|
if(csp.length>0){node.csp=[];for(var c=0;c<csp.length;c++){node.csp.push((csp[c].content||'').substring(0,200));}}
|
||||||
|
}catch(e){}
|
||||||
|
|
||||||
|
try{
|
||||||
|
var allEls=doc.querySelectorAll('*');
|
||||||
|
node.totalElements=allEls.length;
|
||||||
|
// Buttons
|
||||||
|
var btns=doc.querySelectorAll('button');
|
||||||
|
for(var i=0;i<btns.length;i++){
|
||||||
|
var b=btns[i];
|
||||||
|
var txt=(b.textContent||'').trim().substring(0,80);
|
||||||
|
if(!txt)continue;
|
||||||
|
var cls=(b.className||'').substring(0,60);
|
||||||
|
var disabled=b.disabled;
|
||||||
|
var hidden=b.hidden||false;
|
||||||
|
try{if(!b.offsetParent&&b.style.display!=='fixed')hidden=true;}catch(e){}
|
||||||
|
var aria=b.getAttribute('aria-label')||'';
|
||||||
|
var ttl=b.getAttribute('title')||'';
|
||||||
|
node.buttons.push({text:txt,class:cls,disabled:disabled,hidden:hidden,aria:aria,title:ttl});
|
||||||
|
}
|
||||||
|
// role=button
|
||||||
|
var rbs=doc.querySelectorAll('[role="button"]');
|
||||||
|
for(var r=0;r<rbs.length;r++){
|
||||||
|
if(rbs[r].tagName==='BUTTON')continue;
|
||||||
|
var rtxt=(rbs[r].textContent||'').trim().substring(0,60);
|
||||||
|
node.roleBtns.push({tag:rbs[r].tagName.toLowerCase(),text:rtxt});
|
||||||
|
}
|
||||||
|
// Shadow DOMs
|
||||||
|
for(var s=0;s<allEls.length;s++){
|
||||||
|
var sr=allEls[s].shadowRoot;
|
||||||
|
if(sr){node.shadowDOMs++;inspectDoc(sr,depth+1,'shadow(<'+allEls[s].tagName.toLowerCase()+' class="'+(allEls[s].className||'').substring(0,30)+'">)');}
|
||||||
|
}
|
||||||
|
// Iframes
|
||||||
|
var ifs=doc.querySelectorAll('iframe');
|
||||||
|
for(var fi=0;fi<ifs.length;fi++){
|
||||||
|
var f=ifs[fi];
|
||||||
|
var finfo={index:fi,class:(f.className||'').substring(0,60),src:(f.src||'').substring(0,150),id:f.id||'',sandbox:f.getAttribute('sandbox')||'',allow:f.getAttribute('allow')||'',accessible:false,cwExists:false,cwFrames:0};
|
||||||
|
try{
|
||||||
|
var idoc=f.contentDocument||(f.contentWindow&&f.contentWindow.document);
|
||||||
|
if(idoc){finfo.accessible=true;inspectDoc(idoc,depth+1,'iframe#'+fi+'('+finfo.class.substring(0,30)+')');
|
||||||
|
}else{finfo.error='contentDocument=null';}
|
||||||
|
}catch(e){finfo.error=e.message.substring(0,80);}
|
||||||
|
try{var cw=f.contentWindow;if(cw){finfo.cwExists=true;finfo.cwFrames=cw.length;try{finfo.cwLocation=cw.location.href;}catch(e2){finfo.cwLocation='blocked: '+e2.message.substring(0,40);}}}
|
||||||
|
catch(e){}
|
||||||
|
node.iframes.push(finfo);
|
||||||
|
}
|
||||||
|
// Webviews
|
||||||
|
var wvs=doc.querySelectorAll('webview');
|
||||||
|
for(var wi=0;wi<wvs.length;wi++){
|
||||||
|
var wv=wvs[wi];
|
||||||
|
var winfo={index:wi,src:(wv.src||'').substring(0,150),class:(wv.className||'').substring(0,60),partition:wv.getAttribute('partition')||'',preload:wv.getAttribute('preload')||'',nodeintegration:wv.getAttribute('nodeintegration')||'',webpreferences:wv.getAttribute('webpreferences')||'',hasExecJS:typeof wv.executeJavaScript==='function',contentDocAccessible:false};
|
||||||
|
try{var wdoc=wv.contentDocument;if(wdoc){winfo.contentDocAccessible=true;inspectDoc(wdoc,depth+1,'webview#'+wi+'.contentDocument');}}catch(e){winfo.contentDocError=e.message.substring(0,60);}
|
||||||
|
node.webviews.push(winfo);
|
||||||
|
}
|
||||||
|
}catch(e){node.error=e.message;}
|
||||||
|
result.nodes.push(node);
|
||||||
|
return node;
|
||||||
|
}
|
||||||
|
|
||||||
|
inspectDoc(document,0,'MainDocument');
|
||||||
|
|
||||||
|
// Webview executeJavaScript probe (async)
|
||||||
|
var webviews=document.querySelectorAll('webview');
|
||||||
|
var probesPending=webviews.length;
|
||||||
|
result.webviewProbes=[];
|
||||||
|
if(probesPending===0)postResults();
|
||||||
|
for(var pw=0;pw<webviews.length;pw++){
|
||||||
|
(function(wv,idx){
|
||||||
|
if(typeof wv.executeJavaScript!=='function'){result.webviewProbes.push({index:idx,error:'executeJavaScript not available'});probesPending--;if(probesPending<=0)postResults();return;}
|
||||||
|
try{
|
||||||
|
wv.executeJavaScript('(function(){var btns=document.querySelectorAll("button");var allEls=document.querySelectorAll("*");var ifs=document.querySelectorAll("iframe");var wvs=document.querySelectorAll("webview");var btnArr=[];for(var i=0;i<btns.length;i++){var b=btns[i];var txt=(b.textContent||"").trim();var cls=(b.className||"").substring(0,50);var dis=b.disabled;var hid=b.hidden||!b.offsetParent;btnArr.push({text:txt.substring(0,60),class:cls,disabled:dis,hidden:hid,aria:b.getAttribute("aria-label")||"",title:b.getAttribute("title")||""});}var rbs=document.querySelectorAll("[role=button]");var rbArr=[];for(var j=0;j<rbs.length;j++){if(rbs[j].tagName!=="BUTTON")rbArr.push({tag:rbs[j].tagName.toLowerCase(),text:(rbs[j].textContent||"").trim().substring(0,40)});}var sc=0;for(var k=0;k<allEls.length;k++){if(allEls[k].shadowRoot)sc++;}return JSON.stringify({url:document.URL,title:document.title,totalElements:allEls.length,buttons:btnArr,roleBtns:rbArr,iframes:ifs.length,webviews:wvs.length,shadowDOMs:sc});})()')
|
||||||
|
.then(function(r){
|
||||||
|
try{var d=JSON.parse(r);result.webviewProbes.push({index:idx,success:true,data:d});log('DEEP-INSPECT: webview#'+idx+' probe OK: '+d.buttons.length+' buttons, '+d.totalElements+' elements');}catch(e){result.webviewProbes.push({index:idx,parseError:e.message,raw:r});}
|
||||||
|
probesPending--;if(probesPending<=0)postResults();
|
||||||
|
})
|
||||||
|
.catch(function(e){
|
||||||
|
result.webviewProbes.push({index:idx,execError:e.message});
|
||||||
|
log('DEEP-INSPECT: webview#'+idx+' execJS error: '+e.message);
|
||||||
|
probesPending--;if(probesPending<=0)postResults();
|
||||||
|
});
|
||||||
|
}catch(e){
|
||||||
|
result.webviewProbes.push({index:idx,callError:e.message});
|
||||||
|
probesPending--;if(probesPending<=0)postResults();
|
||||||
|
}
|
||||||
|
})(webviews[pw],pw);
|
||||||
|
}
|
||||||
|
|
||||||
|
function postResults(){
|
||||||
|
var summary='nodes='+result.nodes.length;
|
||||||
|
var totalBtns=0;for(var n=0;n<result.nodes.length;n++)totalBtns+=result.nodes[n].buttons.length;
|
||||||
|
summary+=' totalButtons='+totalBtns+' webviewProbes='+result.webviewProbes.length;
|
||||||
|
log('DEEP-INSPECT complete: '+summary);
|
||||||
|
// Also log buttons from each node
|
||||||
|
for(var n2=0;n2<result.nodes.length;n2++){
|
||||||
|
var nd=result.nodes[n2];
|
||||||
|
if(nd.buttons.length>0){
|
||||||
|
log(' '+nd.label+': '+nd.buttons.length+' buttons');
|
||||||
|
for(var bi=0;bi<Math.min(15,nd.buttons.length);bi++){
|
||||||
|
log(' ['+bi+'] "'+nd.buttons[bi].text+'"'+(nd.buttons[bi].disabled?' DISABLED':'')+(nd.buttons[bi].hidden?' HIDDEN':''));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// POST to bridge
|
||||||
|
fetch(BASE+'/deep-inspect-result',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(result)})
|
||||||
|
.then(function(){log('DEEP-INSPECT results posted to bridge');})
|
||||||
|
.catch(function(e){log('DEEP-INSPECT post error: '+e.message);});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Auto-dump on startup (3s delay)
|
||||||
|
function dumpDOMStructure(){runDeepInspect();}
|
||||||
|
|
||||||
|
// ── Port Discovery: async fetch()-based (sync XHR blocked in Electron renderer) ──
|
||||||
|
var HARDCODED_PORT=${_port};
|
||||||
|
|
||||||
|
function tryPingAsync(port){
|
||||||
|
return fetch('http://127.0.0.1:'+port+'/ping?t='+Date.now(),{signal:AbortSignal.timeout(2000)})
|
||||||
|
.then(function(r){return r.text();})
|
||||||
|
.then(function(t){return t==='pong';})
|
||||||
|
.catch(function(){return false;});
|
||||||
|
}
|
||||||
|
|
||||||
|
function discoverPort(cb){
|
||||||
|
log('Trying hardcoded port '+HARDCODED_PORT+'...');
|
||||||
|
tryPingAsync(HARDCODED_PORT).then(function(ok){
|
||||||
|
if(ok){log('Port discovered (hardcoded): '+HARDCODED_PORT);cb(HARDCODED_PORT);return;}
|
||||||
|
log('Hardcoded port failed, retrying with backoff...');
|
||||||
|
|
||||||
|
var attempts=0;
|
||||||
|
var timer=setInterval(function(){
|
||||||
|
attempts++;
|
||||||
|
if(attempts>60){clearInterval(timer);log('Port discovery timeout after 2min');return;}
|
||||||
|
tryPingAsync(HARDCODED_PORT).then(function(ok2){
|
||||||
|
if(ok2){clearInterval(timer);log('Port discovered (retry #'+attempts+'): '+HARDCODED_PORT);cb(HARDCODED_PORT);}
|
||||||
|
});
|
||||||
|
},2000);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
discoverPort(function(port){
|
||||||
|
BASE='http://127.0.0.1:'+port;
|
||||||
|
fetch(BASE+'/ping').then(function(r){return r.text();}).then(function(t){
|
||||||
|
if(t==='pong'){log('Bridge connected on port '+port);_ready=true;startObserver();setTimeout(dumpDOMStructure,3000);}
|
||||||
|
else log('Bridge ping failed: '+t);
|
||||||
|
}).catch(function(e){log('Bridge unreachable: '+e.message);});
|
||||||
|
});
|
||||||
|
|
||||||
|
// ── Button patterns to detect (order matters: first match wins per scan) ──
|
||||||
|
// ONLY positive triggers should initiate a pending request group.
|
||||||
|
// Negative/secondary buttons (Deny, Reject, Dismiss) will be collected as siblings.
|
||||||
|
var PATS=[
|
||||||
|
{re:/^Run/i, type:'terminal_command'},
|
||||||
|
{re:/^Accept all$/i, type:'diff_review'},
|
||||||
|
{re:/^Accept$/i, type:'agent_step'},
|
||||||
|
{re:/^Allow/i, type:'permission'},
|
||||||
|
{re:/^Approve/i, type:'agent_step'},
|
||||||
|
{re:/^Retry$/i, type:'error_recovery'},
|
||||||
|
];
|
||||||
|
|
||||||
|
// ALL actionable button patterns (for grouping siblings in same container)
|
||||||
|
var ALL_ACTION_RE=[/^Run/i,/^Accept/i,/^Reject/i,/^Allow/i,/^Deny/i,/^Approve/i,/^Cancel$/i,/^Retry$/i,/^Dismiss$/i,/^Stop$/i,/^Decline$/i];
|
||||||
|
|
||||||
|
// Reject button patterns for finding the counterpart
|
||||||
|
var REJECT_RE=[/^reject$/i,/^reject all$/i,/^cancel$/i,/^deny$/i,/^stop$/i,/^decline$/i,/^dismiss$/i];
|
||||||
|
|
||||||
|
// ── Stable button fingerprint (no getBoundingClientRect — scroll-safe) ──
|
||||||
|
function btnId(b,type){
|
||||||
|
// Use: type + button text + parent's first 40 chars of text content
|
||||||
|
var txt=(b.textContent||'').trim();
|
||||||
|
var parent=b.parentElement;
|
||||||
|
var pctx=parent?(parent.textContent||'').substring(0,40).replace(/\\s+/g,' '):'';
|
||||||
|
// Also use DOM position: nth-child among sibling buttons
|
||||||
|
var idx=0;
|
||||||
|
if(parent){
|
||||||
|
var siblings=parent.querySelectorAll('button');
|
||||||
|
for(var i=0;i<siblings.length;i++){if(siblings[i]===b){idx=i;break;}}
|
||||||
|
}
|
||||||
|
return type+'|'+txt+'|'+idx+'|'+pctx.substring(0,20);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Context extraction — walk up DOM to find command/code description ──
|
||||||
|
function extractContext(b){
|
||||||
|
// Strategy 1: Look for code/pre/terminal blocks near the button
|
||||||
|
var container=b.closest('[class*="step"]')
|
||||||
|
||b.closest('[class*="action"]')
|
||||||
|
||b.closest('[class*="tool"]')
|
||||||
|
||b.closest('[class*="cascade"]')
|
||||||
|
||b.closest('[class*="message"]');
|
||||||
|
if(!container)container=b.parentElement;
|
||||||
|
if(!container)return '';
|
||||||
|
|
||||||
|
// Look for code blocks
|
||||||
|
var codeEl=container.querySelector('pre,code,[class*="command"],[class*="terminal"],[class*="code-block"]');
|
||||||
|
if(codeEl){
|
||||||
|
var codeText=(codeEl.textContent||'').trim();
|
||||||
|
if(codeText.length>0)return codeText.substring(0,500);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strategy 2: Get surrounding text (exclude button text itself)
|
||||||
|
var full=(container.textContent||'');
|
||||||
|
var btnText=(b.textContent||'');
|
||||||
|
var desc=full.replace(btnText,'').trim();
|
||||||
|
// Trim to reasonable length
|
||||||
|
return desc.substring(0,500);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Find common container of related buttons ──
|
||||||
|
function findButtonContainer(btn){
|
||||||
|
return btn.closest('[class*="step"]')
|
||||||
|
||btn.closest('[class*="action"]')
|
||||||
|
||btn.closest('[class*="tool"]')
|
||||||
|
||btn.closest('[class*="cascade"]')
|
||||||
|
||btn.closest('[class*="message"]')
|
||||||
|
||btn.closest('[class*="dialog"]')
|
||||||
|
||btn.closest('[class*="notification"]')
|
||||||
|
||btn.parentElement;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Collect all actionable sibling buttons from a container ──
|
||||||
|
function collectSiblingButtons(container,triggerBtn){
|
||||||
|
if(!container)return [];
|
||||||
|
var siblings=container.querySelectorAll('button');
|
||||||
|
var result=[];
|
||||||
|
for(var i=0;i<siblings.length;i++){
|
||||||
|
var sb=siblings[i];
|
||||||
|
if(sb.disabled||sb.hidden)continue;
|
||||||
|
try{if(!sb.offsetParent&&sb.style.display!=='fixed')continue;}catch(e){}
|
||||||
|
var stxt=(sb.textContent||'').trim();
|
||||||
|
stxt=stxt.replace(/(Alt|Ctrl|Shift|Meta)\+.*/,'').trim();
|
||||||
|
if(!stxt)continue;
|
||||||
|
// Check if this button matches any actionable pattern
|
||||||
|
var isAction=false;
|
||||||
|
for(var a=0;a<ALL_ACTION_RE.length;a++){
|
||||||
|
if(ALL_ACTION_RE[a].test(stxt)){isAction=true;break;}
|
||||||
|
}
|
||||||
|
if(!isAction)continue;
|
||||||
|
result.push({btn:sb,text:stxt,isPrimary:(sb===triggerBtn)});
|
||||||
|
}
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Find the React app container (Antigravity's main UI root) ──
|
||||||
|
function findPanel(){
|
||||||
|
// Priority order of panel selectors (most specific first)
|
||||||
|
var selectors=[
|
||||||
|
'.antigravity-agent-side-panel',
|
||||||
|
'#jetski-agent-panel',
|
||||||
|
'.react-app-container',
|
||||||
|
'[class*="agent-panel"]',
|
||||||
|
'[class*="agentPanel"]',
|
||||||
|
];
|
||||||
|
for(var i=0;i<selectors.length;i++){
|
||||||
|
var el=document.querySelector(selectors[i]);
|
||||||
|
if(el)return el;
|
||||||
|
}
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Core scan — finds actionable buttons and reports to bridge ──
|
||||||
|
// Groups related buttons from same container into a single pending
|
||||||
|
function scan(){
|
||||||
|
if(!_ready)return;
|
||||||
|
var now=Date.now();
|
||||||
|
|
||||||
|
var panel=findPanel();
|
||||||
|
// Expand search: panel-scoped first, then full body for review bars
|
||||||
|
var searchRoots=[];
|
||||||
|
if(panel)searchRoots.push(panel);
|
||||||
|
// Always also scan body for diff review bar (Accept all/Reject all)
|
||||||
|
// which lives outside the agent panel in the editor notification area
|
||||||
|
if(document.body)searchRoots.push(document.body);
|
||||||
|
if(!searchRoots.length)return;
|
||||||
|
|
||||||
|
var seen={}; // dedupe buttons across search roots
|
||||||
|
for(var r=0;r<searchRoots.length;r++){
|
||||||
|
var allBtns=searchRoots[r].querySelectorAll('button');
|
||||||
|
if(!allBtns.length)continue;
|
||||||
|
|
||||||
|
for(var j=0;j<allBtns.length;j++){
|
||||||
|
var b=allBtns[j];
|
||||||
|
if(b.disabled||b.hidden)continue;
|
||||||
|
// Check visibility (offsetParent null = hidden via CSS)
|
||||||
|
if(!b.offsetParent&&b.style.display!=='fixed')continue;
|
||||||
|
|
||||||
|
var txt=(b.textContent||'').trim();
|
||||||
|
if(!txt)continue;
|
||||||
|
// Strip keyboard shortcut suffixes (e.g. "RunAlt+↵" → "Run")
|
||||||
|
txt=txt.replace(/(Alt|Ctrl|Shift|Meta)\+.*/,'').trim();
|
||||||
|
if(!txt)continue;
|
||||||
|
|
||||||
|
// Match against patterns
|
||||||
|
var matchedType=null;
|
||||||
|
for(var p=0;p<PATS.length;p++){
|
||||||
|
if(PATS[p].re.test(txt)){matchedType=PATS[p].type;break;}
|
||||||
|
}
|
||||||
|
if(!matchedType)continue;
|
||||||
|
|
||||||
|
// Generate stable ID for the GROUP (use container-based key)
|
||||||
|
var container=findButtonContainer(b);
|
||||||
|
var groupKey=matchedType+'|group|'+(container?(container.textContent||'').substring(0,40).replace(/\\s+/g,' '):'none');
|
||||||
|
if(_sent[groupKey])continue;
|
||||||
|
|
||||||
|
// Collect ALL related buttons from the same container
|
||||||
|
var siblings=collectSiblingButtons(container,b);
|
||||||
|
if(siblings.length===0)siblings=[{btn:b,text:txt,isPrimary:true}];
|
||||||
|
|
||||||
|
// Build buttons array for multi-choice support
|
||||||
|
var buttonsArr=[];
|
||||||
|
var btnRefs=[];
|
||||||
|
var bidList=[];
|
||||||
|
for(var si=0;si<siblings.length;si++){
|
||||||
|
var sb=siblings[si];
|
||||||
|
var sbid=btnId(sb.btn,matchedType);
|
||||||
|
buttonsArr.push({text:sb.text,index:si,is_primary:sb.isPrimary});
|
||||||
|
btnRefs.push(sb.btn);
|
||||||
|
bidList.push(sbid);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Extract context from trigger button
|
||||||
|
var desc=extractContext(b);
|
||||||
|
var rid=now.toString()+'_'+Math.random().toString(36).substring(2,6);
|
||||||
|
|
||||||
|
// Mark entire group as sent
|
||||||
|
_sent[groupKey]={rid:rid,ts:now};
|
||||||
|
for(var mk=0;mk<bidList.length;mk++){_sent[bidList[mk]]={rid:rid,ts:now};}
|
||||||
|
|
||||||
|
log('DETECTED GROUP '+matchedType+': '+buttonsArr.map(function(x){return '"'+x.text+'"';}).join(', ')+' → pending to bridge');
|
||||||
|
|
||||||
|
// Send to bridge (closure to capture refs)
|
||||||
|
(function(rid2,btnRefs2,bidList2,groupKey2,txt2,desc2,type2,buttonsArr2){
|
||||||
|
var payload={
|
||||||
|
request_id:rid2,
|
||||||
|
command:txt2,
|
||||||
|
description:desc2,
|
||||||
|
step_type:type2,
|
||||||
|
buttons:buttonsArr2
|
||||||
|
};
|
||||||
|
fetch(BASE+'/pending',{
|
||||||
|
method:'POST',
|
||||||
|
headers:{'Content-Type':'application/json'},
|
||||||
|
body:JSON.stringify(payload)
|
||||||
|
}).then(function(r){return r.json();}).then(function(d){
|
||||||
|
log('Pending created: '+d.request_id+' for group ['+buttonsArr2.map(function(x){return x.text;}).join(', ')+']');
|
||||||
|
pollResponseGroup(d.request_id,btnRefs2,bidList2,groupKey2);
|
||||||
|
}).catch(function(e){
|
||||||
|
log('POST error: '+e.message);
|
||||||
|
delete _sent[groupKey2];
|
||||||
|
for(var di=0;di<bidList2.length;di++){delete _sent[bidList2[di]];}
|
||||||
|
});
|
||||||
|
})(rid,btnRefs,bidList,groupKey,txt,desc,matchedType,buttonsArr);
|
||||||
|
|
||||||
|
// Process ONE button GROUP per scan cycle (avoid flooding)
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
} // end searchRoots loop
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Poll for Discord response (multi-button group aware) ──
|
||||||
|
function pollResponseGroup(rid,btnRefs,bidList,groupKey){
|
||||||
|
var polls=0;
|
||||||
|
var maxPolls=200; // 5 minutes at 1500ms interval
|
||||||
|
var timer=setInterval(function(){
|
||||||
|
polls++;
|
||||||
|
// Check if ANY button in the group is still in DOM
|
||||||
|
var anyAlive=false;
|
||||||
|
for(var ai=0;ai<btnRefs.length;ai++){
|
||||||
|
if(document.body.contains(btnRefs[ai])){anyAlive=true;break;}
|
||||||
|
}
|
||||||
|
if(!anyAlive){
|
||||||
|
log('All buttons removed from DOM — stopping poll for '+rid);
|
||||||
|
clearInterval(timer);
|
||||||
|
delete _sent[groupKey];
|
||||||
|
for(var ci=0;ci<bidList.length;ci++){delete _sent[bidList[ci]];}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if(polls>maxPolls){
|
||||||
|
log('Poll timeout for '+rid);
|
||||||
|
clearInterval(timer);
|
||||||
|
delete _sent[groupKey];
|
||||||
|
for(var ti=0;ti<bidList.length;ti++){delete _sent[bidList[ti]];}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
fetch(BASE+'/response/'+rid).then(function(r){return r.json();}).then(function(d){
|
||||||
|
if(d.waiting)return;
|
||||||
|
clearInterval(timer);
|
||||||
|
var btnIdx=(typeof d.button_index==='number'&&d.button_index>=0)?d.button_index:-1;
|
||||||
|
if(btnIdx>=0&&btnIdx<btnRefs.length){
|
||||||
|
// Multi-choice: click specific button by index
|
||||||
|
var targetBtn=btnRefs[btnIdx];
|
||||||
|
var targetTxt=(targetBtn.textContent||'').trim();
|
||||||
|
log((d.approved?'✅':'❌')+' CHOICE '+rid+' → clicking button['+btnIdx+'] "'+targetTxt+'"');
|
||||||
|
targetBtn.click();
|
||||||
|
} else if(d.approved){
|
||||||
|
// Legacy single-button: click first (primary) button
|
||||||
|
var primaryBtn=btnRefs[0];
|
||||||
|
log('✅ APPROVED '+rid+' → clicking "'+((primaryBtn.textContent||'').trim())+'"');
|
||||||
|
primaryBtn.click();
|
||||||
|
} else {
|
||||||
|
// Legacy reject: find and click reject/deny button
|
||||||
|
log('❌ REJECTED '+rid+' → finding reject button');
|
||||||
|
clickRejectButton(btnRefs[0]);
|
||||||
|
}
|
||||||
|
delete _sent[groupKey];
|
||||||
|
for(var ri=0;ri<bidList.length;ri++){delete _sent[bidList[ri]];}
|
||||||
|
}).catch(function(){});
|
||||||
|
},1500);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Legacy pollResponse for backward compatibility (single button)
|
||||||
|
function pollResponse(rid,btn,bid){
|
||||||
|
pollResponseGroup(rid,[btn],[bid],bid);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Find and click the reject/cancel counterpart button ──
|
||||||
|
function clickRejectButton(approveBtn){
|
||||||
|
// Walk up to find the container, then search for reject buttons
|
||||||
|
var container=approveBtn.closest('[class*="step"]')
|
||||||
|
||approveBtn.closest('[class*="action"]')
|
||||||
|
||approveBtn.closest('[class*="tool"]')
|
||||||
|
||approveBtn.parentElement;
|
||||||
|
if(!container){log('No container for reject');return;}
|
||||||
|
|
||||||
|
var siblings=container.querySelectorAll('button');
|
||||||
|
for(var i=0;i<siblings.length;i++){
|
||||||
|
var t=(siblings[i].textContent||'').trim();
|
||||||
|
for(var r=0;r<REJECT_RE.length;r++){
|
||||||
|
if(REJECT_RE[r].test(t)){
|
||||||
|
log('Clicking reject: "'+t+'"');
|
||||||
|
siblings[i].click();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
log('No reject button found near approve button');
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Throttled scan — leading-edge: fires immediately, then locks ──
|
||||||
|
function scheduleScan(){
|
||||||
|
if(!_ready)return;
|
||||||
|
var now=Date.now();
|
||||||
|
if(now-_lastScanTs>=THROTTLE_MS){
|
||||||
|
_lastScanTs=now;
|
||||||
|
scan();
|
||||||
|
} else if(!_scanScheduled){
|
||||||
|
_scanScheduled=true;
|
||||||
|
setTimeout(function(){
|
||||||
|
_scanScheduled=false;
|
||||||
|
_lastScanTs=Date.now();
|
||||||
|
scan();
|
||||||
|
},THROTTLE_MS-(now-_lastScanTs));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Periodic cleanup of stale _sent entries ──
|
||||||
|
setInterval(function(){
|
||||||
|
var now=Date.now();
|
||||||
|
var keys=Object.keys(_sent);
|
||||||
|
for(var i=0;i<keys.length;i++){
|
||||||
|
var entry=_sent[keys[i]];
|
||||||
|
if(entry&&entry.ts&&(now-entry.ts)>CLEANUP_MS){
|
||||||
|
log('Cleanup stale entry: '+keys[i]);
|
||||||
|
delete _sent[keys[i]];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},60000);
|
||||||
|
|
||||||
|
// ── Start observation ──
|
||||||
|
function startObserver(){
|
||||||
|
if(_obs)return;
|
||||||
|
// PRIMARY: MutationObserver — reacts instantly to DOM changes
|
||||||
|
new MutationObserver(function(mutations){
|
||||||
|
// Only scan if mutations contain added nodes (new buttons potentially)
|
||||||
|
for(var i=0;i<mutations.length;i++){
|
||||||
|
if(mutations[i].addedNodes.length>0){
|
||||||
|
scheduleScan();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}).observe(document.body,{childList:true,subtree:true});
|
||||||
|
|
||||||
|
// FALLBACK: periodic scan every 3s for any missed mutations
|
||||||
|
setInterval(scheduleScan,3000);
|
||||||
|
|
||||||
|
// ── Adaptive idle detection for HTTP polls ──
|
||||||
|
var _lastActivity=Date.now();
|
||||||
|
var _idleThreshold=60000; // 60s without DOM changes → slow mode
|
||||||
|
new MutationObserver(function(){_lastActivity=Date.now();}).observe(document.body,{childList:true,subtree:true,attributes:true});
|
||||||
|
function getAdaptiveInterval(){return (Date.now()-_lastActivity>_idleThreshold)?10000:2000;}
|
||||||
|
|
||||||
|
// ── DEEP-INSPECT POLLING: curl→Bridge→Renderer→Results ──
|
||||||
|
(function pollDeepInspect(){
|
||||||
|
if(_ready&&BASE){
|
||||||
|
fetch(BASE+'/deep-inspect-trigger?t='+Date.now()).then(function(r){return r.json();}).then(function(d){
|
||||||
|
if(d.inspect){log('🔍 Deep inspect triggered via HTTP');runDeepInspect();}
|
||||||
|
}).catch(function(){});
|
||||||
|
}
|
||||||
|
setTimeout(pollDeepInspect,getAdaptiveInterval());
|
||||||
|
})();
|
||||||
|
|
||||||
|
// ── TRIGGER-CLICK: Extension→Renderer bridge for programmatic button clicks ──
|
||||||
|
// Extension sets clickTrigger via tryApprovalStrategies → renderer polls and clicks
|
||||||
|
// v3: uses deepFindButtons() to traverse iframes, webviews, shadow DOMs
|
||||||
|
(function pollTriggerClick(){
|
||||||
|
if(_ready&&BASE){
|
||||||
|
fetch(BASE+'/trigger-click?t='+Date.now()).then(function(r){return r.json();}).then(function(d){
|
||||||
|
if(!d.action)return;
|
||||||
|
log('🔔 TRIGGER-CLICK received: action='+d.action);
|
||||||
|
|
||||||
|
var approveRe=[/^Run$/i,/^Run /i,/^Accept/i,/^Allow/i,/^Approve/i,/^Continue$/i,/^Proceed$/i,/^Retry$/i];
|
||||||
|
var rejectRe=[/^Reject/i,/^Cancel$/i,/^Deny$/i,/^Stop$/i,/^Decline$/i,/^Dismiss$/i];
|
||||||
|
var patterns=(d.action==='approve')?approveRe:rejectRe;
|
||||||
|
var emoji=(d.action==='approve')?'✅':'❌';
|
||||||
|
|
||||||
|
// Phase 1: deepFindButtons in main doc + accessible iframes + shadow DOMs
|
||||||
|
var found=deepFindButtons(patterns);
|
||||||
|
if(found.length>0){
|
||||||
|
log(emoji+' TRIGGER-CLICK: clicking "'+found[0].text+'" from '+found[0].source);
|
||||||
|
found[0].btn.click();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 2: Try <webview>.executeJavaScript for inaccessible webviews
|
||||||
|
var webviews=document.querySelectorAll('webview');
|
||||||
|
if(webviews.length>0){
|
||||||
|
log('TRIGGER-CLICK: trying '+webviews.length+' webview(s) via executeJavaScript...');
|
||||||
|
var patternsStr=patterns.map(function(re){return re.source;}).join('|');
|
||||||
|
var clickScript='(function(){'+
|
||||||
|
'var re=new RegExp("'+patternsStr+'","i");'+
|
||||||
|
'var btns=document.querySelectorAll("button");'+
|
||||||
|
'for(var i=0;i<btns.length;i++){'+
|
||||||
|
'var b=btns[i];if(b.disabled||b.hidden)continue;'+
|
||||||
|
'var t=(b.textContent||"").trim();'+
|
||||||
|
'if(re.test(t)){b.click();return "CLICKED:"+t;}'+
|
||||||
|
'}'+
|
||||||
|
'return "NOT_FOUND:"+btns.length+"_buttons";'+
|
||||||
|
'})()';
|
||||||
|
for(var w=0;w<webviews.length;w++){
|
||||||
|
(function(wv,idx){
|
||||||
|
try{
|
||||||
|
if(typeof wv.executeJavaScript==='function'){
|
||||||
|
wv.executeJavaScript(clickScript).then(function(result){
|
||||||
|
log(emoji+' TRIGGER-CLICK webview#'+idx+': '+result);
|
||||||
|
}).catch(function(e){
|
||||||
|
log('TRIGGER-CLICK webview#'+idx+' execJS error: '+e.message);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}catch(e){
|
||||||
|
log('TRIGGER-CLICK webview#'+idx+' error: '+e.message);
|
||||||
|
}
|
||||||
|
})(webviews[w],w);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 3: Try iframes via postMessage (cross-origin fallback)
|
||||||
|
var iframes=document.querySelectorAll('iframe');
|
||||||
|
if(iframes.length>0){
|
||||||
|
log('TRIGGER-CLICK: trying '+iframes.length+' iframe(s) — checking accessibility...');
|
||||||
|
var clickedAny=false;
|
||||||
|
for(var fi=0;fi<iframes.length;fi++){
|
||||||
|
try{
|
||||||
|
var idoc=iframes[fi].contentDocument||iframes[fi].contentWindow.document;
|
||||||
|
if(!idoc)continue;
|
||||||
|
var ibtns=idoc.querySelectorAll('button');
|
||||||
|
for(var bi=0;bi<ibtns.length;bi++){
|
||||||
|
var ib=ibtns[bi];
|
||||||
|
if(ib.disabled||ib.hidden)continue;
|
||||||
|
var itxt=(ib.textContent||'').trim();
|
||||||
|
for(var pi=0;pi<patterns.length;pi++){
|
||||||
|
if(patterns[pi].test(itxt)){
|
||||||
|
log(emoji+' TRIGGER-CLICK iframe#'+fi+': clicking "'+itxt+'"');
|
||||||
|
ib.click();
|
||||||
|
clickedAny=true;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}catch(e){}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if(!found.length){
|
||||||
|
// Log what we DID find for debugging
|
||||||
|
var allBtns=document.querySelectorAll('button');
|
||||||
|
var btnTexts=[];
|
||||||
|
for(var di=0;di<Math.min(10,allBtns.length);di++){
|
||||||
|
btnTexts.push('"'+((allBtns[di].textContent||'').trim()).substring(0,30)+'"');
|
||||||
|
}
|
||||||
|
log('⚠️ TRIGGER-CLICK: no '+d.action+' button found. Main DOM has '+allBtns.length+' btns: ['+btnTexts.join(',')+']');
|
||||||
|
log('⚠️ iframes='+document.querySelectorAll('iframe').length+' webviews='+document.querySelectorAll('webview').length);
|
||||||
|
}
|
||||||
|
}).catch(function(){});
|
||||||
|
}
|
||||||
|
setTimeout(pollTriggerClick,getAdaptiveInterval());
|
||||||
|
})();
|
||||||
|
|
||||||
|
_obs=true;
|
||||||
|
log('v3 Observer active — deep DOM traversal + MutationObserver + trigger-click polling');
|
||||||
|
}
|
||||||
|
})();
|
||||||
|
`;
|
||||||
|
}
|
||||||
|
|
||||||
1435
extension/src/step-probe.ts
Normal file
1435
extension/src/step-probe.ts
Normal file
File diff suppressed because it is too large
Load Diff
114
extension/src/step-utils.ts
Normal file
114
extension/src/step-utils.ts
Normal file
@@ -0,0 +1,114 @@
|
|||||||
|
/**
|
||||||
|
* Step Utilities — pure functions for parsing step data, planner responses,
|
||||||
|
* and tool call information. No external state dependencies.
|
||||||
|
*
|
||||||
|
* Extracted from extension.ts for maintainability.
|
||||||
|
*/
|
||||||
|
|
||||||
|
export function extractPlannerText(step: any): string | null {
|
||||||
|
if (!step) { return null; }
|
||||||
|
|
||||||
|
// Fields to SKIP — not user-facing content
|
||||||
|
const SKIP_FIELDS = new Set([
|
||||||
|
'thinking', 'thinkingSignature', 'stopReason', 'type', 'status', 'metadata',
|
||||||
|
'ephemeralMessage', 'generatorModel', 'requestedModel',
|
||||||
|
'executionId', 'sourceTrajectoryStepInfo', 'stepIndex',
|
||||||
|
'viewableAt', 'createdAt', 'finishedGeneratingAt',
|
||||||
|
'lastCompletedChunkAt', 'source', 'stepGenerationVersion'
|
||||||
|
]);
|
||||||
|
|
||||||
|
// plannerResponse can be string or object
|
||||||
|
const pr = step.plannerResponse;
|
||||||
|
if (typeof pr === 'string' && pr.length > 10) {
|
||||||
|
return filterEphemeral(pr);
|
||||||
|
}
|
||||||
|
if (pr && typeof pr === 'object') {
|
||||||
|
// Try known content fields first (NOT thinking/stopReason)
|
||||||
|
const text = pr.content || pr.text || pr.summary || pr.message || pr.response || pr.output;
|
||||||
|
if (typeof text === 'string' && text.length > 10) {
|
||||||
|
return filterEphemeral(text);
|
||||||
|
}
|
||||||
|
// Search other fields, but skip non-content ones
|
||||||
|
for (const key of Object.keys(pr)) {
|
||||||
|
if (SKIP_FIELDS.has(key)) continue;
|
||||||
|
const val = pr[key];
|
||||||
|
if (typeof val === 'string' && val.length > 50) { // Higher threshold
|
||||||
|
const filtered = filterEphemeral(val);
|
||||||
|
if (filtered) {
|
||||||
|
console.log(`Gravity Bridge: [DEBUG] planner text in plannerResponse.${key} (${filtered.length} chars)`);
|
||||||
|
return filtered;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try other step fields (skip known non-content)
|
||||||
|
for (const key of Object.keys(step)) {
|
||||||
|
if (SKIP_FIELDS.has(key) || key === 'plannerResponse') continue;
|
||||||
|
const val = step[key];
|
||||||
|
if (typeof val === 'string' && val.length > 50) {
|
||||||
|
const filtered = filterEphemeral(val);
|
||||||
|
if (filtered) {
|
||||||
|
console.log(`Gravity Bridge: [DEBUG] planner text in step.${key}`);
|
||||||
|
return filtered;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Filter out system ephemeral messages and non-content strings. */
|
||||||
|
export function filterEphemeral(text: string): string | null {
|
||||||
|
if (!text || text.length < 10) { return null; }
|
||||||
|
// Skip system prompt metadata
|
||||||
|
if (text.includes('<EPHEMERAL_MESSAGE>') || text.includes('<ephemeral_message>')) { return null; }
|
||||||
|
if (text.includes('artifact_reminder') || text.includes('active_task_reminder')) { return null; }
|
||||||
|
if (text.includes('no_active_task_reminder')) { return null; }
|
||||||
|
// Skip base64/crypto strings (no spaces, mostly alphanumeric)
|
||||||
|
if (!text.includes(' ') && /^[A-Za-z0-9+/=_\-]{50,}$/.test(text)) { return null; }
|
||||||
|
return text;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Extract human-readable command from a tool call step's data. */
|
||||||
|
export function extractToolCommand(stepData: any): string {
|
||||||
|
// Try common step data shapes from protobuf
|
||||||
|
if (stepData.runCommand) {
|
||||||
|
return stepData.runCommand.commandLine || stepData.runCommand.command || 'Run Command';
|
||||||
|
}
|
||||||
|
if (stepData.writeToFile) {
|
||||||
|
const target = stepData.writeToFile.targetFile || stepData.writeToFile.filePath || 'file';
|
||||||
|
return `Write: ${target.split(/[\\/]/).pop()}`;
|
||||||
|
}
|
||||||
|
if (stepData.codeAction) {
|
||||||
|
const fp = stepData.codeAction.filePath || '';
|
||||||
|
return `Edit: ${fp.split(/[\\/]/).pop() || 'file'}`;
|
||||||
|
}
|
||||||
|
if (stepData.replaceFileContent || stepData.multiReplaceFileContent) {
|
||||||
|
const d = stepData.replaceFileContent || stepData.multiReplaceFileContent;
|
||||||
|
const fp = d.targetFile || d.filePath || '';
|
||||||
|
return `Edit: ${fp.split(/[\\/]/).pop() || 'file'}`;
|
||||||
|
}
|
||||||
|
if (stepData.sendCommandInput) {
|
||||||
|
return `Send Input: ${(stepData.sendCommandInput.input || '').substring(0, 50)}`;
|
||||||
|
}
|
||||||
|
// Generic fallback: use first key name
|
||||||
|
const keys = Object.keys(stepData).filter(k => k !== 'status' && k !== 'stepStatus');
|
||||||
|
return keys.length > 0 ? keys[0] : 'Unknown tool call';
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Extract description from a tool call step for Discord display. */
|
||||||
|
export function extractToolDescription(stepData: any, sessionTitle: string, stepIndex: number): string {
|
||||||
|
const parts = [`Step #${stepIndex}`, `Session: "${sessionTitle}"`];
|
||||||
|
// Try to get code/command content for context
|
||||||
|
if (stepData.runCommand) {
|
||||||
|
const cmd = stepData.runCommand.commandLine || stepData.runCommand.command || '';
|
||||||
|
if (cmd) parts.push(`Command: ${cmd.substring(0, 200)}`);
|
||||||
|
}
|
||||||
|
if (stepData.writeToFile?.targetFile) {
|
||||||
|
parts.push(`File: ${stepData.writeToFile.targetFile}`);
|
||||||
|
}
|
||||||
|
if (stepData.codeAction?.filePath) {
|
||||||
|
parts.push(`File: ${stepData.codeAction.filePath}`);
|
||||||
|
}
|
||||||
|
return parts.join('\n');
|
||||||
|
}
|
||||||
505
extension/src/ws-client.ts
Normal file
505
extension/src/ws-client.ts
Normal file
@@ -0,0 +1,505 @@
|
|||||||
|
/**
|
||||||
|
* WebSocket Bridge Client — connects Extension to the Hub server.
|
||||||
|
*
|
||||||
|
* Replaces file-based IPC for:
|
||||||
|
* - Pending approvals (Extension → Hub → Bot → Discord)
|
||||||
|
* - User responses (Discord → Bot → Hub → Extension)
|
||||||
|
* - Chat snapshots (Extension → Hub → Bot → Discord)
|
||||||
|
* - Commands (Discord → Bot → Hub → Extension)
|
||||||
|
* - Session registration
|
||||||
|
* - Auto-resolve notifications
|
||||||
|
*
|
||||||
|
* Features:
|
||||||
|
* - Exponential backoff + jitter reconnection
|
||||||
|
* - Message queue (survives reconnection)
|
||||||
|
* - Heartbeat ping/pong
|
||||||
|
* - First-message JWT authentication
|
||||||
|
*/
|
||||||
|
|
||||||
|
import * as vscode from 'vscode';
|
||||||
|
|
||||||
|
// ─── Types ───
|
||||||
|
|
||||||
|
export interface WSMessage {
|
||||||
|
type: string;
|
||||||
|
data?: any;
|
||||||
|
msg_id?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WSAuthMessage {
|
||||||
|
type: 'auth';
|
||||||
|
token?: string;
|
||||||
|
registration_code?: string;
|
||||||
|
project: string;
|
||||||
|
pc: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WSAuthOkResponse {
|
||||||
|
type: 'auth_ok';
|
||||||
|
conn_id: string;
|
||||||
|
instance_number: number;
|
||||||
|
session_token: string;
|
||||||
|
active_count: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WSPendingData {
|
||||||
|
request_id: string;
|
||||||
|
command: string;
|
||||||
|
description?: string;
|
||||||
|
step_type?: string;
|
||||||
|
status?: string;
|
||||||
|
buttons?: Array<{ text: string; index: number }>;
|
||||||
|
project_name?: string;
|
||||||
|
// diff_review metadata
|
||||||
|
edit_step_indices?: number[];
|
||||||
|
modified_files?: string[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WSResponseData {
|
||||||
|
request_id: string;
|
||||||
|
approved: boolean;
|
||||||
|
button_index?: number;
|
||||||
|
step_type?: string;
|
||||||
|
project_name?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WSCommandData {
|
||||||
|
text: string;
|
||||||
|
project_name?: string;
|
||||||
|
action?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WSChatData {
|
||||||
|
content: string;
|
||||||
|
attached_files?: Array<{ name: string; content: string }>;
|
||||||
|
conversation_id?: string;
|
||||||
|
project_name?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WSRegisterData {
|
||||||
|
conversation_id: string;
|
||||||
|
project_name: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Event Handlers ───
|
||||||
|
|
||||||
|
export interface WSBridgeHandlers {
|
||||||
|
onResponse?: (data: WSResponseData) => void;
|
||||||
|
onCommand?: (data: WSCommandData) => void;
|
||||||
|
onInstanceUpdate?: (activeCount: number, instances: Array<{ instance_number: number; pc: string }>) => void;
|
||||||
|
onConnected?: (connId: string, instanceNumber: number, sessionToken: string) => void;
|
||||||
|
onDisconnected?: () => void;
|
||||||
|
onError?: (error: string) => void;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Constants ───
|
||||||
|
|
||||||
|
const INITIAL_RECONNECT_DELAY = 1000; // 1s
|
||||||
|
const MAX_RECONNECT_DELAY = 60000; // 60s
|
||||||
|
const RECONNECT_JITTER = 0.3; // ±30%
|
||||||
|
const HEARTBEAT_INTERVAL = 25000; // 25s (server expects 30s)
|
||||||
|
const MAX_QUEUE_SIZE = 200;
|
||||||
|
const AUTH_TIMEOUT = 10000; // 10s
|
||||||
|
|
||||||
|
// ─── WSBridgeClient ───
|
||||||
|
|
||||||
|
export class WSBridgeClient {
|
||||||
|
private ws: any = null; // WebSocket instance (Node.js ws module)
|
||||||
|
private hubUrl: string;
|
||||||
|
private registrationCode: string;
|
||||||
|
private project: string;
|
||||||
|
private pcName: string;
|
||||||
|
private handlers: WSBridgeHandlers;
|
||||||
|
private logFn: (msg: string) => void;
|
||||||
|
|
||||||
|
// Connection state
|
||||||
|
private connected = false;
|
||||||
|
private authenticated = false;
|
||||||
|
private connId = '';
|
||||||
|
private instanceNumber = 0;
|
||||||
|
private sessionToken = '';
|
||||||
|
private shouldReconnect = true;
|
||||||
|
private reconnectDelay = INITIAL_RECONNECT_DELAY;
|
||||||
|
private reconnectTimer: NodeJS.Timeout | null = null;
|
||||||
|
private heartbeatTimer: NodeJS.Timeout | null = null;
|
||||||
|
private authTimer: NodeJS.Timeout | null = null;
|
||||||
|
|
||||||
|
// Message queue (survives reconnection)
|
||||||
|
private messageQueue: WSMessage[] = [];
|
||||||
|
private msgIdCounter = 0;
|
||||||
|
|
||||||
|
constructor(
|
||||||
|
hubUrl: string,
|
||||||
|
registrationCode: string,
|
||||||
|
project: string,
|
||||||
|
pcName: string,
|
||||||
|
handlers: WSBridgeHandlers,
|
||||||
|
logFn: (msg: string) => void,
|
||||||
|
) {
|
||||||
|
this.hubUrl = hubUrl;
|
||||||
|
this.registrationCode = registrationCode;
|
||||||
|
this.project = project;
|
||||||
|
this.pcName = pcName;
|
||||||
|
this.handlers = handlers;
|
||||||
|
this.logFn = logFn;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Public API ───
|
||||||
|
|
||||||
|
/** Start the WebSocket connection. */
|
||||||
|
async connect(): Promise<void> {
|
||||||
|
if (!this.hubUrl) {
|
||||||
|
this.logFn('[WS] No hub URL configured — WS disabled');
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
this.shouldReconnect = true;
|
||||||
|
await this._connect();
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Gracefully disconnect. */
|
||||||
|
disconnect(): void {
|
||||||
|
this.shouldReconnect = false;
|
||||||
|
this._cleanup();
|
||||||
|
this.logFn('[WS] Disconnected (intentional)');
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Check if connected and authenticated. */
|
||||||
|
isConnected(): boolean {
|
||||||
|
return this.connected && this.authenticated;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Get the instance number assigned by the Hub. */
|
||||||
|
getInstanceNumber(): number {
|
||||||
|
return this.instanceNumber;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Send a pending approval to the Hub. */
|
||||||
|
sendPending(data: WSPendingData): boolean {
|
||||||
|
return this._send({ type: 'pending', data });
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Send a chat snapshot to the Hub. */
|
||||||
|
sendChat(data: WSChatData): boolean {
|
||||||
|
return this._send({ type: 'chat', data });
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Send a session registration. */
|
||||||
|
sendRegister(data: WSRegisterData): boolean {
|
||||||
|
return this._send({ type: 'register', data });
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Send an auto_resolve notification. */
|
||||||
|
sendAutoResolve(requestId: string): boolean {
|
||||||
|
return this._send({ type: 'auto_resolve', data: { request_id: requestId } });
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Send a brain event. */
|
||||||
|
sendBrainEvent(data: any): boolean {
|
||||||
|
return this._send({ type: 'brain_event', data });
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Internal Connection ───
|
||||||
|
|
||||||
|
private async _connect(): Promise<void> {
|
||||||
|
try {
|
||||||
|
// Dynamic import of ws module (Node.js built-in or npm package)
|
||||||
|
const WebSocket = await this._getWebSocketClass();
|
||||||
|
if (!WebSocket) {
|
||||||
|
this.logFn('[WS] WebSocket module not available');
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
this.logFn(`[WS] Connecting to ${this.hubUrl}...`);
|
||||||
|
const ws = new WebSocket(this.hubUrl);
|
||||||
|
|
||||||
|
ws.on('open', () => {
|
||||||
|
this.logFn('[WS] Connection opened, authenticating...');
|
||||||
|
this.ws = ws;
|
||||||
|
this.connected = true;
|
||||||
|
this._authenticate();
|
||||||
|
});
|
||||||
|
|
||||||
|
ws.on('message', (raw: Buffer | string) => {
|
||||||
|
try {
|
||||||
|
const data = JSON.parse(typeof raw === 'string' ? raw : raw.toString('utf-8'));
|
||||||
|
this._handleMessage(data);
|
||||||
|
} catch (e: any) {
|
||||||
|
this.logFn(`[WS] Parse error: ${e.message}`);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
ws.on('close', (code: number, reason: Buffer) => {
|
||||||
|
const reasonStr = reason ? reason.toString('utf-8') : '';
|
||||||
|
this.logFn(`[WS] Connection closed: code=${code} reason=${reasonStr}`);
|
||||||
|
this._onDisconnect();
|
||||||
|
});
|
||||||
|
|
||||||
|
ws.on('error', (err: Error) => {
|
||||||
|
this.logFn(`[WS] Error: ${err.message}`);
|
||||||
|
// close event will follow
|
||||||
|
});
|
||||||
|
|
||||||
|
ws.on('pong', () => {
|
||||||
|
// Server responded to our ping — connection is alive
|
||||||
|
});
|
||||||
|
|
||||||
|
} catch (e: any) {
|
||||||
|
this.logFn(`[WS] Connect failed: ${e.message}`);
|
||||||
|
this._scheduleReconnect();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private async _getWebSocketClass(): Promise<any> {
|
||||||
|
try {
|
||||||
|
// Try Node.js built-in WebSocket (v21+)
|
||||||
|
if (typeof globalThis.WebSocket !== 'undefined') {
|
||||||
|
return globalThis.WebSocket;
|
||||||
|
}
|
||||||
|
// Try require('ws') — should be available in VS Code's Node.js
|
||||||
|
const ws = require('ws');
|
||||||
|
return ws;
|
||||||
|
} catch {
|
||||||
|
// ws module not available
|
||||||
|
try {
|
||||||
|
// Fallback: try the built-in undici WebSocket
|
||||||
|
const { WebSocket } = require('undici');
|
||||||
|
return WebSocket;
|
||||||
|
} catch {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Authentication ───
|
||||||
|
|
||||||
|
private _authenticate(): void {
|
||||||
|
if (!this.ws) return;
|
||||||
|
|
||||||
|
const authMsg: WSAuthMessage = {
|
||||||
|
type: 'auth',
|
||||||
|
project: this.project,
|
||||||
|
pc: this.pcName,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Use session token if available (from previous connection)
|
||||||
|
if (this.sessionToken) {
|
||||||
|
authMsg.token = this.sessionToken;
|
||||||
|
} else if (this.registrationCode) {
|
||||||
|
authMsg.registration_code = this.registrationCode;
|
||||||
|
}
|
||||||
|
|
||||||
|
this._sendRaw(authMsg);
|
||||||
|
|
||||||
|
// Timeout for auth response
|
||||||
|
this.authTimer = setTimeout(() => {
|
||||||
|
if (!this.authenticated) {
|
||||||
|
this.logFn('[WS] Auth timeout — closing connection');
|
||||||
|
this._cleanup();
|
||||||
|
this._scheduleReconnect();
|
||||||
|
}
|
||||||
|
}, AUTH_TIMEOUT);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Message Handling ───
|
||||||
|
|
||||||
|
private _handleMessage(msg: WSMessage): void {
|
||||||
|
switch (msg.type) {
|
||||||
|
case 'auth_ok': {
|
||||||
|
const authOk = msg as unknown as WSAuthOkResponse;
|
||||||
|
this.authenticated = true;
|
||||||
|
this.connId = authOk.conn_id;
|
||||||
|
this.instanceNumber = authOk.instance_number;
|
||||||
|
this.sessionToken = authOk.session_token;
|
||||||
|
this.reconnectDelay = INITIAL_RECONNECT_DELAY;
|
||||||
|
|
||||||
|
if (this.authTimer) {
|
||||||
|
clearTimeout(this.authTimer);
|
||||||
|
this.authTimer = null;
|
||||||
|
}
|
||||||
|
|
||||||
|
this.logFn(`[WS] Authenticated: conn=${this.connId} instance=#${this.instanceNumber} active=${authOk.active_count}`);
|
||||||
|
this._startHeartbeat();
|
||||||
|
this._flushQueue();
|
||||||
|
this.handlers.onConnected?.(this.connId, this.instanceNumber, this.sessionToken);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
case 'auth_fail': {
|
||||||
|
const reason = (msg as any).reason || 'Unknown';
|
||||||
|
this.logFn(`[WS] Auth failed: ${reason}`);
|
||||||
|
// Clear session token if it was rejected
|
||||||
|
this.sessionToken = '';
|
||||||
|
this._cleanup();
|
||||||
|
// Don't reconnect on auth failure (needs manual fix)
|
||||||
|
this.handlers.onError?.(`Auth failed: ${reason}`);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
case 'response': {
|
||||||
|
const data = msg.data as WSResponseData;
|
||||||
|
if (data) {
|
||||||
|
this.logFn(`[WS] Response received: ${data.request_id?.substring(0, 12)} approved=${data.approved}`);
|
||||||
|
this.handlers.onResponse?.(data);
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
case 'command': {
|
||||||
|
const data = msg.data as WSCommandData;
|
||||||
|
if (data) {
|
||||||
|
this.logFn(`[WS] Command received: ${data.text?.substring(0, 50)}`);
|
||||||
|
this.handlers.onCommand?.(data);
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
case 'instance_update': {
|
||||||
|
const activeCount = (msg as any).active_count || 0;
|
||||||
|
const instances = (msg as any).instances || [];
|
||||||
|
this.logFn(`[WS] Instance update: ${activeCount} active`);
|
||||||
|
this.handlers.onInstanceUpdate?.(activeCount, instances);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
case 'error': {
|
||||||
|
const error = (msg as any).error || 'Unknown error';
|
||||||
|
this.logFn(`[WS] Server error: ${error}`);
|
||||||
|
this.handlers.onError?.(error);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
default:
|
||||||
|
this.logFn(`[WS] Unknown message type: ${msg.type}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Send ───
|
||||||
|
|
||||||
|
private _send(msg: WSMessage): boolean {
|
||||||
|
// Add unique message ID for dedup
|
||||||
|
msg.msg_id = `${this.project}-${Date.now()}-${++this.msgIdCounter}`;
|
||||||
|
|
||||||
|
if (this.isConnected()) {
|
||||||
|
return this._sendRaw(msg);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Queue for later
|
||||||
|
if (this.messageQueue.length >= MAX_QUEUE_SIZE) {
|
||||||
|
// Drop oldest
|
||||||
|
this.messageQueue.shift();
|
||||||
|
this.logFn('[WS] Queue full — dropped oldest message');
|
||||||
|
}
|
||||||
|
this.messageQueue.push(msg);
|
||||||
|
this.logFn(`[WS] Queued message (type=${msg.type}, queue=${this.messageQueue.length})`);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
private _sendRaw(msg: any): boolean {
|
||||||
|
try {
|
||||||
|
if (this.ws && this.connected) {
|
||||||
|
this.ws.send(JSON.stringify(msg));
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
return false;
|
||||||
|
} catch (e: any) {
|
||||||
|
this.logFn(`[WS] Send error: ${e.message}`);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private _flushQueue(): void {
|
||||||
|
if (this.messageQueue.length === 0) return;
|
||||||
|
this.logFn(`[WS] Flushing ${this.messageQueue.length} queued messages`);
|
||||||
|
const queue = [...this.messageQueue];
|
||||||
|
this.messageQueue = [];
|
||||||
|
for (const msg of queue) {
|
||||||
|
this._sendRaw(msg);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Heartbeat ───
|
||||||
|
|
||||||
|
private _startHeartbeat(): void {
|
||||||
|
this._stopHeartbeat();
|
||||||
|
this.heartbeatTimer = setInterval(() => {
|
||||||
|
if (this.ws && this.connected) {
|
||||||
|
try {
|
||||||
|
this.ws.ping();
|
||||||
|
} catch {
|
||||||
|
// ping failure will trigger close event
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}, HEARTBEAT_INTERVAL);
|
||||||
|
}
|
||||||
|
|
||||||
|
private _stopHeartbeat(): void {
|
||||||
|
if (this.heartbeatTimer) {
|
||||||
|
clearInterval(this.heartbeatTimer);
|
||||||
|
this.heartbeatTimer = null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Reconnection ───
|
||||||
|
|
||||||
|
private _onDisconnect(): void {
|
||||||
|
const wasAuthenticated = this.authenticated;
|
||||||
|
this.connected = false;
|
||||||
|
this.authenticated = false;
|
||||||
|
this.ws = null;
|
||||||
|
|
||||||
|
this._stopHeartbeat();
|
||||||
|
if (this.authTimer) {
|
||||||
|
clearTimeout(this.authTimer);
|
||||||
|
this.authTimer = null;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (wasAuthenticated) {
|
||||||
|
this.handlers.onDisconnected?.();
|
||||||
|
}
|
||||||
|
|
||||||
|
if (this.shouldReconnect) {
|
||||||
|
this._scheduleReconnect();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private _scheduleReconnect(): void {
|
||||||
|
if (this.reconnectTimer) return;
|
||||||
|
|
||||||
|
// Exponential backoff with jitter
|
||||||
|
const jitter = 1 + (Math.random() * 2 - 1) * RECONNECT_JITTER;
|
||||||
|
const delay = Math.min(this.reconnectDelay * jitter, MAX_RECONNECT_DELAY);
|
||||||
|
this.logFn(`[WS] Reconnecting in ${Math.round(delay)}ms...`);
|
||||||
|
|
||||||
|
this.reconnectTimer = setTimeout(() => {
|
||||||
|
this.reconnectTimer = null;
|
||||||
|
this.reconnectDelay = Math.min(this.reconnectDelay * 2, MAX_RECONNECT_DELAY);
|
||||||
|
this._connect();
|
||||||
|
}, delay);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── Cleanup ───
|
||||||
|
|
||||||
|
private _cleanup(): void {
|
||||||
|
this._stopHeartbeat();
|
||||||
|
|
||||||
|
if (this.authTimer) {
|
||||||
|
clearTimeout(this.authTimer);
|
||||||
|
this.authTimer = null;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (this.reconnectTimer) {
|
||||||
|
clearTimeout(this.reconnectTimer);
|
||||||
|
this.reconnectTimer = null;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (this.ws) {
|
||||||
|
try {
|
||||||
|
this.ws.close();
|
||||||
|
} catch { }
|
||||||
|
this.ws = null;
|
||||||
|
}
|
||||||
|
|
||||||
|
this.connected = false;
|
||||||
|
this.authenticated = false;
|
||||||
|
}
|
||||||
|
}
|
||||||
47
gateway.py
47
gateway.py
@@ -1,10 +1,12 @@
|
|||||||
"""Gateway HTTP API — receives data from remote Collectors and routes to Discord bot.
|
"""Gateway HTTP API + WebSocket Hub — receives data from Collectors and Extensions.
|
||||||
|
|
||||||
Runs alongside the Discord bot in the server Docker container.
|
Runs alongside the Discord bot in the server Docker container.
|
||||||
Collectors (local PCs) push pending approvals, chat snapshots, and registrations
|
Supports both:
|
||||||
to this API, and poll for responses.
|
- REST API: for legacy Collectors (HTTP polling)
|
||||||
|
- WebSocket: for direct Extension connections (real-time)
|
||||||
|
|
||||||
Endpoints:
|
Endpoints:
|
||||||
|
GET /ws — WebSocket endpoint (Extension direct connection)
|
||||||
POST /api/pending — Collector pushes a new approval request
|
POST /api/pending — Collector pushes a new approval request
|
||||||
GET /api/pending — List all pending requests (for diagnostics)
|
GET /api/pending — List all pending requests (for diagnostics)
|
||||||
POST /api/response/{rid} — Collector polls for response (or Gateway pushes)
|
POST /api/response/{rid} — Collector polls for response (or Gateway pushes)
|
||||||
@@ -14,6 +16,7 @@ Endpoints:
|
|||||||
POST /api/command — Gateway pushes command to specific collector
|
POST /api/command — Gateway pushes command to specific collector
|
||||||
GET /api/commands/{project} — Collector polls for commands
|
GET /api/commands/{project} — Collector polls for commands
|
||||||
GET /health — Health check
|
GET /health — Health check
|
||||||
|
GET /hub/status — WebSocket Hub diagnostics
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
@@ -34,13 +37,15 @@ COMMAND_TTL = 1800 # 30 min — stale commands auto-deleted
|
|||||||
|
|
||||||
|
|
||||||
class GatewayAPI:
|
class GatewayAPI:
|
||||||
"""HTTP API server for Collector ↔ Gateway communication."""
|
"""HTTP API + WebSocket Hub server."""
|
||||||
|
|
||||||
def __init__(self, bot, host: str = "0.0.0.0", port: int = 8585, api_key: str = ""):
|
def __init__(self, bot, host: str = "0.0.0.0", port: int = 8585,
|
||||||
|
api_key: str = "", hub=None):
|
||||||
self.bot = bot
|
self.bot = bot
|
||||||
self.host = host
|
self.host = host
|
||||||
self.port = port
|
self.port = port
|
||||||
self.api_key = api_key
|
self.api_key = api_key
|
||||||
|
self.hub = hub # WSHub instance (None = WS disabled)
|
||||||
self.app = web.Application(
|
self.app = web.Application(
|
||||||
middlewares=[self._auth_middleware],
|
middlewares=[self._auth_middleware],
|
||||||
client_max_size=1024 * 1024, # Security: 1MB max request body
|
client_max_size=1024 * 1024, # Security: 1MB max request body
|
||||||
@@ -52,6 +57,10 @@ class GatewayAPI:
|
|||||||
self._rate_limits: dict[str, list[float]] = defaultdict(list) # IP → [timestamps]
|
self._rate_limits: dict[str, list[float]] = defaultdict(list) # IP → [timestamps]
|
||||||
|
|
||||||
def _setup_routes(self):
|
def _setup_routes(self):
|
||||||
|
# WebSocket endpoint (no auth middleware — Hub handles its own auth)
|
||||||
|
self.app.router.add_get("/ws", self._ws_handler)
|
||||||
|
self.app.router.add_get("/hub/status", self._hub_status)
|
||||||
|
# Legacy REST endpoints (Collector compatibility)
|
||||||
self.app.router.add_get("/health", self._health)
|
self.app.router.add_get("/health", self._health)
|
||||||
self.app.router.add_post("/api/pending", self._post_pending)
|
self.app.router.add_post("/api/pending", self._post_pending)
|
||||||
self.app.router.add_get("/api/pending", self._list_pending)
|
self.app.router.add_get("/api/pending", self._list_pending)
|
||||||
@@ -61,13 +70,29 @@ class GatewayAPI:
|
|||||||
self.app.router.add_get("/api/commands/{project}", self._get_commands)
|
self.app.router.add_get("/api/commands/{project}", self._get_commands)
|
||||||
self.app.router.add_post("/api/event", self._post_event)
|
self.app.router.add_post("/api/event", self._post_event)
|
||||||
|
|
||||||
|
# ─── WebSocket Handler ───
|
||||||
|
|
||||||
|
async def _ws_handler(self, request: web.Request) -> web.WebSocketResponse:
|
||||||
|
"""WebSocket endpoint for direct Extension connections."""
|
||||||
|
if not self.hub:
|
||||||
|
return web.json_response(
|
||||||
|
{"error": "WebSocket Hub not enabled"}, status=503
|
||||||
|
)
|
||||||
|
return await self.hub.handle_ws(request)
|
||||||
|
|
||||||
|
async def _hub_status(self, request: web.Request) -> web.Response:
|
||||||
|
"""WebSocket Hub diagnostics."""
|
||||||
|
if not self.hub:
|
||||||
|
return web.json_response({"hub": "disabled"})
|
||||||
|
return web.json_response(self.hub.get_status())
|
||||||
|
|
||||||
# ─── Auth Middleware ───
|
# ─── Auth Middleware ───
|
||||||
|
|
||||||
@web.middleware
|
@web.middleware
|
||||||
async def _auth_middleware(self, request: web.Request, handler):
|
async def _auth_middleware(self, request: web.Request, handler):
|
||||||
"""Reject requests without valid API key on /api/* routes."""
|
"""Reject requests without valid API key on /api/* routes."""
|
||||||
# Health endpoint is public
|
# WebSocket and public endpoints skip API key auth
|
||||||
if request.path == "/health":
|
if request.path in ("/health", "/ws", "/hub/status"):
|
||||||
return await handler(request)
|
return await handler(request)
|
||||||
|
|
||||||
# All /api/* routes require auth + rate limit
|
# All /api/* routes require auth + rate limit
|
||||||
@@ -109,11 +134,15 @@ class GatewayAPI:
|
|||||||
# ─── Health ───
|
# ─── Health ───
|
||||||
|
|
||||||
async def _health(self, request: web.Request) -> web.Response:
|
async def _health(self, request: web.Request) -> web.Response:
|
||||||
return web.json_response({
|
status = {
|
||||||
"status": "ok",
|
"status": "ok",
|
||||||
"bot_ready": self.bot.is_ready() if self.bot else False,
|
"bot_ready": self.bot.is_ready() if self.bot else False,
|
||||||
"timestamp": time.time(),
|
"timestamp": time.time(),
|
||||||
})
|
"hub_enabled": self.hub is not None,
|
||||||
|
}
|
||||||
|
if self.hub:
|
||||||
|
status["hub_connections"] = len(self.hub.connections)
|
||||||
|
return web.json_response(status)
|
||||||
|
|
||||||
# ─── Pending Approvals (Collector → Gateway → Discord) ───
|
# ─── Pending Approvals (Collector → Gateway → Discord) ───
|
||||||
|
|
||||||
|
|||||||
580
hub.py
Normal file
580
hub.py
Normal file
@@ -0,0 +1,580 @@
|
|||||||
|
"""WebSocket Hub — real-time bidirectional communication between Extensions and Bot.
|
||||||
|
|
||||||
|
Replaces file-based IPC (bridge/) and HTTP polling (Collector/Gateway) with
|
||||||
|
persistent WebSocket connections.
|
||||||
|
|
||||||
|
Architecture:
|
||||||
|
Extension ↔ WSS ↔ Hub ↔ Bot (in-process) ↔ Discord
|
||||||
|
|
||||||
|
Each Extension connects via WebSocket, authenticates, and receives a unique
|
||||||
|
instance number within its project. Messages are routed by project.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from enum import Enum
|
||||||
|
from typing import Any, Callable, Coroutine
|
||||||
|
|
||||||
|
from aiohttp import web, WSMsgType
|
||||||
|
|
||||||
|
from auth import TokenManager
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# ─── Constants ───
|
||||||
|
|
||||||
|
HEARTBEAT_INTERVAL = 30.0 # seconds between server→client pings
|
||||||
|
HEARTBEAT_TIMEOUT = 10.0 # seconds to wait for pong before disconnect
|
||||||
|
CLIENT_QUEUE_SIZE = 100 # max queued messages per client (backpressure)
|
||||||
|
MAX_MSG_SIZE = 1024 * 1024 # 1MB max WebSocket message size
|
||||||
|
PER_CONN_RATE_LIMIT = 60 # max messages per 10s window per connection
|
||||||
|
RATE_WINDOW = 10.0 # seconds
|
||||||
|
|
||||||
|
|
||||||
|
class MsgType(str, Enum):
|
||||||
|
"""WebSocket protocol message types."""
|
||||||
|
# Extension → Hub (upstream)
|
||||||
|
AUTH = "auth"
|
||||||
|
PENDING = "pending"
|
||||||
|
CHAT = "chat"
|
||||||
|
REGISTER = "register"
|
||||||
|
HEARTBEAT = "heartbeat"
|
||||||
|
AUTO_RESOLVE = "auto_resolve"
|
||||||
|
BRAIN_EVENT = "brain_event"
|
||||||
|
|
||||||
|
# Hub → Extension (downstream)
|
||||||
|
AUTH_OK = "auth_ok"
|
||||||
|
AUTH_FAIL = "auth_fail"
|
||||||
|
RESPONSE = "response"
|
||||||
|
COMMAND = "command"
|
||||||
|
INSTANCE_UPDATE = "instance_update"
|
||||||
|
ERROR = "error"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class WSConnection:
|
||||||
|
"""Represents a connected Extension instance."""
|
||||||
|
conn_id: str
|
||||||
|
ws: web.WebSocketResponse
|
||||||
|
project: str = ""
|
||||||
|
pc_name: str = ""
|
||||||
|
instance_number: int = 0
|
||||||
|
connected_at: float = field(default_factory=time.time)
|
||||||
|
last_heartbeat: float = field(default_factory=time.time)
|
||||||
|
authenticated: bool = False
|
||||||
|
send_queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=CLIENT_QUEUE_SIZE))
|
||||||
|
_sender_task: asyncio.Task | None = field(default=None, repr=False)
|
||||||
|
# Rate limiting
|
||||||
|
_msg_timestamps: list[float] = field(default_factory=list, repr=False)
|
||||||
|
|
||||||
|
|
||||||
|
class WSHub:
|
||||||
|
"""WebSocket Hub for routing messages between Extensions and Bot.
|
||||||
|
|
||||||
|
Responsibilities:
|
||||||
|
- Connection lifecycle (auth, heartbeat, disconnect)
|
||||||
|
- Message routing (project-scoped, instance-targeted)
|
||||||
|
- Instance number management (auto-assign, reassign on disconnect)
|
||||||
|
- Rate limiting per connection
|
||||||
|
- Backpressure via per-client queues
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, token_manager: TokenManager):
|
||||||
|
self.token_manager = token_manager
|
||||||
|
self.connections: dict[str, WSConnection] = {} # conn_id → connection
|
||||||
|
self.project_connections: dict[str, set[str]] = {} # project → {conn_ids}
|
||||||
|
self.pending_owners: dict[str, str] = {} # request_id → conn_id
|
||||||
|
self._recent_msg_ids: dict[str, float] = {} # msg_id → timestamp (dedup)
|
||||||
|
self._msg_id_ttl = 300.0 # 5 min dedup window
|
||||||
|
|
||||||
|
# Bot callbacks — set by bot.py during initialization
|
||||||
|
self._on_pending: Callable[..., Coroutine] | None = None
|
||||||
|
self._on_chat: Callable[..., Coroutine] | None = None
|
||||||
|
self._on_register: Callable[..., Coroutine] | None = None
|
||||||
|
self._on_auto_resolve: Callable[..., Coroutine] | None = None
|
||||||
|
self._on_brain_event: Callable[..., Coroutine] | None = None
|
||||||
|
|
||||||
|
# ─── Bot Integration ───
|
||||||
|
|
||||||
|
def set_bot_handlers(
|
||||||
|
self,
|
||||||
|
on_pending: Callable | None = None,
|
||||||
|
on_chat: Callable | None = None,
|
||||||
|
on_register: Callable | None = None,
|
||||||
|
on_auto_resolve: Callable | None = None,
|
||||||
|
on_brain_event: Callable | None = None,
|
||||||
|
):
|
||||||
|
"""Register bot callback functions for incoming Extension messages."""
|
||||||
|
self._on_pending = on_pending
|
||||||
|
self._on_chat = on_chat
|
||||||
|
self._on_register = on_register
|
||||||
|
self._on_auto_resolve = on_auto_resolve
|
||||||
|
self._on_brain_event = on_brain_event
|
||||||
|
|
||||||
|
# ─── Connection Management ───
|
||||||
|
|
||||||
|
async def handle_ws(self, request: web.Request) -> web.WebSocketResponse:
|
||||||
|
"""Handle a new WebSocket connection.
|
||||||
|
|
||||||
|
Protocol:
|
||||||
|
1. Client connects
|
||||||
|
2. Client sends {type:"auth", token:"...", project:"...", pc:"..."}
|
||||||
|
OR {type:"auth", registration_code:"...", project:"...", pc:"..."}
|
||||||
|
3. Server responds {type:"auth_ok", conn_id, instance_number}
|
||||||
|
OR {type:"auth_fail", reason:"..."}
|
||||||
|
4. Bidirectional message exchange
|
||||||
|
"""
|
||||||
|
ws = web.WebSocketResponse(
|
||||||
|
max_msg_size=MAX_MSG_SIZE,
|
||||||
|
heartbeat=HEARTBEAT_INTERVAL,
|
||||||
|
)
|
||||||
|
await ws.prepare(request)
|
||||||
|
|
||||||
|
conn_id = uuid.uuid4().hex[:12]
|
||||||
|
conn = WSConnection(conn_id=conn_id, ws=ws)
|
||||||
|
remote = request.remote or "unknown"
|
||||||
|
logger.info(f"[HUB] New WS connection: {conn_id} from {remote}")
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Wait for auth message (first message must be auth)
|
||||||
|
auth_ok = await self._handle_auth(conn, timeout=10.0)
|
||||||
|
if not auth_ok:
|
||||||
|
return ws
|
||||||
|
|
||||||
|
# Register connection
|
||||||
|
self._register_connection(conn)
|
||||||
|
|
||||||
|
# Start sender task (per-client queue → ws.send)
|
||||||
|
conn._sender_task = asyncio.create_task(
|
||||||
|
self._sender_loop(conn), name=f"sender-{conn_id}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Message loop
|
||||||
|
await self._message_loop(conn)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[HUB] Connection {conn_id} error: {e}")
|
||||||
|
finally:
|
||||||
|
await self._disconnect(conn)
|
||||||
|
|
||||||
|
return ws
|
||||||
|
|
||||||
|
async def _handle_auth(self, conn: WSConnection, timeout: float = 10.0) -> bool:
|
||||||
|
"""Wait for and process the auth message."""
|
||||||
|
try:
|
||||||
|
msg = await asyncio.wait_for(conn.ws.receive(), timeout=timeout)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_FAIL, "reason": "Auth timeout"
|
||||||
|
})
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
|
||||||
|
if msg.type != WSMsgType.TEXT:
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
|
||||||
|
try:
|
||||||
|
data = json.loads(msg.data)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_FAIL, "reason": "Invalid JSON"
|
||||||
|
})
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
|
||||||
|
if data.get("type") != MsgType.AUTH:
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_FAIL, "reason": "First message must be auth"
|
||||||
|
})
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
|
||||||
|
project = data.get("project", "")
|
||||||
|
pc_name = data.get("pc", "unknown")
|
||||||
|
|
||||||
|
if not project:
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_FAIL, "reason": "Project name required"
|
||||||
|
})
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Try token auth first, then registration code
|
||||||
|
token = data.get("token", "")
|
||||||
|
reg_code = data.get("registration_code", "")
|
||||||
|
|
||||||
|
if token:
|
||||||
|
payload = self.token_manager.verify_token(token)
|
||||||
|
if not payload:
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_FAIL, "reason": "Invalid or expired token"
|
||||||
|
})
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
# Token is valid — use project from token (overrides client)
|
||||||
|
project = payload.get("project", project)
|
||||||
|
pc_name = payload.get("pc", pc_name)
|
||||||
|
elif reg_code:
|
||||||
|
if not self.token_manager.validate_registration_code(reg_code):
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_FAIL, "reason": "Invalid registration code"
|
||||||
|
})
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
else:
|
||||||
|
# No auth provided — check if registration code is configured
|
||||||
|
if self.token_manager.registration_code:
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_FAIL, "reason": "Auth required"
|
||||||
|
})
|
||||||
|
await conn.ws.close()
|
||||||
|
return False
|
||||||
|
# No registration code configured → allow (dev mode)
|
||||||
|
|
||||||
|
# Auth successful
|
||||||
|
conn.project = project
|
||||||
|
conn.pc_name = pc_name
|
||||||
|
conn.authenticated = True
|
||||||
|
|
||||||
|
# Issue a session token (for reconnection)
|
||||||
|
session_token = self.token_manager.create_token(project, pc_name)
|
||||||
|
|
||||||
|
# Assign instance number
|
||||||
|
instance_number = self._assign_instance_number(project, conn.conn_id)
|
||||||
|
conn.instance_number = instance_number
|
||||||
|
|
||||||
|
await self._send_direct(conn.ws, {
|
||||||
|
"type": MsgType.AUTH_OK,
|
||||||
|
"conn_id": conn.conn_id,
|
||||||
|
"instance_number": instance_number,
|
||||||
|
"session_token": session_token,
|
||||||
|
"active_count": self.get_active_count(project) + 1, # +1 for self (not yet registered)
|
||||||
|
})
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"[HUB] Auth OK: {conn.conn_id} project={project} pc={pc_name} "
|
||||||
|
f"instance=#{instance_number}"
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _register_connection(self, conn: WSConnection):
|
||||||
|
"""Add connection to tracking structures."""
|
||||||
|
self.connections[conn.conn_id] = conn
|
||||||
|
if conn.project not in self.project_connections:
|
||||||
|
self.project_connections[conn.project] = set()
|
||||||
|
self.project_connections[conn.project].add(conn.conn_id)
|
||||||
|
|
||||||
|
# Broadcast instance update to all project connections
|
||||||
|
asyncio.create_task(self._broadcast_instance_update(conn.project))
|
||||||
|
|
||||||
|
async def _disconnect(self, conn: WSConnection):
|
||||||
|
"""Clean up after a connection closes."""
|
||||||
|
conn_id = conn.conn_id
|
||||||
|
project = conn.project
|
||||||
|
|
||||||
|
# Cancel sender task
|
||||||
|
if conn._sender_task and not conn._sender_task.done():
|
||||||
|
conn._sender_task.cancel()
|
||||||
|
try:
|
||||||
|
await conn._sender_task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Remove from tracking
|
||||||
|
self.connections.pop(conn_id, None)
|
||||||
|
if project in self.project_connections:
|
||||||
|
self.project_connections[project].discard(conn_id)
|
||||||
|
if not self.project_connections[project]:
|
||||||
|
del self.project_connections[project]
|
||||||
|
|
||||||
|
# Clean up pending ownership
|
||||||
|
stale = [rid for rid, cid in self.pending_owners.items() if cid == conn_id]
|
||||||
|
for rid in stale:
|
||||||
|
del self.pending_owners[rid]
|
||||||
|
|
||||||
|
# Close WebSocket if still open
|
||||||
|
if not conn.ws.closed:
|
||||||
|
await conn.ws.close()
|
||||||
|
|
||||||
|
logger.info(f"[HUB] Disconnected: {conn_id} project={project}")
|
||||||
|
|
||||||
|
# Broadcast instance update (remaining connections get notified)
|
||||||
|
if project:
|
||||||
|
await self._broadcast_instance_update(project)
|
||||||
|
|
||||||
|
# ─── Instance Number Management ───
|
||||||
|
|
||||||
|
def _assign_instance_number(self, project: str, conn_id: str) -> int:
|
||||||
|
"""Assign the next available instance number for a project."""
|
||||||
|
used = set()
|
||||||
|
for cid in self.project_connections.get(project, set()):
|
||||||
|
c = self.connections.get(cid)
|
||||||
|
if c:
|
||||||
|
used.add(c.instance_number)
|
||||||
|
|
||||||
|
# Find lowest available number starting from 1
|
||||||
|
num = 1
|
||||||
|
while num in used:
|
||||||
|
num += 1
|
||||||
|
return num
|
||||||
|
|
||||||
|
def get_active_count(self, project: str) -> int:
|
||||||
|
"""Get the number of active connections for a project."""
|
||||||
|
return len(self.project_connections.get(project, set()))
|
||||||
|
|
||||||
|
async def _broadcast_instance_update(self, project: str):
|
||||||
|
"""Notify all project connections of instance count change."""
|
||||||
|
count = self.get_active_count(project)
|
||||||
|
msg = {
|
||||||
|
"type": MsgType.INSTANCE_UPDATE,
|
||||||
|
"active_count": count,
|
||||||
|
"instances": [],
|
||||||
|
}
|
||||||
|
# Include instance list for UI
|
||||||
|
for cid in self.project_connections.get(project, set()):
|
||||||
|
c = self.connections.get(cid)
|
||||||
|
if c:
|
||||||
|
msg["instances"].append({
|
||||||
|
"instance_number": c.instance_number,
|
||||||
|
"pc": c.pc_name,
|
||||||
|
})
|
||||||
|
await self.broadcast_to_project(project, msg)
|
||||||
|
|
||||||
|
# ─── Message Routing ───
|
||||||
|
|
||||||
|
async def broadcast_to_project(self, project: str, message: dict):
|
||||||
|
"""Send a message to all connections in a project."""
|
||||||
|
conn_ids = self.project_connections.get(project, set()).copy()
|
||||||
|
for cid in conn_ids:
|
||||||
|
conn = self.connections.get(cid)
|
||||||
|
if conn and conn.authenticated:
|
||||||
|
await self._queue_send(conn, message)
|
||||||
|
|
||||||
|
async def send_to_connection(self, conn_id: str, message: dict):
|
||||||
|
"""Send a message to a specific connection."""
|
||||||
|
conn = self.connections.get(conn_id)
|
||||||
|
if conn and conn.authenticated:
|
||||||
|
await self._queue_send(conn, message)
|
||||||
|
|
||||||
|
async def send_to_instance(self, project: str, instance_number: int, message: dict):
|
||||||
|
"""Send a message to a specific instance number within a project."""
|
||||||
|
for cid in self.project_connections.get(project, set()):
|
||||||
|
conn = self.connections.get(cid)
|
||||||
|
if conn and conn.instance_number == instance_number:
|
||||||
|
await self._queue_send(conn, message)
|
||||||
|
return True
|
||||||
|
logger.warning(
|
||||||
|
f"[HUB] Instance #{instance_number} not found for project={project}"
|
||||||
|
)
|
||||||
|
return False
|
||||||
|
|
||||||
|
async def send_response_to_pending_owner(self, request_id: str, message: dict):
|
||||||
|
"""Route a response to the Extension that created the pending request."""
|
||||||
|
conn_id = self.pending_owners.get(request_id)
|
||||||
|
if conn_id:
|
||||||
|
await self.send_to_connection(conn_id, message)
|
||||||
|
# Clean up after response delivered
|
||||||
|
self.pending_owners.pop(request_id, None)
|
||||||
|
return True
|
||||||
|
logger.warning(f"[HUB] No owner for pending {request_id[:12]}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
# ─── Per-Client Send Queue (backpressure) ───
|
||||||
|
|
||||||
|
async def _queue_send(self, conn: WSConnection, message: dict):
|
||||||
|
"""Queue a message for sending via the per-client sender loop."""
|
||||||
|
try:
|
||||||
|
conn.send_queue.put_nowait(message)
|
||||||
|
except asyncio.QueueFull:
|
||||||
|
logger.warning(
|
||||||
|
f"[HUB] Queue full for {conn.conn_id} — dropping oldest message"
|
||||||
|
)
|
||||||
|
# Drop oldest to make room
|
||||||
|
try:
|
||||||
|
conn.send_queue.get_nowait()
|
||||||
|
except asyncio.QueueEmpty:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
conn.send_queue.put_nowait(message)
|
||||||
|
except asyncio.QueueFull:
|
||||||
|
pass
|
||||||
|
|
||||||
|
async def _sender_loop(self, conn: WSConnection):
|
||||||
|
"""Dedicated sender coroutine for a single connection.
|
||||||
|
|
||||||
|
Reads from the per-client queue and sends via WebSocket.
|
||||||
|
This prevents slow clients from blocking message routing.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
while not conn.ws.closed:
|
||||||
|
try:
|
||||||
|
msg = await asyncio.wait_for(
|
||||||
|
conn.send_queue.get(), timeout=HEARTBEAT_INTERVAL
|
||||||
|
)
|
||||||
|
await conn.ws.send_json(msg)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
continue # No message to send — loop continues
|
||||||
|
except ConnectionResetError:
|
||||||
|
break
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[HUB] Send error {conn.conn_id}: {e}")
|
||||||
|
break
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
async def _send_direct(self, ws: web.WebSocketResponse, message: dict):
|
||||||
|
"""Send directly (bypasses queue — for auth messages only)."""
|
||||||
|
try:
|
||||||
|
await ws.send_json(message)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[HUB] Direct send error: {e}")
|
||||||
|
|
||||||
|
# ─── Message Loop ───
|
||||||
|
|
||||||
|
async def _message_loop(self, conn: WSConnection):
|
||||||
|
"""Main message receive loop for an authenticated connection."""
|
||||||
|
async for msg in conn.ws:
|
||||||
|
if msg.type == WSMsgType.TEXT:
|
||||||
|
# Rate limit check
|
||||||
|
if not self._check_rate_limit(conn):
|
||||||
|
await self._queue_send(conn, {
|
||||||
|
"type": MsgType.ERROR,
|
||||||
|
"error": "Rate limited — too many messages",
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
|
||||||
|
try:
|
||||||
|
data = json.loads(msg.data)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
continue
|
||||||
|
|
||||||
|
msg_type = data.get("type", "")
|
||||||
|
|
||||||
|
# Dedup check
|
||||||
|
msg_id = data.get("msg_id")
|
||||||
|
if msg_id and self._is_duplicate(msg_id):
|
||||||
|
continue
|
||||||
|
|
||||||
|
await self._handle_message(conn, msg_type, data)
|
||||||
|
|
||||||
|
elif msg.type == WSMsgType.ERROR:
|
||||||
|
logger.error(
|
||||||
|
f"[HUB] WS error {conn.conn_id}: {conn.ws.exception()}"
|
||||||
|
)
|
||||||
|
break
|
||||||
|
|
||||||
|
elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
|
||||||
|
break
|
||||||
|
|
||||||
|
async def _handle_message(self, conn: WSConnection, msg_type: str, data: dict):
|
||||||
|
"""Route an incoming message to the appropriate handler."""
|
||||||
|
conn.last_heartbeat = time.time() # Any message counts as heartbeat
|
||||||
|
|
||||||
|
if msg_type == MsgType.PENDING:
|
||||||
|
payload = data.get("data", {})
|
||||||
|
request_id = payload.get("request_id", "")
|
||||||
|
if request_id:
|
||||||
|
# Track ownership for response routing
|
||||||
|
self.pending_owners[request_id] = conn.conn_id
|
||||||
|
# Add source metadata
|
||||||
|
payload["_conn_id"] = conn.conn_id
|
||||||
|
payload["_instance_number"] = conn.instance_number
|
||||||
|
payload["_pc_name"] = conn.pc_name
|
||||||
|
payload.setdefault("project_name", conn.project)
|
||||||
|
if self._on_pending:
|
||||||
|
await self._on_pending(conn.project, payload)
|
||||||
|
|
||||||
|
elif msg_type == MsgType.CHAT:
|
||||||
|
payload = data.get("data", {})
|
||||||
|
payload.setdefault("project_name", conn.project)
|
||||||
|
payload["_instance_number"] = conn.instance_number
|
||||||
|
payload["_pc_name"] = conn.pc_name
|
||||||
|
if self._on_chat:
|
||||||
|
await self._on_chat(conn.project, payload)
|
||||||
|
|
||||||
|
elif msg_type == MsgType.REGISTER:
|
||||||
|
payload = data.get("data", {})
|
||||||
|
payload.setdefault("project_name", conn.project)
|
||||||
|
if self._on_register:
|
||||||
|
await self._on_register(payload)
|
||||||
|
|
||||||
|
elif msg_type == MsgType.AUTO_RESOLVE:
|
||||||
|
payload = data.get("data", {})
|
||||||
|
request_id = payload.get("request_id", "")
|
||||||
|
if request_id:
|
||||||
|
self.pending_owners.pop(request_id, None)
|
||||||
|
if self._on_auto_resolve:
|
||||||
|
await self._on_auto_resolve(conn.project, payload)
|
||||||
|
|
||||||
|
elif msg_type == MsgType.BRAIN_EVENT:
|
||||||
|
payload = data.get("data", {})
|
||||||
|
payload.setdefault("project_name", conn.project)
|
||||||
|
if self._on_brain_event:
|
||||||
|
await self._on_brain_event(conn.project, payload)
|
||||||
|
|
||||||
|
elif msg_type == MsgType.HEARTBEAT:
|
||||||
|
pass # last_heartbeat already updated above
|
||||||
|
|
||||||
|
else:
|
||||||
|
logger.warning(f"[HUB] Unknown message type: {msg_type} from {conn.conn_id}")
|
||||||
|
|
||||||
|
# ─── Rate Limiting ───
|
||||||
|
|
||||||
|
def _check_rate_limit(self, conn: WSConnection) -> bool:
|
||||||
|
"""Per-connection rate limiting. Returns True if allowed."""
|
||||||
|
now = time.time()
|
||||||
|
conn._msg_timestamps = [
|
||||||
|
t for t in conn._msg_timestamps if now - t < RATE_WINDOW
|
||||||
|
]
|
||||||
|
if len(conn._msg_timestamps) >= PER_CONN_RATE_LIMIT:
|
||||||
|
logger.warning(f"[HUB] Rate limited: {conn.conn_id}")
|
||||||
|
return False
|
||||||
|
conn._msg_timestamps.append(now)
|
||||||
|
return True
|
||||||
|
|
||||||
|
# ─── Deduplication ───
|
||||||
|
|
||||||
|
def _is_duplicate(self, msg_id: str) -> bool:
|
||||||
|
"""Check if a message ID was recently seen."""
|
||||||
|
now = time.time()
|
||||||
|
# Cleanup old entries
|
||||||
|
if len(self._recent_msg_ids) > 10000:
|
||||||
|
cutoff = now - self._msg_id_ttl
|
||||||
|
self._recent_msg_ids = {
|
||||||
|
k: v for k, v in self._recent_msg_ids.items() if v > cutoff
|
||||||
|
}
|
||||||
|
if msg_id in self._recent_msg_ids:
|
||||||
|
return True
|
||||||
|
self._recent_msg_ids[msg_id] = now
|
||||||
|
return False
|
||||||
|
|
||||||
|
# ─── Diagnostics ───
|
||||||
|
|
||||||
|
def get_status(self) -> dict:
|
||||||
|
"""Return hub status for health/diagnostics."""
|
||||||
|
projects = {}
|
||||||
|
for project, conn_ids in self.project_connections.items():
|
||||||
|
instances = []
|
||||||
|
for cid in conn_ids:
|
||||||
|
c = self.connections.get(cid)
|
||||||
|
if c:
|
||||||
|
instances.append({
|
||||||
|
"conn_id": c.conn_id,
|
||||||
|
"instance": c.instance_number,
|
||||||
|
"pc": c.pc_name,
|
||||||
|
"connected": int(time.time() - c.connected_at),
|
||||||
|
})
|
||||||
|
projects[project] = {
|
||||||
|
"count": len(conn_ids),
|
||||||
|
"instances": instances,
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
"total_connections": len(self.connections),
|
||||||
|
"projects": projects,
|
||||||
|
"pending_owners": len(self.pending_owners),
|
||||||
|
}
|
||||||
21
main.py
21
main.py
@@ -102,14 +102,29 @@ async def main():
|
|||||||
else:
|
else:
|
||||||
logger.info("Gateway mode — watcher disabled (data via HTTP API)")
|
logger.info("Gateway mode — watcher disabled (data via HTTP API)")
|
||||||
|
|
||||||
# Start Gateway HTTP API (gateway mode)
|
# Start Gateway HTTP API + WebSocket Hub (gateway mode)
|
||||||
if Config.BOT_MODE == 'gateway':
|
if Config.BOT_MODE == 'gateway':
|
||||||
from gateway import GatewayAPI
|
from gateway import GatewayAPI
|
||||||
|
from hub import WSHub
|
||||||
|
from auth import TokenManager
|
||||||
|
|
||||||
|
# Initialize Hub
|
||||||
|
token_mgr = TokenManager(
|
||||||
|
secret=Config.GRAVITY_HUB_SECRET,
|
||||||
|
registration_code=Config.GRAVITY_REGISTRATION_CODE,
|
||||||
|
)
|
||||||
|
hub = WSHub(token_mgr)
|
||||||
|
|
||||||
gateway_port = int(os.environ.get('GATEWAY_PORT', '8585'))
|
gateway_port = int(os.environ.get('GATEWAY_PORT', '8585'))
|
||||||
gateway = GatewayAPI(bot, port=gateway_port, api_key=Config.GATEWAY_API_KEY)
|
gateway = GatewayAPI(
|
||||||
|
bot, port=gateway_port,
|
||||||
|
api_key=Config.GATEWAY_API_KEY,
|
||||||
|
hub=hub,
|
||||||
|
)
|
||||||
bot.gateway = gateway # Enable _write_command → gateway.push_command
|
bot.gateway = gateway # Enable _write_command → gateway.push_command
|
||||||
|
bot.hub = hub # Enable Hub-based message routing
|
||||||
await gateway.start()
|
await gateway.start()
|
||||||
logger.info(f"Gateway API running on port {gateway_port}")
|
logger.info(f"Gateway API + WS Hub running on port {gateway_port}")
|
||||||
|
|
||||||
# Run Discord bot (blocks until bot disconnects)
|
# Run Discord bot (blocks until bot disconnects)
|
||||||
await bot.start(Config.DISCORD_TOKEN)
|
await bot.start(Config.DISCORD_TOKEN)
|
||||||
|
|||||||
28
tests/test_syntax.py
Normal file
28
tests/test_syntax.py
Normal file
@@ -0,0 +1,28 @@
|
|||||||
|
"""Quick syntax check for all modified/new Python files."""
|
||||||
|
import ast
|
||||||
|
import sys
|
||||||
|
|
||||||
|
files = [
|
||||||
|
'auth.py',
|
||||||
|
'hub.py',
|
||||||
|
'config.py',
|
||||||
|
'gateway.py',
|
||||||
|
'main.py',
|
||||||
|
'bot.py',
|
||||||
|
]
|
||||||
|
|
||||||
|
errors = []
|
||||||
|
for f in files:
|
||||||
|
try:
|
||||||
|
with open(f'../{f}', encoding='utf-8') as fh:
|
||||||
|
ast.parse(fh.read())
|
||||||
|
print(f' OK: {f}')
|
||||||
|
except SyntaxError as e:
|
||||||
|
print(f' FAIL: {f} - {e}')
|
||||||
|
errors.append(f)
|
||||||
|
|
||||||
|
if errors:
|
||||||
|
print(f'\nFailed: {errors}')
|
||||||
|
sys.exit(1)
|
||||||
|
else:
|
||||||
|
print('\nAll files parse OK')
|
||||||
74
tests/test_ws_hub.py
Normal file
74
tests/test_ws_hub.py
Normal file
@@ -0,0 +1,74 @@
|
|||||||
|
"""Quick test for WebSocket Hub connection and auth protocol."""
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import urllib.request
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_ws():
|
||||||
|
import websockets
|
||||||
|
|
||||||
|
reg_code = os.getenv("GRAVITY_REGISTRATION_CODE")
|
||||||
|
uri = "ws://localhost:8586/ws"
|
||||||
|
|
||||||
|
print(f"1. Connecting to {uri}...")
|
||||||
|
async with websockets.connect(uri) as ws:
|
||||||
|
# Correct auth protocol: type=auth with registration_code
|
||||||
|
auth_msg = {
|
||||||
|
"type": "auth",
|
||||||
|
"registration_code": reg_code,
|
||||||
|
"project": "gravity_control",
|
||||||
|
"pc": "test_pc",
|
||||||
|
}
|
||||||
|
await ws.send(json.dumps(auth_msg))
|
||||||
|
print("2. Sent auth message")
|
||||||
|
|
||||||
|
resp = await asyncio.wait_for(ws.recv(), timeout=5)
|
||||||
|
resp_data = json.loads(resp)
|
||||||
|
print(f"3. Auth response: {json.dumps(resp_data, indent=2)}")
|
||||||
|
|
||||||
|
if resp_data.get("type") == "auth_ok":
|
||||||
|
print("4. AUTH SUCCESS!")
|
||||||
|
conn_id = resp_data.get("conn_id", "")
|
||||||
|
instance = resp_data.get("instance_number", "")
|
||||||
|
print(f" conn_id={conn_id}, instance=#{instance}")
|
||||||
|
|
||||||
|
# Send a chat message (Hub expects type='chat', not 'chat_snapshot')
|
||||||
|
snap = {
|
||||||
|
"type": "chat",
|
||||||
|
"data": {
|
||||||
|
"content": "[TEST] Hub WS test message",
|
||||||
|
"project_name": "gravity_control",
|
||||||
|
"conversation_id": "test-ws-hub-123",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
await ws.send(json.dumps(snap))
|
||||||
|
print("5. Sent chat_snapshot via WS")
|
||||||
|
|
||||||
|
# Check hub status via HTTP
|
||||||
|
r = urllib.request.urlopen("http://localhost:8586/hub/status", timeout=3)
|
||||||
|
status = json.loads(r.read())
|
||||||
|
print(f"6. Hub status: {json.dumps(status, indent=2)}")
|
||||||
|
|
||||||
|
# Send register message
|
||||||
|
reg = {
|
||||||
|
"type": "register",
|
||||||
|
"conversation_id": "test-ws-hub-123",
|
||||||
|
"project_name": "gravity_control",
|
||||||
|
}
|
||||||
|
await ws.send(json.dumps(reg))
|
||||||
|
print("7. Sent register")
|
||||||
|
|
||||||
|
else:
|
||||||
|
reason = resp_data.get("reason", "unknown")
|
||||||
|
print(f"4. AUTH FAILED: {reason}")
|
||||||
|
|
||||||
|
await ws.close()
|
||||||
|
print("8. Done!")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(test_ws())
|
||||||
Reference in New Issue
Block a user